diff --git a/changelog_entry.yaml b/changelog_entry.yaml index c31757f84..913a00554 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -17,3 +17,4 @@ - Raised a clear error, never an omission, when a UK constituency or local-authority breakdown is built over a dataset without area codes - Bound UK constituency and local-authority weight matrices to the certified data release of the selected dataset instead of the unversioned bucket copy, and stopped runs before they start when the matrix lacks the run year or household dimension - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error + - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py index 4241c4082..99c32bfc4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py @@ -413,8 +413,11 @@ def validate_spm_provenance(self) -> AggregateReportArtifactPayload: if "spm_config" in self.result: raise ValueError("aggregate result must not include legacy spm_config") - if "spm_provenance" in self.result: - SPMComparisonProvenance.model_validate(self.result["spm_provenance"]) + # Reports without US SPM calculations serialize this optional field + # as null. Validate a supplied record, not the absence of one. + provenance = self.result.get("spm_provenance") + if provenance is not None: + SPMComparisonProvenance.model_validate(provenance) return self diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py index 7c8a326a0..ad8ca1de1 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_execution.py @@ -212,6 +212,29 @@ def test_simulation_artifact_schema_uses_the_receipt_without_a_wrapper() -> None assert "SPMCalculationProvenance" not in schema_text +@pytest.mark.parametrize("result", [{}, {"spm_provenance": None}]) +def test_aggregate_artifact_accepts_absent_spm_provenance(result) -> None: + artifact = AggregateReportArtifactPayload( + evaluation_id=UUID("00000000-0000-0000-0000-000000000001"), + requested_aggregates=(ReportAggregate.BUDGET,), + bundle=_bundle(), + result=result, + ) + + assert artifact.model_dump(mode="json")["result"] == result + + +@pytest.mark.parametrize("provenance", [{}, "invalid", [], 0]) +def test_aggregate_artifact_rejects_invalid_present_spm_provenance(provenance) -> None: + with pytest.raises(ValueError): + AggregateReportArtifactPayload( + evaluation_id=UUID("00000000-0000-0000-0000-000000000001"), + requested_aggregates=(ReportAggregate.BUDGET,), + bundle=_bundle(), + result={"spm_provenance": provenance}, + ) + + def test_aggregate_artifact_accepts_only_canonical_spm_provenance() -> None: comparison = _spm_comparison() artifact = AggregateReportArtifactPayload( diff --git a/projects/policyengine-simulation-executor/fixtures/fake_object_store.py b/projects/policyengine-simulation-executor/fixtures/fake_object_store.py new file mode 100644 index 000000000..eb29be88b --- /dev/null +++ b/projects/policyengine-simulation-executor/fixtures/fake_object_store.py @@ -0,0 +1,24 @@ +"""Shared in-memory stand-in for ``ArtifactStore``'s byte interface. + +The Stage 12 artifact tests and the Stage 12 runtime tests both drive the real +``Stage12ArtifactStore`` over this fake, so its write-once behaviour has one +definition instead of a per-file copy. +""" + + +class InMemoryObjectStore: + """Keep uploaded objects by path, refusing to replace an existing one.""" + + def __init__(self) -> None: + self.objects: dict[str, bytes] = {} + + def upload_bytes(self, path, payload, *, content_type): + # The real store uploads with ``if_generation_match=0``: an existing + # object is left in place and the call reports that it created nothing. + if path in self.objects: + return False + self.objects[path] = payload + return True + + def read_bytes(self, path): + return self.objects.get(path) diff --git a/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py b/projects/policyengine-simulation-executor/fixtures/macro_output.py similarity index 93% rename from projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py rename to projects/policyengine-simulation-executor/fixtures/macro_output.py index a86d9b830..228a5ec58 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py +++ b/projects/policyengine-simulation-executor/fixtures/macro_output.py @@ -1,4 +1,9 @@ -"""Fixtures for simulation API contract tests.""" +"""A complete current single-year macro result and its key sets. + +The API contract, output builder, standalone app and Stage 12 runtime tests +all need a valid ``SingleYearMacroOutput`` payload. This module is its one +shared home, so no test lane imports another lane's fixture module. +""" CURRENT_SINGLE_YEAR_MACRO_KEYS = { "spm_provenance", diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py b/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py index bd249e401..88539c90d 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_api_contracts.py @@ -11,7 +11,7 @@ SingleYearMacroOutput, ) -from fixtures.test_simulation_api_contracts import ( +from fixtures.macro_output import ( CURRENT_REQUIRED_BUDGET_KEYS, CURRENT_SINGLE_YEAR_MACRO_KEYS, CURRENT_SINGLE_YEAR_MACRO_RESULT, diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 2e200af21..a4de52b4d 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -14,7 +14,7 @@ installed_spm_selection, spm_capable_simulation, ) -from fixtures.test_simulation_api_contracts import ( +from fixtures.macro_output import ( CURRENT_SINGLE_YEAR_MACRO_KEYS, CURRENT_SINGLE_YEAR_MACRO_RESULT, ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py index cacc5c332..5632ea606 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py @@ -24,6 +24,7 @@ ) from pydantic import ValidationError +from fixtures.fake_object_store import InMemoryObjectStore from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, comparison_run_prefix, @@ -193,22 +194,8 @@ def test_private_layout_is_environment_and_date_scoped() -> None: ) -class FakeObjectStore: - def __init__(self): - self.values = {} - - def upload_bytes(self, path, payload, *, content_type): - if path in self.values: - return False - self.values[path] = payload - return True - - def read_bytes(self, path): - return self.values.get(path) - - def test_immutable_retry_accepts_identical_bytes_and_rejects_other_bytes() -> None: - objects = FakeObjectStore() + objects = InMemoryObjectStore() store = Stage12ArtifactStore("private-stage12", store=objects) first = store._write_immutable("path", b"same", content_type="application/json") @@ -220,7 +207,7 @@ def test_immutable_retry_accepts_identical_bytes_and_rejects_other_bytes() -> No def test_aggregate_write_rejects_an_undeclared_payload_shape() -> None: - store = Stage12ArtifactStore("private-stage12", store=FakeObjectStore()) + store = Stage12ArtifactStore("private-stage12", store=InMemoryObjectStore()) with pytest.raises(ValueError, match="evaluation_id"): store.write_aggregate( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index c7c9431a2..cddfa82dd 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -23,7 +23,9 @@ StdoutLogDestination, configure, ) +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput from policyengine_simulation_contract.stage12_execution import ( + AggregateReportArtifactPayload, ArtifactMediaType, ArtifactReference, BundleProvenance, @@ -61,7 +63,10 @@ UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) +from fixtures.fake_object_store import InMemoryObjectStore +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_executor.stage12_artifacts import ( + Stage12ArtifactStore, canonical_json_bytes, serialize_simulation_frames, ) @@ -202,6 +207,21 @@ def _report() -> ReportExecutionInput: ) +def _aggregate_payload(result=None) -> dict: + """Return a stand-in aggregate that satisfies the artifact contract. + + Aggregators that must not run return one too, so an unexpected run shows + up as a written aggregate instead of a rejected one. + """ + report = _report() + return AggregateReportArtifactPayload( + evaluation_id=report.evaluation_id, + requested_aggregates=report.requested_aggregates, + bundle=report.baseline.bundle, + result={} if result is None else result, + ).model_dump(mode="json") + + def coordinate_report(*args, **kwargs): kwargs.setdefault("output_plan_resolver", lambda _: _output_plan()) return coordinate_report_impl(*args, **kwargs) @@ -381,6 +401,9 @@ def __init__(self): self.aggregate_writes = [] self.comparison_writes = [] self.input_writes = [] + # Aggregates go through the real store, so every coordinator test + # enforces the aggregate artifact contract that production enforces. + self.report_store = Stage12ArtifactStore("private", store=InMemoryObjectStore()) def write_input(self, *, prefix, simulation): self.input_writes.append((prefix, simulation.role)) @@ -440,20 +463,14 @@ def write_simulation( ) def read(self, uri): - return self.payloads[uri] + if uri in self.payloads: + return self.payloads[uri] + return self.report_store.read(uri) def write_aggregate(self, *, prefix, payload): - encoded = canonical_json_bytes(payload) - digest = sha256(encoded).hexdigest() - uri = "gs://private/aggregate.json" - self.payloads[uri] = encoded + reference = self.report_store.write_aggregate(prefix=prefix, payload=payload) self.aggregate_writes.append(payload) - return ArtifactReference( - uri=uri, - media_type=ArtifactMediaType.JSON, - content_sha256=digest, - size_bytes=len(encoded), - ) + return reference def write_comparison(self, *, prefix, payload): encoded = canonical_json_bytes(payload) @@ -987,6 +1004,87 @@ def serialize(self): assert value.end_date == datetime.strptime(stop, "%Y-%m-%d"), role +def _uk_report() -> ReportExecutionInput: + uk_bundle = _bundle().model_copy( + update={ + "country_package_name": "policyengine-uk", + "country_package_version": "2.102.3", + } + ) + report_payload = _report().model_dump(mode="json") + for role, policy in ( + ("baseline", {}), + ( + "reform", + { + "gov.hmrc.income_tax.charges.CB_HITC.phase_out_start": { + "2026-01-01.2100-12-31": 70_000 + } + }, + ), + ): + report_payload[role].update( + { + "policy": policy, + "geography": {"country": "uk", "region": "uk"}, + "bundle": uk_bundle.model_dump(mode="json"), + } + ) + return ReportExecutionInput.model_validate(report_payload) + + +def test_uk_aggregate_without_spm_provenance_is_written(monkeypatch) -> None: + from policyengine_simulation_executor import simulation_output_builder + + report = _uk_report() + assert report.baseline.options.get("spm") is None + uk_output = { + **{ + key: value + for key, value in CURRENT_SINGLE_YEAR_MACRO_RESULT.items() + if key != "spm_provenance" + }, + "congressional_district_impact": None, + } + + class UKOutputBuilder: + def __init__(self, **_values): + pass + + def serialize(self): + # The real builder dumps SingleYearMacroOutput, which serializes + # the unset optional SPM receipt as an explicit null. + return SingleYearMacroOutput.model_validate(uk_output).model_dump( + mode="json" + ) + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", UKOutputBuilder + ) + uk_frames = { + entity: pd.DataFrame({f"{entity}_id": [1]}) + for entity in ("person", "benunit", "household") + } + artifacts = FakeArtifacts() + + aggregate = build_aggregate_report( + report=report, + baseline_frames=uk_frames, + reform_frames=uk_frames, + baseline_descriptor=artifacts.add_simulation(report.baseline), + reform_descriptor=artifacts.add_simulation(report.reform), + ) + assert aggregate["result"]["spm_provenance"] is None + + aggregate_reference = artifacts.write_aggregate( + prefix=_context().artifact_prefix, payload=aggregate + ) + + stored_aggregate = artifacts.read(aggregate_reference.uri) + assert aggregate_reference.content_sha256 == sha256(stored_aggregate).hexdigest() + assert json.loads(stored_aggregate)["result"]["spm_provenance"] is None + + def test_single_worker_rejects_combined_baseline_and_reform_input() -> None: payload = _planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") payload["baseline"] = {} @@ -1122,7 +1220,7 @@ def test_coordinator_starts_both_children_before_waiting_and_aggregates() -> Non def aggregate(**values): observed.update(values) - return {"result": "complete"} + return _aggregate_payload() context = _context().model_copy( update={ @@ -1160,7 +1258,9 @@ def resolver(report): latest_start = max(interval[0] for interval in invoker.intervals.values()) earliest_end = min(interval[1] for interval in invoker.intervals.values()) assert latest_start < earliest_end - assert result["artifact"]["uri"] == "gs://private/aggregate.json" + assert result["artifact"]["uri"] == ( + f"gs://private/{context.artifact_prefix}/reports/aggregate.json" + ) assert observed["baseline_frames"]["household"][ "household_net_income" ].tolist() == [100.0] @@ -1189,7 +1289,7 @@ def test_duplicate_coordinator_submission_does_not_start_duplicate_children() -> store=store, artifacts=artifacts, invoker=first_invoker, - aggregator=lambda **_: {"result": "complete"}, + aggregator=lambda **_: _aggregate_payload(), ) duplicate_invoker = ConcurrentInvoker(artifacts) duplicate = coordinate_report( @@ -1201,7 +1301,7 @@ def test_duplicate_coordinator_submission_does_not_start_duplicate_children() -> store=store, artifacts=artifacts, invoker=duplicate_invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) assert duplicate == { @@ -1226,7 +1326,7 @@ def test_coordinator_measures_output_and_child_input_planning() -> None: store=store, artifacts=artifacts, invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: {"result": "complete"}, + aggregator=lambda **_: _aggregate_payload(), runtime=runtime, ) @@ -1292,7 +1392,7 @@ def test_coordinator_planning_operations_preserve_observability_id() -> None: store=FakeStore(), artifacts=artifacts, invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: {"result": "complete"}, + aggregator=lambda **_: _aggregate_payload(), runtime=runtime, ) finally: @@ -1337,7 +1437,7 @@ def test_coordinator_compares_automatic_run_after_successful_aggregation() -> No store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"result": stage12_result}, + aggregator=lambda **_: _aggregate_payload(stage12_result), ) assert invoker.events[-1] == "restore:production" @@ -1372,7 +1472,7 @@ def test_comparison_failure_does_not_change_successful_stage12_report() -> None: store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"result": {"budget": 100}}, + aggregator=lambda **_: _aggregate_payload({"budget": 100}), ) assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED @@ -1400,7 +1500,7 @@ def test_repeated_submission_does_not_retry_failed_comparison() -> None: store=store, artifacts=artifacts, invoker=first_invoker, - aggregator=lambda **_: {"result": {"budget": {"total": 100.0}}}, + aggregator=lambda **_: _aggregate_payload({"budget": {"total": 100.0}}), ) failed_at = store.parent.comparison_completed_at failed_summary = store.parent.comparison_error_summary @@ -1418,7 +1518,7 @@ def test_repeated_submission_does_not_retry_failed_comparison() -> None: store=store, artifacts=artifacts, invoker=second_invoker, - aggregator=lambda **_: {"result": {"budget": {"total": 100.0}}}, + aggregator=lambda **_: _aggregate_payload({"budget": {"total": 100.0}}), ) assert result == { @@ -1451,7 +1551,7 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) assert "sensitive" not in str(error.value) @@ -1489,7 +1589,7 @@ def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( store=store, artifacts=artifacts, invoker=invoker(artifacts), - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) assert artifacts.aggregate_writes == [] @@ -1542,7 +1642,7 @@ def record_store_runtime(runtime): coordinator_invocation_id="coordinator-1", artifacts=artifacts, invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: {"result": "complete"}, + aggregator=lambda **_: _aggregate_payload(), runtime=runtime, ) @@ -1778,7 +1878,7 @@ def get_simulation(self, simulation_execution_id): store=store, artifacts=artifacts, invoker=ConcurrentInvoker(artifacts, fail_role="baseline"), - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), runtime=TelemetryRecordingRuntime(ordered_events), ) @@ -1827,7 +1927,7 @@ def test_coordinator_reports_a_child_dispatch_failure_once() -> None: store=store, artifacts=FakeArtifacts(), invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), runtime=TelemetryRecordingRuntime(ordered_events), ) @@ -1866,7 +1966,7 @@ def test_coordinator_rejects_incompatible_row_identity() -> None: store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) assert artifacts.aggregate_writes == [] @@ -1900,7 +2000,7 @@ def test_coordinator_reuses_matching_successful_child() -> None: store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"result": "complete"}, + aggregator=lambda **_: _aggregate_payload(), ) assert invoker.events == ["spawn:reform"] @@ -1930,7 +2030,7 @@ def test_coordinator_resumes_a_matching_running_child_invocation() -> None: store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"result": "complete"}, + aggregator=lambda **_: _aggregate_payload(), ) assert invoker.events == ["restore:baseline", "spawn:reform"] @@ -1974,7 +2074,7 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) assert "sensitive" not in str(error.value) diff --git a/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py b/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py index bba09b1bf..d3da3dc57 100644 --- a/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py +++ b/projects/policyengine-simulation-executor/tests/test_standalone_simulation_contract.py @@ -5,7 +5,7 @@ from fastapi.testclient import TestClient -from fixtures.test_simulation_api_contracts import CURRENT_SINGLE_YEAR_MACRO_RESULT +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_simulation_executor.main import app