Remove the None presets and checks around Stage 12 span and operation blocks

changelog_entry.yaml

@@ -17,6 +17,7 @@
       - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's
     changed:
       - Required policyengine-observability 3.0.4 (>=3.0.4,<4) in the observability library, the entrypoint, the gateway and the executor, and locked it in the six locks that include it, where it is the only package that moves; 3.0.4 annotates its span and operation scopes' __exit__ and __aexit__ as returning Literal[False], so the services now open spans and operations with runtime.span and runtime.operation directly, and the policyengine_simulation_observability.scopes module, which only retyped those scopes for Pyright, is removed with its tests; a new observability library test fails if an older release is installed, Pyright still reports no errors in any package, and 3.0.3's fix, which keeps the recorded exception on the request.completed record and server span of a request that ends with an error outcome but passes no exception to end_request, is included
+      - Removed the None presets and None checks around span and operation blocks that only satisfied Pyright before policyengine-observability 3.0.4; the Stage 12 Modal functions run_single_simulation_segment_us, run_single_simulation_us, run_single_simulation_uk and coordinate_report return their result from inside the operation block, and the Stage 12 segment runner no longer sets the dispatched call, the collected payloads and the merged calculation to None before their spans or checks them afterwards, and returns the merge from inside its span; 3.0.4's scopes never suppress an exception, and each value comes from a function that never returns None, so the removed RuntimeErrors (returned no result, no call, no payloads or no calculation) could not be raised; new tests run each of those blocks through the scopes of the runtime the services build
       - Enforced the file-length rule sim.quality.python-file-length, now that every tracked Python file is within its limit: its empty baseline and its downgrade section are gone, and its tier is metric rather than mechanical, because it measures a number against a threshold, with the at-the-limit and one-over fixtures as its below and above proof
       - Moved the entrypoint's backend calls, response telemetry and comparison dispatch out of app.py (925 lines, now 583) into entry_runtime.py, which app.py re-exports and which still logs under the policyengine_simulation_entry.app logger; the routes stay in app.py, so FastAPI and ruff read their annotations as before; the 2,251-line deployment asset test file is split into seven test files and a support module named test_deployment_assets_support.py, the rules registry names each moved test by its new file, and its guard-test discovery now covers tests/test_deployment_assets*.py
       - Moved the worker's credential setup, request year and policy parsing, and region resolution out of simulation_runtime.py (1,111 lines, now 651) into gcp_credentials.py, request_policy.py and region_resolution.py; simulation_runtime.py keeps dataset selection and loading, simulation building, local-area breakdowns and the run itself, which tests replace by name there, and re-exports the moved names, so every caller still imports them from simulation_runtime; the shared helpers _parse_year, _normalise_policy, _compile_interval_policy, _runtime_policy (now wrapper_policy), _resolve_region, _resolve_dataset_reference and _require_local_area_breakdown are now public, and the credential messages still log under the simulation_runtime logger; the 3,048-line output builder test file is split into seven test files and a module of shared fakes, all named test_simulation_output_*.py so the weighted-frame rule still checks them

projects/policyengine-simulation-executor/src/modal/v2_app.py

