diff --git a/changelog.d/report-run-recording.added.md b/changelog.d/report-run-recording.added.md new file mode 100644 index 000000000..71d024e82 --- /dev/null +++ b/changelog.d/report-run-recording.added.md @@ -0,0 +1 @@ +The v2 Modal report coordinator and its US and UK simulation workers can record a report's run in the Public API's run tables: when a call names a report run and its baseline and reform simulation runs, they write those runs' status, output file and economy results with the Modal worker role's credentials, from the Modal secret `policyengine-modal-worker-database`. The report run is marked succeeded in the same transaction as its result rows. On pre-emption the coordinator of such a call cancels its child calls and resumes when Modal restarts it. Comparison runs behave as before. diff --git a/docs/engineering/skills/database-access.md b/docs/engineering/skills/database-access.md index 1ed83a886..292f29d4a 100644 --- a/docs/engineering/skills/database-access.md +++ b/docs/engineering/skills/database-access.md @@ -34,6 +34,11 @@ It builds on the Stage 12 rules: (`sim.stage12.no-migration-chain`). This repository's code reads and writes their rows, and deployment validation compares their live schema with its SQLAlchemy mappings; nothing here changes them. +- policyengine-api also owns the run tables the v2 Modal coordinator and + workers write as the Modal worker role: `report_runs`, `simulation_runs`, + `report_run_simulation_runs`, `report_run_comparisons` and the result + tables. Their mappings in `policyengine_stage12_persistence.report_run_tables` + and `report_run_results` declare only the columns those writes touch. - **Once #847 merges**, this repository's own Alembic chain at `migrations/` owns the `simulation` schema, and the Stage 12 tables stay with policyengine-api. 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 index d1a1f477f..6006c63ea 100644 --- 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 @@ -38,6 +38,7 @@ ComparisonSimulationRecord, ErrorCode, ErrorSummary, + ReportRunIds, ResultComparisonStatus, Stage12InvocationContext, ) @@ -95,6 +96,7 @@ "ReportAggregate", "ReportExecutionInput", "ReportOutputRequirements", + "ReportRunIds", "ResultComparisonArtifactPayload", "ResultComparisonStatus", "ResultDifference", 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 index ea49b5e2d..31682cdd6 100644 --- 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 @@ -23,13 +23,49 @@ ) +class ReportRunIds(StrictContractModel): + """The Public API's report run and simulation runs one report call records on. + + The simulation service creates the three rows before it starts the + coordinator. The coordinator records the report run's status and results, + and each worker records its own simulation run's. + """ + + report_run_id: UUID + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID + + @model_validator(mode="after") + def validate_distinct_runs(self) -> ReportRunIds: + """Require the report run and both simulation runs to be different rows.""" + run_ids = { + self.report_run_id, + self.baseline_simulation_run_id, + self.reform_simulation_run_id, + } + if len(run_ids) != 3: # noqa: PLR2004 # one report run and two simulation runs + msg = "report run IDs must name three different runs" + raise ValueError(msg) + return self + + def simulation_run_id(self, role: SimulationRole) -> UUID: + """Return the simulation run of ``role``; a standalone role has none.""" + if role is SimulationRole.BASELINE: + return self.baseline_simulation_run_id + if role is SimulationRole.REFORM: + return self.reform_simulation_run_id + msg = "a standalone simulation has no simulation run in a report run" + raise ValueError(msg) + + class Stage12InvocationContext(StrictContractModel): """Deployment and retention context for one comparison-run invocation. Records the originating request, the environments, the worker version and callable, the manifest digests and the artifact prefix. Automatic runs also name the production function call they are compared with. Artifacts are kept - for at most 30 days after ``created_at``. + for at most 30 days after ``created_at``. A call that also records its + outcome on the Public API's run tables names those runs in ``report_run``. """ contract_version: Literal[1] = 1 @@ -51,6 +87,8 @@ class Stage12InvocationContext(StrictContractModel): production_function_call_id: ContractText | None = None created_at: datetime retention_expires_at: datetime + # Comparison runs leave it unset and record only the Stage 12 tables. + report_run: ReportRunIds | None = None @model_validator(mode="after") def validate_retention(self) -> Stage12InvocationContext: @@ -58,6 +96,14 @@ def validate_retention(self) -> Stage12InvocationContext: _validate_retention(self.created_at, self.retention_expires_at) return self + @model_validator(mode="after") + def validate_report_run_is_direct(self) -> Stage12InvocationContext: + """Reject report runs on a shadow run compared with a production call.""" + if self.report_run is not None and self.production_function_call_id: + msg = "a run compared with a production function call has no report run" + raise ValueError(msg) + return self + ErrorCode = Annotated[str, Field(min_length=1, max_length=64)] ErrorSummary = Annotated[str, Field(min_length=1, max_length=512)] diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py b/libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py new file mode 100644 index 000000000..3596581dd --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_stage12_report_run_ids.py @@ -0,0 +1,96 @@ +"""The report run and simulation runs a Stage 12 call records its outcome on.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Any +from uuid import UUID + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ReportRunIds, + SimulationRole, + Stage12InvocationContext, +) +from pydantic import ValidationError + +NOW = datetime(2026, 11, 3, tzinfo=UTC) +REPORT_RUN_ID = UUID("00000000-0000-0000-0000-0000000000a1") +BASELINE_RUN_ID = UUID("00000000-0000-0000-0000-0000000000b1") +REFORM_RUN_ID = UUID("00000000-0000-0000-0000-0000000000c1") + + +def _context_fields(**overrides: Any) -> dict[str, Any]: + return { + "request_id": "request-1", + "environment": "staging", + "modal_environment": "staging", + "worker_version": "6.12.0", + "modal_application": "policyengine-simulation-v2-py6-12-0", + "simulation_callable": "run_single_simulation_us", + "version_manifest_sha256": "a" * 64, + "bundle_manifest_sha256": "b" * 64, + "artifact_prefix": "stage-12-runs/staging/2026/11/report", + "created_at": NOW, + "retention_expires_at": NOW + timedelta(days=30), + **overrides, + } + + +def _run_ids() -> ReportRunIds: + return ReportRunIds( + report_run_id=REPORT_RUN_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_run_id=REFORM_RUN_ID, + ) + + +def test_context_without_run_ids_keeps_its_comparison_fields() -> None: + context = Stage12InvocationContext(**_context_fields()) + + assert context.report_run is None + assert "report_run" not in context.model_dump(mode="json", exclude_none=True) + + +def test_context_carries_the_report_run_and_both_simulation_runs() -> None: + context = Stage12InvocationContext( + **_context_fields(report_run=_run_ids().model_dump(mode="json")) + ) + + assert context.report_run == _run_ids() + restored = Stage12InvocationContext.model_validate(context.model_dump(mode="json")) + assert restored.report_run == _run_ids() + + +def test_each_role_names_its_own_simulation_run() -> None: + run_ids = _run_ids() + + assert run_ids.simulation_run_id(SimulationRole.BASELINE) == BASELINE_RUN_ID + assert run_ids.simulation_run_id(SimulationRole.REFORM) == REFORM_RUN_ID + with pytest.raises(ValueError, match="standalone"): + run_ids.simulation_run_id(SimulationRole.STANDALONE) + + +@pytest.mark.parametrize( + "repeated_ids", + [ + {"baseline_simulation_run_id": REPORT_RUN_ID}, + {"reform_simulation_run_id": REPORT_RUN_ID}, + {"reform_simulation_run_id": BASELINE_RUN_ID}, + ], +) +def test_run_ids_must_name_three_different_runs(repeated_ids) -> None: + fields = {**_run_ids().model_dump(), **repeated_ids} + + with pytest.raises(ValidationError, match="three different runs"): + ReportRunIds(**fields) + + +def test_a_shadow_comparison_run_cannot_record_report_runs() -> None: + with pytest.raises(ValidationError, match="production function call"): + Stage12InvocationContext( + **_context_fields( + production_function_call_id="fc-production", + report_run=_run_ids().model_dump(mode="json"), + ) + ) diff --git a/libs/policyengine-stage12-persistence/README.md b/libs/policyengine-stage12-persistence/README.md index cca735900..b1caeb034 100644 --- a/libs/policyengine-stage12-persistence/README.md +++ b/libs/policyengine-stage12-persistence/README.md @@ -1,6 +1,26 @@ # PolicyEngine Stage 12 persistence -Temporary SQLAlchemy Core mappings and transactional persistence for the two -Stage 12 comparison tables. `policyengine-api` remains the sole owner of the -SQLModel definitions, Alembic revisions, and live schema. This package performs -DML only and is removed with the temporary comparison path during Stage 14. +SQLAlchemy Core mappings and transactional DML for two sets of tables, both +owned by `policyengine-api`, which keeps the SQLModel definitions, Alembic +revisions and live schema. This package never changes a schema. + +- **The Stage 12 comparison tables** (`Stage12PersistenceStore`): temporary + state of the comparison path, removed with it. +- **The Public API's run tables** (`ReportRunStore`): `report_runs`, + `simulation_runs`, `report_run_simulation_runs`, the result tables, and + `report_run_comparisons`, which keeps a compared pair's results that no + result table holds. The v2 Modal coordinator and workers write them as the + Modal worker role (`policyengine_v2_modal_worker`), for a call that names a + report run. The mappings declare only the columns those writes touch. Every + status write is one `UPDATE` conditional on the run's status; a report + run's completion inserts its result rows in the same transaction. + +## Tests + +The report run tests need a disposable Postgres migrated by the Public API's +v2 chain, with its declared privileges applied. They skip unless both URLs +are set: + +- `REPORT_RUN_TEST_DATABASE_URL`: logs in as `policyengine_v2_modal_worker`; +- `REPORT_RUN_TEST_SEED_DATABASE_URL`: logs in as a role that may insert the + rows the simulation service would create. diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index e1344e788..669f4ceb5 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -1,5 +1,15 @@ -"""Direct DML for temporary Stage 12 comparison state.""" +"""Direct DML for temporary Stage 12 comparison state and for report runs.""" +from policyengine_stage12_persistence.report_run_results import ( + DistrictCodes, + ReportRunResult, + report_run_result, +) +from policyengine_stage12_persistence.report_run_store import ( + ReportRunStore, + RunStatus, + RunWrite, +) from policyengine_stage12_persistence.store import ( ConnectionOpener, Stage12PersistenceStore, @@ -13,9 +23,15 @@ __all__ = [ "ConnectionOpener", + "DistrictCodes", + "ReportRunResult", + "ReportRunStore", + "RunStatus", + "RunWrite", "Stage12PersistenceStore", "Stage12RuntimeAccessError", "create_stage12_engine", "normalized_database_url", + "report_run_result", "validate_runtime_database", ] diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py new file mode 100644 index 000000000..5e6e86052 --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_results.py @@ -0,0 +1,611 @@ +"""A report run's economy output as rows of the Public API's result tables. + +The output is the single-year economy result Stage 12's coordinator builds, +in the format API v1 and v3 read today. Each part with a lossless place in a +result table becomes rows: per-report parts keyed by the report run and its +two simulation runs, poverty and inequality keyed by each simulation run. The +labour supply response, cliff impact and SPM receipts, which no result table +holds, are stored whole with the data version, in one ``report_run_comparisons`` +row for the report run and its two simulation runs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, NamedTuple +from uuid import uuid4 + +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from sqlalchemy import JSON, Boolean, Column, Float, Integer, String, Table +from sqlalchemy.dialects.postgresql import ENUM, UUID + +from policyengine_stage12_persistence.report_run_tables import metadata + +if TYPE_CHECKING: + from collections.abc import Mapping + + from policyengine_simulation_contract.macro_output import ( + AgePovertyOutput, + BaselineReformValue, + CongressionalDistrictImpactRecord, + ConstituencyImpactRecord, + DecileOutput, + GenderPovertyOutput, + GeographicImpactOutput, + IntraDecileOutput, + LocalAuthorityImpactRecord, + RacePovertyOutput, + ) + from policyengine_simulation_contract.stage12_execution import ReportRunIds + +decile_type_type = ENUM("income", "wealth", name="v2_decile_type", create_type=False) + + +def _report_run_keys() -> tuple[Column[Any], ...]: + return ( + Column("id", UUID(as_uuid=True), primary_key=True), + Column("report_run_id", UUID(as_uuid=True), nullable=False), + Column("baseline_simulation_run_id", UUID(as_uuid=True), nullable=False), + Column("reform_simulation_run_id", UUID(as_uuid=True), nullable=False), + ) + + +def _simulation_run_keys() -> tuple[Column[Any], ...]: + return ( + Column("id", UUID(as_uuid=True), primary_key=True), + Column("simulation_run_id", UUID(as_uuid=True), nullable=False), + ) + + +report_run_comparisons = Table( + "report_run_comparisons", + metadata, + *_report_run_keys(), + Column("labor_supply_result", JSON(none_as_null=True)), + Column("cliff_impact_result", JSON(none_as_null=True)), + Column("spm_provenance", JSON(none_as_null=True)), + Column("data_version", String(128)), +) +program_statistics = Table( + "program_statistics", + metadata, + *_report_run_keys(), + Column("program_name", String(512), nullable=False), + Column("entity", String(128)), + Column("is_tax", Boolean), + Column("baseline_total", Float), + Column("reform_total", Float), + Column("change", Float), +) + + +def _area_values() -> tuple[Column[Any], ...]: + return ( + Column("average_household_income_change", Float, nullable=False), + Column("relative_household_income_change", Float, nullable=False), + Column("population", Float, nullable=False), + ) + + +def _disclosure_values() -> tuple[Column[Any], ...]: + return ( + Column("disclosure_tier", String(16)), + Column("average_household_income_change_lower", Float), + Column("average_household_income_change_upper", Float), + Column("relative_household_income_change_lower", Float), + Column("relative_household_income_change_upper", Float), + ) + + +congressional_district_impacts = Table( + "congressional_district_impacts", + metadata, + *_report_run_keys(), + *_area_values(), + Column("district_geoid", Integer, nullable=False), + Column("state_fips", Integer, nullable=False), + Column("district_number", Integer, nullable=False), + Column("winner_percentage", Float, nullable=False), + Column("loser_percentage", Float, nullable=False), + Column("no_change_percentage", Float, nullable=False), +) +constituency_impacts = Table( + "constituency_impacts", + metadata, + *_report_run_keys(), + *_area_values(), + *_disclosure_values(), + Column("constituency_code", String(64), nullable=False), + Column("constituency_name", String(255), nullable=False), +) +local_authority_impacts = Table( + "local_authority_impacts", + metadata, + *_report_run_keys(), + *_area_values(), + *_disclosure_values(), + Column("local_authority_code", String(64), nullable=False), + Column("local_authority_name", String(255), nullable=False), +) +budget_summary = Table( + "budget_summary", + metadata, + *_report_run_keys(), + Column("variable_name", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("baseline_total", Float), + Column("reform_total", Float), + Column("change", Float), +) +decile_impacts = Table( + "decile_impacts", + metadata, + *_report_run_keys(), + Column("income_variable", String(512), nullable=False), + Column("decile_variable", String(512)), + Column("entity", String(128), nullable=False), + Column("decile", Integer, nullable=False), + Column("quantiles", Integer, nullable=False), + Column("absolute_change", Float), + Column("relative_change", Float), +) +intra_decile_impacts = Table( + "intra_decile_impacts", + metadata, + *_report_run_keys(), + Column("decile_type", decile_type_type, nullable=False), + Column("decile", Integer, nullable=False), + Column("lose_more_than_5pct", Float), + Column("lose_less_than_5pct", Float), + Column("no_change", Float), + Column("gain_less_than_5pct", Float), + Column("gain_more_than_5pct", Float), +) +poverty = Table( + "poverty", + metadata, + *_simulation_run_keys(), + Column("poverty_type", String(128), nullable=False), + Column("entity", String(128), nullable=False), + Column("filter_variable", String(512)), + Column("filter_group", String(64)), + Column("rate", Float), +) +inequality = Table( + "inequality", + metadata, + *_simulation_run_keys(), + Column("income_variable", String(512), nullable=False), + Column("entity", String(128), nullable=False), + Column("gini", Float), + Column("top_10_share", Float), + Column("top_1_share", Float), +) + +# The household variable each headline budget figure totals, and the column +# its figure goes in. ``budgetary_impact`` is tax minus benefits, as the +# economy output computes it, so it needs no row of its own. +BUDGET_FIGURES = ( + ("tax_revenue_impact", "household_tax", "change"), + ("benefit_spending_impact", "household_benefits", "change"), + ("state_tax_revenue_impact", "household_state_income_tax", "change"), + ("baseline_net_income", "household_net_income", "baseline_total"), + ("households", "household_weight", "baseline_total"), +) +# .py's poverty types behind the output's "poverty" and "deep_poverty". +POVERTY_TYPES = { + "us": {"poverty": "spm", "deep_poverty": "spm_deep"}, + "uk": {"poverty": "relative_bhc", "deep_poverty": "absolute_bhc"}, +} +# .py's poverty groups (``policyengine.outputs.poverty``), by the variable +# that selects them, named as the output's parts name them. +POVERTY_GROUPS = { + "age": ("child", "adult", "senior"), + "is_male": ("male", "female"), + "race": ("white", "black", "hispanic", "other"), +} +# The variable .py groups UK wealth deciles' household net income by. +WEALTH_DECILE_VARIABLE = "household_wealth_decile" +INEQUALITY_INCOME_VARIABLES = { + "us": "household_net_income", + "uk": "equiv_hbai_household_net_income", +} +INTRA_DECILE_COLUMNS = { + "Lose more than 5%": "lose_more_than_5pct", + "Lose less than 5%": "lose_less_than_5pct", + "No change": "no_change", + "Gain less than 5%": "gain_less_than_5pct", + "Gain more than 5%": "gain_more_than_5pct", +} +DISCLOSURE_FIELDS = ( + "disclosure_tier", + "average_household_income_change_lower", + "average_household_income_change_upper", + "relative_household_income_change_lower", + "relative_household_income_change_upper", +) + + +@dataclass(frozen=True) +class DistrictCodes: + """US states' FIPS codes, and the states with one at-large district. + + Today's output names a district as ``"ST-NN"``, numbering an at-large + district ``01``. Its row stores the Census GEOID parts instead: the + state's FIPS code, and the district number, ``0`` for an at-large state. + """ + + state_fips: Mapping[str, int] + at_large_states: frozenset[str] + + +@dataclass(frozen=True) +class ResultRows: + """The rows one result table receives. + + Rows ``shared_by_reports`` belong to a simulation run, which other report + runs may share; a row such a run already holds is kept as stored. + """ + + table: Table + rows: tuple[dict[str, Any], ...] + shared_by_reports: bool + + +@dataclass(frozen=True) +class ReportRunResult: + """Everything a report run's economy output is stored as.""" + + rows: tuple[ResultRows, ...] + + +def report_run_result( + output: Mapping[str, Any], + *, + run_ids: ReportRunIds, + country: str, + district_codes: DistrictCodes | None = None, +) -> ReportRunResult: + """Validate a single-year economy output and return the rows it is stored as. + + Raises ValueError for an output that is not a single-year economy result, + a country without poverty types, or a congressional district whose state + ``district_codes`` doesn't name. + """ + economy = SingleYearMacroOutput.model_validate(output) + if country not in POVERTY_TYPES: + msg = f"no poverty types are known for country {country!r}" + raise ValueError(msg) + report_keys = { + "report_run_id": run_ids.report_run_id, + "baseline_simulation_run_id": run_ids.baseline_simulation_run_id, + "reform_simulation_run_id": run_ids.reform_simulation_run_id, + } + dumped = economy.model_dump(mode="json") + comparison = { + "id": uuid4(), + **report_keys, + "labor_supply_result": dumped["labor_supply_response"], + "cliff_impact_result": dumped["cliff_impact"], + "spm_provenance": dumped["spm_provenance"], + "data_version": economy.data_version, + } + return ReportRunResult( + rows=( + ResultRows(report_run_comparisons, (comparison,), shared_by_reports=False), + ResultRows( + budget_summary, + _budget_rows(economy, report_keys, country), + shared_by_reports=False, + ), + ResultRows( + decile_impacts, + _decile_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + intra_decile_impacts, + _intra_decile_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + poverty, + _poverty_rows(economy, run_ids, country), + shared_by_reports=True, + ), + ResultRows( + inequality, + _inequality_rows(economy, run_ids, country), + shared_by_reports=True, + ), + ResultRows( + program_statistics, + _program_rows(economy, report_keys), + shared_by_reports=False, + ), + ResultRows( + congressional_district_impacts, + _district_rows(economy, report_keys, district_codes), + shared_by_reports=False, + ), + ResultRows( + constituency_impacts, + _area_rows(economy.constituency_impact, report_keys, "constituency"), + shared_by_reports=False, + ), + ResultRows( + local_authority_impacts, + _area_rows( + economy.local_authority_impact, report_keys, "local_authority" + ), + shared_by_reports=False, + ), + ), + ) + + +def _program_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + # The output names no program's entity, so every program's is null. + return tuple( + { + "id": uuid4(), + **keys, + "program_name": program_name, + "entity": None, + "is_tax": program.is_tax, + "baseline_total": program.baseline, + "reform_total": program.reform, + "change": program.difference, + } + for program_name, program in economy.detailed_budget.root.items() + ) + + +def _district_rows( + economy: SingleYearMacroOutput, + keys: Mapping[str, Any], + district_codes: DistrictCodes | None, +) -> tuple[dict[str, Any], ...]: + if economy.congressional_district_impact is None: + return () + districts = economy.congressional_district_impact.districts + if districts and district_codes is None: + msg = "congressional districts need the US states' FIPS codes" + raise ValueError(msg) + return tuple( + _district_row(district, keys, district_codes) + for district in districts + if district_codes is not None + ) + + +def _district_row( + district: CongressionalDistrictImpactRecord, + keys: Mapping[str, Any], + district_codes: DistrictCodes, +) -> dict[str, Any]: + state, _, number = district.district.partition("-") + if state not in district_codes.state_fips or not number.isdigit(): + msg = f"congressional district {district.district!r} names no known state" + raise ValueError(msg) + state_fips = district_codes.state_fips[state] + district_number = 0 if state in district_codes.at_large_states else int(number) + return { + "id": uuid4(), + **keys, + "district_geoid": state_fips * 100 + district_number, + "state_fips": state_fips, + "district_number": district_number, + "average_household_income_change": district.average_household_income_change, + "relative_household_income_change": district.relative_household_income_change, + "winner_percentage": district.winner_percentage, + "loser_percentage": district.loser_percentage, + "no_change_percentage": district.no_change_percentage, + "population": district.population, + } + + +def _area_rows( + areas: GeographicImpactOutput | None, + keys: Mapping[str, Any], + kind: Literal["constituency", "local_authority"], +) -> tuple[dict[str, Any], ...]: + if areas is None: + return () + records: list[ConstituencyImpactRecord | LocalAuthorityImpactRecord] = areas.root + return tuple( + { + "id": uuid4(), + **keys, + f"{kind}_code": getattr(record, f"{kind}_code"), + f"{kind}_name": getattr(record, f"{kind}_name"), + "average_household_income_change": record.average_household_income_change, + "relative_household_income_change": ( + record.relative_household_income_change + ), + "population": record.population, + **{field: getattr(record, field) for field in DISCLOSURE_FIELDS}, + } + for record in records + ) + + +def _budget_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any], country: str +) -> tuple[dict[str, Any], ...]: + """Return the headline budget figures; state income tax only for the US. + + Today's output reports 0.0 state income tax for a country without it, so a + reader without that row reports 0.0 again. + """ + return tuple( + { + "id": uuid4(), + **keys, + "variable_name": variable, + "entity": "household", + column: getattr(economy.budget, figure), + } + for figure, variable, column in BUDGET_FIGURES + if variable != "household_state_income_tax" or country == "us" + ) + + +def _decile_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + rows = _decile_type_rows(economy.decile, None, keys) + if economy.wealth_decile is not None: + rows += _decile_type_rows(economy.wealth_decile, WEALTH_DECILE_VARIABLE, keys) + return rows + + +def _decile_type_rows( + deciles: DecileOutput, decile_variable: str | None, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + """Return the change in household net income in each decile. + + Income deciles rank households by that income, so they name no separate + ``decile_variable``; wealth deciles name the variable that ranks them. + """ + return tuple( + { + "id": uuid4(), + **keys, + "income_variable": "household_net_income", + "decile_variable": decile_variable, + "entity": "household", + "decile": int(decile), + "quantiles": 10, + "absolute_change": average, + # .py's rows hold a percentage; today's output, a fraction. + "relative_change": deciles.relative[decile] * 100, + } + for decile, average in deciles.average.items() + ) + + +def _intra_decile_rows( + economy: SingleYearMacroOutput, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + rows = _intra_decile_type_rows(economy.intra_decile, "income", keys) + if economy.intra_wealth_decile is not None: + rows += _intra_decile_type_rows(economy.intra_wealth_decile, "wealth", keys) + return rows + + +def _intra_decile_type_rows( + shares: IntraDecileOutput, decile_type: str, keys: Mapping[str, Any] +) -> tuple[dict[str, Any], ...]: + """Return one row per decile, 1 to 10. + + Row 0 is .py's share across the whole population, which the output + doesn't carry: its ``all`` is the mean of the decile shares, which a + reader computes again from these rows. + """ + decile_count = max((len(values) for values in shares.deciles.values()), default=0) + return tuple( + { + "id": uuid4(), + **keys, + "decile_type": decile_type, + "decile": index + 1, + **{ + column: shares.deciles[label][index] + for label, column in INTRA_DECILE_COLUMNS.items() + if index < len(shares.deciles.get(label, ())) + }, + } + for index in range(decile_count) + ) + + +def _poverty_rows( + economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str +) -> tuple[dict[str, Any], ...]: + """Return each simulation run's poverty rates, for everyone and by group. + + As in .py's rows, ``filter_variable`` names the variable that selects a + group and ``filter_group`` the group; both are null for everyone. Today's + output reports 0.0 for a rate .py didn't report and for every rate of a + breakdown it failed to calculate, so a stored 0.0 may be such a rate. + """ + poverty_types = POVERTY_TYPES[country] + return tuple( + { + "id": uuid4(), + "simulation_run_id": simulation_run_id, + "poverty_type": poverty_types[rate.measure], + "entity": "person", + "filter_variable": rate.filter_variable, + "filter_group": rate.filter_group, + "rate": getattr(rate.rates, side), + } + for side, simulation_run_id in _sides(run_ids) + for rate in _poverty_rates(economy) + ) + + +class _PovertyRate(NamedTuple): + """One measure's rates for everyone, or for the group a variable selects.""" + + measure: str + filter_variable: str | None + filter_group: str | None + rates: BaselineReformValue + + +def _poverty_rates(economy: SingleYearMacroOutput) -> list[_PovertyRate]: + """Return each measure's rates for everyone and by group. + + Ages and genders have both measures; the output breaks only poverty, not + deep poverty, down by race, and only in the US. + """ + rates: list[_PovertyRate] = [] + for measure in ("poverty", "deep_poverty"): + by_age = getattr(economy.poverty, measure) + rates.append(_PovertyRate(measure, None, None, by_age.all)) + rates += _group_rates(measure, "age", by_age) + rates += _group_rates( + measure, "is_male", getattr(economy.poverty_by_gender, measure) + ) + if economy.poverty_by_race is not None: + rates += _group_rates("poverty", "race", economy.poverty_by_race.poverty) + return rates + + +def _group_rates( + measure: str, + filter_variable: str, + rates_by_group: AgePovertyOutput | GenderPovertyOutput | RacePovertyOutput, +) -> list[_PovertyRate]: + return [ + _PovertyRate(measure, filter_variable, group, getattr(rates_by_group, group)) + for group in POVERTY_GROUPS[filter_variable] + ] + + +def _inequality_rows( + economy: SingleYearMacroOutput, run_ids: ReportRunIds, country: str +) -> tuple[dict[str, Any], ...]: + return tuple( + { + "id": uuid4(), + "simulation_run_id": simulation_run_id, + "income_variable": INEQUALITY_INCOME_VARIABLES[country], + "entity": "household", + "gini": getattr(economy.inequality.gini, side), + "top_10_share": getattr(economy.inequality.top_10_pct_share, side), + "top_1_share": getattr(economy.inequality.top_1_pct_share, side), + } + for side, simulation_run_id in _sides(run_ids) + ) + + +def _sides(run_ids: ReportRunIds) -> tuple[tuple[str, Any], ...]: + return ( + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + ) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py new file mode 100644 index 000000000..b30aa0887 --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -0,0 +1,349 @@ +"""Conditional writes of report run and simulation run status, as the Modal worker role. + +Every status write is one ``UPDATE`` conditional on the run's expected +status, so a call that Modal restarts after pre-emption, or a second writer, +changes nothing that is already final. A report run's completion also inserts +its result rows, in the same transaction and only when its update changed +the run. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING, Any + +from sqlalchemy import and_, func, insert, or_, select, update +from sqlalchemy.dialects.postgresql import insert as postgresql_insert + +from policyengine_stage12_persistence.report_run_tables import ( + report_run_simulation_runs, + report_runs, + simulation_runs, +) +from policyengine_stage12_persistence.store import ( + ConnectionOpener, + create_stage12_engine, + open_connection_once, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_contract.stage12_execution import ReportRunIds + from sqlalchemy import Engine, Table, Update + from sqlalchemy.engine import Connection + + from policyengine_stage12_persistence.report_run_results import ReportRunResult + + +class RunStatus(StrEnum): + """A report run's or simulation run's status.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +ACTIVE_STATUSES = (RunStatus.PENDING.value, RunStatus.RUNNING.value) + + +@dataclass(frozen=True) +class RunWrite: + """One conditional write's outcome: whether it changed the run, and its status.""" + + changed: bool + status: RunStatus + + +class ReportRunStore: + """Status writes to the Public API's report runs and simulation runs.""" + + def __init__( + self, + database_url: str | None = None, + *, + engine: Engine | None = None, + open_connection: ConnectionOpener = open_connection_once, + ) -> None: + """Use ``engine`` if given, else build one from ``database_url``. + + Raises ValueError when neither is given or the URL is invalid. + """ + if engine is None: + if database_url is None: + msg = "report run database URL is required" + raise ValueError(msg) + engine = create_stage12_engine(database_url) + self.engine = engine + self._open_connection = open_connection + + @contextmanager + def _transaction(self, operation_name: str) -> Iterator[Connection]: + # As in the Stage 12 store, the opener only opens the connection; every + # statement and the COMMIT run after it returns. + connection = self._open_connection( + self.engine.connect, + operation_name=operation_name, + ) + with connection, connection.begin(): + yield connection + + def require_simulation_runs(self, run_ids: ReportRunIds) -> None: + """Raise ValueError unless the report run combines both simulation runs.""" + statement = select( + report_run_simulation_runs.c.role, + report_run_simulation_runs.c.simulation_run_id, + ).where(report_run_simulation_runs.c.report_run_id == run_ids.report_run_id) + with self._transaction("require_simulation_runs") as connection: + members = { + (row.role, row.simulation_run_id) + for row in connection.execute(statement) + } + expected = { + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + } + if not expected <= members: + msg = "the report run does not combine the simulation runs its call names" + raise ValueError(msg) + + def claim_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + claimed_at: datetime, + ) -> RunWrite: + """Mark the report run running under ``modal_call_id``. + + A pending run is claimed, and so is a running one that no other call + holds: the simulation service may mark it running before the call + starts, and Modal restarts a pre-empted call under the same call ID. + """ + statement = ( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + or_( + report_runs.c.status == RunStatus.PENDING.value, + and_( + report_runs.c.status == RunStatus.RUNNING.value, + or_( + report_runs.c.modal_call_id.is_(None), + report_runs.c.modal_call_id == modal_call_id, + ), + ), + ), + ) + .values( + status=RunStatus.RUNNING.value, + modal_call_id=modal_call_id, + started_at=func.coalesce(report_runs.c.started_at, claimed_at), + updated_at=claimed_at, + ) + ) + return self._write(report_runs, report_run_id, statement, "claim_report_run") + + def complete_report_run( + self, + run_ids: ReportRunIds, + *, + modal_call_id: str, + completed_at: datetime, + result: ReportRunResult | None = None, + ) -> RunWrite: + """Mark the report run succeeded while ``modal_call_id`` holds it running. + + ``result``'s rows are inserted in the same transaction, only when the + status changed, so a repeated completion stores nothing twice and a + failed insert leaves the run running. A row of a simulation run that + already has one for the same measure, from a report run sharing that + simulation run, keeps the stored row. + """ + statement = ( + update(report_runs) + .where( + report_runs.c.id == run_ids.report_run_id, + report_runs.c.status == RunStatus.RUNNING.value, + report_runs.c.modal_call_id == modal_call_id, + ) + .values( + status=RunStatus.SUCCEEDED.value, + error_message=None, + completed_at=completed_at, + updated_at=completed_at, + ) + ) + return self._write( + report_runs, + run_ids.report_run_id, + statement, + "complete_report_run", + after_change=None if result is None else _result_inserts(result), + ) + + def fail_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark the report run failed while ``modal_call_id`` holds it running.""" + statement = ( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + report_runs.c.status == RunStatus.RUNNING.value, + report_runs.c.modal_call_id == modal_call_id, + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write(report_runs, report_run_id, statement, "fail_report_run") + + def claim_simulation_run( + self, + simulation_run_id: UUID, + *, + modal_call_id: str, + claimed_at: datetime, + ) -> RunWrite: + """Mark a pending or running simulation run running under ``modal_call_id``. + + A running one moves to the newest call, because the coordinator starts + a simulation again only after stopping its earlier call. + """ + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + simulation_runs.c.status.in_(ACTIVE_STATUSES), + ) + .values( + status=RunStatus.RUNNING.value, + modal_call_id=modal_call_id, + started_at=func.coalesce(simulation_runs.c.started_at, claimed_at), + updated_at=claimed_at, + ) + ) + return self._write( + simulation_runs, simulation_run_id, statement, "claim_simulation_run" + ) + + def complete_simulation_run( + self, + simulation_run_id: UUID, + *, + output_uri: str, + output_sha256: str, + completed_at: datetime, + ) -> RunWrite: + """Record a running simulation run's output file and mark it succeeded.""" + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + simulation_runs.c.status == RunStatus.RUNNING.value, + ) + .values( + status=RunStatus.SUCCEEDED.value, + output_uri=output_uri, + output_sha256=output_sha256, + error_message=None, + completed_at=completed_at, + updated_at=completed_at, + ) + ) + return self._write( + simulation_runs, simulation_run_id, statement, "complete_simulation_run" + ) + + def fail_simulation_run( + self, + simulation_run_id: UUID, + *, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark a pending or running simulation run failed with ``error_message``.""" + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + simulation_runs.c.status.in_(ACTIVE_STATUSES), + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write( + simulation_runs, simulation_run_id, statement, "fail_simulation_run" + ) + + def _write( + self, + table: Table, + run_id: UUID, + statement: Update, + operation_name: str, + *, + after_change: Callable[[Connection], None] | None = None, + ) -> RunWrite: + """Run a conditional update; return whether it changed the run, and its status. + + ``after_change`` runs in the same transaction, only after the update + changed the run. Raises LookupError when the run does not exist. + """ + with self._transaction(operation_name) as connection: + changed = connection.execute( + statement.returning(table.c.status) + ).scalar_one_or_none() + if changed is not None: + if after_change is not None: + after_change(connection) + return RunWrite(changed=True, status=RunStatus(changed)) + current = connection.execute( + select(table.c.status).where(table.c.id == run_id) + ).scalar_one_or_none() + if current is None: + noun = "report run" if table is report_runs else "simulation run" + msg = f"{noun} {run_id} does not exist" + raise LookupError(msg) + return RunWrite(changed=False, status=RunStatus(current)) + + +def _result_inserts(result: ReportRunResult) -> Callable[[Connection], None]: + """Return the step that inserts ``result``'s rows on an open transaction.""" + + def insert_rows(connection: Connection) -> None: + for result_rows in result.rows: + if not result_rows.rows: + continue + statement = ( + postgresql_insert(result_rows.table).on_conflict_do_nothing() + if result_rows.shared_by_reports + else insert(result_rows.table) + ) + connection.execute(statement, _uniform_rows(result_rows.rows)) + + return insert_rows + + +def _uniform_rows(rows: tuple[dict[str, Any], ...]) -> list[dict[str, Any]]: + """Give every row the same columns, unset ones null, for one batched insert.""" + columns = {name for row in rows for name in row} + return [dict.fromkeys(columns) | row for row in rows] diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py new file mode 100644 index 000000000..04792d663 --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -0,0 +1,67 @@ +"""SQLAlchemy mappings of the Public API's run tables that report calls update. + +``policyengine-api`` owns these tables and every schema migration. The +mappings declare only the columns the Modal worker role reads or writes, so +the coordinator and workers can issue typed DML; runtime code never emits DDL. +""" + +from __future__ import annotations + +from sqlalchemy import Column, DateTime, MetaData, String, Table +from sqlalchemy.dialects.postgresql import ENUM, UUID + +RUN_STATUS_VALUES = ("pending", "running", "succeeded", "failed") +SIMULATION_ROLE_VALUES = ("baseline", "reform") + +report_run_status_type = ENUM( + *RUN_STATUS_VALUES, + name="v2_report_run_status", + create_type=False, +) +simulation_run_status_type = ENUM( + *RUN_STATUS_VALUES, + name="v2_simulation_status", + create_type=False, +) +simulation_role_type = ENUM( + *SIMULATION_ROLE_VALUES, + name="v2_simulation_role", + create_type=False, +) + +metadata = MetaData(schema="public") + +report_runs = Table( + "report_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("status", report_run_status_type, nullable=False), + Column("modal_call_id", String(255)), + Column("error_message", String), + Column("started_at", DateTime(timezone=True)), + Column("completed_at", DateTime(timezone=True)), + Column("updated_at", DateTime(timezone=True), nullable=False), +) + +simulation_runs = Table( + "simulation_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("status", simulation_run_status_type, nullable=False), + Column("modal_call_id", String(255)), + Column("error_message", String), + Column("started_at", DateTime(timezone=True)), + Column("completed_at", DateTime(timezone=True)), + Column("updated_at", DateTime(timezone=True), nullable=False), + Column("output_uri", String(2048)), + Column("output_sha256", String(64)), +) + +report_run_simulation_runs = Table( + "report_run_simulation_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("report_run_id", UUID(as_uuid=True), nullable=False), + Column("simulation_run_id", UUID(as_uuid=True), nullable=False), + Column("role", simulation_role_type, nullable=False), +) diff --git a/libs/policyengine-stage12-persistence/tests/conftest.py b/libs/policyengine-stage12-persistence/tests/conftest.py new file mode 100644 index 000000000..84d6ceb0b --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/conftest.py @@ -0,0 +1,39 @@ +"""Fixtures for the report run tests against a disposable Postgres.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import pytest +from report_run_database import ( + SEED_URL_ENV, + ReportRunSeeder, + SeededReportRun, + worker_database_url, +) +from sqlalchemy import create_engine + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@pytest.fixture +def seeder() -> Iterator[ReportRunSeeder]: + """Yield a seeder on the seed role's connection, or skip without a database.""" + worker_database_url() + engine = create_engine(os.environ[SEED_URL_ENV]) + try: + yield ReportRunSeeder(engine) + finally: + engine.dispose() + + +@pytest.fixture +def seeded_run(seeder: ReportRunSeeder) -> Iterator[SeededReportRun]: + """Yield one freshly seeded report run, removed again after the test.""" + seeded = seeder.seed_report_run() + try: + yield seeded + finally: + seeder.remove(seeded) diff --git a/libs/policyengine-stage12-persistence/tests/economy_output.py b/libs/policyengine-stage12-persistence/tests/economy_output.py new file mode 100644 index 000000000..33b342b72 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/economy_output.py @@ -0,0 +1,231 @@ +"""A complete single-year economy output with a distinct value in each field.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_comparison_provenance, + build_spm_provenance, +) + +US_ECONOMY_OUTPUT: dict[str, Any] = { + "spm_provenance": None, + "model_version": "1.715.2", + "data_version": "1.115.5", + "budget": { + "budgetary_impact": 70.25, + "tax_revenue_impact": 100.5, + "state_tax_revenue_impact": 20.125, + "benefit_spending_impact": 30.25, + "households": 2.0, + "baseline_net_income": 1000.75, + }, + "detailed_budget": { + "income_tax": { + "baseline": 100.0, + "reform": 125.0, + "difference": 25.0, + "is_tax": True, + } + }, + "decile": { + "relative": {"1": 0.01, "2": 0.02}, + "average": {"1": 10.0, "2": 20.0}, + }, + "inequality": { + "gini": {"baseline": 0.4, "reform": 0.39}, + "top_10_pct_share": {"baseline": 0.3, "reform": 0.29}, + "top_1_pct_share": {"baseline": 0.1, "reform": 0.09}, + }, + "poverty": { + "poverty": { + "adult": {"baseline": 0.081, "reform": 0.071}, + "all": {"baseline": 0.1, "reform": 0.09}, + "child": {"baseline": 0.12, "reform": 0.11}, + "senior": {"baseline": 0.061, "reform": 0.051}, + }, + "deep_poverty": { + "adult": {"baseline": 0.021, "reform": 0.011}, + "all": {"baseline": 0.031, "reform": 0.022}, + "child": {"baseline": 0.04, "reform": 0.03}, + "senior": {"baseline": 0.015, "reform": 0.012}, + }, + }, + "poverty_by_gender": { + "poverty": { + "male": {"baseline": 0.08, "reform": 0.07}, + "female": {"baseline": 0.09, "reform": 0.085}, + }, + "deep_poverty": { + "male": {"baseline": 0.02, "reform": 0.019}, + "female": {"baseline": 0.025, "reform": 0.024}, + }, + }, + "poverty_by_race": { + "poverty": { + "black": {"baseline": 0.15, "reform": 0.14}, + "hispanic": {"baseline": 0.13, "reform": 0.125}, + "other": {"baseline": 0.11, "reform": 0.105}, + "white": {"baseline": 0.06, "reform": 0.05}, + }, + }, + "intra_decile": { + "all": { + "Gain less than 5%": 0.30000000000000004, + "Gain more than 5%": 0.3, + "Lose less than 5%": 0.15000000000000002, + "Lose more than 5%": 0.05, + "No change": 0.44999999999999996, + }, + "deciles": { + "Gain less than 5%": [0.4, 0.2], + "Gain more than 5%": [0.5, 0.1], + "Lose less than 5%": [0.2, 0.1], + "Lose more than 5%": [0.1, 0.0], + "No change": [0.3, 0.6], + }, + }, + "wealth_decile": None, + "intra_wealth_decile": None, + "labor_supply_response": { + "substitution_lsr": 1.5, + "income_lsr": -0.5, + "relative_lsr": {"income": -0.01, "substitution": 0.02}, + "total_change": 1.0, + "revenue_change": 3.0, + "decile": { + "average": {"income": {"1": -0.1}, "substitution": {"1": 0.2}}, + "relative": {"income": {"1": -0.001}, "substitution": {"1": 0.002}}, + }, + "hours": { + "baseline": 100.0, + "reform": 101.0, + "change": 1.0, + "income_effect": -0.5, + "substitution_effect": 1.5, + }, + }, + "constituency_impact": None, + "local_authority_impact": None, + "congressional_district_impact": None, + "cliff_impact": { + "baseline": {"cliff_gap": 10.0, "cliff_share": 0.1}, + "reform": {"cliff_gap": 9.0, "cliff_share": 0.09}, + }, +} + + +def uk_economy_output() -> dict[str, Any]: + """Return the output as the UK reports it: no race breakdown, with wealth parts.""" + output = deepcopy(US_ECONOMY_OUTPUT) + output["poverty_by_race"] = None + output["budget"]["state_tax_revenue_impact"] = 0.0 + output["wealth_decile"] = { + "relative": {"1": 0.03, "2": 0.04}, + "average": {"1": 30.0, "2": 40.0}, + } + output["intra_wealth_decile"] = deepcopy(output["intra_decile"]) + wealth = output["intra_wealth_decile"] + wealth["deciles"]["No change"] = [0.4, 0.6] + # Today's builder makes each overall share the mean of the decile shares. + wealth["all"]["No change"] = (0.4 + 0.6) / 2 + return output + + +US_STATE_FIPS = {"AL": 1, "AK": 2, "CA": 6} +US_AT_LARGE_STATES = frozenset({"AK"}) + + +def us_output_with_every_part() -> dict[str, Any]: + """Return the US output with districts, SPM receipts and an unflagged program.""" + receipt = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="a" * 64, + selection=SPMSelection( + forecast_content_sha256="a" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ), + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "0.0.0-test-policyengine", + "policyengine-core": "0.0.0-test-policyengine-core", + "policyengine-us": "0.0.0-test-policyengine-us", + "spm-calculator": "0.0.0-test-spm-calculator", + } + ), + ) + output = deepcopy(US_ECONOMY_OUTPUT) + output["spm_provenance"] = build_spm_comparison_provenance( + baseline_receipts=[receipt], reform_receipts=[receipt] + ).model_dump(mode="json", by_alias=True) + output["detailed_budget"]["snap"] = { + "baseline": 80.0, + "reform": 70.0, + "difference": -10.0, + "is_tax": None, + } + output["congressional_district_impact"] = { + "districts": [ + { + "district": district, + "average_household_income_change": 10.0 + index, + "relative_household_income_change": 0.01 * (index + 1), + "winner_percentage": 0.6 - index / 10, + "loser_percentage": 0.3, + "no_change_percentage": 0.1 + index / 10, + "population": 1000.0 + index, + } + for index, district in enumerate(("AL-01", "AL-07", "AK-01", "CA-12")) + ] + } + return output + + +def uk_output_with_areas() -> dict[str, Any]: + """Return the UK output with constituencies and local authorities.""" + output = uk_economy_output() + output["constituency_impact"] = [ + { + "constituency_code": "E14001063", + "constituency_name": "Aldershot", + "average_household_income_change": 120.0, + "relative_household_income_change": 0.004, + "population": 45000.0, + "disclosure_tier": "reported", + "average_household_income_change_lower": 100.0, + "average_household_income_change_upper": 140.0, + "relative_household_income_change_lower": 0.003, + "relative_household_income_change_upper": 0.005, + }, + { + "constituency_code": "W07000081", + "constituency_name": "Aberafan Maesteg", + "average_household_income_change": 90.0, + "relative_household_income_change": 0.003, + "population": 40000.0, + }, + ] + output["local_authority_impact"] = [ + { + "local_authority_code": "E06000001", + "local_authority_name": "Hartlepool", + "average_household_income_change": 80.0, + "relative_household_income_change": 0.002, + "population": 42000.0, + "disclosure_tier": "approximate", + "average_household_income_change_lower": 50.0, + "average_household_income_change_upper": 110.0, + "relative_household_income_change_lower": 0.001, + "relative_household_income_change_upper": 0.003, + } + ] + return output diff --git a/libs/policyengine-stage12-persistence/tests/report_run_database.py b/libs/policyengine-stage12-persistence/tests/report_run_database.py new file mode 100644 index 000000000..11c575f44 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/report_run_database.py @@ -0,0 +1,233 @@ +"""A disposable Postgres holding the Public API's v2 tables, for report run tests. + +The tests run only when both URLs name a disposable database that the Public +API's v2 migrations created and whose declared privileges were applied: + +- ``REPORT_RUN_TEST_DATABASE_URL`` logs in as the Modal worker role, which the + store under test uses; +- ``REPORT_RUN_TEST_SEED_DATABASE_URL`` logs in as a role that may insert the + rows the simulation service would create, so each test can set them up. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any +from uuid import UUID, uuid4 + +import pytest +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from sqlalchemy import Connection, Engine, MetaData, Table, delete, insert, select + +WORKER_URL_ENV = "REPORT_RUN_TEST_DATABASE_URL" +SEED_URL_ENV = "REPORT_RUN_TEST_SEED_DATABASE_URL" +# The catalog rows are needed only while simulations and reports keep their +# foreign keys to the catalog; without those keys they are harmless. +SEEDED_TABLES = ( + "tax_benefit_models", + "tax_benefit_model_versions", + "datasets", + "simulations", + "reports", + "report_runs", + "simulation_runs", + "report_run_simulation_runs", + "report_run_comparisons", + "budget_summary", + "decile_impacts", + "intra_decile_impacts", + "poverty", + "inequality", + "program_statistics", + "congressional_district_impacts", + "constituency_impacts", + "local_authority_impacts", +) +NOW = datetime(2026, 11, 4, tzinfo=UTC) +VERSIONS = {"country_package_version": "1.0.0", "policyengine_version": "6.12.0"} + + +def worker_database_url() -> str: + """Return the Modal worker role's test URL, or skip when none is set.""" + url = os.environ.get(WORKER_URL_ENV) + if not url or not os.environ.get(SEED_URL_ENV): + pytest.skip(f"{WORKER_URL_ENV} and {SEED_URL_ENV} name no test database") + return url + + +@dataclass(frozen=True) +class SeededReportRun: + """A report run and its two simulation runs, as the simulation service makes.""" + + run_ids: ReportRunIds + simulation_ids: tuple[UUID, UUID] + model_id: UUID + + +class ReportRunSeeder: + """Insert and read back report runs as the seed role.""" + + def __init__(self, engine: Engine) -> None: + """Reflect the seeded tables through ``engine``; this reads, never alters.""" + self.engine = engine + self.metadata = MetaData(schema="public") + self.metadata.reflect(engine, only=list(SEEDED_TABLES)) + + def table(self, name: str) -> Table: + """Return the reflected table ``name``.""" + return self.metadata.tables[f"public.{name}"] + + def seed_report_run(self) -> SeededReportRun: + """Insert a pending report run whose baseline and reform runs are pending.""" + model_id, version_id, dataset_id = uuid4(), uuid4(), uuid4() + simulation_ids = (uuid4(), uuid4()) + report_id = uuid4() + run_ids = ReportRunIds( + report_run_id=uuid4(), + baseline_simulation_run_id=uuid4(), + reform_simulation_run_id=uuid4(), + ) + with self.engine.begin() as connection: + self._insert_catalog(connection, model_id, version_id, dataset_id) + for simulation_id in simulation_ids: + connection.execute( + insert(self.table("simulations")).values( + id=simulation_id, + simulation_type="economy", + dataset_id=dataset_id, + tax_benefit_model_version_id=version_id, + year=2026, + ) + ) + connection.execute( + insert(self.table("reports")).values( + id=report_id, + label="Report run test", + country="us", + tax_benefit_model_id=model_id, + baseline_simulation_id=simulation_ids[0], + reform_simulation_id=simulation_ids[1], + year=2026, + inputs={}, + ) + ) + connection.execute( + insert(self.table("report_runs")).values( + id=run_ids.report_run_id, + report_id=report_id, + status="pending", + trigger="initial", + **VERSIONS, + ) + ) + for run_id, simulation_id, role in ( + (run_ids.baseline_simulation_run_id, simulation_ids[0], "baseline"), + (run_ids.reform_simulation_run_id, simulation_ids[1], "reform"), + ): + self._insert_simulation_run( + connection, run_ids.report_run_id, run_id, simulation_id, role + ) + return SeededReportRun( + run_ids=run_ids, simulation_ids=simulation_ids, model_id=model_id + ) + + def _insert_catalog( + self, connection: Connection, model_id: UUID, version_id: UUID, dataset_id: UUID + ) -> None: + connection.execute( + insert(self.table("tax_benefit_models")).values( + id=model_id, name=f"test-{model_id.hex[:24]}" + ) + ) + connection.execute( + insert(self.table("tax_benefit_model_versions")).values( + id=version_id, + model_id=model_id, + version="1.0.0", + current_law_id=1, + metadata_time_periods={}, + ) + ) + connection.execute( + insert(self.table("datasets")).values( + id=dataset_id, + name="report-run-test", + year=2026, + is_output_dataset=False, + tax_benefit_model_version_id=version_id, + ) + ) + + def _insert_simulation_run( + self, + connection: Connection, + report_run_id: UUID, + simulation_run_id: UUID, + simulation_id: UUID, + role: str, + ) -> None: + version_id = connection.execute( + select(self.table("simulations").c.tax_benefit_model_version_id).where( + self.table("simulations").c.id == simulation_id + ) + ).scalar_one() + connection.execute( + insert(self.table("simulation_runs")).values( + id=simulation_run_id, + simulation_id=simulation_id, + tax_benefit_model_version_id=version_id, + status="pending", + **VERSIONS, + ) + ) + connection.execute( + insert(self.table("report_run_simulation_runs")).values( + id=uuid4(), + report_run_id=report_run_id, + simulation_run_id=simulation_run_id, + role=role, + ) + ) + + def row(self, table_name: str, row_id: UUID) -> dict[str, Any]: + """Return the row ``row_id`` of ``table_name`` as a dictionary.""" + table = self.table(table_name) + with self.engine.connect() as connection: + row = connection.execute(select(table).where(table.c.id == row_id)).one() + return dict(row._mapping) + + def rows(self, table_name: str, **filters: object) -> list[dict[str, Any]]: + """Return the rows of ``table_name`` whose columns equal ``filters``.""" + table = self.table(table_name) + statement = select(table).where( + *(table.c[name] == value for name, value in filters.items()) + ) + with self.engine.connect() as connection: + return [dict(row._mapping) for row in connection.execute(statement)] + + def remove(self, seeded: SeededReportRun) -> None: + """Delete the seeded rows; the runs and their results cascade.""" + reports = self.table("reports") + simulations = self.table("simulations") + datasets = self.table("datasets") + versions = self.table("tax_benefit_model_versions") + models = self.table("tax_benefit_models") + version_ids = select(versions.c.id).where( + versions.c.model_id == seeded.model_id + ) + with self.engine.begin() as connection: + connection.execute( + delete(reports).where(reports.c.tax_benefit_model_id == seeded.model_id) + ) + connection.execute( + delete(simulations).where(simulations.c.id.in_(seeded.simulation_ids)) + ) + connection.execute( + delete(datasets).where( + datasets.c.tax_benefit_model_version_id.in_(version_ids) + ) + ) + # Deleting the model removes its versions. + connection.execute(delete(models).where(models.c.id == seeded.model_id)) diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_results.py b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py new file mode 100644 index 000000000..fa4b1a2f1 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_results.py @@ -0,0 +1,465 @@ +"""A report run's economy output, stored as rows with the run's success.""" + +from __future__ import annotations + +from datetime import timedelta +from uuid import uuid4 + +import pytest +from economy_output import ( + US_AT_LARGE_STATES, + US_ECONOMY_OUTPUT, + US_STATE_FIPS, + uk_economy_output, + uk_output_with_areas, + us_output_with_every_part, +) +from policyengine_stage12_persistence import ( + DistrictCodes, + ReportRunStore, + RunStatus, + RunWrite, + report_run_result, +) +from policyengine_stage12_persistence.report_run_results import ResultRows +from report_run_database import NOW, SeededReportRun, worker_database_url +from sqlalchemy import insert +from sqlalchemy.exc import IntegrityError + +LATER = NOW + timedelta(minutes=5) +INTRA_DECILE_LABELS = { + "Lose more than 5%": "lose_more_than_5pct", + "Lose less than 5%": "lose_less_than_5pct", + "No change": "no_change", + "Gain less than 5%": "gain_less_than_5pct", + "Gain more than 5%": "gain_more_than_5pct", +} +US_POVERTY_TYPES = {"poverty": "spm", "deep_poverty": "spm_deep"} +UK_POVERTY_TYPES = {"poverty": "relative_bhc", "deep_poverty": "absolute_bhc"} +# The output part each of .py's poverty filter variables rebuilds; a row +# with none is everyone's rate, the part's "all". +POVERTY_PARTS = { + None: "poverty", + "age": "poverty", + "is_male": "poverty_by_gender", + "race": "poverty_by_race", +} +POVERTY_PART_NAMES = ("poverty", "poverty_by_gender", "poverty_by_race") + + +@pytest.fixture +def store(): + run_store = ReportRunStore(worker_database_url()) + try: + yield run_store + finally: + run_store.engine.dispose() + + +DISTRICT_CODES = DistrictCodes( + state_fips=US_STATE_FIPS, at_large_states=US_AT_LARGE_STATES +) +STATE_BY_FIPS = {fips: state for state, fips in US_STATE_FIPS.items()} +AREA_FIELDS = ( + "average_household_income_change", + "relative_household_income_change", + "population", + "disclosure_tier", + "average_household_income_change_lower", + "average_household_income_change_upper", + "relative_household_income_change_lower", + "relative_household_income_change_upper", +) + + +def _complete(store, seeded_run, output=US_ECONOMY_OUTPUT, country="us"): + store.claim_report_run( + seeded_run.run_ids.report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + return store.complete_report_run( + seeded_run.run_ids, + modal_call_id="fc-coordinator", + completed_at=LATER, + result=report_run_result( + output, + run_ids=seeded_run.run_ids, + country=country, + district_codes=DISTRICT_CODES, + ), + ) + + +def _rebuilt_district(row): + """Rebuild today's district code as the output builder formats it.""" + state = STATE_BY_FIPS[row["state_fips"]] + number = 1 if state in US_AT_LARGE_STATES else row["district_number"] + return f"{state}-{number:02d}" + + +def _rebuilt_budget(seeder, report_run_id): + rows = { + row["variable_name"]: row + for row in seeder.rows("budget_summary", report_run_id=report_run_id) + } + tax = rows["household_tax"]["change"] + benefits = rows["household_benefits"]["change"] + # Only the US models state income tax; today's output reports 0.0 elsewhere. + state_tax = rows.get("household_state_income_tax", {"change": 0.0})["change"] + return { + "budgetary_impact": tax - benefits, + "tax_revenue_impact": tax, + "state_tax_revenue_impact": state_tax, + "benefit_spending_impact": benefits, + "households": rows["household_weight"]["baseline_total"], + "baseline_net_income": rows["household_net_income"]["baseline_total"], + } + + +def _rebuilt_intra_decile(rows): + """Rebuild today's shares; each overall one is the deciles' mean, as today.""" + by_decile = {row["decile"]: row for row in rows} + deciles = sorted(by_decile) + shares = { + label: [by_decile[decile][column] for decile in deciles] + for label, column in INTRA_DECILE_LABELS.items() + } + return { + "all": { + label: sum(values) / len(values) if values else 0.0 + for label, values in shares.items() + }, + "deciles": shares, + } + + +def _rebuilt_poverty(seeder, run_ids, poverty_types): + """Rebuild today's three poverty parts from both simulation runs' rows.""" + measures = { + poverty_type: measure for measure, poverty_type in poverty_types.items() + } + parts = {} + for side, run_id in ( + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + ): + for row in seeder.rows("poverty", simulation_run_id=run_id): + assert row["entity"] == "person" + # A group needs the variable that selects it, and everyone has neither. + assert (row["filter_variable"] is None) == (row["filter_group"] is None) + part = parts.setdefault(POVERTY_PARTS[row["filter_variable"]], {}) + groups = part.setdefault(measures[row["poverty_type"]], {}) + groups.setdefault(row["filter_group"] or "all", {})[side] = row["rate"] + return {name: parts.get(name) for name in POVERTY_PART_NAMES} + + +def _rebuilt_deciles(rows): + """Rebuild today's decile part; .py's relative change is a percentage.""" + return { + "average": {str(row["decile"]): row["absolute_change"] for row in rows}, + "relative": {str(row["decile"]): row["relative_change"] / 100 for row in rows}, + } + + +def test_a_report_run_stores_its_economy_output_as_rows( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + + assert _complete(store, seeded_run) == RunWrite( + changed=True, status=RunStatus.SUCCEEDED + ) + + output = US_ECONOMY_OUTPUT + assert _rebuilt_budget(seeder, run_ids.report_run_id) == output["budget"] + deciles = seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id) + assert {str(row["decile"]): row["absolute_change"] for row in deciles} == output[ + "decile" + ]["average"] + # .py's rows hold the relative change as a percentage; today's output + # holds a fraction, so a reader divides by 100. + assert { + str(row["decile"]): row["relative_change"] / 100 for row in deciles + } == pytest.approx(output["decile"]["relative"], rel=1e-12) + assert sorted(row["relative_change"] for row in deciles) == pytest.approx( + [1.0, 2.0] + ) + assert {row["income_variable"] for row in deciles} == {"household_net_income"} + # Income deciles name no separate variable to group households by. + assert {row["decile_variable"] for row in deciles} == {None} + assert all( + row["baseline_simulation_run_id"] == run_ids.baseline_simulation_run_id + and row["reform_simulation_run_id"] == run_ids.reform_simulation_run_id + for row in deciles + ) + intra = seeder.rows("intra_decile_impacts", report_run_id=run_ids.report_run_id) + # Row 0 is .py's share across the whole population, which today's output + # doesn't carry, so none is written. + assert sorted(row["decile"] for row in intra) == [1, 2] + assert _rebuilt_intra_decile(intra) == output["intra_decile"] + # Everyone's rates, and those by age, gender and race, under .py's groups. + assert _rebuilt_poverty(seeder, run_ids, US_POVERTY_TYPES) == { + part: output[part] for part in POVERTY_PART_NAMES + } + for side, run_id in ( + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + ): + (inequality,) = seeder.rows("inequality", simulation_run_id=run_id) + assert inequality["income_variable"] == "household_net_income" + assert inequality["entity"] == "household" + assert ( + inequality["gini"], + inequality["top_10_share"], + inequality["top_1_share"], + ) == ( + output["inequality"]["gini"][side], + output["inequality"]["top_10_pct_share"][side], + output["inequality"]["top_1_pct_share"][side], + ) + # The parts with no result table are kept whole for the compared pair. + (comparison,) = seeder.rows( + "report_run_comparisons", report_run_id=run_ids.report_run_id + ) + assert ( + comparison["baseline_simulation_run_id"], + comparison["reform_simulation_run_id"], + ) == (run_ids.baseline_simulation_run_id, run_ids.reform_simulation_run_id) + assert comparison["labor_supply_result"] == output["labor_supply_response"] + assert comparison["cliff_impact_result"] == output["cliff_impact"] + + +def test_the_uk_output_uses_its_own_poverty_types_and_wealth_deciles( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + output = uk_economy_output() + + _complete(store, seeded_run, output=output, country="uk") + + # No race breakdown outside the US. + assert _rebuilt_poverty(seeder, run_ids, UK_POVERTY_TYPES) == { + part: output[part] for part in POVERTY_PART_NAMES + } + deciles = seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id) + # .py groups wealth deciles' household net income by the wealth decile. + wealth_deciles = [ + row for row in deciles if row["decile_variable"] == "household_wealth_decile" + ] + assert {row["income_variable"] for row in wealth_deciles} == { + "household_net_income" + } + assert {row["entity"] for row in wealth_deciles} == {"household"} + income_deciles = [row for row in deciles if row["decile_variable"] is None] + for rows, part in ((wealth_deciles, "wealth_decile"), (income_deciles, "decile")): + rebuilt = _rebuilt_deciles(rows) + assert rebuilt["average"] == output[part]["average"] + assert rebuilt["relative"] == pytest.approx(output[part]["relative"], rel=1e-12) + assert "household_state_income_tax" not in { + row["variable_name"] + for row in seeder.rows("budget_summary", report_run_id=run_ids.report_run_id) + } + assert ( + _rebuilt_budget(seeder, run_ids.report_run_id) + == (uk_economy_output()["budget"]) + ) + (inequality,) = seeder.rows( + "inequality", simulation_run_id=run_ids.reform_simulation_run_id + ) + assert inequality["income_variable"] == "equiv_hbai_household_net_income" + wealth = [ + row + for row in seeder.rows( + "intra_decile_impacts", report_run_id=run_ids.report_run_id + ) + if row["decile_type"] == "wealth" + ] + assert _rebuilt_intra_decile(wealth) == uk_economy_output()["intra_wealth_decile"] + + +def test_a_repeated_completion_stores_no_second_result( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + _complete(store, seeded_run) + + repeated = store.complete_report_run( + run_ids, + modal_call_id="fc-coordinator", + completed_at=LATER, + result=report_run_result(US_ECONOMY_OUTPUT, run_ids=run_ids, country="us"), + ) + + assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert len(seeder.rows("budget_summary", report_run_id=run_ids.report_run_id)) == 5 + assert len(seeder.rows("decile_impacts", report_run_id=run_ids.report_run_id)) == 2 + assert ( + len(seeder.rows("report_run_comparisons", report_run_id=run_ids.report_run_id)) + == 1 + ) + + +def test_a_result_that_cannot_be_stored_leaves_the_report_run_running( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + store.claim_report_run( + run_ids.report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + result = report_run_result(US_ECONOMY_OUTPUT, run_ids=run_ids, country="us") + budget = next(rows for rows in result.rows if rows.table.name == "budget_summary") + # The same row twice breaks the table's unique key in the last insert. + duplicated = ResultRows( + table=budget.table, + rows=(*budget.rows, {**budget.rows[0], "id": uuid4()}), + shared_by_reports=False, + ) + broken = result.__class__( + rows=(*[rows for rows in result.rows if rows is not budget], duplicated), + ) + + with pytest.raises(IntegrityError): + store.complete_report_run( + run_ids, modal_call_id="fc-coordinator", completed_at=LATER, result=broken + ) + + report_run = seeder.row("report_runs", run_ids.report_run_id) + assert report_run["status"] == "running" + for table_name in ( + "report_run_comparisons", + "decile_impacts", + "budget_summary", + "intra_decile_impacts", + ): + assert seeder.rows(table_name, report_run_id=run_ids.report_run_id) == [] + assert ( + seeder.rows("poverty", simulation_run_id=run_ids.baseline_simulation_run_id) + == [] + ) + + +def test_a_shared_simulation_run_keeps_the_results_first_stored_for_it( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + # An earlier report run that shares this baseline stored its poverty first. + with seeder.engine.begin() as connection: + connection.execute( + insert(seeder.table("poverty")).values( + id=uuid4(), + simulation_run_id=run_ids.baseline_simulation_run_id, + poverty_type="spm", + entity="person", + filter_variable=None, + filter_group=None, + rate=0.5, + ) + ) + + _complete(store, seeded_run) + + stored = seeder.rows( + "poverty", + simulation_run_id=run_ids.baseline_simulation_run_id, + poverty_type="spm", + filter_variable=None, + ) + assert [row["rate"] for row in stored] == [0.5] + # Everyone, three ages and two genders under both poverty types, and four + # races under SPM poverty; the baseline gains all but the stored one. + baseline = seeder.rows( + "poverty", simulation_run_id=run_ids.baseline_simulation_run_id + ) + reform = seeder.rows("poverty", simulation_run_id=run_ids.reform_simulation_run_id) + assert len(reform) == len(baseline) == 16 + + +def test_the_detailed_budget_districts_and_receipts_are_stored_losslessly( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + output = us_output_with_every_part() + + _complete(store, seeded_run, output=output) + + programs = seeder.rows("program_statistics", report_run_id=report_run_id) + assert { + row["program_name"]: { + "baseline": row["baseline_total"], + "reform": row["reform_total"], + "difference": row["change"], + "is_tax": row["is_tax"], + } + for row in programs + } == output["detailed_budget"] + assert {row["entity"] for row in programs} == {None} + districts = seeder.rows( + "congressional_district_impacts", report_run_id=report_run_id + ) + assert sorted( + ( + { + "district": _rebuilt_district(row), + **{ + field: row[field] + for field in ( + "average_household_income_change", + "relative_household_income_change", + "winner_percentage", + "loser_percentage", + "no_change_percentage", + "population", + ) + }, + } + for row in districts + ), + key=lambda district: district["district"], + ) == sorted( + output["congressional_district_impact"]["districts"], + key=lambda district: district["district"], + ) + # Census GEOIDs: an at-large district is numbered 00. + assert {row["district_geoid"] for row in districts} == {101, 107, 200, 612} + (comparison,) = seeder.rows("report_run_comparisons", report_run_id=report_run_id) + assert comparison["spm_provenance"] == output["spm_provenance"] + assert comparison["data_version"] == output["data_version"] + + +def test_the_uk_area_impacts_are_stored_with_their_disclosure_fields( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + output = uk_output_with_areas() + + _complete(store, seeded_run, output=output, country="uk") + + for table_name, kind in ( + ("constituency_impacts", "constituency"), + ("local_authority_impacts", "local_authority"), + ): + stored = { + row[f"{kind}_code"]: { + f"{kind}_code": row[f"{kind}_code"], + f"{kind}_name": row[f"{kind}_name"], + **{field: row[field] for field in AREA_FIELDS}, + } + for row in seeder.rows(table_name, report_run_id=report_run_id) + } + expected = { + record[f"{kind}_code"]: dict.fromkeys(AREA_FIELDS) | record + for record in output[f"{kind}_impact"] + } + assert stored == expected + + +def test_a_district_in_an_unknown_state_is_refused(seeded_run: SeededReportRun) -> None: + output = us_output_with_every_part() + output["congressional_district_impact"]["districts"][0]["district"] = "ZZ-01" + + with pytest.raises(ValueError, match="ZZ"): + report_run_result( + output, + run_ids=seeded_run.run_ids, + country="us", + district_codes=DISTRICT_CODES, + ) diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_store.py b/libs/policyengine-stage12-persistence/tests/test_report_run_store.py new file mode 100644 index 000000000..6d2545e74 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_store.py @@ -0,0 +1,249 @@ +"""The Modal worker role's conditional writes to report runs and simulation runs. + +These run against a disposable Postgres that the Public API's v2 migrations +created, connected as the Modal worker role (see ``report_run_database``). +""" + +from __future__ import annotations + +from datetime import timedelta +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_stage12_persistence import ( + ReportRunStore, + RunStatus, + RunWrite, +) +from report_run_database import NOW, SeededReportRun, worker_database_url +from sqlalchemy import func, insert, select +from sqlalchemy.exc import ProgrammingError + +LATER = NOW + timedelta(minutes=5) + + +@pytest.fixture +def store(): + run_store = ReportRunStore(worker_database_url()) + try: + yield run_store + finally: + run_store.engine.dispose() + + +def test_the_store_connects_as_the_modal_worker_role(store) -> None: + with store.engine.connect() as connection: + role = connection.execute(select(func.current_user())).scalar_one() + + assert role == "policyengine_v2_modal_worker" + + +def test_the_modal_worker_role_cannot_create_a_report_run(store, seeded_run) -> None: + from policyengine_stage12_persistence.report_run_tables import report_runs + + with ( + pytest.raises(ProgrammingError, match="permission denied"), + store.engine.begin() as connection, + ): + connection.execute( + insert(report_runs).values(id=uuid4(), status="pending", updated_at=NOW) + ) + + +def test_a_report_run_is_claimed_once_and_resumed_by_the_same_call( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + + first = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + other_call = store.claim_report_run( + report_run_id, modal_call_id="fc-other", claimed_at=LATER + ) + restarted = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=LATER + ) + + assert first == RunWrite(changed=True, status=RunStatus.RUNNING) + assert other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert restarted == RunWrite(changed=True, status=RunStatus.RUNNING) + row = seeder.row("report_runs", report_run_id) + assert row["modal_call_id"] == "fc-coordinator" + assert row["started_at"] == NOW + assert row["updated_at"] == LATER + + +def test_a_report_run_the_simulation_service_marked_running_is_claimed( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + report_runs = seeder.table("report_runs") + with seeder.engine.begin() as connection: + connection.execute( + report_runs.update() + .where(report_runs.c.id == report_run_id) + .values(status="running", started_at=NOW) + ) + + claimed = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=LATER + ) + + assert claimed == RunWrite(changed=True, status=RunStatus.RUNNING) + assert seeder.row("report_runs", report_run_id)["started_at"] == NOW + + +def test_a_report_run_completes_only_while_its_call_holds_it( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + + by_other_call = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-other", completed_at=LATER + ) + completed = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER + ) + repeated = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER + ) + reclaimed = store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=LATER + ) + + assert by_other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert completed == RunWrite(changed=True, status=RunStatus.SUCCEEDED) + assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert reclaimed == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + row = seeder.row("report_runs", report_run_id) + assert row["status"] == "succeeded" + assert row["completed_at"] == LATER + assert row["error_message"] is None + + +def test_a_report_run_fails_once_and_then_never_completes( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + store.claim_report_run( + report_run_id, modal_call_id="fc-coordinator", claimed_at=NOW + ) + + failed = store.fail_report_run( + report_run_id, + modal_call_id="fc-coordinator", + error_message="simulation_execution_failed: the reform failed", + failed_at=LATER, + ) + again = store.fail_report_run( + report_run_id, + modal_call_id="fc-coordinator", + error_message="another failure", + failed_at=LATER, + ) + completed = store.complete_report_run( + seeded_run.run_ids, modal_call_id="fc-coordinator", completed_at=LATER + ) + + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + assert again == RunWrite(changed=False, status=RunStatus.FAILED) + assert completed == RunWrite(changed=False, status=RunStatus.FAILED) + row = seeder.row("report_runs", report_run_id) + assert row["error_message"] == "simulation_execution_failed: the reform failed" + assert row["completed_at"] == LATER + + +def test_a_simulation_run_moves_to_the_latest_call_until_it_finishes( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_id = seeded_run.run_ids.baseline_simulation_run_id + + first = store.claim_simulation_run( + run_id, modal_call_id="fc-worker-1", claimed_at=NOW + ) + # The coordinator started this simulation again after a restart. + second = store.claim_simulation_run( + run_id, modal_call_id="fc-worker-2", claimed_at=LATER + ) + completed = store.complete_simulation_run( + run_id, + output_uri="gs://private/simulations/baseline.parquet", + output_sha256="a" * 64, + completed_at=LATER, + ) + repeated = store.complete_simulation_run( + run_id, + output_uri="gs://private/simulations/baseline-other.parquet", + output_sha256="b" * 64, + completed_at=LATER, + ) + failed_after = store.fail_simulation_run( + run_id, error_message="cancelled", failed_at=LATER + ) + claimed_after = store.claim_simulation_run( + run_id, modal_call_id="fc-worker-3", claimed_at=LATER + ) + + assert first == RunWrite(changed=True, status=RunStatus.RUNNING) + assert second == RunWrite(changed=True, status=RunStatus.RUNNING) + assert completed == RunWrite(changed=True, status=RunStatus.SUCCEEDED) + assert repeated == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert failed_after == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + assert claimed_after == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + row = seeder.row("simulation_runs", run_id) + assert row["modal_call_id"] == "fc-worker-2" + assert row["started_at"] == NOW + assert row["output_uri"] == "gs://private/simulations/baseline.parquet" + assert row["output_sha256"] == "a" * 64 + + +def test_a_pending_or_running_simulation_run_can_fail_but_not_resume( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_id = seeded_run.run_ids.reform_simulation_run_id + + failed = store.fail_simulation_run( + run_id, error_message="simulation_cancelled: stopped", failed_at=NOW + ) + claimed = store.claim_simulation_run( + run_id, modal_call_id="fc-worker", claimed_at=LATER + ) + completed = store.complete_simulation_run( + run_id, + output_uri="gs://private/x.parquet", + output_sha256="c" * 64, + completed_at=LATER, + ) + + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + assert claimed == RunWrite(changed=False, status=RunStatus.FAILED) + assert completed == RunWrite(changed=False, status=RunStatus.FAILED) + row = seeder.row("simulation_runs", run_id) + assert row["error_message"] == "simulation_cancelled: stopped" + assert row["output_uri"] is None + + +def test_a_report_run_must_combine_the_simulation_runs_its_call_names( + store, seeded_run: SeededReportRun +) -> None: + store.require_simulation_runs(seeded_run.run_ids) + + swapped = ReportRunIds( + report_run_id=seeded_run.run_ids.report_run_id, + baseline_simulation_run_id=seeded_run.run_ids.reform_simulation_run_id, + reform_simulation_run_id=seeded_run.run_ids.baseline_simulation_run_id, + ) + with pytest.raises(ValueError, match="does not combine"): + store.require_simulation_runs(swapped) + + +def test_a_missing_run_raises_lookup_error(store, seeded_run) -> None: + with pytest.raises(LookupError, match="report run"): + store.claim_report_run(uuid4(), modal_call_id="fc-coordinator", claimed_at=NOW) + with pytest.raises(LookupError, match="simulation run"): + store.fail_simulation_run(uuid4(), error_message="failed", failed_at=NOW) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 69eedc493..e9c97852e 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -25,6 +25,10 @@ ) from policyengine_simulation_executor.stage12_cache.models import CacheManifest from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore +from policyengine_simulation_executor.stage12_runtime.report_run_access import ( + MODAL_WORKER_DATABASE_SECRET_NAME, + MODAL_WORKER_DATABASE_URL_ENV, +) from policyengine_simulation_executor.stage12_timeouts import ( COORDINATOR_TIMEOUT_SECONDS, UK_SIMULATION_TIMEOUT_SECONDS, @@ -130,6 +134,13 @@ def _external_assertions(environment: dict[str, str]) -> dict[str, str]: hf_secret, comparison_runtime_secret, ] +# Only the functions that record a report run's or simulation run's status +# and results hold the Modal worker role's database credentials. +modal_worker_database_secret = modal.Secret.from_name( + MODAL_WORKER_DATABASE_SECRET_NAME, + required_keys=[MODAL_WORKER_DATABASE_URL_ENV], +) +report_run_writer_secrets = [*worker_secrets, modal_worker_database_secret] def _country_bundle(country: CountryId) -> Stage12CountryBundle: @@ -304,7 +315,7 @@ def run_single_simulation_segment_us( timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, - secrets=worker_secrets, + secrets=report_run_writer_secrets, ) def run_single_simulation_us( payload: dict, @@ -349,6 +360,7 @@ def segmented_calculator( segmented_calculator=segmented_calculator, runtime=runtime, ), + modal_call_id=modal.current_function_call_id(), ) finally: runtime.shutdown() @@ -361,7 +373,7 @@ def segmented_calculator( timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, - secrets=worker_secrets, + secrets=report_run_writer_secrets, ) def run_single_simulation_uk( payload: dict, @@ -393,6 +405,7 @@ def run_single_simulation_uk( dependencies=stage12_runtime.SimulationWorkerDependencies( runtime=runtime ), + modal_call_id=modal.current_function_call_id(), ) finally: runtime.shutdown() @@ -407,7 +420,7 @@ def run_single_simulation_uk( timeout=COORDINATOR_TIMEOUT_SECONDS, retries=0, max_containers=10, - secrets=worker_secrets, + secrets=report_run_writer_secrets, ) def coordinate_report( payload: dict, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index 71877ee60..c2a11d555 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -72,19 +72,35 @@ def comparison_run_prefix( ) -def input_path(*, prefix: str, role: str) -> str: - """Return the key of the planned execution input JSON for one simulation role.""" - return f"{prefix}/inputs/{role}.json" +def _content_named(name: str, content_sha256: str | None) -> str: + """Return ``name``, followed by ``-`` when a digest is given.""" + return name if content_sha256 is None else f"{name}-{content_sha256}" -def simulation_path(*, prefix: str, role: str) -> str: - """Return the key of the Parquet output for one simulation role.""" - return f"{prefix}/simulations/{role}.parquet" +def input_path(*, prefix: str, role: str, content_sha256: str | None = None) -> str: + """Return the key of the planned execution input JSON for one simulation role. + + With ``content_sha256``, the key also names the content's SHA-256. + """ + return f"{prefix}/inputs/{_content_named(role, content_sha256)}.json" + + +def simulation_path( + *, prefix: str, role: str, content_sha256: str | None = None +) -> str: + """Return the key of the Parquet output for one simulation role. + + With ``content_sha256``, the key also names the content's SHA-256. + """ + return f"{prefix}/simulations/{_content_named(role, content_sha256)}.parquet" -def aggregate_path(*, prefix: str) -> str: - """Return the key of the run's aggregate report JSON.""" - return f"{prefix}/reports/aggregate.json" +def aggregate_path(*, prefix: str, content_sha256: str | None = None) -> str: + """Return the key of the run's aggregate report JSON. + + With ``content_sha256``, the key also names the content's SHA-256. + """ + return f"{prefix}/reports/{_content_named('aggregate', content_sha256)}.json" def comparison_path(*, prefix: str) -> str: @@ -273,16 +289,30 @@ class Stage12ArtifactStore: """Write-once storage for one private bucket's Stage 12 run artifacts. A write to a path already holding identical bytes succeeds; different - bytes raise ``RuntimeError``. + bytes raise ``RuntimeError``. A ``content_addressed`` store names each + input, simulation output and aggregate's SHA-256 in its key, so a call + Modal restarts can store other bytes beside its earlier attempt's, and + identical bytes reuse the earlier object. """ - def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None) -> None: + def __init__( + self, + bucket_name: str, + *, + store: ArtifactStore | None = None, + content_addressed: bool = False, + ) -> None: """Bind to ``bucket_name``, which must be non-empty (there is no default).""" if not bucket_name: msg = "Stage 12 requires an explicitly configured private bucket" raise ValueError(msg) self.bucket_name = bucket_name self._store = store or ArtifactStore(bucket_name) + self._content_addressed = content_addressed + + def _key_digest(self, payload: bytes) -> str | None: + """Return the SHA-256 a content-addressed key names, else None.""" + return sha256(payload).hexdigest() if self._content_addressed else None def _write_immutable( self, @@ -316,9 +346,14 @@ def write_input( simulation: PlannedSimulationExecutionInput, ) -> ArtifactReference: """Store a simulation's planned execution input as canonical JSON.""" + payload = canonical_json_bytes(simulation.model_dump(mode="json")) return self._write_immutable( - input_path(prefix=prefix, role=simulation.role.value), - canonical_json_bytes(simulation.model_dump(mode="json")), + input_path( + prefix=prefix, + role=simulation.role.value, + content_sha256=self._key_digest(payload), + ), + payload, content_type=REPORT_MEDIA_TYPE, ) @@ -338,7 +373,11 @@ def write_simulation( uk_local_authority_metadata=uk_local_authority_metadata, ) artifact = self._write_immutable( - simulation_path(prefix=prefix, role=simulation.role.value), + simulation_path( + prefix=prefix, + role=simulation.role.value, + content_sha256=self._key_digest(payload), + ), payload, content_type=SIMULATION_MEDIA_TYPE, ) @@ -361,9 +400,13 @@ def write_aggregate( ) -> ArtifactReference: """Validate and store the run's aggregate report.""" validated = AggregateReportArtifactPayload.model_validate(payload) + encoded = canonical_json_bytes(validated.model_dump(mode="json")) return self._write_immutable( - aggregate_path(prefix=prefix), - canonical_json_bytes(validated.model_dump(mode="json")), + aggregate_path( + prefix=prefix, + content_sha256=self._key_digest(encoded), + ), + encoded, content_type=REPORT_MEDIA_TYPE, ) 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 d64f9fd2f..fe095ecb9 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 @@ -24,14 +24,20 @@ SIMULATION_WAIT_TIMEOUT_SECONDS, ) -from .coordinated_simulations import SIMULATION_CANCELLED_ERROR_CODE, logger +from .coordinated_simulations import ( + SIMULATION_CANCELLED_ERROR_CODE, + cancel_pending_simulations, + cancellation, + logger, +) from .dependencies import ModalChildInvoker, artifact_store, runtime_store from .failure_causes import StageProgress +from .failures import Stage12ExecutionError, Stage12FailureDetail from .parent_claim import ( - claim_parent, validate_dispatch_context, validate_submitted_parent, ) +from .report_claims import claim_report from .report_coordination import ( SIMULATION_POLL_BACKOFF_FACTOR, SIMULATION_POLL_INITIAL_SECONDS, @@ -40,6 +46,13 @@ CoordinatorInvocation, ReportCoordination, ) +from .report_run_access import report_run_store +from .report_run_recording import ReportRunRecorder, SimulationRunRecordingStore + +CANCELLED_DETAIL = Stage12FailureDetail( + error_code="report_cancelled", + error_summary="Modal cancelled the report's coordinator call.", +) __all__ = [ "SIMULATION_CANCELLED_ERROR_CODE", @@ -85,12 +98,24 @@ def coordinate_report( context=context, parent=submitted_parent, ) + persistence = dependencies.store or runtime_store(dependencies.runtime) + report_run = None + if context.report_run is not None: + report_run = ReportRunRecorder( + writer=dependencies.report_runs or report_run_store(dependencies.runtime), + run_ids=context.report_run, + modal_call_id=invocation.invocation_id, + ) + persistence = SimulationRunRecordingStore( + persistence, report_run.writer, context.report_run + ) coordination = ReportCoordination( report=report, context=context, invocation=invocation, - persistence=dependencies.store or runtime_store(dependencies.runtime), - artifact_storage=dependencies.artifacts or artifact_store(), + persistence=persistence, + artifact_storage=dependencies.artifacts + or artifact_store(content_addressed=report_run is not None), child_invoker=dependencies.invoker or ModalChildInvoker(), dependencies=dependencies, stage_plan=( @@ -100,24 +125,63 @@ def coordinate_report( ), progress=StageProgress(), child_wait_deadline=child_wait_deadline, + report_run=report_run, ) with coordination.span(Stage.STAGE12_COORDINATOR_CLAIM): - parent, should_execute = claim_parent( - submitted=submitted_parent, - coordinator_invocation_id=invocation.invocation_id, - store=coordination.persistence, - ) - if not should_execute: - return { - "deduplicated": True, - "evaluation_id": str(parent.evaluation_id), - "status": parent.status.value, - } + parent, summary = claim_report(coordination, submitted=submitted_parent) + if parent is None: + return summary or {} try: return coordination.run(parent) # Persist a redacted cause for any coordinator, child-call, aggregation, or # artifact failure before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 # every failure is persisted before a stable exception returns to Modal - coordination.fail(parent, error) + _fail_unless_cancelled(coordination, parent, error) msg = "Stage 12 report coordination failed" raise RuntimeError(msg) from None + except KeyboardInterrupt: + # Modal pre-empts a call with SIGINT and runs it again under the same + # call ID. A report run's coordinator stops its unfinished children + # and records nothing, so the restarted call starts them again. + if report_run is not None: + coordination.stop_children() + raise + except BaseException: + # Modal cancels a call with InputCancellation, which is not an + # Exception, and never runs it again: a report run's coordinator + # records its report, report run and unfinished children as failed. + if report_run is not None: + _record_cancellation(coordination, parent) + raise + + +def _fail_unless_cancelled( + coordination: ReportCoordination, + parent: ComparisonReportRecord, + error: Exception, +) -> None: + """Record ``error``; if Modal cancels the call meanwhile, still fail the run.""" + try: + coordination.fail(parent, error) + except KeyboardInterrupt: + raise + except BaseException: + if coordination.report_run is not None: + coordination.report_run.fail_quietly( + CANCELLED_DETAIL, runtime=coordination.runtime + ) + raise + + +def _record_cancellation( + coordination: ReportCoordination, parent: ComparisonReportRecord +) -> None: + """Stop and fail the unfinished children, then fail the report and its run.""" + cancel_pending_simulations( + calls=coordination.calls, + simulations=coordination.simulations, + roles=set(coordination.calls) - set(coordination.descriptors), + store=coordination.persistence, + detail=cancellation("the coordinator's call was cancelled"), + ) + coordination.fail(parent, Stage12ExecutionError(CANCELLED_DETAIL)) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index 747eb839d..120d3946c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -177,6 +177,12 @@ def runtime_store( ) -def artifact_store() -> Stage12ArtifactStore: - """Build the store for ``STAGE12_ARTIFACT_BUCKET``; ValueError if it is unset.""" - return Stage12ArtifactStore(os.environ.get("STAGE12_ARTIFACT_BUCKET", "")) +def artifact_store(*, content_addressed: bool = False) -> Stage12ArtifactStore: + """Build the store for ``STAGE12_ARTIFACT_BUCKET``; ValueError if it is unset. + + A ``content_addressed`` store names each artifact's digest in its key. + """ + return Stage12ArtifactStore( + os.environ.get("STAGE12_ARTIFACT_BUCKET", ""), + content_addressed=content_addressed, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py index 966331f48..570810913 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py @@ -111,7 +111,9 @@ def run(self) -> SimulationCalculation: A segment whose result cannot be read twice in a row raises ``Stage12ExecutionError``, and a requested cancellation raises ``Stage12CancellationError``. Either way, every started segment is - cancelled first. + cancelled first. So it is when Modal pre-empts or cancels this worker, + which it signals with a ``BaseException`` that is not an ``Exception`` + and which is then re-raised: Modal does not cancel spawned calls. """ groups = self._groups() segments = build_segment_inputs(self.simulation, groups=groups) @@ -130,15 +132,23 @@ def run(self) -> SimulationCalculation: segment.model_dump(mode="python"), observability_context=self.runtime.capture_context(), ) - handles.append((segment.segment_index, call)) - except Exception: + # Kept inside the span, so an interruption as it closes + # still finds the call to cancel. + handles.append((segment.segment_index, call)) + except BaseException: self._cancel_all(handles) raise - with self.runtime.span( - STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) - ): - payloads = self._collect(handles) + try: + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) + ): + payloads = self._collect(handles) + except BaseException as interruption: + # Collection cancels the segments itself where a failure needs it. + if not isinstance(interruption, Exception): + self._cancel_all(handles) + raise with self.runtime.span( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) ): 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 index dba19ed4f..1f4a88265 100644 --- 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 @@ -124,8 +124,14 @@ def claim_parent( submitted: ComparisonReportRecord, coordinator_invocation_id: str, store: ComparisonStore, + resume_own_claim: bool = False, ) -> tuple[ComparisonReportRecord, bool]: - """Create the parent in Modal or claim a retry; reject duplicate execution.""" + """Create the parent in Modal or claim a retry; reject duplicate execution. + + With ``resume_own_claim``, a running parent that this coordinator call + claimed before is resumed: Modal runs a pre-empted call again under the + same call ID. + """ coordinating_at = datetime.now(UTC) candidate = submitted.model_copy( update={ @@ -146,6 +152,12 @@ def claim_parent( if parent.evaluation_id != submitted.evaluation_id: msg = "existing comparison identity has a different report identifier" raise ValueError(msg) + if ( + resume_own_claim + and parent.status is ComparisonRunLifecycleStatus.RUNNING + and parent.coordinator_invocation_id == coordinator_invocation_id + ): + return parent, True if parent.status in { ComparisonRunLifecycleStatus.RUNNING, ComparisonRunLifecycleStatus.SUCCEEDED, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py new file mode 100644 index 000000000..285dc0a34 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_claims.py @@ -0,0 +1,148 @@ +"""Claiming a call's Stage 12 report, and the report run its call names. + +A comparison run claims only its Stage 12 report, as before. A call that names +a report run claims that run first, under its Modal call ID, then its Stage 12 +report. Each claim either lets the call run the report or gives the summary +the call returns instead. +""" + +from __future__ import annotations + +from hashlib import sha256 +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_contract.stage12_execution import ( + AggregateReportArtifactPayload, + ComparisonRunLifecycleStatus, +) +from policyengine_stage12_persistence import RunStatus + +from .failures import Stage12FailureDetail +from .parent_claim import claim_parent + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_execution import ( + ComparisonReportRecord, + ) + + from .report_coordination import ReportCoordination + from .report_run_recording import ReportRunRecorder + + +def claim_report( + coordination: ReportCoordination, + *, + submitted: ComparisonReportRecord, +) -> tuple[ComparisonReportRecord | None, dict[str, Any] | None]: + """Return the Stage 12 report to run, or the summary to return instead. + + A report run another call holds, or one already finished, returns its + status. A call Modal restarted after its Stage 12 report succeeded, but + before its report run did, completes the report run from the stored + aggregate and returns that it succeeded. Once this call holds the report + run, any error or cancellation before the report starts fails the run, + then propagates; a pre-emption doesn't, because Modal runs the call again. + """ + report_run = coordination.report_run + if report_run is None: + parent, should_execute = claim_parent( + submitted=submitted, + coordinator_invocation_id=coordination.invocation.invocation_id, + store=coordination.persistence, + ) + return (parent, None) if should_execute else (None, _stage12_summary(parent)) + try: + claimed = report_run.claim() + if not claimed.changed: + return None, _report_run_summary(report_run, claimed.status) + return _claim_held_report(coordination, report_run, submitted=submitted) + except KeyboardInterrupt: + # Pre-emption: Modal runs the call again, and it resumes this claim. + raise + except BaseException as error: + # The claim may have committed before the error, as a cancellation + # during its commit; failing a run this call doesn't hold changes + # nothing. + report_run.fail_quietly(_claim_failure(error), runtime=coordination.runtime) + raise + + +def _claim_failure(error: BaseException) -> Stage12FailureDetail: + """Describe an error, or Modal's cancellation, that stopped the claim.""" + if not isinstance(error, Exception): + return Stage12FailureDetail( + error_code="report_cancelled", + error_summary="Modal cancelled the report's coordinator call.", + ) + return Stage12FailureDetail( + error_code="report_claim_failed", + error_summary=f"Claiming the Stage 12 report failed ({type(error).__name__}).", + ) + + +def _claim_held_report( + coordination: ReportCoordination, + report_run: ReportRunRecorder, + *, + submitted: ComparisonReportRecord, +) -> tuple[ComparisonReportRecord | None, dict[str, Any] | None]: + """Claim the Stage 12 report of a report run this call holds.""" + invocation_id = coordination.invocation.invocation_id + parent, should_execute = claim_parent( + submitted=submitted, + coordinator_invocation_id=invocation_id, + store=coordination.persistence, + resume_own_claim=True, + ) + if should_execute: + return parent, None + if ( + parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + and parent.coordinator_invocation_id == invocation_id + ): + coordination.complete_report_run(_stored_aggregate(coordination, parent)) + return None, _report_run_summary(report_run, RunStatus.SUCCEEDED) + report_run.fail( + Stage12FailureDetail( + error_code="report_already_claimed", + error_summary=( + f"This report's Stage 12 record is {parent.status.value} " + "under another call." + ), + ) + ) + return None, _stage12_summary(parent) + + +def _stored_aggregate( + coordination: ReportCoordination, parent: ComparisonReportRecord +) -> dict[str, Any]: + """Read and digest-check the aggregate a succeeded Stage 12 report names.""" + if parent.aggregate_output_uri is None or parent.aggregate_output_sha256 is None: + msg = "a succeeded Stage 12 report names no aggregate" + raise ValueError(msg) + payload = coordination.artifact_storage.read(parent.aggregate_output_uri) + if sha256(payload).hexdigest() != parent.aggregate_output_sha256: + msg = "stored aggregate artifact digest mismatch" + raise ValueError(msg) + return AggregateReportArtifactPayload.model_validate_json(payload).model_dump( + mode="json" + ) + + +def _report_run_summary( + report_run: ReportRunRecorder, status: RunStatus +) -> dict[str, Any]: + return { + "deduplicated": True, + "report_run_id": str(report_run.run_ids.report_run_id), + "status": status.value, + } + + +def _stage12_summary(parent: ComparisonReportRecord) -> dict[str, Any]: + return { + "deduplicated": True, + "evaluation_id": str(parent.evaluation_id), + "status": parent.status.value, + } 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 index 15c7ab7de..e19053376 100644 --- 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 @@ -25,6 +25,7 @@ Stage12OutputPlan, ) from policyengine_simulation_observability.stages import Stage, StagePlan +from policyengine_stage12_persistence import RunStatus from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, @@ -48,6 +49,7 @@ cancellation, child_failure_detail, child_record, + logger, record_child_failure, simulation_failure_attributes, ) @@ -70,6 +72,7 @@ resolve_report_output_plan, validate_output_frames, ) +from .report_run_recording import SimulationRunRecordingStore from .simulation import descriptor_from_record if TYPE_CHECKING: @@ -79,6 +82,9 @@ from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + from .report_run_access import ReportRunWriter + from .report_run_recording import ReportRunRecorder + SIMULATION_POLL_INITIAL_SECONDS = 0.25 SIMULATION_POLL_MAX_SECONDS = 1.0 SIMULATION_POLL_BACKOFF_FACTOR = 2.0 @@ -116,8 +122,9 @@ class CoordinatorDependencies: The defaults are the Stage 12 database store, the private artifact bucket, Modal child invocation, the standard aggregate builder and output - plan resolution, no telemetry runtime, the real clock and sleep, and the - standard child wait window. + plan resolution, no telemetry runtime, the real clock and sleep, the + standard child wait window, and the report run store, which only a call + that names a report run uses. """ store: ComparisonStore | None = None @@ -131,6 +138,7 @@ class CoordinatorDependencies: sleep: Callable[[float], None] = time.sleep monotonic: Callable[[], float] = time.monotonic wait_timeout_seconds: float = SIMULATION_WAIT_TIMEOUT_SECONDS + report_runs: ReportRunWriter | None = None @dataclass @@ -158,6 +166,8 @@ class ReportCoordination: descriptors: dict[str, SimulationArtifactDescriptor] = dataclasses.field( default_factory=dict ) + # The report run the call names, when it records one. + report_run: ReportRunRecorder | None = None @property def runtime(self) -> ObservabilityRuntime | None: @@ -271,7 +281,13 @@ def _start_children( 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) + if self.report_run is None: + calls[role] = self.child_invoker.restore(child.modal_invocation_id) + else: + # A report run's coordinator cancelled this child when it + # was pre-empted; make sure it stopped, then start it again. + self._cancel_earlier_child(child) + calls[role] = self._dispatch_child(child, simulation) else: calls[role] = self._dispatch_child(child, simulation) return descriptors, calls @@ -300,6 +316,38 @@ def _retained_descriptor( } ) + def _cancel_earlier_child(self, child: ComparisonSimulationRecord) -> None: + """Cancel the child call an earlier attempt of this call started.""" + if child.modal_invocation_id is None: + return + try: + self.child_invoker.restore(child.modal_invocation_id).cancel() + except Exception as error: # noqa: BLE001 # cancelling an earlier attempt's child is best effort; the new child starts regardless + logger.warning( + "Stage 12 could not cancel an earlier %s simulation call in report " + "%s (%s)", + child.role.value, + child.evaluation_id, + type(error).__name__, + ) + + def stop_children(self) -> None: + """Cancel every started, unfinished child call, recording nothing. + + Each child's record stays running, so the call Modal restarts starts + it again. A cancellation error is logged, never raised. + """ + for role in sorted(set(self.calls) - set(self.descriptors)): + try: + self.calls[role].cancel() + except Exception as error: # noqa: BLE001 # every child must be asked to stop; a failed cancellation is logged + logger.warning( + "Stage 12 could not cancel the %s simulation of report %s (%s)", + role, + self.report.evaluation_id, + type(error).__name__, + ) + def _dispatch_child( self, child: ComparisonSimulationRecord, @@ -340,7 +388,16 @@ def _dispatch_child( function_name=self.context.simulation_callable, environment=self.context.modal_environment, simulation=simulation.model_dump(mode="json"), - context=self.context.model_dump(mode="json"), + # A comparison run's children receive the context + # fields they received before report runs existed. + context=self.context.model_dump( + mode="json", + exclude=( + {"report_run"} + if self.context.report_run is None + else None + ), + ), observability_context=child_observability_context, ) ) @@ -561,6 +618,10 @@ def _aggregate_and_complete( } ) ) + # After the Stage 12 report: a call Modal restarts between the two + # writes finds that report succeeded and completes the report run + # from its stored aggregate. + self.complete_report_run(aggregate) compare_completed_report( parent=completed_parent, aggregate=aggregate, @@ -575,6 +636,70 @@ def _aggregate_and_complete( ) return descriptor.model_dump(mode="json") + def complete_report_run(self, aggregate: Mapping[str, Any]) -> None: + """Store the aggregate's result with the report run's success. + + A report run this call no longer holds is logged and left as it is. + A call that names no report run records nothing. Both simulation runs + are given their child records' outcomes first, so a succeeded report + run never combines a simulation run still running. + """ + if self.report_run is None: + return + self.reconcile_simulation_runs() + written = self.report_run.complete( + aggregate["result"], country=self.report.baseline.geography.country + ) + if not written.changed and written.status is not RunStatus.SUCCEEDED: + logger.warning( + "Stage 12 report %s finished, but its report run %s is %s", + self.report.evaluation_id, + self.report_run.run_ids.report_run_id, + written.status.value, + ) + + def reconcile_simulation_runs(self) -> None: + """Give each simulation run of the report its child record's outcome.""" + if not isinstance(self.persistence, SimulationRunRecordingStore): + return + for simulation in (self.report.baseline, self.report.reform): + self.persistence.reconcile_simulation_run( + simulation.simulation_execution_id + ) + + def fail_running_children(self, detail: Stage12FailureDetail) -> None: + """Fail every child record of a report run that is still running. + + A restarted call that fails before it starts a child again leaves + that child's record, and so its simulation run, running under the + call its pre-empted attempt cancelled. Each write is conditional on + the record still running, and is mirrored onto the simulation run. + """ + if not isinstance(self.persistence, SimulationRunRecordingStore): + return + failed_at = datetime.now(UTC) + for simulation in (self.report.baseline, self.report.reform): + try: + record = self.persistence.get_simulation( + simulation.simulation_execution_id + ) + except LookupError: + continue + if record.status is not ComparisonRunLifecycleStatus.RUNNING: + continue + self.persistence.replace_simulation_if_status( + record.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": failed_at, + "completed_at": failed_at, + } + ), + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: """Record a failed run once, stop its unfinished children, fail the parent. @@ -583,7 +708,10 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: 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. + is then stopped so none keeps running for a failed report. A call that + names a report run then fails the report's child records still + running, gives both simulation runs their records' outcomes, and fails + the report run, after the Stage 12 report. """ if isinstance(error, Stage12ExecutionError): detail = error.detail @@ -624,3 +752,24 @@ def fail(self, parent: ComparisonReportRecord, error: Exception) -> None: runtime=self.runtime, failure_scope="report_coordination", ) + # After the Stage 12 report, as on success: a call Modal restarts + # between the two writes finds that report failed and runs it again. + if self.report_run is not None: + persist_failure_record( + partial( + self.fail_running_children, + cancellation("the report failed"), + ), + runtime=self.runtime, + failure_scope="report_coordination", + ) + persist_failure_record( + self.reconcile_simulation_runs, + runtime=self.runtime, + failure_scope="report_coordination", + ) + persist_failure_record( + partial(self.report_run.fail, detail), + runtime=self.runtime, + failure_scope="report_coordination", + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py new file mode 100644 index 000000000..dad36dd64 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py @@ -0,0 +1,107 @@ +"""The Modal worker role's database credentials, which report runs are recorded with. + +An operator creates the Modal secret in each Modal environment. The v2 Modal +application attaches it to the coordinator and the simulation workers, which +record the status and results of the report run and simulation runs their +call names. The secret must exist in a Modal environment before the v2 +application deploys there: ``modal deploy`` checks its required key and fails +without it. +""" + +from __future__ import annotations + +import os +from functools import partial +from typing import TYPE_CHECKING, Protocol + +from policyengine_stage12_persistence import ReportRunStore + +from .database_retry import open_connection_with_retry + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.stage12_execution import ReportRunIds + from policyengine_stage12_persistence import ReportRunResult, RunWrite + +# The Modal secret and the one key it must hold: a postgresql:// URL for the +# Modal worker role on the database that holds the Public API's run tables. +MODAL_WORKER_DATABASE_SECRET_NAME = "policyengine-modal-worker-database" # noqa: S105 # the Modal secret's name, not its value +MODAL_WORKER_DATABASE_URL_ENV = "MODAL_WORKER_DATABASE_URL" + + +class ReportRunWriter(Protocol): + """Report run and simulation run writes; ``ReportRunStore`` implements it. + + Each write is conditional on the run's status and returns whether it + changed the run, and the run's status afterwards. + """ + + def require_simulation_runs(self, run_ids: ReportRunIds) -> None: + """Raise ValueError unless the report run combines both simulation runs.""" + ... + + def claim_report_run( + self, report_run_id: UUID, *, modal_call_id: str, claimed_at: datetime + ) -> RunWrite: + """Mark a pending report run, or one no other call holds, running.""" + ... + + def complete_report_run( + self, + run_ids: ReportRunIds, + *, + modal_call_id: str, + completed_at: datetime, + result: ReportRunResult | None = None, + ) -> RunWrite: + """Mark the report run succeeded with ``result`` while this call holds it.""" + ... + + def fail_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark the report run failed while ``modal_call_id`` holds it.""" + ... + + def claim_simulation_run( + self, simulation_run_id: UUID, *, modal_call_id: str, claimed_at: datetime + ) -> RunWrite: + """Mark a pending or running simulation run running under this call.""" + ... + + def complete_simulation_run( + self, + simulation_run_id: UUID, + *, + output_uri: str, + output_sha256: str, + completed_at: datetime, + ) -> RunWrite: + """Record a running simulation run's output file and mark it succeeded.""" + ... + + def fail_simulation_run( + self, simulation_run_id: UUID, *, error_message: str, failed_at: datetime + ) -> RunWrite: + """Mark a pending or running simulation run failed.""" + ... + + +def report_run_store(runtime: ObservabilityRuntime | None = None) -> ReportRunStore: + """Build the report run store from ``MODAL_WORKER_DATABASE_URL``. + + Raises ``ValueError`` when the variable is unset or not a complete + ``postgresql://`` URL. + """ + return ReportRunStore( + os.environ.get(MODAL_WORKER_DATABASE_URL_ENV, ""), + open_connection=partial(open_connection_with_retry, runtime=runtime), + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py new file mode 100644 index 000000000..f5be374b2 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_recording.py @@ -0,0 +1,303 @@ +"""Recording a report call's report run and simulation runs beside its Stage 12 records. + +A call that names a report run keeps the comparison path's Stage 12 records, +which coordinate its workers. ``SimulationRunRecordingStore`` mirrors every +conditional claim, success and failure of a child record onto that child's +simulation run, as the Modal worker role; ``ReportRunRecorder`` holds the +coordinator's writes to the report run. A mirrored write can be lost, so +before it completes or fails a report run the coordinator also gives both +simulation runs their child records' outcomes, and a failed report fails the +child records still running. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from functools import partial +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from policyengine_stage12_persistence import ( + DistrictCodes, + RunStatus, + report_run_result, +) + +from .failure_recording import persist_failure_record + +if TYPE_CHECKING: + from collections.abc import Mapping + from typing import Any + from uuid import UUID + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.stage12_execution import ( + ComparisonReportPersistenceResult, + ComparisonReportRecord, + ComparisonSimulationPersistenceResult, + ComparisonSimulationRecord, + ReportRunIds, + ) + from policyengine_stage12_persistence import RunWrite + + from .dependencies import ComparisonStore + from .failures import Stage12FailureDetail + from .report_run_access import ReportRunWriter + + +logger = logging.getLogger(__name__) + + +def us_district_codes() -> DistrictCodes: + """Return .py's US state FIPS codes and at-large states, for district rows.""" + # .py's US data loads only for a US report's result. + from policyengine.countries.us.data import ( + AT_LARGE_STATES, + US_STATE_FIPS, + ) + + return DistrictCodes( + state_fips={state: int(fips) for state, fips in US_STATE_FIPS.items()}, + at_large_states=frozenset(AT_LARGE_STATES), + ) + + +def run_error_message(error_code: str | None, error_summary: str | None) -> str: + """Return the redacted failure a run records: its code, then its summary.""" + return f"{error_code or 'simulation_failed'}: {error_summary or 'no summary'}" + + +@dataclass(frozen=True) +class ReportRunRecorder: + """The coordinator's writes to the report run its call names, under its call ID.""" + + writer: ReportRunWriter + run_ids: ReportRunIds + modal_call_id: str + + def claim(self) -> RunWrite: + """Check the run combines the call's simulation runs, then claim it.""" + self.writer.require_simulation_runs(self.run_ids) + return self.writer.claim_report_run( + self.run_ids.report_run_id, + modal_call_id=self.modal_call_id, + claimed_at=datetime.now(UTC), + ) + + def complete(self, economy_output: Mapping[str, Any], *, country: str) -> RunWrite: + """Mark the report run succeeded with its economy output as result rows. + + The rows are written in the same transaction, only while this call + holds the run. An output that cannot be stored raises ValueError. + """ + return self.writer.complete_report_run( + self.run_ids, + modal_call_id=self.modal_call_id, + completed_at=datetime.now(UTC), + result=report_run_result( + economy_output, + run_ids=self.run_ids, + country=country, + district_codes=us_district_codes() if country == "us" else None, + ), + ) + + def fail_quietly( + self, detail: Stage12FailureDetail, *, runtime: ObservabilityRuntime | None + ) -> None: + """Fail the report run, logging rather than raising a failed write.""" + persist_failure_record( + partial(self.fail, detail), runtime=runtime, failure_scope="report_run" + ) + + def fail(self, detail: Stage12FailureDetail) -> RunWrite: + """Mark the report run failed with ``detail`` while this call holds it.""" + return self.writer.fail_report_run( + self.run_ids.report_run_id, + modal_call_id=self.modal_call_id, + error_message=run_error_message(detail.error_code, detail.error_summary), + failed_at=datetime.now(UTC), + ) + + +class SimulationRunRecordingStore: + """A ``ComparisonStore`` that also records each child's simulation run. + + ``modal_call_id`` is the worker's own call, recorded on its simulation run + when it claims it; the coordinator never claims one and passes None. + """ + + def __init__( + self, + store: ComparisonStore, + runs: ReportRunWriter, + run_ids: ReportRunIds, + *, + modal_call_id: str | None = None, + ) -> None: + """Record on ``run_ids``'s simulation runs through ``runs``.""" + self._store = store + self._runs = runs + self._run_ids = run_ids + self._modal_call_id = modal_call_id + + def create_or_resolve_report( + self, record: ComparisonReportRecord + ) -> ComparisonReportPersistenceResult: + """Delegate to the Stage 12 store.""" + return self._store.create_or_resolve_report(record) + + def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Delegate to the Stage 12 store.""" + return self._store.get_report(evaluation_id) + + def replace_report(self, record: ComparisonReportRecord) -> ComparisonReportRecord: + """Delegate to the Stage 12 store.""" + return self._store.replace_report(record) + + def replace_report_result_comparison( + self, record: ComparisonReportRecord + ) -> ComparisonReportRecord: + """Delegate to the Stage 12 store.""" + return self._store.replace_report_result_comparison(record) + + def create_or_resolve_simulation( + self, record: ComparisonSimulationRecord + ) -> ComparisonSimulationPersistenceResult: + """Delegate to the Stage 12 store.""" + return self._store.create_or_resolve_simulation(record) + + def get_simulation( + self, simulation_execution_id: UUID + ) -> ComparisonSimulationRecord: + """Delegate to the Stage 12 store.""" + return self._store.get_simulation(simulation_execution_id) + + def replace_simulation( + self, record: ComparisonSimulationRecord + ) -> ComparisonSimulationRecord: + """Delegate to the Stage 12 store; a dispatch placeholder records nothing.""" + return self._store.replace_simulation(record) + + def attach_simulation_invocation( + self, + simulation_execution_id: UUID, + *, + expected_placeholder: str, + modal_invocation_id: str, + updated_at: datetime, + ) -> ComparisonSimulationRecord: + """Delegate to the Stage 12 store.""" + return self._store.attach_simulation_invocation( + simulation_execution_id, + expected_placeholder=expected_placeholder, + modal_invocation_id=modal_invocation_id, + updated_at=updated_at, + ) + + def replace_simulation_if_status( + self, + record: ComparisonSimulationRecord, + *, + expected_status: ComparisonRunLifecycleStatus, + ) -> tuple[ComparisonSimulationRecord, bool]: + """Write the child record, then its simulation run once the record changed. + + A worker Modal restarts after its record succeeded, but before its + simulation run did, completes the run from that record (see + ``complete_from_succeeded_child``), so both name the same output + file. A simulation run that is already failed refuses the child's + claim, which its worker treats as a cancellation. + """ + simulation_run_id = self._run_ids.simulation_run_id(record.role) + written_at = record.updated_at or datetime.now(UTC) + stored, replaced = self._store.replace_simulation_if_status( + record, expected_status=expected_status + ) + if not replaced: + return stored, replaced + if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: + # The child's result is recorded; if this second write is lost, + # the coordinator completes the run from the child's record. + try: + self._complete(simulation_run_id, record, written_at) + except Exception as error: # noqa: BLE001 # a lost second write is completed again by the coordinator + logger.warning( + "Could not complete simulation run %s (%s)", + simulation_run_id, + type(error).__name__, + ) + elif record.status is ComparisonRunLifecycleStatus.RUNNING: + claimed = self._runs.claim_simulation_run( + simulation_run_id, + modal_call_id=self._require_modal_call_id(), + claimed_at=written_at, + ) + if claimed.status is RunStatus.FAILED: + return stored, False + elif record.status is ComparisonRunLifecycleStatus.FAILED: + self._runs.fail_simulation_run( + simulation_run_id, + error_message=run_error_message( + record.error_code, record.error_summary + ), + failed_at=record.completed_at or written_at, + ) + return stored, replaced + + def reconcile_simulation_run(self, simulation_execution_id: UUID) -> None: + """Give a child's simulation run its finished Stage 12 record's outcome. + + A worker or coordinator that lost its second write leaves the run + behind its record; each write here is conditional, so a run that + already has its outcome is left as it is. An unfinished child, or one + never created, changes nothing. + """ + try: + record = self._store.get_simulation(simulation_execution_id) + except LookupError: + return + if record.status is ComparisonRunLifecycleStatus.SUCCEEDED: + self.complete_from_succeeded_child(record) + elif record.status is ComparisonRunLifecycleStatus.FAILED: + self._runs.fail_simulation_run( + self._run_ids.simulation_run_id(record.role), + error_message=run_error_message( + record.error_code, record.error_summary + ), + failed_at=record.completed_at or datetime.now(UTC), + ) + + def complete_from_succeeded_child(self, record: ComparisonSimulationRecord) -> None: + """Complete a child's simulation run from its succeeded Stage 12 record.""" + self._complete( + self._run_ids.simulation_run_id(record.role), + record, + record.updated_at or datetime.now(UTC), + ) + + def _complete( + self, + simulation_run_id: UUID, + record: ComparisonSimulationRecord, + written_at: datetime, + ) -> RunStatus: + if record.output_uri is None or record.output_sha256 is None: + msg = "a succeeded simulation must name its output file" + raise ValueError(msg) + return self._runs.complete_simulation_run( + simulation_run_id, + output_uri=record.output_uri, + output_sha256=record.output_sha256, + completed_at=record.completed_at or written_at, + ).status + + def _require_modal_call_id(self) -> str: + if self._modal_call_id is None: + msg = "only a worker with its Modal call ID claims a simulation run" + raise ValueError(msg) + return self._modal_call_id diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 1975fc8e0..bbed1923a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -57,6 +57,8 @@ ) from .output_planning import apply_output_plan, validate_output_frames from .partition import should_segment_simulation +from .report_run_access import report_run_store +from .report_run_recording import SimulationRunRecordingStore if TYPE_CHECKING: from contextlib import AbstractContextManager @@ -74,6 +76,8 @@ RegionResolution, ) + from .report_run_access import ReportRunWriter + logger = logging.getLogger(__name__) @@ -473,7 +477,8 @@ class SimulationWorkerDependencies: The defaults are the Stage 12 database store, the private artifact bucket, the standard country-package calculation, the deployed region-group - segment runner and no telemetry runtime. + segment runner, no telemetry runtime, and the report run store, which only + a call that names a report run uses. """ store: ComparisonStore | None = None @@ -490,6 +495,7 @@ class SimulationWorkerDependencies: | None ) = None runtime: ObservabilityRuntime | None = None + report_runs: ReportRunWriter | None = None def run_single_simulation( @@ -498,12 +504,16 @@ def run_single_simulation( *, required_country: CountryId, dependencies: SimulationWorkerDependencies | None = None, + modal_call_id: str | None = None, ) -> dict[str, Any]: """Calculate one planned simulation, store its artifact, and return its descriptor. A child already marked succeeded returns its stored descriptor without recalculating. A failure after the run is marked running marks it FAILED - and raises a ``RuntimeError`` that carries no calculation detail. + and raises a ``RuntimeError`` that carries no calculation detail. A call + whose context names a report run also records its simulation run's + claim, output file and failure under ``modal_call_id``, and names its + artifacts by their digest, so a call Modal restarts can store them again. """ dependencies = dependencies or SimulationWorkerDependencies() runtime = dependencies.runtime @@ -511,9 +521,36 @@ def run_single_simulation( context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) persistence = dependencies.store or runtime_store(runtime) - artifact_storage = dependencies.artifacts or artifact_store() + recording = None + if context.report_run is not None: + if modal_call_id is None: + msg = "a report run's worker needs its Modal call ID" + raise ValueError(msg) + recording = SimulationRunRecordingStore( + persistence, + dependencies.report_runs or report_run_store(runtime), + context.report_run, + modal_call_id=modal_call_id, + ) + persistence = recording + artifact_storage = dependencies.artifacts or artifact_store( + content_addressed=context.report_run is not None + ) child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: + if recording is not None: + # An earlier attempt stopped between the child's and the run's + # success; the run's completion is conditional, so this repeats + # nothing that was recorded, and if it fails the coordinator + # completes the run from the child's record. + try: + recording.complete_from_succeeded_child(child) + except Exception as error: # noqa: BLE001 # the coordinator completes the run again + logger.warning( + "Could not complete the simulation run of %s (%s)", + simulation.simulation_execution_id, + type(error).__name__, + ) return descriptor_from_record(child, simulation).model_dump(mode="json") running = _claim_simulation(persistence, child) cancellation_requested = _cancellation_check(persistence, simulation) diff --git a/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py b/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py new file mode 100644 index 000000000..90b0fd03a --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/stage12_fake_report_runs.py @@ -0,0 +1,158 @@ +"""An in-memory stand-in for the report run store, for Stage 12 runtime tests. + +It applies the same conditions as ``ReportRunStore``'s conditional updates, +which the persistence library tests against Postgres. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from uuid import UUID + +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_stage12_persistence import RunStatus, RunWrite +from stage12_runtime_support import aggregate_payload, invocation_context + +REPORT_RUN_ID = UUID("00000000-0000-0000-0000-0000000000a1") +BASELINE_RUN_ID = UUID("00000000-0000-0000-0000-0000000000b1") +REFORM_RUN_ID = UUID("00000000-0000-0000-0000-0000000000c1") +REPORT_RUN_IDS = ReportRunIds( + report_run_id=REPORT_RUN_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_run_id=REFORM_RUN_ID, +) +ACTIVE = (RunStatus.PENDING, RunStatus.RUNNING) + + +def report_run_context(): + """Return the shared invocation context, naming the report run and its runs.""" + return invocation_context().model_copy(update={"report_run": REPORT_RUN_IDS}) + + +def economy_aggregate(**_values): + """Aggregate the report into a complete single-year economy output.""" + return aggregate_payload(result=CURRENT_SINGLE_YEAR_MACRO_RESULT) + + +@dataclass +class FakeRun: + """One report run's or simulation run's mutable columns.""" + + status: RunStatus = RunStatus.PENDING + modal_call_id: str | None = None + error_message: str | None = None + output_uri: str | None = None + output_sha256: str | None = None + result: object | None = None + writes: list[str] = field(default_factory=list) + + +class FakeReportRuns: + """Report runs and simulation runs held in memory, written conditionally.""" + + def __init__(self, events: list[str] | None = None) -> None: + self.report_runs = {REPORT_RUN_ID: FakeRun()} + self.simulation_runs = {BASELINE_RUN_ID: FakeRun(), REFORM_RUN_ID: FakeRun()} + self.links = {("baseline", BASELINE_RUN_ID), ("reform", REFORM_RUN_ID)} + # Shared with other fakes when a test checks the order of writes. + self.events = events if events is not None else [] + + def require_simulation_runs(self, run_ids: ReportRunIds) -> None: + expected = { + ("baseline", run_ids.baseline_simulation_run_id), + ("reform", run_ids.reform_simulation_run_id), + } + if not expected <= self.links: + msg = "the report run does not combine the simulation runs its call names" + raise ValueError(msg) + + def claim_report_run(self, report_run_id, *, modal_call_id, claimed_at): + run = self._run(self.report_runs, report_run_id, "report run") + claimable = run.status is RunStatus.PENDING or ( + run.status is RunStatus.RUNNING + and run.modal_call_id in {None, modal_call_id} + ) + return self._apply( + run, + "claim_report_run", + changed=claimable, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + ) + + def complete_report_run(self, run_ids, *, modal_call_id, completed_at, result=None): + run = self._run(self.report_runs, run_ids.report_run_id, "report run") + holds_it = self._holds(run, modal_call_id) + return self._apply( + run, + "complete_report_run", + changed=holds_it, + status=RunStatus.SUCCEEDED, + result=result, + ) + + def fail_report_run( + self, report_run_id, *, modal_call_id, error_message, failed_at + ): + run = self._run(self.report_runs, report_run_id, "report run") + holds_it = self._holds(run, modal_call_id) + return self._apply( + run, + "fail_report_run", + changed=holds_it, + status=RunStatus.FAILED, + error_message=error_message, + ) + + def claim_simulation_run(self, simulation_run_id, *, modal_call_id, claimed_at): + run = self._run(self.simulation_runs, simulation_run_id, "simulation run") + return self._apply( + run, + "claim_simulation_run", + changed=run.status in ACTIVE, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + ) + + def complete_simulation_run( + self, simulation_run_id, *, output_uri, output_sha256, completed_at + ): + run = self._run(self.simulation_runs, simulation_run_id, "simulation run") + return self._apply( + run, + "complete_simulation_run", + changed=run.status is RunStatus.RUNNING, + status=RunStatus.SUCCEEDED, + output_uri=output_uri, + output_sha256=output_sha256, + ) + + def fail_simulation_run(self, simulation_run_id, *, error_message, failed_at): + run = self._run(self.simulation_runs, simulation_run_id, "simulation run") + return self._apply( + run, + "fail_simulation_run", + changed=run.status in ACTIVE, + status=RunStatus.FAILED, + error_message=error_message, + ) + + @staticmethod + def _holds(run: FakeRun, modal_call_id: str) -> bool: + return run.status is RunStatus.RUNNING and run.modal_call_id == modal_call_id + + @staticmethod + def _run(runs: dict[UUID, FakeRun], run_id: UUID, noun: str) -> FakeRun: + if run_id not in runs: + msg = f"{noun} {run_id} does not exist" + raise LookupError(msg) + return runs[run_id] + + def _apply(self, run: FakeRun, operation: str, *, changed: bool, **values): + run.writes.append(operation) + self.events.append(operation) + if changed: + for name, value in values.items(): + setattr(run, name, value) + return RunWrite(changed=changed, status=run.status) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index e511ee05e..f81aa8868 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -88,6 +88,15 @@ def model_dump(self, *, mode): observability_context=_observability_context(), ) + # The US and UK workers pass their own Modal call ID to the runtime, so the + # stand-in ``modal`` module needs one; a test that set its own keeps it. + if not hasattr(module.modal, "current_function_call_id"): + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "call-1", + raising=False, + ) monkeypatch.setattr(stage12_runtime, "run_single_simulation", implementation) return getattr(module, function_name)( {"role": "baseline"}, @@ -344,6 +353,60 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert "asgi_app" not in vars(module) +def test_only_report_run_writers_receive_the_modal_worker_database_secret( + monkeypatch, +) -> None: + module = _load(monkeypatch) + functions = dict(module.app.function_calls) + worker_database_secret = { + "args": ("policyengine-modal-worker-database",), + "kwargs": {"required_keys": ["MODAL_WORKER_DATABASE_URL"]}, + } + + assert module.modal_worker_database_secret == worker_database_secret + for writer in ( + "coordinate_report", + "run_single_simulation_us", + "run_single_simulation_uk", + ): + assert functions[writer]["secrets"] == [ + *module.worker_secrets, + worker_database_secret, + ] + # Segments return their result to their worker, and validation records + # nothing, so neither holds the database credentials. + for other in ( + "run_single_simulation_segment_us", + "validate_worker_us", + "validate_worker_uk", + ): + assert functions[other]["secrets"] == module.worker_secrets + + +@pytest.mark.parametrize( + "function_name", ["run_single_simulation_us", "run_single_simulation_uk"] +) +def test_simulation_workers_pass_their_own_modal_call_id( + monkeypatch, function_name +) -> None: + module = _load(monkeypatch) + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "fc-simulation-worker", + raising=False, + ) + received = {} + + def implementation(*_args, **kwargs): + received.update(kwargs) + return {"role": "baseline"} + + _invoke(monkeypatch, module, function_name, implementation) + + assert received["modal_call_id"] == "fc-simulation-worker" + + @pytest.mark.parametrize( "function_name", [ diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py new file mode 100644 index 000000000..3957db241 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_artifacts.py @@ -0,0 +1,81 @@ +"""Artifact keys that name their content, for calls Modal may restart.""" + +from __future__ import annotations + +from hashlib import sha256 + +import pytest +from fixtures.fake_object_store import InMemoryObjectStore +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore +from stage12_fake_stores import simulation_frames +from stage12_runtime_support import aggregate_payload, planned_simulation + +PREFIX = "stage-12-runs/staging/2026/11/report-run" + + +def _store(*, content_addressed: bool = True) -> Stage12ArtifactStore: + return Stage12ArtifactStore( + "private-stage12", + store=InMemoryObjectStore(), + content_addressed=content_addressed, + ) + + +def test_content_named_simulation_outputs_of_one_role_are_kept_apart() -> None: + store = _store() + baseline = planned_simulation(SimulationRole.BASELINE) + + first = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.0) + ) + # A restarted call can calculate different last bits for the same role. + second = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.000001) + ) + repeated = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.0) + ) + + digest = first.artifact.content_sha256 + assert first.artifact.uri == ( + f"gs://private-stage12/{PREFIX}/simulations/baseline-{digest}.parquet" + ) + assert second.artifact.uri != first.artifact.uri + assert repeated.artifact == first.artifact + assert store.read(second.artifact.uri) != store.read(first.artifact.uri) + + +def test_content_named_inputs_and_aggregates_name_their_digest() -> None: + store = _store() + reform = planned_simulation(SimulationRole.REFORM) + + planned_input = store.write_input(prefix=PREFIX, simulation=reform) + aggregate = store.write_aggregate(prefix=PREFIX, payload=aggregate_payload()) + + assert planned_input.uri == ( + f"gs://private-stage12/{PREFIX}/inputs/reform-" + f"{planned_input.content_sha256}.json" + ) + assert aggregate.uri == ( + f"gs://private-stage12/{PREFIX}/reports/aggregate-" + f"{aggregate.content_sha256}.json" + ) + assert sha256(store.read(aggregate.uri)).hexdigest() == aggregate.content_sha256 + + +def test_comparison_runs_keep_their_fixed_write_once_keys() -> None: + store = _store(content_addressed=False) + baseline = planned_simulation(SimulationRole.BASELINE) + + first = store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.0) + ) + + assert first.artifact.uri == ( + f"gs://private-stage12/{PREFIX}/simulations/baseline.parquet" + ) + with pytest.raises(RuntimeError, match="other data"): + store.write_simulation( + prefix=PREFIX, simulation=baseline, frames=simulation_frames(100.000001) + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py new file mode 100644 index 000000000..5217f3277 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_coordinator.py @@ -0,0 +1,447 @@ +"""The Stage 12 coordinator records the status of the report run its call names.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor + +import pytest +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationWorkerDependencies, + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from policyengine_stage12_persistence import RunStatus +from stage12_child_call_support import ( + ConcurrentInvoker, + FutureCall, + InjectedChildFaults, +) +from stage12_fake_report_runs import ( + BASELINE_RUN_ID, + REFORM_RUN_ID, + REPORT_RUN_ID, + FakeReportRuns, + economy_aggregate, + report_run_context, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + child_record, + coordinate_report, + parent_record, + planned_simulation, + report_execution_input, +) + +COORDINATOR_CALL_ID = "fc-coordinator" + + +class RecordingInvoker(ConcurrentInvoker): + """Run both children concurrently, keeping the context each one received.""" + + def __init__(self, artifacts, **options): + super().__init__(artifacts, **options) + self.contexts = [] + + def spawn(self, request): + self.contexts.append(request.context) + return super().spawn(request) + + +def _coordinate(store, report_runs, invoker, artifacts, *, context=None): + context = context or report_run_context() + return coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=economy_aggregate, + report_runs=report_runs, + ) + + +def _recording_report_writes(store, events): + replace_report = store.replace_report + + def recording_replace_report(record): + events.append(f"stage12:{record.status.value}") + return replace_report(record) + + store.replace_report = recording_replace_report + return store + + +def test_coordinator_claims_its_report_run_first_and_completes_it_after_stage12() -> ( + None +): + events: list[str] = [] + store = _recording_report_writes(FakeStore(), events) + report_runs = FakeReportRuns(events) + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + result = _coordinate(store, report_runs, invoker, artifacts) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == COORDINATOR_CALL_ID + assert result["evaluation_id"] == str(report_execution_input().evaluation_id) + # The report run is claimed first and completed last, so a call Modal + # restarts between the two success writes finds its Stage 12 report + # succeeded, and completes the report run from the stored aggregate. + assert events == ["claim_report_run", "stage12:succeeded", "complete_report_run"] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert [context["report_run"] for context in invoker.contexts] == [ + report_run_context().report_run.model_dump(mode="json") + ] * 2 + + +def test_coordinator_completes_its_report_run_with_the_economy_output_as_rows() -> None: + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + + _coordinate(FakeStore(), report_runs, RecordingInvoker(artifacts), artifacts) + + result = report_runs.report_runs[REPORT_RUN_ID].result + rows = {rows.table.name: rows.rows for rows in result.rows} + budget = {row["variable_name"]: row for row in rows["budget_summary"]} + assert budget["household_tax"]["change"] == 100.0 + assert budget["household_tax"]["report_run_id"] == REPORT_RUN_ID + # The shared report is a US one, so its poverty rows use the SPM types. + assert {row["poverty_type"] for row in rows["poverty"]} == {"spm", "spm_deep"} + (comparison,) = rows["report_run_comparisons"] + assert comparison["report_run_id"] == REPORT_RUN_ID + assert ( + comparison["labor_supply_result"] + == (CURRENT_SINGLE_YEAR_MACRO_RESULT["labor_supply_response"]) + ) + # Alabama's first district, by .py's state FIPS codes. + (district,) = rows["congressional_district_impacts"] + assert (district["state_fips"], district["district_number"]) == (1, 1) + assert district["district_geoid"] == 101 + + +def test_comparison_children_receive_the_context_they_received_before() -> None: + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + report_runs = FakeReportRuns() + + _coordinate( + FakeStore(), + report_runs, + invoker, + artifacts, + context=report_run_context().model_copy(update={"report_run": None}), + ) + + assert all("report_run" not in context for context in invoker.contexts) + assert report_runs.events == [] + + +@pytest.mark.parametrize( + ("status", "modal_call_id"), + [ + (RunStatus.RUNNING, "fc-another-coordinator"), + (RunStatus.SUCCEEDED, COORDINATOR_CALL_ID), + (RunStatus.FAILED, COORDINATOR_CALL_ID), + ], +) +def test_a_report_run_that_cannot_be_claimed_starts_nothing( + status, modal_call_id +) -> None: + store = FakeStore() + report_runs = FakeReportRuns() + run = report_runs.report_runs[REPORT_RUN_ID] + run.status, run.modal_call_id = status, modal_call_id + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + result = _coordinate(store, report_runs, invoker, artifacts) + + assert result == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": status.value, + } + assert invoker.events == [] + assert store.parent is None + assert run.writes == ["claim_report_run"] + + +def test_a_report_run_must_combine_the_simulation_runs_its_call_names() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + report_runs.links = {("baseline", REFORM_RUN_ID)} + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + with pytest.raises(ValueError, match="does not combine"): + _coordinate(store, report_runs, invoker, artifacts) + + assert invoker.events == [] + assert store.parent is None + assert report_runs.report_runs[REPORT_RUN_ID].status is RunStatus.PENDING + + +def test_a_failed_report_fails_its_report_run_and_the_failed_simulation_run() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + invoker = RecordingInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises(RuntimeError, match="coordination failed"): + _coordinate(store, report_runs, invoker, artifacts) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("simulation_invocation_failed: ") + reform = report_runs.simulation_runs[REFORM_RUN_ID] + assert reform.status is RunStatus.FAILED + assert reform.error_message == run.error_message + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + + +def test_a_failed_report_records_its_stage12_failure_before_its_report_run() -> None: + events: list[str] = [] + store = _recording_report_writes(FakeStore(), events) + report_runs = FakeReportRuns(events) + artifacts = FakeArtifacts() + invoker = RecordingInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises(RuntimeError, match="coordination failed"): + _coordinate(store, report_runs, invoker, artifacts) + + assert [ + event for event in events if event in {"stage12:failed", "fail_report_run"} + ] == ["stage12:failed", "fail_report_run"] + + +def test_a_stage12_claim_that_fails_fails_the_report_run() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + def unreachable(_record): + msg = "the Stage 12 database refused the connection" + raise ConnectionError(msg) + + store.create_or_resolve_report = unreachable + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + with pytest.raises(ConnectionError): + _coordinate(store, report_runs, invoker, artifacts) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_claim_failed: ") + assert "refused the connection" not in run.error_message + assert invoker.events == [] + + +def test_a_stage12_report_another_call_finished_fails_the_report_run() -> None: + store = FakeStore() + store.parent = parent_record().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "coordinator_invocation_id": "fc-another-coordinator", + } + ) + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + invoker = RecordingInvoker(artifacts) + + result = _coordinate(store, report_runs, invoker, artifacts) + + assert result["deduplicated"] is True + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message == ( + "report_already_claimed: This report's Stage 12 record is succeeded " + "under another call." + ) + assert invoker.events == [] + + +class LosingOneFailureWrite(FakeReportRuns): + """Lose the first write that fails a simulation run, as a dropped connection.""" + + def __init__(self): + super().__init__() + self.lost_failure_writes = 0 + + def fail_simulation_run(self, simulation_run_id, *, error_message, failed_at): + if not self.lost_failure_writes: + self.lost_failure_writes += 1 + msg = "the connection dropped" + raise ConnectionError(msg) + return super().fail_simulation_run( + simulation_run_id, error_message=error_message, failed_at=failed_at + ) + + +def test_a_lost_simulation_run_failure_is_written_when_the_report_fails() -> None: + store = FakeStore() + report_runs = LosingOneFailureWrite() + artifacts = FakeArtifacts() + invoker = RecordingInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises(RuntimeError, match="coordination failed"): + _coordinate(store, report_runs, invoker, artifacts) + + assert report_runs.lost_failure_writes == 1 + reform = report_runs.simulation_runs[REFORM_RUN_ID] + assert reform.status is RunStatus.FAILED + assert reform.error_message.startswith("simulation_invocation_failed: ") + + +def test_a_reused_succeeded_child_completes_its_simulation_run() -> None: + # A worker marked its child record succeeded, then stopped before it + # completed its simulation run; a later call reuses that child's output. + store = FakeStore() + artifacts = FakeArtifacts() + baseline = planned_simulation(SimulationRole.BASELINE) + descriptor = artifacts.add_simulation(baseline) + store.children[baseline.simulation_execution_id] = child_record( + baseline + ).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": descriptor.artifact.uri, + "output_sha256": descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": descriptor.row_identity.identifier_columns, + "row_count": descriptor.row_identity.row_count, + "row_identity_sha256": descriptor.row_identity.identity_sha256, + "completed_at": child_record(baseline).created_at, + } + ) + report_runs = FakeReportRuns() + report_runs.simulation_runs[BASELINE_RUN_ID].status = RunStatus.RUNNING + invoker = RecordingInvoker(artifacts) + + _coordinate(store, report_runs, invoker, artifacts) + + assert invoker.events == ["spawn:reform"] + run = report_runs.simulation_runs[BASELINE_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.output_uri == descriptor.artifact.uri + # The simulation runs are recorded before the report run succeeds. + assert report_runs.events.index("complete_simulation_run") < ( + report_runs.events.index("complete_report_run") + ) + + +class WorkerInvoker: + """Run each child through the real worker code, sharing the coordinator's stores.""" + + def __init__(self, store, artifacts, report_runs): + self.store = store + self.artifacts = artifacts + self.report_runs = report_runs + self.executor = ThreadPoolExecutor(max_workers=2) + + def spawn(self, request): + role = request.simulation["role"] + frames = simulation_frames(100.0 if role == "baseline" else 120.0) + future = self.executor.submit( + run_single_simulation, + request.simulation, + request.context, + required_country="us", + modal_call_id=f"fc-{role}-worker", + dependencies=SimulationWorkerDependencies( + store=self.store, + artifacts=self.artifacts, + calculator=lambda _simulation: frames, + report_runs=self.report_runs, + ), + ) + return FutureCall(f"fc-{role}-worker", future) + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + +def test_a_report_run_through_real_workers_records_every_run_once() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + report_runs = FakeReportRuns() + + report = report_execution_input() + # A reform the real worker accepts: an effective period with dates. + report = report.model_copy( + update={ + "reform": report.reform.model_copy( + update={"policy": {"reform": {"2026-01-01.2026-12-31": 1}}} + ) + } + ) + + coordinate_report( + report.model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=report_run_context().modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=store, + artifacts=artifacts, + invoker=WorkerInvoker(store, artifacts, report_runs), + aggregator=economy_aggregate, + report_runs=report_runs, + ) + + assert report_runs.report_runs[REPORT_RUN_ID].status is RunStatus.SUCCEEDED + for role, run_id in (("baseline", BASELINE_RUN_ID), ("reform", REFORM_RUN_ID)): + run = report_runs.simulation_runs[run_id] + (child,) = [child for child in store.children.values() if child.role == role] + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == f"fc-{role}-worker" + assert run.output_uri == child.output_uri + assert run.writes.count("complete_simulation_run") >= 1 + assert report_runs.report_runs[REPORT_RUN_ID].writes == [ + "claim_report_run", + "complete_report_run", + ] + + +def test_report_run_coordinator_names_its_aggregate_by_digest(monkeypatch) -> None: + built = [] + artifacts = FakeArtifacts() + + def recording_artifact_store(*, content_addressed=False): + built.append(content_addressed) + return artifacts + + monkeypatch.setattr(coordination_module, "artifact_store", recording_artifact_store) + + coordinate_report( + report_execution_input().model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=report_run_context().modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=FakeStore(), + invoker=RecordingInvoker(artifacts), + aggregator=economy_aggregate, + report_runs=FakeReportRuns(), + ) + + assert built == [True] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py new file mode 100644 index 000000000..87c3a028d --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_preemption.py @@ -0,0 +1,425 @@ +"""A report call's coordinator under pre-emption: stop the children, then resume. + +Modal pre-empts a call with SIGINT, which a synchronous function receives as +``KeyboardInterrupt``, and later runs the same call again under the same call +ID. The coordinator of a report run cancels its unfinished children, lets +Modal restart it, and the restarted call starts them again and records the +report run once. +""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from policyengine_stage12_persistence import RunStatus +from stage12_child_call_support import ( + ConcurrentInvoker, + InjectedChildFaults, + TimeoutInvoker, +) +from stage12_fake_report_runs import ( + REPORT_RUN_ID, + FakeReportRuns, + economy_aggregate, + report_run_context, +) +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( + coordinate_report, + invocation_context, + parent_record, + report_execution_input, +) + +COORDINATOR_CALL_ID = "fc-coordinator" + + +class CancelledCall: + def __init__(self, object_id): + self.object_id = object_id + self.cancelled = 0 + + def get(self, *, timeout=None): + msg = "a cancelled call has no result" + raise AssertionError(msg) + + def cancel(self): + self.cancelled += 1 + + +class RestartInvoker(ConcurrentInvoker): + """Run new children; restoring an earlier attempt's child only allows cancelling.""" + + def __init__(self, artifacts): + super().__init__(artifacts) + self.restored: dict[str, CancelledCall] = {} + + def restore(self, invocation_id): + self.events.append(f"restore:{invocation_id}") + call = CancelledCall(invocation_id) + self.restored[invocation_id] = call + return call + + +def _preempt(_seconds): + raise KeyboardInterrupt + + +def _coordinate(store, report_runs, invoker, *, context, **seams): + return coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id=COORDINATOR_CALL_ID, + store=store, + artifacts=seams.pop("artifacts", FakeArtifacts()), + invoker=invoker, + aggregator=economy_aggregate, + report_runs=report_runs, + **seams, + ) + + +def _preempted_first_attempt(store, report_runs, *, context): + invoker = TimeoutInvoker() + calls = [] + spawn = invoker.spawn + + def keeping_spawn(request): + call = spawn(request) + calls.append(call) + return call + + invoker.spawn = keeping_spawn + with pytest.raises(KeyboardInterrupt): + _coordinate(store, report_runs, invoker, context=context, sleep=_preempt) + return calls + + +def test_preempted_coordinator_cancels_its_children_and_records_no_failure() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + calls = _preempted_first_attempt(store, report_runs, context=report_run_context()) + + assert [call.cancelled for call in calls] == [1, 1] + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.RUNNING + assert run.writes == ["claim_report_run"] + assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING + assert all( + child.status is ComparisonRunLifecycleStatus.RUNNING + for child in store.children.values() + ) + assert all( + simulation_run.status is RunStatus.PENDING + for simulation_run in report_runs.simulation_runs.values() + ) + + +def test_restarted_coordinator_resumes_its_claim_and_starts_its_children_again() -> ( + None +): + store = FakeStore() + report_runs = FakeReportRuns() + _preempted_first_attempt(store, report_runs, context=report_run_context()) + artifacts = FakeArtifacts() + invoker = RestartInvoker(artifacts) + + result = _coordinate( + store, report_runs, invoker, context=report_run_context(), artifacts=artifacts + ) + + assert result["evaluation_id"] == str(report_execution_input().evaluation_id) + assert invoker.events == [ + "restore:call-baseline", + "spawn:baseline", + "restore:call-reform", + "spawn:reform", + ] + assert all(call.cancelled == 1 for call in invoker.restored.values()) + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.writes == ["claim_report_run", "claim_report_run", "complete_report_run"] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def _preempted_around_completion(report_runs, *, after_writing): + """Make the report run's completion pre-empt the call, before or after writing.""" + complete = report_runs.complete_report_run + + def preempted_completion(*args, **kwargs): + if after_writing: + complete(*args, **kwargs) + raise KeyboardInterrupt + + report_runs.complete_report_run = preempted_completion + return complete + + +def _first_attempt_preempted_at_completion(store, report_runs, *, after_writing): + complete = _preempted_around_completion(report_runs, after_writing=after_writing) + artifacts = FakeArtifacts() + with pytest.raises(KeyboardInterrupt): + _coordinate( + store, + report_runs, + ConcurrentInvoker(artifacts), + context=report_run_context(), + artifacts=artifacts, + ) + report_runs.complete_report_run = complete + return artifacts + + +def test_a_call_restarted_between_its_two_success_writes_completes_the_run() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + artifacts = _first_attempt_preempted_at_completion( + store, report_runs, after_writing=False + ) + restart_invoker = RestartInvoker(artifacts) + + result = _coordinate( + store, + report_runs, + restart_invoker, + context=report_run_context(), + artifacts=artifacts, + ) + + assert result == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": "succeeded", + } + # Nothing is calculated again: the result comes from the stored aggregate. + assert restart_invoker.events == [] + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + rows = {rows.table.name: rows.rows for rows in run.result.rows} + assert {row["variable_name"] for row in rows["budget_summary"]} >= {"household_tax"} + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def test_a_call_restarted_after_its_report_run_succeeded_records_nothing_again() -> ( + None +): + store = FakeStore() + report_runs = FakeReportRuns() + _first_attempt_preempted_at_completion(store, report_runs, after_writing=True) + restart_invoker = RestartInvoker(FakeArtifacts()) + + result = _coordinate( + store, report_runs, restart_invoker, context=report_run_context() + ) + + assert result == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": "succeeded", + } + assert restart_invoker.events == [] + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.writes == ["claim_report_run", "complete_report_run", "claim_report_run"] + # Both records agree: the Stage 12 report succeeded before the report run. + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +class InputCancelledForTest(BaseException): + """Stands in for Modal's InputCancellation, a BaseException but not an Exception.""" + + +def _cancel(_seconds): + raise InputCancelledForTest + + +def test_a_cancelled_coordinator_records_its_report_run_and_children_failed() -> None: + # Modal does not run a cancelled call again, so nothing would ever finish + # these runs if the coordinator left them running. + store = FakeStore() + report_runs = FakeReportRuns() + + with pytest.raises(InputCancelledForTest): + _coordinate( + store, + report_runs, + TimeoutInvoker(), + context=report_run_context(), + sleep=_cancel, + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_cancelled: ") + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + simulation_run.status is RunStatus.FAILED + for simulation_run in report_runs.simulation_runs.values() + ) + + +@pytest.mark.parametrize( + ("interruption", "expected_status"), + [ + # Modal runs a pre-empted call again, which resumes the claim. + (KeyboardInterrupt, RunStatus.RUNNING), + # Modal never runs a cancelled call again. + (InputCancelledForTest, RunStatus.FAILED), + ], +) +def test_a_coordinator_stopped_while_claiming_its_report( + interruption, expected_status +) -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + def stopped(_record): + raise interruption + + store.create_or_resolve_report = stopped + + with pytest.raises(interruption): + _coordinate( + store, + report_runs, + RestartInvoker(FakeArtifacts()), + context=report_run_context(), + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is expected_status + if expected_status is RunStatus.FAILED: + assert run.error_message.startswith("report_cancelled: ") + + +class FailingRestartInvoker(RestartInvoker): + """Restore an earlier attempt's children; refuse to start the baseline again.""" + + def spawn(self, request): + if request.simulation["role"] == "baseline": + msg = "Modal refused the spawn" + raise ConnectionError(msg) + return super().spawn(request) + + +class CancelledWhileRestoringInvoker(RestartInvoker): + """Receive Modal's cancellation while cancelling the earlier attempt's child.""" + + def restore(self, invocation_id): + raise InputCancelledForTest + + +def _claimed_by_the_preempted_attempt(report_runs): + # In production the pre-empted attempt's workers claimed their runs. + for simulation_run in report_runs.simulation_runs.values(): + simulation_run.status = RunStatus.RUNNING + + +@pytest.mark.parametrize( + ("restart_invoker", "raised"), + [ + (FailingRestartInvoker, RuntimeError), + (CancelledWhileRestoringInvoker, InputCancelledForTest), + ], +) +def test_a_restart_that_stops_before_restarting_its_children_fails_them( + restart_invoker, raised +) -> None: + store = FakeStore() + report_runs = FakeReportRuns() + _preempted_first_attempt(store, report_runs, context=report_run_context()) + _claimed_by_the_preempted_attempt(report_runs) + + with pytest.raises(raised): + _coordinate( + store, + report_runs, + restart_invoker(FakeArtifacts()), + context=report_run_context(), + ) + + assert report_runs.report_runs[REPORT_RUN_ID].status is RunStatus.FAILED + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + simulation_run.status is RunStatus.FAILED + for simulation_run in report_runs.simulation_runs.values() + ) + + +def test_a_cancellation_as_the_report_run_claim_commits_fails_the_run() -> None: + report_runs = FakeReportRuns() + claim = report_runs.claim_report_run + + def cancelled_after_commit(*args, **kwargs): + claim(*args, **kwargs) + raise InputCancelledForTest + + report_runs.claim_report_run = cancelled_after_commit + + with pytest.raises(InputCancelledForTest): + _coordinate( + FakeStore(), + report_runs, + RestartInvoker(FakeArtifacts()), + context=report_run_context(), + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_cancelled: ") + + +def test_a_cancellation_while_recording_a_failure_still_fails_the_run() -> None: + store = FakeStore() + replace_report = store.replace_report + + def cancelled_while_failing(record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + raise InputCancelledForTest + return replace_report(record) + + store.replace_report = cancelled_while_failing + report_runs = FakeReportRuns() + artifacts = FakeArtifacts() + + with pytest.raises(InputCancelledForTest): + _coordinate( + store, + report_runs, + ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ), + context=report_run_context(), + artifacts=artifacts, + ) + + run = report_runs.report_runs[REPORT_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("report_cancelled: ") + + +def test_a_preempted_comparison_run_behaves_as_before() -> None: + store = FakeStore() + report_runs = FakeReportRuns() + + calls = _preempted_first_attempt(store, report_runs, context=invocation_context()) + restart_invoker = RestartInvoker(FakeArtifacts()) + result = _coordinate( + store, report_runs, restart_invoker, context=invocation_context() + ) + + assert [call.cancelled for call in calls] == [0, 0] + assert result["deduplicated"] is True + assert restart_invoker.events == [] + assert report_runs.events == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py new file mode 100644 index 000000000..e6b792825 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_report_run_worker.py @@ -0,0 +1,282 @@ +"""A Stage 12 worker records its simulation run's status and output file.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationWorkerDependencies, + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + simulation as simulation_module, +) +from policyengine_stage12_persistence import RunStatus +from stage12_fake_report_runs import ( + BASELINE_RUN_ID, + FakeReportRuns, + report_run_context, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + child_record, + invocation_context, + planned_simulation, +) + + +def _seeded_store(simulation, events=None): + store = FakeStore() + store.children[simulation.simulation_execution_id] = child_record(simulation) + if events is not None: + replace = store.replace_simulation_if_status + + def recording_replace(record, *, expected_status): + events.append(f"stage12:{record.status.value}") + return replace(record, expected_status=expected_status) + + store.replace_simulation_if_status = recording_replace + return store + + +def _run_worker(simulation, store, report_runs, *, context=None, calculator=None): + return run_single_simulation( + simulation.model_dump(mode="json"), + (context or report_run_context()).model_dump(mode="json"), + required_country="us", + modal_call_id="fc-baseline-worker", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=calculator or (lambda _: simulation_frames()), + report_runs=report_runs, + ), + ) + + +def test_worker_records_its_simulation_run_after_its_stage12_record() -> None: + events: list[str] = [] + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation, events) + report_runs = FakeReportRuns(events) + + result = _run_worker(simulation, store, report_runs) + + run = report_runs.simulation_runs[BASELINE_RUN_ID] + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == "fc-baseline-worker" + assert run.output_uri == result["artifact"]["uri"] + assert run.output_sha256 == result["artifact"]["content_sha256"] + # The simulation run succeeds last, so a worker Modal restarts between + # the two writes finds its Stage 12 record succeeded and completes the + # simulation run from it, with the same output file. + assert events == [ + "stage12:running", + "claim_simulation_run", + "stage12:succeeded", + "complete_simulation_run", + ] + + +def test_worker_failure_marks_its_simulation_run_failed() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + + def failing_calculator(_simulation): + msg = "calculation failed" + raise ValueError(msg) + + with pytest.raises(RuntimeError, match="execution failed"): + _run_worker(simulation, store, report_runs, calculator=failing_calculator) + + run = report_runs.simulation_runs[BASELINE_RUN_ID] + assert run.status is RunStatus.FAILED + assert run.error_message.startswith("simulation_execution_failed: ValueError") + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.FAILED + + +def test_worker_refuses_a_simulation_run_that_was_already_failed() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + report_runs.simulation_runs[BASELINE_RUN_ID].status = RunStatus.FAILED + calculated = [] + + with pytest.raises(RuntimeError, match="cancelled"): + _run_worker( + simulation, + store, + report_runs, + calculator=lambda value: calculated.append(value) or simulation_frames(), + ) + + assert calculated == [] + assert report_runs.simulation_runs[BASELINE_RUN_ID].output_uri is None + + +def test_restarted_worker_completes_its_run_from_its_succeeded_stage12_record() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + # The earlier attempt marked its Stage 12 record succeeded, then Modal + # pre-empted it before it completed its simulation run. + store.children[simulation.simulation_execution_id] = store.children[ + simulation.simulation_execution_id + ].model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": "gs://private/earlier-attempt.parquet", + "output_sha256": "e" * 64, + "output_schema_version": 1, + "row_identity_columns": ("household.household_id",), + "row_count": 1, + "row_identity_sha256": "f" * 64, + "completed_at": child_record(simulation).created_at, + } + ) + run = report_runs.simulation_runs[BASELINE_RUN_ID] + run.status, run.modal_call_id = RunStatus.RUNNING, "fc-baseline-worker" + calculated = [] + + result = _run_worker( + simulation, + store, + report_runs, + calculator=lambda value: calculated.append(value) or simulation_frames(), + ) + + assert calculated == [] + assert result["artifact"]["uri"] == "gs://private/earlier-attempt.parquet" + assert run.status is RunStatus.SUCCEEDED + assert run.output_uri == "gs://private/earlier-attempt.parquet" + assert run.output_sha256 == "e" * 64 + + +def test_a_restarted_worker_returns_its_result_if_completing_its_run_fails() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + store.children[simulation.simulation_execution_id] = store.children[ + simulation.simulation_execution_id + ].model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": "gs://private/earlier-attempt.parquet", + "output_sha256": "e" * 64, + "output_schema_version": 1, + "row_identity_columns": ("household.household_id",), + "row_count": 1, + "row_identity_sha256": "f" * 64, + "completed_at": child_record(simulation).created_at, + } + ) + report_runs = FakeReportRuns() + report_runs.simulation_runs[BASELINE_RUN_ID].status = RunStatus.RUNNING + + def dropped(*_args, **_kwargs): + msg = "the connection dropped" + raise ConnectionError(msg) + + report_runs.complete_simulation_run = dropped + + result = _run_worker(simulation, store, report_runs) + + assert result["artifact"]["uri"] == "gs://private/earlier-attempt.parquet" + + +def test_a_simulation_run_another_report_recorded_keeps_its_output() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + # Another report run, sharing this simulation run, already recorded it. + store.children[simulation.simulation_execution_id] = store.children[ + simulation.simulation_execution_id + ].model_copy(update={"status": ComparisonRunLifecycleStatus.RUNNING}) + earlier = report_runs.simulation_runs[BASELINE_RUN_ID] + earlier.status = RunStatus.SUCCEEDED + earlier.output_uri = "gs://private/earlier-attempt.parquet" + earlier.output_sha256 = "e" * 64 + + result = _run_worker(simulation, store, report_runs) + + assert earlier.output_uri == "gs://private/earlier-attempt.parquet" + assert earlier.writes == ["claim_simulation_run", "complete_simulation_run"] + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert child.output_uri == result["artifact"]["uri"] + + +def test_a_lost_simulation_run_success_still_returns_the_workers_result() -> None: + # The child record already succeeded; the coordinator completes the run + # from it, so a failed second write must not fail the worker. + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + + def dropped(*_args, **_kwargs): + msg = "the connection dropped" + raise ConnectionError(msg) + + report_runs.complete_simulation_run = dropped + + result = _run_worker(simulation, store, report_runs) + + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert child.output_uri == result["artifact"]["uri"] + assert report_runs.simulation_runs[BASELINE_RUN_ID].status is RunStatus.RUNNING + + +def test_comparison_worker_records_no_simulation_run() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + store = _seeded_store(simulation) + report_runs = FakeReportRuns() + + _run_worker(simulation, store, report_runs, context=invocation_context()) + + assert report_runs.events == [] + + +def test_report_run_worker_needs_its_modal_call_id() -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + + with pytest.raises(ValueError, match="Modal call ID"): + run_single_simulation( + simulation.model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=_seeded_store(simulation), + artifacts=FakeArtifacts(), + report_runs=FakeReportRuns(), + ), + ) + + +def test_report_run_worker_names_its_artifacts_by_digest(monkeypatch) -> None: + simulation = planned_simulation(SimulationRole.BASELINE) + built = [] + + def recording_artifact_store(*, content_addressed=False): + built.append(content_addressed) + return FakeArtifacts() + + monkeypatch.setattr(simulation_module, "artifact_store", recording_artifact_store) + + run_single_simulation( + simulation.model_dump(mode="json"), + report_run_context().model_dump(mode="json"), + required_country="us", + modal_call_id="fc-baseline-worker", + dependencies=SimulationWorkerDependencies( + store=_seeded_store(simulation), + calculator=lambda _: simulation_frames(), + report_runs=FakeReportRuns(), + ), + ) + + assert built == [True] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py index 7e37fef45..6506c0a8f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -224,6 +224,43 @@ def test_dispatch_error_cancels_segments_that_already_started() -> None: assert all(call.cancelled == 1 for call in function.calls) +class InputCancelledForTest(BaseException): + """Stands in for Modal's InputCancellation, a BaseException but not an Exception.""" + + +@pytest.mark.parametrize("interruption", [KeyboardInterrupt, InputCancelledForTest]) +def test_preempted_or_cancelled_worker_cancels_every_started_segment( + interruption, +) -> None: + # Modal pre-empts a call with SIGINT (KeyboardInterrupt) and cancels one + # with InputCancellation; it does not cancel the calls the worker spawned. + runner, function, _, _ = _runner(errors={3: [interruption()]}) + + with pytest.raises(interruption): + runner.run() + + assert len(function.calls) == 20 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_worker_preempted_during_dispatch_cancels_the_segments_it_started() -> None: + runner, function, _, _ = _runner() + spawn = function.spawn + + def preempt_during_dispatch(payload, *, observability_context): + if payload["segment_index"] == 5: + raise KeyboardInterrupt + return spawn(payload, observability_context=observability_context) + + function.spawn = preempt_during_dispatch + + with pytest.raises(KeyboardInterrupt): + runner.run() + + assert len(function.calls) == 5 + assert all(call.cancelled == 1 for call in function.calls) + + def test_persisted_parent_cancellation_cancels_all_segments() -> None: runner, function, _, sleeps = _runner( errors={index: [TimeoutError()] for index in range(20)}