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) 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()