From 7e6bb3085b814f5dd50f61785be5e1955ac93866 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:33:18 +0400 Subject: [PATCH 1/2] Add ReportExecutionResult and a result delivery field to the report context Add a typed outcome for a report coordinator that returns its result to the caller instead of writing Stage 12 rows: - ReportExecutionResult carries the aggregate inline, validated as SingleYearMacroOutput, with its canonical JSON SHA-256, the output plan that declares the artifacts' columns, the baseline and reform simulation artifact descriptors, and provenance (bundle, worker version, Modal environment, application, simulation function, version manifest digest and coordinator call ID). Its validator rejects artifacts that do not form the report and a digest that does not match the aggregate. - ReportExecutionFailure and ReportExecutionError carry the redacted error code and summary a Stage 12 row would record. The error pickles with its failure, so Modal can raise it in the caller. - Stage12InvocationContext gains result_delivery. It defaults to stage12_database, so existing callers keep the Stage 12 behavior; a return_to_caller context cannot request a production comparison. The Stage 12 coordinator leaves the new field out of the context it sends to its child simulations, so the child payload keeps exactly the fields it had before and a worker deployed without the field still accepts it. Nothing else reads the field yet. Co-Authored-By: Claude Opus 5.5 --- .../report_execution.py | 187 +++++++ .../stage12_execution.py | 37 +- .../tests/test_report_execution.py | 512 ++++++++++++++++++ .../stage12_runtime/coordination.py | 16 +- .../tests/test_stage12_runtime.py | 36 ++ 5 files changed, 786 insertions(+), 2 deletions(-) create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py create mode 100644 libs/policyengine-simulation-contract/tests/test_report_execution.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py new file mode 100644 index 000000000..2246c1e57 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py @@ -0,0 +1,187 @@ +"""The typed outcome a report coordinator returns to its caller. + +A report coordinator dispatched with ``result_delivery`` set to +``ReportResultDelivery.RETURN_TO_CALLER`` writes no database rows. When the +report succeeds it returns a ``ReportExecutionResult``, which carries the +aggregate inline, references to the baseline and reform simulation artifacts, +and the software, data and Modal deployment that produced them. When the report +fails it raises ``ReportExecutionError``, which carries the redacted error code +and summary that a Stage 12 report row would record. +""" + +from __future__ import annotations + +import json +from datetime import datetime +from hashlib import sha256 +from typing import Annotated, Literal +from uuid import UUID + +from pydantic import Field, field_validator, model_validator + +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + ContractText, + ErrorCode, + ErrorSummary, + ReportAggregate, + Sha256Digest, + SimulationArtifactDescriptor, + SimulationRole, + Stage12OutputPlan, + StrictContractModel, + stage12_output_plan_sha256, +) + +REPORT_EXECUTION_ERROR_MESSAGE = "Report execution failed" + + +def report_aggregate_sha256(aggregate: SingleYearMacroOutput) -> str: + """Return the SHA-256 of the aggregate's canonical JSON encoding. + + The encoding is compact, key-sorted UTF-8 JSON of the aggregate's JSON-mode + dump, so an aggregate has the same digest before and after a JSON round + trip. Raises ``ValueError`` when the aggregate contains NaN or infinity. + """ + try: + payload = json.dumps( + aggregate.model_dump(mode="json"), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + except ValueError: + msg = "aggregate contains a non-finite number" + raise ValueError(msg) from None + return sha256(payload).hexdigest() + + +class ReportExecutionProvenance(StrictContractModel): + """The software, data and Modal deployment that produced one report. + + ``bundle`` names the PolicyEngine.py version, the country model package, the + dataset and the bundle manifest they came from. The other fields name the + Modal environment, application and simulation function that ran, the + digest of the version manifest they were resolved from, and the Modal + function call ID of the report coordinator. + """ + + bundle: BundleProvenance + worker_version: ContractText + modal_environment: ContractText + modal_application: ContractText + simulation_callable: ContractText + version_manifest_sha256: Sha256Digest + coordinator_invocation_id: ContractText + + +class ReportExecutionResult(StrictContractModel): + """A completed report, as its coordinator returns it to the caller. + + ``aggregate`` is the single-year macro result, and ``aggregate_sha256`` the + digest ``report_aggregate_sha256`` gives for it. ``baseline`` and + ``reform`` reference the two simulation artifacts the aggregate was + computed from; each follows ``output_plan``, which declares the columns + the artifacts hold. Both artifacts must belong to this report, carry the + baseline and reform roles, cover the same rows and use the bundle named in + ``provenance``. + """ + + contract_version: Literal[1] = 1 + evaluation_id: UUID + requested_aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] + aggregate_schema_version: Literal[1] = 1 + aggregate: SingleYearMacroOutput + aggregate_sha256: Sha256Digest + output_plan: Stage12OutputPlan + baseline: SimulationArtifactDescriptor + reform: SimulationArtifactDescriptor + provenance: ReportExecutionProvenance + completed_at: datetime + + @field_validator("requested_aggregates") + @classmethod + def require_unique_aggregates( + cls, + value: tuple[ReportAggregate, ...], + ) -> tuple[ReportAggregate, ...]: + """Reject repeated aggregate names.""" + if len(value) != len(set(value)): + msg = "requested aggregates must be unique" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def require_consistent_result(self) -> ReportExecutionResult: + """Check the artifacts form this report and the digest fits the aggregate. + + Also require a timezone-aware ``completed_at``. + """ + if self.completed_at.tzinfo is None: + msg = "completed_at must include a timezone" + raise ValueError(msg) + if self.baseline.role is not SimulationRole.BASELINE: + msg = "the baseline artifact must use the baseline role" + raise ValueError(msg) + if self.reform.role is not SimulationRole.REFORM: + msg = "the reform artifact must use the reform role" + raise ValueError(msg) + expected_plan_sha256 = stage12_output_plan_sha256(self.output_plan) + for descriptor in (self.baseline, self.reform): + if descriptor.evaluation_id != self.evaluation_id: + msg = "simulation artifacts must belong to this report" + raise ValueError(msg) + if descriptor.bundle != self.provenance.bundle: + msg = "simulation artifacts must use the report's bundle" + raise ValueError(msg) + if descriptor.output_plan_sha256 != expected_plan_sha256: + msg = "simulation artifacts must follow the report's output plan" + raise ValueError(msg) + if self.baseline.simulation_execution_id == self.reform.simulation_execution_id: + msg = ( + "baseline and reform artifacts must have different simulation " + "identifiers" + ) + raise ValueError(msg) + if self.baseline.row_identity != self.reform.row_identity: + msg = "baseline and reform artifacts must cover the same rows" + raise ValueError(msg) + if report_aggregate_sha256(self.aggregate) != self.aggregate_sha256: + msg = "aggregate_sha256 does not match the aggregate" + raise ValueError(msg) + return self + + +class ReportExecutionFailure(StrictContractModel): + """A failed report or simulation: its redacted error code and summary. + + The values are the ones a Stage 12 row records for the same failure, which + are safe to store and to show to the caller. + """ + + contract_version: Literal[1] = 1 + evaluation_id: UUID + error_code: ErrorCode + error_summary: ErrorSummary + + +class ReportExecutionError(RuntimeError): + """Raised by a coordinator or simulation whose result goes to its caller. + + The message is fixed; the failure is in ``failure``. The error pickles + with its failure, so ``FunctionCall.get`` raises it in the caller as this + type with the same failure. + """ + + def __init__(self, failure: ReportExecutionFailure) -> None: + """Carry ``failure`` under the fixed message.""" + self.failure = failure + super().__init__(REPORT_EXECUTION_ERROR_MESSAGE) + + def __reduce__( + self, + ) -> tuple[type[ReportExecutionError], tuple[ReportExecutionFailure]]: + """Rebuild the error from its failure when it is unpickled.""" + return (type(self), (self.failure,)) 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 7ef082277..e22d15d7d 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 @@ -587,13 +587,29 @@ def validate_differences(self) -> ResultComparisonArtifactPayload: return self +class ReportResultDelivery(StrEnum): + """How a report coordinator and its simulations deliver their outcome. + + ``STAGE12_DATABASE``, the default, writes the report's and each + simulation's state to the Stage 12 tables in Postgres, and the coordinator + returns an aggregate artifact descriptor. ``RETURN_TO_CALLER`` writes no + database rows: the coordinator returns a ``ReportExecutionResult`` (see + ``report_execution``) or raises ``ReportExecutionError``. + """ + + STAGE12_DATABASE = "stage12_database" + RETURN_TO_CALLER = "return_to_caller" + + 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``. ``result_delivery`` selects how the + coordinator and its simulations deliver their outcome; a returned report + cannot be compared with a production job. """ contract_version: Literal[1] = 1 @@ -615,6 +631,8 @@ class Stage12InvocationContext(StrictContractModel): production_function_call_id: ContractText | None = None created_at: datetime retention_expires_at: datetime + # Callers that omit it keep the Stage 12 behavior. + result_delivery: ReportResultDelivery = ReportResultDelivery.STAGE12_DATABASE @model_validator(mode="after") def validate_retention(self) -> Stage12InvocationContext: @@ -622,6 +640,23 @@ def validate_retention(self) -> Stage12InvocationContext: _validate_retention(self.created_at, self.retention_expires_at) return self + @model_validator(mode="after") + def validate_result_delivery(self) -> Stage12InvocationContext: + """Reject a production comparison for a report returned to its caller. + + The comparison is recorded on the Stage 12 report row, which a + returned report never writes. + """ + if ( + self.result_delivery is ReportResultDelivery.RETURN_TO_CALLER + and self.production_function_call_id is not None + ): + msg = ( + "a report returned to its caller cannot request a production comparison" + ) + 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_report_execution.py b/libs/policyengine-simulation-contract/tests/test_report_execution.py new file mode 100644 index 000000000..c040e5679 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_report_execution.py @@ -0,0 +1,512 @@ +"""Tests for the typed result a report coordinator returns to its caller.""" + +from __future__ import annotations + +import json +import pickle +from datetime import UTC, datetime, timedelta +from hashlib import sha256 +from typing import Any +from uuid import UUID + +import pytest +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, + ReportExecutionProvenance, + ReportExecutionResult, + report_aggregate_sha256, +) +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_comparison_provenance, + build_spm_provenance, +) +from policyengine_simulation_contract.stage12_execution import ( + ArtifactMediaType, + ArtifactReference, + BundleProvenance, + DatasetProvenance, + EntityOutputPlan, + ReportAggregate, + ReportOutputRequirements, + ReportResultDelivery, + RowIdentity, + SimulationArtifactDescriptor, + SimulationRole, + Stage12InvocationContext, + Stage12OutputPlan, + stage12_output_plan_sha256, +) +from pydantic import ValidationError + +NOW = datetime(2026, 10, 20, tzinfo=UTC) +EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") +BASELINE_ID = UUID("00000000-0000-0000-0000-000000000002") +REFORM_ID = UUID("00000000-0000-0000-0000-000000000003") + + +def _bundle() -> BundleProvenance: + return BundleProvenance( + policyengine_version="6.7.5", + country_package_name="policyengine-us", + country_package_version="2.2.1", + dataset=DatasetProvenance( + identity="populace_us_2024", + uri="hf://policyengine/data/populace_us_2024.h5@revision", + artifact_revision="revision", + data_package_name="populace-data", + data_package_version="0.1.0", + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _output_plan() -> Stage12OutputPlan: + return Stage12OutputPlan( + country="us", + requirements=ReportOutputRequirements( + aggregates=tuple(ReportAggregate), + include_cliff_impacts=False, + labor_supply_response_active=False, + ), + entities=( + EntityOutputPlan( + entity="household", + materialized_variables=("household_id", "household_net_income"), + ), + ), + ) + + +def _descriptor(role: SimulationRole) -> SimulationArtifactDescriptor: + return SimulationArtifactDescriptor( + evaluation_id=EVALUATION_ID, + simulation_execution_id=( + BASELINE_ID if role is SimulationRole.BASELINE else REFORM_ID + ), + role=role, + artifact=ArtifactReference( + uri=f"gs://private/stage-12-runs/staging/{role.value}.parquet", + media_type=ArtifactMediaType.PARQUET, + content_sha256=("1" if role is SimulationRole.BASELINE else "2") * 64, + size_bytes=2048, + ), + output_plan_sha256=stage12_output_plan_sha256(_output_plan()), + row_identity=RowIdentity( + identifier_columns=("household.household_id",), + row_count=1, + identity_sha256="c" * 64, + ), + bundle=_bundle(), + ) + + +def _baseline_reform(baseline: float, reform: float) -> dict[str, float]: + return {"baseline": baseline, "reform": reform} + + +def _age_poverty() -> dict[str, dict[str, float]]: + return { + group: _baseline_reform(0.1, 0.09) + for group in ("child", "adult", "senior", "all") + } + + +def _aggregate() -> dict[str, Any]: + """Return a complete single-year macro result with US-only sections set.""" + return { + "spm_provenance": None, + "model_version": "2.2.1", + "data_version": "revision", + "budget": { + "tax_revenue_impact": 100.0, + "state_tax_revenue_impact": 20.0, + "benefit_spending_impact": 30.0, + "budgetary_impact": 70.0, + "households": 2.0, + "baseline_net_income": 1000.0, + }, + "detailed_budget": { + "income_tax": {"baseline": 100.0, "reform": 125.0, "difference": 25.0} + }, + "decile": { + "average": {"1": 10.0, "10": 20.0, "2": 15.0}, + "relative": {"1": 0.01, "10": 0.02, "2": 0.015}, + }, + "inequality": { + "gini": _baseline_reform(0.4, 0.39), + "top_10_pct_share": _baseline_reform(0.3, 0.29), + "top_1_pct_share": _baseline_reform(0.1, 0.09), + }, + "poverty": {"poverty": _age_poverty(), "deep_poverty": _age_poverty()}, + "poverty_by_gender": { + "poverty": { + "male": _baseline_reform(0.08, 0.07), + "female": _baseline_reform(0.09, 0.08), + }, + "deep_poverty": { + "male": _baseline_reform(0.02, 0.02), + "female": _baseline_reform(0.03, 0.02), + }, + }, + "poverty_by_race": { + "poverty": { + race: _baseline_reform(0.06, 0.05) + for race in ("white", "black", "hispanic", "other") + } + }, + "intra_decile": { + "deciles": {"Gain less than 5%": [0.4, 0.2], "No change": [0.6, 0.8]}, + "all": {"Gain less than 5%": 0.3, "No change": 0.7}, + }, + "wealth_decile": None, + "intra_wealth_decile": None, + "labor_supply_response": None, + "constituency_impact": None, + "local_authority_impact": None, + "congressional_district_impact": { + "districts": [ + { + "district": "AL-01", + "average_household_income_change": 10.0, + "relative_household_income_change": 0.01, + "winner_percentage": 0.6, + "loser_percentage": 0.3, + "no_change_percentage": 0.1, + "population": 1000.0, + } + ] + }, + "cliff_impact": None, + } + + +def _spm_comparison() -> dict[str, Any]: + selection = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + receipt = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="a" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "6.7.5", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + } + ), + ) + return build_spm_comparison_provenance( + baseline_receipts=(receipt,), + reform_receipts=(receipt,), + ).model_dump(mode="json", by_alias=True) + + +def _provenance() -> ReportExecutionProvenance: + return ReportExecutionProvenance( + bundle=_bundle(), + worker_version="6.7.5", + modal_environment="staging", + modal_application="policyengine-simulation-v2-py6-7-5", + simulation_callable="run_single_simulation_us", + version_manifest_sha256="d" * 64, + coordinator_invocation_id="fc-coordinator", + ) + + +def _result_payload(**overrides: Any) -> dict[str, Any]: + aggregate = overrides.pop("aggregate", _aggregate()) + payload: dict[str, Any] = { + "evaluation_id": str(EVALUATION_ID), + "requested_aggregates": [item.value for item in ReportAggregate], + "aggregate": aggregate, + "aggregate_sha256": report_aggregate_sha256( + SingleYearMacroOutput.model_validate(aggregate) + ), + "output_plan": _output_plan().model_dump(mode="json"), + "baseline": _descriptor(SimulationRole.BASELINE).model_dump(mode="json"), + "reform": _descriptor(SimulationRole.REFORM).model_dump(mode="json"), + "provenance": _provenance().model_dump(mode="json"), + "completed_at": NOW.isoformat(), + } + payload.update(overrides) + return payload + + +def _context_payload(**overrides: Any) -> dict[str, Any]: + """Return the dispatch context exactly as the Stage 12 entrypoint sends it.""" + payload: dict[str, Any] = { + "request_id": "request-1", + "environment": "staging", + "modal_environment": "staging", + "worker_version": "6.7.5", + "modal_application": "policyengine-simulation-v2-py6-7-5", + "simulation_callable": "run_single_simulation_us", + "version_manifest_sha256": "d" * 64, + "bundle_manifest_sha256": "b" * 64, + "artifact_prefix": f"stage-12-runs/staging/2026/10/{EVALUATION_ID}", + "production_function_call_id": None, + "created_at": NOW.isoformat(), + "retention_expires_at": (NOW + timedelta(days=30)).isoformat(), + } + payload.update(overrides) + return payload + + +def test_existing_dispatch_context_keeps_stage12_database_delivery() -> None: + context = Stage12InvocationContext.model_validate(_context_payload()) + + assert context.result_delivery is ReportResultDelivery.STAGE12_DATABASE + + +def test_dispatch_context_selects_return_to_caller_delivery() -> None: + context = Stage12InvocationContext.model_validate( + _context_payload(result_delivery="return_to_caller") + ) + + assert context.result_delivery is ReportResultDelivery.RETURN_TO_CALLER + + +def test_returned_report_cannot_request_a_production_comparison() -> None: + with pytest.raises(ValidationError, match="production comparison"): + Stage12InvocationContext.model_validate( + _context_payload( + result_delivery="return_to_caller", + production_function_call_id="fc-production", + ) + ) + + +def test_dispatch_context_rejects_an_unknown_delivery() -> None: + with pytest.raises(ValidationError): + Stage12InvocationContext.model_validate( + _context_payload(result_delivery="object_storage") + ) + + +def test_result_round_trips_through_json() -> None: + result = ReportExecutionResult.model_validate(_result_payload()) + + restored = ReportExecutionResult.model_validate( + json.loads(json.dumps(result.model_dump(mode="json"))) + ) + + assert restored == result + assert restored.aggregate.budget.budgetary_impact == 70.0 + assert restored.baseline.role is SimulationRole.BASELINE + assert restored.provenance.bundle == _bundle() + + +def test_result_carries_the_aggregate_as_a_single_year_macro_output() -> None: + result = ReportExecutionResult.model_validate(_result_payload()) + + assert isinstance(result.aggregate, SingleYearMacroOutput) + assert result.aggregate.model_dump(mode="json") == _aggregate() + + +@pytest.mark.parametrize( + "change", + [ + lambda aggregate: aggregate.pop("budget"), + lambda aggregate: aggregate.update({"unexpected_section": {}}), + lambda aggregate: aggregate["budget"].update({"budgetary_impact": "large"}), + ], + ids=["missing-section", "unknown-section", "non-numeric-value"], +) +def test_result_rejects_an_aggregate_outside_the_macro_output_schema(change) -> None: + aggregate = _aggregate() + change(aggregate) + payload = _result_payload() + payload["aggregate"] = aggregate + + with pytest.raises(ValidationError): + ReportExecutionResult.model_validate(payload) + + +def test_result_rejects_an_aggregate_digest_that_does_not_match() -> None: + with pytest.raises(ValidationError, match="aggregate_sha256"): + ReportExecutionResult.model_validate(_result_payload(aggregate_sha256="e" * 64)) + + +def test_result_rejects_a_non_finite_aggregate_value() -> None: + aggregate = _aggregate() + aggregate["budget"]["budgetary_impact"] = float("nan") + payload = _result_payload(aggregate=_aggregate()) + payload["aggregate"] = aggregate + + with pytest.raises(ValidationError, match="finite"): + ReportExecutionResult.model_validate(payload) + + +def test_aggregate_digest_is_the_canonical_json_digest() -> None: + aggregate = SingleYearMacroOutput.model_validate(_aggregate()) + reordered = SingleYearMacroOutput.model_validate( + dict(reversed(list(_aggregate().items()))) + ) + expected = sha256( + json.dumps( + _aggregate(), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + ).hexdigest() + + assert report_aggregate_sha256(aggregate) == expected + assert report_aggregate_sha256(reordered) == expected + + +def test_aggregate_digest_survives_a_json_round_trip_with_spm_provenance() -> None: + aggregate_payload = {**_aggregate(), "spm_provenance": _spm_comparison()} + aggregate = SingleYearMacroOutput.model_validate(aggregate_payload) + restored = SingleYearMacroOutput.model_validate( + json.loads(json.dumps(aggregate.model_dump(mode="json"))) + ) + + assert aggregate.model_dump(mode="json") == aggregate_payload + assert report_aggregate_sha256(restored) == report_aggregate_sha256(aggregate) + ReportExecutionResult.model_validate(_result_payload(aggregate=aggregate_payload)) + + +@pytest.mark.parametrize( + ("field", "replacement", "message"), + [ + ( + "baseline", + lambda: _descriptor(SimulationRole.REFORM).model_copy( + update={"simulation_execution_id": BASELINE_ID} + ), + "baseline artifact must use the baseline role", + ), + ( + "reform", + lambda: _descriptor(SimulationRole.BASELINE).model_copy( + update={"simulation_execution_id": REFORM_ID} + ), + "reform artifact must use the reform role", + ), + ( + "reform", + lambda: _descriptor(SimulationRole.REFORM).model_copy( + update={"evaluation_id": UUID(int=99)} + ), + "must belong to this report", + ), + ( + "reform", + lambda: _descriptor(SimulationRole.REFORM).model_copy( + update={"simulation_execution_id": BASELINE_ID} + ), + "different simulation identifiers", + ), + ( + "reform", + lambda: _descriptor(SimulationRole.REFORM).model_copy( + update={ + "bundle": _bundle().model_copy( + update={"country_package_version": "2.2.2"} + ) + } + ), + "must use the report's bundle", + ), + ( + "baseline", + lambda: _descriptor(SimulationRole.BASELINE).model_copy( + update={"output_plan_sha256": "f" * 64} + ), + "must follow the report's output plan", + ), + ( + "reform", + lambda: _descriptor(SimulationRole.REFORM).model_copy( + update={ + "row_identity": RowIdentity( + identifier_columns=("household.household_id",), + row_count=2, + identity_sha256="c" * 64, + ) + } + ), + "must cover the same rows", + ), + ], + ids=[ + "baseline-role", + "reform-role", + "other-report", + "same-simulation", + "other-bundle", + "other-output-plan", + "other-rows", + ], +) +def test_result_rejects_artifacts_that_do_not_form_this_report( + field, replacement, message +) -> None: + payload = _result_payload(**{field: replacement().model_dump(mode="json")}) + + with pytest.raises(ValidationError, match=message): + ReportExecutionResult.model_validate(payload) + + +def test_result_requires_a_timezone_aware_completion_time() -> None: + with pytest.raises(ValidationError, match="timezone"): + ReportExecutionResult.model_validate( + _result_payload(completed_at="2026-10-20T00:00:00") + ) + + +def test_result_rejects_repeated_requested_aggregates() -> None: + with pytest.raises(ValidationError, match="unique"): + ReportExecutionResult.model_validate( + _result_payload(requested_aggregates=["budget", "budget"]) + ) + + +def test_result_rejects_unknown_fields() -> None: + with pytest.raises(ValidationError): + ReportExecutionResult.model_validate( + _result_payload(aggregate_output_uri="gs://private/aggregate.json") + ) + + +def test_failure_bounds_its_code_and_summary() -> None: + with pytest.raises(ValidationError): + ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="c" * 65, + error_summary="Simulation failed", + ) + with pytest.raises(ValidationError): + ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="simulation_execution_failed", + error_summary="s" * 513, + ) + + +def test_execution_error_keeps_a_fixed_message_and_its_failure_through_pickle() -> None: + failure = ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="simulation_execution_failed", + error_summary="ValueError in dataset_load (baseline): bad input", + ) + + restored = pickle.loads(pickle.dumps(ReportExecutionError(failure))) + + assert isinstance(restored, ReportExecutionError) + assert restored.failure == failure + assert str(restored) == "Report execution failed" 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 49f44d07a..7f81ded06 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 @@ -18,6 +18,7 @@ ComparisonSimulationRecord, PlannedSimulationExecutionInput, ReportExecutionInput, + ReportResultDelivery, ResultComparisonStatus, SimulationArtifactDescriptor, Stage12InvocationContext, @@ -526,6 +527,19 @@ def _require_artifact_digest( raise ValueError(mismatch_message) +def _child_context_payload(context: Stage12InvocationContext) -> dict[str, Any]: + """Return the dispatch context a child simulation receives. + + A Stage 12 child receives the fields it received before + ``result_delivery`` existed, so a worker deployed without that field still + accepts the context. A returned report's child needs the field, so that + it writes no rows either. + """ + if context.result_delivery is ReportResultDelivery.STAGE12_DATABASE: + return context.model_dump(mode="json", exclude={"result_delivery"}) + return context.model_dump(mode="json") + + def coordinate_report( payload: object, context_payload: object, @@ -713,7 +727,7 @@ def coordinate_report( function_name=function_name, environment=context.modal_environment, simulation=simulation.model_dump(mode="json"), - context=context.model_dump(mode="json"), + context=_child_context_payload(context), observability_context=child_observability_context, ) except Exception as error: # noqa: BLE001 # any dispatch failure is recorded on the child and stops its started peers diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 8a0d2644e..405d12663 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1574,6 +1574,42 @@ def test_duplicate_coordinator_submission_does_not_start_duplicate_children() -> assert duplicate_invoker.events == [] +class ContextRecordingInvoker(ConcurrentInvoker): + """Record the dispatch context each child simulation receives.""" + + def __init__(self, artifacts, **options): + super().__init__(artifacts, **options) + self.child_contexts = [] + + def spawn(self, *, context, **values): + self.child_contexts.append(context) + return super().spawn(**values) + + +def test_stage12_children_receive_the_context_without_a_delivery_field() -> None: + artifacts = FakeArtifacts() + invoker = ContextRecordingInvoker(artifacts) + + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=FakeStore(), + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: _aggregate_payload(), + ) + + # Workers deployed before the field existed reject unknown fields. + expected_fields = set(_context().model_dump(mode="json")) - {"result_delivery"} + assert [set(context) for context in invoker.child_contexts] == [ + expected_fields, + expected_fields, + ] + + def test_coordinator_measures_output_and_child_input_planning() -> None: store = FakeStore() artifacts = FakeArtifacts() From c7787766cdeb2107a2148e25f9bb591840a915c8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:08:46 +0400 Subject: [PATCH 2/2] Reject a report result whose requested aggregates differ from its output plan ReportExecutionResult carries the requested aggregates twice: in requested_aggregates and in output_plan.requirements.aggregates, which the coordinator builds from the same report. The validator accepted a result in which the two disagreed, so a caller storing requested_aggregates could keep a value the output plan contradicts. The validator now requires them to be equal, in the same order. Co-Authored-By: Claude Opus 5.5 --- .../report_execution.py | 11 +++++++++-- .../tests/test_report_execution.py | 9 +++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py index 2246c1e57..636de8a2e 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py @@ -86,7 +86,8 @@ class ReportExecutionResult(StrictContractModel): computed from; each follows ``output_plan``, which declares the columns the artifacts hold. Both artifacts must belong to this report, carry the baseline and reform roles, cover the same rows and use the bundle named in - ``provenance``. + ``provenance``. ``requested_aggregates`` must equal the aggregates the + output plan was built for. """ contract_version: Literal[1] = 1 @@ -117,11 +118,17 @@ def require_unique_aggregates( def require_consistent_result(self) -> ReportExecutionResult: """Check the artifacts form this report and the digest fits the aggregate. - Also require a timezone-aware ``completed_at``. + Also require a timezone-aware ``completed_at`` and requested aggregates + that match the output plan's. """ if self.completed_at.tzinfo is None: msg = "completed_at must include a timezone" raise ValueError(msg) + # The coordinator builds the output plan from the report's requested + # aggregates, so the two must be the same values in the same order. + if self.requested_aggregates != self.output_plan.requirements.aggregates: + msg = "requested aggregates must match the output plan's aggregates" + raise ValueError(msg) if self.baseline.role is not SimulationRole.BASELINE: msg = "the baseline artifact must use the baseline role" raise ValueError(msg) diff --git a/libs/policyengine-simulation-contract/tests/test_report_execution.py b/libs/policyengine-simulation-contract/tests/test_report_execution.py index c040e5679..4946c4024 100644 --- a/libs/policyengine-simulation-contract/tests/test_report_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_report_execution.py @@ -476,6 +476,15 @@ def test_result_rejects_repeated_requested_aggregates() -> None: ) +def test_result_rejects_requested_aggregates_that_differ_from_the_output_plan() -> None: + # The output plan records the aggregates the report requested, so a result + # naming other aggregates contradicts its own plan. + with pytest.raises(ValidationError, match="requested aggregates must match"): + ReportExecutionResult.model_validate( + _result_payload(requested_aggregates=["budget"]) + ) + + def test_result_rejects_unknown_fields() -> None: with pytest.raises(ValidationError): ReportExecutionResult.model_validate(