@@ -279,7 +279,6 @@ def run_single_simulation_segment_us(
 
     try:
         propagated = normalize_observability_context(observability_context)
-        result = None
         with runtime.operation(
             STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION),
             attributes={
@@ -289,13 +288,9 @@ def run_single_simulation_segment_us(
             },
             remote_context=propagated,
         ):
-            result = stage12_runtime.calculate_segment(
+            return stage12_runtime.calculate_segment(
                 payload, runtime=runtime
             ).model_dump(mode="python")
-        if result is None:
-            msg = "Stage 12 segment calculation returned no result"
-            raise RuntimeError(msg)
-        return result
     finally:
         runtime.shutdown()
 
@@ -336,7 +331,6 @@ def segmented_calculator(
 
     try:
         propagated = normalize_observability_context(observability_context)
-        result = None
         with runtime.operation(
             STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION),
             attributes={
@@ -345,7 +339,7 @@ def segmented_calculator(
             },
             remote_context=propagated,
         ):
-            result = stage12_runtime.run_single_simulation(
+            return stage12_runtime.run_single_simulation(
                 payload,
                 context,
                 required_country="us",
@@ -354,10 +348,6 @@ def segmented_calculator(
                     runtime=runtime,
                 ),
             )
-        if result is None:
-            msg = "Stage 12 US simulation returned no result"
-            raise RuntimeError(msg)
-        return result
     finally:
         runtime.shutdown()
 
@@ -386,7 +376,6 @@ def run_single_simulation_uk(
     )
     try:
         propagated = normalize_observability_context(observability_context)
-        result = None
         with runtime.operation(
             STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION),
             attributes={
@@ -395,18 +384,14 @@ def run_single_simulation_uk(
             },
             remote_context=propagated,
         ):
-            result = stage12_runtime.run_single_simulation(
+            return stage12_runtime.run_single_simulation(
                 payload,
                 context,
                 required_country="uk",
                 dependencies=stage12_runtime.SimulationWorkerDependencies(
                     runtime=runtime
                 ),
             )
-        if result is None:
-            msg = "Stage 12 UK simulation returned no result"
-            raise RuntimeError(msg)
-        return result
     finally:
         runtime.shutdown()
 
@@ -452,7 +437,6 @@ def coordinate_report(
     )
     try:
         propagated = normalize_observability_context(observability_context)
-        result = None
         with runtime.operation(
             stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION),
             attributes={
@@ -465,7 +449,7 @@ def coordinate_report(
             },
             remote_context=propagated,
         ):
-            result = stage12_runtime.coordinate_report(
+            return stage12_runtime.coordinate_report(
                 payload,
                 context,
                 parent,
@@ -475,9 +459,5 @@ def coordinate_report(
                 ),
                 dependencies=stage12_runtime.CoordinatorDependencies(runtime=runtime),
             )
-        if result is None:
-            msg = "Stage 12 report coordinator returned no result"
-            raise RuntimeError(msg)
-        return result
     finally:
         runtime.shutdown()

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py

@@ -42,14 +42,6 @@
 CONSECUTIVE_POLL_ERROR_LIMIT = 2
 
 
-def _require_call(call: modal.FunctionCall | None) -> modal.FunctionCall:
-    """Return ``call``, raising ``RuntimeError`` when a dispatch returned none."""
-    if call is None:
-        msg = "Stage 12 segment dispatch returned no call"
-        raise RuntimeError(msg)
-    return call
-
-
 def _next_backoff(current: float, *, maximum: float) -> float:
     return min(current * POLL_INTERVAL_BACKOFF_FACTOR, maximum)
 
@@ -130,7 +122,6 @@ def run(self) -> SimulationCalculation:
         handles: list[tuple[int, Any]] = []
         try:
             for segment in segments:
-                call = None
                 with self.runtime.span(
                     STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH),
                     attributes={"segment_index": segment.segment_index},
@@ -139,32 +130,23 @@ def run(self) -> SimulationCalculation:
                         segment.model_dump(mode="python"),
                         observability_context=self.runtime.capture_context(),
                     )
-                handles.append((segment.segment_index, _require_call(call)))
+                handles.append((segment.segment_index, call))
         except Exception:
             self._cancel_all(handles)
             raise
 
-        payloads: list[object] | None = None
         with self.runtime.span(
             STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT)
         ):
             payloads = self._collect(handles)
-        if payloads is None:
-            msg = "Stage 12 segment wait returned no payloads"
-            raise RuntimeError(msg)
-        merged: SimulationCalculation | None = None
         with self.runtime.span(
             STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE)
         ):
-            merged = merge_segment_results(
+            return merge_segment_results(
                 self.simulation,
                 payloads,
                 groups=groups,
             )
-        if merged is None:
-            msg = "Stage 12 segment merge returned no calculation"
-            raise RuntimeError(msg)
-        return merged
 
     def _collect(self, handles: list[tuple[int, Any]]) -> list[object]:
         results: list[object | None] = [None] * len(handles)

projects/policyengine-simulation-executor/tests/conftest.py

@@ -50,6 +50,29 @@ def observability_runtime():
     return NoOpObservabilityRuntime()
 
 
+@pytest.fixture
+def real_observability_runtime(monkeypatch):
+    """Build the observability runtime the services build, exporting nothing.
+
+    Its span and operation scopes are the ones deployed code opens, so a test
+    can check what a block inside them returns or raises. OpenTelemetry is
+    disabled and Google Cloud logging is not configured.
+    """
+    from policyengine_simulation_observability.observability import build_runtime
+
+    monkeypatch.setenv("OTEL_SDK_DISABLED", "true")
+    monkeypatch.delenv("OBSERVABILITY_LOGGING_PROJECT_ID", raising=False)
+    monkeypatch.delenv("OBSERVABILITY_LOG_NAME", raising=False)
+    runtime = build_runtime(
+        service_name="policyengine-simulation-executor-test",
+        service_role="test",
+        platform="local",
+        environment="test",
+    )
+    yield runtime
+    runtime.shutdown()
+
+
 @pytest.fixture(autouse=True)
 def _forget_spm_runtime_caches():
     """Keep the SPM runtime memoization out of every other test's way.

projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py

@@ -398,6 +398,57 @@ def fail(*_, **__):
     assert runtime.shutdown_calls == 1
 
 
+@pytest.mark.parametrize(
+    "function_name",
+    [
+        "run_single_simulation_us",
+        "run_single_simulation_uk",
+        "run_single_simulation_segment_us",
+        "coordinate_report",
+    ],
+)
+def test_stage12_functions_return_the_result_from_the_real_operation_scope(
+    monkeypatch,
+    real_observability_runtime,
+    function_name,
+) -> None:
+    module = _load(monkeypatch)
+    monkeypatch.setattr(
+        module, "init_process_observability", lambda **_: real_observability_runtime
+    )
+
+    result = _invoke(monkeypatch, module, function_name, lambda *_, **__: {"ok": True})
+
+    assert result == {"ok": True}
+
+
+@pytest.mark.parametrize(
+    "function_name",
+    [
+        "run_single_simulation_us",
+        "run_single_simulation_uk",
+        "run_single_simulation_segment_us",
+        "coordinate_report",
+    ],
+)
+def test_stage12_functions_raise_the_application_error_from_the_real_operation_scope(
+    monkeypatch,
+    real_observability_runtime,
+    function_name,
+) -> None:
+    module = _load(monkeypatch)
+    monkeypatch.setattr(
+        module, "init_process_observability", lambda **_: real_observability_runtime
+    )
+
+    def fail(*_, **__):
+        msg = "application failure"
+        raise LookupError(msg)
+
+    with pytest.raises(LookupError, match="application failure"):
+        _invoke(monkeypatch, module, function_name, fail)
+
+
 def test_external_package_values_are_assertions_only(monkeypatch) -> None:
     install_fake_modal(monkeypatch)
     monkeypatch.setenv("STAGE12_EXPECT_US_VERSION", "not-the-bundled-version")