From 324e20f9bd86730fd800f81ab846966294341c9f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:36:27 +0400 Subject: [PATCH 1/2] Test the Stage 12 span and operation blocks through the real runtime scopes The Stage 12 Modal functions in v2_app.py and the segment runner in stage12_runtime/fanout.py set a variable to None before a span or operation block and check it after the block. The next commit removes those presets and checks, which only satisfied Pyright before policyengine-observability 3.0.4. These tests run every block involved and pass on the current code. - The existing v2_app tests run the four functions with a recording runtime whose operation is a generator context manager. Two new parametrized tests run them with the observability runtime the services build (OpenTelemetry disabled), through its real operation scope: each function returns its calculation's result, and an error raised in the block reaches the caller unchanged, never as the "returned no result" RuntimeError. - The segment runner already had tests for its dispatch and wait blocks succeeding and failing. A new test makes the merge block fail, and four tests run the runner with the real runtime: a successful run, and an error raised in each of the dispatch, wait and merge blocks, which reaches the caller unchanged. - conftest.py gains a real_observability_runtime fixture that builds that runtime, with Google Cloud logging unset, and shuts it down afterwards. The dispatch-failure setup moves into a helper the old and new dispatch tests share. Co-Authored-By: Claude Opus 5.5 --- .../tests/conftest.py | 23 ++++ .../tests/test_stage12_modal_app.py | 51 +++++++++ .../tests/test_stage12_segment_fanout.py | 105 +++++++++++++++--- 3 files changed, 166 insertions(+), 13 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/conftest.py b/projects/policyengine-simulation-executor/tests/conftest.py index 2d00ec40c..a43259d61 100644 --- a/projects/policyengine-simulation-executor/tests/conftest.py +++ b/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. diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index ccf4a1918..671aa4d12 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/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") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py index a21794b09..7e37fef45 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -21,6 +21,10 @@ from policyengine_simulation_executor.stage12_runtime.partition import ( US_REGION_GROUPS, ) +from policyengine_simulation_observability.stages import ( + STAGE12_SIMULATION_STAGES, + Stage, +) from stage12_runtime_support import planned_simulation if TYPE_CHECKING: @@ -106,7 +110,11 @@ def _results(simulation) -> list[dict]: ] -def _runner(*, errors: dict[int, list[Exception]] | None = None): +def _runner( + *, + errors: dict[int, list[Exception]] | None = None, + runtime: object | None = None, +): simulation = planned_simulation(SimulationRole.BASELINE) function = FakeFunction(_results(simulation), errors or {}) @@ -117,12 +125,12 @@ def from_name(app_name, function_name): assert function_name == "run_single_simulation_segment_us" return function - runtime = FakeRuntime() + span_runtime = runtime if runtime is not None else FakeRuntime() sleeps: list[float] = [] runner = Stage12SegmentRunner( simulation, app_name="stage12-app", - runtime=cast("ObservabilityRuntime", runtime), + runtime=cast("ObservabilityRuntime", span_runtime), seams=SegmentRunnerSeams( modal_module=SimpleNamespace(Function=FunctionLookup), groups=US_REGION_GROUPS, @@ -131,7 +139,24 @@ def from_name(app_name, function_name): poll_interval_max_seconds=2.0, ), ) - return runner, function, runtime, sleeps + return runner, function, span_runtime, sleeps + + +def _fail_dispatch_of_segment_five(function: FakeFunction) -> None: + spawn = function.spawn + + def fail_during_dispatch(payload, *, observability_context): + if payload["segment_index"] == 5: + msg = "dispatch failed" + raise ConnectionError(msg) + return spawn(payload, observability_context=observability_context) + + function.spawn = fail_during_dispatch + + +def _return_segment_zero_for_segment_one(function: FakeFunction) -> None: + # The merge then reads segment index 0 twice and refuses the result set. + function.results[1] = function.results[0] def test_starts_twenty_segments_and_merges_them_in_partition_order() -> None: @@ -190,15 +215,7 @@ def test_second_poll_error_cancels_all_started_segments() -> None: def test_dispatch_error_cancels_segments_that_already_started() -> None: runner, function, _, _ = _runner() - spawn = function.spawn - - def fail_during_dispatch(payload, *, observability_context): - if payload["segment_index"] == 5: - msg = "dispatch failed" - raise ConnectionError(msg) - return spawn(payload, observability_context=observability_context) - - function.spawn = fail_during_dispatch + _fail_dispatch_of_segment_five(function) with pytest.raises(ConnectionError, match="dispatch failed"): runner.run() @@ -219,3 +236,65 @@ def test_persisted_parent_cancellation_cancels_all_segments() -> None: assert sleeps == [0.5] assert all(call.cancelled == 1 for call in function.calls) + + +def test_merge_error_is_raised_from_the_merge_span() -> None: + runner, function, runtime, _ = _runner() + _return_segment_zero_for_segment_one(function) + + with pytest.raises(ValueError, match="duplicate index"): + runner.run() + + assert runtime.spans[-1][0] == STAGE12_SIMULATION_STAGES.name( + Stage.STAGE12_SEGMENT_MERGE + ) + assert not any(call.cancelled for call in function.calls) + + +def test_segments_are_dispatched_collected_and_merged_in_the_real_spans( + real_observability_runtime, +) -> None: + runner, function, _, _ = _runner(runtime=real_observability_runtime) + + result = runner.run() + + assert len(function.spawns) == 20 + assert result.frames["household"]["household_id"].tolist() == list(range(20)) + + +def test_dispatch_error_is_raised_from_the_real_dispatch_span( + real_observability_runtime, +) -> None: + runner, function, _, _ = _runner(runtime=real_observability_runtime) + _fail_dispatch_of_segment_five(function) + + with pytest.raises(ConnectionError, match="dispatch failed"): + runner.run() + + assert len(function.calls) == 5 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_poll_error_is_raised_from_the_real_wait_span( + real_observability_runtime, +) -> None: + runner, function, _, _ = _runner( + errors={0: [ConnectionError("first"), RuntimeError("second")]}, + runtime=real_observability_runtime, + ) + + with pytest.raises(Stage12ExecutionError) as error: + runner.run() + + assert error.value.detail.error_code == "segment_execution_failed" + assert all(call.cancelled == 1 for call in function.calls) + + +def test_merge_error_is_raised_from_the_real_merge_span( + real_observability_runtime, +) -> None: + runner, function, _, _ = _runner(runtime=real_observability_runtime) + _return_segment_zero_for_segment_one(function) + + with pytest.raises(ValueError, match="duplicate index"): + runner.run() From 8f87202f7ca6250ede83ac71a53e3263cd784b0a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:40:01 +0400 Subject: [PATCH 2/2] Remove the None presets and checks around Stage 12 span and operation blocks Before policyengine-observability 3.0.4, Pyright read the span and operation scopes as possibly suppressing the exception raised in their block. #712 therefore set each value a block assigns to None before the block and raised a RuntimeError after it when the value was still None. 3.0.4 annotates the scopes' __exit__ and __aexit__ as returning Literal[False], and they return False on every path, so the block either assigns its value or its exception reaches the caller. None of the assigned values can be None either: - v2_app.py: run_single_simulation_segment_us returns calculate_segment(...).model_dump(mode="python"), a dict; run_single_simulation_us and run_single_simulation_uk return run_single_simulation(...), and coordinate_report returns coordinate_report(...), both typed dict[str, Any], and every path of each returns a dict or raises. Each function now returns from inside its operation block (ruff's RET504 rejects assigning and then returning); the operation still ends, and the runtime still shuts down, before the function returns. - fanout.py: the dispatched call is the FunctionCall that modal 1.4.2's Function.spawn creates and returns on every path (it raises for a generator function, and the segment function is not one), so _require_call, the only check of it, is removed with the preset; _collect returns a list or raises; and merge_segment_results returns a SimulationCalculation, now from inside the merge span. The removed RuntimeErrors could not be raised, so behaviour is unchanged and there is no failing test to write: the tests added in the previous commit, which run each block through the real runtime scopes, pass before and after. Pyright still reports no errors in the executor. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + .../src/modal/v2_app.py | 28 +++---------------- .../stage12_runtime/fanout.py | 22 ++------------- 3 files changed, 7 insertions(+), 44 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 41efb7fd9..bcc8a4ad3 100644 --- a/changelog_entry.yaml +++ b/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 diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 92642d673..b371403c2 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/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,7 +384,7 @@ 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", @@ -403,10 +392,6 @@ def run_single_simulation_uk( 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() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py index 83194b9f5..966331f48 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py +++ b/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)