From 25d25ece7c5a6b55e6ed41d0514c95e1cd1a2066 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:37:19 +0400 Subject: [PATCH 1/9] Let the report run store find and fail runs whose Modal call ended The store can now list the report runs and simulation runs that are running under a recorded Modal call and started before a given time, and mark one failed only while it is still running under the call it was listed with. A run that finished, or that a restarted coordinator moved to a new call since it was listed, is left as it is, so two sweeps at once fail each run once. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 4 + .../report_run_store.py | 78 +++++++ .../tests/test_report_run_sweep_store.py | 197 ++++++++++++++++++ 3 files changed, 279 insertions(+) create mode 100644 libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py 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..aec46b6c8 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -7,6 +7,8 @@ ) from policyengine_stage12_persistence.report_run_store import ( ReportRunStore, + RunKind, + RunningRun, RunStatus, RunWrite, ) @@ -26,8 +28,10 @@ "DistrictCodes", "ReportRunResult", "ReportRunStore", + "RunKind", "RunStatus", "RunWrite", + "RunningRun", "Stage12PersistenceStore", "Stage12RuntimeAccessError", "create_stage12_engine", 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 62e92d5c2..8329f6381 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 @@ -52,6 +52,25 @@ 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 RunningRun: + """A run recorded as running, and the Modal call recorded as running it.""" + + kind: RunKind + run_id: UUID + modal_call_id: str + + @dataclass(frozen=True) class RunWrite: """One conditional write's outcome: whether it changed the run, and its status.""" @@ -304,6 +323,65 @@ def fail_simulation_run( simulation_runs, simulation_run_id, statement, "fail_simulation_run" ) + def list_running_runs( + self, *, started_before: datetime, limit: int = 500 + ) -> tuple[RunningRun, ...]: + """Return up to ``limit`` runs of each kind running since before a time. + + Only runs with a recorded Modal call are returned, oldest first; a run + without a recorded start counts from its last update. + """ + running: list[RunningRun] = [] + with self._transaction("list_running_runs") as connection: + for kind, table in RUN_TABLES.items(): + started = func.coalesce(table.c.started_at, table.c.updated_at) + statement = ( + select(table.c.id, table.c.modal_call_id) + .where( + table.c.status == RunStatus.RUNNING.value, + table.c.modal_call_id.is_not(None), + started < started_before, + ) + .order_by(started, table.c.id) + .limit(limit) + ) + running.extend( + RunningRun( + kind=kind, run_id=row.id, modal_call_id=row.modal_call_id + ) + for row in connection.execute(statement) + ) + return tuple(running) + + def fail_run_of_ended_call( + self, + run: RunningRun, + *, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark ``run`` failed, only while it is running under ``run.modal_call_id``. + + A run that finished, or that moved to another call since it was + listed, is left as it is, so concurrent sweeps change it once. + """ + table = RUN_TABLES[run.kind] + statement = ( + update(table) + .where( + table.c.id == run.run_id, + table.c.status == RunStatus.RUNNING.value, + table.c.modal_call_id == run.modal_call_id, + ) + .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_run_of_ended_call") + def _write( self, table: Table, 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..c00929e87 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py @@ -0,0 +1,197 @@ +"""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 + +import pytest +from policyengine_stage12_persistence import ( + ReportRunStore, + RunKind, + RunningRun, + RunStatus, + RunWrite, +) +from report_run_database import NOW, SeededReportRun, worker_database_url + +CUTOFF = NOW + timedelta(minutes=30) +ENDED_MESSAGE = "modal_call_ended: the run's Modal call ended without recording it" + + +@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 ( + RunningRun( + kind=RunKind.REPORT_RUN, + run_id=run_ids.report_run_id, + modal_call_id="fc-coordinator", + ), + RunningRun( + kind=RunKind.SIMULATION_RUN, + run_id=run_ids.baseline_simulation_run_id, + modal_call_id="fc-baseline", + ), + ) + + +def _listed_for(store, seeded_run): + 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 store.list_running_runs(started_before=CUTOFF) + if run.run_id in ids + } + + +def test_the_sweep_lists_runs_running_under_a_call_since_before_the_cutoff( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run, baseline_run = _running(seeder, seeded_run) + # The reform run is running, but its call has no ID recorded yet. + _set_run( + seeder, + "simulation_runs", + seeded_run.run_ids.reform_simulation_run_id, + status="running", + started_at=NOW, + ) + + assert _listed_for(store, seeded_run) == {report_run, baseline_run} + assert _listed_for_cutoff(store, seeded_run, NOW) == set() + + +def _listed_for_cutoff(store, seeded_run, cutoff): + ids = { + seeded_run.run_ids.report_run_id, + seeded_run.run_ids.baseline_simulation_run_id, + } + return { + run + for run in store.list_running_runs(started_before=cutoff) + if run.run_id in ids + } + + +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_run_of_ended_call( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + again = store.fail_run_of_ended_call( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + simulation = store.fail_run_of_ended_call( + 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_run_of_ended_call( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + finished = store.fail_run_of_ended_call( + 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_two_sweeps_at_once_fail_each_run_exactly_once( + seeder, seeded_run: SeededReportRun +) -> None: + runs = _running(seeder, seeded_run) + 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_run_of_ended_call( + 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] + assert seeder.row("report_runs", seeded_run.run_ids.report_run_id)["status"] == ( + "failed" + ) From 9b61d552718ce647afe215e3b5f5769601f85131 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:39:02 +0400 Subject: [PATCH 2/9] Add a sweep that fails runs left running after their Modal call ended The sweep lists the report runs and simulation runs that have been running under a recorded Modal call for at least ten minutes, asks Modal for each call's call graph once, and marks a run failed only when every input of its call has finished: succeeded, failed, terminated, timed out or failed to start. A call with a pending input is running; Modal reports a pre-empted call pending until it runs again under the same call ID. A call Modal cannot describe, or any error asking, leaves the run as it is. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/run_sweep.py | 144 +++++++++++++++ .../tests/test_stage12_run_sweep.py | 168 ++++++++++++++++++ 2 files changed, 312 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/run_sweep.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py 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..c4b58d5b5 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/run_sweep.py @@ -0,0 +1,144 @@ +"""The sweep that fails report runs and simulation runs whose Modal call ended. + +A coordinator or worker records the Modal call running each run. If that call +ends without recording the run's outcome, for example because its container +ran out of memory, the run stays running. The sweep asks Modal for the status +of each recorded call and marks the run failed only once the call has ended. +Modal restarts a pre-empted call under the same call ID, and reports it +pending meanwhile, so a pre-empted run is never failed. +""" + +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 + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_stage12_persistence import RunningRun, RunWrite + +logger = logging.getLogger(__name__) + +# A run must have been running at least this long before the sweep checks +# its call, so it never races the dispatch that records the call. +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." +) +# Modal's input statuses (``modal.call_graph.InputStatus``) by name; only a +# pending input is still running or waiting to run again. +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_running_runs(self, *, started_before: datetime) -> tuple[RunningRun, ...]: + """Return the runs running under a recorded call since before a time.""" + ... + + def fail_run_of_ended_call( + self, run: RunningRun, *, error_message: str, failed_at: datetime + ) -> RunWrite: + """Mark ``run`` failed while it is still running under its listed call.""" + ... + + +@dataclass(frozen=True) +class SweepSummary: + """How many listed runs the sweep checked, failed, and left for each reason.""" + + checked: int + failed: int + running: int + unknown: int + + +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. + + A call whose inputs are all finished has ended; one with any pending input + is running. A call Modal cannot describe, 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: + inputs = modal.FunctionCall.from_id(modal_call_id).get_call_graph() + statuses = {input_info.status.name for input_info in inputs} + except Exception as error: # noqa: BLE001 # any failure to ask leaves the run as it is + logger.warning( + "Could not read the status of Modal call %s (%s)", + modal_call_id, + type(error).__name__, + ) + 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 sweep_stuck_runs( + store: StuckRunStore, + *, + call_state: Callable[[str], CallState] = modal_call_state, + now: datetime | None = None, + minimum_age: timedelta = DEFAULT_MINIMUM_AGE, +) -> SweepSummary: + """Fail every listed run whose recorded Modal call has ended. + + Each call is asked about once. A run whose call is running, or whose state + is unknown, is left as it is; so is one that finished or moved to another + call since it was listed, which makes two sweeps at once safe. + """ + now = now or datetime.now(UTC) + runs = store.list_running_runs(started_before=now - minimum_age) + states: dict[str, CallState] = {} + failed = running = unknown = 0 + for run in runs: + if run.modal_call_id not in states: + states[run.modal_call_id] = _safe_call_state(call_state, run.modal_call_id) + state = states[run.modal_call_id] + if state is CallState.RUNNING: + running += 1 + elif state is CallState.UNKNOWN: + unknown += 1 + elif store.fail_run_of_ended_call( + run, error_message=ENDED_CALL_MESSAGE, failed_at=now + ).changed: + failed += 1 + return SweepSummary( + checked=len(runs), failed=failed, running=running, unknown=unknown + ) + + +def _safe_call_state( + call_state: Callable[[str], CallState], modal_call_id: str +) -> CallState: + try: + return call_state(modal_call_id) + except Exception as error: # noqa: BLE001 # a failed lookup must never fail the run + logger.warning( + "Could not read the status of Modal call %s (%s)", + modal_call_id, + type(error).__name__, + ) + return CallState.UNKNOWN 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..7a00dbc96 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -0,0 +1,168 @@ +"""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 ( + CallState, + SweepSummary, + modal_call_state, + sweep_stuck_runs, +) +from policyengine_stage12_persistence import RunKind, RunningRun, 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): + self.runs = list(runs) + self.cutoffs = [] + self.failed = [] + + def list_running_runs(self, *, started_before): + self.cutoffs.append(started_before) + return tuple(self.runs) + + def fail_run_of_ended_call(self, run, *, error_message, failed_at): + self.failed.append((run, error_message, failed_at)) + return RunWrite(changed=True, status=RunStatus.FAILED) + + +def _run(call_id, kind=RunKind.SIMULATION_RUN): + return RunningRun(kind=kind, run_id=uuid4(), modal_call_id=call_id) + + +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) + + +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)] + + +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) + + +def test_a_run_another_sweep_failed_first_is_not_counted_again() -> None: + run = _run("fc-ended") + + class AlreadyFailedStore(FakeSweepStore): + def fail_run_of_ended_call(self, run, *, error_message, failed_at): + super().fail_run_of_ended_call( + 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 + + +def _fake_modal(statuses=None, error=None): + def from_id(call_id): + if error is not None: + raise error + assert call_id == "fc-1" + return SimpleNamespace( + get_call_graph=lambda: [ + SimpleNamespace(status=status) for status in (statuses or []) + ] + ) + + 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 From 95e8efbf514f219439ec801cbeb962212b511876 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:40:27 +0400 Subject: [PATCH 3/9] Run the stuck-run sweep every ten minutes as its own Modal app The new Modal application policyengine-simulation-run-sweep runs the sweep on a cron schedule every ten minutes, one container at a time, with the Modal worker role's database secret and the locked packages of the v2 images, without the bundle's data. It is not versioned with the .py bundle, so one copy runs per Modal environment; an operator deploys it. A sweep that overlaps another is still safe. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/run_sweep_app.py | 78 +++++++++++++++++++ .../tests/test_stage12_run_sweep_app.py | 76 ++++++++++++++++++ 2 files changed, 154 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/modal/run_sweep_app.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py 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..9f652e031 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -0,0 +1,78 @@ +"""A scheduled Modal function that fails runs left running after their call ended. + +Every ten minutes it lists the Public API's report runs and simulation runs +that have been running under a recorded Modal call for at least ten minutes, +asks Modal for each call's status, and marks a run failed once its call has +ended. 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. + +Deploy it, once per Modal environment, from this project's directory: + + uv run --frozen modal deploy --env=staging src/modal/run_sweep_app.py + uv run --frozen modal deploy --env=main src/modal/run_sweep_app.py +""" + +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( + "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 whose Modal call ended; return how many were checked 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/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..99ca0bc39 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py @@ -0,0 +1,76 @@ +"""Declaration tests for the Modal app that runs the stuck-run sweep on a schedule.""" + +from __future__ import annotations + +import importlib +import sys + +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 + + +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_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" + assert calls["add_local_python_source"][0] == ( + "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) + + monkeypatch.setattr(run_sweep, "sweep_stuck_runs", recording_sweep) + + assert module.sweep_stuck_runs() == { + "checked": 3, + "failed": 1, + "running": 1, + "unknown": 1, + } + assert swept == [store] From 42860c708162dea41b5bf4855e03a9a4022c0b6f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 00:40:27 +0400 Subject: [PATCH 4/9] Add a changelog fragment for the stuck-run sweep Co-Authored-By: Claude Opus 5.5 --- changelog.d/stuck-run-sweep.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/stuck-run-sweep.added.md diff --git a/changelog.d/stuck-run-sweep.added.md b/changelog.d/stuck-run-sweep.added.md new file mode 100644 index 000000000..286f915b7 --- /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 they are still running but the Modal call recorded for them has ended. It checks every ten minutes, never fails a run whose call is still running or whose state Modal cannot report, and is safe to run twice at once. An operator deploys it once per Modal environment. From 857d9f275543a55c8e64852480b65f5ff654b95b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:03:58 +0400 Subject: [PATCH 5/9] Mount src.modal in the stuck-run sweep's image The sweep's app module imports src.modal at module level, and Modal imports that module again in the container, so the image now adds the src.modal sources as the project's other Modal apps do. Without them every scheduled run failed with No module named 'src'. Co-Authored-By: Claude Opus 5.5 --- .../policyengine-simulation-executor/src/modal/run_sweep_app.py | 2 ++ .../tests/test_stage12_run_sweep_app.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py index 9f652e031..c32655659 100644 --- a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -48,6 +48,8 @@ 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", 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 index 99ca0bc39..908da9459 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py @@ -47,7 +47,9 @@ def test_the_sweep_image_installs_the_locked_group_and_the_runtime_sources( 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", From 13d32b0b59c9e4dada8a5799ff66d2ec154df9f1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:03:59 +0400 Subject: [PATCH 6/9] Judge a recorded call by its own inputs, and a child after its report Modal returns the whole call graph a call belongs to, and nests a call spawned inside another call's input under that input, so the sweep now reads only the inputs whose call ID is the recorded one, wherever they are, and treats a call missing from its graph as unknown. Before, a worker's simulation run was judged by its coordinator's status. A simulation run is now listed only while no report run that combines it is running: that report run's coordinator records its children's outcomes, and after a pre-emption it cancels them and starts them again under new calls, so a sweep in that gap could fail a run that would have finished. The sweep handles report runs first, so a dead coordinator's children are reached in the same sweep. Listing has no row limit, so runs that stay unknown can't hide newer ones. Co-Authored-By: Claude Opus 5.5 --- .../report_run_store.py | 57 ++++++------ .../tests/test_report_run_sweep_store.py | 31 ++++--- .../stage12_runtime/run_sweep.py | 85 ++++++++++++++---- .../tests/test_stage12_run_sweep.py | 86 +++++++++++++++++-- 4 files changed, 198 insertions(+), 61 deletions(-) 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 8329f6381..f2739b154 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 and_, func, insert, or_, select, update +from sqlalchemy import and_, exists, func, insert, or_, select, update from sqlalchemy.dialects.postgresql import insert as postgresql_insert from policyengine_stage12_persistence.report_run_tables import ( @@ -324,34 +324,41 @@ def fail_simulation_run( ) def list_running_runs( - self, *, started_before: datetime, limit: int = 500 + self, kind: RunKind, *, started_before: datetime ) -> tuple[RunningRun, ...]: - """Return up to ``limit`` runs of each kind running since before a time. + """Return the runs of ``kind`` running under a recorded call since a time. - Only runs with a recorded Modal call are returned, oldest first; a run - without a recorded start counts from its last update. + Runs are oldest first; a run without a recorded start counts from its + last update. A simulation run is listed only while no report run that + combines it is running: that report run's coordinator records its + children's outcomes, and stops and starts them again when Modal + restarts it, so their calls are its to judge until it has finished. """ - running: list[RunningRun] = [] - with self._transaction("list_running_runs") as connection: - for kind, table in RUN_TABLES.items(): - started = func.coalesce(table.c.started_at, table.c.updated_at) - statement = ( - select(table.c.id, table.c.modal_call_id) - .where( - table.c.status == RunStatus.RUNNING.value, - table.c.modal_call_id.is_not(None), - started < started_before, - ) - .order_by(started, table.c.id) - .limit(limit) - ) - running.extend( - RunningRun( - kind=kind, run_id=row.id, modal_call_id=row.modal_call_id - ) - for row in connection.execute(statement) + table = RUN_TABLES[kind] + started = func.coalesce(table.c.started_at, table.c.updated_at) + conditions = [ + table.c.status == RunStatus.RUNNING.value, + table.c.modal_call_id.is_not(None), + started < started_before, + ] + if kind is RunKind.SIMULATION_RUN: + conditions.append( + ~exists().where( + report_run_simulation_runs.c.simulation_run_id == table.c.id, + report_run_simulation_runs.c.report_run_id == report_runs.c.id, + report_runs.c.status == RunStatus.RUNNING.value, ) - return tuple(running) + ) + statement = ( + select(table.c.id, table.c.modal_call_id) + .where(*conditions) + .order_by(started, table.c.id) + ) + with self._transaction("list_running_runs") as connection: + return tuple( + RunningRun(kind=kind, run_id=row.id, modal_call_id=row.modal_call_id) + for row in connection.execute(statement) + ) def fail_run_of_ended_call( self, 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 index c00929e87..2b7dcfbf4 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py @@ -72,23 +72,36 @@ def _running(seeder, seeded_run: SeededReportRun): ) +def _listed(store, cutoff): + return ( + *store.list_running_runs(RunKind.REPORT_RUN, started_before=cutoff), + *store.list_running_runs(RunKind.SIMULATION_RUN, started_before=cutoff), + ) + + def _listed_for(store, seeded_run): 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 store.list_running_runs(started_before=CUTOFF) - if run.run_id in ids - } + return {run for run in _listed(store, CUTOFF) if run.run_id in ids} def test_the_sweep_lists_runs_running_under_a_call_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. _set_run( seeder, @@ -98,7 +111,7 @@ def test_the_sweep_lists_runs_running_under_a_call_since_before_the_cutoff( started_at=NOW, ) - assert _listed_for(store, seeded_run) == {report_run, baseline_run} + assert _listed_for(store, seeded_run) == {baseline_run} assert _listed_for_cutoff(store, seeded_run, NOW) == set() @@ -107,11 +120,7 @@ def _listed_for_cutoff(store, seeded_run, cutoff): seeded_run.run_ids.report_run_id, seeded_run.run_ids.baseline_simulation_run_id, } - return { - run - for run in store.list_running_runs(started_before=cutoff) - if run.run_id in ids - } + return {run for run in _listed(store, cutoff) if run.run_id in ids} def test_the_sweep_fails_a_run_once_while_its_ended_call_still_holds_it( 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 index c4b58d5b5..cf1c0f951 100644 --- 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 @@ -4,8 +4,11 @@ ends without recording the run's outcome, for example because its container ran out of memory, the run stays running. The sweep asks Modal for the status of each recorded call and marks the run failed only once the call has ended. -Modal restarts a pre-empted call under the same call ID, and reports it -pending meanwhile, so a pre-empted run is never failed. + +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, so the sweep +judges a simulation run only once no report run combining it is running. """ from __future__ import annotations @@ -17,6 +20,8 @@ from importlib import import_module from typing import TYPE_CHECKING, Any, Protocol +from policyengine_stage12_persistence import RunKind + if TYPE_CHECKING: from collections.abc import Callable @@ -24,15 +29,17 @@ logger = logging.getLogger(__name__) -# A run must have been running at least this long before the sweep checks -# its call, so it never races the dispatch that records the call. +# A run must have been running at least this long, counted from its first +# claim, before the sweep checks its call. 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." ) -# Modal's input statuses (``modal.call_graph.InputStatus``) by name; only a -# pending input is still running or waiting to run again. +# 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"} @@ -50,8 +57,10 @@ class CallState(StrEnum): class StuckRunStore(Protocol): """The sweep's reads and writes; ``ReportRunStore`` implements them.""" - def list_running_runs(self, *, started_before: datetime) -> tuple[RunningRun, ...]: - """Return the runs running under a recorded call since before a time.""" + def list_running_runs( + self, kind: RunKind, *, started_before: datetime + ) -> tuple[RunningRun, ...]: + """Return the runs of ``kind`` running under a recorded call since a time.""" ... def fail_run_of_ended_call( @@ -74,14 +83,19 @@ class SweepSummary: 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. - A call whose inputs are all finished has ended; one with any pending input - is running. A call Modal cannot describe, or any error asking, is unknown, - so the sweep leaves its run alone. + 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, 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: - inputs = modal.FunctionCall.from_id(modal_call_id).get_call_graph() - statuses = {input_info.status.name for input_info in inputs} + 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 logger.warning( "Could not read the status of Modal call %s (%s)", @@ -96,6 +110,18 @@ def modal_call_state(modal_call_id: str, *, modal_module: Any = None) -> CallSta 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 + + def sweep_stuck_runs( store: StuckRunStore, *, @@ -105,13 +131,38 @@ def sweep_stuck_runs( ) -> SweepSummary: """Fail every listed run whose recorded Modal call has ended. - Each call is asked about once. A run whose call is running, or whose state - is unknown, is left as it is; so is one that finished or moved to another - call since it was listed, which makes two sweeps at once safe. + Report runs are swept first, so a dead coordinator's simulation runs are + listed in the same sweep. Each call is asked about once. A run whose call + is running, or whose state is unknown, is left as it is; so is one that + finished or moved to another call since it was listed, which makes two + sweeps at once safe. """ now = now or datetime.now(UTC) - runs = store.list_running_runs(started_before=now - minimum_age) states: dict[str, CallState] = {} + checked = failed = running = unknown = 0 + for kind in (RunKind.REPORT_RUN, RunKind.SIMULATION_RUN): + runs = store.list_running_runs(kind, started_before=now - minimum_age) + checked += len(runs) + outcome = _sweep_runs( + store, runs, states=states, call_state=call_state, now=now + ) + failed += outcome.failed + running += outcome.running + unknown += outcome.unknown + return SweepSummary( + checked=checked, failed=failed, running=running, unknown=unknown + ) + + +def _sweep_runs( + store: StuckRunStore, + runs: tuple[RunningRun, ...], + *, + states: dict[str, CallState], + call_state: Callable[[str], CallState], + now: datetime, +) -> SweepSummary: + """Fail each of ``runs`` whose call has ended, asking about each call once.""" failed = running = unknown = 0 for run in runs: if run.modal_call_id not in states: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py index 7a00dbc96..8eae345c5 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -35,13 +35,16 @@ def __init__(self, runs): self.runs = list(runs) self.cutoffs = [] self.failed = [] + self.events = [] - def list_running_runs(self, *, started_before): + def list_running_runs(self, kind, *, started_before): self.cutoffs.append(started_before) - return tuple(self.runs) + self.events.append(f"list:{kind.value}") + return tuple(run for run in self.runs if run.kind is kind) def fail_run_of_ended_call(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) @@ -79,7 +82,25 @@ def test_the_sweep_lists_runs_older_than_its_minimum_age() -> None: minimum_age=timedelta(minutes=15), ) - assert store.cutoffs == [NOW - 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 combining 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:report_run", + "fail:report_run", + "list:simulation_run", + "fail:simulation_run", + ] def test_a_call_running_several_runs_is_asked_once() -> None: @@ -129,16 +150,29 @@ def fail_run_of_ended_call(self, run, *, error_message, failed_at): assert summary.failed == 0 -def _fake_modal(statuses=None, error=None): +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" - return SimpleNamespace( - get_call_graph=lambda: [ - SimpleNamespace(status=status) for status in (statuses or []) - ] + 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)) @@ -166,3 +200,39 @@ def test_modal_call_state_is_unknown_when_modal_cannot_answer() -> None: ) 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 + ) From c6438d13e08409947c4a23a5963deaf1846aa42a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:28:19 +0400 Subject: [PATCH 7/9] Fail runs still running a day after their first claim Modal reports a call it won't retry after an internal failure as pending, stops describing old calls, and a run can be marked running with no call recorded; none of these could be decided, so such a run, and the simulation runs it combines, stayed running for good. The sweep now also fails every run still running 24 hours after its first claim, whatever Modal reports and with or without a recorded call, a time no attempt of the v2 functions can approach. A run listed without a call is failed only while it still has none. The summary now also counts runs left because they finished or moved to another call, and runs failed for their age, and the listing of simulation runs is tested with pending, several and no report runs combining them. Co-Authored-By: Claude Opus 5.5 --- .../report_run_store.py | 42 +++-- .../tests/test_report_run_sweep_store.py | 92 ++++++++++ .../stage12_runtime/run_sweep.py | 157 +++++++++++------- .../tests/test_stage12_run_sweep.py | 59 ++++++- .../tests/test_stage12_run_sweep_app.py | 6 +- 5 files changed, 277 insertions(+), 79 deletions(-) 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 f2739b154..4e0f81aa3 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 @@ -64,11 +64,14 @@ class RunKind(StrEnum): @dataclass(frozen=True) class RunningRun: - """A run recorded as running, and the Modal call recorded as running it.""" + """A run recorded as running, and the Modal call recorded as running it. + + ``modal_call_id`` is None for a run marked running with no call recorded. + """ kind: RunKind run_id: UUID - modal_call_id: str + modal_call_id: str | None @dataclass(frozen=True) @@ -324,23 +327,30 @@ def fail_simulation_run( ) def list_running_runs( - self, kind: RunKind, *, started_before: datetime + self, + kind: RunKind, + *, + started_before: datetime, + include_runs_without_call: bool = False, ) -> tuple[RunningRun, ...]: - """Return the runs of ``kind`` running under a recorded call since a time. - - Runs are oldest first; a run without a recorded start counts from its - last update. A simulation run is listed only while no report run that - combines it is running: that report run's coordinator records its - children's outcomes, and stops and starts them again when Modal - restarts it, so their calls are its to judge until it has finished. + """Return the runs of ``kind`` running since before a time. + + Only runs with a recorded Modal call are returned, unless + ``include_runs_without_call``. Runs are oldest first; a run without a + recorded start counts from its last update. A simulation run is + listed only while no report run that combines it is running: that + report run's coordinator records its children's outcomes, and stops + and starts them again when Modal restarts it, so their calls are its + to judge until it has finished. """ table = RUN_TABLES[kind] started = func.coalesce(table.c.started_at, table.c.updated_at) conditions = [ table.c.status == RunStatus.RUNNING.value, - table.c.modal_call_id.is_not(None), started < started_before, ] + if not include_runs_without_call: + conditions.append(table.c.modal_call_id.is_not(None)) if kind is RunKind.SIMULATION_RUN: conditions.append( ~exists().where( @@ -370,15 +380,21 @@ def fail_run_of_ended_call( """Mark ``run`` failed, only while it is running under ``run.modal_call_id``. A run that finished, or that moved to another call since it was - listed, is left as it is, so concurrent sweeps change it once. + listed, is left as it is, so concurrent sweeps change it once. A run + listed with no call is failed only while it still has none. """ 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 == RunStatus.RUNNING.value, - table.c.modal_call_id == run.modal_call_id, + same_call, ) .values( status=RunStatus.FAILED.value, 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 index 2b7dcfbf4..ebb2ecdd7 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py @@ -9,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from datetime import timedelta from threading import Barrier +from uuid import uuid4 import pytest from policyengine_stage12_persistence import ( @@ -19,6 +20,7 @@ RunWrite, ) from report_run_database import NOW, SeededReportRun, worker_database_url +from sqlalchemy import delete, insert CUTOFF = NOW + timedelta(minutes=30) ENDED_MESSAGE = "modal_call_ended: the run's Modal call ended without recording it" @@ -204,3 +206,93 @@ def sweep(run_store): assert seeder.row("report_runs", seeded_run.run_ids.report_run_id)["status"] == ( "failed" ) + + +def _link(seeder, report_run_id, simulation_run_id, role="baseline"): + with seeder.engine.begin() as connection: + connection.execute( + insert(seeder.table("report_run_simulation_runs")).values( + id=uuid4(), + report_run_id=report_run_id, + simulation_run_id=simulation_run_id, + role=role, + ) + ) + + +def _second_report_run(seeder, seeded_run: SeededReportRun, status): + """Insert another report run of the same report, sharing the baseline run.""" + 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"], + status=status, + trigger="manual", + country_package_version=first["country_package_version"], + policyengine_version=first["policyengine_version"], + **({"completed_at": CUTOFF} if status == "succeeded" else {}), + ) + ) + _link(seeder, second_id, seeded_run.run_ids.baseline_simulation_run_id) + return second_id + + +def _baseline_listed(store, seeded_run): + baseline_id = seeded_run.run_ids.baseline_simulation_run_id + return any( + run.run_id == baseline_id + for run in store.list_running_runs( + RunKind.SIMULATION_RUN, started_before=CUTOFF + ) + ) + + +def test_a_simulation_run_is_listed_unless_a_report_run_combining_it_is_running( + 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) + + _set_run(seeder, "report_runs", seeded_run.run_ids.report_run_id, status="pending") + 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) + + links = seeder.table("report_run_simulation_runs") + with seeder.engine.begin() as connection: + connection.execute( + delete(links).where( + links.c.simulation_run_id + == seeded_run.run_ids.baseline_simulation_run_id + ) + ) + # A simulation run no report run combines is the sweep's to judge. + assert _baseline_listed(store, seeded_run) + + +def test_runs_past_the_maximum_age_are_listed_and_failed_without_a_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, status="running", started_at=NOW) + + with_calls = store.list_running_runs(RunKind.REPORT_RUN, started_before=CUTOFF) + every_run = store.list_running_runs( + RunKind.REPORT_RUN, started_before=CUTOFF, include_runs_without_call=True + ) + + assert report_run_id not in {run.run_id for run in with_calls} + (unrecorded,) = [run for run in every_run if run.run_id == report_run_id] + assert unrecorded.modal_call_id is None + failed = store.fail_run_of_ended_call( + unrecorded, error_message="run_exceeded_maximum_age: too old", failed_at=CUTOFF + ) + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) 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 index cf1c0f951..7bd96be54 100644 --- 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 @@ -36,6 +36,17 @@ "modal_call_ended: The Modal call recorded for this run ended without " "recording the run's outcome." ) +# A run still running this long after its first claim 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 running 24 hours after it " + "was first claimed, 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 @@ -58,9 +69,13 @@ class StuckRunStore(Protocol): """The sweep's reads and writes; ``ReportRunStore`` implements them.""" def list_running_runs( - self, kind: RunKind, *, started_before: datetime + self, + kind: RunKind, + *, + started_before: datetime, + include_runs_without_call: bool = False, ) -> tuple[RunningRun, ...]: - """Return the runs of ``kind`` running under a recorded call since a time.""" + """Return the runs of ``kind`` running since before a time.""" ... def fail_run_of_ended_call( @@ -72,12 +87,20 @@ def fail_run_of_ended_call( @dataclass(frozen=True) class SweepSummary: - """How many listed runs the sweep checked, failed, and left for each reason.""" + """What the sweep did with the runs it listed. + + ``checked`` runs had their call asked about: ``failed`` because it ended, + left because it is ``running`` or its state is ``unknown``, or ``left`` + because the run finished or moved to another call after it was listed. + ``failed_for_age`` runs were failed for exceeding the maximum age. + """ checked: int failed: int running: int unknown: int + left: int = 0 + failed_for_age: 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 @@ -87,8 +110,8 @@ def modal_call_state(modal_call_id: str, *, modal_module: Any = None) -> CallSta 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, or any error asking, is unknown, so the - sweep leaves its run alone. + 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: @@ -97,11 +120,7 @@ def modal_call_state(modal_call_id: str, *, modal_module: Any = None) -> CallSta 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 - logger.warning( - "Could not read the status of Modal call %s (%s)", - modal_call_id, - type(error).__name__, - ) + _log_unreadable_call(modal_call_id, error) return CallState.UNKNOWN if statuses & RUNNING_INPUT_STATUSES: return CallState.RUNNING @@ -122,74 +141,90 @@ def _inputs_of(graph: list[Any], modal_call_id: str) -> list[Any]: return found -def sweep_stuck_runs( +@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 + + +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. Report runs are swept first, so a dead coordinator's simulation runs are listed in the same sweep. Each call is asked about once. A run whose call - is running, or whose state is unknown, is left as it is; so is one that - finished or moved to another call since it was listed, which makes two - sweeps at once safe. + is running, or whose state is unknown, is left as it is, unless it has + been running longer than ``maximum_age``; so is one that finished or + moved to another call since it was listed, which makes two sweeps at + once safe. """ now = now or datetime.now(UTC) states: dict[str, CallState] = {} - checked = failed = running = unknown = 0 + tally = _Tally() for kind in (RunKind.REPORT_RUN, RunKind.SIMULATION_RUN): - runs = store.list_running_runs(kind, started_before=now - minimum_age) - checked += len(runs) - outcome = _sweep_runs( - store, runs, states=states, call_state=call_state, now=now - ) - failed += outcome.failed - running += outcome.running - unknown += outcome.unknown - return SweepSummary( - checked=checked, failed=failed, running=running, unknown=unknown - ) - - -def _sweep_runs( + for run in store.list_running_runs(kind, started_before=now - minimum_age): + _sweep_run( + store, run, states=states, call_state=call_state, now=now, tally=tally + ) + for run in store.list_running_runs( + kind, started_before=now - maximum_age, include_runs_without_call=True + ): + if store.fail_run_of_ended_call( + run, error_message=EXCEEDED_AGE_MESSAGE, failed_at=now + ).changed: + tally.failed_for_age += 1 + return SweepSummary(**vars(tally)) + + +def _sweep_run( # noqa: PLR0913 # one run, the store, and the sweep's shared state store: StuckRunStore, - runs: tuple[RunningRun, ...], + run: RunningRun, *, states: dict[str, CallState], call_state: Callable[[str], CallState], now: datetime, -) -> SweepSummary: - """Fail each of ``runs`` whose call has ended, asking about each call once.""" - failed = running = unknown = 0 - for run in runs: - if run.modal_call_id not in states: - states[run.modal_call_id] = _safe_call_state(call_state, run.modal_call_id) - state = states[run.modal_call_id] - if state is CallState.RUNNING: - running += 1 - elif state is CallState.UNKNOWN: - unknown += 1 - elif store.fail_run_of_ended_call( - run, error_message=ENDED_CALL_MESSAGE, failed_at=now - ).changed: - failed += 1 - return SweepSummary( - checked=len(runs), failed=failed, running=running, unknown=unknown + tally: _Tally, +) -> None: + """Fail ``run`` if its call has ended, asking about each call once.""" + tally.checked += 1 + modal_call_id = run.modal_call_id + if modal_call_id is None: + tally.unknown += 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_run_of_ended_call( + 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__, ) - - -def _safe_call_state( - call_state: Callable[[str], CallState], modal_call_id: str -) -> CallState: - try: - return call_state(modal_call_id) - except Exception as error: # noqa: BLE001 # a failed lookup must never fail the run - logger.warning( - "Could not read the status of Modal call %s (%s)", - modal_call_id, - type(error).__name__, - ) - return CallState.UNKNOWN diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py index 8eae345c5..2976cf815 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -31,13 +31,20 @@ class FakeInputStatus(Enum): class FakeSweepStore: - def __init__(self, runs): + def __init__(self, runs, *, expired=()): self.runs = list(runs) + # Runs older than the maximum age, listed with or without a call. + self.expired = list(expired) self.cutoffs = [] self.failed = [] self.events = [] - def list_running_runs(self, kind, *, started_before): + def list_running_runs( + self, kind, *, started_before, include_runs_without_call=False + ): + if include_runs_without_call: + self.events.append(f"list-expired:{kind.value}") + return tuple(run for run in self.expired if run.kind is kind) self.cutoffs.append(started_before) self.events.append(f"list:{kind.value}") return tuple(run for run in self.runs if run.kind is kind) @@ -69,7 +76,9 @@ def test_only_runs_whose_call_ended_are_failed() -> None: (_, 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) + 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: @@ -98,8 +107,10 @@ def test_report_runs_are_swept_before_simulation_runs_are_listed() -> None: assert store.events == [ "list:report_run", "fail:report_run", + "list-expired:report_run", "list:simulation_run", "fail:simulation_run", + "list-expired:simulation_run", ] @@ -130,7 +141,9 @@ def call_state(_call_id): 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) + 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: @@ -148,6 +161,8 @@ def fail_run_of_ended_call(self, run, *, error_message, failed_at): ) 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=()): @@ -236,3 +251,39 @@ def test_a_call_missing_from_its_call_graph_is_unknown() -> None: 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) + without_call = RunningRun( + kind=RunKind.REPORT_RUN, run_id=uuid4(), modal_call_id=None + ) + store = FakeSweepStore([reported_running], expired=[reported_running, 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, without_call] + assert all( + message.startswith("run_exceeded_maximum_age: ") + for _, message, _ in store.failed + ) + assert summary.failed_for_age == 2 + assert summary.running == 1 + + +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 + ) 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 index 908da9459..f3d47ef13 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py @@ -65,7 +65,9 @@ def test_the_scheduled_function_sweeps_with_the_report_run_store(monkeypatch) -> def recording_sweep(received_store): swept.append(received_store) - return run_sweep.SweepSummary(checked=3, failed=1, running=1, unknown=1) + 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) @@ -74,5 +76,7 @@ def recording_sweep(received_store): "failed": 1, "running": 1, "unknown": 1, + "left": 0, + "failed_for_age": 2, } assert swept == [store] From 10a3d3215b0eb43eb95bf7620db43be5a28e4f95 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:43:01 +0400 Subject: [PATCH 8/9] Check the age limit first and keep a combined child out of the sweep The sweep now fails runs past the maximum age before it asks Modal about any call, so lookups that stall can't use up the scheduled function's time before the check that exists for undecidable runs. The update that fails a simulation run repeats the condition that no running report run combines it, so a report run that starts combining it after the listing keeps it. The failing write is renamed fail_stuck_run, the age message no longer names a fixed age, and the changelog fragment and docstrings say that a run past the maximum age is failed whatever Modal reports. Co-Authored-By: Claude Opus 5.5 --- changelog.d/stuck-run-sweep.added.md | 2 +- .../report_run_store.py | 38 ++++++++++++------- .../tests/test_report_run_sweep_store.py | 37 ++++++++++++++---- .../src/modal/run_sweep_app.py | 6 ++- .../stage12_runtime/run_sweep.py | 34 +++++++++-------- .../tests/test_stage12_run_sweep.py | 35 ++++++++++++++--- 6 files changed, 108 insertions(+), 44 deletions(-) diff --git a/changelog.d/stuck-run-sweep.added.md b/changelog.d/stuck-run-sweep.added.md index 286f915b7..6f5418b29 100644 --- a/changelog.d/stuck-run-sweep.added.md +++ b/changelog.d/stuck-run-sweep.added.md @@ -1 +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 they are still running but the Modal call recorded for them has ended. It checks every ten minutes, never fails a run whose call is still running or whose state Modal cannot report, and is safe to run twice at once. An operator deploys it once per Modal environment. +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 they are still running but the Modal call recorded for them has ended. 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 has been running for 24 hours, longer than any of its calls can run; then it fails the run whatever Modal reports. An operator deploys it once per Modal environment. 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 4e0f81aa3..0bea98294 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 @@ -34,7 +34,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 @@ -352,13 +352,7 @@ def list_running_runs( if not include_runs_without_call: conditions.append(table.c.modal_call_id.is_not(None)) if kind is RunKind.SIMULATION_RUN: - conditions.append( - ~exists().where( - report_run_simulation_runs.c.simulation_run_id == table.c.id, - report_run_simulation_runs.c.report_run_id == report_runs.c.id, - report_runs.c.status == RunStatus.RUNNING.value, - ) - ) + conditions.append(_no_running_report_run_combines(table)) statement = ( select(table.c.id, table.c.modal_call_id) .where(*conditions) @@ -370,18 +364,20 @@ def list_running_runs( for row in connection.execute(statement) ) - def fail_run_of_ended_call( + def fail_stuck_run( self, run: RunningRun, *, error_message: str, failed_at: datetime, ) -> RunWrite: - """Mark ``run`` failed, only while it is running under ``run.modal_call_id``. + """Mark a stuck ``run`` failed, while it is as it was listed. - A run that finished, or that moved to another call since it was - listed, is left as it is, so concurrent sweeps change it once. A run - listed with no call is failed only while it still has none. + It must still be running under ``run.modal_call_id``, or with no call + if it was listed with none, and a simulation run must still have no + running report run combining it. A run that finished, moved to another + call, or that a report run started combining since it was listed, is + left as it is, so concurrent sweeps change it once. """ table = RUN_TABLES[run.kind] same_call = ( @@ -395,6 +391,11 @@ def fail_run_of_ended_call( table.c.id == run.run_id, table.c.status == RunStatus.RUNNING.value, same_call, + *( + (_no_running_report_run_combines(table),) + if run.kind is RunKind.SIMULATION_RUN + else () + ), ) .values( status=RunStatus.FAILED.value, @@ -403,7 +404,7 @@ def fail_run_of_ended_call( updated_at=failed_at, ) ) - return self._write(table, run.run_id, statement, "fail_run_of_ended_call") + return self._write(table, run.run_id, statement, "fail_stuck_run") def _write( self, @@ -458,3 +459,12 @@ 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_running_report_run_combines(table: Table) -> ColumnElement[bool]: + """Return the condition that no running report run combines the run.""" + return ~exists().where( + report_run_simulation_runs.c.simulation_run_id == table.c.id, + report_run_simulation_runs.c.report_run_id == report_runs.c.id, + report_runs.c.status == RunStatus.RUNNING.value, + ) 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 index ebb2ecdd7..3edb21886 100644 --- a/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py @@ -130,13 +130,13 @@ def test_the_sweep_fails_a_run_once_while_its_ended_call_still_holds_it( ) -> None: report_run, baseline_run = _running(seeder, seeded_run) - first = store.fail_run_of_ended_call( + first = store.fail_stuck_run( report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF ) - again = store.fail_run_of_ended_call( + again = store.fail_stuck_run( report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF ) - simulation = store.fail_run_of_ended_call( + simulation = store.fail_stuck_run( baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF ) @@ -168,10 +168,10 @@ def test_the_sweep_leaves_a_run_that_moved_to_another_call_or_finished( completed_at=CUTOFF, ) - moved = store.fail_run_of_ended_call( + moved = store.fail_stuck_run( baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF ) - finished = store.fail_run_of_ended_call( + finished = store.fail_stuck_run( report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF ) @@ -189,7 +189,7 @@ def test_two_sweeps_at_once_fail_each_run_exactly_once( def sweep(run_store): start_together.wait() return [ - run_store.fail_run_of_ended_call( + run_store.fail_stuck_run( run, error_message=ENDED_MESSAGE, failed_at=CUTOFF ).changed for run in runs @@ -292,7 +292,30 @@ def test_runs_past_the_maximum_age_are_listed_and_failed_without_a_call( assert report_run_id not in {run.run_id for run in with_calls} (unrecorded,) = [run for run in every_run if run.run_id == report_run_id] assert unrecorded.modal_call_id is None - failed = store.fail_run_of_ended_call( + failed = store.fail_stuck_run( unrecorded, error_message="run_exceeded_maximum_age: too old", failed_at=CUTOFF ) assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + + +def test_a_simulation_run_a_report_run_is_now_running_is_not_failed( + store, seeder, seeded_run: SeededReportRun +) -> 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 combining it starts running; its + # coordinator now judges the simulation run. + _second_report_run(seeder, seeded_run, "running") + + write = store.fail_stuck_run( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert write == RunWrite(changed=False, status=RunStatus.RUNNING) diff --git a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py index c32655659..793714002 100644 --- a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -3,8 +3,10 @@ Every ten minutes it lists the Public API's report runs and simulation runs that have been running under a recorded Modal call for at least ten minutes, asks Modal for each call's status, and marks a run failed once its call has -ended. 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 +ended. It also fails any run still running 24 hours after it 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. Deploy it, once per Modal environment, from this project's directory: 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 index 7bd96be54..fc6351195 100644 --- 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 @@ -3,7 +3,10 @@ A coordinator or worker records the Modal call running each run. If that call ends without recording the run's outcome, for example because its container ran out of memory, the run stays running. The sweep asks Modal for the status -of each recorded call and marks the run failed only once the call has ended. +of each recorded call and marks the run failed once the call has ended. It +also fails a run still running long after it 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 running 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. @@ -44,8 +47,8 @@ # 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 running 24 hours after it " - "was first claimed, longer than any Modal call of it can run." + "run_exceeded_maximum_age: The run was still running long after it " + "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 @@ -78,7 +81,7 @@ def list_running_runs( """Return the runs of ``kind`` running since before a time.""" ... - def fail_run_of_ended_call( + def fail_stuck_run( self, run: RunningRun, *, error_message: str, failed_at: datetime ) -> RunWrite: """Mark ``run`` failed while it is still running under its listed call.""" @@ -164,27 +167,28 @@ def sweep_stuck_runs( # the store, then keyword-only seams and limits """Fail every listed run whose recorded Modal call has ended. Report runs are swept first, so a dead coordinator's simulation runs are - listed in the same sweep. Each call is asked about once. A run whose call - is running, or whose state is unknown, is left as it is, unless it has - been running longer than ``maximum_age``; so is one that finished or - moved to another call since it was listed, which makes two sweeps at - once safe. + listed in the same sweep. For each kind, runs older than ``maximum_age`` + are failed first, before any Modal lookup can use up the scheduled + function's time. Each call is then asked about once. A run whose call is + running, or whose state is unknown, is left as it is; so is one that + finished or moved to another call since it was listed, 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_running_runs(kind, started_before=now - minimum_age): - _sweep_run( - store, run, states=states, call_state=call_state, now=now, tally=tally - ) for run in store.list_running_runs( kind, started_before=now - maximum_age, include_runs_without_call=True ): - if store.fail_run_of_ended_call( + if store.fail_stuck_run( run, error_message=EXCEEDED_AGE_MESSAGE, failed_at=now ).changed: tally.failed_for_age += 1 + for run in store.list_running_runs(kind, started_before=now - minimum_age): + _sweep_run( + store, run, states=states, call_state=call_state, now=now, tally=tally + ) return SweepSummary(**vars(tally)) @@ -214,7 +218,7 @@ def _sweep_run( # noqa: PLR0913 # one run, the store, and the sweep's shared s tally.running += 1 elif state is CallState.UNKNOWN: tally.unknown += 1 - elif store.fail_run_of_ended_call( + elif store.fail_stuck_run( run, error_message=ENDED_CALL_MESSAGE, failed_at=now ).changed: tally.failed += 1 diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py index 2976cf815..1b4ddcc59 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -36,6 +36,7 @@ def __init__(self, runs, *, expired=()): # Runs older than the maximum age, listed with or without a call. self.expired = list(expired) self.cutoffs = [] + self.expired_cutoffs = [] self.failed = [] self.events = [] @@ -43,13 +44,14 @@ def list_running_runs( self, kind, *, started_before, include_runs_without_call=False ): if include_runs_without_call: + self.expired_cutoffs.append(started_before) self.events.append(f"list-expired:{kind.value}") return tuple(run for run in self.expired if run.kind is kind) self.cutoffs.append(started_before) self.events.append(f"list:{kind.value}") return tuple(run for run in self.runs if run.kind is kind) - def fail_run_of_ended_call(self, run, *, error_message, failed_at): + 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) @@ -105,12 +107,12 @@ def test_report_runs_are_swept_before_simulation_runs_are_listed() -> None: sweep_stuck_runs(store, call_state=lambda _: CallState.ENDED, now=NOW) assert store.events == [ + "list-expired:report_run", "list:report_run", "fail:report_run", - "list-expired:report_run", + "list-expired:simulation_run", "list:simulation_run", "fail:simulation_run", - "list-expired:simulation_run", ] @@ -150,8 +152,8 @@ def test_a_run_another_sweep_failed_first_is_not_counted_again() -> None: run = _run("fc-ended") class AlreadyFailedStore(FakeSweepStore): - def fail_run_of_ended_call(self, run, *, error_message, failed_at): - super().fail_run_of_ended_call( + 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) @@ -287,3 +289,26 @@ def test_an_unrecognised_input_status_leaves_the_call_unknown() -> None: 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 + assert asked == [("fc-1", ["list-expired:report_run", "list:report_run"])] From 5cf46f9e41e3597a593d0799b906229b29a4dafd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:52:49 +0400 Subject: [PATCH 9/9] Fail both kinds' overage runs before the sweep asks Modal anything The sweep checked simulation runs' age only after every report run's Modal lookup, so lookups that stalled could still keep simulation runs past the maximum age from being failed. It now fails report runs, then simulation runs, past the maximum age before any lookup, then asks about each kind's calls, report runs first. A refused failure in the age check counts as left, and the docstrings say a write is refused also when a running report run started combining the simulation run. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/run_sweep.py | 32 ++++++++++++------- .../tests/test_stage12_run_sweep.py | 29 +++++++++++++---- 2 files changed, 43 insertions(+), 18 deletions(-) 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 index fc6351195..d5f386d7a 100644 --- 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 @@ -84,7 +84,12 @@ def list_running_runs( def fail_stuck_run( self, run: RunningRun, *, error_message: str, failed_at: datetime ) -> RunWrite: - """Mark ``run`` failed while it is still running under its listed call.""" + """Mark ``run`` failed while it is as it was listed. + + It must still be running under its listed call, or with no call if it + was listed with none; a simulation run must still have no running + report run combining it. + """ ... @@ -93,9 +98,11 @@ class SweepSummary: """What the sweep did with the runs it listed. ``checked`` runs had their call asked about: ``failed`` because it ended, - left because it is ``running`` or its state is ``unknown``, or ``left`` - because the run finished or moved to another call after it was listed. - ``failed_for_age`` runs were failed for exceeding the maximum age. + or left because it is ``running`` or its state is ``unknown``. + ``failed_for_age`` runs were failed for exceeding the maximum age. ``left`` + counts listed runs whose conditional failure changed nothing, because + they finished, moved to another call, or a running report run started + combining them after they were listed. """ checked: int @@ -166,13 +173,13 @@ def sweep_stuck_runs( # the store, then keyword-only seams and limits ) -> SweepSummary: """Fail every listed run whose recorded Modal call has ended. - Report runs are swept first, so a dead coordinator's simulation runs are - listed in the same sweep. For each kind, runs older than ``maximum_age`` - are failed first, before any Modal lookup can use up the scheduled - function's time. Each call is then asked about once. A run whose call is - running, or whose state is unknown, is left as it is; so is one that - finished or moved to another call since it was listed, which makes two - sweeps at once safe. + 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 each recorded call is asked about once, report + runs first, so a dead coordinator's simulation runs are listed in the + same sweep. A run whose call is running, or whose state is unknown, is + left as it is; 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] = {} @@ -185,6 +192,9 @@ def sweep_stuck_runs( # the store, then keyword-only seams and limits 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_running_runs(kind, started_before=now - minimum_age): _sweep_run( store, run, states=states, call_state=call_state, now=now, tally=tally diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py index 1b4ddcc59..a574a187c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -43,13 +43,17 @@ def __init__(self, runs, *, expired=()): def list_running_runs( self, kind, *, started_before, include_runs_without_call=False ): + # As the real store, a run this fake failed is no longer listed. + failed = {failed_run for failed_run, _, _ in self.failed} if include_runs_without_call: self.expired_cutoffs.append(started_before) self.events.append(f"list-expired:{kind.value}") - return tuple(run for run in self.expired if run.kind is kind) - self.cutoffs.append(started_before) - self.events.append(f"list:{kind.value}") - return tuple(run for run in self.runs if run.kind is kind) + listed = self.expired + else: + self.cutoffs.append(started_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)) @@ -108,9 +112,9 @@ def test_report_runs_are_swept_before_simulation_runs_are_listed() -> None: assert store.events == [ "list-expired:report_run", + "list-expired:simulation_run", "list:report_run", "fail:report_run", - "list-expired:simulation_run", "list:simulation_run", "fail:simulation_run", ] @@ -278,7 +282,8 @@ def test_a_run_past_the_maximum_age_is_failed_whatever_modal_reports() -> None: for _, message, _ in store.failed ) assert summary.failed_for_age == 2 - assert summary.running == 1 + # 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: @@ -311,4 +316,14 @@ def call_state(call_id): assert store.expired_cutoffs == [NOW - timedelta(hours=24)] * 2 assert store.cutoffs == [NOW - timedelta(minutes=10)] * 2 - assert asked == [("fc-1", ["list-expired:report_run", "list:report_run"])] + # Both kinds' age checks come before any lookup. + assert asked == [ + ( + "fc-1", + [ + "list-expired:report_run", + "list-expired:simulation_run", + "list:report_run", + ], + ) + ]