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..636de8a2e --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py @@ -0,0 +1,194 @@ +"""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``. ``requested_aggregates`` must equal the aggregates the + output plan was built for. + """ + + 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`` 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) + 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..4946c4024 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_report_execution.py @@ -0,0 +1,521 @@ +"""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_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( + _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()