From 2341fd1e07ef0359fb9595362640e7b6bf37e24d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:20:42 +0400 Subject: [PATCH 01/39] Add simulation runs and report runs, found or created per catalog version and year Revision 0003 adds simulation.simulation_runs and simulation.report_runs. A simulation run is one execution of a simulation for a catalog version and a year. It records the region code and dataset row it resolved in that version (economy runs only), its status, Modal call ID, timings, error, trigger (initial, manual or system) and attempt number, unique per simulation, catalog version, year and attempt. Composite foreign keys keep it on its own version's rows: (catalog version, region code) through the Public API's uq_regions_model_version_code, and (dataset, catalog version) through uq_datasets_id_model_version, both ON DELETE RESTRICT. A key on (simulation, kind) lets the database check that only an economy run names a region and dataset. A report run joins one run of each of the report's simulations (no reform run for a single-simulation household report). Composite keys on (run, catalog version, year) make both runs share the report run's version and year. An economy report run carries its comparison step's own status and Modal call; a household report run has none, since one Modal call computes both households. Checks keep each status consistent with its Modal call, completion time, error and next check time, for the conditional status changes that follow. RunStore.start_simulation_run reuses the newest attempt while it is unfinished, or when it succeeded within the last 28 days, so its output files (kept 30 days) outlive the comparison; a failed run is never reused. Otherwise, and always when forced, it adds the next attempt, and concurrent starts settle on one row through the unique attempt number. create_report_run checks the runs belong to the report's simulations and takes the version and year from them; latest_report_run gives the report's current run. The runtime role may read, insert and update both tables but not delete them. Nothing writes them yet. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 32 +- .../__init__.py | 22 + .../public_api.py | 29 +- .../reports.py | 4 + .../run_store.py | 356 +++++++++++++ .../runs.py | 481 +++++++++++++++++ .../schema.py | 2 +- .../simulations.py | 3 + .../tests/public_api_support.py | 78 ++- .../tests/run_support.py | 85 +++ .../tests/test_runs.py | 495 ++++++++++++++++++ .../tests/test_schema.py | 7 +- .../0003_simulation_runs_and_report_runs.py | 328 ++++++++++++ 13 files changed, 1914 insertions(+), 8 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py create mode 100644 libs/policyengine-simulation-persistence/tests/run_support.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_runs.py create mode 100644 migrations/versions/0003_simulation_runs_and_report_runs.py diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 1436930e8..0a60b73e3 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -14,6 +14,8 @@ writes them with DML only. | --- | --- | | `simulations` | What to run, independent of model versions: kind (`economy` or `household`), country, policy (null for current law), dynamic, and either a household or a region code (`state/ca`) with an optional dataset name (null means the region's default dataset in each catalog version) | | `reports` | Which two simulations a report compares (an economy report needs a reform; a household report may have only a baseline), both of the report's own country and kind, with its year and options | +| `simulation_runs` | One execution of a simulation for a catalog version and year: the region code and dataset row it resolved in that version (economy runs only), its status, Modal call ID, timings, error, trigger (`initial`, `manual` or `system`) and attempt number | +| `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call | Both are definitions: the runtime role may read and insert them but not update or delete them. Rows migrated from the Public API's v1 database keep their v1 @@ -22,9 +24,33 @@ the stores' `create_migrated` methods return the existing row when the same v1 ID is migrated again. Foreign keys into the Public API's tables in the same database (`policies` and -`households`, each together with its country, and `dynamics`) are -`ON DELETE RESTRICT`. Creating them needs `REFERENCES` on those tables for the -migrator role; the runtime role needs no privilege on them. +`households`, each together with its country, `dynamics`, +`tax_benefit_model_versions`, and `regions` and `datasets` each together with +its catalog version) are `ON DELETE RESTRICT`. Creating them needs +`REFERENCES` on those tables for the migrator role; the runtime role needs no +privilege on them. + +### Runs + +A simulation run's setup is its catalog version (one policyengine.py bundle +version for one country, which pins the country package, data release and +datasets) and its year. Its region and dataset keys include its catalog +version, through the Public API's unique constraints +`uq_regions_model_version_code` and `uq_datasets_id_model_version`, so a run +can use only its own version's rows; the caller resolves them in that version +(the simulation's region code, and the dataset it names or the region's +default dataset). A report run's simulation run keys include its catalog +version and year, so both of its runs share them. + +`RunStore.start_simulation_run` reuses the newest attempt for the same +simulation, catalog version and year while it is unfinished, or when it +succeeded within the last 28 days (output files are kept 30); a failed run is +never reused. Otherwise, and always when forced, it adds the next attempt. +Concurrent starts settle on one row through the unique attempt number. +`RunStore.create_report_run` joins one run of each of the report's simulations +and takes the catalog version and year from them; a report's status is its +newest report run's. The runtime role may read, insert and update runs but +not delete them. ## Ownership rules diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 5e87ac336..a83a0d3e5 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -15,6 +15,19 @@ ReportRecord, ReportStore, ) +from policyengine_simulation_persistence.run_store import ( + SUCCEEDED_RUN_REUSE_WINDOW, + RunStore, +) +from policyengine_simulation_persistence.runs import ( + NewReportRun, + NewSimulationRun, + ReportRunRecord, + RunStatus, + RunTrigger, + SimulationRunRecord, + StartedSimulationRun, +) from policyengine_simulation_persistence.schema import ( ALEMBIC_VERSION_TABLE, HEAD_REVISION, @@ -34,14 +47,23 @@ "HEAD_REVISION", "RUNTIME_ROLE", "SIMULATION_SCHEMA", + "SUCCEEDED_RUN_REUSE_WINDOW", "NewReport", + "NewReportRun", "NewSimulation", + "NewSimulationRun", "ReportKind", "ReportRecord", + "ReportRunRecord", "ReportStore", + "RunStatus", + "RunStore", + "RunTrigger", "SimulationKind", "SimulationRecord", + "SimulationRunRecord", "SimulationStore", + "StartedSimulationRun", "create_simulation_engine", "normalized_database_url", "read_schema_revision", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py index 65f709a1f..5dfd974fd 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/public_api.py @@ -20,7 +20,7 @@ from __future__ import annotations -from sqlalchemy import Column, MetaData, String, Table, Uuid +from sqlalchemy import Column, MetaData, String, Table, UniqueConstraint, Uuid public_api_metadata = MetaData() @@ -41,3 +41,30 @@ public_api_metadata, Column("id", Uuid(as_uuid=True), primary_key=True), ) +# A catalog version: one published set of model metadata for one country and +# one policyengine.py version. Each publication creates new region and dataset +# rows, unique by code or name within their version. +tax_benefit_model_versions = Table( + "tax_benefit_model_versions", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), +) +regions = Table( + "regions", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("code", String(255), nullable=False), + UniqueConstraint( + "tax_benefit_model_version_id", "code", name="uq_regions_model_version_code" + ), +) +datasets = Table( + "datasets", + public_api_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + UniqueConstraint( + "id", "tax_benefit_model_version_id", name="uq_datasets_id_model_version" + ), +) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py index 99faccf40..ec911e10f 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py @@ -29,6 +29,7 @@ Integer, Table, Text, + UniqueConstraint, Uuid, func, select, @@ -106,6 +107,9 @@ class ReportKind(StrEnum): name="fk_reports_reform_simulation", ondelete="RESTRICT", ), + # Report runs refer to a report together with its kind, so only an economy + # report run has a comparison step. + UniqueConstraint("id", "kind"), ) Index("ix_reports_baseline_simulation_id", reports.c.baseline_simulation_id) Index("ix_reports_reform_simulation_id", reports.c.reform_simulation_id) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py new file mode 100644 index 000000000..e04547fc4 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py @@ -0,0 +1,356 @@ +"""Find, create and read simulation runs and report runs. + +Starting a simulation run reuses the newest attempt for the same simulation, +catalog version and year while it is unfinished, or succeeded within the +reuse window; otherwise, and always when forced, it adds the next attempt. +Concurrent starts settle on one row through the unique attempt number: the +loser of the insert returns the winner's run. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +from sqlalchemy import select +from sqlalchemy.dialects.postgresql import insert as postgres_insert + +from policyengine_simulation_persistence.reports import ReportKind, reports +from policyengine_simulation_persistence.runs import ( + UNFINISHED_STATUSES, + NewReportRun, + NewSimulationRun, + ReportRunRecord, + RunStatus, + SimulationRunRecord, + StartedSimulationRun, + report_runs, + simulation_runs, +) +from policyengine_simulation_persistence.simulations import ( + SimulationKind, + simulations, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy import ColumnElement, Connection, Engine, RowMapping + +# Output files are kept 30 days. A succeeded run is reused only if it finished +# within the last 28, so its files outlive the comparison step that reads them. +SUCCEEDED_RUN_REUSE_WINDOW = timedelta(days=28) + + +def utc_now() -> datetime: + """Return the current time in UTC.""" + return datetime.now(UTC) + + +def _check_region_code(simulation: RowMapping, new_run: NewSimulationRun) -> None: + """Reject an economy run whose region code is not its simulation's. + + A missing region or dataset, or one on a household run, is left to the + database's population check. + """ + if ( + simulation["kind"] == SimulationKind.ECONOMY.value + and new_run.region_code is not None + and new_run.region_code != simulation["region_code"] + ): + msg = "an economy run's region_code must be its simulation's region code" + raise ValueError(msg) + + +def _check_report_run_simulations( + report: RowMapping, + baseline: RowMapping | None, + reform: RowMapping | None, + new_report_run: NewReportRun, +) -> RowMapping: + """Require one run of each of the report's simulations, and no other. + + Return the baseline run's key, from which the report run takes its catalog + version and year. + """ + expected_reform_run = report["reform_simulation_id"] is not None + if ( + baseline is None + or baseline["simulation_id"] != report["baseline_simulation_id"] + or expected_reform_run != (new_report_run.reform_simulation_run_id is not None) + or ( + expected_reform_run + and ( + reform is None + or reform["simulation_id"] != report["reform_simulation_id"] + ) + ) + ): + msg = ( + "a report run needs one existing run of each of the report's " + "simulations: its baseline and, when it has one, its reform" + ) + raise ValueError(msg) + return baseline + + +class RunStore: + """Typed SQLAlchemy Core reads and inserts for runs.""" + + def __init__( + self, + engine: Engine, + *, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Use ``engine`` (the runtime role) and ``clock`` for every timestamp.""" + self.engine = engine + self._clock = clock + + # --- Simulation runs --------------------------------------------------- + + def start_simulation_run( + self, new_run: NewSimulationRun, *, force: bool = False + ) -> StartedSimulationRun: + """Return the reusable run for this setup, or create the next attempt. + + A run is reused while it is unfinished, or when it succeeded within + ``SUCCEEDED_RUN_REUSE_WINDOW``; a failed run never is. ``force`` always + adds the next attempt, unless another caller added it at the same + time, in which case that run is returned. Raise LookupError when the + simulation does not exist. + """ + now = self._clock() + with self.engine.begin() as connection: + simulation = ( + connection.execute( + select(simulations.c.kind, simulations.c.region_code).where( + simulations.c.id == new_run.simulation_id + ) + ) + .mappings() + .first() + ) + if simulation is None: + msg = f"no simulation {new_run.simulation_id}" + raise LookupError(msg) + _check_region_code(simulation, new_run) + latest = self._latest_simulation_run( + connection, + new_run.simulation_id, + new_run.tax_benefit_model_version_id, + new_run.year, + ) + if latest is not None and not force and self._reusable(latest, now): + return StartedSimulationRun(run=latest, created=False) + attempt = 1 if latest is None else latest.attempt + 1 + created = ( + connection.execute( + postgres_insert(simulation_runs) + .values( + id=uuid4(), + simulation_id=new_run.simulation_id, + simulation_kind=simulation["kind"], + tax_benefit_model_version_id=( + new_run.tax_benefit_model_version_id + ), + year=new_run.year, + attempt=attempt, + region_code=new_run.region_code, + dataset_id=new_run.dataset_id, + trigger=new_run.trigger.value, + status=RunStatus.PENDING.value, + created_at=now, + updated_at=now, + ) + .on_conflict_do_nothing( + index_elements=[ + simulation_runs.c.simulation_id, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.year, + simulation_runs.c.attempt, + ] + ) + .returning(*simulation_runs.c) + ) + .mappings() + .first() + ) + if created is not None: + return StartedSimulationRun( + run=SimulationRunRecord.from_row(created), created=True + ) + # Another caller added this attempt after the read above; the + # insert waited for it to commit, so it is visible now. + concurrent = self._latest_simulation_run( + connection, + new_run.simulation_id, + new_run.tax_benefit_model_version_id, + new_run.year, + ) + if concurrent is None: # pragma: no cover - the conflict proves a row + msg = "the conflicting simulation run disappeared" + raise RuntimeError(msg) + return StartedSimulationRun(run=concurrent, created=False) + + @staticmethod + def _reusable(run: SimulationRunRecord, now: datetime) -> bool: + if run.status in UNFINISHED_STATUSES: + return True + return ( + run.status is RunStatus.SUCCEEDED + and run.completed_at is not None + and run.completed_at >= now - SUCCEEDED_RUN_REUSE_WINDOW + ) + + @staticmethod + def _latest_simulation_run( + connection: Connection, + simulation_id: UUID, + catalog_version_id: UUID, + year: int, + ) -> SimulationRunRecord | None: + row = ( + connection.execute( + select(simulation_runs) + .where( + simulation_runs.c.simulation_id == simulation_id, + simulation_runs.c.tax_benefit_model_version_id + == catalog_version_id, + simulation_runs.c.year == year, + ) + .order_by(simulation_runs.c.attempt.desc()) + .limit(1) + ) + .mappings() + .first() + ) + return None if row is None else SimulationRunRecord.from_row(row) + + def get_simulation_run(self, run_id: UUID) -> SimulationRunRecord | None: + """Return the simulation run with this ID, or None.""" + with self.engine.connect() as connection: + row = ( + connection.execute( + select(simulation_runs).where(simulation_runs.c.id == run_id) + ) + .mappings() + .first() + ) + return None if row is None else SimulationRunRecord.from_row(row) + + def latest_simulation_run( + self, simulation_id: UUID, catalog_version_id: UUID, year: int + ) -> SimulationRunRecord | None: + """Return the newest attempt for this simulation, catalog version and year.""" + with self.engine.connect() as connection: + return self._latest_simulation_run( + connection, simulation_id, catalog_version_id, year + ) + + # --- Report runs ------------------------------------------------------- + + def create_report_run(self, new_report_run: NewReportRun) -> ReportRunRecord: + """Add a report run joining two existing simulation runs. + + The report run takes its catalog version and year from its baseline + run; the database refuses a reform run of another version or year. + An economy report run starts with a pending comparison step. Raise + LookupError when the report does not exist, and ValueError when the + runs are not runs of the report's simulations. + """ + now = self._clock() + with self.engine.begin() as connection: + report = ( + connection.execute( + select( + reports.c.kind, + reports.c.baseline_simulation_id, + reports.c.reform_simulation_id, + ).where(reports.c.id == new_report_run.report_id) + ) + .mappings() + .first() + ) + if report is None: + msg = f"no report {new_report_run.report_id}" + raise LookupError(msg) + baseline = self._run_key( + connection, new_report_run.baseline_simulation_run_id + ) + reform = ( + None + if new_report_run.reform_simulation_run_id is None + else self._run_key(connection, new_report_run.reform_simulation_run_id) + ) + baseline = _check_report_run_simulations( + report, baseline, reform, new_report_run + ) + economy = report["kind"] == ReportKind.ECONOMY.value + row = ( + connection.execute( + postgres_insert(report_runs) + .values( + id=uuid4(), + report_id=new_report_run.report_id, + report_kind=report["kind"], + tax_benefit_model_version_id=( + baseline["tax_benefit_model_version_id"] + ), + year=baseline["year"], + baseline_simulation_run_id=( + new_report_run.baseline_simulation_run_id + ), + reform_simulation_run_id=( + new_report_run.reform_simulation_run_id + ), + trigger=new_report_run.trigger.value, + comparison_status=( + RunStatus.PENDING.value if economy else None + ), + created_at=now, + updated_at=now, + ) + .returning(*report_runs.c) + ) + .mappings() + .one() + ) + return ReportRunRecord.from_row(row) + + @staticmethod + def _run_key(connection: Connection, run_id: UUID) -> RowMapping | None: + return ( + connection.execute( + select( + simulation_runs.c.simulation_id, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.year, + ).where(simulation_runs.c.id == run_id) + ) + .mappings() + .first() + ) + + def get_report_run(self, report_run_id: UUID) -> ReportRunRecord | None: + """Return the report run with this ID, or None.""" + return self._one_report_run(report_runs.c.id == report_run_id) + + def latest_report_run(self, report_id: UUID) -> ReportRunRecord | None: + """Return the report's newest report run, whose status is the report's.""" + return self._one_report_run(report_runs.c.report_id == report_id) + + def _one_report_run(self, condition: ColumnElement[bool]) -> ReportRunRecord | None: + with self.engine.connect() as connection: + row = ( + connection.execute( + select(report_runs) + .where(condition) + .order_by(report_runs.c.created_at.desc(), report_runs.c.id.desc()) + .limit(1) + ) + .mappings() + .first() + ) + return None if row is None else ReportRunRecord.from_row(row) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py new file mode 100644 index 000000000..d69b42012 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py @@ -0,0 +1,481 @@ +"""``simulation_runs`` and ``report_runs``: executions of simulations and reports. + +A simulation run is one execution of a simulation for a catalog version and a +year. It records the region and dataset rows it resolved in that catalog +version, its own status and Modal call, its timings and error, what started it +(its trigger) and its attempt number. Runs are unique per simulation, catalog +version, year and attempt, so a run that already exists is reused, and a forced +re-run adds the next attempt. + +A report run is one calculation of a report: a baseline and a reform +simulation run of the same catalog version and year (the reform is null for a +single-simulation household report), plus, for an economy report, the +comparison step's own status and Modal call. A household report has no +comparison step: one Modal call computes both households, so its simulation +runs share that call. + +Each status moves ``pending`` -> ``dispatching`` -> ``running`` -> ``succeeded`` +or ``failed``. The region and dataset keys include the catalog version, through +the Public API's unique constraints ``uq_regions_model_version_code`` and +``uq_datasets_id_model_version``, so a run can use only its own version's +rows. A report run's simulation run keys include its catalog version and year, +so both runs share them. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from enum import StrEnum +from typing import TYPE_CHECKING, Any, Self + +from sqlalchemy import ( + CheckConstraint, + Column, + DateTime, + ForeignKeyConstraint, + Index, + Integer, + String, + Table, + Text, + UniqueConstraint, + Uuid, + text, +) + +from policyengine_simulation_persistence import public_api +from policyengine_simulation_persistence.reports import ( + FIRST_YEAR, + LAST_YEAR, + ReportKind, + reports, +) +from policyengine_simulation_persistence.schema import simulation_metadata +from policyengine_simulation_persistence.simulations import ( + CATALOG_KEY_LENGTH, + SimulationKind, + simulations, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from sqlalchemy import RowMapping + +MAX_ERROR_MESSAGE_LENGTH = 2000 + + +class RunStatus(StrEnum): + """Where a simulation run, or a report run's comparison step, stands.""" + + PENDING = "pending" + DISPATCHING = "dispatching" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +UNFINISHED_STATUSES = (RunStatus.PENDING, RunStatus.DISPATCHING, RunStatus.RUNNING) +FINISHED_STATUSES = (RunStatus.SUCCEEDED, RunStatus.FAILED) + + +class RunTrigger(StrEnum): + """What started a run. + + ``initial`` is a report's or calculation's first submission, ``manual`` a + user's re-run, and ``system`` a run sim or a migration started itself, for + example the first opening of a migrated report. + """ + + INITIAL = "initial" + MANUAL = "manual" + SYSTEM = "system" + + +def _values(enum: type[StrEnum]) -> str: + return ", ".join(f"'{member.value}'" for member in enum) + + +def _status_checks(prefix: str, *, optional: bool) -> tuple[CheckConstraint, ...]: + """Return the checks that keep one status and its columns consistent. + + ``prefix`` is empty for a simulation run and ``comparison_`` for a report + run's comparison step, whose status is null (``optional``) when the report + run has no comparison step. + """ + status = f"{prefix}status" + unless_absent = f"{status} IS NULL OR " if optional else "" + return ( + CheckConstraint( + f"{unless_absent}{status} IN ({_values(RunStatus)})", + name=f"{prefix}status", + ), + # A Modal call exists once dispatching succeeded, never before; a + # failed run may have one or not. + CheckConstraint( + f"{unless_absent}" + f"({status} IN ('pending', 'dispatching') " + f"AND {prefix}modal_call_id IS NULL) " + f"OR ({status} IN ('running', 'succeeded') " + f"AND {prefix}modal_call_id IS NOT NULL) " + f"OR {status} = 'failed'", + name=f"{prefix}modal_call", + ), + CheckConstraint( + f"{unless_absent}" + f"({status} IN ('succeeded', 'failed')) " + f"= ({prefix}completed_at IS NOT NULL)", + name=f"{prefix}completed_at", + ), + CheckConstraint( + f"{unless_absent}({status} = 'failed') = ({prefix}error_code IS NOT NULL)", + name=f"{prefix}error", + ), + # Only a running call is checked, and every running call is due a check. + CheckConstraint( + f"{unless_absent}" + f"({status} = 'running') = ({prefix}next_check_at IS NOT NULL)", + name=f"{prefix}next_check_at", + ), + CheckConstraint(f"{prefix}check_count >= 0", name=f"{prefix}check_count"), + ) + + +def _execution_columns(prefix: str, *, required: bool) -> list[Column[Any]]: + """Return the status, Modal call, error and timing columns of one execution.""" + return [ + Column(f"{prefix}status", Text, nullable=not required), + Column(f"{prefix}modal_call_id", Text, nullable=True), + Column(f"{prefix}error_code", Text, nullable=True), + Column(f"{prefix}error_message", Text, nullable=True), + Column( + f"{prefix}check_count", + Integer, + nullable=False, + server_default=text("0"), + ), + Column(f"{prefix}dispatched_at", DateTime(timezone=True), nullable=True), + Column(f"{prefix}next_check_at", DateTime(timezone=True), nullable=True), + Column(f"{prefix}completed_at", DateTime(timezone=True), nullable=True), + ] + + +simulation_runs = Table( + "simulation_runs", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("simulation_id", Uuid(as_uuid=True), nullable=False), + Column("simulation_kind", Text, nullable=False), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("year", Integer, nullable=False), + Column("attempt", Integer, nullable=False), + Column("region_code", String(CATALOG_KEY_LENGTH), nullable=True), + Column("dataset_id", Uuid(as_uuid=True), nullable=True), + Column("trigger", Text, nullable=False), + *_execution_columns("", required=True), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), + CheckConstraint("attempt >= 1", name="attempt"), + CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"), + CheckConstraint( + "(simulation_kind = 'household' AND region_code IS NULL " + "AND dataset_id IS NULL) " + "OR (simulation_kind = 'economy' AND region_code IS NOT NULL " + "AND dataset_id IS NOT NULL)", + name="population_matches_kind", + ), + *_status_checks("", optional=False), + UniqueConstraint( + "simulation_id", + "tax_benefit_model_version_id", + "year", + "attempt", + name="uq_simulation_runs_attempt", + ), + # Report runs refer to a simulation run with its catalog version and year, + # so both runs of a report run share them. + UniqueConstraint( + "id", + "tax_benefit_model_version_id", + "year", + name="uq_simulation_runs_id_catalog_version_year", + ), + ForeignKeyConstraint( + ["simulation_id", "simulation_kind"], + [simulations.c.id, simulations.c.kind], + name="fk_simulation_runs_simulation", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + [public_api.tax_benefit_model_versions.c.id], + name="fk_simulation_runs_catalog_version", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["tax_benefit_model_version_id", "region_code"], + [ + public_api.regions.c.tax_benefit_model_version_id, + public_api.regions.c.code, + ], + name="fk_simulation_runs_region", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["dataset_id", "tax_benefit_model_version_id"], + [ + public_api.datasets.c.id, + public_api.datasets.c.tax_benefit_model_version_id, + ], + name="fk_simulation_runs_dataset", + ondelete="RESTRICT", + ), +) +# A Public API delete of a catalog version, region or dataset checks these for +# references. +Index( + "ix_simulation_runs_catalog_version_region", + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.region_code, +) +Index("ix_simulation_runs_dataset_id", simulation_runs.c.dataset_id) +# Every run of one Modal call (a household report's two runs share one). +Index( + "ix_simulation_runs_modal_call_id", + simulation_runs.c.modal_call_id, + postgresql_where=simulation_runs.c.modal_call_id.is_not(None), +) +# The recording loop reads only running runs that are due a check, and +# dispatches that stopped. +Index( + "ix_simulation_runs_due_checks", + simulation_runs.c.next_check_at, + postgresql_where=simulation_runs.c.status == RunStatus.RUNNING.value, +) +Index( + "ix_simulation_runs_dispatching", + simulation_runs.c.updated_at, + postgresql_where=simulation_runs.c.status == RunStatus.DISPATCHING.value, +) + +_SIMULATION_RUN_KEY = [ + simulation_runs.c.id, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.year, +] + +report_runs = Table( + "report_runs", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("report_id", Uuid(as_uuid=True), nullable=False), + Column("report_kind", Text, nullable=False), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("year", Integer, nullable=False), + Column("baseline_simulation_run_id", Uuid(as_uuid=True), nullable=False), + Column("reform_simulation_run_id", Uuid(as_uuid=True), nullable=True), + Column("trigger", Text, nullable=False), + *_execution_columns("comparison_", required=False), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), + CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"), + # An economy report run compares two simulation runs in its own comparison + # step; a household report run has no comparison step. + CheckConstraint( + "(report_kind = 'economy' AND reform_simulation_run_id IS NOT NULL " + "AND comparison_status IS NOT NULL) " + "OR (report_kind = 'household' AND comparison_status IS NULL " + "AND comparison_modal_call_id IS NULL " + "AND comparison_error_code IS NULL " + "AND comparison_dispatched_at IS NULL " + "AND comparison_next_check_at IS NULL " + "AND comparison_completed_at IS NULL " + "AND comparison_check_count = 0)", + name="comparison_matches_kind", + ), + *_status_checks("comparison_", optional=True), + ForeignKeyConstraint( + ["report_id", "report_kind"], + [reports.c.id, reports.c.kind], + name="fk_report_runs_report", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + [ + "baseline_simulation_run_id", + "tax_benefit_model_version_id", + "year", + ], + _SIMULATION_RUN_KEY, + name="fk_report_runs_baseline_simulation_run", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["reform_simulation_run_id", "tax_benefit_model_version_id", "year"], + _SIMULATION_RUN_KEY, + name="fk_report_runs_reform_simulation_run", + ondelete="RESTRICT", + ), +) +# A report's status is its newest report run's. +Index( + "ix_report_runs_report_id_created_at", + report_runs.c.report_id, + report_runs.c.created_at, +) +Index( + "ix_report_runs_baseline_simulation_run_id", + report_runs.c.baseline_simulation_run_id, +) +Index( + "ix_report_runs_reform_simulation_run_id", + report_runs.c.reform_simulation_run_id, +) +Index( + "ix_report_runs_due_comparison_checks", + report_runs.c.comparison_next_check_at, + postgresql_where=report_runs.c.comparison_status == RunStatus.RUNNING.value, +) +Index( + "ix_report_runs_comparisons_dispatching", + report_runs.c.updated_at, + postgresql_where=report_runs.c.comparison_status == RunStatus.DISPATCHING.value, +) + + +def bounded_error_message(message: str) -> str: + """Return ``message`` cut to the length the run tables store.""" + return message[:MAX_ERROR_MESSAGE_LENGTH] + + +@dataclass(frozen=True) +class NewSimulationRun: + """A simulation run to find or create. + + An economy run names the region and dataset rows its caller resolved in + the run's catalog version: the simulation's region code, and the dataset + the simulation names or, when it names none, the region's default dataset + in that version. The database checks that both belong to that version. A + household run names neither. + """ + + simulation_id: UUID + tax_benefit_model_version_id: UUID + year: int + trigger: RunTrigger + region_code: str | None = None + dataset_id: UUID | None = None + + def __post_init__(self) -> None: + """Reject a year outside the supported range without touching the database.""" + if not FIRST_YEAR <= self.year <= LAST_YEAR: + msg = f"year must be between {FIRST_YEAR} and {LAST_YEAR}" + raise ValueError(msg) + + +@dataclass(frozen=True) +class NewReportRun: + """A report run to create from two simulation runs that already exist. + + The baseline run must be a run of the report's baseline simulation, and + the reform run one of its reform simulation; an economy report run needs + both. They must share a catalog version and year, which the report run + takes from them. + """ + + report_id: UUID + baseline_simulation_run_id: UUID + trigger: RunTrigger + reform_simulation_run_id: UUID | None = None + + +def _record_values(cls: type, row: RowMapping) -> dict[str, Any]: + return {field.name: row[field.name] for field in fields(cls)} + + +@dataclass(frozen=True) +class SimulationRunRecord: + """One ``simulation.simulation_runs`` row.""" + + id: UUID + simulation_id: UUID + simulation_kind: SimulationKind + tax_benefit_model_version_id: UUID + year: int + attempt: int + region_code: str | None + dataset_id: UUID | None + trigger: RunTrigger + status: RunStatus + modal_call_id: str | None + error_code: str | None + error_message: str | None + check_count: int + dispatched_at: datetime | None + next_check_at: datetime | None + completed_at: datetime | None + created_at: datetime + updated_at: datetime + + @property + def finished(self) -> bool: + """Whether the run has succeeded or failed.""" + return self.status in FINISHED_STATUSES + + @classmethod + def from_row(cls, row: RowMapping) -> Self: + """Build the record from a selected ``simulation_runs`` row.""" + values = _record_values(cls, row) + values["simulation_kind"] = SimulationKind(values["simulation_kind"]) + values["trigger"] = RunTrigger(values["trigger"]) + values["status"] = RunStatus(values["status"]) + return cls(**values) + + +@dataclass(frozen=True) +class ReportRunRecord: + """One ``simulation.report_runs`` row. + + The comparison fields are None for a household report run, which has no + comparison step. + """ + + id: UUID + report_id: UUID + report_kind: ReportKind + tax_benefit_model_version_id: UUID + year: int + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + trigger: RunTrigger + comparison_status: RunStatus | None + comparison_modal_call_id: str | None + comparison_error_code: str | None + comparison_error_message: str | None + comparison_check_count: int + comparison_dispatched_at: datetime | None + comparison_next_check_at: datetime | None + comparison_completed_at: datetime | None + created_at: datetime + updated_at: datetime + + @classmethod + def from_row(cls, row: RowMapping) -> Self: + """Build the record from a selected ``report_runs`` row.""" + values = _record_values(cls, row) + values["report_kind"] = ReportKind(values["report_kind"]) + values["trigger"] = RunTrigger(values["trigger"]) + if values["comparison_status"] is not None: + values["comparison_status"] = RunStatus(values["comparison_status"]) + return cls(**values) + + +@dataclass(frozen=True) +class StartedSimulationRun: + """The simulation run a start resolved to, and whether this call created it.""" + + run: SimulationRunRecord + created: bool diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py index 7e4080638..7ef7081c6 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py @@ -17,7 +17,7 @@ # The newest revision in migrations/versions. Each migration PR updates it, and # tests/test_schema.py fails when it differs from the chain's head. The # readiness check reports it next to the database's revision. -HEAD_REVISION = "0002" +HEAD_REVISION = "0003" NAMING_CONVENTION = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py index 0ad0e7629..76f134c26 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/simulations.py @@ -106,6 +106,9 @@ class SimulationKind(StrEnum): # Reports refer to a simulation together with its country and kind, so a # report can only compare simulations of its own country and kind. UniqueConstraint("id", "country", "kind"), + # Simulation runs refer to a simulation together with its kind, so a run + # names a region and dataset exactly when its simulation is an economy. + UniqueConstraint("id", "kind"), ForeignKeyConstraint( ["policy_id", "country"], [public_api.policies.c.id, public_api.policies.c.country_id], diff --git a/libs/policyengine-simulation-persistence/tests/public_api_support.py b/libs/policyengine-simulation-persistence/tests/public_api_support.py index cf144c68d..2ace7b50b 100644 --- a/libs/policyengine-simulation-persistence/tests/public_api_support.py +++ b/libs/policyengine-simulation-persistence/tests/public_api_support.py @@ -21,6 +21,24 @@ from sqlalchemy import Engine PUBLIC_API_TABLE_STATEMENTS = ( + "CREATE TABLE public.tax_benefit_model_versions (id uuid PRIMARY KEY)", + "CREATE TABLE public.datasets (" + " id uuid PRIMARY KEY," + " name varchar(255) NOT NULL," + " tax_benefit_model_version_id uuid NOT NULL" + " REFERENCES public.tax_benefit_model_versions (id)," + " CONSTRAINT uq_datasets_model_version_name" + " UNIQUE (tax_benefit_model_version_id, name)," + " CONSTRAINT uq_datasets_id_model_version" + " UNIQUE (id, tax_benefit_model_version_id))", + "CREATE TABLE public.regions (" + " id uuid PRIMARY KEY," + " code varchar(255) NOT NULL," + " tax_benefit_model_version_id uuid NOT NULL" + " REFERENCES public.tax_benefit_model_versions (id)," + " default_dataset_id uuid NOT NULL," + " CONSTRAINT uq_regions_model_version_code" + " UNIQUE (tax_benefit_model_version_id, code))", "CREATE TABLE public.policies (" " id uuid PRIMARY KEY," " country_id varchar(2) NOT NULL," @@ -31,7 +49,14 @@ " CONSTRAINT uq_households_id_country UNIQUE (id, country_id))", "CREATE TABLE public.dynamics (id uuid PRIMARY KEY)", ) -REFERENCED_TABLES = ("policies", "households", "dynamics") +REFERENCED_TABLES = ( + "tax_benefit_model_versions", + "datasets", + "regions", + "policies", + "households", + "dynamics", +) def superuser_engine(superuser_url: str, **options: object) -> Engine: @@ -65,13 +90,57 @@ def create_public_api_tables(superuser_url: str, *, grant_references: bool) -> N @dataclass(frozen=True) class PublicApiRows: - """Public API rows a simulation can refer to.""" + """Public API rows a simulation or run can refer to. + + Two catalog versions each have California (``state/ca``) with its own + default dataset, as each publication creates new rows. + """ us_policy_id: UUID uk_policy_id: UUID us_household_id: UUID uk_household_id: UUID dynamic_id: UUID + catalog_version_id: UUID + newer_catalog_version_id: UUID + dataset_id: UUID + newer_dataset_id: UUID + + +def _catalog_statements(rows: PublicApiRows) -> tuple[tuple[str, dict], ...]: + return ( + ( + "INSERT INTO public.tax_benefit_model_versions (id) VALUES " + "(:version), (:newer_version)", + { + "version": rows.catalog_version_id, + "newer_version": rows.newer_catalog_version_id, + }, + ), + ( + "INSERT INTO public.datasets (id, name, tax_benefit_model_version_id) " + "VALUES (:dataset, 'populace_us_2024', :version), " + "(:newer_dataset, 'populace_us_2024_acs_local', :newer_version)", + { + "dataset": rows.dataset_id, + "newer_dataset": rows.newer_dataset_id, + "version": rows.catalog_version_id, + "newer_version": rows.newer_catalog_version_id, + }, + ), + ( + "INSERT INTO public.regions " + "(id, code, tax_benefit_model_version_id, default_dataset_id) VALUES " + "(gen_random_uuid(), 'state/ca', :version, :dataset), " + "(gen_random_uuid(), 'state/ca', :newer_version, :newer_dataset)", + { + "version": rows.catalog_version_id, + "newer_version": rows.newer_catalog_version_id, + "dataset": rows.dataset_id, + "newer_dataset": rows.newer_dataset_id, + }, + ), + ) def insert_public_api_rows(superuser_url: str) -> PublicApiRows: @@ -82,8 +151,13 @@ def insert_public_api_rows(superuser_url: str) -> PublicApiRows: us_household_id=uuid4(), uk_household_id=uuid4(), dynamic_id=uuid4(), + catalog_version_id=uuid4(), + newer_catalog_version_id=uuid4(), + dataset_id=uuid4(), + newer_dataset_id=uuid4(), ) statements = ( + *_catalog_statements(rows), ( "INSERT INTO public.policies (id, country_id) VALUES " "(:us_policy, 'us'), (:uk_policy, 'uk')", diff --git a/libs/policyengine-simulation-persistence/tests/run_support.py b/libs/policyengine-simulation-persistence/tests/run_support.py new file mode 100644 index 000000000..99a2427bb --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/run_support.py @@ -0,0 +1,85 @@ +"""Builders and direct status changes for the run tests.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_persistence import ( + NewSimulationRun, + RunStore, + RunTrigger, + create_simulation_engine, +) +from public_api_support import superuser_engine +from sqlalchemy import text + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_persistence import SimulationRecord + from public_api_support import PublicApiRows + + +def run_store(database) -> RunStore: + """Return a store that connects as the runtime role.""" + return RunStore(create_simulation_engine(database.runtime_url)) + + +def economy_run( + simulation: SimulationRecord, rows: PublicApiRows, **overrides: object +) -> NewSimulationRun: + """Return a 2026 run of an economy simulation on the first catalog version.""" + values: dict[str, object] = { + "simulation_id": simulation.id, + "tax_benefit_model_version_id": rows.catalog_version_id, + "year": 2026, + "trigger": RunTrigger.INITIAL, + "region_code": "state/ca", + "dataset_id": rows.dataset_id, + } + values.update(overrides) + return NewSimulationRun(**values) # type: ignore[arg-type] + + +def household_run( + simulation: SimulationRecord, rows: PublicApiRows, **overrides: object +) -> NewSimulationRun: + """Return a 2026 run of a household simulation on the first catalog version.""" + values: dict[str, object] = { + "simulation_id": simulation.id, + "tax_benefit_model_version_id": rows.catalog_version_id, + "year": 2026, + "trigger": RunTrigger.INITIAL, + } + values.update(overrides) + return NewSimulationRun(**values) # type: ignore[arg-type] + + +def set_simulation_run_status( + database, run_id: UUID, status: str, *, completed_days_ago: int = 0 +) -> None: + """Move a run to a status directly, as the superuser, for read tests.""" + finished = status in {"succeeded", "failed"} + engine = superuser_engine(database.superuser_url) + try: + with engine.begin() as connection: + connection.execute( + text( + "UPDATE simulation.simulation_runs SET status = :status, " + "modal_call_id = CASE WHEN :status IN ('running', 'succeeded') " + "THEN 'fc-test' ELSE modal_call_id END, " + "error_code = CASE WHEN :status = 'failed' " + "THEN 'test_failure' ELSE NULL END, " + "completed_at = CASE WHEN :finished " + "THEN now() - make_interval(days => :days) ELSE NULL END " + "WHERE id = :id" + ), + { + "status": status, + "finished": finished, + "days": completed_days_ago, + "id": run_id, + }, + ) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_runs.py b/libs/policyengine-simulation-persistence/tests/test_runs.py new file mode 100644 index 000000000..6ad236bf5 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_runs.py @@ -0,0 +1,495 @@ +"""simulation_runs and report_runs: executions of simulations and reports.""" + +from __future__ import annotations + +from datetime import timedelta +from uuid import uuid4 + +import pytest +from policyengine_simulation_persistence import ( + SUCCEEDED_RUN_REUSE_WINDOW, + NewReportRun, + ReportKind, + ReportRunRecord, + RunStatus, + RunTrigger, + SimulationRunRecord, + create_simulation_engine, +) +from public_api_support import delete_public_api_row +from report_support import ( + economy_pair, + economy_report, + household_baseline, + report_store, +) +from run_support import ( + economy_run, + household_run, + run_store, + set_simulation_run_status, +) +from simulation_support import ( + economy_simulation, + household_simulation, + simulation_store, +) +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError, ProgrammingError + +pytestmark = pytest.mark.postgres + + +def test_an_economy_run_records_the_region_and_dataset_it_resolved( + migrated_database, public_api_rows +) -> None: + _, reform = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + + started = store.start_simulation_run(economy_run(reform, public_api_rows)) + + run = started.run + assert started.created is True + assert isinstance(run, SimulationRunRecord) + assert run.simulation_id == reform.id + assert run.simulation_kind == "economy" + assert run.tax_benefit_model_version_id == public_api_rows.catalog_version_id + assert run.year == 2026 + assert run.attempt == 1 + assert run.region_code == "state/ca" + assert run.dataset_id == public_api_rows.dataset_id + assert run.status is RunStatus.PENDING + assert run.trigger is RunTrigger.INITIAL + assert run.modal_call_id is None + assert run.completed_at is None + assert store.get_simulation_run(run.id) == run + + +def test_a_run_that_exists_is_reused(migrated_database, public_api_rows) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + again = store.start_simulation_run( + economy_run(baseline, public_api_rows, trigger=RunTrigger.MANUAL) + ) + + assert again.created is False + assert again.run == first.run + + +def test_a_forced_rerun_adds_the_next_attempt( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + forced = store.start_simulation_run( + economy_run(baseline, public_api_rows, trigger=RunTrigger.MANUAL), force=True + ) + + assert forced.created is True + assert forced.run.attempt == first.run.attempt + 1 + assert forced.run.trigger is RunTrigger.MANUAL + latest = store.latest_simulation_run( + baseline.id, public_api_rows.catalog_version_id, 2026 + ) + assert latest == forced.run + + +def test_a_failed_run_is_not_reused(migrated_database, public_api_rows) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + set_simulation_run_status(migrated_database, first.run.id, "failed") + + again = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + assert again.created is True + assert again.run.attempt == 2 + + +@pytest.mark.parametrize( + ("completed_days_ago", "reused"), + [(27, True), (29, False)], + ids=["within-28-days", "older-than-28-days"], +) +def test_a_succeeded_run_is_reused_only_within_28_days( + migrated_database, public_api_rows, completed_days_ago, reused +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + set_simulation_run_status( + migrated_database, + first.run.id, + "succeeded", + completed_days_ago=completed_days_ago, + ) + + again = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + assert again.created is not reused + assert again.run.attempt == (1 if reused else 2) + assert timedelta(days=28) == SUCCEEDED_RUN_REUSE_WINDOW + + +def test_each_catalog_version_and_year_has_its_own_runs( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + store = run_store(migrated_database) + first = store.start_simulation_run(economy_run(baseline, public_api_rows)) + + newer = store.start_simulation_run( + economy_run( + baseline, + public_api_rows, + tax_benefit_model_version_id=public_api_rows.newer_catalog_version_id, + dataset_id=public_api_rows.newer_dataset_id, + ) + ) + next_year = store.start_simulation_run( + economy_run(baseline, public_api_rows, year=2027) + ) + + assert len({first.run.id, newer.run.id, next_year.run.id}) == 3 + assert newer.created is True + assert next_year.created is True + assert newer.run.attempt == next_year.run.attempt == 1 + + +@pytest.mark.parametrize( + ("overrides", "constraint"), + [ + ({"dataset_attribute": "newer_dataset_id"}, "fk_simulation_runs_dataset"), + ({"region_code": None}, "population_matches_kind"), + ({"dataset_attribute": None}, "population_matches_kind"), + ], + ids=[ + "dataset-of-another-version", + "economy-without-region", + "economy-without-dataset", + ], +) +def test_an_economy_run_uses_only_its_own_versions_dataset( + migrated_database, public_api_rows, overrides, constraint +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + values = dict(overrides) + if "dataset_attribute" in values: + attribute = values.pop("dataset_attribute") + values["dataset_id"] = ( + None if attribute is None else getattr(public_api_rows, attribute) + ) + + with pytest.raises(IntegrityError, match=constraint): + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows, **values) + ) + + +def test_an_economy_run_needs_its_region_in_its_catalog_version( + migrated_database, public_api_rows +) -> None: + new_york = simulation_store(migrated_database).create( + economy_simulation(public_api_rows, region_code="state/ny") + ) + + with pytest.raises(IntegrityError, match="fk_simulation_runs_region"): + run_store(migrated_database).start_simulation_run( + economy_run(new_york, public_api_rows, region_code="state/ny") + ) + + +def test_an_economy_run_uses_its_simulations_region( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(ValueError, match="region_code"): + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows, region_code="state/ny") + ) + + +def test_a_household_run_has_no_region_or_dataset( + migrated_database, public_api_rows +) -> None: + simulation = household_baseline(migrated_database, public_api_rows) + store = run_store(migrated_database) + + started = store.start_simulation_run(household_run(simulation, public_api_rows)) + assert started.run.simulation_kind == "household" + assert started.run.region_code is None + + with pytest.raises(IntegrityError, match="population_matches_kind"): + store.start_simulation_run( + household_run( + simulation, + public_api_rows, + year=2027, + region_code="state/ca", + dataset_id=public_api_rows.dataset_id, + ) + ) + + +def test_a_run_of_an_unknown_simulation_is_refused( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + + with pytest.raises(LookupError, match="simulation"): + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows, simulation_id=uuid4()) + ) + + +def test_the_public_api_cannot_delete_a_dataset_a_run_used( + migrated_database, public_api_rows +) -> None: + baseline, _ = economy_pair(migrated_database, public_api_rows) + run_store(migrated_database).start_simulation_run( + economy_run(baseline, public_api_rows) + ) + + with pytest.raises(IntegrityError, match="violates foreign key constraint"): + delete_public_api_row( + migrated_database.superuser_url, "datasets", public_api_rows.dataset_id + ) + + +def _economy_report_with_runs(database, rows): + baseline, reform = economy_pair(database, rows) + report = report_store(database).create(economy_report(baseline, reform)) + store = run_store(database) + baseline_run = store.start_simulation_run(economy_run(baseline, rows)).run + reform_run = store.start_simulation_run(economy_run(reform, rows)).run + return report, baseline_run, reform_run + + +def test_an_economy_report_run_joins_two_runs_and_waits_for_its_comparison( + migrated_database, public_api_rows +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + store = run_store(migrated_database) + + report_run = store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + assert isinstance(report_run, ReportRunRecord) + assert report_run.report_kind == "economy" + assert report_run.tax_benefit_model_version_id == ( + public_api_rows.catalog_version_id + ) + assert report_run.year == 2026 + assert report_run.comparison_status is RunStatus.PENDING + assert report_run.comparison_modal_call_id is None + assert store.get_report_run(report_run.id) == report_run + assert store.latest_report_run(report.id) == report_run + + +def test_both_runs_of_a_report_run_share_its_version_and_year( + migrated_database, public_api_rows +) -> None: + report, baseline_run, _ = _economy_report_with_runs( + migrated_database, public_api_rows + ) + store = run_store(migrated_database) + reform = report_store(migrated_database).get(report.id) + assert reform is not None + later_reform_run = store.start_simulation_run( + economy_run( + simulation_store(migrated_database).get(reform.reform_simulation_id), + public_api_rows, + year=2027, + ) + ).run + + with pytest.raises(IntegrityError, match="fk_report_runs_reform_simulation_run"): + store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=later_reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +def test_a_report_run_uses_only_runs_of_the_reports_simulations( + migrated_database, public_api_rows +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + + with pytest.raises(ValueError, match="report's simulations"): + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=reform_run.id, + reform_simulation_run_id=baseline_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +def test_an_economy_report_run_needs_its_reform_run( + migrated_database, public_api_rows +) -> None: + report, baseline_run, _ = _economy_report_with_runs( + migrated_database, public_api_rows + ) + + with pytest.raises(ValueError, match="report's simulations"): + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +def test_a_household_report_run_has_no_comparison_step( + migrated_database, public_api_rows +) -> None: + store = simulation_store(migrated_database) + baseline = household_baseline(migrated_database, public_api_rows) + reform = store.create( + household_simulation(public_api_rows, policy_id=public_api_rows.us_policy_id) + ) + report = report_store(migrated_database).create( + economy_report(baseline, reform, kind=ReportKind.HOUSEHOLD) + ) + runs = run_store(migrated_database) + baseline_run = runs.start_simulation_run( + household_run(baseline, public_api_rows) + ).run + reform_run = runs.start_simulation_run(household_run(reform, public_api_rows)).run + + report_run = runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.SYSTEM, + ) + ) + + assert report_run.report_kind == "household" + assert report_run.comparison_status is None + assert report_run.trigger is RunTrigger.SYSTEM + + +def test_a_single_simulation_household_report_run_has_no_reform_run( + migrated_database, public_api_rows +) -> None: + baseline = household_baseline(migrated_database, public_api_rows) + report = report_store(migrated_database).create( + economy_report( + baseline, baseline, kind=ReportKind.HOUSEHOLD, reform_simulation_id=None + ) + ) + runs = run_store(migrated_database) + baseline_run = runs.start_simulation_run( + household_run(baseline, public_api_rows) + ).run + + report_run = runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + assert report_run.reform_simulation_run_id is None + + +def test_the_newest_report_run_is_the_reports_current_run( + migrated_database, public_api_rows +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + store = run_store(migrated_database) + new_report_run = NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + store.create_report_run(new_report_run) + + rerun = store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.MANUAL, + ) + ) + + assert store.latest_report_run(report.id) == rerun + assert store.latest_report_run(uuid4()) is None + + +def test_a_report_run_of_an_unknown_report_is_refused( + migrated_database, public_api_rows +) -> None: + _, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + + with pytest.raises(LookupError, match="report"): + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=uuid4(), + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + +@pytest.mark.parametrize("table_name", ["simulation_runs", "report_runs"]) +def test_the_runtime_role_updates_runs_but_cannot_delete_them( + migrated_database, public_api_rows, table_name +) -> None: + report, baseline_run, reform_run = _economy_report_with_runs( + migrated_database, public_api_rows + ) + run_store(migrated_database).create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + engine = create_simulation_engine(migrated_database.runtime_url) + try: + with engine.begin() as connection: + updated = connection.execute( + text(f"UPDATE simulation.{table_name} SET updated_at = now()") # noqa: S608 # fixed table names + ) + assert updated.rowcount >= 1 + with ( + pytest.raises(ProgrammingError, match="permission denied"), + engine.begin() as connection, + ): + connection.execute(text(f"DELETE FROM simulation.{table_name}")) # noqa: S608 # fixed table names + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_schema.py b/libs/policyengine-simulation-persistence/tests/test_schema.py index 8ac5b7b05..4a6f42383 100644 --- a/libs/policyengine-simulation-persistence/tests/test_schema.py +++ b/libs/policyengine-simulation-persistence/tests/test_schema.py @@ -116,4 +116,9 @@ def test_importing_the_package_registers_every_table() -> None: importlib.import_module(f"policyengine_simulation_persistence.{module.name}") assert set(simulation_metadata.tables) == package_tables - assert {"simulation.simulations", "simulation.reports"} <= package_tables + assert { + "simulation.simulations", + "simulation.reports", + "simulation.simulation_runs", + "simulation.report_runs", + } <= package_tables diff --git a/migrations/versions/0003_simulation_runs_and_report_runs.py b/migrations/versions/0003_simulation_runs_and_report_runs.py new file mode 100644 index 000000000..1225bb99f --- /dev/null +++ b/migrations/versions/0003_simulation_runs_and_report_runs.py @@ -0,0 +1,328 @@ +"""Add simulation runs and report runs. + +``simulation_runs`` records each execution of a simulation for a catalog +version and year: the region and dataset rows it resolved in that version, its +status, Modal call, timings, error, trigger and attempt number, unique per +simulation, catalog version, year and attempt. Composite foreign keys through +the Public API's ``uq_regions_model_version_code`` and +``uq_datasets_id_model_version`` keep a run on its own version's rows, and need +``REFERENCES`` on ``tax_benefit_model_versions``, ``regions`` and ``datasets`` +for the migrator role. ``report_runs`` records each calculation of a report: +two simulation runs of one catalog version and year and, for an economy +report, the comparison step's status and Modal call. The runtime role may read, +insert and update both tables. Nothing writes them yet. + +Revision ID: 0003 +Revises: 0002 +Create Date: 2026-10-22 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "0003" +down_revision = "0002" +branch_labels = None +depends_on = None + +SCHEMA = "simulation" +STATUSES = "'pending', 'dispatching', 'running', 'succeeded', 'failed'" +TRIGGERS = "'initial', 'manual', 'system'" + + +def _public_api(table_and_column: str) -> str: + return f"public.{table_and_column}" + + +def _execution_columns(prefix: str, *, required: bool) -> list[sa.Column]: + return [ + sa.Column(f"{prefix}status", sa.Text(), nullable=not required), + sa.Column(f"{prefix}modal_call_id", sa.Text(), nullable=True), + sa.Column(f"{prefix}error_code", sa.Text(), nullable=True), + sa.Column(f"{prefix}error_message", sa.Text(), nullable=True), + sa.Column( + f"{prefix}check_count", + sa.Integer(), + server_default=sa.text("0"), + nullable=False, + ), + sa.Column(f"{prefix}dispatched_at", sa.DateTime(timezone=True), nullable=True), + sa.Column(f"{prefix}next_check_at", sa.DateTime(timezone=True), nullable=True), + sa.Column(f"{prefix}completed_at", sa.DateTime(timezone=True), nullable=True), + ] + + +def _status_checks( + table: str, prefix: str, *, optional: bool +) -> list[sa.CheckConstraint]: + status = f"{prefix}status" + unless_absent = f"{status} IS NULL OR " if optional else "" + return [ + sa.CheckConstraint( + f"{unless_absent}{status} IN ({STATUSES})", + name=op.f(f"ck_{table}_{prefix}status"), + ), + sa.CheckConstraint( + f"{unless_absent}" + f"({status} IN ('pending', 'dispatching') " + f"AND {prefix}modal_call_id IS NULL) " + f"OR ({status} IN ('running', 'succeeded') " + f"AND {prefix}modal_call_id IS NOT NULL) " + f"OR {status} = 'failed'", + name=op.f(f"ck_{table}_{prefix}modal_call"), + ), + sa.CheckConstraint( + f"{unless_absent}" + f"({status} IN ('succeeded', 'failed')) " + f"= ({prefix}completed_at IS NOT NULL)", + name=op.f(f"ck_{table}_{prefix}completed_at"), + ), + sa.CheckConstraint( + f"{unless_absent}({status} = 'failed') = ({prefix}error_code IS NOT NULL)", + name=op.f(f"ck_{table}_{prefix}error"), + ), + sa.CheckConstraint( + f"{unless_absent}" + f"({status} = 'running') = ({prefix}next_check_at IS NOT NULL)", + name=op.f(f"ck_{table}_{prefix}next_check_at"), + ), + sa.CheckConstraint( + f"{prefix}check_count >= 0", name=op.f(f"ck_{table}_{prefix}check_count") + ), + ] + + +def _create_simulation_runs() -> None: + op.create_unique_constraint( + "uq_simulations_id_kind", "simulations", ["id", "kind"], schema=SCHEMA + ) + op.create_table( + "simulation_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("simulation_id", sa.Uuid(), nullable=False), + sa.Column("simulation_kind", sa.Text(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column("year", sa.Integer(), nullable=False), + sa.Column("attempt", sa.Integer(), nullable=False), + sa.Column("region_code", sa.String(length=255), nullable=True), + sa.Column("dataset_id", sa.Uuid(), nullable=True), + sa.Column("trigger", sa.Text(), nullable=False), + *_execution_columns("", required=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", name=op.f("ck_simulation_runs_year") + ), + sa.CheckConstraint("attempt >= 1", name=op.f("ck_simulation_runs_attempt")), + sa.CheckConstraint( + f"trigger IN ({TRIGGERS})", name=op.f("ck_simulation_runs_trigger") + ), + sa.CheckConstraint( + "(simulation_kind = 'household' AND region_code IS NULL " + "AND dataset_id IS NULL) " + "OR (simulation_kind = 'economy' AND region_code IS NOT NULL " + "AND dataset_id IS NOT NULL)", + name=op.f("ck_simulation_runs_population_matches_kind"), + ), + *_status_checks("simulation_runs", "", optional=False), + sa.ForeignKeyConstraint( + ["simulation_id", "simulation_kind"], + [f"{SCHEMA}.simulations.id", f"{SCHEMA}.simulations.kind"], + name="fk_simulation_runs_simulation", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + [_public_api("tax_benefit_model_versions.id")], + name="fk_simulation_runs_catalog_version", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id", "region_code"], + [ + _public_api("regions.tax_benefit_model_version_id"), + _public_api("regions.code"), + ], + name="fk_simulation_runs_region", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["dataset_id", "tax_benefit_model_version_id"], + [ + _public_api("datasets.id"), + _public_api("datasets.tax_benefit_model_version_id"), + ], + name="fk_simulation_runs_dataset", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_simulation_runs"), + sa.UniqueConstraint( + "simulation_id", + "tax_benefit_model_version_id", + "year", + "attempt", + name="uq_simulation_runs_attempt", + ), + sa.UniqueConstraint( + "id", + "tax_benefit_model_version_id", + "year", + name="uq_simulation_runs_id_catalog_version_year", + ), + schema=SCHEMA, + ) + op.create_index( + "ix_simulation_runs_catalog_version_region", + "simulation_runs", + ["tax_benefit_model_version_id", "region_code"], + schema=SCHEMA, + ) + op.create_index( + "ix_simulation_runs_dataset_id", + "simulation_runs", + ["dataset_id"], + schema=SCHEMA, + ) + op.create_index( + "ix_simulation_runs_modal_call_id", + "simulation_runs", + ["modal_call_id"], + schema=SCHEMA, + postgresql_where=sa.text("modal_call_id IS NOT NULL"), + ) + op.create_index( + "ix_simulation_runs_due_checks", + "simulation_runs", + ["next_check_at"], + schema=SCHEMA, + postgresql_where=sa.text("status = 'running'"), + ) + op.create_index( + "ix_simulation_runs_dispatching", + "simulation_runs", + ["updated_at"], + schema=SCHEMA, + postgresql_where=sa.text("status = 'dispatching'"), + ) + + +def _create_report_runs() -> None: + op.create_unique_constraint( + "uq_reports_id_kind", "reports", ["id", "kind"], schema=SCHEMA + ) + simulation_run_key = [ + f"{SCHEMA}.simulation_runs.id", + f"{SCHEMA}.simulation_runs.tax_benefit_model_version_id", + f"{SCHEMA}.simulation_runs.year", + ] + op.create_table( + "report_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("report_id", sa.Uuid(), nullable=False), + sa.Column("report_kind", sa.Text(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column("year", sa.Integer(), nullable=False), + sa.Column("baseline_simulation_run_id", sa.Uuid(), nullable=False), + sa.Column("reform_simulation_run_id", sa.Uuid(), nullable=True), + sa.Column("trigger", sa.Text(), nullable=False), + *_execution_columns("comparison_", required=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + "year BETWEEN 1900 AND 2200", name=op.f("ck_report_runs_year") + ), + sa.CheckConstraint( + f"trigger IN ({TRIGGERS})", name=op.f("ck_report_runs_trigger") + ), + sa.CheckConstraint( + "(report_kind = 'economy' AND reform_simulation_run_id IS NOT NULL " + "AND comparison_status IS NOT NULL) " + "OR (report_kind = 'household' AND comparison_status IS NULL " + "AND comparison_modal_call_id IS NULL " + "AND comparison_error_code IS NULL " + "AND comparison_dispatched_at IS NULL " + "AND comparison_next_check_at IS NULL " + "AND comparison_completed_at IS NULL " + "AND comparison_check_count = 0)", + name=op.f("ck_report_runs_comparison_matches_kind"), + ), + *_status_checks("report_runs", "comparison_", optional=True), + sa.ForeignKeyConstraint( + ["report_id", "report_kind"], + [f"{SCHEMA}.reports.id", f"{SCHEMA}.reports.kind"], + name="fk_report_runs_report", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["baseline_simulation_run_id", "tax_benefit_model_version_id", "year"], + simulation_run_key, + name="fk_report_runs_baseline_simulation_run", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["reform_simulation_run_id", "tax_benefit_model_version_id", "year"], + simulation_run_key, + name="fk_report_runs_reform_simulation_run", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_report_runs"), + schema=SCHEMA, + ) + op.create_index( + "ix_report_runs_report_id_created_at", + "report_runs", + ["report_id", "created_at"], + schema=SCHEMA, + ) + for column in ("baseline_simulation_run_id", "reform_simulation_run_id"): + op.create_index( + f"ix_report_runs_{column}", "report_runs", [column], schema=SCHEMA + ) + op.create_index( + "ix_report_runs_due_comparison_checks", + "report_runs", + ["comparison_next_check_at"], + schema=SCHEMA, + postgresql_where=sa.text("comparison_status = 'running'"), + ) + op.create_index( + "ix_report_runs_comparisons_dispatching", + "report_runs", + ["updated_at"], + schema=SCHEMA, + postgresql_where=sa.text("comparison_status = 'dispatching'"), + ) + + +def upgrade() -> None: + """Create both tables; the default privileges let the runtime role use them.""" + _create_simulation_runs() + _create_report_runs() + + +def downgrade() -> None: + """Drop the tables ``upgrade`` created, with their indexes and keys.""" + for index in ( + "ix_report_runs_comparisons_dispatching", + "ix_report_runs_due_comparison_checks", + "ix_report_runs_reform_simulation_run_id", + "ix_report_runs_baseline_simulation_run_id", + "ix_report_runs_report_id_created_at", + ): + op.drop_index(index, table_name="report_runs", schema=SCHEMA) + op.drop_table("report_runs", schema=SCHEMA) + op.drop_constraint("uq_reports_id_kind", "reports", schema=SCHEMA, type_="unique") + for index in ( + "ix_simulation_runs_dispatching", + "ix_simulation_runs_due_checks", + "ix_simulation_runs_modal_call_id", + "ix_simulation_runs_dataset_id", + "ix_simulation_runs_catalog_version_region", + ): + op.drop_index(index, table_name="simulation_runs", schema=SCHEMA) + op.drop_table("simulation_runs", schema=SCHEMA) + op.drop_constraint( + "uq_simulations_id_kind", "simulations", schema=SCHEMA, type_="unique" + ) From b48db24321b0154ada5f366bb982cc69be0ff7ec Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:26:53 +0400 Subject: [PATCH 02/39] Change run statuses only with conditional updates, so a second writer changes nothing RunTransitions makes every status change of a simulation run, or of an economy report run's comparison step, one UPDATE conditional on the expected status. When two Simulation Entrypoint instances act on the same run, PostgreSQL makes the second wait for the first's row lock, re-checks the condition and changes nothing; the methods return whether they changed anything. - claim_simulation_dispatch moves runs from pending to dispatching; only the instance that wins starts the Modal call. record_simulation_dispatched records the call ID and the first check time (running), and record_simulation_dispatch_failure fails a run whose call Modal refused. - The runs one Modal call computes, such as a household report's baseline and reform, are passed together and change all together or not at all. - claim_due_simulation_calls and claim_due_comparisons claim running calls that are due a check by moving their next check time one lease ahead (SELECT ... FOR UPDATE SKIP LOCKED), so instances do not check the same call at once, and a call becomes due again if the claiming instance stops. - complete_simulation_call and complete_comparison mark the call's runs succeeded and run the caller's result writer in the same transaction; if the writer raises, nothing is recorded. fail_simulation_call and fail_comparison record an error code and a message bounded to 2,000 characters. - A comparison step can be claimed only once both of its simulation runs have succeeded. - fail_abandoned_dispatches fails runs and comparison steps left in dispatching since a given time (dispatch_abandoned), so they can be submitted again. Tests cover each change, a repeated change, eight concurrent dispatch claims (exactly one wins) and four instances claiming due calls (disjoint sets), on PostgreSQL. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 12 + .../run_transitions.py | 566 ++++++++++++++++++ .../tests/test_run_transitions.py | 491 +++++++++++++++ .../tests/transition_support.py | 49 ++ 4 files changed, 1118 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_run_transitions.py create mode 100644 libs/policyengine-simulation-persistence/tests/transition_support.py diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index a83a0d3e5..be715f3e5 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -19,6 +19,13 @@ SUCCEEDED_RUN_REUSE_WINDOW, RunStore, ) +from policyengine_simulation_persistence.run_transitions import ( + AbandonedDispatches, + DueSimulationCall, + RunError, + RunErrorCode, + RunTransitions, +) from policyengine_simulation_persistence.runs import ( NewReportRun, NewSimulationRun, @@ -48,6 +55,8 @@ "RUNTIME_ROLE", "SIMULATION_SCHEMA", "SUCCEEDED_RUN_REUSE_WINDOW", + "AbandonedDispatches", + "DueSimulationCall", "NewReport", "NewReportRun", "NewSimulation", @@ -56,8 +65,11 @@ "ReportRecord", "ReportRunRecord", "ReportStore", + "RunError", + "RunErrorCode", "RunStatus", "RunStore", + "RunTransitions", "RunTrigger", "SimulationKind", "SimulationRecord", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py new file mode 100644 index 000000000..39c6ee2c8 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -0,0 +1,566 @@ +"""Conditional status changes for simulation runs and comparison steps. + +Every change is one ``UPDATE ... WHERE = ``. When two +Simulation Entrypoint instances act on the same run at once, PostgreSQL makes +the second wait for the first's row lock, then re-checks the condition and +changes nothing. So no lease table or advisory lock is needed: + +- only the instance that moves a run from ``pending`` to ``dispatching`` + starts its Modal call, and it then records the call ID (``running``); +- the runs that share one Modal call, such as a household report's baseline + and reform, change status together, all or none; +- the recording loop claims a running call's next check by moving its + ``next_check_at`` forward, so instances do not check the same call at once, + and only one of them records its outcome; +- a call's results are written in the same transaction that marks its runs + ``succeeded``, so a run never succeeds without its results; +- a run left in ``dispatching`` (the instance stopped between claiming it and + recording its call) is failed after a while, so it can be submitted again. + +An economy report run's comparison step follows the same statuses on its own +``comparison_*`` columns, and may start only once both of its simulation runs +have succeeded. +""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING + +from sqlalchemy import and_, exists, select, update + +from policyengine_simulation_persistence.run_store import utc_now +from policyengine_simulation_persistence.runs import ( + ReportRunRecord, + RunStatus, + SimulationRunRecord, + bounded_error_message, + report_runs, + simulation_runs, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Collection + from datetime import datetime, timedelta + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, Engine + + type SimulationResultWriter = Callable[ + [Connection, tuple[SimulationRunRecord, ...]], None + ] + type ComparisonResultWriter = Callable[[Connection, ReportRunRecord], None] + + +class RunErrorCode(StrEnum): + """Why a run or comparison step failed.""" + + # Modal refused the call or never acknowledged it. + DISPATCH_FAILED = "dispatch_failed" + # The instance stopped between claiming the dispatch and recording the call. + DISPATCH_ABANDONED = "dispatch_abandoned" + # The Modal function raised. + WORKER_ERROR = "worker_error" + WORKER_TIMEOUT = "worker_timeout" + # Modal no longer holds the output, or does not know the call. + RESULT_EXPIRED = "result_expired" + CALL_NOT_FOUND = "call_not_found" + # The output failed validation against its result model. + INVALID_RESULT = "invalid_result" + # No recorder handles results of this kind of run. + NO_RESULT_RECORDER = "no_result_recorder" + + +@dataclass(frozen=True) +class RunError: + """A failure to record on a run; the message must carry no user data.""" + + code: RunErrorCode + message: str + + +@dataclass(frozen=True) +class DueSimulationCall: + """A running Modal call whose runs this instance claimed to check now.""" + + modal_call_id: str + runs: tuple[SimulationRunRecord, ...] + + +@dataclass(frozen=True) +class AbandonedDispatches: + """Runs and comparison steps failed because their dispatch never finished.""" + + simulation_run_ids: tuple[UUID, ...] + report_run_ids: tuple[UUID, ...] + + +class RunTransitions: + """Conditional status writes; each returns whether it changed anything.""" + + def __init__( + self, + engine: Engine, + *, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Use ``engine`` (the runtime role) and ``clock`` for every timestamp.""" + self.engine = engine + self._clock = clock + + # --- Simulation runs --------------------------------------------------- + + def _change_simulation_runs( + self, + run_ids: Collection[UUID], + *, + expected: RunStatus, + values: dict[str, object], + ) -> bool: + """Change every run in ``run_ids`` from ``expected``, or none of them.""" + wanted = set(run_ids) + if not wanted: + msg = "a status change needs at least one run" + raise ValueError(msg) + with self.engine.connect() as connection: + transaction = connection.begin() + changed = connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.id.in_(wanted), + simulation_runs.c.status == expected.value, + ) + .values(updated_at=self._clock(), **values) + .returning(simulation_runs.c.id) + ).scalars() + if set(changed) != wanted: + transaction.rollback() + return False + transaction.commit() + return True + + def claim_simulation_dispatch(self, run_ids: Collection[UUID]) -> bool: + """Move pending runs to ``dispatching``; only the winner starts the call. + + ``run_ids`` are the runs one Modal call computes: one run, or a + household report's baseline and reform. All must be pending. + """ + return self._change_simulation_runs( + run_ids, + expected=RunStatus.PENDING, + values={"status": RunStatus.DISPATCHING.value}, + ) + + def record_simulation_dispatched( + self, + run_ids: Collection[UUID], + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the acknowledged Modal call and move the runs to ``running``.""" + return self._change_simulation_runs( + run_ids, + expected=RunStatus.DISPATCHING, + values={ + "status": RunStatus.RUNNING.value, + "modal_call_id": modal_call_id, + "dispatched_at": self._clock(), + "next_check_at": first_check_at, + }, + ) + + def record_simulation_dispatch_failure( + self, run_ids: Collection[UUID], error: RunError + ) -> bool: + """Fail dispatching runs whose Modal call was refused.""" + return self._change_simulation_runs( + run_ids, + expected=RunStatus.DISPATCHING, + values={ + "status": RunStatus.FAILED.value, + "error_code": error.code.value, + "error_message": bounded_error_message(error.message), + "completed_at": self._clock(), + }, + ) + + def claim_due_simulation_calls( + self, *, limit: int, lease: timedelta + ) -> list[DueSimulationCall]: + """Claim up to ``limit`` running runs that are due a check, by call. + + Each claimed run's next check moves ``lease`` ahead, so no other + instance checks it until the lease ends; if this instance stops, the + run becomes due again then. Rows another instance holds are skipped. + Two runs of one call are normally claimed together; if the limit + splits them, two instances may ask Modal about the same call, and the + conditional writes still record its outcome once. + """ + now = self._clock() + due = ( + select(simulation_runs.c.id) + .where( + simulation_runs.c.status == RunStatus.RUNNING.value, + simulation_runs.c.next_check_at <= now, + ) + .order_by(simulation_runs.c.next_check_at) + .limit(limit) + .with_for_update(skip_locked=True) + .scalar_subquery() + ) + with self.engine.begin() as connection: + rows = ( + connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.id.in_(due), + simulation_runs.c.status == RunStatus.RUNNING.value, + simulation_runs.c.next_check_at <= now, + ) + .values( + next_check_at=now + lease, + check_count=simulation_runs.c.check_count + 1, + updated_at=now, + ) + .returning(*simulation_runs.c) + ) + .mappings() + .all() + ) + runs_by_call: dict[str, list[SimulationRunRecord]] = defaultdict(list) + for row in rows: + run = SimulationRunRecord.from_row(row) + if run.modal_call_id is not None: + runs_by_call[run.modal_call_id].append(run) + return [ + DueSimulationCall(modal_call_id=call_id, runs=tuple(runs)) + for call_id, runs in sorted(runs_by_call.items()) + ] + + def _change_simulation_call( + self, + connection: Connection, + modal_call_id: str, + values: dict[str, object], + ) -> tuple[SimulationRunRecord, ...]: + rows = ( + connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.modal_call_id == modal_call_id, + simulation_runs.c.status == RunStatus.RUNNING.value, + ) + .values(updated_at=self._clock(), **values) + .returning(*simulation_runs.c) + ) + .mappings() + .all() + ) + return tuple(SimulationRunRecord.from_row(row) for row in rows) + + def record_simulation_call_running( + self, modal_call_id: str, *, next_check_at: datetime + ) -> bool: + """Set the next check of a call's runs that are still running.""" + with self.engine.begin() as connection: + return bool( + self._change_simulation_call( + connection, modal_call_id, {"next_check_at": next_check_at} + ) + ) + + def complete_simulation_call( + self, modal_call_id: str, write_results: SimulationResultWriter + ) -> bool: + """Mark a call's running runs succeeded and write their results. + + ``write_results`` receives the transaction and the succeeded runs, and + writes their result rows. If it raises, nothing is recorded and the + exception propagates. Return False, writing nothing, when no run of + the call is still running. + """ + with self.engine.begin() as connection: + runs = self._change_simulation_call( + connection, + modal_call_id, + { + "status": RunStatus.SUCCEEDED.value, + "completed_at": self._clock(), + "next_check_at": None, + }, + ) + if not runs: + return False + write_results(connection, runs) + return True + + def fail_simulation_call(self, modal_call_id: str, error: RunError) -> bool: + """Mark a call's running runs failed; False if none is still running.""" + with self.engine.begin() as connection: + return bool( + self._change_simulation_call( + connection, + modal_call_id, + { + "status": RunStatus.FAILED.value, + "error_code": error.code.value, + "error_message": bounded_error_message(error.message), + "completed_at": self._clock(), + "next_check_at": None, + }, + ) + ) + + # --- Comparison steps of economy report runs ---------------------------- + + def _change_comparison( + self, + connection: Connection, + report_run_id: UUID, + *, + expected: RunStatus, + values: dict[str, object], + conditions: tuple[ColumnElement[bool], ...] = (), + ) -> ReportRunRecord | None: + row = ( + connection.execute( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + report_runs.c.comparison_status == expected.value, + *conditions, + ) + .values(updated_at=self._clock(), **values) + .returning(*report_runs.c) + ) + .mappings() + .first() + ) + return None if row is None else ReportRunRecord.from_row(row) + + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: + """Move a pending comparison step to ``dispatching``. + + Only the winner starts the comparison's Modal call. The step can start + only once both of its simulation runs have succeeded. + """ + baseline = simulation_runs.alias("baseline_run") + reform = simulation_runs.alias("reform_run") + both_runs_succeeded = ( + exists().where( + and_( + baseline.c.id == report_runs.c.baseline_simulation_run_id, + baseline.c.status == RunStatus.SUCCEEDED.value, + ) + ), + exists().where( + and_( + reform.c.id == report_runs.c.reform_simulation_run_id, + reform.c.status == RunStatus.SUCCEEDED.value, + ) + ), + ) + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.PENDING, + values={"comparison_status": RunStatus.DISPATCHING.value}, + conditions=both_runs_succeeded, + ) + is not None + ) + + def record_comparison_dispatched( + self, + report_run_id: UUID, + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the comparison's acknowledged Modal call; it is now running.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.DISPATCHING, + values={ + "comparison_status": RunStatus.RUNNING.value, + "comparison_modal_call_id": modal_call_id, + "comparison_dispatched_at": self._clock(), + "comparison_next_check_at": first_check_at, + }, + ) + is not None + ) + + def record_comparison_dispatch_failure( + self, report_run_id: UUID, error: RunError + ) -> bool: + """Fail a dispatching comparison step whose Modal call was refused.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.DISPATCHING, + values=self._comparison_failure(error), + ) + is not None + ) + + def claim_due_comparisons( + self, *, limit: int, lease: timedelta + ) -> list[ReportRunRecord]: + """Claim up to ``limit`` running comparison steps that are due a check.""" + now = self._clock() + due = ( + select(report_runs.c.id) + .where( + report_runs.c.comparison_status == RunStatus.RUNNING.value, + report_runs.c.comparison_next_check_at <= now, + ) + .order_by(report_runs.c.comparison_next_check_at) + .limit(limit) + .with_for_update(skip_locked=True) + .scalar_subquery() + ) + with self.engine.begin() as connection: + rows = ( + connection.execute( + update(report_runs) + .where( + report_runs.c.id.in_(due), + report_runs.c.comparison_status == RunStatus.RUNNING.value, + report_runs.c.comparison_next_check_at <= now, + ) + .values( + comparison_next_check_at=now + lease, + comparison_check_count=report_runs.c.comparison_check_count + 1, + updated_at=now, + ) + .returning(*report_runs.c) + ) + .mappings() + .all() + ) + return [ReportRunRecord.from_row(row) for row in rows] + + def record_comparison_running( + self, report_run_id: UUID, *, next_check_at: datetime + ) -> bool: + """Set the next check of a comparison step that is still running.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.RUNNING, + values={"comparison_next_check_at": next_check_at}, + ) + is not None + ) + + def complete_comparison( + self, report_run_id: UUID, write_results: ComparisonResultWriter + ) -> bool: + """Mark a running comparison step succeeded and write its results. + + ``write_results`` runs in the same transaction; if it raises, nothing + is recorded and the exception propagates. + """ + with self.engine.begin() as connection: + report_run = self._change_comparison( + connection, + report_run_id, + expected=RunStatus.RUNNING, + values={ + "comparison_status": RunStatus.SUCCEEDED.value, + "comparison_completed_at": self._clock(), + "comparison_next_check_at": None, + }, + ) + if report_run is None: + return False + write_results(connection, report_run) + return True + + def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: + """Mark a running comparison step failed; False if it is not running.""" + with self.engine.begin() as connection: + return ( + self._change_comparison( + connection, + report_run_id, + expected=RunStatus.RUNNING, + values=self._comparison_failure(error), + ) + is not None + ) + + def _comparison_failure(self, error: RunError) -> dict[str, object]: + return { + "comparison_status": RunStatus.FAILED.value, + "comparison_error_code": error.code.value, + "comparison_error_message": bounded_error_message(error.message), + "comparison_completed_at": self._clock(), + "comparison_next_check_at": None, + } + + # --- Dispatches that never finished ------------------------------------ + + def fail_abandoned_dispatches( + self, *, claimed_before: datetime + ) -> AbandonedDispatches: + """Fail runs and comparison steps left in ``dispatching``. + + A run stays ``dispatching`` only if the instance that claimed it + stopped before recording its Modal call. Failing it, with + ``dispatch_abandoned``, lets the same simulation be run again. + """ + now = self._clock() + message = ( + "The instance that claimed the dispatch stopped before it recorded " + "the Modal call." + ) + with self.engine.begin() as connection: + failed_runs = connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.status == RunStatus.DISPATCHING.value, + simulation_runs.c.updated_at < claimed_before, + ) + .values( + status=RunStatus.FAILED.value, + error_code=RunErrorCode.DISPATCH_ABANDONED.value, + error_message=message, + completed_at=now, + updated_at=now, + ) + .returning(simulation_runs.c.id) + ).scalars() + simulation_run_ids = tuple(failed_runs) + report_run_ids = tuple( + connection.execute( + update(report_runs) + .where( + report_runs.c.comparison_status == RunStatus.DISPATCHING.value, + report_runs.c.updated_at < claimed_before, + ) + .values( + updated_at=now, + **self._comparison_failure( + RunError(RunErrorCode.DISPATCH_ABANDONED, message) + ), + ) + .returning(report_runs.c.id) + ).scalars() + ) + return AbandonedDispatches( + simulation_run_ids=simulation_run_ids, report_run_ids=report_run_ids + ) diff --git a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py new file mode 100644 index 000000000..689683849 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py @@ -0,0 +1,491 @@ +"""Conditional status changes: a second writer never changes anything.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier + +import pytest +from policyengine_simulation_persistence import ( + NewReportRun, + ReportKind, + RunError, + RunErrorCode, + RunStatus, + RunTransitions, + RunTrigger, + create_simulation_engine, +) +from report_support import economy_pair, economy_report, household_baseline +from report_support import report_store as definition_report_store +from run_support import economy_run, household_run +from simulation_support import household_simulation, simulation_store +from transition_support import ManualClock, clocked_stores + +pytestmark = pytest.mark.postgres + +LEASE = timedelta(minutes=1) +FIRST_CHECK = timedelta(seconds=2) +WORKER_ERROR = RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + + +def _pending_run(stores, database, rows): + baseline, _ = economy_pair(database, rows) + return stores.runs.start_simulation_run(economy_run(baseline, rows)).run + + +def _running_run(stores, database, rows, call_id="fc-economy"): + run = _pending_run(stores, database, rows) + assert stores.transitions.claim_simulation_dispatch([run.id]) + assert stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=call_id, first_check_at=stores.clock() + FIRST_CHECK + ) + return stores.runs.get_simulation_run(run.id) + + +def _household_pair(stores, database, rows): + baseline = household_baseline(database, rows) + reform = simulation_store(database).create( + household_simulation(rows, policy_id=rows.us_policy_id) + ) + return ( + stores.runs.start_simulation_run(household_run(baseline, rows)).run, + stores.runs.start_simulation_run(household_run(reform, rows)).run, + ) + + +def _no_results(_connection, _records) -> None: + return None + + +def test_only_one_claim_moves_a_run_to_dispatching( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _pending_run(stores, migrated_database, public_api_rows) + + assert stores.transitions.claim_simulation_dispatch([run.id]) is True + assert stores.transitions.claim_simulation_dispatch([run.id]) is False + + claimed = stores.runs.get_simulation_run(run.id) + assert claimed.status is RunStatus.DISPATCHING + assert claimed.modal_call_id is None + + +def test_concurrent_claims_start_one_modal_call( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _pending_run(stores, migrated_database, public_api_rows) + instances = 8 + barrier = Barrier(instances) + + def claim(_instance: int) -> bool: + transitions = RunTransitions( + create_simulation_engine(migrated_database.runtime_url) + ) + try: + barrier.wait() + return transitions.claim_simulation_dispatch([run.id]) + finally: + transitions.engine.dispose() + + with ThreadPoolExecutor(max_workers=instances) as pool: + outcomes = list(pool.map(claim, range(instances))) + + assert outcomes.count(True) == 1 + + +def test_a_dispatched_run_records_its_call_and_first_check( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + + assert run.status is RunStatus.RUNNING + assert run.modal_call_id == "fc-economy" + assert run.dispatched_at == stores.clock() + assert run.next_check_at == stores.clock() + FIRST_CHECK + assert not stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id="fc-other", first_check_at=stores.clock() + ) + + +def test_a_refused_dispatch_fails_the_run_and_frees_the_setup( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _pending_run(stores, migrated_database, public_api_rows) + stores.transitions.claim_simulation_dispatch([run.id]) + + assert stores.transitions.record_simulation_dispatch_failure( + [run.id], RunError(RunErrorCode.DISPATCH_FAILED, "Modal refused the call.") + ) + + failed = stores.runs.get_simulation_run(run.id) + assert failed.status is RunStatus.FAILED + assert failed.error_code == "dispatch_failed" + assert failed.completed_at == stores.clock() + simulation = simulation_store(migrated_database).get(run.simulation_id) + retry = stores.runs.start_simulation_run(economy_run(simulation, public_api_rows)) + assert retry.created is True + assert retry.run.attempt == 2 + + +def test_a_household_pair_shares_one_call_and_changes_together( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + baseline_run, reform_run = _household_pair( + stores, migrated_database, public_api_rows + ) + pair = [baseline_run.id, reform_run.id] + written = [] + + assert stores.transitions.claim_simulation_dispatch(pair) + assert stores.transitions.record_simulation_dispatched( + pair, modal_call_id="fc-household", first_check_at=stores.clock() + ) + assert stores.transitions.complete_simulation_call( + "fc-household", lambda _connection, runs: written.extend(runs) + ) + + assert {run.id for run in written} == set(pair) + for run_id in pair: + run = stores.runs.get_simulation_run(run_id) + assert run.status is RunStatus.SUCCEEDED + assert run.modal_call_id == "fc-household" + assert run.next_check_at is None + + +def test_a_group_claim_changes_nothing_unless_every_run_is_pending( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + baseline_run, reform_run = _household_pair( + stores, migrated_database, public_api_rows + ) + assert stores.transitions.claim_simulation_dispatch([baseline_run.id]) + + assert not stores.transitions.claim_simulation_dispatch( + [baseline_run.id, reform_run.id] + ) + + assert stores.runs.get_simulation_run(reform_run.id).status is RunStatus.PENDING + with pytest.raises(ValueError, match="at least one run"): + stores.transitions.claim_simulation_dispatch([]) + + +def test_due_calls_are_claimed_for_one_lease_at_a_time( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + transitions = stores.transitions + + assert transitions.claim_due_simulation_calls(limit=10, lease=LEASE) == [] + stores.clock.advance(FIRST_CHECK) + claimed = transitions.claim_due_simulation_calls(limit=10, lease=LEASE) + assert [call.modal_call_id for call in claimed] == ["fc-economy"] + assert claimed[0].runs[0].id == run.id + assert claimed[0].runs[0].check_count == 1 + assert claimed[0].runs[0].next_check_at == stores.clock() + LEASE + assert transitions.claim_due_simulation_calls(limit=10, lease=LEASE) == [] + + # The lease ends when the claiming instance stops before recording anything. + stores.clock.advance(LEASE) + reclaimed = transitions.claim_due_simulation_calls(limit=10, lease=LEASE) + assert reclaimed[0].runs[0].check_count == 2 + + +def test_concurrent_instances_claim_disjoint_calls( + migrated_database, public_api_rows +) -> None: + clock = ManualClock() + stores = clocked_stores(migrated_database, clock) + simulations = simulation_store(migrated_database) + run_ids = set() + for index in range(12): + simulation = simulations.create( + household_simulation(public_api_rows, legacy_simulation_id=index + 1) + ) + run = stores.runs.start_simulation_run( + household_run(simulation, public_api_rows) + ).run + stores.transitions.claim_simulation_dispatch([run.id]) + stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=f"fc-{index}", first_check_at=clock() + ) + run_ids.add(run.id) + instances = 4 + barrier = Barrier(instances) + + def claim(_instance: int) -> list[str]: + transitions = RunTransitions( + create_simulation_engine(migrated_database.runtime_url), clock=clock + ) + try: + barrier.wait() + return [ + call.modal_call_id + for call in transitions.claim_due_simulation_calls(limit=5, lease=LEASE) + ] + finally: + transitions.engine.dispose() + + with ThreadPoolExecutor(max_workers=instances) as pool: + claims = list(pool.map(claim, range(instances))) + + claimed_calls = [call for instance_claims in claims for call in instance_claims] + assert len(claimed_calls) == len(set(claimed_calls)) == len(run_ids) + + +def test_a_running_call_gets_its_next_check(migrated_database, public_api_rows) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + later = stores.clock() + timedelta(seconds=30) + + assert stores.transitions.record_simulation_call_running( + "fc-economy", next_check_at=later + ) + + assert stores.runs.get_simulation_run(run.id).next_check_at == later + assert not stores.transitions.record_simulation_call_running( + "fc-unknown", next_check_at=later + ) + + +def test_results_are_written_in_the_transaction_that_marks_success( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + + def failing_writer(_connection, _runs) -> None: + msg = "the result rows could not be written" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="result rows"): + stores.transitions.complete_simulation_call("fc-economy", failing_writer) + + assert stores.runs.get_simulation_run(run.id).status is RunStatus.RUNNING + assert stores.transitions.complete_simulation_call("fc-economy", _no_results) + succeeded = stores.runs.get_simulation_run(run.id) + assert succeeded.status is RunStatus.SUCCEEDED + assert succeeded.completed_at == stores.clock() + + +def test_a_finished_call_is_recorded_once(migrated_database, public_api_rows) -> None: + stores = clocked_stores(migrated_database) + _running_run(stores, migrated_database, public_api_rows) + calls = [] + + assert stores.transitions.complete_simulation_call("fc-economy", _no_results) + + assert not stores.transitions.complete_simulation_call( + "fc-economy", lambda _connection, runs: calls.append(runs) + ) + assert calls == [] + assert not stores.transitions.fail_simulation_call("fc-economy", WORKER_ERROR) + + +def test_a_failed_call_records_a_bounded_error( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + run = _running_run(stores, migrated_database, public_api_rows) + + assert stores.transitions.fail_simulation_call( + "fc-economy", RunError(RunErrorCode.WORKER_ERROR, "x" * 5000) + ) + + failed = stores.runs.get_simulation_run(run.id) + assert failed.status is RunStatus.FAILED + assert failed.error_code == "worker_error" + assert len(failed.error_message) == 2000 + assert failed.next_check_at is None + + +def _economy_report_run(stores, database, rows): + baseline, reform = economy_pair(database, rows) + report = definition_report_store(database).create(economy_report(baseline, reform)) + runs = [ + stores.runs.start_simulation_run(economy_run(simulation, rows)).run + for simulation in (baseline, reform) + ] + report_run = stores.runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=runs[0].id, + reform_simulation_run_id=runs[1].id, + trigger=RunTrigger.INITIAL, + ) + ) + return report_run, runs + + +def _succeed(stores, run, call_id) -> None: + stores.transitions.claim_simulation_dispatch([run.id]) + stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=call_id, first_check_at=stores.clock() + ) + assert stores.transitions.complete_simulation_call(call_id, _no_results) + + +def test_a_comparison_starts_only_after_both_runs_succeed( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + _succeed(stores, baseline_run, "fc-baseline") + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + _succeed(stores, reform_run, "fc-reform") + + assert stores.transitions.claim_comparison_dispatch(report_run.id) + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + assert ( + stores.runs.get_report_run(report_run.id).comparison_status + is RunStatus.DISPATCHING + ) + + +def test_a_comparison_runs_and_records_its_results_once( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + transitions = stores.transitions + transitions.claim_comparison_dispatch(report_run.id) + written = [] + + assert transitions.record_comparison_dispatched( + report_run.id, modal_call_id="fc-comparison", first_check_at=stores.clock() + ) + claimed = transitions.claim_due_comparisons(limit=10, lease=LEASE) + assert [due.id for due in claimed] == [report_run.id] + assert claimed[0].comparison_check_count == 1 + assert transitions.claim_due_comparisons(limit=10, lease=LEASE) == [] + assert transitions.record_comparison_running( + report_run.id, next_check_at=stores.clock() + LEASE + ) + assert transitions.complete_comparison( + report_run.id, lambda _connection, record: written.append(record) + ) + + assert [record.id for record in written] == [report_run.id] + finished = stores.runs.get_report_run(report_run.id) + assert finished.comparison_status is RunStatus.SUCCEEDED + assert finished.comparison_modal_call_id == "fc-comparison" + assert finished.comparison_next_check_at is None + assert not transitions.complete_comparison(report_run.id, _no_results) + assert not transitions.fail_comparison(report_run.id, WORKER_ERROR) + + +def test_a_comparison_failure_is_recorded(migrated_database, public_api_rows) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + transitions = stores.transitions + transitions.claim_comparison_dispatch(report_run.id) + transitions.record_comparison_dispatched( + report_run.id, modal_call_id="fc-comparison", first_check_at=stores.clock() + ) + + assert transitions.fail_comparison(report_run.id, WORKER_ERROR) + + failed = stores.runs.get_report_run(report_run.id) + assert failed.comparison_status is RunStatus.FAILED + assert failed.comparison_error_code == "worker_error" + assert failed.comparison_completed_at == stores.clock() + + +def test_a_refused_comparison_dispatch_fails_the_step( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + stores.transitions.claim_comparison_dispatch(report_run.id) + + assert stores.transitions.record_comparison_dispatch_failure( + report_run.id, RunError(RunErrorCode.DISPATCH_FAILED, "Modal refused it.") + ) + + failed = stores.runs.get_report_run(report_run.id) + assert failed.comparison_status is RunStatus.FAILED + assert failed.comparison_error_code == "dispatch_failed" + assert failed.comparison_modal_call_id is None + + +def test_household_report_runs_have_no_comparison_to_claim( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + baseline_run, reform_run = _household_pair( + stores, migrated_database, public_api_rows + ) + baseline = simulation_store(migrated_database).get(baseline_run.simulation_id) + reform = simulation_store(migrated_database).get(reform_run.simulation_id) + report = definition_report_store(migrated_database).create( + economy_report(baseline, reform, kind=ReportKind.HOUSEHOLD) + ) + report_run = stores.runs.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + ) + ) + + assert not stores.transitions.claim_comparison_dispatch(report_run.id) + + +def test_dispatches_left_for_ten_minutes_are_failed( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + stores.transitions.claim_comparison_dispatch(report_run.id) + stalled = _pending_run(stores, migrated_database, public_api_rows) + stores.transitions.claim_simulation_dispatch([stalled.id]) + stores.clock.advance(timedelta(minutes=10, seconds=1)) + recent = household_run( + household_baseline(migrated_database, public_api_rows), public_api_rows + ) + recent_run = stores.runs.start_simulation_run(recent).run + stores.transitions.claim_simulation_dispatch([recent_run.id]) + + abandoned = stores.transitions.fail_abandoned_dispatches( + claimed_before=stores.clock() - timedelta(minutes=10) + ) + + assert abandoned.simulation_run_ids == (stalled.id,) + assert abandoned.report_run_ids == (report_run.id,) + failed_run = stores.runs.get_simulation_run(stalled.id) + assert failed_run.status is RunStatus.FAILED + assert failed_run.error_code == "dispatch_abandoned" + assert ( + stores.runs.get_report_run(report_run.id).comparison_error_code + == "dispatch_abandoned" + ) + assert stores.runs.get_simulation_run(recent_run.id).status is RunStatus.DISPATCHING diff --git a/libs/policyengine-simulation-persistence/tests/transition_support.py b/libs/policyengine-simulation-persistence/tests/transition_support.py new file mode 100644 index 000000000..426007b8f --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/transition_support.py @@ -0,0 +1,49 @@ +"""A clock the tests move by hand, and stores that share it.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta + +from policyengine_simulation_persistence import ( + RunStore, + RunTransitions, + create_simulation_engine, +) + + +@dataclass +class ManualClock: + """Return a fixed time until the test advances it.""" + + now: datetime = field( + default_factory=lambda: datetime(2026, 10, 22, 12, 0, tzinfo=UTC) + ) + + def __call__(self) -> datetime: + """Return the current fake time.""" + return self.now + + def advance(self, interval: timedelta) -> None: + """Move the fake time forward by ``interval``.""" + self.now += interval + + +@dataclass(frozen=True) +class ClockedStores: + """A run store and transitions on one runtime engine and one clock.""" + + runs: RunStore + transitions: RunTransitions + clock: ManualClock + + +def clocked_stores(database, clock: ManualClock | None = None) -> ClockedStores: + """Return stores that connect as the runtime role and share ``clock``.""" + clock = clock or ManualClock() + engine = create_simulation_engine(database.runtime_url) + return ClockedStores( + runs=RunStore(engine, clock=clock), + transitions=RunTransitions(engine, clock=clock), + clock=clock, + ) From ca6f2022a0c1972c035b7ad6d320b048398cfb3a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:31:52 +0400 Subject: [PATCH 03/39] Start each run's Modal call exactly once from the Simulation Entrypoint RunDispatchService starts the Modal call of simulation runs, or of an economy report run's comparison step, with three conditional writes around Modal's spawn: claim the runs (pending to dispatching), spawn the call and wait only for Modal's acknowledgement, then record the call ID and the first check two seconds later (running). Only the caller whose claim changed the runs spawns; every other caller gets NOT_CLAIMED. A household report's baseline and reform runs are started together under one call. A call Modal refuses fails the runs with dispatch_failed and the error's type only, never its message. A comparison can be claimed only once both of its simulation runs have succeeded. modal_calls holds the main path's own Modal adapters, separate from the Stage 12 comparison path's: ModalCallDispatcher spawns a deployed function with the observability context as its separate keyword argument, and ModalCallInspector asks about a call with get(timeout=0), which returns at once. No output yet (including a call Modal is re-running after pre-emption) reads as running; a timed-out, expired, unknown or raising call reads as failed with its error type only; a problem reaching Modal raises CallInspectionError so the call is asked about again later. No route starts calls yet. Co-Authored-By: Claude Opus 5.5 --- .../modal_calls.py | 154 ++++++++++++ .../run_dispatch.py | 198 +++++++++++++++ .../tests/run_fakes.py | 221 +++++++++++++++++ .../tests/test_run_dispatch.py | 233 ++++++++++++++++++ 4 files changed, 806 insertions(+) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py create mode 100644 projects/policyengine-simulation-entry/tests/run_fakes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_run_dispatch.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py new file mode 100644 index 000000000..92f79f0d0 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py @@ -0,0 +1,154 @@ +"""Start Modal calls and ask about them without waiting, for simulation runs. + +The Simulation Entrypoint starts each run's Modal call with ``spawn`` and later +asks Modal about it with ``get(timeout=0)``, which returns at once. A call +Modal pre-empted is re-run on the same input under the same call ID, so it +reads as still running and finishes later: pre-emption is never a failure. + +These adapters belong to the main path. The Stage 12 comparison path keeps its +own Modal code; the two share nothing. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Protocol + +import modal +from policyengine_simulation_persistence import RunError, RunErrorCode + + +@dataclass(frozen=True) +class ModalTarget: + """A deployed Modal function: its application and function names.""" + + application: str + function: str + + +@dataclass(frozen=True) +class ModalRequest: + """One call to start: the target, its positional arguments and telemetry. + + ``arguments`` hold calculation inputs and functional metadata only. + ``observability_context`` travels as Modal's separate keyword argument. + """ + + target: ModalTarget + arguments: tuple[object, ...] + observability_context: dict[str, Any] | None = None + + +@dataclass(frozen=True) +class CallRunning: + """The call has produced no output yet, or Modal is re-running it.""" + + +@dataclass(frozen=True) +class CallSucceeded: + """The call returned ``output``.""" + + output: object + + +@dataclass(frozen=True) +class CallFailed: + """The call failed; the error carries no calculation data.""" + + error: RunError + + +type CallOutcome = CallRunning | CallSucceeded | CallFailed + + +class CallInspectionError(RuntimeError): + """Modal could not be asked about a call; ask again later.""" + + +class CallDispatcher(Protocol): + """Start a Modal call without waiting for it.""" + + async def spawn(self, request: ModalRequest) -> str: + """Return the call ID once Modal acknowledges the call.""" + ... + + +class CallInspector(Protocol): + """Ask Modal about a call without waiting for it.""" + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the call's outcome; raise CallInspectionError if unknown.""" + ... + + +class ModalCallDispatcher: + """Spawn deployed Modal functions of one Modal environment.""" + + def __init__(self, environment: str) -> None: + """Look functions up in ``environment``.""" + self._environment = environment + + async def spawn(self, request: ModalRequest) -> str: + """Spawn the request's function and return Modal's call ID. + + Raise RuntimeError when Modal acknowledges the call without an ID. + """ + function = modal.Function.from_name( + request.target.application, + request.target.function, + environment_name=self._environment, + ) + call = await function.spawn.aio( + *request.arguments, + observability_context=request.observability_context, + ) + modal_call_id = getattr(call, "object_id", None) + if not isinstance(modal_call_id, str) or not modal_call_id: + msg = "Modal acknowledged the call without a call ID" + raise RuntimeError(msg) + return modal_call_id + + +# Problems reaching Modal: the call's state is unknown, so ask again later. +TRANSIENT_MODAL_ERRORS: tuple[type[Exception], ...] = ( + modal.exception.ConnectionError, + modal.exception.AuthError, + modal.exception.ServiceError, + modal.exception.ResourceExhaustedError, + modal.exception.InternalError, +) + + +def _failed(code: RunErrorCode, message: str) -> CallFailed: + return CallFailed(RunError(code, message)) + + +class ModalCallInspector: + """Read a Modal call's output if it has one, without waiting.""" + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the call's outcome, or raise CallInspectionError.""" + try: + call = modal.FunctionCall.from_id(modal_call_id) + output = await call.get.aio(timeout=0) + # Modal raises the built-in TimeoutError while the call has no output. + except TimeoutError: + return CallRunning() + except modal.exception.OutputExpiredError: + return _failed( + RunErrorCode.RESULT_EXPIRED, "Modal no longer holds the call's output." + ) + except modal.exception.FunctionTimeoutError: + return _failed(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + except modal.exception.NotFoundError: + return _failed(RunErrorCode.CALL_NOT_FOUND, "Modal does not know the call.") + except TRANSIENT_MODAL_ERRORS as error: + raise CallInspectionError(type(error).__name__) from error + # The call re-raises whatever the worker raised. Record only its type: + # the message may contain calculation inputs. + except Exception as error: # noqa: BLE001 # every worker failure fails the run + return _failed( + RunErrorCode.WORKER_ERROR, + f"The Modal call raised {type(error).__name__}.", + ) + return CallSucceeded(output) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py new file mode 100644 index 000000000..76eb398bb --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -0,0 +1,198 @@ +"""Start each simulation run's, or comparison step's, Modal call exactly once. + +Starting a call takes three conditional writes around Modal's ``spawn``: + +1. claim: move the runs from ``pending`` to ``dispatching``. Only the instance + whose claim changes them goes on; every other caller stops here. +2. spawn the Modal call and wait only for Modal's acknowledgement. +3. record the call ID, moving the runs to ``running`` with their first check + due shortly; or, when Modal refused the call, fail them. + +If the instance stops between 1 and 3, the runs stay ``dispatching`` until the +recording loop fails them as abandoned. No route starts calls yet. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from datetime import timedelta +from enum import StrEnum +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_persistence import RunError, RunErrorCode +from policyengine_simulation_persistence.run_store import utc_now + +if TYPE_CHECKING: + from collections.abc import Callable, Collection + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_entry.modal_calls import CallDispatcher, ModalRequest + +logger = logging.getLogger(__name__) + +# The first check of a new call; the recording loop backs off from there. +FIRST_CHECK_DELAY = timedelta(seconds=2) + + +class DispatchOutcome(StrEnum): + """What one attempt to start a call did.""" + + # This caller claimed the runs and Modal acknowledged the call. + STARTED = "started" + # Another caller claimed the runs first, or they were no longer pending. + NOT_CLAIMED = "not_claimed" + # Modal refused the call; the runs are failed. + DISPATCH_FAILED = "dispatch_failed" + # Modal acknowledged the call, but the runs had stopped dispatching before + # the call ID could be recorded. + NOT_RECORDED = "not_recorded" + + +class DispatchTransitions(Protocol): + """The ``RunTransitions`` writes that starting a call needs.""" + + def claim_simulation_dispatch(self, run_ids: Collection[UUID]) -> bool: + """Move pending runs to dispatching.""" + ... + + def record_simulation_dispatched( + self, + run_ids: Collection[UUID], + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the call and move the runs to running.""" + ... + + def record_simulation_dispatch_failure( + self, run_ids: Collection[UUID], error: RunError + ) -> bool: + """Fail dispatching runs.""" + ... + + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: + """Move a pending comparison step to dispatching.""" + ... + + def record_comparison_dispatched( + self, + report_run_id: UUID, + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the comparison's call and move it to running.""" + ... + + def record_comparison_dispatch_failure( + self, report_run_id: UUID, error: RunError + ) -> bool: + """Fail a dispatching comparison step.""" + ... + + +@dataclass(frozen=True) +class _DispatchSteps: + """The three writes around one spawn, bound to the runs they change.""" + + claim: Callable[[], bool] + record: Callable[[str, datetime], bool] + fail: Callable[[RunError], bool] + + +@dataclass +class RunDispatchService: + """Start Modal calls for simulation runs and comparison steps.""" + + transitions: DispatchTransitions + dispatcher: CallDispatcher + clock: Callable[[], datetime] = utc_now + + async def start_simulation_runs( + self, run_ids: Collection[UUID], request: ModalRequest + ) -> DispatchOutcome: + """Start the one Modal call that computes ``run_ids``. + + Pass a household report's baseline and reform runs together: one call + computes both, and they change status together. + """ + runs = frozenset(run_ids) + return await self._start( + _DispatchSteps( + claim=lambda: self.transitions.claim_simulation_dispatch(runs), + record=lambda call_id, first_check_at: ( + self.transitions.record_simulation_dispatched( + runs, modal_call_id=call_id, first_check_at=first_check_at + ) + ), + fail=lambda error: self.transitions.record_simulation_dispatch_failure( + runs, error + ), + ), + request, + subject={"simulation_run_ids": sorted(str(run_id) for run_id in runs)}, + ) + + async def start_comparison( + self, report_run_id: UUID, request: ModalRequest + ) -> DispatchOutcome: + """Start an economy report run's comparison step. + + The claim succeeds only once both of its simulation runs have + succeeded. + """ + return await self._start( + _DispatchSteps( + claim=lambda: self.transitions.claim_comparison_dispatch(report_run_id), + record=lambda call_id, first_check_at: ( + self.transitions.record_comparison_dispatched( + report_run_id, + modal_call_id=call_id, + first_check_at=first_check_at, + ) + ), + fail=lambda error: self.transitions.record_comparison_dispatch_failure( + report_run_id, error + ), + ), + request, + subject={"report_run_id": str(report_run_id)}, + ) + + async def _start( + self, + steps: _DispatchSteps, + request: ModalRequest, + *, + subject: dict[str, object], + ) -> DispatchOutcome: + if not await asyncio.to_thread(steps.claim): + return DispatchOutcome.NOT_CLAIMED + try: + modal_call_id = await self.dispatcher.spawn(request) + # Any failure to start the call fails the runs; they can be run again. + except Exception as error: # noqa: BLE001 # recorded on the runs + logger.warning( + "simulation_run_dispatch_failed", + extra={**subject, "error_type": type(error).__name__}, + ) + await asyncio.to_thread( + steps.fail, + RunError( + RunErrorCode.DISPATCH_FAILED, + f"Modal did not accept the call ({type(error).__name__}).", + ), + ) + return DispatchOutcome.DISPATCH_FAILED + recorded = await asyncio.to_thread( + steps.record, modal_call_id, self.clock() + FIRST_CHECK_DELAY + ) + logger.info( + "simulation_run_dispatched", + extra={**subject, "modal_call_id": modal_call_id, "recorded": recorded}, + ) + return DispatchOutcome.STARTED if recorded else DispatchOutcome.NOT_RECORDED diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py new file mode 100644 index 000000000..1af6983b4 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -0,0 +1,221 @@ +"""In-memory stand-ins for the run transitions and Modal, for the run tests. + +``FakeRunTransitions`` keeps each simulation run's and comparison step's status +and applies the same conditional rules as ``RunTransitions``: a change applies +only from the expected status, and the runs of one call change together. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from datetime import UTC, datetime, timedelta +from threading import Lock +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from policyengine_simulation_entry.modal_calls import ModalRequest, ModalTarget +from policyengine_simulation_persistence import RunError, RunStatus + +if TYPE_CHECKING: + from collections.abc import Collection + + +@dataclass +class ManualClock: + """Return a fixed time until a test advances it.""" + + now: datetime = field( + default_factory=lambda: datetime(2026, 10, 22, 12, 0, tzinfo=UTC) + ) + + def __call__(self) -> datetime: + """Return the current fake time.""" + return self.now + + def advance(self, interval: timedelta) -> None: + """Move the fake time forward.""" + self.now += interval + + +@dataclass +class FakeExecution: + """One simulation run, or one comparison step, as the fake keeps it.""" + + id: UUID + kind: str + status: RunStatus = RunStatus.PENDING + modal_call_id: str | None = None + next_check_at: datetime | None = None + check_count: int = 0 + error: RunError | None = None + updated_at: datetime | None = None + written_results: list[Any] = field(default_factory=list) + + +class FakeRunTransitions: + """Conditional status changes over in-memory runs and comparison steps.""" + + def __init__(self, clock: ManualClock) -> None: + """Start with no runs.""" + self.clock = clock + # Callers run these methods in worker threads, as with the database. + self.lock = Lock() + self.runs: dict[UUID, FakeExecution] = {} + self.comparisons: dict[UUID, FakeExecution] = {} + self.comparison_runs_succeeded: dict[UUID, bool] = {} + + # --- Building state -------------------------------------------------- + + def add_run(self, kind: str = "economy", **values: Any) -> UUID: + """Add a simulation run and return its ID.""" + run = FakeExecution(id=uuid4(), kind=kind, **values) + self.runs[run.id] = run + return run.id + + def add_comparison(self, *, runs_succeeded: bool = True, **values: Any) -> UUID: + """Add an economy report run's comparison step and return its ID.""" + comparison = FakeExecution(id=uuid4(), kind="economy", **values) + self.comparisons[comparison.id] = comparison + self.comparison_runs_succeeded[comparison.id] = runs_succeeded + return comparison.id + + def running_call(self, call_id: str, *kinds: str) -> list[UUID]: + """Add runs already running under ``call_id``, due a check now.""" + return [ + self.add_run( + kind, + status=RunStatus.RUNNING, + modal_call_id=call_id, + next_check_at=self.clock(), + ) + for kind in kinds or ("economy",) + ] + + # --- Dispatch ---------------------------------------------------------- + + def _change_runs( + self, run_ids: Collection[UUID], expected: RunStatus, **values: Any + ) -> bool: + with self.lock: + runs = [self.runs[run_id] for run_id in run_ids] + if not runs or any(run.status is not expected for run in runs): + return False + for run in runs: + for name, value in values.items(): + setattr(run, name, value) + run.updated_at = self.clock() + return True + + def claim_simulation_dispatch(self, run_ids: Collection[UUID]) -> bool: + """Move pending runs to dispatching.""" + return self._change_runs( + run_ids, RunStatus.PENDING, status=RunStatus.DISPATCHING + ) + + def record_simulation_dispatched( + self, + run_ids: Collection[UUID], + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the call and move the runs to running.""" + return self._change_runs( + run_ids, + RunStatus.DISPATCHING, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + next_check_at=first_check_at, + ) + + def record_simulation_dispatch_failure( + self, run_ids: Collection[UUID], error: RunError + ) -> bool: + """Fail dispatching runs.""" + return self._change_runs( + run_ids, RunStatus.DISPATCHING, status=RunStatus.FAILED, error=error + ) + + def _change_comparison( + self, report_run_id: UUID, expected: RunStatus, **values: Any + ) -> FakeExecution | None: + with self.lock: + comparison = self.comparisons[report_run_id] + if comparison.status is not expected: + return None + for name, value in values.items(): + setattr(comparison, name, value) + comparison.updated_at = self.clock() + return comparison + + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: + """Move a pending comparison step whose runs succeeded to dispatching.""" + if not self.comparison_runs_succeeded[report_run_id]: + return False + return ( + self._change_comparison( + report_run_id, RunStatus.PENDING, status=RunStatus.DISPATCHING + ) + is not None + ) + + def record_comparison_dispatched( + self, + report_run_id: UUID, + *, + modal_call_id: str, + first_check_at: datetime, + ) -> bool: + """Record the comparison's call.""" + return ( + self._change_comparison( + report_run_id, + RunStatus.DISPATCHING, + status=RunStatus.RUNNING, + modal_call_id=modal_call_id, + next_check_at=first_check_at, + ) + is not None + ) + + def record_comparison_dispatch_failure( + self, report_run_id: UUID, error: RunError + ) -> bool: + """Fail a dispatching comparison step.""" + return ( + self._change_comparison( + report_run_id, + RunStatus.DISPATCHING, + status=RunStatus.FAILED, + error=error, + ) + is not None + ) + + +class FakeDispatcher: + """Record spawned requests and return call IDs, or fail like Modal.""" + + def __init__(self, failure: Exception | None = None) -> None: + """Fail every spawn with ``failure`` when it is set.""" + self.failure = failure + self.requests: list[ModalRequest] = [] + + async def spawn(self, request: ModalRequest) -> str: + """Return a new call ID, or raise the configured failure.""" + self.requests.append(request) + if self.failure is not None: + raise self.failure + return f"fc-{len(self.requests)}" + + +SAMPLE_REQUEST = ModalRequest( + target=ModalTarget(application="policyengine-simulation-v2", function="run"), + arguments=({"country": "us"},), + observability_context={"traceparent": "00-test"}, +) + + +def execution(transitions: FakeRunTransitions, run_id: UUID) -> FakeExecution: + """Return a copy of a run as the fake keeps it.""" + return replace(transitions.runs[run_id]) diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py new file mode 100644 index 000000000..f4378c4a6 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -0,0 +1,233 @@ +"""Starting a run's Modal call exactly once, and reading Modal's answers.""" + +from __future__ import annotations + +import asyncio + +import modal +import pytest +from policyengine_simulation_entry import modal_calls +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallRunning, + CallSucceeded, + ModalCallDispatcher, + ModalCallInspector, +) +from policyengine_simulation_entry.run_dispatch import ( + FIRST_CHECK_DELAY, + DispatchOutcome, + RunDispatchService, +) +from policyengine_simulation_persistence import RunErrorCode, RunStatus +from run_fakes import ( + SAMPLE_REQUEST, + FakeDispatcher, + FakeRunTransitions, + ManualClock, + execution, +) + + +def _service(dispatcher=None): + clock = ManualClock() + transitions = FakeRunTransitions(clock) + service = RunDispatchService( + transitions=transitions, + dispatcher=dispatcher or FakeDispatcher(), + clock=clock, + ) + return service, transitions, clock + + +@pytest.mark.asyncio +async def test_a_claimed_run_starts_its_call_and_records_it() -> None: + service, transitions, clock = _service() + run_id = transitions.add_run() + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + run = execution(transitions, run_id) + assert run.status is RunStatus.RUNNING + assert run.modal_call_id == "fc-1" + assert run.next_check_at == clock() + FIRST_CHECK_DELAY + assert service.dispatcher.requests == [SAMPLE_REQUEST] + + +@pytest.mark.asyncio +async def test_concurrent_starts_spawn_one_call() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + outcomes = await asyncio.gather( + *(service.start_simulation_runs([run_id], SAMPLE_REQUEST) for _ in range(5)) + ) + + assert outcomes.count(DispatchOutcome.STARTED) == 1 + assert outcomes.count(DispatchOutcome.NOT_CLAIMED) == 4 + assert len(service.dispatcher.requests) == 1 + + +@pytest.mark.asyncio +async def test_a_household_pair_starts_one_shared_call() -> None: + service, transitions, _ = _service() + pair = [transitions.add_run("household"), transitions.add_run("household")] + + outcome = await service.start_simulation_runs(pair, SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + assert {execution(transitions, run_id).modal_call_id for run_id in pair} == {"fc-1"} + + +@pytest.mark.asyncio +async def test_a_refused_call_fails_the_runs_without_raising() -> None: + service, transitions, _ = _service( + FakeDispatcher(failure=modal.exception.NotFoundError("no such function")) + ) + run_id = transitions.add_run() + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.DISPATCH_FAILED + run = execution(transitions, run_id) + assert run.status is RunStatus.FAILED + assert run.error.code is RunErrorCode.DISPATCH_FAILED + assert "NotFoundError" in run.error.message + assert "no such function" not in run.error.message + + +@pytest.mark.asyncio +async def test_a_call_whose_runs_stopped_dispatching_is_reported() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + class SlowDispatcher(FakeDispatcher): + async def spawn(self, request): + # The recording loop gave up on the dispatch in the meantime. + transitions.runs[run_id].status = RunStatus.FAILED + return await super().spawn(request) + + service.dispatcher = SlowDispatcher() + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert execution(transitions, run_id).modal_call_id is None + + +@pytest.mark.asyncio +async def test_a_comparison_starts_only_once_its_runs_succeeded() -> None: + service, transitions, _ = _service() + waiting = transitions.add_comparison(runs_succeeded=False) + ready = transitions.add_comparison() + + assert ( + await service.start_comparison(waiting, SAMPLE_REQUEST) + is DispatchOutcome.NOT_CLAIMED + ) + assert ( + await service.start_comparison(ready, SAMPLE_REQUEST) is DispatchOutcome.STARTED + ) + assert transitions.comparisons[ready].modal_call_id == "fc-1" + assert len(service.dispatcher.requests) == 1 + + +@pytest.mark.asyncio +async def test_a_refused_comparison_call_fails_the_step() -> None: + service, transitions, _ = _service(FakeDispatcher(failure=RuntimeError("down"))) + comparison = transitions.add_comparison() + + outcome = await service.start_comparison(comparison, SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.DISPATCH_FAILED + assert transitions.comparisons[comparison].status is RunStatus.FAILED + + +class _FakeCall: + def __init__(self, behaviour): + self.behaviour = behaviour + self.get = self + + async def aio(self, *, timeout): + assert timeout == 0 + if isinstance(self.behaviour, BaseException): + raise self.behaviour + return self.behaviour + + +def _inspect(monkeypatch, behaviour): + monkeypatch.setattr( + modal_calls.modal.FunctionCall, + "from_id", + lambda call_id: _FakeCall(behaviour), + ) + return asyncio.run(ModalCallInspector().inspect("fc-test")) + + +def test_a_call_without_output_is_running(monkeypatch) -> None: + assert _inspect(monkeypatch, TimeoutError()) == CallRunning() + + +def test_a_finished_call_returns_its_output(monkeypatch) -> None: + assert _inspect(monkeypatch, {"budget": 1.0}) == CallSucceeded({"budget": 1.0}) + + +@pytest.mark.parametrize( + ("raised", "code"), + [ + (modal.exception.FunctionTimeoutError("slow"), RunErrorCode.WORKER_TIMEOUT), + (modal.exception.OutputExpiredError(), RunErrorCode.RESULT_EXPIRED), + (modal.exception.NotFoundError("gone"), RunErrorCode.CALL_NOT_FOUND), + (KeyError("household secret"), RunErrorCode.WORKER_ERROR), + (modal.exception.InternalFailure("lost"), RunErrorCode.WORKER_ERROR), + ], + ids=["timeout", "expired", "not-found", "worker-error", "internal-failure"], +) +def test_a_failed_call_records_only_the_failure_type(monkeypatch, raised, code) -> None: + outcome = _inspect(monkeypatch, raised) + + assert isinstance(outcome, CallFailed) + assert outcome.error.code is code + assert "secret" not in outcome.error.message + + +@pytest.mark.parametrize( + "raised", + [ + modal.exception.ConnectionError("unreachable"), + modal.exception.ServiceError("unavailable"), + ], + ids=["connection", "service"], +) +def test_a_modal_problem_is_retried_later(monkeypatch, raised) -> None: + with pytest.raises(CallInspectionError): + _inspect(monkeypatch, raised) + + +def test_the_dispatcher_passes_observability_as_its_own_argument(monkeypatch) -> None: + spawned = {} + + class _FakeSpawn: + async def aio(self, *arguments, observability_context): + spawned.update(arguments=arguments, context=observability_context) + return type("Call", (), {"object_id": "fc-spawned"})() + + class _FakeFunction: + spawn = _FakeSpawn() + + def from_name(application, function, *, environment_name): + spawned.update(target=(application, function, environment_name)) + return _FakeFunction() + + monkeypatch.setattr(modal_calls.modal.Function, "from_name", from_name) + + call_id = asyncio.run(ModalCallDispatcher("staging").spawn(SAMPLE_REQUEST)) + + assert call_id == "fc-spawned" + assert spawned == { + "target": ("policyengine-simulation-v2", "run", "staging"), + "arguments": ({"country": "us"},), + "context": {"traceparent": "00-test"}, + } From 92a72952085bb5d50a4e218671d73f58d43209fe Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:37:37 +0400 Subject: [PATCH 04/39] Record Modal call outcomes on runs in a background loop in each instance Each Simulation Entrypoint instance with a configured database now runs RunRecordingLoop for its lifetime (started and stopped in the app's lifespan). On each pass, RunRecorder: 1. fails runs and comparison steps left dispatching for ten minutes (dispatch_abandoned), so they can be submitted again; 2. claims up to 50 due simulation calls and 50 due comparison steps, moving each next check one minute ahead, so instances do not check the same call at once and a call is checked again if this instance stops; 3. asks Modal about each claimed call without waiting, ten at a time, and records what it found: - still running (including a call Modal is re-running after pre-emption): the next check, backing off from 2 to 60 seconds; - failed or timed out: the runs fail with the error code and type; - finished: the output is validated by the result recorder registered for the kind of run, then the recorder writes the result rows in the transaction that marks the runs succeeded. A kind with no recorder, or an invalid output, fails the runs (no_result_recorder, invalid_result) without repeating the output. A household report's two runs share one call and are written once, together. If Modal cannot be reached, or writing the results fails, nothing is recorded and the call is checked again after its claim ends. The loop waits five seconds between passes, one second while passes find work, and logs a failed pass without stopping. Every write is conditional on the expected status, so an outcome another instance recorded first is left alone. No kind of run has a result recorder yet (item 1.07 adds them) and no route creates runs, so the loop finds nothing to do until stage 2. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/app.py | 9 + .../entry_runtime.py | 5 + .../run_recording.py | 469 ++++++++++++++++++ .../simulation_database.py | 22 + .../tests/run_fakes.py | 229 ++++++++- .../tests/test_run_recording.py | 366 ++++++++++++++ .../tests/test_simulation_database.py | 3 + 7 files changed, 1101 insertions(+), 2 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py create mode 100644 projects/policyengine-simulation-entry/tests/test_run_recording.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 894394d50..c08e0de2d 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -578,9 +578,18 @@ def create_app( async def lifespan(_: FastAPI) -> AsyncIterator[None]: runtime_settings.validate() await runtime_backend.start() + # Each instance records the outcomes of due Modal calls in the + # background; Cloud Run keeps its CPU allocated between requests. + recording_loop = ( + None if runtime_database is None else runtime_database.recording_loop() + ) + if recording_loop is not None: + recording_loop.start() try: yield finally: + if recording_loop is not None: + await recording_loop.stop() await runtime_backend.close() if runtime_database is not None: runtime_database.dispose() diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py index 4898f758f..1353a6cfc 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py @@ -41,6 +41,7 @@ ) from policyengine_simulation_entry.config import Settings + from policyengine_simulation_entry.run_recording import RunRecordingLoop # Moved out of app.py; this code still logs under that module's name. logger = logging.getLogger("policyengine_simulation_entry.app") @@ -220,6 +221,10 @@ async def read_revision(self) -> str | None: """Return the applied migration revision, or None before the first one.""" ... + def recording_loop(self) -> RunRecordingLoop | None: + """Return the loop that records Modal call outcomes on runs, if any.""" + ... + def dispose(self) -> None: """Close every pooled connection.""" ... diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py new file mode 100644 index 000000000..126392609 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -0,0 +1,469 @@ +"""The background loop that records each Modal call's outcome on its runs. + +Each Simulation Entrypoint instance runs one loop. On every pass it: + +1. fails runs and comparison steps left ``dispatching`` for ten minutes, so + they can be submitted again; +2. claims the running calls that are due a check, by moving their next check + one lease ahead, so instances do not check the same call at once; +3. asks Modal about each claimed call without waiting, and records what it + found with conditional writes: a finished call's result is validated and + its result rows are written in the transaction that marks its runs + ``succeeded``; a failed or timed-out call is marked ``failed`` with its + error; a call still running gets a later check, backing off. + +Pre-emption is never a failure: Modal re-runs a pre-empted call under the +same call ID, so it reads as still running. No caller polls, and Modal holds +no database credentials: this loop writes every outcome. Cloud Run keeps the +instance's CPU allocated between requests, so the loop runs then too. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +from dataclasses import dataclass, field +from datetime import timedelta +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_persistence import RunError, RunErrorCode +from policyengine_simulation_persistence.run_store import utc_now + +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallRunning, +) + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Mapping + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_persistence import ( + AbandonedDispatches, + DueSimulationCall, + ReportKind, + ReportRunRecord, + SimulationKind, + SimulationRunRecord, + ) + from sqlalchemy import Connection + + from policyengine_simulation_entry.modal_calls import CallInspector, CallOutcome + +logger = logging.getLogger(__name__) + +# A claimed call is not checked by another instance for this long; if the +# claiming instance stops, the call becomes due again afterwards. +CHECK_LEASE = timedelta(minutes=1) +# A run left dispatching this long was claimed by an instance that stopped. +DISPATCH_ABANDONED_AFTER = timedelta(minutes=10) +FIRST_BACKOFF = timedelta(seconds=2) +LONGEST_BACKOFF = timedelta(seconds=60) +CLAIM_BATCH_SIZE = 50 +INSPECTION_CONCURRENCY = 10 +BUSY_PASS_INTERVAL_SECONDS = 1.0 +IDLE_PASS_INTERVAL_SECONDS = 5.0 + + +def next_check_delay(check_count: int) -> timedelta: + """Return the wait before the next check of a call checked ``check_count`` times.""" + return min(FIRST_BACKOFF * 2 ** max(check_count - 1, 0), LONGEST_BACKOFF) + + +class SimulationResultRecorder(Protocol): + """Validate one kind of simulation call's output and write its result rows.""" + + def validate(self, output: object) -> object: + """Return the validated result; raise ValueError when it is invalid.""" + ... + + def write( + self, + connection: Connection, + runs: tuple[SimulationRunRecord, ...], + result: object, + ) -> None: + """Write the result rows of ``runs`` in the transaction that succeeds them.""" + ... + + +class ComparisonResultRecorder(Protocol): + """Validate one kind of comparison's output and write its result rows.""" + + def validate(self, output: object) -> object: + """Return the validated result; raise ValueError when it is invalid.""" + ... + + def write( + self, connection: Connection, report_run: ReportRunRecord, result: object + ) -> None: + """Write the comparison's result rows in the transaction that succeeds it.""" + ... + + +@dataclass(frozen=True) +class ResultRecorders: + """The result recorders by kind of run; a kind with none fails its runs.""" + + simulations: Mapping[SimulationKind, SimulationResultRecorder] = field( + default_factory=dict + ) + comparisons: Mapping[ReportKind, ComparisonResultRecorder] = field( + default_factory=dict + ) + + +class RecordingTransitions(Protocol): + """The ``RunTransitions`` writes the loop needs.""" + + def claim_due_simulation_calls( + self, *, limit: int, lease: timedelta + ) -> list[DueSimulationCall]: + """Claim running simulation calls that are due a check.""" + ... + + def record_simulation_call_running( + self, modal_call_id: str, *, next_check_at: datetime + ) -> bool: + """Set a running call's next check.""" + ... + + def complete_simulation_call( + self, + modal_call_id: str, + write_results: Callable[[Connection, tuple[SimulationRunRecord, ...]], None], + ) -> bool: + """Mark a call's runs succeeded and write their results.""" + ... + + def fail_simulation_call(self, modal_call_id: str, error: RunError) -> bool: + """Mark a call's runs failed.""" + ... + + def claim_due_comparisons( + self, *, limit: int, lease: timedelta + ) -> list[ReportRunRecord]: + """Claim running comparison steps that are due a check.""" + ... + + def record_comparison_running( + self, report_run_id: UUID, *, next_check_at: datetime + ) -> bool: + """Set a running comparison step's next check.""" + ... + + def complete_comparison( + self, + report_run_id: UUID, + write_results: Callable[[Connection, ReportRunRecord], None], + ) -> bool: + """Mark a comparison step succeeded and write its results.""" + ... + + def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: + """Mark a comparison step failed.""" + ... + + def fail_abandoned_dispatches( + self, *, claimed_before: datetime + ) -> AbandonedDispatches: + """Fail runs and comparison steps left dispatching.""" + ... + + +@dataclass +class RecordingPass: + """What one pass found and recorded.""" + + checked: int = 0 + succeeded: int = 0 + failed: int = 0 + still_running: int = 0 + not_reached: int = 0 + abandoned: int = 0 + + @property + def found_work(self) -> bool: + """Whether the pass checked or failed anything.""" + return bool(self.checked or self.abandoned) + + +@dataclass(frozen=True) +class _ResultSteps: + """Validate a finished call's output, then record it with its result rows.""" + + validate: Callable[[object], object] + complete: Callable[[object], bool] + + +@dataclass(frozen=True) +class _CallHandlers: + """How to record one claimed call's outcome, bound to its runs. + + ``result_steps`` is None when no recorder handles the runs' kind. + """ + + subject: dict[str, object] + check_count: int + result_steps: _ResultSteps | None + fail: Callable[[RunError], bool] + keep_running: Callable[[datetime], bool] + + +def _invalid_result(error: ValueError) -> RunError: + count = getattr(error, "error_count", None) + detail = f" with {count()} errors" if callable(count) else "" + return RunError( + RunErrorCode.INVALID_RESULT, + f"The Modal call's output failed validation{detail} ({type(error).__name__}).", + ) + + +@dataclass +class RunRecorder: + """Record the outcomes of due Modal calls, one pass at a time.""" + + transitions: RecordingTransitions + inspector: CallInspector + recorders: ResultRecorders = field(default_factory=ResultRecorders) + clock: Callable[[], datetime] = utc_now + + async def record_due_outcomes(self) -> RecordingPass: + """Run one pass: fail abandoned dispatches, then check every due call.""" + recording = RecordingPass() + abandoned = await asyncio.to_thread( + lambda: self.transitions.fail_abandoned_dispatches( + claimed_before=self.clock() - DISPATCH_ABANDONED_AFTER + ) + ) + recording.abandoned = len(abandoned.simulation_run_ids) + len( + abandoned.report_run_ids + ) + if recording.abandoned: + logger.warning( + "simulation_run_dispatches_abandoned", + extra={ + "simulation_run_ids": [ + str(i) for i in abandoned.simulation_run_ids + ], + "report_run_ids": [str(i) for i in abandoned.report_run_ids], + }, + ) + calls = await asyncio.to_thread( + lambda: self.transitions.claim_due_simulation_calls( + limit=CLAIM_BATCH_SIZE, lease=CHECK_LEASE + ) + ) + comparisons = await asyncio.to_thread( + lambda: self.transitions.claim_due_comparisons( + limit=CLAIM_BATCH_SIZE, lease=CHECK_LEASE + ) + ) + handlers = [ + *(self._simulation_handlers(call) for call in calls), + *(self._comparison_handlers(report_run) for report_run in comparisons), + ] + limit = asyncio.Semaphore(INSPECTION_CONCURRENCY) + + async def check(modal_call_id: str, call_handlers: _CallHandlers) -> None: + async with limit: + await self._check(modal_call_id, call_handlers, recording) + + await asyncio.gather( + *( + check(modal_call_id, call_handlers) + for modal_call_id, call_handlers in handlers + ) + ) + return recording + + def _simulation_handlers( + self, call: DueSimulationCall + ) -> tuple[str, _CallHandlers]: + kinds = {run.simulation_kind for run in call.runs} + recorder = ( + self.recorders.simulations.get(next(iter(kinds))) + if len(kinds) == 1 + else None + ) + call_id = call.modal_call_id + result_steps = None + if recorder is not None: + writer = recorder + result_steps = _ResultSteps( + validate=writer.validate, + complete=lambda result: self.transitions.complete_simulation_call( + call_id, + lambda connection, runs: writer.write(connection, runs, result), + ), + ) + return call_id, _CallHandlers( + subject={ + "modal_call_id": call_id, + "simulation_run_ids": [str(run.id) for run in call.runs], + }, + check_count=max(run.check_count for run in call.runs), + result_steps=result_steps, + fail=lambda error: self.transitions.fail_simulation_call(call_id, error), + keep_running=lambda next_check_at: ( + self.transitions.record_simulation_call_running( + call_id, next_check_at=next_check_at + ) + ), + ) + + def _comparison_handlers( + self, report_run: ReportRunRecord + ) -> tuple[str, _CallHandlers]: + recorder = self.recorders.comparisons.get(report_run.report_kind) + report_run_id = report_run.id + call_id = report_run.comparison_modal_call_id or "" + result_steps = None + if recorder is not None: + writer = recorder + result_steps = _ResultSteps( + validate=writer.validate, + complete=lambda result: self.transitions.complete_comparison( + report_run_id, + lambda connection, record: writer.write(connection, record, result), + ), + ) + return call_id, _CallHandlers( + subject={"modal_call_id": call_id, "report_run_id": str(report_run_id)}, + check_count=report_run.comparison_check_count, + result_steps=result_steps, + fail=lambda error: self.transitions.fail_comparison(report_run_id, error), + keep_running=lambda next_check_at: ( + self.transitions.record_comparison_running( + report_run_id, next_check_at=next_check_at + ) + ), + ) + + async def _check( + self, modal_call_id: str, handlers: _CallHandlers, recording: RecordingPass + ) -> None: + recording.checked += 1 + try: + outcome = await self.inspector.inspect(modal_call_id) + except CallInspectionError as error: + # The claim's lease makes the call due again; nothing to record. + recording.not_reached += 1 + logger.warning( + "simulation_run_check_failed", + extra={**handlers.subject, "error_type": str(error)}, + ) + return + await self._record(outcome, handlers, recording) + + async def _record( + self, outcome: CallOutcome, handlers: _CallHandlers, recording: RecordingPass + ) -> None: + if isinstance(outcome, CallRunning): + next_check_at = self.clock() + next_check_delay(handlers.check_count) + await asyncio.to_thread(handlers.keep_running, next_check_at) + recording.still_running += 1 + return + if isinstance(outcome, CallFailed): + await self._fail(handlers, outcome.error, recording) + return + steps = handlers.result_steps + if steps is None: + await self._fail( + handlers, + RunError( + RunErrorCode.NO_RESULT_RECORDER, + "No result recorder handles this kind of run.", + ), + recording, + ) + return + try: + result = await asyncio.to_thread(steps.validate, outcome.output) + except ValueError as error: + await self._fail(handlers, _invalid_result(error), recording) + return + try: + completed = await asyncio.to_thread(steps.complete, result) + # Writing the result rows failed, for example because the database + # was unreachable: nothing was recorded, and the lease makes the call + # due again. + except Exception as error: + recording.not_reached += 1 + logger.exception( + "simulation_run_result_not_written", + extra={**handlers.subject, "error_type": type(error).__name__}, + ) + return + if completed: + recording.succeeded += 1 + logger.info("simulation_run_succeeded", extra=handlers.subject) + + async def _fail( + self, handlers: _CallHandlers, error: RunError, recording: RecordingPass + ) -> None: + if await asyncio.to_thread(handlers.fail, error): + recording.failed += 1 + logger.warning( + "simulation_run_failed", + extra={**handlers.subject, "error_code": error.code.value}, + ) + + +class RunRecordingLoop: + """Run recording passes in the background until stopped. + + A pass that raises is logged and the loop goes on: recording problems + never stop the instance or fail a request. + """ + + def __init__( + self, + recorder: RunRecorder, + *, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + """Run passes of ``recorder``, waiting between them with ``sleep``.""" + self.recorder = recorder + self._sleep = sleep + self._task: asyncio.Task[None] | None = None + + async def run(self) -> None: + """Record outcomes until cancelled; wait less while there is work.""" + while True: + await self._sleep(IDLE_PASS_INTERVAL_SECONDS) + try: + recording = await self.recorder.record_due_outcomes() + except Exception as error: + logger.exception( + "simulation_run_recording_pass_failed", + extra={"error_type": type(error).__name__}, + ) + continue + while recording.found_work: + await self._sleep(BUSY_PASS_INTERVAL_SECONDS) + try: + recording = await self.recorder.record_due_outcomes() + except Exception as error: + logger.exception( + "simulation_run_recording_pass_failed", + extra={"error_type": type(error).__name__}, + ) + break + + def start(self) -> None: + """Start the loop as a background task of the running event loop.""" + if self._task is None: + self._task = asyncio.create_task(self.run()) + + async def stop(self) -> None: + """Cancel the loop and wait for it to end.""" + if self._task is None: + return + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._task + self._task = None diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index e0ca4f1bd..46280ca2e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -6,10 +6,18 @@ from typing import TYPE_CHECKING from policyengine_simulation_persistence import ( + RunTransitions, create_simulation_engine, read_schema_revision, ) +from policyengine_simulation_entry.modal_calls import ModalCallInspector +from policyengine_simulation_entry.run_recording import ( + ResultRecorders, + RunRecorder, + RunRecordingLoop, +) + if TYPE_CHECKING: from sqlalchemy import Engine @@ -37,6 +45,20 @@ async def read_revision(self) -> str | None: """Return the applied migration revision, or None before the first one.""" return await asyncio.to_thread(self._read_revision) + def recording_loop(self) -> RunRecordingLoop: + """Return the loop that records Modal call outcomes on runs. + + No kind of run has a result recorder yet, so a finished call would + fail its runs with ``no_result_recorder``; no route creates runs yet. + """ + return RunRecordingLoop( + RunRecorder( + transitions=RunTransitions(self.engine), + inspector=ModalCallInspector(), + recorders=ResultRecorders(), + ) + ) + def dispose(self) -> None: """Close every pooled connection.""" self.engine.dispose() diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py index 1af6983b4..586693a05 100644 --- a/projects/policyengine-simulation-entry/tests/run_fakes.py +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -13,12 +13,25 @@ from typing import TYPE_CHECKING, Any from uuid import UUID, uuid4 -from policyengine_simulation_entry.modal_calls import ModalRequest, ModalTarget -from policyengine_simulation_persistence import RunError, RunStatus +from policyengine_simulation_entry.modal_calls import ( + CallInspectionError, + CallRunning, + ModalRequest, + ModalTarget, +) +from policyengine_simulation_persistence import ( + AbandonedDispatches, + DueSimulationCall, + RunError, + RunStatus, + SimulationKind, +) if TYPE_CHECKING: from collections.abc import Collection + from policyengine_simulation_entry.modal_calls import CallOutcome + @dataclass class ManualClock: @@ -52,6 +65,26 @@ class FakeExecution: written_results: list[Any] = field(default_factory=list) +@dataclass(frozen=True) +class FakeRunRecord: + """The run fields the recorder reads from a claimed run.""" + + id: UUID + simulation_kind: SimulationKind + modal_call_id: str | None + check_count: int + + +@dataclass(frozen=True) +class FakeReportRunRecord: + """The report run fields the recorder reads from a claimed comparison.""" + + id: UUID + report_kind: str + comparison_modal_call_id: str | None + comparison_check_count: int + + class FakeRunTransitions: """Conditional status changes over in-memory runs and comparison steps.""" @@ -192,6 +225,181 @@ def record_comparison_dispatch_failure( is not None ) + # --- Recording ------------------------------------------------------- + + def claim_due_simulation_calls( + self, *, limit: int, lease: timedelta + ) -> list[DueSimulationCall]: + """Claim running runs that are due, grouped by call.""" + now = self.clock() + due = [ + run + for run in self.runs.values() + if run.status is RunStatus.RUNNING + and run.next_check_at is not None + and run.next_check_at <= now + ][:limit] + calls: dict[str, list[FakeRunRecord]] = {} + for run in due: + run.next_check_at = now + lease + run.check_count += 1 + assert run.modal_call_id is not None + calls.setdefault(run.modal_call_id, []).append( + FakeRunRecord( + id=run.id, + simulation_kind=SimulationKind(run.kind), + modal_call_id=run.modal_call_id, + check_count=run.check_count, + ) + ) + return [ + DueSimulationCall(modal_call_id=call_id, runs=tuple(runs)) # type: ignore[arg-type] + for call_id, runs in sorted(calls.items()) + ] + + def _call_runs(self, modal_call_id: str) -> list[FakeExecution]: + return [ + run + for run in self.runs.values() + if run.modal_call_id == modal_call_id and run.status is RunStatus.RUNNING + ] + + def record_simulation_call_running( + self, modal_call_id: str, *, next_check_at: datetime + ) -> bool: + """Set the next check of a call's running runs.""" + runs = self._call_runs(modal_call_id) + for run in runs: + run.next_check_at = next_check_at + return bool(runs) + + def complete_simulation_call(self, modal_call_id: str, write_results) -> bool: + """Mark a call's runs succeeded after the writer succeeds.""" + runs = self._call_runs(modal_call_id) + if not runs: + return False + records = tuple( + FakeRunRecord( + id=run.id, + simulation_kind=SimulationKind(run.kind), + modal_call_id=run.modal_call_id, + check_count=run.check_count, + ) + for run in runs + ) + connection = object() + write_results(connection, records) + for run in runs: + run.status = RunStatus.SUCCEEDED + run.next_check_at = None + run.written_results.append(connection) + return True + + def fail_simulation_call(self, modal_call_id: str, error: RunError) -> bool: + """Fail a call's running runs.""" + runs = self._call_runs(modal_call_id) + for run in runs: + run.status = RunStatus.FAILED + run.error = error + run.next_check_at = None + return bool(runs) + + def claim_due_comparisons( + self, *, limit: int, lease: timedelta + ) -> list[FakeReportRunRecord]: + """Claim running comparison steps that are due.""" + now = self.clock() + due = [ + comparison + for comparison in self.comparisons.values() + if comparison.status is RunStatus.RUNNING + and comparison.next_check_at is not None + and comparison.next_check_at <= now + ][:limit] + for comparison in due: + comparison.next_check_at = now + lease + comparison.check_count += 1 + return [ + FakeReportRunRecord( + id=comparison.id, + report_kind="economy", + comparison_modal_call_id=comparison.modal_call_id, + comparison_check_count=comparison.check_count, + ) + for comparison in due + ] + + def record_comparison_running( + self, report_run_id: UUID, *, next_check_at: datetime + ) -> bool: + """Set a running comparison step's next check.""" + return ( + self._change_comparison( + report_run_id, RunStatus.RUNNING, next_check_at=next_check_at + ) + is not None + ) + + def complete_comparison(self, report_run_id: UUID, write_results) -> bool: + """Mark a running comparison step succeeded after the writer succeeds.""" + comparison = self.comparisons[report_run_id] + if comparison.status is not RunStatus.RUNNING: + return False + connection = object() + write_results( + connection, + FakeReportRunRecord( + id=comparison.id, + report_kind="economy", + comparison_modal_call_id=comparison.modal_call_id, + comparison_check_count=comparison.check_count, + ), + ) + comparison.status = RunStatus.SUCCEEDED + comparison.next_check_at = None + comparison.written_results.append(connection) + return True + + def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: + """Fail a running comparison step.""" + return ( + self._change_comparison( + report_run_id, + RunStatus.RUNNING, + status=RunStatus.FAILED, + error=error, + next_check_at=None, + ) + is not None + ) + + def fail_abandoned_dispatches( + self, *, claimed_before: datetime + ) -> AbandonedDispatches: + """Fail runs and comparison steps dispatching since before the cutoff.""" + failed_runs = [] + for run in self.runs.values(): + if ( + run.status is RunStatus.DISPATCHING + and run.updated_at is not None + and run.updated_at < claimed_before + ): + run.status = RunStatus.FAILED + failed_runs.append(run.id) + failed_comparisons = [] + for comparison in self.comparisons.values(): + if ( + comparison.status is RunStatus.DISPATCHING + and comparison.updated_at is not None + and comparison.updated_at < claimed_before + ): + comparison.status = RunStatus.FAILED + failed_comparisons.append(comparison.id) + return AbandonedDispatches( + simulation_run_ids=tuple(failed_runs), + report_run_ids=tuple(failed_comparisons), + ) + class FakeDispatcher: """Record spawned requests and return call IDs, or fail like Modal.""" @@ -209,6 +417,23 @@ async def spawn(self, request: ModalRequest) -> str: return f"fc-{len(self.requests)}" +class FakeInspector: + """Return a scripted outcome per call ID, or running when none is set.""" + + def __init__(self) -> None: + """Start with every call running.""" + self.outcomes: dict[str, CallOutcome | CallInspectionError] = {} + self.inspected: list[str] = [] + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the scripted outcome, or raise the scripted error.""" + self.inspected.append(modal_call_id) + outcome = self.outcomes.get(modal_call_id, CallRunning()) + if isinstance(outcome, CallInspectionError): + raise outcome + return outcome + + SAMPLE_REQUEST = ModalRequest( target=ModalTarget(application="policyengine-simulation-v2", function="run"), arguments=({"country": "us"},), diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py new file mode 100644 index 000000000..687b36222 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -0,0 +1,366 @@ +"""The background loop that records Modal call outcomes on runs.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta + +import pytest +from conftest import FakeBackend, make_settings +from fastapi.testclient import TestClient +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallSucceeded, +) +from policyengine_simulation_entry.run_recording import ( + BUSY_PASS_INTERVAL_SECONDS, + CHECK_LEASE, + IDLE_PASS_INTERVAL_SECONDS, + RecordingPass, + ResultRecorders, + RunRecorder, + RunRecordingLoop, + next_check_delay, +) +from policyengine_simulation_persistence import ( + HEAD_REVISION, + ReportKind, + RunError, + RunErrorCode, + RunStatus, + SimulationKind, +) +from pydantic import BaseModel +from run_fakes import FakeInspector, FakeRunTransitions, ManualClock + + +class BudgetResult(BaseModel): + budgetary_impact: float + + +class RecordingRecorder: + """Validate outputs as ``BudgetResult`` and remember what it wrote.""" + + def __init__(self, failure: Exception | None = None) -> None: + self.failure = failure + self.writes = [] + + def validate(self, output): + return BudgetResult.model_validate(output) + + def write(self, connection, records, result) -> None: + if self.failure is not None: + raise self.failure + self.writes.append((connection, records, result)) + + +def _recorder(recorders=None): + clock = ManualClock() + transitions = FakeRunTransitions(clock) + inspector = FakeInspector() + recorder = RunRecorder( + transitions=transitions, + inspector=inspector, + recorders=recorders or ResultRecorders(), + clock=clock, + ) + return recorder, transitions, inspector, clock + + +def _run(recorder) -> RecordingPass: + return asyncio.run(recorder.record_due_outcomes()) + + +@pytest.mark.parametrize( + ("check_count", "seconds"), + [(0, 2), (1, 2), (2, 4), (3, 8), (5, 32), (6, 60), (40, 60)], +) +def test_checks_back_off_up_to_a_minute(check_count, seconds) -> None: + assert next_check_delay(check_count) == timedelta(seconds=seconds) + + +def test_a_running_call_is_checked_again_later() -> None: + recorder, transitions, inspector, clock = _recorder() + (run_id,) = transitions.running_call("fc-1") + + recording = _run(recorder) + + assert inspector.inspected == ["fc-1"] + assert recording.still_running == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.check_count == 1 + assert run.next_check_at == clock() + timedelta(seconds=2) + assert _run(recorder).checked == 0 + + +def test_a_failed_call_fails_its_runs() -> None: + recorder, transitions, inspector, _ = _recorder() + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallFailed( + RunError(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + ) + + recording = _run(recorder) + + assert recording.failed == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.FAILED + assert run.error.code is RunErrorCode.WORKER_TIMEOUT + + +def test_a_finished_call_is_validated_and_written_with_its_success() -> None: + budget = RecordingRecorder() + recorder, transitions, inspector, _ = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: budget}) + ) + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": -1.5}) + + recording = _run(recorder) + + assert recording.succeeded == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.SUCCEEDED + ((connection, records, result),) = budget.writes + assert [record.id for record in records] == [run_id] + assert result == BudgetResult(budgetary_impact=-1.5) + assert run.written_results == [connection] + + +def test_a_household_pair_is_written_once_for_both_runs() -> None: + household = RecordingRecorder() + recorder, transitions, inspector, _ = _recorder( + ResultRecorders(simulations={SimulationKind.HOUSEHOLD: household}) + ) + pair = transitions.running_call("fc-pair", "household", "household") + inspector.outcomes["fc-pair"] = CallSucceeded({"budgetary_impact": 0}) + + recording = _run(recorder) + + assert recording.checked == 1 + assert inspector.inspected == ["fc-pair"] + ((_, records, _),) = household.writes + assert {record.id for record in records} == set(pair) + assert {transitions.runs[run_id].status for run_id in pair} == {RunStatus.SUCCEEDED} + + +def test_a_kind_without_a_recorder_fails_its_runs() -> None: + recorder, transitions, inspector, _ = _recorder() + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": 1}) + + _run(recorder) + + assert transitions.runs[run_id].error.code is RunErrorCode.NO_RESULT_RECORDER + + +def test_an_invalid_output_fails_its_runs_without_repeating_it() -> None: + recorder, transitions, inspector, _ = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: RecordingRecorder()}) + ) + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": "household 42"}) + + recording = _run(recorder) + + assert recording.failed == 1 + error = transitions.runs[run_id].error + assert error.code is RunErrorCode.INVALID_RESULT + assert "1 errors" in error.message + assert "household 42" not in error.message + + +def test_an_unreachable_modal_leaves_the_call_for_after_its_lease() -> None: + recorder, transitions, inspector, clock = _recorder() + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallInspectionError("ConnectionError") + + recording = _run(recorder) + + assert recording.not_reached == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.next_check_at == clock() + CHECK_LEASE + clock.advance(CHECK_LEASE) + assert _run(recorder).checked == 1 + + +def test_a_failed_result_write_records_nothing_and_is_retried() -> None: + recorder, transitions, inspector, clock = _recorder( + ResultRecorders( + simulations={ + SimulationKind.ECONOMY: RecordingRecorder( + failure=ConnectionError("database unreachable") + ) + } + ) + ) + (run_id,) = transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": 1}) + + recording = _run(recorder) + + assert recording.not_reached == 1 + assert recording.succeeded == 0 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.next_check_at == clock() + CHECK_LEASE + + +def test_an_outcome_another_instance_recorded_first_is_not_counted() -> None: + recorder, transitions, _, _ = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: RecordingRecorder()}) + ) + (run_id,) = transitions.running_call("fc-1") + + class RacingInspector(FakeInspector): + async def inspect(self, modal_call_id): + transitions.runs[run_id].status = RunStatus.SUCCEEDED + return CallSucceeded({"budgetary_impact": 1}) + + recorder.inspector = RacingInspector() + + recording = _run(recorder) + + assert recording.checked == 1 + assert recording.succeeded == 0 + + +def test_comparison_steps_are_checked_and_recorded() -> None: + comparisons = RecordingRecorder() + recorder, transitions, inspector, clock = _recorder( + ResultRecorders(comparisons={ReportKind.ECONOMY: comparisons}) + ) + finished = transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-done", next_check_at=clock() + ) + waiting = transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-wait", next_check_at=clock() + ) + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 3}) + + recording = _run(recorder) + + assert recording.succeeded == 1 + assert recording.still_running == 1 + assert transitions.comparisons[finished].status is RunStatus.SUCCEEDED + assert [record.id for _, record, _ in comparisons.writes] == [finished] + assert transitions.comparisons[waiting].next_check_at == clock() + timedelta( + seconds=2 + ) + + +def test_dispatches_left_for_ten_minutes_are_failed_by_a_pass() -> None: + recorder, transitions, _, clock = _recorder() + stalled = transitions.add_run(status=RunStatus.DISPATCHING, updated_at=clock()) + clock.advance(timedelta(minutes=10, seconds=1)) + recent = transitions.add_run(status=RunStatus.DISPATCHING, updated_at=clock()) + + recording = _run(recorder) + + assert recording.abandoned == 1 + assert transitions.runs[stalled].status is RunStatus.FAILED + assert transitions.runs[recent].status is RunStatus.DISPATCHING + + +class _ScriptedRecorder: + """Return scripted passes, or raise, and count how often it ran.""" + + def __init__(self, passes) -> None: + self.passes = list(passes) + self.calls = 0 + + async def record_due_outcomes(self) -> RecordingPass: + self.calls += 1 + outcome = self.passes.pop(0) if self.passes else RecordingPass() + if isinstance(outcome, Exception): + raise outcome + return outcome + + +def test_the_loop_waits_less_while_there_is_work_and_survives_failures() -> None: + scripted = _ScriptedRecorder( + [RecordingPass(checked=3), RecordingPass(), RuntimeError("database down")] + ) + waits = [] + + async def sleep(seconds: float) -> None: + waits.append(seconds) + if len(waits) == 5: + raise asyncio.CancelledError + + loop = RunRecordingLoop(scripted, sleep=sleep) # type: ignore[arg-type] + + with pytest.raises(asyncio.CancelledError): + asyncio.run(loop.run()) + + assert waits == [ + IDLE_PASS_INTERVAL_SECONDS, + BUSY_PASS_INTERVAL_SECONDS, + IDLE_PASS_INTERVAL_SECONDS, + IDLE_PASS_INTERVAL_SECONDS, + IDLE_PASS_INTERVAL_SECONDS, + ] + assert scripted.calls == 4 + + +def test_the_loop_starts_and_stops_as_a_background_task() -> None: + scripted = _ScriptedRecorder([]) + + async def lifetime() -> None: + loop = RunRecordingLoop(scripted, sleep=lambda _seconds: asyncio.sleep(0)) # type: ignore[arg-type] + loop.start() + await asyncio.sleep(0.01) + await loop.stop() + await loop.stop() + + asyncio.run(lifetime()) + + assert scripted.calls > 0 + + +class _TrackedLoop: + def __init__(self) -> None: + self.started = False + self.stopped = False + + def start(self) -> None: + self.started = True + + async def stop(self) -> None: + self.stopped = True + + +class _DatabaseWithLoop: + def __init__(self) -> None: + self.loop = _TrackedLoop() + + async def read_revision(self): + return HEAD_REVISION + + def recording_loop(self): + return self.loop + + def dispose(self) -> None: + return None + + +def test_each_instance_runs_the_loop_for_its_lifetime() -> None: + database = _DatabaseWithLoop() + app = create_app( + settings=make_settings( + simulation_database_url="postgresql://runtime:secret@db.example/postgres" + ), + backend=FakeBackend(), + auth_dependency=lambda: None, + simulation_database=database, + ) + + with TestClient(app): + assert database.loop.started is True + assert database.loop.stopped is False + + assert database.loop.stopped is True diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_database.py b/projects/policyengine-simulation-entry/tests/test_simulation_database.py index 95b1defae..6a7e55547 100644 --- a/projects/policyengine-simulation-entry/tests/test_simulation_database.py +++ b/projects/policyengine-simulation-entry/tests/test_simulation_database.py @@ -33,6 +33,9 @@ async def read_revision(self): raise self.failure return self.revision + def recording_loop(self): + return None + def dispose(self): self.disposed = True From 8cfbd401b09bc9d5b4dbba3e978467643a7dd779 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:59:31 +0400 Subject: [PATCH 05/39] Keep CPU allocated and always mount sim's Modal token for the recording loop The recording loop runs between requests and asks Modal about runs' calls, so the Simulation Entrypoint deploy now: - passes --no-cpu-throttling to gcloud run deploy in staging and production, so Cloud Run keeps the instance's CPU allocated outside requests. Production already keeps one instance running; staging scales to zero and records outcomes only while an instance runs. - mounts MODAL_TOKEN_ID and MODAL_TOKEN_SECRET on every deploy, from the Secret Manager secrets named by the beta and prod environment variables MODAL_TOKEN_ID_SECRET_NAME and MODAL_TOKEN_SECRET_SECRET_NAME, pinned to their newest enabled versions. Before, they were mounted only when DEPLOY_STAGE12_V2 was true; the Stage 12 block no longer mounts them, and its Modal calls use the same token. A deploy without either name stops before gcloud run deploy. The operator step that creates sim's own Modal token sets those two variables to the new secret names. DEPLOYMENT.md describes both settings, and the persistence README names the REFERENCES grants revision 0003 needs. Co-Authored-By: Claude Opus 5.5 --- .../deploy-cloud-run-simulation-entry.sh | 12 ++- .../README.md | 7 +- .../DEPLOYMENT.md | 19 ++++ .../tests/test_run_recording_deployment.py | 88 +++++++++++++++++++ 4 files changed, 118 insertions(+), 8 deletions(-) create mode 100644 projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index 0956ae9a5..05836e82f 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -100,6 +100,8 @@ required_variables=( OLD_GATEWAY_AUTH_CLIENT_ID_VALUE OLD_GATEWAY_AUTH_CLIENT_SECRET_SECRET_NAME SIMULATION_DATABASE_URL_SECRET_NAME + MODAL_TOKEN_ID_SECRET_NAME + MODAL_TOKEN_SECRET_SECRET_NAME OBSERVABILITY_SERVICE_NAMESPACE OBSERVABILITY_TRACE_PROJECT_ID OTEL_EXPORTER_OTLP_ENDPOINT @@ -143,13 +145,16 @@ add_runtime_secret OLD_GATEWAY_AUTH_CLIENT_SECRET \ # The runtime role's URL for the entrypoint's own simulation schema. The # migration job, which uses a separate migrator URL, runs before this deploy. add_runtime_secret SIMULATION_DATABASE_URL "${SIMULATION_DATABASE_URL_SECRET_NAME}" +# sim's own Modal token: the entrypoint starts runs' Modal calls and its +# background loop asks Modal about them, on every deploy. Until the Stage 12 +# comparison path is removed, its Modal calls use the same token. +add_runtime_secret MODAL_TOKEN_ID "${MODAL_TOKEN_ID_SECRET_NAME}" +add_runtime_secret MODAL_TOKEN_SECRET "${MODAL_TOKEN_SECRET_SECRET_NAME}" if [[ "${DEPLOY_STAGE12_V2}" == "true" ]]; then stage12_variables=( STAGE12_ENABLED_VALUE STAGE12_ARTIFACT_BUCKET_VALUE - MODAL_TOKEN_ID_SECRET_NAME - MODAL_TOKEN_SECRET_SECRET_NAME STAGE12_DATABASE_URL_SECRET_NAME ) for variable_name in "${stage12_variables[@]}"; do @@ -161,8 +166,6 @@ if [[ "${DEPLOY_STAGE12_V2}" == "true" ]]; then || fail "STAGE12_ARTIFACT_BUCKET_VALUE is invalid" validate_secret_reference "${STAGE12_DATABASE_URL_SECRET_NAME}" - add_stage12_runtime_secret MODAL_TOKEN_ID "${MODAL_TOKEN_ID_SECRET_NAME}" - add_stage12_runtime_secret MODAL_TOKEN_SECRET "${MODAL_TOKEN_SECRET_SECRET_NAME}" add_stage12_runtime_secret STAGE12_DATABASE_URL "${STAGE12_DATABASE_URL_SECRET_NAME}" fi @@ -211,6 +214,7 @@ gcloud run deploy "${ENTRYPOINT_SERVICE}" \ --env-vars-file "${runtime_environment_file}" \ --set-secrets "${runtime_secrets_csv}" \ --cpu 1 \ + --no-cpu-throttling \ --memory 512Mi \ --concurrency 80 \ --timeout 60 \ diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 0a60b73e3..05394f178 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -91,10 +91,9 @@ privileges on tables the migrator creates later; revision 0002 revokes `UPDATE` on the definition tables, which never change. Both roles must exist before the first migration runs; an operator creates them. Before revision 0002 is deployed, the owner of the Public API's tables must also grant the -migrator `REFERENCES` on `policies`, `households` and `dynamics` (the -operator step grants it on `tax_benefit_model_versions`, `datasets` and -`regions` too, for later revisions); without it the migration fails and -applies nothing. The migrator needs either +migrator `REFERENCES` on `policies`, `households` and `dynamics`, and +before revision 0003 on `tax_benefit_model_versions`, `datasets` and +`regions`; without it the migration fails and applies nothing. The migrator needs either `CREATE` on the database or a `simulation` schema created in advance and owned by it: the environment creates the schema only when it is missing, because PostgreSQL checks `CREATE` on the database even for diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index 87ddd1f5e..901063056 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -56,6 +56,25 @@ rollback checks do not, because an older image may run against a newer schema. The ownership rules and roles are described in `libs/policyengine-simulation-persistence/README.md`. +## Recording run outcomes + +Each Simulation Entrypoint instance runs a background loop that asks Modal +about the runs' calls and records their outcomes in the `simulation` schema +(`policyengine_simulation_entry.run_recording`). Two deploy settings serve it: + +- `gcloud run deploy --no-cpu-throttling` keeps the instance's CPU allocated + outside requests, so the loop runs between them. Production keeps at least + one instance running (`entrypoint_min_instances: 1`); staging scales to + zero, so it records outcomes only while an instance runs, and any request + starts one. +- Every deploy mounts sim's own Modal token as `MODAL_TOKEN_ID` and + `MODAL_TOKEN_SECRET`, whatever the Stage 12 setting: the deploy step + resolves the Secret Manager secrets named by the `beta` and `prod` + environment variables `MODAL_TOKEN_ID_SECRET_NAME` and + `MODAL_TOKEN_SECRET_SECRET_NAME` to their newest enabled versions. The + runtime service account must be able to read both secrets. The Stage 12 + comparison path's Modal calls use the same token until it is removed. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py b/projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py new file mode 100644 index 000000000..ba08a2db6 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py @@ -0,0 +1,88 @@ +"""Deployment settings the background recording loop needs. + +The loop runs between requests, so Cloud Run must keep the instance's CPU +allocated outside them; and it asks Modal about calls, so every deploy mounts +sim's own Modal token, whether or not it also deploys the Stage 12 path. +""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets_cloud_run import cloud_run_deploy_environment +from test_deployment_assets_support import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + workflow_jobs, +) + + +def _run_deploy(tmp_path, **overrides): + environment = cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + arguments_capture = Path(environment["GCLOUD_ARGUMENTS_CAPTURE"]) + arguments = ( + arguments_capture.read_text(encoding="utf-8").splitlines() + if arguments_capture.exists() + else [] + ) + return result, arguments + + +@pytest.mark.parametrize("deploy_stage12", ["true", "false"]) +def test_every_deploy_keeps_cpu_allocated_between_requests( + tmp_path, deploy_stage12 +) -> None: + result, arguments = _run_deploy(tmp_path, DEPLOY_STAGE12_V2=deploy_stage12) + + assert result.returncode == 0, result.stderr + assert "--no-cpu-throttling" in arguments + assert "--cpu-throttling" not in arguments + + +@pytest.mark.parametrize("deploy_stage12", ["true", "false"]) +def test_every_deploy_mounts_sims_own_modal_token(tmp_path, deploy_stage12) -> None: + result, arguments = _run_deploy(tmp_path, DEPLOY_STAGE12_V2=deploy_stage12) + + assert result.returncode == 0, result.stderr + secrets = arguments[arguments.index("--set-secrets") + 1].split(",") + assert "MODAL_TOKEN_ID=modal-token-id:3" in secrets + assert "MODAL_TOKEN_SECRET=modal-token-secret:3" in secrets + assert len(secrets) == len(set(secrets)) + + +@pytest.mark.parametrize( + "variable", ["MODAL_TOKEN_ID_SECRET_NAME", "MODAL_TOKEN_SECRET_SECRET_NAME"] +) +def test_a_deploy_without_the_modal_token_stops(tmp_path, variable) -> None: + result, arguments = _run_deploy( + tmp_path, DEPLOY_STAGE12_V2="false", **{variable: ""} + ) + + assert result.returncode != 0 + assert variable in result.stderr + assert arguments == [] + + +def test_the_deploy_step_always_passes_the_modal_token_secret_names() -> None: + deploy_entrypoint = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + environment = next( + step["env"] + for step in deploy_entrypoint["steps"] + if "deploy-cloud-run-simulation-entry.sh" in step.get("run", "") + ) + + assert environment["MODAL_TOKEN_ID_SECRET_NAME"] == ( + "${{ vars.MODAL_TOKEN_ID_SECRET_NAME }}" + ) + assert environment["MODAL_TOKEN_SECRET_SECRET_NAME"] == ( + "${{ vars.MODAL_TOKEN_SECRET_SECRET_NAME }}" + ) From 71a01c8ffe87c45aa2ad904f17e8aba5cc7d2b10 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:01:21 +0400 Subject: [PATCH 06/39] Load-check the recording loop's conditional writes on PostgreSQL Six simulated instances, each on its own engine, claim due calls and record an outcome for each of 480 running household calls until every run has finished. Some outcomes are recorded only in a later pass, after an instance with nothing to claim has moved the clock past every lease, so another instance can claim the same call and record it first. The test requires every call's results to be written once and every run to finish once, and prints the elapsed time and number of passes as a measurement, with no threshold. With the status condition removed from the call updates, it fails. A staging run of the same check, against the staging database, is left to an operator. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_recording_load.py | 180 ++++++++++++++++++ 1 file changed, 180 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/tests/test_recording_load.py diff --git a/libs/policyengine-simulation-persistence/tests/test_recording_load.py b/libs/policyengine-simulation-persistence/tests/test_recording_load.py new file mode 100644 index 000000000..d3cc76745 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_recording_load.py @@ -0,0 +1,180 @@ +"""A load check of the recording loop's conditional writes on PostgreSQL. + +Several simulated Simulation Entrypoint instances, each on its own engine, +repeatedly claim due calls and record an outcome for each, as the recording +loop does, until every run has finished. Every call must be recorded once: +its results written once and its runs finished once, whatever the +interleaving. The elapsed time and the number of passes are printed, as a +measurement; nothing depends on them. +""" + +from __future__ import annotations + +import random +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier, Lock + +import pytest +from policyengine_simulation_persistence import ( + RunError, + RunErrorCode, + RunStatus, + RunTransitions, + create_simulation_engine, +) +from policyengine_simulation_persistence.runs import simulation_runs +from run_support import household_run +from simulation_support import household_simulation, simulation_store +from sqlalchemy import func, select +from transition_support import ManualClock, clocked_stores + +pytestmark = pytest.mark.postgres + +RUNS = 480 +INSTANCES = 6 +CLAIM_LIMIT = 20 +LEASE = timedelta(minutes=1) +MAXIMUM_PASSES = 400 + + +def _start_running_calls(database, rows, clock) -> int: + """Create household runs, each running under its own Modal call.""" + stores = clocked_stores(database, clock) + simulations = simulation_store(database) + for index in range(RUNS): + simulation = simulations.create( + household_simulation(rows, legacy_simulation_id=index + 1) + ) + run = stores.runs.start_simulation_run(household_run(simulation, rows)).run + stores.transitions.claim_simulation_dispatch([run.id]) + stores.transitions.record_simulation_dispatched( + [run.id], modal_call_id=f"fc-load-{index}", first_check_at=clock() + ) + stores.transitions.engine.dispose() + return RUNS + + +class _Outcomes: + """Remember every recorded outcome, from every instance.""" + + def __init__(self) -> None: + self.lock = Lock() + self.written_calls: list[str] = [] + self.failed_calls: list[str] = [] + + def write(self, call_id: str): + def write_results(_connection, _runs) -> None: + with self.lock: + self.written_calls.append(call_id) + + return write_results + + +def _record(transitions, outcomes, call_id: str, *, succeed: bool) -> None: + if succeed: + transitions.complete_simulation_call(call_id, outcomes.write(call_id)) + elif transitions.fail_simulation_call( + call_id, RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised.") + ): + with outcomes.lock: + outcomes.failed_calls.append(call_id) + + +def _instance(database, clock, outcomes, barrier, seed: int) -> int: + """Claim and record due calls until none is left; return the pass count. + + Some outcomes are recorded only in a later pass, after an instance with + nothing to claim has moved the clock past every lease, so another + instance may claim the same call and record it first: the late write + must then change nothing. + """ + chooser = random.Random(seed) # noqa: S311 # test interleaving, not security + transitions = RunTransitions( + create_simulation_engine(database.runtime_url), clock=clock + ) + late: list[tuple[str, bool]] = [] + barrier.wait() + passes = 0 + try: + while passes < MAXIMUM_PASSES: + passes += 1 + for call_id, succeed in late: + _record(transitions, outcomes, call_id, succeed=succeed) + late.clear() + calls = transitions.claim_due_simulation_calls( + limit=CLAIM_LIMIT, lease=LEASE + ) + if not calls: + if not _unfinished(transitions): + return passes + clock.advance(LEASE) + continue + for call in calls: + outcome = chooser.random() + if outcome < 0.3: + transitions.record_simulation_call_running( + call.modal_call_id, next_check_at=clock() + ) + elif outcome < 0.5: + late.append((call.modal_call_id, outcome < 0.45)) + else: + _record( + transitions, outcomes, call.modal_call_id, succeed=outcome < 0.9 + ) + return passes + finally: + transitions.engine.dispose() + + +def _unfinished(transitions: RunTransitions) -> int: + with transitions.engine.connect() as connection: + return connection.execute( + select(func.count()).where( + simulation_runs.c.status == RunStatus.RUNNING.value + ) + ).scalar_one() + + +def test_concurrent_instances_record_every_call_once( + migrated_database, public_api_rows, capsys +) -> None: + clock = ManualClock() + calls = _start_running_calls(migrated_database, public_api_rows, clock) + outcomes = _Outcomes() + barrier = Barrier(INSTANCES) + started = time.monotonic() + + with ThreadPoolExecutor(max_workers=INSTANCES) as pool: + passes = list( + pool.map( + lambda seed: _instance( + migrated_database, clock, outcomes, barrier, seed + ), + range(INSTANCES), + ) + ) + + elapsed = time.monotonic() - started + recorded = outcomes.written_calls + outcomes.failed_calls + assert len(recorded) == len(set(recorded)) == calls + stores = clocked_stores(migrated_database, clock) + with stores.transitions.engine.connect() as connection: + statuses = dict( + connection.execute( + select(simulation_runs.c.status, func.count()).group_by( + simulation_runs.c.status + ) + ).all() + ) + stores.transitions.engine.dispose() + assert statuses == { + RunStatus.SUCCEEDED.value: len(outcomes.written_calls), + RunStatus.FAILED.value: len(outcomes.failed_calls), + } + with capsys.disabled(): + print( # noqa: T201 # the measurement this check records + f"\nrecording load check: {calls} calls, {INSTANCES} instances, " + f"{sum(passes)} passes, {elapsed:.2f} s" + ) From 1d0861aeff94ed2dfba6b704af952ea0c2b23d20 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:31:40 +0400 Subject: [PATCH 07/39] Cancel a Modal call whose runs cannot record it RunDispatchService handled only a failed spawn. If recording the acknowledged call ID then raised (a dropped connection, or a wait for a pooled connection that ran out), the exception propagated before any log carried the call ID: the runs stayed dispatching until the recording loop failed them as abandoned, with a message saying the instance stopped, while the Modal call ran to its end with nobody collecting its output, and a resubmission started a second call for the same setup. Now a call that cannot be recorded is cancelled (CallDispatcher.cancel, which ModalCallDispatcher implements with FunctionCall.cancel), its runs are failed at once with dispatch_failed when the database allows, and an error log names the call. A call acknowledged after its runs stopped dispatching is cancelled too and logged as a warning. A failed cancel is logged and never raised. Co-Authored-By: Claude Opus 5.5 --- .../modal_calls.py | 10 ++- .../run_dispatch.py | 68 +++++++++++++++--- .../tests/run_fakes.py | 16 ++++- .../tests/test_run_dispatch.py | 71 ++++++++++++++++++- 4 files changed, 151 insertions(+), 14 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py index 92f79f0d0..e5eb7154f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py @@ -66,12 +66,16 @@ class CallInspectionError(RuntimeError): class CallDispatcher(Protocol): - """Start a Modal call without waiting for it.""" + """Start a Modal call without waiting for it, or cancel one.""" async def spawn(self, request: ModalRequest) -> str: """Return the call ID once Modal acknowledges the call.""" ... + async def cancel(self, modal_call_id: str) -> None: + """Ask Modal to stop a call whose runs could not record it.""" + ... + class CallInspector(Protocol): """Ask Modal about a call without waiting for it.""" @@ -108,6 +112,10 @@ async def spawn(self, request: ModalRequest) -> str: raise RuntimeError(msg) return modal_call_id + async def cancel(self, modal_call_id: str) -> None: + """Ask Modal to stop the call; its output will never be collected.""" + await modal.FunctionCall.from_id(modal_call_id).cancel.aio() + # Problems reaching Modal: the call's state is unknown, so ask again later. TRANSIENT_MODAL_ERRORS: tuple[type[Exception], ...] = ( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index 76eb398bb..97115dd9c 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -188,11 +188,63 @@ async def _start( ), ) return DispatchOutcome.DISPATCH_FAILED - recorded = await asyncio.to_thread( - steps.record, modal_call_id, self.clock() + FIRST_CHECK_DELAY - ) - logger.info( - "simulation_run_dispatched", - extra={**subject, "modal_call_id": modal_call_id, "recorded": recorded}, - ) - return DispatchOutcome.STARTED if recorded else DispatchOutcome.NOT_RECORDED + return await self._record(steps, modal_call_id, subject) + + async def _record( + self, steps: _DispatchSteps, modal_call_id: str, subject: dict[str, object] + ) -> DispatchOutcome: + """Record the acknowledged call, or cancel it if it cannot be recorded. + + A call whose runs never record it would run to its end with nobody + collecting its output, and its runs would later be failed as + abandoned; a resubmission would then start a second call. So such a + call is cancelled, and when the record failed, its runs are failed + at once with dispatch_failed if the database allows. + """ + context = {**subject, "modal_call_id": modal_call_id} + try: + recorded = await asyncio.to_thread( + steps.record, modal_call_id, self.clock() + FIRST_CHECK_DELAY + ) + except Exception as error: # the call is cancelled instead + logger.exception( + "simulation_run_call_not_recorded", + extra={**context, "error_type": type(error).__name__}, + ) + await self._cancel(modal_call_id, context) + await self._fail_quietly(steps, context) + return DispatchOutcome.NOT_RECORDED + if not recorded: + # The runs stopped dispatching while Modal acknowledged the call, + # for example because the recording loop failed them as abandoned. + logger.warning("simulation_run_call_not_recorded", extra=context) + await self._cancel(modal_call_id, context) + return DispatchOutcome.NOT_RECORDED + logger.info("simulation_run_dispatched", extra=context) + return DispatchOutcome.STARTED + + async def _cancel(self, modal_call_id: str, context: dict[str, object]) -> None: + try: + await self.dispatcher.cancel(modal_call_id) + except Exception as error: # logged; nothing else can be done + logger.exception( + "simulation_run_call_not_cancelled", + extra={**context, "error_type": type(error).__name__}, + ) + + async def _fail_quietly( + self, steps: _DispatchSteps, context: dict[str, object] + ) -> None: + try: + await asyncio.to_thread( + steps.fail, + RunError( + RunErrorCode.DISPATCH_FAILED, + "The Modal call could not be recorded, so it was cancelled.", + ), + ) + except Exception as error: # the abandoned-dispatch check fails the runs later + logger.exception( + "simulation_run_dispatch_failure_not_recorded", + extra={**context, "error_type": type(error).__name__}, + ) diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py index 586693a05..9e7b111d5 100644 --- a/projects/policyengine-simulation-entry/tests/run_fakes.py +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -404,10 +404,22 @@ def fail_abandoned_dispatches( class FakeDispatcher: """Record spawned requests and return call IDs, or fail like Modal.""" - def __init__(self, failure: Exception | None = None) -> None: - """Fail every spawn with ``failure`` when it is set.""" + def __init__( + self, + failure: Exception | None = None, + cancel_failure: Exception | None = None, + ) -> None: + """Fail spawns with ``failure`` and cancels with ``cancel_failure``.""" self.failure = failure + self.cancel_failure = cancel_failure self.requests: list[ModalRequest] = [] + self.cancelled: list[str] = [] + + async def cancel(self, modal_call_id: str) -> None: + """Record the cancelled call, or raise the configured failure.""" + self.cancelled.append(modal_call_id) + if self.cancel_failure is not None: + raise self.cancel_failure async def spawn(self, request: ModalRequest) -> str: """Return a new call ID, or raise the configured failure.""" diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index f4378c4a6..dc2e24cac 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -99,13 +99,12 @@ async def test_a_refused_call_fails_the_runs_without_raising() -> None: @pytest.mark.asyncio -async def test_a_call_whose_runs_stopped_dispatching_is_reported() -> None: +async def test_a_call_whose_runs_stopped_dispatching_is_cancelled(caplog) -> None: service, transitions, _ = _service() run_id = transitions.add_run() class SlowDispatcher(FakeDispatcher): async def spawn(self, request): - # The recording loop gave up on the dispatch in the meantime. transitions.runs[run_id].status = RunStatus.FAILED return await super().spawn(request) @@ -114,7 +113,73 @@ async def spawn(self, request): outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) assert outcome is DispatchOutcome.NOT_RECORDED - assert execution(transitions, run_id).modal_call_id is None + assert service.dispatcher.cancelled == ["fc-1"] + assert "simulation_run_call_not_recorded" in caplog.text + + +@pytest.mark.asyncio +async def test_a_call_that_cannot_be_recorded_is_cancelled_and_its_runs_failed( + caplog, +) -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + def unreachable_database(*_arguments, **_keywords): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert service.dispatcher.cancelled == ["fc-1"] + run = execution(transitions, run_id) + assert run.status is RunStatus.FAILED + assert run.error.code is RunErrorCode.DISPATCH_FAILED + assert "could not be recorded" in run.error.message + assert any( + getattr(record, "modal_call_id", None) == "fc-1" for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_a_failed_cancel_is_logged_without_raising(caplog) -> None: + service, transitions, _ = _service( + FakeDispatcher(cancel_failure=RuntimeError("modal down")) + ) + run_id = transitions.add_run() + + def unreachable_database(*_arguments, **_keywords): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] + transitions.record_simulation_dispatch_failure = unreachable_database # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert "simulation_run_call_not_cancelled" in caplog.text + + +def test_the_dispatcher_cancels_a_call_by_its_id(monkeypatch) -> None: + cancelled = [] + + class _FakeCancel: + async def aio(self): + cancelled.append("cancelled") + + class _FakeCall: + cancel = _FakeCancel() + + monkeypatch.setattr( + modal_calls.modal.FunctionCall, "from_id", lambda call_id: _FakeCall() + ) + + asyncio.run(ModalCallDispatcher("staging").cancel("fc-orphan")) + + assert cancelled == ["cancelled"] @pytest.mark.asyncio From 890dc4b59b1d768b05cd66a354e53765e803b9e8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:32:59 +0400 Subject: [PATCH 08/39] Fail a comparison step whose simulation run failed A pending comparison step whose baseline or reform simulation run failed could never start and stayed pending, so its report run's status never changed. RunTransitions.fail_comparisons_of_failed_runs fails such steps with simulation_run_failed in one conditional update, and each recording pass runs it next to the abandoned-dispatch check, so a report run reaches a final status without any caller polling. The persistence README now describes the conditional writes and notes that starting a comparison step is left to the stage 2 report path. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 10 +++ .../run_transitions.py | 45 ++++++++++++- .../tests/test_run_transitions.py | 27 ++++++++ .../run_recording.py | 65 +++++++++++++------ .../tests/run_fakes.py | 26 +++++++- .../tests/test_run_recording.py | 16 +++++ 6 files changed, 166 insertions(+), 23 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 05394f178..7e629f5e7 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -52,6 +52,16 @@ and takes the catalog version and year from them; a report's status is its newest report run's. The runtime role may read, insert and update runs but not delete them. +`RunTransitions` makes every status change one conditional update on the +expected status, so a second writer changes nothing: only the caller that +moves runs from `pending` to `dispatching` starts their Modal call, the runs +one call computes change together, the recording loop claims a due call by +moving its next check time, results are written in the transaction that +marks a run succeeded, and dispatches left `dispatching` (and comparison +steps whose simulation run failed) are failed so the work can be submitted +again. A comparison step can start only once both of its simulation runs +have succeeded; starting it is left to the stage 2 report path. + ## Ownership rules These rules are registered in `rules/registry.toml` and checked by diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index 39c6ee2c8..89e4197ff 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -29,7 +29,7 @@ from enum import StrEnum from typing import TYPE_CHECKING -from sqlalchemy import and_, exists, select, update +from sqlalchemy import and_, exists, or_, select, update from policyengine_simulation_persistence.run_store import utc_now from policyengine_simulation_persistence.runs import ( @@ -71,6 +71,8 @@ class RunErrorCode(StrEnum): INVALID_RESULT = "invalid_result" # No recorder handles results of this kind of run. NO_RESULT_RECORDER = "no_result_recorder" + # A comparison step whose baseline or reform simulation run failed. + SIMULATION_RUN_FAILED = "simulation_run_failed" @dataclass(frozen=True) @@ -512,6 +514,47 @@ def _comparison_failure(self, error: RunError) -> dict[str, object]: "comparison_next_check_at": None, } + def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: + """Fail pending comparison steps whose baseline or reform run failed. + + Such a step can never start, so it fails with ``simulation_run_failed`` + and its report run's status says so. Return the report runs changed. + """ + baseline = simulation_runs.alias("failed_baseline_run") + reform = simulation_runs.alias("failed_reform_run") + either_run_failed = or_( + exists().where( + and_( + baseline.c.id == report_runs.c.baseline_simulation_run_id, + baseline.c.status == RunStatus.FAILED.value, + ) + ), + exists().where( + and_( + reform.c.id == report_runs.c.reform_simulation_run_id, + reform.c.status == RunStatus.FAILED.value, + ) + ), + ) + failure = self._comparison_failure( + RunError( + RunErrorCode.SIMULATION_RUN_FAILED, + "A simulation run this comparison needs failed.", + ) + ) + with self.engine.begin() as connection: + return tuple( + connection.execute( + update(report_runs) + .where( + report_runs.c.comparison_status == RunStatus.PENDING.value, + either_run_failed, + ) + .values(updated_at=self._clock(), **failure) + .returning(report_runs.c.id) + ).scalars() + ) + # --- Dispatches that never finished ------------------------------------ def fail_abandoned_dispatches( diff --git a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py index 689683849..d08655317 100644 --- a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py +++ b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py @@ -489,3 +489,30 @@ def test_dispatches_left_for_ten_minutes_are_failed( == "dispatch_abandoned" ) assert stores.runs.get_simulation_run(recent_run.id).status is RunStatus.DISPATCHING + + +def test_a_comparison_of_a_failed_run_fails_too( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + waiting, _ = _economy_report_run(stores, migrated_database, public_api_rows) + _succeed(stores, baseline_run, "fc-baseline") + stores.transitions.claim_simulation_dispatch([reform_run.id]) + stores.transitions.record_simulation_dispatched( + [reform_run.id], modal_call_id="fc-reform", first_check_at=stores.clock() + ) + stores.transitions.fail_simulation_call("fc-reform", WORKER_ERROR) + + failed = stores.transitions.fail_comparisons_of_failed_runs() + + assert failed == (report_run.id,) + record = stores.runs.get_report_run(report_run.id) + assert record.comparison_status is RunStatus.FAILED + assert record.comparison_error_code == "simulation_run_failed" + assert stores.runs.get_report_run(waiting.id).comparison_status is ( + RunStatus.PENDING + ) + assert stores.transitions.fail_comparisons_of_failed_runs() == () diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 126392609..859a04b4c 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -3,7 +3,8 @@ Each Simulation Entrypoint instance runs one loop. On every pass it: 1. fails runs and comparison steps left ``dispatching`` for ten minutes, so - they can be submitted again; + they can be submitted again, and pending comparison steps whose baseline + or reform run failed, since they can never start; 2. claims the running calls that are due a check, by moving their next check one lease ahead, so instances do not check the same call at once; 3. asks Modal about each claimed call without waiting, and records what it @@ -173,6 +174,10 @@ def fail_abandoned_dispatches( """Fail runs and comparison steps left dispatching.""" ... + def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: + """Fail pending comparison steps whose simulation runs failed.""" + ... + @dataclass class RecordingPass: @@ -184,11 +189,12 @@ class RecordingPass: still_running: int = 0 not_reached: int = 0 abandoned: int = 0 + comparisons_without_runs: int = 0 @property def found_work(self) -> bool: """Whether the pass checked or failed anything.""" - return bool(self.checked or self.abandoned) + return bool(self.checked or self.abandoned or self.comparisons_without_runs) @dataclass(frozen=True) @@ -234,24 +240,7 @@ class RunRecorder: async def record_due_outcomes(self) -> RecordingPass: """Run one pass: fail abandoned dispatches, then check every due call.""" recording = RecordingPass() - abandoned = await asyncio.to_thread( - lambda: self.transitions.fail_abandoned_dispatches( - claimed_before=self.clock() - DISPATCH_ABANDONED_AFTER - ) - ) - recording.abandoned = len(abandoned.simulation_run_ids) + len( - abandoned.report_run_ids - ) - if recording.abandoned: - logger.warning( - "simulation_run_dispatches_abandoned", - extra={ - "simulation_run_ids": [ - str(i) for i in abandoned.simulation_run_ids - ], - "report_run_ids": [str(i) for i in abandoned.report_run_ids], - }, - ) + await self._fail_stopped_work(recording) calls = await asyncio.to_thread( lambda: self.transitions.claim_due_simulation_calls( limit=CLAIM_BATCH_SIZE, lease=CHECK_LEASE @@ -280,6 +269,42 @@ async def check(modal_call_id: str, call_handlers: _CallHandlers) -> None: ) return recording + async def _fail_stopped_work(self, recording: RecordingPass) -> None: + """Fail dispatches that stopped, and comparisons that can never start.""" + abandoned = await asyncio.to_thread( + lambda: self.transitions.fail_abandoned_dispatches( + claimed_before=self.clock() - DISPATCH_ABANDONED_AFTER + ) + ) + recording.abandoned = len(abandoned.simulation_run_ids) + len( + abandoned.report_run_ids + ) + if recording.abandoned: + logger.warning( + "simulation_run_dispatches_abandoned", + extra={ + "simulation_run_ids": [ + str(run_id) for run_id in abandoned.simulation_run_ids + ], + "report_run_ids": [ + str(report_run_id) for report_run_id in abandoned.report_run_ids + ], + }, + ) + without_runs = await asyncio.to_thread( + self.transitions.fail_comparisons_of_failed_runs + ) + recording.comparisons_without_runs = len(without_runs) + if without_runs: + logger.warning( + "report_run_comparisons_failed_with_their_runs", + extra={ + "report_run_ids": [ + str(report_run_id) for report_run_id in without_runs + ] + }, + ) + def _simulation_handlers( self, call: DueSimulationCall ) -> tuple[str, _CallHandlers]: diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py index 9e7b111d5..799ed4f0e 100644 --- a/projects/policyengine-simulation-entry/tests/run_fakes.py +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -23,6 +23,7 @@ AbandonedDispatches, DueSimulationCall, RunError, + RunErrorCode, RunStatus, SimulationKind, ) @@ -96,6 +97,7 @@ def __init__(self, clock: ManualClock) -> None: self.runs: dict[UUID, FakeExecution] = {} self.comparisons: dict[UUID, FakeExecution] = {} self.comparison_runs_succeeded: dict[UUID, bool] = {} + self.comparisons_with_failed_runs: set[UUID] = set() # --- Building state -------------------------------------------------- @@ -105,11 +107,17 @@ def add_run(self, kind: str = "economy", **values: Any) -> UUID: self.runs[run.id] = run return run.id - def add_comparison(self, *, runs_succeeded: bool = True, **values: Any) -> UUID: + def add_comparison( + self, *, runs_succeeded: bool = True, runs_failed: bool = False, **values: Any + ) -> UUID: """Add an economy report run's comparison step and return its ID.""" comparison = FakeExecution(id=uuid4(), kind="economy", **values) self.comparisons[comparison.id] = comparison - self.comparison_runs_succeeded[comparison.id] = runs_succeeded + self.comparison_runs_succeeded[comparison.id] = ( + runs_succeeded and not runs_failed + ) + if runs_failed: + self.comparisons_with_failed_runs.add(comparison.id) return comparison.id def running_call(self, call_id: str, *kinds: str) -> list[UUID]: @@ -373,6 +381,20 @@ def fail_comparison(self, report_run_id: UUID, error: RunError) -> bool: is not None ) + def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: + """Fail pending comparison steps whose simulation runs failed.""" + failed = [] + for comparison_id in sorted(self.comparisons_with_failed_runs): + changed = self._change_comparison( + comparison_id, + RunStatus.PENDING, + status=RunStatus.FAILED, + error=RunError(RunErrorCode.SIMULATION_RUN_FAILED, "A run failed."), + ) + if changed is not None: + failed.append(comparison_id) + return tuple(failed) + def fail_abandoned_dispatches( self, *, claimed_before: datetime ) -> AbandonedDispatches: diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index 687b36222..9a856bd19 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -266,6 +266,22 @@ def test_dispatches_left_for_ten_minutes_are_failed_by_a_pass() -> None: assert transitions.runs[recent].status is RunStatus.DISPATCHING +def test_a_pass_fails_comparisons_whose_simulation_run_failed() -> None: + recorder, transitions, _, _ = _recorder() + stuck = transitions.add_comparison(runs_failed=True) + waiting = transitions.add_comparison(runs_succeeded=False) + + recording = _run(recorder) + + assert recording.comparisons_without_runs == 1 + assert recording.found_work + assert transitions.comparisons[stuck].status is RunStatus.FAILED + assert transitions.comparisons[stuck].error.code is ( + RunErrorCode.SIMULATION_RUN_FAILED + ) + assert transitions.comparisons[waiting].status is RunStatus.PENDING + + class _ScriptedRecorder: """Return scripted passes, or raise, and count how often it ran.""" From 0aebd875914f803e3bddca8be791c0c8210a320d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:34:39 +0400 Subject: [PATCH 09/39] Leave a finished call of a kind without a recorder for a later check During a deploy the candidate revision runs beside the current one, and with CPU always allocated an idle older instance keeps running its loop. An older revision with no result recorder for a run's kind failed the run with no_result_recorder, a final status, although the newer revision would have recorded it. Such a call now gets a later check, backing off like a running call, and a warning names it; a revision with the recorder records it, and Modal's output retention bounds the wait (the call then fails with result_expired). The no_result_recorder error code is gone. Co-Authored-By: Claude Opus 5.5 --- .../run_transitions.py | 2 -- .../run_recording.py | 20 +++++++++++-------- .../simulation_database.py | 3 ++- .../tests/test_run_recording.py | 13 ++++++++---- 4 files changed, 23 insertions(+), 15 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index 89e4197ff..ac4c326f1 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -69,8 +69,6 @@ class RunErrorCode(StrEnum): CALL_NOT_FOUND = "call_not_found" # The output failed validation against its result model. INVALID_RESULT = "invalid_result" - # No recorder handles results of this kind of run. - NO_RESULT_RECORDER = "no_result_recorder" # A comparison step whose baseline or reform simulation run failed. SIMULATION_RUN_FAILED = "simulation_run_failed" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 859a04b4c..e0c19fec4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -11,7 +11,8 @@ found with conditional writes: a finished call's result is validated and its result rows are written in the transaction that marks its runs ``succeeded``; a failed or timed-out call is marked ``failed`` with its - error; a call still running gets a later check, backing off. + error; a call still running, or finished but of a kind this revision has + no recorder for, gets a later check, backing off. Pre-emption is never a failure: Modal re-runs a pre-empted call under the same call ID, so it reads as still running. No caller polls, and Modal holds @@ -190,6 +191,7 @@ class RecordingPass: not_reached: int = 0 abandoned: int = 0 comparisons_without_runs: int = 0 + without_recorder: int = 0 @property def found_work(self) -> bool: @@ -397,13 +399,15 @@ async def _record( return steps = handlers.result_steps if steps is None: - await self._fail( - handlers, - RunError( - RunErrorCode.NO_RESULT_RECORDER, - "No result recorder handles this kind of run.", - ), - recording, + # During a deploy an older revision runs beside the new one, and + # with CPU always allocated its loop keeps checking calls. A + # revision without a recorder for this kind leaves the call for + # one that has it; Modal's output retention bounds the wait. + next_check_at = self.clock() + next_check_delay(handlers.check_count) + await asyncio.to_thread(handlers.keep_running, next_check_at) + recording.without_recorder += 1 + logger.warning( + "simulation_run_without_result_recorder", extra=handlers.subject ) return try: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index 46280ca2e..43f595bbb 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -49,7 +49,8 @@ def recording_loop(self) -> RunRecordingLoop: """Return the loop that records Modal call outcomes on runs. No kind of run has a result recorder yet, so a finished call would - fail its runs with ``no_result_recorder``; no route creates runs yet. + stay running until a revision with a recorder records it; no route + creates runs yet. """ return RunRecordingLoop( RunRecorder( diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index 9a856bd19..add632eba 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -147,14 +147,19 @@ def test_a_household_pair_is_written_once_for_both_runs() -> None: assert {transitions.runs[run_id].status for run_id in pair} == {RunStatus.SUCCEEDED} -def test_a_kind_without_a_recorder_fails_its_runs() -> None: - recorder, transitions, inspector, _ = _recorder() +def test_a_kind_without_a_recorder_waits_for_an_instance_with_one() -> None: + """An older revision still running beside a new one must not fail the run.""" + recorder, transitions, inspector, clock = _recorder() (run_id,) = transitions.running_call("fc-1") inspector.outcomes["fc-1"] = CallSucceeded({"budgetary_impact": 1}) - _run(recorder) + recording = _run(recorder) - assert transitions.runs[run_id].error.code is RunErrorCode.NO_RESULT_RECORDER + assert recording.without_recorder == 1 + run = transitions.runs[run_id] + assert run.status is RunStatus.RUNNING + assert run.error is None + assert run.next_check_at == clock() + timedelta(seconds=2) def test_an_invalid_output_fails_its_runs_without_repeating_it() -> None: From 64c78b80fb25cc457a60f9884c15b380a6fb4e6b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:36:41 +0400 Subject: [PATCH 10/39] Leave database connections for requests, and log comparisons as comparisons - The recording loop checked up to ten calls at once, each writing through the runtime engine's pool of two connections plus three overflow, which requests share; a busy pass could leave requests waiting for a connection. It now checks at most three calls at once, and a test keeps that below the pool's size. - Comparison steps were logged under simulation_run_* event names. The dispatch service and the loop now name events for a comparison step report_run_comparison_* (for example report_run_comparison_failed). - Run IDs in log fields are no longer built with a one-letter name. Co-Authored-By: Claude Opus 5.5 --- .../run_dispatch.py | 32 +++++++++++++------ .../run_recording.py | 28 ++++++++++++---- .../tests/test_run_dispatch.py | 5 ++- .../tests/test_run_recording.py | 22 +++++++++++++ 4 files changed, 71 insertions(+), 16 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index 97115dd9c..fd5053501 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -33,6 +33,12 @@ logger = logging.getLogger(__name__) + +def _event_name(prefix: str, outcome: str) -> str: + """Name a log event for a simulation run or a comparison step.""" + return f"{prefix}_{outcome}" + + # The first check of a new call; the recording loop backs off from there. FIRST_CHECK_DELAY = timedelta(seconds=2) @@ -102,6 +108,8 @@ class _DispatchSteps: claim: Callable[[], bool] record: Callable[[str, datetime], bool] fail: Callable[[RunError], bool] + # "simulation_run" or "report_run_comparison": the log event names' start. + event_prefix: str @dataclass @@ -132,6 +140,7 @@ async def start_simulation_runs( fail=lambda error: self.transitions.record_simulation_dispatch_failure( runs, error ), + event_prefix="simulation_run", ), request, subject={"simulation_run_ids": sorted(str(run_id) for run_id in runs)}, @@ -158,6 +167,7 @@ async def start_comparison( fail=lambda error: self.transitions.record_comparison_dispatch_failure( report_run_id, error ), + event_prefix="report_run_comparison", ), request, subject={"report_run_id": str(report_run_id)}, @@ -177,7 +187,7 @@ async def _start( # Any failure to start the call fails the runs; they can be run again. except Exception as error: # noqa: BLE001 # recorded on the runs logger.warning( - "simulation_run_dispatch_failed", + _event_name(steps.event_prefix, "dispatch_failed"), extra={**subject, "error_type": type(error).__name__}, ) await asyncio.to_thread( @@ -208,27 +218,31 @@ async def _record( ) except Exception as error: # the call is cancelled instead logger.exception( - "simulation_run_call_not_recorded", + _event_name(steps.event_prefix, "call_not_recorded"), extra={**context, "error_type": type(error).__name__}, ) - await self._cancel(modal_call_id, context) + await self._cancel(modal_call_id, context, steps.event_prefix) await self._fail_quietly(steps, context) return DispatchOutcome.NOT_RECORDED if not recorded: # The runs stopped dispatching while Modal acknowledged the call, # for example because the recording loop failed them as abandoned. - logger.warning("simulation_run_call_not_recorded", extra=context) - await self._cancel(modal_call_id, context) + logger.warning( + _event_name(steps.event_prefix, "call_not_recorded"), extra=context + ) + await self._cancel(modal_call_id, context, steps.event_prefix) return DispatchOutcome.NOT_RECORDED - logger.info("simulation_run_dispatched", extra=context) + logger.info(_event_name(steps.event_prefix, "dispatched"), extra=context) return DispatchOutcome.STARTED - async def _cancel(self, modal_call_id: str, context: dict[str, object]) -> None: + async def _cancel( + self, modal_call_id: str, context: dict[str, object], event_prefix: str + ) -> None: try: await self.dispatcher.cancel(modal_call_id) except Exception as error: # logged; nothing else can be done logger.exception( - "simulation_run_call_not_cancelled", + _event_name(event_prefix, "call_not_cancelled"), extra={**context, "error_type": type(error).__name__}, ) @@ -245,6 +259,6 @@ async def _fail_quietly( ) except Exception as error: # the abandoned-dispatch check fails the runs later logger.exception( - "simulation_run_dispatch_failure_not_recorded", + _event_name(steps.event_prefix, "dispatch_failure_not_recorded"), extra={**context, "error_type": type(error).__name__}, ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index e0c19fec4..b910a81e7 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -57,6 +57,12 @@ logger = logging.getLogger(__name__) + +def _event_name(prefix: str, outcome: str) -> str: + """Name a log event for a simulation run or a comparison step.""" + return f"{prefix}_{outcome}" + + # A claimed call is not checked by another instance for this long; if the # claiming instance stops, the call becomes due again afterwards. CHECK_LEASE = timedelta(minutes=1) @@ -65,7 +71,10 @@ FIRST_BACKOFF = timedelta(seconds=2) LONGEST_BACKOFF = timedelta(seconds=60) CLAIM_BATCH_SIZE = 50 -INSPECTION_CONCURRENCY = 10 +# Each check writes through the runtime engine's small pool, which requests +# share (policyengine_simulation_persistence.engine: 2 connections plus 3 +# overflow), so the loop uses at most three at once. +INSPECTION_CONCURRENCY = 3 BUSY_PASS_INTERVAL_SECONDS = 1.0 IDLE_PASS_INTERVAL_SECONDS = 5.0 @@ -215,6 +224,8 @@ class _CallHandlers: """ subject: dict[str, object] + # "simulation_run" or "report_run_comparison": the log event names' start. + event_prefix: str check_count: int result_steps: _ResultSteps | None fail: Callable[[RunError], bool] @@ -332,6 +343,7 @@ def _simulation_handlers( "modal_call_id": call_id, "simulation_run_ids": [str(run.id) for run in call.runs], }, + event_prefix="simulation_run", check_count=max(run.check_count for run in call.runs), result_steps=result_steps, fail=lambda error: self.transitions.fail_simulation_call(call_id, error), @@ -360,6 +372,7 @@ def _comparison_handlers( ) return call_id, _CallHandlers( subject={"modal_call_id": call_id, "report_run_id": str(report_run_id)}, + event_prefix="report_run_comparison", check_count=report_run.comparison_check_count, result_steps=result_steps, fail=lambda error: self.transitions.fail_comparison(report_run_id, error), @@ -380,7 +393,7 @@ async def _check( # The claim's lease makes the call due again; nothing to record. recording.not_reached += 1 logger.warning( - "simulation_run_check_failed", + _event_name(handlers.event_prefix, "check_failed"), extra={**handlers.subject, "error_type": str(error)}, ) return @@ -407,7 +420,8 @@ async def _record( await asyncio.to_thread(handlers.keep_running, next_check_at) recording.without_recorder += 1 logger.warning( - "simulation_run_without_result_recorder", extra=handlers.subject + _event_name(handlers.event_prefix, "without_result_recorder"), + extra=handlers.subject, ) return try: @@ -423,13 +437,15 @@ async def _record( except Exception as error: recording.not_reached += 1 logger.exception( - "simulation_run_result_not_written", + _event_name(handlers.event_prefix, "result_not_written"), extra={**handlers.subject, "error_type": type(error).__name__}, ) return if completed: recording.succeeded += 1 - logger.info("simulation_run_succeeded", extra=handlers.subject) + logger.info( + _event_name(handlers.event_prefix, "succeeded"), extra=handlers.subject + ) async def _fail( self, handlers: _CallHandlers, error: RunError, recording: RecordingPass @@ -437,7 +453,7 @@ async def _fail( if await asyncio.to_thread(handlers.fail, error): recording.failed += 1 logger.warning( - "simulation_run_failed", + _event_name(handlers.event_prefix, "failed"), extra={**handlers.subject, "error_code": error.code.value}, ) diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index dc2e24cac..4c01128bf 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -200,7 +200,7 @@ async def test_a_comparison_starts_only_once_its_runs_succeeded() -> None: @pytest.mark.asyncio -async def test_a_refused_comparison_call_fails_the_step() -> None: +async def test_a_refused_comparison_call_fails_the_step(caplog) -> None: service, transitions, _ = _service(FakeDispatcher(failure=RuntimeError("down"))) comparison = transitions.add_comparison() @@ -208,6 +208,9 @@ async def test_a_refused_comparison_call_fails_the_step() -> None: assert outcome is DispatchOutcome.DISPATCH_FAILED assert transitions.comparisons[comparison].status is RunStatus.FAILED + assert [record.message for record in caplog.records] == [ + "report_run_comparison_dispatch_failed" + ] class _FakeCall: diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index add632eba..8d45e2962 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -8,6 +8,7 @@ import pytest from conftest import FakeBackend, make_settings from fastapi.testclient import TestClient +from policyengine_simulation_entry import run_recording from policyengine_simulation_entry.app import create_app from policyengine_simulation_entry.modal_calls import ( CallFailed, @@ -32,6 +33,7 @@ RunStatus, SimulationKind, ) +from policyengine_simulation_persistence.engine import MAX_OVERFLOW, POOL_SIZE from pydantic import BaseModel from run_fakes import FakeInspector, FakeRunTransitions, ManualClock @@ -287,6 +289,26 @@ def test_a_pass_fails_comparisons_whose_simulation_run_failed() -> None: assert transitions.comparisons[waiting].status is RunStatus.PENDING +def test_comparison_outcomes_are_logged_as_comparison_events(caplog) -> None: + recorder, transitions, inspector, clock = _recorder() + transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-failed", next_check_at=clock() + ) + inspector.outcomes["fc-failed"] = CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + ) + + _run(recorder) + + assert [record.message for record in caplog.records] == [ + "report_run_comparison_failed" + ] + + +def test_the_loop_leaves_database_connections_for_requests() -> None: + assert run_recording.INSPECTION_CONCURRENCY < POOL_SIZE + MAX_OVERFLOW + + class _ScriptedRecorder: """Return scripted passes, or raise, and count how often it ran.""" From bc0ae0072386baa3650b5cc1208ae2903bb13998 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:37:55 +0400 Subject: [PATCH 11/39] Load-check household pairs and stalled dispatches too The load check covered only single-run calls. A third of its 480 calls are now household pairs (two runs under one call), and 30 runs are left dispatching as if their instance had stopped; every instance also fails stalled dispatches on each pass, as the recording loop does. It requires each call to be recorded once with all of its runs finished, and each stalled dispatch to be failed once. With the dispatching condition removed from the abandoned-dispatch update, it fails. A run against the staging database remains an operator step. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_recording_load.py | 87 ++++++++++++++----- 1 file changed, 64 insertions(+), 23 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_recording_load.py b/libs/policyengine-simulation-persistence/tests/test_recording_load.py index d3cc76745..efdec7d0b 100644 --- a/libs/policyengine-simulation-persistence/tests/test_recording_load.py +++ b/libs/policyengine-simulation-persistence/tests/test_recording_load.py @@ -1,11 +1,13 @@ """A load check of the recording loop's conditional writes on PostgreSQL. Several simulated Simulation Entrypoint instances, each on its own engine, -repeatedly claim due calls and record an outcome for each, as the recording -loop does, until every run has finished. Every call must be recorded once: -its results written once and its runs finished once, whatever the -interleaving. The elapsed time and the number of passes are printed, as a -measurement; nothing depends on them. +repeatedly fail stalled dispatches, claim due calls and record an outcome for +each, as the recording loop does, until every run has finished. A third of +the calls are household pairs (two runs under one call). Every call must be +recorded once, its results written once and all of its runs finished once, +and every stalled dispatch failed once, whatever the interleaving. The +elapsed time and the number of passes are printed, as a measurement; nothing +depends on them. """ from __future__ import annotations @@ -32,28 +34,51 @@ pytestmark = pytest.mark.postgres -RUNS = 480 +CALLS = 480 +# Every third call is a household report's pair: two runs under one call. +PAIR_EVERY = 3 +# Runs claimed by an instance that stopped before recording their call. +STALLED_DISPATCHES = 30 INSTANCES = 6 CLAIM_LIMIT = 20 LEASE = timedelta(minutes=1) +ABANDONED_AFTER = timedelta(minutes=10) MAXIMUM_PASSES = 400 -def _start_running_calls(database, rows, clock) -> int: - """Create household runs, each running under its own Modal call.""" +def _household_run(stores, simulations, rows, legacy_id: int): + simulation = simulations.create( + household_simulation(rows, legacy_simulation_id=legacy_id) + ) + return stores.runs.start_simulation_run(household_run(simulation, rows)).run + + +def _start_running_calls(database, rows, clock) -> dict[str, int]: + """Create running calls, a third of them household pairs, and stalled runs. + + Return each call's number of runs. + """ stores = clocked_stores(database, clock) simulations = simulation_store(database) - for index in range(RUNS): - simulation = simulations.create( - household_simulation(rows, legacy_simulation_id=index + 1) - ) - run = stores.runs.start_simulation_run(household_run(simulation, rows)).run - stores.transitions.claim_simulation_dispatch([run.id]) + legacy_ids = iter(range(1, 10_000)) + runs_per_call = {} + for index in range(CALLS): + count = 2 if index % PAIR_EVERY == 0 else 1 + run_ids = [ + _household_run(stores, simulations, rows, next(legacy_ids)).id + for _ in range(count) + ] + call_id = f"fc-load-{index}" + stores.transitions.claim_simulation_dispatch(run_ids) stores.transitions.record_simulation_dispatched( - [run.id], modal_call_id=f"fc-load-{index}", first_check_at=clock() + run_ids, modal_call_id=call_id, first_check_at=clock() ) + runs_per_call[call_id] = count + for _ in range(STALLED_DISPATCHES): + run = _household_run(stores, simulations, rows, next(legacy_ids)) + stores.transitions.claim_simulation_dispatch([run.id]) stores.transitions.engine.dispose() - return RUNS + return runs_per_call class _Outcomes: @@ -63,6 +88,7 @@ def __init__(self) -> None: self.lock = Lock() self.written_calls: list[str] = [] self.failed_calls: list[str] = [] + self.abandoned_runs: list[object] = [] def write(self, call_id: str): def write_results(_connection, _runs) -> None: @@ -103,11 +129,19 @@ def _instance(database, clock, outcomes, barrier, seed: int) -> int: for call_id, succeed in late: _record(transitions, outcomes, call_id, succeed=succeed) late.clear() + abandoned = transitions.fail_abandoned_dispatches( + claimed_before=clock() - ABANDONED_AFTER + ) + with outcomes.lock: + outcomes.abandoned_runs.extend(abandoned.simulation_run_ids) calls = transitions.claim_due_simulation_calls( limit=CLAIM_LIMIT, lease=LEASE ) if not calls: - if not _unfinished(transitions): + if ( + not _unfinished(transitions) + and len(outcomes.abandoned_runs) == STALLED_DISPATCHES + ): return passes clock.advance(LEASE) continue @@ -141,7 +175,7 @@ def test_concurrent_instances_record_every_call_once( migrated_database, public_api_rows, capsys ) -> None: clock = ManualClock() - calls = _start_running_calls(migrated_database, public_api_rows, clock) + runs_per_call = _start_running_calls(migrated_database, public_api_rows, clock) outcomes = _Outcomes() barrier = Barrier(INSTANCES) started = time.monotonic() @@ -158,7 +192,9 @@ def test_concurrent_instances_record_every_call_once( elapsed = time.monotonic() - started recorded = outcomes.written_calls + outcomes.failed_calls - assert len(recorded) == len(set(recorded)) == calls + assert sorted(recorded) == sorted(runs_per_call) + assert len(outcomes.abandoned_runs) == len(set(outcomes.abandoned_runs)) + assert len(outcomes.abandoned_runs) == STALLED_DISPATCHES stores = clocked_stores(migrated_database, clock) with stores.transitions.engine.connect() as connection: statuses = dict( @@ -170,11 +206,16 @@ def test_concurrent_instances_record_every_call_once( ) stores.transitions.engine.dispose() assert statuses == { - RunStatus.SUCCEEDED.value: len(outcomes.written_calls), - RunStatus.FAILED.value: len(outcomes.failed_calls), + RunStatus.SUCCEEDED.value: sum( + runs_per_call[call_id] for call_id in outcomes.written_calls + ), + RunStatus.FAILED.value: STALLED_DISPATCHES + + sum(runs_per_call[call_id] for call_id in outcomes.failed_calls), } with capsys.disabled(): print( # noqa: T201 # the measurement this check records - f"\nrecording load check: {calls} calls, {INSTANCES} instances, " - f"{sum(passes)} passes, {elapsed:.2f} s" + f"\nrecording load check: {len(runs_per_call)} calls " + f"({sum(runs_per_call.values())} runs), {STALLED_DISPATCHES} stalled " + f"dispatches, {INSTANCES} instances, {sum(passes)} passes, " + f"{elapsed:.2f} s" ) From f2f88b14257bb8fc1537640a893b8358ae0aa2b4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:58:51 +0400 Subject: [PATCH 12/39] Contain a failed check within its pass, and retry client-side Modal errors - RunRecorder gathered a pass's checks without catching their errors. A database error in one check's write (for example a dropped connection) made the pass raise while its other checks kept running, so the next pass started beside them and the three-check limit no longer held, with their exceptions never collected. Each check now ends on any error: it is logged as check_failed, counted as not reached, and the call is due again when its claim's lease ends. A pass returns only once all of its checks have. - ModalCallInspector treated only five Modal exception types as problems reaching Modal; any other error, including a permission or version error, a closed client or a grpclib stream or status error from the client itself, failed the runs as worker_error while their call might still be running. Those are now retried later too. Errors the worker raised still fail the runs. - Comparison steps failed as abandoned are logged as report_run_comparison_dispatches_abandoned, apart from simulation runs. Co-Authored-By: Claude Opus 5.5 --- .../modal_calls.py | 10 ++++- .../run_recording.py | 40 +++++++++--------- .../tests/test_run_dispatch.py | 17 +++++++- .../tests/test_run_recording.py | 41 +++++++++++++++++++ 4 files changed, 87 insertions(+), 21 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py index e5eb7154f..a4ec22e3d 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py @@ -14,6 +14,7 @@ from dataclasses import dataclass from typing import Any, Protocol +import grpclib.exceptions import modal from policyengine_simulation_persistence import RunError, RunErrorCode @@ -117,13 +118,20 @@ async def cancel(self, modal_call_id: str) -> None: await modal.FunctionCall.from_id(modal_call_id).cancel.aio() -# Problems reaching Modal: the call's state is unknown, so ask again later. +# Problems between this client and Modal: the call's state is unknown, so ask +# again later. Errors the worker raised come back as their own types (or +# RemoteError, InternalFailure, FunctionTimeoutError) and fail the run. TRANSIENT_MODAL_ERRORS: tuple[type[Exception], ...] = ( modal.exception.ConnectionError, modal.exception.AuthError, + modal.exception.PermissionDeniedError, modal.exception.ServiceError, modal.exception.ResourceExhaustedError, modal.exception.InternalError, + modal.exception.VersionError, + modal.exception.ClientClosed, + grpclib.exceptions.GRPCError, + grpclib.exceptions.StreamTerminatedError, ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index b910a81e7..25df018b9 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -34,7 +34,6 @@ from policyengine_simulation_entry.modal_calls import ( CallFailed, - CallInspectionError, CallRunning, ) @@ -117,7 +116,11 @@ def write( @dataclass(frozen=True) class ResultRecorders: - """The result recorders by kind of run; a kind with none fails its runs.""" + """The result recorders by kind of run. + + A finished call of a kind with no recorder here is checked again later, + so that a revision which has one records it. + """ simulations: Mapping[SimulationKind, SimulationResultRecorder] = field( default_factory=dict @@ -292,18 +295,15 @@ async def _fail_stopped_work(self, recording: RecordingPass) -> None: recording.abandoned = len(abandoned.simulation_run_ids) + len( abandoned.report_run_ids ) - if recording.abandoned: - logger.warning( - "simulation_run_dispatches_abandoned", - extra={ - "simulation_run_ids": [ - str(run_id) for run_id in abandoned.simulation_run_ids - ], - "report_run_ids": [ - str(report_run_id) for report_run_id in abandoned.report_run_ids - ], - }, - ) + for event_prefix, field_name, ids in ( + ("simulation_run", "simulation_run_ids", abandoned.simulation_run_ids), + ("report_run_comparison", "report_run_ids", abandoned.report_run_ids), + ): + if ids: + logger.warning( + _event_name(event_prefix, "dispatches_abandoned"), + extra={field_name: [str(identifier) for identifier in ids]}, + ) without_runs = await asyncio.to_thread( self.transitions.fail_comparisons_of_failed_runs ) @@ -389,15 +389,17 @@ async def _check( recording.checked += 1 try: outcome = await self.inspector.inspect(modal_call_id) - except CallInspectionError as error: - # The claim's lease makes the call due again; nothing to record. + await self._record(outcome, handlers, recording) + # Modal could not be reached, or a database write failed (for example + # a dropped connection). Each check ends here, so a pass returns only + # once all of its checks have, and its concurrency limit holds. The + # claim's lease makes the call due again; nothing was recorded. + except Exception as error: # noqa: BLE001 # retried after the lease recording.not_reached += 1 logger.warning( _event_name(handlers.event_prefix, "check_failed"), - extra={**handlers.subject, "error_type": str(error)}, + extra={**handlers.subject, "error_type": type(error).__name__}, ) - return - await self._record(outcome, handlers, recording) async def _record( self, outcome: CallOutcome, handlers: _CallHandlers, recording: RecordingPass diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index 4c01128bf..3d6427e0a 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -4,6 +4,8 @@ import asyncio +import grpclib.const +import grpclib.exceptions import modal import pytest from policyengine_simulation_entry import modal_calls @@ -266,8 +268,21 @@ def test_a_failed_call_records_only_the_failure_type(monkeypatch, raised, code) [ modal.exception.ConnectionError("unreachable"), modal.exception.ServiceError("unavailable"), + modal.exception.PermissionDeniedError("token revoked"), + modal.exception.VersionError("client too old"), + modal.exception.ClientClosed("closed"), + grpclib.exceptions.StreamTerminatedError("stream reset"), + grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), + ], + ids=[ + "connection", + "service", + "permission", + "version", + "client-closed", + "stream-terminated", + "grpc", ], - ids=["connection", "service"], ) def test_a_modal_problem_is_retried_later(monkeypatch, raised) -> None: with pytest.raises(CallInspectionError): diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index 8d45e2962..aa5ae740c 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -309,6 +309,47 @@ def test_the_loop_leaves_database_connections_for_requests() -> None: assert run_recording.INSPECTION_CONCURRENCY < POOL_SIZE + MAX_OVERFLOW +def test_a_database_error_in_one_check_ends_with_its_pass(caplog) -> None: + """Every check finishes before the pass returns, so the limit holds.""" + recorder, transitions, inspector, clock = _recorder() + (broken_run,) = transitions.running_call("fc-broken") + (failed_run,) = transitions.running_call("fc-failed") + inspector.outcomes["fc-failed"] = CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + ) + keep_running = transitions.record_simulation_call_running + + def drop_one_connection(modal_call_id, *, next_check_at): + if modal_call_id == "fc-broken": + msg = "connection lost" + raise ConnectionError(msg) + return keep_running(modal_call_id, next_check_at=next_check_at) + + transitions.record_simulation_call_running = drop_one_connection # type: ignore[method-assign] + + recording = _run(recorder) + + assert recording.not_reached == 1 + assert recording.failed == 1 + assert transitions.runs[failed_run].status is RunStatus.FAILED + assert transitions.runs[broken_run].next_check_at == clock() + CHECK_LEASE + assert "simulation_run_check_failed" in [ + record.message for record in caplog.records + ] + + +def test_abandoned_comparison_steps_are_logged_as_comparison_events(caplog) -> None: + recorder, transitions, _, clock = _recorder() + transitions.add_comparison(status=RunStatus.DISPATCHING, updated_at=clock()) + clock.advance(timedelta(minutes=10, seconds=1)) + + _run(recorder) + + assert [record.message for record in caplog.records] == [ + "report_run_comparison_dispatches_abandoned" + ] + + class _ScriptedRecorder: """Return scripted passes, or raise, and count how often it ran.""" From 1e1aba3ea41942288f8565fc1332dece4f059d6f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:00:43 +0400 Subject: [PATCH 13/39] Index pending comparison steps for the recording pass Each recording pass of each instance looks for pending comparison steps whose simulation run failed, and report_runs had partial indexes only for running and dispatching steps, so that query scanned the whole table. Revision 0003 now adds a partial index on pending comparison steps. The dispatch_failed code's comment also covers a call cancelled because its ID could not be recorded. Co-Authored-By: Claude Opus 5.5 --- .../run_transitions.py | 3 ++- .../src/policyengine_simulation_persistence/runs.py | 6 ++++++ .../versions/0003_simulation_runs_and_report_runs.py | 8 ++++++++ 3 files changed, 16 insertions(+), 1 deletion(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index ac4c326f1..8d442eeed 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -57,7 +57,8 @@ class RunErrorCode(StrEnum): """Why a run or comparison step failed.""" - # Modal refused the call or never acknowledged it. + # Modal refused the call or never acknowledged it, or the acknowledged + # call could not be recorded and was cancelled. DISPATCH_FAILED = "dispatch_failed" # The instance stopped between claiming the dispatch and recording the call. DISPATCH_ABANDONED = "dispatch_abandoned" diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py index d69b42012..4f886d3e1 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py @@ -339,6 +339,12 @@ def _execution_columns(prefix: str, *, required: bool) -> list[Column[Any]]: report_runs.c.comparison_next_check_at, postgresql_where=report_runs.c.comparison_status == RunStatus.RUNNING.value, ) +# Each recording pass looks for pending comparison steps whose run failed. +Index( + "ix_report_runs_pending_comparisons", + report_runs.c.id, + postgresql_where=report_runs.c.comparison_status == RunStatus.PENDING.value, +) Index( "ix_report_runs_comparisons_dispatching", report_runs.c.updated_at, diff --git a/migrations/versions/0003_simulation_runs_and_report_runs.py b/migrations/versions/0003_simulation_runs_and_report_runs.py index 1225bb99f..f7a2b5d43 100644 --- a/migrations/versions/0003_simulation_runs_and_report_runs.py +++ b/migrations/versions/0003_simulation_runs_and_report_runs.py @@ -287,6 +287,13 @@ def _create_report_runs() -> None: schema=SCHEMA, postgresql_where=sa.text("comparison_status = 'running'"), ) + op.create_index( + "ix_report_runs_pending_comparisons", + "report_runs", + ["id"], + schema=SCHEMA, + postgresql_where=sa.text("comparison_status = 'pending'"), + ) op.create_index( "ix_report_runs_comparisons_dispatching", "report_runs", @@ -305,6 +312,7 @@ def upgrade() -> None: def downgrade() -> None: """Drop the tables ``upgrade`` created, with their indexes and keys.""" for index in ( + "ix_report_runs_pending_comparisons", "ix_report_runs_comparisons_dispatching", "ix_report_runs_due_comparison_checks", "ix_report_runs_reform_simulation_run_id", From 3fd4dc7e9275f40b235b80722756d298d96c7789 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:40:08 +0400 Subject: [PATCH 14/39] Read a call's outcome by where Modal raised it, not by the exception's type Modal re-raises in this process whatever the worker raised, so a worker's exception can share a type with errors from the entrypoint's own request to Modal. A worker that raised the built-in TimeoutError read as still running; one that raised a Modal or grpclib client error (from its own Modal calls) read as a problem reaching Modal and was asked about again until Modal dropped the output; and a worker's SystemExit or other BaseException escaped the inspector entirely. ModalCallInspector now decides by origin: an exception raised from the finished call's result (Modal's result processing) is the call's outcome, whatever its type, so it fails the run (a call timeout as worker_timeout, anything else as worker_error with its type only). Any other exception came from this process's own request: no output yet still reads as running, including while Modal re-runs a pre-empted call; an expired output or an unknown call fails the run; every other request error is asked about again later. A signal to this process itself (cancellation, KeyboardInterrupt, SystemExit, GeneratorExit) propagates. The tests raise worker exceptions through Modal's own result processing, as a real finished call does. Co-Authored-By: Claude Opus 5.5 --- .../modal_calls.py | 113 ++++++++++------ .../tests/test_run_dispatch.py | 121 +++++++++++++++++- 2 files changed, 188 insertions(+), 46 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py index a4ec22e3d..7e546344e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py @@ -1,9 +1,10 @@ """Start Modal calls and ask about them without waiting, for simulation runs. The Simulation Entrypoint starts each run's Modal call with ``spawn`` and later -asks Modal about it with ``get(timeout=0)``, which returns at once. A call -Modal pre-empted is re-run on the same input under the same call ID, so it -reads as still running and finishes later: pre-emption is never a failure. +asks Modal about it with ``get(timeout=0)``, which returns at once. Modal stops +a pre-empted worker with KeyboardInterrupt, records no result for it and +re-runs the same input under the same call ID, so the call reads as still +running and finishes later: pre-emption is never a failure. These adapters belong to the main path. The Stage 12 comparison path keeps its own Modal code; the two share nothing. @@ -11,10 +12,11 @@ from __future__ import annotations +import asyncio +import traceback from dataclasses import dataclass from typing import Any, Protocol -import grpclib.exceptions import modal from policyengine_simulation_persistence import RunError, RunErrorCode @@ -118,23 +120,30 @@ async def cancel(self, modal_call_id: str) -> None: await modal.FunctionCall.from_id(modal_call_id).cancel.aio() -# Problems between this client and Modal: the call's state is unknown, so ask -# again later. Errors the worker raised come back as their own types (or -# RemoteError, InternalFailure, FunctionTimeoutError) and fail the run. -TRANSIENT_MODAL_ERRORS: tuple[type[Exception], ...] = ( - modal.exception.ConnectionError, - modal.exception.AuthError, - modal.exception.PermissionDeniedError, - modal.exception.ServiceError, - modal.exception.ResourceExhaustedError, - modal.exception.InternalError, - modal.exception.VersionError, - modal.exception.ClientClosed, - grpclib.exceptions.GRPCError, - grpclib.exceptions.StreamTerminatedError, +# Where Modal raises a finished call's own outcome in this process: the +# worker's exception (whatever its type), the call's timeout, or a failure +# Modal records for the call. Anything raised elsewhere came from this +# process's own request to Modal. +_CALL_RESULT_MODULE = "modal._utils.function_utils" +_CALL_RESULT_FUNCTION = "_process_result" +# Signals to this process itself, which must never be taken for an outcome. +OWN_PROCESS_SIGNALS: tuple[type[BaseException], ...] = ( + asyncio.CancelledError, + KeyboardInterrupt, + SystemExit, + GeneratorExit, ) +def raised_from_call_result(error: BaseException) -> bool: + """Return whether Modal raised ``error`` from a finished call's result.""" + return any( + frame.f_globals.get("__name__") == _CALL_RESULT_MODULE + and frame.f_code.co_name == _CALL_RESULT_FUNCTION + for frame, _ in traceback.walk_tb(error.__traceback__) + ) + + def _failed(code: RunErrorCode, message: str) -> CallFailed: return CallFailed(RunError(code, message)) @@ -143,28 +152,54 @@ class ModalCallInspector: """Read a Modal call's output if it has one, without waiting.""" async def inspect(self, modal_call_id: str) -> CallOutcome: - """Return the call's outcome, or raise CallInspectionError.""" + """Return the call's outcome, or raise CallInspectionError. + + A worker's exception is re-raised here by Modal and can be of any + type, including one Modal's client also raises (a built-in + TimeoutError, a Modal or grpclib error) or a BaseException such as + SystemExit. So the outcome is decided by where the exception was + raised: from the finished call's result, it is the call's outcome; + otherwise it came from this process's own request to Modal, and a + signal to this process propagates. + """ try: call = modal.FunctionCall.from_id(modal_call_id) output = await call.get.aio(timeout=0) - # Modal raises the built-in TimeoutError while the call has no output. - except TimeoutError: - return CallRunning() - except modal.exception.OutputExpiredError: - return _failed( - RunErrorCode.RESULT_EXPIRED, "Modal no longer holds the call's output." - ) - except modal.exception.FunctionTimeoutError: - return _failed(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") - except modal.exception.NotFoundError: - return _failed(RunErrorCode.CALL_NOT_FOUND, "Modal does not know the call.") - except TRANSIENT_MODAL_ERRORS as error: - raise CallInspectionError(type(error).__name__) from error - # The call re-raises whatever the worker raised. Record only its type: - # the message may contain calculation inputs. - except Exception as error: # noqa: BLE001 # every worker failure fails the run - return _failed( - RunErrorCode.WORKER_ERROR, - f"The Modal call raised {type(error).__name__}.", - ) + except BaseException as error: + if raised_from_call_result(error): + return _call_failure(error) + if isinstance(error, OWN_PROCESS_SIGNALS) or not isinstance( + error, Exception + ): + raise + return _request_outcome(error) return CallSucceeded(output) + + +def _call_failure(error: BaseException) -> CallFailed: + """Turn a finished call's own outcome into a run failure.""" + if isinstance(error, modal.exception.FunctionTimeoutError): + return _failed(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + # Record only the type: the message may contain calculation inputs. + return _failed( + RunErrorCode.WORKER_ERROR, f"The Modal call raised {type(error).__name__}." + ) + + +def _request_outcome(error: Exception) -> CallOutcome: + """Read an error from this process's own request to Modal. + + Raise CallInspectionError for a problem reaching Modal, so the call is + asked about again later. + """ + # Modal raises the built-in TimeoutError while the call has no output, + # including while it re-runs a pre-empted call. + if type(error) is TimeoutError: + return CallRunning() + if isinstance(error, modal.exception.OutputExpiredError): + return _failed( + RunErrorCode.RESULT_EXPIRED, "Modal no longer holds the call's output." + ) + if isinstance(error, modal.exception.NotFoundError): + return _failed(RunErrorCode.CALL_NOT_FOUND, "Modal does not know the call.") + raise CallInspectionError(type(error).__name__) from error diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index 3d6427e0a..1bb3ee01f 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -3,11 +3,15 @@ from __future__ import annotations import asyncio +import builtins import grpclib.const import grpclib.exceptions import modal import pytest +from modal._serialization import serialize +from modal._utils.function_utils import _process_result +from modal_proto import api_pb2 from policyengine_simulation_entry import modal_calls from policyengine_simulation_entry.modal_calls import ( CallFailed, @@ -22,7 +26,7 @@ DispatchOutcome, RunDispatchService, ) -from policyengine_simulation_persistence import RunErrorCode, RunStatus +from policyengine_simulation_persistence import RunError, RunErrorCode, RunStatus from run_fakes import ( SAMPLE_REQUEST, FakeDispatcher, @@ -247,20 +251,26 @@ def test_a_finished_call_returns_its_output(monkeypatch) -> None: @pytest.mark.parametrize( ("raised", "code"), [ - (modal.exception.FunctionTimeoutError("slow"), RunErrorCode.WORKER_TIMEOUT), (modal.exception.OutputExpiredError(), RunErrorCode.RESULT_EXPIRED), (modal.exception.NotFoundError("gone"), RunErrorCode.CALL_NOT_FOUND), - (KeyError("household secret"), RunErrorCode.WORKER_ERROR), - (modal.exception.InternalFailure("lost"), RunErrorCode.WORKER_ERROR), ], - ids=["timeout", "expired", "not-found", "worker-error", "internal-failure"], + ids=["expired", "not-found"], ) -def test_a_failed_call_records_only_the_failure_type(monkeypatch, raised, code) -> None: +def test_modal_answers_about_the_call_itself_fail_the_run( + monkeypatch, raised, code +) -> None: outcome = _inspect(monkeypatch, raised) assert isinstance(outcome, CallFailed) assert outcome.error.code is code - assert "secret" not in outcome.error.message + + +def test_a_failed_call_records_only_the_failure_type(monkeypatch) -> None: + outcome = _inspect_finished(monkeypatch, KeyError("household secret")) + + assert outcome == CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + ) @pytest.mark.parametrize( @@ -273,6 +283,7 @@ def test_a_failed_call_records_only_the_failure_type(monkeypatch, raised, code) modal.exception.ClientClosed("closed"), grpclib.exceptions.StreamTerminatedError("stream reset"), grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), + RuntimeError("an unexpected client problem"), ], ids=[ "connection", @@ -282,6 +293,7 @@ def test_a_failed_call_records_only_the_failure_type(monkeypatch, raised, code) "client-closed", "stream-terminated", "grpc", + "other-client-error", ], ) def test_a_modal_problem_is_retried_later(monkeypatch, raised) -> None: @@ -289,6 +301,101 @@ def test_a_modal_problem_is_retried_later(monkeypatch, raised) -> None: _inspect(monkeypatch, raised) +class _FinishedCall: + """A call whose output Modal holds: the worker's exception, or a timeout. + + Reading it goes through Modal's own result processing, which raises the + worker's exception in this process, as a real finished call does. + """ + + def __init__(self, worker_exception, status=None): + self.result = api_pb2.GenericResult( + status=status or api_pb2.GenericResult.GENERIC_STATUS_FAILURE, + data=serialize(worker_exception), + exception=repr(worker_exception), + ) + self.get = self + + async def aio(self, *, timeout): + assert timeout == 0 + return await _process_result( + self.result, api_pb2.DATA_FORMAT_PICKLE, None, None + ) + + +def _inspect_finished(monkeypatch, worker_exception, status=None): + monkeypatch.setattr( + modal_calls.modal.FunctionCall, + "from_id", + lambda call_id: _FinishedCall(worker_exception, status), + ) + return asyncio.run(ModalCallInspector().inspect("fc-test")) + + +@pytest.mark.parametrize( + "worker_exception", + [ + SystemExit(2), + builtins.BaseExceptionGroup("worker", [ValueError("one")]), + KeyboardInterrupt(), + TimeoutError("a child call timed out"), + modal.exception.ConnectionError("the worker's own Modal call failed"), + grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), + ], + ids=[ + "system-exit", + "exception-group", + "keyboard-interrupt", + "timeout", + "modal-connection", + "grpc", + ], +) +def test_whatever_the_worker_raised_fails_the_run( + monkeypatch, worker_exception +) -> None: + """The worker's own exception, even of a type the client also raises, is final.""" + outcome = _inspect_finished(monkeypatch, worker_exception) + + assert isinstance(outcome, CallFailed) + assert outcome.error.code is RunErrorCode.WORKER_ERROR + assert type(worker_exception).__name__ in outcome.error.message + + +def test_a_call_modal_failed_internally_fails_the_run(monkeypatch) -> None: + outcome = _inspect_finished( + monkeypatch, + RuntimeError("ignored"), + status=api_pb2.GenericResult.GENERIC_STATUS_INTERNAL_FAILURE, + ) + + assert isinstance(outcome, CallFailed) + assert outcome.error.code is RunErrorCode.WORKER_ERROR + + +def test_a_call_modal_timed_out_fails_as_a_timeout(monkeypatch) -> None: + outcome = _inspect_finished( + monkeypatch, + RuntimeError("ignored"), + status=api_pb2.GenericResult.GENERIC_STATUS_TIMEOUT, + ) + + assert outcome == CallFailed( + RunError(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + ) + + +@pytest.mark.parametrize( + "signal", + [asyncio.CancelledError(), KeyboardInterrupt(), SystemExit(0)], + ids=["cancelled", "keyboard-interrupt", "system-exit"], +) +def test_a_signal_to_this_process_is_not_swallowed(monkeypatch, signal) -> None: + """Only an exception that came from the call's result is the worker's.""" + with pytest.raises(type(signal)): + _inspect(monkeypatch, signal) + + def test_the_dispatcher_passes_observability_as_its_own_argument(monkeypatch) -> None: spawned = {} From 2fc10abd59a3ed4e30b2080c172e614ed9fd2862 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:40:08 +0400 Subject: [PATCH 15/39] Keep the recording loop running whatever a check or pass raises The checks and the loop caught only Exception. Anything else (a BaseExceptionGroup, or another BaseException that is not a signal) ended a pass with its other checks still running, or ended the loop task with no log. Each check and each pass now catch BaseException: the error is logged, the check counts as not reached and its call is due again after the claim's lease, and the loop goes on. Only a signal to this process (cancellation at shutdown, KeyboardInterrupt, SystemExit, GeneratorExit) passes through, so a Cloud Run shutdown still stops the loop. Co-Authored-By: Claude Opus 5.5 --- .../run_recording.py | 49 ++++++++------- .../tests/test_run_recording.py | 63 +++++++++++++++++-- 2 files changed, 87 insertions(+), 25 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 25df018b9..c71df51a6 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -33,6 +33,7 @@ from policyengine_simulation_persistence.run_store import utc_now from policyengine_simulation_entry.modal_calls import ( + OWN_PROCESS_SIGNALS, CallFailed, CallRunning, ) @@ -391,10 +392,13 @@ async def _check( outcome = await self.inspector.inspect(modal_call_id) await self._record(outcome, handlers, recording) # Modal could not be reached, or a database write failed (for example - # a dropped connection). Each check ends here, so a pass returns only - # once all of its checks have, and its concurrency limit holds. The - # claim's lease makes the call due again; nothing was recorded. - except Exception as error: # noqa: BLE001 # retried after the lease + # a dropped connection). Each check ends here, whatever was raised + # (only a signal to this process passes), so a pass returns only once + # all of its checks have, and its concurrency limit holds. The claim's + # lease makes the call due again; nothing was recorded. + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise recording.not_reached += 1 logger.warning( _event_name(handlers.event_prefix, "check_failed"), @@ -482,24 +486,27 @@ async def run(self) -> None: """Record outcomes until cancelled; wait less while there is work.""" while True: await self._sleep(IDLE_PASS_INTERVAL_SECONDS) - try: - recording = await self.recorder.record_due_outcomes() - except Exception as error: - logger.exception( - "simulation_run_recording_pass_failed", - extra={"error_type": type(error).__name__}, - ) - continue - while recording.found_work: + recording = await self._pass() + while recording is not None and recording.found_work: await self._sleep(BUSY_PASS_INTERVAL_SECONDS) - try: - recording = await self.recorder.record_due_outcomes() - except Exception as error: - logger.exception( - "simulation_run_recording_pass_failed", - extra={"error_type": type(error).__name__}, - ) - break + recording = await self._pass() + + async def _pass(self) -> RecordingPass | None: + """Run one pass; log and return None if it raised anything. + + Only a signal to this process, such as the cancellation that stops + the loop at shutdown, passes through. + """ + try: + return await self.recorder.record_due_outcomes() + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.exception( + "simulation_run_recording_pass_failed", + extra={"error_type": type(error).__name__}, + ) + return None def start(self) -> None: """Start the loop as a background task of the running event loop.""" diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index aa5ae740c..e41bec099 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import builtins from datetime import timedelta import pytest @@ -350,6 +351,42 @@ def test_abandoned_comparison_steps_are_logged_as_comparison_events(caplog) -> N ] +class _NotAnException(BaseException): + """An error that derives from BaseException only, like a worker's exit.""" + + +def test_a_non_exception_error_in_one_check_ends_with_its_pass(caplog) -> None: + recorder, transitions, _, clock = _recorder() + (run_id,) = transitions.running_call("fc-broken") + + def raise_a_group(modal_call_id, *, next_check_at): + message = "writes" + raise builtins.BaseExceptionGroup(message, [_NotAnException()]) + + transitions.record_simulation_call_running = raise_a_group # type: ignore[method-assign] + + recording = _run(recorder) + + assert recording.not_reached == 1 + assert transitions.runs[run_id].next_check_at == clock() + CHECK_LEASE + assert "simulation_run_check_failed" in [ + record.message for record in caplog.records + ] + + +def test_a_signal_to_this_process_stops_a_check() -> None: + recorder, transitions, _, _ = _recorder() + transitions.running_call("fc-1") + + def interrupted(modal_call_id, *, next_check_at): + raise KeyboardInterrupt + + transitions.record_simulation_call_running = interrupted # type: ignore[method-assign] + + with pytest.raises(KeyboardInterrupt): + _run(recorder) + + class _ScriptedRecorder: """Return scripted passes, or raise, and count how often it ran.""" @@ -360,20 +397,25 @@ def __init__(self, passes) -> None: async def record_due_outcomes(self) -> RecordingPass: self.calls += 1 outcome = self.passes.pop(0) if self.passes else RecordingPass() - if isinstance(outcome, Exception): + if isinstance(outcome, BaseException): raise outcome return outcome def test_the_loop_waits_less_while_there_is_work_and_survives_failures() -> None: scripted = _ScriptedRecorder( - [RecordingPass(checked=3), RecordingPass(), RuntimeError("database down")] + [ + RecordingPass(checked=3), + RecordingPass(), + RuntimeError("database down"), + builtins.BaseExceptionGroup("pass", [_NotAnException()]), + ] ) waits = [] async def sleep(seconds: float) -> None: waits.append(seconds) - if len(waits) == 5: + if len(waits) == 6: raise asyncio.CancelledError loop = RunRecordingLoop(scripted, sleep=sleep) # type: ignore[arg-type] @@ -387,8 +429,21 @@ async def sleep(seconds: float) -> None: IDLE_PASS_INTERVAL_SECONDS, IDLE_PASS_INTERVAL_SECONDS, IDLE_PASS_INTERVAL_SECONDS, + IDLE_PASS_INTERVAL_SECONDS, ] - assert scripted.calls == 4 + assert scripted.calls == 5 + + +def test_a_signal_to_this_process_stops_the_loop() -> None: + scripted = _ScriptedRecorder([KeyboardInterrupt()]) + + async def sleep(_seconds: float) -> None: + return None + + loop = RunRecordingLoop(scripted, sleep=sleep) # type: ignore[arg-type] + + with pytest.raises(KeyboardInterrupt): + asyncio.run(loop.run()) def test_the_loop_starts_and_stops_as_a_background_task() -> None: From 729b1c9c9b7b85031cc559d31d0459b63ad2fb86 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:28:06 +0400 Subject: [PATCH 16/39] Retry a failed output download instead of failing the run From the review of the rebased 1.06: Modal stores any spawned call's output over 8 KiB as a blob, so its result processing (_process_result) first downloads nearly every real result with this process's own requests (a BlobGet request, then an HTTP download). The inspector took any exception with that function in its traceback for the call's outcome, so a dropped connection or a failed download during that step failed a call that had succeeded, and its output was never recorded. Now an exception is the call's outcome only when that function raised it itself, or re-raised the worker's exception (whose own frames follow it in the traceback). When the frame it called is Modal's blob download, the error is this process's own: the call is asked about again later. Tests read a blob-stored output, fail its request and its download, and re-raise a worker exception whose own frames include a blob download. Co-Authored-By: Claude Opus 5.5 --- .../modal_calls.py | 36 ++++-- .../tests/test_run_dispatch.py | 121 +++++++++++++++++- 2 files changed, 147 insertions(+), 10 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py index 7e546344e..96ec899bf 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_calls.py @@ -121,11 +121,16 @@ async def cancel(self, modal_call_id: str) -> None: # Where Modal raises a finished call's own outcome in this process: the -# worker's exception (whatever its type), the call's timeout, or a failure -# Modal records for the call. Anything raised elsewhere came from this -# process's own request to Modal. +# worker's exception (whatever its type, followed in the traceback by the +# worker's own frames), the call's timeout, a failure Modal records for the +# call, or an output that cannot be decoded. Anything raised elsewhere came +# from this process's own request to Modal. _CALL_RESULT_MODULE = "modal._utils.function_utils" _CALL_RESULT_FUNCTION = "_process_result" +# Before raising anything, that function fetches an output Modal stored as a +# blob (any spawned call's output over 8 KiB) with this process's own +# requests; an error raised in them is not the call's outcome. +_OUTPUT_DOWNLOAD_MODULE = "modal._utils.blob_utils" # Signals to this process itself, which must never be taken for an outcome. OWN_PROCESS_SIGNALS: tuple[type[BaseException], ...] = ( asyncio.CancelledError, @@ -136,12 +141,25 @@ async def cancel(self, modal_call_id: str) -> None: def raised_from_call_result(error: BaseException) -> bool: - """Return whether Modal raised ``error`` from a finished call's result.""" - return any( - frame.f_globals.get("__name__") == _CALL_RESULT_MODULE - and frame.f_code.co_name == _CALL_RESULT_FUNCTION - for frame, _ in traceback.walk_tb(error.__traceback__) - ) + """Return whether Modal raised ``error`` as a finished call's outcome. + + That is so when Modal's result processing raised it, or re-raised the + worker's exception, but not when the error came from the output download + that processing makes first. + """ + frames = [frame for frame, _ in traceback.walk_tb(error.__traceback__)] + for position, frame in enumerate(frames): + if ( + frame.f_globals.get("__name__") == _CALL_RESULT_MODULE + and frame.f_code.co_name == _CALL_RESULT_FUNCTION + ): + # The frame it called, if any: a download, or the worker's own. + callee = frames[position + 1] if position + 1 < len(frames) else None + return ( + callee is None + or callee.f_globals.get("__name__") != _OUTPUT_DOWNLOAD_MODULE + ) + return False def _failed(code: RunErrorCode, message: str) -> CallFailed: diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index 1bb3ee01f..04c03117b 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -9,7 +9,8 @@ import grpclib.exceptions import modal import pytest -from modal._serialization import serialize +from modal._serialization import pickle_traceback, serialize +from modal._utils import blob_utils from modal._utils.function_utils import _process_result from modal_proto import api_pb2 from policyengine_simulation_entry import modal_calls @@ -385,6 +386,124 @@ def test_a_call_modal_timed_out_fails_as_a_timeout(monkeypatch) -> None: ) +BLOB_DOWNLOAD_URL = "https://blob.test/output" + + +class _BlobStub: + """Modal's API as its result processing uses it to fetch a stored output.""" + + def __init__(self, failure=None): + self.failure = failure + + async def BlobGet(self, request): # noqa: N802 # Modal's method name + assert request.blob_id == "bl-test" + if self.failure is not None: + raise self.failure + return api_pb2.BlobGetResponse(download_url=BLOB_DOWNLOAD_URL) + + +class _BlobStoredCall: + """A finished call whose output Modal stores as a blob. + + Modal does so for any spawned call's output over 8 KiB, so nearly every + real result is read this way: Modal's result processing first makes this + process's own requests for the blob, then decodes it. + """ + + def __init__(self, stub): + self.result = api_pb2.GenericResult( + status=api_pb2.GenericResult.GENERIC_STATUS_SUCCESS, + data_blob_id="bl-test", + ) + self.stub = stub + self.get = self + + async def aio(self, *, timeout): + assert timeout == 0 + return await _process_result( + self.result, api_pb2.DATA_FORMAT_PICKLE, self.stub, None + ) + + +def _inspect_blob_stored(monkeypatch, stub, download): + async def fake_download(download_url): + assert download_url == BLOB_DOWNLOAD_URL + if isinstance(download, BaseException): + raise download + return download + + monkeypatch.setattr(blob_utils, "_download_from_url", fake_download) + monkeypatch.setattr( + modal_calls.modal.FunctionCall, + "from_id", + lambda call_id: _BlobStoredCall(stub), + ) + return asyncio.run(ModalCallInspector().inspect("fc-test")) + + +def test_an_output_stored_as_a_blob_is_read(monkeypatch) -> None: + outcome = _inspect_blob_stored(monkeypatch, _BlobStub(), serialize({"budget": 1.0})) + + assert outcome == CallSucceeded({"budget": 1.0}) + + +@pytest.mark.parametrize( + ("blob_request_failure", "download_failure"), + [ + (modal.exception.ExecutionError("blob request failed"), None), + (grpclib.exceptions.GRPCError(grpclib.const.Status.UNAVAILABLE, "down"), None), + (modal.exception.ConnectionError("unreachable"), None), + (None, modal.exception.ExecutionError("Get from url failed with status 503")), + ], + ids=["blob-request-execution-error", "blob-request-grpc", "connection", "download"], +) +def test_a_failed_output_download_is_retried_later( + monkeypatch, blob_request_failure, download_failure +) -> None: + """Requests for a stored output are the entrypoint's own, not the worker's.""" + with pytest.raises(CallInspectionError): + _inspect_blob_stored( + monkeypatch, _BlobStub(blob_request_failure), download_failure + ) + + +def test_an_output_download_that_times_out_is_checked_again(monkeypatch) -> None: + outcome = _inspect_blob_stored(monkeypatch, _BlobStub(), TimeoutError()) + + assert outcome == CallRunning() + + +def _worker_exception_from_its_own_download(): + """Return what a worker raises when its own Modal blob download fails.""" + + async def worker(): + await blob_utils.blob_download( + "bl-test", _BlobStub(modal.exception.ConnectionError("worker's download")) + ) + + try: + asyncio.run(worker()) + except modal.exception.ConnectionError as error: + return error + raise AssertionError + + +def test_a_worker_exception_with_its_traceback_fails_the_run(monkeypatch) -> None: + """The worker's frames follow Modal's result processing, whatever they hold.""" + worker_exception = _worker_exception_from_its_own_download() + serialized_tb, tb_line_cache = pickle_traceback(worker_exception, "ta-test") + call = _FinishedCall(worker_exception) + call.result.serialized_tb = serialized_tb + call.result.tb_line_cache = tb_line_cache + monkeypatch.setattr(modal_calls.modal.FunctionCall, "from_id", lambda call_id: call) + + outcome = asyncio.run(ModalCallInspector().inspect("fc-test")) + + assert outcome == CallFailed( + RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised ConnectionError.") + ) + + @pytest.mark.parametrize( "signal", [asyncio.CancelledError(), KeyboardInterrupt(), SystemExit(0)], From 5be776944657750cc0e3a77842c92d09e6d618a6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:32:50 +0400 Subject: [PATCH 17/39] Keep a call whose record committed although its reply was lost From the review of the rebased 1.06: when recording an acknowledged call raised, the dispatcher assumed nothing had been written and cancelled the call. If the connection dropped after PostgreSQL committed the record, the runs were already running under the call, so a healthy call was cancelled and its runs later failed as if the worker had failed. Now the dispatcher first fails the runs if they are still dispatching, which only an uncommitted record leaves, and then cancels the call. Otherwise it reads whether the runs hold the call (RunTransitions.simulation_runs_hold_call and comparison_holds_call): if they do, the call started; if not (they were failed meanwhile as abandoned), the call is cancelled. When the database cannot answer either, the call is left running and the new outcome RECORD_UNKNOWN says so: the recording loop collects it if its record committed, and the runs are failed as abandoned otherwise. Co-Authored-By: Claude Opus 5.5 --- .../run_transitions.py | 37 +++++++- .../tests/test_run_transitions.py | 41 +++++++++ .../run_dispatch.py | 74 ++++++++++++---- .../tests/run_fakes.py | 14 +++ .../tests/test_run_dispatch.py | 86 ++++++++++++++++++- 5 files changed, 235 insertions(+), 17 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index 8d442eeed..26f30ce4e 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -29,7 +29,7 @@ from enum import StrEnum from typing import TYPE_CHECKING -from sqlalchemy import and_, exists, or_, select, update +from sqlalchemy import and_, exists, func, or_, select, update from policyengine_simulation_persistence.run_store import utc_now from policyengine_simulation_persistence.runs import ( @@ -188,6 +188,26 @@ def record_simulation_dispatch_failure( }, ) + def simulation_runs_hold_call( + self, run_ids: Collection[UUID], modal_call_id: str + ) -> bool: + """Return whether every one of the runs recorded ``modal_call_id``. + + A dispatcher whose record raised reads this to learn whether the + record committed all the same, its reply having been lost. + """ + wanted = set(run_ids) + with self.engine.connect() as connection: + holding = connection.execute( + select(func.count()) + .select_from(simulation_runs) + .where( + simulation_runs.c.id.in_(wanted), + simulation_runs.c.modal_call_id == modal_call_id, + ) + ).scalar_one() + return bool(wanted) and holding == len(wanted) + def claim_due_simulation_calls( self, *, limit: int, lease: timedelta ) -> list[DueSimulationCall]: @@ -415,6 +435,21 @@ def record_comparison_dispatch_failure( is not None ) + def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool: + """Return whether the comparison step recorded ``modal_call_id``. + + As ``simulation_runs_hold_call``, for a comparison step. + """ + with self.engine.connect() as connection: + return connection.execute( + select( + exists().where( + report_runs.c.id == report_run_id, + report_runs.c.comparison_modal_call_id == modal_call_id, + ) + ) + ).scalar_one() + def claim_due_comparisons( self, *, limit: int, lease: timedelta ) -> list[ReportRunRecord]: diff --git a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py index d08655317..ad1da794a 100644 --- a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py +++ b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py @@ -112,6 +112,25 @@ def test_a_dispatched_run_records_its_call_and_first_check( ) +def test_runs_say_whether_they_recorded_a_call( + migrated_database, public_api_rows +) -> None: + """A dispatcher whose record raised reads whether it committed all the same.""" + stores = clocked_stores(migrated_database) + running = _running_run(stores, migrated_database, public_api_rows) + dispatching = _pending_run(stores, migrated_database, public_api_rows) + stores.transitions.claim_simulation_dispatch([dispatching.id]) + transitions = stores.transitions + + assert transitions.simulation_runs_hold_call([running.id], "fc-economy") + assert not transitions.simulation_runs_hold_call([running.id], "fc-other") + assert not transitions.simulation_runs_hold_call([dispatching.id], "fc-economy") + assert not transitions.simulation_runs_hold_call( + [running.id, dispatching.id], "fc-economy" + ) + assert not transitions.simulation_runs_hold_call([], "fc-economy") + + def test_a_refused_dispatch_fails_the_run_and_frees_the_setup( migrated_database, public_api_rows ) -> None: @@ -411,6 +430,28 @@ def test_a_comparison_failure_is_recorded(migrated_database, public_api_rows) -> assert failed.comparison_completed_at == stores.clock() +def test_a_comparison_says_whether_it_recorded_a_call( + migrated_database, public_api_rows +) -> None: + stores = clocked_stores(migrated_database) + report_run, (baseline_run, reform_run) = _economy_report_run( + stores, migrated_database, public_api_rows + ) + _succeed(stores, baseline_run, "fc-baseline") + _succeed(stores, reform_run, "fc-reform") + transitions = stores.transitions + transitions.claim_comparison_dispatch(report_run.id) + + assert not transitions.comparison_holds_call(report_run.id, "fc-comparison") + + transitions.record_comparison_dispatched( + report_run.id, modal_call_id="fc-comparison", first_check_at=stores.clock() + ) + + assert transitions.comparison_holds_call(report_run.id, "fc-comparison") + assert not transitions.comparison_holds_call(report_run.id, "fc-other") + + def test_a_refused_comparison_dispatch_fails_the_step( migrated_database, public_api_rows ) -> None: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index fd5053501..75c2b5d20 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -53,8 +53,13 @@ class DispatchOutcome(StrEnum): # Modal refused the call; the runs are failed. DISPATCH_FAILED = "dispatch_failed" # Modal acknowledged the call, but the runs had stopped dispatching before - # the call ID could be recorded. + # the call ID could be recorded, or recording it failed; the call is + # cancelled. NOT_RECORDED = "not_recorded" + # Recording the call raised and the database could not say whether the + # record committed; the call is left running. If it did, the recording + # loop collects the call; if not, the runs are failed as abandoned. + RECORD_UNKNOWN = "record_unknown" class DispatchTransitions(Protocol): @@ -80,6 +85,12 @@ def record_simulation_dispatch_failure( """Fail dispatching runs.""" ... + def simulation_runs_hold_call( + self, run_ids: Collection[UUID], modal_call_id: str + ) -> bool: + """Return whether every one of the runs recorded the call.""" + ... + def claim_comparison_dispatch(self, report_run_id: UUID) -> bool: """Move a pending comparison step to dispatching.""" ... @@ -100,14 +111,19 @@ def record_comparison_dispatch_failure( """Fail a dispatching comparison step.""" ... + def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool: + """Return whether the comparison step recorded the call.""" + ... + @dataclass(frozen=True) class _DispatchSteps: - """The three writes around one spawn, bound to the runs they change.""" + """The writes around one spawn, and a read, bound to the runs they change.""" claim: Callable[[], bool] record: Callable[[str, datetime], bool] fail: Callable[[RunError], bool] + holds_call: Callable[[str], bool] # "simulation_run" or "report_run_comparison": the log event names' start. event_prefix: str @@ -140,6 +156,9 @@ async def start_simulation_runs( fail=lambda error: self.transitions.record_simulation_dispatch_failure( runs, error ), + holds_call=lambda call_id: self.transitions.simulation_runs_hold_call( + runs, call_id + ), event_prefix="simulation_run", ), request, @@ -167,6 +186,9 @@ async def start_comparison( fail=lambda error: self.transitions.record_comparison_dispatch_failure( report_run_id, error ), + holds_call=lambda call_id: self.transitions.comparison_holds_call( + report_run_id, call_id + ), event_prefix="report_run_comparison", ), request, @@ -208,22 +230,19 @@ async def _record( A call whose runs never record it would run to its end with nobody collecting its output, and its runs would later be failed as abandoned; a resubmission would then start a second call. So such a - call is cancelled, and when the record failed, its runs are failed - at once with dispatch_failed if the database allows. + call is cancelled. """ context = {**subject, "modal_call_id": modal_call_id} try: recorded = await asyncio.to_thread( steps.record, modal_call_id, self.clock() + FIRST_CHECK_DELAY ) - except Exception as error: # the call is cancelled instead + except Exception as error: # settled below logger.exception( - _event_name(steps.event_prefix, "call_not_recorded"), + _event_name(steps.event_prefix, "call_record_failed"), extra={**context, "error_type": type(error).__name__}, ) - await self._cancel(modal_call_id, context, steps.event_prefix) - await self._fail_quietly(steps, context) - return DispatchOutcome.NOT_RECORDED + return await self._settle_failed_record(steps, modal_call_id, context) if not recorded: # The runs stopped dispatching while Modal acknowledged the call, # for example because the recording loop failed them as abandoned. @@ -246,19 +265,44 @@ async def _cancel( extra={**context, "error_type": type(error).__name__}, ) - async def _fail_quietly( - self, steps: _DispatchSteps, context: dict[str, object] - ) -> None: + async def _settle_failed_record( + self, + steps: _DispatchSteps, + modal_call_id: str, + context: dict[str, object], + ) -> DispatchOutcome: + """Settle a call whose record raised: it may have committed all the same. + + A record that committed but whose reply was lost leaves the runs + running under the call, which must keep running. Only an uncommitted + record leaves them dispatching, so they are failed if they still are, + and then the call is cancelled. Otherwise the runs are read: holding + the call, it started; not holding it (failed meanwhile as abandoned), + the call is cancelled. If the database cannot answer, the call is left + running, since cancelling it could stop a recorded call. + """ try: - await asyncio.to_thread( + failed = await asyncio.to_thread( steps.fail, RunError( RunErrorCode.DISPATCH_FAILED, "The Modal call could not be recorded, so it was cancelled.", ), ) - except Exception as error: # the abandoned-dispatch check fails the runs later + holds_call = not failed and await asyncio.to_thread( + steps.holds_call, modal_call_id + ) + except Exception as error: # the call is left running; see above logger.exception( - _event_name(steps.event_prefix, "dispatch_failure_not_recorded"), + _event_name(steps.event_prefix, "call_record_unknown"), extra={**context, "error_type": type(error).__name__}, ) + return DispatchOutcome.RECORD_UNKNOWN + if holds_call: + logger.info(_event_name(steps.event_prefix, "dispatched"), extra=context) + return DispatchOutcome.STARTED + logger.warning( + _event_name(steps.event_prefix, "call_not_recorded"), extra=context + ) + await self._cancel(modal_call_id, context, steps.event_prefix) + return DispatchOutcome.NOT_RECORDED diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py index 799ed4f0e..232e0eb0c 100644 --- a/projects/policyengine-simulation-entry/tests/run_fakes.py +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -233,6 +233,20 @@ def record_comparison_dispatch_failure( is not None ) + def simulation_runs_hold_call( + self, run_ids: Collection[UUID], modal_call_id: str + ) -> bool: + """Return whether every run recorded ``modal_call_id``.""" + with self.lock: + return all( + self.runs[run_id].modal_call_id == modal_call_id for run_id in run_ids + ) + + def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool: + """Return whether the comparison step recorded ``modal_call_id``.""" + with self.lock: + return self.comparisons[report_run_id].modal_call_id == modal_call_id + # --- Recording ------------------------------------------------------- def claim_due_simulation_calls( diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index 04c03117b..f9ef02a7b 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -162,7 +162,6 @@ def unreachable_database(*_arguments, **_keywords): raise ConnectionError(msg) transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] - transitions.record_simulation_dispatch_failure = unreachable_database # type: ignore[method-assign] outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) @@ -170,6 +169,91 @@ def unreachable_database(*_arguments, **_keywords): assert "simulation_run_call_not_cancelled" in caplog.text +def _commit_then_lose_acknowledgement(record): + """Return ``record`` changed to commit and then raise, as a lost reply does.""" + + def committed_without_acknowledgement(*arguments, **keywords): + assert record(*arguments, **keywords) + msg = "connection lost after the commit" + raise ConnectionError(msg) + + return committed_without_acknowledgement + + +@pytest.mark.asyncio +async def test_an_unacknowledged_record_keeps_its_call() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + transitions.record_simulation_dispatched = _commit_then_lose_acknowledgement( # type: ignore[method-assign] + transitions.record_simulation_dispatched + ) + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + assert service.dispatcher.cancelled == [] + run = execution(transitions, run_id) + assert run.status is RunStatus.RUNNING + assert run.modal_call_id == "fc-1" + + +@pytest.mark.asyncio +async def test_an_unacknowledged_comparison_record_keeps_its_call() -> None: + service, transitions, _ = _service() + comparison = transitions.add_comparison() + transitions.record_comparison_dispatched = _commit_then_lose_acknowledgement( # type: ignore[method-assign] + transitions.record_comparison_dispatched + ) + + outcome = await service.start_comparison(comparison, SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.STARTED + assert service.dispatcher.cancelled == [] + assert transitions.comparisons[comparison].status is RunStatus.RUNNING + + +@pytest.mark.asyncio +async def test_a_failed_record_of_runs_failed_meanwhile_cancels_the_call() -> None: + service, transitions, _ = _service() + run_id = transitions.add_run() + + def abandoned_meanwhile(*_arguments, **_keywords): + transitions.runs[run_id].status = RunStatus.FAILED + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = abandoned_meanwhile # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.NOT_RECORDED + assert service.dispatcher.cancelled == ["fc-1"] + + +@pytest.mark.asyncio +async def test_a_call_whose_record_cannot_be_read_is_left_running(caplog) -> None: + """With the database unreachable, the record may have committed all the same.""" + service, transitions, _ = _service() + run_id = transitions.add_run() + + def unreachable_database(*_arguments, **_keywords): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.record_simulation_dispatched = unreachable_database # type: ignore[method-assign] + transitions.record_simulation_dispatch_failure = unreachable_database # type: ignore[method-assign] + + outcome = await service.start_simulation_runs([run_id], SAMPLE_REQUEST) + + assert outcome is DispatchOutcome.RECORD_UNKNOWN + assert service.dispatcher.cancelled == [] + assert any( + record.message == "simulation_run_call_record_unknown" + and getattr(record, "modal_call_id", None) == "fc-1" + for record in caplog.records + ) + + def test_the_dispatcher_cancels_a_call_by_its_id(monkeypatch) -> None: cancelled = [] From 90d44a4cc04236b6d62b711dc79920373d6f1ca3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:35:10 +0400 Subject: [PATCH 18/39] Pass on a shutdown that arrives while the recording loop stops From the review of the rebased 1.06: RunRecordingLoop.stop suppressed every CancelledError while waiting for the loop to end, including one cancelling the task that called stop, so a shutdown arriving then was swallowed. It now suppresses only the loop's own cancellation and re-raises the caller's. Two descriptions are corrected: the persistence README says simulations and reports (not every table listed above it) are definitions the runtime role cannot update, and the recording loop's docstring says a finished call without a recorder is checked again until Modal drops its output and the run fails as result_expired. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 4 +-- .../run_recording.py | 16 ++++++--- .../simulation_database.py | 7 ++-- .../tests/test_run_recording.py | 34 +++++++++++++++++++ 4 files changed, 52 insertions(+), 9 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 7e629f5e7..4a2ec59f0 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -17,8 +17,8 @@ writes them with DML only. | `simulation_runs` | One execution of a simulation for a catalog version and year: the region code and dataset row it resolved in that version (economy runs only), its status, Modal call ID, timings, error, trigger (`initial`, `manual` or `system`) and attempt number | | `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call | -Both are definitions: the runtime role may read and insert them but not update -or delete them. Rows migrated from the Public API's v1 database keep their v1 +`simulations` and `reports` are definitions: the runtime role may read and +insert them but not update or delete them. Rows migrated from the Public API's v1 database keep their v1 IDs in the unique `legacy_simulation_id` and `legacy_report_id` columns, and the stores' `create_migrated` methods return the existing row when the same v1 ID is migrated again. diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index c71df51a6..9dd41ab45 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -23,7 +23,6 @@ from __future__ import annotations import asyncio -import contextlib import logging from dataclasses import dataclass, field from datetime import timedelta @@ -514,10 +513,19 @@ def start(self) -> None: self._task = asyncio.create_task(self.run()) async def stop(self) -> None: - """Cancel the loop and wait for it to end.""" + """Cancel the loop and wait for it to end. + + A cancellation of the task calling this, such as a shutdown arriving + while the loop ends, is passed on rather than taken for the loop's. + """ if self._task is None: return self._task.cancel() - with contextlib.suppress(asyncio.CancelledError): + try: await self._task - self._task = None + except asyncio.CancelledError: + caller = asyncio.current_task() + if caller is not None and caller.cancelling(): + raise + finally: + self._task = None diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index 43f595bbb..28bffdac0 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -48,9 +48,10 @@ async def read_revision(self) -> str | None: def recording_loop(self) -> RunRecordingLoop: """Return the loop that records Modal call outcomes on runs. - No kind of run has a result recorder yet, so a finished call would - stay running until a revision with a recorder records it; no route - creates runs yet. + No kind of run has a result recorder yet, so a finished call is + checked again until a revision with a recorder records it, or until + Modal no longer holds its output and the run fails as result_expired; + no route creates runs yet. """ return RunRecordingLoop( RunRecorder( diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index e41bec099..c1654dcfc 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -4,6 +4,7 @@ import asyncio import builtins +import contextlib from datetime import timedelta import pytest @@ -461,6 +462,39 @@ async def lifetime() -> None: assert scripted.calls > 0 +class _SlowToStopRecorder: + """A pass that takes a while to end once cancelled, as a closing write does.""" + + def __init__(self) -> None: + self.passing = asyncio.Event() + + async def record_due_outcomes(self) -> RecordingPass: + self.passing.set() + try: + await asyncio.sleep(60) + finally: + await asyncio.shield(asyncio.sleep(0.05)) + return RecordingPass() + + +def test_stopping_passes_on_a_cancellation_of_its_caller() -> None: + """Stopping the loop never swallows a shutdown of the task that stops it.""" + + async def shutdown() -> bool: + recorder = _SlowToStopRecorder() + loop = RunRecordingLoop(recorder, sleep=lambda _seconds: asyncio.sleep(0)) # type: ignore[arg-type] + loop.start() + await recorder.passing.wait() + stopper = asyncio.create_task(loop.stop()) + await asyncio.sleep(0.01) + stopper.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stopper + return stopper.cancelled() + + assert asyncio.run(shutdown()) + + class _TrackedLoop: def __init__(self) -> None: self.started = False From 9d26a6b7ab401318d1e7198f755a885a12550228 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:57:37 +0400 Subject: [PATCH 19/39] Keep backing off for a call checked more than 46 times From the fifth review of 1.06: next_check_delay doubled the first wait before capping it at a minute, so from a call's 47th check (about 42 minutes after dispatch) the product overflowed timedelta and raised. The check was then logged as failed every minute, which would hide real Modal or database errors, and the claim's lease replaced the backoff. The exponent is now capped first; the test covers the 47th and the 10,000th check. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_entry/run_recording.py | 8 +++++++- .../tests/test_run_recording.py | 12 +++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 9dd41ab45..caeac59fe 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -78,9 +78,15 @@ def _event_name(prefix: str, outcome: str) -> str: IDLE_PASS_INTERVAL_SECONDS = 5.0 +# Past this many doublings the wait is LONGEST_BACKOFF anyway; capping the +# exponent keeps the product within timedelta's range for a long call. +_MOST_DOUBLINGS = 16 + + def next_check_delay(check_count: int) -> timedelta: """Return the wait before the next check of a call checked ``check_count`` times.""" - return min(FIRST_BACKOFF * 2 ** max(check_count - 1, 0), LONGEST_BACKOFF) + doublings = min(max(check_count - 1, 0), _MOST_DOUBLINGS) + return min(FIRST_BACKOFF * 2**doublings, LONGEST_BACKOFF) class SimulationResultRecorder(Protocol): diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index c1654dcfc..67c1cda8e 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -79,7 +79,17 @@ def _run(recorder) -> RecordingPass: @pytest.mark.parametrize( ("check_count", "seconds"), - [(0, 2), (1, 2), (2, 4), (3, 8), (5, 32), (6, 60), (40, 60)], + [ + (0, 2), + (1, 2), + (2, 4), + (3, 8), + (5, 32), + (6, 60), + (40, 60), + (47, 60), + (10_000, 60), + ], ) def test_checks_back_off_up_to_a_minute(check_count, seconds) -> None: assert next_check_delay(check_count) == timedelta(seconds=seconds) From 9953a94263213f641008dede723ac79b2f00fa12 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:58:42 +0400 Subject: [PATCH 20/39] Store a failed dispatch's error without claiming the call was cancelled From the fifth review of 1.06: when recording an acknowledged call failed, the runs were failed with "The Modal call could not be recorded, so it was cancelled." before the cancel was tried, so the message was false when the cancel then failed. It now says only that the call could not be recorded; whether the cancel worked is logged (call_not_cancelled on failure). Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_persistence/run_transitions.py | 2 +- .../src/policyengine_simulation_entry/run_dispatch.py | 2 +- .../policyengine-simulation-entry/tests/test_run_dispatch.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index 26f30ce4e..af74e5bce 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -58,7 +58,7 @@ class RunErrorCode(StrEnum): """Why a run or comparison step failed.""" # Modal refused the call or never acknowledged it, or the acknowledged - # call could not be recorded and was cancelled. + # call could not be recorded (the entrypoint then tries to cancel it). DISPATCH_FAILED = "dispatch_failed" # The instance stopped between claiming the dispatch and recording the call. DISPATCH_ABANDONED = "dispatch_abandoned" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index 75c2b5d20..b7730cdeb 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -286,7 +286,7 @@ async def _settle_failed_record( steps.fail, RunError( RunErrorCode.DISPATCH_FAILED, - "The Modal call could not be recorded, so it was cancelled.", + "The Modal call could not be recorded.", ), ) holds_call = not failed and await asyncio.to_thread( diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index f9ef02a7b..0309ace38 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -144,7 +144,7 @@ def unreachable_database(*_arguments, **_keywords): run = execution(transitions, run_id) assert run.status is RunStatus.FAILED assert run.error.code is RunErrorCode.DISPATCH_FAILED - assert "could not be recorded" in run.error.message + assert run.error.message == "The Modal call could not be recorded." assert any( getattr(record, "modal_call_id", None) == "fc-1" for record in caplog.records ) From 8e9bfa3cd291ba2b1cd26dfce8583d244b3cc2c0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:54:44 +0400 Subject: [PATCH 21/39] Stop recording outcomes in a background loop in each instance Decided 2026-10-27: Modal reports when a call ends, and a scheduled Cloud Run job is the safety check. The background loop is removed: Cloud Run keeps a minimum instance for every revision with a traffic tag, so a rejected or rolled-back production revision kept running it and recording production outcomes with its own code. The app's lifespan no longer starts or stops a loop, SimulationDatabase no longer builds one, and the deploy no longer passes --no-cpu-throttling: the web service records outcomes only within requests. RunRecorder, which runs one pass of due checks, stays for the safety check. Every deploy still mounts sim's own Modal token. Co-Authored-By: Claude Opus 5.5 --- .../deploy-cloud-run-simulation-entry.sh | 7 +- .../README.md | 2 +- .../run_transitions.py | 6 +- .../runs.py | 2 +- .../tests/test_recording_load.py | 4 +- .../DEPLOYMENT.md | 25 ++- .../src/policyengine_simulation_entry/app.py | 9 -- .../entry_runtime.py | 5 - .../run_dispatch.py | 8 +- .../run_recording.py | 90 ++--------- .../simulation_database.py | 24 --- ...ment.py => test_modal_token_deployment.py} | 14 +- .../tests/test_run_recording.py | 146 +----------------- .../tests/test_simulation_database.py | 3 - 14 files changed, 46 insertions(+), 299 deletions(-) rename projects/policyengine-simulation-entry/tests/{test_run_recording_deployment.py => test_modal_token_deployment.py} (84%) diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index 05836e82f..36b1899ff 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -145,9 +145,9 @@ add_runtime_secret OLD_GATEWAY_AUTH_CLIENT_SECRET \ # The runtime role's URL for the entrypoint's own simulation schema. The # migration job, which uses a separate migrator URL, runs before this deploy. add_runtime_secret SIMULATION_DATABASE_URL "${SIMULATION_DATABASE_URL_SECRET_NAME}" -# sim's own Modal token: the entrypoint starts runs' Modal calls and its -# background loop asks Modal about them, on every deploy. Until the Stage 12 -# comparison path is removed, its Modal calls use the same token. +# sim's own Modal token, on every deploy: the entrypoint starts runs' Modal +# calls and asks Modal about a call when Modal reports that it ended. Until +# the Stage 12 comparison path is removed, its Modal calls use the same token. add_runtime_secret MODAL_TOKEN_ID "${MODAL_TOKEN_ID_SECRET_NAME}" add_runtime_secret MODAL_TOKEN_SECRET "${MODAL_TOKEN_SECRET_SECRET_NAME}" @@ -214,7 +214,6 @@ gcloud run deploy "${ENTRYPOINT_SERVICE}" \ --env-vars-file "${runtime_environment_file}" \ --set-secrets "${runtime_secrets_csv}" \ --cpu 1 \ - --no-cpu-throttling \ --memory 512Mi \ --concurrency 80 \ --timeout 60 \ diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 4a2ec59f0..85b878aae 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -55,7 +55,7 @@ not delete them. `RunTransitions` makes every status change one conditional update on the expected status, so a second writer changes nothing: only the caller that moves runs from `pending` to `dispatching` starts their Modal call, the runs -one call computes change together, the recording loop claims a due call by +one call computes change together, a recording pass claims a due call by moving its next check time, results are written in the transaction that marks a run succeeded, and dispatches left `dispatching` (and comparison steps whose simulation run failed) are failed so the work can be submitted diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index af74e5bce..0d0a2e5aa 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -9,9 +9,9 @@ starts its Modal call, and it then records the call ID (``running``); - the runs that share one Modal call, such as a household report's baseline and reform, change status together, all or none; -- the recording loop claims a running call's next check by moving its - ``next_check_at`` forward, so instances do not check the same call at once, - and only one of them records its outcome; +- a recording pass claims a running call's next check by moving its + ``next_check_at`` forward, so two passes do not check the same call at + once, and only one of them records its outcome; - a call's results are written in the same transaction that marks its runs ``succeeded``, so a run never succeeds without its results; - a run left in ``dispatching`` (the instance stopped between claiming it and diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py index 4f886d3e1..3ad05ea99 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py @@ -247,7 +247,7 @@ def _execution_columns(prefix: str, *, required: bool) -> list[Column[Any]]: simulation_runs.c.modal_call_id, postgresql_where=simulation_runs.c.modal_call_id.is_not(None), ) -# The recording loop reads only running runs that are due a check, and +# A recording pass reads only running runs that are due a check, and # dispatches that stopped. Index( "ix_simulation_runs_due_checks", diff --git a/libs/policyengine-simulation-persistence/tests/test_recording_load.py b/libs/policyengine-simulation-persistence/tests/test_recording_load.py index efdec7d0b..048151da5 100644 --- a/libs/policyengine-simulation-persistence/tests/test_recording_load.py +++ b/libs/policyengine-simulation-persistence/tests/test_recording_load.py @@ -1,8 +1,8 @@ -"""A load check of the recording loop's conditional writes on PostgreSQL. +"""A load check of the recording passes' conditional writes on PostgreSQL. Several simulated Simulation Entrypoint instances, each on its own engine, repeatedly fail stalled dispatches, claim due calls and record an outcome for -each, as the recording loop does, until every run has finished. A third of +each, as recording passes do, until every run has finished. A third of the calls are household pairs (two runs under one call). Every call must be recorded once, its results written once and all of its runs finished once, and every stalled dispatch failed once, whatever the interleaving. The diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index 901063056..dbe4d5461 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -58,22 +58,15 @@ The ownership rules and roles are described in ## Recording run outcomes -Each Simulation Entrypoint instance runs a background loop that asks Modal -about the runs' calls and records their outcomes in the `simulation` schema -(`policyengine_simulation_entry.run_recording`). Two deploy settings serve it: - -- `gcloud run deploy --no-cpu-throttling` keeps the instance's CPU allocated - outside requests, so the loop runs between them. Production keeps at least - one instance running (`entrypoint_min_instances: 1`); staging scales to - zero, so it records outcomes only while an instance runs, and any request - starts one. -- Every deploy mounts sim's own Modal token as `MODAL_TOKEN_ID` and - `MODAL_TOKEN_SECRET`, whatever the Stage 12 setting: the deploy step - resolves the Secret Manager secrets named by the `beta` and `prod` - environment variables `MODAL_TOKEN_ID_SECRET_NAME` and - `MODAL_TOKEN_SECRET_SECRET_NAME` to their newest enabled versions. The - runtime service account must be able to read both secrets. The Stage 12 - comparison path's Modal calls use the same token until it is removed. +The Simulation Entrypoint records a run's outcome in the `simulation` schema +within a request; nothing runs between requests, so Cloud Run keeps no CPU +allocated outside them. Every deploy mounts sim's own Modal token as +`MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`, whatever the Stage 12 setting: the +deploy step resolves the Secret Manager secrets named by the `beta` and `prod` +environment variables `MODAL_TOKEN_ID_SECRET_NAME` and +`MODAL_TOKEN_SECRET_SECRET_NAME` to their newest enabled versions. The runtime +service account must be able to read both secrets. The Stage 12 comparison +path's Modal calls use the same token until it is removed. ## Candidate qualification diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index c08e0de2d..894394d50 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -578,18 +578,9 @@ def create_app( async def lifespan(_: FastAPI) -> AsyncIterator[None]: runtime_settings.validate() await runtime_backend.start() - # Each instance records the outcomes of due Modal calls in the - # background; Cloud Run keeps its CPU allocated between requests. - recording_loop = ( - None if runtime_database is None else runtime_database.recording_loop() - ) - if recording_loop is not None: - recording_loop.start() try: yield finally: - if recording_loop is not None: - await recording_loop.stop() await runtime_backend.close() if runtime_database is not None: runtime_database.dispose() diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py index 1353a6cfc..4898f758f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/entry_runtime.py @@ -41,7 +41,6 @@ ) from policyengine_simulation_entry.config import Settings - from policyengine_simulation_entry.run_recording import RunRecordingLoop # Moved out of app.py; this code still logs under that module's name. logger = logging.getLogger("policyengine_simulation_entry.app") @@ -221,10 +220,6 @@ async def read_revision(self) -> str | None: """Return the applied migration revision, or None before the first one.""" ... - def recording_loop(self) -> RunRecordingLoop | None: - """Return the loop that records Modal call outcomes on runs, if any.""" - ... - def dispose(self) -> None: """Close every pooled connection.""" ... diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index b7730cdeb..7bb3524ea 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -8,8 +8,8 @@ 3. record the call ID, moving the runs to ``running`` with their first check due shortly; or, when Modal refused the call, fail them. -If the instance stops between 1 and 3, the runs stay ``dispatching`` until the -recording loop fails them as abandoned. No route starts calls yet. +If the instance stops between 1 and 3, the runs stay ``dispatching`` until a +recording pass fails them as abandoned. No route starts calls yet. """ from __future__ import annotations @@ -39,7 +39,7 @@ def _event_name(prefix: str, outcome: str) -> str: return f"{prefix}_{outcome}" -# The first check of a new call; the recording loop backs off from there. +# The first check of a new call; a recording pass backs off from there. FIRST_CHECK_DELAY = timedelta(seconds=2) @@ -245,7 +245,7 @@ async def _record( return await self._settle_failed_record(steps, modal_call_id, context) if not recorded: # The runs stopped dispatching while Modal acknowledged the call, - # for example because the recording loop failed them as abandoned. + # for example because a recording pass failed them as abandoned. logger.warning( _event_name(steps.event_prefix, "call_not_recorded"), extra=context ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index caeac59fe..be934c947 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -1,23 +1,22 @@ -"""The background loop that records each Modal call's outcome on its runs. +"""Record each Modal call's outcome on its runs, one pass of due checks at a time. -Each Simulation Entrypoint instance runs one loop. On every pass it: +A pass: 1. fails runs and comparison steps left ``dispatching`` for ten minutes, so they can be submitted again, and pending comparison steps whose baseline or reform run failed, since they can never start; 2. claims the running calls that are due a check, by moving their next check - one lease ahead, so instances do not check the same call at once; + one lease ahead, so two passes do not check the same call at once; 3. asks Modal about each claimed call without waiting, and records what it found with conditional writes: a finished call's result is validated and its result rows are written in the transaction that marks its runs ``succeeded``; a failed or timed-out call is marked ``failed`` with its error; a call still running, or finished but of a kind this revision has - no recorder for, gets a later check, backing off. + no recorder for, gets a later check. Pre-emption is never a failure: Modal re-runs a pre-empted call under the same call ID, so it reads as still running. No caller polls, and Modal holds -no database credentials: this loop writes every outcome. Cloud Run keeps the -instance's CPU allocated between requests, so the loop runs then too. +no database credentials: sim writes every outcome. """ from __future__ import annotations @@ -38,7 +37,7 @@ ) if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Mapping + from collections.abc import Callable, Mapping from datetime import datetime from uuid import UUID @@ -72,10 +71,8 @@ def _event_name(prefix: str, outcome: str) -> str: CLAIM_BATCH_SIZE = 50 # Each check writes through the runtime engine's small pool, which requests # share (policyengine_simulation_persistence.engine: 2 connections plus 3 -# overflow), so the loop uses at most three at once. +# overflow), so a pass uses at most three at once. INSPECTION_CONCURRENCY = 3 -BUSY_PASS_INTERVAL_SECONDS = 1.0 -IDLE_PASS_INTERVAL_SECONDS = 5.0 # Past this many doublings the wait is LONGEST_BACKOFF anyway; capping the @@ -137,7 +134,7 @@ class ResultRecorders: class RecordingTransitions(Protocol): - """The ``RunTransitions`` writes the loop needs.""" + """The ``RunTransitions`` writes a recording pass needs.""" def claim_due_simulation_calls( self, *, limit: int, lease: timedelta @@ -423,8 +420,7 @@ async def _record( return steps = handlers.result_steps if steps is None: - # During a deploy an older revision runs beside the new one, and - # with CPU always allocated its loop keeps checking calls. A + # During a deploy an older revision runs beside the new one. A # revision without a recorder for this kind leaves the call for # one that has it; Modal's output retention bounds the wait. next_check_at = self.clock() + next_check_delay(handlers.check_count) @@ -467,71 +463,3 @@ async def _fail( _event_name(handlers.event_prefix, "failed"), extra={**handlers.subject, "error_code": error.code.value}, ) - - -class RunRecordingLoop: - """Run recording passes in the background until stopped. - - A pass that raises is logged and the loop goes on: recording problems - never stop the instance or fail a request. - """ - - def __init__( - self, - recorder: RunRecorder, - *, - sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, - ) -> None: - """Run passes of ``recorder``, waiting between them with ``sleep``.""" - self.recorder = recorder - self._sleep = sleep - self._task: asyncio.Task[None] | None = None - - async def run(self) -> None: - """Record outcomes until cancelled; wait less while there is work.""" - while True: - await self._sleep(IDLE_PASS_INTERVAL_SECONDS) - recording = await self._pass() - while recording is not None and recording.found_work: - await self._sleep(BUSY_PASS_INTERVAL_SECONDS) - recording = await self._pass() - - async def _pass(self) -> RecordingPass | None: - """Run one pass; log and return None if it raised anything. - - Only a signal to this process, such as the cancellation that stops - the loop at shutdown, passes through. - """ - try: - return await self.recorder.record_due_outcomes() - except BaseException as error: - if isinstance(error, OWN_PROCESS_SIGNALS): - raise - logger.exception( - "simulation_run_recording_pass_failed", - extra={"error_type": type(error).__name__}, - ) - return None - - def start(self) -> None: - """Start the loop as a background task of the running event loop.""" - if self._task is None: - self._task = asyncio.create_task(self.run()) - - async def stop(self) -> None: - """Cancel the loop and wait for it to end. - - A cancellation of the task calling this, such as a shutdown arriving - while the loop ends, is passed on rather than taken for the loop's. - """ - if self._task is None: - return - self._task.cancel() - try: - await self._task - except asyncio.CancelledError: - caller = asyncio.current_task() - if caller is not None and caller.cancelling(): - raise - finally: - self._task = None diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index 28bffdac0..e0ca4f1bd 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -6,18 +6,10 @@ from typing import TYPE_CHECKING from policyengine_simulation_persistence import ( - RunTransitions, create_simulation_engine, read_schema_revision, ) -from policyengine_simulation_entry.modal_calls import ModalCallInspector -from policyengine_simulation_entry.run_recording import ( - ResultRecorders, - RunRecorder, - RunRecordingLoop, -) - if TYPE_CHECKING: from sqlalchemy import Engine @@ -45,22 +37,6 @@ async def read_revision(self) -> str | None: """Return the applied migration revision, or None before the first one.""" return await asyncio.to_thread(self._read_revision) - def recording_loop(self) -> RunRecordingLoop: - """Return the loop that records Modal call outcomes on runs. - - No kind of run has a result recorder yet, so a finished call is - checked again until a revision with a recorder records it, or until - Modal no longer holds its output and the run fails as result_expired; - no route creates runs yet. - """ - return RunRecordingLoop( - RunRecorder( - transitions=RunTransitions(self.engine), - inspector=ModalCallInspector(), - recorders=ResultRecorders(), - ) - ) - def dispose(self) -> None: """Close every pooled connection.""" self.engine.dispose() diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py b/projects/policyengine-simulation-entry/tests/test_modal_token_deployment.py similarity index 84% rename from projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py rename to projects/policyengine-simulation-entry/tests/test_modal_token_deployment.py index ba08a2db6..24563081a 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording_deployment.py +++ b/projects/policyengine-simulation-entry/tests/test_modal_token_deployment.py @@ -1,8 +1,9 @@ -"""Deployment settings the background recording loop needs. +"""The entrypoint's Modal token, and no CPU kept outside requests. -The loop runs between requests, so Cloud Run must keep the instance's CPU -allocated outside them; and it asks Modal about calls, so every deploy mounts -sim's own Modal token, whether or not it also deploys the Stage 12 path. +The entrypoint starts runs' Modal calls and asks Modal about a call when +Modal reports that it ended, so every deploy mounts sim's own Modal token, +whether or not it also deploys the Stage 12 path. It records outcomes only +within requests, so Cloud Run keeps no CPU allocated between them. """ from __future__ import annotations @@ -38,14 +39,13 @@ def _run_deploy(tmp_path, **overrides): @pytest.mark.parametrize("deploy_stage12", ["true", "false"]) -def test_every_deploy_keeps_cpu_allocated_between_requests( +def test_no_deploy_keeps_cpu_allocated_outside_requests( tmp_path, deploy_stage12 ) -> None: result, arguments = _run_deploy(tmp_path, DEPLOY_STAGE12_V2=deploy_stage12) assert result.returncode == 0, result.stderr - assert "--no-cpu-throttling" in arguments - assert "--cpu-throttling" not in arguments + assert "--no-cpu-throttling" not in arguments @pytest.mark.parametrize("deploy_stage12", ["true", "false"]) diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index 67c1cda8e..5f9dde2aa 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -1,10 +1,9 @@ -"""The background loop that records Modal call outcomes on runs.""" +"""Recording each Modal call's outcome on its runs, one pass at a time.""" from __future__ import annotations import asyncio import builtins -import contextlib from datetime import timedelta import pytest @@ -18,13 +17,10 @@ CallSucceeded, ) from policyengine_simulation_entry.run_recording import ( - BUSY_PASS_INTERVAL_SECONDS, CHECK_LEASE, - IDLE_PASS_INTERVAL_SECONDS, RecordingPass, ResultRecorders, RunRecorder, - RunRecordingLoop, next_check_delay, ) from policyengine_simulation_persistence import ( @@ -398,152 +394,24 @@ def interrupted(modal_call_id, *, next_check_at): _run(recorder) -class _ScriptedRecorder: - """Return scripted passes, or raise, and count how often it ran.""" - - def __init__(self, passes) -> None: - self.passes = list(passes) - self.calls = 0 - - async def record_due_outcomes(self) -> RecordingPass: - self.calls += 1 - outcome = self.passes.pop(0) if self.passes else RecordingPass() - if isinstance(outcome, BaseException): - raise outcome - return outcome - - -def test_the_loop_waits_less_while_there_is_work_and_survives_failures() -> None: - scripted = _ScriptedRecorder( - [ - RecordingPass(checked=3), - RecordingPass(), - RuntimeError("database down"), - builtins.BaseExceptionGroup("pass", [_NotAnException()]), - ] - ) - waits = [] - - async def sleep(seconds: float) -> None: - waits.append(seconds) - if len(waits) == 6: - raise asyncio.CancelledError - - loop = RunRecordingLoop(scripted, sleep=sleep) # type: ignore[arg-type] - - with pytest.raises(asyncio.CancelledError): - asyncio.run(loop.run()) - - assert waits == [ - IDLE_PASS_INTERVAL_SECONDS, - BUSY_PASS_INTERVAL_SECONDS, - IDLE_PASS_INTERVAL_SECONDS, - IDLE_PASS_INTERVAL_SECONDS, - IDLE_PASS_INTERVAL_SECONDS, - IDLE_PASS_INTERVAL_SECONDS, - ] - assert scripted.calls == 5 - - -def test_a_signal_to_this_process_stops_the_loop() -> None: - scripted = _ScriptedRecorder([KeyboardInterrupt()]) - - async def sleep(_seconds: float) -> None: - return None - - loop = RunRecordingLoop(scripted, sleep=sleep) # type: ignore[arg-type] - - with pytest.raises(KeyboardInterrupt): - asyncio.run(loop.run()) - - -def test_the_loop_starts_and_stops_as_a_background_task() -> None: - scripted = _ScriptedRecorder([]) - - async def lifetime() -> None: - loop = RunRecordingLoop(scripted, sleep=lambda _seconds: asyncio.sleep(0)) # type: ignore[arg-type] - loop.start() - await asyncio.sleep(0.01) - await loop.stop() - await loop.stop() - - asyncio.run(lifetime()) - - assert scripted.calls > 0 - - -class _SlowToStopRecorder: - """A pass that takes a while to end once cancelled, as a closing write does.""" - - def __init__(self) -> None: - self.passing = asyncio.Event() - - async def record_due_outcomes(self) -> RecordingPass: - self.passing.set() - try: - await asyncio.sleep(60) - finally: - await asyncio.shield(asyncio.sleep(0.05)) - return RecordingPass() - - -def test_stopping_passes_on_a_cancellation_of_its_caller() -> None: - """Stopping the loop never swallows a shutdown of the task that stops it.""" - - async def shutdown() -> bool: - recorder = _SlowToStopRecorder() - loop = RunRecordingLoop(recorder, sleep=lambda _seconds: asyncio.sleep(0)) # type: ignore[arg-type] - loop.start() - await recorder.passing.wait() - stopper = asyncio.create_task(loop.stop()) - await asyncio.sleep(0.01) - stopper.cancel() - with contextlib.suppress(asyncio.CancelledError): - await stopper - return stopper.cancelled() - - assert asyncio.run(shutdown()) - - -class _TrackedLoop: - def __init__(self) -> None: - self.started = False - self.stopped = False - - def start(self) -> None: - self.started = True - - async def stop(self) -> None: - self.stopped = True - - -class _DatabaseWithLoop: - def __init__(self) -> None: - self.loop = _TrackedLoop() - +class _DatabaseThatOnlyReadsItsRevision: async def read_revision(self): return HEAD_REVISION - def recording_loop(self): - return self.loop - def dispose(self) -> None: return None -def test_each_instance_runs_the_loop_for_its_lifetime() -> None: - database = _DatabaseWithLoop() +def test_an_instance_runs_nothing_in_the_background() -> None: + """Outcomes are recorded when Modal reports them, within a request.""" app = create_app( settings=make_settings( simulation_database_url="postgresql://runtime:secret@db.example/postgres" ), backend=FakeBackend(), auth_dependency=lambda: None, - simulation_database=database, + simulation_database=_DatabaseThatOnlyReadsItsRevision(), ) - with TestClient(app): - assert database.loop.started is True - assert database.loop.stopped is False - - assert database.loop.stopped is True + with TestClient(app) as client: + assert client.get("/health").status_code == 200 diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_database.py b/projects/policyengine-simulation-entry/tests/test_simulation_database.py index 6a7e55547..95b1defae 100644 --- a/projects/policyengine-simulation-entry/tests/test_simulation_database.py +++ b/projects/policyengine-simulation-entry/tests/test_simulation_database.py @@ -33,9 +33,6 @@ async def read_revision(self): raise self.failure return self.revision - def recording_loop(self): - return None - def dispose(self): self.disposed = True From 79b92df22c47234d083a35910ef5c953c722afb4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:56:51 +0400 Subject: [PATCH 22/39] Define the report a Modal call sends sim when it ends A new contract module, call_reports, holds what the reporting route and its sender share: CallEndedReport (a Modal call ID with either the one or two simulation runs it computes or the report run whose comparison it computes, and nothing else), the route's path, the name of the Modal function that sends reports, and CallReportAnswer with the statuses the route answers. The statuses in RETRY_LATER_STATUSES (not dispatched yet, still running, Modal or the database unreachable, no result recorder, not configured) tell the sender to report again after a wait. Co-Authored-By: Claude Opus 5.5 --- .../call_reports.py | 119 ++++++++++++++++++ .../tests/test_call_reports.py | 99 +++++++++++++++ 2 files changed, 218 insertions(+) create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py create mode 100644 libs/policyengine-simulation-contract/tests/test_call_reports.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py new file mode 100644 index 000000000..823457b70 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py @@ -0,0 +1,119 @@ +"""The report a Modal call sends the Simulation Entrypoint when it ends. + +The last step of each top-level v2 Modal function (the economy coordinator +and the household functions) reports that its call ended: it starts the small +Modal function ``CALL_REPORT_FUNCTION_NAME``, which posts a +``CallEndedReport`` to the entrypoint's ``CALL_ENDED_PATH`` with the Modal +identity token of its container as a bearer token. The report carries the +call ID and the run IDs only, never a result: the entrypoint checks that the +request comes from that Modal function, asks Modal for the call's outcome +itself, and records it. + +The entrypoint answers with a ``CallReportAnswer``. While it cannot record +the outcome yet (the call has not finished, its ID is not stored yet, Modal +or the database cannot be reached), it answers 503 with a status in +``RETRY_LATER_STATUSES``, and the sender reports again after a wait. A +repeated report changes nothing. +""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Self +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +# The entrypoint's route for call-ended reports. It is not in the public +# OpenAPI document. +CALL_ENDED_PATH = "/internal/modal-calls/ended" +# The Modal function that sends reports; the entrypoint accepts identity +# tokens of this function only. +CALL_REPORT_FUNCTION_NAME = "report_call_ended" +# Modal sets this variable to the container's identity token in every +# function container. +MODAL_IDENTITY_ENVIRONMENT_VARIABLE = "MODAL_IDENTITY_TOKEN" +# A Modal function call ID, as Modal prints it. +MODAL_CALL_ID_PATTERN = r"^fc-[A-Za-z0-9]{1,64}$" +# A call computes one simulation run, or a household report's baseline and +# reform runs together. +MOST_SIMULATION_RUNS_PER_CALL = 2 + + +class CallEndedReport(BaseModel): + """A Modal call's report that it ended: its ID and the runs it computes. + + Exactly one of ``simulation_run_ids`` (a simulation call: one run, or a + household report's two) and ``report_run_id`` (an economy comparison + step's call) is set. + """ + + model_config = ConfigDict(frozen=True, extra="forbid") + + modal_call_id: str = Field(pattern=MODAL_CALL_ID_PATTERN) + simulation_run_ids: tuple[UUID, ...] | None = Field( + default=None, min_length=1, max_length=MOST_SIMULATION_RUNS_PER_CALL + ) + report_run_id: UUID | None = None + + @model_validator(mode="after") + def _names_one_kind_of_run(self) -> Self: + if (self.simulation_run_ids is None) == (self.report_run_id is None): + msg = "name either simulation_run_ids or report_run_id" + raise ValueError(msg) + if self.simulation_run_ids is not None and len( + set(self.simulation_run_ids) + ) != len(self.simulation_run_ids): + msg = "simulation_run_ids names a run twice" + raise ValueError(msg) + return self + + +class CallReportStatus(StrEnum): + """What the entrypoint did with a report.""" + + # The call's outcome is now recorded on its runs. + RECORDED = "recorded" + # The runs already hold the call's outcome; nothing changed. + ALREADY_RECORDED = "already_recorded" + # The runs are still being dispatched and do not hold the call ID yet. + NOT_DISPATCHED_YET = "not_dispatched_yet" + # Modal holds no output for the call yet. + STILL_RUNNING = "still_running" + # Modal could not be asked about the call. + MODAL_UNREACHABLE = "modal_unreachable" + # The serving revision cannot record this kind of result yet. + NO_RESULT_RECORDER = "no_result_recorder" + # The database could not be read or written. + NOT_WRITTEN = "not_written" + # The entrypoint is not configured to accept reports. + NOT_CONFIGURED = "not_configured" + # No run, or not every run, the report names exists. + UNKNOWN_RUNS = "unknown_runs" + # The runs hold another Modal call. + OTHER_CALL = "other_call" + + +RETRY_LATER_STATUSES = frozenset( + { + CallReportStatus.NOT_DISPATCHED_YET, + CallReportStatus.STILL_RUNNING, + CallReportStatus.MODAL_UNREACHABLE, + CallReportStatus.NO_RESULT_RECORDER, + CallReportStatus.NOT_WRITTEN, + CallReportStatus.NOT_CONFIGURED, + } +) + + +class CallReportAnswer(BaseModel): + """The entrypoint's answer to a report.""" + + model_config = ConfigDict(frozen=True) + + status: CallReportStatus + + @property + def retry_later(self) -> bool: + """Whether the sender should report again after a wait.""" + return self.status in RETRY_LATER_STATUSES diff --git a/libs/policyengine-simulation-contract/tests/test_call_reports.py b/libs/policyengine-simulation-contract/tests/test_call_reports.py new file mode 100644 index 000000000..8ccce497b --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_call_reports.py @@ -0,0 +1,99 @@ +"""The report a Modal function's call sends sim when it ends, and sim's answer.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.call_reports import ( + RETRY_LATER_STATUSES, + CallEndedReport, + CallReportAnswer, + CallReportStatus, +) +from pydantic import ValidationError + +CALL_ID = "fc-01JB3W7X9Z8Q2R5T6Y4U1I0O9P" + + +def test_a_report_names_the_simulation_runs_its_call_computes() -> None: + baseline, reform = uuid4(), uuid4() + + report = CallEndedReport( + modal_call_id=CALL_ID, simulation_run_ids=(baseline, reform) + ) + + assert report.simulation_run_ids == (baseline, reform) + assert report.report_run_id is None + + +def test_a_report_names_the_report_run_whose_comparison_it_computes() -> None: + report_run_id = uuid4() + + report = CallEndedReport(modal_call_id=CALL_ID, report_run_id=report_run_id) + + assert report.report_run_id == report_run_id + assert report.simulation_run_ids is None + + +def test_a_report_round_trips_as_json() -> None: + report = CallEndedReport(modal_call_id=CALL_ID, simulation_run_ids=(uuid4(),)) + + assert CallEndedReport.model_validate_json(report.model_dump_json()) == report + + +@pytest.mark.parametrize( + "fields", + [ + {}, + {"simulation_run_ids": (uuid4(),), "report_run_id": uuid4()}, + {"simulation_run_ids": ()}, + {"simulation_run_ids": (uuid4(), uuid4(), uuid4())}, + ], + ids=["neither", "both", "no-runs", "three-runs"], +) +def test_a_report_names_either_one_or_two_simulation_runs_or_one_report_run( + fields, +) -> None: + with pytest.raises(ValidationError): + CallEndedReport(modal_call_id=CALL_ID, **fields) + + +def test_a_report_names_each_simulation_run_once() -> None: + run_id = uuid4() + + with pytest.raises(ValidationError, match="twice"): + CallEndedReport(modal_call_id=CALL_ID, simulation_run_ids=(run_id, run_id)) + + +@pytest.mark.parametrize( + "call_id", ["", "fc-", "in-01JB3W7X9Z", "fc-01JB/../x", "fc-" + "A" * 80] +) +def test_a_report_names_a_modal_function_call(call_id) -> None: + with pytest.raises(ValidationError): + CallEndedReport(modal_call_id=call_id, report_run_id=uuid4()) + + +def test_a_report_carries_nothing_else() -> None: + with pytest.raises(ValidationError): + CallEndedReport.model_validate( + {"modal_call_id": CALL_ID, "report_run_id": str(uuid4()), "result": {}} + ) + + +def test_the_sender_retries_only_while_sim_cannot_record_yet() -> None: + assert { + CallReportStatus.NOT_DISPATCHED_YET, + CallReportStatus.STILL_RUNNING, + CallReportStatus.MODAL_UNREACHABLE, + CallReportStatus.NO_RESULT_RECORDER, + CallReportStatus.NOT_WRITTEN, + CallReportStatus.NOT_CONFIGURED, + } == RETRY_LATER_STATUSES + assert CallReportStatus.RECORDED not in RETRY_LATER_STATUSES + assert CallReportStatus.UNKNOWN_RUNS not in RETRY_LATER_STATUSES + + +def test_an_answer_says_whether_to_report_again() -> None: + assert CallReportAnswer(status=CallReportStatus.STILL_RUNNING).retry_later + assert not CallReportAnswer(status=CallReportStatus.ALREADY_RECORDED).retry_later From 5fc728f9b7f15d3f9cbff9482b7b8f22d115dbac Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:59:57 +0400 Subject: [PATCH 23/39] Check that a request comes from sim's call-report Modal function Modal gives every function container an RS256 identity token from https://oidc.modal.com whose claims name its workspace, environment, app and function. ModalIdentityVerifier checks a token against Modal's published keys (its JWKS), its issuer, audience and expiry, and accepts it only when it names the configured workspace and environment, a v2 app and the report_call_ended function, so no other Modal function, such as a worker running a user's policy, can report a call's end. When Modal's keys cannot be fetched it raises ModalKeysUnavailableError instead of refusing the token, so the sender reports again later. The expected workspace and environment come from SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID and SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT, set together or not at all. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_entry/config.py | 48 ++++++ .../modal_identity.py | 149 ++++++++++++++++++ .../tests/test_config.py | 42 ++++- .../tests/test_modal_identity.py | 149 ++++++++++++++++++ 4 files changed, 387 insertions(+), 1 deletion(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py create mode 100644 projects/policyengine-simulation-entry/tests/test_modal_identity.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index 5d23e8c54..6ac11436f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import re from dataclasses import dataclass from urllib.parse import urlparse @@ -12,10 +13,15 @@ assert_separate_manifest_names, ) +from policyengine_simulation_entry.modal_identity import ExpectedModalCaller + PRODUCTION_ENVIRONMENTS = frozenset({"main", "prod", "production"}) MODAL_GATEWAY_HOST_SUFFIX = ".modal.run" # Stage 12 comparison artifacts and rows are kept for exactly this many days. REQUIRED_STAGE12_RETENTION_DAYS = 30 +# Modal workspace IDs and environment names, as Modal prints them. +MODAL_WORKSPACE_ID_PATTERN = re.compile(r"ac-[A-Za-z0-9]+") +MODAL_ENVIRONMENT_NAME_PATTERN = re.compile(r"[a-z0-9][a-z0-9-]{0,63}") class ConfigurationError(RuntimeError): @@ -96,6 +102,10 @@ class Settings: # outside deployed environments; nothing on the old-gateway or Stage 12 # paths reads it. simulation_database_url: str = "" + # The Modal workspace and environment whose call-report function may + # report that calls ended. Both empty: the reporting route answers 503. + call_report_modal_workspace_id: str = "" + call_report_modal_environment: str = "" @classmethod def from_env(cls) -> Settings: @@ -147,6 +157,12 @@ def from_env(cls) -> Settings: stage12_artifact_bucket=os.getenv("STAGE12_ARTIFACT_BUCKET", ""), stage12_retention_days=int(os.getenv("STAGE12_RETENTION_DAYS", "30")), simulation_database_url=os.getenv("SIMULATION_DATABASE_URL", ""), + call_report_modal_workspace_id=os.getenv( + "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", "" + ), + call_report_modal_environment=os.getenv( + "SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", "" + ), ) @property @@ -154,6 +170,18 @@ def production(self) -> bool: """Whether this is a production environment, where auth must stay on.""" return self.environment in PRODUCTION_ENVIRONMENTS + @property + def expected_modal_caller(self) -> ExpectedModalCaller | None: + """The Modal caller whose call-ended reports are accepted, if configured.""" + if not ( + self.call_report_modal_workspace_id and self.call_report_modal_environment + ): + return None + return ExpectedModalCaller( + workspace_id=self.call_report_modal_workspace_id, + environment_name=self.call_report_modal_environment, + ) + @property def stage12_resources_configured(self) -> bool: """Whether the Stage 12 manifest environment, database and bucket are set.""" @@ -176,6 +204,7 @@ def validate(self) -> None: self._validate_old_gateway() self._validate_stage12() self._validate_simulation_database() + self._validate_call_reports() def _validate_caller_auth(self) -> None: """Require caller authentication in production, with an HTTPS issuer.""" @@ -281,3 +310,22 @@ def _validate_simulation_database(self) -> None: ): msg = "SIMULATION_DATABASE_URL must use the canonical postgresql:// scheme." raise ConfigurationError(msg) + + def _validate_call_reports(self) -> None: + """Require one Modal workspace ID and environment name, or neither.""" + workspace = self.call_report_modal_workspace_id + environment = self.call_report_modal_environment + if not workspace and not environment: + return + if not MODAL_WORKSPACE_ID_PATTERN.fullmatch(workspace): + msg = ( + "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID must be a Modal " + "workspace ID." + ) + raise ConfigurationError(msg) + if not MODAL_ENVIRONMENT_NAME_PATTERN.fullmatch(environment): + msg = ( + "SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT must be a Modal " + "environment name." + ) + raise ConfigurationError(msg) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py new file mode 100644 index 000000000..cbd971194 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py @@ -0,0 +1,149 @@ +"""Check that a request comes from sim's call-report Modal function. + +Modal gives every function container an OIDC identity token +(``MODAL_IDENTITY_TOKEN``), signed with RS256 by ``MODAL_OIDC_ISSUER``, whose +claims name the container's workspace, environment, app and function. A +call-ended report carries it as a bearer token. The entrypoint accepts it +only when Modal's published key verifies its signature, its issuer, audience +and expiry are valid, and it names the expected workspace and environment, a +v2 app (``V2_APPLICATION_PREFIX``) and the call-report function +(``CALL_REPORT_FUNCTION_NAME``). A token of any other Modal function, such as +a worker running a user's policy, is refused. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +import jwt +from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME +from policyengine_simulation_contract.stage12_manifest import V2_APPLICATION_PREFIX + +if TYPE_CHECKING: + from collections.abc import Mapping + +MODAL_OIDC_ISSUER = "https://oidc.modal.com" +MODAL_OIDC_AUDIENCE = "oidc.modal.com" +MODAL_JWKS_URL = "https://oidc.modal.com/.well-known/jwks.json" +# Modal lists RS256 as its only signing algorithm. +MODAL_SIGNING_ALGORITHMS = ("RS256",) +# Tolerated clock difference between Modal and this instance. +CLOCK_LEEWAY_SECONDS = 30 + + +class ModalIdentityError(Exception): + """The token does not prove that the request comes from the expected caller.""" + + +class ModalKeysUnavailableError(Exception): + """Modal's published keys could not be fetched, so no token can be checked.""" + + +class SigningKeySource(Protocol): + """Where the key that signed a token is found (``jwt.PyJWKClient``).""" + + def get_signing_key_from_jwt(self, token: str) -> jwt.PyJWK: + """Return the published key that signed ``token``.""" + ... + + +@dataclass(frozen=True) +class ExpectedModalCaller: + """The only Modal caller whose reports are accepted.""" + + workspace_id: str + environment_name: str + app_name_prefix: str = V2_APPLICATION_PREFIX + function_name: str = CALL_REPORT_FUNCTION_NAME + + def __post_init__(self) -> None: + """Refuse an expectation that any workspace or environment would meet.""" + if not self.workspace_id: + msg = "the expected Modal workspace is not set" + raise ValueError(msg) + if not self.environment_name: + msg = "the expected Modal environment is not set" + raise ValueError(msg) + + +@dataclass(frozen=True) +class ModalIdentity: + """The verified claims of a Modal container's identity token.""" + + app_name: str + function_name: str + container_id: str + + +def _claim(claims: Mapping[str, object], name: str) -> str: + value = claims.get(name) + return value if isinstance(value, str) else "" + + +class ModalIdentityVerifier: + """Verify Modal identity tokens against Modal's published keys.""" + + def __init__( + self, + expected: ExpectedModalCaller, + *, + signing_keys: SigningKeySource | None = None, + ) -> None: + """Expect ``expected``; fetch keys from ``MODAL_JWKS_URL`` unless given.""" + self.expected = expected + self.signing_keys: SigningKeySource = signing_keys or jwt.PyJWKClient( + MODAL_JWKS_URL + ) + + def verify(self, token: str) -> ModalIdentity: + """Return the token's identity, or raise ModalIdentityError. + + Raise ModalKeysUnavailableError when Modal's keys cannot be fetched. + Fetching them makes a network request the first time and whenever + the key cache expires, so call this in a worker thread. + """ + try: + signing_key = self.signing_keys.get_signing_key_from_jwt(token).key + except jwt.PyJWKClientConnectionError as error: + msg = "Modal's signing keys could not be fetched" + raise ModalKeysUnavailableError(msg) from error + except jwt.PyJWTError as error: + msg = f"the token names no key Modal publishes: {type(error).__name__}" + raise ModalIdentityError(msg) from error + try: + claims = jwt.decode( + token, + signing_key, + algorithms=list(MODAL_SIGNING_ALGORITHMS), + audience=MODAL_OIDC_AUDIENCE, + issuer=MODAL_OIDC_ISSUER, + leeway=CLOCK_LEEWAY_SECONDS, + options={"require": ["exp", "iat", "iss", "aud"]}, + ) + except jwt.PyJWTError as error: + msg = f"the token is not valid: {type(error).__name__}" + raise ModalIdentityError(msg) from error + return self._caller(claims) + + def _caller(self, claims: Mapping[str, object]) -> ModalIdentity: + expected = self.expected + if _claim(claims, "workspace_id") != expected.workspace_id: + msg = "the token names another Modal workspace" + raise ModalIdentityError(msg) + if _claim(claims, "environment_name") != expected.environment_name: + msg = "the token names another Modal environment" + raise ModalIdentityError(msg) + app_name = _claim(claims, "app_name") + if not app_name.startswith(expected.app_name_prefix): + msg = "the token names another Modal app" + raise ModalIdentityError(msg) + function_name = _claim(claims, "function_name") + if function_name != expected.function_name: + msg = "the token names another Modal function" + raise ModalIdentityError(msg) + return ModalIdentity( + app_name=app_name, + function_name=function_name, + container_id=_claim(claims, "container_id"), + ) diff --git a/projects/policyengine-simulation-entry/tests/test_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index 23ed45320..eb8a6b72e 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -2,7 +2,7 @@ import pytest from conftest import make_settings -from policyengine_simulation_entry.config import ConfigurationError +from policyengine_simulation_entry.config import ConfigurationError, Settings def test_production_refuses_disabled_auth(): @@ -335,3 +335,43 @@ def test_the_simulation_database_url_must_be_canonical(database_url): make_settings(simulation_database_url=database_url).validate() assert "secret" not in str(error.value) + + +def test_call_reports_are_off_without_a_modal_workspace(monkeypatch): + monkeypatch.delenv("SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", raising=False) + monkeypatch.delenv("SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", raising=False) + + settings = Settings.from_env() + + assert settings.expected_modal_caller is None + + +def test_call_reports_expect_the_configured_workspace_and_environment(monkeypatch): + monkeypatch.setenv("SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", "ac-workspace0001") + monkeypatch.setenv("SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", "staging") + + expected = Settings.from_env().expected_modal_caller + + assert expected is not None + assert expected.workspace_id == "ac-workspace0001" + assert expected.environment_name == "staging" + + +@pytest.mark.parametrize( + ("workspace", "environment"), + [ + ("ac-workspace0001", ""), + ("", "staging"), + ("workspace0001", "staging"), + ("ac-workspace0001", "Not An Environment"), + ], + ids=["no-environment", "no-workspace", "not-a-workspace-id", "bad-environment"], +) +def test_call_report_settings_must_name_one_modal_workspace_and_environment( + workspace, environment +): + with pytest.raises(ConfigurationError, match="SIMULATION_CALL_REPORT"): + make_settings( + call_report_modal_workspace_id=workspace, + call_report_modal_environment=environment, + ).validate() diff --git a/projects/policyengine-simulation-entry/tests/test_modal_identity.py b/projects/policyengine-simulation-entry/tests/test_modal_identity.py new file mode 100644 index 000000000..5cd0eade2 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_modal_identity.py @@ -0,0 +1,149 @@ +"""Checking that a request comes from sim's call-report Modal function.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import rsa +from policyengine_simulation_entry.modal_identity import ( + MODAL_OIDC_AUDIENCE, + MODAL_OIDC_ISSUER, + ExpectedModalCaller, + ModalIdentityError, + ModalIdentityVerifier, + ModalKeysUnavailableError, +) + +WORKSPACE_ID = "ac-workspace0001" +EXPECTED = ExpectedModalCaller(workspace_id=WORKSPACE_ID, environment_name="staging") +SIGNING_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +OTHER_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) + + +@dataclass +class _KeyFromJwks: + """Stands in for jwt.PyJWKClient: the public key Modal publishes.""" + + public_key: object + fetches: int = 0 + + def get_signing_key_from_jwt(self, _token: str): + self.fetches += 1 + return jwt.PyJWK.from_dict( + { + **jwt.algorithms.RSAAlgorithm.to_jwk(self.public_key, as_dict=True), + "kid": "modal-key", + "alg": "RS256", + } + ) + + +def _claims(**overrides): + now = int(time.time()) + claims = { + "iss": MODAL_OIDC_ISSUER, + "aud": MODAL_OIDC_AUDIENCE, + "sub": "modal:workspace_id:ac-workspace0001:...", + "iat": now, + "exp": now + 3600, + "workspace_id": WORKSPACE_ID, + "environment_name": "staging", + "app_name": "policyengine-simulation-v2-py6-10-0", + "function_name": "report_call_ended", + "container_id": "ta-container0001", + } + claims.update(overrides) + return {name: value for name, value in claims.items() if value is not None} + + +def _token(key=SIGNING_KEY, **overrides) -> str: + return jwt.encode( + _claims(**overrides), key, algorithm="RS256", headers={"kid": "modal-key"} + ) + + +def _verifier() -> ModalIdentityVerifier: + return ModalIdentityVerifier( + EXPECTED, signing_keys=_KeyFromJwks(SIGNING_KEY.public_key()) + ) + + +def test_a_token_of_the_call_report_function_is_accepted() -> None: + identity = _verifier().verify(_token()) + + assert identity.app_name == "policyengine-simulation-v2-py6-10-0" + assert identity.function_name == "report_call_ended" + assert identity.container_id == "ta-container0001" + + +@pytest.mark.parametrize( + ("overrides", "reason"), + [ + ({"workspace_id": "ac-someoneelse"}, "workspace"), + ({"environment_name": "main"}, "environment"), + ({"app_name": "policyengine-simulation-py6-10-0"}, "app"), + ({"function_name": "coordinate_report"}, "function"), + ({"workspace_id": None}, "workspace"), + ({"function_name": None}, "function"), + ], + ids=[ + "other-workspace", + "other-environment", + "main-path-app", + "other-function", + "no-workspace", + "no-function", + ], +) +def test_a_token_of_any_other_modal_caller_is_refused(overrides, reason) -> None: + with pytest.raises(ModalIdentityError, match=reason): + _verifier().verify(_token(**overrides)) + + +@pytest.mark.parametrize( + "token", + [ + _token(key=OTHER_KEY), + _token(iss="https://oidc.example.com"), + _token(aud="sim"), + _token(exp=int(time.time()) - 3600), + "not-a-token", + ], + ids=["wrong-key", "wrong-issuer", "wrong-audience", "expired", "malformed"], +) +def test_a_token_modal_did_not_sign_for_us_is_refused(token) -> None: + with pytest.raises(ModalIdentityError): + _verifier().verify(token) + + +def test_a_token_signed_with_another_algorithm_is_refused() -> None: + token = jwt.encode( + _claims(), "shared-secret-of-at-least-32-bytes!", algorithm="HS256" + ) + + with pytest.raises(ModalIdentityError): + _verifier().verify(token) + + +def test_the_expected_caller_needs_a_workspace_and_an_environment() -> None: + with pytest.raises(ValueError, match="workspace"): + ExpectedModalCaller(workspace_id="", environment_name="staging") + with pytest.raises(ValueError, match="environment"): + ExpectedModalCaller(workspace_id=WORKSPACE_ID, environment_name="") + + +class _UnreachableJwks: + def get_signing_key_from_jwt(self, _token: str): + msg = "Fail to fetch data from the url" + raise jwt.PyJWKClientConnectionError(msg) + + +def test_unreachable_modal_keys_are_not_a_refusal() -> None: + """The sender should report again later, not give up on its report.""" + verifier = ModalIdentityVerifier(EXPECTED, signing_keys=_UnreachableJwks()) + + with pytest.raises(ModalKeysUnavailableError): + verifier.verify(_token()) From 72387dabe4c5d3d80eed468eb3b4d0143714df84 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:04:26 +0400 Subject: [PATCH 24/39] Check a running call again when its claim's lease ends, without backoff The per-call backoff served a loop that asked Modal about running calls every few seconds. With Modal reporting when a call ends, a recording pass is only the safety check for calls that did not report, so a call that is still running, or finished but of a kind with no recorder here, is now left as claimed: it is due again when its claim's lease ends, and nothing is written. next_check_delay and the two writes that set a later check (record_simulation_call_running, record_comparison_running) are gone; the load check leaves such calls for after the lease too. Co-Authored-By: Claude Opus 5.5 --- .../run_transitions.py | 26 -------- .../tests/test_recording_load.py | 7 +-- .../tests/test_run_transitions.py | 18 ------ .../run_recording.py | 50 ++------------- .../tests/run_fakes.py | 20 ------ .../tests/test_run_recording.py | 61 +++++++------------ 6 files changed, 30 insertions(+), 152 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index 0d0a2e5aa..7686d639b 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -282,17 +282,6 @@ def _change_simulation_call( ) return tuple(SimulationRunRecord.from_row(row) for row in rows) - def record_simulation_call_running( - self, modal_call_id: str, *, next_check_at: datetime - ) -> bool: - """Set the next check of a call's runs that are still running.""" - with self.engine.begin() as connection: - return bool( - self._change_simulation_call( - connection, modal_call_id, {"next_check_at": next_check_at} - ) - ) - def complete_simulation_call( self, modal_call_id: str, write_results: SimulationResultWriter ) -> bool: @@ -487,21 +476,6 @@ def claim_due_comparisons( ) return [ReportRunRecord.from_row(row) for row in rows] - def record_comparison_running( - self, report_run_id: UUID, *, next_check_at: datetime - ) -> bool: - """Set the next check of a comparison step that is still running.""" - with self.engine.begin() as connection: - return ( - self._change_comparison( - connection, - report_run_id, - expected=RunStatus.RUNNING, - values={"comparison_next_check_at": next_check_at}, - ) - is not None - ) - def complete_comparison( self, report_run_id: UUID, write_results: ComparisonResultWriter ) -> bool: diff --git a/libs/policyengine-simulation-persistence/tests/test_recording_load.py b/libs/policyengine-simulation-persistence/tests/test_recording_load.py index 048151da5..00df05421 100644 --- a/libs/policyengine-simulation-persistence/tests/test_recording_load.py +++ b/libs/policyengine-simulation-persistence/tests/test_recording_load.py @@ -148,10 +148,9 @@ def _instance(database, clock, outcomes, barrier, seed: int) -> int: for call in calls: outcome = chooser.random() if outcome < 0.3: - transitions.record_simulation_call_running( - call.modal_call_id, next_check_at=clock() - ) - elif outcome < 0.5: + # Still running: checked again once the claim's lease ends. + continue + if outcome < 0.5: late.append((call.modal_call_id, outcome < 0.45)) else: _record( diff --git a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py index ad1da794a..eaec97a5d 100644 --- a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py +++ b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py @@ -260,21 +260,6 @@ def claim(_instance: int) -> list[str]: assert len(claimed_calls) == len(set(claimed_calls)) == len(run_ids) -def test_a_running_call_gets_its_next_check(migrated_database, public_api_rows) -> None: - stores = clocked_stores(migrated_database) - run = _running_run(stores, migrated_database, public_api_rows) - later = stores.clock() + timedelta(seconds=30) - - assert stores.transitions.record_simulation_call_running( - "fc-economy", next_check_at=later - ) - - assert stores.runs.get_simulation_run(run.id).next_check_at == later - assert not stores.transitions.record_simulation_call_running( - "fc-unknown", next_check_at=later - ) - - def test_results_are_written_in_the_transaction_that_marks_success( migrated_database, public_api_rows ) -> None: @@ -393,9 +378,6 @@ def test_a_comparison_runs_and_records_its_results_once( assert [due.id for due in claimed] == [report_run.id] assert claimed[0].comparison_check_count == 1 assert transitions.claim_due_comparisons(limit=10, lease=LEASE) == [] - assert transitions.record_comparison_running( - report_run.id, next_check_at=stores.clock() + LEASE - ) assert transitions.complete_comparison( report_run.id, lambda _connection, record: written.append(record) ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index be934c947..15331ec24 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -12,7 +12,7 @@ its result rows are written in the transaction that marks its runs ``succeeded``; a failed or timed-out call is marked ``failed`` with its error; a call still running, or finished but of a kind this revision has - no recorder for, gets a later check. + no recorder for, is checked again once its claim's lease ends. Pre-emption is never a failure: Modal re-runs a pre-empted call under the same call ID, so it reads as still running. No caller polls, and Modal holds @@ -66,8 +66,6 @@ def _event_name(prefix: str, outcome: str) -> str: CHECK_LEASE = timedelta(minutes=1) # A run left dispatching this long was claimed by an instance that stopped. DISPATCH_ABANDONED_AFTER = timedelta(minutes=10) -FIRST_BACKOFF = timedelta(seconds=2) -LONGEST_BACKOFF = timedelta(seconds=60) CLAIM_BATCH_SIZE = 50 # Each check writes through the runtime engine's small pool, which requests # share (policyengine_simulation_persistence.engine: 2 connections plus 3 @@ -75,17 +73,6 @@ def _event_name(prefix: str, outcome: str) -> str: INSPECTION_CONCURRENCY = 3 -# Past this many doublings the wait is LONGEST_BACKOFF anyway; capping the -# exponent keeps the product within timedelta's range for a long call. -_MOST_DOUBLINGS = 16 - - -def next_check_delay(check_count: int) -> timedelta: - """Return the wait before the next check of a call checked ``check_count`` times.""" - doublings = min(max(check_count - 1, 0), _MOST_DOUBLINGS) - return min(FIRST_BACKOFF * 2**doublings, LONGEST_BACKOFF) - - class SimulationResultRecorder(Protocol): """Validate one kind of simulation call's output and write its result rows.""" @@ -142,12 +129,6 @@ def claim_due_simulation_calls( """Claim running simulation calls that are due a check.""" ... - def record_simulation_call_running( - self, modal_call_id: str, *, next_check_at: datetime - ) -> bool: - """Set a running call's next check.""" - ... - def complete_simulation_call( self, modal_call_id: str, @@ -166,12 +147,6 @@ def claim_due_comparisons( """Claim running comparison steps that are due a check.""" ... - def record_comparison_running( - self, report_run_id: UUID, *, next_check_at: datetime - ) -> bool: - """Set a running comparison step's next check.""" - ... - def complete_comparison( self, report_run_id: UUID, @@ -232,10 +207,8 @@ class _CallHandlers: subject: dict[str, object] # "simulation_run" or "report_run_comparison": the log event names' start. event_prefix: str - check_count: int result_steps: _ResultSteps | None fail: Callable[[RunError], bool] - keep_running: Callable[[datetime], bool] def _invalid_result(error: ValueError) -> RunError: @@ -347,14 +320,8 @@ def _simulation_handlers( "simulation_run_ids": [str(run.id) for run in call.runs], }, event_prefix="simulation_run", - check_count=max(run.check_count for run in call.runs), result_steps=result_steps, fail=lambda error: self.transitions.fail_simulation_call(call_id, error), - keep_running=lambda next_check_at: ( - self.transitions.record_simulation_call_running( - call_id, next_check_at=next_check_at - ) - ), ) def _comparison_handlers( @@ -376,14 +343,8 @@ def _comparison_handlers( return call_id, _CallHandlers( subject={"modal_call_id": call_id, "report_run_id": str(report_run_id)}, event_prefix="report_run_comparison", - check_count=report_run.comparison_check_count, result_steps=result_steps, fail=lambda error: self.transitions.fail_comparison(report_run_id, error), - keep_running=lambda next_check_at: ( - self.transitions.record_comparison_running( - report_run_id, next_check_at=next_check_at - ) - ), ) async def _check( @@ -410,9 +371,9 @@ async def _check( async def _record( self, outcome: CallOutcome, handlers: _CallHandlers, recording: RecordingPass ) -> None: + # A call still running, or of a kind with no recorder here, is checked + # again once its claim's lease ends; nothing is written. if isinstance(outcome, CallRunning): - next_check_at = self.clock() + next_check_delay(handlers.check_count) - await asyncio.to_thread(handlers.keep_running, next_check_at) recording.still_running += 1 return if isinstance(outcome, CallFailed): @@ -420,11 +381,8 @@ async def _record( return steps = handlers.result_steps if steps is None: - # During a deploy an older revision runs beside the new one. A - # revision without a recorder for this kind leaves the call for + # A revision without a recorder for this kind leaves the call for # one that has it; Modal's output retention bounds the wait. - next_check_at = self.clock() + next_check_delay(handlers.check_count) - await asyncio.to_thread(handlers.keep_running, next_check_at) recording.without_recorder += 1 logger.warning( _event_name(handlers.event_prefix, "without_result_recorder"), diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py index 232e0eb0c..de01244a9 100644 --- a/projects/policyengine-simulation-entry/tests/run_fakes.py +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -286,15 +286,6 @@ def _call_runs(self, modal_call_id: str) -> list[FakeExecution]: if run.modal_call_id == modal_call_id and run.status is RunStatus.RUNNING ] - def record_simulation_call_running( - self, modal_call_id: str, *, next_check_at: datetime - ) -> bool: - """Set the next check of a call's running runs.""" - runs = self._call_runs(modal_call_id) - for run in runs: - run.next_check_at = next_check_at - return bool(runs) - def complete_simulation_call(self, modal_call_id: str, write_results) -> bool: """Mark a call's runs succeeded after the writer succeeds.""" runs = self._call_runs(modal_call_id) @@ -351,17 +342,6 @@ def claim_due_comparisons( for comparison in due ] - def record_comparison_running( - self, report_run_id: UUID, *, next_check_at: datetime - ) -> bool: - """Set a running comparison step's next check.""" - return ( - self._change_comparison( - report_run_id, RunStatus.RUNNING, next_check_at=next_check_at - ) - is not None - ) - def complete_comparison(self, report_run_id: UUID, write_results) -> bool: """Mark a running comparison step succeeded after the writer succeeds.""" comparison = self.comparisons[report_run_id] diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index 5f9dde2aa..a47f93b73 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -21,7 +21,6 @@ RecordingPass, ResultRecorders, RunRecorder, - next_check_delay, ) from policyengine_simulation_persistence import ( HEAD_REVISION, @@ -40,6 +39,9 @@ class BudgetResult(BaseModel): budgetary_impact: float +WORKER_ERROR = RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") + + class RecordingRecorder: """Validate outputs as ``BudgetResult`` and remember what it wrote.""" @@ -73,25 +75,7 @@ def _run(recorder) -> RecordingPass: return asyncio.run(recorder.record_due_outcomes()) -@pytest.mark.parametrize( - ("check_count", "seconds"), - [ - (0, 2), - (1, 2), - (2, 4), - (3, 8), - (5, 32), - (6, 60), - (40, 60), - (47, 60), - (10_000, 60), - ], -) -def test_checks_back_off_up_to_a_minute(check_count, seconds) -> None: - assert next_check_delay(check_count) == timedelta(seconds=seconds) - - -def test_a_running_call_is_checked_again_later() -> None: +def test_a_running_call_is_checked_again_after_its_claims_lease() -> None: recorder, transitions, inspector, clock = _recorder() (run_id,) = transitions.running_call("fc-1") @@ -102,8 +86,10 @@ def test_a_running_call_is_checked_again_later() -> None: run = transitions.runs[run_id] assert run.status is RunStatus.RUNNING assert run.check_count == 1 - assert run.next_check_at == clock() + timedelta(seconds=2) + assert run.next_check_at == clock() + CHECK_LEASE assert _run(recorder).checked == 0 + clock.advance(CHECK_LEASE) + assert _run(recorder).checked == 1 def test_a_failed_call_fails_its_runs() -> None: @@ -169,7 +155,7 @@ def test_a_kind_without_a_recorder_waits_for_an_instance_with_one() -> None: run = transitions.runs[run_id] assert run.status is RunStatus.RUNNING assert run.error is None - assert run.next_check_at == clock() + timedelta(seconds=2) + assert run.next_check_at == clock() + CHECK_LEASE def test_an_invalid_output_fails_its_runs_without_repeating_it() -> None: @@ -263,9 +249,7 @@ def test_comparison_steps_are_checked_and_recorded() -> None: assert recording.still_running == 1 assert transitions.comparisons[finished].status is RunStatus.SUCCEEDED assert [record.id for _, record, _ in comparisons.writes] == [finished] - assert transitions.comparisons[waiting].next_check_at == clock() + timedelta( - seconds=2 - ) + assert transitions.comparisons[waiting].next_check_at == clock() + CHECK_LEASE def test_dispatches_left_for_ten_minutes_are_failed_by_a_pass() -> None: @@ -322,18 +306,17 @@ def test_a_database_error_in_one_check_ends_with_its_pass(caplog) -> None: recorder, transitions, inspector, clock = _recorder() (broken_run,) = transitions.running_call("fc-broken") (failed_run,) = transitions.running_call("fc-failed") - inspector.outcomes["fc-failed"] = CallFailed( - RunError(RunErrorCode.WORKER_ERROR, "The Modal call raised KeyError.") - ) - keep_running = transitions.record_simulation_call_running + for call_id in ("fc-broken", "fc-failed"): + inspector.outcomes[call_id] = CallFailed(WORKER_ERROR) + fail_call = transitions.fail_simulation_call - def drop_one_connection(modal_call_id, *, next_check_at): + def drop_one_connection(modal_call_id, error): if modal_call_id == "fc-broken": msg = "connection lost" raise ConnectionError(msg) - return keep_running(modal_call_id, next_check_at=next_check_at) + return fail_call(modal_call_id, error) - transitions.record_simulation_call_running = drop_one_connection # type: ignore[method-assign] + transitions.fail_simulation_call = drop_one_connection # type: ignore[method-assign] recording = _run(recorder) @@ -363,14 +346,15 @@ class _NotAnException(BaseException): def test_a_non_exception_error_in_one_check_ends_with_its_pass(caplog) -> None: - recorder, transitions, _, clock = _recorder() + recorder, transitions, inspector, clock = _recorder() (run_id,) = transitions.running_call("fc-broken") + inspector.outcomes["fc-broken"] = CallFailed(WORKER_ERROR) - def raise_a_group(modal_call_id, *, next_check_at): + def raise_a_group(modal_call_id, error): message = "writes" raise builtins.BaseExceptionGroup(message, [_NotAnException()]) - transitions.record_simulation_call_running = raise_a_group # type: ignore[method-assign] + transitions.fail_simulation_call = raise_a_group # type: ignore[method-assign] recording = _run(recorder) @@ -382,13 +366,14 @@ def raise_a_group(modal_call_id, *, next_check_at): def test_a_signal_to_this_process_stops_a_check() -> None: - recorder, transitions, _, _ = _recorder() + recorder, transitions, inspector, _ = _recorder() transitions.running_call("fc-1") + inspector.outcomes["fc-1"] = CallFailed(WORKER_ERROR) - def interrupted(modal_call_id, *, next_check_at): + def interrupted(modal_call_id, error): raise KeyboardInterrupt - transitions.record_simulation_call_running = interrupted # type: ignore[method-assign] + transitions.fail_simulation_call = interrupted # type: ignore[method-assign] with pytest.raises(KeyboardInterrupt): _run(recorder) From 53e2a46a0ef336f5ed39c0d2bfe1475777faabaa Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:08:04 +0400 Subject: [PATCH 25/39] Record a call's outcome when its Modal function reports that it ended RunRecorder.record_reported_call takes a CallEndedReport, reads the runs (or the comparison step) it names, asks Modal about the call without waiting, and records the outcome with the same conditional writes as a recording pass, returning a CallReportStatus. A report is only a hint, so it records nothing it cannot confirm: - runs that hold the call and are running: the outcome is recorded; - runs that already finished with the call: already_recorded, and Modal is not asked again, so a repeated report changes nothing; - runs still being dispatched, whose call ID is not stored yet (a fast household call can report first): not_dispatched_yet; - a call Modal still runs, Modal or the database unreachable, or a kind without a recorder: still_running, modal_unreachable, not_written or no_result_recorder, with nothing written; - unknown runs, or runs holding another call or finished without one: unknown_runs or other_call. Only a signal to this process propagates. Recording one call now returns what it did, which a pass counts and a report maps to its status, and log records name the error an error was raised from (error_cause_type), so a CallInspectionError no longer hides the Modal error behind it. Co-Authored-By: Claude Opus 5.5 --- .../run_recording.py | 240 ++++++++++++--- .../tests/run_fakes.py | 30 ++ .../tests/test_call_report_recording.py | 280 ++++++++++++++++++ .../tests/test_run_recording.py | 1 + 4 files changed, 515 insertions(+), 36 deletions(-) create mode 100644 projects/policyengine-simulation-entry/tests/test_call_report_recording.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 15331ec24..72a64d83d 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -25,9 +25,16 @@ import logging from dataclasses import dataclass, field from datetime import timedelta +from enum import StrEnum from typing import TYPE_CHECKING, Protocol -from policyengine_simulation_persistence import RunError, RunErrorCode +from policyengine_simulation_contract.call_reports import CallReportStatus +from policyengine_simulation_persistence import ( + DueSimulationCall, + RunError, + RunErrorCode, + RunStatus, +) from policyengine_simulation_persistence.run_store import utc_now from policyengine_simulation_entry.modal_calls import ( @@ -41,9 +48,9 @@ from datetime import datetime from uuid import UUID + from policyengine_simulation_contract.call_reports import CallEndedReport from policyengine_simulation_persistence import ( AbandonedDispatches, - DueSimulationCall, ReportKind, ReportRunRecord, SimulationKind, @@ -170,6 +177,42 @@ def fail_comparisons_of_failed_runs(self) -> tuple[UUID, ...]: ... +class RunReader(Protocol): + """The ``RunStore`` reads a reported call needs.""" + + def get_simulation_run(self, run_id: UUID) -> SimulationRunRecord | None: + """Return the simulation run with this ID, or None.""" + ... + + def get_report_run(self, report_run_id: UUID) -> ReportRunRecord | None: + """Return the report run with this ID, or None.""" + ... + + +class _Recorded(StrEnum): + """What recording one call's outcome did.""" + + SUCCEEDED = "succeeded" + FAILED = "failed" + STILL_RUNNING = "still_running" + WITHOUT_RECORDER = "without_recorder" + NOT_WRITTEN = "not_written" + # Another writer recorded the outcome first. + UNCHANGED = "unchanged" + + +_REPORT_STATUSES = { + _Recorded.SUCCEEDED: CallReportStatus.RECORDED, + _Recorded.FAILED: CallReportStatus.RECORDED, + _Recorded.STILL_RUNNING: CallReportStatus.STILL_RUNNING, + _Recorded.WITHOUT_RECORDER: CallReportStatus.NO_RESULT_RECORDER, + _Recorded.NOT_WRITTEN: CallReportStatus.NOT_WRITTEN, + _Recorded.UNCHANGED: CallReportStatus.ALREADY_RECORDED, +} +# A run or comparison step in these statuses may not hold its call ID yet. +_NOT_YET_DISPATCHED = frozenset({RunStatus.PENDING, RunStatus.DISPATCHING}) + + @dataclass class RecordingPass: """What one pass found and recorded.""" @@ -188,6 +231,19 @@ def found_work(self) -> bool: """Whether the pass checked or failed anything.""" return bool(self.checked or self.abandoned or self.comparisons_without_runs) + def count(self, recorded: _Recorded) -> None: + """Count what recording one checked call did.""" + counters = { + _Recorded.SUCCEEDED: "succeeded", + _Recorded.FAILED: "failed", + _Recorded.STILL_RUNNING: "still_running", + _Recorded.WITHOUT_RECORDER: "without_recorder", + _Recorded.NOT_WRITTEN: "not_reached", + } + counter = counters.get(recorded) + if counter is not None: + setattr(self, counter, getattr(self, counter) + 1) + @dataclass(frozen=True) class _ResultSteps: @@ -222,9 +278,10 @@ def _invalid_result(error: ValueError) -> RunError: @dataclass class RunRecorder: - """Record the outcomes of due Modal calls, one pass at a time.""" + """Record Modal call outcomes: one reported call, or a pass of due checks.""" transitions: RecordingTransitions + runs: RunReader inspector: CallInspector recorders: ResultRecorders = field(default_factory=ResultRecorders) clock: Callable[[], datetime] = utc_now @@ -353,7 +410,7 @@ async def _check( recording.checked += 1 try: outcome = await self.inspector.inspect(modal_call_id) - await self._record(outcome, handlers, recording) + recorded = await self._record(outcome, handlers) # Modal could not be reached, or a database write failed (for example # a dropped connection). Each check ends here, whatever was raised # (only a signal to this process passes), so a pass returns only once @@ -365,59 +422,170 @@ async def _check( recording.not_reached += 1 logger.warning( _event_name(handlers.event_prefix, "check_failed"), - extra={**handlers.subject, "error_type": type(error).__name__}, + extra={**handlers.subject, **_error_fields(error)}, ) + return + recording.count(recorded) + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + """Record the outcome of a call whose Modal function reported its end. + + The report is only a hint: the runs are read, Modal is asked about the + call without waiting, and the outcome is recorded with the same + conditional writes as a recording pass, so a repeated report changes + nothing. While the outcome cannot be recorded yet, the answer is one + of ``RETRY_LATER_STATUSES`` and nothing is written; the sender + reports again, and the safety check finds the call otherwise. Only a + signal to this process propagates. + """ + try: + found = await asyncio.to_thread(self._reported_call, report) + except Exception as error: # the database could not be read + logger.exception( + "call_report_not_read", + extra={"modal_call_id": report.modal_call_id, **_error_fields(error)}, + ) + return CallReportStatus.NOT_WRITTEN + if isinstance(found, CallReportStatus): + return found + try: + outcome = await self.inspector.inspect(report.modal_call_id) + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.warning( + _event_name(found.event_prefix, "check_failed"), + extra={**found.subject, **_error_fields(error)}, + ) + return CallReportStatus.MODAL_UNREACHABLE + try: + recorded = await self._record(outcome, found) + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.warning( + _event_name(found.event_prefix, "outcome_not_written"), + extra={**found.subject, **_error_fields(error)}, + ) + return CallReportStatus.NOT_WRITTEN + return _REPORT_STATUSES[recorded] + + def _reported_call( + self, report: CallEndedReport + ) -> CallReportStatus | _CallHandlers: + """Return how to record the reported call, or why it cannot be yet.""" + call_id = report.modal_call_id + if report.simulation_run_ids is not None: + runs = [ + self.runs.get_simulation_run(run_id) + for run_id in report.simulation_run_ids + ] + known = [run for run in runs if run is not None] + if len(known) != len(runs): + return CallReportStatus.UNKNOWN_RUNS + state = _reported_state( + call_id, [(run.modal_call_id, run.status) for run in known] + ) + if state is not None: + return state + due = DueSimulationCall(modal_call_id=call_id, runs=tuple(known)) + return self._simulation_handlers(due)[1] + report_run = ( + None + if report.report_run_id is None + else self.runs.get_report_run(report.report_run_id) + ) + if report_run is None or report_run.comparison_status is None: + return CallReportStatus.UNKNOWN_RUNS + state = _reported_state( + call_id, + [(report_run.comparison_modal_call_id, report_run.comparison_status)], + ) + if state is not None: + return state + return self._comparison_handlers(report_run)[1] - async def _record( - self, outcome: CallOutcome, handlers: _CallHandlers, recording: RecordingPass - ) -> None: + async def _record(self, outcome: CallOutcome, handlers: _CallHandlers) -> _Recorded: # A call still running, or of a kind with no recorder here, is checked # again once its claim's lease ends; nothing is written. if isinstance(outcome, CallRunning): - recording.still_running += 1 - return + return _Recorded.STILL_RUNNING if isinstance(outcome, CallFailed): - await self._fail(handlers, outcome.error, recording) - return + return await self._fail(handlers, outcome.error) steps = handlers.result_steps if steps is None: # A revision without a recorder for this kind leaves the call for # one that has it; Modal's output retention bounds the wait. - recording.without_recorder += 1 logger.warning( _event_name(handlers.event_prefix, "without_result_recorder"), extra=handlers.subject, ) - return + return _Recorded.WITHOUT_RECORDER + return await self._record_result(outcome.output, steps, handlers) + + async def _record_result( + self, output: object, steps: _ResultSteps, handlers: _CallHandlers + ) -> _Recorded: + """Validate a finished call's output and write it with its success.""" try: - result = await asyncio.to_thread(steps.validate, outcome.output) + result = await asyncio.to_thread(steps.validate, output) except ValueError as error: - await self._fail(handlers, _invalid_result(error), recording) - return + return await self._fail(handlers, _invalid_result(error)) try: completed = await asyncio.to_thread(steps.complete, result) # Writing the result rows failed, for example because the database - # was unreachable: nothing was recorded, and the lease makes the call - # due again. + # was unreachable: nothing was recorded, and the call is checked again. except Exception as error: - recording.not_reached += 1 logger.exception( _event_name(handlers.event_prefix, "result_not_written"), - extra={**handlers.subject, "error_type": type(error).__name__}, - ) - return - if completed: - recording.succeeded += 1 - logger.info( - _event_name(handlers.event_prefix, "succeeded"), extra=handlers.subject + extra={**handlers.subject, **_error_fields(error)}, ) + return _Recorded.NOT_WRITTEN + if not completed: + return _Recorded.UNCHANGED + logger.info( + _event_name(handlers.event_prefix, "succeeded"), extra=handlers.subject + ) + return _Recorded.SUCCEEDED + + async def _fail(self, handlers: _CallHandlers, error: RunError) -> _Recorded: + if not await asyncio.to_thread(handlers.fail, error): + return _Recorded.UNCHANGED + logger.warning( + _event_name(handlers.event_prefix, "failed"), + extra={**handlers.subject, "error_code": error.code.value}, + ) + return _Recorded.FAILED - async def _fail( - self, handlers: _CallHandlers, error: RunError, recording: RecordingPass - ) -> None: - if await asyncio.to_thread(handlers.fail, error): - recording.failed += 1 - logger.warning( - _event_name(handlers.event_prefix, "failed"), - extra={**handlers.subject, "error_code": error.code.value}, - ) + +def _reported_state( + modal_call_id: str, executions: list[tuple[str | None, RunStatus]] +) -> CallReportStatus | None: + """Return why a reported call cannot be recorded, or None if it can. + + ``executions`` holds each reported run's (or the comparison step's) call + ID and status. + """ + held_call_ids = {held for held, _ in executions} + if held_call_ids == {modal_call_id}: + if any(status is RunStatus.RUNNING for _, status in executions): + return None + return CallReportStatus.ALREADY_RECORDED + if held_call_ids - {modal_call_id, None}: + return CallReportStatus.OTHER_CALL + # Some runs hold no call ID yet. While they are being dispatched, the + # dispatcher is about to store it; a finished run without it (failed as + # abandoned, say) never will. + if all( + status in _NOT_YET_DISPATCHED for held, status in executions if held is None + ): + return CallReportStatus.NOT_DISPATCHED_YET + return CallReportStatus.OTHER_CALL + + +def _error_fields(error: BaseException) -> dict[str, str]: + """Name an error, and the error it was raised from, for a log record.""" + fields = {"error_type": type(error).__name__} + if error.__cause__ is not None: + fields["error_cause_type"] = type(error.__cause__).__name__ + return fields diff --git a/projects/policyengine-simulation-entry/tests/run_fakes.py b/projects/policyengine-simulation-entry/tests/run_fakes.py index de01244a9..acee770ed 100644 --- a/projects/policyengine-simulation-entry/tests/run_fakes.py +++ b/projects/policyengine-simulation-entry/tests/run_fakes.py @@ -74,6 +74,7 @@ class FakeRunRecord: simulation_kind: SimulationKind modal_call_id: str | None check_count: int + status: RunStatus = RunStatus.RUNNING @dataclass(frozen=True) @@ -84,6 +85,7 @@ class FakeReportRunRecord: report_kind: str comparison_modal_call_id: str | None comparison_check_count: int + comparison_status: RunStatus | None = RunStatus.RUNNING class FakeRunTransitions: @@ -247,6 +249,34 @@ def comparison_holds_call(self, report_run_id: UUID, modal_call_id: str) -> bool with self.lock: return self.comparisons[report_run_id].modal_call_id == modal_call_id + # --- Reading ------------------------------------------------------- + + def get_simulation_run(self, run_id: UUID) -> FakeRunRecord | None: + """Return the run as a record, or None when there is no such run.""" + run = self.runs.get(run_id) + if run is None: + return None + return FakeRunRecord( + id=run.id, + simulation_kind=SimulationKind(run.kind), + modal_call_id=run.modal_call_id, + check_count=run.check_count, + status=run.status, + ) + + def get_report_run(self, report_run_id: UUID) -> FakeReportRunRecord | None: + """Return the comparison step's report run, or None when there is none.""" + comparison = self.comparisons.get(report_run_id) + if comparison is None: + return None + return FakeReportRunRecord( + id=comparison.id, + report_kind=comparison.kind, + comparison_modal_call_id=comparison.modal_call_id, + comparison_check_count=comparison.check_count, + comparison_status=comparison.status, + ) + # --- Recording ------------------------------------------------------- def claim_due_simulation_calls( diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_recording.py b/projects/policyengine-simulation-entry/tests/test_call_report_recording.py new file mode 100644 index 000000000..76ef79bd0 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_call_report_recording.py @@ -0,0 +1,280 @@ +"""Recording a call's outcome when its Modal function reports that it ended.""" + +from __future__ import annotations + +import asyncio +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.call_reports import ( + CallEndedReport, + CallReportStatus, +) +from policyengine_simulation_entry.modal_calls import ( + CallFailed, + CallInspectionError, + CallSucceeded, +) +from policyengine_simulation_entry.run_recording import ResultRecorders, RunRecorder +from policyengine_simulation_persistence import ( + ReportKind, + RunError, + RunErrorCode, + RunStatus, + SimulationKind, +) +from run_fakes import FakeInspector, FakeRunTransitions, ManualClock +from test_run_recording import RecordingRecorder + +WORKER_TIMEOUT = RunError(RunErrorCode.WORKER_TIMEOUT, "The Modal call timed out.") + + +def _recorder(recorders=None): + clock = ManualClock() + transitions = FakeRunTransitions(clock) + inspector = FakeInspector() + recorder = RunRecorder( + transitions=transitions, + runs=transitions, + inspector=inspector, + recorders=recorders or ResultRecorders(), + clock=clock, + ) + return recorder, transitions, inspector + + +def _economy_recorder(): + budget = RecordingRecorder() + recorder, transitions, inspector = _recorder( + ResultRecorders(simulations={SimulationKind.ECONOMY: budget}) + ) + return recorder, transitions, inspector, budget + + +def _report(recorder, **fields) -> CallReportStatus: + return asyncio.run(recorder.record_reported_call(CallEndedReport(**fields))) + + +def test_a_finished_call_is_recorded_when_it_reports() -> None: + recorder, transitions, inspector, budget = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 2.5}) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.RECORDED + assert transitions.runs[run_id].status is RunStatus.SUCCEEDED + assert len(budget.writes) == 1 + + +def test_a_repeated_report_changes_nothing() -> None: + recorder, transitions, inspector, budget = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 2.5}) + _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.ALREADY_RECORDED + assert len(budget.writes) == 1 + assert inspector.inspected == ["fc-done"] + + +def test_a_household_pair_is_recorded_from_one_report() -> None: + household = RecordingRecorder() + recorder, transitions, inspector = _recorder( + ResultRecorders(simulations={SimulationKind.HOUSEHOLD: household}) + ) + pair = transitions.running_call("fc-pair", "household", "household") + inspector.outcomes["fc-pair"] = CallSucceeded({"budgetary_impact": 0}) + + status = _report(recorder, modal_call_id="fc-pair", simulation_run_ids=tuple(pair)) + + assert status is CallReportStatus.RECORDED + assert {transitions.runs[run_id].status for run_id in pair} == {RunStatus.SUCCEEDED} + assert len(household.writes) == 1 + + +def test_a_failed_call_is_recorded_as_failed() -> None: + recorder, transitions, inspector = _recorder() + (run_id,) = transitions.running_call("fc-failed") + inspector.outcomes["fc-failed"] = CallFailed(WORKER_TIMEOUT) + + status = _report(recorder, modal_call_id="fc-failed", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.RECORDED + assert transitions.runs[run_id].status is RunStatus.FAILED + assert transitions.runs[run_id].error.code is RunErrorCode.WORKER_TIMEOUT + + +def test_a_call_modal_still_runs_is_left_for_another_report() -> None: + recorder, transitions, _, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-running") + next_check_at = transitions.runs[run_id].next_check_at + + status = _report(recorder, modal_call_id="fc-running", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.STILL_RUNNING + assert transitions.runs[run_id].status is RunStatus.RUNNING + assert transitions.runs[run_id].next_check_at == next_check_at + + +def test_a_call_whose_id_is_not_stored_yet_is_left_for_another_report() -> None: + """A fast household call can report before its dispatcher stores its ID.""" + recorder, transitions, inspector, _ = _economy_recorder() + run_id = transitions.add_run(status=RunStatus.DISPATCHING) + + status = _report(recorder, modal_call_id="fc-fast", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NOT_DISPATCHED_YET + assert inspector.inspected == [] + assert transitions.runs[run_id].status is RunStatus.DISPATCHING + + +def test_a_report_naming_an_unknown_run_is_refused() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + + status = _report( + recorder, modal_call_id="fc-done", simulation_run_ids=(run_id, uuid4()) + ) + + assert status is CallReportStatus.UNKNOWN_RUNS + assert inspector.inspected == [] + + +@pytest.mark.parametrize( + "run_state", + [ + {"status": RunStatus.RUNNING, "modal_call_id": "fc-other"}, + {"status": RunStatus.SUCCEEDED, "modal_call_id": "fc-other"}, + # Failed as abandoned before the call's ID was stored. + {"status": RunStatus.FAILED, "modal_call_id": None}, + ], + ids=["running-another-call", "finished-another-call", "failed-without-call"], +) +def test_a_report_for_runs_that_do_not_hold_its_call_is_refused(run_state) -> None: + recorder, transitions, inspector, _ = _economy_recorder() + run_id = transitions.add_run(**run_state) + + status = _report(recorder, modal_call_id="fc-late", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.OTHER_CALL + assert inspector.inspected == [] + assert transitions.runs[run_id].status is run_state["status"] + + +def test_an_unreachable_modal_is_left_for_another_report() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallInspectionError("ConnectionError") + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.MODAL_UNREACHABLE + assert transitions.runs[run_id].status is RunStatus.RUNNING + + +def test_a_kind_without_a_recorder_is_left_for_a_revision_with_one() -> None: + recorder, transitions, inspector = _recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallSucceeded({"budgetary_impact": 1}) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NO_RESULT_RECORDER + assert transitions.runs[run_id].status is RunStatus.RUNNING + + +def test_a_failed_read_is_left_for_another_report() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + + def unreachable_database(_run_id): + msg = "connection lost" + raise ConnectionError(msg) + + transitions.get_simulation_run = unreachable_database # type: ignore[method-assign] + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NOT_WRITTEN + assert inspector.inspected == [] + + +@pytest.mark.parametrize( + ("outcome", "write"), + [ + (CallSucceeded({"budgetary_impact": 2.5}), "complete_simulation_call"), + (CallFailed(WORKER_TIMEOUT), "fail_simulation_call"), + ], + ids=["result", "failure"], +) +def test_a_failed_write_is_left_for_another_report(outcome, write) -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = outcome + + def unreachable_database(*_arguments): + msg = "connection lost" + raise ConnectionError(msg) + + setattr(transitions, write, unreachable_database) + + status = _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + assert status is CallReportStatus.NOT_WRITTEN + assert transitions.runs[run_id].status is RunStatus.RUNNING + + +def test_a_signal_to_this_process_is_not_swallowed() -> None: + recorder, transitions, inspector, _ = _economy_recorder() + (run_id,) = transitions.running_call("fc-done") + inspector.outcomes["fc-done"] = CallFailed(WORKER_TIMEOUT) + + def interrupted(*_arguments): + raise KeyboardInterrupt + + transitions.fail_simulation_call = interrupted # type: ignore[method-assign] + + with pytest.raises(KeyboardInterrupt): + _report(recorder, modal_call_id="fc-done", simulation_run_ids=(run_id,)) + + +def test_a_comparison_call_is_recorded_when_it_reports() -> None: + comparisons = RecordingRecorder() + recorder, transitions, inspector = _recorder( + ResultRecorders(comparisons={ReportKind.ECONOMY: comparisons}) + ) + report_run_id = transitions.add_comparison( + status=RunStatus.RUNNING, modal_call_id="fc-comparison" + ) + inspector.outcomes["fc-comparison"] = CallSucceeded({"budgetary_impact": 3}) + + status = _report( + recorder, modal_call_id="fc-comparison", report_run_id=report_run_id + ) + + assert status is CallReportStatus.RECORDED + assert transitions.comparisons[report_run_id].status is RunStatus.SUCCEEDED + assert [record.id for _, record, _ in comparisons.writes] == [report_run_id] + + +def test_a_comparison_report_for_an_unknown_report_run_is_refused() -> None: + recorder, _, inspector, _ = _economy_recorder() + + status = _report(recorder, modal_call_id="fc-comparison", report_run_id=uuid4()) + + assert status is CallReportStatus.UNKNOWN_RUNS + assert inspector.inspected == [] + + +def test_a_comparison_whose_call_id_is_not_stored_yet_waits() -> None: + recorder, transitions, _, _ = _economy_recorder() + report_run_id = transitions.add_comparison(status=RunStatus.DISPATCHING) + + status = _report( + recorder, modal_call_id="fc-comparison", report_run_id=report_run_id + ) + + assert status is CallReportStatus.NOT_DISPATCHED_YET diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index a47f93b73..f646cfca9 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -64,6 +64,7 @@ def _recorder(recorders=None): inspector = FakeInspector() recorder = RunRecorder( transitions=transitions, + runs=transitions, inspector=inspector, recorders=recorders or ResultRecorders(), clock=clock, From 591a4acd60b792c669520eaaf891ef612be82f89 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:14:23 +0400 Subject: [PATCH 26/39] Accept call-ended reports from sim's call-report Modal function POST /internal/modal-calls/ended, not in the public OpenAPI document, takes a CallEndedReport with the Modal identity token of the report_call_ended function as a bearer token. The token is checked first (ModalIdentityVerifier, built from the expected Modal caller), so a caller without it gets 401 and learns nothing about the body; then the body is parsed (422 when it is not a report) and the outcome recorded with RunRecorder.record_reported_call, which SimulationDatabase now builds. No caller permission applies: the Modal token is the only proof. The answer is a CallReportAnswer: 200 for recorded or already_recorded; 503 with Retry-After for a status the sender should report again after (including unreachable Modal keys, and an instance without the expected Modal caller or a database); 404 for unknown runs and 409 for runs that hold another call, which the sender does not repeat. create_app's database seam and the route's two seams are grouped as RunServices, so the factory keeps five arguments. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/app.py | 37 +++- .../call_report_route.py | 147 ++++++++++++++ .../simulation_database.py | 19 ++ .../tests/test_call_report_route.py | 181 ++++++++++++++++++ .../tests/test_run_recording.py | 4 +- .../tests/test_simulation_database.py | 4 +- 6 files changed, 386 insertions(+), 6 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py create mode 100644 projects/policyengine-simulation-entry/tests/test_call_report_route.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 894394d50..82d360a01 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -10,6 +10,7 @@ import logging import uuid from contextlib import asynccontextmanager +from dataclasses import dataclass from typing import TYPE_CHECKING from uuid import UUID @@ -51,6 +52,11 @@ CallerAuthorizer, ) from policyengine_simulation_entry.backend import OldGatewayBackend, SimulationBackend +from policyengine_simulation_entry.call_report_route import ( + CallReportRecorder, + ModalCallerVerifier, + register_call_report_route, +) from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.entry_runtime import ( BACKEND_RESPONSE_HEADER, @@ -65,6 +71,7 @@ route_template, set_observability_header, ) +from policyengine_simulation_entry.modal_identity import ModalIdentityVerifier from policyengine_simulation_entry.schemas import ( SimulationDatabaseReadiness, SimulationEntrypointReadinessResponse, @@ -88,6 +95,7 @@ "STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS", "ComparisonBackend", "ResponseIdentifier", + "RunServices", "SimulationDatabaseConnection", "app", "create_app", @@ -97,6 +105,20 @@ SIMULATION_DATABASE_READINESS_TIMEOUT_SECONDS = 3.0 +@dataclass(frozen=True) +class RunServices: + """The simulation database and the call-ended route's seams, for tests. + + ``create_app`` builds each one it is not given: the database from + ``SIMULATION_DATABASE_URL``, the call recorder from that database, and the + Modal identity verifier from the expected Modal caller. + """ + + database: SimulationDatabaseConnection | None = None + call_recorder: CallReportRecorder | None = None + modal_identity_verifier: ModalCallerVerifier | None = None + + def _temporary_submission_payload( report: ComparisonReportRecord, ) -> TemporaryStage12SubmissionResponse: @@ -564,15 +586,23 @@ def create_app( backend: SimulationBackend | None = None, auth_dependency: AuthenticationDependency | None = None, comparison_backend: ComparisonBackend | None = None, - simulation_database: SimulationDatabaseConnection | None = None, + run_services: RunServices | None = None, ) -> FastAPI: """Build the app with injectable auth/backend seams for hermetic tests.""" runtime_settings = settings or Settings.from_env() - runtime_database = simulation_database + services = run_services or RunServices() + runtime_database = services.database if runtime_database is None and runtime_settings.simulation_database_url: runtime_database = SimulationDatabase.from_url( runtime_settings.simulation_database_url ) + call_recorder = services.call_recorder + if call_recorder is None and isinstance(runtime_database, SimulationDatabase): + call_recorder = runtime_database.call_recorder() + modal_identity_verifier = services.modal_identity_verifier + expected_modal_caller = runtime_settings.expected_modal_caller + if modal_identity_verifier is None and expected_modal_caller is not None: + modal_identity_verifier = ModalIdentityVerifier(expected_modal_caller) @asynccontextmanager async def lifespan(_: FastAPI) -> AsyncIterator[None]: @@ -635,6 +665,9 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: reader=simulation_reader, ) _register_service_routes(app, entry) + register_call_report_route( + app, verifier=modal_identity_verifier, recorder=call_recorder + ) return app diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py new file mode 100644 index 000000000..1673b25cd --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py @@ -0,0 +1,147 @@ +"""The route on which sim's call-report Modal function reports that a call ended. + +``POST /internal/modal-calls/ended`` (``CALL_ENDED_PATH``), not in the public +OpenAPI document. The request carries the Modal identity token of the +``report_call_ended`` function as a bearer token, and a ``CallEndedReport`` +naming the call and its runs, never a result. The token is checked first, so +a caller without it learns nothing about the body. Then the reported call's +outcome is recorded, if Modal has one (``RunRecorder.record_reported_call``). + +Answers, each a ``CallReportAnswer``: + +- 200: ``recorded`` or ``already_recorded``; a repeated report changes + nothing. +- 503 with ``Retry-After``: a status in ``RETRY_LATER_STATUSES``; the sender + reports again after a wait, and the safety check covers a call whose + sender gives up. +- 404 ``unknown_runs`` and 409 ``other_call``: the report cannot be recorded; + the sender stops. +- 401: the token is missing or refused; 422: the body is not a report. + +Only the revision that serves traffic receives these requests, since the +sender posts to sim's main URL. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import TYPE_CHECKING, Protocol + +from fastapi import Request, status +from fastapi.responses import JSONResponse +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + CallEndedReport, + CallReportAnswer, + CallReportStatus, +) +from pydantic import ValidationError + +from policyengine_simulation_entry.modal_identity import ( + ModalIdentityError, + ModalKeysUnavailableError, +) + +if TYPE_CHECKING: + from fastapi import FastAPI + + from policyengine_simulation_entry.modal_identity import ModalIdentity + +logger = logging.getLogger(__name__) + +# How long the sender waits before reporting again, in seconds. +RETRY_AFTER_SECONDS = 5 +_REFUSED_STATUSES = { + CallReportStatus.UNKNOWN_RUNS: status.HTTP_404_NOT_FOUND, + CallReportStatus.OTHER_CALL: status.HTTP_409_CONFLICT, +} +_BEARER_PREFIX = "bearer " + + +class CallReportRecorder(Protocol): + """Record a reported call's outcome (``RunRecorder``).""" + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + """Record the call's outcome if Modal has one, and say what was done.""" + ... + + +class ModalCallerVerifier(Protocol): + """Check a Modal identity token (``ModalIdentityVerifier``).""" + + def verify(self, token: str) -> ModalIdentity: + """Return the token's identity, or raise ModalIdentityError.""" + ... + + +def _answer(report_status: CallReportStatus) -> JSONResponse: + answer = CallReportAnswer(status=report_status) + if answer.retry_later: + return JSONResponse( + answer.model_dump(mode="json"), + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + headers={"Retry-After": str(RETRY_AFTER_SECONDS)}, + ) + return JSONResponse( + answer.model_dump(mode="json"), + status_code=_REFUSED_STATUSES.get(report_status, status.HTTP_200_OK), + ) + + +def _unauthorized() -> JSONResponse: + return JSONResponse( + {"detail": "A Modal identity token of the call-report function is required."}, + status_code=status.HTTP_401_UNAUTHORIZED, + headers={"WWW-Authenticate": "Bearer"}, + ) + + +def _bearer_token(request: Request) -> str | None: + header = request.headers.get("authorization", "") + if not header.lower().startswith(_BEARER_PREFIX): + return None + return header[len(_BEARER_PREFIX) :].strip() or None + + +def register_call_report_route( + app: FastAPI, + *, + verifier: ModalCallerVerifier | None, + recorder: CallReportRecorder | None, +) -> None: + """Add the call-ended route; without both seams it answers 503.""" + + @app.post(CALL_ENDED_PATH, include_in_schema=False) + async def report_call_ended(request: Request) -> JSONResponse: + if verifier is None or recorder is None: + return _answer(CallReportStatus.NOT_CONFIGURED) + token = _bearer_token(request) + if token is None: + return _unauthorized() + try: + identity = await asyncio.to_thread(verifier.verify, token) + except ModalKeysUnavailableError: + logger.warning("call_report_modal_keys_unavailable") + return _answer(CallReportStatus.MODAL_UNREACHABLE) + except ModalIdentityError as error: + logger.warning("call_report_refused", extra={"reason": str(error)}) + return _unauthorized() + try: + report = CallEndedReport.model_validate_json(await request.body()) + except ValidationError as error: + return JSONResponse( + {"detail": error.errors(include_url=False, include_input=False)}, + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ) + report_status = await recorder.record_reported_call(report) + logger.info( + "call_report_answered", + extra={ + "modal_call_id": report.modal_call_id, + "call_report_status": report_status.value, + "modal_app_name": identity.app_name, + "modal_container_id": identity.container_id, + }, + ) + return _answer(report_status) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index e0ca4f1bd..4401f8b7b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -6,10 +6,15 @@ from typing import TYPE_CHECKING from policyengine_simulation_persistence import ( + RunStore, + RunTransitions, create_simulation_engine, read_schema_revision, ) +from policyengine_simulation_entry.modal_calls import ModalCallInspector +from policyengine_simulation_entry.run_recording import ResultRecorders, RunRecorder + if TYPE_CHECKING: from sqlalchemy import Engine @@ -37,6 +42,20 @@ async def read_revision(self) -> str | None: """Return the applied migration revision, or None before the first one.""" return await asyncio.to_thread(self._read_revision) + def call_recorder(self) -> RunRecorder: + """Return what records a reported call's outcome on its runs. + + No kind of run has a result recorder yet (stage 2 registers them), so + a finished call is answered no_result_recorder and left for the + safety check; no route creates runs yet. + """ + return RunRecorder( + transitions=RunTransitions(self.engine), + runs=RunStore(self.engine), + inspector=ModalCallInspector(), + recorders=ResultRecorders(), + ) + def dispose(self) -> None: """Close every pooled connection.""" self.engine.dispose() diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_route.py b/projects/policyengine-simulation-entry/tests/test_call_report_route.py new file mode 100644 index 000000000..326d4dbf2 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_call_report_route.py @@ -0,0 +1,181 @@ +"""The route on which sim's call-report Modal function reports that a call ended.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from uuid import uuid4 + +import pytest +from conftest import FakeBackend, make_settings +from fastapi import HTTPException +from fastapi.testclient import TestClient +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + RETRY_LATER_STATUSES, + CallEndedReport, + CallReportStatus, +) +from policyengine_simulation_entry.app import RunServices, create_app +from policyengine_simulation_entry.modal_identity import ( + ModalIdentity, + ModalIdentityError, + ModalKeysUnavailableError, +) + +GOOD_TOKEN = "token-of-the-call-report-function" +RUN_ID = uuid4() +REPORT = {"modal_call_id": "fc-ended", "simulation_run_ids": [str(RUN_ID)]} + + +@dataclass +class _Verifier: + """Accept one token; refuse or fail on the others.""" + + keys_unavailable: bool = False + verified: list[str] = field(default_factory=list) + + def verify(self, token: str) -> ModalIdentity: + self.verified.append(token) + if self.keys_unavailable: + raise ModalKeysUnavailableError + if token != GOOD_TOKEN: + msg = "the token names another Modal function" + raise ModalIdentityError(msg) + return ModalIdentity( + app_name="policyengine-simulation-v2-py6-10-0", + function_name="report_call_ended", + container_id="ta-container0001", + ) + + +@dataclass +class _Recorder: + """Answer every report with one status, and keep the reports.""" + + status: CallReportStatus = CallReportStatus.RECORDED + reports: list[CallEndedReport] = field(default_factory=list) + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + self.reports.append(report) + return self.status + + +def _reject_every_caller(): + raise HTTPException(status_code=403) + + +def _client(verifier=None, recorder=None) -> TestClient: + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=_reject_every_caller, + run_services=RunServices( + call_recorder=recorder, modal_identity_verifier=verifier + ), + ) + ) + + +def _post(client, body=REPORT, token=GOOD_TOKEN): + headers = {} if token is None else {"Authorization": f"Bearer {token}"} + return client.post(CALL_ENDED_PATH, json=body, headers=headers) + + +def test_a_report_from_the_call_report_function_is_recorded() -> None: + verifier, recorder = _Verifier(), _Recorder() + + response = _post(_client(verifier, recorder)) + + assert response.status_code == 200 + assert response.json() == {"status": "recorded"} + assert verifier.verified == [GOOD_TOKEN] + assert recorder.reports == [CallEndedReport.model_validate(REPORT)] + + +def test_a_repeated_report_is_answered_as_already_recorded() -> None: + recorder = _Recorder(status=CallReportStatus.ALREADY_RECORDED) + + response = _post(_client(_Verifier(), recorder)) + + assert response.status_code == 200 + assert response.json() == {"status": "already_recorded"} + + +@pytest.mark.parametrize("status", sorted(RETRY_LATER_STATUSES - {"not_configured"})) +def test_a_report_sim_cannot_record_yet_is_answered_with_a_wait(status) -> None: + response = _post(_client(_Verifier(), _Recorder(status=CallReportStatus(status)))) + + assert response.status_code == 503 + assert response.json() == {"status": status} + assert int(response.headers["Retry-After"]) > 0 + + +@pytest.mark.parametrize( + ("status", "http_status"), + [(CallReportStatus.UNKNOWN_RUNS, 404), (CallReportStatus.OTHER_CALL, 409)], +) +def test_a_report_that_cannot_be_recorded_is_refused(status, http_status) -> None: + response = _post(_client(_Verifier(), _Recorder(status=status))) + + assert response.status_code == http_status + assert response.json() == {"status": status.value} + assert "Retry-After" not in response.headers + + +@pytest.mark.parametrize("token", [None, "token-of-a-policy-worker"]) +def test_a_request_without_the_functions_token_is_refused(token) -> None: + recorder = _Recorder() + + response = _post(_client(_Verifier(), recorder), token=token) + + assert response.status_code == 401 + assert recorder.reports == [] + + +def test_an_unauthenticated_request_learns_nothing_about_its_body() -> None: + response = _post(_client(_Verifier(), _Recorder()), body={"x": 1}, token=None) + + assert response.status_code == 401 + + +def test_a_malformed_report_is_refused_after_authentication() -> None: + recorder = _Recorder() + + response = _post( + _client(_Verifier(), recorder), + body={"modal_call_id": "fc-ended", "simulation_run_ids": [], "result": {}}, + ) + + assert response.status_code == 422 + assert recorder.reports == [] + + +def test_unreachable_modal_keys_are_answered_with_a_wait() -> None: + recorder = _Recorder() + + response = _post(_client(_Verifier(keys_unavailable=True), recorder)) + + assert response.status_code == 503 + assert response.json() == {"status": "modal_unreachable"} + assert recorder.reports == [] + + +@pytest.mark.parametrize( + ("verifier", "recorder"), + [(None, _Recorder()), (_Verifier(), None)], + ids=["no-modal-caller", "no-database"], +) +def test_an_instance_not_configured_for_reports_answers_with_a_wait( + verifier, recorder +) -> None: + response = _post(_client(verifier, recorder)) + + assert response.status_code == 503 + assert response.json() == {"status": "not_configured"} + + +def test_the_route_is_not_in_the_public_openapi_document() -> None: + document = _client(_Verifier(), _Recorder()).get("/openapi.json").json() + + assert CALL_ENDED_PATH not in document["paths"] diff --git a/projects/policyengine-simulation-entry/tests/test_run_recording.py b/projects/policyengine-simulation-entry/tests/test_run_recording.py index f646cfca9..a2408ea8a 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_recording.py +++ b/projects/policyengine-simulation-entry/tests/test_run_recording.py @@ -10,7 +10,7 @@ from conftest import FakeBackend, make_settings from fastapi.testclient import TestClient from policyengine_simulation_entry import run_recording -from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.app import RunServices, create_app from policyengine_simulation_entry.modal_calls import ( CallFailed, CallInspectionError, @@ -396,7 +396,7 @@ def test_an_instance_runs_nothing_in_the_background() -> None: ), backend=FakeBackend(), auth_dependency=lambda: None, - simulation_database=_DatabaseThatOnlyReadsItsRevision(), + run_services=RunServices(database=_DatabaseThatOnlyReadsItsRevision()), ) with TestClient(app) as client: diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_database.py b/projects/policyengine-simulation-entry/tests/test_simulation_database.py index 95b1defae..c5fd00f91 100644 --- a/projects/policyengine-simulation-entry/tests/test_simulation_database.py +++ b/projects/policyengine-simulation-entry/tests/test_simulation_database.py @@ -9,7 +9,7 @@ from fastapi.testclient import TestClient from policyengine_simulation_entry import app as app_module from policyengine_simulation_entry import simulation_database as database_module -from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.app import RunServices, create_app from policyengine_simulation_entry.simulation_database import SimulationDatabase from policyengine_simulation_persistence import HEAD_REVISION from sqlalchemy.exc import OperationalError @@ -43,7 +43,7 @@ def _client(backend, database): settings=make_settings(simulation_database_url=DATABASE_URL), backend=backend, auth_dependency=lambda: None, - simulation_database=database, + run_services=RunServices(database=database), ) ) From 08c0bd47a569d9fcb267fce2c45efaaafd89a018 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:17:03 +0400 Subject: [PATCH 27/39] Add the safety check: one recording pass per execution of a job A Modal call reports when it ends; the safety check covers a call whose report never arrived. python -m policyengine_simulation_entry.safety_check runs one RunRecorder pass and exits: it fails runs and comparison steps left dispatching for ten minutes, and records what Modal reports for running calls dispatched, or last checked, at least one interval ago. It exits 0 after a pass, 1 when the pass raised (the next execution tries again) and 2 without SIMULATION_DATABASE_URL; only a signal propagates. SAFETY_CHECK_INTERVAL (five minutes) is the schedule's interval, the claim's lease, so each execution checks a running call once, and the wait before a new call's first check, which the dispatcher now sets. Co-Authored-By: Claude Opus 5.5 --- .../run_dispatch.py | 8 +- .../run_recording.py | 11 +- .../safety_check.py | 90 ++++++++++++++ .../tests/test_safety_check.py | 110 ++++++++++++++++++ 4 files changed, 213 insertions(+), 6 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py create mode 100644 projects/policyengine-simulation-entry/tests/test_safety_check.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index 7bb3524ea..af392ffd8 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -17,13 +17,14 @@ import asyncio import logging from dataclasses import dataclass -from datetime import timedelta from enum import StrEnum from typing import TYPE_CHECKING, Protocol from policyengine_simulation_persistence import RunError, RunErrorCode from policyengine_simulation_persistence.run_store import utc_now +from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL + if TYPE_CHECKING: from collections.abc import Callable, Collection from datetime import datetime @@ -39,8 +40,9 @@ def _event_name(prefix: str, outcome: str) -> str: return f"{prefix}_{outcome}" -# The first check of a new call; a recording pass backs off from there. -FIRST_CHECK_DELAY = timedelta(seconds=2) +# A call reports when it ends; the safety check first looks at one that has +# not reported after this long. +FIRST_CHECK_DELAY = SAFETY_CHECK_INTERVAL class DispatchOutcome(StrEnum): diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 72a64d83d..395f23153 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -68,9 +68,14 @@ def _event_name(prefix: str, outcome: str) -> str: return f"{prefix}_{outcome}" -# A claimed call is not checked by another instance for this long; if the -# claiming instance stops, the call becomes due again afterwards. -CHECK_LEASE = timedelta(minutes=1) +# Cloud Scheduler starts the safety check this often (the deploy's +# schedule): a call that has not reported is checked about this long after +# its dispatch, then once per interval while it runs. +SAFETY_CHECK_INTERVAL = timedelta(minutes=5) +# A claimed call is not checked by another pass for this long, so each +# execution of the safety check checks a running call once; if the pass +# stops, the call becomes due again afterwards. +CHECK_LEASE = SAFETY_CHECK_INTERVAL # A run left dispatching this long was claimed by an instance that stopped. DISPATCH_ABANDONED_AFTER = timedelta(minutes=10) CLAIM_BATCH_SIZE = 50 diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py new file mode 100644 index 000000000..378b123ef --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py @@ -0,0 +1,90 @@ +"""The safety check: one recording pass, run as a Cloud Run job, then exit. + +A Modal call reports when it ends, and sim records its outcome then +(``call_report_route``). The safety check covers a call whose report never +arrived: Cloud Scheduler starts this job every ``SAFETY_CHECK_INTERVAL``, and +each execution runs one ``RunRecorder`` pass, which + +- fails runs and comparison steps left ``dispatching`` for ten minutes (sim + stopped between claiming them and recording their Modal call), so they can + be submitted again, and pending comparison steps whose run failed; +- asks Modal about running calls first dispatched, or last checked, at least + one interval ago, and records what it reports, with the same conditional + writes as a report. + +The job runs only the image of the revision that serves traffic: the deploy +updates it when that revision is promoted, and a rollback moves it back, so +a rejected candidate never records anything. Run it with +``python -m policyengine_simulation_entry.safety_check``; it reads +``SIMULATION_DATABASE_URL`` and sim's Modal token (``MODAL_TOKEN_ID``, +``MODAL_TOKEN_SECRET``). +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import sys +from dataclasses import asdict +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_entry.modal_calls import OWN_PROCESS_SIGNALS +from policyengine_simulation_entry.simulation_database import SimulationDatabase + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from policyengine_simulation_entry.run_recording import RecordingPass + +logger = logging.getLogger(__name__) + +EXIT_PASSED = 0 +EXIT_PASS_FAILED = 1 +EXIT_NOT_CONFIGURED = 2 + + +class DueOutcomeRecorder(Protocol): + """Run one recording pass (``RunRecorder``).""" + + async def record_due_outcomes(self) -> RecordingPass: + """Fail abandoned dispatches, then check every due call.""" + ... + + +def _database_recorder(database_url: str) -> DueOutcomeRecorder: + return SimulationDatabase.from_url(database_url).call_recorder() + + +def main( + *, + environment: Mapping[str, str] | None = None, + build_recorder: Callable[[str], DueOutcomeRecorder] = _database_recorder, +) -> int: + """Run one pass and return the process exit status. + + Only a signal to this process propagates; any other error fails the + execution, and the next scheduled execution tries again. + """ + values = os.environ if environment is None else environment + database_url = values.get("SIMULATION_DATABASE_URL", "") + if not database_url: + logger.error("safety_check_not_configured") + print("SIMULATION_DATABASE_URL is not set.", file=sys.stderr) # noqa: T201 # the job's own error output + return EXIT_NOT_CONFIGURED + try: + recording = asyncio.run(build_recorder(database_url).record_due_outcomes()) + except BaseException as error: + if isinstance(error, OWN_PROCESS_SIGNALS): + raise + logger.exception( + "safety_check_pass_failed", extra={"error_type": type(error).__name__} + ) + return EXIT_PASS_FAILED + logger.info("safety_check_pass", extra=asdict(recording)) + return EXIT_PASSED + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + sys.exit(main()) diff --git a/projects/policyengine-simulation-entry/tests/test_safety_check.py b/projects/policyengine-simulation-entry/tests/test_safety_check.py new file mode 100644 index 000000000..20810acbd --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_safety_check.py @@ -0,0 +1,110 @@ +"""The safety check: one recording pass per execution of a Cloud Run job.""" + +from __future__ import annotations + +import subprocess +import sys +from datetime import timedelta + +import pytest +from policyengine_simulation_entry import safety_check +from policyengine_simulation_entry.run_dispatch import FIRST_CHECK_DELAY +from policyengine_simulation_entry.run_recording import ( + CHECK_LEASE, + SAFETY_CHECK_INTERVAL, + RecordingPass, +) + +DATABASE_URL = "postgresql://runtime:secret@db.example/postgres" + + +class _OnePass: + """A recorder whose pass returns a scripted result, or raises.""" + + def __init__(self, result) -> None: + self.result = result + self.passes = 0 + + async def record_due_outcomes(self) -> RecordingPass: + self.passes += 1 + if isinstance(self.result, BaseException): + raise self.result + return self.result + + +def _run(environment, recorder): + built_for = [] + + def build(database_url): + built_for.append(database_url) + return recorder + + status = safety_check.main(environment=environment, build_recorder=build) + return status, built_for + + +def test_each_execution_runs_one_pass_and_exits() -> None: + recorder = _OnePass(RecordingPass(checked=4, succeeded=1, failed=1)) + + status, built_for = _run({"SIMULATION_DATABASE_URL": DATABASE_URL}, recorder) + + assert status == 0 + assert recorder.passes == 1 + assert built_for == [DATABASE_URL] + + +def test_the_pass_is_logged_with_what_it_found(caplog) -> None: + caplog.set_level("INFO") + recorder = _OnePass(RecordingPass(checked=4, succeeded=1, abandoned=2)) + + _run({"SIMULATION_DATABASE_URL": DATABASE_URL}, recorder) + + (record,) = [ + record for record in caplog.records if record.message == "safety_check_pass" + ] + assert record.checked == 4 + assert record.succeeded == 1 + assert record.abandoned == 2 + + +def test_an_execution_without_the_database_url_fails_without_a_pass() -> None: + recorder = _OnePass(RecordingPass()) + + status, built_for = _run({}, recorder) + + assert status == safety_check.EXIT_NOT_CONFIGURED + assert built_for == [] + + +def test_a_pass_that_raises_fails_the_execution() -> None: + status, _ = _run( + {"SIMULATION_DATABASE_URL": DATABASE_URL}, _OnePass(RuntimeError("down")) + ) + + assert status == safety_check.EXIT_PASS_FAILED + + +def test_a_signal_to_the_job_is_not_swallowed() -> None: + with pytest.raises(KeyboardInterrupt): + _run({"SIMULATION_DATABASE_URL": DATABASE_URL}, _OnePass(KeyboardInterrupt())) + + +def test_the_job_runs_the_module_as_its_command() -> None: + """The deploy runs ``python -m policyengine_simulation_entry.safety_check``.""" + result = subprocess.run( + [sys.executable, "-m", "policyengine_simulation_entry.safety_check"], + check=False, + capture_output=True, + text=True, + env={"PATH": "/usr/bin:/bin"}, + ) + + assert result.returncode == safety_check.EXIT_NOT_CONFIGURED + assert "SIMULATION_DATABASE_URL" in result.stderr + + +def test_a_call_is_first_checked_one_interval_after_dispatch() -> None: + """A call reports when it ends; the safety check looks only after a while.""" + assert timedelta(minutes=5) == SAFETY_CHECK_INTERVAL + assert FIRST_CHECK_DELAY == SAFETY_CHECK_INTERVAL + assert CHECK_LEASE == SAFETY_CHECK_INTERVAL From 38e8d997a7c1d1c1eb7153c52df67f3d37c91d6a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:19:39 +0400 Subject: [PATCH 28/39] Deploy the entrypoint with the Modal caller its reports must come from The deploy step passes the beta/prod GitHub environment variable MODAL_WORKSPACE_ID to the deploy script, which requires a Modal workspace ID and writes it, with the deployment's own Modal environment (MODAL_ENVIRONMENT: staging or main), to the runtime as SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID and SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT. Every deploy stops until an operator sets the variable. The deployment notes say how the route checks a report and that Modal needs no Google credentials, since the service already allows unauthenticated requests through Cloud Run. Co-Authored-By: Claude Opus 5.5 --- .../deploy-cloud-run-simulation-entry.sh | 7 ++ .../workflows/simulation-deploy.reusable.yml | 1 + .../DEPLOYMENT.md | 13 ++++ .../tests/test_call_report_deployment.py | 76 +++++++++++++++++++ .../tests/test_deployment_assets_cloud_run.py | 1 + 5 files changed, 98 insertions(+) create mode 100644 projects/policyengine-simulation-entry/tests/test_call_report_deployment.py diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index 36b1899ff..45784c03a 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -102,6 +102,7 @@ required_variables=( SIMULATION_DATABASE_URL_SECRET_NAME MODAL_TOKEN_ID_SECRET_NAME MODAL_TOKEN_SECRET_SECRET_NAME + MODAL_WORKSPACE_ID_VALUE OBSERVABILITY_SERVICE_NAMESPACE OBSERVABILITY_TRACE_PROJECT_ID OTEL_EXPORTER_OTLP_ENDPOINT @@ -128,6 +129,10 @@ esac || fail "TAG is invalid" [[ "${ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ || fail "ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT is invalid" +# The Modal workspace whose call-report function may report that calls +# ended (the entrypoint also checks the environment, MODAL_ENVIRONMENT). +[[ "${MODAL_WORKSPACE_ID_VALUE}" =~ ^ac-[A-Za-z0-9]+$ ]] \ + || fail "MODAL_WORKSPACE_ID_VALUE must be a Modal workspace ID" [[ "${ENTRYPOINT_MIN_INSTANCES}" =~ ^[0-9]+$ ]] \ || fail "ENTRYPOINT_MIN_INSTANCES must be a non-negative integer" [[ "${ENTRYPOINT_MAX_INSTANCES}" =~ ^[0-9]+$ ]] \ @@ -191,6 +196,8 @@ jq -n ' OTEL_METRICS_EXPORTER: "otlp", OTEL_TRACES_SAMPLER_ARG: "1.0", POLICYENGINE_OTEL_GOOGLE_AUDIENCE: env.POLICYENGINE_OTEL_GOOGLE_AUDIENCE, + SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID: env.MODAL_WORKSPACE_ID_VALUE, + SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT: env.MODAL_ENVIRONMENT, STAGE12_ENABLED: ( if env.DEPLOY_STAGE12_V2 == "true" then env.STAGE12_ENABLED_VALUE else "0" end ) diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index ea4f48585..495ad9d8a 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -279,6 +279,7 @@ jobs: STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} MODAL_TOKEN_ID_SECRET_NAME: ${{ vars.MODAL_TOKEN_ID_SECRET_NAME }} MODAL_TOKEN_SECRET_SECRET_NAME: ${{ vars.MODAL_TOKEN_SECRET_SECRET_NAME }} + MODAL_WORKSPACE_ID_VALUE: ${{ vars.MODAL_WORKSPACE_ID }} run: .github/scripts/deploy-cloud-run-simulation-entry.sh - name: Resolve exact candidate metadata diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index dbe4d5461..140087530 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -68,6 +68,19 @@ environment variables `MODAL_TOKEN_ID_SECRET_NAME` and service account must be able to read both secrets. The Stage 12 comparison path's Modal calls use the same token until it is removed. +A Modal call reports that it ended on `POST /internal/modal-calls/ended` +(`policyengine_simulation_entry.call_report_route`), with the Modal identity +token of the `report_call_ended` function as a bearer token. The entrypoint +checks the token against Modal's published keys and accepts only that +function, in a v2 app, of the workspace in the `beta`/`prod` environment +variable `MODAL_WORKSPACE_ID` and of the deployment's own Modal environment; +the deploy passes both to the runtime as +`SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID` and +`SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT`, and stops when the workspace ID is +missing. The service already allows unauthenticated requests through Cloud +Run (`--allow-unauthenticated`): callers prove who they are to the +entrypoint, so Modal needs no Google credentials. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py b/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py new file mode 100644 index 000000000..9b428dfbb --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py @@ -0,0 +1,76 @@ +"""Deploy settings for the route on which Modal reports that a call ended. + +The entrypoint accepts reports only from the call-report function of one +Modal workspace and environment. The workspace ID comes from the `beta` and +`prod` GitHub environment variable `MODAL_WORKSPACE_ID`; the environment is +the deployment's own Modal environment. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets_cloud_run import cloud_run_deploy_environment +from test_deployment_assets_support import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + workflow_jobs, +) + + +def _deploy(tmp_path, **overrides): + environment = cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + capture = Path(environment["GCLOUD_ENV_FILE_CAPTURE"]) + runtime = ( + json.loads(capture.read_text(encoding="utf-8")) if capture.exists() else {} + ) + return result, runtime + + +@pytest.mark.parametrize( + ("app_environment", "modal_environment"), + [("staging", "staging"), ("production", "main")], +) +def test_the_entrypoint_expects_reports_from_its_own_modal_environment( + tmp_path, app_environment, modal_environment +) -> None: + result, runtime = _deploy( + tmp_path, + APP_ENVIRONMENT=app_environment, + MODAL_ENVIRONMENT=modal_environment, + MODAL_WORKSPACE_ID_VALUE="ac-workspace0001", + ) + + assert result.returncode == 0, result.stderr + assert runtime["SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID"] == "ac-workspace0001" + assert runtime["SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT"] == modal_environment + + +@pytest.mark.parametrize("workspace_id", ["", "workspace0001", "ac-", "ac-a b"]) +def test_a_deploy_without_a_modal_workspace_id_stops(tmp_path, workspace_id) -> None: + result, runtime = _deploy(tmp_path, MODAL_WORKSPACE_ID_VALUE=workspace_id) + + assert result.returncode != 0 + assert "MODAL_WORKSPACE_ID_VALUE" in result.stderr + assert runtime == {} + + +def test_the_deploy_step_passes_the_modal_workspace_id() -> None: + deploy_entrypoint = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + environment = next( + step["env"] + for step in deploy_entrypoint["steps"] + if "deploy-cloud-run-simulation-entry.sh" in step.get("run", "") + ) + + assert environment["MODAL_WORKSPACE_ID_VALUE"] == "${{ vars.MODAL_WORKSPACE_ID }}" diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py index a6d870841..bc251481f 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets_cloud_run.py @@ -116,6 +116,7 @@ def cloud_run_deploy_environment(tmp_path: Path) -> dict[str, str]: "STAGE12_DATABASE_URL_SECRET_NAME": "stage12-staging-database-url", "MODAL_TOKEN_ID_SECRET_NAME": "modal-token-id", "MODAL_TOKEN_SECRET_SECRET_NAME": "modal-token-secret", + "MODAL_WORKSPACE_ID_VALUE": "ac-workspace0001", "SIMULATION_DATABASE_URL_SECRET_NAME": ( "simulation-staging-runtime-database-url" ), From 91e0f99f19ad6169bb2ddd31b4ad7fff33bb842f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:23:40 +0400 Subject: [PATCH 29/39] Deploy the safety check from the promoted revision, on a schedule deploy-cloud-run-simulation-safety-check.sh points the Cloud Run job -safety-check at one entrypoint revision: it copies the revision's image digest and service account, the secret versions it pins for SIMULATION_DATABASE_URL, MODAL_TOKEN_ID and MODAL_TOKEN_SECRET, and its plain telemetry settings (not the old-gateway secret), runs python -m policyengine_simulation_entry.safety_check as one task with no retries, and creates or updates a Cloud Scheduler entry that starts it every five minutes with the scheduler's own service account. A revision of another service is refused; a revision without the simulation database predates the safety check, so the schedule is paused. Only the promotion job (after the promoted revision passes every check, and back to the previous revision if this step fails) and the rollback workflow (after it restores traffic) call it, so a rejected candidate never records anything. The scheduler's service account comes from the new beta/prod variable SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT. The deployment notes list the one-time operator steps; staging has no safety check until it promotes its candidates (plan item 1.39). Co-Authored-By: Claude Opus 5.5 --- ...eploy-cloud-run-simulation-safety-check.sh | 108 +++++++ .../workflows/simulation-deploy.reusable.yml | 16 + .github/workflows/simulation-rollback.yml | 7 + .../DEPLOYMENT.md | 38 +++ .../tests/test_safety_check_deployment.py | 288 ++++++++++++++++++ 5 files changed, 457 insertions(+) create mode 100755 .github/scripts/deploy-cloud-run-simulation-safety-check.sh create mode 100644 projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py diff --git a/.github/scripts/deploy-cloud-run-simulation-safety-check.sh b/.github/scripts/deploy-cloud-run-simulation-safety-check.sh new file mode 100755 index 000000000..c50e3d7d3 --- /dev/null +++ b/.github/scripts/deploy-cloud-run-simulation-safety-check.sh @@ -0,0 +1,108 @@ +#!/usr/bin/env bash +# Point the safety-check Cloud Run job at the entrypoint revision that serves +# traffic, and keep its Cloud Scheduler entry. +# +# The job runs `python -m policyengine_simulation_entry.safety_check`: one +# recording pass per execution. It is updated only from a revision that was +# promoted (or restored by a rollback), so a rejected candidate never records +# anything. It copies the revision's image digest, service account, Modal +# token and database URL (at the secret versions the revision pins) and its +# plain telemetry settings; nothing else. A revision without the simulation +# database URL predates the safety check, so its schedule is paused. +# +# Environment: PROJECT_ID, REGION, SERVICE, SAFETY_CHECK_REVISION, +# SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT, RUNNER_TEMP. +set -euo pipefail + +fail() { + printf '%s\n' "$*" >&2 + exit 1 +} + +# Every five minutes: SAFETY_CHECK_INTERVAL in +# policyengine_simulation_entry.run_recording. +SCHEDULE="*/5 * * * *" +SECRETS=(SIMULATION_DATABASE_URL MODAL_TOKEN_ID MODAL_TOKEN_SECRET) +PLAIN_VARIABLES_FILTER='^(APP_ENVIRONMENT|CLOUD_RUN_REGION|OBSERVABILITY_[A-Z_]+|OTEL_[A-Z_]+|POLICYENGINE_OTEL_GOOGLE_AUDIENCE)$' + +for variable_name in PROJECT_ID REGION SERVICE SAFETY_CHECK_REVISION \ + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT RUNNER_TEMP; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" +done +[[ "${PROJECT_ID}" =~ ^[a-z][a-z0-9-]{4,28}[a-z0-9]$ ]] || fail "PROJECT_ID is invalid" +[[ "${REGION}" =~ ^[a-z]+-[a-z]+[0-9]+$ ]] || fail "REGION is invalid" +[[ "${SERVICE}" =~ ^[a-z]([a-z0-9-]{0,47}[a-z0-9])?$ ]] || fail "SERVICE is invalid" +[[ "${SAFETY_CHECK_REVISION}" =~ ^[a-z]([a-z0-9-]{0,61}[a-z0-9])?$ ]] \ + || fail "SAFETY_CHECK_REVISION is invalid" +[[ "${SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ + || fail "SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT is invalid" + +job="${SERVICE}-safety-check" +location=(--project "${PROJECT_ID}" --region "${REGION}") +scheduler_location=(--project "${PROJECT_ID}" --location "${REGION}") + +revision_json="$(gcloud run revisions describe "${SAFETY_CHECK_REVISION}" "${location[@]}" --format=json)" +revision_service="$(jq -r '.metadata.labels["serving.knative.dev/service"] // ""' <<<"${revision_json}")" +[[ "${revision_service}" == "${SERVICE}" ]] \ + || fail "${SAFETY_CHECK_REVISION} belongs to ${revision_service:-no service}, not ${SERVICE}" + +if ! jq -e '.spec.containers[0].env // [] | any(.name == "SIMULATION_DATABASE_URL")' >/dev/null <<<"${revision_json}"; then + if gcloud scheduler jobs describe "${job}" "${scheduler_location[@]}" >/dev/null 2>&1; then + gcloud scheduler jobs pause "${job}" "${scheduler_location[@]}" --quiet + fi + printf '%s has no simulation database, so the safety check is paused.\n' "${SAFETY_CHECK_REVISION}" + exit 0 +fi + +image="$(jq -er '.status.imageDigest // .spec.containers[0].image' <<<"${revision_json}")" +service_account="$(jq -er '.spec.serviceAccountName' <<<"${revision_json}")" + +job_secrets=() +for secret_variable in "${SECRETS[@]}"; do + reference="$(jq -er --arg name "${secret_variable}" ' + [.spec.containers[0].env[] | select(.name == $name) | .valueFrom.secretKeyRef + | select(.name != null and .key != null) | "\(.name):\(.key)"] + | if length == 1 then .[0] else error("not one secret reference") end + ' <<<"${revision_json}")" || fail "${SAFETY_CHECK_REVISION} has no pinned secret for ${secret_variable}" + job_secrets+=("${secret_variable}=${reference}") +done +job_secrets_csv="$(IFS=,; printf '%s' "${job_secrets[*]}")" + +environment_file="${RUNNER_TEMP}/simulation-safety-check-environment.yaml" +trap 'rm -f "${environment_file}"' EXIT +jq --arg filter "${PLAIN_VARIABLES_FILTER}" ' + [.spec.containers[0].env[] | select(.value != null and (.name | test($filter))) + | {(.name): .value}] | add // {} +' <<<"${revision_json}" > "${environment_file}" + +gcloud run jobs deploy "${job}" "${location[@]}" \ + --image "${image}" \ + --service-account "${service_account}" \ + --set-secrets "${job_secrets_csv}" \ + --env-vars-file "${environment_file}" \ + --command /opt/venv/bin/python \ + --args=-m,policyengine_simulation_entry.safety_check \ + --tasks 1 \ + --max-retries 0 \ + --task-timeout 600s \ + --cpu 1 \ + --memory 512Mi \ + --labels "serving-revision=${SAFETY_CHECK_REVISION}" \ + --quiet + +schedule_action=create +if gcloud scheduler jobs describe "${job}" "${scheduler_location[@]}" >/dev/null 2>&1; then + schedule_action=update +fi +gcloud scheduler jobs "${schedule_action}" http "${job}" "${scheduler_location[@]}" \ + --schedule "${SCHEDULE}" \ + --time-zone Etc/UTC \ + --uri "https://${REGION}-run.googleapis.com/apis/run.googleapis.com/v1/namespaces/${PROJECT_ID}/jobs/${job}:run" \ + --http-method POST \ + --oauth-service-account-email "${SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT}" \ + --attempt-deadline 60s \ + --quiet +if [[ "${schedule_action}" == "update" ]]; then + gcloud scheduler jobs resume "${job}" "${scheduler_location[@]}" --quiet +fi +printf 'The safety check runs %s from %s.\n' "${image}" "${SAFETY_CHECK_REVISION}" diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 495ad9d8a..24d2b898b 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -1012,6 +1012,15 @@ jobs: SIMULATION_ENTRYPOINT_TEST_ENFORCE_PERMISSIONS: ${{ vars.SIMULATION_ENTRYPOINT_ENFORCE_PERMISSIONS }} run: uv run --frozen pytest authenticated_tests/ -v + # Only a verified, promoted revision records outcomes in the safety + # check, so a rejected candidate never records anything. + - name: Point the safety check at the promoted revision + id: safety_check + env: + SAFETY_CHECK_REVISION: ${{ env.TARGET_REVISION }} + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + - name: Restore previous revision after failed verification if: ${{ failure() && steps.promote.outcome == 'success' }} env: @@ -1022,6 +1031,13 @@ jobs: SIMULATION_ENTRYPOINT_EXPECTED_CURRENT_REVISION: ${{ env.TARGET_REVISION }} run: .github/scripts/set-cloud-run-simulation-entry-revision.sh + - name: Point the safety check back at the previous revision + if: ${{ failure() && steps.safety_check.outcome == 'failure' && env.PREVIOUS_REVISION != '' }} + env: + SAFETY_CHECK_REVISION: ${{ env.PREVIOUS_REVISION }} + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + finalize_routing: name: Finalize the Modal routing publication # Finish the durable journal after every gate, including failed uploads diff --git a/.github/workflows/simulation-rollback.yml b/.github/workflows/simulation-rollback.yml index 10054ea5d..46c165372 100644 --- a/.github/workflows/simulation-rollback.yml +++ b/.github/workflows/simulation-rollback.yml @@ -218,6 +218,13 @@ jobs: .github/scripts/set-cloud-run-simulation-entry-revision.sh fi + - name: Point the safety check at the restored revision + if: ${{ inputs.entrypoint_revision != '' }} + env: + SAFETY_CHECK_REVISION: ${{ inputs.entrypoint_revision }} + SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + - name: Verify the stable service serves the restored release if: ${{ !inputs.resolve_pending_publication }} env: diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index 140087530..e910794df 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -81,6 +81,44 @@ missing. The service already allows unauthenticated requests through Cloud Run (`--allow-unauthenticated`): callers prove who they are to the entrypoint, so Modal needs no Google credentials. +## Safety check + +A call whose report never arrives is recorded by the safety check, the Cloud +Run job `-safety-check`. Each execution runs one recording +pass (`python -m policyengine_simulation_entry.safety_check`) and exits: it +fails runs left `dispatching` for ten minutes and records what Modal reports +for calls dispatched, or last checked, at least five minutes ago. Cloud +Scheduler starts it every five minutes (`*/5 * * * *`, UTC): a lost report is +recorded within about ten minutes, and each execution is a single short +container start, so the job costs little when every call reports. + +`.github/scripts/deploy-cloud-run-simulation-safety-check.sh` points the job +at one entrypoint revision. It copies that revision's image digest, service +account, the secret versions it pins for `SIMULATION_DATABASE_URL`, +`MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`, and its plain telemetry settings, +then creates or updates the schedule. Only two places call it, so the job +only ever runs a revision that serves traffic and a rejected candidate never +records anything: + +- the promotion job, after the promoted revision passes every check; if this + step fails, the previous revision is restored and the job is pointed back + at it; +- the rollback workflow, after it restores entrypoint traffic. A restored + revision without the simulation database predates the safety check, so the + schedule is paused. + +`beta` does not promote its candidate yet, so staging has no safety check +until it does (plan item 1.39). + +One-time operator steps, in each environment: enable the Cloud Scheduler +API; create the scheduler's service account, give it `roles/run.invoker` +on the job (or the project), and set its email in the GitHub environment +variable `SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT`; give the deploy +service account permission to deploy Cloud Run jobs and manage Cloud +Scheduler jobs, and `iam.serviceAccounts.actAs` on the scheduler's service +account. The job uses the entrypoint's runtime service account, which can +already read its secrets. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py b/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py new file mode 100644 index 000000000..f031bf68c --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py @@ -0,0 +1,288 @@ +"""Deploying the safety-check job from the revision that serves traffic. + +The job runs only the image of the promoted (or restored) entrypoint +revision, so a rejected candidate never records anything. The deploy copies +that revision's image, service account, database URL and Modal token, and +keeps a Cloud Scheduler entry that starts the job every five minutes. +""" + +from __future__ import annotations + +import json +from datetime import timedelta + +import pytest +from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL +from test_deployment_assets_support import ( + REUSABLE_DEPLOY_WORKFLOW, + ROLLBACK_WORKFLOW, + workflow_jobs, +) +from workflow_script_support import FakeCommands, run_script + +SCRIPT = "deploy-cloud-run-simulation-safety-check.sh" +SERVICE = "policyengine-simulation-entry" +REVISION = "policyengine-simulation-entry-p-123" +IMAGE = ( + "us-central1-docker.pkg.dev/policyengine/policyengine-simulation-entry/" + "policyengine-simulation-entry@sha256:" + "a" * 64 +) +RUNTIME_ACCOUNT = "sim-entry-prod-runtime@policyengine.iam.gserviceaccount.com" +SCHEDULER_ACCOUNT = "sim-safety-check-prod@policyengine.iam.gserviceaccount.com" +JOB = f"{SERVICE}-safety-check" + + +def _secret(name, secret, version): + return { + "name": name, + "valueFrom": {"secretKeyRef": {"name": secret, "key": version}}, + } + + +def _revision(*, service=SERVICE, environment=None): + return { + "apiVersion": "serving.knative.dev/v1", + "kind": "Revision", + "metadata": { + "name": REVISION, + "labels": {"serving.knative.dev/service": service}, + }, + "spec": { + "serviceAccountName": RUNTIME_ACCOUNT, + "containers": [ + { + "image": "us-central1-docker.pkg.dev/policyengine/x/y:sha", + "env": environment + if environment is not None + else [ + {"name": "APP_ENVIRONMENT", "value": "production"}, + {"name": "OLD_GATEWAY_URL", "value": "https://legacy.example"}, + {"name": "OBSERVABILITY_SERVICE_NAMESPACE", "value": "sim"}, + {"name": "CLOUD_RUN_REGION", "value": "us-central1"}, + _secret( + "SIMULATION_DATABASE_URL", "simulation-runtime-url", "4" + ), + _secret("MODAL_TOKEN_ID", "modal-token-id", "3"), + _secret("MODAL_TOKEN_SECRET", "modal-token-secret", "3"), + _secret( + "OLD_GATEWAY_AUTH_CLIENT_SECRET", "old-client-secret", "2" + ), + ], + } + ], + }, + "status": {"imageDigest": IMAGE}, + } + + +def _fakes(tmp_path, revision, *, schedule_exists): + fakes = FakeCommands.create(tmp_path) + revision_file = tmp_path / "revision.json" + revision_file.write_text(json.dumps(revision), encoding="utf-8") + environment_capture = tmp_path / "job-environment.json" + fakes.add( + "gcloud", + f""" + case "$1 $2 $3" in + "run revisions describe") cat "{revision_file}" ;; + "run jobs deploy") + while (($#)); do + if [[ "$1" == "--env-vars-file" ]]; then + cp "$2" "{environment_capture}" + fi + shift + done + ;; + "scheduler jobs describe") {"exit 0" if schedule_exists else "exit 1"} ;; + *) ;; + esac + """, + ) + return fakes, environment_capture + + +def _deploy(tmp_path, revision=None, *, schedule_exists=False, **overrides): + fakes, environment_capture = _fakes( + tmp_path, revision or _revision(), schedule_exists=schedule_exists + ) + settings = { + "PROJECT_ID": "policyengine", + "REGION": "us-central1", + "SERVICE": SERVICE, + "SAFETY_CHECK_REVISION": REVISION, + "SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT": SCHEDULER_ACCOUNT, + "RUNNER_TEMP": str(tmp_path), + } + result = run_script(SCRIPT, env=fakes.env(**(settings | overrides))) + return result, fakes, environment_capture + + +def _call(fakes, *prefix): + calls = [ + call + for call in fakes.calls("gcloud") + if call[1 : 1 + len(prefix)] == list(prefix) + ] + assert len(calls) == 1, calls + return calls[0] + + +def _flag(call, name): + return call[call.index(name) + 1] + + +def test_the_job_runs_the_serving_revisions_exact_image_and_identity(tmp_path) -> None: + result, fakes, _ = _deploy(tmp_path) + + assert result.returncode == 0, result.stderr + deploy = _call(fakes, "run", "jobs", "deploy") + assert deploy[4] == JOB + assert _flag(deploy, "--image") == IMAGE + assert _flag(deploy, "--service-account") == RUNTIME_ACCOUNT + assert _flag(deploy, "--command") == "/opt/venv/bin/python" + assert "--args=-m,policyengine_simulation_entry.safety_check" in deploy + assert _flag(deploy, "--tasks") == "1" + assert _flag(deploy, "--max-retries") == "0" + assert _flag(deploy, "--project") == "policyengine" + assert _flag(deploy, "--region") == "us-central1" + + +def test_the_job_gets_only_the_secrets_it_reads_at_the_revisions_versions( + tmp_path, +) -> None: + result, fakes, environment_capture = _deploy(tmp_path) + + assert result.returncode == 0, result.stderr + secrets = _flag(_call(fakes, "run", "jobs", "deploy"), "--set-secrets").split(",") + assert sorted(secrets) == [ + "MODAL_TOKEN_ID=modal-token-id:3", + "MODAL_TOKEN_SECRET=modal-token-secret:3", + "SIMULATION_DATABASE_URL=simulation-runtime-url:4", + ] + environment = json.loads(environment_capture.read_text(encoding="utf-8")) + assert environment == { + "APP_ENVIRONMENT": "production", + "OBSERVABILITY_SERVICE_NAMESPACE": "sim", + "CLOUD_RUN_REGION": "us-central1", + } + + +@pytest.mark.parametrize("schedule_exists", [False, True]) +def test_the_schedule_starts_the_job_every_interval(tmp_path, schedule_exists) -> None: + result, fakes, _ = _deploy(tmp_path, schedule_exists=schedule_exists) + + assert result.returncode == 0, result.stderr + action = "update" if schedule_exists else "create" + schedule = _call(fakes, "scheduler", "jobs", action, "http") + assert schedule[5] == JOB + assert _flag(schedule, "--location") == "us-central1" + assert _flag(schedule, "--schedule") == "*/5 * * * *" + assert _flag(schedule, "--uri") == ( + "https://us-central1-run.googleapis.com/apis/run.googleapis.com/v1/" + f"namespaces/policyengine/jobs/{JOB}:run" + ) + assert _flag(schedule, "--http-method") == "POST" + assert _flag(schedule, "--oauth-service-account-email") == SCHEDULER_ACCOUNT + + +def test_the_schedule_matches_the_safety_check_interval() -> None: + assert timedelta(minutes=5) == SAFETY_CHECK_INTERVAL + + +def test_a_revision_of_another_service_is_refused(tmp_path) -> None: + result, fakes, _ = _deploy(tmp_path, _revision(service="another-service")) + + assert result.returncode != 0 + assert "another-service" in result.stderr + assert [call[1:3] for call in fakes.calls("gcloud")] == [["run", "revisions"]] + + +def test_a_revision_without_the_simulation_database_pauses_the_schedule( + tmp_path, +) -> None: + """An older revision restored by a rollback has no safety check.""" + revision = _revision( + environment=[{"name": "APP_ENVIRONMENT", "value": "production"}] + ) + + result, fakes, _ = _deploy(tmp_path, revision, schedule_exists=True) + + assert result.returncode == 0, result.stderr + assert [call[1:4] for call in fakes.calls("gcloud")] == [ + ["run", "revisions", "describe"], + ["scheduler", "jobs", "describe"], + ["scheduler", "jobs", "pause"], + ] + + +@pytest.mark.parametrize( + ("variable", "value"), + [ + ("SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT", ""), + ("SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT", "not-an-account"), + ("SAFETY_CHECK_REVISION", ""), + ("SAFETY_CHECK_REVISION", "Bad Revision"), + ("SERVICE", ""), + ], +) +def test_incomplete_settings_stop_before_any_change(tmp_path, variable, value) -> None: + result, fakes, _ = _deploy(tmp_path, **{variable: value}) + + assert result.returncode != 0 + assert variable in result.stderr + assert fakes.calls("gcloud") == [] + + +def _steps(job): + return job["steps"] + + +def _step_named(steps, name): + return next(index for index, step in enumerate(steps) if step.get("name") == name) + + +def test_promotion_points_the_job_at_the_revision_after_it_is_verified() -> None: + steps = _steps(workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["promote_entrypoint"]) + promote = _step_named(steps, "Promote exact tested revision") + safety = _step_named(steps, "Point the safety check at the promoted revision") + restore = _step_named(steps, "Restore previous revision after failed verification") + restore_safety = _step_named( + steps, "Point the safety check back at the previous revision" + ) + + verification = [ + index + for index, step in enumerate(steps) + if index > promote and step.get("name", "").startswith("Verify") + ] + assert verification + assert max(verification) < safety < restore < restore_safety + assert SCRIPT in steps[safety]["run"] + assert steps[safety]["env"]["SAFETY_CHECK_REVISION"] == "${{ env.TARGET_REVISION }}" + assert steps[safety]["env"]["SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT"] == ( + "${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }}" + ) + assert "steps.safety_check.outcome == 'failure'" in steps[restore_safety]["if"] + assert steps[restore_safety]["env"]["SAFETY_CHECK_REVISION"] == ( + "${{ env.PREVIOUS_REVISION }}" + ) + + +def test_no_candidate_deploy_touches_the_job() -> None: + jobs = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW) + for name, job in jobs.items(): + runs = " ".join(step.get("run", "") for step in job.get("steps", [])) + if SCRIPT in runs: + assert name == "promote_entrypoint" + + +def test_a_rollback_points_the_job_at_the_restored_revision() -> None: + steps = _steps(workflow_jobs(ROLLBACK_WORKFLOW)["rollback"]) + restore = _step_named(steps, "Restore entrypoint traffic") + safety = _step_named(steps, "Point the safety check at the restored revision") + + assert restore < safety + assert steps[safety]["if"] == "${{ inputs.entrypoint_revision != '' }}" + assert steps[safety]["env"]["SAFETY_CHECK_REVISION"] == ( + "${{ inputs.entrypoint_revision }}" + ) From 27efca96d6d2482fbc0fb66177bdeaa5cfe73130 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:27:03 +0400 Subject: [PATCH 30/39] Add the sender a v2 Modal call uses to report that it ended src.modal.call_reports is v2 code only (the main path's run_simulation does not use it). Its interface for the v2 functions' last step (items 2.01 and 2.03): reporting_call_end wraps a top-level function's work, and when the block returns or raises an Exception it starts a separate call of the report_call_ended Modal function with the call ID and run IDs; the block's result or exception goes on unchanged, and a failure to start the report is logged, never raised. A signal sends nothing: Modal delivers pre-emption as KeyboardInterrupt, re-runs the call under its ID, and the re-run reports. The report goes through a separate call because sim asks Modal for the outcome without waiting, and a call's own output exists only once it has returned. send_call_ended_report, which that function runs, posts the CallEndedReport to the entrypoint's main URL with the container's Modal identity token, reports again after a wait (Retry-After, or a backoff from 1 to 30 seconds) while sim answers 429/502/503/504 or cannot be reached, stops on any other answer, and gives up after six attempts, leaving the call to the safety check. It never raises an Exception. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/call_reports.py | 268 ++++++++++++++++++ .../tests/test_call_reports.py | 237 ++++++++++++++++ 2 files changed, 505 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/modal/call_reports.py create mode 100644 projects/policyengine-simulation-executor/tests/test_call_reports.py diff --git a/projects/policyengine-simulation-executor/src/modal/call_reports.py b/projects/policyengine-simulation-executor/src/modal/call_reports.py new file mode 100644 index 000000000..81ef33db2 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/call_reports.py @@ -0,0 +1,268 @@ +"""Report to the Simulation Entrypoint that a v2 Modal call ended. + +Interface for the v2 functions' last step (the economy coordinator, item +2.01, and the household functions, item 2.03). Wrap the work of a top-level +function whose call sim started:: + + with reporting_call_end( + modal_call_id=modal.current_function_call_id(), + simulation_run_ids=run_ids, # or report_run_id=report_run_id + start_report=report_call_ended.spawn, + ): + result = calculate(...) + return result + +When the block ends, normally or by raising an ``Exception``, the report is +started as a separate call of the ``report_call_ended`` Modal function and +the block's own result or exception goes on unchanged. Starting it never +raises. A signal sends nothing: Modal delivers pre-emption as +``KeyboardInterrupt``, re-runs the call under the same ID, and the re-run +reports. Segment calls report to their coordinator, not to sim. + +Why a separate call: this function's own output exists only once it has +returned, and sim asks Modal for the outcome itself, without waiting, so a +report sent from inside the call would always find it still running. +``report_call_ended`` runs ``send_call_ended_report_from_container``: it +posts the ``CallEndedReport`` (the call ID and run IDs, never a result) to +``SIMULATION_ENTRY_URL`` + ``CALL_ENDED_PATH`` with the container's Modal +identity token, and reports again after a wait while sim answers 503 or +cannot be reached, up to ``CALL_REPORT_ATTEMPTS`` times. A report that never +arrives is recorded by sim's safety check within about ten minutes. + +This module is v2 code only: the main path's ``run_simulation`` does not use +it. +""" + +from __future__ import annotations + +import json +import logging +import os +import time +import urllib.error +import urllib.request +from contextlib import contextmanager +from enum import StrEnum +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + MODAL_IDENTITY_ENVIRONMENT_VARIABLE, + CallEndedReport, +) +from pydantic import ValidationError + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Mapping, Sequence + from uuid import UUID + +logger = logging.getLogger(__name__) + +# The Simulation Entrypoint's main URL, which the deploy bakes into the +# report_call_ended function's environment. +ENTRY_URL_VARIABLE = "SIMULATION_ENTRY_URL" +CALL_REPORT_ATTEMPTS = 6 +FIRST_WAIT_SECONDS = 1.0 +LONGEST_WAIT_SECONDS = 30.0 +REQUEST_TIMEOUT_SECONDS = 15.0 +# sim answers these when it cannot record the outcome yet. +_RETRY_LATER_HTTP_STATUSES = frozenset({429, 502, 503, 504}) + + +class ReportDelivery(StrEnum): + """What sending one report came to.""" + + # sim recorded the outcome, or already had it. + DELIVERED = "delivered" + # sim refused the report; sending it again would not help. + REFUSED = "refused" + # sim could not record it after every attempt; the safety check will. + GAVE_UP = "gave_up" + # The container has no URL for sim or no identity token. + NOT_CONFIGURED = "not_configured" + + +class Poster(Protocol): + """Send one HTTP POST; return the status, headers and body.""" + + def __call__( + self, + url: str, + body: bytes, + headers: Mapping[str, str], + timeout: float, + ) -> tuple[int, Mapping[str, str], bytes]: + """Post ``body`` to ``url``; raise OSError when nothing came back.""" + ... + + +def post_with_urllib( + url: str, body: bytes, headers: Mapping[str, str], timeout: float +) -> tuple[int, Mapping[str, str], bytes]: + """Post with the standard library; an HTTP error status is an answer.""" + request = urllib.request.Request( # noqa: S310 # an https URL the deploy sets + url, data=body, headers=dict(headers), method="POST" + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: # noqa: S310 # as above + return response.status, dict(response.headers), response.read() + except urllib.error.HTTPError as error: + return error.code, dict(error.headers), error.read() + + +def _wait_seconds(attempt: int, headers: Mapping[str, str]) -> float: + retry_after = {name.lower(): value for name, value in headers.items()}.get( + "retry-after", "" + ) + backoff = min(FIRST_WAIT_SECONDS * 2**attempt, LONGEST_WAIT_SECONDS) + if retry_after.isdigit(): + return min(max(float(retry_after), backoff), LONGEST_WAIT_SECONDS) + return backoff + + +def _answer_status(body: bytes) -> str: + try: + answer = json.loads(body) + except ValueError: + return "" + return str(answer.get("status", "")) if isinstance(answer, dict) else "" + + +def send_call_ended_report( + report: CallEndedReport, + *, + entry_url: str, + identity_token: str, + post: Poster = post_with_urllib, + sleep: Callable[[float], None] = time.sleep, +) -> ReportDelivery: + """Post ``report`` to sim, reporting again while sim cannot record it yet. + + Never raises an ``Exception``: a report that cannot be delivered is + logged, and sim's safety check records the call instead. + """ + if not entry_url or not identity_token: + logger.warning( + "call_report_not_configured", + extra={"modal_call_id": report.modal_call_id}, + ) + return ReportDelivery.NOT_CONFIGURED + url = entry_url.rstrip("/") + CALL_ENDED_PATH + body = report.model_dump_json().encode() + headers = { + "Authorization": f"Bearer {identity_token}", + "Content-Type": "application/json", + } + answer_headers: Mapping[str, str] = {} + for attempt in range(CALL_REPORT_ATTEMPTS): + if attempt: + sleep(_wait_seconds(attempt - 1, answer_headers)) + try: + status, answer_headers, answer_body = post( + url, body, headers, REQUEST_TIMEOUT_SECONDS + ) + except OSError as error: # no answer: unreachable, reset or timed out + answer_headers = {} + logger.warning( + "call_report_not_answered", + extra={ + "modal_call_id": report.modal_call_id, + "attempt": attempt + 1, + "error_type": type(error).__name__, + }, + ) + continue + fields = { + "modal_call_id": report.modal_call_id, + "attempt": attempt + 1, + "http_status": status, + "call_report_status": _answer_status(answer_body), + } + if status < 300: # noqa: PLR2004 # any 2xx answer + logger.info("call_report_delivered", extra=fields) + return ReportDelivery.DELIVERED + if status not in _RETRY_LATER_HTTP_STATUSES: + logger.warning("call_report_refused", extra=fields) + return ReportDelivery.REFUSED + logger.info("call_report_not_recorded_yet", extra=fields) + logger.warning( + "call_report_given_up", + extra={"modal_call_id": report.modal_call_id, "attempts": CALL_REPORT_ATTEMPTS}, + ) + return ReportDelivery.GAVE_UP + + +def send_call_ended_report_from_container( + report: Mapping[str, object], + *, + post: Poster = post_with_urllib, + sleep: Callable[[float], None] = time.sleep, +) -> str: + """Run in ``report_call_ended``: send ``report`` with this container's settings. + + Return the ``ReportDelivery`` value. A malformed report is refused. + """ + try: + call_report = CallEndedReport.model_validate(report) + except ValidationError: + logger.warning("call_report_malformed") + return ReportDelivery.REFUSED.value + return send_call_ended_report( + call_report, + entry_url=os.environ.get(ENTRY_URL_VARIABLE, ""), + identity_token=os.environ.get(MODAL_IDENTITY_ENVIRONMENT_VARIABLE, ""), + post=post, + sleep=sleep, + ).value + + +def _start_report( + report: CallEndedReport, start_report: Callable[[dict[str, object]], object] +) -> None: + try: + start_report(report.model_dump(mode="json")) + except Exception as error: # a failed report must never fail the calculation + logger.exception( + "call_report_not_started", + extra={ + "modal_call_id": report.modal_call_id, + "error_type": type(error).__name__, + }, + ) + + +@contextmanager +def reporting_call_end( + *, + modal_call_id: str | None, + start_report: Callable[[dict[str, object]], object], + simulation_run_ids: Sequence[UUID] | None = None, + report_run_id: UUID | None = None, +) -> Iterator[None]: + """Report the call's end when the block returns or raises an Exception. + + ``start_report`` starts the ``report_call_ended`` function (its + ``spawn``). Without a call ID (outside Modal), or with run IDs the report + cannot carry, nothing is reported. + """ + report: CallEndedReport | None = None + try: + if modal_call_id is not None: + report = CallEndedReport( + modal_call_id=modal_call_id, + simulation_run_ids=( + None if simulation_run_ids is None else tuple(simulation_run_ids) + ), + report_run_id=report_run_id, + ) + except ValidationError: + logger.warning("call_report_malformed", extra={"modal_call_id": modal_call_id}) + if report is None: + yield + return + try: + yield + except Exception: + _start_report(report, start_report) + raise + _start_report(report, start_report) diff --git a/projects/policyengine-simulation-executor/tests/test_call_reports.py b/projects/policyengine-simulation-executor/tests/test_call_reports.py new file mode 100644 index 000000000..4d33d5219 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_call_reports.py @@ -0,0 +1,237 @@ +"""Reporting to the Simulation Entrypoint that a v2 Modal call ended.""" + +from __future__ import annotations + +import json +import urllib.error +from dataclasses import dataclass, field +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.call_reports import ( + CALL_ENDED_PATH, + CallEndedReport, +) +from src.modal.call_reports import ( + CALL_REPORT_ATTEMPTS, + ReportDelivery, + reporting_call_end, + send_call_ended_report, + send_call_ended_report_from_container, +) + +ENTRY_URL = "https://policyengine-simulation-entry-abc.a.run.app" +RUN_ID = uuid4() +REPORT = CallEndedReport(modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,)) + + +@dataclass +class _Entrypoint: + """Answer each request with the next scripted (status, headers, body) or error.""" + + answers: list + requests: list = field(default_factory=list) + + def __call__(self, url, body, headers, timeout): + self.requests.append({"url": url, "body": body, "headers": headers}) + answer = self.answers.pop(0) + if isinstance(answer, BaseException): + raise answer + status, answer_headers, answer_body = answer + return status, answer_headers, json.dumps(answer_body).encode() + + +@dataclass +class _Waits: + seconds: list = field(default_factory=list) + + def __call__(self, seconds: float) -> None: + self.seconds.append(seconds) + + +def _send(answers, **settings): + entrypoint, waits = _Entrypoint(list(answers)), _Waits() + delivery = send_call_ended_report( + REPORT, + entry_url=settings.get("entry_url", ENTRY_URL), + identity_token=settings.get("identity_token", "modal-identity-token"), + post=entrypoint, + sleep=waits, + ) + return delivery, entrypoint, waits + + +def test_a_report_is_posted_with_the_containers_identity_token() -> None: + delivery, entrypoint, waits = _send([(200, {}, {"status": "recorded"})]) + + assert delivery is ReportDelivery.DELIVERED + (request,) = entrypoint.requests + assert request["url"] == ENTRY_URL + CALL_ENDED_PATH + assert request["headers"]["Authorization"] == "Bearer modal-identity-token" + assert request["headers"]["Content-Type"] == "application/json" + assert CallEndedReport.model_validate_json(request["body"]) == REPORT + assert waits.seconds == [] + + +def test_a_report_sim_cannot_record_yet_is_sent_again_after_its_wait() -> None: + delivery, entrypoint, waits = _send( + [ + (503, {"Retry-After": "3"}, {"status": "still_running"}), + (503, {}, {"status": "not_dispatched_yet"}), + (200, {}, {"status": "already_recorded"}), + ] + ) + + assert delivery is ReportDelivery.DELIVERED + assert len(entrypoint.requests) == 3 + assert waits.seconds[0] == 3 + assert waits.seconds[1] > 0 + + +@pytest.mark.parametrize( + "failure", + [urllib.error.URLError("unreachable"), TimeoutError(), ConnectionResetError()], + ids=["unreachable", "timeout", "reset"], +) +def test_a_report_that_does_not_arrive_is_sent_again(failure) -> None: + delivery, entrypoint, _ = _send([failure, (200, {}, {"status": "recorded"})]) + + assert delivery is ReportDelivery.DELIVERED + assert len(entrypoint.requests) == 2 + + +@pytest.mark.parametrize("status", [401, 404, 409, 422]) +def test_a_refused_report_is_not_sent_again(status) -> None: + delivery, entrypoint, waits = _send([(status, {}, {"status": "other_call"})]) + + assert delivery is ReportDelivery.REFUSED + assert len(entrypoint.requests) == 1 + assert waits.seconds == [] + + +def test_the_sender_gives_up_after_a_few_attempts_with_growing_waits() -> None: + answers = [(503, {}, {"status": "still_running"})] * CALL_REPORT_ATTEMPTS + + delivery, entrypoint, waits = _send(answers) + + assert delivery is ReportDelivery.GAVE_UP + assert len(entrypoint.requests) == CALL_REPORT_ATTEMPTS + assert len(waits.seconds) == CALL_REPORT_ATTEMPTS - 1 + assert waits.seconds == sorted(waits.seconds) + assert sum(waits.seconds) < 120 + + +def test_an_unexpected_answer_is_not_a_crash() -> None: + delivery, _, _ = _send([(200, {}, "not an answer")]) + + assert delivery is ReportDelivery.DELIVERED + + +def test_a_report_without_sims_url_or_a_token_is_not_sent() -> None: + delivery, entrypoint, _ = _send([], entry_url="") + assert delivery is ReportDelivery.NOT_CONFIGURED + delivery, entrypoint, _ = _send([], identity_token="") + assert delivery is ReportDelivery.NOT_CONFIGURED + assert entrypoint.requests == [] + + +def test_the_report_function_reads_its_settings_from_the_container( + monkeypatch, +) -> None: + monkeypatch.setenv("SIMULATION_ENTRY_URL", ENTRY_URL) + monkeypatch.setenv("MODAL_IDENTITY_TOKEN", "modal-identity-token") + entrypoint = _Entrypoint([(200, {}, {"status": "recorded"})]) + + delivery = send_call_ended_report_from_container( + REPORT.model_dump(mode="json"), post=entrypoint, sleep=_Waits() + ) + + assert delivery == "delivered" + assert entrypoint.requests[0]["url"] == ENTRY_URL + CALL_ENDED_PATH + + +@dataclass +class _ReportStarter: + """Stands in for ``report_call_ended.spawn``.""" + + failure: BaseException | None = None + started: list = field(default_factory=list) + + def __call__(self, report: dict) -> None: + if self.failure is not None: + raise self.failure + self.started.append(report) + + +def test_a_call_that_returns_reports_its_end() -> None: + starter = _ReportStarter() + + with reporting_call_end( + modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,), start_report=starter + ): + pass + + assert [CallEndedReport.model_validate(report) for report in starter.started] == [ + REPORT + ] + + +def _failing_comparison(report_run_id, starter) -> None: + with reporting_call_end( + modal_call_id="fc-failed", report_run_id=report_run_id, start_report=starter + ): + msg = "calculation failed" + raise ValueError(msg) + + +def test_a_call_that_fails_reports_its_end_and_still_fails() -> None: + starter = _ReportStarter() + report_run_id = uuid4() + + with pytest.raises(ValueError, match="calculation failed"): + _failing_comparison(report_run_id, starter) + + assert starter.started == [ + { + "modal_call_id": "fc-failed", + "simulation_run_ids": None, + "report_run_id": str(report_run_id), + } + ] + + +@pytest.mark.parametrize("signal", [KeyboardInterrupt(), SystemExit(1)]) +def test_a_pre_empted_call_reports_nothing(signal) -> None: + """Modal re-runs a pre-empted call under its ID, and the re-run reports.""" + starter = _ReportStarter() + + with ( + pytest.raises(type(signal)), + reporting_call_end( + modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,), start_report=starter + ), + ): + raise signal + + assert starter.started == [] + + +def test_a_report_that_cannot_start_never_fails_the_calculation() -> None: + starter = _ReportStarter(failure=RuntimeError("Modal unreachable")) + + with reporting_call_end( + modal_call_id="fc-ended", simulation_run_ids=(RUN_ID,), start_report=starter + ): + pass + + +def test_a_call_without_a_modal_call_id_reports_nothing() -> None: + """Run outside Modal (a local test), the call has no ID to report.""" + starter = _ReportStarter() + + with reporting_call_end( + modal_call_id=None, simulation_run_ids=(RUN_ID,), start_report=starter + ): + pass + + assert starter.started == [] From 260300417da5bd2e46a6ce355234ea9fed15dc05 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:30:43 +0400 Subject: [PATCH 31/39] Declare report_call_ended in the v2 Modal app with sim's main URL The v2 app (src/modal/v2_app.py) gains report_call_ended, the function the v2 functions' last step starts and the only Modal function whose identity token sim accepts reports from; a check keeps its name equal to the contract's CALL_REPORT_FUNCTION_NAME. It runs send_call_ended_report_from_container on the coordinator's existing image with a quarter CPU, no retries, a five-minute timeout covering every attempt, and no secrets: its own identity token proves who it is. Its SIMULATION_ENTRY_URL is sim's main URL, set at deploy time: the v2 deploy job reads the entrypoint service's stable URL (status.url) with the new simulation-entry-main-url.sh, which refuses anything but an https URL, and passes it to modal deploy; v2_app refuses a URL that is not https. Nothing calls the function yet (items 2.01 and 2.03 do). Co-Authored-By: Claude Opus 5.5 --- .github/scripts/simulation-entry-main-url.sh | 24 +++++++ .../workflows/simulation-deploy.reusable.yml | 10 +++ .../DEPLOYMENT.md | 10 +++ .../tests/test_call_report_deployment.py | 63 +++++++++++++++++++ .../src/modal/v2_app.py | 47 ++++++++++++++ .../tests/test_call_report_function.py | 53 ++++++++++++++++ .../tests/test_stage12_modal_app.py | 1 + 7 files changed, 208 insertions(+) create mode 100755 .github/scripts/simulation-entry-main-url.sh create mode 100644 projects/policyengine-simulation-executor/tests/test_call_report_function.py diff --git a/.github/scripts/simulation-entry-main-url.sh b/.github/scripts/simulation-entry-main-url.sh new file mode 100755 index 000000000..8e86373d8 --- /dev/null +++ b/.github/scripts/simulation-entry-main-url.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Write the Simulation Entrypoint's main URL (its stable Cloud Run URL, which +# reaches the revision that serves traffic) to $GITHUB_OUTPUT as `url`, for +# the v2 Modal app's report_call_ended function. +# +# Environment: PROJECT_ID, REGION, SERVICE, GITHUB_OUTPUT. +set -euo pipefail + +fail() { + printf '%s\n' "$*" >&2 + exit 1 +} + +for variable_name in PROJECT_ID REGION SERVICE GITHUB_OUTPUT; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" +done + +url="$(gcloud run services describe "${SERVICE}" \ + --project "${PROJECT_ID}" \ + --region "${REGION}" \ + --format='value(status.url)')" +[[ "${url}" =~ ^https://[a-z0-9.-]+$ ]] \ + || fail "${SERVICE} has no https main URL: ${url:-none}" +printf 'url=%s\n' "${url}" >> "${GITHUB_OUTPUT}" diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 24d2b898b..aaee397f0 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -671,12 +671,22 @@ jobs: STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} run: ../../.github/scripts/stage12-cache-precompute.sh "${{ inputs.modal_environment }}" + # The v2 functions report to sim's main URL that a call ended. + - name: Find the Simulation Entrypoint's main URL + id: entry_url + env: + PROJECT_ID: ${{ vars.SIMULATION_ENTRYPOINT_GCP_PROJECT_ID }} + REGION: ${{ vars.SIMULATION_ENTRYPOINT_GCP_REGION || 'us-central1' }} + SERVICE: ${{ inputs.entrypoint_service }} + run: .github/scripts/simulation-entry-main-url.sh + - name: Deploy the additional versioned v2 Modal application working-directory: projects/policyengine-simulation-executor env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} MODAL_APP_NAME: ${{ steps.application.outputs.name }} + SIMULATION_ENTRY_URL: ${{ steps.entry_url.outputs.url }} STAGE12_EXPECT_POLICYENGINE_VERSION: ${{ needs.prepare.outputs.policyengine_version }} STAGE12_EXPECT_US_VERSION: ${{ needs.prepare.outputs.us_version }} STAGE12_EXPECT_UK_VERSION: ${{ needs.prepare.outputs.uk_version }} diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index e910794df..00b7a1ebf 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -81,6 +81,16 @@ missing. The service already allows unauthenticated requests through Cloud Run (`--allow-unauthenticated`): callers prove who they are to the entrypoint, so Modal needs no Google credentials. +The reports come from the v2 Modal app's `report_call_ended` function, which +the v2 functions start as their last step (`src/modal/call_reports.py` in the +executor). It posts to the entrypoint's main URL: the v2 deploy job reads the +service's stable URL (`status.url`) with +`.github/scripts/simulation-entry-main-url.sh` and passes it to `modal deploy` +as `SIMULATION_ENTRY_URL`. Only the revision that serves traffic receives +these requests. `beta` does not promote its candidate yet, so in staging the +main URL reaches an older revision until it does (plan item 1.39), and +staging calls are recorded by the safety check only once it exists there. + ## Safety check A call whose report never arrives is recorded by the safety check, the Cloud diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py b/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py index 9b428dfbb..cb09bfad3 100644 --- a/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py +++ b/projects/policyengine-simulation-entry/tests/test_call_report_deployment.py @@ -19,6 +19,7 @@ REUSABLE_DEPLOY_WORKFLOW, workflow_jobs, ) +from workflow_script_support import FakeCommands, github_output, run_script def _deploy(tmp_path, **overrides): @@ -74,3 +75,65 @@ def test_the_deploy_step_passes_the_modal_workspace_id() -> None: ) assert environment["MODAL_WORKSPACE_ID_VALUE"] == "${{ vars.MODAL_WORKSPACE_ID }}" + + +URL_SCRIPT = "simulation-entry-main-url.sh" + + +def _find_url(tmp_path, service_url, **overrides): + fakes = FakeCommands.create(tmp_path) + fakes.add("gcloud", f"printf '%s\\n' '{service_url}'\n") + output = tmp_path / "github-output" + settings = { + "PROJECT_ID": "policyengine", + "REGION": "us-central1", + "SERVICE": "policyengine-simulation-entry", + "GITHUB_OUTPUT": str(output), + } + result = run_script(URL_SCRIPT, env=fakes.env(**(settings | overrides))) + return result, fakes, github_output(output) + + +def test_the_modal_functions_report_to_the_entrypoints_main_url(tmp_path) -> None: + url = "https://policyengine-simulation-entry-abc-uc.a.run.app" + + result, fakes, outputs = _find_url(tmp_path, url) + + assert result.returncode == 0, result.stderr + assert outputs == {"url": url} + (describe,) = fakes.calls("gcloud") + assert describe[1:5] == [ + "run", + "services", + "describe", + "policyengine-simulation-entry", + ] + assert "--format=value(status.url)" in describe + + +@pytest.mark.parametrize("service_url", ["", "http://insecure.example", "not a url"]) +def test_a_main_url_that_is_not_https_stops_the_deploy(tmp_path, service_url) -> None: + result, _, outputs = _find_url(tmp_path, service_url) + + assert result.returncode != 0 + assert outputs == {} + + +def test_the_v2_deploy_passes_the_main_url_to_the_modal_app() -> None: + job = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_stage12_v2"] + steps = job["steps"] + find = next( + index for index, step in enumerate(steps) if URL_SCRIPT in step.get("run", "") + ) + deploy = next( + index + for index, step in enumerate(steps) + if "src/modal/v2_app.py" in step.get("run", "") + ) + + assert find < deploy + assert steps[find]["id"] == "entry_url" + assert steps[find]["env"]["SERVICE"] == "${{ inputs.entrypoint_service }}" + assert steps[deploy]["env"]["SIMULATION_ENTRY_URL"] == ( + "${{ steps.entry_url.outputs.url }}" + ) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 69eedc493..3789a940f 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import TYPE_CHECKING +from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME from policyengine_simulation_contract.stage12_manifest import v2_application_name from policyengine_simulation_executor import stage12_runtime from policyengine_simulation_executor.hf_access_validation import ( @@ -94,6 +95,23 @@ def _deploy_time_cache_inputs() -> tuple[str, dict | None]: STAGE12_CACHE_BUCKET, STAGE12_CACHE_MANIFEST = _deploy_time_cache_inputs() +def _deploy_time_entry_url() -> str: + """Return sim's main URL, which the deploy passes, for report_call_ended.""" + if not modal.is_local(): + return "" + url = os.environ.get("SIMULATION_ENTRY_URL", "") + if url and (url != url.strip() or not url.startswith("https://")): + msg = "SIMULATION_ENTRY_URL must be an https URL" + raise RuntimeError(msg) + return url + + +SIMULATION_ENTRY_URL = _deploy_time_entry_url() +# Every attempt to report, each with its request timeout, and every wait +# between them (src/modal/call_reports.py), with a margin. +CALL_REPORT_TIMEOUT_SECONDS = 300 + + def _external_assertions(environment: dict[str, str]) -> dict[str, str]: mapping = { "STAGE12_EXPECT_POLICYENGINE_VERSION": "policyengine_version", @@ -463,3 +481,32 @@ def coordinate_report( ) finally: runtime.shutdown() + + +@app.function( + image=coordinator_image, + cpu=0.25, + memory=512, + timeout=CALL_REPORT_TIMEOUT_SECONDS, + retries=0, + max_containers=20, + # No secrets: sim accepts this function's own Modal identity token. + env={"SIMULATION_ENTRY_URL": SIMULATION_ENTRY_URL}, +) +def report_call_ended(report: dict) -> str: + """Tell sim that the call in ``report`` ended; return how that went. + + The v2 functions start this as their last step + (``src.modal.call_reports.reporting_call_end``), so the reported call + has returned by the time sim asks Modal for its outcome. + """ + from src.modal.call_reports import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + send_call_ended_report_from_container, + ) + + return send_call_ended_report_from_container(report) + + +if report_call_ended.__name__ != CALL_REPORT_FUNCTION_NAME: + msg = "report_call_ended must keep the name sim accepts reports from" + raise RuntimeError(msg) diff --git a/projects/policyengine-simulation-executor/tests/test_call_report_function.py b/projects/policyengine-simulation-executor/tests/test_call_report_function.py new file mode 100644 index 000000000..cf9937db0 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_call_report_function.py @@ -0,0 +1,53 @@ +"""The v2 app's report_call_ended function, which tells sim that a call ended.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME +from src.modal import call_reports +from test_stage12_modal_app import _load + +ENTRY_URL = "https://policyengine-simulation-entry-abc.a.run.app" + + +def test_the_v2_app_declares_the_function_sim_accepts_reports_from( + monkeypatch, +) -> None: + monkeypatch.setenv("SIMULATION_ENTRY_URL", ENTRY_URL) + module = _load(monkeypatch) + + declaration = dict(module.app.function_calls)[CALL_REPORT_FUNCTION_NAME] + + assert declaration["image"] is module.coordinator_image + assert declaration["env"] == {"SIMULATION_ENTRY_URL": ENTRY_URL} + assert declaration["retries"] == 0 + # Every attempt, each with its request timeout, and every wait between them. + assert declaration["timeout"] >= call_reports.CALL_REPORT_ATTEMPTS * ( + call_reports.REQUEST_TIMEOUT_SECONDS + call_reports.LONGEST_WAIT_SECONDS + ) + # It needs no secret: its own identity token proves who it is. + assert "secrets" not in declaration + + +def test_the_function_sends_the_report_with_its_containers_settings( + monkeypatch, +) -> None: + module = _load(monkeypatch) + sent = [] + monkeypatch.setattr( + call_reports, + "send_call_ended_report_from_container", + lambda report: sent.append(report) or "delivered", + ) + report = {"modal_call_id": "fc-ended", "report_run_id": "x"} + + assert module.report_call_ended(report) == "delivered" + assert sent == [report] + + +@pytest.mark.parametrize("url", ["http://entry.example", "entry.example", " https://x"]) +def test_a_deploy_with_a_url_that_is_not_https_stops(monkeypatch, url) -> None: + monkeypatch.setenv("SIMULATION_ENTRY_URL", url) + + with pytest.raises(RuntimeError, match="SIMULATION_ENTRY_URL"): + _load(monkeypatch) 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 671aa4d12..1850f4bda 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -326,6 +326,7 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( "run_single_simulation_us", "run_single_simulation_uk", "coordinate_report", + "report_call_ended", } assert functions["run_single_simulation_us"]["image"] is module.us_worker_image assert functions["run_single_simulation_uk"]["image"] is module.uk_worker_image From 2ef2b88b31396d80ebff0a01eefc50ddb8db41f7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:34:03 +0400 Subject: [PATCH 32/39] Store dispatch errors that say only what is known From the third review of 1.06: two stored error messages could be false. A spawn that raised stored "Modal did not accept the call", but Modal may have created the call before the error; it now stores "Starting the Modal call raised ." And a dispatch failed as abandoned stored that its instance had stopped, which is also stored when the instance kept running but could not tell whether its record committed (RECORD_UNKNOWN); it now stores "No Modal call was recorded for the dispatch in time." Co-Authored-By: Claude Opus 5.5 --- .../run_transitions.py | 12 +++++------- .../tests/test_run_transitions.py | 5 +++++ .../policyengine_simulation_entry/run_dispatch.py | 2 +- .../tests/test_run_dispatch.py | 3 +++ 4 files changed, 14 insertions(+), 8 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py index 7686d639b..9d83327f2 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_transitions.py @@ -60,7 +60,7 @@ class RunErrorCode(StrEnum): # Modal refused the call or never acknowledged it, or the acknowledged # call could not be recorded (the entrypoint then tries to cancel it). DISPATCH_FAILED = "dispatch_failed" - # The instance stopped between claiming the dispatch and recording the call. + # No Modal call was recorded in time after the dispatch was claimed. DISPATCH_ABANDONED = "dispatch_abandoned" # The Modal function raised. WORKER_ERROR = "worker_error" @@ -570,15 +570,13 @@ def fail_abandoned_dispatches( ) -> AbandonedDispatches: """Fail runs and comparison steps left in ``dispatching``. - A run stays ``dispatching`` only if the instance that claimed it - stopped before recording its Modal call. Failing it, with + A run stays ``dispatching`` when no Modal call was recorded for it: + the instance that claimed it stopped, or could not tell whether its + record committed and left the call running. Failing it, with ``dispatch_abandoned``, lets the same simulation be run again. """ now = self._clock() - message = ( - "The instance that claimed the dispatch stopped before it recorded " - "the Modal call." - ) + message = "No Modal call was recorded for the dispatch in time." with self.engine.begin() as connection: failed_runs = connection.execute( update(simulation_runs) diff --git a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py index eaec97a5d..e54fce729 100644 --- a/libs/policyengine-simulation-persistence/tests/test_run_transitions.py +++ b/libs/policyengine-simulation-persistence/tests/test_run_transitions.py @@ -507,6 +507,11 @@ def test_dispatches_left_for_ten_minutes_are_failed( failed_run = stores.runs.get_simulation_run(stalled.id) assert failed_run.status is RunStatus.FAILED assert failed_run.error_code == "dispatch_abandoned" + # The claimer may have stopped, or may not know whether its record + # committed; the message says only what is known. + assert failed_run.error_message == ( + "No Modal call was recorded for the dispatch in time." + ) assert ( stores.runs.get_report_run(report_run.id).comparison_error_code == "dispatch_abandoned" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index af392ffd8..07a691b9c 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -218,7 +218,7 @@ async def _start( steps.fail, RunError( RunErrorCode.DISPATCH_FAILED, - f"Modal did not accept the call ({type(error).__name__}).", + f"Starting the Modal call raised {type(error).__name__}.", ), ) return DispatchOutcome.DISPATCH_FAILED diff --git a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py index 0309ace38..a90f420bc 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_dispatch.py +++ b/projects/policyengine-simulation-entry/tests/test_run_dispatch.py @@ -103,6 +103,9 @@ async def test_a_refused_call_fails_the_runs_without_raising() -> None: assert run.error.code is RunErrorCode.DISPATCH_FAILED assert "NotFoundError" in run.error.message assert "no such function" not in run.error.message + # Modal may have created the call before the error, so the message says + # only that starting it raised. + assert run.error.message == "Starting the Modal call raised NotFoundError." @pytest.mark.asyncio From fe581b9ce5a02247b3784b93baf535bbb31adec5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:34:30 +0400 Subject: [PATCH 33/39] Describe the load check and the conditional writes for reported calls The load check's late writes stand for a call's report arriving while a safety-check pass checks the same call; its description says so, and that the same check of the reporting route against staging is an operator step. The persistence README says a reported call needs no claim, since its outcome is written conditionally too. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 7 +++--- .../tests/test_recording_load.py | 23 +++++++++++-------- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 85b878aae..faeff7e29 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -55,9 +55,10 @@ not delete them. `RunTransitions` makes every status change one conditional update on the expected status, so a second writer changes nothing: only the caller that moves runs from `pending` to `dispatching` starts their Modal call, the runs -one call computes change together, a recording pass claims a due call by -moving its next check time, results are written in the transaction that -marks a run succeeded, and dispatches left `dispatching` (and comparison +one call computes change together, a safety-check pass claims a due call by +moving its next check time (a call's own report needs no claim, since its +outcome is written conditionally too), results are written in the +transaction that marks a run succeeded, and dispatches left `dispatching` (and comparison steps whose simulation run failed) are failed so the work can be submitted again. A comparison step can start only once both of its simulation runs have succeeded; starting it is left to the stage 2 report path. diff --git a/libs/policyengine-simulation-persistence/tests/test_recording_load.py b/libs/policyengine-simulation-persistence/tests/test_recording_load.py index 00df05421..28b096c68 100644 --- a/libs/policyengine-simulation-persistence/tests/test_recording_load.py +++ b/libs/policyengine-simulation-persistence/tests/test_recording_load.py @@ -1,13 +1,16 @@ -"""A load check of the recording passes' conditional writes on PostgreSQL. - -Several simulated Simulation Entrypoint instances, each on its own engine, -repeatedly fail stalled dispatches, claim due calls and record an outcome for -each, as recording passes do, until every run has finished. A third of -the calls are household pairs (two runs under one call). Every call must be -recorded once, its results written once and all of its runs finished once, -and every stalled dispatch failed once, whatever the interleaving. The -elapsed time and the number of passes are printed, as a measurement; nothing -depends on them. +"""A load check of the conditional writes that record outcomes, on PostgreSQL. + +Several simulated writers, each on its own engine, repeatedly fail stalled +dispatches, claim due calls and record an outcome for each, as safety-check +passes do, until every run has finished. Some outcomes are written late, +after another writer could claim the same call and record it, as a report +from Modal can arrive while a pass checks the same call. A third of the calls +are household pairs (two runs under one call). Every call must be recorded +once, its results written once and all of its runs finished once, and every +stalled dispatch failed once, whatever the interleaving. The elapsed time and +the number of passes are printed, as a measurement; nothing depends on them. +The same check of the reporting route against the staging database is an +operator step. """ from __future__ import annotations From ef93571496d686feee9d8ea2d6ed07d14f6d294a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:06:54 +0400 Subject: [PATCH 34/39] Answer a report with a wait whenever Modal's keys cannot be read From the extra review of the reworked 1.06: ModalIdentityVerifier mapped only PyJWKClientConnectionError to unavailable keys, but PyJWKClient wraps only URLError and TimeoutError. A key set that arrived truncated, reset or not as JSON escaped the route's handlers and became a 500, which the sender takes as a refusal, so the report was dropped and the outcome waited for the safety check. Any failure to fetch or read the keys, other than a token naming no published key, now raises ModalKeysUnavailableError, and the route answers any other error from the token check with 503 modal_unreachable as well. The check moved into its own helper, so the route keeps six returns. Co-Authored-By: Claude Opus 5.5 --- .../call_report_route.py | 36 +++++++++++++------ .../modal_identity.py | 6 ++++ .../tests/test_call_report_route.py | 17 +++++++++ .../tests/test_modal_identity.py | 36 +++++++++++++++++++ 4 files changed, 84 insertions(+), 11 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py index 1673b25cd..f21e56325 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/call_report_route.py @@ -104,6 +104,28 @@ def _bearer_token(request: Request) -> str | None: return header[len(_BEARER_PREFIX) :].strip() or None +async def _caller( + request: Request, verifier: ModalCallerVerifier +) -> ModalIdentity | JSONResponse: + """Return the verified Modal caller, or the answer that refuses or defers it.""" + token = _bearer_token(request) + if token is None: + return _unauthorized() + try: + return await asyncio.to_thread(verifier.verify, token) + except ModalKeysUnavailableError: + logger.warning("call_report_modal_keys_unavailable") + return _answer(CallReportStatus.MODAL_UNREACHABLE) + except ModalIdentityError as error: + logger.warning("call_report_refused", extra={"reason": str(error)}) + return _unauthorized() + # Anything else is this instance's own problem, not a refusal: the sender + # reports again after the wait. + except Exception: + logger.exception("call_report_token_not_checked") + return _answer(CallReportStatus.MODAL_UNREACHABLE) + + def register_call_report_route( app: FastAPI, *, @@ -116,17 +138,9 @@ def register_call_report_route( async def report_call_ended(request: Request) -> JSONResponse: if verifier is None or recorder is None: return _answer(CallReportStatus.NOT_CONFIGURED) - token = _bearer_token(request) - if token is None: - return _unauthorized() - try: - identity = await asyncio.to_thread(verifier.verify, token) - except ModalKeysUnavailableError: - logger.warning("call_report_modal_keys_unavailable") - return _answer(CallReportStatus.MODAL_UNREACHABLE) - except ModalIdentityError as error: - logger.warning("call_report_refused", extra={"reason": str(error)}) - return _unauthorized() + identity = await _caller(request, verifier) + if isinstance(identity, JSONResponse): + return identity try: report = CallEndedReport.model_validate_json(await request.body()) except ValidationError as error: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py index cbd971194..cecb7190c 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/modal_identity.py @@ -111,6 +111,12 @@ def verify(self, token: str) -> ModalIdentity: except jwt.PyJWTError as error: msg = f"the token names no key Modal publishes: {type(error).__name__}" raise ModalIdentityError(msg) from error + # PyJWKClient wraps only URLError and TimeoutError; a key set that + # arrives truncated, reset or not as JSON raises something else, and + # is no more the caller's fault. + except Exception as error: + msg = f"Modal's signing keys could not be read: {type(error).__name__}" + raise ModalKeysUnavailableError(msg) from error try: claims = jwt.decode( token, diff --git a/projects/policyengine-simulation-entry/tests/test_call_report_route.py b/projects/policyengine-simulation-entry/tests/test_call_report_route.py index 326d4dbf2..ebdeb996b 100644 --- a/projects/policyengine-simulation-entry/tests/test_call_report_route.py +++ b/projects/policyengine-simulation-entry/tests/test_call_report_route.py @@ -179,3 +179,20 @@ def test_the_route_is_not_in_the_public_openapi_document() -> None: document = _client(_Verifier(), _Recorder()).get("/openapi.json").json() assert CALL_ENDED_PATH not in document["paths"] + + +class _FailingVerifier: + def verify(self, token: str) -> ModalIdentity: + msg = "an error the verifier did not expect" + raise RuntimeError(msg) + + +def test_an_unexpected_verifier_failure_is_answered_with_a_wait() -> None: + """The sender reports again, instead of taking a 500 as a refusal.""" + recorder = _Recorder() + + response = _post(_client(_FailingVerifier(), recorder)) + + assert response.status_code == 503 + assert response.json() == {"status": "modal_unreachable"} + assert recorder.reports == [] diff --git a/projects/policyengine-simulation-entry/tests/test_modal_identity.py b/projects/policyengine-simulation-entry/tests/test_modal_identity.py index 5cd0eade2..7f4c542ae 100644 --- a/projects/policyengine-simulation-entry/tests/test_modal_identity.py +++ b/projects/policyengine-simulation-entry/tests/test_modal_identity.py @@ -147,3 +147,39 @@ def test_unreachable_modal_keys_are_not_a_refusal() -> None: with pytest.raises(ModalKeysUnavailableError): verifier.verify(_token()) + + +class _BrokenJwks: + """Modal's key set answered, but not with a key set PyJWT could read.""" + + def __init__(self, failure: Exception) -> None: + self.failure = failure + + def get_signing_key_from_jwt(self, _token: str): + raise self.failure + + +@pytest.mark.parametrize( + "failure", + [ + ValueError("Expecting value: line 1 column 1 (char 0)"), + ConnectionResetError("reset while reading the key set"), + OSError("truncated response"), + ], + ids=["not-json", "reset", "truncated"], +) +def test_any_failure_to_fetch_modal_keys_is_not_a_refusal(failure) -> None: + verifier = ModalIdentityVerifier(EXPECTED, signing_keys=_BrokenJwks(failure)) + + with pytest.raises(ModalKeysUnavailableError): + verifier.verify(_token()) + + +def test_a_token_naming_no_published_key_is_still_refused() -> None: + verifier = ModalIdentityVerifier( + EXPECTED, + signing_keys=_BrokenJwks(jwt.PyJWKClientError("Unable to find a signing key")), + ) + + with pytest.raises(ModalIdentityError): + verifier.verify(_token()) From c7842cfcffaf0e96b5ce092e5366b77f0423048e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:09:41 +0400 Subject: [PATCH 35/39] Pause the safety check rather than run an image without it From the extra review of the reworked 1.06: - Pointing the job at a revision built from 1.04 or 1.05 deployed it: such a revision mounts the database URL and, with Stage 12, the Modal token, but its image has no safety_check module, so every execution would have failed. A revision now needs SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID, which every deploy has required since the safety check exists; otherwise the schedule is paused. - A failed step in the rollback workflow left the job on the image traffic had just moved away from. The script now takes `pause`, which only pauses the schedule; the rollback workflow runs it when pointing the job at the restored revision fails, and the promotion job when pointing it back at the previous revision fails. Co-Authored-By: Claude Opus 5.5 --- ...eploy-cloud-run-simulation-safety-check.sh | 54 +++++++++++---- .../workflows/simulation-deploy.reusable.yml | 5 ++ .github/workflows/simulation-rollback.yml | 6 ++ .../DEPLOYMENT.md | 12 ++-- .../tests/test_safety_check_deployment.py | 69 +++++++++++++++++++ 5 files changed, 127 insertions(+), 19 deletions(-) diff --git a/.github/scripts/deploy-cloud-run-simulation-safety-check.sh b/.github/scripts/deploy-cloud-run-simulation-safety-check.sh index c50e3d7d3..54d830beb 100755 --- a/.github/scripts/deploy-cloud-run-simulation-safety-check.sh +++ b/.github/scripts/deploy-cloud-run-simulation-safety-check.sh @@ -7,11 +7,15 @@ # promoted (or restored by a rollback), so a rejected candidate never records # anything. It copies the revision's image digest, service account, Modal # token and database URL (at the secret versions the revision pins) and its -# plain telemetry settings; nothing else. A revision without the simulation -# database URL predates the safety check, so its schedule is paused. +# plain telemetry settings; nothing else. A revision deployed before the +# safety check existed has no job code, so its schedule is paused instead. # -# Environment: PROJECT_ID, REGION, SERVICE, SAFETY_CHECK_REVISION, -# SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT, RUNNER_TEMP. +# With the argument `pause`, only pause the schedule: a workflow does this +# when pointing the job at a revision failed part way, so the job never keeps +# running the image traffic just moved away from. +# +# Environment: PROJECT_ID, REGION, SERVICE, and unless pausing, +# SAFETY_CHECK_REVISION, SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT, RUNNER_TEMP. set -euo pipefail fail() { @@ -19,38 +23,58 @@ fail() { exit 1 } +pause_schedule() { + if gcloud scheduler jobs describe "${job}" "${scheduler_location[@]}" >/dev/null 2>&1; then + gcloud scheduler jobs pause "${job}" "${scheduler_location[@]}" --quiet + fi +} + # Every five minutes: SAFETY_CHECK_INTERVAL in # policyengine_simulation_entry.run_recording. SCHEDULE="*/5 * * * *" SECRETS=(SIMULATION_DATABASE_URL MODAL_TOKEN_ID MODAL_TOKEN_SECRET) PLAIN_VARIABLES_FILTER='^(APP_ENVIRONMENT|CLOUD_RUN_REGION|OBSERVABILITY_[A-Z_]+|OTEL_[A-Z_]+|POLICYENGINE_OTEL_GOOGLE_AUDIENCE)$' -for variable_name in PROJECT_ID REGION SERVICE SAFETY_CHECK_REVISION \ - SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT RUNNER_TEMP; do +for variable_name in PROJECT_ID REGION SERVICE; do [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" done [[ "${PROJECT_ID}" =~ ^[a-z][a-z0-9-]{4,28}[a-z0-9]$ ]] || fail "PROJECT_ID is invalid" [[ "${REGION}" =~ ^[a-z]+-[a-z]+[0-9]+$ ]] || fail "REGION is invalid" [[ "${SERVICE}" =~ ^[a-z]([a-z0-9-]{0,47}[a-z0-9])?$ ]] || fail "SERVICE is invalid" -[[ "${SAFETY_CHECK_REVISION}" =~ ^[a-z]([a-z0-9-]{0,61}[a-z0-9])?$ ]] \ - || fail "SAFETY_CHECK_REVISION is invalid" -[[ "${SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ - || fail "SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT is invalid" job="${SERVICE}-safety-check" location=(--project "${PROJECT_ID}" --region "${REGION}") scheduler_location=(--project "${PROJECT_ID}" --location "${REGION}") +if [[ "${1:-}" == "pause" ]]; then + pause_schedule + printf 'The safety check schedule for %s is paused.\n' "${SERVICE}" + exit 0 +fi + +for variable_name in SAFETY_CHECK_REVISION SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT RUNNER_TEMP; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} is required" +done +[[ "${SAFETY_CHECK_REVISION}" =~ ^[a-z]([a-z0-9-]{0,61}[a-z0-9])?$ ]] \ + || fail "SAFETY_CHECK_REVISION is invalid" +[[ "${SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ + || fail "SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT is invalid" + revision_json="$(gcloud run revisions describe "${SAFETY_CHECK_REVISION}" "${location[@]}" --format=json)" revision_service="$(jq -r '.metadata.labels["serving.knative.dev/service"] // ""' <<<"${revision_json}")" [[ "${revision_service}" == "${SERVICE}" ]] \ || fail "${SAFETY_CHECK_REVISION} belongs to ${revision_service:-no service}, not ${SERVICE}" -if ! jq -e '.spec.containers[0].env // [] | any(.name == "SIMULATION_DATABASE_URL")' >/dev/null <<<"${revision_json}"; then - if gcloud scheduler jobs describe "${job}" "${scheduler_location[@]}" >/dev/null 2>&1; then - gcloud scheduler jobs pause "${job}" "${scheduler_location[@]}" --quiet - fi - printf '%s has no simulation database, so the safety check is paused.\n' "${SAFETY_CHECK_REVISION}" +# Every revision deployed with the safety check has the simulation database +# and the expected Modal caller (SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID, +# which the deploy has required since then); an older one has no job code. +if ! jq -e ' + .spec.containers[0].env // [] + | any(.name == "SIMULATION_DATABASE_URL") + and any(.name == "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID") +' >/dev/null <<<"${revision_json}"; then + pause_schedule + printf '%s predates the safety check, so its schedule is paused.\n' "${SAFETY_CHECK_REVISION}" exit 0 fi diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index aaee397f0..7324468fe 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -1042,12 +1042,17 @@ jobs: run: .github/scripts/set-cloud-run-simulation-entry-revision.sh - name: Point the safety check back at the previous revision + id: safety_check_restore if: ${{ failure() && steps.safety_check.outcome == 'failure' && env.PREVIOUS_REVISION != '' }} env: SAFETY_CHECK_REVISION: ${{ env.PREVIOUS_REVISION }} SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + - name: Pause the safety check after a failed restore + if: ${{ failure() && steps.safety_check_restore.outcome == 'failure' }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh pause + finalize_routing: name: Finalize the Modal routing publication # Finish the durable journal after every gate, including failed uploads diff --git a/.github/workflows/simulation-rollback.yml b/.github/workflows/simulation-rollback.yml index 46c165372..c51ef0f8e 100644 --- a/.github/workflows/simulation-rollback.yml +++ b/.github/workflows/simulation-rollback.yml @@ -219,12 +219,18 @@ jobs: fi - name: Point the safety check at the restored revision + id: safety_check if: ${{ inputs.entrypoint_revision != '' }} env: SAFETY_CHECK_REVISION: ${{ inputs.entrypoint_revision }} SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT: ${{ vars.SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT }} run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh + # The job must not keep running the image traffic just moved away from. + - name: Pause the safety check after a failed update + if: ${{ failure() && steps.safety_check.outcome == 'failure' }} + run: .github/scripts/deploy-cloud-run-simulation-safety-check.sh pause + - name: Verify the stable service serves the restored release if: ${{ !inputs.resolve_pending_publication }} env: diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index 00b7a1ebf..4f445fd98 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -112,10 +112,14 @@ records anything: - the promotion job, after the promoted revision passes every check; if this step fails, the previous revision is restored and the job is pointed back - at it; -- the rollback workflow, after it restores entrypoint traffic. A restored - revision without the simulation database predates the safety check, so the - schedule is paused. + at it, and if that fails too, the schedule is paused; +- the rollback workflow, after it restores entrypoint traffic; if this step + fails, the schedule is paused, so the job never keeps running the image + traffic just moved away from. + +A revision deployed before the safety check existed (without +`SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID`) has no job code, so pointing the +job at it pauses the schedule instead. `beta` does not promote its candidate yet, so staging has no safety check until it does (plan item 1.39). diff --git a/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py b/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py index f031bf68c..01b009b51 100644 --- a/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py +++ b/projects/policyengine-simulation-entry/tests/test_safety_check_deployment.py @@ -56,6 +56,10 @@ def _revision(*, service=SERVICE, environment=None): if environment is not None else [ {"name": "APP_ENVIRONMENT", "value": "production"}, + { + "name": "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID", + "value": "ac-workspace0001", + }, {"name": "OLD_GATEWAY_URL", "value": "https://legacy.example"}, {"name": "OBSERVABILITY_SERVICE_NAMESPACE", "value": "sim"}, {"name": "CLOUD_RUN_REGION", "value": "us-central1"}, @@ -286,3 +290,68 @@ def test_a_rollback_points_the_job_at_the_restored_revision() -> None: assert steps[safety]["env"]["SAFETY_CHECK_REVISION"] == ( "${{ inputs.entrypoint_revision }}" ) + + +def test_a_revision_from_before_the_safety_check_pauses_the_schedule(tmp_path) -> None: + """A 1.04 or 1.05 revision mounts the database and token but has no job code.""" + environment = [ + entry + for entry in _revision()["spec"]["containers"][0]["env"] + if entry["name"] != "SIMULATION_CALL_REPORT_MODAL_WORKSPACE_ID" + ] + + result, fakes, _ = _deploy( + tmp_path, _revision(environment=environment), schedule_exists=True + ) + + assert result.returncode == 0, result.stderr + assert [call[1:4] for call in fakes.calls("gcloud")] == [ + ["run", "revisions", "describe"], + ["scheduler", "jobs", "describe"], + ["scheduler", "jobs", "pause"], + ] + + +@pytest.mark.parametrize("schedule_exists", [True, False]) +def test_pausing_stops_the_schedule_without_a_revision( + tmp_path, schedule_exists +) -> None: + fakes, _ = _fakes(tmp_path, _revision(), schedule_exists=schedule_exists) + settings = { + "PROJECT_ID": "policyengine", + "REGION": "us-central1", + "SERVICE": SERVICE, + } + + result = run_script(SCRIPT, "pause", env=fakes.env(**settings)) + + assert result.returncode == 0, result.stderr + expected = [["scheduler", "jobs", "describe"]] + if schedule_exists: + expected.append(["scheduler", "jobs", "pause"]) + assert [call[1:4] for call in fakes.calls("gcloud")] == expected + + +def test_a_failed_update_in_a_rollback_pauses_the_safety_check() -> None: + """The job must not keep running the image traffic just moved away from.""" + steps = _steps(workflow_jobs(ROLLBACK_WORKFLOW)["rollback"]) + safety = _step_named(steps, "Point the safety check at the restored revision") + pause = _step_named(steps, "Pause the safety check after a failed update") + + assert safety < pause + assert steps[safety]["id"] == "safety_check" + assert steps[pause]["if"] == ( + "${{ failure() && steps.safety_check.outcome == 'failure' }}" + ) + assert steps[pause]["run"] == f".github/scripts/{SCRIPT} pause" + + +def test_a_failed_restore_in_a_promotion_pauses_the_safety_check() -> None: + steps = _steps(workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["promote_entrypoint"]) + restore = _step_named(steps, "Point the safety check back at the previous revision") + pause = _step_named(steps, "Pause the safety check after a failed restore") + + assert restore < pause + assert steps[restore]["id"] == "safety_check_restore" + assert "steps.safety_check_restore.outcome == 'failure'" in steps[pause]["if"] + assert steps[pause]["run"] == f".github/scripts/{SCRIPT} pause" From 2133796c36672b75f30a23745add87965dce4314 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:12:33 +0400 Subject: [PATCH 36/39] Log the safety check's passes through the observability runtime From the extra review of the reworked 1.06: the job configured logging with logging.basicConfig, whose plain format drops every field of a log record, so a pass's counts never reached the job's logs. The job now starts the process observability runtime, as the Modal workers do, logs JSON like the web service, and shuts the runtime down on exit. The pass summary carries its counts as policyengine_attributes, the only record fields the observability log handler keeps. Co-Authored-By: Claude Opus 5.5 --- .../safety_check.py | 42 ++++++++++++++++-- .../tests/test_safety_check.py | 44 ++++++++++++++++--- 2 files changed, 78 insertions(+), 8 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py index 378b123ef..d630ee44e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py @@ -29,6 +29,10 @@ from dataclasses import asdict from typing import TYPE_CHECKING, Protocol +from policyengine_simulation_observability.observability import ( + init_process_observability, +) + from policyengine_simulation_entry.modal_calls import OWN_PROCESS_SIGNALS from policyengine_simulation_entry.simulation_database import SimulationDatabase @@ -44,6 +48,23 @@ EXIT_NOT_CONFIGURED = 2 +class Telemetry(Protocol): + """The observability runtime's part the job uses.""" + + def shutdown(self) -> None: + """Flush and stop telemetry.""" + ... + + +def _start_observability(environment: str) -> Telemetry: + return init_process_observability( + service_name="policyengine-simulation-safety-check", + service_role="simulation_safety_check", + platform="google_cloud_run", + environment=environment, + ) + + class DueOutcomeRecorder(Protocol): """Run one recording pass (``RunRecorder``).""" @@ -60,13 +81,25 @@ def main( *, environment: Mapping[str, str] | None = None, build_recorder: Callable[[str], DueOutcomeRecorder] = _database_recorder, + start_observability: Callable[[str], Telemetry] = _start_observability, ) -> int: """Run one pass and return the process exit status. - Only a signal to this process propagates; any other error fails the + Logs go through the observability runtime, as the web service's do. Only + a signal to this process propagates; any other error fails the execution, and the next scheduled execution tries again. """ values = os.environ if environment is None else environment + telemetry = start_observability(values.get("APP_ENVIRONMENT", "local")) + try: + return _run_pass(values, build_recorder) + finally: + telemetry.shutdown() + + +def _run_pass( + values: Mapping[str, str], build_recorder: Callable[[str], DueOutcomeRecorder] +) -> int: database_url = values.get("SIMULATION_DATABASE_URL", "") if not database_url: logger.error("safety_check_not_configured") @@ -81,10 +114,13 @@ def main( "safety_check_pass_failed", extra={"error_type": type(error).__name__} ) return EXIT_PASS_FAILED - logger.info("safety_check_pass", extra=asdict(recording)) + # The observability log handler keeps a record's policyengine_attributes. + logger.info( + "safety_check_pass", extra={"policyengine_attributes": asdict(recording)} + ) return EXIT_PASSED if __name__ == "__main__": - logging.basicConfig(level=logging.INFO) + logging.getLogger().setLevel(logging.INFO) sys.exit(main()) diff --git a/projects/policyengine-simulation-entry/tests/test_safety_check.py b/projects/policyengine-simulation-entry/tests/test_safety_check.py index 20810acbd..18b49b0ef 100644 --- a/projects/policyengine-simulation-entry/tests/test_safety_check.py +++ b/projects/policyengine-simulation-entry/tests/test_safety_check.py @@ -32,14 +32,34 @@ async def record_due_outcomes(self) -> RecordingPass: return self.result -def _run(environment, recorder): +class _Telemetry: + """Stands in for the observability runtime the job starts and shuts down.""" + + def __init__(self) -> None: + self.environments = [] + self.shut_down = False + + def start(self, environment): + self.environments.append(environment) + return self + + def shutdown(self) -> None: + self.shut_down = True + + +def _run(environment, recorder, telemetry=None): built_for = [] def build(database_url): built_for.append(database_url) return recorder - status = safety_check.main(environment=environment, build_recorder=build) + telemetry = telemetry or _Telemetry() + status = safety_check.main( + environment=environment, + build_recorder=build, + start_observability=telemetry.start, + ) return status, built_for @@ -62,9 +82,23 @@ def test_the_pass_is_logged_with_what_it_found(caplog) -> None: (record,) = [ record for record in caplog.records if record.message == "safety_check_pass" ] - assert record.checked == 4 - assert record.succeeded == 1 - assert record.abandoned == 2 + # The observability log handler keeps only these attributes of a record. + assert record.policyengine_attributes["checked"] == 4 + assert record.policyengine_attributes["succeeded"] == 1 + assert record.policyengine_attributes["abandoned"] == 2 + + +def test_the_job_logs_through_the_observability_runtime_and_shuts_it_down() -> None: + telemetry = _Telemetry() + + _run( + {"SIMULATION_DATABASE_URL": DATABASE_URL, "APP_ENVIRONMENT": "production"}, + _OnePass(RecordingPass()), + telemetry, + ) + + assert telemetry.environments == ["production"] + assert telemetry.shut_down def test_an_execution_without_the_database_url_fails_without_a_pass() -> None: From 2c8c69215e006613ce942e5a66415982eddf06ed Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:12:53 +0400 Subject: [PATCH 37/39] Describe the dispatch, the report function and other_call as they work From the extra review of the reworked 1.06: run_dispatch still said a new call's first check is due shortly (it is the safety check's, one interval later) and that the recording loop collects a call whose record is unknown; report_call_ended's docstring said the reported call has always returned when sim asks Modal (usually it has; otherwise sim answers still_running and the function reports again); and other_call is also the answer for runs that finished without holding any call. Co-Authored-By: Claude Opus 5.5 --- .../call_reports.py | 3 ++- .../policyengine_simulation_entry/run_dispatch.py | 14 ++++++++------ .../src/modal/v2_app.py | 6 ++++-- 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py index 823457b70..4b9a32cf9 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/call_reports.py @@ -90,7 +90,8 @@ class CallReportStatus(StrEnum): NOT_CONFIGURED = "not_configured" # No run, or not every run, the report names exists. UNKNOWN_RUNS = "unknown_runs" - # The runs hold another Modal call. + # The runs hold another Modal call, or finished without holding one (for + # example failed as abandoned before the call's ID was stored). OTHER_CALL = "other_call" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py index 07a691b9c..acdc426a6 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_dispatch.py @@ -5,11 +5,12 @@ 1. claim: move the runs from ``pending`` to ``dispatching``. Only the instance whose claim changes them goes on; every other caller stops here. 2. spawn the Modal call and wait only for Modal's acknowledgement. -3. record the call ID, moving the runs to ``running`` with their first check - due shortly; or, when Modal refused the call, fail them. +3. record the call ID, moving the runs to ``running``, with the safety + check's first look due one interval later (the call itself reports when + it ends); or, when starting the call raised, fail them. -If the instance stops between 1 and 3, the runs stay ``dispatching`` until a -recording pass fails them as abandoned. No route starts calls yet. +If the instance stops between 1 and 3, the runs stay ``dispatching`` until the +safety check fails them as abandoned. No route starts calls yet. """ from __future__ import annotations @@ -59,8 +60,9 @@ class DispatchOutcome(StrEnum): # cancelled. NOT_RECORDED = "not_recorded" # Recording the call raised and the database could not say whether the - # record committed; the call is left running. If it did, the recording - # loop collects the call; if not, the runs are failed as abandoned. + # record committed; the call is left running. If it did, the call's report + # or the safety check records its outcome; if not, the safety check fails + # the runs as abandoned. RECORD_UNKNOWN = "record_unknown" diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 3789a940f..b0f1960d2 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -497,8 +497,10 @@ def report_call_ended(report: dict) -> str: """Tell sim that the call in ``report`` ended; return how that went. The v2 functions start this as their last step - (``src.modal.call_reports.reporting_call_end``), so the reported call - has returned by the time sim asks Modal for its outcome. + (``src.modal.call_reports.reporting_call_end``) and then return, so the + reported call has usually returned by the time sim asks Modal for its + outcome; when it has not, sim answers still_running and this reports + again after a wait. """ from src.modal.call_reports import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time send_call_ended_report_from_container, From 6e45ba341a1218f8caf2fab8c211ad230e66d889 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 28 Oct 2026 00:55:20 +0400 Subject: [PATCH 38/39] Let modal deploy import the v2 app: drop the check of report_call_ended's name v2_app checked report_call_ended.__name__ at import, after Modal's decorator had run. With the real modal package a declared function is a modal.Function, which has no __name__, so importing the module raised AttributeError and every `modal deploy` of the v2 app would have failed; the unit tests load the app with a fake Modal whose decorator returns the plain function, so they never saw it. The check is gone. The unit tests still require report_call_ended to be declared under the contract's CALL_REPORT_FUNCTION_NAME, and a new test imports the module with the real Modal package, as `modal deploy` does. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 6 ---- .../tests/test_call_report_function.py | 31 +++++++++++++++++++ 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index b0f1960d2..ceb8ce9d2 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -12,7 +12,6 @@ from pathlib import Path from typing import TYPE_CHECKING -from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME from policyengine_simulation_contract.stage12_manifest import v2_application_name from policyengine_simulation_executor import stage12_runtime from policyengine_simulation_executor.hf_access_validation import ( @@ -507,8 +506,3 @@ def report_call_ended(report: dict) -> str: ) return send_call_ended_report_from_container(report) - - -if report_call_ended.__name__ != CALL_REPORT_FUNCTION_NAME: - msg = "report_call_ended must keep the name sim accepts reports from" - raise RuntimeError(msg) diff --git a/projects/policyengine-simulation-executor/tests/test_call_report_function.py b/projects/policyengine-simulation-executor/tests/test_call_report_function.py index cf9937db0..9ce5cd33d 100644 --- a/projects/policyengine-simulation-executor/tests/test_call_report_function.py +++ b/projects/policyengine-simulation-executor/tests/test_call_report_function.py @@ -2,12 +2,18 @@ from __future__ import annotations +import os +import subprocess +import sys +from pathlib import Path + import pytest from policyengine_simulation_contract.call_reports import CALL_REPORT_FUNCTION_NAME from src.modal import call_reports from test_stage12_modal_app import _load ENTRY_URL = "https://policyengine-simulation-entry-abc.a.run.app" +PROJECT_ROOT = Path(__file__).resolve().parents[1] def test_the_v2_app_declares_the_function_sim_accepts_reports_from( @@ -51,3 +57,28 @@ def test_a_deploy_with_a_url_that_is_not_https_stops(monkeypatch, url) -> None: with pytest.raises(RuntimeError, match="SIMULATION_ENTRY_URL"): _load(monkeypatch) + + +def test_the_v2_app_imports_with_the_real_modal_package() -> None: + """`modal deploy` imports the module with Modal's own decorators. + + There a declared function is a ``modal.Function``, not the decorated + Python function, so nothing at import may rely on its Python attributes. + """ + environment = { + name: value + for name, value in os.environ.items() + if not name.startswith(("MODAL_", "STAGE12_", "SIMULATION_ENTRY_")) + } + + result = subprocess.run( + [sys.executable, "-c", "import src.modal.v2_app"], + cwd=PROJECT_ROOT, + env=environment, + capture_output=True, + text=True, + check=False, + timeout=300, + ) + + assert result.returncode == 0, result.stderr[-2000:] From 628572009767f72a57b024287895fb21e95d0de6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:43:15 +0400 Subject: [PATCH 39/39] Add changelog fragments for simulation runs, report runs and call reports Co-Authored-By: Claude Opus 5.5 --- changelog.d/947.added.md | 1 + changelog.d/947.changed.md | 3 +++ 2 files changed, 4 insertions(+) create mode 100644 changelog.d/947.added.md create mode 100644 changelog.d/947.changed.md diff --git a/changelog.d/947.added.md b/changelog.d/947.added.md new file mode 100644 index 000000000..20ff473cc --- /dev/null +++ b/changelog.d/947.added.md @@ -0,0 +1 @@ +Added simulation runs and report runs, in migration revision `0003` of the `simulation` schema. Each run has its own status and Modal call, every status change is one conditional update, and only the instance that claims a run starts its Modal call. When a v2 Modal call ends, the new `report_call_ended` Modal function reports it on `POST /internal/modal-calls/ended`, which accepts only a Modal identity token from that function; sim then reads the call's outcome from Modal and records it on the runs. A new Cloud Run job, `-safety-check`, runs every five minutes on Cloud Scheduler and records calls that ended without a report. No route creates runs yet. diff --git a/changelog.d/947.changed.md b/changelog.d/947.changed.md new file mode 100644 index 000000000..a3b61fcdc --- /dev/null +++ b/changelog.d/947.changed.md @@ -0,0 +1,3 @@ +Every Simulation Entrypoint deploy now mounts sim's own Modal token as `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`, no longer only when `DEPLOY_STAGE12_V2` is `true`. + +Every deploy now also needs, in both `beta` and `prod`, the `MODAL_TOKEN_ID_SECRET_NAME` and `MODAL_TOKEN_SECRET_SECRET_NAME` environment variables naming sim's own Modal token secrets, the `MODAL_WORKSPACE_ID` environment variable, and the `SIMULATION_SAFETY_CHECK_SCHEDULER_SERVICE_ACCOUNT` environment variable with a scheduler service account allowed to run the safety-check job. Without the first three the entrypoint deploy stops; without the last, the promotion's safety-check step fails and traffic goes back to the previous revision.