diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 667b222db..cd7d8e162 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -377,6 +377,17 @@ jobs: POLICYENGINE_OTEL_GOOGLE_AUDIENCE: ${{ vars.OBSERVABILITY_OTLP_GOOGLE_AUDIENCE }} run: uv run --frozen modal deploy --env="${{ inputs.modal_environment }}" src/modal/app.py + # The scheduled sweep that fails report runs and simulation runs nothing + # will finish. Its database URL comes from the Modal secret + # policyengine-modal-worker-database, so this step fails until an + # operator has created that secret in the Modal environment. + - name: Deploy the stuck-run sweep + working-directory: projects/policyengine-simulation-executor + env: + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + run: uv run --frozen modal deploy --env="${{ inputs.modal_environment }}" src/modal/run_sweep_app.py + - name: Run executor unit tests working-directory: projects/policyengine-simulation-executor run: uv run --frozen pytest tests/ -v diff --git a/changelog.d/stuck-run-sweep.added.md b/changelog.d/stuck-run-sweep.added.md new file mode 100644 index 000000000..063aa164f --- /dev/null +++ b/changelog.d/stuck-run-sweep.added.md @@ -0,0 +1 @@ +A scheduled Modal function, in the new Modal application `policyengine-simulation-run-sweep`, marks the Public API's report runs and simulation runs failed when nothing will finish them: a pending or running run whose recorded Modal call has ended, or a pending run that still has no Modal call recorded ten minutes after its creation. It checks every ten minutes and is safe to run twice at once. It leaves a run alone while Modal reports its call running, or can't report it, until the run is 24 hours old, counted from its first claim or, if it was never claimed, from its creation, longer than any of its calls can run; then it fails the run whatever Modal reports. The deployment workflow deploys it with the executor, and fails until the Modal secret `policyengine-modal-worker-database` exists in the Modal environment. diff --git a/libs/policyengine-stage12-persistence/README.md b/libs/policyengine-stage12-persistence/README.md index 40cb5aee3..a0bf80d21 100644 --- a/libs/policyengine-stage12-persistence/README.md +++ b/libs/policyengine-stage12-persistence/README.md @@ -11,9 +11,10 @@ revisions and live schema. This package never changes a schema. SPM receipts and data version, `simulation_runs`, and the result tables, the two labour supply tables included. The v2 Modal coordinator and workers write them as the Modal worker role (`policyengine_v2_modal_worker`), for a - call that names a report run. The mappings declare only the columns the - coordinator's and workers' statements read or write, and each `UPDATE` - sets only columns the role may update. Every status write is one `UPDATE` + call that names a report run, and the scheduled stuck-run sweep fails the + runs nothing will finish, as the same role. The mappings declare only the + columns those statements read or write, and each `UPDATE` sets only + columns the role may update. Every status write is one `UPDATE` conditional on the run's status; a report run's completion sets its own result columns and inserts its result rows in the same transaction. diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index 669f4ceb5..5f86c8e0b 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -6,7 +6,9 @@ report_run_result, ) from policyengine_stage12_persistence.report_run_store import ( + ActiveRun, ReportRunStore, + RunKind, RunStatus, RunWrite, ) @@ -22,10 +24,12 @@ ) __all__ = [ + "ActiveRun", "ConnectionOpener", "DistrictCodes", "ReportRunResult", "ReportRunStore", + "RunKind", "RunStatus", "RunWrite", "Stage12PersistenceStore", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index da0eeaeb4..9e5edeb1c 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -14,7 +14,7 @@ from enum import StrEnum from typing import TYPE_CHECKING, Any -from sqlalchemy import func, insert, or_, select, update +from sqlalchemy import exists, func, insert, or_, select, update from sqlalchemy.dialects.postgresql import insert as postgresql_insert from policyengine_stage12_persistence.report_run_tables import ( @@ -33,7 +33,7 @@ from uuid import UUID from policyengine_simulation_contract.stage12_execution import ReportRunIds - from sqlalchemy import Engine, Table, Update + from sqlalchemy import ColumnElement, Engine, Table, Update from sqlalchemy.engine import Connection from policyengine_stage12_persistence.report_run_results import ReportRunResult @@ -51,6 +51,29 @@ class RunStatus(StrEnum): ACTIVE_STATUSES = (RunStatus.PENDING.value, RunStatus.RUNNING.value) +class RunKind(StrEnum): + """Which run table a run is in.""" + + REPORT_RUN = "report_run" + SIMULATION_RUN = "simulation_run" + + +RUN_TABLES = {RunKind.REPORT_RUN: report_runs, RunKind.SIMULATION_RUN: simulation_runs} + + +@dataclass(frozen=True) +class ActiveRun: + """A pending or running run, and the Modal call recorded on it. + + ``modal_call_id`` is None for a run with no call recorded. + """ + + kind: RunKind + run_id: UUID + status: RunStatus + modal_call_id: str | None + + @dataclass(frozen=True) class RunWrite: """One conditional write's outcome: whether it changed the run, and its status.""" @@ -298,6 +321,92 @@ def fail_simulation_run( simulation_runs, simulation_run_id, statement, "fail_simulation_run" ) + def list_active_runs( + self, + kind: RunKind, + *, + active_since_before: datetime, + ) -> tuple[ActiveRun, ...]: + """Return the pending and running runs of ``kind`` active since before a time. + + A run counts from its first claim, or from its creation while it is + pending. Runs are oldest first, with or without a recorded Modal + call. A simulation run is listed only while no pending or running + report run names it as its baseline or reform: that report run's + coordinator starts its children's calls, records their outcomes, and + stops and starts them again when Modal restarts it, so they are its + to judge until it has finished. + """ + table = RUN_TABLES[kind] + active_since = func.coalesce(table.c.started_at, table.c.created_at) + conditions = [ + table.c.status.in_(ACTIVE_STATUSES), + active_since < active_since_before, + ] + if kind is RunKind.SIMULATION_RUN: + conditions.append(_no_active_report_run_names(table)) + statement = ( + select(table.c.id, table.c.status, table.c.modal_call_id) + .where(*conditions) + .order_by(active_since, table.c.id) + ) + with self._transaction("list_active_runs") as connection: + return tuple( + ActiveRun( + kind=kind, + run_id=row.id, + status=RunStatus(row.status), + modal_call_id=row.modal_call_id, + ) + for row in connection.execute(statement) + ) + + def fail_stuck_run( + self, + run: ActiveRun, + *, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark a stuck ``run`` failed, while it is as it was listed. + + It must still have its listed status and ``run.modal_call_id``, or no + call if it was listed with none, and a simulation run must still have + no pending or running report run naming it. A run that was claimed, + finished, or moved to another call, that the simulation service + recorded a call on, or that a pending or running report run has named + since it was listed, is left as it is, so concurrent sweeps change it + once. A report run naming it counts once its insert has committed, + before this update starts: the database doesn't make an insert in + progress wait for this update, or this update for it. + """ + table = RUN_TABLES[run.kind] + same_call = ( + table.c.modal_call_id.is_(None) + if run.modal_call_id is None + else table.c.modal_call_id == run.modal_call_id + ) + statement = ( + update(table) + .where( + table.c.id == run.run_id, + table.c.status == run.status.value, + same_call, + *( + (_no_active_report_run_names(table),) + if run.kind is RunKind.SIMULATION_RUN + else () + ), + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write(table, run.run_id, statement, "fail_stuck_run") + def _write( self, table: Table, @@ -351,3 +460,17 @@ def _uniform_rows(rows: tuple[dict[str, Any], ...]) -> list[dict[str, Any]]: """Give every row the same columns, unset ones null, for one batched insert.""" columns = {name for row in rows for name in row} return [dict.fromkeys(columns) | row for row in rows] + + +def _no_active_report_run_names(table: Table) -> ColumnElement[bool]: + """Return the condition that no pending or running report run names the run. + + A report run names its simulation runs as its baseline and its reform. + """ + return ~exists().where( + or_( + report_runs.c.baseline_simulation_run_id == table.c.id, + report_runs.c.reform_simulation_run_id == table.c.id, + ), + report_runs.c.status.in_(ACTIVE_STATUSES), + ) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py index 91728e95c..99cb71418 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -1,9 +1,9 @@ """SQLAlchemy mappings of the Public API's run tables that report calls update. ``policyengine-api`` owns these tables and every schema migration. The -mappings declare only the columns the coordinator and workers read or write -as the Modal worker role, so they can issue typed DML; runtime code never -emits DDL. +mappings declare only the columns the coordinator, the workers and the +stuck-run sweep read or write as the Modal worker role, so they can issue +typed DML; runtime code never emits DDL. """ from __future__ import annotations @@ -33,6 +33,7 @@ "report_runs", metadata, Column("id", UUID(as_uuid=True), primary_key=True), + Column("created_at", DateTime(timezone=True), nullable=False), Column("baseline_simulation_run_id", UUID(as_uuid=True), nullable=False), Column("reform_simulation_run_id", UUID(as_uuid=True)), Column("status", report_run_status_type, nullable=False), @@ -53,6 +54,7 @@ "simulation_runs", metadata, Column("id", UUID(as_uuid=True), primary_key=True), + Column("created_at", DateTime(timezone=True), nullable=False), Column("status", simulation_run_status_type, nullable=False), Column("modal_call_id", String(255)), Column("error_message", String), diff --git a/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py new file mode 100644 index 000000000..492090f74 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py @@ -0,0 +1,616 @@ +"""The stuck-run sweep's reads and writes, as the Modal worker role. + +These run against a disposable Postgres that the Public API's v2 migrations +created (see ``report_run_database``). +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier +from uuid import uuid4 + +import pytest +from policyengine_stage12_persistence import ( + ActiveRun, + ReportRunStore, + RunKind, + RunStatus, + RunWrite, +) +from report_run_database import NOW, SeededReportRun, worker_database_url +from sqlalchemy import insert + +CUTOFF = NOW + timedelta(minutes=30) +ENDED_MESSAGE = "modal_call_ended: the run's Modal call ended without recording it" +NOT_STARTED_MESSAGE = "run_not_started: no Modal call was recorded for the run" + + +@pytest.fixture +def store(): + run_store = ReportRunStore(worker_database_url()) + try: + yield run_store + finally: + run_store.engine.dispose() + + +def _set_run(seeder, table_name, run_id, **values): + table = seeder.table(table_name) + with seeder.engine.begin() as connection: + connection.execute(table.update().where(table.c.id == run_id).values(**values)) + + +def _running(seeder, seeded_run: SeededReportRun): + """Mark the report run and baseline run running under calls started at NOW.""" + run_ids = seeded_run.run_ids + _set_run( + seeder, + "report_runs", + run_ids.report_run_id, + status="running", + modal_call_id="fc-coordinator", + started_at=NOW, + ) + _set_run( + seeder, + "simulation_runs", + run_ids.baseline_simulation_run_id, + status="running", + modal_call_id="fc-baseline", + started_at=NOW, + ) + return ( + ActiveRun( + kind=RunKind.REPORT_RUN, + run_id=run_ids.report_run_id, + status=RunStatus.RUNNING, + modal_call_id="fc-coordinator", + ), + ActiveRun( + kind=RunKind.SIMULATION_RUN, + run_id=run_ids.baseline_simulation_run_id, + status=RunStatus.RUNNING, + modal_call_id="fc-baseline", + ), + ) + + +def _listed(store, cutoff): + return ( + *store.list_active_runs(RunKind.REPORT_RUN, active_since_before=cutoff), + *store.list_active_runs(RunKind.SIMULATION_RUN, active_since_before=cutoff), + ) + + +def _listed_for(store, seeded_run, cutoff=CUTOFF): + ids = { + seeded_run.run_ids.report_run_id, + seeded_run.run_ids.baseline_simulation_run_id, + seeded_run.run_ids.reform_simulation_run_id, + } + return {run for run in _listed(store, cutoff) if run.run_id in ids} + + +def test_the_sweep_lists_runs_running_since_before_the_cutoff( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run, baseline_run = _running(seeder, seeded_run) + # Only once its report run is no longer running is a simulation run's + # call the sweep's to check. + assert _listed_for(store, seeded_run) == {report_run} + _set_run( + seeder, + "report_runs", + seeded_run.run_ids.report_run_id, + status="failed", + completed_at=CUTOFF, + ) + # The reform run is running, but its call has no ID recorded yet. + reform_id = seeded_run.run_ids.reform_simulation_run_id + _set_run(seeder, "simulation_runs", reform_id, status="running", started_at=NOW) + + assert _listed_for(store, seeded_run) == { + baseline_run, + ActiveRun( + kind=RunKind.SIMULATION_RUN, + run_id=reform_id, + status=RunStatus.RUNNING, + modal_call_id=None, + ), + } + assert _listed_for(store, seeded_run, NOW) == set() + + +def test_a_pending_run_is_listed_from_its_creation( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + _set_run(seeder, "report_runs", report_run_id, created_at=NOW) + unstarted = ActiveRun( + kind=RunKind.REPORT_RUN, + run_id=report_run_id, + status=RunStatus.PENDING, + modal_call_id=None, + ) + # The simulation runs are pending too, but the pending report run's + # coordinator will start them. + for simulation_run_id in ( + seeded_run.run_ids.baseline_simulation_run_id, + seeded_run.run_ids.reform_simulation_run_id, + ): + _set_run(seeder, "simulation_runs", simulation_run_id, created_at=NOW) + + assert _listed_for(store, seeded_run) == {unstarted} + assert _listed_for(store, seeded_run, NOW) == set() + + # The simulation service recorded the call it started. + _set_run(seeder, "report_runs", report_run_id, modal_call_id="fc-coordinator") + assert _listed_for(store, seeded_run) == { + ActiveRun( + kind=RunKind.REPORT_RUN, + run_id=report_run_id, + status=RunStatus.PENDING, + modal_call_id="fc-coordinator", + ) + } + + +def test_never_started_simulation_runs_are_listed_once_their_report_run_finished( + store, seeder, seeded_run: SeededReportRun +) -> None: + # A coordinator that died before starting its children leaves them + # pending with no call; once its report run has failed, they are listed. + run_ids = seeded_run.run_ids + for simulation_run_id in ( + run_ids.baseline_simulation_run_id, + run_ids.reform_simulation_run_id, + ): + _set_run(seeder, "simulation_runs", simulation_run_id, created_at=NOW) + _set_run( + seeder, + "report_runs", + run_ids.report_run_id, + status="failed", + completed_at=CUTOFF, + ) + + assert _listed_for(store, seeded_run) == { + ActiveRun( + kind=RunKind.SIMULATION_RUN, + run_id=simulation_run_id, + status=RunStatus.PENDING, + modal_call_id=None, + ) + for simulation_run_id in ( + run_ids.baseline_simulation_run_id, + run_ids.reform_simulation_run_id, + ) + } + + +def test_the_sweep_fails_a_run_once_while_its_ended_call_still_holds_it( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run, baseline_run = _running(seeder, seeded_run) + + first = store.fail_stuck_run( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + again = store.fail_stuck_run( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + simulation = store.fail_stuck_run( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert first == RunWrite(changed=True, status=RunStatus.FAILED) + assert again == RunWrite(changed=False, status=RunStatus.FAILED) + assert simulation == RunWrite(changed=True, status=RunStatus.FAILED) + row = seeder.row("report_runs", seeded_run.run_ids.report_run_id) + assert row["error_message"] == ENDED_MESSAGE + assert row["completed_at"] == CUTOFF + + +def test_the_sweep_leaves_a_run_that_moved_to_another_call_or_finished( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run, baseline_run = _running(seeder, seeded_run) + # After the sweep listed them, Modal restarted the baseline's worker + # under a new call, and the report run succeeded. + _set_run( + seeder, + "simulation_runs", + seeded_run.run_ids.baseline_simulation_run_id, + modal_call_id="fc-baseline-restarted", + ) + _set_run( + seeder, + "report_runs", + seeded_run.run_ids.report_run_id, + status="succeeded", + completed_at=CUTOFF, + ) + + moved = store.fail_stuck_run( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + finished = store.fail_stuck_run( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert moved == RunWrite(changed=False, status=RunStatus.RUNNING) + assert finished == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + + +def test_a_pending_run_is_failed_only_while_pending_under_its_listed_call( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + _set_run(seeder, "report_runs", report_run_id, created_at=NOW) + (listed_without_call,) = _listed_for(store, seeded_run) + # After the listing, the simulation service recorded the call it started. + _set_run(seeder, "report_runs", report_run_id, modal_call_id="fc-coordinator") + late_call = store.fail_stuck_run( + listed_without_call, error_message=NOT_STARTED_MESSAGE, failed_at=CUTOFF + ) + (listed_with_call,) = _listed_for(store, seeded_run) + # After the listing, the call claimed the run. + _set_run( + seeder, + "report_runs", + report_run_id, + status="running", + started_at=CUTOFF, + ) + claimed = store.fail_stuck_run( + listed_with_call, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert listed_without_call.modal_call_id is None + assert late_call == RunWrite(changed=False, status=RunStatus.PENDING) + assert listed_with_call.modal_call_id == "fc-coordinator" + assert claimed == RunWrite(changed=False, status=RunStatus.RUNNING) + assert seeder.row("report_runs", report_run_id)["error_message"] is None + + +def test_two_sweeps_at_once_fail_each_run_exactly_once( + seeder, seeded_run: SeededReportRun +) -> None: + reform_id = seeded_run.run_ids.reform_simulation_run_id + _set_run(seeder, "simulation_runs", reform_id, created_at=NOW) + runs = ( + *_running(seeder, seeded_run), + ActiveRun( + kind=RunKind.SIMULATION_RUN, + run_id=reform_id, + status=RunStatus.PENDING, + modal_call_id=None, + ), + ) + stores = [ReportRunStore(worker_database_url()) for _ in range(2)] + start_together = Barrier(len(stores)) + + def sweep(run_store): + start_together.wait() + return [ + run_store.fail_stuck_run( + run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ).changed + for run in runs + ] + + try: + with ThreadPoolExecutor(max_workers=len(stores)) as pool: + changes = list(pool.map(sweep, stores)) + finally: + for run_store in stores: + run_store.engine.dispose() + + assert [sum(column) for column in zip(*changes, strict=True)] == [1, 1, 1] + assert seeder.row("report_runs", seeded_run.run_ids.report_run_id)["status"] == ( + "failed" + ) + assert seeder.row("simulation_runs", reform_id)["status"] == "failed" + + +_SAME_RUN = object() + + +def _second_report_run( + seeder, + seeded_run: SeededReportRun, + status, + *, + baseline_simulation_run_id=_SAME_RUN, + reform_simulation_run_id=_SAME_RUN, +): + """Insert another report run of the same report. + + It names the first report run's simulation runs unless given others; a + null reform run is a report with only a baseline. + """ + first = seeder.row("report_runs", seeded_run.run_ids.report_run_id) + second_id = uuid4() + with seeder.engine.begin() as connection: + connection.execute( + insert(seeder.table("report_runs")).values( + id=second_id, + report_id=first["report_id"], + baseline_simulation_run_id=( + first["baseline_simulation_run_id"] + if baseline_simulation_run_id is _SAME_RUN + else baseline_simulation_run_id + ), + reform_simulation_run_id=( + first["reform_simulation_run_id"] + if reform_simulation_run_id is _SAME_RUN + else reform_simulation_run_id + ), + status=status, + trigger="manual", + country_package_version=first["country_package_version"], + policyengine_version=first["policyengine_version"], + **( + {"completed_at": CUTOFF} + if status in {"succeeded", "failed"} + else {} + ), + ) + ) + return second_id + + +# What another simulation run of the seeded baseline's simulation copies. +SIMULATION_RUN_DEFINITION = ( + "simulation_id", + "year", + "tax_benefit_model_version_id", + "dataset_id", + "country_package_version", + "policyengine_version", +) + + +def _another_simulation_run(seeder, seeded_run: SeededReportRun, **state): + """Insert another simulation run of the seeded baseline's simulation.""" + baseline = seeder.row( + "simulation_runs", seeded_run.run_ids.baseline_simulation_run_id + ) + run_id = uuid4() + with seeder.engine.begin() as connection: + connection.execute( + insert(seeder.table("simulation_runs")).values( + id=run_id, + **{name: baseline[name] for name in SIMULATION_RUN_DEFINITION}, + **state, + ) + ) + return run_id + + +def _finished_simulation_run(seeder, seeded_run: SeededReportRun): + """Insert a succeeded simulation run of the seeded baseline's simulation.""" + return _another_simulation_run( + seeder, seeded_run, status="succeeded", started_at=NOW, completed_at=NOW + ) + + +def _simulation_run_listed(store, simulation_run_id): + return any( + run.run_id == simulation_run_id + for run in store.list_active_runs( + RunKind.SIMULATION_RUN, active_since_before=CUTOFF + ) + ) + + +def _baseline_listed(store, seeded_run): + return _simulation_run_listed(store, seeded_run.run_ids.baseline_simulation_run_id) + + +def test_a_simulation_run_is_listed_unless_a_pending_or_running_report_run_names_it( + store, seeder, seeded_run: SeededReportRun +) -> None: + _running(seeder, seeded_run) + # The first report run is running, so its coordinator judges the run. + assert not _baseline_listed(store, seeded_run) + + # A pending report run's coordinator is still to start it. + _set_run(seeder, "report_runs", seeded_run.run_ids.report_run_id, status="pending") + assert not _baseline_listed(store, seeded_run) + + _set_run( + seeder, + "report_runs", + seeded_run.run_ids.report_run_id, + status="failed", + completed_at=CUTOFF, + ) + assert _baseline_listed(store, seeded_run) + + second = _second_report_run(seeder, seeded_run, "running") + assert not _baseline_listed(store, seeded_run) + + _set_run(seeder, "report_runs", second, status="succeeded", completed_at=CUTOFF) + assert _baseline_listed(store, seeded_run) + + +def test_a_report_run_keeps_its_reform_run_out_of_the_sweep_until_it_finishes( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + reform_id = seeded_run.run_ids.reform_simulation_run_id + _set_run( + seeder, + "report_runs", + report_run_id, + status="running", + modal_call_id="fc-coordinator", + started_at=NOW, + ) + _set_run( + seeder, + "simulation_runs", + reform_id, + status="running", + modal_call_id="fc-reform", + started_at=NOW, + ) + # The running report run names it as its reform, so its coordinator + # judges the run. + assert not _simulation_run_listed(store, reform_id) + + _set_run(seeder, "report_runs", report_run_id, status="failed", completed_at=CUTOFF) + assert _simulation_run_listed(store, reform_id) + + +def test_a_run_with_no_call_is_listed_and_failed_while_it_still_has_none( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run_id = seeded_run.run_ids.report_run_id + _set_run(seeder, "report_runs", report_run_id, status="running", started_at=NOW) + + (unrecorded,) = [ + run + for run in store.list_active_runs( + RunKind.REPORT_RUN, active_since_before=CUTOFF + ) + if run.run_id == report_run_id + ] + failed = store.fail_stuck_run( + unrecorded, error_message="run_exceeded_maximum_age: too old", failed_at=CUTOFF + ) + + assert unrecorded.modal_call_id is None + assert unrecorded.status is RunStatus.RUNNING + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + + +@pytest.mark.parametrize("second_status", ["pending", "running"]) +def test_a_simulation_run_a_report_run_now_names_is_not_failed( + store, seeder, seeded_run: SeededReportRun, second_status +) -> None: + _, baseline_run = _running(seeder, seeded_run) + _set_run( + seeder, + "report_runs", + seeded_run.run_ids.report_run_id, + status="failed", + completed_at=CUTOFF, + ) + assert _baseline_listed(store, seeded_run) + # After the listing, another report run naming it is created or starts + # running; its coordinator now judges the simulation run. + _second_report_run(seeder, seeded_run, second_status) + + write = store.fail_stuck_run( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert write == RunWrite(changed=False, status=RunStatus.RUNNING) + + +@pytest.mark.parametrize("second_status", ["pending", "running"]) +def test_a_reform_run_a_report_run_now_names_is_not_failed( + store, seeder, seeded_run: SeededReportRun, second_status +) -> None: + run_ids = seeded_run.run_ids + _set_run( + seeder, + "simulation_runs", + run_ids.reform_simulation_run_id, + status="running", + modal_call_id="fc-reform", + started_at=NOW, + ) + _set_run( + seeder, + "report_runs", + run_ids.report_run_id, + status="failed", + completed_at=CUTOFF, + ) + (reform_run,) = [ + run + for run in store.list_active_runs( + RunKind.SIMULATION_RUN, active_since_before=CUTOFF + ) + if run.run_id == run_ids.reform_simulation_run_id + ] + # After the listing, another report run names it, as its reform only. + _second_report_run( + seeder, + seeded_run, + second_status, + baseline_simulation_run_id=_finished_simulation_run(seeder, seeded_run), + ) + + write = store.fail_stuck_run( + reform_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert write == RunWrite(changed=False, status=RunStatus.RUNNING) + + +def test_a_report_run_with_only_a_baseline_keeps_only_its_baseline_out( + store, seeder, seeded_run: SeededReportRun +) -> None: + run_ids = seeded_run.run_ids + _running(seeder, seeded_run) + _set_run( + seeder, + "simulation_runs", + run_ids.reform_simulation_run_id, + status="running", + modal_call_id="fc-reform", + started_at=NOW, + ) + _set_run( + seeder, + "report_runs", + run_ids.report_run_id, + status="failed", + completed_at=CUTOFF, + ) + _second_report_run(seeder, seeded_run, "running", reform_simulation_run_id=None) + + assert not _baseline_listed(store, seeded_run) + assert _simulation_run_listed(store, run_ids.reform_simulation_run_id) + + +def test_a_simulation_run_no_report_run_names_is_listed_and_failed( + store, seeder, seeded_run: SeededReportRun +) -> None: + # A simulation run the simulation service starts on its own, with no + # report run naming it, is the sweep's to judge from the start. + unnamed_id = _another_simulation_run( + seeder, + seeded_run, + status="running", + modal_call_id="fc-unnamed", + started_at=NOW, + ) + + listed = [ + run + for run in store.list_active_runs( + RunKind.SIMULATION_RUN, active_since_before=CUTOFF + ) + if run.run_id == unnamed_id + ] + failed = [ + store.fail_stuck_run(run, error_message=ENDED_MESSAGE, failed_at=CUTOFF) + for run in listed + ] + + assert listed == [ + ActiveRun( + kind=RunKind.SIMULATION_RUN, + run_id=unnamed_id, + status=RunStatus.RUNNING, + modal_call_id="fc-unnamed", + ) + ] + assert failed == [RunWrite(changed=True, status=RunStatus.FAILED)] diff --git a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py new file mode 100644 index 000000000..98dd27d45 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -0,0 +1,83 @@ +"""A scheduled Modal function that fails runs nothing will finish. + +Every ten minutes it lists the Public API's report runs and simulation runs +that have been pending or running for at least ten minutes, counted from +their first claim or, while pending, from their creation. It asks Modal for +the status of each recorded call and marks a run failed once its call has +ended, and fails a pending run that has no call recorded, because nothing +will run it. It also fails any run still unfinished 24 hours after it was +created or started, longer than any Modal call of it can run, whatever Modal +reports. It connects as the Modal worker role, from the same Modal secret the +v2 coordinator and workers use. This application is not versioned with the +.py bundle: one copy runs per Modal environment. + +The deployment workflow's executor job deploys it to each Modal environment, +after the versioned executor (``simulation-deploy.reusable.yml``). Modal +refuses the deployment until the Modal secret exists in that environment. +""" + +from __future__ import annotations + +from dataclasses import asdict +from pathlib import Path + +from policyengine_simulation_executor.stage12_runtime.report_run_access import ( + MODAL_WORKER_DATABASE_SECRET_NAME, + MODAL_WORKER_DATABASE_URL_ENV, +) +from src.modal.ca_certificates import GRPC_CA_BUNDLE_ENV, OPERATING_SYSTEM_CA_PACKAGE + +import modal + +APP_NAME = "policyengine-simulation-run-sweep" +_UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." + +app = modal.App(APP_NAME) +modal_worker_database_secret = modal.Secret.from_name( + MODAL_WORKER_DATABASE_SECRET_NAME, + required_keys=[MODAL_WORKER_DATABASE_URL_ENV], +) +sweep_image = ( + modal.Image.debian_slim(python_version="3.13") + # The same CA store and locked packages as the v2 images, without the + # bundle's data, which the sweep never reads. + .apt_install(OPERATING_SYSTEM_CA_PACKAGE) + .env(GRPC_CA_BUNDLE_ENV) + .uv_sync( + uv_project_dir=_UV_PROJECT_DIR, + frozen=True, + extra_options="--only-group modal-simulation-image", + # The repository's uv release, from .uv-version. + uv_version="0.13.6", + ) + .add_local_python_source( + # This module imports src.modal, so the container needs it too. + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + copy=True, + ) +) + + +@app.function( + image=sweep_image, + schedule=modal.Cron("*/10 * * * *"), + cpu=0.25, + memory=512, + timeout=600, + retries=0, + max_containers=1, + secrets=[modal_worker_database_secret], +) +def sweep_stuck_runs() -> dict: + """Fail the runs nothing will finish; return how many were judged and why.""" + from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the container, not at deploy time + report_run_access, + run_sweep, + ) + + summary = run_sweep.sweep_stuck_runs(report_run_access.report_run_store()) + return asdict(summary) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/run_sweep.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/run_sweep.py new file mode 100644 index 000000000..a4e1ada45 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/run_sweep.py @@ -0,0 +1,267 @@ +"""The sweep that fails report runs and simulation runs nothing will finish. + +A coordinator or worker records the Modal call running each run it claims, +and the simulation service is to record the call it starts on a pending +report run. If that call ends without recording the run's outcome, for +example because its container ran out of memory, the run stays pending or +running. The sweep asks Modal for the status of each recorded call and marks +the run failed once the call has ended. A pending run with no call recorded +soon after its creation is failed without asking Modal: usually the +simulation service stopped before starting its call. If it started the call +but failed to record it, that call's claim is refused once the run has +failed, so it finishes nothing. The sweep also fails a run +still unfinished long after it was created or started, longer than any Modal +call of it can run, whatever Modal reports and with or without a recorded +call, so a run Modal can't decide doesn't stay unfinished for good. + +Modal restarts a pre-empted call under the same call ID and reports it +pending meanwhile, so a pre-empted coordinator's report run is never failed. +Its children are cancelled and started again under new calls, and it starts +them only once Modal runs it, so the sweep judges a simulation run only once +no pending or running report run names it as its baseline or reform. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from enum import StrEnum +from importlib import import_module +from typing import TYPE_CHECKING, Any, Protocol + +from policyengine_stage12_persistence import RunKind, RunStatus + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_stage12_persistence import ActiveRun, RunWrite + +logger = logging.getLogger(__name__) + +# A run must have been unfinished at least this long, counted from its first +# claim or, while it is pending, from its creation, before the sweep judges +# it; a pending run with no call recorded by then is failed. +DEFAULT_MINIMUM_AGE = timedelta(minutes=10) +ENDED_CALL_MESSAGE = ( + "modal_call_ended: The Modal call recorded for this run ended without " + "recording the run's outcome." +) +NOT_STARTED_MESSAGE = ( + "run_not_started: The run stayed pending with no Modal call recorded for it." +) +# A run still unfinished this long after its first claim, or its creation if +# it was never claimed, is failed whatever Modal reports, or with no call +# recorded at all. No attempt of the v2 +# functions can run anywhere near this long (the coordinator's own time +# limit is 90 minutes), however often Modal pre-empts and restarts it; it +# stops Modal's pending status for an internal failure it won't retry, or +# a call Modal no longer describes, from keeping a run running for good. +DEFAULT_MAXIMUM_AGE = timedelta(hours=24) +EXCEEDED_AGE_MESSAGE = ( + "run_exceeded_maximum_age: The run was still unfinished long after it " + "was created or started, longer than any Modal call of it can run." +) +# Modal's input statuses (``modal.call_graph.InputStatus``) by name. A pending +# input is running or waiting to run again; Modal reports an input pre-empted +# and waiting to restart (internally INTERNAL_FAILURE) as pending. Any other +# name leaves the call's state unknown. +RUNNING_INPUT_STATUSES = frozenset({"PENDING"}) +ENDED_INPUT_STATUSES = frozenset( + {"SUCCESS", "FAILURE", "TERMINATED", "TIMEOUT", "INIT_FAILURE"} +) + + +class CallState(StrEnum): + """What Modal reports about a recorded call.""" + + RUNNING = "running" + ENDED = "ended" + UNKNOWN = "unknown" + + +class StuckRunStore(Protocol): + """The sweep's reads and writes; ``ReportRunStore`` implements them.""" + + def list_active_runs( + self, kind: RunKind, *, active_since_before: datetime + ) -> tuple[ActiveRun, ...]: + """Return the pending and running runs of ``kind`` older than a time.""" + ... + + def fail_stuck_run( + self, run: ActiveRun, *, error_message: str, failed_at: datetime + ) -> RunWrite: + """Mark ``run`` failed while it is as it was listed. + + It must still have its listed status and call, or no call if it was + listed with none; a simulation run must still have no pending or + running report run naming it. + """ + ... + + +@dataclass(frozen=True) +class SweepSummary: + """What the sweep did with the runs it listed. + + ``checked`` runs were judged after the minimum age: ``failed`` because + their call ended, ``failed_without_call`` because they stayed pending with + no call recorded, or left because their call is ``running`` or its state + is ``unknown``, as is a running run with no call recorded. + ``failed_for_age`` runs were failed for exceeding the maximum age. ``left`` + counts listed runs whose conditional failure changed nothing, because + they were claimed, finished, moved to another call, had a call recorded, + or a pending or running report run named them after they were listed. + """ + + checked: int + failed: int + running: int + unknown: int + left: int = 0 + failed_for_age: int = 0 + failed_without_call: int = 0 + + +def modal_call_state(modal_call_id: str, *, modal_module: Any = None) -> CallState: # noqa: ANN401 # the modal module, or a stand-in for it + """Return whether Modal reports the call as running or ended. + + Modal returns the whole call graph the call is part of, where a call + spawned inside another call's input is nested under it, so the call's own + inputs are found by their call ID anywhere in the graph. A call whose + inputs are all finished has ended; one with any pending input is running. + A call missing from its graph, a status this code doesn't know, or any + error asking, is unknown, so the sweep leaves its run alone. + """ + modal = modal_module if modal_module is not None else import_module("modal") + try: + graph = modal.FunctionCall.from_id(modal_call_id).get_call_graph() + statuses = { + input_info.status.name for input_info in _inputs_of(graph, modal_call_id) + } + except Exception as error: # noqa: BLE001 # any failure to ask leaves the run as it is + _log_unreadable_call(modal_call_id, error) + return CallState.UNKNOWN + if statuses & RUNNING_INPUT_STATUSES: + return CallState.RUNNING + if statuses and statuses <= ENDED_INPUT_STATUSES: + return CallState.ENDED + return CallState.UNKNOWN + + +def _inputs_of(graph: list[Any], modal_call_id: str) -> list[Any]: + """Return every input in ``graph`` that belongs to call ``modal_call_id``.""" + found: list[Any] = [] + pending = list(graph) + while pending: + input_info = pending.pop() + if input_info.function_call_id == modal_call_id: + found.append(input_info) + pending.extend(input_info.children) + return found + + +@dataclass +class _Tally: + """Running counts of one sweep.""" + + checked: int = 0 + failed: int = 0 + running: int = 0 + unknown: int = 0 + left: int = 0 + failed_for_age: int = 0 + failed_without_call: int = 0 + + +def sweep_stuck_runs( # the store, then keyword-only seams and limits + store: StuckRunStore, + *, + call_state: Callable[[str], CallState] = modal_call_state, + now: datetime | None = None, + minimum_age: timedelta = DEFAULT_MINIMUM_AGE, + maximum_age: timedelta = DEFAULT_MAXIMUM_AGE, +) -> SweepSummary: + """Fail every listed run whose recorded Modal call has ended or was never started. + + Runs older than ``maximum_age`` are failed first, report runs then + simulation runs, before any Modal lookup can use up the scheduled + function's time. Then, report runs first, so a dead coordinator's + simulation runs are listed in the same sweep, each run older than + ``minimum_age`` is judged: a pending run with no call recorded is failed, + and each recorded call is asked about once. A run whose call is running, + or whose state is unknown, is left as it is, as is a running run with no + call recorded, until the maximum age; so is one whose conditional failure + changes nothing, which makes two sweeps at once safe. + """ + now = now or datetime.now(UTC) + states: dict[str, CallState] = {} + tally = _Tally() + for kind in (RunKind.REPORT_RUN, RunKind.SIMULATION_RUN): + for run in store.list_active_runs(kind, active_since_before=now - maximum_age): + if store.fail_stuck_run( + run, error_message=EXCEEDED_AGE_MESSAGE, failed_at=now + ).changed: + tally.failed_for_age += 1 + else: + tally.left += 1 + for kind in (RunKind.REPORT_RUN, RunKind.SIMULATION_RUN): + for run in store.list_active_runs(kind, active_since_before=now - minimum_age): + _sweep_run( + store, run, states=states, call_state=call_state, now=now, tally=tally + ) + return SweepSummary(**vars(tally)) + + +def _sweep_run( # noqa: PLR0913 # one run, the store, and the sweep's shared state + store: StuckRunStore, + run: ActiveRun, + *, + states: dict[str, CallState], + call_state: Callable[[str], CallState], + now: datetime, + tally: _Tally, +) -> None: + """Fail ``run`` if its call has ended or was never started. + + Each call is asked about once. + """ + tally.checked += 1 + modal_call_id = run.modal_call_id + if modal_call_id is None: + if run.status is not RunStatus.PENDING: + # Nothing to ask Modal about; the maximum age decides it. + tally.unknown += 1 + elif store.fail_stuck_run( + run, error_message=NOT_STARTED_MESSAGE, failed_at=now + ).changed: + tally.failed_without_call += 1 + else: + tally.left += 1 + return + if modal_call_id not in states: + try: + states[modal_call_id] = call_state(modal_call_id) + except Exception as error: # noqa: BLE001 # a failed lookup must never fail the run + _log_unreadable_call(modal_call_id, error) + states[modal_call_id] = CallState.UNKNOWN + state = states[modal_call_id] + if state is CallState.RUNNING: + tally.running += 1 + elif state is CallState.UNKNOWN: + tally.unknown += 1 + elif store.fail_stuck_run( + run, error_message=ENDED_CALL_MESSAGE, failed_at=now + ).changed: + tally.failed += 1 + else: + tally.left += 1 + + +def _log_unreadable_call(modal_call_id: str, error: Exception) -> None: + logger.warning( + "Could not read the status of Modal call %s (%s)", + modal_call_id, + type(error).__name__, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py new file mode 100644 index 000000000..4c03fabdd --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -0,0 +1,383 @@ +"""The sweep that fails runs left running after their Modal call ended.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from enum import Enum +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from policyengine_simulation_executor.stage12_runtime.run_sweep import ( + DEFAULT_MAXIMUM_AGE, + CallState, + SweepSummary, + modal_call_state, + sweep_stuck_runs, +) +from policyengine_stage12_persistence import ActiveRun, RunKind, RunStatus, RunWrite + +NOW = datetime(2026, 11, 4, 12, tzinfo=UTC) + + +class FakeInputStatus(Enum): + """Stands in for ``modal.call_graph.InputStatus``; the sweep reads names.""" + + PENDING = 0 + SUCCESS = 1 + FAILURE = 2 + TERMINATED = 3 + TIMEOUT = 4 + INIT_FAILURE = 5 + + +class FakeSweepStore: + def __init__(self, runs, *, expired=(), maximum_age=DEFAULT_MAXIMUM_AGE): + self.runs = list(runs) + # Runs older than the maximum age, which the sweep lists first, at + # exactly this cutoff; every other listing is for the minimum age. + self.expired = list(expired) + self.maximum_age_cutoff = NOW - maximum_age + self.cutoffs = [] + self.expired_cutoffs = [] + self.failed = [] + self.events = [] + + def list_active_runs(self, kind, *, active_since_before): + # As the real store, a run this fake failed is no longer listed. + failed = {failed_run for failed_run, _, _ in self.failed} + if active_since_before == self.maximum_age_cutoff: + self.expired_cutoffs.append(active_since_before) + self.events.append(f"list-expired:{kind.value}") + listed = self.expired + else: + self.cutoffs.append(active_since_before) + self.events.append(f"list:{kind.value}") + listed = self.runs + return tuple(run for run in listed if run.kind is kind and run not in failed) + + def fail_stuck_run(self, run, *, error_message, failed_at): + self.failed.append((run, error_message, failed_at)) + self.events.append(f"fail:{run.kind.value}") + return RunWrite(changed=True, status=RunStatus.FAILED) + + +def _run(call_id, kind=RunKind.SIMULATION_RUN, status=RunStatus.RUNNING): + return ActiveRun(kind=kind, run_id=uuid4(), status=status, modal_call_id=call_id) + + +def test_a_pending_run_with_no_call_is_failed_without_asking_modal() -> None: + # The simulation service stopped between inserting the run and starting + # its call, so nothing will run it. + never_started = _run(None, RunKind.REPORT_RUN, RunStatus.PENDING) + # A running run with no call recorded is left to the maximum age. + running_without_call = _run(None) + store = FakeSweepStore([never_started, running_without_call]) + asked = [] + + summary = sweep_stuck_runs(store, call_state=asked.append, now=NOW) + + assert asked == [] + assert [run for run, _, _ in store.failed] == [never_started] + (_, message, failed_at) = store.failed[0] + assert message.startswith("run_not_started: ") + assert failed_at == NOW + assert summary == SweepSummary( + checked=2, + failed=0, + running=0, + unknown=1, + left=0, + failed_for_age=0, + failed_without_call=1, + ) + + +def test_a_pending_run_with_a_call_is_judged_by_its_call() -> None: + # The simulation service recorded the call it started; Modal may still + # be queueing it, or it may have ended without claiming the run. + queued = _run("fc-queued", RunKind.REPORT_RUN, RunStatus.PENDING) + ended = _run("fc-ended", RunKind.REPORT_RUN, RunStatus.PENDING) + store = FakeSweepStore([queued, ended]) + states = {"fc-queued": CallState.RUNNING, "fc-ended": CallState.ENDED} + + summary = sweep_stuck_runs(store, call_state=states.__getitem__, now=NOW) + + assert [run for run, _, _ in store.failed] == [ended] + assert store.failed[0][1].startswith("modal_call_ended: ") + assert (summary.checked, summary.failed, summary.running) == (2, 1, 1) + + +def test_only_runs_whose_call_ended_are_failed() -> None: + ended = _run("fc-ended") + alive = _run("fc-alive") + unknown = _run("fc-unknown") + store = FakeSweepStore([ended, alive, unknown]) + states = { + "fc-ended": CallState.ENDED, + "fc-alive": CallState.RUNNING, + "fc-unknown": CallState.UNKNOWN, + } + + summary = sweep_stuck_runs(store, call_state=states.__getitem__, now=NOW) + + assert [run for run, _, _ in store.failed] == [ended] + (_, message, failed_at) = store.failed[0] + assert message.startswith("modal_call_ended: ") + assert failed_at == NOW + assert summary == SweepSummary( + checked=3, failed=1, running=1, unknown=1, left=0, failed_for_age=0 + ) + + +def test_the_sweep_lists_runs_older_than_its_minimum_age() -> None: + store = FakeSweepStore([]) + + sweep_stuck_runs( + store, + call_state=lambda _: CallState.ENDED, + now=NOW, + minimum_age=timedelta(minutes=15), + ) + + assert store.cutoffs == [NOW - timedelta(minutes=15)] * 2 + + +def test_report_runs_are_swept_before_simulation_runs_are_listed() -> None: + # A simulation run is listed only once no report run naming it is + # running, so failing a dead coordinator's report run first lets the same + # sweep reach its children. + report_run = _run("fc-coordinator", RunKind.REPORT_RUN) + simulation_run = _run("fc-worker") + store = FakeSweepStore([simulation_run, report_run]) + + sweep_stuck_runs(store, call_state=lambda _: CallState.ENDED, now=NOW) + + assert store.events == [ + "list-expired:report_run", + "list-expired:simulation_run", + "list:report_run", + "fail:report_run", + "list:simulation_run", + "fail:simulation_run", + ] + + +def test_a_call_running_several_runs_is_asked_once() -> None: + report_run = _run("fc-coordinator", RunKind.REPORT_RUN) + other = _run("fc-coordinator", RunKind.REPORT_RUN) + asked = [] + + def call_state(call_id): + asked.append(call_id) + return CallState.ENDED + + sweep_stuck_runs( + FakeSweepStore([report_run, other]), call_state=call_state, now=NOW + ) + + assert asked == ["fc-coordinator"] + + +def test_a_failed_status_lookup_never_fails_the_run() -> None: + run = _run("fc-error") + store = FakeSweepStore([run]) + + def call_state(_call_id): + msg = "Modal is unreachable" + raise ConnectionError(msg) + + summary = sweep_stuck_runs(store, call_state=call_state, now=NOW) + + assert store.failed == [] + assert summary == SweepSummary( + checked=1, failed=0, running=0, unknown=1, left=0, failed_for_age=0 + ) + + +def test_a_run_another_sweep_failed_first_is_not_counted_again() -> None: + run = _run("fc-ended") + + class AlreadyFailedStore(FakeSweepStore): + def fail_stuck_run(self, run, *, error_message, failed_at): + super().fail_stuck_run( + run, error_message=error_message, failed_at=failed_at + ) + return RunWrite(changed=False, status=RunStatus.FAILED) + + summary = sweep_stuck_runs( + AlreadyFailedStore([run]), call_state=lambda _: CallState.ENDED, now=NOW + ) + + assert summary.failed == 0 + # It finished, or moved to another call, after it was listed. + assert summary.left == 1 + + +def _node(call_id, status, children=()): + return SimpleNamespace( + function_call_id=call_id, status=status, children=list(children) + ) + + +def _fake_modal(statuses=None, error=None, *, graph=None): + """Return a modal stand-in whose call graph holds call fc-1's inputs. + + By default fc-1 is a top-level call; ``graph`` gives the whole graph a + spawned call is part of instead. + """ + + def from_id(call_id): + if error is not None: + raise error + assert call_id == "fc-1" + roots = ( + graph + if graph is not None + else [_node("fc-1", status) for status in (statuses or [])] + ) + return SimpleNamespace(get_call_graph=lambda: roots) + + return SimpleNamespace(FunctionCall=SimpleNamespace(from_id=from_id)) + + +@pytest.mark.parametrize( + ("statuses", "expected"), + [ + ([FakeInputStatus.PENDING], CallState.RUNNING), + ([FakeInputStatus.SUCCESS], CallState.ENDED), + ([FakeInputStatus.FAILURE], CallState.ENDED), + ([FakeInputStatus.TERMINATED], CallState.ENDED), + ([FakeInputStatus.TIMEOUT], CallState.ENDED), + ([FakeInputStatus.INIT_FAILURE], CallState.ENDED), + ([FakeInputStatus.FAILURE, FakeInputStatus.PENDING], CallState.RUNNING), + ([], CallState.UNKNOWN), + ], +) +def test_modal_call_state_reads_the_call_graph(statuses, expected) -> None: + assert modal_call_state("fc-1", modal_module=_fake_modal(statuses)) is expected + + +def test_modal_call_state_is_unknown_when_modal_cannot_answer() -> None: + state = modal_call_state( + "fc-1", modal_module=_fake_modal(error=ConnectionError("unreachable")) + ) + + assert state is CallState.UNKNOWN + + +@pytest.mark.parametrize( + ("coordinator_status", "worker_status", "expected"), + [ + # A coordinator that ran out of memory left its worker running. + (FakeInputStatus.FAILURE, FakeInputStatus.PENDING, CallState.RUNNING), + (FakeInputStatus.PENDING, FakeInputStatus.TERMINATED, CallState.ENDED), + ], +) +def test_modal_call_state_reads_the_recorded_calls_own_inputs( + coordinator_status, worker_status, expected +) -> None: + # A worker is spawned inside its coordinator's input, so Modal nests its + # inputs under the coordinator's in the call graph. + graph = [ + _node( + "fc-coordinator", + coordinator_status, + children=[ + _node("fc-1", worker_status, children=[_node("fc-segment", "x")]), + _node("fc-other-worker", FakeInputStatus.SUCCESS), + ], + ) + ] + + assert modal_call_state("fc-1", modal_module=_fake_modal(graph=graph)) is expected + + +def test_a_call_missing_from_its_call_graph_is_unknown() -> None: + graph = [_node("fc-coordinator", FakeInputStatus.SUCCESS)] + + assert ( + modal_call_state("fc-1", modal_module=_fake_modal(graph=graph)) + is CallState.UNKNOWN + ) + + +def test_a_run_past_the_maximum_age_is_failed_whatever_modal_reports() -> None: + # Modal reports a call it won't retry after an internal failure as + # pending, and no longer describes old calls; neither may keep a run + # running for good, and nor may a run with no recorded call. + reported_running = _run("fc-reported-running", RunKind.REPORT_RUN) + # A pending run counts from its creation; its call may still be queued. + still_queued = _run("fc-queued", RunKind.REPORT_RUN, RunStatus.PENDING) + without_call = ActiveRun( + kind=RunKind.REPORT_RUN, + run_id=uuid4(), + status=RunStatus.RUNNING, + modal_call_id=None, + ) + store = FakeSweepStore( + [reported_running], expired=[reported_running, still_queued, without_call] + ) + + summary = sweep_stuck_runs( + store, + call_state=lambda _: CallState.RUNNING, + now=NOW, + maximum_age=timedelta(hours=24), + ) + + assert [run for run, _, _ in store.failed] == [ + reported_running, + still_queued, + without_call, + ] + assert all( + message.startswith("run_exceeded_maximum_age: ") + for _, message, _ in store.failed + ) + assert summary.failed_for_age == 3 + # Failed for its age first, it is no longer listed for a lookup. + assert summary.running == 0 + + +def test_an_unrecognised_input_status_leaves_the_call_unknown() -> None: + unrecognised = SimpleNamespace(name="SOMETHING_NEW") + graph = [_node("fc-1", unrecognised)] + + assert ( + modal_call_state("fc-1", modal_module=_fake_modal(graph=graph)) + is CallState.UNKNOWN + ) + + +def test_the_age_limit_is_checked_first_with_its_own_cutoff() -> None: + # Before any Modal lookup, so lookups that stall can't keep it from + # running within the scheduled function's time limit. + store = FakeSweepStore([_run("fc-1", RunKind.REPORT_RUN)]) + asked = [] + + def call_state(call_id): + asked.append((call_id, list(store.events))) + return CallState.RUNNING + + sweep_stuck_runs( + store, + call_state=call_state, + now=NOW, + minimum_age=timedelta(minutes=10), + maximum_age=timedelta(hours=24), + ) + + assert store.expired_cutoffs == [NOW - timedelta(hours=24)] * 2 + assert store.cutoffs == [NOW - timedelta(minutes=10)] * 2 + # Both kinds' age checks come before any lookup. + assert asked == [ + ( + "fc-1", + [ + "list-expired:report_run", + "list-expired:simulation_run", + "list:report_run", + ], + ) + ] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py new file mode 100644 index 000000000..91a04fa96 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py @@ -0,0 +1,117 @@ +"""Declaration tests for the Modal app that runs the stuck-run sweep on a schedule.""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path + +import yaml +from fixtures.fake_modal import install_fake_modal +from policyengine_simulation_executor.stage12_runtime import ( + report_run_access as report_run_access_module, +) +from policyengine_simulation_executor.stage12_runtime import run_sweep + +REUSABLE_DEPLOY_WORKFLOW = ( + Path(__file__).resolve().parents[3] + / ".github" + / "workflows" + / "simulation-deploy.reusable.yml" +) + + +def _load(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.run_sweep_app", None) + return importlib.import_module("src.modal.run_sweep_app") + + +def test_the_executor_deployment_deploys_the_sweep_with_the_modal_tokens() -> None: + workflow = yaml.safe_load(REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8")) + steps = workflow["jobs"]["deploy_executor"]["steps"] + step_names = [step.get("name") for step in steps] + (sweep_step,) = [ + step for step in steps if "src/modal/run_sweep_app.py" in step.get("run", "") + ] + + assert sweep_step["run"] == ( + 'uv run --frozen modal deploy --env="${{ inputs.modal_environment }}" ' + "src/modal/run_sweep_app.py" + ) + assert ( + sweep_step["working-directory"] == "projects/policyengine-simulation-executor" + ) + # Only Modal's tokens: the database URL reaches the scheduled function from + # its Modal secret, which must exist before this step can deploy it. + assert sweep_step["env"] == { + "MODAL_TOKEN_ID": "${{ secrets.MODAL_TOKEN_ID }}", + "MODAL_TOKEN_SECRET": "${{ secrets.MODAL_TOKEN_SECRET }}", + } + assert step_names.index(sweep_step["name"]) > step_names.index( + "Deploy versioned Modal executor" + ) + + +def test_the_sweep_runs_every_ten_minutes_with_the_worker_database_secret( + monkeypatch, +) -> None: + module = _load(monkeypatch) + ((name, options),) = module.app.function_calls + + assert module.app.name == "policyengine-simulation-run-sweep" + assert name == "sweep_stuck_runs" + assert options["schedule"] == {"schedule": "*/10 * * * *"} + assert options["secrets"] == [ + { + "args": ("policyengine-modal-worker-database",), + "kwargs": {"required_keys": ["MODAL_WORKER_DATABASE_URL"]}, + } + ] + # One sweep at a time in this app; two at once would still be safe. + assert options["max_containers"] == 1 + assert options["retries"] == 0 + assert options["image"] is module.sweep_image + + +def test_the_sweep_image_installs_the_locked_group_and_the_runtime_sources( + monkeypatch, +) -> None: + module = _load(monkeypatch) + calls = {call[0]: call[1:] for call in module.sweep_image.calls} + + assert calls["uv_sync"][1]["frozen"] is True + assert calls["uv_sync"][1]["extra_options"] == "--only-group modal-simulation-image" + # The app module imports src.modal at module level, in the container too. + assert calls["add_local_python_source"][0] == ( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + ) + + +def test_the_scheduled_function_sweeps_with_the_report_run_store(monkeypatch) -> None: + module = _load(monkeypatch) + store = object() + swept = [] + monkeypatch.setattr(report_run_access_module, "report_run_store", lambda: store) + + def recording_sweep(received_store): + swept.append(received_store) + return run_sweep.SweepSummary( + checked=3, failed=1, running=1, unknown=1, left=0, failed_for_age=2 + ) + + monkeypatch.setattr(run_sweep, "sweep_stuck_runs", recording_sweep) + + assert module.sweep_stuck_runs() == { + "checked": 3, + "failed": 1, + "running": 1, + "unknown": 1, + "left": 0, + "failed_for_age": 2, + } + assert swept == [store]