From 4afab74021b4264d95575e0e48cf085a0331eda7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:06:43 +0400 Subject: [PATCH] Split Stage 12 coordination and the execution contracts into smaller modules stage12_runtime/coordination.py (1,091 lines) and the contract's stage12_execution.py (929 lines) were over the file-length limit. The code moves unchanged. coordination.py keeps coordinate_report and re-exports every public name. The coordinator class, ReportCoordination, its dataclasses and the polling constants move to report_coordination.py, the child simulations' records, failures and cancellation to coordinated_simulations.py, and the parent report's checks and claim to parent_claim.py. Child failures still log under the coordination module's logger, now named explicitly. Private names another module uses are public: _ReportCoordination, _cancel_pending_simulations, _cancellation, _child_failure_detail, _child_record and _simulation_failure_attributes (both with a new docstring), _record_child_failure, _validate_dispatch_context, _validate_submitted_parent and _claim_parent. stage12_execution becomes a package with the same import path, whose __init__.py re-exports every name: base.py holds the field types, the strict contract base and the simulation execution input; report_inputs.py the report input and output plans; result_artifacts.py the artifact descriptors and payloads; comparison_runs.py the invocation context and the comparison-run records. The models in the other modules subclass StrictContractModel and SimulationExecutionInput through the package path, which ruff.toml already lists as runtime-evaluated, so ruff keeps their annotation imports at runtime with no settings change; base.py sorts first, so the package binds those bases before the other modules import them. Every model's JSON schema is unchanged. The coordination tests' fixed clock patches datetime in each of the three modules that read it, and uuid4 and resolve_report_output_plan in report_coordination.py, where the coordinator looks them up. The file-length baseline drops both files (count 10 to 8). Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + docs/generated/rules.md | 2 +- .../stage12_execution.py | 929 --------------- .../stage12_execution/__init__.py | 114 ++ .../stage12_execution/base.py | 220 ++++ .../stage12_execution/comparison_runs.py | 365 ++++++ .../stage12_execution/report_inputs.py | 208 ++++ .../stage12_execution/result_artifacts.py | 185 +++ .../coordinated_simulations.py | 263 +++++ .../stage12_runtime/coordination.py | 1040 +---------------- .../stage12_runtime/parent_claim.py | 169 +++ .../stage12_runtime/report_coordination.py | 626 ++++++++++ .../tests/test_stage12_coordination.py | 12 +- rules/baselines/python-file-length.txt | 2 - rules/registry.toml | 2 +- 15 files changed, 2198 insertions(+), 1940 deletions(-) delete mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/base.py create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/report_inputs.py create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/result_artifacts.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordinated_simulations.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 586bf1a99..43e2c2382 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -16,6 +16,7 @@ - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: + - Split Stage 12 report coordination (stage12_runtime/coordination.py, 1,091 lines) into modules for the report coordinator class, its child simulations' records, failures and cancellation, and the parent report's checks and claim, with coordinate_report and every public name still in coordination.py and child failures still logged under its logger; and turned the contract's stage12_execution.py (929 lines) into a stage12_execution package of modules for the shared field types and simulation input, report inputs and output plans, result artifacts, and comparison-run records, which re-exports every name from the same import path, with unchanged JSON schemas; models subclass the base models through the package path ruff.toml lists, so no ruff setting changes - Split the gateway's endpoints.py (1,454 lines) into modules for the shared router and request helpers, route resolution from the active routing state, Modal reads of job metadata and routing, bundle payloads, and the job submission, job status and version routes, which register their routes in the same order, so the OpenAPI document is unchanged; endpoints.py re-exports every public name, the routes keep logging under the policyengine_simulation_gateway.endpoints logger, and the private helpers the new modules share are now public; its 1,801-line test file is split into five test files, with the shared expected-response helpers in fixtures/gateway_responses.py - Split the routing publisher (src/modal/utils/update_version_registry.py, 1,584 lines) into a src/modal/routing_publisher package of modules for the bundle manifest, the routing state, routing plans and their publication record, the check that only serialized workflows write routing, Modal Dict access, plan publication, the legacy seed and the command-line options, and its 2,728-line test file into a tests/version_registry directory of nine test files with shared fixtures; update_version_registry.py keeps the command line, still runs with python -m src.modal.utils.update_version_registry and re-exports every public name; behaviour and output are unchanged, and the private helpers the new modules share are now public - Removed the UK constituency and local-authority weight-matrix region fallback (_build_uk_weight_replacement_region, _require_certified_uk_weight_matrix and its Stage 12 call); each deployment rebuilds only the app for the current policyengine.py version, now 6.10.0, which routes those regions to the local-area dataset, and apps deployed for older versions are not rebuilt and keep their own code; on a bundle that does not route them, an area region is now an unsupported region diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 43721ef3c..78e27a08c 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -66,7 +66,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` -- `sim.quality.python-file-length`: 10 remaining in `rules/baselines/python-file-length.txt` +- `sim.quality.python-file-length`: 8 remaining in `rules/baselines/python-file-length.txt` ### Exceptions and review dates diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py deleted file mode 100644 index cc41fb2a4..000000000 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py +++ /dev/null @@ -1,929 +0,0 @@ -"""Versioned Stage 12 report, simulation, and artifact contracts.""" - -from __future__ import annotations - -import json -from dataclasses import dataclass -from datetime import datetime, timedelta -from enum import StrEnum -from hashlib import sha256 -from typing import Annotated, Literal -from uuid import UUID - -from pydantic import ( - BaseModel, - ConfigDict, - Field, - JsonValue, - field_validator, - model_validator, -) - -from policyengine_simulation_contract.spm import ( - SPMComparisonProvenance, - SPMProvenance, -) -from policyengine_simulation_contract.uk_geography import ( - UKLocalAuthorityBoundaryVersion as UKLocalAuthorityBoundaryVersion, # noqa: PLC0414 # explicit re-export: callers import it from stage12_execution -) -from policyengine_simulation_contract.uk_geography import ( - UKLocalAuthorityMetadata as UKLocalAuthorityMetadata, # noqa: PLC0414 # explicit re-export: callers import it from stage12_execution -) - -ContractText = Annotated[str, Field(min_length=1, max_length=255)] -ObservabilityId = Annotated[str, Field(min_length=1, max_length=36)] -# ``evaluation_id`` is the already-deployed physical database and wire field -# for the temporary comparison-run identifier. Stage 14 owns its removal. -Sha256Digest = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")] -StorageUri = Annotated[ - str, - Field(min_length=6, max_length=2048, pattern=r"^gs://[^/]+/.+$"), -] -DatasetArtifactUri = Annotated[ - str, - Field( - min_length=6, - max_length=2048, - pattern=r"^(?:gs|hf)://[^/]+/.+$", - ), -] -CountryId = Annotated[ - Literal["us", "uk"], - Field(description="Supported PolicyEngine country ID"), -] - - -class StrictContractModel(BaseModel): - """Base for Stage 12 contracts; rejects unknown fields and non-finite numbers.""" - - model_config = ConfigDict(extra="forbid", allow_inf_nan=False, frozen=True) - - -class SimulationRole(StrEnum): - """Role of a simulation: baseline or reform of a comparison, or a standalone run.""" - - BASELINE = "baseline" - REFORM = "reform" - STANDALONE = "standalone" - - -class ArtifactMediaType(StrEnum): - """Media type of a Stage 12 output artifact.""" - - JSON = "application/json" - PARQUET = "application/vnd.apache.parquet" - - -class SimulationParquetPayloadContract(StrictContractModel): - """Physical schema rules for a Stage 12 simulation artifact.""" - - payload_schema_version: Literal[1] = 1 - media_type: Literal["application/vnd.apache.parquet"] = ( - "application/vnd.apache.parquet" - ) - compression: Literal["zstd"] = "zstd" - parquet_version: Literal["2.6"] = "2.6" - data_page_version: Literal["2.0"] = "2.0" - entity_column: Literal["__entity__"] = "__entity__" - row_order_column: Literal["__row_order__"] = "__row_order__" - identifier_column_template: Literal["{entity}_id"] = "{entity}_id" - column_order: Literal["system columns, then lexicographic"] = ( - "system columns, then lexicographic" - ) - row_order: Literal["entity name, then entity identifier"] = ( - "entity name, then entity identifier" - ) - schema_version_metadata_key: Literal["policyengine.stage12.schema_version"] = ( - "policyengine.stage12.schema_version" - ) - dtype_metadata_key: Literal["policyengine.stage12.dtypes"] = ( - "policyengine.stage12.dtypes" - ) - calculation_provenance_metadata_key: Literal[ - "policyengine.stage12.calculation_provenance" - ] = "policyengine.stage12.calculation_provenance" - uk_local_authority_metadata_key: Literal[ - "policyengine.stage12.uk_local_authority_metadata" - ] = "policyengine.stage12.uk_local_authority_metadata" - - -SIMULATION_PARQUET_PAYLOAD_CONTRACT = SimulationParquetPayloadContract() - - -class DatasetArtifactMediaType(StrEnum): - """Media type of an input dataset artifact.""" - - HDF5 = "application/x-hdf5" - JSON = "application/json" - PARQUET = "application/vnd.apache.parquet" - - -class DatasetProvenance(StrictContractModel): - """The dataset a simulation used: identity, URI, revision and data package.""" - - identity: ContractText - uri: Annotated[str, Field(min_length=1, max_length=2048)] - artifact_revision: ContractText - data_package_name: ContractText - data_package_version: ContractText - - -class BundleProvenance(StrictContractModel): - """The pinned software and data a simulation ran with. - - Covers the PolicyEngine.py version, the country model package, the dataset, - and the SHA-256 of the bundle manifest they came from. - """ - - policyengine_version: ContractText - country_package_name: ContractText - country_package_version: ContractText - dataset: DatasetProvenance - bundle_manifest_sha256: Sha256Digest - - -class ArtifactReference(StrictContractModel): - """An artifact in Cloud Storage, with its media type, SHA-256 and size in bytes.""" - - uri: StorageUri - media_type: ArtifactMediaType - content_sha256: Sha256Digest - size_bytes: Annotated[int, Field(ge=0)] - - -class DatasetArtifactReference(StrictContractModel): - """A dataset artifact in Cloud Storage (``gs://``) or on Hugging Face (``hf://``). - - ``size_bytes`` may be omitted when the size is not known. - """ - - uri: DatasetArtifactUri - media_type: DatasetArtifactMediaType - content_sha256: Sha256Digest - size_bytes: Annotated[int, Field(ge=0)] | None = None - - -class DatasetPopulationInput(StrictContractModel): - """Population input that simulates the households in a dataset artifact.""" - - kind: Literal["dataset"] = "dataset" - artifact: DatasetArtifactReference - - -class HouseholdPopulationInput(StrictContractModel): - """Population input that simulates one household from a situation document.""" - - kind: Literal["household"] = "household" - document: dict[str, JsonValue] - - -PopulationInput = Annotated[ - DatasetPopulationInput | HouseholdPopulationInput, - Field(discriminator="kind"), -] - - -class GeographySelection(StrictContractModel): - """Country and region to simulate, optionally filtered on one dataset field. - - ``filter_field`` and ``filter_value`` must be given together, and - ``filter_strategy`` requires both. - """ - - country: CountryId - region: ContractText - filter_field: ContractText | None = None - filter_value: ContractText | None = None - filter_strategy: ContractText | None = None - - @model_validator(mode="after") - def require_complete_filter(self) -> GeographySelection: - """Require the filter field and value together, and both for a strategy.""" - if (self.filter_field is None) != (self.filter_value is None): - msg = "filter_field and filter_value must be supplied together" - raise ValueError(msg) - if self.filter_strategy is not None and self.filter_field is None: - msg = "filter_strategy requires filter_field and filter_value" - raise ValueError(msg) - return self - - -class SimulationExecutionInput(StrictContractModel): - """Everything needed to run one simulation of a comparison run. - - ``evaluation_id`` identifies the comparison run and - ``simulation_execution_id`` this simulation within it. The input fixes the - policy, population, year, geography, options and pinned bundle; a baseline - and reform pair is described by ``ReportExecutionInput`` instead. - """ - - contract_version: Literal[1] = 1 - evaluation_id: UUID - simulation_execution_id: UUID - role: SimulationRole - policy: dict[str, JsonValue] - population: PopulationInput - year: Annotated[int, Field(ge=1900, le=2200)] - geography: GeographySelection - options: dict[str, JsonValue] = Field(default_factory=dict) - bundle: BundleProvenance - - @model_validator(mode="before") - @classmethod - def reject_combined_policy_input(cls, value: object) -> object: - """Reject report-shaped input that carries ``baseline`` or ``reform`` keys. - - One execution input describes one simulation; the pair belongs in - ``ReportExecutionInput``. - """ - if isinstance(value, dict) and ({"baseline", "reform"} & value.keys()): - msg = "single-simulation input cannot contain baseline or reform fields" - raise ValueError(msg) - return value - - -class ReportAggregate(StrEnum): - """A family of aggregate results that a comparison report computes.""" - - BUDGET = "budget" - POVERTY = "poverty" - INEQUALITY = "inequality" - DISTRIBUTIONAL = "distributional" - WINNERS_AND_LOSERS = "winners_and_losers" - GEOGRAPHIC = "geographic" - PROGRAM_STATISTICS = "program_statistics" - - -class ReportOutputRequirements(StrictContractModel): - """Report features that determine which simulation columns must exist.""" - - schema_version: Literal[1] = 1 - aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] - include_cliff_impacts: bool - labor_supply_response_active: bool - - @field_validator("aggregates") - @classmethod - def require_complete_aggregate_profile( - cls, - value: tuple[ReportAggregate, ...], - ) -> tuple[ReportAggregate, ...]: - """Require every ``ReportAggregate``, once each and in declaration order.""" - if len(value) != len(set(value)): - msg = "report output aggregates must be unique" - raise ValueError(msg) - if value != tuple(ReportAggregate): - msg = "Stage 12 currently requires the complete aggregate profile" - raise ValueError(msg) - return value - - -class EntityOutputPlan(StrictContractModel): - """Required materialized columns for one country-model entity.""" - - entity: ContractText - materialized_variables: Annotated[ - tuple[ContractText, ...], - Field(min_length=1), - ] - additional_variables: tuple[ContractText, ...] = () - dataset_variables: tuple[ContractText, ...] = () - - @field_validator( - "materialized_variables", - "additional_variables", - "dataset_variables", - ) - @classmethod - def require_canonical_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]: - """Require variable names to be unique and sorted, so plans hash stably.""" - if len(value) != len(set(value)): - msg = "output-plan variables must be unique" - raise ValueError(msg) - if value != tuple(sorted(value)): - msg = "output-plan variables must use canonical order" - raise ValueError(msg) - return value - - @model_validator(mode="after") - def require_variable_subsets(self) -> EntityOutputPlan: - """Require the additional and dataset variables not to overlap. - - Both must also be subsets of ``materialized_variables``. - """ - if not set(self.additional_variables).issubset(self.materialized_variables): - msg = "additional output-plan variables must be a materialized subset" - raise ValueError(msg) - if not set(self.dataset_variables).issubset(self.materialized_variables): - msg = "dataset output-plan variables must be a materialized subset" - raise ValueError(msg) - if set(self.additional_variables).intersection(self.dataset_variables): - msg = ( - "calculated additional variables and dataset variables must be disjoint" - ) - raise ValueError(msg) - return self - - -class Stage12OutputPlan(StrictContractModel): - """One immutable output schema shared by both report simulations.""" - - schema_version: Literal[1] = 1 - country: CountryId - requirements: ReportOutputRequirements - entities: Annotated[tuple[EntityOutputPlan, ...], Field(min_length=1)] - - @field_validator("entities") - @classmethod - def require_canonical_entities( - cls, - value: tuple[EntityOutputPlan, ...], - ) -> tuple[EntityOutputPlan, ...]: - """Require entity plans to be unique and sorted by entity name.""" - names = tuple(entity.entity for entity in value) - if len(names) != len(set(names)): - msg = "output-plan entities must be unique" - raise ValueError(msg) - if names != tuple(sorted(names)): - msg = "output-plan entities must use canonical order" - raise ValueError(msg) - return value - - -def stage12_output_plan_sha256(plan: Stage12OutputPlan) -> str: - """Return the stable digest recorded by each Stage 12 child artifact.""" - payload = json.dumps( - plan.model_dump(mode="json"), - sort_keys=True, - separators=(",", ":"), - ensure_ascii=False, - allow_nan=False, - ).encode("utf-8") - return sha256(payload).hexdigest() - - -class PlannedSimulationExecutionInput(SimulationExecutionInput): - """Internal coordinator-to-worker input with a resolved output schema.""" - - output_plan: Stage12OutputPlan - - @model_validator(mode="after") - def require_output_plan_country(self) -> PlannedSimulationExecutionInput: - """Reject an output plan built for a different country than the geography.""" - if self.output_plan.country != self.geography.country: - msg = "output plan country must match simulation geography" - raise ValueError(msg) - return self - - -class RowIdentity(StrictContractModel): - """Fingerprint of the rows in a simulation artifact. - - Names the identifier columns and gives the row count and a SHA-256 over the - row identifiers, so artifacts can be checked to cover the same rows. - """ - - schema_version: Literal[1] = 1 - identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)] - row_count: Annotated[int, Field(ge=0)] - identity_sha256: Sha256Digest - - @field_validator("identifier_columns") - @classmethod - def require_unique_columns(cls, value: tuple[str, ...]) -> tuple[str, ...]: - """Reject repeated identifier column names.""" - if len(value) != len(set(value)): - msg = "row identifier columns must be unique" - raise ValueError(msg) - return value - - -class SimulationArtifactDescriptor(StrictContractModel): - """The stored output of one simulation in a comparison run. - - Points to the artifact and records the digest of the output plan it follows, - its row identity, its bundle provenance and, for SPM calculations, the - calculation receipt. - """ - - contract_version: Literal[1] = 1 - evaluation_id: UUID - simulation_execution_id: UUID - role: SimulationRole - artifact: ArtifactReference - output_schema_version: Literal[1] = 1 - output_plan_sha256: Sha256Digest - row_identity: RowIdentity - bundle: BundleProvenance - calculation_provenance: SPMProvenance | None = None - - -class ReportExecutionInput(StrictContractModel): - """A baseline and reform simulation pair to run and aggregate into one report. - - Both simulations must belong to this report, carry the baseline and reform - roles and distinct execution IDs, and differ only in their policy. - """ - - contract_version: Literal[1] = 1 - evaluation_id: UUID - baseline: SimulationExecutionInput - reform: SimulationExecutionInput - requested_aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] - - @model_validator(mode="after") - def require_aligned_simulations(self) -> ReportExecutionInput: - """Check the two simulations form one comparable report. - - Roles must be baseline and reform, both must belong to this report with - distinct execution ids, and everything but the policy must match. - """ - if self.baseline.role is not SimulationRole.BASELINE: - msg = "baseline input must use the baseline role" - raise ValueError(msg) - if self.reform.role is not SimulationRole.REFORM: - msg = "reform input must use the reform role" - raise ValueError(msg) - if self.baseline.evaluation_id != self.evaluation_id: - msg = "baseline evaluation_id must match the report" - raise ValueError(msg) - if self.reform.evaluation_id != self.evaluation_id: - msg = "reform evaluation_id must match the report" - raise ValueError(msg) - if self.baseline.simulation_execution_id == self.reform.simulation_execution_id: - msg = "baseline and reform simulation identifiers must differ" - raise ValueError(msg) - for field_name in ( - "population", - "year", - "geography", - "options", - "bundle", - ): - if getattr(self.baseline, field_name) != getattr(self.reform, field_name): - msg = f"baseline and reform {field_name} must match before execution" - raise ValueError(msg) - if len(self.requested_aggregates) != len(set(self.requested_aggregates)): - msg = "requested aggregates must be unique" - raise ValueError(msg) - return self - - -class AggregateReportArtifactDescriptor(StrictContractModel): - """The stored aggregate results of a comparison run. - - Records the digests of the baseline and reform artifacts the results were - computed from, and the bundle provenance. - """ - - contract_version: Literal[1] = 1 - evaluation_id: UUID - artifact: ArtifactReference - aggregate_schema_version: Literal[1] = 1 - baseline_artifact_sha256: Sha256Digest - reform_artifact_sha256: Sha256Digest - bundle: BundleProvenance - - -class AggregateReportArtifactPayload(StrictContractModel): - """Complete JSON payload stored in one aggregate report artifact.""" - - contract_version: Literal[1] = 1 - aggregate_schema_version: Literal[1] = 1 - evaluation_id: UUID - requested_aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] - bundle: BundleProvenance - result: dict[str, JsonValue] - - @field_validator("requested_aggregates") - @classmethod - def require_unique_aggregates( - cls, - value: tuple[ReportAggregate, ...], - ) -> tuple[ReportAggregate, ...]: - """Reject repeated aggregate names.""" - if len(value) != len(set(value)): - msg = "requested aggregates must be unique" - raise ValueError(msg) - return value - - @model_validator(mode="after") - def validate_spm_provenance(self) -> AggregateReportArtifactPayload: - """Reject duplicate or noncanonical SPM data in aggregate artifacts.""" - if "spm_config" in self.result: - msg = "aggregate result must not include legacy spm_config" - raise ValueError(msg) - # Reports without US SPM calculations serialize this optional field - # as null. Validate a supplied record, not the absence of one. - provenance = self.result.get("spm_provenance") - if provenance is not None: - SPMComparisonProvenance.model_validate(provenance) - return self - - -class ResultDifference(StrictContractModel): - """One exact aggregate-result difference identified by JSON Pointer.""" - - path: Annotated[str, Field(max_length=4096)] - production_present: bool - stage12_present: bool - production_value: JsonValue = None - stage12_value: JsonValue = None - absolute_delta: float | None = None - relative_delta: float | None = None - - @model_validator(mode="after") - def validate_presence(self) -> ResultDifference: - """Check presence flags against values, and that deltas come as a pair. - - An absent side must be null, and a relative delta needs an absolute one; - only an absolute delta may stand alone (a zero production value). - """ - if not self.production_present and not self.stage12_present: - msg = "a result difference must be present in at least one result" - raise ValueError(msg) - if not self.production_present and self.production_value is not None: - msg = "an absent production value must be null" - raise ValueError(msg) - if not self.stage12_present and self.stage12_value is not None: - msg = "an absent Stage 12 value must be null" - raise ValueError(msg) - # A zero production value has no finite relative delta, so an absolute - # delta with no relative delta is the only permitted partial pair. - if (self.absolute_delta is None) != (self.relative_delta is None) and ( - self.absolute_delta is None or self.relative_delta is not None - ): - msg = "numeric deltas must be supplied together" - raise ValueError(msg) - return self - - -class ResultComparisonArtifactPayload(StrictContractModel): - """Private comparison receipt without duplicate complete result objects.""" - - schema_version: Literal[1] = 1 - evaluation_id: UUID - production_job_id: ContractText - compared_at: datetime - status: Literal["matched", "different"] - production_result_sha256: Sha256Digest - stage12_result_sha256: Sha256Digest - difference_count: Annotated[int, Field(ge=0)] - differences: tuple[ResultDifference, ...] - - @model_validator(mode="after") - def validate_differences(self) -> ResultComparisonArtifactPayload: - """Require a timezone-aware timestamp and a count and status that match.""" - if self.compared_at.tzinfo is None: - msg = "comparison timestamp must include a timezone" - raise ValueError(msg) - if self.difference_count != len(self.differences): - msg = "difference_count must equal the number of differences" - raise ValueError(msg) - expected_status = "matched" if not self.differences else "different" - if self.status != expected_status: - msg = "comparison status does not match its differences" - raise ValueError(msg) - return self - - -class Stage12InvocationContext(StrictContractModel): - """Deployment and retention context for one comparison-run invocation. - - Records the originating request, the environments, the worker version and - callable, the manifest digests and the artifact prefix. Automatic runs also - name the production function call they are compared with. Artifacts are kept - for at most 30 days after ``created_at``. - """ - - contract_version: Literal[1] = 1 - request_id: ContractText - # The logical environment labels durable records and artifact paths. - environment: ContractText - # Modal uses ``main`` for production, so it cannot be inferred from the - # logical ``production`` label. - modal_environment: ContractText - worker_version: ContractText - modal_application: ContractText - simulation_callable: ContractText - version_manifest_sha256: Sha256Digest - bundle_manifest_sha256: Sha256Digest - artifact_prefix: Annotated[str, Field(min_length=1, max_length=2048)] - # Automatic runs retain the existing Modal function-call identifier so the - # coordinator can retrieve the production aggregate after Stage 12 finishes. - # Direct Stage 12-only runs have no associated production call. - production_function_call_id: ContractText | None = None - created_at: datetime - retention_expires_at: datetime - - @model_validator(mode="after") - def validate_retention(self) -> Stage12InvocationContext: - """Require aware timestamps and a retention window of at most 30 days.""" - _validate_retention(self.created_at, self.retention_expires_at) - return self - - -ErrorCode = Annotated[str, Field(min_length=1, max_length=64)] -ErrorSummary = Annotated[str, Field(min_length=1, max_length=512)] -MAX_COMPARISON_RUN_ARTIFACT_RETENTION = timedelta(days=30) - - -def _validate_retention(created_at: datetime, retention_expires_at: datetime) -> None: - if created_at.tzinfo is None or retention_expires_at.tzinfo is None: - msg = "comparison retention timestamps must include a timezone" - raise ValueError(msg) - retention = retention_expires_at - created_at - if retention <= timedelta(0): - msg = "comparison retention must be greater than zero" - raise ValueError(msg) - if retention > MAX_COMPARISON_RUN_ARTIFACT_RETENTION: - msg = "comparison retention must not exceed 30 days" - raise ValueError(msg) - - -class ComparisonRunLifecycleStatus(StrEnum): - """Overall status of a comparison run.""" - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" - INCOMPLETE = "incomplete" - SKIPPED = "skipped" - - -class ComparisonRunAggregationStatus(StrEnum): - """Status of aggregating a comparison run's simulations into a report.""" - - NOT_STARTED = "not_started" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" - - -class ResultComparisonStatus(StrEnum): - """Outcome of checking a comparison run's aggregate result against production.""" - - NOT_REQUESTED = "not_requested" - PENDING = "pending" - RUNNING = "running" - MATCHED = "matched" - DIFFERENT = "different" - FAILED = "failed" - - -class ComparisonReportRecord(StrictContractModel): - """Complete producer/consumer contract for one temporary parent row.""" - - contract_version: Literal[1] = 1 - evaluation_id: UUID - status: ComparisonRunLifecycleStatus - aggregation_status: ComparisonRunAggregationStatus - environment: ContractText - calculation_flow: ContractText - originating_request_id: ContractText - observability_id: ObservabilityId | None = None - production_identity: ContractText - incumbent_execution_id: ContractText | None = None - worker_version: ContractText - modal_application: ContractText - report_coordinator_callable: ContractText - version_manifest_sha256: Sha256Digest - policyengine_version: ContractText - country_package_name: ContractText - country_package_version: ContractText - country: CountryId - dataset_identity: ContractText - dataset_uri: Annotated[str, Field(min_length=1, max_length=2048)] - data_package_name: ContractText - data_package_version: ContractText - data_artifact_revision: ContractText - coordinator_invocation_id: ContractText | None = None - error_code: ErrorCode | None = None - error_summary: ErrorSummary | None = None - aggregate_output_uri: StorageUri | None = None - aggregate_output_sha256: Sha256Digest | None = None - aggregate_schema_version: Annotated[int, Field(ge=1)] | None = None - comparison_status: ResultComparisonStatus = ResultComparisonStatus.NOT_REQUESTED - comparison_output_uri: StorageUri | None = None - comparison_output_sha256: Sha256Digest | None = None - comparison_schema_version: Annotated[int, Field(ge=1)] | None = None - comparison_completed_at: datetime | None = None - comparison_error_code: ErrorCode | None = None - comparison_error_summary: ErrorSummary | None = None - created_at: datetime - updated_at: datetime - started_at: datetime | None = None - completed_at: datetime | None = None - retention_expires_at: datetime - - @model_validator(mode="after") - def validate_state(self) -> ComparisonReportRecord: - """Check that the row's fields are consistent with its lifecycle states. - - Aggregate and comparison output fields come all-or-nothing, a success - needs its outputs and ``completed_at``, failures and skips need an - ``error_code``, and the retention window is at most 30 days. - """ - _validate_retention(self.created_at, self.retention_expires_at) - self._validate_report_state() - self._validate_comparison_state() - return self - - def _validate_report_state(self) -> None: - """Check the aggregate output fields against the report's status.""" - output_fields = ( - self.aggregate_output_uri, - self.aggregate_output_sha256, - self.aggregate_schema_version, - ) - if any(value is not None for value in output_fields) and not all( - value is not None for value in output_fields - ): - msg = "aggregate output fields must be supplied together" - raise ValueError(msg) - if self.status is ComparisonRunLifecycleStatus.SUCCEEDED: - if self.aggregation_status is not ComparisonRunAggregationStatus.SUCCEEDED: - msg = "a successful report requires successful aggregation" - raise ValueError(msg) - if not all(value is not None for value in output_fields): - msg = "a successful report requires an aggregate output" - raise ValueError(msg) - if self.completed_at is None: - msg = "a successful report requires completed_at" - raise ValueError(msg) - if ( - self.status - in { - ComparisonRunLifecycleStatus.FAILED, - ComparisonRunLifecycleStatus.SKIPPED, - } - and self.error_code is None - ): - msg = "a failed or skipped report requires an error_code" - raise ValueError(msg) - - def _validate_comparison_state(self) -> None: - """Check the result comparison fields against the comparison's status.""" - if self.comparison_status in { - ResultComparisonStatus.NOT_REQUESTED, - ResultComparisonStatus.PENDING, - ResultComparisonStatus.RUNNING, - }: - self._validate_incomplete_comparison() - elif self.comparison_status in { - ResultComparisonStatus.MATCHED, - ResultComparisonStatus.DIFFERENT, - }: - self._validate_completed_comparison() - else: - self._validate_failed_comparison() - - @property - def _comparison_output_fields(self) -> tuple[object, ...]: - return ( - self.comparison_output_uri, - self.comparison_output_sha256, - self.comparison_schema_version, - ) - - def _validate_incomplete_comparison(self) -> None: - """Reject outputs, completion and errors on a comparison still to finish.""" - if any(value is not None for value in self._comparison_output_fields): - msg = "an incomplete comparison cannot reference an output artifact" - raise ValueError(msg) - if self.comparison_completed_at is not None: - msg = "an incomplete comparison cannot have completed_at" - raise ValueError(msg) - if self.comparison_error_code is not None: - msg = "an incomplete comparison cannot have an error_code" - raise ValueError(msg) - if self.comparison_error_summary is not None: - msg = "an incomplete comparison cannot have an error_summary" - raise ValueError(msg) - - def _validate_completed_comparison(self) -> None: - """Require an output and completion, and no error, on a finished comparison.""" - if not all(value is not None for value in self._comparison_output_fields): - msg = "a completed comparison requires an output artifact" - raise ValueError(msg) - if self.comparison_completed_at is None: - msg = "a completed comparison requires completed_at" - raise ValueError(msg) - if self.comparison_error_code is not None: - msg = "a completed comparison cannot have an error_code" - raise ValueError(msg) - if self.comparison_error_summary is not None: - msg = "a completed comparison cannot have an error_summary" - raise ValueError(msg) - - def _validate_failed_comparison(self) -> None: - """Require completion and an error code, but no output, on a failure.""" - if any(value is not None for value in self._comparison_output_fields): - msg = "a failed comparison cannot reference an output artifact" - raise ValueError(msg) - if self.comparison_completed_at is None: - msg = "a failed comparison requires completed_at" - raise ValueError(msg) - if self.comparison_error_code is None: - msg = "a failed comparison requires an error_code" - raise ValueError(msg) - - -class ComparisonSimulationRecord(StrictContractModel): - """Complete producer/consumer contract for one temporary child row.""" - - contract_version: Literal[1] = 1 - simulation_execution_id: UUID - evaluation_id: UUID - role: SimulationRole - input_sha256: Sha256Digest - worker_version: ContractText - modal_application: ContractText - simulation_callable: ContractText - version_manifest_sha256: Sha256Digest - modal_invocation_id: ContractText | None = None - status: ComparisonRunLifecycleStatus - error_code: ErrorCode | None = None - error_summary: ErrorSummary | None = None - output_uri: StorageUri | None = None - output_sha256: Sha256Digest | None = None - output_schema_version: Annotated[int, Field(ge=1)] | None = None - row_identity_columns: tuple[ContractText, ...] | None = None - row_count: Annotated[int, Field(ge=0)] | None = None - row_identity_sha256: Sha256Digest | None = None - created_at: datetime - updated_at: datetime - started_at: datetime | None = None - completed_at: datetime | None = None - retention_expires_at: datetime - - @model_validator(mode="after") - def validate_state(self) -> ComparisonSimulationRecord: - """Check that the row's fields are consistent with its lifecycle status. - - Output and row-identity fields come all-or-nothing, a success needs - them and ``completed_at``, failures and skips need an ``error_code``, - and the retention window is at most 30 days. - """ - _validate_retention(self.created_at, self.retention_expires_at) - output_fields = ( - self.output_uri, - self.output_sha256, - self.output_schema_version, - self.row_identity_columns, - self.row_count, - self.row_identity_sha256, - ) - if any(value is not None for value in output_fields) and not all( - value is not None for value in output_fields - ): - msg = "simulation output fields must be supplied together" - raise ValueError(msg) - if self.row_identity_columns is not None: - if not self.row_identity_columns: - msg = "row_identity_columns must not be empty" - raise ValueError(msg) - if len(self.row_identity_columns) != len(set(self.row_identity_columns)): - msg = "row_identity_columns must be unique" - raise ValueError(msg) - if self.status is ComparisonRunLifecycleStatus.SUCCEEDED: - if not all(value is not None for value in output_fields): - msg = "a successful simulation requires an output" - raise ValueError(msg) - if self.completed_at is None: - msg = "a successful simulation requires completed_at" - raise ValueError(msg) - if ( - self.status - in { - ComparisonRunLifecycleStatus.FAILED, - ComparisonRunLifecycleStatus.SKIPPED, - } - and self.error_code is None - ): - msg = "a failed or skipped simulation requires an error_code" - raise ValueError(msg) - return self - - -@dataclass(frozen=True) -class ComparisonReportPersistenceResult: - """The stored parent row after an idempotent create. - - ``created`` is False when an identical row already existed and was reused. - """ - - record: ComparisonReportRecord - created: bool - - -@dataclass(frozen=True) -class ComparisonSimulationPersistenceResult: - """The stored child row after an idempotent create. - - ``created`` is False when an identical row already existed and was reused. - """ - - record: ComparisonSimulationRecord - created: bool diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py new file mode 100644 index 000000000..d1a1f477f --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py @@ -0,0 +1,114 @@ +"""Versioned Stage 12 report, simulation, and artifact contracts. + +The contracts live in this package's modules, and this module re-exports +them. base.py comes first: the other modules subclass its models through +this package's path. +""" + +from policyengine_simulation_contract.stage12_execution.base import ( + SIMULATION_PARQUET_PAYLOAD_CONTRACT, + ArtifactMediaType, + ArtifactReference, + BundleProvenance, + ContractText, + CountryId, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetArtifactUri, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + HouseholdPopulationInput, + ObservabilityId, + PopulationInput, + Sha256Digest, + SimulationExecutionInput, + SimulationParquetPayloadContract, + SimulationRole, + StorageUri, + StrictContractModel, +) +from policyengine_simulation_contract.stage12_execution.comparison_runs import ( + MAX_COMPARISON_RUN_ARTIFACT_RETENTION, + ComparisonReportPersistenceResult, + ComparisonReportRecord, + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, + ComparisonSimulationPersistenceResult, + ComparisonSimulationRecord, + ErrorCode, + ErrorSummary, + ResultComparisonStatus, + Stage12InvocationContext, +) +from policyengine_simulation_contract.stage12_execution.report_inputs import ( + EntityOutputPlan, + PlannedSimulationExecutionInput, + ReportAggregate, + ReportExecutionInput, + ReportOutputRequirements, + Stage12OutputPlan, + stage12_output_plan_sha256, +) +from policyengine_simulation_contract.stage12_execution.result_artifacts import ( + AggregateReportArtifactDescriptor, + AggregateReportArtifactPayload, + ResultComparisonArtifactPayload, + ResultDifference, + RowIdentity, + SimulationArtifactDescriptor, +) +from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, +) + +__all__ = [ + "MAX_COMPARISON_RUN_ARTIFACT_RETENTION", + "SIMULATION_PARQUET_PAYLOAD_CONTRACT", + "AggregateReportArtifactDescriptor", + "AggregateReportArtifactPayload", + "ArtifactMediaType", + "ArtifactReference", + "BundleProvenance", + "ComparisonReportPersistenceResult", + "ComparisonReportRecord", + "ComparisonRunAggregationStatus", + "ComparisonRunLifecycleStatus", + "ComparisonSimulationPersistenceResult", + "ComparisonSimulationRecord", + "ContractText", + "CountryId", + "DatasetArtifactMediaType", + "DatasetArtifactReference", + "DatasetArtifactUri", + "DatasetPopulationInput", + "DatasetProvenance", + "EntityOutputPlan", + "ErrorCode", + "ErrorSummary", + "GeographySelection", + "HouseholdPopulationInput", + "ObservabilityId", + "PlannedSimulationExecutionInput", + "PopulationInput", + "ReportAggregate", + "ReportExecutionInput", + "ReportOutputRequirements", + "ResultComparisonArtifactPayload", + "ResultComparisonStatus", + "ResultDifference", + "RowIdentity", + "Sha256Digest", + "SimulationArtifactDescriptor", + "SimulationExecutionInput", + "SimulationParquetPayloadContract", + "SimulationRole", + "Stage12InvocationContext", + "Stage12OutputPlan", + "StorageUri", + "StrictContractModel", + "UKLocalAuthorityBoundaryVersion", + "UKLocalAuthorityMetadata", + "stage12_output_plan_sha256", +] diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/base.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/base.py new file mode 100644 index 000000000..c24727dcb --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/base.py @@ -0,0 +1,220 @@ +"""Field types, the strict contract base, and one simulation's execution input.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Annotated, Literal +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator + +ContractText = Annotated[str, Field(min_length=1, max_length=255)] +ObservabilityId = Annotated[str, Field(min_length=1, max_length=36)] +# ``evaluation_id`` is the already-deployed physical database and wire field +# for the temporary comparison-run identifier. Stage 14 owns its removal. +Sha256Digest = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")] +StorageUri = Annotated[ + str, + Field(min_length=6, max_length=2048, pattern=r"^gs://[^/]+/.+$"), +] +DatasetArtifactUri = Annotated[ + str, + Field( + min_length=6, + max_length=2048, + pattern=r"^(?:gs|hf)://[^/]+/.+$", + ), +] +CountryId = Annotated[ + Literal["us", "uk"], + Field(description="Supported PolicyEngine country ID"), +] + + +class StrictContractModel(BaseModel): + """Base for Stage 12 contracts; rejects unknown fields and non-finite numbers.""" + + model_config = ConfigDict(extra="forbid", allow_inf_nan=False, frozen=True) + + +class SimulationRole(StrEnum): + """Role of a simulation: baseline or reform of a comparison, or a standalone run.""" + + BASELINE = "baseline" + REFORM = "reform" + STANDALONE = "standalone" + + +class ArtifactMediaType(StrEnum): + """Media type of a Stage 12 output artifact.""" + + JSON = "application/json" + PARQUET = "application/vnd.apache.parquet" + + +class SimulationParquetPayloadContract(StrictContractModel): + """Physical schema rules for a Stage 12 simulation artifact.""" + + payload_schema_version: Literal[1] = 1 + media_type: Literal["application/vnd.apache.parquet"] = ( + "application/vnd.apache.parquet" + ) + compression: Literal["zstd"] = "zstd" + parquet_version: Literal["2.6"] = "2.6" + data_page_version: Literal["2.0"] = "2.0" + entity_column: Literal["__entity__"] = "__entity__" + row_order_column: Literal["__row_order__"] = "__row_order__" + identifier_column_template: Literal["{entity}_id"] = "{entity}_id" + column_order: Literal["system columns, then lexicographic"] = ( + "system columns, then lexicographic" + ) + row_order: Literal["entity name, then entity identifier"] = ( + "entity name, then entity identifier" + ) + schema_version_metadata_key: Literal["policyengine.stage12.schema_version"] = ( + "policyengine.stage12.schema_version" + ) + dtype_metadata_key: Literal["policyengine.stage12.dtypes"] = ( + "policyengine.stage12.dtypes" + ) + calculation_provenance_metadata_key: Literal[ + "policyengine.stage12.calculation_provenance" + ] = "policyengine.stage12.calculation_provenance" + uk_local_authority_metadata_key: Literal[ + "policyengine.stage12.uk_local_authority_metadata" + ] = "policyengine.stage12.uk_local_authority_metadata" + + +SIMULATION_PARQUET_PAYLOAD_CONTRACT = SimulationParquetPayloadContract() + + +class DatasetArtifactMediaType(StrEnum): + """Media type of an input dataset artifact.""" + + HDF5 = "application/x-hdf5" + JSON = "application/json" + PARQUET = "application/vnd.apache.parquet" + + +class DatasetProvenance(StrictContractModel): + """The dataset a simulation used: identity, URI, revision and data package.""" + + identity: ContractText + uri: Annotated[str, Field(min_length=1, max_length=2048)] + artifact_revision: ContractText + data_package_name: ContractText + data_package_version: ContractText + + +class BundleProvenance(StrictContractModel): + """The pinned software and data a simulation ran with. + + Covers the PolicyEngine.py version, the country model package, the dataset, + and the SHA-256 of the bundle manifest they came from. + """ + + policyengine_version: ContractText + country_package_name: ContractText + country_package_version: ContractText + dataset: DatasetProvenance + bundle_manifest_sha256: Sha256Digest + + +class ArtifactReference(StrictContractModel): + """An artifact in Cloud Storage, with its media type, SHA-256 and size in bytes.""" + + uri: StorageUri + media_type: ArtifactMediaType + content_sha256: Sha256Digest + size_bytes: Annotated[int, Field(ge=0)] + + +class DatasetArtifactReference(StrictContractModel): + """A dataset artifact in Cloud Storage (``gs://``) or on Hugging Face (``hf://``). + + ``size_bytes`` may be omitted when the size is not known. + """ + + uri: DatasetArtifactUri + media_type: DatasetArtifactMediaType + content_sha256: Sha256Digest + size_bytes: Annotated[int, Field(ge=0)] | None = None + + +class DatasetPopulationInput(StrictContractModel): + """Population input that simulates the households in a dataset artifact.""" + + kind: Literal["dataset"] = "dataset" + artifact: DatasetArtifactReference + + +class HouseholdPopulationInput(StrictContractModel): + """Population input that simulates one household from a situation document.""" + + kind: Literal["household"] = "household" + document: dict[str, JsonValue] + + +PopulationInput = Annotated[ + DatasetPopulationInput | HouseholdPopulationInput, + Field(discriminator="kind"), +] + + +class GeographySelection(StrictContractModel): + """Country and region to simulate, optionally filtered on one dataset field. + + ``filter_field`` and ``filter_value`` must be given together, and + ``filter_strategy`` requires both. + """ + + country: CountryId + region: ContractText + filter_field: ContractText | None = None + filter_value: ContractText | None = None + filter_strategy: ContractText | None = None + + @model_validator(mode="after") + def require_complete_filter(self) -> GeographySelection: + """Require the filter field and value together, and both for a strategy.""" + if (self.filter_field is None) != (self.filter_value is None): + msg = "filter_field and filter_value must be supplied together" + raise ValueError(msg) + if self.filter_strategy is not None and self.filter_field is None: + msg = "filter_strategy requires filter_field and filter_value" + raise ValueError(msg) + return self + + +class SimulationExecutionInput(StrictContractModel): + """Everything needed to run one simulation of a comparison run. + + ``evaluation_id`` identifies the comparison run and + ``simulation_execution_id`` this simulation within it. The input fixes the + policy, population, year, geography, options and pinned bundle; a baseline + and reform pair is described by ``ReportExecutionInput`` instead. + """ + + contract_version: Literal[1] = 1 + evaluation_id: UUID + simulation_execution_id: UUID + role: SimulationRole + policy: dict[str, JsonValue] + population: PopulationInput + year: Annotated[int, Field(ge=1900, le=2200)] + geography: GeographySelection + options: dict[str, JsonValue] = Field(default_factory=dict) + bundle: BundleProvenance + + @model_validator(mode="before") + @classmethod + def reject_combined_policy_input(cls, value: object) -> object: + """Reject report-shaped input that carries ``baseline`` or ``reform`` keys. + + One execution input describes one simulation; the pair belongs in + ``ReportExecutionInput``. + """ + if isinstance(value, dict) and ({"baseline", "reform"} & value.keys()): + msg = "single-simulation input cannot contain baseline or reform fields" + raise ValueError(msg) + return value diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py new file mode 100644 index 000000000..ea49b5e2d --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/comparison_runs.py @@ -0,0 +1,365 @@ +"""A comparison run's invocation context, statuses, and its parent and child records.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime, timedelta +from enum import StrEnum +from typing import Annotated, Literal +from uuid import UUID + +from pydantic import Field, model_validator + +# Models subclass these bases through the package's path, which ruff.toml +# lists as runtime-evaluated; base.py binds them in the package first. +from policyengine_simulation_contract.stage12_execution import StrictContractModel +from policyengine_simulation_contract.stage12_execution.base import ( + ContractText, + CountryId, + ObservabilityId, + Sha256Digest, + SimulationRole, + StorageUri, +) + + +class Stage12InvocationContext(StrictContractModel): + """Deployment and retention context for one comparison-run invocation. + + Records the originating request, the environments, the worker version and + callable, the manifest digests and the artifact prefix. Automatic runs also + name the production function call they are compared with. Artifacts are kept + for at most 30 days after ``created_at``. + """ + + contract_version: Literal[1] = 1 + request_id: ContractText + # The logical environment labels durable records and artifact paths. + environment: ContractText + # Modal uses ``main`` for production, so it cannot be inferred from the + # logical ``production`` label. + modal_environment: ContractText + worker_version: ContractText + modal_application: ContractText + simulation_callable: ContractText + version_manifest_sha256: Sha256Digest + bundle_manifest_sha256: Sha256Digest + artifact_prefix: Annotated[str, Field(min_length=1, max_length=2048)] + # Automatic runs retain the existing Modal function-call identifier so the + # coordinator can retrieve the production aggregate after Stage 12 finishes. + # Direct Stage 12-only runs have no associated production call. + production_function_call_id: ContractText | None = None + created_at: datetime + retention_expires_at: datetime + + @model_validator(mode="after") + def validate_retention(self) -> Stage12InvocationContext: + """Require aware timestamps and a retention window of at most 30 days.""" + _validate_retention(self.created_at, self.retention_expires_at) + return self + + +ErrorCode = Annotated[str, Field(min_length=1, max_length=64)] +ErrorSummary = Annotated[str, Field(min_length=1, max_length=512)] +MAX_COMPARISON_RUN_ARTIFACT_RETENTION = timedelta(days=30) + + +def _validate_retention(created_at: datetime, retention_expires_at: datetime) -> None: + if created_at.tzinfo is None or retention_expires_at.tzinfo is None: + msg = "comparison retention timestamps must include a timezone" + raise ValueError(msg) + retention = retention_expires_at - created_at + if retention <= timedelta(0): + msg = "comparison retention must be greater than zero" + raise ValueError(msg) + if retention > MAX_COMPARISON_RUN_ARTIFACT_RETENTION: + msg = "comparison retention must not exceed 30 days" + raise ValueError(msg) + + +class ComparisonRunLifecycleStatus(StrEnum): + """Overall status of a comparison run.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + INCOMPLETE = "incomplete" + SKIPPED = "skipped" + + +class ComparisonRunAggregationStatus(StrEnum): + """Status of aggregating a comparison run's simulations into a report.""" + + NOT_STARTED = "not_started" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +class ResultComparisonStatus(StrEnum): + """Outcome of checking a comparison run's aggregate result against production.""" + + NOT_REQUESTED = "not_requested" + PENDING = "pending" + RUNNING = "running" + MATCHED = "matched" + DIFFERENT = "different" + FAILED = "failed" + + +class ComparisonReportRecord(StrictContractModel): + """Complete producer/consumer contract for one temporary parent row.""" + + contract_version: Literal[1] = 1 + evaluation_id: UUID + status: ComparisonRunLifecycleStatus + aggregation_status: ComparisonRunAggregationStatus + environment: ContractText + calculation_flow: ContractText + originating_request_id: ContractText + observability_id: ObservabilityId | None = None + production_identity: ContractText + incumbent_execution_id: ContractText | None = None + worker_version: ContractText + modal_application: ContractText + report_coordinator_callable: ContractText + version_manifest_sha256: Sha256Digest + policyengine_version: ContractText + country_package_name: ContractText + country_package_version: ContractText + country: CountryId + dataset_identity: ContractText + dataset_uri: Annotated[str, Field(min_length=1, max_length=2048)] + data_package_name: ContractText + data_package_version: ContractText + data_artifact_revision: ContractText + coordinator_invocation_id: ContractText | None = None + error_code: ErrorCode | None = None + error_summary: ErrorSummary | None = None + aggregate_output_uri: StorageUri | None = None + aggregate_output_sha256: Sha256Digest | None = None + aggregate_schema_version: Annotated[int, Field(ge=1)] | None = None + comparison_status: ResultComparisonStatus = ResultComparisonStatus.NOT_REQUESTED + comparison_output_uri: StorageUri | None = None + comparison_output_sha256: Sha256Digest | None = None + comparison_schema_version: Annotated[int, Field(ge=1)] | None = None + comparison_completed_at: datetime | None = None + comparison_error_code: ErrorCode | None = None + comparison_error_summary: ErrorSummary | None = None + created_at: datetime + updated_at: datetime + started_at: datetime | None = None + completed_at: datetime | None = None + retention_expires_at: datetime + + @model_validator(mode="after") + def validate_state(self) -> ComparisonReportRecord: + """Check that the row's fields are consistent with its lifecycle states. + + Aggregate and comparison output fields come all-or-nothing, a success + needs its outputs and ``completed_at``, failures and skips need an + ``error_code``, and the retention window is at most 30 days. + """ + _validate_retention(self.created_at, self.retention_expires_at) + self._validate_report_state() + self._validate_comparison_state() + return self + + def _validate_report_state(self) -> None: + """Check the aggregate output fields against the report's status.""" + output_fields = ( + self.aggregate_output_uri, + self.aggregate_output_sha256, + self.aggregate_schema_version, + ) + if any(value is not None for value in output_fields) and not all( + value is not None for value in output_fields + ): + msg = "aggregate output fields must be supplied together" + raise ValueError(msg) + if self.status is ComparisonRunLifecycleStatus.SUCCEEDED: + if self.aggregation_status is not ComparisonRunAggregationStatus.SUCCEEDED: + msg = "a successful report requires successful aggregation" + raise ValueError(msg) + if not all(value is not None for value in output_fields): + msg = "a successful report requires an aggregate output" + raise ValueError(msg) + if self.completed_at is None: + msg = "a successful report requires completed_at" + raise ValueError(msg) + if ( + self.status + in { + ComparisonRunLifecycleStatus.FAILED, + ComparisonRunLifecycleStatus.SKIPPED, + } + and self.error_code is None + ): + msg = "a failed or skipped report requires an error_code" + raise ValueError(msg) + + def _validate_comparison_state(self) -> None: + """Check the result comparison fields against the comparison's status.""" + if self.comparison_status in { + ResultComparisonStatus.NOT_REQUESTED, + ResultComparisonStatus.PENDING, + ResultComparisonStatus.RUNNING, + }: + self._validate_incomplete_comparison() + elif self.comparison_status in { + ResultComparisonStatus.MATCHED, + ResultComparisonStatus.DIFFERENT, + }: + self._validate_completed_comparison() + else: + self._validate_failed_comparison() + + @property + def _comparison_output_fields(self) -> tuple[object, ...]: + return ( + self.comparison_output_uri, + self.comparison_output_sha256, + self.comparison_schema_version, + ) + + def _validate_incomplete_comparison(self) -> None: + """Reject outputs, completion and errors on a comparison still to finish.""" + if any(value is not None for value in self._comparison_output_fields): + msg = "an incomplete comparison cannot reference an output artifact" + raise ValueError(msg) + if self.comparison_completed_at is not None: + msg = "an incomplete comparison cannot have completed_at" + raise ValueError(msg) + if self.comparison_error_code is not None: + msg = "an incomplete comparison cannot have an error_code" + raise ValueError(msg) + if self.comparison_error_summary is not None: + msg = "an incomplete comparison cannot have an error_summary" + raise ValueError(msg) + + def _validate_completed_comparison(self) -> None: + """Require an output and completion, and no error, on a finished comparison.""" + if not all(value is not None for value in self._comparison_output_fields): + msg = "a completed comparison requires an output artifact" + raise ValueError(msg) + if self.comparison_completed_at is None: + msg = "a completed comparison requires completed_at" + raise ValueError(msg) + if self.comparison_error_code is not None: + msg = "a completed comparison cannot have an error_code" + raise ValueError(msg) + if self.comparison_error_summary is not None: + msg = "a completed comparison cannot have an error_summary" + raise ValueError(msg) + + def _validate_failed_comparison(self) -> None: + """Require completion and an error code, but no output, on a failure.""" + if any(value is not None for value in self._comparison_output_fields): + msg = "a failed comparison cannot reference an output artifact" + raise ValueError(msg) + if self.comparison_completed_at is None: + msg = "a failed comparison requires completed_at" + raise ValueError(msg) + if self.comparison_error_code is None: + msg = "a failed comparison requires an error_code" + raise ValueError(msg) + + +class ComparisonSimulationRecord(StrictContractModel): + """Complete producer/consumer contract for one temporary child row.""" + + contract_version: Literal[1] = 1 + simulation_execution_id: UUID + evaluation_id: UUID + role: SimulationRole + input_sha256: Sha256Digest + worker_version: ContractText + modal_application: ContractText + simulation_callable: ContractText + version_manifest_sha256: Sha256Digest + modal_invocation_id: ContractText | None = None + status: ComparisonRunLifecycleStatus + error_code: ErrorCode | None = None + error_summary: ErrorSummary | None = None + output_uri: StorageUri | None = None + output_sha256: Sha256Digest | None = None + output_schema_version: Annotated[int, Field(ge=1)] | None = None + row_identity_columns: tuple[ContractText, ...] | None = None + row_count: Annotated[int, Field(ge=0)] | None = None + row_identity_sha256: Sha256Digest | None = None + created_at: datetime + updated_at: datetime + started_at: datetime | None = None + completed_at: datetime | None = None + retention_expires_at: datetime + + @model_validator(mode="after") + def validate_state(self) -> ComparisonSimulationRecord: + """Check that the row's fields are consistent with its lifecycle status. + + Output and row-identity fields come all-or-nothing, a success needs + them and ``completed_at``, failures and skips need an ``error_code``, + and the retention window is at most 30 days. + """ + _validate_retention(self.created_at, self.retention_expires_at) + output_fields = ( + self.output_uri, + self.output_sha256, + self.output_schema_version, + self.row_identity_columns, + self.row_count, + self.row_identity_sha256, + ) + if any(value is not None for value in output_fields) and not all( + value is not None for value in output_fields + ): + msg = "simulation output fields must be supplied together" + raise ValueError(msg) + if self.row_identity_columns is not None: + if not self.row_identity_columns: + msg = "row_identity_columns must not be empty" + raise ValueError(msg) + if len(self.row_identity_columns) != len(set(self.row_identity_columns)): + msg = "row_identity_columns must be unique" + raise ValueError(msg) + if self.status is ComparisonRunLifecycleStatus.SUCCEEDED: + if not all(value is not None for value in output_fields): + msg = "a successful simulation requires an output" + raise ValueError(msg) + if self.completed_at is None: + msg = "a successful simulation requires completed_at" + raise ValueError(msg) + if ( + self.status + in { + ComparisonRunLifecycleStatus.FAILED, + ComparisonRunLifecycleStatus.SKIPPED, + } + and self.error_code is None + ): + msg = "a failed or skipped simulation requires an error_code" + raise ValueError(msg) + return self + + +@dataclass(frozen=True) +class ComparisonReportPersistenceResult: + """The stored parent row after an idempotent create. + + ``created`` is False when an identical row already existed and was reused. + """ + + record: ComparisonReportRecord + created: bool + + +@dataclass(frozen=True) +class ComparisonSimulationPersistenceResult: + """The stored child row after an idempotent create. + + ``created`` is False when an identical row already existed and was reused. + """ + + record: ComparisonSimulationRecord + created: bool diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/report_inputs.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/report_inputs.py new file mode 100644 index 000000000..8d4a4ea42 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/report_inputs.py @@ -0,0 +1,208 @@ +"""A report's simulation pair, its required outputs, and the planned output schema.""" + +from __future__ import annotations + +import json +from enum import StrEnum +from hashlib import sha256 +from typing import Annotated, Literal +from uuid import UUID + +from pydantic import Field, field_validator, model_validator + +# Models subclass these bases through the package's path, which ruff.toml +# lists as runtime-evaluated; base.py binds them in the package first. +from policyengine_simulation_contract.stage12_execution import ( + SimulationExecutionInput, + StrictContractModel, +) +from policyengine_simulation_contract.stage12_execution.base import ( + ContractText, + CountryId, + SimulationRole, +) + + +class ReportAggregate(StrEnum): + """A family of aggregate results that a comparison report computes.""" + + BUDGET = "budget" + POVERTY = "poverty" + INEQUALITY = "inequality" + DISTRIBUTIONAL = "distributional" + WINNERS_AND_LOSERS = "winners_and_losers" + GEOGRAPHIC = "geographic" + PROGRAM_STATISTICS = "program_statistics" + + +class ReportOutputRequirements(StrictContractModel): + """Report features that determine which simulation columns must exist.""" + + schema_version: Literal[1] = 1 + aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] + include_cliff_impacts: bool + labor_supply_response_active: bool + + @field_validator("aggregates") + @classmethod + def require_complete_aggregate_profile( + cls, + value: tuple[ReportAggregate, ...], + ) -> tuple[ReportAggregate, ...]: + """Require every ``ReportAggregate``, once each and in declaration order.""" + if len(value) != len(set(value)): + msg = "report output aggregates must be unique" + raise ValueError(msg) + if value != tuple(ReportAggregate): + msg = "Stage 12 currently requires the complete aggregate profile" + raise ValueError(msg) + return value + + +class EntityOutputPlan(StrictContractModel): + """Required materialized columns for one country-model entity.""" + + entity: ContractText + materialized_variables: Annotated[ + tuple[ContractText, ...], + Field(min_length=1), + ] + additional_variables: tuple[ContractText, ...] = () + dataset_variables: tuple[ContractText, ...] = () + + @field_validator( + "materialized_variables", + "additional_variables", + "dataset_variables", + ) + @classmethod + def require_canonical_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Require variable names to be unique and sorted, so plans hash stably.""" + if len(value) != len(set(value)): + msg = "output-plan variables must be unique" + raise ValueError(msg) + if value != tuple(sorted(value)): + msg = "output-plan variables must use canonical order" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def require_variable_subsets(self) -> EntityOutputPlan: + """Require the additional and dataset variables not to overlap. + + Both must also be subsets of ``materialized_variables``. + """ + if not set(self.additional_variables).issubset(self.materialized_variables): + msg = "additional output-plan variables must be a materialized subset" + raise ValueError(msg) + if not set(self.dataset_variables).issubset(self.materialized_variables): + msg = "dataset output-plan variables must be a materialized subset" + raise ValueError(msg) + if set(self.additional_variables).intersection(self.dataset_variables): + msg = ( + "calculated additional variables and dataset variables must be disjoint" + ) + raise ValueError(msg) + return self + + +class Stage12OutputPlan(StrictContractModel): + """One immutable output schema shared by both report simulations.""" + + schema_version: Literal[1] = 1 + country: CountryId + requirements: ReportOutputRequirements + entities: Annotated[tuple[EntityOutputPlan, ...], Field(min_length=1)] + + @field_validator("entities") + @classmethod + def require_canonical_entities( + cls, + value: tuple[EntityOutputPlan, ...], + ) -> tuple[EntityOutputPlan, ...]: + """Require entity plans to be unique and sorted by entity name.""" + names = tuple(entity.entity for entity in value) + if len(names) != len(set(names)): + msg = "output-plan entities must be unique" + raise ValueError(msg) + if names != tuple(sorted(names)): + msg = "output-plan entities must use canonical order" + raise ValueError(msg) + return value + + +def stage12_output_plan_sha256(plan: Stage12OutputPlan) -> str: + """Return the stable digest recorded by each Stage 12 child artifact.""" + payload = json.dumps( + plan.model_dump(mode="json"), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + return sha256(payload).hexdigest() + + +class PlannedSimulationExecutionInput(SimulationExecutionInput): + """Internal coordinator-to-worker input with a resolved output schema.""" + + output_plan: Stage12OutputPlan + + @model_validator(mode="after") + def require_output_plan_country(self) -> PlannedSimulationExecutionInput: + """Reject an output plan built for a different country than the geography.""" + if self.output_plan.country != self.geography.country: + msg = "output plan country must match simulation geography" + raise ValueError(msg) + return self + + +class ReportExecutionInput(StrictContractModel): + """A baseline and reform simulation pair to run and aggregate into one report. + + Both simulations must belong to this report, carry the baseline and reform + roles and distinct execution IDs, and differ only in their policy. + """ + + contract_version: Literal[1] = 1 + evaluation_id: UUID + baseline: SimulationExecutionInput + reform: SimulationExecutionInput + requested_aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] + + @model_validator(mode="after") + def require_aligned_simulations(self) -> ReportExecutionInput: + """Check the two simulations form one comparable report. + + Roles must be baseline and reform, both must belong to this report with + distinct execution ids, and everything but the policy must match. + """ + if self.baseline.role is not SimulationRole.BASELINE: + msg = "baseline input must use the baseline role" + raise ValueError(msg) + if self.reform.role is not SimulationRole.REFORM: + msg = "reform input must use the reform role" + raise ValueError(msg) + if self.baseline.evaluation_id != self.evaluation_id: + msg = "baseline evaluation_id must match the report" + raise ValueError(msg) + if self.reform.evaluation_id != self.evaluation_id: + msg = "reform evaluation_id must match the report" + raise ValueError(msg) + if self.baseline.simulation_execution_id == self.reform.simulation_execution_id: + msg = "baseline and reform simulation identifiers must differ" + raise ValueError(msg) + for field_name in ( + "population", + "year", + "geography", + "options", + "bundle", + ): + if getattr(self.baseline, field_name) != getattr(self.reform, field_name): + msg = f"baseline and reform {field_name} must match before execution" + raise ValueError(msg) + if len(self.requested_aggregates) != len(set(self.requested_aggregates)): + msg = "requested aggregates must be unique" + raise ValueError(msg) + return self diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/result_artifacts.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/result_artifacts.py new file mode 100644 index 000000000..855eab6c9 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/result_artifacts.py @@ -0,0 +1,185 @@ +"""The stored simulation, report and comparison artifacts of a comparison run.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Annotated, Literal +from uuid import UUID + +from pydantic import Field, JsonValue, field_validator, model_validator + +from policyengine_simulation_contract.spm import SPMComparisonProvenance, SPMProvenance + +# Models subclass these bases through the package's path, which ruff.toml +# lists as runtime-evaluated; base.py binds them in the package first. +from policyengine_simulation_contract.stage12_execution import StrictContractModel +from policyengine_simulation_contract.stage12_execution.base import ( + ArtifactReference, + BundleProvenance, + ContractText, + Sha256Digest, + SimulationRole, +) +from policyengine_simulation_contract.stage12_execution.report_inputs import ( + ReportAggregate, +) + + +class RowIdentity(StrictContractModel): + """Fingerprint of the rows in a simulation artifact. + + Names the identifier columns and gives the row count and a SHA-256 over the + row identifiers, so artifacts can be checked to cover the same rows. + """ + + schema_version: Literal[1] = 1 + identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)] + row_count: Annotated[int, Field(ge=0)] + identity_sha256: Sha256Digest + + @field_validator("identifier_columns") + @classmethod + def require_unique_columns(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Reject repeated identifier column names.""" + if len(value) != len(set(value)): + msg = "row identifier columns must be unique" + raise ValueError(msg) + return value + + +class SimulationArtifactDescriptor(StrictContractModel): + """The stored output of one simulation in a comparison run. + + Points to the artifact and records the digest of the output plan it follows, + its row identity, its bundle provenance and, for SPM calculations, the + calculation receipt. + """ + + contract_version: Literal[1] = 1 + evaluation_id: UUID + simulation_execution_id: UUID + role: SimulationRole + artifact: ArtifactReference + output_schema_version: Literal[1] = 1 + output_plan_sha256: Sha256Digest + row_identity: RowIdentity + bundle: BundleProvenance + calculation_provenance: SPMProvenance | None = None + + +class AggregateReportArtifactDescriptor(StrictContractModel): + """The stored aggregate results of a comparison run. + + Records the digests of the baseline and reform artifacts the results were + computed from, and the bundle provenance. + """ + + contract_version: Literal[1] = 1 + evaluation_id: UUID + artifact: ArtifactReference + aggregate_schema_version: Literal[1] = 1 + baseline_artifact_sha256: Sha256Digest + reform_artifact_sha256: Sha256Digest + bundle: BundleProvenance + + +class AggregateReportArtifactPayload(StrictContractModel): + """Complete JSON payload stored in one aggregate report artifact.""" + + contract_version: Literal[1] = 1 + aggregate_schema_version: Literal[1] = 1 + evaluation_id: UUID + requested_aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)] + bundle: BundleProvenance + result: dict[str, JsonValue] + + @field_validator("requested_aggregates") + @classmethod + def require_unique_aggregates( + cls, + value: tuple[ReportAggregate, ...], + ) -> tuple[ReportAggregate, ...]: + """Reject repeated aggregate names.""" + if len(value) != len(set(value)): + msg = "requested aggregates must be unique" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def validate_spm_provenance(self) -> AggregateReportArtifactPayload: + """Reject duplicate or noncanonical SPM data in aggregate artifacts.""" + if "spm_config" in self.result: + msg = "aggregate result must not include legacy spm_config" + raise ValueError(msg) + # Reports without US SPM calculations serialize this optional field + # as null. Validate a supplied record, not the absence of one. + provenance = self.result.get("spm_provenance") + if provenance is not None: + SPMComparisonProvenance.model_validate(provenance) + return self + + +class ResultDifference(StrictContractModel): + """One exact aggregate-result difference identified by JSON Pointer.""" + + path: Annotated[str, Field(max_length=4096)] + production_present: bool + stage12_present: bool + production_value: JsonValue = None + stage12_value: JsonValue = None + absolute_delta: float | None = None + relative_delta: float | None = None + + @model_validator(mode="after") + def validate_presence(self) -> ResultDifference: + """Check presence flags against values, and that deltas come as a pair. + + An absent side must be null, and a relative delta needs an absolute one; + only an absolute delta may stand alone (a zero production value). + """ + if not self.production_present and not self.stage12_present: + msg = "a result difference must be present in at least one result" + raise ValueError(msg) + if not self.production_present and self.production_value is not None: + msg = "an absent production value must be null" + raise ValueError(msg) + if not self.stage12_present and self.stage12_value is not None: + msg = "an absent Stage 12 value must be null" + raise ValueError(msg) + # A zero production value has no finite relative delta, so an absolute + # delta with no relative delta is the only permitted partial pair. + if (self.absolute_delta is None) != (self.relative_delta is None) and ( + self.absolute_delta is None or self.relative_delta is not None + ): + msg = "numeric deltas must be supplied together" + raise ValueError(msg) + return self + + +class ResultComparisonArtifactPayload(StrictContractModel): + """Private comparison receipt without duplicate complete result objects.""" + + schema_version: Literal[1] = 1 + evaluation_id: UUID + production_job_id: ContractText + compared_at: datetime + status: Literal["matched", "different"] + production_result_sha256: Sha256Digest + stage12_result_sha256: Sha256Digest + difference_count: Annotated[int, Field(ge=0)] + differences: tuple[ResultDifference, ...] + + @model_validator(mode="after") + def validate_differences(self) -> ResultComparisonArtifactPayload: + """Require a timezone-aware timestamp and a count and status that match.""" + if self.compared_at.tzinfo is None: + msg = "comparison timestamp must include a timezone" + raise ValueError(msg) + if self.difference_count != len(self.differences): + msg = "difference_count must equal the number of differences" + raise ValueError(msg) + expected_status = "matched" if not self.differences else "different" + if self.status != expected_status: + msg = "comparison status does not match its differences" + raise ValueError(msg) + return self diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordinated_simulations.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordinated_simulations.py new file mode 100644 index 000000000..86016e19b --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordinated_simulations.py @@ -0,0 +1,263 @@ +"""The coordinator's child simulations: their records, failures and cancellation.""" + +from __future__ import annotations + +import logging +from datetime import UTC, datetime +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + ComparisonSimulationRecord, + PlannedSimulationExecutionInput, + Stage12InvocationContext, +) + +from .failure_recording import FailureSite, report_failure +from .failures import Stage12FailureDetail, failure_detail_from_record +from .partition import should_segment_simulation +from .simulation import simulation_input_sha256 + +if TYPE_CHECKING: + from policyengine_observability import ObservabilityRuntime + + from .dependencies import ChildCall, ComparisonStore + from .failure_causes import StageProgress + +# Moved out of coordination.py; child failures still log under its name. +logger = logging.getLogger( + "policyengine_simulation_executor.stage12_runtime.coordination" +) + + +def _log_child_operation_failure( + level: int, + message: str, + *, + simulation: PlannedSimulationExecutionInput, + error: Exception, +) -> None: + """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 + 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, + 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 + + +SIMULATION_CANCELLED_ERROR_CODE = "simulation_cancelled" + + +def cancellation(reason: str) -> Stage12FailureDetail: + """Describe a child the coordinator stopped; ``reason`` ends the summary.""" + return Stage12FailureDetail( + error_code=SIMULATION_CANCELLED_ERROR_CODE, + error_summary=f"Cancelled after {reason}.", + ) + + +def cancel_pending_simulations( + *, + calls: dict[str, ChildCall], + simulations: dict[str, PlannedSimulationExecutionInput], + roles: set[str], + store: ComparisonStore, + detail: Stage12FailureDetail, +) -> None: + """Record each started, unfinished child as failed, then stop it. + + 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: + _replace_running_simulation( + store=store, + simulation=simulation, + status=ComparisonRunLifecycleStatus.FAILED, + detail=detail, + ) + except Exception as error: # noqa: BLE001 # a failed cancellation write is logged and must not replace the failure that caused it + _log_child_operation_failure( + logging.ERROR, + "Stage 12 could not record cancellation of %s simulation %s " + "in report %s (%s)", + simulation=simulation, + error=error, + ) + 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() + except Exception as error: # noqa: BLE001 # a failed Modal cancellation is logged and must not replace the failure that caused it + _log_child_operation_failure( + logging.WARNING, + "Stage 12 could not cancel %s simulation %s in report %s (%s)", + simulation=simulation, + error=error, + ) + + +def simulation_failure_attributes( + simulation: PlannedSimulationExecutionInput, +) -> dict[str, object]: + """Return the telemetry attributes that identify a failed simulation.""" + 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( + *, + simulation: PlannedSimulationExecutionInput, + 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, + ) + + +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 # a failed state write is logged and must not replace the child's failure + _log_child_state_failure(simulation=simulation, error=state_error) + + +def child_failure_detail( + *, + error: Exception, + simulation: PlannedSimulationExecutionInput, + store: ComparisonStore, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> 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, + FailureSite( + scope="child_wait", + default_code="simulation_invocation_failed", + attributes=simulation_failure_attributes(simulation), + ), + runtime=runtime, + progress=progress, + ) + try: + child = store.get_simulation(simulation.simulation_execution_id) + except Exception as state_error: # noqa: BLE001 # if the child's record cannot be read, its call failure is reported instead + _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 + + +def child_record( + *, + simulation: PlannedSimulationExecutionInput, + context: Stage12InvocationContext, + function_name: str, +) -> ComparisonSimulationRecord: + """Return the pending child record for ``simulation`` in this invocation.""" + return ComparisonSimulationRecord( + simulation_execution_id=simulation.simulation_execution_id, + evaluation_id=simulation.evaluation_id, + role=simulation.role, + input_sha256=simulation_input_sha256(simulation), + worker_version=context.worker_version, + modal_application=context.modal_application, + simulation_callable=function_name, + version_manifest_sha256=context.version_manifest_sha256, + status=ComparisonRunLifecycleStatus.PENDING, + created_at=context.created_at, + updated_at=context.created_at, + retention_expires_at=context.retention_expires_at, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 280467162..d64f9fd2f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -1,1025 +1,57 @@ -"""Coordination of a Stage 12 report's baseline and reform simulations.""" +"""Coordination of a Stage 12 report's baseline and reform simulations. + +coordinate_report runs one report; report_coordination, +coordinated_simulations and parent_claim hold its steps, and this module +re-exports their public names. +""" from __future__ import annotations -import dataclasses -import logging -import time -from dataclasses import dataclass -from datetime import UTC, datetime -from functools import partial -from hashlib import sha256 -from typing import TYPE_CHECKING, Any -from uuid import uuid4 +from typing import Any from policyengine_simulation_contract.stage12_execution import ( - AggregateReportArtifactDescriptor, ComparisonReportRecord, - ComparisonRunAggregationStatus, - ComparisonRunLifecycleStatus, - ComparisonSimulationRecord, - PlannedSimulationExecutionInput, ReportExecutionInput, - ResultComparisonStatus, - SimulationArtifactDescriptor, - SimulationRole, Stage12InvocationContext, - Stage12OutputPlan, ) from policyengine_simulation_observability.stages import ( STAGE12_CANONICAL_REPORT_STAGES, STAGE12_SHADOW_REPORT_STAGES, Stage, - StagePlan, ) -from policyengine_simulation_executor.stage12_artifacts import ( - Stage12ArtifactStore, - deserialize_calculation_provenance, - deserialize_simulation_frames, - deserialize_uk_local_authority_metadata, -) from policyengine_simulation_executor.stage12_timeouts import ( SIMULATION_WAIT_TIMEOUT_SECONDS, ) -from .aggregation import ( - RetainedSimulationOutput, - build_aggregate_report, - validate_aligned_outputs, - validate_uk_local_authority_metadata, -) -from .comparison import ComparisonServices, compare_completed_report -from .dependencies import ( - ChildCall, - ChildInvoker, - ChildSpawnRequest, - ComparisonStore, - ModalChildInvoker, - artifact_store, - runtime_store, -) -from .failure_causes import ( - StageProgress, - StageTracking, - record_failure_cause, - tracked_stage, -) -from .failure_recording import ( - FailureSite, - persist_failure_record, - report_failure, - report_original_failure, +from .coordinated_simulations import SIMULATION_CANCELLED_ERROR_CODE, logger +from .dependencies import ModalChildInvoker, artifact_store, runtime_store +from .failure_causes import StageProgress +from .parent_claim import ( + claim_parent, + validate_dispatch_context, + validate_submitted_parent, ) -from .failures import ( - Stage12ExecutionError, - Stage12FailureDetail, - failure_detail_from_record, +from .report_coordination import ( + SIMULATION_POLL_BACKOFF_FACTOR, + SIMULATION_POLL_INITIAL_SECONDS, + SIMULATION_POLL_MAX_SECONDS, + CoordinatorDependencies, + CoordinatorInvocation, + ReportCoordination, ) -from .output_planning import ( - plan_simulation_input, - resolve_report_output_plan, - validate_output_frames, -) -from .partition import should_segment_simulation -from .simulation import descriptor_from_record, simulation_input_sha256 - -if TYPE_CHECKING: - from collections.abc import Callable, Mapping - from contextlib import AbstractContextManager - - from policyengine_observability import ObservabilityRuntime - from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata - -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 _log_child_operation_failure( - level: int, - message: str, - *, - simulation: PlannedSimulationExecutionInput, - error: Exception, -) -> None: - """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 - 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, - 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 - -SIMULATION_CANCELLED_ERROR_CODE = "simulation_cancelled" - - -def _cancellation(reason: str) -> Stage12FailureDetail: - """Describe a child the coordinator stopped; ``reason`` ends the summary.""" - return Stage12FailureDetail( - error_code=SIMULATION_CANCELLED_ERROR_CODE, - error_summary=f"Cancelled after {reason}.", - ) - - -def _cancel_pending_simulations( - *, - calls: dict[str, ChildCall], - simulations: dict[str, PlannedSimulationExecutionInput], - roles: set[str], - store: ComparisonStore, - detail: Stage12FailureDetail, -) -> None: - """Record each started, unfinished child as failed, then stop it. - - 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: - _replace_running_simulation( - store=store, - simulation=simulation, - status=ComparisonRunLifecycleStatus.FAILED, - detail=detail, - ) - except Exception as error: # noqa: BLE001 # a failed cancellation write is logged and must not replace the failure that caused it - _log_child_operation_failure( - logging.ERROR, - "Stage 12 could not record cancellation of %s simulation %s " - "in report %s (%s)", - simulation=simulation, - error=error, - ) - 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() - except Exception as error: # noqa: BLE001 # a failed Modal cancellation is logged and must not replace the failure that caused it - _log_child_operation_failure( - logging.WARNING, - "Stage 12 could not cancel %s simulation %s in report %s (%s)", - simulation=simulation, - error=error, - ) - - -def _simulation_failure_attributes( - 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( - *, - simulation: PlannedSimulationExecutionInput, - 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, - ) - - -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 # a failed state write is logged and must not replace the child's failure - _log_child_state_failure(simulation=simulation, error=state_error) - - -def _child_failure_detail( - *, - error: Exception, - simulation: PlannedSimulationExecutionInput, - store: ComparisonStore, - runtime: ObservabilityRuntime | None, - progress: StageProgress, -) -> 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, - FailureSite( - scope="child_wait", - default_code="simulation_invocation_failed", - attributes=_simulation_failure_attributes(simulation), - ), - runtime=runtime, - progress=progress, - ) - try: - child = store.get_simulation(simulation.simulation_execution_id) - except Exception as state_error: # noqa: BLE001 # if the child's record cannot be read, its call failure is reported instead - _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 - - -def _child_record( - *, - simulation: PlannedSimulationExecutionInput, - context: Stage12InvocationContext, - function_name: str, -) -> ComparisonSimulationRecord: - return ComparisonSimulationRecord( - simulation_execution_id=simulation.simulation_execution_id, - evaluation_id=simulation.evaluation_id, - role=simulation.role, - input_sha256=simulation_input_sha256(simulation), - worker_version=context.worker_version, - modal_application=context.modal_application, - simulation_callable=function_name, - version_manifest_sha256=context.version_manifest_sha256, - status=ComparisonRunLifecycleStatus.PENDING, - created_at=context.created_at, - updated_at=context.created_at, - retention_expires_at=context.retention_expires_at, - ) - - -def _validate_dispatch_context( - report: ReportExecutionInput, - context: Stage12InvocationContext, - application_name: str, -) -> None: - """Reject a dispatch context naming another application, worker or bundle.""" - if application_name != context.modal_application: - msg = "report coordinator application differs from dispatch context" - raise ValueError(msg) - if report.baseline.bundle.policyengine_version != context.worker_version: - msg = "report worker version differs from dispatch context" - raise ValueError(msg) - if report.baseline.bundle.bundle_manifest_sha256 != context.bundle_manifest_sha256: - msg = "report bundle digest differs from dispatch context" - raise ValueError(msg) - - -def _validate_submitted_parent( - *, - report: ReportExecutionInput, - context: Stage12InvocationContext, - parent: ComparisonReportRecord, -) -> None: - """Reject dispatch metadata that does not describe this exact report.""" - _validate_parent_matches_invocation(report=report, context=context, parent=parent) - _validate_parent_is_unclaimed(parent) - _validate_parent_production_link(report=report, context=context, parent=parent) - - -def _validate_parent_matches_invocation( - *, - report: ReportExecutionInput, - context: Stage12InvocationContext, - parent: ComparisonReportRecord, -) -> None: - """Require every parent field the dispatch fixes to match the report and context.""" - bundle = report.baseline.bundle - expected = { - "evaluation_id": report.evaluation_id, - "environment": context.environment, - "originating_request_id": context.request_id, - "worker_version": context.worker_version, - "modal_application": context.modal_application, - "version_manifest_sha256": context.version_manifest_sha256, - "policyengine_version": bundle.policyengine_version, - "country_package_name": bundle.country_package_name, - "country_package_version": bundle.country_package_version, - "country": report.baseline.geography.country, - "dataset_identity": bundle.dataset.identity, - "dataset_uri": bundle.dataset.uri, - "data_package_name": bundle.dataset.data_package_name, - "data_package_version": bundle.dataset.data_package_version, - "data_artifact_revision": bundle.dataset.artifact_revision, - "created_at": context.created_at, - "retention_expires_at": context.retention_expires_at, - } - for field, value in expected.items(): - if getattr(parent, field) != value: - msg = f"submitted parent {field} does not match invocation" - raise ValueError(msg) - - -def _validate_parent_is_unclaimed(parent: ComparisonReportRecord) -> None: - """Require a pending parent that no coordinator has started.""" - if parent.status is not ComparisonRunLifecycleStatus.PENDING: - msg = "submitted parent must be pending" - raise ValueError(msg) - if parent.aggregation_status is not ComparisonRunAggregationStatus.NOT_STARTED: - msg = "submitted parent aggregation must not be started" - raise ValueError(msg) - if parent.coordinator_invocation_id is not None: - msg = "submitted parent must not contain an invocation identifier" - raise ValueError(msg) - - -def _validate_parent_production_link( - *, - report: ReportExecutionInput, - context: Stage12InvocationContext, - parent: ComparisonReportRecord, -) -> None: - """Require a direct parent to stand alone and a shadow parent to name its job.""" - production_job_id = context.production_function_call_id - if production_job_id is None: - if parent.production_identity != f"direct:{report.evaluation_id}": - msg = "direct parent identity does not match its report" - raise ValueError(msg) - if parent.incumbent_execution_id is not None: - msg = "direct parent must not name a production execution" - raise ValueError(msg) - if parent.comparison_status is not ResultComparisonStatus.NOT_REQUESTED: - msg = "direct parent must not request result comparison" - raise ValueError(msg) - elif ( - parent.production_identity != production_job_id - or parent.incumbent_execution_id != production_job_id - ): - msg = "automatic parent does not match the production execution" - raise ValueError(msg) - - -def _claim_parent( - *, - submitted: ComparisonReportRecord, - coordinator_invocation_id: str, - store: ComparisonStore, -) -> tuple[ComparisonReportRecord, bool]: - """Create the parent in Modal or claim a retry; reject duplicate execution.""" - coordinating_at = datetime.now(UTC) - candidate = submitted.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "aggregation_status": ComparisonRunAggregationStatus.RUNNING, - "coordinator_invocation_id": coordinator_invocation_id, - "error_code": None, - "error_summary": None, - "started_at": coordinating_at, - "updated_at": coordinating_at, - "completed_at": None, - } - ) - persistence = store.create_or_resolve_report(candidate) - parent = persistence.record - if persistence.created: - return parent, True - if parent.evaluation_id != submitted.evaluation_id: - msg = "existing comparison identity has a different report identifier" - raise ValueError(msg) - if parent.status in { - ComparisonRunLifecycleStatus.RUNNING, - ComparisonRunLifecycleStatus.SUCCEEDED, - ComparisonRunLifecycleStatus.SKIPPED, - }: - return parent, False - claimed = store.replace_report( - parent.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "aggregation_status": ComparisonRunAggregationStatus.RUNNING, - "coordinator_invocation_id": coordinator_invocation_id, - "error_code": None, - "error_summary": None, - "started_at": parent.started_at or coordinating_at, - "updated_at": coordinating_at, - "completed_at": None, - } - ) - ) - return claimed, True - - -def _require_artifact_digest( - payload: bytes, - expected_sha256: str, - *, - mismatch_message: str, -) -> None: - """Raise ``ValueError(mismatch_message)`` unless ``payload`` matches its digest.""" - if sha256(payload).hexdigest() != expected_sha256: - raise ValueError(mismatch_message) - - -@dataclass(frozen=True) -class CoordinatorInvocation: - """The deployed coordinator call running a report: its Modal app and call ID.""" - - application_name: str - invocation_id: str - - -@dataclass(frozen=True) -class CoordinatorDependencies: - """Seams the coordinator uses; an unset seam selects the deployed default. - - The defaults are the Stage 12 database store, the private artifact - bucket, Modal child invocation, the standard aggregate builder and output - plan resolution, no telemetry runtime, the real clock and sleep, and the - standard child wait window. - """ - - store: ComparisonStore | None = None - artifacts: Stage12ArtifactStore | None = None - invoker: ChildInvoker | None = None - aggregator: Callable[..., dict[str, Any]] = build_aggregate_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 - - -@dataclass -class _ReportCoordination: - """One claimed report's inputs and resolved seams, run through each phase.""" - - report: ReportExecutionInput - context: Stage12InvocationContext - invocation: CoordinatorInvocation - persistence: ComparisonStore - artifact_storage: Stage12ArtifactStore - child_invoker: ChildInvoker - dependencies: CoordinatorDependencies - stage_plan: StagePlan - # The last stage entered and whether it finished, for failure causes. - progress: StageProgress - # The shared child-wait window closes at this monotonic time. - child_wait_deadline: float - # Each child by role: its plan, its started call and its descriptor once - # finished, so a failure at any point can stop the unfinished children. - simulations: dict[str, PlannedSimulationExecutionInput] = dataclasses.field( - default_factory=dict - ) - calls: dict[str, ChildCall] = dataclasses.field(default_factory=dict) - descriptors: dict[str, SimulationArtifactDescriptor] = dataclasses.field( - default_factory=dict - ) - - @property - def runtime(self) -> ObservabilityRuntime | None: - """The telemetry runtime, or None when the coordinator records nothing.""" - return self.dependencies.runtime - - @property - def stage_tracking(self) -> StageTracking: - """How this report's stages are named, recorded and timed.""" - return StageTracking(self.stage_plan, self.progress, self.runtime) - - def span( - self, - stage: Stage, - attributes: Mapping[str, object] | None = None, - ) -> AbstractContextManager[object]: - """Track ``stage`` for failure causes and time it as a span.""" - return tracked_stage(stage, attributes, tracking=self.stage_tracking) - - def stage_scope( - self, - stage: Stage, - attributes: Mapping[str, object] | None = None, - ) -> AbstractContextManager[object]: - """Track ``stage`` for failure causes and record it as an operation.""" - return tracked_stage( - stage, attributes, tracking=self.stage_tracking, operation=True - ) - - def run(self, parent: ComparisonReportRecord) -> dict[str, Any]: - """Plan, start and await both children, then aggregate and complete.""" - output_plan, planned_simulations = self._prepare(parent) - self.simulations = { - simulation.role.value: simulation for simulation in planned_simulations - } - children = self._create_children(planned_simulations) - descriptors, calls = self._start_children(planned_simulations, children) - # Every required call has been started before the coordinator observes - # baseline and reform together through nonblocking result probes. - self._wait_for_children(calls, descriptors) - return self._aggregate_and_complete(parent, output_plan, descriptors) - - def _prepare( - self, parent: ComparisonReportRecord - ) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: - """Bind the run's artifact prefix, then plan the outputs and child inputs.""" - with self.stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): - self.context = self.context.model_copy( - update={ - "artifact_prefix": ( - f"stage-12-runs/{parent.environment}/" - f"{parent.created_at:%Y}/{parent.created_at:%m}/" - f"{parent.evaluation_id}" - ), - "created_at": parent.created_at, - "retention_expires_at": parent.retention_expires_at, - } - ) - with self.stage_scope( - Stage.STAGE12_OUTPUT_PLANNING, - {"country": self.report.baseline.geography.country}, - ): - resolver = self.dependencies.output_plan_resolver - if resolver is None: - output_plan = resolve_report_output_plan( - self.report, - stage_scope=self.stage_scope, - ) - else: - output_plan = resolver(self.report) - simulations = [] - for simulation in (self.report.baseline, self.report.reform): - with self.stage_scope( - Stage.STAGE12_CHILD_INPUT_PLANNING, - {"simulation_role": simulation.role.value}, - ): - simulations.append(plan_simulation_input(simulation, output_plan)) - return output_plan, tuple(simulations) - - def _create_children( - self, planned_simulations: tuple[PlannedSimulationExecutionInput, ...] - ) -> dict[SimulationRole, ComparisonSimulationRecord]: - """Create each child's row, or resolve the row an earlier attempt made.""" - with self.span(Stage.STAGE12_CHILD_STATE_CREATE): - return { - simulation.role: self.persistence.create_or_resolve_simulation( - _child_record( - simulation=simulation, - context=self.context, - function_name=self.context.simulation_callable, - ) - ).record - for simulation in planned_simulations - } - - def _start_children( - self, - planned_simulations: tuple[PlannedSimulationExecutionInput, ...], - children: dict[SimulationRole, ComparisonSimulationRecord], - ) -> tuple[dict[str, SimulationArtifactDescriptor], dict[str, ChildCall]]: - """Reuse finished children, reattach running ones and dispatch the rest.""" - descriptors = self.descriptors - calls = self.calls - for simulation in planned_simulations: - child = children[simulation.role] - role = simulation.role.value - if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: - descriptors[role] = self._retained_descriptor(child, simulation) - elif ( - child.status is ComparisonRunLifecycleStatus.RUNNING - and child.modal_invocation_id is not None - and not child.modal_invocation_id.startswith("dispatch-pending-") - ): - calls[role] = self.child_invoker.restore(child.modal_invocation_id) - else: - calls[role] = self._dispatch_child(child, simulation) - return descriptors, calls - - def _retained_descriptor( - self, - child: ComparisonSimulationRecord, - simulation: PlannedSimulationExecutionInput, - ) -> SimulationArtifactDescriptor: - """Describe a succeeded child's digest-checked retained artifact.""" - descriptor = descriptor_from_record(child, simulation) - retained = self.artifact_storage.read(descriptor.artifact.uri) - _require_artifact_digest( - retained, - descriptor.artifact.content_sha256, - mismatch_message="retained child artifact digest mismatch", - ) - return descriptor.model_copy( - update={ - "artifact": descriptor.artifact.model_copy( - update={"size_bytes": len(retained)} - ), - "calculation_provenance": ( - deserialize_calculation_provenance(retained) - ), - } - ) - - def _dispatch_child( - self, - child: ComparisonSimulationRecord, - simulation: PlannedSimulationExecutionInput, - ) -> ChildCall: - """Mark the child dispatching, spawn it, and attach its Modal call ID. - - A spawn failure is reported and recorded on the child (unless it already - finished), every started peer is stopped, and Stage12ExecutionError is - raised with the failure. - """ - dispatching_at = datetime.now(UTC) - dispatch_placeholder = f"dispatch-pending-{uuid4()}" - self.persistence.replace_simulation( - child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "modal_invocation_id": dispatch_placeholder, - "started_at": child.started_at or dispatching_at, - "updated_at": dispatching_at, - "error_code": None, - "error_summary": None, - "completed_at": None, - } - ) - ) - try: - with self.span( - Stage.STAGE12_CHILD_DISPATCH, - {"simulation_role": simulation.role.value}, - ): - child_observability_context = ( - self.runtime.capture_context() if self.runtime is not None else None - ) - call = self.child_invoker.spawn( - ChildSpawnRequest( - application_name=self.invocation.application_name, - function_name=self.context.simulation_callable, - environment=self.context.modal_environment, - simulation=simulation.model_dump(mode="json"), - context=self.context.model_dump(mode="json"), - observability_context=child_observability_context, - ) - ) - except Exception as error: # noqa: BLE001 # any dispatch failure is recorded on the child and stops its started peers - detail = report_failure( - error, - FailureSite( - scope="child_dispatch", - default_code="simulation_dispatch_failed", - attributes=_simulation_failure_attributes(simulation), - ), - runtime=self.runtime, - progress=self.progress, - ) - _record_child_failure( - store=self.persistence, - simulation=simulation, - detail=detail, - ) - _cancel_pending_simulations( - calls=self.calls, - simulations=self.simulations, - roles=set(self.calls), - store=self.persistence, - detail=_cancellation(f"the {simulation.role.value} simulation failed"), - ) - raise Stage12ExecutionError(detail) from None - # Record the call before attaching its ID, so a failed attach still - # stops it. - self.calls[simulation.role.value] = call - self.persistence.attach_simulation_invocation( - simulation.simulation_execution_id, - expected_placeholder=dispatch_placeholder, - modal_invocation_id=call.object_id, - updated_at=datetime.now(UTC), - ) - return call - - def _wait_for_children( - self, - calls: dict[str, ChildCall], - descriptors: dict[str, SimulationArtifactDescriptor], - ) -> None: - """Observe every started child together until each has a descriptor. - - A child's failure, or the shared wait window closing, stops the - unfinished children and raises Stage12ExecutionError. - """ - with self.span(Stage.STAGE12_CHILD_WAIT): - self._observe_children(calls, descriptors) - - def _observe_children( - self, - calls: dict[str, ChildCall], - descriptors: dict[str, SimulationArtifactDescriptor], - ) -> None: - """Probe every unfinished child, without blocking, until each finishes. - - Between passes the coordinator sleeps, backing off while no child - finishes, and never past the shared wait window. - """ - wait_stage_name = self.stage_plan.name(Stage.STAGE12_CHILD_WAIT) - timeout_seconds = self.dependencies.wait_timeout_seconds - pending = set(calls) - poll_interval = SIMULATION_POLL_INITIAL_SECONDS - - while pending: - progress_made = False - for role in sorted(pending): - call = calls[role] - try: - # The probe records the child's role, so a failure names it. - with self.progress.stage(wait_stage_name, role=role): - result = call.get(timeout=0) - descriptor = SimulationArtifactDescriptor.model_validate(result) - except TimeoutError: - continue - except Exception as error: # noqa: BLE001 # a child call's failure can be any exception type; it is recorded and stops the other child - detail = _child_failure_detail( - error=error, - simulation=self.simulations[role], - store=self.persistence, - runtime=self.runtime, - progress=self.progress, - ) - _cancel_pending_simulations( - calls=calls, - simulations=self.simulations, - roles=pending - {role}, - store=self.persistence, - detail=_cancellation(f"the {role} simulation failed"), - ) - raise Stage12ExecutionError(detail) from None - descriptors[role] = descriptor - pending.remove(role) - progress_made = True - - if not pending: - return - remaining_seconds = self.child_wait_deadline - self.dependencies.monotonic() - if remaining_seconds <= 0: - detail = Stage12FailureDetail( - error_code="simulation_wait_timeout", - error_summary=( - "Stage 12 simulations did not finish within " - f"{timeout_seconds:g} seconds of the report coordinator " - "starting." - ), - ) - wait_timeout = Stage12ExecutionError(detail) - report_original_failure( - wait_timeout, - runtime=self.runtime, - failure_scope="child_wait_timeout", - ) - _cancel_pending_simulations( - calls=calls, - simulations=self.simulations, - roles=pending, - store=self.persistence, - detail=_cancellation( - f"the {timeout_seconds:g}-second wait window closed" - ), - ) - raise wait_timeout - # Never sleep past the deadline, so the window closes on time. - self.dependencies.sleep(min(poll_interval, remaining_seconds)) - poll_interval = ( - SIMULATION_POLL_INITIAL_SECONDS - if progress_made - else _next_poll_interval(poll_interval) - ) - - def _read_retained_outputs( - self, - output_plan: Stage12OutputPlan, - baseline: SimulationArtifactDescriptor, - reform: SimulationArtifactDescriptor, - ) -> tuple[ - RetainedSimulationOutput, - RetainedSimulationOutput, - UKLocalAuthorityMetadata | None, - ]: - """Read, digest-check and validate both children's retained artifacts.""" - validate_aligned_outputs(self.report, output_plan, baseline, reform) - with self.span(Stage.STAGE12_ARTIFACT_READ): - baseline_payload = self.artifact_storage.read(baseline.artifact.uri) - reform_payload = self.artifact_storage.read(reform.artifact.uri) - _require_artifact_digest( - baseline_payload, - baseline.artifact.content_sha256, - mismatch_message="baseline artifact digest mismatch", - ) - _require_artifact_digest( - reform_payload, - reform.artifact.content_sha256, - mismatch_message="reform artifact digest mismatch", - ) - baseline_frames = deserialize_simulation_frames(baseline_payload) - reform_frames = deserialize_simulation_frames(reform_payload) - uk_local_authority_metadata = validate_uk_local_authority_metadata( - self.report.baseline.geography.country, - deserialize_uk_local_authority_metadata(baseline_payload), - deserialize_uk_local_authority_metadata(reform_payload), - region_code=self.report.baseline.geography.region, - ) - validate_output_frames(baseline_frames, output_plan) - validate_output_frames(reform_frames, output_plan) - return ( - RetainedSimulationOutput(baseline_frames, baseline), - RetainedSimulationOutput(reform_frames, reform), - uk_local_authority_metadata, - ) - - def _aggregate_and_complete( - self, - parent: ComparisonReportRecord, - output_plan: Stage12OutputPlan, - descriptors: dict[str, SimulationArtifactDescriptor], - ) -> dict[str, Any]: - """Aggregate both outputs, store the aggregate and complete the parent. - - A shadow run's result is then compared with the production result. - """ - baseline, reform, uk_local_authority_metadata = self._read_retained_outputs( - output_plan, descriptors["baseline"], descriptors["reform"] - ) - with self.span(Stage.STAGE12_AGGREGATION): - aggregate = self.dependencies.aggregator( - report=self.report, - baseline=baseline, - reform=reform, - uk_local_authority_metadata=uk_local_authority_metadata, - ) - with self.span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): - aggregate_artifact = self.artifact_storage.write_aggregate( - prefix=self.context.artifact_prefix, - payload=aggregate, - ) - descriptor = AggregateReportArtifactDescriptor( - evaluation_id=self.report.evaluation_id, - artifact=aggregate_artifact, - baseline_artifact_sha256=baseline.descriptor.artifact.content_sha256, - reform_artifact_sha256=reform.descriptor.artifact.content_sha256, - bundle=self.report.baseline.bundle, - ) - completed = datetime.now(UTC) - completed_parent = self.persistence.replace_report( - parent.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.SUCCEEDED, - "aggregation_status": ComparisonRunAggregationStatus.SUCCEEDED, - "aggregate_output_uri": aggregate_artifact.uri, - "aggregate_output_sha256": aggregate_artifact.content_sha256, - "aggregate_schema_version": descriptor.aggregate_schema_version, - "updated_at": completed, - "completed_at": completed, - } - ) - ) - compare_completed_report( - parent=completed_parent, - aggregate=aggregate, - context=self.context, - services=ComparisonServices( - store=self.persistence, - artifacts=self.artifact_storage, - invoker=self.child_invoker, - runtime=self.runtime, - stage_plan=self.stage_plan, - ), - ) - return descriptor.model_dump(mode="json") - - def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: - """Record a failed run once, stop its unfinished children, fail the parent. - - Child, dispatch and wait-timeout failures raise Stage12ExecutionError - only after reporting the original failure and stopping their unfinished - peers. Any other failure is reported here, once, before any database - write. If it occurred after a child started, such as a failed - placeholder write or invocation attach, every started, unfinished child - is then stopped so none keeps running for a failed report. - """ - if isinstance(error, Stage12ExecutionError): - detail = error.detail - else: - detail = report_failure( - error, - FailureSite( - scope="report_coordination", - default_code="report_coordination_failed", - attributes={"evaluation_id": str(self.report.evaluation_id)}, - ), - runtime=self.runtime, - progress=self.progress, - ) - _cancel_pending_simulations( - calls=self.calls, - simulations=self.simulations, - roles=set(self.calls) - set(self.descriptors), - store=self.persistence, - detail=_cancellation("the coordinator failed"), - ) - record_failure_cause(self.runtime, error, detail.error_summary) - failed_at = datetime.now(UTC) - persist_failure_record( - partial( - self.persistence.replace_report, - parent.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "aggregation_status": ComparisonRunAggregationStatus.FAILED, - "error_code": detail.error_code, - "error_summary": detail.error_summary, - "updated_at": failed_at, - "completed_at": failed_at, - } - ), - ), - runtime=self.runtime, - failure_scope="report_coordination", - ) +__all__ = [ + "SIMULATION_CANCELLED_ERROR_CODE", + "SIMULATION_POLL_BACKOFF_FACTOR", + "SIMULATION_POLL_INITIAL_SECONDS", + "SIMULATION_POLL_MAX_SECONDS", + "SIMULATION_WAIT_TIMEOUT_SECONDS", + "CoordinatorDependencies", + "CoordinatorInvocation", + "coordinate_report", + "logger", +] def coordinate_report( @@ -1047,13 +79,13 @@ def coordinate_report( report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) - _validate_dispatch_context(report, context, invocation.application_name) - _validate_submitted_parent( + validate_dispatch_context(report, context, invocation.application_name) + validate_submitted_parent( report=report, context=context, parent=submitted_parent, ) - coordination = _ReportCoordination( + coordination = ReportCoordination( report=report, context=context, invocation=invocation, @@ -1070,7 +102,7 @@ def coordinate_report( child_wait_deadline=child_wait_deadline, ) with coordination.span(Stage.STAGE12_COORDINATOR_CLAIM): - parent, should_execute = _claim_parent( + parent, should_execute = claim_parent( submitted=submitted_parent, coordinator_invocation_id=invocation.invocation_id, store=coordination.persistence, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py new file mode 100644 index 000000000..dba19ed4f --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/parent_claim.py @@ -0,0 +1,169 @@ +"""Check a dispatched report against its parent record, then claim the parent.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.stage12_execution import ( + ComparisonReportRecord, + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, + ReportExecutionInput, + ResultComparisonStatus, + Stage12InvocationContext, +) + +if TYPE_CHECKING: + from .dependencies import ComparisonStore + + +def validate_dispatch_context( + report: ReportExecutionInput, + context: Stage12InvocationContext, + application_name: str, +) -> None: + """Reject a dispatch context naming another application, worker or bundle.""" + if application_name != context.modal_application: + msg = "report coordinator application differs from dispatch context" + raise ValueError(msg) + if report.baseline.bundle.policyengine_version != context.worker_version: + msg = "report worker version differs from dispatch context" + raise ValueError(msg) + if report.baseline.bundle.bundle_manifest_sha256 != context.bundle_manifest_sha256: + msg = "report bundle digest differs from dispatch context" + raise ValueError(msg) + + +def validate_submitted_parent( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + parent: ComparisonReportRecord, +) -> None: + """Reject dispatch metadata that does not describe this exact report.""" + _validate_parent_matches_invocation(report=report, context=context, parent=parent) + _validate_parent_is_unclaimed(parent) + _validate_parent_production_link(report=report, context=context, parent=parent) + + +def _validate_parent_matches_invocation( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + parent: ComparisonReportRecord, +) -> None: + """Require every parent field the dispatch fixes to match the report and context.""" + bundle = report.baseline.bundle + expected = { + "evaluation_id": report.evaluation_id, + "environment": context.environment, + "originating_request_id": context.request_id, + "worker_version": context.worker_version, + "modal_application": context.modal_application, + "version_manifest_sha256": context.version_manifest_sha256, + "policyengine_version": bundle.policyengine_version, + "country_package_name": bundle.country_package_name, + "country_package_version": bundle.country_package_version, + "country": report.baseline.geography.country, + "dataset_identity": bundle.dataset.identity, + "dataset_uri": bundle.dataset.uri, + "data_package_name": bundle.dataset.data_package_name, + "data_package_version": bundle.dataset.data_package_version, + "data_artifact_revision": bundle.dataset.artifact_revision, + "created_at": context.created_at, + "retention_expires_at": context.retention_expires_at, + } + for field, value in expected.items(): + if getattr(parent, field) != value: + msg = f"submitted parent {field} does not match invocation" + raise ValueError(msg) + + +def _validate_parent_is_unclaimed(parent: ComparisonReportRecord) -> None: + """Require a pending parent that no coordinator has started.""" + if parent.status is not ComparisonRunLifecycleStatus.PENDING: + msg = "submitted parent must be pending" + raise ValueError(msg) + if parent.aggregation_status is not ComparisonRunAggregationStatus.NOT_STARTED: + msg = "submitted parent aggregation must not be started" + raise ValueError(msg) + if parent.coordinator_invocation_id is not None: + msg = "submitted parent must not contain an invocation identifier" + raise ValueError(msg) + + +def _validate_parent_production_link( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + parent: ComparisonReportRecord, +) -> None: + """Require a direct parent to stand alone and a shadow parent to name its job.""" + production_job_id = context.production_function_call_id + if production_job_id is None: + if parent.production_identity != f"direct:{report.evaluation_id}": + msg = "direct parent identity does not match its report" + raise ValueError(msg) + if parent.incumbent_execution_id is not None: + msg = "direct parent must not name a production execution" + raise ValueError(msg) + if parent.comparison_status is not ResultComparisonStatus.NOT_REQUESTED: + msg = "direct parent must not request result comparison" + raise ValueError(msg) + elif ( + parent.production_identity != production_job_id + or parent.incumbent_execution_id != production_job_id + ): + msg = "automatic parent does not match the production execution" + raise ValueError(msg) + + +def claim_parent( + *, + submitted: ComparisonReportRecord, + coordinator_invocation_id: str, + store: ComparisonStore, +) -> tuple[ComparisonReportRecord, bool]: + """Create the parent in Modal or claim a retry; reject duplicate execution.""" + coordinating_at = datetime.now(UTC) + candidate = submitted.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "aggregation_status": ComparisonRunAggregationStatus.RUNNING, + "coordinator_invocation_id": coordinator_invocation_id, + "error_code": None, + "error_summary": None, + "started_at": coordinating_at, + "updated_at": coordinating_at, + "completed_at": None, + } + ) + persistence = store.create_or_resolve_report(candidate) + parent = persistence.record + if persistence.created: + return parent, True + if parent.evaluation_id != submitted.evaluation_id: + msg = "existing comparison identity has a different report identifier" + raise ValueError(msg) + if parent.status in { + ComparisonRunLifecycleStatus.RUNNING, + ComparisonRunLifecycleStatus.SUCCEEDED, + ComparisonRunLifecycleStatus.SKIPPED, + }: + return parent, False + claimed = store.replace_report( + parent.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "aggregation_status": ComparisonRunAggregationStatus.RUNNING, + "coordinator_invocation_id": coordinator_invocation_id, + "error_code": None, + "error_summary": None, + "started_at": parent.started_at or coordinating_at, + "updated_at": coordinating_at, + "completed_at": None, + } + ) + ) + return claimed, True diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py new file mode 100644 index 000000000..15c7ab7de --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_coordination.py @@ -0,0 +1,626 @@ +"""One claimed report's coordination: start, wait for and aggregate its simulations.""" + +from __future__ import annotations + +import dataclasses +import time +from dataclasses import dataclass +from datetime import UTC, datetime +from functools import partial +from hashlib import sha256 +from typing import TYPE_CHECKING, Any +from uuid import uuid4 + +from policyengine_simulation_contract.stage12_execution import ( + AggregateReportArtifactDescriptor, + ComparisonReportRecord, + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, + ComparisonSimulationRecord, + PlannedSimulationExecutionInput, + ReportExecutionInput, + SimulationArtifactDescriptor, + SimulationRole, + Stage12InvocationContext, + Stage12OutputPlan, +) +from policyengine_simulation_observability.stages import Stage, StagePlan + +from policyengine_simulation_executor.stage12_artifacts import ( + Stage12ArtifactStore, + deserialize_calculation_provenance, + deserialize_simulation_frames, + deserialize_uk_local_authority_metadata, +) +from policyengine_simulation_executor.stage12_timeouts import ( + SIMULATION_WAIT_TIMEOUT_SECONDS, +) + +from .aggregation import ( + RetainedSimulationOutput, + build_aggregate_report, + validate_aligned_outputs, + validate_uk_local_authority_metadata, +) +from .comparison import ComparisonServices, compare_completed_report +from .coordinated_simulations import ( + cancel_pending_simulations, + cancellation, + child_failure_detail, + child_record, + record_child_failure, + simulation_failure_attributes, +) +from .dependencies import ChildCall, ChildInvoker, ChildSpawnRequest, ComparisonStore +from .failure_causes import ( + StageProgress, + StageTracking, + record_failure_cause, + tracked_stage, +) +from .failure_recording import ( + FailureSite, + persist_failure_record, + report_failure, + report_original_failure, +) +from .failures import Stage12ExecutionError, Stage12FailureDetail +from .output_planning import ( + plan_simulation_input, + resolve_report_output_plan, + validate_output_frames, +) +from .simulation import descriptor_from_record + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + from contextlib import AbstractContextManager + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + +SIMULATION_POLL_INITIAL_SECONDS = 0.25 +SIMULATION_POLL_MAX_SECONDS = 1.0 +SIMULATION_POLL_BACKOFF_FACTOR = 2.0 + + +def _next_poll_interval(current: float) -> float: + return min( + current * SIMULATION_POLL_BACKOFF_FACTOR, + SIMULATION_POLL_MAX_SECONDS, + ) + + +def _require_artifact_digest( + payload: bytes, + expected_sha256: str, + *, + mismatch_message: str, +) -> None: + """Raise ``ValueError(mismatch_message)`` unless ``payload`` matches its digest.""" + if sha256(payload).hexdigest() != expected_sha256: + raise ValueError(mismatch_message) + + +@dataclass(frozen=True) +class CoordinatorInvocation: + """The deployed coordinator call running a report: its Modal app and call ID.""" + + application_name: str + invocation_id: str + + +@dataclass(frozen=True) +class CoordinatorDependencies: + """Seams the coordinator uses; an unset seam selects the deployed default. + + The defaults are the Stage 12 database store, the private artifact + bucket, Modal child invocation, the standard aggregate builder and output + plan resolution, no telemetry runtime, the real clock and sleep, and the + standard child wait window. + """ + + store: ComparisonStore | None = None + artifacts: Stage12ArtifactStore | None = None + invoker: ChildInvoker | None = None + aggregator: Callable[..., dict[str, Any]] = build_aggregate_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 + + +@dataclass +class ReportCoordination: + """One claimed report's inputs and resolved seams, run through each phase.""" + + report: ReportExecutionInput + context: Stage12InvocationContext + invocation: CoordinatorInvocation + persistence: ComparisonStore + artifact_storage: Stage12ArtifactStore + child_invoker: ChildInvoker + dependencies: CoordinatorDependencies + stage_plan: StagePlan + # The last stage entered and whether it finished, for failure causes. + progress: StageProgress + # The shared child-wait window closes at this monotonic time. + child_wait_deadline: float + # Each child by role: its plan, its started call and its descriptor once + # finished, so a failure at any point can stop the unfinished children. + simulations: dict[str, PlannedSimulationExecutionInput] = dataclasses.field( + default_factory=dict + ) + calls: dict[str, ChildCall] = dataclasses.field(default_factory=dict) + descriptors: dict[str, SimulationArtifactDescriptor] = dataclasses.field( + default_factory=dict + ) + + @property + def runtime(self) -> ObservabilityRuntime | None: + """The telemetry runtime, or None when the coordinator records nothing.""" + return self.dependencies.runtime + + @property + def stage_tracking(self) -> StageTracking: + """How this report's stages are named, recorded and timed.""" + return StageTracking(self.stage_plan, self.progress, self.runtime) + + def span( + self, + stage: Stage, + attributes: Mapping[str, object] | None = None, + ) -> AbstractContextManager[object]: + """Track ``stage`` for failure causes and time it as a span.""" + return tracked_stage(stage, attributes, tracking=self.stage_tracking) + + def stage_scope( + self, + stage: Stage, + attributes: Mapping[str, object] | None = None, + ) -> AbstractContextManager[object]: + """Track ``stage`` for failure causes and record it as an operation.""" + return tracked_stage( + stage, attributes, tracking=self.stage_tracking, operation=True + ) + + def run(self, parent: ComparisonReportRecord) -> dict[str, Any]: + """Plan, start and await both children, then aggregate and complete.""" + output_plan, planned_simulations = self._prepare(parent) + self.simulations = { + simulation.role.value: simulation for simulation in planned_simulations + } + children = self._create_children(planned_simulations) + descriptors, calls = self._start_children(planned_simulations, children) + # Every required call has been started before the coordinator observes + # baseline and reform together through nonblocking result probes. + self._wait_for_children(calls, descriptors) + return self._aggregate_and_complete(parent, output_plan, descriptors) + + def _prepare( + self, parent: ComparisonReportRecord + ) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: + """Bind the run's artifact prefix, then plan the outputs and child inputs.""" + with self.stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): + self.context = self.context.model_copy( + update={ + "artifact_prefix": ( + f"stage-12-runs/{parent.environment}/" + f"{parent.created_at:%Y}/{parent.created_at:%m}/" + f"{parent.evaluation_id}" + ), + "created_at": parent.created_at, + "retention_expires_at": parent.retention_expires_at, + } + ) + with self.stage_scope( + Stage.STAGE12_OUTPUT_PLANNING, + {"country": self.report.baseline.geography.country}, + ): + resolver = self.dependencies.output_plan_resolver + if resolver is None: + output_plan = resolve_report_output_plan( + self.report, + stage_scope=self.stage_scope, + ) + else: + output_plan = resolver(self.report) + simulations = [] + for simulation in (self.report.baseline, self.report.reform): + with self.stage_scope( + Stage.STAGE12_CHILD_INPUT_PLANNING, + {"simulation_role": simulation.role.value}, + ): + simulations.append(plan_simulation_input(simulation, output_plan)) + return output_plan, tuple(simulations) + + def _create_children( + self, planned_simulations: tuple[PlannedSimulationExecutionInput, ...] + ) -> dict[SimulationRole, ComparisonSimulationRecord]: + """Create each child's row, or resolve the row an earlier attempt made.""" + with self.span(Stage.STAGE12_CHILD_STATE_CREATE): + return { + simulation.role: self.persistence.create_or_resolve_simulation( + child_record( + simulation=simulation, + context=self.context, + function_name=self.context.simulation_callable, + ) + ).record + for simulation in planned_simulations + } + + def _start_children( + self, + planned_simulations: tuple[PlannedSimulationExecutionInput, ...], + children: dict[SimulationRole, ComparisonSimulationRecord], + ) -> tuple[dict[str, SimulationArtifactDescriptor], dict[str, ChildCall]]: + """Reuse finished children, reattach running ones and dispatch the rest.""" + descriptors = self.descriptors + calls = self.calls + for simulation in planned_simulations: + child = children[simulation.role] + role = simulation.role.value + if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: + descriptors[role] = self._retained_descriptor(child, simulation) + elif ( + child.status is ComparisonRunLifecycleStatus.RUNNING + and child.modal_invocation_id is not None + and not child.modal_invocation_id.startswith("dispatch-pending-") + ): + calls[role] = self.child_invoker.restore(child.modal_invocation_id) + else: + calls[role] = self._dispatch_child(child, simulation) + return descriptors, calls + + def _retained_descriptor( + self, + child: ComparisonSimulationRecord, + simulation: PlannedSimulationExecutionInput, + ) -> SimulationArtifactDescriptor: + """Describe a succeeded child's digest-checked retained artifact.""" + descriptor = descriptor_from_record(child, simulation) + retained = self.artifact_storage.read(descriptor.artifact.uri) + _require_artifact_digest( + retained, + descriptor.artifact.content_sha256, + mismatch_message="retained child artifact digest mismatch", + ) + return descriptor.model_copy( + update={ + "artifact": descriptor.artifact.model_copy( + update={"size_bytes": len(retained)} + ), + "calculation_provenance": ( + deserialize_calculation_provenance(retained) + ), + } + ) + + def _dispatch_child( + self, + child: ComparisonSimulationRecord, + simulation: PlannedSimulationExecutionInput, + ) -> ChildCall: + """Mark the child dispatching, spawn it, and attach its Modal call ID. + + A spawn failure is reported and recorded on the child (unless it already + finished), every started peer is stopped, and Stage12ExecutionError is + raised with the failure. + """ + dispatching_at = datetime.now(UTC) + dispatch_placeholder = f"dispatch-pending-{uuid4()}" + self.persistence.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": dispatch_placeholder, + "started_at": child.started_at or dispatching_at, + "updated_at": dispatching_at, + "error_code": None, + "error_summary": None, + "completed_at": None, + } + ) + ) + try: + with self.span( + Stage.STAGE12_CHILD_DISPATCH, + {"simulation_role": simulation.role.value}, + ): + child_observability_context = ( + self.runtime.capture_context() if self.runtime is not None else None + ) + call = self.child_invoker.spawn( + ChildSpawnRequest( + application_name=self.invocation.application_name, + function_name=self.context.simulation_callable, + environment=self.context.modal_environment, + simulation=simulation.model_dump(mode="json"), + context=self.context.model_dump(mode="json"), + observability_context=child_observability_context, + ) + ) + except Exception as error: # noqa: BLE001 # any dispatch failure is recorded on the child and stops its started peers + detail = report_failure( + error, + FailureSite( + scope="child_dispatch", + default_code="simulation_dispatch_failed", + attributes=simulation_failure_attributes(simulation), + ), + runtime=self.runtime, + progress=self.progress, + ) + record_child_failure( + store=self.persistence, + simulation=simulation, + detail=detail, + ) + cancel_pending_simulations( + calls=self.calls, + simulations=self.simulations, + roles=set(self.calls), + store=self.persistence, + detail=cancellation(f"the {simulation.role.value} simulation failed"), + ) + raise Stage12ExecutionError(detail) from None + # Record the call before attaching its ID, so a failed attach still + # stops it. + self.calls[simulation.role.value] = call + self.persistence.attach_simulation_invocation( + simulation.simulation_execution_id, + expected_placeholder=dispatch_placeholder, + modal_invocation_id=call.object_id, + updated_at=datetime.now(UTC), + ) + return call + + def _wait_for_children( + self, + calls: dict[str, ChildCall], + descriptors: dict[str, SimulationArtifactDescriptor], + ) -> None: + """Observe every started child together until each has a descriptor. + + A child's failure, or the shared wait window closing, stops the + unfinished children and raises Stage12ExecutionError. + """ + with self.span(Stage.STAGE12_CHILD_WAIT): + self._observe_children(calls, descriptors) + + def _observe_children( + self, + calls: dict[str, ChildCall], + descriptors: dict[str, SimulationArtifactDescriptor], + ) -> None: + """Probe every unfinished child, without blocking, until each finishes. + + Between passes the coordinator sleeps, backing off while no child + finishes, and never past the shared wait window. + """ + wait_stage_name = self.stage_plan.name(Stage.STAGE12_CHILD_WAIT) + timeout_seconds = self.dependencies.wait_timeout_seconds + pending = set(calls) + poll_interval = SIMULATION_POLL_INITIAL_SECONDS + + while pending: + progress_made = False + for role in sorted(pending): + call = calls[role] + try: + # The probe records the child's role, so a failure names it. + with self.progress.stage(wait_stage_name, role=role): + result = call.get(timeout=0) + descriptor = SimulationArtifactDescriptor.model_validate(result) + except TimeoutError: + continue + except Exception as error: # noqa: BLE001 # a child call's failure can be any exception type; it is recorded and stops the other child + detail = child_failure_detail( + error=error, + simulation=self.simulations[role], + store=self.persistence, + runtime=self.runtime, + progress=self.progress, + ) + cancel_pending_simulations( + calls=calls, + simulations=self.simulations, + roles=pending - {role}, + store=self.persistence, + detail=cancellation(f"the {role} simulation failed"), + ) + raise Stage12ExecutionError(detail) from None + descriptors[role] = descriptor + pending.remove(role) + progress_made = True + + if not pending: + return + remaining_seconds = self.child_wait_deadline - self.dependencies.monotonic() + if remaining_seconds <= 0: + detail = Stage12FailureDetail( + error_code="simulation_wait_timeout", + error_summary=( + "Stage 12 simulations did not finish within " + f"{timeout_seconds:g} seconds of the report coordinator " + "starting." + ), + ) + wait_timeout = Stage12ExecutionError(detail) + report_original_failure( + wait_timeout, + runtime=self.runtime, + failure_scope="child_wait_timeout", + ) + cancel_pending_simulations( + calls=calls, + simulations=self.simulations, + roles=pending, + store=self.persistence, + detail=cancellation( + f"the {timeout_seconds:g}-second wait window closed" + ), + ) + raise wait_timeout + # Never sleep past the deadline, so the window closes on time. + self.dependencies.sleep(min(poll_interval, remaining_seconds)) + poll_interval = ( + SIMULATION_POLL_INITIAL_SECONDS + if progress_made + else _next_poll_interval(poll_interval) + ) + + def _read_retained_outputs( + self, + output_plan: Stage12OutputPlan, + baseline: SimulationArtifactDescriptor, + reform: SimulationArtifactDescriptor, + ) -> tuple[ + RetainedSimulationOutput, + RetainedSimulationOutput, + UKLocalAuthorityMetadata | None, + ]: + """Read, digest-check and validate both children's retained artifacts.""" + validate_aligned_outputs(self.report, output_plan, baseline, reform) + with self.span(Stage.STAGE12_ARTIFACT_READ): + baseline_payload = self.artifact_storage.read(baseline.artifact.uri) + reform_payload = self.artifact_storage.read(reform.artifact.uri) + _require_artifact_digest( + baseline_payload, + baseline.artifact.content_sha256, + mismatch_message="baseline artifact digest mismatch", + ) + _require_artifact_digest( + reform_payload, + reform.artifact.content_sha256, + mismatch_message="reform artifact digest mismatch", + ) + baseline_frames = deserialize_simulation_frames(baseline_payload) + reform_frames = deserialize_simulation_frames(reform_payload) + uk_local_authority_metadata = validate_uk_local_authority_metadata( + self.report.baseline.geography.country, + deserialize_uk_local_authority_metadata(baseline_payload), + deserialize_uk_local_authority_metadata(reform_payload), + region_code=self.report.baseline.geography.region, + ) + validate_output_frames(baseline_frames, output_plan) + validate_output_frames(reform_frames, output_plan) + return ( + RetainedSimulationOutput(baseline_frames, baseline), + RetainedSimulationOutput(reform_frames, reform), + uk_local_authority_metadata, + ) + + def _aggregate_and_complete( + self, + parent: ComparisonReportRecord, + output_plan: Stage12OutputPlan, + descriptors: dict[str, SimulationArtifactDescriptor], + ) -> dict[str, Any]: + """Aggregate both outputs, store the aggregate and complete the parent. + + A shadow run's result is then compared with the production result. + """ + baseline, reform, uk_local_authority_metadata = self._read_retained_outputs( + output_plan, descriptors["baseline"], descriptors["reform"] + ) + with self.span(Stage.STAGE12_AGGREGATION): + aggregate = self.dependencies.aggregator( + report=self.report, + baseline=baseline, + reform=reform, + uk_local_authority_metadata=uk_local_authority_metadata, + ) + with self.span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): + aggregate_artifact = self.artifact_storage.write_aggregate( + prefix=self.context.artifact_prefix, + payload=aggregate, + ) + descriptor = AggregateReportArtifactDescriptor( + evaluation_id=self.report.evaluation_id, + artifact=aggregate_artifact, + baseline_artifact_sha256=baseline.descriptor.artifact.content_sha256, + reform_artifact_sha256=reform.descriptor.artifact.content_sha256, + bundle=self.report.baseline.bundle, + ) + completed = datetime.now(UTC) + completed_parent = self.persistence.replace_report( + parent.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "aggregation_status": ComparisonRunAggregationStatus.SUCCEEDED, + "aggregate_output_uri": aggregate_artifact.uri, + "aggregate_output_sha256": aggregate_artifact.content_sha256, + "aggregate_schema_version": descriptor.aggregate_schema_version, + "updated_at": completed, + "completed_at": completed, + } + ) + ) + compare_completed_report( + parent=completed_parent, + aggregate=aggregate, + context=self.context, + services=ComparisonServices( + store=self.persistence, + artifacts=self.artifact_storage, + invoker=self.child_invoker, + runtime=self.runtime, + stage_plan=self.stage_plan, + ), + ) + return descriptor.model_dump(mode="json") + + def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: + """Record a failed run once, stop its unfinished children, fail the parent. + + Child, dispatch and wait-timeout failures raise Stage12ExecutionError + only after reporting the original failure and stopping their unfinished + peers. Any other failure is reported here, once, before any database + write. If it occurred after a child started, such as a failed + placeholder write or invocation attach, every started, unfinished child + is then stopped so none keeps running for a failed report. + """ + if isinstance(error, Stage12ExecutionError): + detail = error.detail + else: + detail = report_failure( + error, + FailureSite( + scope="report_coordination", + default_code="report_coordination_failed", + attributes={"evaluation_id": str(self.report.evaluation_id)}, + ), + runtime=self.runtime, + progress=self.progress, + ) + cancel_pending_simulations( + calls=self.calls, + simulations=self.simulations, + roles=set(self.calls) - set(self.descriptors), + store=self.persistence, + detail=cancellation("the coordinator failed"), + ) + record_failure_cause(self.runtime, error, detail.error_summary) + failed_at = datetime.now(UTC) + persist_failure_record( + partial( + self.persistence.replace_report, + parent.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "aggregation_status": ComparisonRunAggregationStatus.FAILED, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": failed_at, + "completed_at": failed_at, + } + ), + ), + runtime=self.runtime, + failure_scope="report_coordination", + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index 933f5a435..3f547dcb9 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -16,7 +16,10 @@ from policyengine_simulation_executor.stage12_runtime import ( CoordinatorDependencies, CoordinatorInvocation, + coordinated_simulations, coordination, + parent_claim, + report_coordination, ) from test_stage12_runtime import ( BASELINE_ID, @@ -64,8 +67,11 @@ def now(cls, tz=None): @pytest.fixture def fixed_clock(monkeypatch): - monkeypatch.setattr(coordination, "datetime", FixedClock) - monkeypatch.setattr(coordination, "uuid4", lambda: PLACEHOLDER_ID) + # Every module that reads the clock while it coordinates a report. + monkeypatch.setattr(coordinated_simulations, "datetime", FixedClock) + monkeypatch.setattr(parent_claim, "datetime", FixedClock) + monkeypatch.setattr(report_coordination, "datetime", FixedClock) + monkeypatch.setattr(report_coordination, "uuid4", lambda: PLACEHOLDER_ID) class TracingRuntime: @@ -414,7 +420,7 @@ def resolve(report, *, stage_scope): resolved.append(report.evaluation_id) return _output_plan() - monkeypatch.setattr(coordination, "resolve_report_output_plan", resolve) + monkeypatch.setattr(report_coordination, "resolve_report_output_plan", resolve) coordinate_report_impl( _report().model_dump(mode="json"), diff --git a/rules/baselines/python-file-length.txt b/rules/baselines/python-file-length.txt index 506efdcb9..a93a18dba 100644 --- a/rules/baselines/python-file-length.txt +++ b/rules/baselines/python-file-length.txt @@ -4,13 +4,11 @@ # Python file that was already over its line limit (800 lines, or 1,500 for # test code) when the rule was added. Remove an entry, and lower the rule's # baseline count, in the pull request that splits the file. -libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py projects/policyengine-simulation-entry/tests/test_deployment_assets.py projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py projects/policyengine-simulation-executor/tests/test_precompute.py projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py projects/policyengine-simulation-executor/tests/test_stage12_runtime.py diff --git a/rules/registry.toml b/rules/registry.toml index dc8dc8849..8add82684 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1190,7 +1190,7 @@ job = "quality-guards" invocation = "run_quality_guards.py" [rule.baseline] file = "rules/baselines/python-file-length.txt" -count = 10 +count = 8 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.python-file-length/bad_files.toml" good_fixture = "rules/fixtures/sim.quality.python-file-length/good_files.toml"