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..a25c4c29a --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py @@ -0,0 +1,226 @@ +"""The typed outcome a report coordinator returns to its caller. + +A report coordinator dispatched with ``result_delivery`` set to +``ReportResultDelivery.RETURN_TO_CALLER`` reads and writes no Stage 12 rows. +When the report succeeds it returns a ``ReportExecutionResult``, which carries +the report's 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 for +the same failure. +""" + +from __future__ import annotations + +import json +from datetime import datetime +from hashlib import sha256 +from typing import TYPE_CHECKING, Annotated, Any, 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, +) + +if TYPE_CHECKING: + from collections.abc import Callable + +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 report's single-year macro result: its totals, inline. + ``aggregate_sha256`` is the digest ``report_aggregate_sha256`` gives for it. + ``baseline`` and ``reform`` reference the two simulation artifacts the + aggregate was computed from, and ``output_plan`` declares the columns they + hold. Both artifacts must belong to this report, carry the baseline and + reform roles, follow the output plan, cover the same rows and use the + bundle named in ``provenance``. ``requested_aggregates`` must be the + aggregates the output plan was built for, in the same order. + """ + + 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 + + @field_validator("completed_at") + @classmethod + def require_aware_completion_time(cls, value: datetime) -> datetime: + """Reject a completion time without a timezone.""" + if value.tzinfo is None: + msg = "completed_at must include a timezone" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def require_aggregates_of_the_output_plan(self) -> ReportExecutionResult: + """Require the requested aggregates the output plan was built for. + + The coordinator builds the output plan from the report's requested + aggregates, so a result naming others contradicts its own plan. + """ + if self.requested_aggregates != self.output_plan.requirements.aggregates: + msg = "requested aggregates must match the output plan's aggregates" + raise ValueError(msg) + return self + + @model_validator(mode="after") + def require_artifacts_of_this_report(self) -> ReportExecutionResult: + """Require a baseline and a reform artifact that form this report.""" + 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) + for descriptor in (self.baseline, self.reform): + self._require_artifact_of_this_report(descriptor) + 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) + return self + + def _require_artifact_of_this_report( + self, + descriptor: SimulationArtifactDescriptor, + ) -> None: + """Require one artifact to belong to this report, bundle and output plan.""" + 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 != stage12_output_plan_sha256( + self.output_plan + ): + msg = "simulation artifacts must follow the report's output plan" + raise ValueError(msg) + + @model_validator(mode="after") + def require_the_aggregate_digest(self) -> ReportExecutionResult: + """Require ``aggregate_sha256`` to be the digest of ``aggregate``.""" + 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 as plain JSON values, so Modal's ``FunctionCall.get`` + raises it in the caller as this type with the same failure, provided the + caller can import this module, whatever pydantic version each side runs. + """ + + 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[Callable[[dict[str, Any]], ReportExecutionError], tuple[dict[str, Any]]]: + """Rebuild the error from its failure's JSON values when it is unpickled.""" + return (_report_execution_error, (self.failure.model_dump(mode="json"),)) + + +def _report_execution_error(failure: dict[str, Any]) -> ReportExecutionError: + """Rebuild a pickled ``ReportExecutionError`` from its failure's JSON values.""" + return ReportExecutionError(ReportExecutionFailure.model_validate(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 cc41fb2a4..1e4c10269 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, records the report's and each + simulation's state in the Stage 12 tables, and the coordinator returns an + aggregate artifact descriptor. ``RETURN_TO_CALLER`` reads and writes no + Stage 12 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 report returned to + its caller cannot be compared with a production job. """ contract_version: Literal[1] = 1 @@ -615,6 +631,9 @@ class Stage12InvocationContext(StrictContractModel): production_function_call_id: ContractText | None = None created_at: datetime retention_expires_at: datetime + # Callers that leave it out, such as the Simulation Entrypoint today, keep + # the Stage 12 database behavior. + result_delivery: ReportResultDelivery = ReportResultDelivery.STAGE12_DATABASE @model_validator(mode="after") def validate_retention(self) -> Stage12InvocationContext: @@ -622,6 +641,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..9e881f5ea --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_report_execution.py @@ -0,0 +1,489 @@ +"""Tests for the typed outcome a report coordinator returns to its caller.""" + +from __future__ import annotations + +import json +import pickle +from datetime import UTC, datetime +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, + RowIdentity, + SimulationArtifactDescriptor, + SimulationRole, + Stage12OutputPlan, + stage12_output_plan_sha256, +) +from pydantic import ValidationError + +COMPLETED_AT = datetime(2026, 10, 22, 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: + is_baseline = role is SimulationRole.BASELINE + return SimulationArtifactDescriptor( + evaluation_id=EVALUATION_ID, + simulation_execution_id=BASELINE_ID if is_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 is_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_and_reform(baseline: float, reform: float) -> dict[str, float]: + return {"baseline": baseline, "reform": reform} + + +def _poverty_by_age() -> dict[str, dict[str, float]]: + return { + group: _baseline_and_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 the 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_and_reform(0.4, 0.39), + "top_10_pct_share": _baseline_and_reform(0.3, 0.29), + "top_1_pct_share": _baseline_and_reform(0.1, 0.09), + }, + "poverty": {"poverty": _poverty_by_age(), "deep_poverty": _poverty_by_age()}, + "poverty_by_gender": { + "poverty": { + "male": _baseline_and_reform(0.08, 0.07), + "female": _baseline_and_reform(0.09, 0.08), + }, + "deep_poverty": { + "male": _baseline_and_reform(0.02, 0.02), + "female": _baseline_and_reform(0.03, 0.02), + }, + }, + "poverty_by_race": { + "poverty": { + race: _baseline_and_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": [ + aggregate_name.value for aggregate_name 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": COMPLETED_AT.isoformat(), + } + payload.update(overrides) + return payload + + +def _failure() -> ReportExecutionFailure: + return ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="simulation_execution_failed", + error_summary="ValueError in dataset_load (baseline): bad input", + ) + + +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_totals_inline_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() + payload["aggregate"] = aggregate + + with pytest.raises(ValidationError, match="non-finite"): + ReportExecutionResult.model_validate(payload) + + +def test_aggregate_digest_is_the_digest_of_its_canonical_json() -> 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-22T00: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 coordinator builds the output plan from 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_requested_aggregates_in_another_order_than_the_plan() -> None: + reordered = [aggregate_name.value for aggregate_name in reversed(ReportAggregate)] + + with pytest.raises(ValidationError, match="requested aggregates must match"): + ReportExecutionResult.model_validate( + _result_payload(requested_aggregates=reordered) + ) + + +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_error_pickles_its_failure_as_plain_json_values() -> None: + # The error crosses Modal as a pickle. Plain JSON values in it mean the + # container's and the caller's pydantic versions need not agree. + pickled = pickle.dumps(ReportExecutionError(_failure())) + + assert b"ReportExecutionFailure" not in pickled + assert b"pydantic" not in pickled + + +def test_error_keeps_a_fixed_message_and_its_failure_through_pickle() -> None: + restored = pickle.loads(pickle.dumps(ReportExecutionError(_failure()))) + + assert isinstance(restored, ReportExecutionError) + assert restored.failure == _failure() + assert str(restored) == "Report execution failed" diff --git a/libs/policyengine-simulation-contract/tests/test_report_execution_modal_transport.py b/libs/policyengine-simulation-contract/tests/test_report_execution_modal_transport.py new file mode 100644 index 000000000..d92611b45 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_report_execution_modal_transport.py @@ -0,0 +1,55 @@ +"""Tests that Modal raises ``ReportExecutionError`` in the caller as the same type. + +They call Modal's own exception transport, which is private to the Modal +client, so they live apart from the other report execution tests: a Modal +upgrade that moves these functions fails only this module. +""" + +from __future__ import annotations + +import asyncio +from uuid import UUID + +import pytest +from modal._serialization import pickle_exception +from modal._utils.function_utils import _process_result +from modal_proto import api_pb2 +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, +) + +EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") + + +def _failure() -> ReportExecutionFailure: + return ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="simulation_execution_failed", + error_summary="ValueError in dataset_load (baseline): bad input", + ) + + +def test_modal_raises_the_error_in_the_caller_as_the_same_type() -> None: + """Send the error through Modal's own exception transport. + + A function container pickles the exception it raised with + ``pickle_exception``, and ``FunctionCall.get`` in the caller unpickles and + raises it in ``_process_result``. Both are Modal's functions, so this + checks the installed Modal client rather than a copy of its behavior. + """ + error = ReportExecutionError(_failure()) + remote_result = api_pb2.GenericResult( + status=api_pb2.GenericResult.GENERIC_STATUS_FAILURE, + data=pickle_exception(error), + exception=repr(error), + ) + + with pytest.raises(ReportExecutionError) as raised: + asyncio.run( + _process_result(remote_result, api_pb2.DATA_FORMAT_PICKLE, None, None) + ) + + assert type(raised.value) is ReportExecutionError + assert raised.value.failure == _failure() + assert str(raised.value) == "Report execution failed" diff --git a/libs/policyengine-simulation-contract/tests/test_report_result_delivery.py b/libs/policyengine-simulation-contract/tests/test_report_result_delivery.py new file mode 100644 index 000000000..bc96cddb7 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_report_result_delivery.py @@ -0,0 +1,79 @@ +"""Tests for how a report coordinator is told to deliver its outcome.""" + +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 ( + ReportResultDelivery, + Stage12InvocationContext, +) +from pydantic import ValidationError + +CREATED_AT = datetime(2026, 10, 22, tzinfo=UTC) +EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") + + +def _entrypoint_context_payload(**overrides: Any) -> dict[str, Any]: + """Return the dispatch context as the Simulation Entrypoint builds it today.""" + 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": CREATED_AT.isoformat(), + "retention_expires_at": (CREATED_AT + timedelta(days=30)).isoformat(), + } + payload.update(overrides) + return payload + + +def test_a_context_without_a_delivery_keeps_the_comparison_path_storage() -> None: + context = Stage12InvocationContext.model_validate(_entrypoint_context_payload()) + + assert context.result_delivery is ReportResultDelivery.STAGE12_DATABASE + + +def test_a_context_can_ask_for_the_report_to_return_to_its_caller() -> None: + context = Stage12InvocationContext.model_validate( + _entrypoint_context_payload(result_delivery="return_to_caller") + ) + + assert context.result_delivery is ReportResultDelivery.RETURN_TO_CALLER + + +def test_a_returned_report_cannot_request_a_production_comparison() -> None: + # The comparison with a production result is recorded on the comparison + # path's report row, which a returned report never writes. + with pytest.raises(ValidationError, match="production comparison"): + Stage12InvocationContext.model_validate( + _entrypoint_context_payload( + result_delivery="return_to_caller", + production_function_call_id="fc-production", + ) + ) + + +def test_a_comparison_path_context_may_still_name_its_production_job() -> None: + context = Stage12InvocationContext.model_validate( + _entrypoint_context_payload(production_function_call_id="fc-production") + ) + + assert context.production_function_call_id == "fc-production" + assert context.result_delivery is ReportResultDelivery.STAGE12_DATABASE + + +def test_a_context_rejects_an_unknown_delivery() -> None: + with pytest.raises(ValidationError): + Stage12InvocationContext.model_validate( + _entrypoint_context_payload(result_delivery="object_storage") + ) diff --git a/projects/policyengine-simulation-executor/fixtures/fake_modal.py b/projects/policyengine-simulation-executor/fixtures/fake_modal.py index bc64d36f5..d203c0d8f 100644 --- a/projects/policyengine-simulation-executor/fixtures/fake_modal.py +++ b/projects/policyengine-simulation-executor/fixtures/fake_modal.py @@ -102,6 +102,16 @@ def from_name(*args: object, **kwargs: object) -> dict[str, object]: return {"args": args, "kwargs": kwargs} +class InputCancellation(BaseException): + """Stand-in for ``modal.exception.InputCancellation``. + + Modal raises it in a running synchronous input that is cancelled. It + subclasses ``BaseException`` so that ``except Exception`` does not catch + it. The project's ``src/modal`` package shadows the installed Modal package + in unit tests, so the real class cannot be imported there. + """ + + class FakeApp: """Stand-in for ``modal.App`` that records ``function`` decorator options.""" diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 92642d673..d71cdcb16 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import TYPE_CHECKING +from policyengine_simulation_contract.stage12_execution import ReportResultDelivery from policyengine_simulation_contract.stage12_manifest import v2_application_name from policyengine_simulation_executor import stage12_runtime from policyengine_simulation_executor.hf_access_validation import ( @@ -300,6 +301,23 @@ def run_single_simulation_segment_us( runtime.shutdown() +def _returns_to_caller(context: dict) -> bool: + """Report whether a dispatch context asks for the report to return to its caller.""" + return context.get("result_delivery") == ReportResultDelivery.RETURN_TO_CALLER.value + + +def _simulation_operation_attributes(payload: dict, context: dict) -> dict[str, object]: + """Return a worker operation's attributes; a returned report's worker says so.""" + attributes: dict[str, object] = { + "runner_name": "stage12", + "simulation_role": payload.get("role"), + } + if _returns_to_caller(context): + # Telemetry tells a returned report's workers apart from Stage 12 runs. + attributes["result_delivery"] = ReportResultDelivery.RETURN_TO_CALLER.value + return attributes + + @app.function( image=us_worker_image, cpu=8.0, @@ -339,10 +357,7 @@ def segmented_calculator( result = None with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), - attributes={ - "runner_name": "stage12", - "simulation_role": payload.get("role"), - }, + attributes=_simulation_operation_attributes(payload, context), remote_context=propagated, ): result = stage12_runtime.run_single_simulation( @@ -389,10 +404,7 @@ def run_single_simulation_uk( result = None with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), - attributes={ - "runner_name": "stage12", - "simulation_role": payload.get("role"), - }, + attributes=_simulation_operation_attributes(payload, context), remote_context=propagated, ): result = stage12_runtime.run_single_simulation( @@ -411,6 +423,24 @@ def run_single_simulation_uk( runtime.shutdown() +def _coordinator_operation_attributes( + context: dict, + *, + returns_to_caller: bool, +) -> dict[str, str]: + """Return the coordinator operation's attributes; a returned report says so.""" + attributes = { + "runner_name": "stage12", + "execution_mode": ( + "shadow" if context.get("production_function_call_id") else "authoritative" + ), + } + if returns_to_caller: + # Telemetry tells returned reports apart from Stage 12 runs. + attributes["result_delivery"] = ReportResultDelivery.RETURN_TO_CALLER.value + return attributes + + @app.function( image=coordinator_image, cpu=2.0, @@ -425,16 +455,24 @@ def run_single_simulation_uk( def coordinate_report( payload: dict, context: dict, - parent: dict, + parent: dict | None = None, *, observability_context: dict | None = None, ) -> dict: - """Run one Stage 12 report: fan out both simulations, aggregate, and persist. + """Run one report: fan out both simulations, aggregate, and deliver the result. A ``production_function_call_id`` in ``context`` marks a shadow run whose result is compared with that production job's. A report another invocation already claimed returns a ``deduplicated`` marker instead. + + A ``result_delivery`` of ``return_to_caller`` in ``context`` takes no + ``parent`` and reads and writes no Stage 12 rows: the report returns a + ``ReportExecutionResult`` payload or raises ``ReportExecutionError``. """ + returns_to_caller = _returns_to_caller(context) + if returns_to_caller and parent is not None: + msg = "a report returned to its caller takes no Stage 12 parent record" + raise ValueError(msg) coordinator_invocation_id = modal.current_function_call_id() if coordinator_invocation_id is None: msg = "Modal coordinator invocation identifier is unavailable" @@ -450,31 +488,36 @@ def coordinate_report( if context.get("production_function_call_id") else STAGE12_CANONICAL_REPORT_STAGES ) + invocation = stage12_runtime.CoordinatorInvocation( + application_name=APP_NAME, + invocation_id=coordinator_invocation_id, + ) + dependencies = stage12_runtime.CoordinatorDependencies(runtime=runtime) try: propagated = normalize_observability_context(observability_context) result = None with runtime.operation( stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION), - attributes={ - "runner_name": "stage12", - "execution_mode": ( - "shadow" - if context.get("production_function_call_id") - else "authoritative" - ), - }, + attributes=_coordinator_operation_attributes( + context, returns_to_caller=returns_to_caller + ), remote_context=propagated, ): - result = stage12_runtime.coordinate_report( - payload, - context, - parent, - invocation=stage12_runtime.CoordinatorInvocation( - application_name=APP_NAME, - invocation_id=coordinator_invocation_id, - ), - dependencies=stage12_runtime.CoordinatorDependencies(runtime=runtime), - ) + if returns_to_caller: + result = stage12_runtime.coordinate_report_for_caller( + payload, + context, + invocation=invocation, + dependencies=dependencies, + ) + else: + result = stage12_runtime.coordinate_report( + payload, + context, + parent, + invocation=invocation, + dependencies=dependencies, + ) if result is None: msg = "Stage 12 report coordinator returned no result" raise RuntimeError(msg) 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..73804c6d0 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 @@ -77,11 +77,41 @@ def input_path(*, prefix: str, role: str) -> str: return f"{prefix}/inputs/{role}.json" +def content_addressed_input_path( + *, + prefix: str, + role: str, + content_sha256: str, +) -> str: + """Return the key of a planned execution input JSON named by its own SHA-256. + + A restart that plans a simulation differently, for example on a + redeployed application, stores its input under a new key; an identical + input reuses the earlier object. + """ + return f"{prefix}/inputs/{role}-{content_sha256}.json" + + 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 content_addressed_simulation_path( + *, + prefix: str, + role: str, + content_sha256: str, +) -> str: + """Return the key of a Parquet output named by its own SHA-256. + + A restarted simulation whose output differs in any byte stores it under a + new key instead of colliding with the earlier attempt's object, and an + identical output reuses the earlier object. + """ + return f"{prefix}/simulations/{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" @@ -322,6 +352,28 @@ def write_input( content_type=REPORT_MEDIA_TYPE, ) + def write_input_named_by_digest( + self, + *, + prefix: str, + simulation: PlannedSimulationExecutionInput, + ) -> ArtifactReference: + """Store a simulation's planned execution input under a key naming its SHA-256. + + A restarted simulation can always store its input this way (see + ``content_addressed_input_path``). + """ + payload = canonical_json_bytes(simulation.model_dump(mode="json")) + return self._write_immutable( + content_addressed_input_path( + prefix=prefix, + role=simulation.role.value, + content_sha256=sha256(payload).hexdigest(), + ), + payload, + content_type=REPORT_MEDIA_TYPE, + ) + def write_simulation( self, *, @@ -331,14 +383,66 @@ def write_simulation( calculation_provenance: SPMProvenance | None = None, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> SimulationArtifactDescriptor: - """Store a simulation's output frames and describe the artifact for reports.""" + """Store a simulation's output frames under its role's fixed key. + + Returns the artifact's descriptor for reports. + """ payload, row_identity = serialize_simulation_frames( frames, calculation_provenance=calculation_provenance, uk_local_authority_metadata=uk_local_authority_metadata, ) - artifact = self._write_immutable( + return self._store_simulation_payload( simulation_path(prefix=prefix, role=simulation.role.value), + simulation, + payload, + row_identity, + calculation_provenance, + ) + + def write_simulation_named_by_digest( + self, + *, + prefix: str, + simulation: PlannedSimulationExecutionInput, + frames: Mapping[str, pd.DataFrame], + calculation_provenance: SPMProvenance | None = None, + uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, + ) -> SimulationArtifactDescriptor: + """Store a simulation's output frames under a key naming their SHA-256. + + A restarted simulation can always store its output this way (see + ``content_addressed_simulation_path``). Returns the artifact's + descriptor for reports. + """ + payload, row_identity = serialize_simulation_frames( + frames, + calculation_provenance=calculation_provenance, + uk_local_authority_metadata=uk_local_authority_metadata, + ) + return self._store_simulation_payload( + content_addressed_simulation_path( + prefix=prefix, + role=simulation.role.value, + content_sha256=sha256(payload).hexdigest(), + ), + simulation, + payload, + row_identity, + calculation_provenance, + ) + + def _store_simulation_payload( + self, + path: str, + simulation: PlannedSimulationExecutionInput, + payload: bytes, + row_identity: RowIdentity, + calculation_provenance: SPMProvenance | None, + ) -> SimulationArtifactDescriptor: + """Store one encoded simulation output at ``path`` and describe it.""" + artifact = self._write_immutable( + path, payload, content_type=SIMULATION_MEDIA_TYPE, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index 0bc91c889..cf49c79ce 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -28,6 +28,7 @@ stage12_region_groups, stage12_region_groups_for_model, ) +from .returned_coordination import coordinate_report_for_caller from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult from .segmentation import ( build_segment_inputs, @@ -67,6 +68,7 @@ "calculate_segment", "calculate_simulation_frames", "coordinate_report", + "coordinate_report_for_caller", "merge_segment_results", "redacted_failure_detail", "run_segmented_simulation", 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 280467162..9a1835f3c 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 @@ -1,15 +1,22 @@ -"""Coordination of a Stage 12 report's baseline and reform simulations.""" +"""Coordination of a report's baseline and reform simulations. + +``ReportSimulationCoordination`` holds the steps every report coordinator +runs. ``coordinate_report`` runs a Stage 12 report and records each step in +the Stage 12 tables; ``returned_coordination`` runs the same steps for a +report whose result returns to its caller and records nothing. +""" from __future__ import annotations import dataclasses import logging import time +from abc import ABC, abstractmethod from dataclasses import dataclass from datetime import UTC, datetime from functools import partial from hashlib import sha256 -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from uuid import uuid4 from policyengine_simulation_contract.stage12_execution import ( @@ -20,6 +27,7 @@ ComparisonSimulationRecord, PlannedSimulationExecutionInput, ReportExecutionInput, + ReportResultDelivery, ResultComparisonStatus, SimulationArtifactDescriptor, SimulationRole, @@ -209,21 +217,34 @@ def _cancel_pending_simulations( cancellation_recorded = True if cancellation_recorded and should_segment_simulation(simulation): # A segmented worker reads its record, cancels the region-group - # calls it started, and exits. Modal does not cancel spawned calls - # recursively, so cancelling the worker through Modal would leave - # those calls running. That is the fallback only when the record - # could not be written and the worker cannot learn of the - # cancellation. + # calls it started, and exits through its own failure handling. + # Cancelling it through Modal instead raises Modal's + # InputCancellation in it, which also makes it cancel those calls + # (Stage12SegmentRunner.run), but records nothing: that is the + # fallback only when the record could not be written and the + # worker cannot learn of the cancellation. continue - try: - calls[role].cancel() - except Exception as error: # noqa: BLE001 # a failed Modal cancellation is logged and must not replace the failure that caused it - _log_child_operation_failure( - logging.WARNING, - "Stage 12 could not cancel %s simulation %s in report %s (%s)", - simulation=simulation, - error=error, - ) + _cancel_child_call(calls[role], simulation) + + +def _cancel_child_call( + call: ChildCall, + simulation: PlannedSimulationExecutionInput, +) -> None: + """Ask Modal to stop one started child; log a failure rather than raise it. + + A cancellation error must never replace the failure that caused the + cancellation. + """ + try: + call.cancel() + except Exception as error: # noqa: BLE001 # a failed Modal cancellation is logged and must not replace the failure that caused it + _log_child_operation_failure( + logging.WARNING, + "Stage 12 could not cancel %s simulation %s in report %s (%s)", + simulation=simulation, + error=error, + ) def _simulation_failure_attributes( @@ -271,6 +292,26 @@ def _record_child_failure( _log_child_state_failure(simulation=simulation, error=state_error) +def _report_child_wait_failure( + error: Exception, + *, + simulation: PlannedSimulationExecutionInput, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> Stage12FailureDetail: + """Report a failed child call once and return its redacted detail.""" + return report_failure( + error, + FailureSite( + scope="child_wait", + default_code="simulation_invocation_failed", + attributes=_simulation_failure_attributes(simulation), + ), + runtime=runtime, + progress=progress, + ) + + def _child_failure_detail( *, error: Exception, @@ -287,13 +328,9 @@ def _child_failure_detail( the call's own failure is used and still recorded where possible, so a state error never replaces the child's failure. """ - detail = report_failure( + detail = _report_child_wait_failure( error, - FailureSite( - scope="child_wait", - default_code="simulation_invocation_failed", - attributes=_simulation_failure_attributes(simulation), - ), + simulation=simulation, runtime=runtime, progress=progress, ) @@ -336,7 +373,7 @@ def _child_record( ) -def _validate_dispatch_context( +def validate_dispatch_context( report: ReportExecutionInput, context: Stage12InvocationContext, application_name: str, @@ -498,6 +535,18 @@ 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 exactly the fields it received before + ``result_delivery`` existed. A child of a report returned to its caller + receives the field, so that it reads and writes no Stage 12 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") + + @dataclass(frozen=True) class CoordinatorInvocation: """The deployed coordinator call running a report: its Modal app and call ID.""" @@ -513,7 +562,8 @@ 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. + standard child wait window. A coordinator whose report returns to its + caller never uses ``store``. """ store: ComparisonStore | None = None @@ -529,14 +579,20 @@ class CoordinatorDependencies: wait_timeout_seconds: float = SIMULATION_WAIT_TIMEOUT_SECONDS -@dataclass -class _ReportCoordination: - """One claimed report's inputs and resolved seams, run through each phase.""" +@dataclass(kw_only=True) +class ReportSimulationCoordination(ABC): + """The steps a report coordinator runs however it delivers the report. + + It plans the report's output and both child inputs, starts each child, + waits for both within one shared window, and reads and checks their + artifacts. Subclasses decide what a failed or stopped child records: the + Stage 12 path writes its rows, and a report returned to its caller has no + rows to write. + """ report: ReportExecutionInput context: Stage12InvocationContext invocation: CoordinatorInvocation - persistence: ComparisonStore artifact_storage: Stage12ArtifactStore child_invoker: ChildInvoker dependencies: CoordinatorDependencies @@ -583,145 +639,58 @@ def stage_scope( stage, attributes, tracking=self.stage_tracking, operation=True ) - def run(self, parent: ComparisonReportRecord) -> dict[str, Any]: - """Plan, start and await both children, then aggregate and complete.""" - output_plan, planned_simulations = self._prepare(parent) - self.simulations = { - simulation.role.value: simulation for simulation in planned_simulations - } - children = self._create_children(planned_simulations) - descriptors, calls = self._start_children(planned_simulations, children) - # Every required call has been started before the coordinator observes - # baseline and reform together through nonblocking result probes. - self._wait_for_children(calls, descriptors) - return self._aggregate_and_complete(parent, output_plan, descriptors) + @abstractmethod + def _record_dispatch_failure( + self, + simulation: PlannedSimulationExecutionInput, + detail: Stage12FailureDetail, + ) -> None: + """Record that ``simulation`` could not be started, where there is a record.""" - def _prepare( - self, parent: ComparisonReportRecord - ) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: - """Bind the run's artifact prefix, then plan the outputs and child inputs.""" - with self.stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): - self.context = self.context.model_copy( - update={ - "artifact_prefix": ( - f"stage-12-runs/{parent.environment}/" - f"{parent.created_at:%Y}/{parent.created_at:%m}/" - f"{parent.evaluation_id}" - ), - "created_at": parent.created_at, - "retention_expires_at": parent.retention_expires_at, - } - ) - with self.stage_scope( - Stage.STAGE12_OUTPUT_PLANNING, - {"country": self.report.baseline.geography.country}, - ): - resolver = self.dependencies.output_plan_resolver - if resolver is None: - output_plan = resolve_report_output_plan( - self.report, - stage_scope=self.stage_scope, - ) - else: - output_plan = resolver(self.report) - simulations = [] - for simulation in (self.report.baseline, self.report.reform): - with self.stage_scope( - Stage.STAGE12_CHILD_INPUT_PLANNING, - {"simulation_role": simulation.role.value}, - ): - simulations.append(plan_simulation_input(simulation, output_plan)) - return output_plan, tuple(simulations) + @abstractmethod + def _child_call_failed(self, role: str, error: Exception) -> Stage12FailureDetail: + """Report a failed child call and return the failure its report records.""" - def _create_children( - self, planned_simulations: tuple[PlannedSimulationExecutionInput, ...] - ) -> dict[SimulationRole, ComparisonSimulationRecord]: - """Create each child's row, or resolve the row an earlier attempt made.""" - with self.span(Stage.STAGE12_CHILD_STATE_CREATE): - return { - simulation.role: self.persistence.create_or_resolve_simulation( - _child_record( - simulation=simulation, - context=self.context, - function_name=self.context.simulation_callable, - ) - ).record - for simulation in planned_simulations - } + @abstractmethod + def _stop_children(self, roles: set[str], *, reason: str) -> None: + """Stop each started, unfinished child in ``roles``. - def _start_children( - self, - planned_simulations: tuple[PlannedSimulationExecutionInput, ...], - children: dict[SimulationRole, ComparisonSimulationRecord], - ) -> tuple[dict[str, SimulationArtifactDescriptor], dict[str, ChildCall]]: - """Reuse finished children, reattach running ones and dispatch the rest.""" - descriptors = self.descriptors - calls = self.calls - for simulation in planned_simulations: - child = children[simulation.role] - role = simulation.role.value - if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: - descriptors[role] = self._retained_descriptor(child, simulation) - elif ( - child.status is ComparisonRunLifecycleStatus.RUNNING - 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) - else: - calls[role] = self._dispatch_child(child, simulation) - return descriptors, calls + ``reason`` ends the summary a recorded cancellation carries. + """ - def _retained_descriptor( + def _plan_report_simulations( self, - child: ComparisonSimulationRecord, - simulation: PlannedSimulationExecutionInput, - ) -> SimulationArtifactDescriptor: - """Describe a succeeded child's digest-checked retained artifact.""" - descriptor = descriptor_from_record(child, simulation) - retained = self.artifact_storage.read(descriptor.artifact.uri) - _require_artifact_digest( - retained, - descriptor.artifact.content_sha256, - mismatch_message="retained child artifact digest mismatch", - ) - return descriptor.model_copy( - update={ - "artifact": descriptor.artifact.model_copy( - update={"size_bytes": len(retained)} - ), - "calculation_provenance": ( - deserialize_calculation_provenance(retained) - ), - } - ) + ) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: + """Resolve the report's shared output plan and attach it to both inputs.""" + with self.stage_scope( + Stage.STAGE12_OUTPUT_PLANNING, + {"country": self.report.baseline.geography.country}, + ): + resolver = self.dependencies.output_plan_resolver + if resolver is None: + output_plan = resolve_report_output_plan( + self.report, + stage_scope=self.stage_scope, + ) + else: + output_plan = resolver(self.report) + simulations = [] + for simulation in (self.report.baseline, self.report.reform): + with self.stage_scope( + Stage.STAGE12_CHILD_INPUT_PLANNING, + {"simulation_role": simulation.role.value}, + ): + simulations.append(plan_simulation_input(simulation, output_plan)) + return output_plan, tuple(simulations) - def _dispatch_child( - self, - child: ComparisonSimulationRecord, - simulation: PlannedSimulationExecutionInput, - ) -> ChildCall: - """Mark the child dispatching, spawn it, and attach its Modal call ID. + def _spawn_child(self, simulation: PlannedSimulationExecutionInput) -> ChildCall: + """Start one child simulation, keep its call in ``calls`` and return it. - A spawn failure is reported and recorded on the child (unless it already - finished), every started peer is stopped, and Stage12ExecutionError is - raised with the failure. + The call is kept before the dispatch span closes, so a stop while the + span closes still finds it. A spawn failure is reported, recorded + through ``_record_dispatch_failure``, every started child is stopped, + and Stage12ExecutionError is raised with the failure. """ - dispatching_at = datetime.now(UTC) - dispatch_placeholder = f"dispatch-pending-{uuid4()}" - self.persistence.replace_simulation( - child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "modal_invocation_id": dispatch_placeholder, - "started_at": child.started_at or dispatching_at, - "updated_at": dispatching_at, - "error_code": None, - "error_summary": None, - "completed_at": None, - } - ) - ) try: with self.span( Stage.STAGE12_CHILD_DISPATCH, @@ -736,10 +705,11 @@ 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"), + context=_child_context_payload(self.context), observability_context=child_observability_context, ) ) + self.calls[simulation.role.value] = call except Exception as error: # noqa: BLE001 # any dispatch failure is recorded on the child and stops its started peers detail = report_failure( error, @@ -751,28 +721,12 @@ def _dispatch_child( runtime=self.runtime, progress=self.progress, ) - _record_child_failure( - store=self.persistence, - simulation=simulation, - detail=detail, - ) - _cancel_pending_simulations( - calls=self.calls, - simulations=self.simulations, - roles=set(self.calls), - store=self.persistence, - detail=_cancellation(f"the {simulation.role.value} simulation failed"), + self._record_dispatch_failure(simulation, detail) + self._stop_children( + set(self.calls), + reason=f"the {simulation.role.value} simulation failed", ) raise Stage12ExecutionError(detail) from None - # Record the call before attaching its ID, so a failed attach still - # stops it. - self.calls[simulation.role.value] = call - self.persistence.attach_simulation_invocation( - simulation.simulation_execution_id, - expected_placeholder=dispatch_placeholder, - modal_invocation_id=call.object_id, - updated_at=datetime.now(UTC), - ) return call def _wait_for_children( @@ -815,19 +769,10 @@ def _observe_children( except TimeoutError: continue except Exception as error: # noqa: BLE001 # a child call's failure can be any exception type; it is recorded and stops the other child - detail = _child_failure_detail( - error=error, - simulation=self.simulations[role], - store=self.persistence, - runtime=self.runtime, - progress=self.progress, - ) - _cancel_pending_simulations( - calls=calls, - simulations=self.simulations, - roles=pending - {role}, - store=self.persistence, - detail=_cancellation(f"the {role} simulation failed"), + detail = self._child_call_failed(role, error) + self._stop_children( + pending - {role}, + reason=f"the {role} simulation failed", ) raise Stage12ExecutionError(detail) from None descriptors[role] = descriptor @@ -852,14 +797,9 @@ def _observe_children( runtime=self.runtime, failure_scope="child_wait_timeout", ) - _cancel_pending_simulations( - calls=calls, - simulations=self.simulations, - roles=pending, - store=self.persistence, - detail=_cancellation( - f"the {timeout_seconds:g}-second wait window closed" - ), + self._stop_children( + pending, + reason=f"the {timeout_seconds:g}-second wait window closed", ) raise wait_timeout # Never sleep past the deadline, so the window closes on time. @@ -911,6 +851,229 @@ def _read_retained_outputs( uk_local_authority_metadata, ) + def _cancel_through_modal(self, roles: set[str]) -> None: + """Ask Modal to stop each started child in ``roles``, logging any failure.""" + for role in sorted(roles): + _cancel_child_call(self.calls[role], self.simulations[role]) + + def _report_failed_child_call( + self, + role: str, + error: Exception, + ) -> Stage12FailureDetail: + """Report the failed call of the child in ``role`` once; return its detail.""" + return _report_child_wait_failure( + error, + simulation=self.simulations[role], + runtime=self.runtime, + progress=self.progress, + ) + + def _failure_detail(self, error: Exception) -> Stage12FailureDetail: + """Return the redacted failure of a failed run, reporting it once. + + Child, dispatch and wait-timeout failures raise Stage12ExecutionError + only after reporting the original failure and stopping their unfinished + peers. Any other failure is reported here, before the caller records + or raises it. If it occurred while a child was still running, such as + a failed placeholder write or invocation attach on the Stage 12 path, + every started, unfinished child is then stopped through + ``_stop_children`` so none keeps running for a failed report. + """ + if isinstance(error, Stage12ExecutionError): + return error.detail + detail = report_failure( + error, + FailureSite( + scope="report_coordination", + default_code="report_coordination_failed", + attributes={"evaluation_id": str(self.report.evaluation_id)}, + ), + runtime=self.runtime, + progress=self.progress, + ) + self._stop_children( + set(self.calls) - set(self.descriptors), + reason="the coordinator failed", + ) + return detail + + +@dataclass(kw_only=True) +class _ReportCoordination(ReportSimulationCoordination): + """One claimed Stage 12 report's inputs and resolved seams, run through each phase. + + Every child's state, and the report's, is recorded in ``persistence``. + """ + + persistence: ComparisonStore + + def run(self, parent: ComparisonReportRecord) -> dict[str, Any]: + """Plan, start and await both children, then aggregate and complete.""" + output_plan, planned_simulations = self._prepare(parent) + self.simulations = { + simulation.role.value: simulation for simulation in planned_simulations + } + children = self._create_children(planned_simulations) + descriptors, calls = self._start_children(planned_simulations, children) + # Every required call has been started before the coordinator observes + # baseline and reform together through nonblocking result probes. + self._wait_for_children(calls, descriptors) + return self._aggregate_and_complete(parent, output_plan, descriptors) + + def _prepare( + self, parent: ComparisonReportRecord + ) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: + """Bind the run's artifact prefix, then plan the outputs and child inputs.""" + with self.stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): + self.context = self.context.model_copy( + update={ + "artifact_prefix": ( + f"stage-12-runs/{parent.environment}/" + f"{parent.created_at:%Y}/{parent.created_at:%m}/" + f"{parent.evaluation_id}" + ), + "created_at": parent.created_at, + "retention_expires_at": parent.retention_expires_at, + } + ) + return self._plan_report_simulations() + + def _create_children( + self, planned_simulations: tuple[PlannedSimulationExecutionInput, ...] + ) -> dict[SimulationRole, ComparisonSimulationRecord]: + """Create each child's row, or resolve the row an earlier attempt made.""" + with self.span(Stage.STAGE12_CHILD_STATE_CREATE): + return { + simulation.role: self.persistence.create_or_resolve_simulation( + _child_record( + simulation=simulation, + context=self.context, + function_name=self.context.simulation_callable, + ) + ).record + for simulation in planned_simulations + } + + def _start_children( + self, + planned_simulations: tuple[PlannedSimulationExecutionInput, ...], + children: dict[SimulationRole, ComparisonSimulationRecord], + ) -> tuple[dict[str, SimulationArtifactDescriptor], dict[str, ChildCall]]: + """Reuse finished children, reattach running ones and dispatch the rest.""" + descriptors = self.descriptors + calls = self.calls + for simulation in planned_simulations: + child = children[simulation.role] + role = simulation.role.value + if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: + descriptors[role] = self._retained_descriptor(child, simulation) + elif ( + child.status is ComparisonRunLifecycleStatus.RUNNING + 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) + else: + calls[role] = self._dispatch_child(child, simulation) + return descriptors, calls + + def _retained_descriptor( + self, + child: ComparisonSimulationRecord, + simulation: PlannedSimulationExecutionInput, + ) -> SimulationArtifactDescriptor: + """Describe a succeeded child's digest-checked retained artifact.""" + descriptor = descriptor_from_record(child, simulation) + retained = self.artifact_storage.read(descriptor.artifact.uri) + _require_artifact_digest( + retained, + descriptor.artifact.content_sha256, + mismatch_message="retained child artifact digest mismatch", + ) + return descriptor.model_copy( + update={ + "artifact": descriptor.artifact.model_copy( + update={"size_bytes": len(retained)} + ), + "calculation_provenance": ( + deserialize_calculation_provenance(retained) + ), + } + ) + + def _dispatch_child( + self, + child: ComparisonSimulationRecord, + simulation: PlannedSimulationExecutionInput, + ) -> ChildCall: + """Mark the child dispatching, spawn it, and attach its Modal call ID. + + A spawn failure is reported and recorded on the child (unless it already + finished), every started peer is stopped, and Stage12ExecutionError is + raised with the failure. + """ + dispatching_at = datetime.now(UTC) + dispatch_placeholder = f"dispatch-pending-{uuid4()}" + self.persistence.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": dispatch_placeholder, + "started_at": child.started_at or dispatching_at, + "updated_at": dispatching_at, + "error_code": None, + "error_summary": None, + "completed_at": None, + } + ) + ) + # _spawn_child keeps the call in calls, so a failed attach still stops + # it. + call = self._spawn_child(simulation) + self.persistence.attach_simulation_invocation( + simulation.simulation_execution_id, + expected_placeholder=dispatch_placeholder, + modal_invocation_id=call.object_id, + updated_at=datetime.now(UTC), + ) + return call + + @override + def _record_dispatch_failure( + self, + simulation: PlannedSimulationExecutionInput, + detail: Stage12FailureDetail, + ) -> None: + """Record the dispatch failure on the child, unless it already finished.""" + _record_child_failure( + store=self.persistence, + simulation=simulation, + detail=detail, + ) + + @override + def _child_call_failed(self, role: str, error: Exception) -> Stage12FailureDetail: + """Report the failed call; the failure its worker recorded takes precedence.""" + return _child_failure_detail( + error=error, + simulation=self.simulations[role], + store=self.persistence, + runtime=self.runtime, + progress=self.progress, + ) + + @override + def _stop_children(self, roles: set[str], *, reason: str) -> None: + """Record each unfinished child as cancelled, then stop it.""" + _cancel_pending_simulations( + calls=self.calls, + simulations=self.simulations, + roles=roles, + store=self.persistence, + detail=_cancellation(reason), + ) + def _aggregate_and_complete( self, parent: ComparisonReportRecord, @@ -974,33 +1137,10 @@ def _aggregate_and_complete( def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: """Record a failed run once, stop its unfinished children, fail the parent. - Child, dispatch and wait-timeout failures raise Stage12ExecutionError - only after reporting the original failure and stopping their unfinished - 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. + See ``_failure_detail`` for how the failure is reported and which + children are stopped. """ - if isinstance(error, Stage12ExecutionError): - detail = error.detail - else: - detail = report_failure( - error, - FailureSite( - scope="report_coordination", - default_code="report_coordination_failed", - attributes={"evaluation_id": str(self.report.evaluation_id)}, - ), - runtime=self.runtime, - progress=self.progress, - ) - _cancel_pending_simulations( - calls=self.calls, - simulations=self.simulations, - roles=set(self.calls) - set(self.descriptors), - store=self.persistence, - detail=_cancellation("the coordinator failed"), - ) + detail = self._failure_detail(error) record_failure_cause(self.runtime, error, detail.error_summary) failed_at = datetime.now(UTC) persist_failure_record( @@ -1036,18 +1176,24 @@ def coordinate_report( ``deduplicated`` summary instead), starts or reattaches the baseline and reform child simulations within one shared wait window, and aggregates their digest-checked artifacts; shadow runs are then compared with the - production result. A dispatch context that disagrees with the report - raises ``ValueError`` before the claim; any later failure stops the - unfinished children, marks the report FAILED and raises a - ``RuntimeError`` that carries no calculation detail. + production result. A dispatch context that disagrees with the report, or + asks for the report to return to its caller, raises ``ValueError`` before + the claim; any later failure stops the unfinished children, marks the + report FAILED and raises a ``RuntimeError`` that carries no calculation + detail. """ dependencies = dependencies or CoordinatorDependencies() # Open the shared child-wait window before any setup work. child_wait_deadline = dependencies.monotonic() + dependencies.wait_timeout_seconds report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) + if context.result_delivery is not ReportResultDelivery.STAGE12_DATABASE: + msg = ( + "a report returned to its caller runs through coordinate_report_for_caller" + ) + raise ValueError(msg) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) - _validate_dispatch_context(report, context, invocation.application_name) + validate_dispatch_context(report, context, invocation.application_name) _validate_submitted_parent( report=report, context=context, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py index fd6b36573..ed5b0b537 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py @@ -4,14 +4,21 @@ from collections.abc import Mapping from datetime import date -from typing import Any, NoReturn +from typing import TYPE_CHECKING, Any, NoReturn +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, +) from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_observability.errors import make_correlation_id from pydantic import BaseModel, ConfigDict from .failure_causes import StageProgress, failure_summary +if TYPE_CHECKING: + from uuid import UUID + ERROR_CODE_MAX_LENGTH = 64 ERROR_SUMMARY_MAX_LENGTH = 512 @@ -83,6 +90,24 @@ def failure_detail_from_record( ) +def report_execution_error( + evaluation_id: UUID, + detail: Stage12FailureDetail, +) -> ReportExecutionError: + """Return the typed error that carries ``detail`` to a returned report's caller. + + It holds the error code and summary a Stage 12 row would record for the + same failure. + """ + return ReportExecutionError( + ReportExecutionFailure( + evaluation_id=evaluation_id, + error_code=detail.error_code, + error_summary=detail.error_summary or "Simulation failed", + ) + ) + + def redacted_failure_detail( error: BaseException, *, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py index 83194b9f5..9a407f4eb 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 @@ -29,6 +29,7 @@ PlannedSimulationExecutionInput, ) + from .segment_contracts import Stage12SegmentInput from .simulation import SimulationCalculation logger = logging.getLogger(__name__) @@ -119,7 +120,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. Every started segment is also cancelled when the + run stops for any other reason before the results are collected, + such as a failed dispatch or Modal cancelling this worker's input. """ groups = self._groups() segments = build_segment_inputs(self.simulation, groups=groups) @@ -128,27 +131,24 @@ def run(self) -> SimulationCalculation: stage12_segment_count=len(segments), ) handles: list[tuple[int, Any]] = [] + payloads: list[object] | None = None try: - for segment in segments: - call = None - with self.runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH), - attributes={"segment_index": segment.segment_index}, - ): - call = self.child_function.spawn( - segment.model_dump(mode="python"), - observability_context=self.runtime.capture_context(), - ) - handles.append((segment.segment_index, _require_call(call))) + self._start_segments(segments, handles) + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) + ): + payloads = self._collect(handles) + # An Exception keeps its existing handling: a failed dispatch, a + # segment failure and a requested cancellation each cancel every + # started segment once before they are raised. except Exception: + raise + # Modal stops a cancelled input by raising InputCancellation, a + # BaseException, in it, and does not cancel the calls that input + # spawned, so this worker cancels them before it stops. + except BaseException: self._cancel_all(handles) raise - - payloads: list[object] | None = None - with self.runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) - ): - payloads = self._collect(handles) if payloads is None: msg = "Stage 12 segment wait returned no payloads" raise RuntimeError(msg) @@ -166,6 +166,35 @@ def run(self) -> SimulationCalculation: raise RuntimeError(msg) return merged + def _start_segments( + self, + segments: Sequence[Stage12SegmentInput], + handles: list[tuple[int, Any]], + ) -> None: + """Spawn every segment, adding each call to ``handles`` as soon as it starts. + + A failed dispatch cancels every segment already started, then is raised. + """ + try: + for segment in segments: + call = None + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH), + attributes={"segment_index": segment.segment_index}, + ): + call = self.child_function.spawn( + segment.model_dump(mode="python"), + observability_context=self.runtime.capture_context(), + ) + # Keep the call before its span closes, so a stop while + # the span closes still cancels it. + if call is not None: + handles.append((segment.segment_index, call)) + _require_call(call) + except Exception: + self._cancel_all(handles) + raise + def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: results: list[object | None] = [None] * len(handles) pending = set(range(len(handles))) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/returned_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/returned_coordination.py new file mode 100644 index 000000000..764850938 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/returned_coordination.py @@ -0,0 +1,256 @@ +"""Coordination of a report whose result returns to its caller. + +The coordinator plans, starts, awaits and checks the report's baseline and +reform simulations with the same steps as the Stage 12 coordinator +(``ReportSimulationCoordination``), but reads and writes no Stage 12 rows. It +returns a ``ReportExecutionResult`` with the aggregate inline, or raises +``ReportExecutionError``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, override + +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionProvenance, + ReportExecutionResult, + report_aggregate_sha256, +) +from policyengine_simulation_contract.stage12_execution import ( + AggregateReportArtifactPayload, + ReportExecutionInput, + ReportResultDelivery, + Stage12InvocationContext, +) +from policyengine_simulation_observability.stages import ( + STAGE12_CANONICAL_REPORT_STAGES, + Stage, +) + +from policyengine_simulation_executor.stage12_artifacts import comparison_run_prefix + +from .coordination import ( + CoordinatorDependencies, + CoordinatorInvocation, + ReportSimulationCoordination, + validate_dispatch_context, +) +from .dependencies import ModalChildInvoker, artifact_store +from .failure_causes import StageProgress, record_failure_cause +from .failures import failure_detail_from_record, report_execution_error + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + Stage12OutputPlan, + ) + + from .failures import Stage12FailureDetail + + +@dataclass(kw_only=True) +class _ReturnedReportCoordination(ReportSimulationCoordination): + """One report whose result returns to its caller, run through each phase. + + No Stage 12 row exists for the report or its children, so nothing records + their state: a failed child's own ``ReportExecutionError`` carries its + failure, and a child that must stop is cancelled through Modal. + """ + + def run(self) -> ReportExecutionResult: + """Plan, start and await both children, then aggregate and check the result.""" + with self.stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): + self.context = self.context.model_copy( + update={ + "artifact_prefix": comparison_run_prefix( + environment=self.context.environment, + created_at=self.context.created_at, + evaluation_id=self.report.evaluation_id, + ) + } + ) + output_plan, planned_simulations = self._plan_report_simulations() + self.simulations = { + simulation.role.value: simulation for simulation in planned_simulations + } + for simulation in planned_simulations: + self._spawn_child(simulation) + # Both children have started before the coordinator observes them + # together through nonblocking result probes. + self._wait_for_children(self.calls, self.descriptors) + baseline, reform, uk_local_authority_metadata = self._read_retained_outputs( + output_plan, self.descriptors["baseline"], self.descriptors["reform"] + ) + with self.span(Stage.STAGE12_AGGREGATION): + aggregate = self.dependencies.aggregator( + report=self.report, + baseline=baseline, + reform=reform, + uk_local_authority_metadata=uk_local_authority_metadata, + ) + return self._result(output_plan, aggregate) + + def _result( + self, + output_plan: Stage12OutputPlan, + aggregate: dict[str, Any], + ) -> ReportExecutionResult: + """Check the aggregate and build the result this report delivers. + + The aggregate passes the checks a Stage 12 aggregate artifact write + applies, must describe this report, and its result must be a + ``SingleYearMacroOutput``. + """ + payload = AggregateReportArtifactPayload.model_validate(aggregate) + if ( + payload.evaluation_id != self.report.evaluation_id + or payload.bundle != self.report.baseline.bundle + or payload.requested_aggregates != self.report.requested_aggregates + ): + msg = "aggregate does not describe this report" + raise ValueError(msg) + totals = SingleYearMacroOutput.model_validate(payload.result) + return ReportExecutionResult( + evaluation_id=self.report.evaluation_id, + requested_aggregates=self.report.requested_aggregates, + aggregate=totals, + aggregate_sha256=report_aggregate_sha256(totals), + output_plan=output_plan, + baseline=self.descriptors["baseline"], + reform=self.descriptors["reform"], + provenance=ReportExecutionProvenance( + bundle=self.report.baseline.bundle, + worker_version=self.context.worker_version, + modal_environment=self.context.modal_environment, + modal_application=self.context.modal_application, + simulation_callable=self.context.simulation_callable, + version_manifest_sha256=self.context.version_manifest_sha256, + coordinator_invocation_id=self.invocation.invocation_id, + ), + completed_at=datetime.now(UTC), + ) + + def fail(self, error: Exception) -> ReportExecutionError: + """Return the typed error for a failed run, reporting the failure once. + + See ``_failure_detail`` for how the failure is reported and which + children are stopped. The redacted failure is also recorded as the + operation's cause. + """ + detail = self._failure_detail(error) + record_failure_cause(self.runtime, error, detail.error_summary) + return report_execution_error(self.report.evaluation_id, detail) + + def stop_unfinished_children(self) -> None: + """Cancel every started, unfinished child through Modal. + + The coordinator does this when it stops without a failure to report, + as when Modal cancels its own input: Modal does not cancel the calls + that input spawned. + """ + self._stop_children( + set(self.calls) - set(self.descriptors), + reason="the coordinator stopped", + ) + + @override + def _record_dispatch_failure( + self, + simulation: PlannedSimulationExecutionInput, + detail: Stage12FailureDetail, + ) -> None: + """Record nothing: a returned report's children have no record.""" + + @override + def _child_call_failed(self, role: str, error: Exception) -> Stage12FailureDetail: + """Report the failed call; the child's own ``ReportExecutionError`` wins. + + That error carries the failure its worker reported, as a Stage 12 + worker's recorded failure does. + """ + detail = self._report_failed_child_call(role, error) + if isinstance(error, ReportExecutionError): + return failure_detail_from_record( + error_code=error.failure.error_code, + error_summary=error.failure.error_summary, + ) + return detail + + @override + def _stop_children(self, roles: set[str], *, reason: str) -> None: + """Cancel each child through Modal: no record can tell it to stop. + + Modal raises its ``InputCancellation`` in the child, and a segmented US + child then cancels the region-group segment calls it started + (``Stage12SegmentRunner.run``), since Modal does not cancel them. + """ + self._cancel_through_modal(roles) + + +def coordinate_report_for_caller( + payload: object, + context_payload: object, + *, + invocation: CoordinatorInvocation, + dependencies: CoordinatorDependencies | None = None, +) -> dict[str, Any]: + """Run one report whose result returns to its caller; return the result payload. + + The payload is a ``ReportExecutionResult`` in JSON mode. No Stage 12 row is + read or written and ``dependencies.store`` is never used. Nothing claims + the report, so a repeated dispatch runs it again, and the caller + deduplicates. Artifacts go under the Stage 12 layout for the context's + environment, creation time and report, whatever prefix the caller sent. + Both children start with the same delivery, so they write no rows either, + and each names its input and output by their digests. A repeated dispatch, + or a call Modal restarts after pre-empting it, therefore stores its + artifacts beside any an earlier attempt stored, even where it plans or + calculates differently. No aggregate artifact is written: the aggregate is + returned inline. + + A context that does not ask for the report to return to its caller, or + that disagrees with the report, raises ``ValueError`` before any child + starts. Any later failure stops the unfinished children through Modal, is + reported once, and is raised as ``ReportExecutionError`` with the redacted + code and summary a Stage 12 report row would record. When the coordinator + stops on a ``BaseException`` that is not an ``Exception``, such as the + ``InputCancellation`` Modal raises in a cancelled input, it cancels the + unfinished children through Modal and re-raises it unreported. + """ + dependencies = dependencies or CoordinatorDependencies() + # Open the shared child-wait window before any setup work. + child_wait_deadline = dependencies.monotonic() + dependencies.wait_timeout_seconds + report = ReportExecutionInput.model_validate(payload) + context = Stage12InvocationContext.model_validate(context_payload) + if context.result_delivery is not ReportResultDelivery.RETURN_TO_CALLER: + msg = "coordinate_report_for_caller requires result_delivery return_to_caller" + raise ValueError(msg) + validate_dispatch_context(report, context, invocation.application_name) + coordination = _ReturnedReportCoordination( + report=report, + context=context, + invocation=invocation, + artifact_storage=dependencies.artifacts or artifact_store(), + child_invoker=dependencies.invoker or ModalChildInvoker(), + dependencies=dependencies, + # A returned report cannot be compared with a production job. + stage_plan=STAGE12_CANONICAL_REPORT_STAGES, + progress=StageProgress(), + child_wait_deadline=child_wait_deadline, + ) + try: + result = coordination.run() + # Report a redacted cause for any coordinator, child-call, aggregation or + # artifact failure, and return it to the caller as a typed error. + except Exception as error: # noqa: BLE001 # every failure reaches the caller as a typed error with its redacted detail + raise coordination.fail(error) from None + # Modal stops a cancelled coordinator by raising InputCancellation, a + # BaseException, in it. + except BaseException: + coordination.stop_unfinished_children() + raise + return result.model_dump(mode="json") 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 728331273..df042f14f 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 @@ -19,6 +19,7 @@ ComparisonRunLifecycleStatus, ComparisonSimulationRecord, PlannedSimulationExecutionInput, + ReportResultDelivery, RowIdentity, SimulationArtifactDescriptor, SimulationExecutionInput, @@ -53,6 +54,7 @@ Stage12CancellationError, Stage12ExecutionError, Stage12FailureDetail, + report_execution_error, validate_policy_periods, ) from .output_planning import apply_output_plan, validate_output_frames @@ -485,6 +487,72 @@ class SimulationWorkerDependencies: runtime: ObservabilityRuntime | None = None +@dataclass(frozen=True) +class _SimulationRun: + """One planned simulation with its context, artifact store, seams and progress. + + ``progress`` records the stage the run is in, so a failure names it. + """ + + simulation: PlannedSimulationExecutionInput + context: Stage12InvocationContext + artifact_storage: Stage12ArtifactStore + dependencies: SimulationWorkerDependencies + progress: StageProgress + + def calculate_and_store( + self, + cancellation_requested: Callable[[], bool], + ) -> SimulationArtifactDescriptor: + """Check the policy, store the input, calculate, and store the output artifact. + + An eligible US simulation goes to the segmented calculator, which + checks ``cancellation_requested`` while it waits for its segments. + """ + simulation = self.simulation + stage_span = partial( + tracked_stage, + tracking=StageTracking( + STAGE12_SIMULATION_STAGES, self.progress, self.dependencies.runtime + ), + ) + # Modal restarts a pre-empted call, possibly on another machine whose + # calculation differs in its last bits, or on a redeployed application + # that plans the simulation differently. A returned simulation names + # its input and output by their digests, so it never needs a key that + # already holds other data. + returned = self.context.result_delivery is ReportResultDelivery.RETURN_TO_CALLER + write_input = ( + self.artifact_storage.write_input_named_by_digest + if returned + else self.artifact_storage.write_input + ) + write_simulation = ( + self.artifact_storage.write_simulation_named_by_digest + if returned + else self.artifact_storage.write_simulation + ) + validate_policy_periods(simulation.policy) + with stage_span(Stage.STAGE12_INPUT_WRITE): + write_input(prefix=self.context.artifact_prefix, simulation=simulation) + calculation = _calculate( + simulation, + self.dependencies, + progress=self.progress, + stage_span=stage_span, + cancellation_requested=cancellation_requested, + ) + validate_output_frames(calculation.frames, simulation.output_plan) + with stage_span(Stage.STAGE12_SIMULATION_ARTIFACT_WRITE): + return write_simulation( + prefix=self.context.artifact_prefix, + simulation=simulation, + frames=calculation.frames, + calculation_provenance=calculation.calculation_provenance, + uk_local_authority_metadata=calculation.uk_local_authority_metadata, + ) + + def run_single_simulation( payload: object, context_payload: object, @@ -497,12 +565,27 @@ def run_single_simulation( 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. + + When the context's ``result_delivery`` returns the report to its caller, + no Stage 12 row is read or written and ``store`` is not used; a failure + raises ``ReportExecutionError`` with its redacted detail instead (see + ``_run_simulation_for_caller``). """ dependencies = dependencies or SimulationWorkerDependencies() runtime = dependencies.runtime simulation = PlannedSimulationExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) + if context.result_delivery is ReportResultDelivery.RETURN_TO_CALLER: + return _run_simulation_for_caller( + _SimulationRun( + simulation=simulation, + context=context, + artifact_storage=dependencies.artifacts or artifact_store(), + dependencies=dependencies, + progress=StageProgress(role=simulation.role.value), + ) + ) persistence = dependencies.store or runtime_store(runtime) artifact_storage = dependencies.artifacts or artifact_store() child = persistence.get_simulation(simulation.simulation_execution_id) @@ -510,34 +593,15 @@ def run_single_simulation( return descriptor_from_record(child, simulation).model_dump(mode="json") running = _claim_simulation(persistence, child) cancellation_requested = _cancellation_check(persistence, simulation) - progress = StageProgress(role=simulation.role.value) - stage_span = partial( - tracked_stage, - tracking=StageTracking(STAGE12_SIMULATION_STAGES, progress, runtime), + run = _SimulationRun( + simulation=simulation, + context=context, + artifact_storage=artifact_storage, + dependencies=dependencies, + progress=StageProgress(role=simulation.role.value), ) try: - validate_policy_periods(simulation.policy) - with stage_span(Stage.STAGE12_INPUT_WRITE): - artifact_storage.write_input( - prefix=context.artifact_prefix, - simulation=simulation, - ) - calculation = _calculate( - simulation, - dependencies, - progress=progress, - stage_span=stage_span, - cancellation_requested=cancellation_requested, - ) - validate_output_frames(calculation.frames, simulation.output_plan) - with stage_span(Stage.STAGE12_SIMULATION_ARTIFACT_WRITE): - descriptor = artifact_storage.write_simulation( - prefix=context.artifact_prefix, - simulation=simulation, - frames=calculation.frames, - calculation_provenance=calculation.calculation_provenance, - uk_local_authority_metadata=calculation.uk_local_authority_metadata, - ) + descriptor = run.calculate_and_store(cancellation_requested) _record_success(persistence, running, descriptor) return descriptor.model_dump(mode="json") except Stage12CancellationError: @@ -551,12 +615,44 @@ def run_single_simulation( simulation, persistence=persistence, runtime=runtime, - progress=progress, + progress=run.progress, ) msg = "Stage 12 simulation execution failed" raise RuntimeError(msg) from None +def _never_cancelled() -> bool: + """Report no cancellation: nothing records one for a returned simulation.""" + return False + + +def _run_simulation_for_caller(run: _SimulationRun) -> dict[str, Any]: + """Calculate one simulation of a report returned to its caller. + + No Stage 12 row is read or written. The simulation is always calculated: + no record can carry a cancellation, so the coordinator stops this call + through Modal instead. Its input and output are stored under keys naming + their digests, so a call Modal restarts after pre-empting it can always + store its own. + A failure is reported once, recorded as the operation's cause, and raised + as ``ReportExecutionError`` with the redacted detail a Stage 12 row would + record. + """ + runtime = run.dependencies.runtime + try: + descriptor = run.calculate_and_store(_never_cancelled) + except Exception as error: # noqa: BLE001 # every calculation failure reaches the caller as a typed error with its redacted detail + detail = _simulation_failure_detail( + error, + run.simulation, + runtime=runtime, + progress=run.progress, + ) + record_failure_cause(runtime, error, detail.error_summary) + raise report_execution_error(run.simulation.evaluation_id, detail) from None + return descriptor.model_dump(mode="json") + + def _claim_simulation( persistence: ComparisonStore, child: ComparisonSimulationRecord, @@ -669,6 +765,34 @@ def _record_success( _require_completion_recorded(replaced=replaced) +def _simulation_failure_detail( + error: Exception, + simulation: PlannedSimulationExecutionInput, + *, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> Stage12FailureDetail: + """Report a calculation failure once and return its redacted detail.""" + if isinstance(error, Stage12ExecutionError): + # A segmented calculation reports its segment failure where it + # occurs, before raising this error. + return error.detail + return report_failure( + error, + FailureSite( + scope="simulation_execution", + default_code="simulation_execution_failed", + attributes={ + "simulation_role": simulation.role.value, + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + }, + ), + runtime=runtime, + progress=progress, + ) + + def _record_failure( error: Exception, simulation: PlannedSimulationExecutionInput, @@ -678,25 +802,12 @@ def _record_failure( progress: StageProgress, ) -> None: """Report a calculation failure once, then record it if still running.""" - if isinstance(error, Stage12ExecutionError): - # A segmented calculation reports its segment failure where it - # occurs, before raising this error. - detail = error.detail - else: - detail = report_failure( - error, - FailureSite( - scope="simulation_execution", - default_code="simulation_execution_failed", - attributes={ - "simulation_role": simulation.role.value, - "evaluation_id": str(simulation.evaluation_id), - "simulation_execution_id": str(simulation.simulation_execution_id), - }, - ), - runtime=runtime, - progress=progress, - ) + detail = _simulation_failure_detail( + error, + simulation, + runtime=runtime, + progress=progress, + ) record_failure_cause(runtime, error, detail.error_summary) failed_at = datetime.now(UTC) persist_failure_record( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py index 2a1d6aa06..73390a8a6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_artifacts.py @@ -20,6 +20,7 @@ validate_spm_calculation_provenance, ) from policyengine_simulation_contract.stage12_execution import ( + SimulationRole, UKLocalAuthorityBoundaryVersion, UKLocalAuthorityMetadata, ) @@ -32,6 +33,7 @@ serialize_simulation_frames, ) from pydantic import ValidationError +from test_stage12_runtime import _planned_simulation def _frames(): @@ -205,6 +207,59 @@ def test_immutable_retry_accepts_identical_bytes_and_rejects_other_bytes() -> No store._write_immutable("path", b"different", content_type="application/json") +def test_a_simulation_named_by_digest_gets_one_key_per_distinct_output() -> None: + store = Stage12ArtifactStore("private-stage12", store=InMemoryObjectStore()) + simulation = _planned_simulation(SimulationRole.BASELINE) + other_frames = _frames() + other_frames["household"]["household_net_income"] = np.array( + [20, 11], dtype=np.float32 + ) + + first = store.write_simulation_named_by_digest( + prefix="run", simulation=simulation, frames=_frames() + ) + repeated = store.write_simulation_named_by_digest( + prefix="run", simulation=simulation, frames=_frames() + ) + other = store.write_simulation_named_by_digest( + prefix="run", simulation=simulation, frames=other_frames + ) + fixed_key = store.write_simulation( + prefix="run", simulation=simulation, frames=_frames() + ) + + assert first.artifact.uri == ( + "gs://private-stage12/run/simulations/" + f"baseline-{first.artifact.content_sha256}.parquet" + ) + assert repeated == first + assert other.artifact.uri != first.artifact.uri + # The comparison path keeps its fixed key. + assert fixed_key.artifact.uri == ( + "gs://private-stage12/run/simulations/baseline.parquet" + ) + + +def test_an_input_named_by_digest_gets_one_key_per_distinct_input() -> None: + objects = InMemoryObjectStore() + store = Stage12ArtifactStore("private-stage12", store=objects) + simulation = _planned_simulation(SimulationRole.BASELINE) + replanned = simulation.model_copy( + update={"policy": {"reform": {"2026-01-01.2026-12-31": 1}}} + ) + + first = store.write_input_named_by_digest(prefix="run", simulation=simulation) + repeated = store.write_input_named_by_digest(prefix="run", simulation=simulation) + other = store.write_input_named_by_digest(prefix="run", simulation=replanned) + + assert first.uri == ( + f"gs://private-stage12/run/inputs/baseline-{first.content_sha256}.json" + ) + assert repeated == first + assert other.uri != first.uri + assert len(objects.objects) == 2 + + def test_aggregate_write_rejects_an_undeclared_payload_shape() -> None: store = Stage12ArtifactStore("private-stage12", store=InMemoryObjectStore()) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index ccf4a1918..4334c5528 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -7,9 +7,14 @@ import tomllib from contextlib import contextmanager from pathlib import Path +from uuid import UUID import pytest from fixtures.fake_modal import install_fake_modal +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, +) from policyengine_simulation_executor.release_bundle import ( get_bundled_package_version, ) @@ -398,6 +403,192 @@ def fail(*_, **__): assert runtime.shutdown_calls == 1 +def _invoke_coordinator( + monkeypatch, module, context, *parent, returned_error=None, runtime=None +): + """Call the Modal coordinator with both runtime coordinators replaced. + + The returned-report coordinator raises ``returned_error`` when one is given. + """ + from policyengine_simulation_executor import stage12_runtime + + runtime = runtime or RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "call-1", + raising=False, + ) + coordinator_calls = [] + + def stage12_coordinator( + _payload, _context_payload, parent_payload, *, invocation, dependencies + ): + coordinator_calls.append( + { + "coordinator": "stage12", + "parent": parent_payload, + "invocation_id": invocation.invocation_id, + "runtime": dependencies.runtime, + } + ) + return {"delivered_by": "stage12"} + + def returning_coordinator(_payload, _context_payload, *, invocation, dependencies): + coordinator_calls.append( + { + "coordinator": "return_to_caller", + "invocation_id": invocation.invocation_id, + "runtime": dependencies.runtime, + } + ) + if returned_error is not None: + raise returned_error + return {"delivered_by": "return_to_caller"} + + monkeypatch.setattr(stage12_runtime, "coordinate_report", stage12_coordinator) + monkeypatch.setattr( + stage12_runtime, "coordinate_report_for_caller", returning_coordinator + ) + result = module.coordinate_report( + {}, + context, + *parent, + observability_context=_observability_context(), + ) + return result, coordinator_calls, runtime + + +def test_coordinator_runs_a_returned_report_without_a_parent(monkeypatch) -> None: + module = _load(monkeypatch) + + result, coordinator_calls, runtime = _invoke_coordinator( + monkeypatch, module, {"result_delivery": "return_to_caller"} + ) + + assert result == {"delivered_by": "return_to_caller"} + assert coordinator_calls == [ + { + "coordinator": "return_to_caller", + "invocation_id": "call-1", + "runtime": runtime, + } + ] + # Telemetry tells a returned report apart from a Stage 12 run. + assert runtime.operations[0]["attributes"] == { + "runner_name": "stage12", + "execution_mode": "authoritative", + "result_delivery": "return_to_caller", + } + assert runtime.shutdown_calls == 1 + + +def test_coordinator_raises_a_returned_reports_error_unchanged(monkeypatch) -> None: + module = _load(monkeypatch) + runtime = RecordingRuntime() + error = ReportExecutionError( + ReportExecutionFailure( + evaluation_id=UUID("00000000-0000-0000-0000-000000000001"), + error_code="simulation_execution_failed", + error_summary="ValueError in dataset_load (baseline): bad input", + ) + ) + + with pytest.raises(ReportExecutionError) as raised: + _invoke_coordinator( + monkeypatch, + module, + {"result_delivery": "return_to_caller"}, + returned_error=error, + runtime=runtime, + ) + + # Modal pickles exactly what the function raises for the caller. + assert raised.value is error + assert runtime.shutdown_calls == 1 + + +def test_coordinator_refuses_a_parent_record_for_a_returned_report( + monkeypatch, +) -> None: + module = _load(monkeypatch) + + with pytest.raises(ValueError, match="no Stage 12 parent record"): + _invoke_coordinator( + monkeypatch, + module, + {"result_delivery": "return_to_caller"}, + {"evaluation_id": "parent"}, + ) + + +def test_coordinator_runs_a_stage12_report_with_unchanged_attributes( + monkeypatch, +) -> None: + module = _load(monkeypatch) + + result, coordinator_calls, runtime = _invoke_coordinator( + monkeypatch, module, {}, {"evaluation_id": "parent"} + ) + + assert result == {"delivered_by": "stage12"} + assert coordinator_calls == [ + { + "coordinator": "stage12", + "parent": {"evaluation_id": "parent"}, + "invocation_id": "call-1", + "runtime": runtime, + } + ] + assert runtime.operations[0]["attributes"] == { + "runner_name": "stage12", + "execution_mode": "authoritative", + } + + +@pytest.mark.parametrize( + "function_name", ["run_single_simulation_us", "run_single_simulation_uk"] +) +@pytest.mark.parametrize( + ("context", "expected_attributes"), + [ + ({}, {"runner_name": "stage12", "simulation_role": "baseline"}), + ( + {"result_delivery": "return_to_caller"}, + { + "runner_name": "stage12", + "simulation_role": "baseline", + "result_delivery": "return_to_caller", + }, + ), + ], +) +def test_worker_operation_says_when_its_report_returns_to_the_caller( + monkeypatch, + function_name, + context, + expected_attributes, +) -> None: + from policyengine_simulation_executor import stage12_runtime + + module = _load(monkeypatch) + runtime = RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + monkeypatch.setattr( + stage12_runtime, "run_single_simulation", lambda *_, **__: {"ok": True} + ) + + getattr(module, function_name)( + {"role": "baseline"}, + context, + observability_context=_observability_context(), + ) + + # Telemetry tells a returned report's workers apart from Stage 12 runs. + assert runtime.operations[0]["attributes"] == expected_attributes + + def test_external_package_values_are_assertions_only(monkeypatch) -> None: install_fake_modal(monkeypatch) monkeypatch.setenv("STAGE12_EXPECT_US_VERSION", "not-the-bundled-version") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_returned_report.py b/projects/policyengine-simulation-executor/tests/test_stage12_returned_report.py new file mode 100644 index 000000000..45294732f --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_returned_report.py @@ -0,0 +1,1004 @@ +"""Tests for reports and simulations whose result returns to their caller. + +With ``result_delivery`` set to ``return_to_caller``, the report coordinator +and its simulation workers read and write no Stage 12 rows. They return typed +values, and a failure raises ``ReportExecutionError`` with the redacted detail +a Stage 12 row would record. +""" + +from __future__ import annotations + +import logging +import math +from contextlib import contextmanager +from hashlib import sha256 +from uuid import uuid4 + +import pytest +from fixtures.fake_modal import InputCancellation +from fixtures.fake_object_store import InMemoryObjectStore +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +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 SPMInputError +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + ReportAggregate, + ReportResultDelivery, + SimulationArtifactDescriptor, + SimulationRole, + Stage12InvocationContext, +) +from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore +from policyengine_simulation_executor.stage12_runtime import ( + CoordinatorDependencies, + CoordinatorInvocation, + SimulationCalculation, + SimulationWorkerDependencies, + Stage12ExecutionError, + Stage12FailureDetail, + coordinate_report_for_caller, + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from policyengine_simulation_executor.stage12_runtime import ( + simulation as simulation_module, +) +from policyengine_simulation_observability.stages import ( + STAGE12_CANONICAL_REPORT_STAGES, + Stage, +) +from pydantic import ValidationError +from test_stage12_runtime import ( + BASELINE_ID, + EVALUATION_ID, + REFORM_ID, + ConcurrentInvoker, + ContextRecordingInvoker, + FakeArtifacts, + FakeStore, + InjectedChildFaults, + ReformDispatchFailureInvoker, + ScriptedChildInvoker, + _aggregate_payload, + _assert_log_identifies_child, + _bundle, + _child, + _context, + _coordination_log, + _failure_cause, + _frames, + _output_plan, + _parent, + _planned_simulation, + _report, + _report_with_segmentation, + coordinate_report, +) + + +class StageRecordingRuntime: + """Record the stages a run opens and the failure causes it attaches.""" + + def __init__(self) -> None: + self.spans: list[str] = [] + self.operations: list[str] = [] + self.recorded_exceptions: list[Exception] = [] + # Raise InputCancellation as the span with this name and these + # attributes closes, as Modal's cancellation can arrive at any point. + self.stop_on_close: tuple[str, dict] | None = None + + @contextmanager + def span(self, name, *_args, attributes=None, **_kwargs): + self.spans.append(name) + yield + if self.stop_on_close == (name, attributes): + msg = "Input was cancelled by user" + raise InputCancellation(msg) + + @contextmanager + def operation(self, name, *_args, **_kwargs): + self.operations.append(name) + yield + + def capture_context(self): + return {} + + def record_exception(self, error, *, handled, status_code=None): + assert handled + self.recorded_exceptions.append(error) + + +class NoDatabaseStore: + """A store that fails the test on any use: a returned report writes no rows.""" + + def __getattr__(self, name): + pytest.fail(f"a returned report used the Stage 12 store ({name})") + + +@pytest.fixture +def forbid_database(monkeypatch): + """Fail the test if a returned report builds a Stage 12 database store.""" + + def refuse_store(*_args, **_kwargs): + pytest.fail("a returned report built a Stage 12 database store") + + monkeypatch.setattr(simulation_module, "runtime_store", refuse_store) + monkeypatch.setattr(coordination_module, "runtime_store", refuse_store) + + +def _returning_context() -> Stage12InvocationContext: + return _context().model_copy( + update={"result_delivery": ReportResultDelivery.RETURN_TO_CALLER} + ) + + +def _run_returned_simulation(simulation, **seams): + seams.setdefault("artifacts", FakeArtifacts()) + return run_single_simulation( + simulation.model_dump(mode="json"), + _returning_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies(**seams), + ) + + +def test_returned_simulation_stores_its_artifact_without_the_database( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + artifacts = FakeArtifacts() + + result = _run_returned_simulation( + simulation, + artifacts=artifacts, + calculator=lambda _: _frames(), + ) + + descriptor = SimulationArtifactDescriptor.model_validate(result) + assert descriptor.role is SimulationRole.BASELINE + assert descriptor.simulation_execution_id == BASELINE_ID + stored = artifacts.read(descriptor.artifact.uri) + assert sha256(stored).hexdigest() == descriptor.artifact.content_sha256 + assert artifacts.input_writes == [ + (_context().artifact_prefix, SimulationRole.BASELINE) + ] + + +def test_returned_simulation_never_uses_a_store_it_was_given() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + + result = _run_returned_simulation( + simulation, + store=NoDatabaseStore(), + calculator=lambda _: _frames(), + ) + + assert result["role"] == "baseline" + + +def test_returned_simulation_times_the_same_stages_as_a_stage12_simulation( + forbid_database, +) -> None: + runtime = StageRecordingRuntime() + + _run_returned_simulation( + _planned_simulation(SimulationRole.BASELINE), + calculator=lambda _: _frames(), + runtime=runtime, + ) + + assert runtime.spans == [ + "stage12_input_write", + "stage12_calculation", + "stage12_simulation_artifact_write", + ] + + +def test_returned_simulation_failure_raises_its_redacted_detail( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + runtime = StageRecordingRuntime() + + def fail(_simulation): + msg = "could not reach postgresql://runtime:database-secret@db.example/postgres" + raise RuntimeError(msg) + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation(simulation, calculator=fail, runtime=runtime) + + failure = raised.value.failure + assert failure.evaluation_id == EVALUATION_ID + assert failure.error_code == "simulation_execution_failed" + assert _failure_cause(failure.error_summary) == ( + "RuntimeError in stage12_calculation (baseline): " + "could not reach postgresql://[REDACTED]@db.example/postgres" + ) + assert str(raised.value) == "Report execution failed" + # The original error, with its unredacted message, stays in telemetry. + assert raised.value.__cause__ is None + assert raised.value.__suppress_context__ is True + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == failure.error_summary + + +def test_returned_simulation_keeps_a_typed_failure(forbid_database) -> None: + def fail(_simulation): + msg = "SPM_YEAR_UNAVAILABLE" + raise SPMInputError(msg, "No SPM values for 2026") + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation( + _planned_simulation(SimulationRole.BASELINE), calculator=fail + ) + + assert raised.value.failure.error_code == "SPM_YEAR_UNAVAILABLE" + assert raised.value.failure.error_summary == "No SPM values for 2026" + + +def test_returned_simulation_rejects_an_invalid_policy_period_before_calculating( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={ + "policy": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3_000}} + } + ) + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation( + simulation, + calculator=lambda _: pytest.fail("calculation must not start"), + ) + + assert raised.value.failure.error_code == "invalid_policy_period" + assert raised.value.failure.error_summary == ( + "Policy parameter 'gov.irs.credits.ctc.amount.base[0].amount' has invalid " + "effective period '2026'; expected YYYY-MM-DD or " + "YYYY-MM-DD.YYYY-MM-DD." + ) + + +def test_returned_segmented_simulation_is_never_told_to_stop( + forbid_database, +) -> None: + # No record can carry a cancellation, so the coordinator stops this call + # through Modal instead. + simulation = _planned_simulation(SimulationRole.BASELINE) + cancellation_answers = [] + + def segmented_calculator(planned, cancellation_requested): + cancellation_answers.append((planned, cancellation_requested())) + return SimulationCalculation(frames=_frames()) + + _run_returned_simulation( + simulation, + calculator=lambda _: pytest.fail("the unsegmented calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + assert cancellation_answers == [(simulation, False)] + + +def test_returned_segmented_simulation_keeps_the_segment_failure( + forbid_database, +) -> None: + segment_failure = Stage12FailureDetail( + error_code="segment_execution_failed", + error_summary="Segment 3: MemoryError in stage12_calculation", + ) + + def segmented_calculator(_planned, _cancellation_requested): + raise Stage12ExecutionError(segment_failure) + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation( + _planned_simulation(SimulationRole.BASELINE), + calculator=lambda _: pytest.fail("the unsegmented calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + assert raised.value.failure.error_code == "segment_execution_failed" + assert raised.value.failure.error_summary == ( + "Segment 3: MemoryError in stage12_calculation" + ) + + +def _macro_aggregate(**_inputs): + return _aggregate_payload(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + +def _coordinate_returned_report(invoker, *, context=None, report=None, **seams): + context = context or _returning_context() + seams.setdefault("store", NoDatabaseStore()) + seams.setdefault("artifacts", FakeArtifacts()) + seams.setdefault("aggregator", _macro_aggregate) + seams.setdefault("output_plan_resolver", lambda _: _output_plan()) + return coordinate_report_for_caller( + (report or _report()).model_dump(mode="json"), + context.model_dump(mode="json"), + invocation=CoordinatorInvocation( + application_name=context.modal_application, + invocation_id="coordinator-1", + ), + dependencies=CoordinatorDependencies(invoker=invoker, **seams), + ) + + +def test_returned_report_returns_its_typed_result_without_the_database( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + invoker = ContextRecordingInvoker(artifacts) + aggregator_inputs = {} + + def aggregate(**inputs): + aggregator_inputs.update(inputs) + return _macro_aggregate() + + # The coordinator, not the caller, chooses where the artifacts go. + caller_context = _returning_context().model_copy( + update={"artifact_prefix": f"caller-chosen/{EVALUATION_ID}"} + ) + + returned = _coordinate_returned_report( + invoker, + context=caller_context, + artifacts=artifacts, + aggregator=aggregate, + ) + + result = ReportExecutionResult.model_validate(returned) + assert returned == result.model_dump(mode="json") + assert result.evaluation_id == EVALUATION_ID + assert result.requested_aggregates == tuple(ReportAggregate) + assert result.aggregate == SingleYearMacroOutput.model_validate( + CURRENT_SINGLE_YEAR_MACRO_RESULT + ) + assert result.aggregate_sha256 == report_aggregate_sha256(result.aggregate) + assert result.output_plan == _output_plan() + assert result.baseline.simulation_execution_id == BASELINE_ID + assert result.reform.simulation_execution_id == REFORM_ID + for descriptor in (result.baseline, result.reform): + stored = artifacts.read(descriptor.artifact.uri) + assert sha256(stored).hexdigest() == descriptor.artifact.content_sha256 + assert result.provenance == ReportExecutionProvenance( + bundle=_bundle(), + worker_version="5.2.0", + modal_environment="staging", + modal_application="policyengine-simulation-v2-py5-2-0", + simulation_callable="run_single_simulation_us", + version_manifest_sha256="b" * 64, + coordinator_invocation_id="coordinator-1", + ) + assert result.completed_at.tzinfo is not None + household_income = { + role: aggregator_inputs[role] + .frames["household"]["household_net_income"] + .tolist() + for role in ("baseline", "reform") + } + assert household_income == {"baseline": [100.0], "reform": [120.0]} + # The aggregate is returned inline, never written as an artifact. + assert artifacts.aggregate_writes == [] + assert invoker.events == ["spawn:baseline", "spawn:reform"] + 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 [context["result_delivery"] for context in invoker.child_contexts] == [ + "return_to_caller", + "return_to_caller", + ] + assert [context["artifact_prefix"] for context in invoker.child_contexts] == [ + f"stage-12-runs/staging/2026/09/{EVALUATION_ID}", + f"stage-12-runs/staging/2026/09/{EVALUATION_ID}", + ] + + +def test_returned_report_runs_only_the_stages_it_needs(forbid_database) -> None: + runtime = StageRecordingRuntime() + artifacts = FakeArtifacts() + + _coordinate_returned_report( + ConcurrentInvoker(artifacts), artifacts=artifacts, runtime=runtime + ) + + stage_name = STAGE12_CANONICAL_REPORT_STAGES.name + # No claim, child state or aggregate artifact stages: nothing is stored. + assert runtime.operations == [ + stage_name(Stage.STAGE12_COORDINATOR_PREPARATION), + stage_name(Stage.STAGE12_OUTPUT_PLANNING), + stage_name(Stage.STAGE12_CHILD_INPUT_PLANNING), + stage_name(Stage.STAGE12_CHILD_INPUT_PLANNING), + ] + assert runtime.spans == [ + stage_name(Stage.STAGE12_CHILD_DISPATCH), + stage_name(Stage.STAGE12_CHILD_DISPATCH), + stage_name(Stage.STAGE12_CHILD_WAIT), + stage_name(Stage.STAGE12_ARTIFACT_READ), + stage_name(Stage.STAGE12_AGGREGATION), + ] + + +def test_returned_report_claims_nothing_so_a_repeated_dispatch_runs_again( + forbid_database, +) -> None: + # The caller deduplicates; the coordinator keeps no record of the report. + artifacts = FakeArtifacts() + first_invoker = ConcurrentInvoker(artifacts) + repeated_invoker = ConcurrentInvoker(artifacts) + + first = _coordinate_returned_report(first_invoker, artifacts=artifacts) + repeated = _coordinate_returned_report(repeated_invoker, artifacts=artifacts) + + assert first_invoker.events == ["spawn:baseline", "spawn:reform"] + assert repeated_invoker.events == ["spawn:baseline", "spawn:reform"] + assert repeated["aggregate"] == first["aggregate"] + + +def test_returned_report_checks_the_dispatch_context_before_starting( + forbid_database, +) -> None: + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValueError, match="application differs"): + coordinate_report_for_caller( + _report().model_dump(mode="json"), + _returning_context().model_dump(mode="json"), + invocation=CoordinatorInvocation( + application_name="policyengine-simulation-v2-py9-9-9", + invocation_id="coordinator-1", + ), + dependencies=CoordinatorDependencies( + store=NoDatabaseStore(), + artifacts=FakeArtifacts(), + invoker=invoker, + ), + ) + + assert invoker.events == [] + + +def test_returned_report_coordinator_refuses_a_stage12_context( + forbid_database, +) -> None: + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValueError, match="return_to_caller"): + _coordinate_returned_report(invoker, context=_context()) + + assert invoker.events == [] + + +def test_stage12_coordinator_refuses_a_returned_report_context() -> None: + store = FakeStore() + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValueError, match="coordinate_report_for_caller"): + coordinate_report( + _report().model_dump(mode="json"), + _returning_context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=_macro_aggregate, + ) + + assert store.parent is None + assert invoker.events == [] + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_returned_report_keeps_the_failed_childs_detail_and_stops_its_peer( + forbid_database, + segmented, +) -> None: + child_failure = ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="SPM_YEAR_UNAVAILABLE", + error_summary="No SPM values for 2026", + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": ReportExecutionError(child_failure)} + ) + runtime = StageRecordingRuntime() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + invoker, + report=_report_with_segmentation(segmented=segmented), + runtime=runtime, + ) + + assert raised.value.failure == child_failure + assert raised.value.__cause__ is None + assert raised.value.__suppress_context__ is True + # No record can tell a child to stop, segmented or not, so Modal cancels it. + assert invoker.calls["reform"].cancel_count == 1 + assert invoker.calls["baseline"].cancel_count == 0 + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == "No SPM values for 2026" + + +def test_returned_report_redacts_a_child_call_failure_without_detail( + forbid_database, +) -> None: + invoker = ScriptedChildInvoker( + wait_errors={"reform": LookupError("worker lost with token=opaque-token")} + ) + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report(invoker) + + failure = raised.value.failure + assert failure.evaluation_id == EVALUATION_ID + assert failure.error_code == "simulation_invocation_failed" + assert _failure_cause(failure.error_summary) == ( + "LookupError in stage12_child_wait (reform): worker lost with token=[REDACTED]" + ) + assert invoker.calls["baseline"].cancel_count == 1 + assert invoker.calls["reform"].cancel_count == 0 + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_returned_report_wait_timeout_stops_every_child_through_modal( + forbid_database, + segmented, +) -> None: + invoker = ScriptedChildInvoker() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + invoker, + report=_report_with_segmentation(segmented=segmented), + wait_timeout_seconds=0, + ) + + assert raised.value.failure.error_code == "simulation_wait_timeout" + assert raised.value.failure.error_summary == ( + "Stage 12 simulations did not finish within 0 seconds of the report " + "coordinator starting." + ) + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_returned_report_stops_its_children_when_modal_cancels_it( + forbid_database, + segmented, +) -> None: + # Modal stops the coordinator's own cancelled input by raising + # InputCancellation in it, and does not cancel the calls it spawned. + def cancel_the_coordinator(): + msg = "Input was cancelled by user" + raise InputCancellation(msg) + + invoker = ScriptedChildInvoker(during_wait={"reform": cancel_the_coordinator}) + runtime = StageRecordingRuntime() + + with pytest.raises(InputCancellation): + _coordinate_returned_report( + invoker, + report=_report_with_segmentation(segmented=segmented), + runtime=runtime, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + # A cancellation is not a report failure, so none is reported. + assert runtime.recorded_exceptions == [] + + +@pytest.mark.parametrize( + ("closing_role", "started_roles"), + [("baseline", ["baseline"]), ("reform", ["baseline", "reform"])], +) +def test_returned_report_stops_a_child_whose_dispatch_is_closing_when_cancelled( + forbid_database, + closing_role, + started_roles, +) -> None: + runtime = StageRecordingRuntime() + # The worker has started, but its dispatch span has not closed yet. + runtime.stop_on_close = ( + STAGE12_CANONICAL_REPORT_STAGES.name(Stage.STAGE12_CHILD_DISPATCH), + {"simulation_role": closing_role}, + ) + invoker = ScriptedChildInvoker() + + with pytest.raises(InputCancellation): + _coordinate_returned_report(invoker, runtime=runtime) + + assert list(invoker.calls) == started_roles + assert [call.cancel_count for call in invoker.calls.values()] == [1] * len( + started_roles + ) + + +def test_returned_report_dispatch_failure_stops_the_started_child( + forbid_database, +) -> None: + invoker = ReformDispatchFailureInvoker() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report(invoker) + + failure = raised.value.failure + assert failure.error_code == "simulation_dispatch_failed" + assert _failure_cause(failure.error_summary) == ( + "ConnectionError in stage12_child_dispatch (reform): " + "dispatch contains sensitive details" + ) + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert invoker.calls["baseline"].cancel_count == 1 + + +def test_returned_report_cancellation_error_does_not_replace_the_failure( + forbid_database, + caplog, +) -> None: + invoker = ScriptedChildInvoker( + wait_errors={"reform": LookupError("reform worker failed")}, + cancel_errors={"baseline": ConnectionError("cancel contains sensitive data")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(ReportExecutionError) as raised, + ): + _coordinate_returned_report(invoker) + + assert raised.value.failure.error_code == "simulation_invocation_failed" + assert invoker.calls["baseline"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="baseline", + error_type="ConnectionError", + ) + + +def test_returned_report_aggregation_failure_is_redacted_and_reported_once( + forbid_database, +) -> None: + runtime = StageRecordingRuntime() + artifacts = FakeArtifacts() + + def fail(**_inputs): + msg = "aggregation read postgresql://runtime:database-secret@db.example/x" + raise RuntimeError(msg) + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker(artifacts), + artifacts=artifacts, + aggregator=fail, + runtime=runtime, + ) + + failure = raised.value.failure + assert failure.error_code == "report_coordination_failed" + assert _failure_cause(failure.error_summary) == ( + "RuntimeError in stage12_aggregation: " + "aggregation read postgresql://[REDACTED]@db.example/x" + ) + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == failure.error_summary + + +def test_returned_report_rejects_an_aggregate_outside_the_macro_output_schema( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker(artifacts), + artifacts=artifacts, + aggregator=lambda **_: _aggregate_payload({"budget": {}}), + ) + + failure = raised.value.failure + assert failure.error_code == "report_coordination_failed" + assert _failure_cause(failure.error_summary).startswith( + "ValidationError in stage12_aggregation: " + ) + assert artifacts.aggregate_writes == [] + + +def _aggregate_for_another_evaluation() -> dict: + return {**_macro_aggregate(), "evaluation_id": str(uuid4())} + + +def _aggregate_of_other_requested_aggregates() -> dict: + aggregate = _macro_aggregate() + return {**aggregate, "requested_aggregates": aggregate["requested_aggregates"][:1]} + + +def _aggregate_of_another_bundle() -> dict: + aggregate = _macro_aggregate() + return { + **aggregate, + "bundle": {**aggregate["bundle"], "country_package_version": "9.9.9"}, + } + + +@pytest.mark.parametrize( + "mismatched_aggregate", + [ + _aggregate_for_another_evaluation, + _aggregate_of_other_requested_aggregates, + _aggregate_of_another_bundle, + ], +) +def test_returned_report_rejects_an_aggregate_for_another_report( + forbid_database, + mismatched_aggregate, +) -> None: + artifacts = FakeArtifacts() + + def aggregate_for_another_report(**_inputs): + return mismatched_aggregate() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker(artifacts), + artifacts=artifacts, + aggregator=aggregate_for_another_report, + ) + + assert raised.value.failure.error_code == "report_coordination_failed" + assert "aggregate does not describe this report" in ( + raised.value.failure.error_summary + ) + + +def test_returned_report_rejects_artifacts_that_do_not_follow_the_plan( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(wrong_plan_digest=True) + ), + artifacts=artifacts, + aggregator=lambda **_: pytest.fail("aggregation must not start"), + ) + + assert raised.value.failure.error_code == "report_coordination_failed" + assert "do not satisfy the report output plan" in ( + raised.value.failure.error_summary + ) + + +def test_stage12_report_still_requires_its_parent_record() -> None: + # The Modal coordinator's parent argument is optional now; a Stage 12 + # report without one is rejected before anything is claimed or started. + store = FakeStore() + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValidationError): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + None, + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=_macro_aggregate, + ) + + assert store.parent is None + assert invoker.events == [] + + +# Modal restarts a pre-empted call on another machine, and the restarted call +# starts over. Its calculation can differ in the last bits from the first +# attempt's, so its output must not need to match what that attempt stored. + + +def _write_once_artifacts() -> Stage12ArtifactStore: + """Return an artifact store that, like the bucket, never replaces an object.""" + return Stage12ArtifactStore("private", store=InMemoryObjectStore()) + + +def test_a_restarted_returned_simulation_stores_output_that_differs_from_before( + forbid_database, +) -> None: + artifacts = _write_once_artifacts() + simulation = _planned_simulation(SimulationRole.BASELINE) + + first_attempt = _run_returned_simulation( + simulation, artifacts=artifacts, calculator=lambda _: _frames(100.0) + ) + restarted = _run_returned_simulation( + simulation, + artifacts=artifacts, + calculator=lambda _: _frames(math.nextafter(100.0, math.inf)), + ) + + for result in (first_attempt, restarted): + descriptor = SimulationArtifactDescriptor.model_validate(result) + stored = artifacts.read(descriptor.artifact.uri) + assert sha256(stored).hexdigest() == descriptor.artifact.content_sha256 + assert first_attempt["artifact"]["uri"] != restarted["artifact"]["uri"] + + +def test_a_restarted_returned_simulation_with_the_same_output_reuses_its_artifact( + forbid_database, +) -> None: + artifacts = _write_once_artifacts() + simulation = _planned_simulation(SimulationRole.BASELINE) + + first_attempt = _run_returned_simulation( + simulation, artifacts=artifacts, calculator=lambda _: _frames(100.0) + ) + restarted = _run_returned_simulation( + simulation, artifacts=artifacts, calculator=lambda _: _frames(100.0) + ) + + assert restarted == first_attempt + + +def test_a_restarted_returned_simulation_stores_an_input_that_differs_from_before( + forbid_database, +) -> None: + objects = InMemoryObjectStore() + artifacts = Stage12ArtifactStore("private", store=objects) + simulation = _planned_simulation(SimulationRole.BASELINE) + # A restart that lands on a redeployed application can plan the same + # simulation differently; any change to the planned input stands for it. + replanned = simulation.model_copy( + update={"policy": {"reform": {"2026-01-01.2026-12-31": 1}}} + ) + + _run_returned_simulation( + simulation, artifacts=artifacts, calculator=lambda _: _frames(100.0) + ) + _run_returned_simulation( + replanned, artifacts=artifacts, calculator=lambda _: _frames(100.0) + ) + + stored_inputs = [key for key in objects.objects if "/inputs/" in key] + assert len(stored_inputs) == 2 + + +def test_a_stage12_simulation_keeps_its_fixed_artifact_keys() -> None: + objects = InMemoryObjectStore() + artifacts = Stage12ArtifactStore("private", store=objects) + simulation = _planned_simulation(SimulationRole.BASELINE) + store = FakeStore() + store.children[simulation.simulation_execution_id] = _child(simulation) + + result = run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: _frames(100.0), + ), + ) + + prefix = _context().artifact_prefix + assert result["artifact"]["uri"] == ( + f"gs://private/{prefix}/simulations/baseline.parquet" + ) + assert sorted(objects.objects) == [ + f"{prefix}/inputs/baseline.json", + f"{prefix}/simulations/baseline.parquet", + ] + + +class InProcessWorkerCall: + """A started worker whose outcome was computed when it was started.""" + + def __init__(self, object_id, outcome) -> None: + self.object_id = object_id + self.outcome = outcome + self.cancel_count = 0 + + def get(self, *, timeout=None): + if isinstance(self.outcome, BaseException): + raise self.outcome + return self.outcome + + def cancel(self) -> None: + self.cancel_count += 1 + + +class InProcessWorkerInvoker: + """Run each returned worker in this process, storing into ``artifacts``. + + ``household_income`` gives each role's calculated household net income. + Probing ``interrupted_role``'s result raises ``KeyboardInterrupt``, as the + interrupt Modal sends a pre-empted container does in the coordinator. + """ + + def __init__(self, artifacts, *, household_income, interrupted_role=None): + self.artifacts = artifacts + self.household_income = household_income + self.interrupted_role = interrupted_role + self.calls = {} + + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role.value + if role == self.interrupted_role: + outcome = KeyboardInterrupt() + else: + income = self.household_income[role] + outcome = run_single_simulation( + request.simulation, + request.context, + required_country="us", + dependencies=SimulationWorkerDependencies( + artifacts=self.artifacts, + calculator=lambda _: _frames(income), + ), + ) + call = InProcessWorkerCall(f"call-{role}", outcome) + self.calls[role] = call + return call + + +def _report_with_a_dated_reform_period(): + """Return the test report with a reform period the worker's check accepts.""" + report = _report() + return report.model_copy( + update={ + "reform": report.reform.model_copy( + update={"policy": {"reform": {"2026-01-01.2026-12-31": 1}}} + ) + } + ) + + +def test_a_pre_empted_returned_report_finishes_when_modal_restarts_it( + forbid_database, +) -> None: + artifacts = _write_once_artifacts() + report = _report_with_a_dated_reform_period() + first_attempt = InProcessWorkerInvoker( + artifacts, + household_income={"baseline": 100.0}, + interrupted_role="reform", + ) + + # The baseline worker stores its output before Modal pre-empts the + # coordinator, which then cancels the unfinished reform worker. + with pytest.raises(KeyboardInterrupt): + _coordinate_returned_report(first_attempt, report=report, artifacts=artifacts) + assert first_attempt.calls["reform"].cancel_count == 1 + + restarted = InProcessWorkerInvoker( + artifacts, + household_income={ + "baseline": math.nextafter(100.0, math.inf), + "reform": 120.0, + }, + ) + returned = _coordinate_returned_report( + restarted, report=report, artifacts=artifacts + ) + + result = ReportExecutionResult.model_validate(returned) + for descriptor in (result.baseline, result.reform): + stored = artifacts.read(descriptor.artifact.uri) + assert sha256(stored).hexdigest() == descriptor.artifact.content_sha256 diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index eca8612bf..f18f57f05 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -450,6 +450,8 @@ def __init__(self): def write_input(self, *, prefix, simulation): self.input_writes.append((prefix, simulation.role)) + write_input_named_by_digest = write_input + def add_simulation(self, simulation, frames=None): payload, row_identity = serialize_simulation_frames(frames or _frames()) digest = sha256(payload).hexdigest() @@ -504,6 +506,10 @@ def write_simulation( update={"calculation_provenance": calculation_provenance} ) + # A returned simulation names its input and output by their digests; this + # fake keeps one of each per role either way. + write_simulation_named_by_digest = write_simulation + def read(self, uri): if uri in self.payloads: return self.payloads[uri] @@ -2320,6 +2326,57 @@ 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, request): + self.child_contexts.append(request.context) + return super().spawn(request) + + +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(), + ) + + # The Stage 12 worker payload keeps exactly the fields it had before the + # delivery field existed. + fields_before_the_delivery_field = { + "contract_version", + "request_id", + "environment", + "modal_environment", + "worker_version", + "modal_application", + "simulation_callable", + "version_manifest_sha256", + "bundle_manifest_sha256", + "artifact_prefix", + "production_function_call_id", + "created_at", + "retention_expires_at", + } + assert [set(context) for context in invoker.child_contexts] == [ + fields_before_the_delivery_field, + fields_before_the_delivery_field, + ] + + def test_coordinator_measures_output_and_child_input_planning() -> None: store = FakeStore() artifacts = FakeArtifacts() 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 15284a869..298b18cb2 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -8,6 +8,7 @@ import pandas as pd import pytest +from fixtures.fake_modal import InputCancellation from policyengine_simulation_contract.stage12_execution import SimulationRole from policyengine_simulation_executor.stage12_runtime import ( SegmentRunnerSeams, @@ -21,6 +22,10 @@ from policyengine_simulation_executor.stage12_runtime.partition import ( US_REGION_GROUPS, ) +from policyengine_simulation_observability.stages import ( + STAGE12_SIMULATION_STAGES, + Stage, +) from test_stage12_runtime import _planned_simulation if TYPE_CHECKING: @@ -49,6 +54,10 @@ class FakeRuntime: def __init__(self) -> None: self.context: dict[str, object] = {} self.spans: list[tuple[str, dict | None]] = [] + # Raise this when the span with this name and attributes opens or + # closes, as Modal's cancellation can arrive at any point. + self.stop_on_open: tuple[str, dict | None] | None = None + self.stop_on_close: tuple[str, dict | None] | None = None def set_context(self, **attributes) -> None: self.context.update(attributes) @@ -59,7 +68,13 @@ def capture_context(self) -> dict[str, str]: @contextmanager def span(self, name, *, attributes=None): self.spans.append((name, attributes)) + if self.stop_on_open == (name, attributes): + msg = "Input was cancelled by user" + raise InputCancellation(msg) yield + if self.stop_on_close == (name, attributes): + msg = "Input was cancelled by user" + raise InputCancellation(msg) class FakeCall: @@ -80,7 +95,9 @@ def cancel(self) -> None: class FakeFunction: - def __init__(self, results: list[dict], errors: dict[int, list[Exception]]): + def __init__( + self, results: list[dict], errors: dict[int, list[BaseException]] + ) -> None: self.results = results self.errors = errors self.spawns: list[tuple[dict, dict | None]] = [] @@ -106,7 +123,7 @@ def _results(simulation) -> list[dict]: ] -def _runner(*, errors: dict[int, list[Exception]] | None = None): +def _runner(*, errors: dict[int, list[BaseException]] | None = None): simulation = _planned_simulation(SimulationRole.BASELINE) function = FakeFunction(_results(simulation), errors or {}) @@ -219,3 +236,66 @@ def test_persisted_parent_cancellation_cancels_all_segments() -> None: assert sleeps == [0.5] assert all(call.cancelled == 1 for call in function.calls) + + +def test_modal_cancelling_the_worker_cancels_every_started_segment() -> None: + # Modal stops a cancelled input by raising InputCancellation in it, and does + # not cancel the calls that input spawned. + runner, function, _, _ = _runner( + errors={3: [InputCancellation("Input was cancelled by user")]} + ) + + with pytest.raises(InputCancellation): + runner.run() + + assert len(function.calls) == 20 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_modal_cancelling_the_worker_during_dispatch_cancels_started_segments() -> None: + runner, function, _, _ = _runner() + spawn = function.spawn + + def cancel_during_dispatch(payload, *, observability_context): + if payload["segment_index"] == 5: + msg = "Input was cancelled by user" + raise InputCancellation(msg) + return spawn(payload, observability_context=observability_context) + + function.spawn = cancel_during_dispatch + + with pytest.raises(InputCancellation): + runner.run() + + assert len(function.calls) == 5 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_modal_cancelling_the_worker_as_a_dispatch_span_closes_cancels_it() -> None: + runner, function, runtime, _ = _runner() + # The call has started, but its dispatch span has not closed yet. + runtime.stop_on_close = ( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH), + {"segment_index": 3}, + ) + + with pytest.raises(InputCancellation): + runner.run() + + assert len(function.calls) == 4 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_modal_cancelling_the_worker_before_the_wait_cancels_every_segment() -> None: + runner, function, runtime, _ = _runner() + # Every segment has started; the wait has not begun. + runtime.stop_on_open = ( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT), + None, + ) + + with pytest.raises(InputCancellation): + runner.run() + + assert len(function.calls) == 20 + assert all(call.cancelled == 1 for call in function.calls)