From c15785a11477888749e2d01077896bf7b3227399 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 21:47:31 +0400 Subject: [PATCH 01/25] Add Stage 12 segmentation foundations --- .../stage12_runtime/__init__.py | 6 + .../stage12_runtime/aggregation.py | 13 +- .../stage12_runtime/partition.py | 64 ++++++++++ .../stage12_runtime/segment_contracts.py | 63 ++++++++++ .../test_stage12_segmentation_foundation.py | 113 ++++++++++++++++++ 5 files changed, 256 insertions(+), 3 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py 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 dd61bce17..2b16e317f 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 @@ -2,6 +2,8 @@ from .aggregation import build_aggregate_report, build_spm_result from .coordination import coordinate_report +from .partition import stage12_region_groups, stage12_region_groups_for_model +from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult from .simulation import ( SimulationCalculation, calculate_simulation_frames, @@ -15,9 +17,13 @@ __all__ = [ "SimulationCalculation", + "Stage12SegmentInput", + "Stage12SegmentResult", "build_aggregate_report", "calculate_simulation_frames", "coordinate_report", "run_single_simulation", "simulation_input_sha256", + "stage12_region_groups", + "stage12_region_groups_for_model", ] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index a9dadb23a..b064fccbe 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -6,6 +6,7 @@ from typing import Any import pandas as pd +from policyengine.core import Simulation from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( ReportExecutionInput, @@ -51,6 +52,15 @@ def validate_uk_local_authority_metadata( return None +class PrecomputedSimulation(Simulation): + """A Stage 12 aggregate input whose output dataset is already complete.""" + + def ensure(self) -> None: + """Do not recalculate frames loaded from Stage 12 artifacts.""" + + return + + def validate_aligned_outputs( report: ReportExecutionInput, output_plan: Stage12OutputPlan, @@ -130,9 +140,6 @@ def build_aggregate_report( reform_descriptor: SimulationArtifactDescriptor, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> dict[str, Any]: - from policyengine_simulation_executor.segmented_national_reduce import ( - PrecomputedSimulation, - ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py new file mode 100644 index 000000000..5e7066b80 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py @@ -0,0 +1,64 @@ +"""Stage 12-owned US partition for segmented national calculations.""" + +from __future__ import annotations + +from typing import Any + +# This is an independent snapshot of the measured 20-way partition used by +# the v1 runner when Stage 12 segmentation was introduced. Stage 12 must not +# import the v1 partition: either execution path can now evolve without +# changing the other one implicitly. +US_REGION_GROUPS: tuple[tuple[str, ...], ...] = ( + ("state/hi", "state/ia", "state/wi"), + ("state/ak", "state/ga", "state/la"), + ("state/al", "state/co", "state/ri"), + ("state/mn", "state/mo", "state/nh"), + ("state/az", "state/de", "state/me"), + ("state/ny", "state/wv"), + ("state/in", "state/ky", "state/md"), + ("state/mt", "state/sc", "state/wa"), + ("state/id", "state/or", "state/tn"), + ("state/ar", "state/nc", "state/nd"), + ("state/dc", "state/ks", "state/mi"), + ("state/ct", "state/oh", "state/vt"), + ("state/nv", "state/va", "state/wy"), + ("state/pa", "state/ut"), + ("state/sd", "state/tx"), + ("state/ca",), + ("state/il", "state/ms"), + ("state/fl", "state/nm"), + ("state/ne", "state/nj"), + ("state/ma", "state/ok"), +) + + +def stage12_region_groups(country: str) -> list[list[str]] | None: + """Return a mutable copy of the Stage 12 partition when supported.""" + + if country.lower() != "us": + return None + return [list(group) for group in US_REGION_GROUPS] + + +def stage12_region_groups_for_model(country: str, model: Any) -> list[list[str]]: + """Return the partition, retaining newly registered state regions. + + The static partition is measured for the certified dataset. A later model + can register another state-level region before the partition is rebalanced; + assigning it to the final group prevents silently omitting its households. + """ + + groups = stage12_region_groups(country) + if groups is None: + raise ValueError(f"Stage 12 has no national partition for {country!r}") + registry = getattr(model, "region_registry", None) + if registry is None: + return groups + registered = { + region.code + for region in registry.regions + if getattr(region, "region_type", None) == "state" + } + covered = {code for group in groups for code in group} + groups[-1].extend(sorted(registered - covered)) + return groups diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py new file mode 100644 index 000000000..a22873e0e --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py @@ -0,0 +1,63 @@ +"""Strict internal contracts for Stage 12 segment function calls.""" + +from __future__ import annotations + +from typing import Annotated, Literal + +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + Sha256Digest, + SimulationRole, +) +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + + +class _StrictSegmentModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class Stage12SegmentInput(_StrictSegmentModel): + """One single-policy region-group calculation.""" + + schema_version: Literal[1] = 1 + simulation: PlannedSimulationExecutionInput + segment_index: Annotated[int, Field(ge=0, lt=20)] + region_codes: Annotated[tuple[str, ...], Field(min_length=1)] + + @field_validator("region_codes") + @classmethod + def require_unique_region_codes(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if len(value) != len(set(value)): + raise ValueError("segment region codes must be unique") + if any(not code.startswith("state/") for code in value): + raise ValueError("segment region codes must identify state regions") + return value + + @model_validator(mode="after") + def require_us_simulation(self) -> Stage12SegmentInput: + if self.simulation.geography.country != "us": + raise ValueError("Stage 12 segments currently support only the US") + if self.simulation.role not in { + SimulationRole.BASELINE, + SimulationRole.REFORM, + }: + raise ValueError("Stage 12 segment role must be baseline or reform") + return self + + +class Stage12SegmentResult(_StrictSegmentModel): + """Compressed, dtype-preserving output from one segment calculation.""" + + schema_version: Literal[1] = 1 + segment_index: Annotated[int, Field(ge=0, lt=20)] + role: SimulationRole + region_codes: Annotated[tuple[str, ...], Field(min_length=1)] + parquet_payload: Annotated[bytes, Field(min_length=1)] + payload_sha256: Sha256Digest + + @field_validator("role") + @classmethod + def require_report_role(cls, value: SimulationRole) -> SimulationRole: + if value not in {SimulationRole.BASELINE, SimulationRole.REFORM}: + raise ValueError("Stage 12 segment role must be baseline or reform") + return value diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py new file mode 100644 index 000000000..63cc3b9eb --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py @@ -0,0 +1,113 @@ +"""Foundational tests for Stage 12-owned segmented execution.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, + stage12_region_groups, + stage12_region_groups_for_model, +) +from policyengine_simulation_executor.stage12_runtime.segment_contracts import ( + Stage12SegmentInput, +) +from test_stage12_runtime import _planned_simulation +from policyengine_simulation_contract.stage12_execution import SimulationRole + +PROJECT = Path(__file__).resolve().parents[1] +PACKAGE = PROJECT / "src" / "policyengine_simulation_executor" +FORBIDDEN_IMPORTS = { + "policyengine_simulation_executor.artifact_keys", + "policyengine_simulation_executor.baseline_artifacts", + "policyengine_simulation_executor.national_partition", + "policyengine_simulation_executor.precompute", + "policyengine_simulation_executor.precompute_models", + "policyengine_simulation_executor.segmented_national_reduce", + "src.modal._image_setup", + "src.modal.precompute_app", + "src.modal.segmented_national", +} + + +def test_us_partition_has_twenty_unique_groups_and_complete_state_coverage() -> None: + groups = stage12_region_groups("us") + + assert groups is not None + assert len(groups) == 20 + assert all(groups) + flattened = [code for group in groups for code in group] + assert len(flattened) == len(set(flattened)) == 51 + assert set(flattened) == {code for group in US_REGION_GROUPS for code in group} + assert stage12_region_groups("uk") is None + + +def test_partition_returns_independent_mutable_copies() -> None: + first = stage12_region_groups("us") + second = stage12_region_groups("us") + assert first is not None and second is not None + + first[-1].append("state/test") + + assert "state/test" not in second[-1] + + +def test_new_model_regions_are_added_to_the_final_group_in_canonical_order() -> None: + model = SimpleNamespace( + region_registry=SimpleNamespace( + regions=[ + SimpleNamespace(code="state/zz", region_type="state"), + SimpleNamespace(code="state/aa", region_type="state"), + SimpleNamespace(code="district/1", region_type="district"), + ] + ) + ) + + groups = stage12_region_groups_for_model("us", model) + + assert groups[-1][-2:] == ["state/aa", "state/zz"] + + +def test_segment_input_rejects_duplicate_or_non_state_regions() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + + with pytest.raises(ValidationError, match="must be unique"): + Stage12SegmentInput( + simulation=simulation, + segment_index=0, + region_codes=("state/ca", "state/ca"), + ) + with pytest.raises(ValidationError, match="state regions"): + Stage12SegmentInput( + simulation=simulation, + segment_index=0, + region_codes=("district/1",), + ) + + +def test_stage12_feature_modules_do_not_import_v1_feature_implementations() -> None: + feature_paths = list((PACKAGE / "stage12_runtime").glob("*.py")) + cache_package = PACKAGE / "stage12_cache" + if cache_package.exists(): + feature_paths.extend(cache_package.glob("*.py")) + + violations: list[str] = [] + for path in feature_paths: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + imported = [node.module or ""] + else: + continue + for module in imported: + if module in FORBIDDEN_IMPORTS: + violations.append(f"{path.name}: {module}") + + assert violations == [] From 939f10e2f1bb4e017be616469ac45629eb788788 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 21:50:05 +0400 Subject: [PATCH 02/25] Add Stage 12 segment calculation and merge --- .../stage12_runtime/__init__.py | 12 ++ .../stage12_runtime/segment_contracts.py | 14 ++ .../stage12_runtime/segmentation.py | 188 ++++++++++++++++++ .../stage12_runtime/simulation.py | 33 ++- .../tests/test_stage12_runtime.py | 2 +- .../tests/test_stage12_segment_calculation.py | 170 ++++++++++++++++ 6 files changed, 416 insertions(+), 3 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py 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 2b16e317f..70be17666 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 @@ -4,8 +4,15 @@ from .coordination import coordinate_report from .partition import stage12_region_groups, stage12_region_groups_for_model from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult +from .segmentation import ( + build_segment_inputs, + calculate_segment, + merge_segment_results, + should_segment_simulation, +) from .simulation import ( SimulationCalculation, + build_stage12_simulation, calculate_simulation_frames, run_single_simulation, simulation_input_sha256, @@ -20,10 +27,15 @@ "Stage12SegmentInput", "Stage12SegmentResult", "build_aggregate_report", + "build_segment_inputs", + "build_stage12_simulation", "calculate_simulation_frames", + "calculate_segment", "coordinate_report", "run_single_simulation", + "merge_segment_results", "simulation_input_sha256", "stage12_region_groups", "stage12_region_groups_for_model", + "should_segment_simulation", ] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py index a22873e0e..ba4cfb68d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py @@ -2,6 +2,7 @@ from __future__ import annotations +from hashlib import sha256 from typing import Annotated, Literal from policyengine_simulation_contract.stage12_execution import ( @@ -55,9 +56,22 @@ class Stage12SegmentResult(_StrictSegmentModel): parquet_payload: Annotated[bytes, Field(min_length=1)] payload_sha256: Sha256Digest + @field_validator("region_codes") + @classmethod + def require_unique_region_codes(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if len(value) != len(set(value)): + raise ValueError("segment region codes must be unique") + return value + @field_validator("role") @classmethod def require_report_role(cls, value: SimulationRole) -> SimulationRole: if value not in {SimulationRole.BASELINE, SimulationRole.REFORM}: raise ValueError("Stage 12 segment role must be baseline or reform") return value + + @model_validator(mode="after") + def require_payload_digest(self) -> Stage12SegmentResult: + if sha256(self.parquet_payload).hexdigest() != self.payload_sha256: + raise ValueError("segment payload digest differs from its bytes") + return self diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py new file mode 100644 index 000000000..f3b1408da --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py @@ -0,0 +1,188 @@ +"""Independent Stage 12 segment calculation and exact national merge.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from hashlib import sha256 +from typing import Any + +import pandas as pd +from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, +) + +from policyengine_simulation_executor.stage12_artifacts import ( + deserialize_calculation_provenance, + deserialize_simulation_frames, + serialize_simulation_frames, +) + +from .output_planning import validate_output_frames +from .partition import stage12_region_groups +from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult +from .simulation import SimulationCalculation, calculate_simulation_frames + + +def should_segment_simulation(simulation: PlannedSimulationExecutionInput) -> bool: + """Return whether this logical simulation should use the US partition.""" + + if simulation.geography.country != "us": + return False + if simulation.geography.region.strip().lower() != "us": + return False + if simulation.options.get("segmented") is False: + return False + requirements = simulation.output_plan.requirements + if requirements.include_cliff_impacts: + return False + if requirements.labor_supply_response_active: + return False + return stage12_region_groups("us") is not None + + +def build_segment_inputs( + simulation: PlannedSimulationExecutionInput, + *, + groups: Sequence[Sequence[str]] | None = None, +) -> tuple[Stage12SegmentInput, ...]: + """Build one immutable single-policy input per region group.""" + + selected = groups or stage12_region_groups(simulation.geography.country) + if not selected: + raise ValueError("Stage 12 simulation has no national partition") + if len(selected) != 20: + raise ValueError("Stage 12 national partition must contain 20 groups") + return tuple( + Stage12SegmentInput( + simulation=simulation, + segment_index=index, + region_codes=tuple(group), + ) + for index, group in enumerate(selected) + ) + + +def _scoped_simulation(segment: Stage12SegmentInput) -> PlannedSimulationExecutionInput: + options = { + key: value + for key, value in segment.simulation.options.items() + if key != "segmented" + } + options["region_group"] = list(segment.region_codes) + return segment.simulation.model_copy(update={"options": options}) + + +def calculate_segment( + payload: object, + *, + calculator: Callable[ + [PlannedSimulationExecutionInput], + SimulationCalculation | Mapping[str, pd.DataFrame], + ] = calculate_simulation_frames, + runtime: ObservabilityRuntime | None = None, +) -> Stage12SegmentResult: + """Calculate one segment without persistence or report aggregation.""" + + segment = Stage12SegmentInput.model_validate(payload) + scoped = _scoped_simulation(segment) + if calculator is calculate_simulation_frames: + calculated = calculate_simulation_frames(scoped, runtime=runtime) + else: + calculated = calculator(scoped) + if isinstance(calculated, SimulationCalculation): + frames = calculated.frames + provenance = calculated.calculation_provenance + else: + frames = calculated + provenance = None + validate_output_frames(frames, scoped.output_plan) + parquet_payload, _ = serialize_simulation_frames( + frames, + calculation_provenance=provenance, + ) + return Stage12SegmentResult( + segment_index=segment.segment_index, + role=segment.simulation.role, + region_codes=segment.region_codes, + parquet_payload=parquet_payload, + payload_sha256=sha256(parquet_payload).hexdigest(), + ) + + +def _merge_calculation_provenance( + values: Sequence[dict[str, Any] | None], +) -> dict[str, Any] | None: + present = [value for value in values if value is not None] + if not present: + return None + if len(present) != len(values): + raise ValueError("segment calculation metadata is incomplete") + selection = present[0].get("spm_config") + if any(value.get("spm_config") != selection for value in present[1:]): + raise ValueError("segment calculation selections do not match") + return { + "spm_config": selection, + "spm_provenance": [value.get("spm_provenance") for value in present], + } + + +def merge_segment_results( + simulation: PlannedSimulationExecutionInput, + payloads: Sequence[object], + *, + groups: Sequence[Sequence[str]] | None = None, +) -> SimulationCalculation: + """Validate and merge all segment payloads into one national result.""" + + expected = build_segment_inputs(simulation, groups=groups) + results = [Stage12SegmentResult.model_validate(payload) for payload in payloads] + if len(results) != len(expected): + raise ValueError( + "Stage 12 segmented calculation returned an incomplete result set" + ) + by_index: dict[int, Stage12SegmentResult] = {} + for result in results: + if result.segment_index in by_index: + raise ValueError( + "Stage 12 segmented calculation returned a duplicate index" + ) + by_index[result.segment_index] = result + + decoded: list[dict[str, pd.DataFrame]] = [] + provenance: list[dict[str, Any] | None] = [] + for segment in expected: + result = by_index.get(segment.segment_index) + if result is None: + raise ValueError("Stage 12 segmented calculation omitted a segment") + if result.role is not simulation.role: + raise ValueError("Stage 12 segment result names another simulation role") + if result.region_codes != segment.region_codes: + raise ValueError("Stage 12 segment result names another region group") + decoded.append(deserialize_simulation_frames(result.parquet_payload)) + provenance.append(deserialize_calculation_provenance(result.parquet_payload)) + + expected_entities = set(decoded[0]) + if any(set(frames) != expected_entities for frames in decoded[1:]): + raise ValueError("Stage 12 segment entity schemas do not match") + merged: dict[str, pd.DataFrame] = {} + for entity in sorted(expected_entities): + parts = [frames[entity] for frames in decoded] + first_dtypes = {column: str(dtype) for column, dtype in parts[0].dtypes.items()} + for part in parts[1:]: + if { + column: str(dtype) for column, dtype in part.dtypes.items() + } != first_dtypes: + raise ValueError("Stage 12 segment column dtypes do not match") + frame = pd.concat(parts, ignore_index=True) + identifier = f"{entity}_id" + if identifier not in frame: + raise ValueError(f"Stage 12 segment entity {entity!r} has no identifier") + if frame[identifier].duplicated().any(): + raise ValueError(f"Stage 12 segment entity {entity!r} has duplicate IDs") + merged[entity] = frame + validate_output_frames(merged, simulation.output_plan) + return SimulationCalculation( + frames=merged, + calculation_provenance=_merge_calculation_provenance(provenance), + ) 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 139ce06cd..157eada98 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 @@ -50,6 +50,35 @@ class SimulationCalculation: uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None +def build_stage12_simulation( + params: dict[str, Any], + *, + dataset: Any, + dataset_selection: Any, + country_module: Any, + policy: dict[str, Any] | None, + scoping_strategy: Any, + region_code: str | None, +): + """Construct a Stage 12 simulation without selecting v1 cache code.""" + + from policyengine.core import Simulation + from policyengine_simulation_executor.spm import normalize_runtime_spm + + # These values become part of Stage 12's independent cache eligibility + # and identity in the next layer; retaining them in this constructor keeps + # the monolithic and segmented call sites identical. + del dataset_selection, region_code + selection = normalize_runtime_spm(params) + return Simulation( + **({"spm": selection} if selection is not None else {}), + dataset=dataset, + tax_benefit_model_version=country_module.model, + policy=policy, + scoping_strategy=scoping_strategy, + ) + + def simulation_input_sha256(simulation: SimulationExecutionInput) -> str: normalized = simulation.model_dump( mode="json", @@ -128,7 +157,6 @@ def calculate_simulation_frames( **simulation.options, } from policyengine_simulation_executor.simulation_runtime import ( - _build_simulation, _country_module, _load_dataset, _normalise_policy, @@ -204,10 +232,11 @@ def calculate_simulation_frames( else nullcontext() ) with build_span: - model = _build_simulation( + model = build_stage12_simulation( params, dataset=dataset, dataset_selection=dataset_selection, + country_module=country_module, policy=policy, scoping_strategy=region.scoping_strategy, region_code=region.code, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index cddfa82dd..51732662e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -561,7 +561,7 @@ def build_simulation(params, *, dataset, dataset_selection, **kwargs): resolve_dataset, ) monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) - monkeypatch.setattr(simulation_runtime, "_build_simulation", build_simulation) + monkeypatch.setattr(worker, "build_stage12_simulation", build_simulation) result = worker.calculate_simulation_frames( _planned_simulation(SimulationRole.BASELINE) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py new file mode 100644 index 000000000..ef4159530 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py @@ -0,0 +1,170 @@ +"""Unit tests for Stage 12 segment calculation and merging.""" + +from __future__ import annotations + +from hashlib import sha256 + +import pandas as pd +import pytest + +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_artifacts import ( + serialize_simulation_frames, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + Stage12SegmentResult, + build_segment_inputs, + calculate_segment, + merge_segment_results, + should_segment_simulation, +) +from test_stage12_runtime import _planned_simulation + + +def _frames(index: int) -> dict[str, pd.DataFrame]: + return { + "household": pd.DataFrame( + { + "household_id": pd.Series([index], dtype="int64"), + "household_net_income": pd.Series([index + 0.5], dtype="float32"), + } + ), + "person": pd.DataFrame( + { + "age": pd.Series([30 + index], dtype="int16"), + "household_id": pd.Series([index], dtype="int64"), + "person_id": pd.Series([100 + index], dtype="int64"), + } + ), + } + + +def _results(simulation): + results = [] + for segment in build_segment_inputs(simulation): + payload, _ = serialize_simulation_frames(_frames(segment.segment_index)) + results.append( + Stage12SegmentResult( + segment_index=segment.segment_index, + role=simulation.role, + region_codes=segment.region_codes, + parquet_payload=payload, + payload_sha256=sha256(payload).hexdigest(), + ) + ) + return results + + +def test_eligibility_matches_stage12_supported_national_shape() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + assert should_segment_simulation(simulation) is True + assert ( + should_segment_simulation( + simulation.model_copy(update={"options": {"segmented": False}}) + ) + is False + ) + assert ( + should_segment_simulation( + simulation.model_copy( + update={ + "geography": simulation.geography.model_copy( + update={"region": "state/ca"} + ) + } + ) + ) + is False + ) + assert ( + should_segment_simulation( + simulation.model_copy( + update={ + "output_plan": simulation.output_plan.model_copy( + update={ + "requirements": simulation.output_plan.requirements.model_copy( + update={"include_cliff_impacts": True} + ) + } + ) + } + ) + ) + is False + ) + + +def test_builds_twenty_single_policy_segment_inputs() -> None: + simulation = _planned_simulation(SimulationRole.REFORM) + + segments = build_segment_inputs(simulation) + + assert len(segments) == 20 + assert {segment.segment_index for segment in segments} == set(range(20)) + assert all(segment.simulation.role is SimulationRole.REFORM for segment in segments) + assert all(segment.simulation.policy == simulation.policy for segment in segments) + + +def test_segment_calculation_scopes_one_group_and_preserves_dtypes() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + segment = build_segment_inputs(simulation)[0] + received = [] + + def calculator(scoped): + received.append(scoped) + return SimulationCalculation(frames=_frames(0)) + + result = calculate_segment(segment, calculator=calculator) + merged = merge_segment_results( + simulation, + [ + calculate_segment( + item, + calculator=lambda scoped, index=item.segment_index: ( + SimulationCalculation(frames=_frames(index)) + ), + ) + for item in build_segment_inputs(simulation) + ], + ) + + assert received[0].options["region_group"] == list(segment.region_codes) + assert "segmented" not in received[0].options + assert result.role is SimulationRole.BASELINE + assert str(merged.frames["household"]["household_net_income"].dtype) == "float32" + assert str(merged.frames["person"]["age"].dtype) == "int16" + assert merged.frames["household"]["household_id"].tolist() == list(range(20)) + + +def test_merge_rejects_missing_duplicate_and_cross_role_results() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + results = _results(simulation) + + with pytest.raises(ValueError, match="incomplete"): + merge_segment_results(simulation, results[:-1]) + with pytest.raises(ValueError, match="duplicate index"): + merge_segment_results(simulation, [*results[:-1], results[0]]) + with pytest.raises(ValueError, match="another simulation role"): + merge_segment_results( + simulation, + [ + results[0].model_copy(update={"role": SimulationRole.REFORM}), + *results[1:], + ], + ) + + +def test_merge_rejects_duplicate_entity_ids_across_groups() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + results = _results(simulation) + duplicate_payload, _ = serialize_simulation_frames(_frames(0)) + results[1] = results[1].model_copy( + update={ + "parquet_payload": duplicate_payload, + "payload_sha256": sha256(duplicate_payload).hexdigest(), + } + ) + + with pytest.raises(ValueError, match="duplicate IDs"): + merge_segment_results(simulation, results) From cdf774babd44fdbe06cf5c9eeb0dd33fa93d23a1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 21:55:36 +0400 Subject: [PATCH 03/25] Add Stage 12 cache identity and loading --- .../stage12_cache/__init__.py | 27 ++ .../stage12_cache/keys.py | 177 ++++++++++++ .../stage12_cache/runtime.py | 115 ++++++++ .../stage12_runtime/simulation.py | 29 +- .../tests/test_stage12_cache_runtime.py | 262 ++++++++++++++++++ 5 files changed, 605 insertions(+), 5 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py new file mode 100644 index 000000000..b68b75e70 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py @@ -0,0 +1,27 @@ +"""Independent cache identities and runtime loading for Stage 12.""" + +from .keys import ( + Stage12BaselineCacheIdentity, + Stage12DatasetCacheIdentity, + collect_baseline_identity, + collect_dataset_identity, +) +from .runtime import ( + CACHE_HIT, + CACHE_INCOMPLETE, + CACHE_MISS, + Stage12CachedSimulation, + qualifying_cache_identity, +) + +__all__ = [ + "CACHE_HIT", + "CACHE_INCOMPLETE", + "CACHE_MISS", + "Stage12BaselineCacheIdentity", + "Stage12CachedSimulation", + "Stage12DatasetCacheIdentity", + "collect_baseline_identity", + "collect_dataset_identity", + "qualifying_cache_identity", +] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py new file mode 100644 index 000000000..682263be4 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py @@ -0,0 +1,177 @@ +"""Content-derived Stage 12 cache keys and object paths.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from hashlib import sha256 +import json +from typing import Any + +from policyengine_simulation_contract.stage12_bundle import ( + CountryId, + ResolvedStage12Bundle, +) +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + stage12_output_plan_sha256, +) + +DATASET_SCHEMA = "stage12-ds1" +BASELINE_SCHEMA = "stage12-bl1" +CURRENT_LAW_POLICY = "current-law" +_ID_DIGEST_CHARS = 16 + + +def canonical_digest(payload: Mapping[str, Any]) -> str: + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + return sha256(encoded).hexdigest() + + +@dataclass(frozen=True) +class Stage12DatasetCacheIdentity: + country: CountryId + dataset: str + stem: str + year: int + source_sha256: str + data_artifact_revision: str + data_package_version: str + country_package_version: str + core_package_version: str + policyengine_version: str + bundle_manifest_sha256: str + + @property + def payload(self) -> dict[str, Any]: + return { + "schema": DATASET_SCHEMA, + "country": self.country, + "dataset": self.dataset, + "year": self.year, + "source_sha256": self.source_sha256, + "data_artifact_revision": self.data_artifact_revision, + "data_package_version": self.data_package_version, + "country_package_version": self.country_package_version, + "core_package_version": self.core_package_version, + "policyengine_version": self.policyengine_version, + "bundle_manifest_sha256": self.bundle_manifest_sha256, + } + + @property + def digest(self) -> str: + return canonical_digest(self.payload) + + @property + def filename(self) -> str: + return f"{self.stem}_year_{self.year}.h5" + + @property + def store_path(self) -> str: + return f"datasets/{self.country}/{self.digest}/{self.filename}" + + +@dataclass(frozen=True) +class Stage12BaselineCacheIdentity: + dataset: Stage12DatasetCacheIdentity + region_codes: tuple[str, ...] + scope_key: str + output_plan_sha256: str + optional_configuration: Any = None + + @property + def payload(self) -> dict[str, Any]: + return { + "schema": BASELINE_SCHEMA, + "country": self.dataset.country, + "dataset_digest": self.dataset.digest, + "region_codes": list(self.region_codes), + "scope_key": self.scope_key, + "output_plan_sha256": self.output_plan_sha256, + "optional_configuration": self.optional_configuration, + "policy": CURRENT_LAW_POLICY, + } + + @property + def digest(self) -> str: + return canonical_digest(self.payload) + + @property + def simulation_id(self) -> str: + return f"{BASELINE_SCHEMA}-{self.digest[:_ID_DIGEST_CHARS]}" + + @property + def filename(self) -> str: + return f"{self.simulation_id}.h5" + + @property + def store_path(self) -> str: + return f"baselines/{self.dataset.country}/{self.digest}/{self.filename}" + + +def collect_dataset_identity( + resolved: ResolvedStage12Bundle, + *, + country: CountryId, + year: int, +) -> Stage12DatasetCacheIdentity: + """Collect dataset identity only from the installed Stage 12 bundle.""" + + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + country_bundle = next( + item for item in resolved.bundle.countries if item.country == country + ) + dataset = next( + item + for item in country_bundle.datasets + if item.identity == country_bundle.default_dataset + ) + stem = dataset_logical_name( + resolve_dataset_reference(country, country_bundle.default_dataset) + ) + return Stage12DatasetCacheIdentity( + country=country, + dataset=dataset.identity, + stem=stem, + year=year, + source_sha256=dataset.sha256, + data_artifact_revision=dataset.artifact_revision, + data_package_version=country_bundle.data_package_version, + country_package_version=country_bundle.country_package_version, + core_package_version=resolved.bundle.core_package_version, + policyengine_version=resolved.bundle.policyengine_version, + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + ) + + +def collect_baseline_identity( + simulation: PlannedSimulationExecutionInput, + *, + resolved: ResolvedStage12Bundle, + region_codes: tuple[str, ...], + scope_key: str, +) -> Stage12BaselineCacheIdentity: + """Build the identity shared by the Stage 12 writer and reader.""" + + optional_configuration = simulation.options.get("spm") + return Stage12BaselineCacheIdentity( + dataset=collect_dataset_identity( + resolved, + country=simulation.geography.country, + year=simulation.year, + ), + region_codes=region_codes, + scope_key=scope_key, + output_plan_sha256=stage12_output_plan_sha256(simulation.output_plan), + optional_configuration=optional_configuration, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py new file mode 100644 index 000000000..3151d744d --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py @@ -0,0 +1,115 @@ +"""Validated Stage 12 cache loading with calculation fallback.""" + +from __future__ import annotations + +import logging +from typing import Any + +from policyengine.core import Simulation +from policyengine.core.scoping_strategy import RegionGroupStrategy +from policyengine_simulation_contract.stage12_bundle import ResolvedStage12Bundle +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + SimulationRole, + Stage12OutputPlan, +) +from pydantic import PrivateAttr + +from policyengine_simulation_executor.stage12_runtime.output_planning import ( + validate_output_frames, +) + +from .keys import Stage12BaselineCacheIdentity, collect_baseline_identity + +logger = logging.getLogger(__name__) + +CACHE_HIT = "hit" +CACHE_INCOMPLETE = "incomplete" +CACHE_MISS = "miss" + + +def qualifying_cache_identity( + simulation: PlannedSimulationExecutionInput, + *, + resolved: ResolvedStage12Bundle, + dataset_is_default: bool, + scoping_strategy: Any, +) -> Stage12BaselineCacheIdentity | None: + """Return a safe cache identity, or None for an ineligible simulation.""" + + if simulation.role is not SimulationRole.BASELINE: + return None + if simulation.policy: + return None + if not dataset_is_default: + return None + if not isinstance(scoping_strategy, RegionGroupStrategy): + return None + raw_group = simulation.options.get("region_group") + if not isinstance(raw_group, list) or not raw_group: + return None + if any( + not isinstance(code, str) or not code.startswith("state/") for code in raw_group + ): + return None + scope_key = getattr(scoping_strategy, "cache_key", None) + if not isinstance(scope_key, str) or not scope_key: + return None + return collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=tuple(raw_group), + scope_key=scope_key, + ) + + +class Stage12CachedSimulation(Simulation): + """A deterministic baseline that validates any loaded Stage 12 cache.""" + + _stage12_output_plan: Stage12OutputPlan | None = PrivateAttr(default=None) + _computed_this_process: bool = PrivateAttr(default=False) + _stage12_cache_outcome: str | None = PrivateAttr(default=None) + + @property + def stage12_cache_outcome(self) -> str | None: + return self._stage12_cache_outcome + + def configure_stage12_cache(self, output_plan: Stage12OutputPlan) -> None: + self._stage12_output_plan = output_plan + + def run(self) -> None: + self._computed_this_process = True + super().run() + + def _record_outcome(self, outcome: str) -> None: + if self._stage12_cache_outcome is None: + self._stage12_cache_outcome = outcome + + def ensure(self) -> None: + self._computed_this_process = False + super().ensure() + if self._computed_this_process: + self._record_outcome(CACHE_MISS) + return + plan = self._stage12_output_plan + if plan is None: + raise RuntimeError("Stage 12 cached simulation has no output plan") + data = getattr(getattr(self, "output_dataset", None), "data", None) + entity_data = getattr(data, "entity_data", None) + try: + validate_output_frames(entity_data or {}, plan) + except (TypeError, ValueError): + logger.warning( + "Stage 12 cache %s is incomplete; recalculating", + self.id, + ) + self._record_outcome(CACHE_INCOMPLETE) + self.run() + self.save() + from policyengine.core.simulation import _cache + + storage_id = getattr(self, "storage_id", self.id) + _cache._cache.pop(storage_id, None) + _cache.add(storage_id, self) + return + self._record_outcome(CACHE_HIT) 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 157eada98..82a209c23 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 @@ -59,24 +59,39 @@ def build_stage12_simulation( policy: dict[str, Any] | None, scoping_strategy: Any, region_code: str | None, + execution: PlannedSimulationExecutionInput, ): """Construct a Stage 12 simulation without selecting v1 cache code.""" from policyengine.core import Simulation from policyengine_simulation_executor.spm import normalize_runtime_spm - # These values become part of Stage 12's independent cache eligibility - # and identity in the next layer; retaining them in this constructor keeps - # the monolithic and segmented call sites identical. - del dataset_selection, region_code selection = normalize_runtime_spm(params) - return Simulation( + from policyengine_simulation_executor.stage12_cache.runtime import ( + Stage12CachedSimulation, + qualifying_cache_identity, + ) + + resolved = load_stage12_bundle() + identity = qualifying_cache_identity( + execution, + resolved=resolved, + dataset_is_default=dataset_selection.is_default, + scoping_strategy=scoping_strategy, + ) + simulation_type = Stage12CachedSimulation if identity is not None else Simulation + model = simulation_type( **({"spm": selection} if selection is not None else {}), + **({"id": identity.simulation_id} if identity is not None else {}), dataset=dataset, tax_benefit_model_version=country_module.model, policy=policy, scoping_strategy=scoping_strategy, ) + if isinstance(model, Stage12CachedSimulation): + model.configure_stage12_cache(execution.output_plan) + del region_code + return model def simulation_input_sha256(simulation: SimulationExecutionInput) -> str: @@ -240,6 +255,7 @@ def calculate_simulation_frames( policy=policy, scoping_strategy=region.scoping_strategy, region_code=region.code, + execution=simulation, ) calculation_span = ( runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_CALCULATION)) @@ -249,6 +265,9 @@ def calculate_simulation_frames( with calculation_span: apply_output_plan(model, simulation.output_plan) model.ensure() + cache_outcome = getattr(model, "stage12_cache_outcome", None) + if runtime is not None and cache_outcome is not None: + runtime.set_context(stage12_cache_outcome=cache_outcome) output_data = getattr(getattr(model, "output_dataset", None), "data", None) entity_data = getattr(output_data, "entity_data", None) if not isinstance(entity_data, Mapping): diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py new file mode 100644 index 000000000..6d8df6654 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py @@ -0,0 +1,262 @@ +"""Unit tests for independent Stage 12 cache identities and loading.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +from policyengine.core import Simulation + +from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, +) +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_cache.keys import ( + collect_baseline_identity, + collect_dataset_identity, +) +from policyengine_simulation_executor.stage12_cache.runtime import ( + CACHE_HIT, + CACHE_INCOMPLETE, + CACHE_MISS, + Stage12CachedSimulation, + qualifying_cache_identity, +) +from test_stage12_runtime import _planned_simulation + + +def _resolved() -> ResolvedStage12Bundle: + dataset = Stage12Dataset( + identity="populace_us_2024", + uri="hf://policyengine/data/populace_us_2024.h5@revision", + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + country = Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.0.0", + country_package_requirement="policyengine-us==1.0.0", + data_package_name="policyengine-us-data", + data_package_version="2.0.0", + data_release_version="2.0.0", + data_artifact_revision="revision", + default_dataset=dataset.identity, + default_dataset_uri=dataset.uri, + datasets=(dataset,), + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=(country,), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _strategy(): + from policyengine.core.scoping_strategy import ( + RegionGroupStrategy, + RowFilterStrategy, + ) + + return RegionGroupStrategy( + members=[RowFilterStrategy(variable_name="state_code", variable_value="CA")] + ) + + +def test_cache_writer_and_reader_share_full_identity() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"region_group": ["state/ca"]}} + ) + resolved = _resolved() + strategy = _strategy() + + reader = qualifying_cache_identity( + simulation, + resolved=resolved, + dataset_is_default=True, + scoping_strategy=strategy, + ) + writer = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ca",), + scope_key=strategy.cache_key, + ) + + assert reader == writer + assert reader is not None + assert reader.simulation_id.startswith("stage12-bl1-") + assert reader.store_path.endswith(f"/{reader.simulation_id}.h5") + + +def test_dataset_and_baseline_inputs_rotate_cache_identity() -> None: + resolved = _resolved() + first = collect_dataset_identity(resolved, country="us", year=2026) + later = collect_dataset_identity(resolved, country="us", year=2027) + assert first.digest != later.digest + + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"region_group": ["state/ca"]}} + ) + ca = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ca",), + scope_key="state/ca", + ) + ny = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ny",), + scope_key="state/ny", + ) + first_entity = simulation.output_plan.entities[0] + changed_plan = collect_baseline_identity( + simulation.model_copy( + update={ + "output_plan": simulation.output_plan.model_copy( + update={ + "entities": ( + first_entity.model_copy( + update={ + "materialized_variables": ( + *first_entity.materialized_variables, + "new_cache_variable", + ) + } + ), + *simulation.output_plan.entities[1:], + ) + } + ) + } + ), + resolved=resolved, + region_codes=("state/ca",), + scope_key="state/ca", + ) + assert ca.digest != ny.digest + assert changed_plan.digest != ca.digest + + +def test_only_current_law_default_dataset_baseline_qualifies() -> None: + baseline = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"region_group": ["state/ca"]}} + ) + kwargs = { + "resolved": _resolved(), + "dataset_is_default": True, + "scoping_strategy": _strategy(), + } + assert qualifying_cache_identity(baseline, **kwargs) is not None + assert ( + qualifying_cache_identity( + baseline.model_copy(update={"role": SimulationRole.REFORM}), **kwargs + ) + is None + ) + assert ( + qualifying_cache_identity( + baseline.model_copy(update={"policy": {"reform": {"2026": 1}}}), + **kwargs, + ) + is None + ) + assert ( + qualifying_cache_identity(baseline, **{**kwargs, "dataset_is_default": False}) + is None + ) + + +def _cached_model() -> Stage12CachedSimulation: + model = Stage12CachedSimulation.model_construct( + id="stage12-bl1-test", + output_dataset=None, + dataset=SimpleNamespace(year=2026), + tax_benefit_model_version=None, + policy=None, + scoping_strategy=None, + extra_variables={}, + dynamic=None, + ) + model.configure_stage12_cache( + _planned_simulation(SimulationRole.BASELINE).output_plan + ) + return model + + +def _output(complete: bool = True): + household = { + "household_id": pd.Series([1], dtype="int64"), + "household_net_income": pd.Series([1.0], dtype="float32"), + } + if not complete: + household.pop("household_net_income") + return SimpleNamespace( + data=SimpleNamespace( + entity_data={ + "household": pd.DataFrame(household), + "person": pd.DataFrame( + { + "age": [30], + "household_id": [1], + "person_id": [1], + } + ), + } + ) + ) + + +def test_cached_simulation_reports_hit_and_miss(monkeypatch) -> None: + hit = _cached_model() + monkeypatch.setattr( + Simulation, + "ensure", + lambda self: setattr(self, "output_dataset", _output()), + ) + hit.ensure() + assert hit.stage12_cache_outcome == CACHE_HIT + + miss = _cached_model() + + def ensure_by_running(self): + self.run() + + monkeypatch.setattr(Simulation, "ensure", ensure_by_running) + monkeypatch.setattr( + Simulation, + "run", + lambda self: setattr(self, "output_dataset", _output()), + ) + miss.ensure() + assert miss.stage12_cache_outcome == CACHE_MISS + + +def test_incomplete_cache_recalculates(monkeypatch) -> None: + model = _cached_model() + monkeypatch.setattr( + Simulation, + "ensure", + lambda self: setattr(self, "output_dataset", _output(complete=False)), + ) + monkeypatch.setattr( + Simulation, + "run", + lambda self: setattr(self, "output_dataset", _output()), + ) + monkeypatch.setattr(Simulation, "save", lambda self: None) + + model.ensure() + + assert model.stage12_cache_outcome == CACHE_INCOMPLETE + assert "household_net_income" in model.output_dataset.data.entity_data["household"] From 8a5ae6d24a701a5e2209e3b2c30f34dceab80b86 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 21:59:15 +0400 Subject: [PATCH 04/25] Add Stage 12 cache precompute app --- .../src/modal/stage12_precompute_app.py | 165 ++++++++ .../stage12_cache/models.py | 112 +++++ .../stage12_cache/precompute.py | 393 ++++++++++++++++++ .../stage12_cache/store.py | 164 ++++++++ .../tests/test_stage12_cache_precompute.py | 276 ++++++++++++ .../tests/test_stage12_precompute_app.py | 50 +++ 6 files changed, 1160 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py new file mode 100644 index 000000000..3e40bd5ad --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py @@ -0,0 +1,165 @@ +# pyright: reportAttributeAccessIssue=false +"""Ephemeral Modal application that builds the independent Stage 12 cache.""" + +from __future__ import annotations + +import shlex +from pathlib import Path + +import modal + +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle + +STAGE12_DATA_DIR = "/opt/policyengine/stage12-data" +_UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." +RESOLVED_BUNDLE = load_stage12_bundle() + +app = modal.App("policyengine-stage12-cache-precompute") +gcp_secret = modal.Secret.from_name("stage12-evaluation-gcp-credentials") +data_secret = modal.Secret.from_name("policyengine-data-credentials") +hf_secret = modal.Secret.from_name("huggingface-token", required_keys=["HF_TOKEN"]) +worker_secrets = [gcp_secret, data_secret, hf_secret] + + +def _bundle_install_command() -> str: + parts = [ + "uvx", + "--from", + RESOLVED_BUNDLE.bundle.policyengine_requirement, + "policyengine", + "bundle", + "install", + RESOLVED_BUNDLE.bundle.policyengine_version, + "--venv", + "/.uv/.venv", + "--country", + "us", + "--data-dir", + STAGE12_DATA_DIR, + "--yes", + ] + return " ".join(shlex.quote(part) for part in parts) + + +precompute_image = ( + modal.Image.debian_slim(python_version="3.13") + .uv_sync( + uv_project_dir=_UV_PROJECT_DIR, + frozen=True, + extra_options="--only-group modal-simulation-image", + ) + .run_commands(_bundle_install_command(), secrets=[data_secret, hf_secret]) + .env({"POLICYENGINE_DATA_FOLDER": STAGE12_DATA_DIR}) + .add_local_python_source( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + copy=True, + ) +) + + +@app.function( + image=precompute_image, + cpu=4.0, + memory=16384, + timeout=1800, + secrets=worker_secrets, +) +def plan_artifacts(bucket: str) -> dict: + from policyengine_simulation_executor.stage12_cache.precompute import ( + plan_artifacts_impl, + ) + + return plan_artifacts_impl(bucket).model_dump() + + +@app.function( + image=precompute_image, + cpu=8.0, + memory=65536, + timeout=7200, + secrets=worker_secrets, +) +def build_dataset(bucket: str, expected: dict) -> dict: + from policyengine_simulation_executor.stage12_cache.models import DatasetPlanEntry + from policyengine_simulation_executor.stage12_cache.precompute import ( + build_dataset_impl, + ) + + return build_dataset_impl( + bucket, DatasetPlanEntry.model_validate(expected) + ).model_dump() + + +@app.function( + image=precompute_image, + cpu=8.0, + memory=32768, + timeout=3600, + secrets=worker_secrets, +) +def compute_baseline(bucket: str, expected: dict) -> dict: + from policyengine_simulation_executor.stage12_cache.models import BaselinePlanEntry + from policyengine_simulation_executor.stage12_cache.precompute import ( + compute_baseline_impl, + ) + + return compute_baseline_impl( + bucket, BaselinePlanEntry.model_validate(expected) + ).model_dump() + + +@app.function( + image=precompute_image, + cpu=8.0, + memory=32768, + timeout=3600, + secrets=worker_secrets, +) +def verify_determinism(bucket: str, expected: dict) -> dict: + from policyengine_simulation_executor.stage12_cache.models import BaselinePlanEntry + from policyengine_simulation_executor.stage12_cache.precompute import ( + verify_determinism_impl, + ) + + return verify_determinism_impl( + bucket, BaselinePlanEntry.model_validate(expected) + ).model_dump() + + +@app.function( + image=precompute_image, + cpu=2.0, + memory=8192, + timeout=600, + secrets=worker_secrets, +) +def publish_manifest(bucket: str, plan: dict) -> str: + from policyengine_simulation_executor.stage12_cache.models import PrecomputePlan + from policyengine_simulation_executor.stage12_cache.precompute import ( + publish_manifest_impl, + ) + + return publish_manifest_impl(bucket, PrecomputePlan.model_validate(plan)) + + +@app.local_entrypoint() +def main(force: bool = False): + from policyengine_simulation_executor.stage12_cache.precompute import ( + run_precompute, + ) + from policyengine_simulation_executor.stage12_cache.store import ( + resolve_cache_bucket, + ) + + run_precompute( + resolve_cache_bucket(), + force=force, + plan_artifacts=plan_artifacts, + build_dataset=build_dataset, + compute_baseline=compute_baseline, + verify_determinism=verify_determinism, + publish_manifest=publish_manifest, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py new file mode 100644 index 000000000..47732c3a7 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -0,0 +1,112 @@ +"""Strict models crossing Stage 12 precompute function boundaries.""" + +from __future__ import annotations + +from typing import Annotated, Literal, Protocol, Any + +from pydantic import BaseModel, ConfigDict, Field + + +class StrictCacheModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class CacheBundleIdentity(StrictCacheModel): + bundle_manifest_sha256: str + policyengine_version: str + core_package_version: str + country_package_version: str + data_package_version: str + data_artifact_revision: str + default_dataset: str + + +class DatasetPlanEntry(StrictCacheModel): + year: int + digest: str + path: str + filename: str + exists: bool + + +class BaselinePlanEntry(StrictCacheModel): + year: int + segment_index: Annotated[int, Field(ge=0, lt=20)] + region_codes: Annotated[tuple[str, ...], Field(min_length=1)] + digest: str + path: str + simulation_id: str + output_plan_sha256: str + exists: bool + + +class PrecomputePlan(StrictCacheModel): + schema_version: Literal[1] = 1 + years: tuple[int, ...] + datasets: tuple[DatasetPlanEntry, ...] + baselines: tuple[BaselinePlanEntry, ...] + bundle: CacheBundleIdentity + partition_sha256: str + + +class WorkSelection(StrictCacheModel): + datasets: tuple[DatasetPlanEntry, ...] + baselines: tuple[BaselinePlanEntry, ...] + + +class CacheObject(StrictCacheModel): + path: str + content_sha256: str + size_bytes: Annotated[int, Field(ge=0)] + + +class DatasetBuildResult(CacheObject): + year: int + uploaded: bool + build_seconds: float + + +class BaselineBuildResult(CacheObject): + year: int + segment_index: int + simulation_id: str + cache_outcome: Literal["hit", "incomplete", "miss"] | None + uploaded: bool + compute_seconds: float + + +class DeterminismVerdict(StrictCacheModel): + equal: bool + differences: tuple[str, ...] + + +class ManifestArtifact(StrictCacheModel): + type: Literal["dataset", "baseline"] + path: str + filename: str + year: int + identity_digest: str + content_sha256: str + size_bytes: int + + +class CacheManifest(StrictCacheModel): + manifest_schema: Literal["stage12-mf1"] = Field( + default="stage12-mf1", + alias="schema", + serialization_alias="schema", + ) + country: Literal["us"] = "us" + years: tuple[int, ...] + partition_sha256: str + bundle: CacheBundleIdentity + artifacts: tuple[ManifestArtifact, ...] + + def canonical_payload(self) -> dict[str, Any]: + return self.model_dump(mode="json", by_alias=True) + + +class RemoteFunction(Protocol): + def remote(self, *args: Any) -> Any: ... + + def spawn(self, *args: Any) -> Any: ... diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py new file mode 100644 index 000000000..9cbc39263 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -0,0 +1,393 @@ +"""Build and publish the independent Stage 12 US cache.""" + +from __future__ import annotations + +from pathlib import Path +import time +from typing import Callable +from uuid import NAMESPACE_URL, uuid5 + +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + ReportAggregate, + ReportExecutionInput, + SimulationExecutionInput, + SimulationRole, + stage12_output_plan_sha256, +) + +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.stage12_runtime.output_planning import ( + plan_simulation_input, + resolve_report_output_plan, +) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) + +from .keys import ( + canonical_digest, + collect_baseline_identity, + collect_dataset_identity, +) +from .models import ( + BaselineBuildResult, + BaselinePlanEntry, + CacheBundleIdentity, + CacheManifest, + DatasetBuildResult, + DatasetPlanEntry, + DeterminismVerdict, + ManifestArtifact, + PrecomputePlan, + RemoteFunction, + WorkSelection, +) +from .store import Stage12CacheStore + +PRECOMPUTE_YEARS = (2026, 2027, 2025) +MANIFEST_DIGEST_PREFIX = "STAGE12_CACHE_MANIFEST_DIGEST=" + + +def manifest_digest_line(digest: str) -> str: + return f"{MANIFEST_DIGEST_PREFIX}{digest}" + + +def _bundle_identity() -> CacheBundleIdentity: + resolved = load_stage12_bundle() + country = next(item for item in resolved.bundle.countries if item.country == "us") + return CacheBundleIdentity( + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + policyengine_version=resolved.bundle.policyengine_version, + core_package_version=resolved.bundle.core_package_version, + country_package_version=country.country_package_version, + data_package_version=country.data_package_version, + data_artifact_revision=country.data_artifact_revision, + default_dataset=country.default_dataset, + ) + + +def _planned_pair(year: int): + resolved = load_stage12_bundle() + country = next(item for item in resolved.bundle.countries if item.country == "us") + dataset = next( + item for item in country.datasets if item.identity == country.default_dataset + ) + provenance = BundleProvenance( + policyengine_version=resolved.bundle.policyengine_version, + country_package_name=country.country_package_name, + country_package_version=country.country_package_version, + dataset=DatasetProvenance( + identity=dataset.identity, + uri=dataset.uri, + artifact_revision=dataset.artifact_revision, + data_package_name=country.data_package_name, + data_package_version=country.data_package_version, + ), + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + ) + population = DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256=dataset.sha256, + ) + ) + evaluation_id = uuid5(NAMESPACE_URL, f"stage12-cache:{year}") + + def simulation(role: SimulationRole) -> SimulationExecutionInput: + return SimulationExecutionInput( + evaluation_id=evaluation_id, + simulation_execution_id=uuid5(evaluation_id, role.value), + role=role, + policy={}, + population=population, + year=year, + geography=GeographySelection(country="us", region="us"), + bundle=provenance, + ) + + report = ReportExecutionInput( + evaluation_id=evaluation_id, + baseline=simulation(SimulationRole.BASELINE), + reform=simulation(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), + ) + plan = resolve_report_output_plan(report) + return ( + plan_simulation_input(report.baseline, plan), + plan_simulation_input(report.reform, plan), + ) + + +def _scoped_planned_simulation( + year: int, group: tuple[str, ...], *, reform: bool = False +): + baseline, reform_input = _planned_pair(year) + selected = reform_input if reform else baseline + return selected.model_copy(update={"options": {"region_group": list(group)}}) + + +def _scope_key(year: int, group: tuple[str, ...]) -> str: + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + _resolve_region, + ) + + params = { + "country": "us", + "scope": "macro", + "time_period": str(year), + "region_group": list(group), + } + resolution = _resolve_region( + country_module=_country_module("us"), + country="us", + params=params, + ) + value = getattr(resolution.scoping_strategy, "cache_key", None) + if not isinstance(value, str) or not value: + raise RuntimeError("Stage 12 region group produced no cache identity") + return value + + +def plan_artifacts_impl(bucket: str) -> PrecomputePlan: + resolved = load_stage12_bundle() + store = Stage12CacheStore(bucket) + datasets = [] + baselines = [] + for year in PRECOMPUTE_YEARS: + dataset = collect_dataset_identity(resolved, country="us", year=year) + datasets.append( + DatasetPlanEntry( + year=year, + digest=dataset.digest, + path=dataset.store_path, + filename=dataset.filename, + exists=store.exists(dataset.store_path), + ) + ) + for index, group in enumerate(US_REGION_GROUPS): + simulation = _scoped_planned_simulation(year, group) + identity = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=group, + scope_key=_scope_key(year, group), + ) + baselines.append( + BaselinePlanEntry( + year=year, + segment_index=index, + region_codes=group, + digest=identity.digest, + path=identity.store_path, + simulation_id=identity.simulation_id, + output_plan_sha256=stage12_output_plan_sha256( + simulation.output_plan + ), + exists=store.exists(identity.store_path), + ) + ) + return PrecomputePlan( + years=PRECOMPUTE_YEARS, + datasets=tuple(datasets), + baselines=tuple(baselines), + bundle=_bundle_identity(), + partition_sha256=canonical_digest( + {"groups": [list(group) for group in US_REGION_GROUPS]} + ), + ) + + +def select_work(plan: PrecomputePlan, *, force: bool) -> WorkSelection: + return WorkSelection( + datasets=tuple(entry for entry in plan.datasets if force or not entry.exists), + baselines=tuple(entry for entry in plan.baselines if force or not entry.exists), + ) + + +def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + resolve_data_folder, + ) + + resolved = load_stage12_bundle() + identity = collect_dataset_identity(resolved, country="us", year=expected.year) + if identity.store_path != expected.path: + raise RuntimeError("planned and runtime Stage 12 dataset identities differ") + started = time.monotonic() + _country_module("us").ensure_datasets( + datasets=[identity.dataset], + years=[expected.year], + data_folder=resolve_data_folder(), + ) + local_path = Path(resolve_data_folder()) / identity.filename + uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) + return DatasetBuildResult( + **stored.model_dump(), + year=expected.year, + uploaded=uploaded, + build_seconds=round(time.monotonic() - started, 1), + ) + + +def _ensure_dataset_file(bucket: str, year: int) -> None: + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + + identity = collect_dataset_identity(load_stage12_bundle(), country="us", year=year) + local_path = Path(resolve_data_folder()) / identity.filename + if not local_path.exists(): + Stage12CacheStore(bucket).download_file(identity.store_path, local_path) + + +def compute_baseline_impl( + bucket: str, expected: BaselinePlanEntry +) -> BaselineBuildResult: + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + from policyengine_simulation_executor.stage12_runtime.simulation import ( + calculate_simulation_frames, + ) + + group = tuple(expected.region_codes) + simulation = _scoped_planned_simulation(expected.year, group) + identity = collect_baseline_identity( + simulation, + resolved=load_stage12_bundle(), + region_codes=group, + scope_key=_scope_key(expected.year, group), + ) + if identity.store_path != expected.path: + raise RuntimeError("planned and runtime Stage 12 baseline identities differ") + _ensure_dataset_file(bucket, expected.year) + started = time.monotonic() + calculate_simulation_frames(simulation) + local_path = Path(resolve_data_folder()) / identity.filename + if not local_path.exists(): + raise RuntimeError("Stage 12 baseline calculation produced no cache file") + uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) + return BaselineBuildResult( + **stored.model_dump(), + year=expected.year, + segment_index=expected.segment_index, + simulation_id=expected.simulation_id, + cache_outcome="miss", + uploaded=uploaded, + compute_seconds=round(time.monotonic() - started, 1), + ) + + +def verify_determinism_impl( + bucket: str, expected: BaselinePlanEntry +) -> DeterminismVerdict: + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + from policyengine_simulation_executor.stage12_runtime.simulation import ( + calculate_simulation_frames, + ) + + group = tuple(expected.region_codes) + baseline = _scoped_planned_simulation(expected.year, group) + reform = _scoped_planned_simulation(expected.year, group, reform=True) + _ensure_dataset_file(bucket, expected.year) + local_path = Path(resolve_data_folder()) / f"{expected.simulation_id}.h5" + if not local_path.exists(): + Stage12CacheStore(bucket).download_file(expected.path, local_path) + cached = calculate_simulation_frames(baseline).frames + fresh = calculate_simulation_frames(reform).frames + differences = [] + for entity in sorted(set(cached) | set(fresh)): + left = cached.get(entity) + right = fresh.get(entity) + if left is None or right is None: + differences.append(f"{entity}: missing on one side") + continue + for column in sorted(set(left) | set(right)): + if column not in left or column not in right: + differences.append(f"{entity}.{column}: only on one side") + elif str(left[column].dtype) != str(right[column].dtype): + differences.append(f"{entity}.{column}: dtype differs") + elif not left[column].equals(right[column]): + differences.append(f"{entity}.{column}: values differ") + return DeterminismVerdict( + equal=not differences, + differences=tuple(differences[:50]), + ) + + +def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManifest: + artifacts = [] + for artifact_type, entries in ( + ("dataset", plan.datasets), + ("baseline", plan.baselines), + ): + for entry in entries: + stored = store.describe(entry.path) + artifacts.append( + ManifestArtifact( + type=artifact_type, + path=entry.path, + filename=entry.filename + if artifact_type == "dataset" + else f"{entry.simulation_id}.h5", + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, + ) + ) + return CacheManifest( + years=plan.years, + partition_sha256=plan.partition_sha256, + bundle=plan.bundle, + artifacts=tuple(artifacts), + ) + + +def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: + store = Stage12CacheStore(bucket) + return store.write_manifest(build_manifest(plan, store).canonical_payload()) + + +def run_precompute( + bucket: str, + *, + force: bool, + plan_artifacts: RemoteFunction, + build_dataset: RemoteFunction, + compute_baseline: RemoteFunction, + verify_determinism: RemoteFunction, + publish_manifest: RemoteFunction, + echo: Callable[[str], None] = print, +) -> str: + plan = PrecomputePlan.model_validate(plan_artifacts.remote(bucket)) + work = select_work(plan, force=force) + dataset_calls = [ + build_dataset.spawn(bucket, entry.model_dump()) for entry in work.datasets + ] + for call in dataset_calls: + DatasetBuildResult.model_validate(call.get()) + baseline_calls = [ + compute_baseline.spawn(bucket, entry.model_dump()) for entry in work.baselines + ] + for call in baseline_calls: + BaselineBuildResult.model_validate(call.get()) + if work.baselines: + verdict = DeterminismVerdict.model_validate( + verify_determinism.remote(bucket, work.baselines[0].model_dump()) + ) + if not verdict.equal: + raise SystemExit( + "Stage 12 cache determinism validation failed: " + + "; ".join(verdict.differences) + ) + digest = publish_manifest.remote(bucket, plan.model_dump()) + if not isinstance(digest, str) or not digest: + raise RuntimeError("Stage 12 cache manifest publisher returned no digest") + echo(manifest_digest_line(digest)) + return digest diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py new file mode 100644 index 000000000..f9932bb15 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py @@ -0,0 +1,164 @@ +"""Dedicated write-once object store for Stage 12 cache files.""" + +from __future__ import annotations + +from collections.abc import Mapping +from hashlib import sha256 +import json +import os +from pathlib import Path +from typing import Any + +from google.api_core.exceptions import NotFound, PreconditionFailed + +from .keys import canonical_digest +from .models import CacheObject + +CACHE_BUCKET_ENV = "STAGE12_CACHE_BUCKET" + + +def _file_sha256(path: Path) -> str: + digest = sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def resolve_cache_bucket(explicit: str | None = None) -> str: + bucket = explicit or os.environ.get(CACHE_BUCKET_ENV) + if not bucket: + raise RuntimeError( + f"Stage 12 cache bucket is not configured; set {CACHE_BUCKET_ENV}" + ) + return bucket + + +def _default_client(): + from google.cloud import storage + from policyengine_simulation_executor.simulation_runtime import ( + setup_gcp_credentials, + ) + + with setup_gcp_credentials(): + return storage.Client() + + +class Stage12CacheStore: + def __init__(self, bucket_name: str | None = None, *, client: Any = None): + self.bucket_name = resolve_cache_bucket(bucket_name) + self._client = client + + @property + def client(self): + if self._client is None: + self._client = _default_client() + return self._client + + def _blob(self, path: str): + return self.client.bucket(self.bucket_name).blob(path) + + def exists(self, path: str) -> bool: + return bool(self._blob(path).exists()) + + def upload_file( + self, path: str, local_path: str | Path + ) -> tuple[bool, CacheObject]: + source = Path(local_path) + payload_digest = _file_sha256(source) + size_bytes = source.stat().st_size + blob = self._blob(path) + blob.metadata = { + "sha256": payload_digest, + "size_bytes": str(size_bytes), + } + try: + blob.upload_from_filename(str(source), if_generation_match=0) + uploaded = True + except PreconditionFailed: + uploaded = False + stored = self.describe(path) + if stored.content_sha256 != payload_digest or stored.size_bytes != size_bytes: + raise RuntimeError("immutable Stage 12 cache path contains other data") + return uploaded, stored + + def upload_bytes( + self, + path: str, + payload: bytes, + *, + content_type: str, + ) -> tuple[bool, CacheObject]: + digest = sha256(payload).hexdigest() + blob = self._blob(path) + blob.metadata = {"sha256": digest, "size_bytes": str(len(payload))} + try: + blob.upload_from_string( + payload, + content_type=content_type, + if_generation_match=0, + ) + uploaded = True + except PreconditionFailed: + uploaded = False + stored = self.describe(path) + if stored.content_sha256 != digest or stored.size_bytes != len(payload): + raise RuntimeError("immutable Stage 12 cache path contains other data") + return uploaded, stored + + def describe(self, path: str) -> CacheObject: + blob = self._blob(path) + try: + blob.reload() + except NotFound: + raise FileNotFoundError(path) from None + metadata = blob.metadata or {} + digest = metadata.get("sha256") + raw_size = metadata.get("size_bytes") + if not isinstance(digest, str) or not isinstance(raw_size, str): + payload = blob.download_as_bytes() + digest = sha256(payload).hexdigest() + raw_size = str(len(payload)) + return CacheObject( + path=path, + content_sha256=digest, + size_bytes=int(raw_size), + ) + + def download_file(self, path: str, local_path: str | Path) -> CacheObject: + destination = Path(local_path) + destination.parent.mkdir(parents=True, exist_ok=True) + partial = destination.with_name(destination.name + ".partial") + try: + self._blob(path).download_to_filename(str(partial)) + stored = self.describe(path) + if _file_sha256(partial) != stored.content_sha256: + raise RuntimeError("downloaded Stage 12 cache digest differs") + partial.replace(destination) + return stored + finally: + partial.unlink(missing_ok=True) + + def write_manifest(self, payload: Mapping[str, Any]) -> str: + digest = canonical_digest(payload) + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + self.upload_bytes( + f"manifests/{digest}.json", + encoded, + content_type="application/json", + ) + return digest + + def read_manifest(self, digest: str) -> Mapping[str, Any] | None: + try: + payload = self._blob(f"manifests/{digest}.json").download_as_bytes() + except NotFound: + return None + value = json.loads(payload) + return value if isinstance(value, Mapping) else None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py new file mode 100644 index 000000000..a2578cde0 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -0,0 +1,276 @@ +"""Unit tests for Stage 12 cache planning and orchestration.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from policyengine_simulation_executor.stage12_cache import precompute +from policyengine_simulation_executor.stage12_cache.models import ( + BaselinePlanEntry, + CacheBundleIdentity, + CacheObject, + DatasetPlanEntry, + PrecomputePlan, +) +from policyengine_simulation_executor.stage12_runtime.partition import US_REGION_GROUPS + + +def _bundle() -> CacheBundleIdentity: + return CacheBundleIdentity( + bundle_manifest_sha256="a" * 64, + policyengine_version="5.2.0", + core_package_version="3.0.0", + country_package_version="1.0.0", + data_package_version="2.0.0", + data_artifact_revision="revision", + default_dataset="populace_us_2024", + ) + + +def _plan(*, exists: bool = False) -> PrecomputePlan: + datasets = tuple( + DatasetPlanEntry( + year=year, + digest=f"dataset-{year}", + path=f"datasets/us/dataset-{year}/populace_year_{year}.h5", + filename=f"populace_year_{year}.h5", + exists=exists, + ) + for year in precompute.PRECOMPUTE_YEARS + ) + baselines = tuple( + BaselinePlanEntry( + year=year, + segment_index=index, + region_codes=group, + digest=f"baseline-{year}-{index}", + path=f"baselines/us/baseline-{year}-{index}/stage12-bl1-{index}.h5", + simulation_id=f"stage12-bl1-{index}", + output_plan_sha256="b" * 64, + exists=exists, + ) + for year in precompute.PRECOMPUTE_YEARS + for index, group in enumerate(US_REGION_GROUPS) + ) + return PrecomputePlan( + years=precompute.PRECOMPUTE_YEARS, + datasets=datasets, + baselines=baselines, + bundle=_bundle(), + partition_sha256="c" * 64, + ) + + +def test_plan_shape_matches_v1_year_window_and_twenty_groups() -> None: + plan = _plan() + + assert plan.years == (2026, 2027, 2025) + assert len(plan.datasets) == 3 + assert len(plan.baselines) == 60 + assert [entry.year for entry in plan.datasets] == [2026, 2027, 2025] + assert { + entry.segment_index for entry in plan.baselines if entry.year == 2026 + } == set(range(20)) + + +def test_strict_plan_rejects_unknown_modal_boundary_fields() -> None: + payload = _plan().model_dump() + payload["unexpected"] = True + + with pytest.raises(ValidationError): + PrecomputePlan.model_validate(payload) + + +def test_work_selection_is_idempotent_and_force_selects_everything() -> None: + cold = precompute.select_work(_plan(), force=False) + warm = precompute.select_work(_plan(exists=True), force=False) + forced = precompute.select_work(_plan(exists=True), force=True) + + assert len(cold.datasets) == 3 and len(cold.baselines) == 60 + assert warm.datasets == () and warm.baselines == () + assert len(forced.datasets) == 3 and len(forced.baselines) == 60 + + +class _Call: + def __init__(self, result, events, label): + self.result = result + self.events = events + self.label = label + + def get(self): + self.events.append(f"get:{self.label}") + return self.result + + +class _Function: + def __init__(self, callback, events, label): + self.callback = callback + self.events = events + self.label = label + self.spawn_calls = [] + self.remote_calls = [] + + def spawn(self, *args): + self.spawn_calls.append(args) + self.events.append(f"spawn:{self.label}") + return _Call(self.callback(*args), self.events, self.label) + + def remote(self, *args): + self.remote_calls.append(args) + self.events.append(f"remote:{self.label}") + return self.callback(*args) + + +def test_precompute_finishes_dataset_wave_before_baseline_wave() -> None: + events = [] + plan = _plan() + planner = _Function(lambda _: plan.model_dump(), events, "plan") + dataset = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "d" * 64, + "size_bytes": 1, + "year": entry["year"], + "uploaded": True, + "build_seconds": 1.0, + }, + events, + "dataset", + ) + baseline = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "e" * 64, + "size_bytes": 1, + "year": entry["year"], + "segment_index": entry["segment_index"], + "simulation_id": entry["simulation_id"], + "cache_outcome": "miss", + "uploaded": True, + "compute_seconds": 1.0, + }, + events, + "baseline", + ) + verify = _Function(lambda *_: {"equal": True, "differences": []}, events, "verify") + publish = _Function(lambda *_: "manifest-digest", events, "publish") + lines = [] + + digest = precompute.run_precompute( + "cache-bucket", + force=False, + plan_artifacts=planner, + build_dataset=dataset, + compute_baseline=baseline, + verify_determinism=verify, + publish_manifest=publish, + echo=lines.append, + ) + + first_baseline_spawn = events.index("spawn:baseline") + assert all( + index < first_baseline_spawn + for index, event in enumerate(events) + if event == "get:dataset" + ) + assert len(dataset.spawn_calls) == 3 + assert len(baseline.spawn_calls) == 60 + assert len(verify.remote_calls) == 1 + assert digest == "manifest-digest" + assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" + + +def test_warm_precompute_skips_builds_and_determinism_validation() -> None: + events = [] + plan = _plan(exists=True) + planner = _Function(lambda _: plan.model_dump(), events, "plan") + unused = _Function(lambda *_: pytest.fail("unexpected build"), events, "build") + verify = _Function( + lambda *_: pytest.fail("unexpected validation"), events, "verify" + ) + publish = _Function(lambda *_: "manifest-digest", events, "publish") + + precompute.run_precompute( + "cache-bucket", + force=False, + plan_artifacts=planner, + build_dataset=unused, + compute_baseline=unused, + verify_determinism=verify, + publish_manifest=publish, + ) + + assert unused.spawn_calls == [] + assert verify.remote_calls == [] + + +def test_failed_determinism_validation_prevents_manifest_publication() -> None: + events = [] + plan = _plan() + planner = _Function(lambda _: plan.model_dump(), events, "plan") + dataset = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "d" * 64, + "size_bytes": 1, + "year": entry["year"], + "uploaded": True, + "build_seconds": 1.0, + }, + events, + "dataset", + ) + baseline = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "e" * 64, + "size_bytes": 1, + "year": entry["year"], + "segment_index": entry["segment_index"], + "simulation_id": entry["simulation_id"], + "cache_outcome": "miss", + "uploaded": True, + "compute_seconds": 1.0, + }, + events, + "baseline", + ) + verify = _Function( + lambda *_: {"equal": False, "differences": ["person.age differs"]}, + events, + "verify", + ) + publish = _Function(lambda *_: "unexpected", events, "publish") + + with pytest.raises(SystemExit, match="person.age differs"): + precompute.run_precompute( + "cache-bucket", + force=False, + plan_artifacts=planner, + build_dataset=dataset, + compute_baseline=baseline, + verify_determinism=verify, + publish_manifest=publish, + ) + + assert publish.remote_calls == [] + + +def test_manifest_records_content_hashes_and_sizes() -> None: + plan = _plan(exists=True) + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, + content_sha256="f" * 64, + size_bytes=123, + ) + ) + + manifest = precompute.build_manifest(plan, store) + + assert len(manifest.artifacts) == 63 + assert all(item.content_sha256 == "f" * 64 for item in manifest.artifacts) + assert all(item.size_bytes == 123 for item in manifest.artifacts) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py new file mode 100644 index 000000000..3cb2930b6 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -0,0 +1,50 @@ +"""Declaration tests for the independent Stage 12 precompute app.""" + +from __future__ import annotations + +import importlib +import sys + +from fixtures.fake_modal import install_fake_modal + + +def _load(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.stage12_precompute_app", None) + return importlib.import_module("src.modal.stage12_precompute_app") + + +def test_precompute_app_declares_independent_worker_resources(monkeypatch) -> None: + module = _load(monkeypatch) + functions = dict(module.app.function_calls) + + assert set(functions) == { + "plan_artifacts", + "build_dataset", + "compute_baseline", + "verify_determinism", + "publish_manifest", + } + assert functions["build_dataset"]["cpu"] == 8.0 + assert functions["build_dataset"]["memory"] == 65536 + assert functions["build_dataset"]["timeout"] == 7200 + assert functions["compute_baseline"]["cpu"] == 8.0 + assert functions["compute_baseline"]["memory"] == 32768 + assert functions["compute_baseline"]["timeout"] == 3600 + assert module.app.name == "policyengine-stage12-cache-precompute" + + +def test_precompute_image_installs_only_the_us_bundle(monkeypatch) -> None: + module = _load(monkeypatch) + command = next( + call for call in module.precompute_image.calls if call[0] == "run_commands" + )[1][0] + + assert "--country us" in command + assert "--country uk" not in command + local_source = next( + call + for call in module.precompute_image.calls + if call[0] == "add_local_python_source" + ) + assert "policyengine_simulation_executor" in local_source[1] From 697d93936569066296245a6dc435c461f08a2cd2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 22:34:25 +0400 Subject: [PATCH 05/25] Bake Stage 12 caches into US workers --- .github/scripts/stage12-cache-precompute.sh | 29 +++++ .github/scripts/stage12-sync-modal-secrets.sh | 4 +- .../stage12-validate-infrastructure.sh | 23 +++- .../scripts/stage12-validate-v2-workers.sh | 9 +- .../workflows/simulation-deploy.reusable.yml | 16 ++- .../stage12_manifest.py | 13 +++ .../tests/test_stage12_manifest.py | 16 +++ .../tests/test_deploy_workflow_scripts.py | 28 ++++- .../tests/test_stage12_infrastructure.py | 11 +- .../src/modal/stage12_cache_image.py | 100 ++++++++++++++++ .../modal/utils/update_v2_version_manifest.py | 1 + .../modal/utils/validate_v2_worker_release.py | 20 +++- .../src/modal/v2_app.py | 60 ++++++++-- .../stage12_cache/precompute.py | 43 ++++--- .../stage12_cache/runtime.py | 3 +- .../stage12_runtime/segmentation.py | 1 + .../stage12_runtime/simulation.py | 2 + .../stage12_worker_validation.py | 54 +++++++++ .../tests/test_stage12_cache_image.py | 108 ++++++++++++++++++ .../tests/test_stage12_modal_app.py | 8 ++ .../tests/test_stage12_worker_validation.py | 30 +++++ .../tests/test_validate_v2_worker_release.py | 24 +++- 22 files changed, 556 insertions(+), 47 deletions(-) create mode 100755 .github/scripts/stage12-cache-precompute.sh create mode 100644 projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py diff --git a/.github/scripts/stage12-cache-precompute.sh b/.github/scripts/stage12-cache-precompute.sh new file mode 100755 index 000000000..8e3a88d19 --- /dev/null +++ b/.github/scripts/stage12-cache-precompute.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# Build the independent Stage 12 cache and export its manifest digest. + +set -euo pipefail + +modal_environment="${1:?Modal environment is required}" +force="${2:-false}" +: "${STAGE12_CACHE_BUCKET:?STAGE12_CACHE_BUCKET is required}" +: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" + +command=(uv run modal run --env="${modal_environment}" src/modal/stage12_precompute_app.py) +if [[ "${force}" == "true" || "${force}" == "1" ]]; then + command+=(--force) +fi + +output_file="$(mktemp)" +trap 'rm -f "${output_file}"' EXIT +"${command[@]}" 2>&1 | tee "${output_file}" +digest_line="$(tr -d '\r' <"${output_file}" | grep '^STAGE12_CACHE_MANIFEST_DIGEST=' | tail -n 1 || true)" +if [[ -z "${digest_line}" ]]; then + echo "Stage 12 precompute returned no cache manifest digest" >&2 + exit 1 +fi +digest="${digest_line#STAGE12_CACHE_MANIFEST_DIGEST=}" +if [[ ! "${digest}" =~ ^[0-9a-f]{64}$ ]]; then + echo "Stage 12 precompute returned an invalid cache manifest digest" >&2 + exit 1 +fi +echo "manifest_digest=${digest}" >>"${GITHUB_OUTPUT}" diff --git a/.github/scripts/stage12-sync-modal-secrets.sh b/.github/scripts/stage12-sync-modal-secrets.sh index c6c97b180..5954fd68a 100755 --- a/.github/scripts/stage12-sync-modal-secrets.sh +++ b/.github/scripts/stage12-sync-modal-secrets.sh @@ -9,6 +9,7 @@ modal_environment="${1:?Modal environment is required}" : "${STAGE12_DATABASE_URL_SECRET_NAME:?STAGE12_DATABASE_URL_SECRET_NAME is required}" : "${STAGE12_GCP_CREDENTIALS_SECRET_NAME:?STAGE12_GCP_CREDENTIALS_SECRET_NAME is required}" : "${STAGE12_ARTIFACT_BUCKET:?STAGE12_ARTIFACT_BUCKET is required}" +: "${STAGE12_CACHE_BUCKET:?STAGE12_CACHE_BUCKET is required}" : "${MODAL_TOKEN_ID:?MODAL_TOKEN_ID is required}" : "${MODAL_TOKEN_SECRET:?MODAL_TOKEN_SECRET is required}" @@ -34,7 +35,8 @@ gcloud secrets versions access latest \ jq -n \ --rawfile database_url "${database_url_file}" \ --arg artifact_bucket "${STAGE12_ARTIFACT_BUCKET}" \ - '{STAGE12_DATABASE_URL: ($database_url | rtrimstr("\n")), STAGE12_ARTIFACT_BUCKET: $artifact_bucket}' \ + --arg cache_bucket "${STAGE12_CACHE_BUCKET}" \ + '{STAGE12_DATABASE_URL: ($database_url | rtrimstr("\n")), STAGE12_ARTIFACT_BUCKET: $artifact_bucket, STAGE12_CACHE_BUCKET: $cache_bucket}' \ >"${runtime_modal_file}" jq -n \ --rawfile credentials "${credentials_file}" \ diff --git a/.github/scripts/stage12-validate-infrastructure.sh b/.github/scripts/stage12-validate-infrastructure.sh index d12daabad..a2acf8b6b 100755 --- a/.github/scripts/stage12-validate-infrastructure.sh +++ b/.github/scripts/stage12-validate-infrastructure.sh @@ -7,6 +7,7 @@ set +x : "${STAGE12_ENVIRONMENT:?STAGE12_ENVIRONMENT is required}" : "${STAGE12_GCP_PROJECT_ID:?STAGE12_GCP_PROJECT_ID is required}" : "${STAGE12_ARTIFACT_BUCKET:?STAGE12_ARTIFACT_BUCKET is required}" +: "${STAGE12_CACHE_BUCKET:?STAGE12_CACHE_BUCKET is required}" : "${STAGE12_ENTRYPOINT_SERVICE_ACCOUNT:?STAGE12_ENTRYPOINT_SERVICE_ACCOUNT is required}" : "${STAGE12_DEPLOY_SERVICE_ACCOUNT:?STAGE12_DEPLOY_SERVICE_ACCOUNT is required}" : "${STAGE12_MODAL_SERVICE_ACCOUNT:?STAGE12_MODAL_SERVICE_ACCOUNT is required}" @@ -28,6 +29,11 @@ if [[ ! "${STAGE12_ARTIFACT_BUCKET}" =~ ^[a-z0-9][a-z0-9._-]{2,61}[a-z0-9]$ ]] | echo "STAGE12_ARTIFACT_BUCKET must be valid and identify the environment" >&2 exit 1 fi +if [[ ! "${STAGE12_CACHE_BUCKET}" =~ ^[a-z0-9][a-z0-9._-]{2,61}[a-z0-9]$ ]] || + [[ ! "${STAGE12_CACHE_BUCKET}" =~ ${STAGE12_ENVIRONMENT} ]]; then + echo "STAGE12_CACHE_BUCKET must be valid and identify the environment" >&2 + exit 1 +fi for account in \ "${STAGE12_ENTRYPOINT_SERVICE_ACCOUNT}" \ "${STAGE12_DEPLOY_SERVICE_ACCOUNT}" \ @@ -68,6 +74,7 @@ canary_source_file="$(mktemp)" canary_download_file="$(mktemp)" runtime_gcloud_config="$(mktemp -d)" canary_object="" +cache_canary_object="" runtime_gcloud() { # The GitHub auth action exports credentials for the deployment identity. # Remove them only here so the activated Modal worker account controls this check. @@ -83,6 +90,10 @@ cleanup() { runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ "${canary_object}" --quiet >/dev/null 2>&1 || true fi + if [[ -n "${cache_canary_object}" ]]; then + runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ + "${cache_canary_object}" --quiet >/dev/null 2>&1 || true + fi rm -f \ "${database_url_file}" \ "${credentials_file}" \ @@ -164,6 +175,16 @@ runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ "${canary_object}" --quiet canary_object="" +cache_canary_object="gs://${STAGE12_CACHE_BUCKET}/_deployment-validation/${STAGE12_ENVIRONMENT}/${canary_id}.txt" +runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ + "${canary_source_file}" "${cache_canary_object}" --quiet +runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ + "${cache_canary_object}" "${canary_download_file}" --quiet +cmp "${canary_source_file}" "${canary_download_file}" +runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ + "${cache_canary_object}" --quiet +cache_canary_object="" + for lookup_object in "${uk_geography_lookup_objects[@]}"; do runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ "${lookup_object}" "${canary_download_file}" --quiet @@ -173,4 +194,4 @@ for lookup_object in "${uk_geography_lookup_objects[@]}"; do fi done -echo "Pre-provisioned Stage 12 database, secret, artifact storage, and UK geography lookup access is verified." +echo "Pre-provisioned Stage 12 database, secret, report storage, cache storage, and UK geography lookup access is verified." diff --git a/.github/scripts/stage12-validate-v2-workers.sh b/.github/scripts/stage12-validate-v2-workers.sh index ad0ed6e7b..df813a8be 100755 --- a/.github/scripts/stage12-validate-v2-workers.sh +++ b/.github/scripts/stage12-validate-v2-workers.sh @@ -1,7 +1,8 @@ #!/usr/bin/env bash -# Validate the installed bundles, secrets, datasets and calculations of the -# deployed v2 workers, and keep the result for manifest publication. -# Usage: stage12-validate-v2-workers.sh +# Validate the installed bundles, secrets, datasets, calculations and Stage 12 +# cache of the deployed v2 workers, and keep the result for manifest +# publication. +# Usage: stage12-validate-v2-workers.sh # Run from projects/policyengine-simulation-executor. # Required env vars: RUNNER_TEMP, GITHUB_OUTPUT. # Outputs: file (the validation result path) in GITHUB_OUTPUT. @@ -9,9 +10,11 @@ set -euo pipefail MODAL_ENV="${1:?Modal environment required}" +CACHE_MANIFEST_SHA256="${2:?Stage 12 cache manifest digest required}" output="${RUNNER_TEMP}/stage12-worker-validation.json" uv run python -m src.modal.utils.validate_v2_worker_release \ --environment "${MODAL_ENV}" \ + --cache-manifest-sha256 "${CACHE_MANIFEST_SHA256}" \ --output "${output}" echo "file=${output}" >> "$GITHUB_OUTPUT" diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 11d47de71..657e38277 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -547,6 +547,7 @@ jobs: STAGE12_ENVIRONMENT: ${{ inputs.deployment_environment }} STAGE12_GCP_PROJECT_ID: ${{ vars.SIMULATION_ENTRYPOINT_GCP_PROJECT_ID }} STAGE12_ARTIFACT_BUCKET: ${{ vars.STAGE12_ARTIFACT_BUCKET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} STAGE12_ENTRYPOINT_SERVICE_ACCOUNT: ${{ vars.SIMULATION_ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT }} STAGE12_DEPLOY_SERVICE_ACCOUNT: ${{ secrets.SIMULATION_ENTRYPOINT_GCP_DEPLOY_SERVICE_ACCOUNT }} STAGE12_MODAL_SERVICE_ACCOUNT: ${{ vars.STAGE12_MODAL_SERVICE_ACCOUNT }} @@ -568,6 +569,7 @@ jobs: outputs: application_name: ${{ steps.application.outputs.name }} validation_file: ${{ steps.validation.outputs.file }} + cache_manifest_digest: ${{ steps.stage12_cache.outputs.manifest_digest }} steps: - uses: actions/checkout@v6 @@ -603,11 +605,21 @@ jobs: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} STAGE12_ARTIFACT_BUCKET: ${{ vars.STAGE12_ARTIFACT_BUCKET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} STAGE12_GCP_PROJECT_ID: ${{ vars.SIMULATION_ENTRYPOINT_GCP_PROJECT_ID }} STAGE12_GCP_CREDENTIALS_SECRET_NAME: ${{ vars.STAGE12_GCP_CREDENTIALS_SECRET_NAME }} run: ../../.github/scripts/stage12-sync-modal-secrets.sh "${{ inputs.modal_environment }}" + - name: Build the independent Stage 12 cache + id: stage12_cache + working-directory: projects/policyengine-simulation-executor + env: + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} + run: ../../.github/scripts/stage12-cache-precompute.sh "${{ inputs.modal_environment }}" "${{ inputs.force_recompute }}" + - name: Deploy the additional versioned v2 Modal application working-directory: projects/policyengine-simulation-executor env: @@ -618,6 +630,8 @@ jobs: STAGE12_EXPECT_US_VERSION: ${{ needs.prepare.outputs.us_version }} STAGE12_EXPECT_UK_VERSION: ${{ needs.prepare.outputs.uk_version }} STAGE12_ARTIFACT_BUCKET: ${{ vars.STAGE12_ARTIFACT_BUCKET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} + STAGE12_CACHE_MANIFEST_DIGEST: ${{ steps.stage12_cache.outputs.manifest_digest }} OBSERVABILITY_SERVICE_NAMESPACE: ${{ vars.OBSERVABILITY_SERVICE_NAMESPACE }} OBSERVABILITY_TRACE_PROJECT_ID: ${{ vars.OBSERVABILITY_TRACE_PROJECT_ID }} OBSERVABILITY_LOGGING_PROJECT_ID: ${{ vars.OBSERVABILITY_LOGGING_PROJECT_ID }} @@ -634,7 +648,7 @@ jobs: env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} - run: ../../.github/scripts/stage12-validate-v2-workers.sh "${{ inputs.modal_environment }}" + run: ../../.github/scripts/stage12-validate-v2-workers.sh "${{ inputs.modal_environment }}" "${{ steps.stage12_cache.outputs.manifest_digest }}" - name: Retain worker validation for manifest publication uses: actions/upload-artifact@v7 diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 920aa1251..47873a82c 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -39,6 +39,15 @@ class V2CountryWorker(StrictManifestModel): single_simulation_callable: CallableName +class Stage12CacheRelease(StrictManifestModel): + manifest_sha256: Sha256Digest + partition_sha256: Sha256Digest + years: tuple[int, ...] + segment_count: Literal[20] = 20 + dataset_count: Literal[3] = 3 + baseline_count: Literal[60] = 60 + + class V2WorkerValidation(StrictManifestModel): validated: Literal[True] application_name: ManifestText @@ -46,6 +55,7 @@ class V2WorkerValidation(StrictManifestModel): validated_at: datetime validation_invocation_id: ManifestText country_validation_invocation_ids: dict[CountryId, ManifestText] + cache: Stage12CacheRelease | None = None class V2WorkerVersion(StrictManifestModel): @@ -55,6 +65,7 @@ class V2WorkerVersion(StrictManifestModel): bundle: Stage12BundleManifest bundle_manifest_sha256: Sha256Digest validation: V2WorkerValidation + cache: Stage12CacheRelease | None = None @model_validator(mode="after") def validate_worker(self) -> V2WorkerVersion: @@ -69,6 +80,8 @@ def validate_worker(self) -> V2WorkerVersion: raise ValueError("worker validation names another application") if self.validation.bundle_manifest_sha256 != self.bundle_manifest_sha256: raise ValueError("worker validation names another bundle digest") + if self.validation.cache != self.cache: + raise ValueError("worker cache metadata differs from worker validation") countries = tuple(worker.country for worker in self.countries) bundle_countries = tuple(country.country for country in self.bundle.countries) if len(countries) != len(set(countries)): diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 627a43f5b..944c1814c 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -16,6 +16,7 @@ ACTIVE_MANIFEST_KEY, V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, + Stage12CacheRelease, V2CountryWorker, V2ManifestLoader, V2WorkerValidation, @@ -178,6 +179,21 @@ def test_manifest_rejects_incomplete_or_unvalidated_worker() -> None: with pytest.raises(ValueError): V2WorkerVersion.model_validate(value) + +def test_cache_metadata_is_optional_for_old_workers_and_must_match_validation() -> None: + old = _worker() + assert old.cache is None + + cache = Stage12CacheRelease( + manifest_sha256="c" * 64, + partition_sha256="d" * 64, + years=(2026, 2027, 2025), + ) + value = old.model_dump(mode="json") + value["cache"] = cache.model_dump(mode="json") + with pytest.raises(ValueError, match="cache metadata differs"): + V2WorkerVersion.model_validate(value) + value = _worker().model_dump(mode="json") value["validation"]["validated"] = False with pytest.raises(ValueError): diff --git a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py index 2a3c502b6..5203ebc17 100644 --- a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py @@ -294,6 +294,7 @@ def test_candidate_fails_when_the_tag_is_not_deployed(fakes, tmp_path): V2_APPLICATION_NAME = "stage12-v2-application-name.sh" VALIDATE_V2_WORKERS = "stage12-validate-v2-workers.sh" +CACHE_MANIFEST_SHA256 = "a" * 64 STAGE12_READY = "stage12-deployment-ready.sh" V2_APPLICATION_CODE = ( "from policyengine_simulation_contract.stage12_manifest import " @@ -321,9 +322,9 @@ def test_deploy_workflow_runs_the_stage12_worker_scripts(): validation ) assert ( - f'run: ../../.github/scripts/{VALIDATE_V2_WORKERS} "${{{{ inputs.modal_environment }}}}"' - in validation - ) + f'run: ../../.github/scripts/{VALIDATE_V2_WORKERS} "${{{{ inputs.modal_environment }}}}" ' + '"${{ steps.stage12_cache.outputs.manifest_digest }}"\n' + ) in validation assert "application_name: ${{ steps.application.outputs.name }}" in workflow assert "validation_file: ${{ steps.validation.outputs.file }}" in workflow @@ -368,7 +369,9 @@ def _validate_v2_workers(fakes: FakeCommands, tmp_path: Path, *args: str): def test_validate_v2_workers_records_the_validation_file(fakes, tmp_path): fakes.add("uv") - result, output, validation_file = _validate_v2_workers(fakes, tmp_path, "main") + result, output, validation_file = _validate_v2_workers( + fakes, tmp_path, "main", CACHE_MANIFEST_SHA256 + ) assert result.returncode == 0, result.stderr assert output == f"file={validation_file}\n" @@ -381,6 +384,8 @@ def test_validate_v2_workers_records_the_validation_file(fakes, tmp_path): "src.modal.utils.validate_v2_worker_release", "--environment", "main", + "--cache-manifest-sha256", + CACHE_MANIFEST_SHA256, "--output", str(validation_file), ] @@ -392,7 +397,9 @@ def test_validate_v2_workers_fails_without_writing_when_validation_fails( ): fakes.add("uv", "exit 1\n") - result, output, _ = _validate_v2_workers(fakes, tmp_path, "main") + result, output, _ = _validate_v2_workers( + fakes, tmp_path, "main", CACHE_MANIFEST_SHA256 + ) assert result.returncode == 1 assert output == "" @@ -409,6 +416,17 @@ def test_validate_v2_workers_requires_the_modal_environment(fakes, tmp_path): assert fakes.calls() == [] +def test_validate_v2_workers_requires_the_cache_manifest_digest(fakes, tmp_path): + fakes.add("uv") + + result, output, _ = _validate_v2_workers(fakes, tmp_path, "main") + + assert result.returncode != 0 + assert "Stage 12 cache manifest digest required" in result.stderr + assert output == "" + assert fakes.calls() == [] + + def test_deploy_workflow_checks_stage12_results_with_the_tested_script(): workflow = REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8") step = _step(workflow, "Check Stage 12 deployment result") diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py index 4eb899072..fec6ba12e 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py @@ -29,9 +29,10 @@ def test_infrastructure_validation_is_bounded_and_valid_shell() -> None: source = SCRIPT.read_text(encoding="utf-8") assert "STAGE12_ENVIRONMENT" in source assert "STAGE12_ARTIFACT_BUCKET" in source + assert "STAGE12_CACHE_BUCKET" in source assert ( - "Pre-provisioned Stage 12 database, secret, artifact storage, and UK " - "geography lookup access" in source + "Pre-provisioned Stage 12 database, secret, report storage, cache storage, " + "and UK geography lookup access" in source ) assert "STAGE12_DATABASE_URL_SECRET_NAME" in source assert "policyengine_v2_runtime" in source @@ -60,6 +61,7 @@ def test_infrastructure_validation_is_bounded_and_valid_shell() -> None: assert "stage12-evaluation-runtime" in modal_sync assert "stage12-evaluation-gcp-credentials" in modal_sync assert "STAGE12_DATABASE_URL:" in modal_sync + assert "STAGE12_CACHE_BUCKET:" in modal_sync assert "GOOGLE_APPLICATION_CREDENTIALS_JSON:" in modal_sync @@ -90,8 +92,9 @@ def test_storage_validation_uses_only_the_modal_worker_credentials() -> None: assert "runtime_gcloud auth activate-service-account" in source modal_account = '--account="${STAGE12_MODAL_SERVICE_ACCOUNT}"' - assert source.count(f"runtime_gcloud {modal_account} storage cp") == 3 - assert source.count(f"runtime_gcloud {modal_account} storage rm") == 2 + assert source.count(f"runtime_gcloud {modal_account} storage cp") == 5 + assert source.count(f"runtime_gcloud {modal_account} storage rm") == 4 + assert "gs://${STAGE12_CACHE_BUCKET}/_deployment-validation" in source assert "gs://policyengine-uk-data-private/constituencies_2024.csv" in source assert "gs://policyengine-uk-data-private/local_authorities_2021.csv" not in source assert "STAGE12_UK_GEOGRAPHY_BUCKET" not in source diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py new file mode 100644 index 000000000..71c8b3218 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py @@ -0,0 +1,100 @@ +"""Self-contained Stage 12 cache fetch used during Modal image construction.""" + +from __future__ import annotations + + +def fetch_stage12_cache(bucket: str, manifest: dict | None, *, client=None) -> None: + """Download and verify every manifest-listed cache file into the image.""" + + from hashlib import sha256 + import json + import os + from pathlib import Path + + if not bucket or not manifest: + raise RuntimeError("Stage 12 cache deployment has no bucket or manifest") + if manifest.get("schema") != "stage12-mf1": + raise RuntimeError("Stage 12 cache manifest schema is unsupported") + bundle = manifest.get("bundle") + if not isinstance(bundle, dict) or bundle.get( + "bundle_manifest_sha256" + ) != os.getenv("STAGE12_BUNDLE_MANIFEST_SHA256"): + raise RuntimeError("Stage 12 cache manifest names another installed bundle") + if manifest.get("years") != [2026, 2027, 2025]: + raise RuntimeError("Stage 12 cache manifest has an unexpected year set") + artifacts = manifest.get("artifacts") + if not isinstance(artifacts, list): + raise RuntimeError("Stage 12 cache manifest has no artifact list") + type_counts = { + kind: sum( + item.get("type") == kind for item in artifacts if isinstance(item, dict) + ) + for kind in ("dataset", "baseline") + } + if type_counts != {"dataset": 3, "baseline": 60}: + raise RuntimeError("Stage 12 cache manifest has an unexpected artifact count") + canonical = json.dumps( + manifest, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + manifest_digest = sha256(canonical).hexdigest() + if manifest_digest != os.getenv("STAGE12_CACHE_MANIFEST_DIGEST"): + raise RuntimeError("Stage 12 cache manifest digest differs from deployment") + + if client is None: + credentials = {} + raw_credentials = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON") + if raw_credentials: + decoded = json.loads(raw_credentials) + info = json.loads(decoded) if isinstance(decoded, str) else decoded + from google.oauth2 import service_account + + credentials = { + "credentials": service_account.Credentials.from_service_account_info( + info + ), + "project": info.get("project_id"), + } + from google.cloud import storage + + client = storage.Client(**credentials) + + data_dir = Path(os.environ["POLICYENGINE_DATA_FOLDER"]) + data_dir.mkdir(parents=True, exist_ok=True) + filenames: set[str] = set() + for artifact in artifacts: + if not isinstance(artifact, dict): + raise RuntimeError("Stage 12 cache artifact entry is invalid") + filename = artifact.get("filename") + path = artifact.get("path") + expected_digest = artifact.get("content_sha256") + expected_size = artifact.get("size_bytes") + if ( + not isinstance(filename, str) + or Path(filename).name != filename + or filename in filenames + or not isinstance(path, str) + or not isinstance(expected_digest, str) + or not isinstance(expected_size, int) + ): + raise RuntimeError("Stage 12 cache artifact metadata is invalid") + filenames.add(filename) + destination = data_dir / filename + partial = destination.with_name(filename + ".partial") + try: + client.bucket(bucket).blob(path).download_to_filename(str(partial)) + digest = sha256() + with partial.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + if partial.stat().st_size != expected_size: + raise RuntimeError("Stage 12 cache artifact size differs") + if digest.hexdigest() != expected_digest: + raise RuntimeError("Stage 12 cache artifact digest differs") + partial.replace(destination) + finally: + partial.unlink(missing_ok=True) + (data_dir / ".stage12-cache-manifest.json").write_bytes(canonical) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py index 0d267bb90..97674a136 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py @@ -38,6 +38,7 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: bundle=resolved.bundle, bundle_manifest_sha256=resolved.bundle_manifest_sha256, validation=validation, + cache=validation.cache, ) validate_manifest_has_no_credentials(worker.model_dump(mode="json")) return worker diff --git a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py index 348076fd1..02eda2071 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py @@ -14,6 +14,7 @@ from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_manifest import ( ManifestText, + Stage12CacheRelease, V2WorkerValidation, v2_application_name, ) @@ -39,10 +40,13 @@ def _spawn_validation( return result, invocation_id -def validate_release(*, environment: str) -> V2WorkerValidation: +def validate_release( + *, environment: str, expected_cache_manifest_sha256: str +) -> V2WorkerValidation: resolved = load_stage12_bundle() application_name = v2_application_name(resolved.bundle.policyengine_version) invocation_ids: dict[CountryId, ManifestText] = {} + cache_release: Stage12CacheRelease | None = None for country_bundle in resolved.bundle.countries: function_name = f"validate_worker_{country_bundle.country}" result, invocation_id = _spawn_validation( @@ -71,6 +75,13 @@ def validate_release(*, environment: str) -> V2WorkerValidation: raise RuntimeError( f"{function_name} reported bundle values that differ from deployment" ) + if country_bundle.country == "us": + reported_cache = result.get("cache") + cache_release = Stage12CacheRelease.model_validate(reported_cache) + if cache_release.manifest_sha256 != expected_cache_manifest_sha256: + raise RuntimeError(f"{function_name} reported another cache manifest") + elif result.get("cache") is not None: + raise RuntimeError(f"{function_name} unexpectedly reported a cache") invocation_ids[country_bundle.country] = invocation_id return V2WorkerValidation( validated=True, @@ -79,6 +90,7 @@ def validate_release(*, environment: str) -> V2WorkerValidation: validated_at=datetime.now(timezone.utc), validation_invocation_id=f"stage12-validation-{uuid4()}", country_validation_invocation_ids=invocation_ids, + cache=cache_release, ) @@ -88,8 +100,12 @@ def main() -> None: ) parser.add_argument("--environment", required=True) parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--cache-manifest-sha256", required=True) args = parser.parse_args() - validation = validate_release(environment=args.environment) + validation = validate_release( + environment=args.environment, + expected_cache_manifest_sha256=args.cache_manifest_sha256, + ) args.output.write_text( json.dumps( validation.model_dump(mode="json"), diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 357670a2a..377595948 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -39,6 +39,7 @@ ) from src.modal.bundle_data import bundle_data_install_command from src.modal.static_runtime_files import add_static_runtime_files +from src.modal.stage12_cache_image import fetch_stage12_cache STAGE12_DATA_DIR = "/opt/policyengine/stage12-data" _UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." @@ -47,6 +48,30 @@ APP_NAME = v2_application_name(RESOLVED_BUNDLE.bundle.policyengine_version) +def _deploy_time_cache_inputs() -> tuple[str, dict | None]: + """Resolve and validate the cache manifest on the deploying machine.""" + + if not modal.is_local(): + return "", None + bucket = os.environ.get("STAGE12_CACHE_BUCKET", "") + digest = os.environ.get("STAGE12_CACHE_MANIFEST_DIGEST") + if not digest: + return bucket, None + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore + + payload = Stage12CacheStore(bucket).read_manifest(digest) + if payload is None: + raise RuntimeError("Stage 12 cache manifest is unavailable") + manifest = CacheManifest.model_validate(payload) + if manifest.bundle.bundle_manifest_sha256 != RESOLVED_BUNDLE.bundle_manifest_sha256: + raise RuntimeError("Stage 12 cache manifest names another bundle") + return bucket, manifest.canonical_payload() + + +STAGE12_CACHE_BUCKET, STAGE12_CACHE_MANIFEST = _deploy_time_cache_inputs() + + def _external_assertions(environment: dict[str, str]) -> dict[str, str]: mapping = { "STAGE12_EXPECT_POLICYENGINE_VERSION": "policyengine_version", @@ -90,7 +115,9 @@ def _country_bundle(country: CountryId): ) -def build_v2_image(countries: tuple[CountryId, ...]) -> modal.Image: +def build_v2_image( + countries: tuple[CountryId, ...], *, include_us_cache: bool = False +) -> modal.Image: country_values = { country: _country_bundle(country).model_dump(mode="json") for country in countries @@ -122,21 +149,33 @@ def build_v2_image(countries: tuple[CountryId, ...]) -> modal.Image: sort_keys=True, separators=(",", ":"), ), + "STAGE12_CACHE_MANIFEST_DIGEST": os.environ.get( + "STAGE12_CACHE_MANIFEST_DIGEST", "" + ), } ) - .add_local_python_source( - "src.modal", - "policyengine_simulation_executor", - "policyengine_simulation_observability", - "policyengine_simulation_contract", - "policyengine_stage12_persistence", - copy=True, + ) + if include_us_cache: + image = image.run_function( + fetch_stage12_cache, + args=(STAGE12_CACHE_BUCKET, STAGE12_CACHE_MANIFEST), + secrets=[gcp_secret], + cpu=2.0, + memory=4096, + timeout=1800, ) + image = image.add_local_python_source( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + copy=True, ) return add_static_runtime_files(image, uv_project_dir=_UV_PROJECT_DIR) -us_worker_image = build_v2_image(("us",)) +us_worker_image = build_v2_image(("us",), include_us_cache=True) uk_worker_image = build_v2_image(("uk",)) coordinator_image = build_v2_image(("us", "uk")) @@ -149,6 +188,9 @@ def _validate(country: CountryId) -> dict: return validate_country_worker( country=country, expected_bundle_manifest_sha256=os.environ["STAGE12_BUNDLE_MANIFEST_SHA256"], + expected_cache_manifest_sha256=( + os.environ.get("STAGE12_CACHE_MANIFEST_DIGEST") if country == "us" else None + ), ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index 9cbc39263..0961f3e9c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -322,25 +322,32 @@ def verify_determinism_impl( def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManifest: artifacts = [] - for artifact_type, entries in ( - ("dataset", plan.datasets), - ("baseline", plan.baselines), - ): - for entry in entries: - stored = store.describe(entry.path) - artifacts.append( - ManifestArtifact( - type=artifact_type, - path=entry.path, - filename=entry.filename - if artifact_type == "dataset" - else f"{entry.simulation_id}.h5", - year=entry.year, - identity_digest=entry.digest, - content_sha256=stored.content_sha256, - size_bytes=stored.size_bytes, - ) + for entry in plan.datasets: + stored = store.describe(entry.path) + artifacts.append( + ManifestArtifact( + type="dataset", + path=entry.path, + filename=entry.filename, + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, + ) + ) + for entry in plan.baselines: + stored = store.describe(entry.path) + artifacts.append( + ManifestArtifact( + type="baseline", + path=entry.path, + filename=f"{entry.simulation_id}.h5", + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, ) + ) return CacheManifest( years=plan.years, partition_sha256=plan.partition_sha256, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py index 3151d744d..3659da8f9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py @@ -52,13 +52,14 @@ def qualifying_cache_identity( not isinstance(code, str) or not code.startswith("state/") for code in raw_group ): return None + region_codes = tuple(code for code in raw_group if isinstance(code, str)) scope_key = getattr(scoping_strategy, "cache_key", None) if not isinstance(scope_key, str) or not scope_key: return None return collect_baseline_identity( simulation, resolved=resolved, - region_codes=tuple(raw_group), + region_codes=region_codes, scope_key=scope_key, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py index f3b1408da..8d4c61e1e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py @@ -185,4 +185,5 @@ def merge_segment_results( return SimulationCalculation( frames=merged, calculation_provenance=_merge_calculation_provenance(provenance), + cache_outcome=None, ) 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 82a209c23..3da356fd9 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 @@ -48,6 +48,7 @@ class SimulationCalculation: frames: Mapping[str, pd.DataFrame] calculation_provenance: SPMProvenance | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + cache_outcome: str | None = None def build_stage12_simulation( @@ -289,6 +290,7 @@ def calculate_simulation_frames( frames=frames, calculation_provenance=calculation_provenance, uk_local_authority_metadata=uk_local_authority_metadata, + cache_outcome=cache_outcome, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 008bfae40..a8fbdcff3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -25,6 +25,7 @@ ), ("STAGE12_DATABASE_URL",), ("STAGE12_ARTIFACT_BUCKET",), + ("STAGE12_CACHE_BUCKET",), ) @@ -95,10 +96,58 @@ def _check_uk_local_authority_dataset(dataset_path: str) -> None: detect_uk_local_authority_metadata_from_hdf(dataset_path) +def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: + """Validate the baked manifest and require one real baseline cache hit.""" + + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + from policyengine_simulation_executor.stage12_cache.precompute import ( + _scoped_planned_simulation, + ) + from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, + ) + from policyengine_simulation_executor.stage12_runtime.simulation import ( + calculate_simulation_frames, + ) + + data_folder = Path(resolve_data_folder()) + manifest_path = data_folder / ".stage12-cache-manifest.json" + payload = manifest_path.read_bytes() + if sha256(payload).hexdigest() != expected_manifest_sha256: + raise RuntimeError("installed Stage 12 cache manifest digest differs") + manifest = CacheManifest.model_validate_json(payload) + if manifest.years != (2026, 2027, 2025): + raise RuntimeError("installed Stage 12 cache years differ") + if len(manifest.artifacts) != 63: + raise RuntimeError("installed Stage 12 cache artifact count differs") + probe = next( + item + for item in manifest.artifacts + if item.type == "baseline" and item.year == 2026 + ) + probe_path = data_folder / probe.filename + _check_dataset_access(str(probe_path), probe.content_sha256) + calculation = calculate_simulation_frames( + _scoped_planned_simulation(2026, US_REGION_GROUPS[0]) + ) + if calculation.cache_outcome != "hit": + raise RuntimeError("deployed Stage 12 cache probe did not load a cache hit") + return { + "manifest_sha256": expected_manifest_sha256, + "partition_sha256": manifest.partition_sha256, + "years": list(manifest.years), + "segment_count": 20, + "dataset_count": sum(item.type == "dataset" for item in manifest.artifacts), + "baseline_count": sum(item.type == "baseline" for item in manifest.artifacts), + } + + def validate_country_worker( *, country: CountryId, expected_bundle_manifest_sha256: str, + expected_cache_manifest_sha256: str | None = None, environment: Mapping[str, str] | None = None, dataset_path_resolver: Callable[[CountryId], str | None] = ( _resolve_installed_dataset_path @@ -108,6 +157,7 @@ def validate_country_worker( local_authority_resource_check: Callable[[str], None] = ( _check_uk_local_authority_dataset ), + cache_check: Callable[[str], dict[str, Any]] = _check_installed_cache, ) -> dict[str, Any]: """Validate imports, installed bundle, secrets, data, and calculation.""" @@ -145,6 +195,9 @@ def validate_country_worker( if country == "uk": local_authority_resource_check(installed_dataset_path) calculation_check(country) + cache = None + if country == "us" and expected_cache_manifest_sha256: + cache = cache_check(expected_cache_manifest_sha256) return { "validated": True, "country": country, @@ -155,4 +208,5 @@ def validate_country_worker( "dataset_uri": country_bundle.default_dataset_uri, "data_artifact_revision": country_bundle.data_artifact_revision, "bundle_manifest_sha256": resolved.bundle_manifest_sha256, + "cache": cache, } diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py new file mode 100644 index 000000000..e0c433e01 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py @@ -0,0 +1,108 @@ +"""Tests for the self-contained Stage 12 image cache fetch.""" + +from __future__ import annotations + +from hashlib import sha256 +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from src.modal.stage12_cache_image import fetch_stage12_cache + + +class _Blob: + def __init__(self, payload: bytes): + self.payload = payload + + def download_to_filename(self, filename: str) -> None: + Path(filename).write_bytes(self.payload) + + +class _Client: + def __init__(self, payloads: dict[str, bytes]): + self.payloads = payloads + + def bucket(self, _name: str): + return SimpleNamespace(blob=lambda path: _Blob(self.payloads[path])) + + +def _manifest() -> tuple[dict, dict[str, bytes]]: + artifacts = [] + payloads = {} + for kind, count in (("dataset", 3), ("baseline", 60)): + for index in range(count): + filename = f"{kind}-{index}.h5" + path = f"{kind}s/us/{index}/{filename}" + payload = f"{kind}-{index}".encode() + payloads[path] = payload + artifacts.append( + { + "type": kind, + "path": path, + "filename": filename, + "year": (2026, 2027, 2025)[index % 3], + "identity_digest": f"identity-{kind}-{index}", + "content_sha256": sha256(payload).hexdigest(), + "size_bytes": len(payload), + } + ) + return ( + { + "schema": "stage12-mf1", + "country": "us", + "years": [2026, 2027, 2025], + "partition_sha256": "p" * 64, + "bundle": { + "bundle_manifest_sha256": "b" * 64, + "policyengine_version": "5.2.0", + "core_package_version": "3.0.0", + "country_package_version": "1.0.0", + "data_package_version": "2.0.0", + "data_artifact_revision": "revision", + "default_dataset": "populace_us_2024", + }, + "artifacts": artifacts, + }, + payloads, + ) + + +def _digest(manifest: dict) -> str: + return sha256( + json.dumps( + manifest, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode() + ).hexdigest() + + +def test_fetches_and_verifies_all_cache_files(monkeypatch, tmp_path) -> None: + manifest, payloads = _manifest() + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + monkeypatch.setenv("STAGE12_BUNDLE_MANIFEST_SHA256", "b" * 64) + monkeypatch.setenv("STAGE12_CACHE_MANIFEST_DIGEST", _digest(manifest)) + + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert len(list(tmp_path.glob("*.h5"))) == 63 + assert (tmp_path / ".stage12-cache-manifest.json").exists() + + +def test_rejects_foreign_or_corrupt_cache(monkeypatch, tmp_path) -> None: + manifest, payloads = _manifest() + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + monkeypatch.setenv("STAGE12_BUNDLE_MANIFEST_SHA256", "different") + monkeypatch.setenv("STAGE12_CACHE_MANIFEST_DIGEST", _digest(manifest)) + with pytest.raises(RuntimeError, match="another installed bundle"): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + monkeypatch.setenv("STAGE12_BUNDLE_MANIFEST_SHA256", "b" * 64) + corrupt_path = manifest["artifacts"][0]["path"] + payloads[corrupt_path] = b"corrupt" + with pytest.raises(RuntimeError, match="size differs|digest differs"): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) 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 ade681c12..3fcb76a7e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -182,6 +182,14 @@ def test_v2_static_runtime_file_smoke_uses_exact_uk_worker_image(monkeypatch) -> } +def test_only_us_worker_image_fetches_stage12_cache(monkeypatch) -> None: + module = _load(monkeypatch) + + assert any(call[0] == "run_function" for call in module.us_worker_image.calls) + assert not any(call[0] == "run_function" for call in module.uk_worker_image.calls) + assert not any(call[0] == "run_function" for call in module.coordinator_image.calls) + + def test_v2_app_declares_validation_workers_and_non_http_coordinator( monkeypatch, ) -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index a1abf01d9..4e284b6cd 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -20,6 +20,7 @@ def _environment() -> dict[str, str]: "GOOGLE_APPLICATION_CREDENTIALS_JSON": "{}", "STAGE12_DATABASE_URL": "postgresql://runtime:secret@db.example/postgres", "STAGE12_ARTIFACT_BUCKET": "policyengine-stage12-staging", + "STAGE12_CACHE_BUCKET": "policyengine-stage12-cache-staging", } @@ -35,6 +36,7 @@ def test_validation_checks_dataset_and_non_serving_calculation() -> None: dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", dataset_check=lambda path, digest: datasets.append((path, digest)), calculation_check=countries.append, + cache_check=lambda digest: {"manifest_sha256": digest}, ) assert result["validated"] is True @@ -49,6 +51,7 @@ def test_validation_checks_dataset_and_non_serving_calculation() -> None: ) assert datasets == [("/installed/populace_us_2024.h5", expected_dataset.sha256)] assert countries == ["us"] + assert result["cache"] is None def test_non_serving_us_calculation_supports_national_spm() -> None: @@ -108,6 +111,33 @@ def test_validation_rejects_digest_mismatch_before_dataset_access() -> None: assert accessed == [] +def test_us_validation_reports_the_verified_cache_release() -> None: + resolved = load_stage12_bundle() + seen: list[str] = [] + cache = { + "manifest_sha256": "c" * 64, + "partition_sha256": "d" * 64, + "years": [2026, 2027, 2025], + "segment_count": 20, + "dataset_count": 3, + "baseline_count": 60, + } + + result = validate_country_worker( + country="us", + expected_bundle_manifest_sha256=resolved.bundle_manifest_sha256, + expected_cache_manifest_sha256="c" * 64, + environment=_environment(), + dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", + dataset_check=lambda *_: None, + calculation_check=lambda _: None, + cache_check=lambda digest: seen.append(digest) or cache, + ) + + assert seen == ["c" * 64] + assert result["cache"] == cache + + @pytest.mark.parametrize( "environment", [ diff --git a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py index f3a521c9a..e03a7e7c3 100644 --- a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py @@ -29,13 +29,28 @@ def fake_spawn_validation(*, application_name, function_name, environment): "dataset_uri": bundle.default_dataset_uri, "data_artifact_revision": bundle.data_artifact_revision, "bundle_manifest_sha256": resolved.bundle_manifest_sha256, + "cache": ( + { + "manifest_sha256": "c" * 64, + "partition_sha256": "d" * 64, + "years": [2026, 2027, 2025], + "segment_count": 20, + "dataset_count": 3, + "baseline_count": 60, + } + if country == "us" + else None + ), }, invocation_id, ) monkeypatch.setattr(validation, "_spawn_validation", fake_spawn_validation) - result = validation.validate_release(environment="staging") + result = validation.validate_release( + environment="staging", + expected_cache_manifest_sha256="c" * 64, + ) assert result.validated is True assert result.country_validation_invocation_ids == { @@ -43,6 +58,8 @@ def fake_spawn_validation(*, application_name, function_name, environment): "uk": "call-uk", } assert invocation_ids == ["call-us", "call-uk"] + assert result.cache is not None + assert result.cache.manifest_sha256 == "c" * 64 def test_country_result_mismatch_stops_publication(monkeypatch) -> None: @@ -60,4 +77,7 @@ def test_country_result_mismatch_stops_publication(monkeypatch) -> None: ) with pytest.raises(RuntimeError, match="differ from deployment"): - validation.validate_release(environment="staging") + validation.validate_release( + environment="staging", + expected_cache_manifest_sha256="c" * 64, + ) From dd51dd9141b42b313945c5f6e56b433b957839f7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 29 Sep 2026 22:53:06 +0400 Subject: [PATCH 06/25] Fan out Stage 12 US simulations --- .../stages.py | 28 ++- .../stage12_adapter.py | 11 +- .../tests/test_stage12_adapter.py | 17 +- .../src/modal/v2_app.py | 58 +++++- .../stage12_runtime/__init__.py | 3 + .../stage12_runtime/fanout.py | 188 +++++++++++++++++ .../stage12_runtime/segmentation.py | 7 +- .../stage12_runtime/simulation.py | 36 +++- .../tests/test_stage12_modal_app.py | 27 ++- .../tests/test_stage12_runtime.py | 44 ++++ .../tests/test_stage12_segment_fanout.py | 193 ++++++++++++++++++ 11 files changed, 590 insertions(+), 22 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py index d923c3379..93f4e2627 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -7,10 +7,10 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from enum import StrEnum from types import MappingProxyType -from typing import Mapping class RunConfiguration(StrEnum): @@ -20,6 +20,7 @@ class RunConfiguration(StrEnum): STAGE12_SHADOW_REPORT = "stage12_shadow_report" STAGE12_CANONICAL_REPORT = "stage12_canonical_report" STAGE12_SIMULATION = "stage12_simulation" + STAGE12_SEGMENT = "stage12_segment" class Stage(StrEnum): @@ -91,6 +92,11 @@ class Stage(StrEnum): STAGE12_CHILD_WAIT = "stage12_child_wait" STAGE12_INPUT_WRITE = "stage12_input_write" STAGE12_SIMULATION_EXECUTION = "stage12_simulation_execution" + STAGE12_SEGMENT_EXECUTION = "stage12_segment_execution" + STAGE12_SEGMENT_DISPATCH = "stage12_segment_dispatch" + STAGE12_SEGMENT_WAIT = "stage12_segment_wait" + STAGE12_SEGMENT_MERGE = "stage12_segment_merge" + STAGE12_CACHE_LOOKUP = "stage12_cache_lookup" STAGE12_CALCULATION = "stage12_calculation" STAGE12_SIMULATION_ARTIFACT_WRITE = "stage12_simulation_artifact_write" STAGE12_ARTIFACT_READ = "stage12_artifact_read" @@ -226,6 +232,9 @@ def name(self, stage: Stage) -> str: ( Stage.STAGE12_SIMULATION_EXECUTION, Stage.STAGE12_INPUT_WRITE, + Stage.STAGE12_SEGMENT_DISPATCH, + Stage.STAGE12_SEGMENT_WAIT, + Stage.STAGE12_SEGMENT_MERGE, Stage.CREDENTIAL_SETUP, Stage.COUNTRY_MODULE_LOAD, Stage.REGION_RESOLUTION, @@ -233,10 +242,26 @@ def name(self, stage: Stage) -> str: Stage.DATASET_LOAD, Stage.POLICY_NORMALIZATION, Stage.SIMULATION_BUILD, + Stage.STAGE12_CACHE_LOOKUP, Stage.STAGE12_CALCULATION, Stage.STAGE12_SIMULATION_ARTIFACT_WRITE, ), ), + RunConfiguration.STAGE12_SEGMENT: StagePlan( + RunConfiguration.STAGE12_SEGMENT, + ( + Stage.STAGE12_SEGMENT_EXECUTION, + Stage.CREDENTIAL_SETUP, + Stage.COUNTRY_MODULE_LOAD, + Stage.REGION_RESOLUTION, + Stage.DATASET_RESOLUTION, + Stage.DATASET_LOAD, + Stage.POLICY_NORMALIZATION, + Stage.SIMULATION_BUILD, + Stage.STAGE12_CACHE_LOOKUP, + Stage.STAGE12_CALCULATION, + ), + ), } ) @@ -250,3 +275,4 @@ def name(self, stage: Stage) -> str: RunConfiguration.STAGE12_CANONICAL_REPORT ] STAGE12_SIMULATION_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SIMULATION] +STAGE12_SEGMENT_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SEGMENT] diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index c51e7aa2e..f0873ac4b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -99,8 +99,9 @@ def adapt_annual_comparison( include_cliffs = payload.get("include_cliffs") if include_cliffs is not None and not isinstance(include_cliffs, bool): return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) - if payload.get("segmented") not in {None, False}: - return _skip(ComparisonSkipReason.UNSUPPORTED_OPTIONS) + segmented = payload.get("segmented") + if segmented is not None and not isinstance(segmented, bool): + return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) country = payload.get("country") if not isinstance(country, str) or _country_worker(worker, country) is None: return _skip(ComparisonSkipReason.UNSUPPORTED_COUNTRY) @@ -175,7 +176,11 @@ def adapt_annual_comparison( country=cast(CountryId, country), region=region, ) - options = {key: payload[key] for key in ("spm",) if payload.get(key) is not None} + options = { + key: payload[key] + for key in ("spm", "segmented") + if payload.get(key) is not None + } if include_cliffs is True: options["include_cliffs"] = True diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index a8b2cce35..0c54b3bb5 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -107,6 +107,21 @@ def test_cliff_analysis_is_forwarded_to_both_simulations() -> None: assert result.report.reform.options == {"include_cliffs": True} +@pytest.mark.parametrize("segmented", [True, False]) +def test_segmentation_selection_is_forwarded_to_both_simulations(segmented) -> None: + payload = {**eligible_payload(), "segmented": segmented} + + result = adapt_annual_comparison( + payload, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.report is not None + assert result.report.baseline.options == {"segmented": segmented} + assert result.report.reform.options == {"segmented": segmented} + + @pytest.mark.parametrize( ("update", "reason"), [ @@ -124,7 +139,7 @@ def test_cliff_analysis_is_forwarded_to_both_simulations() -> None: ), ({"time_period": "2026-01"}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), ({"region": None}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), - ({"segmented": True}, ComparisonSkipReason.UNSUPPORTED_OPTIONS), + ({"segmented": "true"}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), ({"region_group": ["state/ca"]}, ComparisonSkipReason.UNSUPPORTED_OPTIONS), ], ) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 377595948..559a5df5d 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -19,6 +19,7 @@ ) from policyengine_simulation_observability.stages import ( STAGE12_CANONICAL_REPORT_STAGES, + STAGE12_SEGMENT_STAGES, STAGE12_SHADOW_REPORT_STAGES, STAGE12_SIMULATION_STAGES, Stage, @@ -222,7 +223,46 @@ def validate_worker_uk() -> dict: image=us_worker_image, cpu=8.0, memory=32768, - timeout=3000, + timeout=3600, + retries=0, + max_containers=300, + secrets=worker_secrets, +) +def run_single_simulation_segment_us( + payload: dict, + *, + observability_context: dict | None = None, +) -> dict: + from policyengine_simulation_executor.stage12_runtime import calculate_segment + + runtime = init_process_observability( + service_name="policyengine-stage12-us-segment-worker", + service_role="stage12_segment_worker", + platform="modal", + environment=os.getenv("MODAL_ENVIRONMENT", "local"), + ) + + try: + propagated = normalize_observability_context(observability_context) + with runtime.operation( + STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION), + attributes={ + "runner_name": "stage12", + "simulation_role": payload.get("simulation", {}).get("role"), + "segment_index": payload.get("segment_index"), + }, + remote_context=propagated, + ): + return calculate_segment(payload, runtime=runtime).model_dump(mode="python") + finally: + runtime.shutdown() + + +@app.function( + image=us_worker_image, + cpu=8.0, + memory=32768, + timeout=3900, retries=0, max_containers=10, secrets=worker_secrets, @@ -234,6 +274,7 @@ def run_single_simulation_us( observability_context: dict | None = None, ) -> dict: from policyengine_simulation_executor.stage12_runtime import ( + run_segmented_simulation, run_single_simulation, ) @@ -243,6 +284,14 @@ def run_single_simulation_us( platform="modal", environment=os.getenv("MODAL_ENVIRONMENT", "local"), ) + + def segmented_calculator(simulation): + return run_segmented_simulation( + simulation, + app_name=APP_NAME, + runtime=runtime, + ) + try: propagated = normalize_observability_context(observability_context) with runtime.operation( @@ -257,6 +306,7 @@ def run_single_simulation_us( payload, context, required_country="us", + segmented_calculator=segmented_calculator, runtime=runtime, ) finally: @@ -312,9 +362,9 @@ def run_single_simulation_uk( image=coordinator_image, cpu=2.0, memory=8192, - # Allow one 50-minute child-calculation window, the subsequent 15-minute - # production-result wait, and bounded aggregation/persistence overhead. - timeout=4500, + # Allow the two logical simulation parents to start and collect all region + # calculations, then retain bounded aggregation/persistence overhead. + timeout=5100, retries=0, max_containers=10, secrets=worker_secrets, 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 70be17666..4074151c6 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 @@ -2,6 +2,7 @@ from .aggregation import build_aggregate_report, build_spm_result from .coordination import coordinate_report +from .fanout import Stage12SegmentRunner, run_segmented_simulation from .partition import stage12_region_groups, stage12_region_groups_for_model from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult from .segmentation import ( @@ -26,6 +27,7 @@ "SimulationCalculation", "Stage12SegmentInput", "Stage12SegmentResult", + "Stage12SegmentRunner", "build_aggregate_report", "build_segment_inputs", "build_stage12_simulation", @@ -33,6 +35,7 @@ "calculate_segment", "coordinate_report", "run_single_simulation", + "run_segmented_simulation", "merge_segment_results", "simulation_input_sha256", "stage12_region_groups", 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 new file mode 100644 index 000000000..4172d5858 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py @@ -0,0 +1,188 @@ +"""Independent Modal orchestration for one Stage 12 segmented simulation.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +import logging +import time +from typing import Any + +import modal +from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, +) +from policyengine_simulation_observability.stages import ( + STAGE12_SIMULATION_STAGES, + Stage, +) + +from .partition import stage12_region_groups_for_model +from .segmentation import build_segment_inputs, merge_segment_results +from .simulation import SimulationCalculation + +logger = logging.getLogger(__name__) + +SEGMENT_FUNCTION_NAME = "run_single_simulation_segment_us" +POLL_INTERVAL_INITIAL_SECONDS = 0.5 +POLL_INTERVAL_MAX_SECONDS = 8.0 +POLL_INTERVAL_BACKOFF_FACTOR = 2.0 + + +def _next_backoff(current: float, *, maximum: float) -> float: + return min(current * POLL_INTERVAL_BACKOFF_FACTOR, maximum) + + +class Stage12SegmentRunner: + """Start, collect, and merge the 20 parts of one logical simulation.""" + + def __init__( + self, + simulation: PlannedSimulationExecutionInput, + *, + app_name: str, + runtime: ObservabilityRuntime, + modal_module: Any = modal, + groups: Sequence[Sequence[str]] | None = None, + sleep: Callable[[float], None] = time.sleep, + poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, + poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, + ) -> None: + self.simulation = simulation + self.runtime = runtime + self.sleep = sleep + self.poll_interval_initial_seconds = poll_interval_seconds + self.poll_interval_max_seconds = poll_interval_max_seconds + self.groups = [list(group) for group in groups] if groups is not None else None + self.child_function = modal_module.Function.from_name( + app_name, + SEGMENT_FUNCTION_NAME, + ) + + def _groups(self) -> list[list[str]]: + if self.groups is not None: + return self.groups + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + ) + + country = self.simulation.geography.country + return stage12_region_groups_for_model( + country, + _country_module(country).model, + ) + + def run(self) -> SimulationCalculation: + groups = self._groups() + segments = build_segment_inputs(self.simulation, groups=groups) + self.runtime.set_context( + stage12_execution="segmented", + stage12_segment_count=len(segments), + ) + handles: list[tuple[int, Any]] = [] + 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(), + ) + if call is None: + raise RuntimeError("Stage 12 segment dispatch returned no call") + handles.append((segment.segment_index, call)) + except Exception: + 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: + raise RuntimeError("Stage 12 segment wait returned no payloads") + merged: SimulationCalculation | None = None + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) + ): + merged = merge_segment_results( + self.simulation, + payloads, + groups=groups, + ) + if merged is None: + raise RuntimeError("Stage 12 segment merge returned no calculation") + return merged + + def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: + results: list[object | None] = [None] * len(handles) + pending = set(range(len(handles))) + poll_errors: dict[int, int] = {} + current_sleep = self.poll_interval_initial_seconds + poll_count = 0 + + while pending: + progress_made = False + for index in sorted(pending): + segment_index, call = handles[index] + poll_count += 1 + try: + results[index] = call.get(timeout=0) + except TimeoutError: + poll_errors.pop(index, None) + continue + except Exception as error: + attempts = poll_errors.get(index, 0) + 1 + poll_errors[index] = attempts + if attempts < 2: + logger.warning( + "Stage 12 segment %d poll failed once; retrying (%s)", + segment_index, + type(error).__name__, + ) + continue + self._cancel_all(handles) + raise RuntimeError( + f"Stage 12 segment {segment_index} failed" + ) from None + pending.discard(index) + poll_errors.pop(index, None) + progress_made = True + self.runtime.set_context(stage12_segment_poll_count=poll_count) + if pending and not progress_made: + self.sleep(current_sleep) + current_sleep = _next_backoff( + current_sleep, + maximum=self.poll_interval_max_seconds, + ) + elif progress_made: + current_sleep = self.poll_interval_initial_seconds + + if any(result is None for result in results): + raise RuntimeError("Stage 12 segment collection is incomplete") + return [result for result in results if result is not None] + + @staticmethod + def _cancel_all(handles: list[tuple[int, Any]]) -> None: + for _, call in handles: + try: + call.cancel() + except Exception: + pass + + +def run_segmented_simulation( + simulation: PlannedSimulationExecutionInput, + *, + app_name: str, + runtime: ObservabilityRuntime, +) -> SimulationCalculation: + return Stage12SegmentRunner( + simulation, + app_name=app_name, + runtime=runtime, + ).run() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py index 8d4c61e1e..5b37426e1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py @@ -11,6 +11,7 @@ from policyengine_simulation_contract.stage12_execution import ( PlannedSimulationExecutionInput, ) +from policyengine_simulation_observability.stages import STAGE12_SEGMENT_STAGES from policyengine_simulation_executor.stage12_artifacts import ( deserialize_calculation_provenance, @@ -87,7 +88,11 @@ def calculate_segment( segment = Stage12SegmentInput.model_validate(payload) scoped = _scoped_simulation(segment) if calculator is calculate_simulation_frames: - calculated = calculate_simulation_frames(scoped, runtime=runtime) + calculated = calculate_simulation_frames( + scoped, + runtime=runtime, + stage_plan=STAGE12_SEGMENT_STAGES, + ) else: calculated = calculator(scoped) if isinstance(calculated, SimulationCalculation): 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 3da356fd9..aa0026155 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 @@ -27,6 +27,7 @@ from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, Stage, + StagePlan, ) from policyengine_simulation_executor.stage12_artifacts import ( @@ -156,6 +157,7 @@ def calculate_simulation_frames( simulation: PlannedSimulationExecutionInput, *, runtime: ObservabilityRuntime | None = None, + stage_plan: StagePlan = STAGE12_SIMULATION_STAGES, ) -> SimulationCalculation: """Run one policy and return the coordinator-planned entity output tables.""" @@ -183,20 +185,20 @@ def calculate_simulation_frames( ) credential_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.CREDENTIAL_SETUP)) + runtime.span(stage_plan.name(Stage.CREDENTIAL_SETUP)) if runtime is not None else nullcontext() ) with credential_span, setup_gcp_credentials(): country_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.COUNTRY_MODULE_LOAD)) + runtime.span(stage_plan.name(Stage.COUNTRY_MODULE_LOAD)) if runtime is not None else nullcontext() ) with country_span: country_module = _country_module(country) region_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.REGION_RESOLUTION)) + runtime.span(stage_plan.name(Stage.REGION_RESOLUTION)) if runtime is not None else nullcontext() ) @@ -207,7 +209,7 @@ def calculate_simulation_frames( params=params, ) dataset_resolution_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.DATASET_RESOLUTION)) + runtime.span(stage_plan.name(Stage.DATASET_RESOLUTION)) if runtime is not None else nullcontext() ) @@ -217,7 +219,7 @@ def calculate_simulation_frames( region_resolution=region, ) dataset_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.DATASET_LOAD)) + runtime.span(stage_plan.name(Stage.DATASET_LOAD)) if runtime is not None else nullcontext() ) @@ -236,14 +238,14 @@ def calculate_simulation_frames( region.scoping_strategy, dataset, dataset_selection ) policy_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.POLICY_NORMALIZATION)) + runtime.span(stage_plan.name(Stage.POLICY_NORMALIZATION)) if runtime is not None else nullcontext() ) with policy_span: policy = _normalise_policy(simulation.policy) build_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.SIMULATION_BUILD)) + runtime.span(stage_plan.name(Stage.SIMULATION_BUILD)) if runtime is not None else nullcontext() ) @@ -259,13 +261,19 @@ def calculate_simulation_frames( execution=simulation, ) calculation_span = ( - runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_CALCULATION)) + runtime.span(stage_plan.name(Stage.STAGE12_CALCULATION)) if runtime is not None else nullcontext() ) with calculation_span: apply_output_plan(model, simulation.output_plan) - model.ensure() + cache_span = ( + runtime.span(stage_plan.name(Stage.STAGE12_CACHE_LOOKUP)) + if runtime is not None and hasattr(model, "stage12_cache_outcome") + else nullcontext() + ) + with cache_span: + model.ensure() cache_outcome = getattr(model, "stage12_cache_outcome", None) if runtime is not None and cache_outcome is not None: runtime.set_context(stage12_cache_outcome=cache_outcome) @@ -351,6 +359,10 @@ def run_single_simulation( [PlannedSimulationExecutionInput], Mapping[str, pd.DataFrame] | SimulationCalculation, ] = calculate_simulation_frames, + segmented_calculator: Callable[ + [PlannedSimulationExecutionInput], SimulationCalculation + ] + | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: simulation = PlannedSimulationExecutionInput.model_validate(payload) @@ -383,7 +395,11 @@ def run_single_simulation( prefix=context.artifact_prefix, simulation=simulation, ) - if calculator is calculate_simulation_frames: + from .segmentation import should_segment_simulation + + if segmented_calculator is not None and should_segment_simulation(simulation): + calculated = segmented_calculator(simulation) + elif calculator is calculate_simulation_frames: calculated = calculate_simulation_frames(simulation, runtime=runtime) else: calculation_span = ( 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 3fcb76a7e..3b6a952a8 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -59,6 +59,21 @@ def _invoke( ): from policyengine_simulation_executor import stage12_runtime + if function_name == "run_single_simulation_segment_us": + + class Result: + def model_dump(self, *, mode): + assert mode == "python" + return implementation() + + monkeypatch.setattr( + stage12_runtime, "calculate_segment", lambda *_, **__: Result() + ) + return module.run_single_simulation_segment_us( + {"simulation": {"role": "baseline"}, "segment_index": 0}, + observability_context=_observability_context(), + ) + if function_name == "coordinate_report": monkeypatch.setattr( module.modal, @@ -199,18 +214,24 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert set(functions) == { "validate_worker_us", "validate_worker_uk", + "run_single_simulation_segment_us", "run_single_simulation_us", "run_single_simulation_uk", "coordinate_report", } assert functions["run_single_simulation_us"]["image"] is module.us_worker_image assert functions["run_single_simulation_uk"]["image"] is module.uk_worker_image - assert functions["run_single_simulation_us"]["timeout"] == 3000 + assert ( + functions["run_single_simulation_segment_us"]["image"] is module.us_worker_image + ) + assert functions["run_single_simulation_segment_us"]["timeout"] == 3600 + assert functions["run_single_simulation_segment_us"]["max_containers"] == 300 + assert functions["run_single_simulation_us"]["timeout"] == 3900 assert functions["run_single_simulation_uk"]["timeout"] == 3000 assert functions["run_single_simulation_us"]["max_containers"] == 10 assert functions["run_single_simulation_uk"]["max_containers"] == 10 assert functions["coordinate_report"]["image"] is module.coordinator_image - assert functions["coordinate_report"]["timeout"] == 4500 + assert functions["coordinate_report"]["timeout"] == 5100 assert functions["coordinate_report"]["max_containers"] == 10 assert "asgi_app" not in vars(module) @@ -220,6 +241,7 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( [ "run_single_simulation_us", "run_single_simulation_uk", + "run_single_simulation_segment_us", "coordinate_report", ], ) @@ -246,6 +268,7 @@ def test_stage12_functions_propagate_identifier_and_close_runtime( [ "run_single_simulation_us", "run_single_simulation_uk", + "run_single_simulation_segment_us", "coordinate_report", ], ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 51732662e..633a1be8b 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -679,6 +679,50 @@ def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: ] +def test_single_worker_routes_eligible_us_calculation_to_segment_runner() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + received = [] + + def segmented_calculator(value): + received.append(value) + return SimulationCalculation(frames=_frames()) + + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + assert received == [simulation] + + +def test_single_worker_honors_explicit_segmentation_opt_out() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"segmented": False}} + ) + store.children[simulation.simulation_execution_id] = _child(simulation) + segmented_calls = [] + + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: _frames(), + segmented_calculator=lambda value: segmented_calls.append(value), + ) + + assert segmented_calls == [] + + def test_single_worker_retains_detached_calculation_provenance() -> None: store = FakeStore() simulation = _planned_simulation(SimulationRole.BASELINE) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py new file mode 100644 index 000000000..d8b39386d --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -0,0 +1,193 @@ +"""Unit tests for independent Stage 12 Modal segment orchestration.""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace + +import pandas as pd +import pytest +from policyengine_simulation_contract.stage12_execution import SimulationRole + +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + Stage12SegmentRunner, + build_segment_inputs, + calculate_segment, +) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) +from test_stage12_runtime import _planned_simulation + + +def _frames(index: int) -> dict[str, pd.DataFrame]: + return { + "household": pd.DataFrame( + { + "household_id": pd.Series([index], dtype="int64"), + "household_net_income": pd.Series([index + 0.5], dtype="float32"), + } + ), + "person": pd.DataFrame( + { + "age": pd.Series([30 + index], dtype="int16"), + "household_id": pd.Series([index], dtype="int64"), + "person_id": pd.Series([100 + index], dtype="int64"), + } + ), + } + + +class FakeRuntime: + def __init__(self) -> None: + self.context: dict[str, object] = {} + self.spans: list[tuple[str, dict | None]] = [] + + def set_context(self, **attributes) -> None: + self.context.update(attributes) + + def capture_context(self) -> dict[str, str]: + return {"observability_id": "00000000-0000-4000-8000-000000000001"} + + @contextmanager + def span(self, name, *, attributes=None): + self.spans.append((name, attributes)) + yield + + +class FakeCall: + def __init__(self, outcomes: list[object]) -> None: + self.outcomes = outcomes + self.get_timeouts: list[float] = [] + self.cancelled = 0 + + def get(self, *, timeout: float): + self.get_timeouts.append(timeout) + outcome = self.outcomes.pop(0) + if isinstance(outcome, BaseException): + raise outcome + return outcome + + def cancel(self) -> None: + self.cancelled += 1 + + +class FakeFunction: + def __init__(self, results: list[dict], errors: dict[int, list[Exception]]): + self.results = results + self.errors = errors + self.spawns: list[tuple[dict, dict | None]] = [] + self.calls: list[FakeCall] = [] + + def spawn(self, payload: dict, *, observability_context: dict | None): + index = payload["segment_index"] + self.spawns.append((payload, observability_context)) + call = FakeCall([*self.errors.get(index, []), self.results[index]]) + self.calls.append(call) + return call + + +def _results(simulation) -> list[dict]: + return [ + calculate_segment( + segment, + calculator=lambda scoped, index=segment.segment_index: ( + SimulationCalculation(frames=_frames(index)) + ), + ).model_dump(mode="python") + for segment in build_segment_inputs(simulation) + ] + + +def _runner(*, errors: dict[int, list[Exception]] | None = None): + simulation = _planned_simulation(SimulationRole.BASELINE) + function = FakeFunction(_results(simulation), errors or {}) + + class FunctionLookup: + @staticmethod + def from_name(app_name, function_name): + assert app_name == "stage12-app" + assert function_name == "run_single_simulation_segment_us" + return function + + runtime = FakeRuntime() + sleeps: list[float] = [] + runner = Stage12SegmentRunner( + simulation, + app_name="stage12-app", + runtime=runtime, + modal_module=SimpleNamespace(Function=FunctionLookup), + groups=US_REGION_GROUPS, + sleep=sleeps.append, + poll_interval_seconds=0.5, + poll_interval_max_seconds=2.0, + ) + return runner, function, runtime, sleeps + + +def test_starts_twenty_segments_and_merges_them_in_partition_order() -> None: + runner, function, runtime, sleeps = _runner() + + result = runner.run() + + assert len(function.spawns) == 20 + assert [payload["segment_index"] for payload, _ in function.spawns] == list( + range(20) + ) + assert all(context is not None for _, context in function.spawns) + assert all(call.get_timeouts == [0] for call in function.calls) + assert result.frames["household"]["household_id"].tolist() == list(range(20)) + assert runtime.context["stage12_segment_count"] == 20 + assert sleeps == [] + + +def test_retries_one_poll_error_without_restarting_the_segment() -> None: + runner, function, _, _ = _runner(errors={0: [ConnectionError("transient")]}) + + runner.run() + + assert len(function.spawns) == 20 + assert function.calls[0].get_timeouts == [0, 0] + assert not any(call.cancelled for call in function.calls) + + +def test_not_ready_polls_use_bounded_backoff() -> None: + runner, function, _, sleeps = _runner( + errors={index: [TimeoutError()] for index in range(20)} + ) + + runner.run() + + assert sleeps == [0.5] + assert all(call.get_timeouts == [0, 0] for call in function.calls) + + +def test_second_poll_error_cancels_all_started_segments() -> None: + runner, function, _, _ = _runner( + errors={0: [ConnectionError("first"), RuntimeError("second")]} + ) + + with pytest.raises(RuntimeError, match="segment 0 failed"): + runner.run() + + assert len(function.calls) == 20 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_dispatch_error_cancels_segments_that_already_started() -> None: + runner, function, _, _ = _runner() + spawn = function.spawn + + def fail_during_dispatch(payload, *, observability_context): + if payload["segment_index"] == 5: + raise ConnectionError("dispatch failed") + return spawn(payload, observability_context=observability_context) + + function.spawn = fail_during_dispatch + + with pytest.raises(ConnectionError, match="dispatch failed"): + runner.run() + + assert len(function.calls) == 5 + assert all(call.cancelled == 1 for call in function.calls) From 15b4cba99da345dc47d2012779a658f860db75ab Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 30 Sep 2026 00:33:26 +0400 Subject: [PATCH 07/25] Qualify and document Stage 12 execution --- .../stage-12-modal-worker-foundation.md | 101 ++++++++++++++++-- docs/operations/api-v1-observability.md | 9 +- .../stage12_manifest.py | 4 +- .../tests/test_stage12_manifest.py | 4 +- .../README.md | 30 ++++++ .../src/modal/v2_app.py | 26 ++++- .../stage12_qualification.py | 67 ++++++++++-- .../tests/test_modal_scripts.py | 69 ++++++++++++ .../tests/test_stage12_qualification.py | 49 ++++++++- .../tests/test_stage12_runtime.py | 6 +- .../tests/test_stage12_segment_fanout.py | 4 +- 11 files changed, 335 insertions(+), 34 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index f5a802179..edfe360bf 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -20,9 +20,17 @@ simulation submission or polling contracts. records as `main`. - One v2 report-coordinator invocation starts baseline and reform as distinct single-simulation calls before awaiting either result. -- Each country-specific single-simulation function and the report-coordinator - function has a Modal `max_containers` value of 10. This is a per-function - limit; the functions do not share one application-wide container quota. +- An eligible US national single-simulation call starts 20 region-group + calculation calls and merges their typed Parquet results. Because baseline + and reform are separate and start before either is awaited, an eligible + report starts 40 region-group calculation calls. These are invocations; + Modal may reuse a warm container for more than one invocation. +- Country-specific single-simulation functions and the report coordinator have + a Modal `max_containers` value of 10. The separate US region-group function + has a value of 300 so calculation work cannot wait behind the parent calls + that are collecting it. These are per-function limits, not one shared quota. +- Stage 12 owns its partition, request and result models, orchestration, merge, + and cache implementation. It does not import the corresponding v1 modules. - The v2 functions write canonical private artifacts and use a shared SQLAlchemy Core package to persist temporary comparison state directly. They do not create production simulations, reports, report runs, or user @@ -67,6 +75,10 @@ with the live API-owned schema. validation downloads that lookup file using the worker credential. Stage 12 local-authority output uses packaged metadata and does not require access to `local_authorities_2021.csv`. +- Stage 12 uses a dedicated environment-specific cache bucket named by + `STAGE12_CACHE_BUCKET`. Its object keys and manifest namespace are separate + from the v1 precompute store. The same restricted worker credential can + create, read, and delete objects in this cache bucket. - `STAGE12_DATABASE_URL` is delivered from the environment-specific Secret Manager resource named by `STAGE12_DATABASE_URL_SECRET_NAME`. It authenticates as the existing shared `policyengine_v2_runtime` account; Stage 12 does not @@ -98,9 +110,10 @@ automation or infrastructure configuration: | Operation | Maintained owner | | --- | --- | | Build and deploy each versioned v2 Modal application | Simulation deployment workflow | +| Build and verify the Stage 12 dataset and baseline cache | Simulation deployment workflow | | Create or update the separate v2 manifest storage object | V2 manifest publisher | | Configure the v2 manifest reader and `STAGE12_ENABLED` | Cloud Run deployment workflow | -| Create the private artifact namespace and its retention policy | Infrastructure configuration | +| Create the private report-artifact and cache buckets and their policies | Infrastructure configuration | | Ensure the existing API v2 runtime identity has comparison-row database access | API infrastructure configuration | | Deliver the existing shared runtime database secret to Cloud Run and Modal | Simulation deployment workflow | | Grant private-object access only to the Modal runtime | Infrastructure configuration | @@ -160,9 +173,10 @@ Before deploying, the simulation workflow obtains the shared runtime URL from Secret Manager, verifies the expected `policyengine_v2_runtime` identity and permissions, inspects the complete temporary schema against the declared SQLAlchemy mappings, and runs a rolled-back insert/read/update/delete canary. It -also verifies required Secret Manager access and performs a create/read/delete -canary in the private artifact bucket. The storage canary is deleted in the -normal path and by exit cleanup after a failure. +also verifies required Secret Manager access and performs separate +create/read/delete checks in the private report-artifact and cache buckets. The +temporary objects are deleted in the normal path and by exit cleanup after a +failure. The current Public API omits `data` for the certified default dataset. The v2 adapter therefore resolves an absent `data` field to the exact default dataset @@ -255,6 +269,63 @@ move to the front end, and Stage 14 or later work should remove these temporary executor resources when the authoritative v2 report architecture supersedes Stage 12. +## US segmentation and baseline cache + +Segmentation applies by default when both simulation inputs describe a US +national society-wide calculation and the shared output plan does not require +cliff analysis or labor-supply responses. `segmented=false` selects the +single-container calculation explicitly. UK calculations, subnational US +calculations, cliff analysis, and labor-supply-response calculations also use +the single-container path. `segmented=true` is accepted, but it does not +override those numerical eligibility conditions. + +The Stage 12 coordinator still persists only two logical simulation records: +one baseline and one reform. Each eligible logical simulation call performs +the following work independently: + +```text +logical baseline or reform simulation + | + +---- start 20 US region-group calculation calls + | + +---- collect and validate all 20 results + | + +---- concatenate entity rows and reject duplicate identifiers + | + +---- write one logical simulation artifact and update one record +``` + +Region-group calls do not write temporary database rows or private report +artifacts. They return compressed, dtype-preserving data to their logical +simulation parent. The parent validates the result role, partition index, +region membership, content digest, entity schemas, column dtypes, and unique +entity identifiers before merging. A missing or failed region-group call fails +that logical simulation. The parent makes a best-effort cancellation request +for every calculation call that it already started. + +Before each v2 deployment, the independent Stage 12 precompute application +builds three annual US dataset files and 60 current-law baseline files: one for +each of 20 region groups in 2025, 2026, and 2027. Cache identities include the +bundle manifest digest, package and data versions, source-data digest, year, +region membership, scoping identity, output-plan digest, and optional SPM +configuration. Objects are content-addressed and written with a create-only +condition; an existing object must have the expected digest and size. + +Precompute publishes a canonical manifest containing all 63 files. The +deployment passes that manifest's SHA-256 digest into the US image build. The +image build downloads and verifies every listed file, then records the exact +cache release in worker-validation metadata and the v2 version manifest. The UK +and coordinator images do not contain these files. Deployment validation runs +one real 2026 region-group baseline and requires it to load the precomputed +file without recalculation. + +Only an empty-policy baseline using the bundle's default dataset and an exact +region-group scope can read a baseline cache file. Reform calculations always +compute. Missing cache files fall back to calculation; a loaded file that does +not satisfy the requested output plan is recalculated. Each region-group log +records whether it loaded a complete file, recalculated an incomplete file, or +computed because no file existed. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section @@ -299,9 +370,9 @@ Modal report coordinator | atomically create or resolve the temporary parent row | stop if another invocation already owns the same logical run | - +------> baseline single-simulation call + +------> baseline logical simulation call ------> 20 region-group calls | - +------> reform single-simulation call + +------> reform logical simulation call --------> 20 region-group calls | +------> write aggregate artifact and durable state @@ -333,11 +404,17 @@ curl -X POST "${SIMULATION_ENTRYPOINT_URL}/internal/stage12/reports" \ "scope": "macro", "region": "us", "time_period": "2026", + "segmented": true, "baseline": {}, "reform": {} }' ``` +The `segmented` field is optional. Omitting it uses segmentation when the +request meets the eligibility conditions above. `false` explicitly selects the +single-container path. `true` requests segmentation but still falls back to +the single-container path when the output requirements are not eligible. + An accepted request returns `202` and `Retry-After: 1` without waiting for the calculations or for the coordinator's first database write: @@ -434,3 +511,9 @@ aggregate, and compares every aggregate field. Differences fail qualification with only a bounded code and field path; calculation values are not included in the diagnostic or receipt. Numerical tolerances default to zero and can be introduced only for explicit field paths through a reviewed JSON file. +For an eligible US national input, the Stage 12 side of qualification executes +and merges all 20 region groups separately for both baseline and reform. This +checks the partition and merge numerically in addition to comparing the final +entity tables and aggregate report with the existing combined executor. Inputs +that are not eligible for segmentation continue through the single-container +Stage 12 calculation. diff --git a/docs/operations/api-v1-observability.md b/docs/operations/api-v1-observability.md index 07cacc8b6..77ed93d53 100644 --- a/docs/operations/api-v1-observability.md +++ b/docs/operations/api-v1-observability.md @@ -54,8 +54,13 @@ The simulation entry service dispatches the current production run and the Stage 12 report with the same `observability_id`. Stage 12 creates separate functional identifiers for the report and its baseline and reform simulations. The coordinator passes captured context to both simulation workers. When Stage -12 becomes the authoritative runner, the same report identifier flow applies; -only the selected functional result changes. +12 runs an eligible US national report, each simulation worker passes that same +context to 20 region-group calculation calls. Baseline and reform therefore +produce 40 region-group invocations with one `observability_id`, two +`simulation_execution_id` values, and no additional persisted simulation +records. The region-group index and cache outcome distinguish their structured +logs. When Stage 12 becomes the authoritative runner, the same report +identifier flow applies; only the selected functional result changes. ## Query model diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 47873a82c..0cea893a1 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -2,12 +2,12 @@ from __future__ import annotations +import json +import re from collections.abc import Mapping from copy import deepcopy from datetime import datetime from hashlib import sha256 -import json -import re from typing import Annotated, Literal, Protocol from packaging.version import InvalidVersion, Version diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 944c1814c..14b9a2b89 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime import pytest @@ -107,7 +107,7 @@ def _worker(version: str = "5.2.0") -> V2WorkerVersion: validated=True, application_name=application_name, bundle_manifest_sha256=resolved.bundle_manifest_sha256, - validated_at=datetime(2026, 9, 14, tzinfo=timezone.utc), + validated_at=datetime(2026, 9, 14, tzinfo=UTC), validation_invocation_id="validation-123", country_validation_invocation_ids={ "us": "validation-us-123", diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 83170934c..3ad9e4150 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -156,6 +156,36 @@ Operational notes: That staleness class predates the artifact pipeline and lives in the bundle install, not the fetch. +## Independent Stage 12 segmentation and cache + +The separately named Stage 12 v2 application has its own implementation under +`stage12_runtime/` and `stage12_cache/`. It does not call the v1 segmented-run +or precompute modules described above. + +For an eligible US national report, the coordinator starts baseline and reform +as separate logical simulations. Each logical simulation starts 20 calls to +`run_single_simulation_segment_us`, then validates and concatenates their +dtype-preserving entity tables. This produces 40 region-group calculation +invocations per report while retaining two simulation records and two private +simulation artifacts. `segmented=false`, US subnational requests, UK requests, +cliff analysis, and labor-supply-response calculations use the +single-container Stage 12 path. + +Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the +dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 +current-law US baseline files and emits +`STAGE12_CACHE_MANIFEST_DIGEST=`. The US worker image downloads and +verifies the manifest's 63 files during image construction. The release +validation checks the recorded manifest, hashes a baseline file, and requires +a real calculation to load that file without recomputing it. The cache manifest +digest is then stored in the separate v2 version manifest. + +The cache is an optimization, not required state for a calculation: an absent +file causes a normal calculation, and a file missing planned output columns is +recalculated. Only empty-policy baseline region-group calculations against the +bundle's default dataset are eligible to read it. Reform calculations never +read a baseline file. + ## Observability `policyengine-observability` 3.x emits structured request, operation, diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 559a5df5d..556b39b80 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -244,6 +244,7 @@ def run_single_simulation_segment_us( try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION), attributes={ @@ -253,7 +254,12 @@ def run_single_simulation_segment_us( }, remote_context=propagated, ): - return calculate_segment(payload, runtime=runtime).model_dump(mode="python") + result = calculate_segment(payload, runtime=runtime).model_dump( + mode="python" + ) + if result is None: + raise RuntimeError("Stage 12 segment calculation returned no result") + return result finally: runtime.shutdown() @@ -294,6 +300,7 @@ def segmented_calculator(simulation): try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ @@ -302,13 +309,16 @@ def segmented_calculator(simulation): }, remote_context=propagated, ): - return run_single_simulation( + result = run_single_simulation( payload, context, required_country="us", segmented_calculator=segmented_calculator, runtime=runtime, ) + if result is None: + raise RuntimeError("Stage 12 US simulation returned no result") + return result finally: runtime.shutdown() @@ -340,6 +350,7 @@ def run_single_simulation_uk( ) try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ @@ -348,12 +359,15 @@ def run_single_simulation_uk( }, remote_context=propagated, ): - return run_single_simulation( + result = run_single_simulation( payload, context, required_country="uk", runtime=runtime, ) + if result is None: + raise RuntimeError("Stage 12 UK simulation returned no result") + return result finally: runtime.shutdown() @@ -396,6 +410,7 @@ def coordinate_report( ) try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION), attributes={ @@ -408,7 +423,7 @@ def coordinate_report( }, remote_context=propagated, ): - return run_report_coordinator( + result = run_report_coordinator( payload, context, parent, @@ -416,5 +431,8 @@ def coordinate_report( coordinator_invocation_id=coordinator_invocation_id, runtime=runtime, ) + if result is None: + raise RuntimeError("Stage 12 report coordinator returned no result") + return result finally: runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index 473d1ea82..9a9c86cfb 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from hashlib import sha256 from typing import Any, Literal @@ -80,6 +80,46 @@ class Stage12ParityReceipt(BaseModel): AggregateBuilder = Callable[..., dict[str, Any]] +def _run_stage12_simulation_for_qualification( + simulation: PlannedSimulationExecutionInput, + *, + groups: Sequence[Sequence[str]] | None = None, + segment_runner: Callable[[object], object] | None = None, + monolithic_runner: SingleSimulationRunner | None = None, +) -> SimulationCalculation: + """Exercise the segmented calculation and merge without Modal transport.""" + + from policyengine_simulation_executor.stage12_runtime import ( + build_segment_inputs, + calculate_segment, + calculate_simulation_frames, + merge_segment_results, + should_segment_simulation, + stage12_region_groups_for_model, + ) + + run_monolithic = monolithic_runner or calculate_simulation_frames + if not should_segment_simulation(simulation): + return _calculation(run_monolithic(simulation)) + selected_groups = groups + if selected_groups is None: + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + ) + + country = simulation.geography.country + selected_groups = stage12_region_groups_for_model( + country, + _country_module(country).model, + ) + run_segment = segment_runner or calculate_segment + results = [ + run_segment(segment) + for segment in build_segment_inputs(simulation, groups=selected_groups) + ] + return merge_segment_results(simulation, results, groups=selected_groups) + + def _local_qualification_runtime() -> ObservabilityRuntime: """Create a local runtime that cannot send qualification data remotely.""" @@ -127,6 +167,7 @@ def _validate_matching_input( "policyengine_version": bundle.policyengine_version, "version": bundle.country_package_version, "spm": baseline.options.get("spm"), + "segmented": baseline.options.get("segmented"), } for field, expected_value in expected.items(): actual_value = existing_request.get(field) @@ -243,12 +284,9 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: return run_simulation_impl(request, runtime=local_runtime) existing_runner = run_existing + default_stage12_runner = single_simulation_runner is None if single_simulation_runner is None: - from policyengine_simulation_executor.stage12_runtime import ( - calculate_simulation_frames, - ) - - single_simulation_runner = calculate_simulation_frames + single_simulation_runner = _run_stage12_simulation_for_qualification if aggregate_builder is None: from policyengine_simulation_executor.stage12_runtime import ( build_aggregate_report, @@ -279,11 +317,18 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: "report.data_version", ) - with ThreadPoolExecutor(max_workers=2) as executor: - baseline_future = executor.submit(single_simulation_runner, baseline_input) - reform_future = executor.submit(single_simulation_runner, reform_input) - baseline_calculation = _calculation(baseline_future.result()) - reform_calculation = _calculation(reform_future.result()) + if default_stage12_runner: + # The production segments run in isolated Modal processes. Run the + # local qualifier sequentially so country-model process caches are not + # accessed concurrently by Python threads. + baseline_calculation = _calculation(single_simulation_runner(baseline_input)) + reform_calculation = _calculation(single_simulation_runner(reform_input)) + else: + with ThreadPoolExecutor(max_workers=2) as executor: + baseline_future = executor.submit(single_simulation_runner, baseline_input) + reform_future = executor.submit(single_simulation_runner, reform_input) + baseline_calculation = _calculation(baseline_future.result()) + reform_calculation = _calculation(reform_future.result()) compare_simulation_frames( existing_baseline, diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index d00fbb7cc..e3bebcd7b 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -755,6 +755,75 @@ def test_precompute_is_gated_by_the_shared_secret_sync(self): assert "needs: prepare" in executor_job +class TestStage12CachePrecompute: + """Tests for the independent Stage 12 cache deployment script.""" + + script = SCRIPTS_DIR / "stage12-cache-precompute.sh" + + def _run_with_fake_uv(self, tmp_path, output: str, *args): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + log_path = tmp_path / "uv-calls.log" + github_output = tmp_path / "github-output.txt" + uv_path = bin_dir / "uv" + uv_path.write_text( + '#!/bin/bash\nprintf \'%s\\n\' "$*" >> "$UV_FAKE_LOG"\nprintf \'%s\\n\' "$UV_FAKE_OUTPUT"\n', + encoding="utf-8", + ) + uv_path.chmod(0o755) + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "UV_FAKE_LOG": str(log_path), + "UV_FAKE_OUTPUT": output, + "GITHUB_OUTPUT": str(github_output), + "STAGE12_CACHE_BUCKET": "policyengine-stage12-cache-staging", + } + result = subprocess.run( + ["bash", str(self.script), "staging", *args], + capture_output=True, + text=True, + env=env, + ) + calls = log_path.read_text(encoding="utf-8") if log_path.exists() else "" + emitted = ( + github_output.read_text(encoding="utf-8") if github_output.exists() else "" + ) + return result, calls, emitted + + def test_validates_syntax_and_exports_exact_manifest_digest(self, tmp_path): + syntax = subprocess.run( + ["bash", "-n", str(self.script)], + capture_output=True, + text=True, + ) + assert syntax.returncode == 0, syntax.stderr + digest = "a" * 64 + + result, calls, emitted = self._run_with_fake_uv( + tmp_path, + f"progress\nSTAGE12_CACHE_MANIFEST_DIGEST={digest}", + ) + + assert result.returncode == 0, result.stderr + assert ( + "run modal run --env=staging src/modal/stage12_precompute_app.py" in calls + ) + assert "--force" not in calls + assert emitted == f"manifest_digest={digest}\n" + + def test_rejects_invalid_digest_and_passes_force(self, tmp_path): + result, calls, emitted = self._run_with_fake_uv( + tmp_path, + "STAGE12_CACHE_MANIFEST_DIGEST=not-a-digest", + "true", + ) + + assert result.returncode != 0 + assert calls.rstrip().endswith("--force") + assert emitted == "" + + class TestModalRecordDeployment: """Tests for modal-record-deployment.sh""" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index ef215fae2..c7f02b75c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -14,11 +14,18 @@ ) from policyengine_simulation_executor.stage12_parity import ParityMismatch from policyengine_simulation_executor.stage12_qualification import ( + _run_stage12_simulation_for_qualification, qualify_report_parity, ) -from policyengine_simulation_executor.stage12_runtime import SimulationCalculation +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + calculate_segment, +) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) -from test_stage12_runtime import _report +from test_stage12_runtime import _planned_simulation, _report def _frames(value: float) -> dict[str, pd.DataFrame]: @@ -62,6 +69,44 @@ def _existing_request() -> dict[str, Any]: } +def test_qualification_runner_exercises_all_segments_and_exact_merge() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + calls = [] + + def run_segment(segment): + calls.append(segment.segment_index) + index = segment.segment_index + frames = { + "household": pd.DataFrame( + { + "household_id": pd.Series([index], dtype="int64"), + "household_net_income": pd.Series([index + 0.5], dtype="float32"), + } + ), + "person": pd.DataFrame( + { + "age": pd.Series([30 + index], dtype="int16"), + "household_id": pd.Series([index], dtype="int64"), + "person_id": pd.Series([100 + index], dtype="int64"), + } + ), + } + return calculate_segment( + segment, + calculator=lambda _: SimulationCalculation(frames=frames), + ) + + result = _run_stage12_simulation_for_qualification( + simulation, + groups=US_REGION_GROUPS, + segment_runner=run_segment, + ) + + assert calls == list(range(20)) + assert result.frames["household"]["household_id"].tolist() == list(range(20)) + assert str(result.frames["person"]["age"].dtype) == "int16" + + def test_qualification_runs_matching_inputs_through_both_complete_paths() -> None: report = _report() baseline = _frames(100.0) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 633a1be8b..b11165fb8 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -710,6 +710,10 @@ def test_single_worker_honors_explicit_segmentation_opt_out() -> None: store.children[simulation.simulation_execution_id] = _child(simulation) segmented_calls = [] + def unexpected_segmented(value): + segmented_calls.append(value) + return SimulationCalculation(frames=_frames()) + run_single_simulation( simulation.model_dump(mode="json"), _context().model_dump(mode="json"), @@ -717,7 +721,7 @@ def test_single_worker_honors_explicit_segmentation_opt_out() -> None: store=store, artifacts=FakeArtifacts(), calculator=lambda _: _frames(), - segmented_calculator=lambda value: segmented_calls.append(value), + segmented_calculator=unexpected_segmented, ) assert segmented_calls == [] 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 d8b39386d..2c64a1e6f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -4,9 +4,11 @@ from contextlib import contextmanager from types import SimpleNamespace +from typing import cast import pandas as pd import pytest +from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import SimulationRole from policyengine_simulation_executor.stage12_runtime import ( @@ -116,7 +118,7 @@ def from_name(app_name, function_name): runner = Stage12SegmentRunner( simulation, app_name="stage12-app", - runtime=runtime, + runtime=cast(ObservabilityRuntime, runtime), modal_module=SimpleNamespace(Function=FunctionLookup), groups=US_REGION_GROUPS, sleep=sleeps.append, From 9ef2f0272ad828514aefb4e2418562fc97266ff6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:49:03 +0400 Subject: [PATCH 08/25] Fix Stage 12 precompute image sources --- .../src/modal/stage12_precompute_app.py | 1 + .../tests/test_stage12_precompute_app.py | 1 + 2 files changed, 2 insertions(+) diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py index 3e40bd5ad..cf2ecda4a 100644 --- a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py @@ -55,6 +55,7 @@ def _bundle_install_command() -> str: "policyengine_simulation_executor", "policyengine_simulation_observability", "policyengine_simulation_contract", + "policyengine_stage12_persistence", copy=True, ) ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py index 3cb2930b6..a993a2c5f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -48,3 +48,4 @@ def test_precompute_image_installs_only_the_us_bundle(monkeypatch) -> None: if call[0] == "add_local_python_source" ) assert "policyengine_simulation_executor" in local_source[1] + assert "policyengine_stage12_persistence" in local_source[1] From 9b13a8be0cf9faca5bc47bb5d4c235ede11552d8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 30 Sep 2026 19:45:57 +0400 Subject: [PATCH 09/25] Validate Stage 12 cache deploy access --- .github/scripts/stage12-validate-infrastructure.sh | 3 +++ .../tests/test_stage12_infrastructure.py | 6 +++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/scripts/stage12-validate-infrastructure.sh b/.github/scripts/stage12-validate-infrastructure.sh index a2acf8b6b..c64413204 100755 --- a/.github/scripts/stage12-validate-infrastructure.sh +++ b/.github/scripts/stage12-validate-infrastructure.sh @@ -181,6 +181,9 @@ runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ "${cache_canary_object}" "${canary_download_file}" --quiet cmp "${canary_source_file}" "${canary_download_file}" +gcloud storage cp \ + "${cache_canary_object}" "${canary_download_file}" --quiet +cmp "${canary_source_file}" "${canary_download_file}" runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ "${cache_canary_object}" --quiet cache_canary_object="" diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py index fec6ba12e..2fd05a021 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py @@ -77,7 +77,7 @@ def test_infrastructure_validation_rejects_missing_configuration() -> None: assert "STAGE12_ENVIRONMENT is required" in result.stderr -def test_storage_validation_uses_only_the_modal_worker_credentials() -> None: +def test_storage_validation_checks_runtime_and_deployment_credentials() -> None: source = SCRIPT.read_text(encoding="utf-8") helper_start = source.index("runtime_gcloud()") helper_end = source.index("\n}\n", helper_start) @@ -94,6 +94,10 @@ def test_storage_validation_uses_only_the_modal_worker_credentials() -> None: modal_account = '--account="${STAGE12_MODAL_SERVICE_ACCOUNT}"' assert source.count(f"runtime_gcloud {modal_account} storage cp") == 5 assert source.count(f"runtime_gcloud {modal_account} storage rm") == 4 + assert ( + 'gcloud storage cp \\\n "${cache_canary_object}" ' + '"${canary_download_file}" --quiet' + ) in source assert "gs://${STAGE12_CACHE_BUCKET}/_deployment-validation" in source assert "gs://policyengine-uk-data-private/constituencies_2024.csv" in source assert "gs://policyengine-uk-data-private/local_authorities_2021.csv" not in source From 42abf0ef2a9ca93f9a8fee9f3d6cab1cb58e45cb Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 30 Sep 2026 20:42:30 +0400 Subject: [PATCH 10/25] Keep Stage 12 deployment validation independent --- .github/workflows/stage12-deploy.yml | 4 +++- .../tests/test_deployment_assets.py | 7 +++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/.github/workflows/stage12-deploy.yml b/.github/workflows/stage12-deploy.yml index c3fb9d1ee..cecb6b097 100644 --- a/.github/workflows/stage12-deploy.yml +++ b/.github/workflows/stage12-deploy.yml @@ -54,7 +54,9 @@ jobs: entrypoint_min_instances: 0 entrypoint_max_instances: 2 promote_entrypoint: false - run_full_integration: true + # This workflow deploys and validates only the independent Stage 12 path. + # The main API deployment continues to run the existing-stack suite. + run_full_integration: false deploy_existing_stack: false deploy_stage12_v2: true secrets: inherit diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index e2351c18e..c07b624fe 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -712,6 +712,13 @@ def test_complete_integration_suite_is_configured_for_beta_only(): assert "run_full_integration:" in reusable_workflow +def test_stage12_only_deployment_skips_existing_stack_integration_suite(): + stage12_workflow = STAGE12_DEPLOY_WORKFLOW.read_text(encoding="utf-8") + + assert stage12_workflow.count("run_full_integration: false") == 2 + assert "run_full_integration: true" not in stage12_workflow + + def test_main_deployment_automatically_deploys_stage12_in_both_environments(): deploy_workflow = DEPLOY_WORKFLOW.read_text(encoding="utf-8") reusable_workflow = REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8") From 09f54a0c039caac0543b538e4fcb75966238e5a7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 00:10:09 +0400 Subject: [PATCH 11/25] Validate Stage 12 reform intervals --- .../tests/stage12_fixtures.py | 4 ++- .../stage12_worker_validation.py | 36 ++++++++++++++++--- .../tests/test_stage12_worker_validation.py | 32 ++++++++++++++++- 3 files changed, 66 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py index 4133ac0c5..46ebe97e7 100644 --- a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py +++ b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py @@ -97,7 +97,9 @@ def eligible_payload() -> dict: "country": "us", "scope": "macro", "baseline": {}, - "reform": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3000}}, + "reform": { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3000} + }, "time_period": "2026", "region": "us", "include_cliffs": False, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index a8fbdcff3..3ef087e41 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -28,6 +28,10 @@ ("STAGE12_CACHE_BUCKET",), ) +US_VALIDATION_REFORM = { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3_000} +} + def _require_secret_environment(environment: Mapping[str, str]) -> None: for alternatives in REQUIRED_SECRET_ALTERNATIVES: @@ -56,16 +60,40 @@ def _check_dataset_access(path: str, expected_sha256: str) -> None: def _run_non_serving_calculation(country: CountryId) -> None: country_module = import_module(f"policyengine.tax_benefit_models.{country}") - spm = {"geography_kind": "national"} if country == "us" else None + people = [{"age": 40, "employment_income": 50_000}] + extra_variables = ["household_net_income"] + options: dict[str, Any] = {} + if country == "us": + from policyengine_simulation_executor.simulation_runtime import ( + _normalise_policy, + ) + + people = [ + { + "age": 40, + "employment_income": 50_000, + "is_tax_unit_head": True, + }, + {"age": 10, "is_tax_unit_dependent": True}, + ] + extra_variables.append("ctc") + options["spm"] = {"geography_kind": "national"} + options["reform"] = _normalise_policy(US_VALIDATION_REFORM) result = country_module.calculate_household( - people=[{"age": 40, "employment_income": 50_000}], - **({"spm": spm} if spm is not None else {}), + people=people, year=2026, - extra_variables=["household_net_income"], + extra_variables=extra_variables, + **options, ) household = getattr(result, "household", None) if not isinstance(household, Mapping) or "household_net_income" not in household: raise RuntimeError("Stage 12 validation calculation returned no net income") + if country == "us": + tax_unit = getattr(result, "tax_unit", None) + if not isinstance(tax_unit, Mapping) or not isinstance( + tax_unit.get("ctc"), (int, float) + ): + raise RuntimeError("Stage 12 validation reform returned no CTC value") def _check_uk_local_authority_resources() -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index 4e284b6cd..d9f32fc04 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -66,7 +66,10 @@ def test_non_serving_us_calculation_selects_national_without_county( def calculate_household(**kwargs): calls.append(kwargs) - return SimpleNamespace(household={"household_net_income": 1}) + return SimpleNamespace( + household={"household_net_income": 1}, + tax_unit={"ctc": 1}, + ) monkeypatch.setattr( "policyengine_simulation_executor.stage12_worker_validation.import_module", @@ -79,6 +82,33 @@ def calculate_household(**kwargs): assert "household" not in calls[0] +def test_us_validation_normalizes_the_frontend_reform_interval(monkeypatch) -> None: + from policyengine_simulation_executor import stage12_worker_validation + + received = {} + + def calculate_household(**kwargs): + received.update(kwargs) + return SimpleNamespace( + household={"household_net_income": 50_000}, + tax_unit={"ctc": 3_000}, + ) + + monkeypatch.setattr( + stage12_worker_validation, + "import_module", + lambda _: SimpleNamespace(calculate_household=calculate_household), + ) + + _run_non_serving_calculation("us") + + assert received["reform"] == { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01": 3_000} + } + assert received["people"][1]["is_tax_unit_dependent"] is True + assert "ctc" in received["extra_variables"] + + def test_uk_validation_loads_packaged_local_authority_resources() -> None: resolved = load_stage12_bundle() checks: list[str] = [] From 582e7828107806e3ef1d55797f8bd9c6778ef8f4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 00:14:23 +0400 Subject: [PATCH 12/25] Preserve Stage 12 failure details --- .../stage12_runtime/__init__.py | 12 ++ .../stage12_runtime/coordination.py | 89 ++++++--- .../stage12_runtime/failures.py | 172 ++++++++++++++++++ .../stage12_runtime/fanout.py | 13 +- .../stage12_runtime/simulation.py | 58 +++++- .../tests/test_stage12_runtime.py | 141 ++++++++++++-- .../tests/test_stage12_segment_fanout.py | 6 +- 7 files changed, 443 insertions(+), 48 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py 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 4074151c6..78e53c41f 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 @@ -2,6 +2,13 @@ from .aggregation import build_aggregate_report, build_spm_result from .coordination import coordinate_report +from .failures import ( + Stage12ExecutionError, + Stage12FailureDetail, + Stage12InputError, + failure_detail_from_exception, + validate_policy_periods, +) from .fanout import Stage12SegmentRunner, run_segmented_simulation from .partition import stage12_region_groups, stage12_region_groups_for_model from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult @@ -28,12 +35,16 @@ "Stage12SegmentInput", "Stage12SegmentResult", "Stage12SegmentRunner", + "Stage12ExecutionError", + "Stage12FailureDetail", + "Stage12InputError", "build_aggregate_report", "build_segment_inputs", "build_stage12_simulation", "calculate_simulation_frames", "calculate_segment", "coordinate_report", + "failure_detail_from_exception", "run_single_simulation", "run_segmented_simulation", "merge_segment_results", @@ -41,4 +52,5 @@ "stage12_region_groups", "stage12_region_groups_for_model", "should_segment_simulation", + "validate_policy_periods", ] 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 5ccf33c70..6f3c279ad 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 @@ -53,6 +53,11 @@ runtime_store, ) from .failure_recording import persist_failure_record, report_original_failure +from .failures import ( + Stage12ExecutionError, + failure_detail_from_exception, + failure_detail_from_record, +) from .output_planning import ( plan_simulation_input, resolve_report_output_plan, @@ -269,9 +274,6 @@ def coordinate_report( } descriptors: dict[str, SimulationArtifactDescriptor] = {} calls: dict[str, ChildCall] = {} - # A child failure is reported where it occurs and then propagates to the - # report handler below, which must not report the same error twice. - reported_child_failure: Exception | None = None try: def stage_scope( @@ -404,19 +406,28 @@ def stage_scope( failure_scope="child_dispatch", attributes={"simulation_role": simulation.role.value}, ) - reported_child_failure = error + detail = failure_detail_from_exception( + error, + runtime=runtime, + scope="stage12_simulation_dispatch", + default_code="simulation_dispatch_failed", + context={ + "evaluation_id": str(report.evaluation_id), + "simulation_role": simulation.role.value, + }, + ) persist_failure_record( partial( _mark_child_failed, persistence, simulation.simulation_execution_id, - error_code="simulation_dispatch_failed", - error_summary=type(error).__name__, + error_code=detail.error_code, + error_summary=detail.error_summary, ), runtime=runtime, failure_scope="child_dispatch", ) - raise + raise Stage12ExecutionError(detail) from None calls[simulation.role.value] = call persistence.attach_simulation_invocation( simulation.simulation_execution_id, @@ -446,22 +457,45 @@ def stage_scope( failure_scope="child_wait", attributes={"simulation_role": role}, ) - reported_child_failure = error simulation = next( item for item in simulations if item.role.value == role ) - persist_failure_record( - partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code="simulation_invocation_failed", - error_summary=type(error).__name__, - ), - runtime=runtime, - failure_scope="child_wait", - ) - raise + child = persistence.get_simulation(simulation.simulation_execution_id) + if ( + child.status is ComparisonRunLifecycleStatus.FAILED + and child.error_code is not None + ): + detail = failure_detail_from_record( + error_code=child.error_code, + error_summary=child.error_summary, + ) + else: + detail = failure_detail_from_exception( + error, + runtime=runtime, + scope="stage12_simulation_invocation", + default_code="simulation_invocation_failed", + context={ + "evaluation_id": str(report.evaluation_id), + "simulation_role": role, + }, + ) + if child.status not in { + ComparisonRunLifecycleStatus.SUCCEEDED, + ComparisonRunLifecycleStatus.FAILED, + }: + persist_failure_record( + partial( + _mark_child_failed, + persistence, + simulation.simulation_execution_id, + error_code=detail.error_code, + error_summary=detail.error_summary, + ), + runtime=runtime, + failure_scope="child_wait", + ) + raise Stage12ExecutionError(detail) from None baseline = descriptors["baseline"] reform = descriptors["reform"] validate_aligned_outputs(report, output_plan, baseline, reform) @@ -546,13 +580,22 @@ def stage_scope( # Persist any coordinator, child-call, aggregation, or artifact failure # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 - if error is not reported_child_failure: + if not isinstance(error, Stage12ExecutionError): + # Child dispatch and wait failures raise Stage12ExecutionError only + # after reporting the original error where it occurred. report_original_failure( error, runtime=runtime, failure_scope="report_coordination", ) failed_at = datetime.now(UTC) + detail = failure_detail_from_exception( + error, + runtime=runtime, + scope="stage12_report_coordination", + default_code="report_coordination_failed", + context={"evaluation_id": str(report.evaluation_id)}, + ) persist_failure_record( partial( persistence.replace_report, @@ -560,8 +603,8 @@ def stage_scope( update={ "status": ComparisonRunLifecycleStatus.FAILED, "aggregation_status": ComparisonRunAggregationStatus.FAILED, - "error_code": "report_coordination_failed", - "error_summary": type(error).__name__, + "error_code": detail.error_code, + "error_summary": detail.error_summary, "updated_at": failed_at, "completed_at": failed_at, } 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 new file mode 100644 index 000000000..132ccd0e9 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py @@ -0,0 +1,172 @@ +"""Safe, structured failure details for Stage 12 execution.""" + +from __future__ import annotations + +from collections.abc import Mapping +from datetime import date +import logging +from typing import Any + +from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.spm import spm_error_detail +from pydantic import BaseModel, ConfigDict + +from policyengine_simulation_observability.errors import ( + GENERIC_JOB_FAILURE_MESSAGE, + log_and_redact_exception, + make_correlation_id, +) + +logger = logging.getLogger(__name__) + +ERROR_CODE_MAX_LENGTH = 64 +ERROR_SUMMARY_MAX_LENGTH = 512 + + +def _bounded(value: str, maximum: int) -> str: + if len(value) <= maximum: + return value + return value[: maximum - 3] + "..." + + +class Stage12FailureDetail(BaseModel): + """Failure values safe to persist and return from the internal status API.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + error_code: str + error_summary: str + correlation_id: str | None = None + segment_index: int | None = None + + +class Stage12InputError(ValueError): + """An input error whose code and message are safe to return to the caller.""" + + def __init__(self, code: str, message: str) -> None: + self.code = code + self.message = message + super().__init__(message) + + +class Stage12ExecutionError(RuntimeError): + """Carry one safe failure between Stage 12 execution layers.""" + + def __init__(self, detail: Stage12FailureDetail) -> None: + self.detail = detail + super().__init__("Stage 12 execution failed") + + +def _safe_detail( + *, + error_code: str, + error_summary: str, + correlation_id: str | None = None, + segment_index: int | None = None, +) -> Stage12FailureDetail: + return Stage12FailureDetail( + error_code=_bounded(error_code, ERROR_CODE_MAX_LENGTH), + error_summary=_bounded(error_summary, ERROR_SUMMARY_MAX_LENGTH), + correlation_id=correlation_id, + segment_index=segment_index, + ) + + +def failure_detail_from_record( + *, + error_code: str, + error_summary: str | None, +) -> Stage12FailureDetail: + """Recreate a propagated failure without replacing its persisted values.""" + + return _safe_detail( + error_code=error_code, + error_summary=error_summary or "Simulation failed", + ) + + +def failure_detail_from_exception( + error: BaseException, + *, + runtime: ObservabilityRuntime | None, + scope: str, + default_code: str, + context: dict[str, Any] | None = None, + segment_index: int | None = None, +) -> Stage12FailureDetail: + """Retain typed public failures and redact unexpected exception messages.""" + + if isinstance(error, Stage12ExecutionError): + return error.detail + if isinstance(error, Stage12InputError): + return _safe_detail( + error_code=error.code, + error_summary=error.message, + segment_index=segment_index, + ) + spm = spm_error_detail(error) + if spm is not None: + return _safe_detail( + error_code=spm.code, + error_summary=spm.message, + segment_index=segment_index, + ) + + if runtime is not None: + summary = log_and_redact_exception( + error, + runtime=runtime, + scope=scope, + context=context, + ) + correlation_id = summary.removeprefix( + f"{GENERIC_JOB_FAILURE_MESSAGE} (correlation_id=" + ).removesuffix(")") + else: + correlation_id = make_correlation_id() + logger.error( + "Stage 12 %s failed (correlation_id=%s)", + scope, + correlation_id, + exc_info=error, + extra={"correlation_id": correlation_id, **(context or {})}, + ) + summary = f"{GENERIC_JOB_FAILURE_MESSAGE} (correlation_id={correlation_id})" + if segment_index is not None: + summary = f"Segment {segment_index}: {summary}" + return _safe_detail( + error_code=default_code, + error_summary=summary, + correlation_id=correlation_id, + segment_index=segment_index, + ) + + +def validate_policy_periods(policy: Mapping[str, Any]) -> None: + """Require effective dates in the wire formats emitted by the frontend.""" + + for parameter, value in policy.items(): + if not isinstance(value, Mapping): + continue + for raw_period in value: + period = str(raw_period) + parts = period.split(".") + if len(parts) not in {1, 2}: + _raise_invalid_period(parameter, period) + parsed = [] + for part in parts: + try: + parsed.append(date.fromisoformat(part)) + except ValueError: + _raise_invalid_period(parameter, period) + if len(parsed) == 2 and parsed[1] < parsed[0]: + _raise_invalid_period(parameter, period) + + +def _raise_invalid_period(parameter: str, period: str) -> None: + raise Stage12InputError( + "invalid_policy_period", + f"Policy parameter {parameter!r} has invalid effective period " + f"{period!r}; expected YYYY-MM-DD or " + "YYYY-MM-DD.YYYY-MM-DD.", + ) 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 4172d5858..333694d07 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 @@ -18,6 +18,7 @@ ) from .partition import stage12_region_groups_for_model +from .failures import Stage12ExecutionError, failure_detail_from_exception from .segmentation import build_segment_inputs, merge_segment_results from .simulation import SimulationCalculation @@ -145,10 +146,16 @@ def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: type(error).__name__, ) continue + detail = failure_detail_from_exception( + error, + runtime=self.runtime, + scope="stage12_segment_call", + default_code="segment_execution_failed", + context={"segment_index": segment_index}, + segment_index=segment_index, + ) self._cancel_all(handles) - raise RuntimeError( - f"Stage 12 segment {segment_index} failed" - ) from None + raise Stage12ExecutionError(detail) from None pending.discard(index) poll_errors.pop(index, None) progress_made = True 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 aa0026155..45932b0de 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 @@ -9,6 +9,7 @@ from functools import partial from hashlib import sha256 from typing import Any +from uuid import UUID import pandas as pd from policyengine_observability import ObservabilityRuntime @@ -41,6 +42,11 @@ from .dependencies import ComparisonStore, artifact_store, runtime_store from .failure_recording import persist_failure_record, report_original_failure +from .failures import ( + Stage12FailureDetail, + failure_detail_from_exception, + validate_policy_periods, +) from .output_planning import apply_output_plan, validate_output_frames @@ -348,6 +354,31 @@ def descriptor_from_record( ) +def _mark_running_simulation_failed( + persistence: ComparisonStore, + simulation_execution_id: UUID, + *, + detail: Stage12FailureDetail, + failed_at: datetime, +) -> None: + """Record a failure only while the simulation is still running.""" + + latest = persistence.get_simulation(simulation_execution_id) + if latest.status is not ComparisonRunLifecycleStatus.RUNNING: + return + persistence.replace_simulation( + latest.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + + def run_single_simulation( payload: object, context_payload: object, @@ -385,6 +416,7 @@ def run_single_simulation( ) persistence.replace_simulation(running) try: + validate_policy_periods(simulation.policy) input_span = ( runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_INPUT_WRITE)) if runtime is not None @@ -460,18 +492,24 @@ def run_single_simulation( attributes={"simulation_role": simulation.role.value}, ) failed_at = datetime.now(UTC) + detail = failure_detail_from_exception( + error, + runtime=runtime, + scope="stage12_simulation_execution", + default_code="simulation_execution_failed", + context={ + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + "simulation_role": simulation.role.value, + }, + ) persist_failure_record( partial( - persistence.replace_simulation, - running.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_execution_failed", - "error_summary": type(error).__name__, - "updated_at": failed_at, - "completed_at": failed_at, - } - ), + _mark_running_simulation_failed, + persistence, + simulation.simulation_execution_id, + detail=detail, + failed_at=failed_at, ), runtime=runtime, failure_scope="simulation_execution", diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index b11165fb8..4660a60aa 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -96,6 +96,7 @@ from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata, ) +from policyengine_simulation_contract.spm import SPMInputError from policyengine_stage12_persistence import store as persistence_store_module from sqlalchemy.exc import OperationalError @@ -776,7 +777,59 @@ def fail(_simulation): child = store.children[simulation.simulation_execution_id] assert child.status is ComparisonRunLifecycleStatus.FAILED assert child.error_code == "simulation_execution_failed" - assert child.error_summary == "RuntimeError" + assert child.error_summary.startswith("Simulation failed (correlation_id=") + assert "sensitive" not in child.error_summary + + +def test_single_worker_preserves_typed_failure_values() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + + def fail(_simulation): + raise SPMInputError("SPM_YEAR_UNAVAILABLE", "No SPM values for 2026") + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_code == "SPM_YEAR_UNAVAILABLE" + assert child.error_summary == "No SPM values for 2026" + + +def test_single_worker_rejects_bare_year_policy_period_before_calculation() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={ + "policy": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3_000}} + } + ) + store.children[simulation.simulation_execution_id] = _child(simulation) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("calculation must not start"), + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_code == "invalid_policy_period" + assert child.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 _connection_limit_error() -> OperationalError: @@ -864,7 +917,8 @@ def fail(_simulation): "stage12_failure_scope": "simulation_execution", "simulation_role": "baseline", } - assert failure_write == ("failure_write", "ValueError") + assert failure_write[0] == "failure_write" + assert failure_write[1].startswith("Simulation failed (correlation_id=") assert persistence_report[0] == "telemetry" assert persistence_report[1][0] == ( "Stage 12 failure record could not be persisted" @@ -1606,12 +1660,67 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None assert artifacts.aggregate_writes == [] assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "RuntimeError" + assert store.parent.error_summary.startswith("Simulation failed (correlation_id=") assert "sensitive" not in store.parent.error_summary failed_child = store.children[REFORM_ID] assert failed_child.status is ComparisonRunLifecycleStatus.FAILED assert failed_child.error_code == "simulation_invocation_failed" - assert failed_child.error_summary == "RuntimeError" + assert failed_child.error_summary == store.parent.error_summary + + +def test_coordinator_propagates_an_already_persisted_child_failure() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + + class PersistedFailureCall: + object_id = "call-reform" + + def get(self, *, timeout=None): + child = store.children[REFORM_ID] + failed_at = datetime.now(UTC) + store.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + raise RuntimeError("remote wrapper must not replace the child failure") + + class PersistedFailureInvoker: + def spawn(self, *, simulation, **_): + parsed = PlannedSimulationExecutionInput.model_validate(simulation) + if parsed.role is SimulationRole.REFORM: + return PersistedFailureCall() + descriptor = artifacts.add_simulation(parsed, _frames(100.0)) + return ImmediateCall( + "call-baseline", descriptor.model_dump(mode="json"), [] + ) + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=store, + artifacts=artifacts, + invoker=PersistedFailureInvoker(), + aggregator=lambda **_: {"must": "not run"}, + ) + + assert store.children[REFORM_ID].error_code == "invalid_policy_period" + assert store.children[REFORM_ID].error_summary == "The reform period is invalid." + assert store.parent.error_code == "invalid_policy_period" + assert store.parent.error_summary == "The reform period is invalid." @pytest.mark.parametrize( @@ -1665,7 +1774,7 @@ def create_or_resolve_simulation(self, record): assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_code == "report_coordination_failed" - assert store.parent.error_summary == "RuntimeError" + assert store.parent.error_summary.startswith("Simulation failed (correlation_id=") def test_coordinator_builds_its_store_with_the_coordinator_runtime( @@ -1898,11 +2007,20 @@ def fail_aggregation(**_): original_report, failure_write, persistence_report = ordered_events assert original_report[1][1] == {"stage12_failure_scope": "report_coordination"} assert original_report[1][2] is aggregation_error - assert failure_write == ("failure_write", "ValueError") + assert failure_write[0] == "failure_write" + assert failure_write[1].startswith("Simulation failed (correlation_id=") assert isinstance(persistence_report[1][2], OperationalError) assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING +@pytest.mark.xfail( + strict=True, + reason=( + "Known open bug: the coordinator reads a failed child's record to " + "choose its failure detail outside the failure-recording guard, so " + "an unreadable record replaces the child's original failure." + ), +) def test_coordinator_reports_a_child_failure_once_when_its_record_is_unwritable() -> ( None ): @@ -1942,7 +2060,7 @@ def get_simulation(self, simulation_execution_id): assert ordered_events[1] == ("child_read", str(BASELINE_ID)) # The report record is still written after the child record could not be. assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "RuntimeError" + assert store.parent.error_summary.startswith("Simulation failed (correlation_id=") class RefusingDispatchInvoker: @@ -1994,9 +2112,9 @@ def test_coordinator_reports_a_child_dispatch_failure_once() -> None: baseline_child = store.children[BASELINE_ID] assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED assert baseline_child.error_code == "simulation_dispatch_failed" - assert baseline_child.error_summary == "ConnectionError" + assert baseline_child.error_summary.startswith("Simulation failed (correlation_id=") assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "ConnectionError" + assert store.parent.error_summary == baseline_child.error_summary def test_coordinator_rejects_incompatible_row_identity() -> None: @@ -2130,7 +2248,8 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: assert invoker.events == ["spawn:baseline", "spawn:reform"] assert artifacts.aggregate_writes == [] assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == "TimeoutError" + assert store.parent.error_summary.startswith("Simulation failed (correlation_id=") timed_out_child = store.children[BASELINE_ID] assert timed_out_child.status is ComparisonRunLifecycleStatus.FAILED - assert timed_out_child.error_summary == "TimeoutError" + assert timed_out_child.error_code == "simulation_invocation_failed" + assert timed_out_child.error_summary == store.parent.error_summary 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 2c64a1e6f..c4637a69c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -13,6 +13,7 @@ from policyengine_simulation_executor.stage12_runtime import ( SimulationCalculation, + Stage12ExecutionError, Stage12SegmentRunner, build_segment_inputs, calculate_segment, @@ -170,11 +171,14 @@ def test_second_poll_error_cancels_all_started_segments() -> None: errors={0: [ConnectionError("first"), RuntimeError("second")]} ) - with pytest.raises(RuntimeError, match="segment 0 failed"): + with pytest.raises(Stage12ExecutionError) as error: runner.run() assert len(function.calls) == 20 assert all(call.cancelled == 1 for call in function.calls) + assert error.value.detail.error_code == "segment_execution_failed" + assert error.value.detail.segment_index == 0 + assert "second" not in error.value.detail.error_summary def test_dispatch_error_cancels_segments_that_already_started() -> None: From 0e63bc5975410445b9b9503bb1809fa35f2ca76e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:36:36 +0400 Subject: [PATCH 13/25] Observe Stage 12 simulations concurrently --- .../policyengine_stage12_persistence/store.py | 49 +++ .../tests/test_store.py | 67 ++++ .../src/modal/v2_app.py | 3 +- .../stage12_runtime/__init__.py | 2 + .../stage12_runtime/coordination.py | 356 +++++++++++++----- .../stage12_runtime/dependencies.py | 10 + .../stage12_runtime/failures.py | 4 + .../stage12_runtime/fanout.py | 32 +- .../stage12_runtime/simulation.py | 47 ++- .../tests/test_stage12_runtime.py | 211 ++++++++++- .../tests/test_stage12_segment_fanout.py | 15 + 11 files changed, 684 insertions(+), 112 deletions(-) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py index f463d0a32..b9e0ce05d 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -492,6 +492,55 @@ def replace_simulation( ) return _simulation(updated) + def replace_simulation_if_status( + self, + record: ComparisonSimulationRecord, + *, + expected_status: ComparisonRunLifecycleStatus, + ) -> tuple[ComparisonSimulationRecord, bool]: + """Replace a simulation only while it remains in ``expected_status``.""" + + values = _record_values(record) + with self._transaction("replace_simulation_if_status") as connection: + row = ( + connection.execute( + select(comparison_simulations) + .where( + comparison_simulations.c.simulation_execution_id + == record.simulation_execution_id + ) + .with_for_update() + ) + .mappings() + .one_or_none() + ) + if row is None: + raise LookupError( + "comparison simulation " + f"{record.simulation_execution_id} does not exist" + ) + existing = _simulation(row) + require_simulation_identity(existing, record) + if existing.status is not expected_status: + return existing, False + require_lifecycle_transition(existing.status, record.status) + updated = ( + connection.execute( + update(comparison_simulations) + .where( + comparison_simulations.c.simulation_execution_id + == record.simulation_execution_id + ) + .values( + **{name: values[name] for name in SIMULATION_MUTABLE_COLUMNS} + ) + .returning(*comparison_simulations.c) + ) + .mappings() + .one() + ) + return _simulation(updated), True + def attach_simulation_invocation( self, simulation_execution_id: UUID, diff --git a/libs/policyengine-stage12-persistence/tests/test_store.py b/libs/policyengine-stage12-persistence/tests/test_store.py index 9df7938b2..685814cac 100644 --- a/libs/policyengine-stage12-persistence/tests/test_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_store.py @@ -370,6 +370,73 @@ def test_invocation_attachment_is_compare_and_set() -> None: assert result == attached +def test_simulation_status_replacement_is_compare_and_set() -> None: + running = _simulation().model_copy( + update={"status": ComparisonRunLifecycleStatus.RUNNING} + ) + failed = running.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_execution_failed", + "completed_at": NOW, + } + ) + engine = FakeEngine( + [ + FakeResult(row=_row(running)), + FakeResult(row=_row(failed)), + ] + ) + opener = RetryingConnectionOpener() + store = Stage12PersistenceStore( + engine=engine, # type: ignore[arg-type] + open_connection=opener, + ) + + result, replaced = store.replace_simulation_if_status( + failed, + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + + assert replaced is True + assert result == failed + assert opener.operation_names == ["replace_simulation_if_status"] + assert len(engine.connection.statements) == 2 + + +def test_simulation_status_replacement_preserves_a_concurrent_cancellation() -> None: + incomplete = _simulation().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.INCOMPLETE, + "error_code": "cancelled_after_peer_failure", + "completed_at": NOW, + } + ) + succeeded = _simulation().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": "gs://private/baseline.parquet", + "output_sha256": "c" * 64, + "output_schema_version": 1, + "row_identity_columns": ("household.household_id",), + "row_count": 1, + "row_identity_sha256": "d" * 64, + "completed_at": NOW, + } + ) + engine = FakeEngine([FakeResult(row=_row(incomplete))]) + store = Stage12PersistenceStore(engine=engine) # type: ignore[arg-type] + + result, replaced = store.replace_simulation_if_status( + succeeded, + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + + assert replaced is False + assert result == incomplete + assert len(engine.connection.statements) == 1 + + def test_database_url_uses_one_canonical_secret_format() -> None: parsed = normalized_database_url( "postgresql://policyengine_v2_runtime:secret@db.example/postgres" diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 556b39b80..92e6a5ab8 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -291,11 +291,12 @@ def run_single_simulation_us( environment=os.getenv("MODAL_ENVIRONMENT", "local"), ) - def segmented_calculator(simulation): + def segmented_calculator(simulation, cancellation_requested): return run_segmented_simulation( simulation, app_name=APP_NAME, runtime=runtime, + cancellation_requested=cancellation_requested, ) try: 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 78e53c41f..29ca966f8 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 @@ -3,6 +3,7 @@ from .aggregation import build_aggregate_report, build_spm_result from .coordination import coordinate_report from .failures import ( + Stage12Cancellation, Stage12ExecutionError, Stage12FailureDetail, Stage12InputError, @@ -38,6 +39,7 @@ "Stage12ExecutionError", "Stage12FailureDetail", "Stage12InputError", + "Stage12Cancellation", "build_aggregate_report", "build_segment_inputs", "build_stage12_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 6f3c279ad..06595c0fc 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 @@ -7,8 +7,10 @@ from datetime import UTC, datetime from functools import partial from hashlib import sha256 +import logging +import time from typing import Any -from uuid import UUID, uuid4 +from uuid import uuid4 from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( @@ -55,6 +57,7 @@ from .failure_recording import persist_failure_record, report_original_failure from .failures import ( Stage12ExecutionError, + Stage12FailureDetail, failure_detail_from_exception, failure_detail_from_record, ) @@ -64,8 +67,228 @@ validate_output_frames, ) from .simulation import descriptor_from_record, simulation_input_sha256 +from .segmentation import should_segment_simulation SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 +SIMULATION_POLL_INITIAL_SECONDS = 0.25 +SIMULATION_POLL_MAX_SECONDS = 1.0 +SIMULATION_POLL_BACKOFF_FACTOR = 2.0 + +logger = logging.getLogger(__name__) + + +def _next_poll_interval(current: float) -> float: + return min( + current * SIMULATION_POLL_BACKOFF_FACTOR, + SIMULATION_POLL_MAX_SECONDS, + ) + + +def _replace_running_simulation( + *, + store: ComparisonStore, + simulation: PlannedSimulationExecutionInput, + status: ComparisonRunLifecycleStatus, + detail: Stage12FailureDetail, +) -> ComparisonSimulationRecord: + changed_at = datetime.now(UTC) + current = store.get_simulation(simulation.simulation_execution_id) + candidate = current.model_copy( + update={ + "status": status, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": changed_at, + "completed_at": changed_at, + } + ) + result, _ = store.replace_simulation_if_status( + candidate, + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + return result + + +def _cancel_pending_simulations( + *, + calls: dict[str, ChildCall], + simulations: dict[str, PlannedSimulationExecutionInput], + roles: set[str], + store: ComparisonStore, + failed_role: str, +) -> None: + detail = Stage12FailureDetail( + error_code="cancelled_after_peer_failure", + error_summary=f"Cancelled after {failed_role} simulation failed.", + ) + for role in sorted(roles): + if role == failed_role: + continue + simulation = simulations[role] + try: + latest = _replace_running_simulation( + store=store, + simulation=simulation, + status=ComparisonRunLifecycleStatus.INCOMPLETE, + detail=detail, + ) + except Exception as error: # noqa: BLE001 + logger.error( + "Stage 12 peer state cancellation failed for %s (%s)", + role, + type(error).__name__, + ) + continue + if latest.status is not ComparisonRunLifecycleStatus.INCOMPLETE: + continue + if should_segment_simulation(simulation): + # The simulation worker reads the persisted state and cancels the + # segment handles it owns. Modal does not guarantee recursive + # cancellation for calls created with spawn(). + continue + try: + calls[role].cancel() + except Exception as error: # noqa: BLE001 + logger.warning( + "Stage 12 peer cancellation failed for %s (%s)", + role, + type(error).__name__, + ) + + +def _child_failure_detail( + *, + error: BaseException, + simulation: PlannedSimulationExecutionInput, + store: ComparisonStore, + runtime: ObservabilityRuntime | None, +) -> Stage12FailureDetail: + child = store.get_simulation(simulation.simulation_execution_id) + if ( + child.status is ComparisonRunLifecycleStatus.FAILED + and child.error_code is not None + ): + return failure_detail_from_record( + error_code=child.error_code, + error_summary=child.error_summary, + ) + detail = failure_detail_from_exception( + error, + runtime=runtime, + scope="stage12_simulation_invocation", + default_code="simulation_invocation_failed", + context={ + "evaluation_id": str(simulation.evaluation_id), + "simulation_role": simulation.role.value, + }, + ) + persist_failure_record( + partial( + _replace_running_simulation, + store=store, + simulation=simulation, + status=ComparisonRunLifecycleStatus.FAILED, + detail=detail, + ), + runtime=runtime, + failure_scope="child_wait", + ) + return detail + + +def _observe_simulations( + *, + calls: dict[str, ChildCall], + simulations: dict[str, PlannedSimulationExecutionInput], + descriptors: dict[str, SimulationArtifactDescriptor], + store: ComparisonStore, + runtime: ObservabilityRuntime | None, + sleep: Callable[[float], None], + monotonic: Callable[[], float], + timeout_seconds: float, +) -> None: + pending = set(calls) + deadline = monotonic() + timeout_seconds + poll_interval = SIMULATION_POLL_INITIAL_SECONDS + + while pending: + progress_made = False + for role in sorted(pending): + call = calls[role] + try: + result = call.get(timeout=0) + descriptor = SimulationArtifactDescriptor.model_validate(result) + except TimeoutError: + continue + except Exception as error: # noqa: BLE001 + report_original_failure( + error, + runtime=runtime, + failure_scope="child_wait", + attributes={"simulation_role": role}, + ) + detail = _child_failure_detail( + error=error, + simulation=simulations[role], + store=store, + runtime=runtime, + ) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=pending, + store=store, + failed_role=role, + ) + raise Stage12ExecutionError(detail) from None + descriptors[role] = descriptor + pending.remove(role) + progress_made = True + + if not pending: + return + if monotonic() >= deadline: + detail = Stage12FailureDetail( + error_code="simulation_wait_timeout", + error_summary=( + "Stage 12 simulations did not finish within " + f"{timeout_seconds:g} seconds." + ), + ) + wait_timeout = Stage12ExecutionError(detail) + report_original_failure( + wait_timeout, + runtime=runtime, + failure_scope="child_wait_timeout", + ) + for role in sorted(pending): + persist_failure_record( + partial( + _replace_running_simulation, + store=store, + simulation=simulations[role], + status=ComparisonRunLifecycleStatus.FAILED, + detail=detail, + ), + runtime=runtime, + failure_scope="child_wait_timeout", + ) + if not should_segment_simulation(simulations[role]): + try: + calls[role].cancel() + except Exception as error: # noqa: BLE001 + logger.warning( + "Stage 12 timeout cancellation failed for %s (%s)", + role, + type(error).__name__, + ) + raise wait_timeout + sleep(poll_interval) + poll_interval = ( + SIMULATION_POLL_INITIAL_SECONDS + if progress_made + else _next_poll_interval(poll_interval) + ) def _child_record( @@ -194,30 +417,6 @@ def _claim_parent( return claimed, True -def _mark_child_failed( - persistence: ComparisonStore, - simulation_execution_id: UUID, - *, - error_code: str, - error_summary: str, -) -> None: - latest = persistence.get_simulation(simulation_execution_id) - if latest.status is ComparisonRunLifecycleStatus.SUCCEEDED: - return - failed_at = datetime.now(UTC) - persistence.replace_simulation( - latest.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": error_code, - "error_summary": error_summary, - "updated_at": failed_at, - "completed_at": failed_at, - } - ) - ) - - def coordinate_report( payload: object, context_payload: object, @@ -232,6 +431,9 @@ def coordinate_report( output_plan_resolver: Callable[[ReportExecutionInput], Stage12OutputPlan] | None = None, runtime: ObservabilityRuntime | None = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, + wait_timeout_seconds: float = SIMULATION_WAIT_TIMEOUT_SECONDS, ) -> dict[str, Any]: report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) @@ -319,6 +521,9 @@ def stage_scope( ): simulations.append(plan_simulation_input(simulation, output_plan)) planned_simulations = tuple(simulations) + simulations = { + simulation.role.value: simulation for simulation in planned_simulations + } child_state_span = ( runtime.span(stage_plan.name(Stage.STAGE12_CHILD_STATE_CREATE)) if runtime is not None @@ -418,15 +623,22 @@ def stage_scope( ) persist_failure_record( partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code=detail.error_code, - error_summary=detail.error_summary, + _replace_running_simulation, + store=persistence, + simulation=simulation, + status=ComparisonRunLifecycleStatus.FAILED, + detail=detail, ), runtime=runtime, failure_scope="child_dispatch", ) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=set(calls), + store=persistence, + failed_role=simulation.role.value, + ) raise Stage12ExecutionError(detail) from None calls[simulation.role.value] = call persistence.attach_simulation_invocation( @@ -435,67 +647,24 @@ def stage_scope( modal_invocation_id=call.object_id, updated_at=datetime.now(UTC), ) - # Every required call has been started before the coordinator waits. - for role, call in calls.items(): - try: - wait_span = ( - runtime.span( - stage_plan.name(Stage.STAGE12_CHILD_WAIT), - attributes={"simulation_role": role}, - ) - if runtime is not None - else nullcontext() - ) - with wait_span: - descriptors[role] = SimulationArtifactDescriptor.model_validate( - call.get(timeout=SIMULATION_WAIT_TIMEOUT_SECONDS) - ) - except Exception as error: - report_original_failure( - error, - runtime=runtime, - failure_scope="child_wait", - attributes={"simulation_role": role}, - ) - simulation = next( - item for item in simulations if item.role.value == role - ) - child = persistence.get_simulation(simulation.simulation_execution_id) - if ( - child.status is ComparisonRunLifecycleStatus.FAILED - and child.error_code is not None - ): - detail = failure_detail_from_record( - error_code=child.error_code, - error_summary=child.error_summary, - ) - else: - detail = failure_detail_from_exception( - error, - runtime=runtime, - scope="stage12_simulation_invocation", - default_code="simulation_invocation_failed", - context={ - "evaluation_id": str(report.evaluation_id), - "simulation_role": role, - }, - ) - if child.status not in { - ComparisonRunLifecycleStatus.SUCCEEDED, - ComparisonRunLifecycleStatus.FAILED, - }: - persist_failure_record( - partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code=detail.error_code, - error_summary=detail.error_summary, - ), - runtime=runtime, - failure_scope="child_wait", - ) - raise Stage12ExecutionError(detail) from None + # Every required call has been started before the coordinator observes + # baseline and reform together through nonblocking result probes. + wait_span = ( + runtime.span(stage_plan.name(Stage.STAGE12_CHILD_WAIT)) + if runtime is not None + else nullcontext() + ) + with wait_span: + _observe_simulations( + calls=calls, + simulations=simulations, + descriptors=descriptors, + store=persistence, + runtime=runtime, + sleep=sleep, + monotonic=monotonic, + timeout_seconds=wait_timeout_seconds, + ) baseline = descriptors["baseline"] reform = descriptors["reform"] validate_aligned_outputs(report, output_plan, baseline, reform) @@ -581,8 +750,9 @@ def stage_scope( # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 if not isinstance(error, Stage12ExecutionError): - # Child dispatch and wait failures raise Stage12ExecutionError only - # after reporting the original error where it occurred. + # Child dispatch, child and wait-timeout failures raise + # Stage12ExecutionError only after reporting the original failure + # where it occurred. report_original_failure( error, runtime=runtime, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index e5848ebe0..3009789b1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -12,6 +12,7 @@ from policyengine_simulation_contract.stage12_execution import ( ComparisonReportPersistenceResult, ComparisonReportRecord, + ComparisonRunLifecycleStatus, ComparisonSimulationPersistenceResult, ComparisonSimulationRecord, ) @@ -53,6 +54,13 @@ def replace_simulation( record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: ... + def replace_simulation_if_status( + self, + record: ComparisonSimulationRecord, + *, + expected_status: ComparisonRunLifecycleStatus, + ) -> tuple[ComparisonSimulationRecord, bool]: ... + def attach_simulation_invocation( self, simulation_execution_id: UUID, @@ -68,6 +76,8 @@ class ChildCall(Protocol): def get(self, *, timeout: float | None = None) -> object: ... + def cancel(self) -> None: ... + class ChildInvoker(Protocol): def spawn( 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 132ccd0e9..d52661a7f 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 @@ -57,6 +57,10 @@ def __init__(self, detail: Stage12FailureDetail) -> None: super().__init__("Stage 12 execution failed") +class Stage12Cancellation(RuntimeError): + """Stop a simulation whose durable record was cancelled by its parent.""" + + def _safe_detail( *, error_code: str, 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 333694d07..012df1c2e 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 @@ -18,7 +18,11 @@ ) from .partition import stage12_region_groups_for_model -from .failures import Stage12ExecutionError, failure_detail_from_exception +from .failures import ( + Stage12Cancellation, + Stage12ExecutionError, + failure_detail_from_exception, +) from .segmentation import build_segment_inputs, merge_segment_results from .simulation import SimulationCalculation @@ -48,12 +52,14 @@ def __init__( sleep: Callable[[float], None] = time.sleep, poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, + cancellation_requested: Callable[[], bool] | None = None, ) -> None: self.simulation = simulation self.runtime = runtime self.sleep = sleep self.poll_interval_initial_seconds = poll_interval_seconds self.poll_interval_max_seconds = poll_interval_max_seconds + self.cancellation_requested = cancellation_requested or (lambda: False) self.groups = [list(group) for group in groups] if groups is not None else None self.child_function = modal_module.Function.from_name( app_name, @@ -127,6 +133,7 @@ def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: poll_count = 0 while pending: + self._raise_if_cancelled(handles) progress_made = False for index in sorted(pending): segment_index, call = handles[index] @@ -161,7 +168,7 @@ def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: progress_made = True self.runtime.set_context(stage12_segment_poll_count=poll_count) if pending and not progress_made: - self.sleep(current_sleep) + self._wait_with_cancellation(current_sleep, handles) current_sleep = _next_backoff( current_sleep, maximum=self.poll_interval_max_seconds, @@ -173,6 +180,25 @@ def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: raise RuntimeError("Stage 12 segment collection is incomplete") return [result for result in results if result is not None] + def _wait_with_cancellation( + self, + seconds: float, + handles: list[tuple[int, Any]], + ) -> None: + remaining = seconds + while remaining > 0: + self._raise_if_cancelled(handles) + delay = min(1.0, remaining) + self.sleep(delay) + remaining -= delay + self._raise_if_cancelled(handles) + + def _raise_if_cancelled(self, handles: list[tuple[int, Any]]) -> None: + if not self.cancellation_requested(): + return + self._cancel_all(handles) + raise Stage12Cancellation("Stage 12 simulation was cancelled") + @staticmethod def _cancel_all(handles: list[tuple[int, Any]]) -> None: for _, call in handles: @@ -187,9 +213,11 @@ def run_segmented_simulation( *, app_name: str, runtime: ObservabilityRuntime, + cancellation_requested: Callable[[], bool] | None = None, ) -> SimulationCalculation: return Stage12SegmentRunner( simulation, app_name=app_name, runtime=runtime, + cancellation_requested=cancellation_requested, ).run() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 45932b0de..f24364987 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 @@ -8,6 +8,7 @@ from datetime import UTC, datetime from functools import partial from hashlib import sha256 +import logging from typing import Any from uuid import UUID @@ -43,12 +44,15 @@ from .dependencies import ComparisonStore, artifact_store, runtime_store from .failure_recording import persist_failure_record, report_original_failure from .failures import ( + Stage12Cancellation, Stage12FailureDetail, failure_detail_from_exception, validate_policy_periods, ) from .output_planning import apply_output_plan, validate_output_frames +logger = logging.getLogger(__name__) + @dataclass(frozen=True) class SimulationCalculation: @@ -366,7 +370,7 @@ def _mark_running_simulation_failed( latest = persistence.get_simulation(simulation_execution_id) if latest.status is not ComparisonRunLifecycleStatus.RUNNING: return - persistence.replace_simulation( + persistence.replace_simulation_if_status( latest.model_copy( update={ "status": ComparisonRunLifecycleStatus.FAILED, @@ -375,7 +379,8 @@ def _mark_running_simulation_failed( "updated_at": failed_at, "completed_at": failed_at, } - ) + ), + expected_status=ComparisonRunLifecycleStatus.RUNNING, ) @@ -391,7 +396,7 @@ def run_single_simulation( Mapping[str, pd.DataFrame] | SimulationCalculation, ] = calculate_simulation_frames, segmented_calculator: Callable[ - [PlannedSimulationExecutionInput], SimulationCalculation + [PlannedSimulationExecutionInput, Callable[[], bool]], SimulationCalculation ] | None = None, runtime: ObservabilityRuntime | None = None, @@ -404,6 +409,11 @@ def run_single_simulation( child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: return descriptor_from_record(child, simulation).model_dump(mode="json") + if child.status not in { + ComparisonRunLifecycleStatus.PENDING, + ComparisonRunLifecycleStatus.RUNNING, + }: + raise RuntimeError("Stage 12 simulation is no longer active") started = datetime.now(UTC) running = child.model_copy( update={ @@ -414,7 +424,25 @@ def run_single_simulation( "error_summary": None, } ) - persistence.replace_simulation(running) + running, claimed = persistence.replace_simulation_if_status( + running, + expected_status=child.status, + ) + if not claimed: + raise RuntimeError("Stage 12 simulation was cancelled") + + def cancellation_requested() -> bool: + try: + latest = persistence.get_simulation(simulation.simulation_execution_id) + except Exception as error: # noqa: BLE001 + logger.warning( + "Stage 12 cancellation check failed for %s (%s)", + simulation.simulation_execution_id, + type(error).__name__, + ) + return False + return latest.status is not ComparisonRunLifecycleStatus.RUNNING + try: validate_policy_periods(simulation.policy) input_span = ( @@ -430,7 +458,7 @@ def run_single_simulation( from .segmentation import should_segment_simulation if segmented_calculator is not None and should_segment_simulation(simulation): - calculated = segmented_calculator(simulation) + calculated = segmented_calculator(simulation, cancellation_requested) elif calculator is calculate_simulation_frames: calculated = calculate_simulation_frames(simulation, runtime=runtime) else: @@ -466,7 +494,7 @@ def run_single_simulation( uk_local_authority_metadata=uk_local_authority_metadata, ) completed = datetime.now(UTC) - persistence.replace_simulation( + _, replaced = persistence.replace_simulation_if_status( running.model_copy( update={ "status": ComparisonRunLifecycleStatus.SUCCEEDED, @@ -479,9 +507,14 @@ def run_single_simulation( "updated_at": completed, "completed_at": completed, } - ) + ), + expected_status=ComparisonRunLifecycleStatus.RUNNING, ) + if not replaced: + raise Stage12Cancellation("Stage 12 simulation was cancelled") return descriptor.model_dump(mode="json") + except Stage12Cancellation: + raise RuntimeError("Stage 12 simulation was cancelled") from None # Persist any country-package calculation failure before returning a # stable exception to Modal. except Exception as error: # noqa: BLE001 diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 4660a60aa..734199f3b 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -375,6 +375,14 @@ def replace_simulation(self, record): self.children[record.simulation_execution_id] = record return record + def replace_simulation_if_status(self, record, *, expected_status): + with self.lock: + current = self.children[record.simulation_execution_id] + if current.status is not expected_status: + return current, False + self.children[record.simulation_execution_id] = record + return record, True + def attach_simulation_invocation( self, simulation_execution_id, @@ -686,8 +694,9 @@ def test_single_worker_routes_eligible_us_calculation_to_segment_runner() -> Non store.children[simulation.simulation_execution_id] = _child(simulation) received = [] - def segmented_calculator(value): + def segmented_calculator(value, cancellation_requested): received.append(value) + assert cancellation_requested() is False return SimulationCalculation(frames=_frames()) run_single_simulation( @@ -711,7 +720,7 @@ def test_single_worker_honors_explicit_segmentation_opt_out() -> None: store.children[simulation.simulation_execution_id] = _child(simulation) segmented_calls = [] - def unexpected_segmented(value): + def unexpected_segmented(value, cancellation_requested): segmented_calls.append(value) return SimulationCalculation(frames=_frames()) @@ -884,11 +893,13 @@ def test_single_worker_reports_failure_before_an_unwritable_failure_record() -> ordered_events: list[tuple[str, object]] = [] class UnwritableFailureStore(FakeStore): - def replace_simulation(self, record): + def replace_simulation_if_status(self, record, *, expected_status): if record.status is ComparisonRunLifecycleStatus.FAILED: ordered_events.append(("failure_write", record.error_summary)) raise _connection_limit_error() - return super().replace_simulation(record) + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) store = UnwritableFailureStore() simulation = _planned_simulation(SimulationRole.BASELINE) @@ -928,6 +939,44 @@ def fail(_simulation): assert child.status is ComparisonRunLifecycleStatus.RUNNING +def test_late_simulation_success_does_not_overwrite_parent_cancellation() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + + def segmented_calculator(value, cancellation_requested): + current = store.children[value.simulation_execution_id] + cancelled_at = datetime.now(UTC) + store.replace_simulation( + current.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.INCOMPLETE, + "error_code": "cancelled_after_peer_failure", + "error_summary": "Cancelled after reform simulation failed.", + "updated_at": cancelled_at, + "completed_at": cancelled_at, + } + ) + ) + assert cancellation_requested() is True + return SimulationCalculation(frames=_frames()) + + with pytest.raises(RuntimeError, match="Stage 12 simulation was cancelled"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.INCOMPLETE + assert child.error_code == "cancelled_after_peer_failure" + + def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> None: store = FakeStore() simulation = _planned_simulation(SimulationRole.BASELINE) @@ -1223,6 +1272,9 @@ def __init__(self, object_id, future): def get(self, *, timeout=None): return self.future.result(timeout=timeout) + def cancel(self): + self.future.cancel() + class ImmediateCall: def __init__(self, object_id, result, observed_timeouts): @@ -1234,6 +1286,9 @@ def get(self, *, timeout=None): self.observed_timeouts.append(timeout) return self.result + def cancel(self): + return None + class ConcurrentInvoker: def __init__( @@ -1668,6 +1723,132 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None assert failed_child.error_summary == store.parent.error_summary +@pytest.mark.parametrize( + ("failed_role", "failed_id", "pending_role", "pending_id"), + [ + ("reform", REFORM_ID, "baseline", BASELINE_ID), + ("baseline", BASELINE_ID, "reform", REFORM_ID), + ], +) +def test_coordinator_observes_one_role_failure_while_the_peer_is_pending( + failed_role, + failed_id, + pending_role, + pending_id, +) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + + class PendingCall: + object_id = f"call-{pending_role}" + + def __init__(self): + self.timeouts = [] + + def get(self, *, timeout=None): + self.timeouts.append(timeout) + raise TimeoutError() + + def cancel(self): + raise AssertionError("segmented cancellation must be cooperative") + + class FailedCall: + object_id = f"call-{failed_role}" + + def __init__(self): + self.timeouts = [] + + def get(self, *, timeout=None): + self.timeouts.append(timeout) + child = store.children[failed_id] + failed_at = datetime.now(UTC) + store.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + raise RuntimeError("remote wrapper") + + def cancel(self): + return None + + baseline_call = PendingCall() + failed_call = FailedCall() + + class ConcurrentFailureInvoker: + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role + return failed_call if role.value == failed_role else baseline_call + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=store, + artifacts=artifacts, + invoker=ConcurrentFailureInvoker(), + aggregator=lambda **_: {"must": "not run"}, + sleep=lambda _: pytest.fail("failure should be observed before sleeping"), + ) + + assert failed_call.timeouts == [0] + assert baseline_call.timeouts == ([0] if pending_role == "baseline" else []) + pending = store.children[pending_id] + assert pending.status is ComparisonRunLifecycleStatus.INCOMPLETE + assert pending.error_code == "cancelled_after_peer_failure" + assert pending.error_summary == f"Cancelled after {failed_role} simulation failed." + assert store.parent.error_code == "invalid_policy_period" + assert store.parent.error_summary == "The reform period is invalid." + assert artifacts.aggregate_writes == [] + + +def test_dispatch_failure_cancels_an_already_started_peer() -> None: + store = FakeStore() + baseline_call = TimeoutCall("call-baseline") + + class DispatchFailureInvoker: + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role + if role is SimulationRole.REFORM: + raise ConnectionError("dispatch contains sensitive details") + return baseline_call + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=store, + artifacts=FakeArtifacts(), + invoker=DispatchFailureInvoker(), + ) + + baseline = store.children[BASELINE_ID] + reform = store.children[REFORM_ID] + assert baseline.status is ComparisonRunLifecycleStatus.INCOMPLETE + assert baseline.error_code == "cancelled_after_peer_failure" + assert reform.status is ComparisonRunLifecycleStatus.FAILED + assert reform.error_code == "simulation_dispatch_failed" + assert store.parent.error_code == "simulation_dispatch_failed" + + def test_coordinator_propagates_an_already_persisted_child_failure() -> None: store = FakeStore() artifacts = FakeArtifacts() @@ -2205,10 +2386,14 @@ def test_coordinator_resumes_a_matching_running_child_invocation() -> None: class TimeoutCall: def __init__(self, object_id): self.object_id = object_id + self.cancelled = 0 def get(self, *, timeout=None): raise TimeoutError("contains sensitive timeout details") + def cancel(self): + self.cancelled += 1 + class TimeoutInvoker: def __init__(self): @@ -2241,6 +2426,7 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: artifacts=artifacts, invoker=invoker, aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + wait_timeout_seconds=0, ) assert "sensitive" not in str(error.value) @@ -2248,8 +2434,15 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: assert invoker.events == ["spawn:baseline", "spawn:reform"] assert artifacts.aggregate_writes == [] assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary.startswith("Simulation failed (correlation_id=") - timed_out_child = store.children[BASELINE_ID] - assert timed_out_child.status is ComparisonRunLifecycleStatus.FAILED - assert timed_out_child.error_code == "simulation_invocation_failed" - assert timed_out_child.error_summary == store.parent.error_summary + assert store.parent.error_code == "simulation_wait_timeout" + assert store.parent.error_summary == ( + "Stage 12 simulations did not finish within 0 seconds." + ) + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + child.error_code == "simulation_wait_timeout" + for child in store.children.values() + ) 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 c4637a69c..63b22d20f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -13,6 +13,7 @@ from policyengine_simulation_executor.stage12_runtime import ( SimulationCalculation, + Stage12Cancellation, Stage12ExecutionError, Stage12SegmentRunner, build_segment_inputs, @@ -197,3 +198,17 @@ def fail_during_dispatch(payload, *, observability_context): assert len(function.calls) == 5 assert all(call.cancelled == 1 for call in function.calls) + + +def test_persisted_parent_cancellation_cancels_all_segments() -> None: + runner, function, _, sleeps = _runner( + errors={index: [TimeoutError()] for index in range(20)} + ) + checks = iter([False, False, True]) + runner.cancellation_requested = lambda: next(checks, True) + + with pytest.raises(Stage12Cancellation): + runner.run() + + assert sleeps == [0.5] + assert all(call.cancelled == 1 for call in function.calls) From 4ee5cdc9929ccf2b8feba8c7e62439b054c2f175 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:38:58 +0400 Subject: [PATCH 14/25] Document Stage 12 failure handling --- .../stage-12-modal-worker-foundation.md | 33 +++++++++++-- .../tests/test_app.py | 48 +++++++++++++++++++ .../README.md | 16 ++++++- 3 files changed, 92 insertions(+), 5 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index edfe360bf..6bce1330d 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -19,7 +19,8 @@ simulation submission or polling contracts. simulations therefore resolve in Modal `main` without relabeling durable records as `main`. - One v2 report-coordinator invocation starts baseline and reform as distinct - single-simulation calls before awaiting either result. + single-simulation calls, then checks both calls with nonblocking result probes + on every polling pass. - An eligible US national single-simulation call starts 20 region-group calculation calls and merges their typed Parquet results. Because baseline and reform are separate and start before either is awaited, an eligible @@ -303,6 +304,21 @@ entity identifiers before merging. A missing or failed region-group call fails that logical simulation. The parent makes a best-effort cancellation request for every calculation call that it already started. +The report coordinator observes baseline and reform together rather than +blocking on one role first. On the first failure, it copies the failed +simulation's `error_code` and safe `error_summary` to the report, marks the +unresolved peer `incomplete` with `cancelled_after_peer_failure`, and completes +the report as failed without waiting for the peer calculation. A segmented +peer checks its persisted simulation status during bounded waits; after seeing +the cancellation state, it cancels all region-group calls it owns and exits. +Atomic status replacement prevents a late calculation result from replacing +that cancellation record. + +Typed input failures, including invalid policy periods and SPM selection +errors, retain their code and bounded message through the segment, simulation, +and report layers. Unexpected exceptions are logged with their complete stack +trace and persisted only as `Simulation failed (correlation_id=...)`. + Before each v2 deployment, the independent Stage 12 precompute application builds three annual US dataset files and 60 current-law baseline files: one for each of 20 region groups in 2025, 2026, and 2027. Cache identities include the @@ -406,10 +422,18 @@ curl -X POST "${SIMULATION_ENTRYPOINT_URL}/internal/stage12/reports" \ "time_period": "2026", "segmented": true, "baseline": {}, - "reform": {} + "reform": { + "gov.irs.credits.ctc.amount.base[0].amount": { + "2026-01-01.2100-12-31": 3000 + } + } }' ``` +Policy values may be scalars, ISO effective-date mappings, or frontend interval +mappings in `YYYY-MM-DD.YYYY-MM-DD` form. A bare year such as `"2026"` is not +a valid effective-period key. + The `segmented` field is optional. Omitting it uses segmentation when the request meets the eligibility conditions above. `false` explicitly selects the single-container path. `true` requests segmentation but still falls back to @@ -436,6 +460,9 @@ While the report is pending or running, GET returns `202` and `Retry-After: 5`; completed, failed, incomplete, or skipped states return `200`. The response contains the complete temporary parent and child execution metadata, including private artifact references and digests, but never returns artifact contents. +For a failed report, the parent record contains the originating safe failure; +the failed simulation contains the same values, and a peer stopped because of +that failure is `incomplete` with `cancelled_after_peer_failure`. Cloud Run therefore requires no additional object-storage permission for this interface. @@ -469,7 +496,7 @@ starting child simulations. Unsupported automatic inputs are logged with a bounded reason and are not persisted. The report coordinator starts baseline and reform as independent Modal calls, -waits for them, and derives the Stage 12 aggregate. It then restores the +observes both until they succeed, and derives the Stage 12 aggregate. It then restores the production Modal function call by the retained production job identifier and waits up to 15 minutes for that production call to finish. It compares both complete aggregate result objects exactly. The private diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index 4c3a43ee2..34b77ba0c 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -488,6 +488,54 @@ def test_temporary_stage12_poll_returns_terminal_metadata_with_200(backend): ) +def test_temporary_stage12_poll_returns_propagated_failure_values(backend): + report = comparison_report(status=ComparisonRunLifecycleStatus.FAILED).model_copy( + update={ + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + } + ) + baseline = comparison_simulation(SimulationRole.BASELINE).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.INCOMPLETE, + "error_code": "cancelled_after_peer_failure", + "error_summary": "Cancelled after reform simulation failed.", + "completed_at": report.completed_at, + } + ) + reform = comparison_simulation(SimulationRole.REFORM).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "completed_at": report.completed_at, + } + ) + comparison = TemporaryComparisonBackend( + report=report, + simulations=(baseline, reform), + ) + app = create_app( + settings=make_settings(), + backend=backend, + auth_dependency=lambda: None, + comparison_backend=comparison, + ) + + from fastapi.testclient import TestClient + + with TestClient(app) as test_client: + result = test_client.get(f"/internal/stage12/reports/{EVALUATION_ID}") + + assert result.status_code == 200 + payload = result.json() + assert payload["report"]["error_code"] == "invalid_policy_period" + assert payload["report"]["error_summary"] == "The reform period is invalid." + assert payload["simulations"][0]["status"] == "incomplete" + assert payload["simulations"][0]["error_code"] == "cancelled_after_peer_failure" + assert payload["simulations"][1]["error_code"] == "invalid_policy_period" + + def test_temporary_stage12_routes_are_excluded_from_openapi(backend): app = create_app( settings=make_settings(), diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 3ad9e4150..e58b83498 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -171,14 +171,26 @@ simulation artifacts. `segmented=false`, US subnational requests, UK requests, cliff analysis, and labor-supply-response calculations use the single-container Stage 12 path. +After dispatching both logical simulations, the coordinator checks both Modal +calls with nonblocking result probes on every polling pass. It does not wait +for baseline before checking reform. The first failed simulation supplies the +report's `error_code` and safe `error_summary`; the coordinator marks the +unresolved peer `incomplete` with `cancelled_after_peer_failure`. A segmented +peer reads that persisted state, cancels the region-group calls it owns, and +exits without overwriting the cancellation record. Expected input failures +retain their typed code and message. Unexpected exceptions are logged in full +and persisted as a generic message with a correlation ID. + Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 current-law US baseline files and emits `STAGE12_CACHE_MANIFEST_DIGEST=`. The US worker image downloads and verifies the manifest's 63 files during image construction. The release validation checks the recorded manifest, hashes a baseline file, and requires -a real calculation to load that file without recomputing it. The cache manifest -digest is then stored in the separate v2 version manifest. +a real calculation to load that file without recomputing it. Worker validation +also calculates a US household under a nonempty CTC reform using the frontend +interval format `2026-01-01.2100-12-31`. The cache manifest digest is then +stored in the separate v2 version manifest. The cache is an optimization, not required state for a calculation: an absent file causes a normal calculation, and a file missing planned output columns is From 0f5b81a09707979417896fb5f4c6cc248fc1ebd3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:32:12 +0400 Subject: [PATCH 15/25] Reconcile segmented Stage 12 execution with current typed provenance --- .../stage12_runtime/coordination.py | 8 +- .../stage12_runtime/segmentation.py | 20 ++--- .../tests/test_stage12_segment_calculation.py | 86 ++++++++++++++++++- .../tests/test_stage12_worker_validation.py | 6 +- 4 files changed, 101 insertions(+), 19 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 06595c0fc..abd8dcfbd 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 @@ -513,14 +513,16 @@ def stage_scope( ) else: output_plan = output_plan_resolver(report) - simulations = [] + planned_inputs: list[PlannedSimulationExecutionInput] = [] for simulation in (report.baseline, report.reform): with stage_scope( Stage.STAGE12_CHILD_INPUT_PLANNING, {"simulation_role": simulation.role.value}, ): - simulations.append(plan_simulation_input(simulation, output_plan)) - planned_simulations = tuple(simulations) + planned_inputs.append( + plan_simulation_input(simulation, output_plan) + ) + planned_simulations = tuple(planned_inputs) simulations = { simulation.role.value: simulation for simulation in planned_simulations } diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py index 5b37426e1..38b0b76f2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py @@ -4,13 +4,13 @@ from collections.abc import Callable, Mapping, Sequence from hashlib import sha256 -from typing import Any import pandas as pd from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( PlannedSimulationExecutionInput, ) +from policyengine_simulation_contract.spm import SPMProvenance from policyengine_simulation_observability.stages import STAGE12_SEGMENT_STAGES from policyengine_simulation_executor.stage12_artifacts import ( @@ -116,20 +116,18 @@ def calculate_segment( def _merge_calculation_provenance( - values: Sequence[dict[str, Any] | None], -) -> dict[str, Any] | None: + values: Sequence[SPMProvenance | None], +) -> SPMProvenance | None: + """Keep one shared compact receipt only when every segment agrees.""" present = [value for value in values if value is not None] if not present: return None if len(present) != len(values): raise ValueError("segment calculation metadata is incomplete") - selection = present[0].get("spm_config") - if any(value.get("spm_config") != selection for value in present[1:]): - raise ValueError("segment calculation selections do not match") - return { - "spm_config": selection, - "spm_provenance": [value.get("spm_provenance") for value in present], - } + receipt = present[0] + if any(value != receipt for value in present[1:]): + raise ValueError("segment calculation provenance does not match") + return receipt def merge_segment_results( @@ -155,7 +153,7 @@ def merge_segment_results( by_index[result.segment_index] = result decoded: list[dict[str, pd.DataFrame]] = [] - provenance: list[dict[str, Any] | None] = [] + provenance: list[SPMProvenance | None] = [] for segment in expected: result = by_index.get(segment.segment_index) if result is None: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py index ef4159530..c0180680a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py @@ -7,7 +7,16 @@ import pandas as pd import pytest -from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + SimulationRole, +) +from policyengine_simulation_contract.spm import ( + SPMProvenance, + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, +) from policyengine_simulation_executor.stage12_artifacts import ( serialize_simulation_frames, ) @@ -40,10 +49,14 @@ def _frames(index: int) -> dict[str, pd.DataFrame]: } -def _results(simulation): +def _results( + simulation: PlannedSimulationExecutionInput, provenance: SPMProvenance | None = None +) -> list[Stage12SegmentResult]: results = [] for segment in build_segment_inputs(simulation): - payload, _ = serialize_simulation_frames(_frames(segment.segment_index)) + payload, _ = serialize_simulation_frames( + _frames(segment.segment_index), calculation_provenance=provenance + ) results.append( Stage12SegmentResult( segment_index=segment.segment_index, @@ -56,6 +69,73 @@ def _results(simulation): return results +def _spm_provenance() -> SPMProvenance: + """Build an explicit test fixture, not installed runtime version metadata.""" + selection = SPMSelection( + forecast_content_sha256="f" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + return build_spm_provenance( + forecast_id="test-forecast", + forecast_sha256="f" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions( + policyengine="test-fixture", + policyengine_core="test-fixture", + policyengine_us="test-fixture", + spm_calculator="test-fixture", + ), + ) + + +def test_merge_preserves_one_identical_compact_spm_receipt() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + receipt = _spm_provenance() + + calculated = merge_segment_results(simulation, _results(simulation, receipt)) + + assert calculated.calculation_provenance == receipt + assert isinstance(calculated.calculation_provenance, SPMProvenance) + + +@pytest.mark.parametrize("difference", ["missing", "scenario", "versions", "years"]) +def test_merge_rejects_missing_or_incompatible_segment_spm_receipts( + difference: str, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + receipt = _spm_provenance() + results = _results(simulation, receipt) + alternative = { + "missing": None, + "scenario": receipt.model_copy(update={"scenario": "alternative"}), + "versions": receipt.model_copy( + update={ + "runtime_versions": receipt.runtime_versions.model_copy( + update={"policyengine": "different-test-fixture"} + ) + } + ), + "years": receipt.model_copy(update={"years": ("2027",)}), + }[difference] + payload, _ = serialize_simulation_frames( + _frames(0), calculation_provenance=alternative + ) + results[0] = results[0].model_copy( + update={ + "parquet_payload": payload, + "payload_sha256": sha256(payload).hexdigest(), + } + ) + + with pytest.raises(ValueError, match="segment calculation"): + merge_segment_results(simulation, results) + + def test_eligibility_matches_stage12_supported_national_shape() -> None: simulation = _planned_simulation(SimulationRole.BASELINE) assert should_segment_simulation(simulation) is True diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index d9f32fc04..0ef8f0bd6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -82,7 +82,7 @@ def calculate_household(**kwargs): assert "household" not in calls[0] -def test_us_validation_normalizes_the_frontend_reform_interval(monkeypatch) -> None: +def test_us_validation_preserves_the_frontend_reform_interval(monkeypatch) -> None: from policyengine_simulation_executor import stage12_worker_validation received = {} @@ -103,7 +103,9 @@ def calculate_household(**kwargs): _run_non_serving_calculation("us") assert received["reform"] == { - "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01": 3_000} + # The current wrapper compiles the complete interval; do not discard + # its end date while passing the front-end reform to calculate_household. + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3_000} } assert received["people"][1]["is_tax_unit_dependent"] is True assert "ctc" in received["extra_variables"] From 0b5edc1bd0a1431ae0c3a77a5fd41c63d1e405fe Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:07:03 +0400 Subject: [PATCH 16/25] Identify the report and child in Stage 12 cancellation logs Cancellation state-update and cancel-request failures previously logged only the simulation role and error type. They now also log the report's evaluation_id and the child's simulation_execution_id, in the message and as record attributes, and still omit the error's own message. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 60 ++++-- .../tests/test_stage12_runtime.py | 198 ++++++++++++++++++ 2 files changed, 246 insertions(+), 12 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index abd8dcfbd..083a52f4f 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 @@ -84,6 +84,37 @@ def _next_poll_interval(current: float) -> float: ) +def _log_cancellation_failure( + level: int, + message: str, + *, + simulation: PlannedSimulationExecutionInput, + error: Exception, +) -> None: + """Log a cancellation failure with the report and child identifiers. + + ``message`` receives the simulation role, simulation execution ID, + evaluation ID and error type, in that order. The error's own message is + never logged. + """ + + error_type = type(error).__name__ + logger.log( + level, + message, + simulation.role.value, + simulation.simulation_execution_id, + simulation.evaluation_id, + error_type, + extra={ + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + "simulation_role": simulation.role.value, + "error_type": error_type, + }, + ) + + def _replace_running_simulation( *, store: ComparisonStore, @@ -133,10 +164,12 @@ def _cancel_pending_simulations( detail=detail, ) except Exception as error: # noqa: BLE001 - logger.error( - "Stage 12 peer state cancellation failed for %s (%s)", - role, - type(error).__name__, + _log_cancellation_failure( + logging.ERROR, + "Stage 12 could not record cancellation of %s simulation %s " + "in report %s (%s)", + simulation=simulation, + error=error, ) continue if latest.status is not ComparisonRunLifecycleStatus.INCOMPLETE: @@ -149,10 +182,11 @@ def _cancel_pending_simulations( try: calls[role].cancel() except Exception as error: # noqa: BLE001 - logger.warning( - "Stage 12 peer cancellation failed for %s (%s)", - role, - type(error).__name__, + _log_cancellation_failure( + logging.WARNING, + "Stage 12 could not cancel %s simulation %s in report %s (%s)", + simulation=simulation, + error=error, ) @@ -277,10 +311,12 @@ def _observe_simulations( try: calls[role].cancel() except Exception as error: # noqa: BLE001 - logger.warning( - "Stage 12 timeout cancellation failed for %s (%s)", - role, - type(error).__name__, + _log_cancellation_failure( + logging.WARNING, + "Stage 12 could not cancel timed-out %s simulation %s " + "in report %s (%s)", + simulation=simulations[role], + error=error, ) raise wait_timeout sleep(poll_interval) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 734199f3b..636924f73 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -4,6 +4,7 @@ import io import json +import logging import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager, nullcontext @@ -2446,3 +2447,200 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: child.error_code == "simulation_wait_timeout" for child in store.children.values() ) + + +class ScriptedChildCall: + """A started child that keeps running unless a wait error is scripted.""" + + def __init__(self, object_id, *, clock=None, wait_error=None, cancel_error=None): + self.object_id = object_id + self.clock = clock + self.wait_error = wait_error + self.cancel_error = cancel_error + self.probed_at = [] + self.cancel_count = 0 + + def get(self, *, timeout=None): + if self.clock is not None: + self.probed_at.append(self.clock.monotonic()) + if self.wait_error is not None: + raise self.wait_error + raise TimeoutError("still running") + + def cancel(self): + self.cancel_count += 1 + if self.cancel_error is not None: + raise self.cancel_error + + +class ScriptedChildInvoker: + def __init__(self, *, clock=None, wait_errors=None, cancel_errors=None): + self.clock = clock + self.wait_errors = wait_errors or {} + self.cancel_errors = cancel_errors or {} + self.events = [] + self.calls = {} + + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + self.events.append(f"spawn:{role}") + return self._call(role) + + def restore(self, invocation_id): + role = invocation_id.removeprefix("call-") + self.events.append(f"restore:{role}") + return self._call(role) + + def _call(self, role): + call = ScriptedChildCall( + f"call-{role}", + clock=self.clock, + wait_error=self.wait_errors.get(role), + cancel_error=self.cancel_errors.get(role), + ) + self.calls[role] = call + return call + + +class CancellationStateFailingStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + if record.status is ComparisonRunLifecycleStatus.INCOMPLETE: + raise ConnectionError("database contains sensitive details") + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + +def _simulation_id(role: str) -> UUID: + return BASELINE_ID if role == "baseline" else REFORM_ID + + +def _report_with_segmentation(*, segmented: bool) -> ReportExecutionInput: + if segmented: + return _report() + + def unsegmented(role: SimulationRole) -> SimulationExecutionInput: + return _simulation(role).model_copy(update={"options": {"segmented": False}}) + + return ReportExecutionInput( + evaluation_id=EVALUATION_ID, + baseline=unsegmented(SimulationRole.BASELINE), + reform=unsegmented(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), + ) + + +def _coordinate_with(report, store, invoker, **coordinator_options): + coordinator_options.setdefault("artifacts", FakeArtifacts()) + return 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=store, + invoker=invoker, + aggregator=lambda **_: {"must": "not run"}, + **coordinator_options, + ) + + +def _coordination_log(caplog, level): + (record,) = [ + record + for record in caplog.records + if record.name == coordination_module.__name__ and record.levelno == level + ] + return record + + +def _assert_log_identifies_child(record, *, role: str, error_type: str) -> None: + simulation_execution_id = str(_simulation_id(role)) + assert record.evaluation_id == str(EVALUATION_ID) + assert record.simulation_execution_id == simulation_execution_id + assert record.simulation_role == role + assert record.error_type == error_type + message = record.getMessage() + assert str(EVALUATION_ID) in message + assert simulation_execution_id in message + assert "sensitive" not in message + assert record.exc_info is None + + +def test_peer_cancellation_failure_log_identifies_the_report_and_child( + caplog, +) -> None: + store = FakeStore() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + cancel_errors={ + "reform": ConnectionError("cancel request contains sensitive details") + }, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID].error_code == "cancelled_after_peer_failure" + assert store.parent.error_code == "simulation_invocation_failed" + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="reform", + error_type="ConnectionError", + ) + + +def test_peer_state_cancellation_failure_log_identifies_the_report_and_child( + caplog, +) -> None: + store = CancellationStateFailingStore() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert store.parent.error_code == "simulation_invocation_failed" + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="reform", + error_type="ConnectionError", + ) + + +def test_timeout_cancellation_failure_log_identifies_the_report_and_child( + caplog, +) -> None: + store = FakeStore() + invoker = ScriptedChildInvoker( + cancel_errors={ + "reform": ConnectionError("cancel request contains sensitive details") + }, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + wait_timeout_seconds=0, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + assert store.parent.error_code == "simulation_wait_timeout" + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="reform", + error_type="ConnectionError", + ) From 352851d13412fbce37d0b18ef54e2ebfbd29b308 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:08:59 +0400 Subject: [PATCH 17/25] Stop started Stage 12 workers when coordinator startup fails Only a failed spawn used to cancel already-started children. Any other coordinator failure after the first spawn, such as a failed dispatch placeholder write, a failed invocation attach, a failed restore, or a retained-artifact error, reached the report failure handler, which marked the report failed but left every started worker running until its own timeout. The handler now stops every started, unfinished child with the existing cancellation path: it records the child as incomplete with cancelled_after_coordinator_failure, cancels an unsegmented call, and lets a segmented call stop itself once it reads that record. Children that already finished are left unchanged. Child, dispatch and wait-timeout failures keep their own peer cancellation, and cancellation or state-update errors are logged without replacing the original failure. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 11 +- .../README.md | 5 +- .../stage12_runtime/coordination.py | 51 +++-- .../tests/test_stage12_runtime.py | 200 ++++++++++++++++++ 4 files changed, 253 insertions(+), 14 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 6bce1330d..5e3cd6546 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -314,6 +314,13 @@ the cancellation state, it cancels all region-group calls it owns and exits. Atomic status replacement prevents a late calculation result from replacing that cancellation record. +If the coordinator itself fails after starting a simulation, for example while +recording a later dispatch or attaching a started call's invocation +identifier, it stops every started, unfinished simulation in the same way and +marks it `incomplete` with `cancelled_after_coordinator_failure`. Simulations +that already finished are left unchanged, and the report keeps the +coordinator's original failure. + Typed input failures, including invalid policy periods and SPM selection errors, retain their code and bounded message through the segment, simulation, and report layers. Unexpected exceptions are logged with their complete stack @@ -462,7 +469,9 @@ contains the complete temporary parent and child execution metadata, including private artifact references and digests, but never returns artifact contents. For a failed report, the parent record contains the originating safe failure; the failed simulation contains the same values, and a peer stopped because of -that failure is `incomplete` with `cancelled_after_peer_failure`. +that failure is `incomplete` with `cancelled_after_peer_failure`. A simulation +stopped because the coordinator failed is `incomplete` with +`cancelled_after_coordinator_failure`. Cloud Run therefore requires no additional object-storage permission for this interface. diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index e58b83498..2731b323b 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -177,7 +177,10 @@ for baseline before checking reform. The first failed simulation supplies the report's `error_code` and safe `error_summary`; the coordinator marks the unresolved peer `incomplete` with `cancelled_after_peer_failure`. A segmented peer reads that persisted state, cancels the region-group calls it owns, and -exits without overwriting the cancellation record. Expected input failures +exits without overwriting the cancellation record. If the coordinator itself +fails after starting a simulation, it stops every started, unfinished +simulation the same way and marks it `incomplete` with +`cancelled_after_coordinator_failure`. Expected input failures retain their typed code and message. Unexpected exceptions are logged in full and persisted as a generic message with a correlation ID. 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 083a52f4f..97a6a02a9 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 @@ -140,21 +140,36 @@ def _replace_running_simulation( return result +COORDINATOR_FAILURE_CANCELLATION = Stage12FailureDetail( + error_code="cancelled_after_coordinator_failure", + error_summary="Cancelled after report coordination failed.", +) + + +def _peer_failure_cancellation(failed_role: str) -> Stage12FailureDetail: + return Stage12FailureDetail( + error_code="cancelled_after_peer_failure", + error_summary=f"Cancelled after {failed_role} simulation failed.", + ) + + def _cancel_pending_simulations( *, calls: dict[str, ChildCall], simulations: dict[str, PlannedSimulationExecutionInput], roles: set[str], store: ComparisonStore, - failed_role: str, + detail: Stage12FailureDetail, ) -> None: - detail = Stage12FailureDetail( - error_code="cancelled_after_peer_failure", - error_summary=f"Cancelled after {failed_role} simulation failed.", - ) + """Record each started, unfinished child as incomplete, then stop it. + + A child whose record is no longer running has already finished or been + stopped, so it is left unchanged. State-update and cancellation errors are + logged rather than raised, so they never replace the failure that caused + the cancellation. + """ + for role in sorted(roles): - if role == failed_role: - continue simulation = simulations[role] try: latest = _replace_running_simulation( @@ -270,9 +285,9 @@ def _observe_simulations( _cancel_pending_simulations( calls=calls, simulations=simulations, - roles=pending, + roles=pending - {role}, store=store, - failed_role=role, + detail=_peer_failure_cancellation(role), ) raise Stage12ExecutionError(detail) from None descriptors[role] = descriptor @@ -512,6 +527,7 @@ def coordinate_report( } descriptors: dict[str, SimulationArtifactDescriptor] = {} calls: dict[str, ChildCall] = {} + simulations: dict[str, PlannedSimulationExecutionInput] = {} try: def stage_scope( @@ -675,7 +691,7 @@ def stage_scope( simulations=simulations, roles=set(calls), store=persistence, - failed_role=simulation.role.value, + detail=_peer_failure_cancellation(simulation.role.value), ) raise Stage12ExecutionError(detail) from None calls[simulation.role.value] = call @@ -788,14 +804,25 @@ def stage_scope( # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 if not isinstance(error, Stage12ExecutionError): - # Child dispatch, child and wait-timeout failures raise + # Child, dispatch and wait-timeout failures raise # Stage12ExecutionError only after reporting the original failure - # where it occurred. + # and stopping their unfinished peers. Any other failure is + # reported here 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. report_original_failure( error, runtime=runtime, failure_scope="report_coordination", ) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=set(calls) - set(descriptors), + store=persistence, + detail=COORDINATOR_FAILURE_CANCELLATION, + ) failed_at = datetime.now(UTC) detail = failure_detail_from_exception( error, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 636924f73..a34c54ed0 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -2644,3 +2644,203 @@ def test_timeout_cancellation_failure_log_identifies_the_report_and_child( role="reform", error_type="ConnectionError", ) + + +class StartupWriteError(Exception): + """A coordinator persistence failure after a child may have started.""" + + +class PlaceholderWriteFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def replace_simulation(self, record): + # The coordinator's only plain replacement is the dispatch placeholder. + if record.simulation_execution_id == self.failing_simulation_id: + raise StartupWriteError("placeholder write contains sensitive details") + return super().replace_simulation(record) + + +class AttachFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def attach_simulation_invocation(self, simulation_execution_id, **values): + if simulation_execution_id == self.failing_simulation_id: + raise StartupWriteError("attach contains sensitive details") + return super().attach_simulation_invocation(simulation_execution_id, **values) + + +def _assert_cancelled_after_coordinator_failure(child) -> None: + assert child.status is ComparisonRunLifecycleStatus.INCOMPLETE + assert child.error_code == "cancelled_after_coordinator_failure" + assert child.error_summary == "Cancelled after report coordination failed." + assert child.completed_at is not None + + +def _assert_original_coordinator_failure(error, store) -> None: + assert isinstance(error.value.__context__, StartupWriteError) + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "report_coordination_failed" + assert "sensitive" not in store.parent.error_summary + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_failed_placeholder_write_stops_the_started_baseline(segmented) -> None: + store = PlaceholderWriteFailingStore(REFORM_ID) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(_report_with_segmentation(segmented=segmented), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == ["spawn:baseline"] + _assert_cancelled_after_coordinator_failure(store.children[BASELINE_ID]) + # A segmented worker stops itself, and its region-group calls, once it + # reads the cancelled record; Modal does not cancel spawned calls + # recursively. + assert invoker.calls["baseline"].cancel_count == (0 if segmented else 1) + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + + +@pytest.mark.parametrize( + ("failing_role", "started_roles"), + [ + ("baseline", ["baseline"]), + ("reform", ["baseline", "reform"]), + ], +) +def test_failed_invocation_attach_stops_every_started_child( + failing_role, + started_roles, +) -> None: + store = AttachFailingStore(_simulation_id(failing_role)) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == [f"spawn:{role}" for role in started_roles] + # The child whose invocation could not be attached is stopped as well. + for role in started_roles: + _assert_cancelled_after_coordinator_failure( + store.children[_simulation_id(role)] + ) + assert invoker.calls[role].cancel_count == 1 + if failing_role == "baseline": + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + + +def test_startup_failure_stops_a_call_restored_from_an_earlier_attempt() -> None: + report = _report_with_segmentation(segmented=False) + store = PlaceholderWriteFailingStore(REFORM_ID) + store.children[BASELINE_ID] = _child(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-baseline", + "started_at": NOW, + } + ) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(report, store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == ["restore:baseline"] + _assert_cancelled_after_coordinator_failure(store.children[BASELINE_ID]) + assert invoker.calls["baseline"].cancel_count == 1 + + +def test_startup_failure_leaves_a_finished_child_unchanged() -> None: + report = _report_with_segmentation(segmented=False) + store = AttachFailingStore(REFORM_ID) + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + finished_baseline = _child(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": baseline_descriptor.artifact.uri, + "output_sha256": baseline_descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": baseline_descriptor.row_identity.identifier_columns, + "row_count": baseline_descriptor.row_identity.row_count, + "row_identity_sha256": baseline_descriptor.row_identity.identity_sha256, + "completed_at": NOW, + } + ) + store.children[BASELINE_ID] = finished_baseline + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(report, store, invoker, artifacts=artifacts) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == ["spawn:reform"] + assert store.children[BASELINE_ID] == finished_baseline + _assert_cancelled_after_coordinator_failure(store.children[REFORM_ID]) + assert invoker.calls["reform"].cancel_count == 1 + + +def test_startup_cancellation_error_does_not_replace_the_original_failure( + caplog, +) -> None: + store = PlaceholderWriteFailingStore(REFORM_ID) + invoker = ScriptedChildInvoker( + cancel_errors={ + "baseline": ConnectionError("cancel request contains sensitive details") + } + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.calls["baseline"].cancel_count == 1 + _assert_cancelled_after_coordinator_failure(store.children[BASELINE_ID]) + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="baseline", + error_type="ConnectionError", + ) + + +def test_startup_cancellation_state_error_does_not_replace_the_original_failure( + caplog, +) -> None: + class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): + pass + + store = FailingStore(REFORM_ID) + invoker = ScriptedChildInvoker() + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + _assert_original_coordinator_failure(error, store) + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="baseline", + error_type="ConnectionError", + ) From b7d3b1cf950c27d62dbd96f7557e8fef52590249 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:10:53 +0400 Subject: [PATCH 18/25] Open the Stage 12 child wait window at coordinator entry The child-wait deadline used to start when observation began, after claiming, output planning (including the country model load), retained-artifact reads and dispatch. Setup time therefore came out of the coordinator's remaining Modal time rather than the children's window. The 3,000 s window now opens when coordinate_report starts, so setup shortens the children's window and the coordinator always keeps 2,100 s of its 5,100 s Modal timeout after the window: up to 900 s for the production-result wait and 1,200 s to read results, aggregate, write artifacts and record the outcome. The timeout summary now says the window is measured from the coordinator's start. A new test runs the coordinator with 0, 400 and 2,000 s of simulated setup, checks that the window closes 3,000 s after entry, and checks the declared Modal timeouts still leave the production wait plus at least 600 s. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 9 ++ .../README.md | 23 ++--- .../stage12_runtime/coordination.py | 24 +++++- .../tests/test_stage12_runtime.py | 86 ++++++++++++++++++- 4 files changed, 129 insertions(+), 13 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 5e3cd6546..127ea110b 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -314,6 +314,15 @@ the cancellation state, it cancels all region-group calls it owns and exits. Atomic status replacement prevents a late calculation result from replacing that cancellation record. +Both simulations share one 3,000-second wait window. It opens when the +coordinator starts, so the coordinator's own setup, including loading the +country model and dispatching both simulations, comes out of that window and +not out of the time it needs afterwards. When the window closes, the +coordinator records each unfinished simulation as `failed` with +`simulation_wait_timeout` and stops it. Its 5,100-second Modal timeout leaves +2,100 seconds after the window: up to 900 seconds for the production-result +wait and 1,200 seconds to aggregate, write artifacts, and record the outcome. + If the coordinator itself fails after starting a simulation, for example while recording a later dispatch or attaching a started call's invocation identifier, it stops every started, unfinished simulation in the same way and diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 2731b323b..702b8dd3b 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -173,16 +173,19 @@ single-container Stage 12 path. After dispatching both logical simulations, the coordinator checks both Modal calls with nonblocking result probes on every polling pass. It does not wait -for baseline before checking reform. The first failed simulation supplies the -report's `error_code` and safe `error_summary`; the coordinator marks the -unresolved peer `incomplete` with `cancelled_after_peer_failure`. A segmented -peer reads that persisted state, cancels the region-group calls it owns, and -exits without overwriting the cancellation record. If the coordinator itself -fails after starting a simulation, it stops every started, unfinished -simulation the same way and marks it `incomplete` with -`cancelled_after_coordinator_failure`. Expected input failures -retain their typed code and message. Unexpected exceptions are logged in full -and persisted as a generic message with a correlation ID. +for baseline before checking reform. Both simulations share one 3,000-second +wait window that opens when the coordinator starts, so coordinator setup comes +out of that window. When it closes, every unfinished simulation is recorded as +`failed` with `simulation_wait_timeout` and stopped. The first failed +simulation supplies the report's `error_code` and safe `error_summary`; the +coordinator marks the unresolved peer `incomplete` with +`cancelled_after_peer_failure`. A segmented peer reads that persisted state, +cancels the region-group calls it owns, and exits without overwriting the +cancellation record. If the coordinator itself fails after starting a +simulation, it stops every started, unfinished simulation the same way and +marks it `incomplete` with `cancelled_after_coordinator_failure`. Expected +input failures retain their typed code and message. Unexpected exceptions are +logged in full and persisted as a generic message with a correlation ID. Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 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 97a6a02a9..8c213b016 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 @@ -69,6 +69,22 @@ from .simulation import descriptor_from_record, simulation_input_sha256 from .segmentation import should_segment_simulation +# One child-wait window opens when `coordinate_report` starts. It covers +# coordinator setup (claiming, output planning including the country model +# load, child-state creation, retained-artifact reads and dispatch) and the +# concurrent wait for both simulations, so setup time shortens the children's +# window instead of the time left after it. The deadline is checked after each +# polling pass, so the coordinator notices a closed window at most one poll +# interval (1 s) late. A simulation worker may have a longer Modal timeout +# (3,900 s for US), but the report stops waiting for it, and cancels it, when +# this window closes. +# +# Budget against the coordinator's 5,100 s Modal timeout in `v2_app.py`: the +# 3,000 s window leaves 2,100 s. Up to 900 s of that is the production-result +# wait (`PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS`), which leaves 1,200 s to read +# and validate both simulation artifacts, aggregate, write the aggregate and +# comparison artifacts, and record the outcome or failure before Modal stops +# the coordinator. SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 SIMULATION_POLL_INITIAL_SECONDS = 0.25 SIMULATION_POLL_MAX_SECONDS = 1.0 @@ -254,10 +270,10 @@ def _observe_simulations( runtime: ObservabilityRuntime | None, sleep: Callable[[float], None], monotonic: Callable[[], float], + deadline: float, timeout_seconds: float, ) -> None: pending = set(calls) - deadline = monotonic() + timeout_seconds poll_interval = SIMULATION_POLL_INITIAL_SECONDS while pending: @@ -301,7 +317,8 @@ def _observe_simulations( error_code="simulation_wait_timeout", error_summary=( "Stage 12 simulations did not finish within " - f"{timeout_seconds:g} seconds." + f"{timeout_seconds:g} seconds of the report coordinator " + "starting." ), ) wait_timeout = Stage12ExecutionError(detail) @@ -486,6 +503,8 @@ def coordinate_report( monotonic: Callable[[], float] = time.monotonic, wait_timeout_seconds: float = SIMULATION_WAIT_TIMEOUT_SECONDS, ) -> dict[str, Any]: + # Open the shared child-wait window before any setup work. + child_wait_deadline = monotonic() + wait_timeout_seconds report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) @@ -717,6 +736,7 @@ def stage_scope( runtime=runtime, sleep=sleep, monotonic=monotonic, + deadline=child_wait_deadline, timeout_seconds=wait_timeout_seconds, ) baseline = descriptors["baseline"] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index a34c54ed0..3b87b71ac 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -2,9 +2,11 @@ from __future__ import annotations +import importlib import io import json import logging +import sys import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager, nullcontext @@ -15,6 +17,7 @@ import pandas as pd import pytest +from fixtures.fake_modal import install_fake_modal from policyengine_observability import ( DeploymentIdentity, LoggingConfig, @@ -91,6 +94,9 @@ from policyengine_simulation_executor.stage12_runtime.aggregation import ( validate_uk_local_authority_metadata, ) +from policyengine_simulation_executor.stage12_runtime.comparison import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, +) from policyengine_simulation_executor.stage12_runtime.dependencies import ( runtime_store, ) @@ -2437,7 +2443,8 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_code == "simulation_wait_timeout" assert store.parent.error_summary == ( - "Stage 12 simulations did not finish within 0 seconds." + "Stage 12 simulations did not finish within 0 seconds of the report " + "coordinator starting." ) assert all( child.status is ComparisonRunLifecycleStatus.FAILED @@ -2844,3 +2851,80 @@ class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): role="baseline", error_type="ConnectionError", ) + + +class SteppedClock: + def __init__(self, start_seconds=10_000.0): + self.current_seconds = start_seconds + + def monotonic(self): + return self.current_seconds + + def advance(self, elapsed_seconds): + self.current_seconds += elapsed_seconds + + +def _declared_modal_timeouts(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.v2_app", None) + modal_app = importlib.import_module("src.modal.v2_app") + return { + name: options.get("timeout") for name, options in modal_app.app.function_calls + } + + +# After the child window closes and the production-result wait ends, the +# coordinator still has to read both simulation artifacts, aggregate, write its +# artifacts, and record the outcome before Modal stops it. +MINIMUM_FINISHING_SECONDS = 600 + + +@pytest.mark.parametrize("setup_seconds", [0.0, 400.0, 2_000.0]) +def test_child_wait_window_opens_when_the_coordinator_starts( + monkeypatch, + setup_seconds, +) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + poll_seconds = coordination_module.SIMULATION_POLL_MAX_SECONDS + clock = SteppedClock() + coordinator_started_at = clock.monotonic() + + def slow_output_planning(_report): + # Loading the country model is coordinator setup, not child time. + clock.advance(setup_seconds) + return _output_plan() + + store = FakeStore() + invoker = ScriptedChildInvoker(clock=clock) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report(), + store, + invoker, + output_plan_resolver=slow_output_planning, + monotonic=clock.monotonic, + sleep=clock.advance, + ) + + probed_at = invoker.calls["baseline"].probed_at + child_waits_began_after = probed_at[0] - coordinator_started_at + window_closed_after = probed_at[-1] - coordinator_started_at + assert child_waits_began_after == setup_seconds + # Setup shortens the children's window rather than the time after it. + assert window_seconds <= window_closed_after < window_seconds + poll_seconds + assert probed_at[-1] - probed_at[0] == pytest.approx( + window_seconds - setup_seconds, abs=poll_seconds + ) + assert ( + declared_timeouts["coordinate_report"] + - window_closed_after + - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + >= MINIMUM_FINISHING_SECONDS + ) + assert store.parent.error_code == "simulation_wait_timeout" + assert all( + child.error_code == "simulation_wait_timeout" + for child in store.children.values() + ) From 48c233ecd69cff21554d455767768c438e73a14c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:03:32 +0400 Subject: [PATCH 19/25] Size the Stage 12 wait window and coordinator timeout from worker timeouts The coordinator stopped waiting 3,000 s after it started, although a US simulation worker may run for 3,900 s and each of its region-group calls for 3,600 s, so US simulations could never use the time Modal allowed them. The coordinator's 5,100 s timeout was not derived from either. Following the rule from #793, a new dependency-free stage12_timeouts module defines the worker timeouts, sets the shared child-wait window to the longest of them (3,900 s), and derives the coordinator's Modal timeout as that window plus the 900 s production-result wait plus 600 s to read results, aggregate, write artifacts and record the outcome: 5,400 s. A complete testing-environment US report took 288 s, so 600 s of finishing time is ample. The Modal app, coordinator and comparison step all read these constants. The last poll no longer sleeps past the deadline, so the window closes on time instead of up to 1 s late and the finishing time is not eroded. Tests read the declared Modal timeouts and check that the window covers every simulation worker's timeout and that the coordinator's timeout leaves the production wait plus at least 600 s. They also check that setup time shortens the window and that an exhausted window gives the children no extra time. Only the Modal timeout values change: no image, secret or environment variable attachment changes. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 16 +++--- .../README.md | 7 ++- .../src/modal/v2_app.py | 15 +++-- .../stage12_runtime/comparison.py | 4 +- .../stage12_runtime/coordination.py | 26 +++------ .../stage12_timeouts.py | 41 +++++++++++++ .../tests/test_stage12_modal_app.py | 2 +- .../tests/test_stage12_runtime.py | 57 +++++++++++++++++-- 8 files changed, 126 insertions(+), 42 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 127ea110b..ce5a58b80 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -314,14 +314,14 @@ the cancellation state, it cancels all region-group calls it owns and exits. Atomic status replacement prevents a late calculation result from replacing that cancellation record. -Both simulations share one 3,000-second wait window. It opens when the -coordinator starts, so the coordinator's own setup, including loading the -country model and dispatching both simulations, comes out of that window and -not out of the time it needs afterwards. When the window closes, the -coordinator records each unfinished simulation as `failed` with -`simulation_wait_timeout` and stops it. Its 5,100-second Modal timeout leaves -2,100 seconds after the window: up to 900 seconds for the production-result -wait and 1,200 seconds to aggregate, write artifacts, and record the outcome. +Both simulations share one 3,900-second wait window, the longest simulation +worker timeout. It opens when the coordinator starts, so the coordinator's own +setup, including loading the country model and dispatching both simulations, +comes out of that window and not out of the time it needs afterwards. When the +window closes, the coordinator records each unfinished simulation as `failed` +with `simulation_wait_timeout` and stops it. The coordinator's 5,400-second +Modal timeout is that window plus up to 900 seconds for the production-result +wait plus 600 seconds to aggregate, write artifacts, and record the outcome. If the coordinator itself fails after starting a simulation, for example while recording a later dispatch or attaching a started call's invocation diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 702b8dd3b..838fa2d22 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -173,9 +173,10 @@ single-container Stage 12 path. After dispatching both logical simulations, the coordinator checks both Modal calls with nonblocking result probes on every polling pass. It does not wait -for baseline before checking reform. Both simulations share one 3,000-second -wait window that opens when the coordinator starts, so coordinator setup comes -out of that window. When it closes, every unfinished simulation is recorded as +for baseline before checking reform. Both simulations share one 3,900-second +wait window, the longest simulation worker timeout, that opens when the +coordinator starts, so coordinator setup comes out of that window. When it +closes, every unfinished simulation is recorded as `failed` with `simulation_wait_timeout` and stopped. The first failed simulation supplies the report's `error_code` and safe `error_summary`; the coordinator marks the unresolved peer `incomplete` with diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 92e6a5ab8..34ca1ccee 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -38,6 +38,11 @@ assertion_values, load_stage12_bundle, ) +from policyengine_simulation_executor.stage12_timeouts import ( + COORDINATOR_TIMEOUT_SECONDS, + UK_SIMULATION_TIMEOUT_SECONDS, + US_SIMULATION_TIMEOUT_SECONDS, +) from src.modal.bundle_data import bundle_data_install_command from src.modal.static_runtime_files import add_static_runtime_files from src.modal.stage12_cache_image import fetch_stage12_cache @@ -268,7 +273,7 @@ def run_single_simulation_segment_us( image=us_worker_image, cpu=8.0, memory=32768, - timeout=3900, + timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, secrets=worker_secrets, @@ -328,7 +333,7 @@ def segmented_calculator(simulation, cancellation_requested): image=uk_worker_image, cpu=8.0, memory=32768, - timeout=3000, + timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, secrets=worker_secrets, @@ -377,9 +382,9 @@ def run_single_simulation_uk( image=coordinator_image, cpu=2.0, memory=8192, - # Allow the two logical simulation parents to start and collect all region - # calculations, then retain bounded aggregation/persistence overhead. - timeout=5100, + # The shared child-wait window (the longest simulation worker timeout), + # the production-result wait, and time to finish and record the report. + timeout=COORDINATOR_TIMEOUT_SECONDS, retries=0, max_containers=10, secrets=worker_secrets, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py index 99055a349..2c11a190d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py @@ -17,10 +17,12 @@ from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore from policyengine_simulation_executor.stage12_result_comparison import compare_results +from policyengine_simulation_executor.stage12_timeouts import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, +) from .dependencies import ChildInvoker, ComparisonStore -PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS = 900 logger = logging.getLogger(__name__) 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 8c213b016..9dcafc053 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 @@ -39,6 +39,9 @@ deserialize_simulation_frames, deserialize_uk_local_authority_metadata, ) +from policyengine_simulation_executor.stage12_timeouts import ( + SIMULATION_WAIT_TIMEOUT_SECONDS, +) from .aggregation import ( build_aggregate_report, @@ -69,23 +72,6 @@ from .simulation import descriptor_from_record, simulation_input_sha256 from .segmentation import should_segment_simulation -# One child-wait window opens when `coordinate_report` starts. It covers -# coordinator setup (claiming, output planning including the country model -# load, child-state creation, retained-artifact reads and dispatch) and the -# concurrent wait for both simulations, so setup time shortens the children's -# window instead of the time left after it. The deadline is checked after each -# polling pass, so the coordinator notices a closed window at most one poll -# interval (1 s) late. A simulation worker may have a longer Modal timeout -# (3,900 s for US), but the report stops waiting for it, and cancels it, when -# this window closes. -# -# Budget against the coordinator's 5,100 s Modal timeout in `v2_app.py`: the -# 3,000 s window leaves 2,100 s. Up to 900 s of that is the production-result -# wait (`PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS`), which leaves 1,200 s to read -# and validate both simulation artifacts, aggregate, write the aggregate and -# comparison artifacts, and record the outcome or failure before Modal stops -# the coordinator. -SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 SIMULATION_POLL_INITIAL_SECONDS = 0.25 SIMULATION_POLL_MAX_SECONDS = 1.0 SIMULATION_POLL_BACKOFF_FACTOR = 2.0 @@ -312,7 +298,8 @@ def _observe_simulations( if not pending: return - if monotonic() >= deadline: + remaining_seconds = deadline - monotonic() + if remaining_seconds <= 0: detail = Stage12FailureDetail( error_code="simulation_wait_timeout", error_summary=( @@ -351,7 +338,8 @@ def _observe_simulations( error=error, ) raise wait_timeout - sleep(poll_interval) + # Never sleep past the deadline, so the window closes on time. + sleep(min(poll_interval, remaining_seconds)) poll_interval = ( SIMULATION_POLL_INITIAL_SECONDS if progress_made diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py new file mode 100644 index 000000000..0f7385517 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py @@ -0,0 +1,41 @@ +"""Modal timeouts and waits shared by the Stage 12 app and report coordinator. + +This module has no imports, so the Modal app declaration can read it without +loading the Stage 12 runtime. +""" + +# Modal timeouts of the simulation workers that the report coordinator starts. +# A US worker starts 20 region-group calls (3,600 s each), then merges their +# results and writes one artifact. +US_SIMULATION_TIMEOUT_SECONDS = 3_900 +UK_SIMULATION_TIMEOUT_SECONDS = 3_000 + +# One child-wait window opens when `coordinate_report` starts. It covers +# coordinator setup (claiming, output planning including the country model +# load, child-state creation, retained-artifact reads and dispatch) and the +# concurrent wait for both simulations, so setup time shortens the children's +# window instead of the time left after it. The window equals the longest +# simulation worker timeout, so the coordinator only stops waiting for a worker +# early by its own setup time and the time the worker spent queued in Modal. +# The coordinator's last poll never sleeps past the end of the window. +SIMULATION_WAIT_TIMEOUT_SECONDS = max( + US_SIMULATION_TIMEOUT_SECONDS, + UK_SIMULATION_TIMEOUT_SECONDS, +) + +# An automatic (shadow) report then waits up to 15 minutes for the production +# result it compares against. +PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS = 900 + +# Time left after both waits to read and validate both simulation artifacts, +# aggregate, write the aggregate and comparison artifacts, and record the +# outcome or failure. A complete testing-environment US report, including all +# 40 region-group calculations, took 288 s. +REPORT_FINISHING_SECONDS = 600 + +# 3,900 s + 900 s + 600 s = 5,400 s. +COORDINATOR_TIMEOUT_SECONDS = ( + SIMULATION_WAIT_TIMEOUT_SECONDS + + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + + REPORT_FINISHING_SECONDS +) 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 3b6a952a8..f5e0c1094 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -231,7 +231,7 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert functions["run_single_simulation_us"]["max_containers"] == 10 assert functions["run_single_simulation_uk"]["max_containers"] == 10 assert functions["coordinate_report"]["image"] is module.coordinator_image - assert functions["coordinate_report"]["timeout"] == 5100 + assert functions["coordinate_report"]["timeout"] == 5400 assert functions["coordinate_report"]["max_containers"] == 10 assert "asgi_app" not in vars(module) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 3b87b71ac..222576e66 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -2879,6 +2879,30 @@ def _declared_modal_timeouts(monkeypatch): MINIMUM_FINISHING_SECONDS = 600 +def test_declared_modal_timeouts_fit_the_shared_wait_window(monkeypatch) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + child_worker_timeouts = { + name: timeout + for name, timeout in declared_timeouts.items() + if name.startswith("run_single_simulation") + } + + assert set(child_worker_timeouts) == { + "run_single_simulation_us", + "run_single_simulation_uk", + "run_single_simulation_segment_us", + } + # The window covers every worker's own Modal timeout. + assert all(window_seconds >= timeout for timeout in child_worker_timeouts.values()) + assert ( + declared_timeouts["coordinate_report"] + >= window_seconds + + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + + MINIMUM_FINISHING_SECONDS + ) + + @pytest.mark.parametrize("setup_seconds", [0.0, 400.0, 2_000.0]) def test_child_wait_window_opens_when_the_coordinator_starts( monkeypatch, @@ -2886,7 +2910,6 @@ def test_child_wait_window_opens_when_the_coordinator_starts( ) -> None: declared_timeouts = _declared_modal_timeouts(monkeypatch) window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS - poll_seconds = coordination_module.SIMULATION_POLL_MAX_SECONDS clock = SteppedClock() coordinator_started_at = clock.monotonic() @@ -2913,10 +2936,8 @@ def slow_output_planning(_report): window_closed_after = probed_at[-1] - coordinator_started_at assert child_waits_began_after == setup_seconds # Setup shortens the children's window rather than the time after it. - assert window_seconds <= window_closed_after < window_seconds + poll_seconds - assert probed_at[-1] - probed_at[0] == pytest.approx( - window_seconds - setup_seconds, abs=poll_seconds - ) + assert window_closed_after == pytest.approx(window_seconds) + assert probed_at[-1] - probed_at[0] == pytest.approx(window_seconds - setup_seconds) assert ( declared_timeouts["coordinate_report"] - window_closed_after @@ -2928,3 +2949,29 @@ def slow_output_planning(_report): child.error_code == "simulation_wait_timeout" for child in store.children.values() ) + + +def test_exhausted_wait_window_gives_the_children_no_extra_time() -> None: + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + clock = SteppedClock() + + def setup_that_uses_the_whole_window(_report): + clock.advance(window_seconds + 60.0) + return _output_plan() + + store = FakeStore() + invoker = ScriptedChildInvoker(clock=clock) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + output_plan_resolver=setup_that_uses_the_whole_window, + monotonic=clock.monotonic, + sleep=lambda _: pytest.fail("an exhausted window must not wait"), + ) + + assert [len(call.probed_at) for call in invoker.calls.values()] == [1, 1] + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + assert store.parent.error_code == "simulation_wait_timeout" From 786b7aaec711d1007c90228ac826390f36d77da1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:05:34 +0400 Subject: [PATCH 20/25] Record cancelled Stage 12 simulations as failed and always stop them Adopt #793's cancellation practice in the concurrent coordinator: - Every simulation the coordinator stops, after a peer failure, a wait timeout or a coordinator failure, is recorded as failed with simulation_cancelled, and its error_summary gives the reason, for example "Cancelled after the baseline simulation failed.". The cancelled_after_peer_failure and cancelled_after_coordinator_failure codes and the incomplete status for cancellations are gone; the wait timeout stays the report's error code. - Only a running record changes, so a simulation that already succeeded or failed keeps its outcome. Workers already treat any status other than running as a cancellation, so a segmented worker still sees the failed record and cancels its own region-group calls. - An unsegmented simulation is always cancelled through Modal, even when its record could not be written. A segmented simulation is cancelled through Modal only when its record could not be written, because Modal does not cancel spawned calls recursively. Cancellation and state-update errors are still logged without replacing the original failure. Carry over #793's tests for a cancelled and recorded reform, a finished baseline left alone, a restored call cancelled on retry, finished siblings left unchanged, and both unfinished children cancelled on timeout. Add tests for the Modal fallback after a failed state write, for segmented and unsegmented simulations. ImmediateCall now counts cancellations, and the entry and persistence fixtures use the new cancellation record. Add a changelog line, as #793 did. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + .../stage-12-modal-worker-foundation.md | 40 ++- .../tests/test_store.py | 11 +- .../tests/test_app.py | 10 +- .../README.md | 25 +- .../stage12_runtime/coordination.py | 85 ++--- .../tests/test_stage12_runtime.py | 326 ++++++++++++++---- 7 files changed, 336 insertions(+), 162 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index a362bd8df..6543cbc11 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -21,3 +21,4 @@ - Bound UK constituency and local-authority weight matrices to the certified data release of the selected dataset instead of the unversioned bucket copy, and stopped runs before they start when the matrix lacks the run year or household dimension - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated + - Gave Stage 12 baseline and reform simulations one shared wait window that opens when the coordinator starts and fits inside the coordinator's Modal timeout, and cancelled every unfinished simulation, recording it as failed with simulation_cancelled, when a simulation fails, the window closes, or the coordinator fails after starting a simulation diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index ce5a58b80..39b0bdc27 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -306,29 +306,34 @@ for every calculation call that it already started. The report coordinator observes baseline and reform together rather than blocking on one role first. On the first failure, it copies the failed -simulation's `error_code` and safe `error_summary` to the report, marks the -unresolved peer `incomplete` with `cancelled_after_peer_failure`, and completes -the report as failed without waiting for the peer calculation. A segmented -peer checks its persisted simulation status during bounded waits; after seeing -the cancellation state, it cancels all region-group calls it owns and exits. -Atomic status replacement prevents a late calculation result from replacing -that cancellation record. +simulation's `error_code` and safe `error_summary` to the report, records the +unresolved peer as `failed` with `simulation_cancelled` and the summary +"Cancelled after the baseline simulation failed." (or reform), and completes +the report as failed without waiting for the peer calculation. An unsegmented +peer is also cancelled through Modal. A segmented peer checks its persisted +simulation status during bounded waits; after seeing the cancellation, it +cancels all region-group calls it owns and exits. Modal does not cancel +spawned calls recursively, so the coordinator cancels a segmented peer through +Modal only if it could not write the cancellation record. Only a `running` +record changes, so a simulation that already succeeded or failed keeps its +outcome, and atomic status replacement prevents a late calculation result from +replacing the cancellation record. Both simulations share one 3,900-second wait window, the longest simulation worker timeout. It opens when the coordinator starts, so the coordinator's own setup, including loading the country model and dispatching both simulations, comes out of that window and not out of the time it needs afterwards. When the -window closes, the coordinator records each unfinished simulation as `failed` -with `simulation_wait_timeout` and stops it. The coordinator's 5,400-second -Modal timeout is that window plus up to 900 seconds for the production-result -wait plus 600 seconds to aggregate, write artifacts, and record the outcome. +window closes, the report fails with `simulation_wait_timeout`, and each +unfinished simulation is cancelled in the same way. The coordinator's +5,400-second Modal timeout is that window plus up to 900 seconds for the +production-result wait plus 600 seconds to aggregate, write artifacts, and +record the outcome. If the coordinator itself fails after starting a simulation, for example while recording a later dispatch or attaching a started call's invocation -identifier, it stops every started, unfinished simulation in the same way and -marks it `incomplete` with `cancelled_after_coordinator_failure`. Simulations -that already finished are left unchanged, and the report keeps the -coordinator's original failure. +identifier, it cancels every started, unfinished simulation in the same way, +with the summary "Cancelled after the coordinator failed.". The report keeps +the coordinator's original failure. Typed input failures, including invalid policy periods and SPM selection errors, retain their code and bounded message through the segment, simulation, @@ -478,9 +483,8 @@ contains the complete temporary parent and child execution metadata, including private artifact references and digests, but never returns artifact contents. For a failed report, the parent record contains the originating safe failure; the failed simulation contains the same values, and a peer stopped because of -that failure is `incomplete` with `cancelled_after_peer_failure`. A simulation -stopped because the coordinator failed is `incomplete` with -`cancelled_after_coordinator_failure`. +that failure is `failed` with `simulation_cancelled`; its `error_summary` +gives the reason. Cloud Run therefore requires no additional object-storage permission for this interface. diff --git a/libs/policyengine-stage12-persistence/tests/test_store.py b/libs/policyengine-stage12-persistence/tests/test_store.py index 685814cac..7cf41b236 100644 --- a/libs/policyengine-stage12-persistence/tests/test_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_store.py @@ -405,10 +405,11 @@ def test_simulation_status_replacement_is_compare_and_set() -> None: def test_simulation_status_replacement_preserves_a_concurrent_cancellation() -> None: - incomplete = _simulation().model_copy( + cancelled = _simulation().model_copy( update={ - "status": ComparisonRunLifecycleStatus.INCOMPLETE, - "error_code": "cancelled_after_peer_failure", + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", "completed_at": NOW, } ) @@ -424,7 +425,7 @@ def test_simulation_status_replacement_preserves_a_concurrent_cancellation() -> "completed_at": NOW, } ) - engine = FakeEngine([FakeResult(row=_row(incomplete))]) + engine = FakeEngine([FakeResult(row=_row(cancelled))]) store = Stage12PersistenceStore(engine=engine) # type: ignore[arg-type] result, replaced = store.replace_simulation_if_status( @@ -433,7 +434,7 @@ def test_simulation_status_replacement_preserves_a_concurrent_cancellation() -> ) assert replaced is False - assert result == incomplete + assert result == cancelled assert len(engine.connection.statements) == 1 diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index 34b77ba0c..1a55e433c 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -497,9 +497,9 @@ def test_temporary_stage12_poll_returns_propagated_failure_values(backend): ) baseline = comparison_simulation(SimulationRole.BASELINE).model_copy( update={ - "status": ComparisonRunLifecycleStatus.INCOMPLETE, - "error_code": "cancelled_after_peer_failure", - "error_summary": "Cancelled after reform simulation failed.", + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", "completed_at": report.completed_at, } ) @@ -531,8 +531,8 @@ def test_temporary_stage12_poll_returns_propagated_failure_values(backend): payload = result.json() assert payload["report"]["error_code"] == "invalid_policy_period" assert payload["report"]["error_summary"] == "The reform period is invalid." - assert payload["simulations"][0]["status"] == "incomplete" - assert payload["simulations"][0]["error_code"] == "cancelled_after_peer_failure" + assert payload["simulations"][0]["status"] == "failed" + assert payload["simulations"][0]["error_code"] == "simulation_cancelled" assert payload["simulations"][1]["error_code"] == "invalid_policy_period" diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 838fa2d22..4e7227d21 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -175,18 +175,19 @@ After dispatching both logical simulations, the coordinator checks both Modal calls with nonblocking result probes on every polling pass. It does not wait for baseline before checking reform. Both simulations share one 3,900-second wait window, the longest simulation worker timeout, that opens when the -coordinator starts, so coordinator setup comes out of that window. When it -closes, every unfinished simulation is recorded as -`failed` with `simulation_wait_timeout` and stopped. The first failed -simulation supplies the report's `error_code` and safe `error_summary`; the -coordinator marks the unresolved peer `incomplete` with -`cancelled_after_peer_failure`. A segmented peer reads that persisted state, -cancels the region-group calls it owns, and exits without overwriting the -cancellation record. If the coordinator itself fails after starting a -simulation, it stops every started, unfinished simulation the same way and -marks it `incomplete` with `cancelled_after_coordinator_failure`. Expected -input failures retain their typed code and message. Unexpected exceptions are -logged in full and persisted as a generic message with a correlation ID. +coordinator starts, so coordinator setup comes out of that window. The first +failed simulation supplies the report's `error_code` and safe `error_summary`. +When a simulation fails, when the window closes, or when the coordinator +itself fails after starting a simulation, the coordinator records every +started, unfinished simulation as `failed` with `simulation_cancelled` and a +summary that gives the reason, then stops it. Only a `running` record changes, +so a simulation that already succeeded or failed keeps its outcome. An +unsegmented simulation is cancelled through Modal. A segmented simulation +reads its cancelled record, cancels the region-group calls it owns, and exits +without overwriting the record; it is cancelled through Modal only if that +record could not be written. Expected input failures retain their typed code +and message. Unexpected exceptions are logged in full and persisted as a +generic message with a correlation ID. Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 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 9dcafc053..993d44dea 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 @@ -142,16 +142,15 @@ def _replace_running_simulation( return result -COORDINATOR_FAILURE_CANCELLATION = Stage12FailureDetail( - error_code="cancelled_after_coordinator_failure", - error_summary="Cancelled after report coordination failed.", -) +SIMULATION_CANCELLED_ERROR_CODE = "simulation_cancelled" + +def _cancellation(reason: str) -> Stage12FailureDetail: + """Describe a child the coordinator stopped; ``reason`` ends the summary.""" -def _peer_failure_cancellation(failed_role: str) -> Stage12FailureDetail: return Stage12FailureDetail( - error_code="cancelled_after_peer_failure", - error_summary=f"Cancelled after {failed_role} simulation failed.", + error_code=SIMULATION_CANCELLED_ERROR_CODE, + error_summary=f"Cancelled after {reason}.", ) @@ -163,21 +162,22 @@ def _cancel_pending_simulations( store: ComparisonStore, detail: Stage12FailureDetail, ) -> None: - """Record each started, unfinished child as incomplete, then stop it. + """Record each started, unfinished child as failed, then stop it. - A child whose record is no longer running has already finished or been - stopped, so it is left unchanged. State-update and cancellation errors are - logged rather than raised, so they never replace the failure that caused - the cancellation. + Only a running record changes, so a child that has already succeeded or + failed keeps its own outcome. An unsegmented child is always cancelled + through Modal, even if its record could not be written. State-update and + cancellation errors are logged rather than raised, so they never replace + the failure that caused the cancellation. """ for role in sorted(roles): simulation = simulations[role] try: - latest = _replace_running_simulation( + _replace_running_simulation( store=store, simulation=simulation, - status=ComparisonRunLifecycleStatus.INCOMPLETE, + status=ComparisonRunLifecycleStatus.FAILED, detail=detail, ) except Exception as error: # noqa: BLE001 @@ -188,13 +188,16 @@ def _cancel_pending_simulations( simulation=simulation, error=error, ) - continue - if latest.status is not ComparisonRunLifecycleStatus.INCOMPLETE: - continue - if should_segment_simulation(simulation): - # The simulation worker reads the persisted state and cancels the - # segment handles it owns. Modal does not guarantee recursive - # cancellation for calls created with spawn(). + cancellation_recorded = False + else: + 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. continue try: calls[role].cancel() @@ -289,7 +292,7 @@ def _observe_simulations( simulations=simulations, roles=pending - {role}, store=store, - detail=_peer_failure_cancellation(role), + detail=_cancellation(f"the {role} simulation failed"), ) raise Stage12ExecutionError(detail) from None descriptors[role] = descriptor @@ -314,29 +317,15 @@ def _observe_simulations( runtime=runtime, failure_scope="child_wait_timeout", ) - for role in sorted(pending): - persist_failure_record( - partial( - _replace_running_simulation, - store=store, - simulation=simulations[role], - status=ComparisonRunLifecycleStatus.FAILED, - detail=detail, - ), - runtime=runtime, - failure_scope="child_wait_timeout", - ) - if not should_segment_simulation(simulations[role]): - try: - calls[role].cancel() - except Exception as error: # noqa: BLE001 - _log_cancellation_failure( - logging.WARNING, - "Stage 12 could not cancel timed-out %s simulation %s " - "in report %s (%s)", - simulation=simulations[role], - error=error, - ) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=pending, + store=store, + detail=_cancellation( + f"the {timeout_seconds:g}-second wait window closed" + ), + ) raise wait_timeout # Never sleep past the deadline, so the window closes on time. sleep(min(poll_interval, remaining_seconds)) @@ -698,7 +687,9 @@ def stage_scope( simulations=simulations, roles=set(calls), store=persistence, - detail=_peer_failure_cancellation(simulation.role.value), + detail=_cancellation( + f"the {simulation.role.value} simulation failed" + ), ) raise Stage12ExecutionError(detail) from None calls[simulation.role.value] = call @@ -829,7 +820,7 @@ def stage_scope( simulations=simulations, roles=set(calls) - set(descriptors), store=persistence, - detail=COORDINATOR_FAILURE_CANCELLATION, + detail=_cancellation("the coordinator failed"), ) failed_at = datetime.now(UTC) detail = failure_detail_from_exception( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 222576e66..b714a2482 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -957,9 +957,9 @@ def segmented_calculator(value, cancellation_requested): store.replace_simulation( current.model_copy( update={ - "status": ComparisonRunLifecycleStatus.INCOMPLETE, - "error_code": "cancelled_after_peer_failure", - "error_summary": "Cancelled after reform simulation failed.", + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", "updated_at": cancelled_at, "completed_at": cancelled_at, } @@ -980,8 +980,8 @@ def segmented_calculator(value, cancellation_requested): ) child = store.children[simulation.simulation_execution_id] - assert child.status is ComparisonRunLifecycleStatus.INCOMPLETE - assert child.error_code == "cancelled_after_peer_failure" + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_cancelled" def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> None: @@ -1288,13 +1288,14 @@ def __init__(self, object_id, result, observed_timeouts): self.object_id = object_id self.result = result self.observed_timeouts = observed_timeouts + self.cancel_count = 0 def get(self, *, timeout=None): self.observed_timeouts.append(timeout) return self.result def cancel(self): - return None + self.cancel_count += 1 class ConcurrentInvoker: @@ -1751,13 +1752,14 @@ class PendingCall: def __init__(self): self.timeouts = [] + self.cancel_count = 0 def get(self, *, timeout=None): self.timeouts.append(timeout) raise TimeoutError() def cancel(self): - raise AssertionError("segmented cancellation must be cooperative") + self.cancel_count += 1 class FailedCall: object_id = f"call-{failed_role}" @@ -1812,10 +1814,15 @@ def restore(self, invocation_id): assert failed_call.timeouts == [0] assert baseline_call.timeouts == ([0] if pending_role == "baseline" else []) + # A segmented peer reads its cancelled record and stops its own + # region-group calls, so it is not cancelled through Modal. + assert baseline_call.cancel_count == 0 pending = store.children[pending_id] - assert pending.status is ComparisonRunLifecycleStatus.INCOMPLETE - assert pending.error_code == "cancelled_after_peer_failure" - assert pending.error_summary == f"Cancelled after {failed_role} simulation failed." + assert pending.status is ComparisonRunLifecycleStatus.FAILED + assert pending.error_code == "simulation_cancelled" + assert pending.error_summary == ( + f"Cancelled after the {failed_role} simulation failed." + ) assert store.parent.error_code == "invalid_policy_period" assert store.parent.error_summary == "The reform period is invalid." assert artifacts.aggregate_writes == [] @@ -1849,8 +1856,10 @@ def restore(self, invocation_id): baseline = store.children[BASELINE_ID] reform = store.children[REFORM_ID] - assert baseline.status is ComparisonRunLifecycleStatus.INCOMPLETE - assert baseline.error_code == "cancelled_after_peer_failure" + assert baseline_call.cancelled == 0 + assert baseline.status is ComparisonRunLifecycleStatus.FAILED + assert baseline.error_code == "simulation_cancelled" + assert baseline.error_summary == "Cancelled after the reform simulation failed." assert reform.status is ComparisonRunLifecycleStatus.FAILED assert reform.error_code == "simulation_dispatch_failed" assert store.parent.error_code == "simulation_dispatch_failed" @@ -2451,27 +2460,46 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: for child in store.children.values() ) assert all( - child.error_code == "simulation_wait_timeout" - for child in store.children.values() + child.error_code == "simulation_cancelled" for child in store.children.values() ) +class StartupWriteError(Exception): + """A coordinator persistence failure after a child may have started.""" + + class ScriptedChildCall: - """A started child that keeps running unless a wait error is scripted.""" + """A started child that keeps running unless the test scripts an outcome.""" - def __init__(self, object_id, *, clock=None, wait_error=None, cancel_error=None): + def __init__( + self, + object_id, + *, + clock=None, + result=None, + wait_error=None, + cancel_error=None, + during_wait=None, + ): self.object_id = object_id self.clock = clock + self.result = result self.wait_error = wait_error self.cancel_error = cancel_error + self.during_wait = during_wait self.probed_at = [] self.cancel_count = 0 def get(self, *, timeout=None): - if self.clock is not None: - self.probed_at.append(self.clock.monotonic()) + self.probed_at.append( + self.clock.monotonic() if self.clock is not None else None + ) + if self.during_wait is not None: + self.during_wait() if self.wait_error is not None: raise self.wait_error + if self.result is not None: + return self.result raise TimeoutError("still running") def cancel(self): @@ -2481,10 +2509,20 @@ def cancel(self): class ScriptedChildInvoker: - def __init__(self, *, clock=None, wait_errors=None, cancel_errors=None): + def __init__( + self, + *, + clock=None, + results=None, + wait_errors=None, + cancel_errors=None, + during_wait=None, + ): self.clock = clock + self.results = results or {} self.wait_errors = wait_errors or {} self.cancel_errors = cancel_errors or {} + self.during_wait = during_wait or {} self.events = [] self.calls = {} @@ -2502,16 +2540,41 @@ def _call(self, role): call = ScriptedChildCall( f"call-{role}", clock=self.clock, + result=self.results.get(role), wait_error=self.wait_errors.get(role), cancel_error=self.cancel_errors.get(role), + during_wait=self.during_wait.get(role), ) self.calls[role] = call return call +class PlaceholderWriteFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def replace_simulation(self, record): + # The coordinator's only plain replacement is the dispatch placeholder. + if record.simulation_execution_id == self.failing_simulation_id: + raise StartupWriteError("placeholder write contains sensitive details") + return super().replace_simulation(record) + + +class AttachFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def attach_simulation_invocation(self, simulation_execution_id, **values): + if simulation_execution_id == self.failing_simulation_id: + raise StartupWriteError("attach contains sensitive details") + return super().attach_simulation_invocation(simulation_execution_id, **values) + + class CancellationStateFailingStore(FakeStore): def replace_simulation_if_status(self, record, *, expected_status): - if record.status is ComparisonRunLifecycleStatus.INCOMPLETE: + if record.error_code == "simulation_cancelled": raise ConnectionError("database contains sensitive details") return super().replace_simulation_if_status( record, expected_status=expected_status @@ -2552,6 +2615,20 @@ def _coordinate_with(report, store, invoker, **coordinator_options): ) +def _assert_cancelled(child, *, after: str) -> None: + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_cancelled" + assert child.error_summary == f"Cancelled after {after}." + assert child.completed_at is not None + + +def _assert_original_coordinator_failure(error, store) -> None: + assert isinstance(error.value.__context__, StartupWriteError) + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "report_coordination_failed" + assert "sensitive" not in store.parent.error_summary + + def _coordination_log(caplog, level): (record,) = [ record @@ -2574,7 +2651,118 @@ def _assert_log_identifies_child(record, *, role: str, error_type: str) -> None: assert record.exc_info is None -def test_peer_cancellation_failure_log_identifies_the_report_and_child( +def test_baseline_failure_cancels_and_records_the_reform_child() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")} + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + artifacts=artifacts, + ) + + reform_call = invoker.calls["reform"] + assert reform_call.cancel_count == 1 + assert reform_call.probed_at == [] + assert artifacts.aggregate_writes == [] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.parent.error_code == "simulation_invocation_failed" + + +def test_reform_failure_leaves_the_finished_baseline_alone() -> None: + report = _report_with_segmentation(segmented=False) + store = FakeStore() + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + invoker = ScriptedChildInvoker( + results={"baseline": baseline_descriptor.model_dump(mode="json")}, + wait_errors={"reform": LookupError("reform worker failed")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(report, store, invoker, artifacts=artifacts) + + assert invoker.calls["baseline"].cancel_count == 0 + assert store.children[BASELINE_ID].status is ComparisonRunLifecycleStatus.RUNNING + assert store.children[BASELINE_ID].error_code is None + assert store.children[REFORM_ID].error_code == "simulation_invocation_failed" + + +def test_baseline_failure_cancels_a_reform_call_restored_on_retry() -> None: + report = _report_with_segmentation(segmented=False) + store = FakeStore() + # An earlier coordinator attempt started the reform child before stopping. + store.children[REFORM_ID] = _child(report.reform).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-reform", + "started_at": NOW, + } + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")} + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(report, store, invoker) + + assert invoker.events == ["spawn:baseline", "restore:reform"] + assert invoker.calls["reform"].cancel_count == 1 + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +@pytest.mark.parametrize( + ("worker_status", "worker_error_code"), + [ + (ComparisonRunLifecycleStatus.SUCCEEDED, None), + (ComparisonRunLifecycleStatus.FAILED, "simulation_execution_failed"), + ], +) +def test_cancellation_leaves_a_sibling_that_already_finished_unchanged( + worker_status, + worker_error_code, +) -> None: + store = FakeStore() + worker_records = [] + + def reform_worker_records_its_outcome(): + # The reform worker finishes while the coordinator probes baseline. + finished_at = NOW + timedelta(minutes=5) + worker_records.append( + store.replace_simulation( + store.get_simulation(REFORM_ID).model_copy( + update={ + "status": worker_status, + "error_code": worker_error_code, + "updated_at": finished_at, + "completed_at": finished_at, + } + ) + ) + ) + + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")}, + during_wait={"baseline": reform_worker_records_its_outcome}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + # Cancelling a finished Modal call is harmless; its record is not touched. + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID] == worker_records[0] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +def test_peer_cancellation_failure_does_not_replace_the_child_failure( caplog, ) -> None: store = FakeStore() @@ -2592,7 +2780,7 @@ def test_peer_cancellation_failure_log_identifies_the_report_and_child( _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) assert invoker.calls["reform"].cancel_count == 1 - assert store.children[REFORM_ID].error_code == "cancelled_after_peer_failure" + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") assert store.parent.error_code == "simulation_invocation_failed" _assert_log_identifies_child( _coordination_log(caplog, logging.WARNING), @@ -2601,8 +2789,10 @@ def test_peer_cancellation_failure_log_identifies_the_report_and_child( ) -def test_peer_state_cancellation_failure_log_identifies_the_report_and_child( +@pytest.mark.parametrize("segmented", [True, False]) +def test_failed_cancellation_state_write_still_cancels_through_modal( caplog, + segmented, ) -> None: store = CancellationStateFailingStore() invoker = ScriptedChildInvoker( @@ -2613,8 +2803,12 @@ def test_peer_state_cancellation_failure_log_identifies_the_report_and_child( caplog.at_level(logging.WARNING, logger=coordination_module.__name__), pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), ): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + _coordinate_with(_report_with_segmentation(segmented=segmented), store, invoker) + # A segmented worker cannot read a cancellation that was never written, so + # it is cancelled through Modal as well. + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.RUNNING assert store.parent.error_code == "simulation_invocation_failed" _assert_log_identifies_child( _coordination_log(caplog, logging.ERROR), @@ -2623,6 +2817,27 @@ def test_peer_state_cancellation_failure_log_identifies_the_report_and_child( ) +def test_wait_timeout_cancels_and_records_every_unfinished_child() -> None: + store = FakeStore() + invoker = ScriptedChildInvoker() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + wait_timeout_seconds=0, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + for role in ("baseline", "reform"): + _assert_cancelled( + store.children[_simulation_id(role)], + after="the 0-second wait window closed", + ) + assert store.parent.error_code == "simulation_wait_timeout" + + def test_timeout_cancellation_failure_log_identifies_the_report_and_child( caplog, ) -> None: @@ -2653,47 +2868,6 @@ def test_timeout_cancellation_failure_log_identifies_the_report_and_child( ) -class StartupWriteError(Exception): - """A coordinator persistence failure after a child may have started.""" - - -class PlaceholderWriteFailingStore(FakeStore): - def __init__(self, failing_simulation_id): - super().__init__() - self.failing_simulation_id = failing_simulation_id - - def replace_simulation(self, record): - # The coordinator's only plain replacement is the dispatch placeholder. - if record.simulation_execution_id == self.failing_simulation_id: - raise StartupWriteError("placeholder write contains sensitive details") - return super().replace_simulation(record) - - -class AttachFailingStore(FakeStore): - def __init__(self, failing_simulation_id): - super().__init__() - self.failing_simulation_id = failing_simulation_id - - def attach_simulation_invocation(self, simulation_execution_id, **values): - if simulation_execution_id == self.failing_simulation_id: - raise StartupWriteError("attach contains sensitive details") - return super().attach_simulation_invocation(simulation_execution_id, **values) - - -def _assert_cancelled_after_coordinator_failure(child) -> None: - assert child.status is ComparisonRunLifecycleStatus.INCOMPLETE - assert child.error_code == "cancelled_after_coordinator_failure" - assert child.error_summary == "Cancelled after report coordination failed." - assert child.completed_at is not None - - -def _assert_original_coordinator_failure(error, store) -> None: - assert isinstance(error.value.__context__, StartupWriteError) - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_code == "report_coordination_failed" - assert "sensitive" not in store.parent.error_summary - - @pytest.mark.parametrize("segmented", [True, False]) def test_failed_placeholder_write_stops_the_started_baseline(segmented) -> None: store = PlaceholderWriteFailingStore(REFORM_ID) @@ -2706,7 +2880,7 @@ def test_failed_placeholder_write_stops_the_started_baseline(segmented) -> None: _assert_original_coordinator_failure(error, store) assert invoker.events == ["spawn:baseline"] - _assert_cancelled_after_coordinator_failure(store.children[BASELINE_ID]) + _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") # A segmented worker stops itself, and its region-group calls, once it # reads the cancelled record; Modal does not cancel spawned calls # recursively. @@ -2737,8 +2911,8 @@ def test_failed_invocation_attach_stops_every_started_child( assert invoker.events == [f"spawn:{role}" for role in started_roles] # The child whose invocation could not be attached is stopped as well. for role in started_roles: - _assert_cancelled_after_coordinator_failure( - store.children[_simulation_id(role)] + _assert_cancelled( + store.children[_simulation_id(role)], after="the coordinator failed" ) assert invoker.calls[role].cancel_count == 1 if failing_role == "baseline": @@ -2764,7 +2938,7 @@ def test_startup_failure_stops_a_call_restored_from_an_earlier_attempt() -> None _assert_original_coordinator_failure(error, store) assert invoker.events == ["restore:baseline"] - _assert_cancelled_after_coordinator_failure(store.children[BASELINE_ID]) + _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") assert invoker.calls["baseline"].cancel_count == 1 @@ -2796,7 +2970,7 @@ def test_startup_failure_leaves_a_finished_child_unchanged() -> None: _assert_original_coordinator_failure(error, store) assert invoker.events == ["spawn:reform"] assert store.children[BASELINE_ID] == finished_baseline - _assert_cancelled_after_coordinator_failure(store.children[REFORM_ID]) + _assert_cancelled(store.children[REFORM_ID], after="the coordinator failed") assert invoker.calls["reform"].cancel_count == 1 @@ -2820,7 +2994,7 @@ def test_startup_cancellation_error_does_not_replace_the_original_failure( _assert_original_coordinator_failure(error, store) assert invoker.calls["baseline"].cancel_count == 1 - _assert_cancelled_after_coordinator_failure(store.children[BASELINE_ID]) + _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") _assert_log_identifies_child( _coordination_log(caplog, logging.WARNING), role="baseline", @@ -2846,6 +3020,7 @@ class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) _assert_original_coordinator_failure(error, store) + assert invoker.calls["baseline"].cancel_count == 1 _assert_log_identifies_child( _coordination_log(caplog, logging.ERROR), role="baseline", @@ -2945,10 +3120,11 @@ def slow_output_planning(_report): >= MINIMUM_FINISHING_SECONDS ) assert store.parent.error_code == "simulation_wait_timeout" - assert all( - child.error_code == "simulation_wait_timeout" - for child in store.children.values() - ) + for role in ("baseline", "reform"): + _assert_cancelled( + store.children[_simulation_id(role)], + after=f"the {window_seconds:g}-second wait window closed", + ) def test_exhausted_wait_window_gives_the_children_no_extra_time() -> None: From 2da84b8ea43ac85026a4ae893521a53579e3f2dd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:28:49 +0400 Subject: [PATCH 21/25] Test that a closed Stage 12 wait window is reported first The rebase onto #789 reports a closed child-wait window to telemetry before cancellation and report failure records are written, as every other Stage 12 failure is. This test makes those writes fail and checks that the timeout still reached telemetry first. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_runtime.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index b714a2482..0696e893e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -2464,6 +2464,61 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: ) +def test_closed_wait_window_is_reported_before_its_failure_records() -> None: + ordered_events: list[tuple[str, object]] = [] + + class UnwritableFailureStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_code)) + raise _connection_limit_error() + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + def replace_report(self, record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_code)) + raise _connection_limit_error() + return super().replace_report(record) + + store = UnwritableFailureStore() + invoker = TimeoutInvoker() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=store, + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=lambda **_: {"must": "not run"}, + runtime=TelemetryRecordingRuntime(ordered_events), + wait_timeout_seconds=0, + ) + + original_report, *failure_writes, persistence_report = ordered_events + event_kind, (message, attributes, reported_error) = original_report + assert event_kind == "telemetry" + assert message == "Stage 12 execution failed" + assert attributes == {"stage12_failure_scope": "child_wait_timeout"} + assert reported_error.detail.error_code == "simulation_wait_timeout" + # Both cancellation records and the report record were attempted after + # the closed window had already been reported. + assert failure_writes == [ + ("failure_write", "simulation_cancelled"), + ("failure_write", "simulation_cancelled"), + ("failure_write", "simulation_wait_timeout"), + ] + assert persistence_report[1][0] == ( + "Stage 12 failure record could not be persisted" + ) + assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING + + class StartupWriteError(Exception): """A coordinator persistence failure after a child may have started.""" From 2a82903d0f802971e48aefe3ad7e56a89a339adc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:09:30 +0400 Subject: [PATCH 22/25] Keep a failed Stage 12 child's failure when its own record errors The coordinator read a failed child's record, to choose its failure detail, outside any guard. One database error on that read reached the report handler: the report was recorded as report_coordination_failed, and the failed child and its peer were both recorded as cancelled "after the coordinator failed". Reading and writing the failed child's own record, in both the child-wait and the dispatch-failure paths, is now non-fatal in the same way as the cancellation writes. The error is logged with the report and simulation identifiers and its type, but not its message. The report keeps the child's failure as its code and cause, the failed child is recorded as failed with its own cause where the record can still be written and is never recorded as cancelled, and its peers are cancelled "after the simulation failed". Each failure is now also reported once. The coordinator, the simulation worker and the segment fan-out reported an unexpected error both through the original-failure report and again through the correlation-ID log. One report now carries the correlation ID that the redacted summary persists, with the report and simulation identifiers. A segment failure is reported where it occurs, and the worker does not report it again. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 13 +- .../README.md | 8 +- .../stage12_runtime/__init__.py | 4 +- .../stage12_runtime/coordination.py | 174 ++++---- .../stage12_runtime/failure_recording.py | 45 ++- .../stage12_runtime/failures.py | 40 +- .../stage12_runtime/fanout.py | 20 +- .../stage12_runtime/simulation.py | 37 +- .../tests/test_stage12_runtime.py | 381 +++++++++++++++++- 9 files changed, 568 insertions(+), 154 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 39b0bdc27..c0dc50f19 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -335,10 +335,21 @@ identifier, it cancels every started, unfinished simulation in the same way, with the summary "Cancelled after the coordinator failed.". The report keeps the coordinator's original failure. +Errors while reading or updating simulation records never replace the failure +that caused them. If the coordinator cannot read or update the failed +simulation's own record, or record a peer's cancellation, it logs that error +with the report and simulation identifiers and the error type, but not the +error's message. The report still records the simulation's failure, the +failed simulation is never recorded as cancelled, and its peers are still +cancelled with the reason "Cancelled after the baseline simulation failed." +(or reform). + Typed input failures, including invalid policy periods and SPM selection errors, retain their code and bounded message through the segment, simulation, and report layers. Unexpected exceptions are logged with their complete stack -trace and persisted only as `Simulation failed (correlation_id=...)`. +trace and persisted only as `Simulation failed (correlation_id=...)`. The +coordinator and simulation workers report each failure once, before any +database write, and that report carries the same correlation ID. Before each v2 deployment, the independent Stage 12 precompute application builds three annual US dataset files and 60 current-law baseline files: one for diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 4e7227d21..cd76ec6fb 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -185,9 +185,11 @@ so a simulation that already succeeded or failed keeps its outcome. An unsegmented simulation is cancelled through Modal. A segmented simulation reads its cancelled record, cancels the region-group calls it owns, and exits without overwriting the record; it is cancelled through Modal only if that -record could not be written. Expected input failures retain their typed code -and message. Unexpected exceptions are logged in full and persisted as a -generic message with a correlation ID. +record could not be written. An error while reading or updating a simulation +record is logged with its identifiers and never replaces the failure that +caused it. Expected input failures retain their typed code and message. +Unexpected exceptions are reported once, in full, and persisted as a generic +message with the same correlation ID. Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 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 29ca966f8..3a876f234 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 @@ -7,7 +7,7 @@ Stage12ExecutionError, Stage12FailureDetail, Stage12InputError, - failure_detail_from_exception, + redacted_failure_detail, validate_policy_periods, ) from .fanout import Stage12SegmentRunner, run_segmented_simulation @@ -46,7 +46,7 @@ "calculate_simulation_frames", "calculate_segment", "coordinate_report", - "failure_detail_from_exception", + "redacted_failure_detail", "run_single_simulation", "run_segmented_simulation", "merge_segment_results", 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 993d44dea..2133ee9dd 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 @@ -57,11 +57,14 @@ artifact_store, runtime_store, ) -from .failure_recording import persist_failure_record, report_original_failure +from .failure_recording import ( + persist_failure_record, + report_failure, + report_original_failure, +) from .failures import ( Stage12ExecutionError, Stage12FailureDetail, - failure_detail_from_exception, failure_detail_from_record, ) from .output_planning import ( @@ -86,14 +89,14 @@ def _next_poll_interval(current: float) -> float: ) -def _log_cancellation_failure( +def _log_child_operation_failure( level: int, message: str, *, simulation: PlannedSimulationExecutionInput, error: Exception, ) -> None: - """Log a cancellation failure with the report and child identifiers. + """Log a failed child state update or cancellation with its identifiers. ``message`` receives the simulation role, simulation execution ID, evaluation ID and error type, in that order. The error's own message is @@ -181,7 +184,7 @@ def _cancel_pending_simulations( detail=detail, ) except Exception as error: # noqa: BLE001 - _log_cancellation_failure( + _log_child_operation_failure( logging.ERROR, "Stage 12 could not record cancellation of %s simulation %s " "in report %s (%s)", @@ -202,7 +205,7 @@ def _cancel_pending_simulations( try: calls[role].cancel() except Exception as error: # noqa: BLE001 - _log_cancellation_failure( + _log_child_operation_failure( logging.WARNING, "Stage 12 could not cancel %s simulation %s in report %s (%s)", simulation=simulation, @@ -210,43 +213,89 @@ def _cancel_pending_simulations( ) -def _child_failure_detail( +def _simulation_failure_attributes( + simulation: PlannedSimulationExecutionInput, +) -> dict[str, object]: + return { + "simulation_role": simulation.role.value, + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + } + + +def _log_child_state_failure( *, - error: BaseException, simulation: PlannedSimulationExecutionInput, - store: ComparisonStore, - runtime: ObservabilityRuntime | None, -) -> Stage12FailureDetail: - child = store.get_simulation(simulation.simulation_execution_id) - if ( - child.status is ComparisonRunLifecycleStatus.FAILED - and child.error_code is not None - ): - return failure_detail_from_record( - error_code=child.error_code, - error_summary=child.error_summary, - ) - detail = failure_detail_from_exception( - error, - runtime=runtime, - scope="stage12_simulation_invocation", - default_code="simulation_invocation_failed", - context={ - "evaluation_id": str(simulation.evaluation_id), - "simulation_role": simulation.role.value, - }, + error: Exception, +) -> None: + _log_child_operation_failure( + logging.ERROR, + "Stage 12 could not record the failure of %s simulation %s in report %s (%s)", + simulation=simulation, + error=error, ) - persist_failure_record( - partial( - _replace_running_simulation, + + +def _record_child_failure( + *, + store: ComparisonStore, + simulation: PlannedSimulationExecutionInput, + detail: Stage12FailureDetail, +) -> None: + """Record a running child as failed with its own cause. + + A child that has already finished keeps its own outcome. A state error is + logged rather than raised, so it never replaces the child's failure. + """ + + try: + _replace_running_simulation( store=store, simulation=simulation, status=ComparisonRunLifecycleStatus.FAILED, detail=detail, - ), + ) + except Exception as state_error: # noqa: BLE001 + _log_child_state_failure(simulation=simulation, error=state_error) + + +def _child_failure_detail( + *, + error: Exception, + simulation: PlannedSimulationExecutionInput, + store: ComparisonStore, + runtime: ObservabilityRuntime | None, +) -> Stage12FailureDetail: + """Report a failed child call and return the failure for its report. + + A failure the child's worker recorded itself takes precedence over the + error its call raised. Reading or writing the child's record can fail, + for example when the database is unavailable. That error is logged, and + 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( + error, runtime=runtime, failure_scope="child_wait", + default_code="simulation_invocation_failed", + attributes=_simulation_failure_attributes(simulation), ) + try: + child = store.get_simulation(simulation.simulation_execution_id) + except Exception as state_error: # noqa: BLE001 + _log_child_state_failure(simulation=simulation, error=state_error) + else: + if ( + child.status is ComparisonRunLifecycleStatus.FAILED + and child.error_code is not None + ): + return failure_detail_from_record( + error_code=child.error_code, + error_summary=child.error_summary, + ) + _record_child_failure(store=store, simulation=simulation, detail=detail) return detail @@ -275,12 +324,6 @@ def _observe_simulations( except TimeoutError: continue except Exception as error: # noqa: BLE001 - report_original_failure( - error, - runtime=runtime, - failure_scope="child_wait", - attributes={"simulation_role": role}, - ) detail = _child_failure_detail( error=error, simulation=simulations[role], @@ -655,32 +698,17 @@ def stage_scope( observability_context=child_observability_context, ) except Exception as error: - report_original_failure( + detail = report_failure( error, runtime=runtime, failure_scope="child_dispatch", - attributes={"simulation_role": simulation.role.value}, - ) - detail = failure_detail_from_exception( - error, - runtime=runtime, - scope="stage12_simulation_dispatch", default_code="simulation_dispatch_failed", - context={ - "evaluation_id": str(report.evaluation_id), - "simulation_role": simulation.role.value, - }, + attributes=_simulation_failure_attributes(simulation), ) - persist_failure_record( - partial( - _replace_running_simulation, - store=persistence, - simulation=simulation, - status=ComparisonRunLifecycleStatus.FAILED, - detail=detail, - ), - runtime=runtime, - failure_scope="child_dispatch", + _record_child_failure( + store=persistence, + simulation=simulation, + detail=detail, ) _cancel_pending_simulations( calls=calls, @@ -802,18 +830,23 @@ def stage_scope( # Persist any coordinator, child-call, aggregation, or artifact failure # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 - if not isinstance(error, Stage12ExecutionError): + if isinstance(error, Stage12ExecutionError): # 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 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. - report_original_failure( + # and stopping their unfinished peers. + detail = error.detail + else: + # 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. + detail = report_failure( error, runtime=runtime, failure_scope="report_coordination", + default_code="report_coordination_failed", + attributes={"evaluation_id": str(report.evaluation_id)}, ) _cancel_pending_simulations( calls=calls, @@ -823,13 +856,6 @@ def stage_scope( detail=_cancellation("the coordinator failed"), ) failed_at = datetime.now(UTC) - detail = failure_detail_from_exception( - error, - runtime=runtime, - scope="stage12_report_coordination", - default_code="report_coordination_failed", - context={"evaluation_id": str(report.evaluation_id)}, - ) persist_failure_record( partial( persistence.replace_report, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py index 342c315f4..409f3fbb5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -8,9 +8,48 @@ from policyengine_observability import ObservabilityRuntime +from .failures import Stage12FailureDetail, redacted_failure_detail + logger = logging.getLogger(__name__) +def report_failure( + error: Exception, + *, + runtime: ObservabilityRuntime | None, + failure_scope: str, + default_code: str, + attributes: Mapping[str, object] | None = None, + segment_index: int | None = None, +) -> Stage12FailureDetail: + """Report a failure once and return the safe detail to persist for it. + + The report is emitted before any database write is attempted, as + ``report_original_failure`` requires. An unexpected error's detail has a + new correlation ID, and the report carries the same ID, so the redacted + summary persisted for callers leads back to this one record of the + original error. + """ + + detail = redacted_failure_detail( + error, + default_code=default_code, + segment_index=segment_index, + ) + correlation = ( + {"correlation_id": detail.correlation_id} + if detail.correlation_id is not None + else {} + ) + report_original_failure( + error, + runtime=runtime, + failure_scope=failure_scope, + attributes={**(attributes or {}), **correlation}, + ) + return detail + + def report_original_failure( error: Exception, *, @@ -66,10 +105,14 @@ def _emit_error( attributes: Mapping[str, object], ) -> None: if runtime is None: + failure_label = str(attributes.get("stage12_failure_scope")) + correlation_id = attributes.get("correlation_id") + if correlation_id is not None: + failure_label = f"{failure_label}; correlation_id={correlation_id}" logger.error( "%s (%s)", message, - attributes.get("stage12_failure_scope"), + failure_label, exc_info=error, ) return 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 d52661a7f..9952a6f4a 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,21 +4,16 @@ from collections.abc import Mapping from datetime import date -import logging from typing import Any -from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.spm import spm_error_detail from pydantic import BaseModel, ConfigDict from policyengine_simulation_observability.errors import ( GENERIC_JOB_FAILURE_MESSAGE, - log_and_redact_exception, make_correlation_id, ) -logger = logging.getLogger(__name__) - ERROR_CODE_MAX_LENGTH = 64 ERROR_SUMMARY_MAX_LENGTH = 512 @@ -89,16 +84,19 @@ def failure_detail_from_record( ) -def failure_detail_from_exception( +def redacted_failure_detail( error: BaseException, *, - runtime: ObservabilityRuntime | None, - scope: str, default_code: str, - context: dict[str, Any] | None = None, segment_index: int | None = None, ) -> Stage12FailureDetail: - """Retain typed public failures and redact unexpected exception messages.""" + """Retain typed public failures and redact unexpected exception messages. + + This only describes the failure; it does not report it. Any other error + is summarized by a generic message and a new correlation ID, and the + caller reports the error once with that ID so the summary leads back to + it. + """ if isinstance(error, Stage12ExecutionError): return error.detail @@ -116,26 +114,8 @@ def failure_detail_from_exception( segment_index=segment_index, ) - if runtime is not None: - summary = log_and_redact_exception( - error, - runtime=runtime, - scope=scope, - context=context, - ) - correlation_id = summary.removeprefix( - f"{GENERIC_JOB_FAILURE_MESSAGE} (correlation_id=" - ).removesuffix(")") - else: - correlation_id = make_correlation_id() - logger.error( - "Stage 12 %s failed (correlation_id=%s)", - scope, - correlation_id, - exc_info=error, - extra={"correlation_id": correlation_id, **(context or {})}, - ) - summary = f"{GENERIC_JOB_FAILURE_MESSAGE} (correlation_id={correlation_id})" + correlation_id = make_correlation_id() + summary = f"{GENERIC_JOB_FAILURE_MESSAGE} (correlation_id={correlation_id})" if segment_index is not None: summary = f"Segment {segment_index}: {summary}" return _safe_detail( 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 012df1c2e..28ab4f57f 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 @@ -18,11 +18,8 @@ ) from .partition import stage12_region_groups_for_model -from .failures import ( - Stage12Cancellation, - Stage12ExecutionError, - failure_detail_from_exception, -) +from .failure_recording import report_failure +from .failures import Stage12Cancellation, Stage12ExecutionError from .segmentation import build_segment_inputs, merge_segment_results from .simulation import SimulationCalculation @@ -153,12 +150,19 @@ def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: type(error).__name__, ) continue - detail = failure_detail_from_exception( + detail = report_failure( error, runtime=self.runtime, - scope="stage12_segment_call", + failure_scope="segment_call", default_code="segment_execution_failed", - context={"segment_index": segment_index}, + attributes={ + "simulation_role": self.simulation.role.value, + "evaluation_id": str(self.simulation.evaluation_id), + "simulation_execution_id": str( + self.simulation.simulation_execution_id + ), + "segment_index": segment_index, + }, segment_index=segment_index, ) self._cancel_all(handles) 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 f24364987..e328e41a8 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 @@ -42,11 +42,11 @@ ) from .dependencies import ComparisonStore, artifact_store, runtime_store -from .failure_recording import persist_failure_record, report_original_failure +from .failure_recording import persist_failure_record, report_failure from .failures import ( Stage12Cancellation, + Stage12ExecutionError, Stage12FailureDetail, - failure_detail_from_exception, validate_policy_periods, ) from .output_planning import apply_output_plan, validate_output_frames @@ -518,24 +518,23 @@ def cancellation_requested() -> bool: # Persist any country-package calculation failure before returning a # stable exception to Modal. except Exception as error: # noqa: BLE001 - report_original_failure( - error, - runtime=runtime, - failure_scope="simulation_execution", - attributes={"simulation_role": simulation.role.value}, - ) + 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, + runtime=runtime, + failure_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), + }, + ) failed_at = datetime.now(UTC) - detail = failure_detail_from_exception( - error, - runtime=runtime, - scope="stage12_simulation_execution", - default_code="simulation_execution_failed", - context={ - "evaluation_id": str(simulation.evaluation_id), - "simulation_execution_id": str(simulation.simulation_execution_id), - "simulation_role": simulation.role.value, - }, - ) persist_failure_record( partial( _mark_running_simulation_failed, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 0696e893e..bae2c5c2c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -13,6 +13,7 @@ from datetime import UTC, datetime, timedelta from hashlib import sha256 from threading import Lock +from types import SimpleNamespace from uuid import UUID, uuid4 import pandas as pd @@ -76,6 +77,7 @@ ) from policyengine_simulation_executor.stage12_runtime import ( SimulationCalculation, + Stage12SegmentRunner, _build_spm_result, build_aggregate_report, database_retry, @@ -97,6 +99,9 @@ from policyengine_simulation_executor.stage12_runtime.comparison import ( PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, ) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) from policyengine_simulation_executor.stage12_runtime.dependencies import ( runtime_store, ) @@ -857,6 +862,15 @@ def _connection_limit_error() -> OperationalError: ) +def _correlation_id(error_summary: str) -> str: + """Return the correlation ID in a redacted failure summary.""" + + prefix = "Simulation failed (correlation_id=" + assert error_summary.startswith(prefix) + assert error_summary.endswith(")") + return error_summary.removeprefix(prefix).removesuffix(")") + + class TelemetryRecordingRuntime(RecordingRuntime): """Record telemetry in the same ordered list as database writes.""" @@ -934,9 +948,11 @@ def fail(_simulation): assert attributes == { "stage12_failure_scope": "simulation_execution", "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": _correlation_id(failure_write[1]), } assert failure_write[0] == "failure_write" - assert failure_write[1].startswith("Simulation failed (correlation_id=") assert persistence_report[0] == "telemetry" assert persistence_report[1][0] == ( "Stage 12 failure record could not be persisted" @@ -2202,22 +2218,17 @@ def fail_aggregation(**_): assert error.value.__cause__ is None original_report, failure_write, persistence_report = ordered_events - assert original_report[1][1] == {"stage12_failure_scope": "report_coordination"} + assert original_report[1][1] == { + "stage12_failure_scope": "report_coordination", + "evaluation_id": str(EVALUATION_ID), + "correlation_id": _correlation_id(failure_write[1]), + } assert original_report[1][2] is aggregation_error assert failure_write[0] == "failure_write" - assert failure_write[1].startswith("Simulation failed (correlation_id=") assert isinstance(persistence_report[1][2], OperationalError) assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING -@pytest.mark.xfail( - strict=True, - reason=( - "Known open bug: the coordinator reads a failed child's record to " - "choose its failure detail outside the failure-recording guard, so " - "an unreadable record replaces the child's original failure." - ), -) def test_coordinator_reports_a_child_failure_once_when_its_record_is_unwritable() -> ( None ): @@ -2251,13 +2262,19 @@ def get_simulation(self, simulation_execution_id): if event[0] == "telemetry" and event[1][0] == "Stage 12 execution failed" ] assert [report[1] for report in failure_reports] == [ - {"stage12_failure_scope": "child_wait", "simulation_role": "baseline"} + { + "stage12_failure_scope": "child_wait", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": _correlation_id(store.parent.error_summary), + } ] assert ordered_events[0][0] == "telemetry" assert ordered_events[1] == ("child_read", str(BASELINE_ID)) # The report record is still written after the child record could not be. assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary.startswith("Simulation failed (correlation_id=") + assert store.parent.error_code == "simulation_invocation_failed" class RefusingDispatchInvoker: @@ -2299,17 +2316,22 @@ def test_coordinator_reports_a_child_dispatch_failure_once() -> None: for event in ordered_events if event[0] == "telemetry" and event[1][0] == "Stage 12 execution failed" ] + baseline_child = store.children[BASELINE_ID] assert [(attributes, error) for _message, attributes, error in failure_reports] == [ ( - {"stage12_failure_scope": "child_dispatch", "simulation_role": "baseline"}, + { + "stage12_failure_scope": "child_dispatch", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": _correlation_id(baseline_child.error_summary), + }, dispatch_error, ) ] assert invoker.events == ["spawn:baseline"] - baseline_child = store.children[BASELINE_ID] assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED assert baseline_child.error_code == "simulation_dispatch_failed" - assert baseline_child.error_summary.startswith("Simulation failed (correlation_id=") assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_summary == baseline_child.error_summary @@ -3083,6 +3105,333 @@ class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): ) +class ChildStateError(Exception): + """A database failure while reading or writing a child's own record.""" + + +class OwnChildStateFailingStore(FakeStore): + """Fail one read, or one failed-status write, of one child's record once. + + Cancellation writes are left alone, so only the failed child's own state + update can fail. + """ + + def __init__(self, *, failing_simulation_id, failing_operation): + super().__init__() + self.failing_simulation_id = failing_simulation_id + self.failing_operation = failing_operation + self.remaining_failures = 1 + + def _fail_once(self, operation, simulation_execution_id): + if ( + operation == self.failing_operation + and simulation_execution_id == self.failing_simulation_id + and self.remaining_failures > 0 + ): + self.remaining_failures -= 1 + raise ChildStateError("database contains sensitive details") + + def get_simulation(self, simulation_execution_id): + self._fail_once("read", simulation_execution_id) + return super().get_simulation(simulation_execution_id) + + def replace_simulation_if_status(self, record, *, expected_status): + if record.error_code != "simulation_cancelled": + self._fail_once("write", record.simulation_execution_id) + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + +@pytest.mark.parametrize( + ("failing_operation", "baseline_status", "baseline_error_code"), + [ + # The failure is still recorded: the read before the conditional + # write succeeds. + ( + "read", + ComparisonRunLifecycleStatus.FAILED, + "simulation_invocation_failed", + ), + ("write", ComparisonRunLifecycleStatus.RUNNING, None), + ], +) +def test_failed_child_state_error_does_not_replace_the_child_failure( + caplog, + failing_operation, + baseline_status, + baseline_error_code, +) -> None: + store = OwnChildStateFailingStore( + failing_simulation_id=BASELINE_ID, + failing_operation=failing_operation, + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert store.parent.error_code == "simulation_invocation_failed" + _correlation_id(store.parent.error_summary) + baseline = store.children[BASELINE_ID] + assert baseline.status is baseline_status + assert baseline.error_code == baseline_error_code + if baseline_error_code is not None: + assert baseline.error_summary == store.parent.error_summary + # The failed child is neither cancelled nor recorded as cancelled. + assert invoker.calls["baseline"].cancel_count == 0 + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert invoker.calls["reform"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="baseline", + error_type="ChildStateError", + ) + + +class ReformDispatchFailureInvoker(ScriptedChildInvoker): + def spawn(self, *, simulation, **values): + role = PlannedSimulationExecutionInput.model_validate(simulation).role + if role is SimulationRole.REFORM: + self.events.append("spawn:reform") + raise ConnectionError("dispatch contains sensitive details") + return super().spawn(simulation=simulation, **values) + + +@pytest.mark.parametrize("failing_operation", ["read", "write"]) +def test_dispatch_failure_state_error_does_not_replace_the_child_failure( + caplog, + failing_operation, +) -> None: + store = OwnChildStateFailingStore( + failing_simulation_id=REFORM_ID, + failing_operation=failing_operation, + ) + invoker = ReformDispatchFailureInvoker() + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert store.parent.error_code == "simulation_dispatch_failed" + _correlation_id(store.parent.error_summary) + # The reform record could not be updated, so it keeps its dispatch + # placeholder rather than a cancellation. + reform = store.children[REFORM_ID] + assert reform.status is ComparisonRunLifecycleStatus.RUNNING + assert reform.error_code is None + _assert_cancelled(store.children[BASELINE_ID], after="the reform simulation failed") + assert invoker.calls["baseline"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="reform", + error_type="ChildStateError", + ) + + +class ExceptionRecordingRuntime(RecordingRuntime): + """Record every telemetry call that can report a failure.""" + + def __init__(self) -> None: + super().__init__() + self.reports: list[tuple[str, object, dict, object]] = [] + + def log(self, message, *, severity="INFO", attributes=None, error=None): + self.reports.append(("log", message, dict(attributes or {}), error)) + + def record_exception(self, error, *, handled, status_code=None): + self.reports.append(("exception", None, {}, error)) + + def event(self, name, *, attributes=None, **_options): + self.reports.append(("event", name, dict(attributes or {}), None)) + + def set_context(self, **_attributes): + return None + + +def _fail_simulation_calculation(runtime): + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + calculation_error = ValueError("calculation contains sensitive details") + + def fail(_simulation): + raise calculation_error + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=runtime, + ) + return calculation_error, store.children[BASELINE_ID].error_summary + + +def _fail_child_call(runtime): + store = FakeStore() + child_error = LookupError("baseline failure contains sensitive details") + invoker = ScriptedChildInvoker(wait_errors={"baseline": child_error}) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + runtime=runtime, + ) + return child_error, store.parent.error_summary + + +def _fail_child_dispatch(runtime): + store = FakeStore() + dispatch_error = ConnectionError("dispatch contains sensitive details") + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report(), + store, + RefusingDispatchInvoker(dispatch_error), + runtime=runtime, + ) + return dispatch_error, store.parent.error_summary + + +def _fail_report_aggregation(runtime): + store = FakeStore() + artifacts = FakeArtifacts() + aggregation_error = ValueError("aggregation contains sensitive details") + + def fail_aggregation(**_values): + raise aggregation_error + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=fail_aggregation, + runtime=runtime, + ) + return aggregation_error, store.parent.error_summary + + +@pytest.mark.parametrize( + "fail", + [ + _fail_simulation_calculation, + _fail_child_call, + _fail_child_dispatch, + _fail_report_aggregation, + ], +) +@pytest.mark.parametrize("with_runtime", [True, False], ids=["runtime", "no-runtime"]) +def test_an_unexpected_failure_is_reported_exactly_once( + caplog, + fail, + with_runtime, +) -> None: + runtime = ExceptionRecordingRuntime() if with_runtime else None + + with caplog.at_level(logging.DEBUG): + original_error, error_summary = fail(runtime) + + correlation_id = _correlation_id(error_summary) + logged_failures = [record for record in caplog.records if record.exc_info] + if runtime is not None: + [(kind, message, attributes, reported_error)] = runtime.reports + assert (kind, message) == ("log", "Stage 12 execution failed") + assert reported_error is original_error + assert attributes["correlation_id"] == correlation_id + assert logged_failures == [] + else: + [logged_failure] = logged_failures + assert logged_failure.exc_info[1] is original_error + assert f"correlation_id={correlation_id}" in logged_failure.getMessage() + + +def test_a_segment_failure_is_reported_exactly_once(caplog) -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + runtime = ExceptionRecordingRuntime() + segment_error = RuntimeError("segment contains sensitive details") + + class SegmentFunction: + @staticmethod + def spawn(payload, *, observability_context): + segment_index = payload["segment_index"] + # The first segment fails on every probe; the others keep running. + return ScriptedChildCall( + f"segment-{segment_index}", + wait_error=segment_error if segment_index == 0 else None, + ) + + class SegmentFunctionLookup: + @staticmethod + def from_name(_app_name, _function_name): + return SegmentFunction + + def segmented_calculator(value, cancellation_requested): + return Stage12SegmentRunner( + value, + app_name="stage12-app", + runtime=runtime, + modal_module=SimpleNamespace(Function=SegmentFunctionLookup), + groups=US_REGION_GROUPS, + sleep=lambda _seconds: None, + cancellation_requested=cancellation_requested, + ).run() + + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"), + ): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + runtime=runtime, + ) + + child = store.children[BASELINE_ID] + assert child.error_code == "segment_execution_failed" + [(kind, message, attributes, reported_error)] = runtime.reports + assert (kind, message) == ("log", "Stage 12 execution failed") + assert reported_error is segment_error + assert attributes == { + "stage12_failure_scope": "segment_call", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "segment_index": 0, + "correlation_id": _correlation_id( + child.error_summary.removeprefix("Segment 0: ") + ), + } + assert [record for record in caplog.records if record.exc_info] == [] + + class SteppedClock: def __init__(self, start_seconds=10_000.0): self.current_seconds = start_seconds From 9428eb8c65ab44d1f75633a6d16cf0c62bf119d2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:12:44 +0400 Subject: [PATCH 23/25] Mount only the read-only Hugging Face secret in the Stage 12 precompute app The Stage 12 cache precompute app, which runs on every Stage 12 deployment, mounted the retiring policyengine-data-credentials and huggingface-token (HF_TOKEN) Modal secrets on its image build and all five functions. The credential runbook removes both from simulation code. Its image build and functions now mount pe-uk-private-hf-read-token, which supplies HUGGING_FACE_TOKEN, declared through the same constants and required key as the executor and v2 apps. The US data the app reads comes only from the public policyengine/populace-us dataset repository, which the v2 US worker image already downloads with this token. The guard test that keeps the retired names out of deployment sources now covers this app. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 18 +++++++++++------- .../src/modal/stage12_precompute_app.py | 16 ++++++++++++---- .../tests/test_modal_scripts.py | 3 +++ .../tests/test_stage12_precompute_app.py | 19 +++++++++++++++++++ 4 files changed, 45 insertions(+), 11 deletions(-) diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index cd76ec6fb..5698f4320 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -194,13 +194,17 @@ message with the same correlation ID. Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 current-law US baseline files and emits -`STAGE12_CACHE_MANIFEST_DIGEST=`. The US worker image downloads and -verifies the manifest's 63 files during image construction. The release -validation checks the recorded manifest, hashes a baseline file, and requires -a real calculation to load that file without recomputing it. Worker validation -also calculates a US household under a nonempty CTC reform using the frontend -interval format `2026-01-01.2100-12-31`. The cache manifest digest is then -stored in the separate v2 version manifest. +`STAGE12_CACHE_MANIFEST_DIGEST=`. Like the v2 workers, its image and +functions mount only the read-only `pe-uk-private-hf-read-token` Hugging Face +secret, which supplies `HUGGING_FACE_TOKEN`, and +`stage12-evaluation-gcp-credentials`; the US data it reads comes from the +public `policyengine/populace-us` dataset repository. The US worker image +downloads and verifies the manifest's 63 files during image construction. The +release validation checks the recorded manifest, hashes a baseline file, and +requires a real calculation to load that file without recomputing it. Worker +validation also calculates a US household under a nonempty CTC reform using +the frontend interval format `2026-01-01.2100-12-31`. The cache manifest +digest is then stored in the separate v2 version manifest. The cache is an optimization, not required state for a calculation: an absent file causes a normal calculation, and a file missing planned output columns is diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py index cf2ecda4a..0ef84c8c1 100644 --- a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py @@ -8,6 +8,10 @@ import modal +from policyengine_simulation_executor.hf_access_validation import ( + HF_MODAL_SECRET_NAME, + HF_RUNTIME_ENV_NAME, +) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle STAGE12_DATA_DIR = "/opt/policyengine/stage12-data" @@ -16,9 +20,13 @@ app = modal.App("policyengine-stage12-cache-precompute") gcp_secret = modal.Secret.from_name("stage12-evaluation-gcp-credentials") -data_secret = modal.Secret.from_name("policyengine-data-credentials") -hf_secret = modal.Secret.from_name("huggingface-token", required_keys=["HF_TOKEN"]) -worker_secrets = [gcp_secret, data_secret, hf_secret] +# The read-only Hugging Face credential the other simulation apps mount. The +# US cache reads only the public policyengine/populace-us dataset repository. +hf_secret = modal.Secret.from_name( + HF_MODAL_SECRET_NAME, + required_keys=[HF_RUNTIME_ENV_NAME], +) +worker_secrets = [gcp_secret, hf_secret] def _bundle_install_command() -> str: @@ -48,7 +56,7 @@ def _bundle_install_command() -> str: frozen=True, extra_options="--only-group modal-simulation-image", ) - .run_commands(_bundle_install_command(), secrets=[data_secret, hf_secret]) + .run_commands(_bundle_install_command(), secrets=[hf_secret]) .env({"POLICYENGINE_DATA_FOLDER": STAGE12_DATA_DIR}) .add_local_python_source( "src.modal", diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index e3bebcd7b..a65f550af 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -225,6 +225,9 @@ def test_deployment_sources_exclude_legacy_hf_secret_names(self): REPO_ROOT / "projects/policyengine-simulation-executor/src/modal/v2_app.py", REPO_ROOT / "projects/policyengine-simulation-executor/src/modal/precompute_app.py", + REPO_ROOT + / "projects/policyengine-simulation-executor/src/modal" + / "stage12_precompute_app.py", ) sources = "\n".join(path.read_text(encoding="utf-8") for path in paths) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py index a993a2c5f..e245af84b 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -49,3 +49,22 @@ def test_precompute_image_installs_only_the_us_bundle(monkeypatch) -> None: ) assert "policyengine_simulation_executor" in local_source[1] assert "policyengine_stage12_persistence" in local_source[1] + + +def test_precompute_app_mounts_only_the_read_only_hugging_face_secret( + monkeypatch, +) -> None: + module = _load(monkeypatch) + read_only_secret = { + "args": ("pe-uk-private-hf-read-token",), + "kwargs": {"required_keys": ["HUGGING_FACE_TOKEN"]}, + } + gcp_secret = {"args": ("stage12-evaluation-gcp-credentials",), "kwargs": {}} + + assert module.hf_secret == read_only_secret + bundle_install = next( + call for call in module.precompute_image.calls if call[0] == "run_commands" + ) + assert bundle_install[2]["secrets"] == [read_only_secret] + for function_name, declaration in module.app.function_calls: + assert declaration["secrets"] == [gcp_secret, read_only_secret], function_name From 249e84c339c3c0dc4e840e86d5010a5d4c4b7e41 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 23:51:42 +0400 Subject: [PATCH 24/25] Register the Stage 12-only integration test in the rules registry The rules registry guard that landed on main with #752 requires every deployment-asset guard test to belong to a rule. Add the test that keeps the Stage 12-only deploy from running the existing stack's complete integration suite to sim.stage12.deployment-scope, widen the statement to say so, and re-render the rules page and the AGENTS.md block. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 2 +- docs/generated/rules.md | 2 +- rules/registry.toml | 3 ++- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a4738d7a7..a0ffb5c55 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -93,7 +93,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. ### Stage 12 persistence -- `sim.stage12.deployment-scope`: Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources; the Stage 12 worker can request a Google OIDC token. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) +- `sim.stage12.deployment-scope`: Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources or runs the existing stack's complete integration suite; the Stage 12 worker can request a Google OIDC token. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) - `sim.stage12.deployment-uses-api-schema`: Stage 12 deployment uses the API-owned schema directly, and the Stage 12 runtime it deploys uses typed DML without raw SQL or schema changes. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) - `sim.stage12.no-migration-chain`: This repository has no schema migration chain (no Alembic); Stage 12 tables are migrated by policyengine-api. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py`) - `sim.stage12.no-production-user-associations`: Stage 12 runtime modules never create production user associations (user_simulations or user_reports). (blocking: pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py`) diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 9882a1c96..17163e4fb 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -89,7 +89,7 @@ flowchart LR | Rule | Statement | Tier | Status | Blocking | Enforced by | CI | Proof | | --- | --- | --- | --- | --- | --- | --- | --- | -| `sim.stage12.deployment-scope` | Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources; the Stage 12 worker can request a Google OIDC token. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py` | `pr.yml:test` | **unproven** | +| `sim.stage12.deployment-scope` | Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources or runs the existing stack's complete integration suite; the Stage 12 worker can request a Google OIDC token. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py` | `pr.yml:test` | **unproven** | | `sim.stage12.deployment-uses-api-schema` | Stage 12 deployment uses the API-owned schema directly, and the Stage 12 runtime it deploys uses typed DML without raw SQL or schema changes. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py` | `pr.yml:test` | **unproven** | | `sim.stage12.no-migration-chain` | This repository has no schema migration chain (no Alembic); Stage 12 tables are migrated by policyengine-api. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py` | `pr.yml:test` | **unproven** | | `sim.stage12.no-production-user-associations` | Stage 12 runtime modules never create production user associations (user_simulations or user_reports). | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py` | `pr.yml:test` | **unproven** | diff --git a/rules/registry.toml b/rules/registry.toml index 3cf4b55eb..0483de2f2 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -197,7 +197,7 @@ invocation = "uv run pytest tests/" [[rule]] id = "sim.stage12.deployment-scope" category = "stage12" -statement = "Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources; the Stage 12 worker can request a Google OIDC token." +statement = "Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources or runs the existing stack's complete integration suite; the Stage 12 worker can request a Google OIDC token." rationale = "docs/operations" scope = [".github/workflows/*.yml"] tier = "mechanical" @@ -210,6 +210,7 @@ nodes = [ "tests/test_deployment_assets.py::test_main_deployment_automatically_deploys_stage12_in_both_environments", "tests/test_deployment_assets.py::test_stage12_only_deployment_runs_candidate_tests_and_requested_promotion", "tests/test_deployment_assets.py::test_stage12_only_deployment_cannot_redeploy_existing_modal_resources", + "tests/test_deployment_assets.py::test_stage12_only_deployment_skips_existing_stack_integration_suite", "tests/test_deployment_assets.py::test_stage12_worker_deployment_can_request_a_google_oidc_token", "tests/test_deployment_assets.py::test_stage12_only_promotion_runs_whatever_the_skipped_stack_jobs", ] From 0cc170d19f2eb041e63a6e95395b9ebc21c2bfc3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:23:56 +0400 Subject: [PATCH 25/25] Return valid aggregates from the Stage 12 test aggregators #787 made the coordinator tests' artifact double write aggregates through the real Stage12ArtifactStore, which rejects anything that is not a valid AggregateReportArtifactPayload. Use _aggregate_payload for this PR's aggregator stubs that must not run, as #787 did for main's, so an unexpected run shows up as a written aggregate instead of as a rejection hidden by the failure the test expects. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_runtime.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index bae2c5c2c..bca5f3377 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1824,7 +1824,7 @@ def restore(self, invocation_id): store=store, artifacts=artifacts, invoker=ConcurrentFailureInvoker(), - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), sleep=lambda _: pytest.fail("failure should be observed before sleeping"), ) @@ -1927,7 +1927,7 @@ def restore(self, invocation_id): store=store, artifacts=artifacts, invoker=PersistedFailureInvoker(), - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), ) assert store.children[REFORM_ID].error_code == "invalid_policy_period" @@ -2517,7 +2517,7 @@ def replace_report(self, record): store=store, artifacts=FakeArtifacts(), invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), runtime=TelemetryRecordingRuntime(ordered_events), wait_timeout_seconds=0, ) @@ -2687,7 +2687,7 @@ def _coordinate_with(report, store, invoker, **coordinator_options): coordinator_invocation_id="coordinator-1", store=store, invoker=invoker, - aggregator=lambda **_: {"must": "not run"}, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), **coordinator_options, )