From 595f04407cd4c63192f3566068301f468b710540 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 3 Nov 2026 23:59:17 +0400 Subject: [PATCH 01/24] Let a Stage 12 call name the report run and simulation runs it records on A report call's invocation context can now carry the IDs of the Public API's report run and its baseline and reform simulation runs. Comparison runs leave the field unset, and a shadow run compared with a production call cannot carry it. Co-Authored-By: Claude Opus 5.5 --- .../stage12_execution/__init__.py | 2 + .../stage12_execution/comparison_runs.py | 48 +++++++++- .../tests/test_stage12_report_run_ids.py | 96 +++++++++++++++++++ 3 files changed, 145 insertions(+), 1 deletion(-) create mode 100644 libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py index d1a1f477f..6006c63ea 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py @@ -38,6 +38,7 @@ ComparisonSimulationRecord, ErrorCode, ErrorSummary, + ReportRunIds, ResultComparisonStatus, Stage12InvocationContext, ) @@ -95,6 +96,7 @@ "ReportAggregate", "ReportExecutionInput", "ReportOutputRequirements", + "ReportRunIds", "ResultComparisonArtifactPayload", "ResultComparisonStatus", "ResultDifference", diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py index ea49b5e2d..31682cdd6 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py @@ -23,13 +23,49 @@ ) +class ReportRunIds(StrictContractModel): + """The Public API's report run and simulation runs one report call records on. + + The simulation service creates the three rows before it starts the + coordinator. The coordinator records the report run's status and results, + and each worker records its own simulation run's. + """ + + report_run_id: UUID + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID + + @model_validator(mode="after") + def validate_distinct_runs(self) -> ReportRunIds: + """Require the report run and both simulation runs to be different rows.""" + run_ids = { + self.report_run_id, + self.baseline_simulation_run_id, + self.reform_simulation_run_id, + } + if len(run_ids) != 3: # noqa: PLR2004 # one report run and two simulation runs + msg = "report run IDs must name three different runs" + raise ValueError(msg) + return self + + def simulation_run_id(self, role: SimulationRole) -> UUID: + """Return the simulation run of ``role``; a standalone role has none.""" + if role is SimulationRole.BASELINE: + return self.baseline_simulation_run_id + if role is SimulationRole.REFORM: + return self.reform_simulation_run_id + msg = "a standalone simulation has no simulation run in a report run" + raise ValueError(msg) + + class Stage12InvocationContext(StrictContractModel): """Deployment and retention context for one comparison-run invocation. Records the originating request, the environments, the worker version and callable, the manifest digests and the artifact prefix. Automatic runs also name the production function call they are compared with. Artifacts are kept - for at most 30 days after ``created_at``. + for at most 30 days after ``created_at``. A call that also records its + outcome on the Public API's run tables names those runs in ``report_run``. """ contract_version: Literal[1] = 1 @@ -51,6 +87,8 @@ class Stage12InvocationContext(StrictContractModel): production_function_call_id: ContractText | None = None created_at: datetime retention_expires_at: datetime + # Comparison runs leave it unset and record only the Stage 12 tables. + report_run: ReportRunIds | None = None @model_validator(mode="after") def validate_retention(self) -> Stage12InvocationContext: @@ -58,6 +96,14 @@ def validate_retention(self) -> Stage12InvocationContext: _validate_retention(self.created_at, self.retention_expires_at) return self + @model_validator(mode="after") + def validate_report_run_is_direct(self) -> Stage12InvocationContext: + """Reject report runs on a shadow run compared with a production call.""" + if self.report_run is not None and self.production_function_call_id: + msg = "a run compared with a production function call has no report run" + raise ValueError(msg) + return self + ErrorCode = Annotated[str, Field(min_length=1, max_length=64)] ErrorSummary = Annotated[str, Field(min_length=1, max_length=512)] diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py b/libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py new file mode 100644 index 000000000..3596581dd --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py @@ -0,0 +1,96 @@ +"""The report run and simulation runs a Stage 12 call records its outcome on.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Any +from uuid import UUID + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ReportRunIds, + SimulationRole, + Stage12InvocationContext, +) +from pydantic import ValidationError + +NOW = datetime(2026, 11, 3, tzinfo=UTC) +REPORT_RUN_ID = UUID("00000000-0000-0000-0000-0000000000a1") +BASELINE_RUN_ID = UUID("00000000-0000-0000-0000-0000000000b1") +REFORM_RUN_ID = UUID("00000000-0000-0000-0000-0000000000c1") + + +def _context_fields(**overrides: Any) -> dict[str, Any]: + return { + "request_id": "request-1", + "environment": "staging", + "modal_environment": "staging", + "worker_version": "6.12.0", + "modal_application": "policyengine-simulation-v2-py6-12-0", + "simulation_callable": "run_single_simulation_us", + "version_manifest_sha256": "a" * 64, + "bundle_manifest_sha256": "b" * 64, + "artifact_prefix": "stage-12-runs/staging/2026/11/report", + "created_at": NOW, + "retention_expires_at": NOW + timedelta(days=30), + **overrides, + } + + +def _run_ids() -> ReportRunIds: + return ReportRunIds( + report_run_id=REPORT_RUN_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_run_id=REFORM_RUN_ID, + ) + + +def test_context_without_run_ids_keeps_its_comparison_fields() -> None: + context = Stage12InvocationContext(**_context_fields()) + + assert context.report_run is None + assert "report_run" not in context.model_dump(mode="json", exclude_none=True) + + +def test_context_carries_the_report_run_and_both_simulation_runs() -> None: + context = Stage12InvocationContext( + **_context_fields(report_run=_run_ids().model_dump(mode="json")) + ) + + assert context.report_run == _run_ids() + restored = Stage12InvocationContext.model_validate(context.model_dump(mode="json")) + assert restored.report_run == _run_ids() + + +def test_each_role_names_its_own_simulation_run() -> None: + run_ids = _run_ids() + + assert run_ids.simulation_run_id(SimulationRole.BASELINE) == BASELINE_RUN_ID + assert run_ids.simulation_run_id(SimulationRole.REFORM) == REFORM_RUN_ID + with pytest.raises(ValueError, match="standalone"): + run_ids.simulation_run_id(SimulationRole.STANDALONE) + + +@pytest.mark.parametrize( + "repeated_ids", + [ + {"baseline_simulation_run_id": REPORT_RUN_ID}, + {"reform_simulation_run_id": REPORT_RUN_ID}, + {"reform_simulation_run_id": BASELINE_RUN_ID}, + ], +) +def test_run_ids_must_name_three_different_runs(repeated_ids) -> None: + fields = {**_run_ids().model_dump(), **repeated_ids} + + with pytest.raises(ValidationError, match="three different runs"): + ReportRunIds(**fields) + + +def test_a_shadow_comparison_run_cannot_record_report_runs() -> None: + with pytest.raises(ValidationError, match="production function call"): + Stage12InvocationContext( + **_context_fields( + production_function_call_id="fc-production", + report_run=_run_ids().model_dump(mode="json"), + ) + ) From 649bb7ae04cc661c93b6de1f251a3fa2bbf1e1e7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:01:13 +0400 Subject: [PATCH 02/24] Give the report run writers the Modal worker role's database secret The v2 Modal application now attaches the Modal secret policyengine-modal-worker-database, which must hold MODAL_WORKER_DATABASE_URL, to the report coordinator and the US and UK simulation workers. Region-group segments and worker validation do not receive it. An operator creates the secret before this deploys. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 17 +++++++++-- .../stage12_runtime/report_run_access.py | 14 +++++++++ .../tests/test_stage12_modal_app.py | 30 +++++++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 69eedc493..623e4334e 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -25,6 +25,10 @@ ) from policyengine_simulation_executor.stage12_cache.models import CacheManifest from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore +from policyengine_simulation_executor.stage12_runtime.report_run_access import ( + MODAL_WORKER_DATABASE_SECRET_NAME, + MODAL_WORKER_DATABASE_URL_ENV, +) from policyengine_simulation_executor.stage12_timeouts import ( COORDINATOR_TIMEOUT_SECONDS, UK_SIMULATION_TIMEOUT_SECONDS, @@ -130,6 +134,13 @@ def _external_assertions(environment: dict[str, str]) -> dict[str, str]: hf_secret, comparison_runtime_secret, ] +# Only the functions that record a report run's or simulation run's status +# and results hold the Modal worker role's database credentials. +modal_worker_database_secret = modal.Secret.from_name( + MODAL_WORKER_DATABASE_SECRET_NAME, + required_keys=[MODAL_WORKER_DATABASE_URL_ENV], +) +report_run_writer_secrets = [*worker_secrets, modal_worker_database_secret] def _country_bundle(country: CountryId) -> Stage12CountryBundle: @@ -304,7 +315,7 @@ def run_single_simulation_segment_us( timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, - secrets=worker_secrets, + secrets=report_run_writer_secrets, ) def run_single_simulation_us( payload: dict, @@ -361,7 +372,7 @@ def segmented_calculator( timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, - secrets=worker_secrets, + secrets=report_run_writer_secrets, ) def run_single_simulation_uk( payload: dict, @@ -407,7 +418,7 @@ def run_single_simulation_uk( timeout=COORDINATOR_TIMEOUT_SECONDS, retries=0, max_containers=10, - secrets=worker_secrets, + secrets=report_run_writer_secrets, ) def coordinate_report( payload: dict, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py new file mode 100644 index 000000000..25b63edc9 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py @@ -0,0 +1,14 @@ +"""The Modal worker role's database credentials, which report runs are recorded with. + +An operator creates the Modal secret in each Modal environment. The v2 Modal +application attaches it to the coordinator and the simulation workers, which +record the status and results of the report run and simulation runs their +call names. +""" + +from __future__ import annotations + +# The Modal secret and the one key it must hold: a postgresql:// URL for the +# Modal worker role on the database that holds the Public API's run tables. +MODAL_WORKER_DATABASE_SECRET_NAME = "policyengine-modal-worker-database" # noqa: S105 # the Modal secret's name, not its value +MODAL_WORKER_DATABASE_URL_ENV = "MODAL_WORKER_DATABASE_URL" 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 e511ee05e..75b3d3e44 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -344,6 +344,36 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert "asgi_app" not in vars(module) +def test_only_report_run_writers_receive_the_modal_worker_database_secret( + monkeypatch, +) -> None: + module = _load(monkeypatch) + functions = dict(module.app.function_calls) + worker_database_secret = { + "args": ("policyengine-modal-worker-database",), + "kwargs": {"required_keys": ["MODAL_WORKER_DATABASE_URL"]}, + } + + assert module.modal_worker_database_secret == worker_database_secret + for writer in ( + "coordinate_report", + "run_single_simulation_us", + "run_single_simulation_uk", + ): + assert functions[writer]["secrets"] == [ + *module.worker_secrets, + worker_database_secret, + ] + # Segments return their result to their worker, and validation records + # nothing, so neither holds the database credentials. + for other in ( + "run_single_simulation_segment_us", + "validate_worker_us", + "validate_worker_uk", + ): + assert functions[other]["secrets"] == module.worker_secrets + + @pytest.mark.parametrize( "function_name", [ From 9b9d7e6e468059779b50080db150cc6040eb0214 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:04:52 +0400 Subject: [PATCH 03/24] Cancel a Stage 12 worker's segment calls when Modal stops the worker Modal pre-empts a call with SIGINT, raised as KeyboardInterrupt, and cancels one with InputCancellation. Both are BaseExceptions that are not Exceptions, and Modal does not cancel the calls a stopped call spawned. The segmented US worker caught only Exceptions, so its region-group segment calls kept running after Modal stopped it. It now cancels every started segment call, then re-raises. Each started call is also recorded inside its dispatch span, so an interruption as the span closes still finds it. No database row or stored file changes. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/fanout.py | 24 ++++++++---- .../tests/test_stage12_segment_fanout.py | 37 +++++++++++++++++++ 2 files changed, 54 insertions(+), 7 deletions(-) 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 966331f48..570810913 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 @@ -111,7 +111,9 @@ def run(self) -> SimulationCalculation: A segment whose result cannot be read twice in a row raises ``Stage12ExecutionError``, and a requested cancellation raises ``Stage12CancellationError``. Either way, every started segment is - cancelled first. + cancelled first. So it is when Modal pre-empts or cancels this worker, + which it signals with a ``BaseException`` that is not an ``Exception`` + and which is then re-raised: Modal does not cancel spawned calls. """ groups = self._groups() segments = build_segment_inputs(self.simulation, groups=groups) @@ -130,15 +132,23 @@ def run(self) -> SimulationCalculation: segment.model_dump(mode="python"), observability_context=self.runtime.capture_context(), ) - handles.append((segment.segment_index, call)) - except Exception: + # Kept inside the span, so an interruption as it closes + # still finds the call to cancel. + handles.append((segment.segment_index, call)) + except BaseException: self._cancel_all(handles) raise - with self.runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) - ): - payloads = self._collect(handles) + try: + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) + ): + payloads = self._collect(handles) + except BaseException as interruption: + # Collection cancels the segments itself where a failure needs it. + if not isinstance(interruption, Exception): + self._cancel_all(handles) + raise with self.runtime.span( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) ): 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 7e37fef45..6506c0a8f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -224,6 +224,43 @@ def test_dispatch_error_cancels_segments_that_already_started() -> None: assert all(call.cancelled == 1 for call in function.calls) +class InputCancelledForTest(BaseException): + """Stands in for Modal's InputCancellation, a BaseException but not an Exception.""" + + +@pytest.mark.parametrize("interruption", [KeyboardInterrupt, InputCancelledForTest]) +def test_preempted_or_cancelled_worker_cancels_every_started_segment( + interruption, +) -> None: + # Modal pre-empts a call with SIGINT (KeyboardInterrupt) and cancels one + # with InputCancellation; it does not cancel the calls the worker spawned. + runner, function, _, _ = _runner(errors={3: [interruption()]}) + + with pytest.raises(interruption): + runner.run() + + assert len(function.calls) == 20 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_worker_preempted_during_dispatch_cancels_the_segments_it_started() -> None: + runner, function, _, _ = _runner() + spawn = function.spawn + + def preempt_during_dispatch(payload, *, observability_context): + if payload["segment_index"] == 5: + raise KeyboardInterrupt + return spawn(payload, observability_context=observability_context) + + function.spawn = preempt_during_dispatch + + with pytest.raises(KeyboardInterrupt): + runner.run() + + assert len(function.calls) == 5 + assert all(call.cancelled == 1 for call in function.calls) + + def test_persisted_parent_cancellation_cancels_all_segments() -> None: runner, function, _, sleeps = _runner( errors={index: [TimeoutError()] for index in range(20)} From 0e65bd001f901efaf905c12afc7ae8e7b040fa0f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:10:51 +0400 Subject: [PATCH 04/24] Let a Stage 12 artifact store name each artifact's digest in its key A content-addressed store writes a simulation's planned input, its Parquet output and the report's aggregate under keys that end in the content's SHA-256. A call that Modal restarts after pre-emption can calculate different last bits, or plan its input on a redeployed copy of the application; its artifacts are then stored beside the earlier attempt's instead of failing the write-once check, and identical bytes reuse the earlier object. Comparison runs keep their fixed keys. Co-Authored-By: Claude Opus 5.5 --- .../stage12_artifacts.py | 75 +++++++++++++---- .../test_stage12_report_run_artifacts.py | 81 +++++++++++++++++++ 2 files changed, 140 insertions(+), 16 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index 71877ee60..c2a11d555 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -72,19 +72,35 @@ def comparison_run_prefix( ) -def input_path(*, prefix: str, role: str) -> str: - """Return the key of the planned execution input JSON for one simulation role.""" - return f"{prefix}/inputs/{role}.json" +def _content_named(name: str, content_sha256: str | None) -> str: + """Return ``name``, followed by ``-`` when a digest is given.""" + return name if content_sha256 is None else f"{name}-{content_sha256}" -def simulation_path(*, prefix: str, role: str) -> str: - """Return the key of the Parquet output for one simulation role.""" - return f"{prefix}/simulations/{role}.parquet" +def input_path(*, prefix: str, role: str, content_sha256: str | None = None) -> str: + """Return the key of the planned execution input JSON for one simulation role. + + With ``content_sha256``, the key also names the content's SHA-256. + """ + return f"{prefix}/inputs/{_content_named(role, content_sha256)}.json" + + +def simulation_path( + *, prefix: str, role: str, content_sha256: str | None = None +) -> str: + """Return the key of the Parquet output for one simulation role. + + With ``content_sha256``, the key also names the content's SHA-256. + """ + return f"{prefix}/simulations/{_content_named(role, content_sha256)}.parquet" -def aggregate_path(*, prefix: str) -> str: - """Return the key of the run's aggregate report JSON.""" - return f"{prefix}/reports/aggregate.json" +def aggregate_path(*, prefix: str, content_sha256: str | None = None) -> str: + """Return the key of the run's aggregate report JSON. + + With ``content_sha256``, the key also names the content's SHA-256. + """ + return f"{prefix}/reports/{_content_named('aggregate', content_sha256)}.json" def comparison_path(*, prefix: str) -> str: @@ -273,16 +289,30 @@ class Stage12ArtifactStore: """Write-once storage for one private bucket's Stage 12 run artifacts. A write to a path already holding identical bytes succeeds; different - bytes raise ``RuntimeError``. + bytes raise ``RuntimeError``. A ``content_addressed`` store names each + input, simulation output and aggregate's SHA-256 in its key, so a call + Modal restarts can store other bytes beside its earlier attempt's, and + identical bytes reuse the earlier object. """ - def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None) -> None: + def __init__( + self, + bucket_name: str, + *, + store: ArtifactStore | None = None, + content_addressed: bool = False, + ) -> None: """Bind to ``bucket_name``, which must be non-empty (there is no default).""" if not bucket_name: msg = "Stage 12 requires an explicitly configured private bucket" raise ValueError(msg) self.bucket_name = bucket_name self._store = store or ArtifactStore(bucket_name) + self._content_addressed = content_addressed + + def _key_digest(self, payload: bytes) -> str | None: + """Return the SHA-256 a content-addressed key names, else None.""" + return sha256(payload).hexdigest() if self._content_addressed else None def _write_immutable( self, @@ -316,9 +346,14 @@ def write_input( simulation: PlannedSimulationExecutionInput, ) -> ArtifactReference: """Store a simulation's planned execution input as canonical JSON.""" + payload = canonical_json_bytes(simulation.model_dump(mode="json")) return self._write_immutable( - input_path(prefix=prefix, role=simulation.role.value), - canonical_json_bytes(simulation.model_dump(mode="json")), + input_path( + prefix=prefix, + role=simulation.role.value, + content_sha256=self._key_digest(payload), + ), + payload, content_type=REPORT_MEDIA_TYPE, ) @@ -338,7 +373,11 @@ def write_simulation( uk_local_authority_metadata=uk_local_authority_metadata, ) artifact = self._write_immutable( - simulation_path(prefix=prefix, role=simulation.role.value), + simulation_path( + prefix=prefix, + role=simulation.role.value, + content_sha256=self._key_digest(payload), + ), payload, content_type=SIMULATION_MEDIA_TYPE, ) @@ -361,9 +400,13 @@ def write_aggregate( ) -> ArtifactReference: """Validate and store the run's aggregate report.""" validated = AggregateReportArtifactPayload.model_validate(payload) + encoded = canonical_json_bytes(validated.model_dump(mode="json")) return self._write_immutable( - aggregate_path(prefix=prefix), - canonical_json_bytes(validated.model_dump(mode="json")), + aggregate_path( + prefix=prefix, + content_sha256=self._key_digest(encoded), + ), + encoded, content_type=REPORT_MEDIA_TYPE, ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py new file mode 100644 index 000000000..3957db241 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py @@ -0,0 +1,81 @@ +"""Artifact keys that name their content, for calls Modal may restart.""" + +from __future__ import annotations + +from hashlib import sha256 + +import pytest +from fixtures.fake_object_store import InMemoryObjectStore +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore +from stage12_fake_stores import simulation_frames +from stage12_runtime_support import aggregate_payload, planned_simulation + +PREFIX = "stage-12-runs/staging/2026/11/report-run" + + +def _store(*, content_addressed: bool = True) -> Stage12ArtifactStore: + return Stage12ArtifactStore( + "private-stage12", + store=InMemoryObjectStore(), + content_addressed=content_addressed, + ) + + +def test_content_named_simulation_outputs_of_one_role_are_kept_apart() -> None: + store = _store() + baseline = planned_simulation(SimulationRole.BASELINE) + + first = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.0) + ) + # A restarted call can calculate different last bits for the same role. + second = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.000001) + ) + repeated = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.0) + ) + + digest = first.artifact.content_sha256 + assert first.artifact.uri == ( + f"gs://private-stage12/{PREFIX}/simulations/baseline-{digest}.parquet" + ) + assert second.artifact.uri != first.artifact.uri + assert repeated.artifact == first.artifact + assert store.read(second.artifact.uri) != store.read(first.artifact.uri) + + +def test_content_named_inputs_and_aggregates_name_their_digest() -> None: + store = _store() + reform = planned_simulation(SimulationRole.REFORM) + + planned_input = store.write_input(prefix=PREFIX, simulation=reform) + aggregate = store.write_aggregate(prefix=PREFIX, payload=aggregate_payload()) + + assert planned_input.uri == ( + f"gs://private-stage12/{PREFIX}/inputs/reform-" + f"{planned_input.content_sha256}.json" + ) + assert aggregate.uri == ( + f"gs://private-stage12/{PREFIX}/reports/aggregate-" + f"{aggregate.content_sha256}.json" + ) + assert sha256(store.read(aggregate.uri)).hexdigest() == aggregate.content_sha256 + + +def test_comparison_runs_keep_their_fixed_write_once_keys() -> None: + store = _store(content_addressed=False) + baseline = planned_simulation(SimulationRole.BASELINE) + + first = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.0) + ) + + assert first.artifact.uri == ( + f"gs://private-stage12/{PREFIX}/simulations/baseline.parquet" + ) + with pytest.raises(RuntimeError, match="other data"): + store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.000001) + ) From c2afdcd19fd9a28fda1c8d6fd1c1d87b7b86ae1b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:21:07 +0400 Subject: [PATCH 05/24] Add the report run store that the Modal worker role writes through The Stage 12 persistence library can now update the Public API's report_runs and simulation_runs rows: claim a report run under a Modal call (a pending one, or a running one no other call holds, which also lets a call Modal restarts resume its claim), mark it succeeded or failed only while that call holds it, claim a simulation run, record its output file as it succeeds, and fail a pending or running one. Each write is one conditional UPDATE and returns whether it changed the run and the run's status, so a repeated write changes nothing that is final. It also checks that a report run combines the simulation runs a call names. The mappings declare only the columns these writes touch. The tests run against a disposable Postgres migrated by the Public API's v2 chain, as the Modal worker role, and skip when no test database is configured. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 10 +- .../report_run_store.py | 304 ++++++++++++++++++ .../report_run_tables.py | 69 ++++ .../tests/conftest.py | 39 +++ .../tests/report_run_database.py | 223 +++++++++++++ .../tests/test_report_run_store.py | 249 ++++++++++++++ 6 files changed, 893 insertions(+), 1 deletion(-) create mode 100644 libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py create mode 100644 libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py create mode 100644 libs/policyengine-stage12-persistence/tests/conftest.py create mode 100644 libs/policyengine-stage12-persistence/tests/report_run_database.py create mode 100644 libs/policyengine-stage12-persistence/tests/test_report_run_store.py diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index e1344e788..e985543e6 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -1,5 +1,10 @@ -"""Direct DML for temporary Stage 12 comparison state.""" +"""Direct DML for temporary Stage 12 comparison state and for report runs.""" +from policyengine_stage12_persistence.report_run_store import ( + ReportRunStore, + RunStatus, + RunWrite, +) from policyengine_stage12_persistence.store import ( ConnectionOpener, Stage12PersistenceStore, @@ -13,6 +18,9 @@ __all__ = [ "ConnectionOpener", + "ReportRunStore", + "RunStatus", + "RunWrite", "Stage12PersistenceStore", "Stage12RuntimeAccessError", "create_stage12_engine", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py new file mode 100644 index 000000000..147730712 --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -0,0 +1,304 @@ +"""Conditional writes of report run and simulation run status, as the Modal worker role. + +Every write is one ``UPDATE`` conditional on the run's expected status, so a +call that Modal restarts after pre-emption, or a second writer, changes +nothing that is already final. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING + +from sqlalchemy import and_, func, or_, select, update + +from policyengine_stage12_persistence.report_run_tables import ( + report_run_simulation_runs, + report_runs, + simulation_runs, +) +from policyengine_stage12_persistence.store import ( + ConnectionOpener, + create_stage12_engine, + open_connection_once, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_contract.stage12_execution import ReportRunIds + from sqlalchemy import Engine, Table, Update + from sqlalchemy.engine import Connection + + +class RunStatus(StrEnum): + """A report run's or simulation run's status.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +ACTIVE_STATUSES = (RunStatus.PENDING.value, RunStatus.RUNNING.value) + + +@dataclass(frozen=True) +class RunWrite: + """One conditional write's outcome: whether it changed the run, and its status.""" + + changed: bool + status: RunStatus + + +class ReportRunStore: + """Status writes to the Public API's report runs and simulation runs.""" + + def __init__( + self, + database_url: str | None = None, + *, + engine: Engine | None = None, + open_connection: ConnectionOpener = open_connection_once, + ) -> None: + """Use ``engine`` if given, else build one from ``database_url``. + + Raises ValueError when neither is given or the URL is invalid. + """ + if engine is None: + if database_url is None: + msg = "report run database URL is required" + raise ValueError(msg) + engine = create_stage12_engine(database_url) + self.engine = engine + self._open_connection = open_connection + + @contextmanager + def _transaction(self, operation_name: str) -> Iterator[Connection]: + # As in the Stage 12 store, the opener only opens the connection; every + # statement and the COMMIT run after it returns. + connection = self._open_connection( + self.engine.connect, + operation_name=operation_name, + ) + with connection, connection.begin(): + yield connection + + def require_simulation_runs(self, run_ids: ReportRunIds) -> None: + """Raise ValueError unless the report run combines both simulation runs.""" + statement = select( + report_run_simulation_runs.c.role, + report_run_simulation_runs.c.simulation_run_id, + ).where(report_run_simulation_runs.c.report_run_id == run_ids.report_run_id) + with self._transaction("require_simulation_runs") as connection: + members = { + (row.role, row.simulation_run_id) + for row in connection.execute(statement) + } + expected = { + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + } + if not expected <= members: + msg = "the report run does not combine the simulation runs its call names" + raise ValueError(msg) + + def claim_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + claimed_at: datetime, + ) -> RunWrite: + """Mark the report run running under ``modal_call_id``. + + A pending run is claimed, and so is a running one that no other call + holds: the simulation service may mark it running before the call + starts, and Modal restarts a pre-empted call under the same call ID. + """ + statement = ( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + or_( + report_runs.c.status == RunStatus.PENDING.value, + and_( + report_runs.c.status == RunStatus.RUNNING.value, + or_( + report_runs.c.modal_call_id.is_(None), + report_runs.c.modal_call_id == modal_call_id, + ), + ), + ), + ) + .values( + status=RunStatus.RUNNING.value, + modal_call_id=modal_call_id, + started_at=func.coalesce(report_runs.c.started_at, claimed_at), + updated_at=claimed_at, + ) + ) + return self._write(report_runs, report_run_id, statement, "claim_report_run") + + def complete_report_run( + self, + run_ids: ReportRunIds, + *, + modal_call_id: str, + completed_at: datetime, + ) -> RunWrite: + """Mark the report run succeeded while ``modal_call_id`` holds it running.""" + statement = ( + update(report_runs) + .where( + report_runs.c.id == run_ids.report_run_id, + report_runs.c.status == RunStatus.RUNNING.value, + report_runs.c.modal_call_id == modal_call_id, + ) + .values( + status=RunStatus.SUCCEEDED.value, + error_message=None, + completed_at=completed_at, + updated_at=completed_at, + ) + ) + return self._write( + report_runs, run_ids.report_run_id, statement, "complete_report_run" + ) + + def fail_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark the report run failed while ``modal_call_id`` holds it running.""" + statement = ( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + report_runs.c.status == RunStatus.RUNNING.value, + report_runs.c.modal_call_id == modal_call_id, + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write(report_runs, report_run_id, statement, "fail_report_run") + + def claim_simulation_run( + self, + simulation_run_id: UUID, + *, + modal_call_id: str, + claimed_at: datetime, + ) -> RunWrite: + """Mark a pending or running simulation run running under ``modal_call_id``. + + A running one moves to the newest call, because the coordinator starts + a simulation again only after stopping its earlier call. + """ + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + simulation_runs.c.status.in_(ACTIVE_STATUSES), + ) + .values( + status=RunStatus.RUNNING.value, + modal_call_id=modal_call_id, + started_at=func.coalesce(simulation_runs.c.started_at, claimed_at), + updated_at=claimed_at, + ) + ) + return self._write( + simulation_runs, simulation_run_id, statement, "claim_simulation_run" + ) + + def complete_simulation_run( + self, + simulation_run_id: UUID, + *, + output_uri: str, + output_sha256: str, + completed_at: datetime, + ) -> RunWrite: + """Record a running simulation run's output file and mark it succeeded.""" + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + simulation_runs.c.status == RunStatus.RUNNING.value, + ) + .values( + status=RunStatus.SUCCEEDED.value, + output_uri=output_uri, + output_sha256=output_sha256, + error_message=None, + completed_at=completed_at, + updated_at=completed_at, + ) + ) + return self._write( + simulation_runs, simulation_run_id, statement, "complete_simulation_run" + ) + + def fail_simulation_run( + self, + simulation_run_id: UUID, + *, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark a pending or running simulation run failed with ``error_message``.""" + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + simulation_runs.c.status.in_(ACTIVE_STATUSES), + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write( + simulation_runs, simulation_run_id, statement, "fail_simulation_run" + ) + + def _write( + self, + table: Table, + run_id: UUID, + statement: Update, + operation_name: str, + ) -> RunWrite: + """Run a conditional update; return whether it changed the run, and its status. + + Raises LookupError when the run does not exist. + """ + with self._transaction(operation_name) as connection: + changed = connection.execute( + statement.returning(table.c.status) + ).scalar_one_or_none() + if changed is not None: + return RunWrite(changed=True, status=RunStatus(changed)) + current = connection.execute( + select(table.c.status).where(table.c.id == run_id) + ).scalar_one_or_none() + if current is None: + noun = "report run" if table is report_runs else "simulation run" + msg = f"{noun} {run_id} does not exist" + raise LookupError(msg) + return RunWrite(changed=False, status=RunStatus(current)) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py new file mode 100644 index 000000000..7ef743a10 --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -0,0 +1,69 @@ +"""SQLAlchemy mappings of the Public API's run tables that report calls update. + +``policyengine-api`` owns these tables and every schema migration. The +mappings declare only the columns the Modal worker role reads or writes, so +the coordinator and workers can issue typed DML; runtime code never emits DDL. +""" + +from __future__ import annotations + +from sqlalchemy import JSON, Column, DateTime, MetaData, String, Table +from sqlalchemy.dialects.postgresql import ENUM, UUID + +RUN_STATUS_VALUES = ("pending", "running", "succeeded", "failed") +SIMULATION_ROLE_VALUES = ("baseline", "reform") + +report_run_status_type = ENUM( + *RUN_STATUS_VALUES, + name="v2_report_run_status", + create_type=False, +) +simulation_run_status_type = ENUM( + *RUN_STATUS_VALUES, + name="v2_simulation_status", + create_type=False, +) +simulation_role_type = ENUM( + *SIMULATION_ROLE_VALUES, + name="v2_simulation_role", + create_type=False, +) + +metadata = MetaData(schema="public") + +report_runs = Table( + "report_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("status", report_run_status_type, nullable=False), + Column("modal_call_id", String(255)), + Column("error_message", String), + Column("started_at", DateTime(timezone=True)), + Column("completed_at", DateTime(timezone=True)), + Column("updated_at", DateTime(timezone=True), nullable=False), + Column("labor_supply_result", JSON(none_as_null=True)), + Column("cliff_impact_result", JSON(none_as_null=True)), +) + +simulation_runs = Table( + "simulation_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("status", simulation_run_status_type, nullable=False), + Column("modal_call_id", String(255)), + Column("error_message", String), + Column("started_at", DateTime(timezone=True)), + Column("completed_at", DateTime(timezone=True)), + Column("updated_at", DateTime(timezone=True), nullable=False), + Column("output_uri", String(2048)), + Column("output_sha256", String(64)), +) + +report_run_simulation_runs = Table( + "report_run_simulation_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("report_run_id", UUID(as_uuid=True), nullable=False), + Column("simulation_run_id", UUID(as_uuid=True), nullable=False), + Column("role", simulation_role_type, nullable=False), +) diff --git a/libs/policyengine-stage12-persistence/tests/conftest.py b/libs/policyengine-stage12-persistence/tests/conftest.py new file mode 100644 index 000000000..84d6ceb0b --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/conftest.py @@ -0,0 +1,39 @@ +"""Fixtures for the report run tests against a disposable Postgres.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import pytest +from report_run_database import ( + SEED_URL_ENV, + ReportRunSeeder, + SeededReportRun, + worker_database_url, +) +from sqlalchemy import create_engine + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@pytest.fixture +def seeder() -> Iterator[ReportRunSeeder]: + """Yield a seeder on the seed role's connection, or skip without a database.""" + worker_database_url() + engine = create_engine(os.environ[SEED_URL_ENV]) + try: + yield ReportRunSeeder(engine) + finally: + engine.dispose() + + +@pytest.fixture +def seeded_run(seeder: ReportRunSeeder) -> Iterator[SeededReportRun]: + """Yield one freshly seeded report run, removed again after the test.""" + seeded = seeder.seed_report_run() + try: + yield seeded + finally: + seeder.remove(seeded) diff --git a/libs/policyengine-stage12-persistence/tests/report_run_database.py b/libs/policyengine-stage12-persistence/tests/report_run_database.py new file mode 100644 index 000000000..b4a39b67e --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/report_run_database.py @@ -0,0 +1,223 @@ +"""A disposable Postgres holding the Public API's v2 tables, for report run tests. + +The tests run only when both URLs name a disposable database that the Public +API's v2 migrations created and whose declared privileges were applied: + +- ``REPORT_RUN_TEST_DATABASE_URL`` logs in as the Modal worker role, which the + store under test uses; +- ``REPORT_RUN_TEST_SEED_DATABASE_URL`` logs in as a role that may insert the + rows the simulation service would create, so each test can set them up. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any +from uuid import UUID, uuid4 + +import pytest +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from sqlalchemy import Connection, Engine, MetaData, Table, delete, insert, select + +WORKER_URL_ENV = "REPORT_RUN_TEST_DATABASE_URL" +SEED_URL_ENV = "REPORT_RUN_TEST_SEED_DATABASE_URL" +# The catalog rows are needed only while simulations and reports keep their +# foreign keys to the catalog; without those keys they are harmless. +SEEDED_TABLES = ( + "tax_benefit_models", + "tax_benefit_model_versions", + "datasets", + "simulations", + "reports", + "report_runs", + "simulation_runs", + "report_run_simulation_runs", +) +NOW = datetime(2026, 11, 4, tzinfo=UTC) +VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "6.12.0"} + + +def worker_database_url() -> str: + """Return the Modal worker role's test URL, or skip when none is set.""" + url = os.environ.get(WORKER_URL_ENV) + if not url or not os.environ.get(SEED_URL_ENV): + pytest.skip(f"{WORKER_URL_ENV} and {SEED_URL_ENV} name no test database") + return url + + +@dataclass(frozen=True) +class SeededReportRun: + """A report run and its two simulation runs, as the simulation service makes.""" + + run_ids: ReportRunIds + simulation_ids: tuple[UUID, UUID] + model_id: UUID + + +class ReportRunSeeder: + """Insert and read back report runs as the seed role.""" + + def __init__(self, engine: Engine) -> None: + """Reflect the seeded tables through ``engine``; this reads, never alters.""" + self.engine = engine + self.metadata = MetaData(schema="public") + self.metadata.reflect(engine, only=list(SEEDED_TABLES)) + + def table(self, name: str) -> Table: + """Return the reflected table ``name``.""" + return self.metadata.tables[f"public.{name}"] + + def seed_report_run(self) -> SeededReportRun: + """Insert a pending report run whose baseline and reform runs are pending.""" + model_id, version_id, dataset_id = uuid4(), uuid4(), uuid4() + simulation_ids = (uuid4(), uuid4()) + report_id = uuid4() + run_ids = ReportRunIds( + report_run_id=uuid4(), + baseline_simulation_run_id=uuid4(), + reform_simulation_run_id=uuid4(), + ) + with self.engine.begin() as connection: + self._insert_catalog(connection, model_id, version_id, dataset_id) + for simulation_id in simulation_ids: + connection.execute( + insert(self.table("simulations")).values( + id=simulation_id, + simulation_type="economy", + dataset_id=dataset_id, + tax_benefit_model_version_id=version_id, + year=2026, + ) + ) + connection.execute( + insert(self.table("reports")).values( + id=report_id, + label="Report run test", + country="us", + tax_benefit_model_id=model_id, + baseline_simulation_id=simulation_ids[0], + reform_simulation_id=simulation_ids[1], + year=2026, + inputs={}, + ) + ) + connection.execute( + insert(self.table("report_runs")).values( + id=run_ids.report_run_id, + report_id=report_id, + status="pending", + trigger="initial", + **VERSIONS, + ) + ) + for run_id, simulation_id, role in ( + (run_ids.baseline_simulation_run_id, simulation_ids[0], "baseline"), + (run_ids.reform_simulation_run_id, simulation_ids[1], "reform"), + ): + self._insert_simulation_run( + connection, run_ids.report_run_id, run_id, simulation_id, role + ) + return SeededReportRun( + run_ids=run_ids, simulation_ids=simulation_ids, model_id=model_id + ) + + def _insert_catalog( + self, connection: Connection, model_id: UUID, version_id: UUID, dataset_id: UUID + ) -> None: + connection.execute( + insert(self.table("tax_benefit_models")).values( + id=model_id, name=f"test-{model_id.hex[:24]}" + ) + ) + connection.execute( + insert(self.table("tax_benefit_model_versions")).values( + id=version_id, + model_id=model_id, + version="1.0.0", + current_law_id=1, + metadata_time_periods={}, + ) + ) + connection.execute( + insert(self.table("datasets")).values( + id=dataset_id, + name="report-run-test", + year=2026, + is_output_dataset=False, + tax_benefit_model_version_id=version_id, + ) + ) + + def _insert_simulation_run( + self, + connection: Connection, + report_run_id: UUID, + simulation_run_id: UUID, + simulation_id: UUID, + role: str, + ) -> None: + version_id = connection.execute( + select(self.table("simulations").c.tax_benefit_model_version_id).where( + self.table("simulations").c.id == simulation_id + ) + ).scalar_one() + connection.execute( + insert(self.table("simulation_runs")).values( + id=simulation_run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=version_id, + status="pending", + **VERSIONS, + ) + ) + connection.execute( + insert(self.table("report_run_simulation_runs")).values( + id=uuid4(), + report_run_id=report_run_id, + simulation_run_id=simulation_run_id, + role=role, + ) + ) + + def row(self, table_name: str, row_id: UUID) -> dict[str, Any]: + """Return the row ``row_id`` of ``table_name`` as a dictionary.""" + table = self.table(table_name) + with self.engine.connect() as connection: + row = connection.execute(select(table).where(table.c.id == row_id)).one() + return dict(row._mapping) + + def rows(self, table_name: str, **filters: object) -> list[dict[str, Any]]: + """Return the rows of ``table_name`` whose columns equal ``filters``.""" + table = self.table(table_name) + statement = select(table).where( + *(table.c[name] == value for name, value in filters.items()) + ) + with self.engine.connect() as connection: + return [dict(row._mapping) for row in connection.execute(statement)] + + def remove(self, seeded: SeededReportRun) -> None: + """Delete the seeded rows; the runs and their results cascade.""" + reports = self.table("reports") + simulations = self.table("simulations") + datasets = self.table("datasets") + versions = self.table("tax_benefit_model_versions") + models = self.table("tax_benefit_models") + version_ids = select(versions.c.id).where( + versions.c.model_id == seeded.model_id + ) + with self.engine.begin() as connection: + connection.execute( + delete(reports).where(reports.c.tax_benefit_model_id == seeded.model_id) + ) + connection.execute( + delete(simulations).where(simulations.c.id.in_(seeded.simulation_ids)) + ) + connection.execute( + delete(datasets).where( + datasets.c.tax_benefit_model_version_id.in_(version_ids) + ) + ) + # Deleting the model removes its versions. + connection.execute(delete(models).where(models.c.id == seeded.model_id)) diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_store.py b/libs/policyengine-stage12-persistence/tests/test_report_run_store.py new file mode 100644 index 000000000..6d2545e74 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_store.py @@ -0,0 +1,249 @@ +"""The Modal worker role's conditional writes to report runs and simulation runs. + +These run against a disposable Postgres that the Public API's v2 migrations +created, connected as the Modal worker role (see ``report_run_database``). +""" + +from __future__ import annotations + +from datetime import timedelta +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_stage12_persistence import ( + ReportRunStore, + RunStatus, + RunWrite, +) +from report_run_database import NOW, SeededReportRun, worker_database_url +from sqlalchemy import func, insert, select +from sqlalchemy.exc import ProgrammingError + +LATER = NOW + timedelta(minutes=5) + + +@pytest.fixture +def store(): + run_store = ReportRunStore(worker_database_url()) + try: + yield run_store + finally: + run_store.engine.dispose() + + +def test_the_store_connects_as_the_modal_worker_role(store) -> None: + with store.engine.connect() as connection: + role = connection.execute(select(func.current_user())).scalar_one() + + assert role == "policyengine_v2_modal_worker" + + +def test_the_modal_worker_role_cannot_create_a_report_run(store, seeded_run) -> None: + from policyengine_stage12_persistence.report_run_tables import report_runs + + with ( + pytest.raises(ProgrammingError, match="permission denied"), + store.engine.begin() as connection, + ): + connection.execute( + insert(report_runs).values(id=uuid4(), status="pending", updated_at=NOW) + ) + + +def test_a_report_run_is_claimed_once_and_resumed_by_the_same_call( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + + first = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + other_call = store.claim_report_run( + report_run_id, modal_call_id="fc-other", claimed_at=LATER + ) + restarted = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=LATER + ) + + assert first == RunWrite(changed=True, status=RunStatus.RUNNING) + assert other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert restarted == RunWrite(changed=True, status=RunStatus.RUNNING) + row = seeder.row("report_runs", report_run_id) + assert row["modal_call_id"] == "fc-coordinator" + assert row["started_at"] == NOW + assert row["updated_at"] == LATER + + +def test_a_report_run_the_simulation_service_marked_running_is_claimed( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + report_runs = seeder.table("report_runs") + with seeder.engine.begin() as connection: + connection.execute( + report_runs.update() + .where(report_runs.c.id == report_run_id) + .values(status="running", started_at=NOW) + ) + + claimed = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=LATER + ) + + assert claimed == RunWrite(changed=True, status=RunStatus.RUNNING) + assert seeder.row("report_runs", report_run_id)["started_at"] == NOW + + +def test_a_report_run_completes_only_while_its_call_holds_it( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + + by_other_call = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-other", completed_at=LATER + ) + completed = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER + ) + repeated = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER + ) + reclaimed = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=LATER + ) + + assert by_other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert completed == RunWrite(changed=True, status=RunStatus.SUCCEEDED) + assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert reclaimed == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + row = seeder.row("report_runs", report_run_id) + assert row["status"] == "succeeded" + assert row["completed_at"] == LATER + assert row["error_message"] is None + + +def test_a_report_run_fails_once_and_then_never_completes( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + + failed = store.fail_report_run( + report_run_id, + modal_call_id="fc-coordinator", + error_message="simulation_execution_failed: the reform failed", + failed_at=LATER, + ) + again = store.fail_report_run( + report_run_id, + modal_call_id="fc-coordinator", + error_message="another failure", + failed_at=LATER, + ) + completed = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER + ) + + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + assert again == RunWrite(changed=False, status=RunStatus.FAILED) + assert completed == RunWrite(changed=False, status=RunStatus.FAILED) + row = seeder.row("report_runs", report_run_id) + assert row["error_message"] == "simulation_execution_failed: the reform failed" + assert row["completed_at"] == LATER + + +def test_a_simulation_run_moves_to_the_latest_call_until_it_finishes( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_id = seeded_run.run_ids.baseline_simulation_run_id + + first = store.claim_simulation_run( + run_id, modal_call_id="fc-worker-1", claimed_at=NOW + ) + # The coordinator started this simulation again after a restart. + second = store.claim_simulation_run( + run_id, modal_call_id="fc-worker-2", claimed_at=LATER + ) + completed = store.complete_simulation_run( + run_id, + output_uri="gs://private/simulations/baseline.parquet", + output_sha256="a" * 64, + completed_at=LATER, + ) + repeated = store.complete_simulation_run( + run_id, + output_uri="gs://private/simulations/baseline-other.parquet", + output_sha256="b" * 64, + completed_at=LATER, + ) + failed_after = store.fail_simulation_run( + run_id, error_message="cancelled", failed_at=LATER + ) + claimed_after = store.claim_simulation_run( + run_id, modal_call_id="fc-worker-3", claimed_at=LATER + ) + + assert first == RunWrite(changed=True, status=RunStatus.RUNNING) + assert second == RunWrite(changed=True, status=RunStatus.RUNNING) + assert completed == RunWrite(changed=True, status=RunStatus.SUCCEEDED) + assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert failed_after == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert claimed_after == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + row = seeder.row("simulation_runs", run_id) + assert row["modal_call_id"] == "fc-worker-2" + assert row["started_at"] == NOW + assert row["output_uri"] == "gs://private/simulations/baseline.parquet" + assert row["output_sha256"] == "a" * 64 + + +def test_a_pending_or_running_simulation_run_can_fail_but_not_resume( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_id = seeded_run.run_ids.reform_simulation_run_id + + failed = store.fail_simulation_run( + run_id, error_message="simulation_cancelled: stopped", failed_at=NOW + ) + claimed = store.claim_simulation_run( + run_id, modal_call_id="fc-worker", claimed_at=LATER + ) + completed = store.complete_simulation_run( + run_id, + output_uri="gs://private/x.parquet", + output_sha256="c" * 64, + completed_at=LATER, + ) + + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + assert claimed == RunWrite(changed=False, status=RunStatus.FAILED) + assert completed == RunWrite(changed=False, status=RunStatus.FAILED) + row = seeder.row("simulation_runs", run_id) + assert row["error_message"] == "simulation_cancelled: stopped" + assert row["output_uri"] is None + + +def test_a_report_run_must_combine_the_simulation_runs_its_call_names( + store, seeded_run: SeededReportRun +) -> None: + store.require_simulation_runs(seeded_run.run_ids) + + swapped = ReportRunIds( + report_run_id=seeded_run.run_ids.report_run_id, + baseline_simulation_run_id=seeded_run.run_ids.reform_simulation_run_id, + reform_simulation_run_id=seeded_run.run_ids.baseline_simulation_run_id, + ) + with pytest.raises(ValueError, match="does not combine"): + store.require_simulation_runs(swapped) + + +def test_a_missing_run_raises_lookup_error(store, seeded_run) -> None: + with pytest.raises(LookupError, match="report run"): + store.claim_report_run(uuid4(), modal_call_id="fc-coordinator", claimed_at=NOW) + with pytest.raises(LookupError, match="simulation run"): + store.fail_simulation_run(uuid4(), error_message="failed", failed_at=NOW) From 915a82810b2c8c0364b91924056cfc316896ed1b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:21:20 +0400 Subject: [PATCH 06/24] Record a report call's simulation runs from its Stage 12 workers When a worker's invocation context names a report run, the worker keeps its Stage 12 child record as before and also records its own simulation run, as the Modal worker role: it claims the run under its Modal call ID after claiming the child record, records the output file and marks the run succeeded before marking the child record succeeded, and marks the run failed after a recorded failure. A wrapper around the comparison store makes these writes, so the coordinator's own failure and cancellation records reach the simulation runs too. A simulation run that is already failed refuses the claim, which the worker treats as a cancellation. Success comes first so a worker that Modal restarts between the two writes finds its child record still running and runs again, and its second output never replaces the first. Such a worker names its artifacts by their digest. The US and UK worker functions now pass their Modal call ID; comparison runs record nothing new. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 2 + .../stage12_runtime/dependencies.py | 12 +- .../stage12_runtime/report_run_access.py | 90 ++++++++ .../stage12_runtime/report_run_recording.py | 181 +++++++++++++++++ .../stage12_runtime/simulation.py | 28 ++- .../tests/stage12_fake_report_runs.py | 150 ++++++++++++++ .../tests/test_stage12_modal_app.py | 24 +++ .../tests/test_stage12_report_run_worker.py | 192 ++++++++++++++++++ 8 files changed, 673 insertions(+), 6 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py create mode 100644 projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 623e4334e..e9c97852e 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -360,6 +360,7 @@ def segmented_calculator( segmented_calculator=segmented_calculator, runtime=runtime, ), + modal_call_id=modal.current_function_call_id(), ) finally: runtime.shutdown() @@ -404,6 +405,7 @@ def run_single_simulation_uk( dependencies=stage12_runtime.SimulationWorkerDependencies( runtime=runtime ), + modal_call_id=modal.current_function_call_id(), ) finally: runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index 747eb839d..120d3946c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -177,6 +177,12 @@ def runtime_store( ) -def artifact_store() -> Stage12ArtifactStore: - """Build the store for ``STAGE12_ARTIFACT_BUCKET``; ValueError if it is unset.""" - return Stage12ArtifactStore(os.environ.get("STAGE12_ARTIFACT_BUCKET", "")) +def artifact_store(*, content_addressed: bool = False) -> Stage12ArtifactStore: + """Build the store for ``STAGE12_ARTIFACT_BUCKET``; ValueError if it is unset. + + A ``content_addressed`` store names each artifact's digest in its key. + """ + return Stage12ArtifactStore( + os.environ.get("STAGE12_ARTIFACT_BUCKET", ""), + content_addressed=content_addressed, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py index 25b63edc9..3926f3972 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py @@ -8,7 +8,97 @@ from __future__ import annotations +import os +from functools import partial +from typing import TYPE_CHECKING, Protocol + +from policyengine_stage12_persistence import ReportRunStore + +from .database_retry import open_connection_with_retry + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.stage12_execution import ReportRunIds + from policyengine_stage12_persistence import RunWrite + # The Modal secret and the one key it must hold: a postgresql:// URL for the # Modal worker role on the database that holds the Public API's run tables. MODAL_WORKER_DATABASE_SECRET_NAME = "policyengine-modal-worker-database" # noqa: S105 # the Modal secret's name, not its value MODAL_WORKER_DATABASE_URL_ENV = "MODAL_WORKER_DATABASE_URL" + + +class ReportRunWriter(Protocol): + """Report run and simulation run writes; ``ReportRunStore`` implements it. + + Each write is conditional on the run's status and returns whether it + changed the run, and the run's status afterwards. + """ + + def require_simulation_runs(self, run_ids: ReportRunIds) -> None: + """Raise ValueError unless the report run combines both simulation runs.""" + ... + + def claim_report_run( + self, report_run_id: UUID, *, modal_call_id: str, claimed_at: datetime + ) -> RunWrite: + """Mark a pending report run, or one no other call holds, running.""" + ... + + def complete_report_run( + self, + run_ids: ReportRunIds, + *, + modal_call_id: str, + completed_at: datetime, + ) -> RunWrite: + """Mark the report run succeeded while ``modal_call_id`` holds it.""" + ... + + def fail_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark the report run failed while ``modal_call_id`` holds it.""" + ... + + def claim_simulation_run( + self, simulation_run_id: UUID, *, modal_call_id: str, claimed_at: datetime + ) -> RunWrite: + """Mark a pending or running simulation run running under this call.""" + ... + + def complete_simulation_run( + self, + simulation_run_id: UUID, + *, + output_uri: str, + output_sha256: str, + completed_at: datetime, + ) -> RunWrite: + """Record a running simulation run's output file and mark it succeeded.""" + ... + + def fail_simulation_run( + self, simulation_run_id: UUID, *, error_message: str, failed_at: datetime + ) -> RunWrite: + """Mark a pending or running simulation run failed.""" + ... + + +def report_run_store(runtime: ObservabilityRuntime | None = None) -> ReportRunStore: + """Build the report run store from ``MODAL_WORKER_DATABASE_URL``. + + Raises ``ValueError`` when the variable is unset or not a complete + ``postgresql://`` URL. + """ + return ReportRunStore( + os.environ.get(MODAL_WORKER_DATABASE_URL_ENV, ""), + open_connection=partial(open_connection_with_retry, runtime=runtime), + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py new file mode 100644 index 000000000..c7a783668 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -0,0 +1,181 @@ +"""Recording a report call's simulation runs beside its Stage 12 child records. + +A call that names a report run keeps the comparison path's Stage 12 records, +which coordinate its workers. Every claim, success and failure of a child +record is also written to that child's simulation run, as the Modal worker +role, so the workers and the coordinator need no other change to record them. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from policyengine_stage12_persistence import RunStatus + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_contract.stage12_execution import ( + ComparisonReportPersistenceResult, + ComparisonReportRecord, + ComparisonSimulationPersistenceResult, + ComparisonSimulationRecord, + ReportRunIds, + ) + + from .dependencies import ComparisonStore + from .report_run_access import ReportRunWriter + + +def run_error_message(error_code: str | None, error_summary: str | None) -> str: + """Return the redacted failure a run records: its code, then its summary.""" + return f"{error_code or 'simulation_failed'}: {error_summary or 'no summary'}" + + +class SimulationRunRecordingStore: + """A ``ComparisonStore`` that also records each child's simulation run. + + ``modal_call_id`` is the worker's own call, recorded on its simulation run + when it claims it; the coordinator never claims one and passes None. + """ + + def __init__( + self, + store: ComparisonStore, + runs: ReportRunWriter, + run_ids: ReportRunIds, + *, + modal_call_id: str | None = None, + ) -> None: + """Record on ``run_ids``'s simulation runs through ``runs``.""" + self._store = store + self._runs = runs + self._run_ids = run_ids + self._modal_call_id = modal_call_id + + def create_or_resolve_report( + self, record: ComparisonReportRecord + ) -> ComparisonReportPersistenceResult: + """Delegate to the Stage 12 store.""" + return self._store.create_or_resolve_report(record) + + def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Delegate to the Stage 12 store.""" + return self._store.get_report(evaluation_id) + + def replace_report(self, record: ComparisonReportRecord) -> ComparisonReportRecord: + """Delegate to the Stage 12 store.""" + return self._store.replace_report(record) + + def replace_report_result_comparison( + self, record: ComparisonReportRecord + ) -> ComparisonReportRecord: + """Delegate to the Stage 12 store.""" + return self._store.replace_report_result_comparison(record) + + def create_or_resolve_simulation( + self, record: ComparisonSimulationRecord + ) -> ComparisonSimulationPersistenceResult: + """Delegate to the Stage 12 store.""" + return self._store.create_or_resolve_simulation(record) + + def get_simulation( + self, simulation_execution_id: UUID + ) -> ComparisonSimulationRecord: + """Delegate to the Stage 12 store.""" + return self._store.get_simulation(simulation_execution_id) + + def replace_simulation( + self, record: ComparisonSimulationRecord + ) -> ComparisonSimulationRecord: + """Delegate to the Stage 12 store; a dispatch placeholder records nothing.""" + return self._store.replace_simulation(record) + + def attach_simulation_invocation( + self, + simulation_execution_id: UUID, + *, + expected_placeholder: str, + modal_invocation_id: str, + updated_at: datetime, + ) -> ComparisonSimulationRecord: + """Delegate to the Stage 12 store.""" + return self._store.attach_simulation_invocation( + simulation_execution_id, + expected_placeholder=expected_placeholder, + modal_invocation_id=modal_invocation_id, + updated_at=updated_at, + ) + + def replace_simulation_if_status( + self, + record: ComparisonSimulationRecord, + *, + expected_status: ComparisonRunLifecycleStatus, + ) -> tuple[ComparisonSimulationRecord, bool]: + """Write the child record and its simulation run, in a restart-safe order. + + A success is recorded on the simulation run first, so a call Modal + restarts between the two writes finds its Stage 12 record still + running and runs again without recording a second output. A claim or + failure is recorded on the simulation run only after the Stage 12 + record changed. A simulation run that is already failed refuses the + child's claim and success, which its worker treats as a cancellation. + """ + simulation_run_id = self._run_ids.simulation_run_id(record.role) + written_at = record.updated_at or datetime.now(UTC) + if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: + completed = self._complete(simulation_run_id, record, written_at) + if completed is not RunStatus.SUCCEEDED: + return self._store.get_simulation(record.simulation_execution_id), False + return self._store.replace_simulation_if_status( + record, expected_status=expected_status + ) + stored, replaced = self._store.replace_simulation_if_status( + record, expected_status=expected_status + ) + if not replaced: + return stored, replaced + if record.status is ComparisonRunLifecycleStatus.RUNNING: + claimed = self._runs.claim_simulation_run( + simulation_run_id, + modal_call_id=self._require_modal_call_id(), + claimed_at=written_at, + ) + if claimed.status is RunStatus.FAILED: + return stored, False + elif record.status is ComparisonRunLifecycleStatus.FAILED: + self._runs.fail_simulation_run( + simulation_run_id, + error_message=run_error_message( + record.error_code, record.error_summary + ), + failed_at=record.completed_at or written_at, + ) + return stored, replaced + + def _complete( + self, + simulation_run_id: UUID, + record: ComparisonSimulationRecord, + written_at: datetime, + ) -> RunStatus: + if record.output_uri is None or record.output_sha256 is None: + msg = "a succeeded simulation must name its output file" + raise ValueError(msg) + return self._runs.complete_simulation_run( + simulation_run_id, + output_uri=record.output_uri, + output_sha256=record.output_sha256, + completed_at=record.completed_at or written_at, + ).status + + def _require_modal_call_id(self) -> str: + if self._modal_call_id is None: + msg = "only a worker with its Modal call ID claims a simulation run" + raise ValueError(msg) + return self._modal_call_id diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 1975fc8e0..c99a9a361 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -57,6 +57,8 @@ ) from .output_planning import apply_output_plan, validate_output_frames from .partition import should_segment_simulation +from .report_run_access import report_run_store +from .report_run_recording import SimulationRunRecordingStore if TYPE_CHECKING: from contextlib import AbstractContextManager @@ -74,6 +76,8 @@ RegionResolution, ) + from .report_run_access import ReportRunWriter + logger = logging.getLogger(__name__) @@ -473,7 +477,8 @@ class SimulationWorkerDependencies: The defaults are the Stage 12 database store, the private artifact bucket, the standard country-package calculation, the deployed region-group - segment runner and no telemetry runtime. + segment runner, no telemetry runtime, and the report run store, which only + a call that names a report run uses. """ store: ComparisonStore | None = None @@ -490,6 +495,7 @@ class SimulationWorkerDependencies: | None ) = None runtime: ObservabilityRuntime | None = None + report_runs: ReportRunWriter | None = None def run_single_simulation( @@ -498,12 +504,16 @@ def run_single_simulation( *, required_country: CountryId, dependencies: SimulationWorkerDependencies | None = None, + modal_call_id: str | None = None, ) -> dict[str, Any]: """Calculate one planned simulation, store its artifact, and return its descriptor. A child already marked succeeded returns its stored descriptor without recalculating. A failure after the run is marked running marks it FAILED - and raises a ``RuntimeError`` that carries no calculation detail. + and raises a ``RuntimeError`` that carries no calculation detail. A call + whose context names a report run also records its simulation run's + claim, output file and failure under ``modal_call_id``, and names its + artifacts by their digest, so a call Modal restarts can store them again. """ dependencies = dependencies or SimulationWorkerDependencies() runtime = dependencies.runtime @@ -511,7 +521,19 @@ def run_single_simulation( context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) persistence = dependencies.store or runtime_store(runtime) - artifact_storage = dependencies.artifacts or artifact_store() + if context.report_run is not None: + if modal_call_id is None: + msg = "a report run's worker needs its Modal call ID" + raise ValueError(msg) + persistence = SimulationRunRecordingStore( + persistence, + dependencies.report_runs or report_run_store(runtime), + context.report_run, + modal_call_id=modal_call_id, + ) + artifact_storage = dependencies.artifacts or artifact_store( + content_addressed=context.report_run is not None + ) child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: return descriptor_from_record(child, simulation).model_dump(mode="json") diff --git a/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py b/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py new file mode 100644 index 000000000..12dfcc9f3 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py @@ -0,0 +1,150 @@ +"""An in-memory stand-in for the report run store, for Stage 12 runtime tests. + +It applies the same conditions as ``ReportRunStore``'s conditional updates, +which the persistence library tests against Postgres. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from uuid import UUID + +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_stage12_persistence import RunStatus, RunWrite +from stage12_runtime_support import invocation_context + +REPORT_RUN_ID = UUID("00000000-0000-0000-0000-0000000000a1") +BASELINE_RUN_ID = UUID("00000000-0000-0000-0000-0000000000b1") +REFORM_RUN_ID = UUID("00000000-0000-0000-0000-0000000000c1") +REPORT_RUN_IDS = ReportRunIds( + report_run_id=REPORT_RUN_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_run_id=REFORM_RUN_ID, +) +ACTIVE = (RunStatus.PENDING, RunStatus.RUNNING) + + +def report_run_context(): + """Return the shared invocation context, naming the report run and its runs.""" + return invocation_context().model_copy(update={"report_run": REPORT_RUN_IDS}) + + +@dataclass +class FakeRun: + """One report run's or simulation run's mutable columns.""" + + status: RunStatus = RunStatus.PENDING + modal_call_id: str | None = None + error_message: str | None = None + output_uri: str | None = None + output_sha256: str | None = None + writes: list[str] = field(default_factory=list) + + +class FakeReportRuns: + """Report runs and simulation runs held in memory, written conditionally.""" + + def __init__(self, events: list[str] | None = None) -> None: + self.report_runs = {REPORT_RUN_ID: FakeRun()} + self.simulation_runs = {BASELINE_RUN_ID: FakeRun(), REFORM_RUN_ID: FakeRun()} + self.links = {("baseline", BASELINE_RUN_ID), ("reform", REFORM_RUN_ID)} + # Shared with other fakes when a test checks the order of writes. + self.events = events if events is not None else [] + + def require_simulation_runs(self, run_ids: ReportRunIds) -> None: + expected = { + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + } + if not expected <= self.links: + msg = "the report run does not combine the simulation runs its call names" + raise ValueError(msg) + + def claim_report_run(self, report_run_id, *, modal_call_id, claimed_at): + run = self._run(self.report_runs, report_run_id, "report run") + claimable = run.status is RunStatus.PENDING or ( + run.status is RunStatus.RUNNING + and run.modal_call_id in {None, modal_call_id} + ) + return self._apply( + run, + "claim_report_run", + changed=claimable, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + ) + + def complete_report_run(self, run_ids, *, modal_call_id, completed_at): + run = self._run(self.report_runs, run_ids.report_run_id, "report run") + holds_it = self._holds(run, modal_call_id) + return self._apply( + run, + "complete_report_run", + changed=holds_it, + status=RunStatus.SUCCEEDED, + ) + + def fail_report_run( + self, report_run_id, *, modal_call_id, error_message, failed_at + ): + run = self._run(self.report_runs, report_run_id, "report run") + holds_it = self._holds(run, modal_call_id) + return self._apply( + run, + "fail_report_run", + changed=holds_it, + status=RunStatus.FAILED, + error_message=error_message, + ) + + def claim_simulation_run(self, simulation_run_id, *, modal_call_id, claimed_at): + run = self._run(self.simulation_runs, simulation_run_id, "simulation run") + return self._apply( + run, + "claim_simulation_run", + changed=run.status in ACTIVE, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + ) + + def complete_simulation_run( + self, simulation_run_id, *, output_uri, output_sha256, completed_at + ): + run = self._run(self.simulation_runs, simulation_run_id, "simulation run") + return self._apply( + run, + "complete_simulation_run", + changed=run.status is RunStatus.RUNNING, + status=RunStatus.SUCCEEDED, + output_uri=output_uri, + output_sha256=output_sha256, + ) + + def fail_simulation_run(self, simulation_run_id, *, error_message, failed_at): + run = self._run(self.simulation_runs, simulation_run_id, "simulation run") + return self._apply( + run, + "fail_simulation_run", + changed=run.status in ACTIVE, + status=RunStatus.FAILED, + error_message=error_message, + ) + + @staticmethod + def _holds(run: FakeRun, modal_call_id: str) -> bool: + return run.status is RunStatus.RUNNING and run.modal_call_id == modal_call_id + + @staticmethod + def _run(runs: dict[UUID, FakeRun], run_id: UUID, noun: str) -> FakeRun: + if run_id not in runs: + msg = f"{noun} {run_id} does not exist" + raise LookupError(msg) + return runs[run_id] + + def _apply(self, run: FakeRun, operation: str, *, changed: bool, **values): + run.writes.append(operation) + self.events.append(operation) + if changed: + for name, value in values.items(): + setattr(run, name, value) + return RunWrite(changed=changed, status=run.status) 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 75b3d3e44..73ad26255 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -374,6 +374,30 @@ def test_only_report_run_writers_receive_the_modal_worker_database_secret( assert functions[other]["secrets"] == module.worker_secrets +@pytest.mark.parametrize( + "function_name", ["run_single_simulation_us", "run_single_simulation_uk"] +) +def test_simulation_workers_pass_their_own_modal_call_id( + monkeypatch, function_name +) -> None: + module = _load(monkeypatch) + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "fc-simulation-worker", + raising=False, + ) + received = {} + + def implementation(*_args, **kwargs): + received.update(kwargs) + return {"role": "baseline"} + + _invoke(monkeypatch, module, function_name, implementation) + + assert received["modal_call_id"] == "fc-simulation-worker" + + @pytest.mark.parametrize( "function_name", [ diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py new file mode 100644 index 000000000..10e7d1687 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py @@ -0,0 +1,192 @@ +"""A Stage 12 worker records its simulation run's status and output file.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationWorkerDependencies, + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + simulation as simulation_module, +) +from policyengine_stage12_persistence import RunStatus +from stage12_fake_report_runs import ( + BASELINE_RUN_ID, + FakeReportRuns, + report_run_context, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + child_record, + invocation_context, + planned_simulation, +) + + +def _seeded_store(simulation, events=None): + store = FakeStore() + store.children[simulation.simulation_execution_id] = child_record(simulation) + if events is not None: + replace = store.replace_simulation_if_status + + def recording_replace(record, *, expected_status): + events.append(f"stage12:{record.status.value}") + return replace(record, expected_status=expected_status) + + store.replace_simulation_if_status = recording_replace + return store + + +def _run_worker(simulation, store, report_runs, *, context=None, calculator=None): + return run_single_simulation( + simulation.model_dump(mode="json"), + (context or report_run_context()).model_dump(mode="json"), + required_country="us", + modal_call_id="fc-baseline-worker", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=calculator or (lambda _: simulation_frames()), + report_runs=report_runs, + ), + ) + + +def test_worker_records_its_simulation_run_before_its_stage12_record() -> None: + events: list[str] = [] + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation, events) + report_runs = FakeReportRuns(events) + + result = _run_worker(simulation, store, report_runs) + + run = report_runs.simulation_runs[BASELINE_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == "fc-baseline-worker" + assert run.output_uri == result["artifact"]["uri"] + assert run.output_sha256 == result["artifact"]["content_sha256"] + # The simulation run succeeds first, so a restart between the two writes + # finds its Stage 12 record still running and records nothing twice. + assert events == [ + "stage12:running", + "claim_simulation_run", + "complete_simulation_run", + "stage12:succeeded", + ] + + +def test_worker_failure_marks_its_simulation_run_failed() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + + def failing_calculator(_simulation): + msg = "calculation failed" + raise ValueError(msg) + + with pytest.raises(RuntimeError, match="execution failed"): + _run_worker(simulation, store, report_runs, calculator=failing_calculator) + + run = report_runs.simulation_runs[BASELINE_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("simulation_execution_failed: ValueError") + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.FAILED + + +def test_worker_refuses_a_simulation_run_that_was_already_failed() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + report_runs.simulation_runs[BASELINE_RUN_ID].status = RunStatus.FAILED + calculated = [] + + with pytest.raises(RuntimeError, match="cancelled"): + _run_worker( + simulation, + store, + report_runs, + calculator=lambda value: calculated.append(value) or simulation_frames(), + ) + + assert calculated == [] + assert report_runs.simulation_runs[BASELINE_RUN_ID].output_uri is None + + +def test_restarted_worker_keeps_the_output_its_earlier_attempt_recorded() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + # The earlier attempt recorded its simulation run, then Modal pre-empted + # it before it marked its Stage 12 record succeeded. + store.children[simulation.simulation_execution_id] = store.children[ + simulation.simulation_execution_id + ].model_copy(update={"status": ComparisonRunLifecycleStatus.RUNNING}) + earlier = report_runs.simulation_runs[BASELINE_RUN_ID] + earlier.status = RunStatus.SUCCEEDED + earlier.output_uri = "gs://private/earlier-attempt.parquet" + earlier.output_sha256 = "e" * 64 + + result = _run_worker(simulation, store, report_runs) + + assert earlier.output_uri == "gs://private/earlier-attempt.parquet" + assert earlier.writes == ["claim_simulation_run", "complete_simulation_run"] + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert child.output_uri == result["artifact"]["uri"] + + +def test_comparison_worker_records_no_simulation_run() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + + _run_worker(simulation, store, report_runs, context=invocation_context()) + + assert report_runs.events == [] + + +def test_report_run_worker_needs_its_modal_call_id() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + + with pytest.raises(ValueError, match="Modal call ID"): + run_single_simulation( + simulation.model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=_seeded_store(simulation), + artifacts=FakeArtifacts(), + report_runs=FakeReportRuns(), + ), + ) + + +def test_report_run_worker_names_its_artifacts_by_digest(monkeypatch) -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + built = [] + + def recording_artifact_store(*, content_addressed=False): + built.append(content_addressed) + return FakeArtifacts() + + monkeypatch.setattr(simulation_module, "artifact_store", recording_artifact_store) + + run_single_simulation( + simulation.model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + required_country="us", + modal_call_id="fc-baseline-worker", + dependencies=SimulationWorkerDependencies( + store=_seeded_store(simulation), + calculator=lambda _: simulation_frames(), + report_runs=FakeReportRuns(), + ), + ) + + assert built == [True] From a5550c835b2df40c4d7aa30c544093a3ca76a6ce Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:23:42 +0400 Subject: [PATCH 07/24] Record a report call's report run from its Stage 12 coordinator When the coordinator's invocation context names a report run, it first checks that the run combines the call's baseline and reform simulation runs and claims it under its own Modal call ID. A run another call holds, or one already finished, makes the call return a deduplicated summary before it starts anything. The coordinator marks the report run succeeded before it marks its Stage 12 report succeeded, and marks it failed, with the redacted failure, whenever it records a Stage 12 failure. Its failure and cancellation records of child simulations also reach their simulation runs, and it names its aggregate by its digest. Comparison runs record nothing new, and their children receive exactly the context fields they received before. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 37 +++- .../stage12_runtime/report_coordination.py | 47 ++++- .../stage12_runtime/report_run_recording.py | 38 ++++ .../test_stage12_report_run_coordinator.py | 196 ++++++++++++++++++ 4 files changed, 313 insertions(+), 5 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index d64f9fd2f..faf0999c8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -27,6 +27,7 @@ from .coordinated_simulations import SIMULATION_CANCELLED_ERROR_CODE, logger from .dependencies import ModalChildInvoker, artifact_store, runtime_store from .failure_causes import StageProgress +from .failures import Stage12FailureDetail from .parent_claim import ( claim_parent, validate_dispatch_context, @@ -40,6 +41,8 @@ CoordinatorInvocation, ReportCoordination, ) +from .report_run_access import report_run_store +from .report_run_recording import ReportRunRecorder, SimulationRunRecordingStore __all__ = [ "SIMULATION_CANCELLED_ERROR_CODE", @@ -85,12 +88,24 @@ def coordinate_report( context=context, parent=submitted_parent, ) + persistence = dependencies.store or runtime_store(dependencies.runtime) + report_run = None + if context.report_run is not None: + report_run = ReportRunRecorder( + writer=dependencies.report_runs or report_run_store(dependencies.runtime), + run_ids=context.report_run, + modal_call_id=invocation.invocation_id, + ) + persistence = SimulationRunRecordingStore( + persistence, report_run.writer, context.report_run + ) coordination = ReportCoordination( report=report, context=context, invocation=invocation, - persistence=dependencies.store or runtime_store(dependencies.runtime), - artifact_storage=dependencies.artifacts or artifact_store(), + persistence=persistence, + artifact_storage=dependencies.artifacts + or artifact_store(content_addressed=report_run is not None), child_invoker=dependencies.invoker or ModalChildInvoker(), dependencies=dependencies, stage_plan=( @@ -100,14 +115,32 @@ def coordinate_report( ), progress=StageProgress(), child_wait_deadline=child_wait_deadline, + report_run=report_run, ) with coordination.span(Stage.STAGE12_COORDINATOR_CLAIM): + if report_run is not None: + # The report run decides whether this call works: another call + # may hold it, or it may already be finished. + claimed = report_run.claim() + if not claimed.changed: + return { + "deduplicated": True, + "report_run_id": str(report_run.run_ids.report_run_id), + "status": claimed.status.value, + } parent, should_execute = claim_parent( submitted=submitted_parent, coordinator_invocation_id=invocation.invocation_id, store=coordination.persistence, ) if not should_execute: + if report_run is not None: + report_run.fail( + Stage12FailureDetail( + error_code="report_already_claimed", + error_summary="Another call holds this report's Stage 12 record.", + ) + ) return { "deduplicated": True, "evaluation_id": str(parent.evaluation_id), diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py index 15c7ab7de..66b20aeed 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -25,6 +25,7 @@ Stage12OutputPlan, ) from policyengine_simulation_observability.stages import Stage, StagePlan +from policyengine_stage12_persistence import RunStatus from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, @@ -48,6 +49,7 @@ cancellation, child_failure_detail, child_record, + logger, record_child_failure, simulation_failure_attributes, ) @@ -79,6 +81,9 @@ from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + from .report_run_access import ReportRunWriter + from .report_run_recording import ReportRunRecorder + SIMULATION_POLL_INITIAL_SECONDS = 0.25 SIMULATION_POLL_MAX_SECONDS = 1.0 SIMULATION_POLL_BACKOFF_FACTOR = 2.0 @@ -116,8 +121,9 @@ class CoordinatorDependencies: The defaults are the Stage 12 database store, the private artifact bucket, Modal child invocation, the standard aggregate builder and output - plan resolution, no telemetry runtime, the real clock and sleep, and the - standard child wait window. + plan resolution, no telemetry runtime, the real clock and sleep, the + standard child wait window, and the report run store, which only a call + that names a report run uses. """ store: ComparisonStore | None = None @@ -131,6 +137,7 @@ class CoordinatorDependencies: sleep: Callable[[float], None] = time.sleep monotonic: Callable[[], float] = time.monotonic wait_timeout_seconds: float = SIMULATION_WAIT_TIMEOUT_SECONDS + report_runs: ReportRunWriter | None = None @dataclass @@ -158,6 +165,8 @@ class ReportCoordination: descriptors: dict[str, SimulationArtifactDescriptor] = dataclasses.field( default_factory=dict ) + # The report run the call names, when it records one. + report_run: ReportRunRecorder | None = None @property def runtime(self) -> ObservabilityRuntime | None: @@ -340,7 +349,16 @@ def _dispatch_child( function_name=self.context.simulation_callable, environment=self.context.modal_environment, simulation=simulation.model_dump(mode="json"), - context=self.context.model_dump(mode="json"), + # A comparison run's children receive the context + # fields they received before report runs existed. + context=self.context.model_dump( + mode="json", + exclude=( + {"report_run"} + if self.context.report_run is None + else None + ), + ), observability_context=child_observability_context, ) ) @@ -547,6 +565,10 @@ def _aggregate_and_complete( reform_artifact_sha256=reform.descriptor.artifact.content_sha256, bundle=self.report.baseline.bundle, ) + if self.report_run is not None: + # Before the Stage 12 report, so a call Modal restarts between the + # two writes finds that report running and resumes it. + self._complete_report_run() completed = datetime.now(UTC) completed_parent = self.persistence.replace_report( parent.model_copy( @@ -575,6 +597,19 @@ def _aggregate_and_complete( ) return descriptor.model_dump(mode="json") + def _complete_report_run(self) -> None: + """Mark the report run succeeded; log a run this call no longer holds.""" + if self.report_run is None: + return + written = self.report_run.complete() + if not written.changed and written.status is not RunStatus.SUCCEEDED: + logger.warning( + "Stage 12 report %s finished, but its report run %s is %s", + self.report.evaluation_id, + self.report_run.run_ids.report_run_id, + written.status.value, + ) + def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: """Record a failed run once, stop its unfinished children, fail the parent. @@ -606,6 +641,12 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: detail=cancellation("the coordinator failed"), ) record_failure_cause(self.runtime, error, detail.error_summary) + if self.report_run is not None: + persist_failure_record( + partial(self.report_run.fail, detail), + runtime=self.runtime, + failure_scope="report_coordination", + ) failed_at = datetime.now(UTC) persist_failure_record( partial( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py index c7a783668..62d0b2090 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -8,6 +8,7 @@ from __future__ import annotations +from dataclasses import dataclass from datetime import UTC, datetime from typing import TYPE_CHECKING @@ -26,8 +27,10 @@ ComparisonSimulationRecord, ReportRunIds, ) + from policyengine_stage12_persistence import RunWrite from .dependencies import ComparisonStore + from .failures import Stage12FailureDetail from .report_run_access import ReportRunWriter @@ -36,6 +39,41 @@ def run_error_message(error_code: str | None, error_summary: str | None) -> str: return f"{error_code or 'simulation_failed'}: {error_summary or 'no summary'}" +@dataclass(frozen=True) +class ReportRunRecorder: + """The coordinator's writes to the report run its call names, under its call ID.""" + + writer: ReportRunWriter + run_ids: ReportRunIds + modal_call_id: str + + def claim(self) -> RunWrite: + """Check the run combines the call's simulation runs, then claim it.""" + self.writer.require_simulation_runs(self.run_ids) + return self.writer.claim_report_run( + self.run_ids.report_run_id, + modal_call_id=self.modal_call_id, + claimed_at=datetime.now(UTC), + ) + + def complete(self) -> RunWrite: + """Mark the report run succeeded while this call holds it.""" + return self.writer.complete_report_run( + self.run_ids, + modal_call_id=self.modal_call_id, + completed_at=datetime.now(UTC), + ) + + def fail(self, detail: Stage12FailureDetail) -> RunWrite: + """Mark the report run failed with ``detail`` while this call holds it.""" + return self.writer.fail_report_run( + self.run_ids.report_run_id, + modal_call_id=self.modal_call_id, + error_message=run_error_message(detail.error_code, detail.error_summary), + failed_at=datetime.now(UTC), + ) + + class SimulationRunRecordingStore: """A ``ComparisonStore`` that also records each child's simulation run. diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py new file mode 100644 index 000000000..4acc7a486 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -0,0 +1,196 @@ +"""The Stage 12 coordinator records the status of the report run its call names.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from policyengine_stage12_persistence import RunStatus +from stage12_child_call_support import ConcurrentInvoker, InjectedChildFaults +from stage12_fake_report_runs import ( + REFORM_RUN_ID, + REPORT_RUN_ID, + FakeReportRuns, + report_run_context, +) +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( + aggregate_payload, + coordinate_report, + parent_record, + report_execution_input, +) + +COORDINATOR_CALL_ID = "fc-coordinator" + + +class RecordingInvoker(ConcurrentInvoker): + """Run both children concurrently, keeping the context each one received.""" + + def __init__(self, artifacts, **options): + super().__init__(artifacts, **options) + self.contexts = [] + + def spawn(self, request): + self.contexts.append(request.context) + return super().spawn(request) + + +def _coordinate(store, report_runs, invoker, artifacts, *, context=None): + context = context or report_run_context() + return coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload(), + report_runs=report_runs, + ) + + +def _recording_report_writes(store, events): + replace_report = store.replace_report + + def recording_replace_report(record): + events.append(f"stage12:{record.status.value}") + return replace_report(record) + + store.replace_report = recording_replace_report + return store + + +def test_coordinator_claims_and_completes_its_report_run_before_stage12() -> None: + events: list[str] = [] + store = _recording_report_writes(FakeStore(), events) + report_runs = FakeReportRuns(events) + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + result = _coordinate(store, report_runs, invoker, artifacts) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == COORDINATOR_CALL_ID + assert result["evaluation_id"] == str(report_execution_input().evaluation_id) + # The report run is claimed before, and completed before, the Stage 12 + # report, so a restart between the two finds the Stage 12 report running. + assert events == ["claim_report_run", "complete_report_run", "stage12:succeeded"] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert [context["report_run"] for context in invoker.contexts] == [ + report_run_context().report_run.model_dump(mode="json") + ] * 2 + + +def test_comparison_children_receive_the_context_they_received_before() -> None: + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + report_runs = FakeReportRuns() + + _coordinate( + FakeStore(), + report_runs, + invoker, + artifacts, + context=report_run_context().model_copy(update={"report_run": None}), + ) + + assert all("report_run" not in context for context in invoker.contexts) + assert report_runs.events == [] + + +@pytest.mark.parametrize( + ("status", "modal_call_id"), + [ + (RunStatus.RUNNING, "fc-another-coordinator"), + (RunStatus.SUCCEEDED, COORDINATOR_CALL_ID), + (RunStatus.FAILED, COORDINATOR_CALL_ID), + ], +) +def test_a_report_run_that_cannot_be_claimed_starts_nothing( + status, modal_call_id +) -> None: + store = FakeStore() + report_runs = FakeReportRuns() + run = report_runs.report_runs[REPORT_RUN_ID] + run.status, run.modal_call_id = status, modal_call_id + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + result = _coordinate(store, report_runs, invoker, artifacts) + + assert result == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": status.value, + } + assert invoker.events == [] + assert store.parent is None + assert run.writes == ["claim_report_run"] + + +def test_a_report_run_must_combine_the_simulation_runs_its_call_names() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + report_runs.links = {("baseline", REFORM_RUN_ID)} + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + with pytest.raises(ValueError, match="does not combine"): + _coordinate(store, report_runs, invoker, artifacts) + + assert invoker.events == [] + assert store.parent is None + assert report_runs.report_runs[REPORT_RUN_ID].status is RunStatus.PENDING + + +def test_a_failed_report_fails_its_report_run_and_the_failed_simulation_run() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + invoker = RecordingInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises(RuntimeError, match="coordination failed"): + _coordinate(store, report_runs, invoker, artifacts) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("simulation_invocation_failed: ") + reform = report_runs.simulation_runs[REFORM_RUN_ID] + assert reform.status is RunStatus.FAILED + assert reform.error_message == run.error_message + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + + +def test_report_run_coordinator_names_its_aggregate_by_digest(monkeypatch) -> None: + built = [] + artifacts = FakeArtifacts() + + def recording_artifact_store(*, content_addressed=False): + built.append(content_addressed) + return artifacts + + monkeypatch.setattr(coordination_module, "artifact_store", recording_artifact_store) + + coordinate_report( + report_execution_input().model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=report_run_context().modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=FakeStore(), + invoker=RecordingInvoker(artifacts), + aggregator=lambda **_: aggregate_payload(), + report_runs=FakeReportRuns(), + ) + + assert built == [True] From dc79bb70169ac4df791c0f2923dee7a204a2b482 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:25:02 +0400 Subject: [PATCH 08/24] Let a report run's coordinator stop its children when Modal pre-empts it Modal pre-empts a call with SIGINT, which the coordinator receives as KeyboardInterrupt, and later runs the same call again under the same call ID; it cancels a call with InputCancellation. Neither is an Exception. A coordinator whose call names a report run now cancels its started, unfinished child calls on either and re-raises, recording no failure. When Modal restarts it, it resumes the Stage 12 claim it holds under that call ID instead of treating the report as a duplicate, cancels the child calls its earlier attempt started, and starts the children again; the report run is claimed again by the same call and marked succeeded once. A restart after the report run succeeded starts nothing. Comparison runs keep their behaviour. Time limits are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 9 + .../stage12_runtime/parent_claim.py | 14 +- .../stage12_runtime/report_coordination.py | 40 +++- .../test_stage12_report_run_preemption.py | 193 ++++++++++++++++++ 4 files changed, 254 insertions(+), 2 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index faf0999c8..35acdc39e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -132,6 +132,7 @@ def coordinate_report( submitted=submitted_parent, coordinator_invocation_id=invocation.invocation_id, store=coordination.persistence, + resume_own_claim=report_run is not None, ) if not should_execute: if report_run is not None: @@ -154,3 +155,11 @@ def coordinate_report( coordination.fail(parent, error) msg = "Stage 12 report coordination failed" raise RuntimeError(msg) from None + except BaseException: + # Modal pre-empts a call with KeyboardInterrupt and cancels one with + # InputCancellation, neither an Exception, and runs a pre-empted call + # again under the same call ID. A report run's coordinator stops its + # unfinished children and records nothing, so that run starts again. + if report_run is not None: + coordination.stop_children() + raise diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py index dba19ed4f..1f4a88265 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py @@ -124,8 +124,14 @@ def claim_parent( submitted: ComparisonReportRecord, coordinator_invocation_id: str, store: ComparisonStore, + resume_own_claim: bool = False, ) -> tuple[ComparisonReportRecord, bool]: - """Create the parent in Modal or claim a retry; reject duplicate execution.""" + """Create the parent in Modal or claim a retry; reject duplicate execution. + + With ``resume_own_claim``, a running parent that this coordinator call + claimed before is resumed: Modal runs a pre-empted call again under the + same call ID. + """ coordinating_at = datetime.now(UTC) candidate = submitted.model_copy( update={ @@ -146,6 +152,12 @@ def claim_parent( if parent.evaluation_id != submitted.evaluation_id: msg = "existing comparison identity has a different report identifier" raise ValueError(msg) + if ( + resume_own_claim + and parent.status is ComparisonRunLifecycleStatus.RUNNING + and parent.coordinator_invocation_id == coordinator_invocation_id + ): + return parent, True if parent.status in { ComparisonRunLifecycleStatus.RUNNING, ComparisonRunLifecycleStatus.SUCCEEDED, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py index 66b20aeed..551614626 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -280,7 +280,13 @@ def _start_children( and child.modal_invocation_id is not None and not child.modal_invocation_id.startswith("dispatch-pending-") ): - calls[role] = self.child_invoker.restore(child.modal_invocation_id) + if self.report_run is None: + calls[role] = self.child_invoker.restore(child.modal_invocation_id) + else: + # A report run's coordinator cancelled this child when it + # was pre-empted; make sure it stopped, then start it again. + self._cancel_earlier_child(child) + calls[role] = self._dispatch_child(child, simulation) else: calls[role] = self._dispatch_child(child, simulation) return descriptors, calls @@ -309,6 +315,38 @@ def _retained_descriptor( } ) + def _cancel_earlier_child(self, child: ComparisonSimulationRecord) -> None: + """Cancel the child call an earlier attempt of this call started.""" + if child.modal_invocation_id is None: + return + try: + self.child_invoker.restore(child.modal_invocation_id).cancel() + except Exception as error: # noqa: BLE001 # cancelling an earlier attempt's child is best effort; the new child starts regardless + logger.warning( + "Stage 12 could not cancel an earlier %s simulation call in report " + "%s (%s)", + child.role.value, + child.evaluation_id, + type(error).__name__, + ) + + def stop_children(self) -> None: + """Cancel every started, unfinished child call, recording nothing. + + Each child's record stays running, so the call Modal restarts starts + it again. A cancellation error is logged, never raised. + """ + for role in sorted(set(self.calls) - set(self.descriptors)): + try: + self.calls[role].cancel() + except Exception as error: # noqa: BLE001 # every child must be asked to stop; a failed cancellation is logged + logger.warning( + "Stage 12 could not cancel the %s simulation of report %s (%s)", + role, + self.report.evaluation_id, + type(error).__name__, + ) + def _dispatch_child( self, child: ComparisonSimulationRecord, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py new file mode 100644 index 000000000..d3fa617c6 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -0,0 +1,193 @@ +"""A report call's coordinator under pre-emption: stop the children, then resume. + +Modal pre-empts a call with SIGINT, which a synchronous function receives as +``KeyboardInterrupt``, and later runs the same call again under the same call +ID. The coordinator of a report run cancels its unfinished children, lets +Modal restart it, and the restarted call starts them again and records the +report run once. +""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from policyengine_stage12_persistence import RunStatus +from stage12_child_call_support import ConcurrentInvoker, TimeoutInvoker +from stage12_fake_report_runs import REPORT_RUN_ID, FakeReportRuns, report_run_context +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( + aggregate_payload, + coordinate_report, + invocation_context, + parent_record, + report_execution_input, +) + +COORDINATOR_CALL_ID = "fc-coordinator" + + +class CancelledCall: + def __init__(self, object_id): + self.object_id = object_id + self.cancelled = 0 + + def get(self, *, timeout=None): + msg = "a cancelled call has no result" + raise AssertionError(msg) + + def cancel(self): + self.cancelled += 1 + + +class RestartInvoker(ConcurrentInvoker): + """Run new children; restoring an earlier attempt's child only allows cancelling.""" + + def __init__(self, artifacts): + super().__init__(artifacts) + self.restored: dict[str, CancelledCall] = {} + + def restore(self, invocation_id): + self.events.append(f"restore:{invocation_id}") + call = CancelledCall(invocation_id) + self.restored[invocation_id] = call + return call + + +def _preempt(_seconds): + raise KeyboardInterrupt + + +def _coordinate(store, report_runs, invoker, *, context, **seams): + return coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=store, + artifacts=seams.pop("artifacts", FakeArtifacts()), + invoker=invoker, + aggregator=lambda **_: aggregate_payload(), + report_runs=report_runs, + **seams, + ) + + +def _preempted_first_attempt(store, report_runs, *, context): + invoker = TimeoutInvoker() + calls = [] + spawn = invoker.spawn + + def keeping_spawn(request): + call = spawn(request) + calls.append(call) + return call + + invoker.spawn = keeping_spawn + with pytest.raises(KeyboardInterrupt): + _coordinate(store, report_runs, invoker, context=context, sleep=_preempt) + return calls + + +def test_preempted_coordinator_cancels_its_children_and_records_no_failure() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + calls = _preempted_first_attempt(store, report_runs, context=report_run_context()) + + assert [call.cancelled for call in calls] == [1, 1] + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.RUNNING + assert run.writes == ["claim_report_run"] + assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING + assert all( + child.status is ComparisonRunLifecycleStatus.RUNNING + for child in store.children.values() + ) + assert all( + simulation_run.status is RunStatus.PENDING + for simulation_run in report_runs.simulation_runs.values() + ) + + +def test_restarted_coordinator_resumes_its_claim_and_starts_its_children_again() -> ( + None +): + store = FakeStore() + report_runs = FakeReportRuns() + _preempted_first_attempt(store, report_runs, context=report_run_context()) + artifacts = FakeArtifacts() + invoker = RestartInvoker(artifacts) + + result = _coordinate( + store, report_runs, invoker, context=report_run_context(), artifacts=artifacts + ) + + assert result["evaluation_id"] == str(report_execution_input().evaluation_id) + assert invoker.events == [ + "restore:call-baseline", + "spawn:baseline", + "restore:call-reform", + "spawn:reform", + ] + assert all(call.cancelled == 1 for call in invoker.restored.values()) + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.writes == ["claim_report_run", "claim_report_run", "complete_report_run"] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def test_a_call_restarted_after_its_report_run_succeeded_records_nothing_again() -> ( + None +): + store = FakeStore() + report_runs = FakeReportRuns() + replace_report = store.replace_report + + def preempted_before_the_stage12_success(record): + if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: + raise KeyboardInterrupt + return replace_report(record) + + store.replace_report = preempted_before_the_stage12_success + artifacts = FakeArtifacts() + with pytest.raises(KeyboardInterrupt): + _coordinate( + store, + report_runs, + ConcurrentInvoker(artifacts), + context=report_run_context(), + artifacts=artifacts, + ) + restart_invoker = RestartInvoker(FakeArtifacts()) + + result = _coordinate( + store, report_runs, restart_invoker, context=report_run_context() + ) + + assert result == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": "succeeded", + } + assert restart_invoker.events == [] + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.writes == ["claim_report_run", "complete_report_run", "claim_report_run"] + + +def test_a_preempted_comparison_run_behaves_as_before() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + calls = _preempted_first_attempt(store, report_runs, context=invocation_context()) + restart_invoker = RestartInvoker(FakeArtifacts()) + result = _coordinate( + store, report_runs, restart_invoker, context=invocation_context() + ) + + assert [call.cancelled for call in calls] == [0, 0] + assert result["deduplicated"] is True + assert restart_invoker.events == [] + assert report_runs.events == [] From 3c76d1cebc4b973946a2f0691caa0d23e898d3d5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:32:23 +0400 Subject: [PATCH 09/24] Store a report run's economy output as result rows with its success The report run store can now write a report's single-year economy output in the same transaction that marks the report run succeeded, only when that status write changes the run. Budget totals, income deciles and intra-decile shares become rows keyed by the report run and its two simulation runs; poverty rates and inequality become rows keyed by each simulation run, under .py's poverty types and income variables for the country; the labour supply response and cliff impact are stored whole on the report run. A repeated completion stores nothing, a failed insert leaves the run running, and a simulation run another report run already stored a row for keeps that row. The geographic impacts, the detailed budget, the UK's wealth deciles, the SPM provenance and the data version have no lossless place in the result tables yet and are not written. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 6 + .../report_run_results.py | 321 ++++++++++++++++++ .../report_run_store.py | 63 +++- .../tests/economy_output.py | 122 +++++++ .../tests/report_run_database.py | 5 + .../tests/test_report_run_results.py | 274 +++++++++++++++ 6 files changed, 785 insertions(+), 6 deletions(-) create mode 100644 libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py create mode 100644 libs/policyengine-stage12-persistence/tests/economy_output.py create mode 100644 libs/policyengine-stage12-persistence/tests/test_report_run_results.py diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index e985543e6..c6b3909a4 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -1,5 +1,9 @@ """Direct DML for temporary Stage 12 comparison state and for report runs.""" +from policyengine_stage12_persistence.report_run_results import ( + ReportRunResult, + report_run_result, +) from policyengine_stage12_persistence.report_run_store import ( ReportRunStore, RunStatus, @@ -18,6 +22,7 @@ __all__ = [ "ConnectionOpener", + "ReportRunResult", "ReportRunStore", "RunStatus", "RunWrite", @@ -25,5 +30,6 @@ "Stage12RuntimeAccessError", "create_stage12_engine", "normalized_database_url", + "report_run_result", "validate_runtime_database", ] diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py new file mode 100644 index 000000000..d689809ba --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -0,0 +1,321 @@ +"""A report run's economy output as rows of the Public API's result tables. + +The output is the single-year economy result Stage 12's coordinator builds, +in the format API v1 and v3 read today. Each part with a lossless place in a +result table becomes rows: per-report parts keyed by the report run and its +two simulation runs, poverty and inequality keyed by each simulation run. The +labour supply response and cliff impact are stored whole on the report run. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any +from uuid import uuid4 + +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from sqlalchemy import Column, Float, Integer, String, Table +from sqlalchemy.dialects.postgresql import ENUM, UUID + +from policyengine_stage12_persistence.report_run_tables import metadata + +if TYPE_CHECKING: + from collections.abc import Mapping + + from policyengine_simulation_contract.macro_output import IntraDecileOutput + from policyengine_simulation_contract.stage12_execution import ReportRunIds + +decile_type_type = ENUM("income", "wealth", name="v2_decile_type", create_type=False) + + +def _report_run_keys() -> tuple[Column[Any], ...]: + return ( + Column("id", UUID(as_uuid=True), primary_key=True), + Column("report_run_id", UUID(as_uuid=True), nullable=False), + Column("baseline_simulation_run_id", UUID(as_uuid=True), nullable=False), + Column("reform_simulation_run_id", UUID(as_uuid=True), nullable=False), + ) + + +def _simulation_run_keys() -> tuple[Column[Any], ...]: + return ( + Column("id", UUID(as_uuid=True), primary_key=True), + Column("simulation_run_id", UUID(as_uuid=True), nullable=False), + ) + + +budget_summary = Table( + "budget_summary", + metadata, + *_report_run_keys(), + Column("variable_name", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("baseline_total", Float), + Column("reform_total", Float), + Column("change", Float), +) +decile_impacts = Table( + "decile_impacts", + metadata, + *_report_run_keys(), + Column("income_variable", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("decile", Integer, nullable=False), + Column("quantiles", Integer, nullable=False), + Column("absolute_change", Float), + Column("relative_change", Float), +) +intra_decile_impacts = Table( + "intra_decile_impacts", + metadata, + *_report_run_keys(), + Column("decile_type", decile_type_type, nullable=False), + Column("decile", Integer, nullable=False), + Column("lose_more_than_5pct", Float), + Column("lose_less_than_5pct", Float), + Column("no_change", Float), + Column("gain_less_than_5pct", Float), + Column("gain_more_than_5pct", Float), +) +poverty = Table( + "poverty", + metadata, + *_simulation_run_keys(), + Column("poverty_type", String(128), nullable=False), + Column("entity", String(128), nullable=False), + Column("filter_variable", String(512)), + Column("rate", Float), +) +inequality = Table( + "inequality", + metadata, + *_simulation_run_keys(), + Column("income_variable", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("gini", Float), + Column("top_10_share", Float), + Column("top_1_share", Float), +) + +# The household variable each headline budget figure totals, and the column +# its figure goes in. ``budgetary_impact`` is tax minus benefits, as the +# economy output computes it, so it needs no row of its own. +BUDGET_FIGURES = ( + ("tax_revenue_impact", "household_tax", "change"), + ("benefit_spending_impact", "household_benefits", "change"), + ("state_tax_revenue_impact", "household_state_income_tax", "change"), + ("baseline_net_income", "household_net_income", "baseline_total"), + ("households", "household_weight", "baseline_total"), +) +# .py's poverty types behind the output's "poverty" and "deep_poverty". +POVERTY_TYPES = { + "us": {"poverty": "spm", "deep_poverty": "spm_deep"}, + "uk": {"poverty": "relative_bhc", "deep_poverty": "absolute_bhc"}, +} +INEQUALITY_INCOME_VARIABLES = { + "us": "household_net_income", + "uk": "equiv_hbai_household_net_income", +} +INTRA_DECILE_COLUMNS = { + "Lose more than 5%": "lose_more_than_5pct", + "Lose less than 5%": "lose_less_than_5pct", + "No change": "no_change", + "Gain less than 5%": "gain_less_than_5pct", + "Gain more than 5%": "gain_more_than_5pct", +} +# The intra-decile row holding each category's share across all deciles. +ALL_DECILES = 0 + + +@dataclass(frozen=True) +class ResultRows: + """The rows one result table receives. + + Rows ``shared_by_reports`` belong to a simulation run, which other report + runs may share; a row such a run already holds is kept as stored. + """ + + table: Table + rows: tuple[dict[str, Any], ...] + shared_by_reports: bool + + +@dataclass(frozen=True) +class ReportRunResult: + """Everything a report run's economy output is stored as.""" + + rows: tuple[ResultRows, ...] + labor_supply_result: dict[str, Any] | None + cliff_impact_result: dict[str, Any] | None + + +def report_run_result( + output: Mapping[str, Any], + *, + run_ids: ReportRunIds, + country: str, +) -> ReportRunResult: + """Validate a single-year economy output and return the rows it is stored as. + + Raises ValueError for an output that is not a single-year economy result + or a country without poverty types. + """ + economy = SingleYearMacroOutput.model_validate(output) + if country not in POVERTY_TYPES: + msg = f"no poverty types are known for country {country!r}" + raise ValueError(msg) + report_keys = { + "report_run_id": run_ids.report_run_id, + "baseline_simulation_run_id": run_ids.baseline_simulation_run_id, + "reform_simulation_run_id": run_ids.reform_simulation_run_id, + } + dumped = economy.model_dump(mode="json") + return ReportRunResult( + rows=( + ResultRows( + budget_summary, + _budget_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + decile_impacts, + _decile_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + intra_decile_impacts, + _intra_decile_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + poverty, + _poverty_rows(economy, run_ids, country), + shared_by_reports=True, + ), + ResultRows( + inequality, + _inequality_rows(economy, run_ids, country), + shared_by_reports=True, + ), + ), + labor_supply_result=dumped["labor_supply_response"], + cliff_impact_result=dumped["cliff_impact"], + ) + + +def _budget_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + return tuple( + { + "id": uuid4(), + **keys, + "variable_name": variable, + "entity": "household", + column: getattr(economy.budget, figure), + } + for figure, variable, column in BUDGET_FIGURES + ) + + +def _decile_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + return tuple( + { + "id": uuid4(), + **keys, + "income_variable": "household_net_income", + "entity": "household", + "decile": int(decile), + "quantiles": 10, + "absolute_change": average, + "relative_change": economy.decile.relative[decile], + } + for decile, average in economy.decile.average.items() + ) + + +def _intra_decile_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + rows = _intra_decile_type_rows(economy.intra_decile, "income", keys) + if economy.intra_wealth_decile is not None: + rows += _intra_decile_type_rows(economy.intra_wealth_decile, "wealth", keys) + return rows + + +def _intra_decile_type_rows( + shares: IntraDecileOutput, decile_type: str, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + decile_count = max((len(values) for values in shares.deciles.values()), default=0) + by_decile: list[dict[str, Any]] = [ + { + column: shares.deciles[label][index] + for label, column in INTRA_DECILE_COLUMNS.items() + if index < len(shares.deciles.get(label, ())) + } + | {"decile": index + 1} + for index in range(decile_count) + ] + overall = { + column: shares.all[label] + for label, column in INTRA_DECILE_COLUMNS.items() + if label in shares.all + } | {"decile": ALL_DECILES} + return tuple( + {"id": uuid4(), **keys, "decile_type": decile_type, **values} + for values in (*by_decile, overall) + ) + + +def _poverty_rows( + economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str +) -> tuple[dict[str, Any], ...]: + poverty_types = POVERTY_TYPES[country] + blocks: list[tuple[str, Mapping[str, Any]]] = [ + ("poverty", economy.poverty.poverty.model_dump()), + ("deep_poverty", economy.poverty.deep_poverty.model_dump()), + ("poverty", economy.poverty_by_gender.poverty.model_dump()), + ("deep_poverty", economy.poverty_by_gender.deep_poverty.model_dump()), + ] + if economy.poverty_by_race is not None: + blocks.append(("poverty", economy.poverty_by_race.poverty.model_dump())) + return tuple( + { + "id": uuid4(), + "simulation_run_id": simulation_run_id, + "poverty_type": poverty_types[measure], + "entity": "person", + # Everyone's rate has no filter; each group's names the group. + "filter_variable": None if group == "all" else group, + "rate": rates[side], + } + for side, simulation_run_id in _sides(run_ids) + for measure, groups in blocks + for group, rates in groups.items() + ) + + +def _inequality_rows( + economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str +) -> tuple[dict[str, Any], ...]: + return tuple( + { + "id": uuid4(), + "simulation_run_id": simulation_run_id, + "income_variable": INEQUALITY_INCOME_VARIABLES[country], + "entity": "household", + "gini": getattr(economy.inequality.gini, side), + "top_10_share": getattr(economy.inequality.top_10_pct_share, side), + "top_1_share": getattr(economy.inequality.top_1_pct_share, side), + } + for side, simulation_run_id in _sides(run_ids) + ) + + +def _sides(run_ids: ReportRunIds) -> tuple[tuple[str, Any], ...]: + return ( + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + ) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index 147730712..bbf9197e1 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -10,9 +10,10 @@ from contextlib import contextmanager from dataclasses import dataclass from enum import StrEnum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any -from sqlalchemy import and_, func, or_, select, update +from sqlalchemy import and_, func, insert, or_, select, update +from sqlalchemy.dialects.postgresql import insert as postgresql_insert from policyengine_stage12_persistence.report_run_tables import ( report_run_simulation_runs, @@ -26,7 +27,7 @@ ) if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator from datetime import datetime from uuid import UUID @@ -34,6 +35,8 @@ from sqlalchemy import Engine, Table, Update from sqlalchemy.engine import Connection + from policyengine_stage12_persistence.report_run_results import ReportRunResult + class RunStatus(StrEnum): """A report run's or simulation run's status.""" @@ -150,8 +153,16 @@ def complete_report_run( *, modal_call_id: str, completed_at: datetime, + result: ReportRunResult | None = None, ) -> RunWrite: - """Mark the report run succeeded while ``modal_call_id`` holds it running.""" + """Mark the report run succeeded while ``modal_call_id`` holds it running. + + ``result``'s rows are inserted in the same transaction, only when the + status changed, so a repeated completion stores nothing twice and a + failed insert leaves the run running. A row of a simulation run that + already has one for the same measure, from a report run sharing that + simulation run, keeps the stored row. + """ statement = ( update(report_runs) .where( @@ -164,10 +175,22 @@ def complete_report_run( error_message=None, completed_at=completed_at, updated_at=completed_at, + **( + { + "labor_supply_result": result.labor_supply_result, + "cliff_impact_result": result.cliff_impact_result, + } + if result is not None + else {} + ), ) ) return self._write( - report_runs, run_ids.report_run_id, statement, "complete_report_run" + report_runs, + run_ids.report_run_id, + statement, + "complete_report_run", + after_change=None if result is None else _result_inserts(result), ) def fail_report_run( @@ -283,16 +306,21 @@ def _write( run_id: UUID, statement: Update, operation_name: str, + *, + after_change: Callable[[Connection], None] | None = None, ) -> RunWrite: """Run a conditional update; return whether it changed the run, and its status. - Raises LookupError when the run does not exist. + ``after_change`` runs in the same transaction, only after the update + changed the run. Raises LookupError when the run does not exist. """ with self._transaction(operation_name) as connection: changed = connection.execute( statement.returning(table.c.status) ).scalar_one_or_none() if changed is not None: + if after_change is not None: + after_change(connection) return RunWrite(changed=True, status=RunStatus(changed)) current = connection.execute( select(table.c.status).where(table.c.id == run_id) @@ -302,3 +330,26 @@ def _write( msg = f"{noun} {run_id} does not exist" raise LookupError(msg) return RunWrite(changed=False, status=RunStatus(current)) + + +def _result_inserts(result: ReportRunResult) -> Callable[[Connection], None]: + """Return the step that inserts ``result``'s rows on an open transaction.""" + + def insert_rows(connection: Connection) -> None: + for result_rows in result.rows: + if not result_rows.rows: + continue + statement = ( + postgresql_insert(result_rows.table).on_conflict_do_nothing() + if result_rows.shared_by_reports + else insert(result_rows.table) + ) + connection.execute(statement, _uniform_rows(result_rows.rows)) + + return insert_rows + + +def _uniform_rows(rows: tuple[dict[str, Any], ...]) -> list[dict[str, Any]]: + """Give every row the same columns, unset ones null, for one batched insert.""" + columns = {name for row in rows for name in row} + return [dict.fromkeys(columns) | row for row in rows] diff --git a/libs/policyengine-stage12-persistence/tests/economy_output.py b/libs/policyengine-stage12-persistence/tests/economy_output.py new file mode 100644 index 000000000..a59e4dc54 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/economy_output.py @@ -0,0 +1,122 @@ +"""A complete single-year economy output with a distinct value in each field.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + +US_ECONOMY_OUTPUT: dict[str, Any] = { + "spm_provenance": None, + "model_version": "1.715.2", + "data_version": "1.115.5", + "budget": { + "budgetary_impact": 70.25, + "tax_revenue_impact": 100.5, + "state_tax_revenue_impact": 20.125, + "benefit_spending_impact": 30.25, + "households": 2.0, + "baseline_net_income": 1000.75, + }, + "detailed_budget": { + "income_tax": { + "baseline": 100.0, + "reform": 125.0, + "difference": 25.0, + "is_tax": True, + } + }, + "decile": { + "relative": {"1": 0.01, "2": 0.02}, + "average": {"1": 10.0, "2": 20.0}, + }, + "inequality": { + "gini": {"baseline": 0.4, "reform": 0.39}, + "top_10_pct_share": {"baseline": 0.3, "reform": 0.29}, + "top_1_pct_share": {"baseline": 0.1, "reform": 0.09}, + }, + "poverty": { + "poverty": { + "adult": {"baseline": 0.081, "reform": 0.071}, + "all": {"baseline": 0.1, "reform": 0.09}, + "child": {"baseline": 0.12, "reform": 0.11}, + "senior": {"baseline": 0.061, "reform": 0.051}, + }, + "deep_poverty": { + "adult": {"baseline": 0.021, "reform": 0.011}, + "all": {"baseline": 0.031, "reform": 0.022}, + "child": {"baseline": 0.04, "reform": 0.03}, + "senior": {"baseline": 0.015, "reform": 0.012}, + }, + }, + "poverty_by_gender": { + "poverty": { + "male": {"baseline": 0.08, "reform": 0.07}, + "female": {"baseline": 0.09, "reform": 0.085}, + }, + "deep_poverty": { + "male": {"baseline": 0.02, "reform": 0.019}, + "female": {"baseline": 0.025, "reform": 0.024}, + }, + }, + "poverty_by_race": { + "poverty": { + "black": {"baseline": 0.15, "reform": 0.14}, + "hispanic": {"baseline": 0.13, "reform": 0.125}, + "other": {"baseline": 0.11, "reform": 0.105}, + "white": {"baseline": 0.06, "reform": 0.05}, + }, + }, + "intra_decile": { + "all": { + "Gain less than 5%": 0.30000000000000004, + "Gain more than 5%": 0.3, + "Lose less than 5%": 0.15000000000000002, + "Lose more than 5%": 0.05, + "No change": 0.44999999999999996, + }, + "deciles": { + "Gain less than 5%": [0.4, 0.2], + "Gain more than 5%": [0.5, 0.1], + "Lose less than 5%": [0.2, 0.1], + "Lose more than 5%": [0.1, 0.0], + "No change": [0.3, 0.6], + }, + }, + "wealth_decile": None, + "intra_wealth_decile": None, + "labor_supply_response": { + "substitution_lsr": 1.5, + "income_lsr": -0.5, + "relative_lsr": {"income": -0.01, "substitution": 0.02}, + "total_change": 1.0, + "revenue_change": 3.0, + "decile": { + "average": {"income": {"1": -0.1}, "substitution": {"1": 0.2}}, + "relative": {"income": {"1": -0.001}, "substitution": {"1": 0.002}}, + }, + "hours": { + "baseline": 100.0, + "reform": 101.0, + "change": 1.0, + "income_effect": -0.5, + "substitution_effect": 1.5, + }, + }, + "constituency_impact": None, + "local_authority_impact": None, + "congressional_district_impact": None, + "cliff_impact": { + "baseline": {"cliff_gap": 10.0, "cliff_share": 0.1}, + "reform": {"cliff_gap": 9.0, "cliff_share": 0.09}, + }, +} + + +def uk_economy_output() -> dict[str, Any]: + """Return the output as the UK reports it: no race breakdown, with wealth parts.""" + output = deepcopy(US_ECONOMY_OUTPUT) + output["poverty_by_race"] = None + output["budget"]["state_tax_revenue_impact"] = 0.0 + output["intra_wealth_decile"] = deepcopy(output["intra_decile"]) + output["intra_wealth_decile"]["all"]["No change"] = 0.5 + return output diff --git a/libs/policyengine-stage12-persistence/tests/report_run_database.py b/libs/policyengine-stage12-persistence/tests/report_run_database.py index b4a39b67e..a16499828 100644 --- a/libs/policyengine-stage12-persistence/tests/report_run_database.py +++ b/libs/policyengine-stage12-persistence/tests/report_run_database.py @@ -34,6 +34,11 @@ "report_runs", "simulation_runs", "report_run_simulation_runs", + "budget_summary", + "decile_impacts", + "intra_decile_impacts", + "poverty", + "inequality", ) NOW = datetime(2026, 11, 4, tzinfo=UTC) VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "6.12.0"} diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py new file mode 100644 index 000000000..63043171d --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -0,0 +1,274 @@ +"""A report run's economy output, stored as rows with the run's success.""" + +from __future__ import annotations + +from datetime import timedelta +from uuid import uuid4 + +import pytest +from economy_output import US_ECONOMY_OUTPUT, uk_economy_output +from policyengine_stage12_persistence import ( + ReportRunStore, + RunStatus, + RunWrite, + report_run_result, +) +from policyengine_stage12_persistence.report_run_results import ResultRows +from report_run_database import NOW, SeededReportRun, worker_database_url +from sqlalchemy import insert +from sqlalchemy.exc import IntegrityError + +LATER = NOW + timedelta(minutes=5) +INTRA_DECILE_LABELS = { + "Lose more than 5%": "lose_more_than_5pct", + "Lose less than 5%": "lose_less_than_5pct", + "No change": "no_change", + "Gain less than 5%": "gain_less_than_5pct", + "Gain more than 5%": "gain_more_than_5pct", +} +AGE_GROUPS = ("child", "adult", "senior") + + +@pytest.fixture +def store(): + run_store = ReportRunStore(worker_database_url()) + try: + yield run_store + finally: + run_store.engine.dispose() + + +def _complete(store, seeded_run, output=US_ECONOMY_OUTPUT, country="us"): + store.claim_report_run( + seeded_run.run_ids.report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + return store.complete_report_run( + seeded_run.run_ids, + modal_call_id="fc-coordinator", + completed_at=LATER, + result=report_run_result(output, run_ids=seeded_run.run_ids, country=country), + ) + + +def _rebuilt_budget(seeder, report_run_id): + rows = { + row["variable_name"]: row + for row in seeder.rows("budget_summary", report_run_id=report_run_id) + } + tax = rows["household_tax"]["change"] + benefits = rows["household_benefits"]["change"] + return { + "budgetary_impact": tax - benefits, + "tax_revenue_impact": tax, + "state_tax_revenue_impact": rows["household_state_income_tax"]["change"], + "benefit_spending_impact": benefits, + "households": rows["household_weight"]["baseline_total"], + "baseline_net_income": rows["household_net_income"]["baseline_total"], + } + + +def _rebuilt_intra_decile(rows): + by_decile = {row["decile"]: row for row in rows} + deciles = sorted(decile for decile in by_decile if decile != 0) + return { + "all": { + label: by_decile[0][column] for label, column in INTRA_DECILE_LABELS.items() + }, + "deciles": { + label: [by_decile[decile][column] for decile in deciles] + for label, column in INTRA_DECILE_LABELS.items() + }, + } + + +def _rebuilt_poverty(seeder, simulation_run_id, *, poverty_type, groups): + rows = { + row["filter_variable"]: row["rate"] + for row in seeder.rows( + "poverty", simulation_run_id=simulation_run_id, poverty_type=poverty_type + ) + } + return {group: rows[None if group == "all" else group] for group in groups} + + +def test_a_report_run_stores_its_economy_output_as_rows( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + + assert _complete(store, seeded_run) == RunWrite( + changed=True, status=RunStatus.SUCCEEDED + ) + + output = US_ECONOMY_OUTPUT + assert _rebuilt_budget(seeder, run_ids.report_run_id) == output["budget"] + deciles = seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id) + assert { + "average": {str(row["decile"]): row["absolute_change"] for row in deciles}, + "relative": {str(row["decile"]): row["relative_change"] for row in deciles}, + } == output["decile"] + assert {row["income_variable"] for row in deciles} == {"household_net_income"} + assert all( + row["baseline_simulation_run_id"] == run_ids.baseline_simulation_run_id + and row["reform_simulation_run_id"] == run_ids.reform_simulation_run_id + for row in deciles + ) + intra = seeder.rows("intra_decile_impacts", report_run_id=run_ids.report_run_id) + assert _rebuilt_intra_decile(intra) == output["intra_decile"] + for side, run_id in ( + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + ): + for poverty_type, key in (("spm", "poverty"), ("spm_deep", "deep_poverty")): + assert _rebuilt_poverty( + seeder, + run_id, + poverty_type=poverty_type, + groups=(*AGE_GROUPS, "all", "male", "female"), + ) == { + **{ + group: output["poverty"][key][group][side] + for group in (*AGE_GROUPS, "all") + }, + **{ + group: output["poverty_by_gender"][key][group][side] + for group in ("male", "female") + }, + } + assert _rebuilt_poverty( + seeder, + run_id, + poverty_type="spm", + groups=("white", "black", "hispanic", "other"), + ) == { + group: value[side] + for group, value in output["poverty_by_race"]["poverty"].items() + } + (inequality,) = seeder.rows("inequality", simulation_run_id=run_id) + assert inequality["income_variable"] == "household_net_income" + assert inequality["entity"] == "household" + assert ( + inequality["gini"], + inequality["top_10_share"], + inequality["top_1_share"], + ) == ( + output["inequality"]["gini"][side], + output["inequality"]["top_10_pct_share"][side], + output["inequality"]["top_1_pct_share"][side], + ) + report_run = seeder.row("report_runs", run_ids.report_run_id) + assert report_run["labor_supply_result"] == output["labor_supply_response"] + assert report_run["cliff_impact_result"] == output["cliff_impact"] + + +def test_the_uk_output_uses_its_own_poverty_types_and_wealth_shares( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + + _complete(store, seeded_run, output=uk_economy_output(), country="uk") + + poverty_types = { + row["poverty_type"] + for row in seeder.rows( + "poverty", simulation_run_id=run_ids.baseline_simulation_run_id + ) + } + assert poverty_types == {"relative_bhc", "absolute_bhc"} + (inequality,) = seeder.rows( + "inequality", simulation_run_id=run_ids.reform_simulation_run_id + ) + assert inequality["income_variable"] == "equiv_hbai_household_net_income" + wealth = [ + row + for row in seeder.rows( + "intra_decile_impacts", report_run_id=run_ids.report_run_id + ) + if row["decile_type"] == "wealth" + ] + assert _rebuilt_intra_decile(wealth) == uk_economy_output()["intra_wealth_decile"] + + +def test_a_repeated_completion_stores_no_second_result( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + _complete(store, seeded_run) + + repeated = store.complete_report_run( + run_ids, + modal_call_id="fc-coordinator", + completed_at=LATER, + result=report_run_result(US_ECONOMY_OUTPUT, run_ids=run_ids, country="us"), + ) + + assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert len(seeder.rows("budget_summary", report_run_id=run_ids.report_run_id)) == 5 + assert len(seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id)) == 2 + + +def test_a_result_that_cannot_be_stored_leaves_the_report_run_running( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + store.claim_report_run( + run_ids.report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + result = report_run_result(US_ECONOMY_OUTPUT, run_ids=run_ids, country="us") + budget = next(rows for rows in result.rows if rows.table.name == "budget_summary") + # The same row twice breaks the table's unique key in the last insert. + duplicated = ResultRows( + table=budget.table, + rows=(*budget.rows, {**budget.rows[0], "id": uuid4()}), + shared_by_reports=False, + ) + broken = result.__class__( + rows=(*[rows for rows in result.rows if rows is not budget], duplicated), + labor_supply_result=result.labor_supply_result, + cliff_impact_result=result.cliff_impact_result, + ) + + with pytest.raises(IntegrityError): + store.complete_report_run( + run_ids, modal_call_id="fc-coordinator", completed_at=LATER, result=broken + ) + + report_run = seeder.row("report_runs", run_ids.report_run_id) + assert report_run["status"] == "running" + assert report_run["labor_supply_result"] is None + for table_name in ("decile_impacts", "budget_summary", "intra_decile_impacts"): + assert seeder.rows(table_name, report_run_id=run_ids.report_run_id) == [] + assert ( + seeder.rows("poverty", simulation_run_id=run_ids.baseline_simulation_run_id) + == [] + ) + + +def test_a_shared_simulation_run_keeps_the_results_first_stored_for_it( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + # An earlier report run that shares this baseline stored its poverty first. + with seeder.engine.begin() as connection: + connection.execute( + insert(seeder.table("poverty")).values( + id=uuid4(), + simulation_run_id=run_ids.baseline_simulation_run_id, + poverty_type="spm", + entity="person", + filter_variable=None, + rate=0.5, + ) + ) + + _complete(store, seeded_run) + + stored = seeder.rows( + "poverty", + simulation_run_id=run_ids.baseline_simulation_run_id, + poverty_type="spm", + ) + assert next(row for row in stored if row["filter_variable"] is None)["rate"] == 0.5 + assert len(stored) == len({row["filter_variable"] for row in stored}) + reform = seeder.rows("poverty", simulation_run_id=run_ids.reform_simulation_run_id) + assert len(reform) == 2 * 6 + 4 From 82c289ee969fd9d6de20661b74404795363f7996 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:33:40 +0400 Subject: [PATCH 10/24] Store a report call's aggregate result with its report run's success A coordinator whose call names a report run now passes the economy output of its aggregate to the report run store, which writes it as result rows in the transaction that marks the run succeeded. An output the store cannot hold fails the report, as any other failure does. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/report_coordination.py | 13 +++++++--- .../stage12_runtime/report_run_access.py | 5 ++-- .../stage12_runtime/report_run_recording.py | 15 ++++++++--- .../tests/stage12_fake_report_runs.py | 12 +++++++-- .../test_stage12_report_run_coordinator.py | 26 ++++++++++++++++--- .../test_stage12_report_run_preemption.py | 10 ++++--- 6 files changed, 64 insertions(+), 17 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py index 551614626..1a71b01a8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -606,7 +606,7 @@ def _aggregate_and_complete( if self.report_run is not None: # Before the Stage 12 report, so a call Modal restarts between the # two writes finds that report running and resumes it. - self._complete_report_run() + self._complete_report_run(aggregate) completed = datetime.now(UTC) completed_parent = self.persistence.replace_report( parent.model_copy( @@ -635,11 +635,16 @@ def _aggregate_and_complete( ) return descriptor.model_dump(mode="json") - def _complete_report_run(self) -> None: - """Mark the report run succeeded; log a run this call no longer holds.""" + def _complete_report_run(self, aggregate: Mapping[str, Any]) -> None: + """Store the aggregate's result with the report run's success. + + A report run this call no longer holds is logged and left as it is. + """ if self.report_run is None: return - written = self.report_run.complete() + written = self.report_run.complete( + aggregate["result"], country=self.report.baseline.geography.country + ) if not written.changed and written.status is not RunStatus.SUCCEEDED: logger.warning( "Stage 12 report %s finished, but its report run %s is %s", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py index 3926f3972..2e3dea9f1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py @@ -22,7 +22,7 @@ from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ReportRunIds - from policyengine_stage12_persistence import RunWrite + from policyengine_stage12_persistence import ReportRunResult, RunWrite # The Modal secret and the one key it must hold: a postgresql:// URL for the # Modal worker role on the database that holds the Public API's run tables. @@ -53,8 +53,9 @@ def complete_report_run( *, modal_call_id: str, completed_at: datetime, + result: ReportRunResult | None = None, ) -> RunWrite: - """Mark the report run succeeded while ``modal_call_id`` holds it.""" + """Mark the report run succeeded with ``result`` while this call holds it.""" ... def fail_report_run( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py index 62d0b2090..f18855d42 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -15,9 +15,11 @@ from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, ) -from policyengine_stage12_persistence import RunStatus +from policyengine_stage12_persistence import RunStatus, report_run_result if TYPE_CHECKING: + from collections.abc import Mapping + from typing import Any from uuid import UUID from policyengine_simulation_contract.stage12_execution import ( @@ -56,12 +58,19 @@ def claim(self) -> RunWrite: claimed_at=datetime.now(UTC), ) - def complete(self) -> RunWrite: - """Mark the report run succeeded while this call holds it.""" + def complete(self, economy_output: Mapping[str, Any], *, country: str) -> RunWrite: + """Mark the report run succeeded with its economy output as result rows. + + The rows are written in the same transaction, only while this call + holds the run. An output that cannot be stored raises ValueError. + """ return self.writer.complete_report_run( self.run_ids, modal_call_id=self.modal_call_id, completed_at=datetime.now(UTC), + result=report_run_result( + economy_output, run_ids=self.run_ids, country=country + ), ) def fail(self, detail: Stage12FailureDetail) -> RunWrite: diff --git a/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py b/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py index 12dfcc9f3..90b0fd03a 100644 --- a/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py +++ b/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py @@ -9,9 +9,10 @@ from dataclasses import dataclass, field from uuid import UUID +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_contract.stage12_execution import ReportRunIds from policyengine_stage12_persistence import RunStatus, RunWrite -from stage12_runtime_support import invocation_context +from stage12_runtime_support import aggregate_payload, invocation_context REPORT_RUN_ID = UUID("00000000-0000-0000-0000-0000000000a1") BASELINE_RUN_ID = UUID("00000000-0000-0000-0000-0000000000b1") @@ -29,6 +30,11 @@ def report_run_context(): return invocation_context().model_copy(update={"report_run": REPORT_RUN_IDS}) +def economy_aggregate(**_values): + """Aggregate the report into a complete single-year economy output.""" + return aggregate_payload(result=CURRENT_SINGLE_YEAR_MACRO_RESULT) + + @dataclass class FakeRun: """One report run's or simulation run's mutable columns.""" @@ -38,6 +44,7 @@ class FakeRun: error_message: str | None = None output_uri: str | None = None output_sha256: str | None = None + result: object | None = None writes: list[str] = field(default_factory=list) @@ -74,7 +81,7 @@ def claim_report_run(self, report_run_id, *, modal_call_id, claimed_at): modal_call_id=modal_call_id, ) - def complete_report_run(self, run_ids, *, modal_call_id, completed_at): + def complete_report_run(self, run_ids, *, modal_call_id, completed_at, result=None): run = self._run(self.report_runs, run_ids.report_run_id, "report run") holds_it = self._holds(run, modal_call_id) return self._apply( @@ -82,6 +89,7 @@ def complete_report_run(self, run_ids, *, modal_call_id, completed_at): "complete_report_run", changed=holds_it, status=RunStatus.SUCCEEDED, + result=result, ) def fail_report_run( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index 4acc7a486..d8503b9e6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -3,6 +3,7 @@ from __future__ import annotations import pytest +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, ) @@ -15,11 +16,11 @@ REFORM_RUN_ID, REPORT_RUN_ID, FakeReportRuns, + economy_aggregate, report_run_context, ) from stage12_fake_stores import FakeArtifacts, FakeStore from stage12_runtime_support import ( - aggregate_payload, coordinate_report, parent_record, report_execution_input, @@ -51,7 +52,7 @@ def _coordinate(store, report_runs, invoker, artifacts, *, context=None): store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: aggregate_payload(), + aggregator=economy_aggregate, report_runs=report_runs, ) @@ -89,6 +90,25 @@ def test_coordinator_claims_and_completes_its_report_run_before_stage12() -> Non ] * 2 +def test_coordinator_completes_its_report_run_with_the_economy_output_as_rows() -> None: + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + + _coordinate(FakeStore(), report_runs, RecordingInvoker(artifacts), artifacts) + + result = report_runs.report_runs[REPORT_RUN_ID].result + rows = {rows.table.name: rows.rows for rows in result.rows} + budget = {row["variable_name"]: row for row in rows["budget_summary"]} + assert budget["household_tax"]["change"] == 100.0 + assert budget["household_tax"]["report_run_id"] == REPORT_RUN_ID + # The shared report is a US one, so its poverty rows use the SPM types. + assert {row["poverty_type"] for row in rows["poverty"]} == {"spm", "spm_deep"} + assert ( + result.labor_supply_result + == (CURRENT_SINGLE_YEAR_MACRO_RESULT["labor_supply_response"]) + ) + + def test_comparison_children_receive_the_context_they_received_before() -> None: artifacts = FakeArtifacts() invoker = RecordingInvoker(artifacts) @@ -189,7 +209,7 @@ def recording_artifact_store(*, content_addressed=False): coordinator_invocation_id=COORDINATOR_CALL_ID, store=FakeStore(), invoker=RecordingInvoker(artifacts), - aggregator=lambda **_: aggregate_payload(), + aggregator=economy_aggregate, report_runs=FakeReportRuns(), ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py index d3fa617c6..607ba3b98 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -15,10 +15,14 @@ ) from policyengine_stage12_persistence import RunStatus from stage12_child_call_support import ConcurrentInvoker, TimeoutInvoker -from stage12_fake_report_runs import REPORT_RUN_ID, FakeReportRuns, report_run_context +from stage12_fake_report_runs import ( + REPORT_RUN_ID, + FakeReportRuns, + economy_aggregate, + report_run_context, +) from stage12_fake_stores import FakeArtifacts, FakeStore from stage12_runtime_support import ( - aggregate_payload, coordinate_report, invocation_context, parent_record, @@ -69,7 +73,7 @@ def _coordinate(store, report_runs, invoker, *, context, **seams): store=store, artifacts=seams.pop("artifacts", FakeArtifacts()), invoker=invoker, - aggregator=lambda **_: aggregate_payload(), + aggregator=economy_aggregate, report_runs=report_runs, **seams, ) From 44f66c4f2174f8cbb7e3ce35c4c1efdbbd498789 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:33:59 +0400 Subject: [PATCH 11/24] Add a changelog fragment for recording report runs from Modal Co-Authored-By: Claude Opus 5.5 --- changelog.d/report-run-recording.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/report-run-recording.added.md diff --git a/changelog.d/report-run-recording.added.md b/changelog.d/report-run-recording.added.md new file mode 100644 index 000000000..71d024e82 --- /dev/null +++ b/changelog.d/report-run-recording.added.md @@ -0,0 +1 @@ +The v2 Modal report coordinator and its US and UK simulation workers can record a report's run in the Public API's run tables: when a call names a report run and its baseline and reform simulation runs, they write those runs' status, output file and economy results with the Modal worker role's credentials, from the Modal secret `policyengine-modal-worker-database`. The report run is marked succeeded in the same transaction as its result rows. On pre-emption the coordinator of such a call cancels its child calls and resumes when Modal restarts it. Comparison runs behave as before. From 637b68623ec682f6c9314ff2c2ab33237e198ea8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:53:46 +0400 Subject: [PATCH 12/24] Store the detailed budget, district and area impacts, and SPM receipts The Public API's result tables now have a place for the parts of the economy output that had none. A report run's result also writes the detailed budget to program_statistics, one row per program with a null entity and the output's is_tax flag; congressional district impacts with their winner, loser and no-change shares, keyed by the Census state FIPS code and district number, an at-large district being 0, from .py's US state data; UK constituency and local authority impacts with their disclosure tier and interval bounds; and the SPM provenance and data version on the report run. A district in a state .py doesn't name is refused. The UK's wealth deciles still have no place and aren't written. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 2 + .../report_run_results.py | 213 +++++++++++++++++- .../report_run_store.py | 2 + .../report_run_tables.py | 2 + .../tests/economy_output.py | 102 +++++++++ .../tests/report_run_database.py | 4 + .../tests/test_report_run_results.py | 132 ++++++++++- .../stage12_runtime/report_run_recording.py | 25 +- .../test_stage12_report_run_coordinator.py | 4 + 9 files changed, 477 insertions(+), 9 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index c6b3909a4..669f4ceb5 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -1,6 +1,7 @@ """Direct DML for temporary Stage 12 comparison state and for report runs.""" from policyengine_stage12_persistence.report_run_results import ( + DistrictCodes, ReportRunResult, report_run_result, ) @@ -22,6 +23,7 @@ __all__ = [ "ConnectionOpener", + "DistrictCodes", "ReportRunResult", "ReportRunStore", "RunStatus", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py index d689809ba..6cbec7913 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -4,7 +4,8 @@ in the format API v1 and v3 read today. Each part with a lossless place in a result table becomes rows: per-report parts keyed by the report run and its two simulation runs, poverty and inequality keyed by each simulation run. The -labour supply response and cliff impact are stored whole on the report run. +labour supply response, cliff impact and SPM receipts are stored whole on the +report run, with the data version. """ from __future__ import annotations @@ -14,7 +15,7 @@ from uuid import uuid4 from policyengine_simulation_contract.macro_output import SingleYearMacroOutput -from sqlalchemy import Column, Float, Integer, String, Table +from sqlalchemy import Boolean, Column, Float, Integer, String, Table from sqlalchemy.dialects.postgresql import ENUM, UUID from policyengine_stage12_persistence.report_run_tables import metadata @@ -22,7 +23,12 @@ if TYPE_CHECKING: from collections.abc import Mapping - from policyengine_simulation_contract.macro_output import IntraDecileOutput + from policyengine_simulation_contract.macro_output import ( + CongressionalDistrictImpactRecord, + ConstituencyImpactRecord, + IntraDecileOutput, + LocalAuthorityImpactRecord, + ) from policyengine_simulation_contract.stage12_execution import ReportRunIds decile_type_type = ENUM("income", "wealth", name="v2_decile_type", create_type=False) @@ -44,6 +50,67 @@ def _simulation_run_keys() -> tuple[Column[Any], ...]: ) +program_statistics = Table( + "program_statistics", + metadata, + *_report_run_keys(), + Column("program_name", String(512), nullable=False), + Column("entity", String(128)), + Column("is_tax", Boolean), + Column("baseline_total", Float), + Column("reform_total", Float), + Column("change", Float), +) + + +def _area_values() -> tuple[Column[Any], ...]: + return ( + Column("average_household_income_change", Float, nullable=False), + Column("relative_household_income_change", Float, nullable=False), + Column("population", Float, nullable=False), + ) + + +def _disclosure_values() -> tuple[Column[Any], ...]: + return ( + Column("disclosure_tier", String(16)), + Column("average_household_income_change_lower", Float), + Column("average_household_income_change_upper", Float), + Column("relative_household_income_change_lower", Float), + Column("relative_household_income_change_upper", Float), + ) + + +congressional_district_impacts = Table( + "congressional_district_impacts", + metadata, + *_report_run_keys(), + *_area_values(), + Column("district_geoid", Integer, nullable=False), + Column("state_fips", Integer, nullable=False), + Column("district_number", Integer, nullable=False), + Column("winner_percentage", Float, nullable=False), + Column("loser_percentage", Float, nullable=False), + Column("no_change_percentage", Float, nullable=False), +) +constituency_impacts = Table( + "constituency_impacts", + metadata, + *_report_run_keys(), + *_area_values(), + *_disclosure_values(), + Column("constituency_code", String(64), nullable=False), + Column("constituency_name", String(255), nullable=False), +) +local_authority_impacts = Table( + "local_authority_impacts", + metadata, + *_report_run_keys(), + *_area_values(), + *_disclosure_values(), + Column("local_authority_code", String(64), nullable=False), + Column("local_authority_name", String(255), nullable=False), +) budget_summary = Table( "budget_summary", metadata, @@ -125,6 +192,26 @@ def _simulation_run_keys() -> tuple[Column[Any], ...]: } # The intra-decile row holding each category's share across all deciles. ALL_DECILES = 0 +DISCLOSURE_FIELDS = ( + "disclosure_tier", + "average_household_income_change_lower", + "average_household_income_change_upper", + "relative_household_income_change_lower", + "relative_household_income_change_upper", +) + + +@dataclass(frozen=True) +class DistrictCodes: + """US states' FIPS codes, and the states with one at-large district. + + Today's output names a district as ``"ST-NN"``, numbering an at-large + district ``01``. Its row stores the Census GEOID parts instead: the + state's FIPS code, and the district number, ``0`` for an at-large state. + """ + + state_fips: Mapping[str, int] + at_large_states: frozenset[str] @dataclass(frozen=True) @@ -147,6 +234,8 @@ class ReportRunResult: rows: tuple[ResultRows, ...] labor_supply_result: dict[str, Any] | None cliff_impact_result: dict[str, Any] | None + spm_provenance: dict[str, Any] | None = None + data_version: str | None = None def report_run_result( @@ -154,11 +243,13 @@ def report_run_result( *, run_ids: ReportRunIds, country: str, + district_codes: DistrictCodes | None = None, ) -> ReportRunResult: """Validate a single-year economy output and return the rows it is stored as. - Raises ValueError for an output that is not a single-year economy result - or a country without poverty types. + Raises ValueError for an output that is not a single-year economy result, + a country without poverty types, or a congressional district whose state + ``district_codes`` doesn't name. """ economy = SingleYearMacroOutput.model_validate(output) if country not in POVERTY_TYPES: @@ -197,9 +288,121 @@ def report_run_result( _inequality_rows(economy, run_ids, country), shared_by_reports=True, ), + ResultRows( + program_statistics, + _program_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + congressional_district_impacts, + _district_rows(economy, report_keys, district_codes), + shared_by_reports=False, + ), + ResultRows( + constituency_impacts, + _area_rows(economy.constituency_impact, report_keys, "constituency"), + shared_by_reports=False, + ), + ResultRows( + local_authority_impacts, + _area_rows( + economy.local_authority_impact, report_keys, "local_authority" + ), + shared_by_reports=False, + ), ), labor_supply_result=dumped["labor_supply_response"], cliff_impact_result=dumped["cliff_impact"], + spm_provenance=dumped["spm_provenance"], + data_version=economy.data_version, + ) + + +def _program_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + # The output names no program's entity, so every program's is null. + return tuple( + { + "id": uuid4(), + **keys, + "program_name": program_name, + "entity": None, + "is_tax": program.is_tax, + "baseline_total": program.baseline, + "reform_total": program.reform, + "change": program.difference, + } + for program_name, program in economy.detailed_budget.root.items() + ) + + +def _district_rows( + economy: SingleYearMacroOutput, + keys: Mapping[str, Any], + district_codes: DistrictCodes | None, +) -> tuple[dict[str, Any], ...]: + if economy.congressional_district_impact is None: + return () + districts = economy.congressional_district_impact.districts + if districts and district_codes is None: + msg = "congressional districts need the US states' FIPS codes" + raise ValueError(msg) + return tuple( + _district_row(district, keys, district_codes) + for district in districts + if district_codes is not None + ) + + +def _district_row( + district: CongressionalDistrictImpactRecord, + keys: Mapping[str, Any], + district_codes: DistrictCodes, +) -> dict[str, Any]: + state, _, number = district.district.partition("-") + if state not in district_codes.state_fips or not number.isdigit(): + msg = f"congressional district {district.district!r} names no known state" + raise ValueError(msg) + state_fips = district_codes.state_fips[state] + district_number = 0 if state in district_codes.at_large_states else int(number) + return { + "id": uuid4(), + **keys, + "district_geoid": state_fips * 100 + district_number, + "state_fips": state_fips, + "district_number": district_number, + "average_household_income_change": district.average_household_income_change, + "relative_household_income_change": district.relative_household_income_change, + "winner_percentage": district.winner_percentage, + "loser_percentage": district.loser_percentage, + "no_change_percentage": district.no_change_percentage, + "population": district.population, + } + + +def _area_rows( + areas: Any, # noqa: ANN401 # a GeographicImpactOutput, or None + keys: Mapping[str, Any], + kind: str, +) -> tuple[dict[str, Any], ...]: + if areas is None: + return () + records: list[ConstituencyImpactRecord | LocalAuthorityImpactRecord] = areas.root + return tuple( + { + "id": uuid4(), + **keys, + f"{kind}_code": getattr(record, f"{kind}_code"), + f"{kind}_name": getattr(record, f"{kind}_name"), + "average_household_income_change": record.average_household_income_change, + "relative_household_income_change": ( + record.relative_household_income_change + ), + "population": record.population, + **{field: getattr(record, field) for field in DISCLOSURE_FIELDS}, + } + for record in records ) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index bbf9197e1..1fc37d414 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -179,6 +179,8 @@ def complete_report_run( { "labor_supply_result": result.labor_supply_result, "cliff_impact_result": result.cliff_impact_result, + "spm_provenance": result.spm_provenance, + "data_version": result.data_version, } if result is not None else {} diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py index 7ef743a10..74cd36802 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -43,6 +43,8 @@ Column("updated_at", DateTime(timezone=True), nullable=False), Column("labor_supply_result", JSON(none_as_null=True)), Column("cliff_impact_result", JSON(none_as_null=True)), + Column("spm_provenance", JSON(none_as_null=True)), + Column("data_version", String(128)), ) simulation_runs = Table( diff --git a/libs/policyengine-stage12-persistence/tests/economy_output.py b/libs/policyengine-stage12-persistence/tests/economy_output.py index a59e4dc54..c02bfaac0 100644 --- a/libs/policyengine-stage12-persistence/tests/economy_output.py +++ b/libs/policyengine-stage12-persistence/tests/economy_output.py @@ -5,6 +5,13 @@ from copy import deepcopy from typing import Any +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_comparison_provenance, + build_spm_provenance, +) + US_ECONOMY_OUTPUT: dict[str, Any] = { "spm_provenance": None, "model_version": "1.715.2", @@ -120,3 +127,98 @@ def uk_economy_output() -> dict[str, Any]: output["intra_wealth_decile"] = deepcopy(output["intra_decile"]) output["intra_wealth_decile"]["all"]["No change"] = 0.5 return output + + +US_STATE_FIPS = {"AL": 1, "AK": 2, "CA": 6} +US_AT_LARGE_STATES = frozenset({"AK"}) + + +def us_output_with_every_part() -> dict[str, Any]: + """Return the US output with districts, SPM receipts and an unflagged program.""" + receipt = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="a" * 64, + selection=SPMSelection( + forecast_content_sha256="a" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ), + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "0.0.0-test-policyengine", + "policyengine-core": "0.0.0-test-policyengine-core", + "policyengine-us": "0.0.0-test-policyengine-us", + "spm-calculator": "0.0.0-test-spm-calculator", + } + ), + ) + output = deepcopy(US_ECONOMY_OUTPUT) + output["spm_provenance"] = build_spm_comparison_provenance( + baseline_receipts=[receipt], reform_receipts=[receipt] + ).model_dump(mode="json", by_alias=True) + output["detailed_budget"]["snap"] = { + "baseline": 80.0, + "reform": 70.0, + "difference": -10.0, + "is_tax": None, + } + output["congressional_district_impact"] = { + "districts": [ + { + "district": district, + "average_household_income_change": 10.0 + index, + "relative_household_income_change": 0.01 * (index + 1), + "winner_percentage": 0.6 - index / 10, + "loser_percentage": 0.3, + "no_change_percentage": 0.1 + index / 10, + "population": 1000.0 + index, + } + for index, district in enumerate(("AL-01", "AL-07", "AK-01", "CA-12")) + ] + } + return output + + +def uk_output_with_areas() -> dict[str, Any]: + """Return the UK output with constituencies and local authorities.""" + output = uk_economy_output() + output["constituency_impact"] = [ + { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "average_household_income_change": 120.0, + "relative_household_income_change": 0.004, + "population": 45000.0, + "disclosure_tier": "reported", + "average_household_income_change_lower": 100.0, + "average_household_income_change_upper": 140.0, + "relative_household_income_change_lower": 0.003, + "relative_household_income_change_upper": 0.005, + }, + { + "constituency_code": "W07000081", + "constituency_name": "Aberafan Maesteg", + "average_household_income_change": 90.0, + "relative_household_income_change": 0.003, + "population": 40000.0, + }, + ] + output["local_authority_impact"] = [ + { + "local_authority_code": "E06000001", + "local_authority_name": "Hartlepool", + "average_household_income_change": 80.0, + "relative_household_income_change": 0.002, + "population": 42000.0, + "disclosure_tier": "approximate", + "average_household_income_change_lower": 50.0, + "average_household_income_change_upper": 110.0, + "relative_household_income_change_lower": 0.001, + "relative_household_income_change_upper": 0.003, + } + ] + return output diff --git a/libs/policyengine-stage12-persistence/tests/report_run_database.py b/libs/policyengine-stage12-persistence/tests/report_run_database.py index a16499828..876b88d66 100644 --- a/libs/policyengine-stage12-persistence/tests/report_run_database.py +++ b/libs/policyengine-stage12-persistence/tests/report_run_database.py @@ -39,6 +39,10 @@ "intra_decile_impacts", "poverty", "inequality", + "program_statistics", + "congressional_district_impacts", + "constituency_impacts", + "local_authority_impacts", ) NOW = datetime(2026, 11, 4, tzinfo=UTC) VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "6.12.0"} diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py index 63043171d..457097aa9 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -6,8 +6,16 @@ from uuid import uuid4 import pytest -from economy_output import US_ECONOMY_OUTPUT, uk_economy_output +from economy_output import ( + US_AT_LARGE_STATES, + US_ECONOMY_OUTPUT, + US_STATE_FIPS, + uk_economy_output, + uk_output_with_areas, + us_output_with_every_part, +) from policyengine_stage12_persistence import ( + DistrictCodes, ReportRunStore, RunStatus, RunWrite, @@ -38,6 +46,22 @@ def store(): run_store.engine.dispose() +DISTRICT_CODES = DistrictCodes( + state_fips=US_STATE_FIPS, at_large_states=US_AT_LARGE_STATES +) +STATE_BY_FIPS = {fips: state for state, fips in US_STATE_FIPS.items()} +AREA_FIELDS = ( + "average_household_income_change", + "relative_household_income_change", + "population", + "disclosure_tier", + "average_household_income_change_lower", + "average_household_income_change_upper", + "relative_household_income_change_lower", + "relative_household_income_change_upper", +) + + def _complete(store, seeded_run, output=US_ECONOMY_OUTPUT, country="us"): store.claim_report_run( seeded_run.run_ids.report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW @@ -46,10 +70,22 @@ def _complete(store, seeded_run, output=US_ECONOMY_OUTPUT, country="us"): seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER, - result=report_run_result(output, run_ids=seeded_run.run_ids, country=country), + result=report_run_result( + output, + run_ids=seeded_run.run_ids, + country=country, + district_codes=DISTRICT_CODES, + ), ) +def _rebuilt_district(row): + """Rebuild today's district code as the output builder formats it.""" + state = STATE_BY_FIPS[row["state_fips"]] + number = 1 if state in US_AT_LARGE_STATES else row["district_number"] + return f"{state}-{number:02d}" + + def _rebuilt_budget(seeder, report_run_id): rows = { row["variable_name"]: row @@ -272,3 +308,95 @@ def test_a_shared_simulation_run_keeps_the_results_first_stored_for_it( assert len(stored) == len({row["filter_variable"] for row in stored}) reform = seeder.rows("poverty", simulation_run_id=run_ids.reform_simulation_run_id) assert len(reform) == 2 * 6 + 4 + + +def test_the_detailed_budget_districts_and_receipts_are_stored_losslessly( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + output = us_output_with_every_part() + + _complete(store, seeded_run, output=output) + + programs = seeder.rows("program_statistics", report_run_id=report_run_id) + assert { + row["program_name"]: { + "baseline": row["baseline_total"], + "reform": row["reform_total"], + "difference": row["change"], + "is_tax": row["is_tax"], + } + for row in programs + } == output["detailed_budget"] + assert {row["entity"] for row in programs} == {None} + districts = seeder.rows( + "congressional_district_impacts", report_run_id=report_run_id + ) + assert sorted( + ( + { + "district": _rebuilt_district(row), + **{ + field: row[field] + for field in ( + "average_household_income_change", + "relative_household_income_change", + "winner_percentage", + "loser_percentage", + "no_change_percentage", + "population", + ) + }, + } + for row in districts + ), + key=lambda district: district["district"], + ) == sorted( + output["congressional_district_impact"]["districts"], + key=lambda district: district["district"], + ) + # Census GEOIDs: an at-large district is numbered 00. + assert {row["district_geoid"] for row in districts} == {101, 107, 200, 612} + report_run = seeder.row("report_runs", report_run_id) + assert report_run["spm_provenance"] == output["spm_provenance"] + assert report_run["data_version"] == output["data_version"] + + +def test_the_uk_area_impacts_are_stored_with_their_disclosure_fields( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + output = uk_output_with_areas() + + _complete(store, seeded_run, output=output, country="uk") + + for table_name, kind in ( + ("constituency_impacts", "constituency"), + ("local_authority_impacts", "local_authority"), + ): + stored = { + row[f"{kind}_code"]: { + f"{kind}_code": row[f"{kind}_code"], + f"{kind}_name": row[f"{kind}_name"], + **{field: row[field] for field in AREA_FIELDS}, + } + for row in seeder.rows(table_name, report_run_id=report_run_id) + } + expected = { + record[f"{kind}_code"]: dict.fromkeys(AREA_FIELDS) | record + for record in output[f"{kind}_impact"] + } + assert stored == expected + + +def test_a_district_in_an_unknown_state_is_refused(seeded_run: SeededReportRun) -> None: + output = us_output_with_every_part() + output["congressional_district_impact"]["districts"][0]["district"] = "ZZ-01" + + with pytest.raises(ValueError, match="ZZ"): + report_run_result( + output, + run_ids=seeded_run.run_ids, + country="us", + district_codes=DISTRICT_CODES, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py index f18855d42..57fc0261e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -15,7 +15,11 @@ from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, ) -from policyengine_stage12_persistence import RunStatus, report_run_result +from policyengine_stage12_persistence import ( + DistrictCodes, + RunStatus, + report_run_result, +) if TYPE_CHECKING: from collections.abc import Mapping @@ -36,6 +40,20 @@ from .report_run_access import ReportRunWriter +def us_district_codes() -> DistrictCodes: + """Return .py's US state FIPS codes and at-large states, for district rows.""" + # .py's US data loads only for a US report's result. + from policyengine.countries.us.data import ( + AT_LARGE_STATES, + US_STATE_FIPS, + ) + + return DistrictCodes( + state_fips={state: int(fips) for state, fips in US_STATE_FIPS.items()}, + at_large_states=frozenset(AT_LARGE_STATES), + ) + + def run_error_message(error_code: str | None, error_summary: str | None) -> str: """Return the redacted failure a run records: its code, then its summary.""" return f"{error_code or 'simulation_failed'}: {error_summary or 'no summary'}" @@ -69,7 +87,10 @@ def complete(self, economy_output: Mapping[str, Any], *, country: str) -> RunWri modal_call_id=self.modal_call_id, completed_at=datetime.now(UTC), result=report_run_result( - economy_output, run_ids=self.run_ids, country=country + economy_output, + run_ids=self.run_ids, + country=country, + district_codes=us_district_codes() if country == "us" else None, ), ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index d8503b9e6..57dd6b72d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -107,6 +107,10 @@ def test_coordinator_completes_its_report_run_with_the_economy_output_as_rows() result.labor_supply_result == (CURRENT_SINGLE_YEAR_MACRO_RESULT["labor_supply_response"]) ) + # Alabama's first district, by .py's state FIPS codes. + (district,) = rows["congressional_district_impacts"] + assert (district["state_fips"], district["district_number"]) == (1, 1) + assert district["district_geoid"] == 101 def test_comparison_children_receive_the_context_they_received_before() -> None: From 5c45699f58b0ba075a24578c8034ec214479be59 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:55:21 +0400 Subject: [PATCH 13/24] Store decile results with the meaning .py's rows give the columns The result tables follow .py's row models. decile_impacts.relative_change holds .py's percentage, while the economy output holds a fraction, so the writer now multiplies it by 100 and a reader divides it again, equal to within floating-point rounding. intra_decile_impacts row 0 is .py's share across the whole population, while the output's overall shares are the mean of the decile shares; the writer no longer puts the mean in row 0, and a reader computes it again from deciles 1 to 10, exactly as the output does. Co-Authored-By: Claude Opus 5.5 --- .../report_run_results.py | 35 ++++++++++--------- .../tests/economy_output.py | 5 ++- .../tests/test_report_run_results.py | 31 +++++++++++----- 3 files changed, 44 insertions(+), 27 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py index 6cbec7913..b6605a25b 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -190,8 +190,6 @@ def _disclosure_values() -> tuple[Column[Any], ...]: "Gain less than 5%": "gain_less_than_5pct", "Gain more than 5%": "gain_more_than_5pct", } -# The intra-decile row holding each category's share across all deciles. -ALL_DECILES = 0 DISCLOSURE_FIELDS = ( "disclosure_tier", "average_household_income_change_lower", @@ -433,7 +431,8 @@ def _decile_rows( "decile": int(decile), "quantiles": 10, "absolute_change": average, - "relative_change": economy.decile.relative[decile], + # .py's rows hold a percentage; today's output, a fraction. + "relative_change": economy.decile.relative[decile] * 100, } for decile, average in economy.decile.average.items() ) @@ -451,24 +450,26 @@ def _intra_decile_rows( def _intra_decile_type_rows( shares: IntraDecileOutput, decile_type: str, keys: Mapping[str, Any] ) -> tuple[dict[str, Any], ...]: + """Return one row per decile, 1 to 10. + + Row 0 is .py's share across the whole population, which the output + doesn't carry: its ``all`` is the mean of the decile shares, which a + reader computes again from these rows. + """ decile_count = max((len(values) for values in shares.deciles.values()), default=0) - by_decile: list[dict[str, Any]] = [ + return tuple( { - column: shares.deciles[label][index] - for label, column in INTRA_DECILE_COLUMNS.items() - if index < len(shares.deciles.get(label, ())) + "id": uuid4(), + **keys, + "decile_type": decile_type, + "decile": index + 1, + **{ + column: shares.deciles[label][index] + for label, column in INTRA_DECILE_COLUMNS.items() + if index < len(shares.deciles.get(label, ())) + }, } - | {"decile": index + 1} for index in range(decile_count) - ] - overall = { - column: shares.all[label] - for label, column in INTRA_DECILE_COLUMNS.items() - if label in shares.all - } | {"decile": ALL_DECILES} - return tuple( - {"id": uuid4(), **keys, "decile_type": decile_type, **values} - for values in (*by_decile, overall) ) diff --git a/libs/policyengine-stage12-persistence/tests/economy_output.py b/libs/policyengine-stage12-persistence/tests/economy_output.py index c02bfaac0..6898079a5 100644 --- a/libs/policyengine-stage12-persistence/tests/economy_output.py +++ b/libs/policyengine-stage12-persistence/tests/economy_output.py @@ -125,7 +125,10 @@ def uk_economy_output() -> dict[str, Any]: output["poverty_by_race"] = None output["budget"]["state_tax_revenue_impact"] = 0.0 output["intra_wealth_decile"] = deepcopy(output["intra_decile"]) - output["intra_wealth_decile"]["all"]["No change"] = 0.5 + wealth = output["intra_wealth_decile"] + wealth["deciles"]["No change"] = [0.4, 0.6] + # Today's builder makes each overall share the mean of the decile shares. + wealth["all"]["No change"] = (0.4 + 0.6) / 2 return output diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py index 457097aa9..9af11ec1d 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -104,16 +104,19 @@ def _rebuilt_budget(seeder, report_run_id): def _rebuilt_intra_decile(rows): + """Rebuild today's shares; each overall one is the deciles' mean, as today.""" by_decile = {row["decile"]: row for row in rows} - deciles = sorted(decile for decile in by_decile if decile != 0) + deciles = sorted(by_decile) + shares = { + label: [by_decile[decile][column] for decile in deciles] + for label, column in INTRA_DECILE_LABELS.items() + } return { "all": { - label: by_decile[0][column] for label, column in INTRA_DECILE_LABELS.items() - }, - "deciles": { - label: [by_decile[decile][column] for decile in deciles] - for label, column in INTRA_DECILE_LABELS.items() + label: sum(values) / len(values) if values else 0.0 + for label, values in shares.items() }, + "deciles": shares, } @@ -139,10 +142,17 @@ def test_a_report_run_stores_its_economy_output_as_rows( output = US_ECONOMY_OUTPUT assert _rebuilt_budget(seeder, run_ids.report_run_id) == output["budget"] deciles = seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id) + assert {str(row["decile"]): row["absolute_change"] for row in deciles} == output[ + "decile" + ]["average"] + # .py's rows hold the relative change as a percentage; today's output + # holds a fraction, so a reader divides by 100. assert { - "average": {str(row["decile"]): row["absolute_change"] for row in deciles}, - "relative": {str(row["decile"]): row["relative_change"] for row in deciles}, - } == output["decile"] + str(row["decile"]): row["relative_change"] / 100 for row in deciles + } == pytest.approx(output["decile"]["relative"], rel=1e-12) + assert sorted(row["relative_change"] for row in deciles) == pytest.approx( + [1.0, 2.0] + ) assert {row["income_variable"] for row in deciles} == {"household_net_income"} assert all( row["baseline_simulation_run_id"] == run_ids.baseline_simulation_run_id @@ -150,6 +160,9 @@ def test_a_report_run_stores_its_economy_output_as_rows( for row in deciles ) intra = seeder.rows("intra_decile_impacts", report_run_id=run_ids.report_run_id) + # Row 0 is .py's share across the whole population, which today's output + # doesn't carry, so none is written. + assert sorted(row["decile"] for row in intra) == [1, 2] assert _rebuilt_intra_decile(intra) == output["intra_decile"] for side, run_id in ( ("baseline", run_ids.baseline_simulation_run_id), From e7288dd296e4ba33c9761529855645244211f54c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:59:32 +0400 Subject: [PATCH 14/24] Write each Stage 12 record before its report or simulation run The coordinator now marks its Stage 12 report succeeded or failed first, and its report run second; each worker marks its Stage 12 child record succeeded first, and its simulation run second. A call Modal restarts between the two writes therefore finds the Stage 12 record finished. A coordinator restarted after its Stage 12 report succeeded under the same call ID reads the stored aggregate, checks its digest, and completes the report run from it without calculating again; one restarted after a failure runs the report again. A worker restarted after its child record succeeded completes its simulation run from that record, so both name the same output file. Before, a restart after the report run succeeded left the Stage 12 report running for good. The claims now live in their own module. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 35 +----- .../stage12_runtime/report_claims.py | 109 ++++++++++++++++++ .../stage12_runtime/report_coordination.py | 25 ++-- .../stage12_runtime/report_run_recording.py | 34 +++--- .../stage12_runtime/simulation.py | 9 +- .../test_stage12_report_run_coordinator.py | 28 ++++- .../test_stage12_report_run_preemption.py | 66 +++++++++-- .../tests/test_stage12_report_run_worker.py | 52 +++++++-- 8 files changed, 277 insertions(+), 81 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 35acdc39e..449a6ef60 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -27,12 +27,11 @@ from .coordinated_simulations import SIMULATION_CANCELLED_ERROR_CODE, logger from .dependencies import ModalChildInvoker, artifact_store, runtime_store from .failure_causes import StageProgress -from .failures import Stage12FailureDetail from .parent_claim import ( - claim_parent, validate_dispatch_context, validate_submitted_parent, ) +from .report_claims import claim_report from .report_coordination import ( SIMULATION_POLL_BACKOFF_FACTOR, SIMULATION_POLL_INITIAL_SECONDS, @@ -118,35 +117,9 @@ def coordinate_report( report_run=report_run, ) with coordination.span(Stage.STAGE12_COORDINATOR_CLAIM): - if report_run is not None: - # The report run decides whether this call works: another call - # may hold it, or it may already be finished. - claimed = report_run.claim() - if not claimed.changed: - return { - "deduplicated": True, - "report_run_id": str(report_run.run_ids.report_run_id), - "status": claimed.status.value, - } - parent, should_execute = claim_parent( - submitted=submitted_parent, - coordinator_invocation_id=invocation.invocation_id, - store=coordination.persistence, - resume_own_claim=report_run is not None, - ) - if not should_execute: - if report_run is not None: - report_run.fail( - Stage12FailureDetail( - error_code="report_already_claimed", - error_summary="Another call holds this report's Stage 12 record.", - ) - ) - return { - "deduplicated": True, - "evaluation_id": str(parent.evaluation_id), - "status": parent.status.value, - } + parent, summary = claim_report(coordination, submitted=submitted_parent) + if parent is None: + return summary or {} try: return coordination.run(parent) # Persist a redacted cause for any coordinator, child-call, aggregation, or diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py new file mode 100644 index 000000000..b97f34bf5 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py @@ -0,0 +1,109 @@ +"""Claiming a call's Stage 12 report, and the report run its call names. + +A comparison run claims only its Stage 12 report, as before. A call that names +a report run claims that run first, under its Modal call ID, then its Stage 12 +report. Each claim either lets the call run the report or gives the summary +the call returns instead. +""" + +from __future__ import annotations + +from hashlib import sha256 +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_contract.stage12_execution import ( + AggregateReportArtifactPayload, + ComparisonRunLifecycleStatus, +) +from policyengine_stage12_persistence import RunStatus + +from .failures import Stage12FailureDetail +from .parent_claim import claim_parent + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_execution import ( + ComparisonReportRecord, + ) + + from .report_coordination import ReportCoordination + from .report_run_recording import ReportRunRecorder + + +def claim_report( + coordination: ReportCoordination, + *, + submitted: ComparisonReportRecord, +) -> tuple[ComparisonReportRecord | None, dict[str, Any] | None]: + """Return the Stage 12 report to run, or the summary to return instead. + + A report run another call holds, or one already finished, returns its + status. A call Modal restarted after its Stage 12 report succeeded, but + before its report run did, completes the report run from the stored + aggregate and returns that it succeeded. + """ + report_run = coordination.report_run + invocation_id = coordination.invocation.invocation_id + if report_run is not None: + claimed = report_run.claim() + if not claimed.changed: + return None, _report_run_summary(report_run, claimed.status) + parent, should_execute = claim_parent( + submitted=submitted, + coordinator_invocation_id=invocation_id, + store=coordination.persistence, + resume_own_claim=report_run is not None, + ) + if should_execute: + return parent, None + if report_run is None: + return None, _stage12_summary(parent) + if ( + parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + and parent.coordinator_invocation_id == invocation_id + ): + coordination.complete_report_run(_stored_aggregate(coordination, parent)) + return None, _report_run_summary(report_run, RunStatus.SUCCEEDED) + report_run.fail( + Stage12FailureDetail( + error_code="report_already_claimed", + error_summary=( + f"This report's Stage 12 record is {parent.status.value} " + "under another call." + ), + ) + ) + return None, _stage12_summary(parent) + + +def _stored_aggregate( + coordination: ReportCoordination, parent: ComparisonReportRecord +) -> dict[str, Any]: + """Read and digest-check the aggregate a succeeded Stage 12 report names.""" + if parent.aggregate_output_uri is None or parent.aggregate_output_sha256 is None: + msg = "a succeeded Stage 12 report names no aggregate" + raise ValueError(msg) + payload = coordination.artifact_storage.read(parent.aggregate_output_uri) + if sha256(payload).hexdigest() != parent.aggregate_output_sha256: + msg = "stored aggregate artifact digest mismatch" + raise ValueError(msg) + return AggregateReportArtifactPayload.model_validate_json(payload).model_dump( + mode="json" + ) + + +def _report_run_summary( + report_run: ReportRunRecorder, status: RunStatus +) -> dict[str, Any]: + return { + "deduplicated": True, + "report_run_id": str(report_run.run_ids.report_run_id), + "status": status.value, + } + + +def _stage12_summary(parent: ComparisonReportRecord) -> dict[str, Any]: + return { + "deduplicated": True, + "evaluation_id": str(parent.evaluation_id), + "status": parent.status.value, + } diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py index 1a71b01a8..ecaa50767 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -603,10 +603,6 @@ def _aggregate_and_complete( reform_artifact_sha256=reform.descriptor.artifact.content_sha256, bundle=self.report.baseline.bundle, ) - if self.report_run is not None: - # Before the Stage 12 report, so a call Modal restarts between the - # two writes finds that report running and resumes it. - self._complete_report_run(aggregate) completed = datetime.now(UTC) completed_parent = self.persistence.replace_report( parent.model_copy( @@ -621,6 +617,10 @@ def _aggregate_and_complete( } ) ) + # After the Stage 12 report: a call Modal restarts between the two + # writes finds that report succeeded and completes the report run + # from its stored aggregate. + self.complete_report_run(aggregate) compare_completed_report( parent=completed_parent, aggregate=aggregate, @@ -635,10 +635,11 @@ def _aggregate_and_complete( ) return descriptor.model_dump(mode="json") - def _complete_report_run(self, aggregate: Mapping[str, Any]) -> None: + def complete_report_run(self, aggregate: Mapping[str, Any]) -> None: """Store the aggregate's result with the report run's success. A report run this call no longer holds is logged and left as it is. + A call that names no report run records nothing. """ if self.report_run is None: return @@ -684,12 +685,6 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: detail=cancellation("the coordinator failed"), ) record_failure_cause(self.runtime, error, detail.error_summary) - if self.report_run is not None: - persist_failure_record( - partial(self.report_run.fail, detail), - runtime=self.runtime, - failure_scope="report_coordination", - ) failed_at = datetime.now(UTC) persist_failure_record( partial( @@ -708,3 +703,11 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: runtime=self.runtime, failure_scope="report_coordination", ) + # After the Stage 12 report, as on success: a call Modal restarts + # between the two writes finds that report failed and runs it again. + if self.report_run is not None: + persist_failure_record( + partial(self.report_run.fail, detail), + runtime=self.runtime, + failure_scope="report_coordination", + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py index 57fc0261e..ab6758794 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -185,30 +185,24 @@ def replace_simulation_if_status( *, expected_status: ComparisonRunLifecycleStatus, ) -> tuple[ComparisonSimulationRecord, bool]: - """Write the child record and its simulation run, in a restart-safe order. - - A success is recorded on the simulation run first, so a call Modal - restarts between the two writes finds its Stage 12 record still - running and runs again without recording a second output. A claim or - failure is recorded on the simulation run only after the Stage 12 - record changed. A simulation run that is already failed refuses the - child's claim and success, which its worker treats as a cancellation. + """Write the child record, then its simulation run once the record changed. + + A worker Modal restarts after its record succeeded, but before its + simulation run did, completes the run from that record (see + ``complete_from_succeeded_child``), so both name the same output + file. A simulation run that is already failed refuses the child's + claim, which its worker treats as a cancellation. """ simulation_run_id = self._run_ids.simulation_run_id(record.role) written_at = record.updated_at or datetime.now(UTC) - if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: - completed = self._complete(simulation_run_id, record, written_at) - if completed is not RunStatus.SUCCEEDED: - return self._store.get_simulation(record.simulation_execution_id), False - return self._store.replace_simulation_if_status( - record, expected_status=expected_status - ) stored, replaced = self._store.replace_simulation_if_status( record, expected_status=expected_status ) if not replaced: return stored, replaced - if record.status is ComparisonRunLifecycleStatus.RUNNING: + if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: + self._complete(simulation_run_id, record, written_at) + elif record.status is ComparisonRunLifecycleStatus.RUNNING: claimed = self._runs.claim_simulation_run( simulation_run_id, modal_call_id=self._require_modal_call_id(), @@ -226,6 +220,14 @@ def replace_simulation_if_status( ) return stored, replaced + def complete_from_succeeded_child(self, record: ComparisonSimulationRecord) -> None: + """Complete a child's simulation run from its succeeded Stage 12 record.""" + self._complete( + self._run_ids.simulation_run_id(record.role), + record, + record.updated_at or datetime.now(UTC), + ) + def _complete( self, simulation_run_id: UUID, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index c99a9a361..badfc26e5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -521,21 +521,28 @@ def run_single_simulation( context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) persistence = dependencies.store or runtime_store(runtime) + recording = None if context.report_run is not None: if modal_call_id is None: msg = "a report run's worker needs its Modal call ID" raise ValueError(msg) - persistence = SimulationRunRecordingStore( + recording = SimulationRunRecordingStore( persistence, dependencies.report_runs or report_run_store(runtime), context.report_run, modal_call_id=modal_call_id, ) + persistence = recording artifact_storage = dependencies.artifacts or artifact_store( content_addressed=context.report_run is not None ) child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: + if recording is not None: + # An earlier attempt stopped between the child's and the run's + # success; the run's completion is conditional, so this repeats + # nothing that was recorded. + recording.complete_from_succeeded_child(child) return descriptor_from_record(child, simulation).model_dump(mode="json") running = _claim_simulation(persistence, child) cancellation_requested = _cancellation_check(persistence, simulation) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index 57dd6b72d..c26499e70 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -68,7 +68,9 @@ def recording_replace_report(record): return store -def test_coordinator_claims_and_completes_its_report_run_before_stage12() -> None: +def test_coordinator_claims_its_report_run_first_and_completes_it_after_stage12() -> ( + None +): events: list[str] = [] store = _recording_report_writes(FakeStore(), events) report_runs = FakeReportRuns(events) @@ -81,9 +83,10 @@ def test_coordinator_claims_and_completes_its_report_run_before_stage12() -> Non assert run.status is RunStatus.SUCCEEDED assert run.modal_call_id == COORDINATOR_CALL_ID assert result["evaluation_id"] == str(report_execution_input().evaluation_id) - # The report run is claimed before, and completed before, the Stage 12 - # report, so a restart between the two finds the Stage 12 report running. - assert events == ["claim_report_run", "complete_report_run", "stage12:succeeded"] + # The report run is claimed first and completed last, so a call Modal + # restarts between the two success writes finds its Stage 12 report + # succeeded, and completes the report run from the stored aggregate. + assert events == ["claim_report_run", "stage12:succeeded", "complete_report_run"] assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED assert [context["report_run"] for context in invoker.contexts] == [ report_run_context().report_run.model_dump(mode="json") @@ -195,6 +198,23 @@ def test_a_failed_report_fails_its_report_run_and_the_failed_simulation_run() -> assert store.parent.status is ComparisonRunLifecycleStatus.FAILED +def test_a_failed_report_records_its_stage12_failure_before_its_report_run() -> None: + events: list[str] = [] + store = _recording_report_writes(FakeStore(), events) + report_runs = FakeReportRuns(events) + artifacts = FakeArtifacts() + invoker = RecordingInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises(RuntimeError, match="coordination failed"): + _coordinate(store, report_runs, invoker, artifacts) + + assert [ + event for event in events if event in {"stage12:failed", "fail_report_run"} + ] == ["stage12:failed", "fail_report_run"] + + def test_report_run_coordinator_names_its_aggregate_by_digest(monkeypatch) -> None: built = [] artifacts = FakeArtifacts() diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py index 607ba3b98..8199845f3 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -143,19 +143,21 @@ def test_restarted_coordinator_resumes_its_claim_and_starts_its_children_again() assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED -def test_a_call_restarted_after_its_report_run_succeeded_records_nothing_again() -> ( - None -): - store = FakeStore() - report_runs = FakeReportRuns() - replace_report = store.replace_report +def _preempted_around_completion(report_runs, *, after_writing): + """Make the report run's completion pre-empt the call, before or after writing.""" + complete = report_runs.complete_report_run - def preempted_before_the_stage12_success(record): - if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: - raise KeyboardInterrupt - return replace_report(record) + def preempted_completion(*args, **kwargs): + if after_writing: + complete(*args, **kwargs) + raise KeyboardInterrupt - store.replace_report = preempted_before_the_stage12_success + report_runs.complete_report_run = preempted_completion + return complete + + +def _first_attempt_preempted_at_completion(store, report_runs, *, after_writing): + complete = _preempted_around_completion(report_runs, after_writing=after_writing) artifacts = FakeArtifacts() with pytest.raises(KeyboardInterrupt): _coordinate( @@ -165,6 +167,46 @@ def preempted_before_the_stage12_success(record): context=report_run_context(), artifacts=artifacts, ) + report_runs.complete_report_run = complete + return artifacts + + +def test_a_call_restarted_between_its_two_success_writes_completes_the_run() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + artifacts = _first_attempt_preempted_at_completion( + store, report_runs, after_writing=False + ) + restart_invoker = RestartInvoker(artifacts) + + result = _coordinate( + store, + report_runs, + restart_invoker, + context=report_run_context(), + artifacts=artifacts, + ) + + assert result == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": "succeeded", + } + # Nothing is calculated again: the result comes from the stored aggregate. + assert restart_invoker.events == [] + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + rows = {rows.table.name: rows.rows for rows in run.result.rows} + assert {row["variable_name"] for row in rows["budget_summary"]} >= {"household_tax"} + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def test_a_call_restarted_after_its_report_run_succeeded_records_nothing_again() -> ( + None +): + store = FakeStore() + report_runs = FakeReportRuns() + _first_attempt_preempted_at_completion(store, report_runs, after_writing=True) restart_invoker = RestartInvoker(FakeArtifacts()) result = _coordinate( @@ -179,6 +221,8 @@ def preempted_before_the_stage12_success(record): assert restart_invoker.events == [] run = report_runs.report_runs[REPORT_RUN_ID] assert run.writes == ["claim_report_run", "complete_report_run", "claim_report_run"] + # Both records agree: the Stage 12 report succeeded before the report run. + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED def test_a_preempted_comparison_run_behaves_as_before() -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py index 10e7d1687..b677af695 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py @@ -57,7 +57,7 @@ def _run_worker(simulation, store, report_runs, *, context=None, calculator=None ) -def test_worker_records_its_simulation_run_before_its_stage12_record() -> None: +def test_worker_records_its_simulation_run_after_its_stage12_record() -> None: events: list[str] = [] simulation = planned_simulation(SimulationRole.BASELINE) store = _seeded_store(simulation, events) @@ -70,13 +70,14 @@ def test_worker_records_its_simulation_run_before_its_stage12_record() -> None: assert run.modal_call_id == "fc-baseline-worker" assert run.output_uri == result["artifact"]["uri"] assert run.output_sha256 == result["artifact"]["content_sha256"] - # The simulation run succeeds first, so a restart between the two writes - # finds its Stage 12 record still running and records nothing twice. + # The simulation run succeeds last, so a worker Modal restarts between + # the two writes finds its Stage 12 record succeeded and completes the + # simulation run from it, with the same output file. assert events == [ "stage12:running", "claim_simulation_run", - "complete_simulation_run", "stage12:succeeded", + "complete_simulation_run", ] @@ -118,12 +119,49 @@ def test_worker_refuses_a_simulation_run_that_was_already_failed() -> None: assert report_runs.simulation_runs[BASELINE_RUN_ID].output_uri is None -def test_restarted_worker_keeps_the_output_its_earlier_attempt_recorded() -> None: +def test_restarted_worker_completes_its_run_from_its_succeeded_stage12_record() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + # The earlier attempt marked its Stage 12 record succeeded, then Modal + # pre-empted it before it completed its simulation run. + store.children[simulation.simulation_execution_id] = store.children[ + simulation.simulation_execution_id + ].model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": "gs://private/earlier-attempt.parquet", + "output_sha256": "e" * 64, + "output_schema_version": 1, + "row_identity_columns": ("household.household_id",), + "row_count": 1, + "row_identity_sha256": "f" * 64, + "completed_at": child_record(simulation).created_at, + } + ) + run = report_runs.simulation_runs[BASELINE_RUN_ID] + run.status, run.modal_call_id = RunStatus.RUNNING, "fc-baseline-worker" + calculated = [] + + result = _run_worker( + simulation, + store, + report_runs, + calculator=lambda value: calculated.append(value) or simulation_frames(), + ) + + assert calculated == [] + assert result["artifact"]["uri"] == "gs://private/earlier-attempt.parquet" + assert run.status is RunStatus.SUCCEEDED + assert run.output_uri == "gs://private/earlier-attempt.parquet" + assert run.output_sha256 == "e" * 64 + + +def test_a_simulation_run_another_report_recorded_keeps_its_output() -> None: simulation = planned_simulation(SimulationRole.BASELINE) store = _seeded_store(simulation) report_runs = FakeReportRuns() - # The earlier attempt recorded its simulation run, then Modal pre-empted - # it before it marked its Stage 12 record succeeded. + # Another report run, sharing this simulation run, already recorded it. store.children[simulation.simulation_execution_id] = store.children[ simulation.simulation_execution_id ].model_copy(update={"status": ComparisonRunLifecycleStatus.RUNNING}) From 98cf80f47f8603f8ce1e8301c792e59949ffa989 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:00:58 +0400 Subject: [PATCH 15/24] Fail a report run its coordinator can no longer finish Once a coordinator holds its report run, an error claiming its Stage 12 report, or completing the run from a stored aggregate, now fails the report run before the error propagates, instead of leaving it running. Modal cancels a call with InputCancellation and never runs it again, so a coordinator whose call is cancelled now records its unfinished children, its Stage 12 report and its report run as failed, with the code report_cancelled; pre-emption, KeyboardInterrupt, still records nothing, because Modal restarts that call. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 46 ++++++++++++++--- .../stage12_runtime/report_claims.py | 49 ++++++++++++++++--- .../test_stage12_report_run_coordinator.py | 46 +++++++++++++++++ .../test_stage12_report_run_preemption.py | 37 ++++++++++++++ 4 files changed, 164 insertions(+), 14 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 449a6ef60..6c365e0c9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -24,9 +24,15 @@ SIMULATION_WAIT_TIMEOUT_SECONDS, ) -from .coordinated_simulations import SIMULATION_CANCELLED_ERROR_CODE, logger +from .coordinated_simulations import ( + SIMULATION_CANCELLED_ERROR_CODE, + cancel_pending_simulations, + cancellation, + logger, +) from .dependencies import ModalChildInvoker, artifact_store, runtime_store from .failure_causes import StageProgress +from .failures import Stage12ExecutionError, Stage12FailureDetail from .parent_claim import ( validate_dispatch_context, validate_submitted_parent, @@ -128,11 +134,39 @@ def coordinate_report( coordination.fail(parent, error) msg = "Stage 12 report coordination failed" raise RuntimeError(msg) from None - except BaseException: - # Modal pre-empts a call with KeyboardInterrupt and cancels one with - # InputCancellation, neither an Exception, and runs a pre-empted call - # again under the same call ID. A report run's coordinator stops its - # unfinished children and records nothing, so that run starts again. + except KeyboardInterrupt: + # Modal pre-empts a call with SIGINT and runs it again under the same + # call ID. A report run's coordinator stops its unfinished children + # and records nothing, so the restarted call starts them again. if report_run is not None: coordination.stop_children() raise + except BaseException: + # Modal cancels a call with InputCancellation, which is not an + # Exception, and never runs it again: a report run's coordinator + # records its report, report run and unfinished children as failed. + if report_run is not None: + _record_cancellation(coordination, parent) + raise + + +def _record_cancellation( + coordination: ReportCoordination, parent: ComparisonReportRecord +) -> None: + """Stop and fail the unfinished children, then fail the report and its run.""" + cancel_pending_simulations( + calls=coordination.calls, + simulations=coordination.simulations, + roles=set(coordination.calls) - set(coordination.descriptors), + store=coordination.persistence, + detail=cancellation("the coordinator's call was cancelled"), + ) + coordination.fail( + parent, + Stage12ExecutionError( + Stage12FailureDetail( + error_code="report_cancelled", + error_summary="Modal cancelled the report's coordinator call.", + ) + ), + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py index b97f34bf5..f26683ec7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py @@ -8,6 +8,7 @@ from __future__ import annotations +from functools import partial from hashlib import sha256 from typing import TYPE_CHECKING, Any @@ -17,6 +18,7 @@ ) from policyengine_stage12_persistence import RunStatus +from .failure_recording import persist_failure_record from .failures import Stage12FailureDetail from .parent_claim import claim_parent @@ -39,24 +41,55 @@ def claim_report( A report run another call holds, or one already finished, returns its status. A call Modal restarted after its Stage 12 report succeeded, but before its report run did, completes the report run from the stored - aggregate and returns that it succeeded. + aggregate and returns that it succeeded. Once this call holds the report + run, any error before the report starts fails the run, then propagates. """ report_run = coordination.report_run + if report_run is None: + parent, should_execute = claim_parent( + submitted=submitted, + coordinator_invocation_id=coordination.invocation.invocation_id, + store=coordination.persistence, + ) + return (parent, None) if should_execute else (None, _stage12_summary(parent)) + claimed = report_run.claim() + if not claimed.changed: + return None, _report_run_summary(report_run, claimed.status) + try: + return _claim_held_report(coordination, report_run, submitted=submitted) + except Exception as error: + persist_failure_record( + partial( + report_run.fail, + Stage12FailureDetail( + error_code="report_claim_failed", + error_summary=( + f"Claiming the Stage 12 report failed ({type(error).__name__})." + ), + ), + ), + runtime=coordination.runtime, + failure_scope="report_claim", + ) + raise + + +def _claim_held_report( + coordination: ReportCoordination, + report_run: ReportRunRecorder, + *, + submitted: ComparisonReportRecord, +) -> tuple[ComparisonReportRecord | None, dict[str, Any] | None]: + """Claim the Stage 12 report of a report run this call holds.""" invocation_id = coordination.invocation.invocation_id - if report_run is not None: - claimed = report_run.claim() - if not claimed.changed: - return None, _report_run_summary(report_run, claimed.status) parent, should_execute = claim_parent( submitted=submitted, coordinator_invocation_id=invocation_id, store=coordination.persistence, - resume_own_claim=report_run is not None, + resume_own_claim=True, ) if should_execute: return parent, None - if report_run is None: - return None, _stage12_summary(parent) if ( parent.status is ComparisonRunLifecycleStatus.SUCCEEDED and parent.coordinator_invocation_id == invocation_id diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index c26499e70..6d1cad08a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -215,6 +215,52 @@ def test_a_failed_report_records_its_stage12_failure_before_its_report_run() -> ] == ["stage12:failed", "fail_report_run"] +def test_a_stage12_claim_that_fails_fails_the_report_run() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + def unreachable(_record): + msg = "the Stage 12 database refused the connection" + raise ConnectionError(msg) + + store.create_or_resolve_report = unreachable + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + with pytest.raises(ConnectionError): + _coordinate(store, report_runs, invoker, artifacts) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_claim_failed: ") + assert "refused the connection" not in run.error_message + assert invoker.events == [] + + +def test_a_stage12_report_another_call_finished_fails_the_report_run() -> None: + store = FakeStore() + store.parent = parent_record().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "coordinator_invocation_id": "fc-another-coordinator", + } + ) + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + result = _coordinate(store, report_runs, invoker, artifacts) + + assert result["deduplicated"] is True + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message == ( + "report_already_claimed: This report's Stage 12 record is succeeded " + "under another call." + ) + assert invoker.events == [] + + def test_report_run_coordinator_names_its_aggregate_by_digest(monkeypatch) -> None: built = [] artifacts = FakeArtifacts() diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py index 8199845f3..47d3517ec 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -225,6 +225,43 @@ def test_a_call_restarted_after_its_report_run_succeeded_records_nothing_again() assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED +class InputCancelledForTest(BaseException): + """Stands in for Modal's InputCancellation, a BaseException but not an Exception.""" + + +def _cancel(_seconds): + raise InputCancelledForTest + + +def test_a_cancelled_coordinator_records_its_report_run_and_children_failed() -> None: + # Modal does not run a cancelled call again, so nothing would ever finish + # these runs if the coordinator left them running. + store = FakeStore() + report_runs = FakeReportRuns() + + with pytest.raises(InputCancelledForTest): + _coordinate( + store, + report_runs, + TimeoutInvoker(), + context=report_run_context(), + sleep=_cancel, + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_cancelled: ") + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + simulation_run.status is RunStatus.FAILED + for simulation_run in report_runs.simulation_runs.values() + ) + + def test_a_preempted_comparison_run_behaves_as_before() -> None: store = FakeStore() report_runs = FakeReportRuns() From b467f671cbe1e2ca8e13e76d54795429cbb911bc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:01:32 +0400 Subject: [PATCH 16/24] Describe the report run store and the Modal secret it needs The persistence library's README and the database-access skill now describe the Public API's run tables that the Modal worker role writes, the store's module docstring says that a completion also inserts its result rows, and the secret's module says that the secret must exist in a Modal environment before the v2 application deploys there. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/database-access.md | 5 ++++ .../README.md | 27 ++++++++++++++++--- .../report_run_store.py | 8 +++--- .../stage12_runtime/report_run_access.py | 4 ++- 4 files changed, 36 insertions(+), 8 deletions(-) diff --git a/docs/engineering/skills/database-access.md b/docs/engineering/skills/database-access.md index 1ed83a886..cb6ac2ce8 100644 --- a/docs/engineering/skills/database-access.md +++ b/docs/engineering/skills/database-access.md @@ -34,6 +34,11 @@ It builds on the Stage 12 rules: (`sim.stage12.no-migration-chain`). This repository's code reads and writes their rows, and deployment validation compares their live schema with its SQLAlchemy mappings; nothing here changes them. +- policyengine-api also owns the run tables the v2 Modal coordinator and + workers write as the Modal worker role: `report_runs`, `simulation_runs`, + `report_run_simulation_runs` and the result tables. Their mappings in + `policyengine_stage12_persistence.report_run_tables` and + `report_run_results` declare only the columns those writes touch. - **Once #847 merges**, this repository's own Alembic chain at `migrations/` owns the `simulation` schema, and the Stage 12 tables stay with policyengine-api. diff --git a/libs/policyengine-stage12-persistence/README.md b/libs/policyengine-stage12-persistence/README.md index cca735900..2b0194d61 100644 --- a/libs/policyengine-stage12-persistence/README.md +++ b/libs/policyengine-stage12-persistence/README.md @@ -1,6 +1,25 @@ # PolicyEngine Stage 12 persistence -Temporary SQLAlchemy Core mappings and transactional persistence for the two -Stage 12 comparison tables. `policyengine-api` remains the sole owner of the -SQLModel definitions, Alembic revisions, and live schema. This package performs -DML only and is removed with the temporary comparison path during Stage 14. +SQLAlchemy Core mappings and transactional DML for two sets of tables, both +owned by `policyengine-api`, which keeps the SQLModel definitions, Alembic +revisions and live schema. This package never changes a schema. + +- **The Stage 12 comparison tables** (`Stage12PersistenceStore`): temporary + state of the comparison path, removed with it. +- **The Public API's run tables** (`ReportRunStore`): `report_runs`, + `simulation_runs`, `report_run_simulation_runs` and the result tables. The + v2 Modal coordinator and workers write them as the Modal worker role + (`policyengine_v2_modal_worker`), for a call that names a report run. The + mappings declare only the columns those writes touch. Every status write is + one `UPDATE` conditional on the run's status; a report run's completion + inserts its result rows in the same transaction. + +## Tests + +The report run tests need a disposable Postgres migrated by the Public API's +v2 chain, with its declared privileges applied. They skip unless both URLs +are set: + +- `REPORT_RUN_TEST_DATABASE_URL`: logs in as `policyengine_v2_modal_worker`; +- `REPORT_RUN_TEST_SEED_DATABASE_URL`: logs in as a role that may insert the + rows the simulation service would create. diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index 1fc37d414..62e92d5c2 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -1,8 +1,10 @@ """Conditional writes of report run and simulation run status, as the Modal worker role. -Every write is one ``UPDATE`` conditional on the run's expected status, so a -call that Modal restarts after pre-emption, or a second writer, changes -nothing that is already final. +Every status write is one ``UPDATE`` conditional on the run's expected +status, so a call that Modal restarts after pre-emption, or a second writer, +changes nothing that is already final. A report run's completion also inserts +its result rows, in the same transaction and only when its update changed +the run. """ from __future__ import annotations diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py index 2e3dea9f1..dad36dd64 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py @@ -3,7 +3,9 @@ An operator creates the Modal secret in each Modal environment. The v2 Modal application attaches it to the coordinator and the simulation workers, which record the status and results of the report run and simulation runs their -call names. +call names. The secret must exist in a Modal environment before the v2 +application deploys there: ``modal deploy`` checks its required key and fails +without it. """ from __future__ import annotations From 501b22f021dc8783893242aa4e898f8aed65700d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:20:15 +0400 Subject: [PATCH 17/24] Store only the poverty rates .py's poverty rows can hold In .py's poverty rows, filter_variable names the model variable a group is filtered on (age, sex or race), with its bounds and group label in columns the table doesn't have, so a group's rate has no column that keeps .py's meaning, and the three age groups would collide on the table's unique key as age. The writer now stores only each simulation run's poverty and deep poverty rate for everyone, with a null filter variable, instead of putting group labels in filter_variable. A UK report no longer writes a household_state_income_tax budget row: the UK model has no such variable, and today's output reports 0.0, which a reader gives again when the row is absent. The area mapping is typed. Co-Authored-By: Claude Opus 5.5 --- .../report_run_results.py | 44 +++++++------ .../tests/test_report_run_results.py | 63 +++++++------------ 2 files changed, 50 insertions(+), 57 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py index b6605a25b..7e2314319 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -11,7 +11,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal from uuid import uuid4 from policyengine_simulation_contract.macro_output import SingleYearMacroOutput @@ -26,6 +26,7 @@ from policyengine_simulation_contract.macro_output import ( CongressionalDistrictImpactRecord, ConstituencyImpactRecord, + GeographicImpactOutput, IntraDecileOutput, LocalAuthorityImpactRecord, ) @@ -263,7 +264,7 @@ def report_run_result( rows=( ResultRows( budget_summary, - _budget_rows(economy, report_keys), + _budget_rows(economy, report_keys, country), shared_by_reports=False, ), ResultRows( @@ -380,9 +381,9 @@ def _district_row( def _area_rows( - areas: Any, # noqa: ANN401 # a GeographicImpactOutput, or None + areas: GeographicImpactOutput | None, keys: Mapping[str, Any], - kind: str, + kind: Literal["constituency", "local_authority"], ) -> tuple[dict[str, Any], ...]: if areas is None: return () @@ -405,8 +406,13 @@ def _area_rows( def _budget_rows( - economy: SingleYearMacroOutput, keys: Mapping[str, Any] + economy: SingleYearMacroOutput, keys: Mapping[str, Any], country: str ) -> tuple[dict[str, Any], ...]: + """Return the headline budget figures; state income tax only for the US. + + Today's output reports 0.0 state income tax for a country without it, so a + reader without that row reports 0.0 again. + """ return tuple( { "id": uuid4(), @@ -416,6 +422,7 @@ def _budget_rows( column: getattr(economy.budget, figure), } for figure, variable, column in BUDGET_FIGURES + if variable != "household_state_income_tax" or country == "us" ) @@ -476,28 +483,29 @@ def _intra_decile_type_rows( def _poverty_rows( economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str ) -> tuple[dict[str, Any], ...]: + """Return each simulation run's poverty and deep poverty rate for everyone. + + In .py's rows a filter variable names the model variable a group is + filtered on (age, sex, race), with its bounds, and the table has columns + for neither the group nor the bounds; the age, gender and race rates + therefore have no lossless column and are not written. + """ poverty_types = POVERTY_TYPES[country] - blocks: list[tuple[str, Mapping[str, Any]]] = [ - ("poverty", economy.poverty.poverty.model_dump()), - ("deep_poverty", economy.poverty.deep_poverty.model_dump()), - ("poverty", economy.poverty_by_gender.poverty.model_dump()), - ("deep_poverty", economy.poverty_by_gender.deep_poverty.model_dump()), - ] - if economy.poverty_by_race is not None: - blocks.append(("poverty", economy.poverty_by_race.poverty.model_dump())) + everyone = { + "poverty": economy.poverty.poverty.all, + "deep_poverty": economy.poverty.deep_poverty.all, + } return tuple( { "id": uuid4(), "simulation_run_id": simulation_run_id, "poverty_type": poverty_types[measure], "entity": "person", - # Everyone's rate has no filter; each group's names the group. - "filter_variable": None if group == "all" else group, - "rate": rates[side], + "filter_variable": None, + "rate": getattr(rates, side), } for side, simulation_run_id in _sides(run_ids) - for measure, groups in blocks - for group, rates in groups.items() + for measure, rates in everyone.items() ) diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py index 9af11ec1d..1bc9d7ca4 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -34,7 +34,6 @@ "Gain less than 5%": "gain_less_than_5pct", "Gain more than 5%": "gain_more_than_5pct", } -AGE_GROUPS = ("child", "adult", "senior") @pytest.fixture @@ -93,10 +92,12 @@ def _rebuilt_budget(seeder, report_run_id): } tax = rows["household_tax"]["change"] benefits = rows["household_benefits"]["change"] + # Only the US models state income tax; today's output reports 0.0 elsewhere. + state_tax = rows.get("household_state_income_tax", {"change": 0.0})["change"] return { "budgetary_impact": tax - benefits, "tax_revenue_impact": tax, - "state_tax_revenue_impact": rows["household_state_income_tax"]["change"], + "state_tax_revenue_impact": state_tax, "benefit_spending_impact": benefits, "households": rows["household_weight"]["baseline_total"], "baseline_net_income": rows["household_net_income"]["baseline_total"], @@ -120,16 +121,6 @@ def _rebuilt_intra_decile(rows): } -def _rebuilt_poverty(seeder, simulation_run_id, *, poverty_type, groups): - rows = { - row["filter_variable"]: row["rate"] - for row in seeder.rows( - "poverty", simulation_run_id=simulation_run_id, poverty_type=poverty_type - ) - } - return {group: rows[None if group == "all" else group] for group in groups} - - def test_a_report_run_stores_its_economy_output_as_rows( store, seeder, seeded_run: SeededReportRun ) -> None: @@ -168,31 +159,18 @@ def test_a_report_run_stores_its_economy_output_as_rows( ("baseline", run_ids.baseline_simulation_run_id), ("reform", run_ids.reform_simulation_run_id), ): - for poverty_type, key in (("spm", "poverty"), ("spm_deep", "deep_poverty")): - assert _rebuilt_poverty( - seeder, - run_id, - poverty_type=poverty_type, - groups=(*AGE_GROUPS, "all", "male", "female"), - ) == { - **{ - group: output["poverty"][key][group][side] - for group in (*AGE_GROUPS, "all") - }, - **{ - group: output["poverty_by_gender"][key][group][side] - for group in ("male", "female") - }, - } - assert _rebuilt_poverty( - seeder, - run_id, - poverty_type="spm", - groups=("white", "black", "hispanic", "other"), - ) == { - group: value[side] - for group, value in output["poverty_by_race"]["poverty"].items() + # Only everyone's rate has a column with .py's meaning: in .py's rows + # a filter variable names the model variable filtered on (age, sex, + # race), and the table has no column for the group or its bounds. + poverty_rows = seeder.rows("poverty", simulation_run_id=run_id) + assert { + (row["poverty_type"], row["filter_variable"]): row["rate"] + for row in poverty_rows + } == { + ("spm", None): output["poverty"]["poverty"]["all"][side], + ("spm_deep", None): output["poverty"]["deep_poverty"]["all"][side], } + assert {row["entity"] for row in poverty_rows} == {"person"} (inequality,) = seeder.rows("inequality", simulation_run_id=run_id) assert inequality["income_variable"] == "household_net_income" assert inequality["entity"] == "household" @@ -224,6 +202,14 @@ def test_the_uk_output_uses_its_own_poverty_types_and_wealth_shares( ) } assert poverty_types == {"relative_bhc", "absolute_bhc"} + assert "household_state_income_tax" not in { + row["variable_name"] + for row in seeder.rows("budget_summary", report_run_id=run_ids.report_run_id) + } + assert ( + _rebuilt_budget(seeder, run_ids.report_run_id) + == (uk_economy_output()["budget"]) + ) (inequality,) = seeder.rows( "inequality", simulation_run_id=run_ids.reform_simulation_run_id ) @@ -317,10 +303,9 @@ def test_a_shared_simulation_run_keeps_the_results_first_stored_for_it( simulation_run_id=run_ids.baseline_simulation_run_id, poverty_type="spm", ) - assert next(row for row in stored if row["filter_variable"] is None)["rate"] == 0.5 - assert len(stored) == len({row["filter_variable"] for row in stored}) + assert [row["rate"] for row in stored] == [0.5] reform = seeder.rows("poverty", simulation_run_id=run_ids.reform_simulation_run_id) - assert len(reform) == 2 * 6 + 4 + assert len(reform) == 2 def test_the_detailed_budget_districts_and_receipts_are_stored_losslessly( From afb0bf9818e25de6bd39165d7058270b5610cf89 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:22:12 +0400 Subject: [PATCH 18/24] Bring a report's simulation runs in line with its child records A worker writes its Stage 12 child record first and its simulation run second, so a lost second write left the simulation run behind: a failure the worker or coordinator could not mirror stayed running, and so did a success when the worker stopped between the writes and a later call reused the child. Before it completes or fails a report run, the coordinator now gives each of the report's two simulation runs its child record's outcome, with the same conditional writes, so a report run never succeeds beside a simulation run still running. A worker whose second write fails after its child record succeeded logs it and returns its result instead of failing. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/report_coordination.py | 20 ++++- .../stage12_runtime/report_run_recording.py | 38 +++++++++- .../test_stage12_report_run_coordinator.py | 75 +++++++++++++++++++ .../tests/test_stage12_report_run_worker.py | 21 ++++++ 4 files changed, 152 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py index ecaa50767..e48a9c437 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -72,6 +72,7 @@ resolve_report_output_plan, validate_output_frames, ) +from .report_run_recording import SimulationRunRecordingStore from .simulation import descriptor_from_record if TYPE_CHECKING: @@ -639,10 +640,13 @@ def complete_report_run(self, aggregate: Mapping[str, Any]) -> None: """Store the aggregate's result with the report run's success. A report run this call no longer holds is logged and left as it is. - A call that names no report run records nothing. + A call that names no report run records nothing. Both simulation runs + are given their child records' outcomes first, so a succeeded report + run never combines a simulation run still running. """ if self.report_run is None: return + self.reconcile_simulation_runs() written = self.report_run.complete( aggregate["result"], country=self.report.baseline.geography.country ) @@ -654,6 +658,15 @@ def complete_report_run(self, aggregate: Mapping[str, Any]) -> None: written.status.value, ) + def reconcile_simulation_runs(self) -> None: + """Give each simulation run of the report its child record's outcome.""" + if not isinstance(self.persistence, SimulationRunRecordingStore): + return + for simulation in (self.report.baseline, self.report.reform): + self.persistence.reconcile_simulation_run( + simulation.simulation_execution_id + ) + def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: """Record a failed run once, stop its unfinished children, fail the parent. @@ -706,6 +719,11 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: # After the Stage 12 report, as on success: a call Modal restarts # between the two writes finds that report failed and runs it again. if self.report_run is not None: + persist_failure_record( + self.reconcile_simulation_runs, + runtime=self.runtime, + failure_scope="report_coordination", + ) persist_failure_record( partial(self.report_run.fail, detail), runtime=self.runtime, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py index ab6758794..73949842e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -8,6 +8,7 @@ from __future__ import annotations +import logging from dataclasses import dataclass from datetime import UTC, datetime from typing import TYPE_CHECKING @@ -40,6 +41,9 @@ from .report_run_access import ReportRunWriter +logger = logging.getLogger(__name__) + + def us_district_codes() -> DistrictCodes: """Return .py's US state FIPS codes and at-large states, for district rows.""" # .py's US data loads only for a US report's result. @@ -201,7 +205,16 @@ def replace_simulation_if_status( if not replaced: return stored, replaced if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: - self._complete(simulation_run_id, record, written_at) + # The child's result is recorded; if this second write is lost, + # the coordinator completes the run from the child's record. + try: + self._complete(simulation_run_id, record, written_at) + except Exception as error: # noqa: BLE001 # a lost second write is completed again by the coordinator + logger.warning( + "Could not complete simulation run %s (%s)", + simulation_run_id, + type(error).__name__, + ) elif record.status is ComparisonRunLifecycleStatus.RUNNING: claimed = self._runs.claim_simulation_run( simulation_run_id, @@ -220,6 +233,29 @@ def replace_simulation_if_status( ) return stored, replaced + def reconcile_simulation_run(self, simulation_execution_id: UUID) -> None: + """Give a child's simulation run its finished Stage 12 record's outcome. + + A worker or coordinator that lost its second write leaves the run + behind its record; each write here is conditional, so a run that + already has its outcome is left as it is. An unfinished child, or one + never created, changes nothing. + """ + try: + record = self._store.get_simulation(simulation_execution_id) + except LookupError: + return + if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: + self.complete_from_succeeded_child(record) + elif record.status is ComparisonRunLifecycleStatus.FAILED: + self._runs.fail_simulation_run( + self._run_ids.simulation_run_id(record.role), + error_message=run_error_message( + record.error_code, record.error_summary + ), + failed_at=record.completed_at or datetime.now(UTC), + ) + def complete_from_succeeded_child(self, record: ComparisonSimulationRecord) -> None: """Complete a child's simulation run from its succeeded Stage 12 record.""" self._complete( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index 6d1cad08a..3fb968643 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -6,6 +6,7 @@ from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, + SimulationRole, ) from policyengine_simulation_executor.stage12_runtime import ( coordination as coordination_module, @@ -13,6 +14,7 @@ from policyengine_stage12_persistence import RunStatus from stage12_child_call_support import ConcurrentInvoker, InjectedChildFaults from stage12_fake_report_runs import ( + BASELINE_RUN_ID, REFORM_RUN_ID, REPORT_RUN_ID, FakeReportRuns, @@ -21,8 +23,10 @@ ) from stage12_fake_stores import FakeArtifacts, FakeStore from stage12_runtime_support import ( + child_record, coordinate_report, parent_record, + planned_simulation, report_execution_input, ) @@ -261,6 +265,77 @@ def test_a_stage12_report_another_call_finished_fails_the_report_run() -> None: assert invoker.events == [] +class LosingOneFailureWrite(FakeReportRuns): + """Lose the first write that fails a simulation run, as a dropped connection.""" + + def __init__(self): + super().__init__() + self.lost_failure_writes = 0 + + def fail_simulation_run(self, simulation_run_id, *, error_message, failed_at): + if not self.lost_failure_writes: + self.lost_failure_writes += 1 + msg = "the connection dropped" + raise ConnectionError(msg) + return super().fail_simulation_run( + simulation_run_id, error_message=error_message, failed_at=failed_at + ) + + +def test_a_lost_simulation_run_failure_is_written_when_the_report_fails() -> None: + store = FakeStore() + report_runs = LosingOneFailureWrite() + artifacts = FakeArtifacts() + invoker = RecordingInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises(RuntimeError, match="coordination failed"): + _coordinate(store, report_runs, invoker, artifacts) + + assert report_runs.lost_failure_writes == 1 + reform = report_runs.simulation_runs[REFORM_RUN_ID] + assert reform.status is RunStatus.FAILED + assert reform.error_message.startswith("simulation_invocation_failed: ") + + +def test_a_reused_succeeded_child_completes_its_simulation_run() -> None: + # A worker marked its child record succeeded, then stopped before it + # completed its simulation run; a later call reuses that child's output. + store = FakeStore() + artifacts = FakeArtifacts() + baseline = planned_simulation(SimulationRole.BASELINE) + descriptor = artifacts.add_simulation(baseline) + store.children[baseline.simulation_execution_id] = child_record( + baseline + ).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": descriptor.artifact.uri, + "output_sha256": descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": descriptor.row_identity.identifier_columns, + "row_count": descriptor.row_identity.row_count, + "row_identity_sha256": descriptor.row_identity.identity_sha256, + "completed_at": child_record(baseline).created_at, + } + ) + report_runs = FakeReportRuns() + report_runs.simulation_runs[BASELINE_RUN_ID].status = RunStatus.RUNNING + invoker = RecordingInvoker(artifacts) + + _coordinate(store, report_runs, invoker, artifacts) + + assert invoker.events == ["spawn:reform"] + run = report_runs.simulation_runs[BASELINE_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.output_uri == descriptor.artifact.uri + # The simulation runs are recorded before the report run succeeds. + assert report_runs.events.index("complete_simulation_run") < ( + report_runs.events.index("complete_report_run") + ) + + def test_report_run_coordinator_names_its_aggregate_by_digest(monkeypatch) -> None: built = [] artifacts = FakeArtifacts() diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py index b677af695..8bd51d804 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py @@ -179,6 +179,27 @@ def test_a_simulation_run_another_report_recorded_keeps_its_output() -> None: assert child.output_uri == result["artifact"]["uri"] +def test_a_lost_simulation_run_success_still_returns_the_workers_result() -> None: + # The child record already succeeded; the coordinator completes the run + # from it, so a failed second write must not fail the worker. + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + + def dropped(*_args, **_kwargs): + msg = "the connection dropped" + raise ConnectionError(msg) + + report_runs.complete_simulation_run = dropped + + result = _run_worker(simulation, store, report_runs) + + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert child.output_uri == result["artifact"]["uri"] + assert report_runs.simulation_runs[BASELINE_RUN_ID].status is RunStatus.RUNNING + + def test_comparison_worker_records_no_simulation_run() -> None: simulation = planned_simulation(SimulationRole.BASELINE) store = _seeded_store(simulation) From fad49eb3a3ef7bc7c35de7278eb085279b4cab05 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:23:00 +0400 Subject: [PATCH 19/24] Fail a report run whose coordinator is cancelled while claiming The coordinator claims its report run, then its Stage 12 report. A cancellation between the two, or while it completes the run from a stored aggregate on a restart, now fails the report run with the code report_cancelled before it propagates; Modal never runs a cancelled call again. A pre-emption there still records nothing, because Modal runs the call again and it resumes its claim. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/report_claims.py | 31 +++++++++++------ .../test_stage12_report_run_preemption.py | 34 +++++++++++++++++++ 2 files changed, 54 insertions(+), 11 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py index f26683ec7..980f37e3c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py @@ -42,7 +42,8 @@ def claim_report( status. A call Modal restarted after its Stage 12 report succeeded, but before its report run did, completes the report run from the stored aggregate and returns that it succeeded. Once this call holds the report - run, any error before the report starts fails the run, then propagates. + run, any error or cancellation before the report starts fails the run, + then propagates; a pre-emption doesn't, because Modal runs the call again. """ report_run = coordination.report_run if report_run is None: @@ -57,23 +58,31 @@ def claim_report( return None, _report_run_summary(report_run, claimed.status) try: return _claim_held_report(coordination, report_run, submitted=submitted) - except Exception as error: + except KeyboardInterrupt: + # Pre-emption: Modal runs the call again, and it resumes this claim. + raise + except BaseException as error: persist_failure_record( - partial( - report_run.fail, - Stage12FailureDetail( - error_code="report_claim_failed", - error_summary=( - f"Claiming the Stage 12 report failed ({type(error).__name__})." - ), - ), - ), + partial(report_run.fail, _claim_failure(error)), runtime=coordination.runtime, failure_scope="report_claim", ) raise +def _claim_failure(error: BaseException) -> Stage12FailureDetail: + """Describe an error, or Modal's cancellation, that stopped the claim.""" + if not isinstance(error, Exception): + return Stage12FailureDetail( + error_code="report_cancelled", + error_summary="Modal cancelled the report's coordinator call.", + ) + return Stage12FailureDetail( + error_code="report_claim_failed", + error_summary=f"Claiming the Stage 12 report failed ({type(error).__name__}).", + ) + + def _claim_held_report( coordination: ReportCoordination, report_run: ReportRunRecorder, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py index 47d3517ec..13a63a471 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -262,6 +262,40 @@ def test_a_cancelled_coordinator_records_its_report_run_and_children_failed() -> ) +@pytest.mark.parametrize( + ("interruption", "expected_status"), + [ + # Modal runs a pre-empted call again, which resumes the claim. + (KeyboardInterrupt, RunStatus.RUNNING), + # Modal never runs a cancelled call again. + (InputCancelledForTest, RunStatus.FAILED), + ], +) +def test_a_coordinator_stopped_while_claiming_its_report( + interruption, expected_status +) -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + def stopped(_record): + raise interruption + + store.create_or_resolve_report = stopped + + with pytest.raises(interruption): + _coordinate( + store, + report_runs, + RestartInvoker(FakeArtifacts()), + context=report_run_context(), + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is expected_status + if expected_status is RunStatus.FAILED: + assert run.error_message.startswith("report_cancelled: ") + + def test_a_preempted_comparison_run_behaves_as_before() -> None: store = FakeStore() report_runs = FakeReportRuns() From df355023d0d43c0c8b4b91d17c21323e2298730c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:24:27 +0400 Subject: [PATCH 20/24] Test a report run end to end through the real worker code The coordinator test runs both children through run_single_simulation, sharing the coordinator's Stage 12 store and report run store, and checks that the report run and both simulation runs are recorded once, each simulation run under its own worker's call and naming its child record's output file. Co-Authored-By: Claude Opus 5.5 --- .../test_stage12_report_run_coordinator.py | 88 ++++++++++++++++++- 1 file changed, 86 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index 3fb968643..d1916c59e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -2,17 +2,27 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor + import pytest from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_contract.stage12_execution import ( ComparisonRunLifecycleStatus, SimulationRole, ) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationWorkerDependencies, + run_single_simulation, +) from policyengine_simulation_executor.stage12_runtime import ( coordination as coordination_module, ) from policyengine_stage12_persistence import RunStatus -from stage12_child_call_support import ConcurrentInvoker, InjectedChildFaults +from stage12_child_call_support import ( + ConcurrentInvoker, + FutureCall, + InjectedChildFaults, +) from stage12_fake_report_runs import ( BASELINE_RUN_ID, REFORM_RUN_ID, @@ -21,7 +31,7 @@ economy_aggregate, report_run_context, ) -from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames from stage12_runtime_support import ( child_record, coordinate_report, @@ -336,6 +346,80 @@ def test_a_reused_succeeded_child_completes_its_simulation_run() -> None: ) +class WorkerInvoker: + """Run each child through the real worker code, sharing the coordinator's stores.""" + + def __init__(self, store, artifacts, report_runs): + self.store = store + self.artifacts = artifacts + self.report_runs = report_runs + self.executor = ThreadPoolExecutor(max_workers=2) + + def spawn(self, request): + role = request.simulation["role"] + frames = simulation_frames(100.0 if role == "baseline" else 120.0) + future = self.executor.submit( + run_single_simulation, + request.simulation, + request.context, + required_country="us", + modal_call_id=f"fc-{role}-worker", + dependencies=SimulationWorkerDependencies( + store=self.store, + artifacts=self.artifacts, + calculator=lambda _simulation: frames, + report_runs=self.report_runs, + ), + ) + return FutureCall(f"fc-{role}-worker", future) + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + +def test_a_report_run_through_real_workers_records_every_run_once() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + report_runs = FakeReportRuns() + + report = report_execution_input() + # A reform the real worker accepts: an effective period with dates. + report = report.model_copy( + update={ + "reform": report.reform.model_copy( + update={"policy": {"reform": {"2026-01-01.2026-12-31": 1}}} + ) + } + ) + + coordinate_report( + report.model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=report_run_context().modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=store, + artifacts=artifacts, + invoker=WorkerInvoker(store, artifacts, report_runs), + aggregator=economy_aggregate, + report_runs=report_runs, + ) + + assert report_runs.report_runs[REPORT_RUN_ID].status is RunStatus.SUCCEEDED + for role, run_id in (("baseline", BASELINE_RUN_ID), ("reform", REFORM_RUN_ID)): + run = report_runs.simulation_runs[run_id] + (child,) = [child for child in store.children.values() if child.role == role] + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == f"fc-{role}-worker" + assert run.output_uri == child.output_uri + assert run.writes.count("complete_simulation_run") >= 1 + assert report_runs.report_runs[REPORT_RUN_ID].writes == [ + "claim_report_run", + "complete_report_run", + ] + + def test_report_run_coordinator_names_its_aggregate_by_digest(monkeypatch) -> None: built = [] artifacts = FakeArtifacts() From ef4ecef6a5d6f7c2a0f879425e045b14473ca5dd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:47:14 +0400 Subject: [PATCH 21/24] Leave no run running when a restarted or cancelled report stops A coordinator restarted after a pre-emption starts its children again; if it failed, or was cancelled, before it had, the children its earlier attempt claimed stayed running under calls that attempt had cancelled. A failed report run's coordinator now fails each child record of the report still running, which also fails its simulation run. A cancellation during the report run claim's commit, or while the coordinator records a failure, now still fails the report run. A restarted worker whose completion of its simulation run fails logs it and returns its result, as a worker does after a fresh success, since the coordinator completes the run from the child's record. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 35 ++++-- .../stage12_runtime/report_claims.py | 17 ++- .../stage12_runtime/report_coordination.py | 46 ++++++- .../stage12_runtime/report_run_recording.py | 24 +++- .../stage12_runtime/simulation.py | 12 +- .../test_stage12_report_run_preemption.py | 115 +++++++++++++++++- .../tests/test_stage12_report_run_worker.py | 31 +++++ 7 files changed, 252 insertions(+), 28 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 6c365e0c9..fe095ecb9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -49,6 +49,11 @@ from .report_run_access import report_run_store from .report_run_recording import ReportRunRecorder, SimulationRunRecordingStore +CANCELLED_DETAIL = Stage12FailureDetail( + error_code="report_cancelled", + error_summary="Modal cancelled the report's coordinator call.", +) + __all__ = [ "SIMULATION_CANCELLED_ERROR_CODE", "SIMULATION_POLL_BACKOFF_FACTOR", @@ -131,7 +136,7 @@ def coordinate_report( # Persist a redacted cause for any coordinator, child-call, aggregation, or # artifact failure before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 # every failure is persisted before a stable exception returns to Modal - coordination.fail(parent, error) + _fail_unless_cancelled(coordination, parent, error) msg = "Stage 12 report coordination failed" raise RuntimeError(msg) from None except KeyboardInterrupt: @@ -150,6 +155,24 @@ def coordinate_report( raise +def _fail_unless_cancelled( + coordination: ReportCoordination, + parent: ComparisonReportRecord, + error: Exception, +) -> None: + """Record ``error``; if Modal cancels the call meanwhile, still fail the run.""" + try: + coordination.fail(parent, error) + except KeyboardInterrupt: + raise + except BaseException: + if coordination.report_run is not None: + coordination.report_run.fail_quietly( + CANCELLED_DETAIL, runtime=coordination.runtime + ) + raise + + def _record_cancellation( coordination: ReportCoordination, parent: ComparisonReportRecord ) -> None: @@ -161,12 +184,4 @@ def _record_cancellation( store=coordination.persistence, detail=cancellation("the coordinator's call was cancelled"), ) - coordination.fail( - parent, - Stage12ExecutionError( - Stage12FailureDetail( - error_code="report_cancelled", - error_summary="Modal cancelled the report's coordinator call.", - ) - ), - ) + coordination.fail(parent, Stage12ExecutionError(CANCELLED_DETAIL)) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py index 980f37e3c..285dc0a34 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py @@ -8,7 +8,6 @@ from __future__ import annotations -from functools import partial from hashlib import sha256 from typing import TYPE_CHECKING, Any @@ -18,7 +17,6 @@ ) from policyengine_stage12_persistence import RunStatus -from .failure_recording import persist_failure_record from .failures import Stage12FailureDetail from .parent_claim import claim_parent @@ -53,20 +51,19 @@ def claim_report( store=coordination.persistence, ) return (parent, None) if should_execute else (None, _stage12_summary(parent)) - claimed = report_run.claim() - if not claimed.changed: - return None, _report_run_summary(report_run, claimed.status) try: + claimed = report_run.claim() + if not claimed.changed: + return None, _report_run_summary(report_run, claimed.status) return _claim_held_report(coordination, report_run, submitted=submitted) except KeyboardInterrupt: # Pre-emption: Modal runs the call again, and it resumes this claim. raise except BaseException as error: - persist_failure_record( - partial(report_run.fail, _claim_failure(error)), - runtime=coordination.runtime, - failure_scope="report_claim", - ) + # The claim may have committed before the error, as a cancellation + # during its commit; failing a run this call doesn't hold changes + # nothing. + report_run.fail_quietly(_claim_failure(error), runtime=coordination.runtime) raise diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py index e48a9c437..e19053376 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -667,6 +667,39 @@ def reconcile_simulation_runs(self) -> None: simulation.simulation_execution_id ) + def fail_running_children(self, detail: Stage12FailureDetail) -> None: + """Fail every child record of a report run that is still running. + + A restarted call that fails before it starts a child again leaves + that child's record, and so its simulation run, running under the + call its pre-empted attempt cancelled. Each write is conditional on + the record still running, and is mirrored onto the simulation run. + """ + if not isinstance(self.persistence, SimulationRunRecordingStore): + return + failed_at = datetime.now(UTC) + for simulation in (self.report.baseline, self.report.reform): + try: + record = self.persistence.get_simulation( + simulation.simulation_execution_id + ) + except LookupError: + continue + if record.status is not ComparisonRunLifecycleStatus.RUNNING: + continue + self.persistence.replace_simulation_if_status( + record.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": failed_at, + "completed_at": failed_at, + } + ), + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: """Record a failed run once, stop its unfinished children, fail the parent. @@ -675,7 +708,10 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: peers. Any other failure is reported here, once, before any database write. If it occurred after a child started, such as a failed placeholder write or invocation attach, every started, unfinished child - is then stopped so none keeps running for a failed report. + is then stopped so none keeps running for a failed report. A call that + names a report run then fails the report's child records still + running, gives both simulation runs their records' outcomes, and fails + the report run, after the Stage 12 report. """ if isinstance(error, Stage12ExecutionError): detail = error.detail @@ -719,6 +755,14 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: # After the Stage 12 report, as on success: a call Modal restarts # between the two writes finds that report failed and runs it again. if self.report_run is not None: + persist_failure_record( + partial( + self.fail_running_children, + cancellation("the report failed"), + ), + runtime=self.runtime, + failure_scope="report_coordination", + ) persist_failure_record( self.reconcile_simulation_runs, runtime=self.runtime, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py index 73949842e..f5be374b2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -1,9 +1,13 @@ -"""Recording a report call's simulation runs beside its Stage 12 child records. +"""Recording a report call's report run and simulation runs beside its Stage 12 records. A call that names a report run keeps the comparison path's Stage 12 records, -which coordinate its workers. Every claim, success and failure of a child -record is also written to that child's simulation run, as the Modal worker -role, so the workers and the coordinator need no other change to record them. +which coordinate its workers. ``SimulationRunRecordingStore`` mirrors every +conditional claim, success and failure of a child record onto that child's +simulation run, as the Modal worker role; ``ReportRunRecorder`` holds the +coordinator's writes to the report run. A mirrored write can be lost, so +before it completes or fails a report run the coordinator also gives both +simulation runs their child records' outcomes, and a failed report fails the +child records still running. """ from __future__ import annotations @@ -11,6 +15,7 @@ import logging from dataclasses import dataclass from datetime import UTC, datetime +from functools import partial from typing import TYPE_CHECKING from policyengine_simulation_contract.stage12_execution import ( @@ -22,11 +27,14 @@ report_run_result, ) +from .failure_recording import persist_failure_record + if TYPE_CHECKING: from collections.abc import Mapping from typing import Any from uuid import UUID + from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( ComparisonReportPersistenceResult, ComparisonReportRecord, @@ -98,6 +106,14 @@ def complete(self, economy_output: Mapping[str, Any], *, country: str) -> RunWri ), ) + def fail_quietly( + self, detail: Stage12FailureDetail, *, runtime: ObservabilityRuntime | None + ) -> None: + """Fail the report run, logging rather than raising a failed write.""" + persist_failure_record( + partial(self.fail, detail), runtime=runtime, failure_scope="report_run" + ) + def fail(self, detail: Stage12FailureDetail) -> RunWrite: """Mark the report run failed with ``detail`` while this call holds it.""" return self.writer.fail_report_run( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index badfc26e5..bbed1923a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -541,8 +541,16 @@ def run_single_simulation( if recording is not None: # An earlier attempt stopped between the child's and the run's # success; the run's completion is conditional, so this repeats - # nothing that was recorded. - recording.complete_from_succeeded_child(child) + # nothing that was recorded, and if it fails the coordinator + # completes the run from the child's record. + try: + recording.complete_from_succeeded_child(child) + except Exception as error: # noqa: BLE001 # the coordinator completes the run again + logger.warning( + "Could not complete the simulation run of %s (%s)", + simulation.simulation_execution_id, + type(error).__name__, + ) return descriptor_from_record(child, simulation).model_dump(mode="json") running = _claim_simulation(persistence, child) cancellation_requested = _cancellation_check(persistence, simulation) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py index 13a63a471..87c3a028d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -14,7 +14,11 @@ ComparisonRunLifecycleStatus, ) from policyengine_stage12_persistence import RunStatus -from stage12_child_call_support import ConcurrentInvoker, TimeoutInvoker +from stage12_child_call_support import ( + ConcurrentInvoker, + InjectedChildFaults, + TimeoutInvoker, +) from stage12_fake_report_runs import ( REPORT_RUN_ID, FakeReportRuns, @@ -296,6 +300,115 @@ def stopped(_record): assert run.error_message.startswith("report_cancelled: ") +class FailingRestartInvoker(RestartInvoker): + """Restore an earlier attempt's children; refuse to start the baseline again.""" + + def spawn(self, request): + if request.simulation["role"] == "baseline": + msg = "Modal refused the spawn" + raise ConnectionError(msg) + return super().spawn(request) + + +class CancelledWhileRestoringInvoker(RestartInvoker): + """Receive Modal's cancellation while cancelling the earlier attempt's child.""" + + def restore(self, invocation_id): + raise InputCancelledForTest + + +def _claimed_by_the_preempted_attempt(report_runs): + # In production the pre-empted attempt's workers claimed their runs. + for simulation_run in report_runs.simulation_runs.values(): + simulation_run.status = RunStatus.RUNNING + + +@pytest.mark.parametrize( + ("restart_invoker", "raised"), + [ + (FailingRestartInvoker, RuntimeError), + (CancelledWhileRestoringInvoker, InputCancelledForTest), + ], +) +def test_a_restart_that_stops_before_restarting_its_children_fails_them( + restart_invoker, raised +) -> None: + store = FakeStore() + report_runs = FakeReportRuns() + _preempted_first_attempt(store, report_runs, context=report_run_context()) + _claimed_by_the_preempted_attempt(report_runs) + + with pytest.raises(raised): + _coordinate( + store, + report_runs, + restart_invoker(FakeArtifacts()), + context=report_run_context(), + ) + + assert report_runs.report_runs[REPORT_RUN_ID].status is RunStatus.FAILED + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + simulation_run.status is RunStatus.FAILED + for simulation_run in report_runs.simulation_runs.values() + ) + + +def test_a_cancellation_as_the_report_run_claim_commits_fails_the_run() -> None: + report_runs = FakeReportRuns() + claim = report_runs.claim_report_run + + def cancelled_after_commit(*args, **kwargs): + claim(*args, **kwargs) + raise InputCancelledForTest + + report_runs.claim_report_run = cancelled_after_commit + + with pytest.raises(InputCancelledForTest): + _coordinate( + FakeStore(), + report_runs, + RestartInvoker(FakeArtifacts()), + context=report_run_context(), + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_cancelled: ") + + +def test_a_cancellation_while_recording_a_failure_still_fails_the_run() -> None: + store = FakeStore() + replace_report = store.replace_report + + def cancelled_while_failing(record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + raise InputCancelledForTest + return replace_report(record) + + store.replace_report = cancelled_while_failing + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + + with pytest.raises(InputCancelledForTest): + _coordinate( + store, + report_runs, + ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ), + context=report_run_context(), + artifacts=artifacts, + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_cancelled: ") + + def test_a_preempted_comparison_run_behaves_as_before() -> None: store = FakeStore() report_runs = FakeReportRuns() diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py index 8bd51d804..e6b792825 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py @@ -157,6 +157,37 @@ def test_restarted_worker_completes_its_run_from_its_succeeded_stage12_record() assert run.output_sha256 == "e" * 64 +def test_a_restarted_worker_returns_its_result_if_completing_its_run_fails() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + store.children[simulation.simulation_execution_id] = store.children[ + simulation.simulation_execution_id + ].model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": "gs://private/earlier-attempt.parquet", + "output_sha256": "e" * 64, + "output_schema_version": 1, + "row_identity_columns": ("household.household_id",), + "row_count": 1, + "row_identity_sha256": "f" * 64, + "completed_at": child_record(simulation).created_at, + } + ) + report_runs = FakeReportRuns() + report_runs.simulation_runs[BASELINE_RUN_ID].status = RunStatus.RUNNING + + def dropped(*_args, **_kwargs): + msg = "the connection dropped" + raise ConnectionError(msg) + + report_runs.complete_simulation_run = dropped + + result = _run_worker(simulation, store, report_runs) + + assert result["artifact"]["uri"] == "gs://private/earlier-attempt.parquet" + + def test_a_simulation_run_another_report_recorded_keeps_its_output() -> None: simulation = planned_simulation(SimulationRole.BASELINE) store = _seeded_store(simulation) From 2f6b737150f699ecf1504dd0de60b6e1a8f54d2f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 02:45:58 +0400 Subject: [PATCH 22/24] Store results by compared pair, wealth decile and poverty group The Public API's run tables changed under this work (PolicyEngine/policyengine-api#4162): the labour supply response, cliff impact, SPM receipts and data version moved from `report_runs` to a `report_run_comparisons` row per report run and compared pair, and the result tables gained the keys that hold the parts left out before. - A report run's completion inserts its `report_run_comparisons` row with its other result rows, in the same transaction as its success. - UK wealth deciles are stored in `decile_impacts`, with `decile_variable` `household_wealth_decile`; income deciles leave it null. - Poverty rates by age, gender and race are stored with .py's groups: `filter_variable` names the variable that selects a group (`age`, `is_male`, `race`) and `filter_group` the group; both stay null for everyone. Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/database-access.md | 6 +- .../README.md | 13 +- .../report_run_results.py | 108 +++++++++++---- .../report_run_store.py | 10 -- .../report_run_tables.py | 6 +- .../tests/economy_output.py | 4 + .../tests/report_run_database.py | 1 + .../tests/test_report_run_results.py | 129 +++++++++++++----- .../test_stage12_report_run_coordinator.py | 4 +- 9 files changed, 201 insertions(+), 80 deletions(-) diff --git a/docs/engineering/skills/database-access.md b/docs/engineering/skills/database-access.md index cb6ac2ce8..292f29d4a 100644 --- a/docs/engineering/skills/database-access.md +++ b/docs/engineering/skills/database-access.md @@ -36,9 +36,9 @@ It builds on the Stage 12 rules: with its SQLAlchemy mappings; nothing here changes them. - policyengine-api also owns the run tables the v2 Modal coordinator and workers write as the Modal worker role: `report_runs`, `simulation_runs`, - `report_run_simulation_runs` and the result tables. Their mappings in - `policyengine_stage12_persistence.report_run_tables` and - `report_run_results` declare only the columns those writes touch. + `report_run_simulation_runs`, `report_run_comparisons` and the result + tables. Their mappings in `policyengine_stage12_persistence.report_run_tables` + and `report_run_results` declare only the columns those writes touch. - **Once #847 merges**, this repository's own Alembic chain at `migrations/` owns the `simulation` schema, and the Stage 12 tables stay with policyengine-api. diff --git a/libs/policyengine-stage12-persistence/README.md b/libs/policyengine-stage12-persistence/README.md index 2b0194d61..b1caeb034 100644 --- a/libs/policyengine-stage12-persistence/README.md +++ b/libs/policyengine-stage12-persistence/README.md @@ -7,12 +7,13 @@ revisions and live schema. This package never changes a schema. - **The Stage 12 comparison tables** (`Stage12PersistenceStore`): temporary state of the comparison path, removed with it. - **The Public API's run tables** (`ReportRunStore`): `report_runs`, - `simulation_runs`, `report_run_simulation_runs` and the result tables. The - v2 Modal coordinator and workers write them as the Modal worker role - (`policyengine_v2_modal_worker`), for a call that names a report run. The - mappings declare only the columns those writes touch. Every status write is - one `UPDATE` conditional on the run's status; a report run's completion - inserts its result rows in the same transaction. + `simulation_runs`, `report_run_simulation_runs`, the result tables, and + `report_run_comparisons`, which keeps a compared pair's results that no + result table holds. The v2 Modal coordinator and workers write them as the + Modal worker role (`policyengine_v2_modal_worker`), for a call that names a + report run. The mappings declare only the columns those writes touch. Every + status write is one `UPDATE` conditional on the run's status; a report + run's completion inserts its result rows in the same transaction. ## Tests diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py index 7e2314319..279d3ca9c 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -4,8 +4,9 @@ in the format API v1 and v3 read today. Each part with a lossless place in a result table becomes rows: per-report parts keyed by the report run and its two simulation runs, poverty and inequality keyed by each simulation run. The -labour supply response, cliff impact and SPM receipts are stored whole on the -report run, with the data version. +labour supply response, cliff impact and SPM receipts, which no result table +holds, are stored whole with the data version, in one ``report_run_comparisons`` +row for the report run and its two simulation runs. """ from __future__ import annotations @@ -15,7 +16,7 @@ from uuid import uuid4 from policyengine_simulation_contract.macro_output import SingleYearMacroOutput -from sqlalchemy import Boolean, Column, Float, Integer, String, Table +from sqlalchemy import JSON, Boolean, Column, Float, Integer, String, Table from sqlalchemy.dialects.postgresql import ENUM, UUID from policyengine_stage12_persistence.report_run_tables import metadata @@ -24,8 +25,10 @@ from collections.abc import Mapping from policyengine_simulation_contract.macro_output import ( + BaselineReformValue, CongressionalDistrictImpactRecord, ConstituencyImpactRecord, + DecileOutput, GeographicImpactOutput, IntraDecileOutput, LocalAuthorityImpactRecord, @@ -51,6 +54,15 @@ def _simulation_run_keys() -> tuple[Column[Any], ...]: ) +report_run_comparisons = Table( + "report_run_comparisons", + metadata, + *_report_run_keys(), + Column("labor_supply_result", JSON(none_as_null=True)), + Column("cliff_impact_result", JSON(none_as_null=True)), + Column("spm_provenance", JSON(none_as_null=True)), + Column("data_version", String(128)), +) program_statistics = Table( "program_statistics", metadata, @@ -127,6 +139,7 @@ def _disclosure_values() -> tuple[Column[Any], ...]: metadata, *_report_run_keys(), Column("income_variable", String(512), nullable=False), + Column("decile_variable", String(512)), Column("entity", String(128), nullable=False), Column("decile", Integer, nullable=False), Column("quantiles", Integer, nullable=False), @@ -152,6 +165,7 @@ def _disclosure_values() -> tuple[Column[Any], ...]: Column("poverty_type", String(128), nullable=False), Column("entity", String(128), nullable=False), Column("filter_variable", String(512)), + Column("filter_group", String(64)), Column("rate", Float), ) inequality = Table( @@ -180,6 +194,13 @@ def _disclosure_values() -> tuple[Column[Any], ...]: "us": {"poverty": "spm", "deep_poverty": "spm_deep"}, "uk": {"poverty": "relative_bhc", "deep_poverty": "absolute_bhc"}, } +# .py's poverty groups (``policyengine.outputs.poverty``), named as the +# output's parts name them; ``age``, ``is_male`` and ``race`` select them. +AGE_GROUPS = ("child", "adult", "senior") +GENDER_GROUPS = ("male", "female") +RACE_GROUPS = ("white", "black", "hispanic", "other") +# The variable .py groups UK wealth deciles' household net income by. +WEALTH_DECILE_VARIABLE = "household_wealth_decile" INEQUALITY_INCOME_VARIABLES = { "us": "household_net_income", "uk": "equiv_hbai_household_net_income", @@ -231,10 +252,6 @@ class ReportRunResult: """Everything a report run's economy output is stored as.""" rows: tuple[ResultRows, ...] - labor_supply_result: dict[str, Any] | None - cliff_impact_result: dict[str, Any] | None - spm_provenance: dict[str, Any] | None = None - data_version: str | None = None def report_run_result( @@ -260,8 +277,17 @@ def report_run_result( "reform_simulation_run_id": run_ids.reform_simulation_run_id, } dumped = economy.model_dump(mode="json") + comparison = { + "id": uuid4(), + **report_keys, + "labor_supply_result": dumped["labor_supply_response"], + "cliff_impact_result": dumped["cliff_impact"], + "spm_provenance": dumped["spm_provenance"], + "data_version": economy.data_version, + } return ReportRunResult( rows=( + ResultRows(report_run_comparisons, (comparison,), shared_by_reports=False), ResultRows( budget_summary, _budget_rows(economy, report_keys, country), @@ -310,10 +336,6 @@ def report_run_result( shared_by_reports=False, ), ), - labor_supply_result=dumped["labor_supply_response"], - cliff_impact_result=dumped["cliff_impact"], - spm_provenance=dumped["spm_provenance"], - data_version=economy.data_version, ) @@ -429,19 +451,34 @@ def _budget_rows( def _decile_rows( economy: SingleYearMacroOutput, keys: Mapping[str, Any] ) -> tuple[dict[str, Any], ...]: + rows = _decile_type_rows(economy.decile, None, keys) + if economy.wealth_decile is not None: + rows += _decile_type_rows(economy.wealth_decile, WEALTH_DECILE_VARIABLE, keys) + return rows + + +def _decile_type_rows( + deciles: DecileOutput, decile_variable: str | None, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + """Return the change in household net income in each decile. + + Income deciles rank households by that income, so they name no separate + ``decile_variable``; wealth deciles name the variable that ranks them. + """ return tuple( { "id": uuid4(), **keys, "income_variable": "household_net_income", + "decile_variable": decile_variable, "entity": "household", "decile": int(decile), "quantiles": 10, "absolute_change": average, # .py's rows hold a percentage; today's output, a fraction. - "relative_change": economy.decile.relative[decile] * 100, + "relative_change": deciles.relative[decile] * 100, } - for decile, average in economy.decile.average.items() + for decile, average in deciles.average.items() ) @@ -483,32 +520,55 @@ def _intra_decile_type_rows( def _poverty_rows( economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str ) -> tuple[dict[str, Any], ...]: - """Return each simulation run's poverty and deep poverty rate for everyone. + """Return each simulation run's poverty rates, for everyone and by group. - In .py's rows a filter variable names the model variable a group is - filtered on (age, sex, race), with its bounds, and the table has columns - for neither the group nor the bounds; the age, gender and race rates - therefore have no lossless column and are not written. + As in .py's rows, ``filter_variable`` names the variable that selects a + group and ``filter_group`` the group; both are null for everyone. """ poverty_types = POVERTY_TYPES[country] - everyone = { - "poverty": economy.poverty.poverty.all, - "deep_poverty": economy.poverty.deep_poverty.all, - } return tuple( { "id": uuid4(), "simulation_run_id": simulation_run_id, "poverty_type": poverty_types[measure], "entity": "person", - "filter_variable": None, + "filter_variable": filter_variable, + "filter_group": filter_group, "rate": getattr(rates, side), } for side, simulation_run_id in _sides(run_ids) - for measure, rates in everyone.items() + for measure, filter_variable, filter_group, rates in _poverty_rates(economy) ) +def _poverty_rates( + economy: SingleYearMacroOutput, +) -> list[tuple[str, str | None, str | None, BaselineReformValue]]: + """Return each measure's rates with the variable and group that select them. + + Ages and genders have both measures; the output breaks only poverty, not + deep poverty, down by race, and only in the US. + """ + rates: list[tuple[str, str | None, str | None, BaselineReformValue]] = [] + for measure in ("poverty", "deep_poverty"): + by_age = getattr(economy.poverty, measure) + by_gender = getattr(economy.poverty_by_gender, measure) + rates.append((measure, None, None, by_age.all)) + rates += [ + (measure, "age", group, getattr(by_age, group)) for group in AGE_GROUPS + ] + rates += [ + (measure, "is_male", group, getattr(by_gender, group)) + for group in GENDER_GROUPS + ] + if economy.poverty_by_race is not None: + by_race = economy.poverty_by_race.poverty + rates += [ + ("poverty", "race", group, getattr(by_race, group)) for group in RACE_GROUPS + ] + return rates + + def _inequality_rows( economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str ) -> tuple[dict[str, Any], ...]: diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index 62e92d5c2..b30aa0887 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -177,16 +177,6 @@ def complete_report_run( error_message=None, completed_at=completed_at, updated_at=completed_at, - **( - { - "labor_supply_result": result.labor_supply_result, - "cliff_impact_result": result.cliff_impact_result, - "spm_provenance": result.spm_provenance, - "data_version": result.data_version, - } - if result is not None - else {} - ), ) ) return self._write( diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py index 74cd36802..04792d663 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -7,7 +7,7 @@ from __future__ import annotations -from sqlalchemy import JSON, Column, DateTime, MetaData, String, Table +from sqlalchemy import Column, DateTime, MetaData, String, Table from sqlalchemy.dialects.postgresql import ENUM, UUID RUN_STATUS_VALUES = ("pending", "running", "succeeded", "failed") @@ -41,10 +41,6 @@ Column("started_at", DateTime(timezone=True)), Column("completed_at", DateTime(timezone=True)), Column("updated_at", DateTime(timezone=True), nullable=False), - Column("labor_supply_result", JSON(none_as_null=True)), - Column("cliff_impact_result", JSON(none_as_null=True)), - Column("spm_provenance", JSON(none_as_null=True)), - Column("data_version", String(128)), ) simulation_runs = Table( diff --git a/libs/policyengine-stage12-persistence/tests/economy_output.py b/libs/policyengine-stage12-persistence/tests/economy_output.py index 6898079a5..33b342b72 100644 --- a/libs/policyengine-stage12-persistence/tests/economy_output.py +++ b/libs/policyengine-stage12-persistence/tests/economy_output.py @@ -124,6 +124,10 @@ def uk_economy_output() -> dict[str, Any]: output = deepcopy(US_ECONOMY_OUTPUT) output["poverty_by_race"] = None output["budget"]["state_tax_revenue_impact"] = 0.0 + output["wealth_decile"] = { + "relative": {"1": 0.03, "2": 0.04}, + "average": {"1": 30.0, "2": 40.0}, + } output["intra_wealth_decile"] = deepcopy(output["intra_decile"]) wealth = output["intra_wealth_decile"] wealth["deciles"]["No change"] = [0.4, 0.6] diff --git a/libs/policyengine-stage12-persistence/tests/report_run_database.py b/libs/policyengine-stage12-persistence/tests/report_run_database.py index 876b88d66..11c575f44 100644 --- a/libs/policyengine-stage12-persistence/tests/report_run_database.py +++ b/libs/policyengine-stage12-persistence/tests/report_run_database.py @@ -34,6 +34,7 @@ "report_runs", "simulation_runs", "report_run_simulation_runs", + "report_run_comparisons", "budget_summary", "decile_impacts", "intra_decile_impacts", diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py index 1bc9d7ca4..ce3b1ed73 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -34,6 +34,17 @@ "Gain less than 5%": "gain_less_than_5pct", "Gain more than 5%": "gain_more_than_5pct", } +US_POVERTY_TYPES = {"poverty": "spm", "deep_poverty": "spm_deep"} +UK_POVERTY_TYPES = {"poverty": "relative_bhc", "deep_poverty": "absolute_bhc"} +# The output part each of .py's poverty filter variables rebuilds; a row +# with none is everyone's rate, the part's "all". +POVERTY_PARTS = { + None: "poverty", + "age": "poverty", + "is_male": "poverty_by_gender", + "race": "poverty_by_race", +} +POVERTY_PART_NAMES = ("poverty", "poverty_by_gender", "poverty_by_race") @pytest.fixture @@ -121,6 +132,34 @@ def _rebuilt_intra_decile(rows): } +def _rebuilt_poverty(seeder, run_ids, poverty_types): + """Rebuild today's three poverty parts from both simulation runs' rows.""" + measures = { + poverty_type: measure for measure, poverty_type in poverty_types.items() + } + parts = {} + for side, run_id in ( + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + ): + for row in seeder.rows("poverty", simulation_run_id=run_id): + assert row["entity"] == "person" + # A group needs the variable that selects it, and everyone has neither. + assert (row["filter_variable"] is None) == (row["filter_group"] is None) + part = parts.setdefault(POVERTY_PARTS[row["filter_variable"]], {}) + groups = part.setdefault(measures[row["poverty_type"]], {}) + groups.setdefault(row["filter_group"] or "all", {})[side] = row["rate"] + return {name: parts.get(name) for name in POVERTY_PART_NAMES} + + +def _rebuilt_deciles(rows): + """Rebuild today's decile part; .py's relative change is a percentage.""" + return { + "average": {str(row["decile"]): row["absolute_change"] for row in rows}, + "relative": {str(row["decile"]): row["relative_change"] / 100 for row in rows}, + } + + def test_a_report_run_stores_its_economy_output_as_rows( store, seeder, seeded_run: SeededReportRun ) -> None: @@ -145,6 +184,8 @@ def test_a_report_run_stores_its_economy_output_as_rows( [1.0, 2.0] ) assert {row["income_variable"] for row in deciles} == {"household_net_income"} + # Income deciles name no separate variable to group households by. + assert {row["decile_variable"] for row in deciles} == {None} assert all( row["baseline_simulation_run_id"] == run_ids.baseline_simulation_run_id and row["reform_simulation_run_id"] == run_ids.reform_simulation_run_id @@ -155,22 +196,14 @@ def test_a_report_run_stores_its_economy_output_as_rows( # doesn't carry, so none is written. assert sorted(row["decile"] for row in intra) == [1, 2] assert _rebuilt_intra_decile(intra) == output["intra_decile"] + # Everyone's rates, and those by age, gender and race, under .py's groups. + assert _rebuilt_poverty(seeder, run_ids, US_POVERTY_TYPES) == { + part: output[part] for part in POVERTY_PART_NAMES + } for side, run_id in ( ("baseline", run_ids.baseline_simulation_run_id), ("reform", run_ids.reform_simulation_run_id), ): - # Only everyone's rate has a column with .py's meaning: in .py's rows - # a filter variable names the model variable filtered on (age, sex, - # race), and the table has no column for the group or its bounds. - poverty_rows = seeder.rows("poverty", simulation_run_id=run_id) - assert { - (row["poverty_type"], row["filter_variable"]): row["rate"] - for row in poverty_rows - } == { - ("spm", None): output["poverty"]["poverty"]["all"][side], - ("spm_deep", None): output["poverty"]["deep_poverty"]["all"][side], - } - assert {row["entity"] for row in poverty_rows} == {"person"} (inequality,) = seeder.rows("inequality", simulation_run_id=run_id) assert inequality["income_variable"] == "household_net_income" assert inequality["entity"] == "household" @@ -183,25 +216,46 @@ def test_a_report_run_stores_its_economy_output_as_rows( output["inequality"]["top_10_pct_share"][side], output["inequality"]["top_1_pct_share"][side], ) - report_run = seeder.row("report_runs", run_ids.report_run_id) - assert report_run["labor_supply_result"] == output["labor_supply_response"] - assert report_run["cliff_impact_result"] == output["cliff_impact"] + # The parts with no result table are kept whole for the compared pair. + (comparison,) = seeder.rows( + "report_run_comparisons", report_run_id=run_ids.report_run_id + ) + assert ( + comparison["baseline_simulation_run_id"], + comparison["reform_simulation_run_id"], + ) == (run_ids.baseline_simulation_run_id, run_ids.reform_simulation_run_id) + assert comparison["labor_supply_result"] == output["labor_supply_response"] + assert comparison["cliff_impact_result"] == output["cliff_impact"] -def test_the_uk_output_uses_its_own_poverty_types_and_wealth_shares( +def test_the_uk_output_uses_its_own_poverty_types_and_wealth_deciles( store, seeder, seeded_run: SeededReportRun ) -> None: run_ids = seeded_run.run_ids + output = uk_economy_output() - _complete(store, seeded_run, output=uk_economy_output(), country="uk") + _complete(store, seeded_run, output=output, country="uk") - poverty_types = { - row["poverty_type"] - for row in seeder.rows( - "poverty", simulation_run_id=run_ids.baseline_simulation_run_id - ) + # No race breakdown outside the US. + assert _rebuilt_poverty(seeder, run_ids, UK_POVERTY_TYPES) == { + part: output[part] for part in POVERTY_PART_NAMES } - assert poverty_types == {"relative_bhc", "absolute_bhc"} + deciles = seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id) + # .py groups wealth deciles' household net income by the wealth decile. + wealth_deciles = [ + row for row in deciles if row["decile_variable"] == "household_wealth_decile" + ] + assert {row["income_variable"] for row in wealth_deciles} == { + "household_net_income" + } + assert {row["entity"] for row in wealth_deciles} == {"household"} + rebuilt_wealth = _rebuilt_deciles(wealth_deciles) + assert rebuilt_wealth["average"] == output["wealth_decile"]["average"] + assert rebuilt_wealth["relative"] == pytest.approx( + output["wealth_decile"]["relative"], rel=1e-12 + ) + income_deciles = [row for row in deciles if row["decile_variable"] is None] + assert _rebuilt_deciles(income_deciles)["average"] == output["decile"]["average"] assert "household_state_income_tax" not in { row["variable_name"] for row in seeder.rows("budget_summary", report_run_id=run_ids.report_run_id) @@ -240,6 +294,10 @@ def test_a_repeated_completion_stores_no_second_result( assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) assert len(seeder.rows("budget_summary", report_run_id=run_ids.report_run_id)) == 5 assert len(seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id)) == 2 + assert ( + len(seeder.rows("report_run_comparisons", report_run_id=run_ids.report_run_id)) + == 1 + ) def test_a_result_that_cannot_be_stored_leaves_the_report_run_running( @@ -259,8 +317,6 @@ def test_a_result_that_cannot_be_stored_leaves_the_report_run_running( ) broken = result.__class__( rows=(*[rows for rows in result.rows if rows is not budget], duplicated), - labor_supply_result=result.labor_supply_result, - cliff_impact_result=result.cliff_impact_result, ) with pytest.raises(IntegrityError): @@ -270,8 +326,12 @@ def test_a_result_that_cannot_be_stored_leaves_the_report_run_running( report_run = seeder.row("report_runs", run_ids.report_run_id) assert report_run["status"] == "running" - assert report_run["labor_supply_result"] is None - for table_name in ("decile_impacts", "budget_summary", "intra_decile_impacts"): + for table_name in ( + "report_run_comparisons", + "decile_impacts", + "budget_summary", + "intra_decile_impacts", + ): assert seeder.rows(table_name, report_run_id=run_ids.report_run_id) == [] assert ( seeder.rows("poverty", simulation_run_id=run_ids.baseline_simulation_run_id) @@ -292,6 +352,7 @@ def test_a_shared_simulation_run_keeps_the_results_first_stored_for_it( poverty_type="spm", entity="person", filter_variable=None, + filter_group=None, rate=0.5, ) ) @@ -302,10 +363,16 @@ def test_a_shared_simulation_run_keeps_the_results_first_stored_for_it( "poverty", simulation_run_id=run_ids.baseline_simulation_run_id, poverty_type="spm", + filter_variable=None, ) assert [row["rate"] for row in stored] == [0.5] + # Everyone, three ages and two genders under both poverty types, and four + # races under SPM poverty; the baseline gains all but the stored one. + baseline = seeder.rows( + "poverty", simulation_run_id=run_ids.baseline_simulation_run_id + ) reform = seeder.rows("poverty", simulation_run_id=run_ids.reform_simulation_run_id) - assert len(reform) == 2 + assert len(reform) == len(baseline) == 16 def test_the_detailed_budget_districts_and_receipts_are_stored_losslessly( @@ -355,9 +422,9 @@ def test_the_detailed_budget_districts_and_receipts_are_stored_losslessly( ) # Census GEOIDs: an at-large district is numbered 00. assert {row["district_geoid"] for row in districts} == {101, 107, 200, 612} - report_run = seeder.row("report_runs", report_run_id) - assert report_run["spm_provenance"] == output["spm_provenance"] - assert report_run["data_version"] == output["data_version"] + (comparison,) = seeder.rows("report_run_comparisons", report_run_id=report_run_id) + assert comparison["spm_provenance"] == output["spm_provenance"] + assert comparison["data_version"] == output["data_version"] def test_the_uk_area_impacts_are_stored_with_their_disclosure_fields( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py index d1916c59e..5217f3277 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -120,8 +120,10 @@ def test_coordinator_completes_its_report_run_with_the_economy_output_as_rows() assert budget["household_tax"]["report_run_id"] == REPORT_RUN_ID # The shared report is a US one, so its poverty rows use the SPM types. assert {row["poverty_type"] for row in rows["poverty"]} == {"spm", "spm_deep"} + (comparison,) = rows["report_run_comparisons"] + assert comparison["report_run_id"] == REPORT_RUN_ID assert ( - result.labor_supply_result + comparison["labor_supply_result"] == (CURRENT_SINGLE_YEAR_MACRO_RESULT["labor_supply_response"]) ) # Alabama's first district, by .py's state FIPS codes. From 11a0b7adc95793821fc6542be44d3ec142f84b80 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 02:57:20 +0400 Subject: [PATCH 23/24] Name the poverty rates' fields, and say a 0.0 rate may be uncalculated Each poverty rate's measure, selecting variable, group and values are now a named tuple, with every group listed under the variable that selects it. The rows' docstring says that today's output reports 0.0 for a rate .py didn't report and for a breakdown it failed to calculate. The UK test also checks the income deciles' relative change. Co-Authored-By: Claude Opus 5.5 --- .../report_run_results.py | 78 ++++++++++++------- .../tests/test_report_run_results.py | 10 +-- 2 files changed, 52 insertions(+), 36 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py index 279d3ca9c..5e6e86052 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -12,7 +12,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, NamedTuple from uuid import uuid4 from policyengine_simulation_contract.macro_output import SingleYearMacroOutput @@ -25,13 +25,16 @@ from collections.abc import Mapping from policyengine_simulation_contract.macro_output import ( + AgePovertyOutput, BaselineReformValue, CongressionalDistrictImpactRecord, ConstituencyImpactRecord, DecileOutput, + GenderPovertyOutput, GeographicImpactOutput, IntraDecileOutput, LocalAuthorityImpactRecord, + RacePovertyOutput, ) from policyengine_simulation_contract.stage12_execution import ReportRunIds @@ -194,11 +197,13 @@ def _disclosure_values() -> tuple[Column[Any], ...]: "us": {"poverty": "spm", "deep_poverty": "spm_deep"}, "uk": {"poverty": "relative_bhc", "deep_poverty": "absolute_bhc"}, } -# .py's poverty groups (``policyengine.outputs.poverty``), named as the -# output's parts name them; ``age``, ``is_male`` and ``race`` select them. -AGE_GROUPS = ("child", "adult", "senior") -GENDER_GROUPS = ("male", "female") -RACE_GROUPS = ("white", "black", "hispanic", "other") +# .py's poverty groups (``policyengine.outputs.poverty``), by the variable +# that selects them, named as the output's parts name them. +POVERTY_GROUPS = { + "age": ("child", "adult", "senior"), + "is_male": ("male", "female"), + "race": ("white", "black", "hispanic", "other"), +} # The variable .py groups UK wealth deciles' household net income by. WEALTH_DECILE_VARIABLE = "household_wealth_decile" INEQUALITY_INCOME_VARIABLES = { @@ -523,52 +528,65 @@ def _poverty_rows( """Return each simulation run's poverty rates, for everyone and by group. As in .py's rows, ``filter_variable`` names the variable that selects a - group and ``filter_group`` the group; both are null for everyone. + group and ``filter_group`` the group; both are null for everyone. Today's + output reports 0.0 for a rate .py didn't report and for every rate of a + breakdown it failed to calculate, so a stored 0.0 may be such a rate. """ poverty_types = POVERTY_TYPES[country] return tuple( { "id": uuid4(), "simulation_run_id": simulation_run_id, - "poverty_type": poverty_types[measure], + "poverty_type": poverty_types[rate.measure], "entity": "person", - "filter_variable": filter_variable, - "filter_group": filter_group, - "rate": getattr(rates, side), + "filter_variable": rate.filter_variable, + "filter_group": rate.filter_group, + "rate": getattr(rate.rates, side), } for side, simulation_run_id in _sides(run_ids) - for measure, filter_variable, filter_group, rates in _poverty_rates(economy) + for rate in _poverty_rates(economy) ) -def _poverty_rates( - economy: SingleYearMacroOutput, -) -> list[tuple[str, str | None, str | None, BaselineReformValue]]: - """Return each measure's rates with the variable and group that select them. +class _PovertyRate(NamedTuple): + """One measure's rates for everyone, or for the group a variable selects.""" + + measure: str + filter_variable: str | None + filter_group: str | None + rates: BaselineReformValue + + +def _poverty_rates(economy: SingleYearMacroOutput) -> list[_PovertyRate]: + """Return each measure's rates for everyone and by group. Ages and genders have both measures; the output breaks only poverty, not deep poverty, down by race, and only in the US. """ - rates: list[tuple[str, str | None, str | None, BaselineReformValue]] = [] + rates: list[_PovertyRate] = [] for measure in ("poverty", "deep_poverty"): by_age = getattr(economy.poverty, measure) - by_gender = getattr(economy.poverty_by_gender, measure) - rates.append((measure, None, None, by_age.all)) - rates += [ - (measure, "age", group, getattr(by_age, group)) for group in AGE_GROUPS - ] - rates += [ - (measure, "is_male", group, getattr(by_gender, group)) - for group in GENDER_GROUPS - ] + rates.append(_PovertyRate(measure, None, None, by_age.all)) + rates += _group_rates(measure, "age", by_age) + rates += _group_rates( + measure, "is_male", getattr(economy.poverty_by_gender, measure) + ) if economy.poverty_by_race is not None: - by_race = economy.poverty_by_race.poverty - rates += [ - ("poverty", "race", group, getattr(by_race, group)) for group in RACE_GROUPS - ] + rates += _group_rates("poverty", "race", economy.poverty_by_race.poverty) return rates +def _group_rates( + measure: str, + filter_variable: str, + rates_by_group: AgePovertyOutput | GenderPovertyOutput | RacePovertyOutput, +) -> list[_PovertyRate]: + return [ + _PovertyRate(measure, filter_variable, group, getattr(rates_by_group, group)) + for group in POVERTY_GROUPS[filter_variable] + ] + + def _inequality_rows( economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str ) -> tuple[dict[str, Any], ...]: diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py index ce3b1ed73..fa4b1a2f1 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -249,13 +249,11 @@ def test_the_uk_output_uses_its_own_poverty_types_and_wealth_deciles( "household_net_income" } assert {row["entity"] for row in wealth_deciles} == {"household"} - rebuilt_wealth = _rebuilt_deciles(wealth_deciles) - assert rebuilt_wealth["average"] == output["wealth_decile"]["average"] - assert rebuilt_wealth["relative"] == pytest.approx( - output["wealth_decile"]["relative"], rel=1e-12 - ) income_deciles = [row for row in deciles if row["decile_variable"] is None] - assert _rebuilt_deciles(income_deciles)["average"] == output["decile"]["average"] + for rows, part in ((wealth_deciles, "wealth_decile"), (income_deciles, "decile")): + rebuilt = _rebuilt_deciles(rows) + assert rebuilt["average"] == output[part]["average"] + assert rebuilt["relative"] == pytest.approx(output[part]["relative"], rel=1e-12) assert "household_state_income_tax" not in { row["variable_name"] for row in seeder.rows("budget_summary", report_run_id=run_ids.report_run_id) From 1899837153594077687d54d481f15a754eff39d0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 14:44:05 +0400 Subject: [PATCH 24/24] Give the Modal app tests' US and UK workers a Modal call ID The US and UK single-simulation functions in the Stage 12 Modal app now pass `modal.current_function_call_id()` to the runtime, so a call that names a report run can record its simulation run under its own call ID. The tests load the app against a stand-in `modal` module, which has no `current_function_call_id`, and the `_invoke` helper set one only for the coordinator. The US and UK cases of the four tests that call each Stage 12 function through `_invoke` therefore raised `AttributeError` before reaching the code they test. `_invoke` now gives the stand-in a call ID for the workers too, unless the test already set its own, so the test that checks a worker passes its own call ID still sees its own value. Production code is unchanged. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_modal_app.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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 73ad26255..f81aa8868 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -88,6 +88,15 @@ def model_dump(self, *, mode): observability_context=_observability_context(), ) + # The US and UK workers pass their own Modal call ID to the runtime, so the + # stand-in ``modal`` module needs one; a test that set its own keeps it. + if not hasattr(module.modal, "current_function_call_id"): + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "call-1", + raising=False, + ) monkeypatch.setattr(stage12_runtime, "run_single_simulation", implementation) return getattr(module, function_name)( {"role": "baseline"},