diff --git a/changelog.d/stuck-run-sweep.added.md b/changelog.d/stuck-run-sweep.added.md new file mode 100644 index 000000000..6f5418b29 --- /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 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/__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..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 @@ -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 ( @@ -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 @@ -52,6 +52,28 @@ 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. + + ``modal_call_id`` is None for a run marked running with no call recorded. + """ + + kind: RunKind + run_id: UUID + modal_call_id: str | None + + @dataclass(frozen=True) class RunWrite: """One conditional write's outcome: whether it changed the run, and its status.""" @@ -304,6 +326,86 @@ def fail_simulation_run( simulation_runs, simulation_run_id, statement, "fail_simulation_run" ) + def list_running_runs( + self, + kind: RunKind, + *, + started_before: datetime, + include_runs_without_call: bool = False, + ) -> tuple[RunningRun, ...]: + """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, + 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(_no_running_report_run_combines(table)) + 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_stuck_run( + self, + run: RunningRun, + *, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark a stuck ``run`` failed, while it is as it was listed. + + 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 = ( + 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, + same_call, + *( + (_no_running_report_run_combines(table),) + if run.kind is RunKind.SIMULATION_RUN + else () + ), + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write(table, run.run_id, statement, "fail_stuck_run") + def _write( self, table: Table, @@ -357,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 new file mode 100644 index 000000000..3edb21886 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_report_run_sweep_store.py @@ -0,0 +1,321 @@ +"""The stuck-run sweep's reads and writes, as the Modal worker role. + +These run against a disposable Postgres that the Public API's v2 migrations +created (see ``report_run_database``). +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier +from uuid import uuid4 + +import pytest +from policyengine_stage12_persistence import ( + ReportRunStore, + RunKind, + RunningRun, + RunStatus, + 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" + + +@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(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 _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, + "simulation_runs", + seeded_run.run_ids.reform_simulation_run_id, + status="running", + started_at=NOW, + ) + + assert _listed_for(store, seeded_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 _listed(store, 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_stuck_run( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + again = store.fail_stuck_run( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + simulation = store.fail_stuck_run( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert first == RunWrite(changed=True, status=RunStatus.FAILED) + assert again == RunWrite(changed=False, status=RunStatus.FAILED) + assert simulation == RunWrite(changed=True, status=RunStatus.FAILED) + row = seeder.row("report_runs", seeded_run.run_ids.report_run_id) + assert row["error_message"] == ENDED_MESSAGE + assert row["completed_at"] == CUTOFF + + +def test_the_sweep_leaves_a_run_that_moved_to_another_call_or_finished( + store, seeder, seeded_run: SeededReportRun +) -> None: + report_run, baseline_run = _running(seeder, seeded_run) + # After the sweep listed them, Modal restarted the baseline's worker + # under a new call, and the report run succeeded. + _set_run( + seeder, + "simulation_runs", + seeded_run.run_ids.baseline_simulation_run_id, + modal_call_id="fc-baseline-restarted", + ) + _set_run( + seeder, + "report_runs", + seeded_run.run_ids.report_run_id, + status="succeeded", + completed_at=CUTOFF, + ) + + moved = store.fail_stuck_run( + baseline_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + finished = store.fail_stuck_run( + report_run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ) + + assert moved == RunWrite(changed=False, status=RunStatus.RUNNING) + assert finished == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + + +def test_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_stuck_run( + run, error_message=ENDED_MESSAGE, failed_at=CUTOFF + ).changed + for run in runs + ] + + try: + with ThreadPoolExecutor(max_workers=len(stores)) as pool: + changes = list(pool.map(sweep, stores)) + finally: + for run_store in stores: + run_store.engine.dispose() + + assert [sum(column) for column in zip(*changes, strict=True)] == [1, 1] + 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_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 new file mode 100644 index 000000000..793714002 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -0,0 +1,82 @@ +"""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 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: + + 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( + # This module imports src.modal, so the container needs it too. + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + copy=True, + ) +) + + +@app.function( + image=sweep_image, + schedule=modal.Cron("*/10 * * * *"), + cpu=0.25, + memory=512, + timeout=600, + retries=0, + max_containers=1, + secrets=[modal_worker_database_secret], +) +def sweep_stuck_runs() -> dict: + """Fail the runs 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/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..d5f386d7a --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/run_sweep.py @@ -0,0 +1,244 @@ +"""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 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. +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 + +import logging +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from enum import StrEnum +from importlib import import_module +from typing import TYPE_CHECKING, Any, Protocol + +from policyengine_stage12_persistence import RunKind + +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, 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." +) +# 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 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 +# and waiting to restart (internally INTERNAL_FAILURE) as pending. Any other +# name leaves the call's state unknown. +RUNNING_INPUT_STATUSES = frozenset({"PENDING"}) +ENDED_INPUT_STATUSES = frozenset( + {"SUCCESS", "FAILURE", "TERMINATED", "TIMEOUT", "INIT_FAILURE"} +) + + +class CallState(StrEnum): + """What Modal reports about a recorded call.""" + + RUNNING = "running" + ENDED = "ended" + UNKNOWN = "unknown" + + +class StuckRunStore(Protocol): + """The sweep's reads and writes; ``ReportRunStore`` implements them.""" + + def list_running_runs( + self, + kind: RunKind, + *, + started_before: datetime, + include_runs_without_call: bool = False, + ) -> tuple[RunningRun, ...]: + """Return the runs of ``kind`` running since before a time.""" + ... + + def fail_stuck_run( + self, run: RunningRun, *, error_message: str, failed_at: datetime + ) -> RunWrite: + """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. + """ + ... + + +@dataclass(frozen=True) +class SweepSummary: + """What the sweep did with the runs it listed. + + ``checked`` runs had their call asked about: ``failed`` because it ended, + 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 + 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 + """Return whether Modal reports the call as running or ended. + + Modal returns the whole call graph the call is part of, where a call + spawned inside another call's input is nested under it, so the call's own + inputs are found by their call ID anywhere in the graph. A call whose + inputs are all finished has ended; one with any pending input is running. + A call missing from its graph, a status this code doesn't know, or any + error asking, is unknown, so the sweep leaves its run alone. + """ + modal = modal_module if modal_module is not None else import_module("modal") + try: + graph = modal.FunctionCall.from_id(modal_call_id).get_call_graph() + statuses = { + input_info.status.name for input_info in _inputs_of(graph, modal_call_id) + } + except Exception as error: # noqa: BLE001 # any failure to ask leaves the run as it is + _log_unreadable_call(modal_call_id, error) + return CallState.UNKNOWN + if statuses & RUNNING_INPUT_STATUSES: + return CallState.RUNNING + if statuses and statuses <= ENDED_INPUT_STATUSES: + return CallState.ENDED + return CallState.UNKNOWN + + +def _inputs_of(graph: list[Any], modal_call_id: str) -> list[Any]: + """Return every input in ``graph`` that belongs to call ``modal_call_id``.""" + found: list[Any] = [] + pending = list(graph) + while pending: + input_info = pending.pop() + if input_info.function_call_id == modal_call_id: + found.append(input_info) + pending.extend(input_info.children) + return found + + +@dataclass +class _Tally: + """Running counts of one sweep.""" + + checked: int = 0 + failed: int = 0 + running: int = 0 + unknown: int = 0 + left: int = 0 + failed_for_age: int = 0 + + +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. + + 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] = {} + tally = _Tally() + for kind in (RunKind.REPORT_RUN, RunKind.SIMULATION_RUN): + for run in store.list_running_runs( + kind, started_before=now - maximum_age, include_runs_without_call=True + ): + if store.fail_stuck_run( + run, error_message=EXCEEDED_AGE_MESSAGE, failed_at=now + ).changed: + tally.failed_for_age += 1 + else: + tally.left += 1 + for kind in (RunKind.REPORT_RUN, RunKind.SIMULATION_RUN): + for run in store.list_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)) + + +def _sweep_run( # noqa: PLR0913 # one run, the store, and the sweep's shared state + store: StuckRunStore, + run: RunningRun, + *, + states: dict[str, CallState], + call_state: Callable[[str], CallState], + now: datetime, + 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_stuck_run( + run, error_message=ENDED_CALL_MESSAGE, failed_at=now + ).changed: + tally.failed += 1 + else: + tally.left += 1 + + +def _log_unreadable_call(modal_call_id: str, error: Exception) -> None: + logger.warning( + "Could not read the status of Modal call %s (%s)", + modal_call_id, + type(error).__name__, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py new file mode 100644 index 000000000..a574a187c --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -0,0 +1,329 @@ +"""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, *, expired=()): + self.runs = list(runs) + # 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 = [] + + 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}") + 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)) + self.events.append(f"fail:{run.kind.value}") + 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, left=0, failed_for_age=0 + ) + + +def test_the_sweep_lists_runs_older_than_its_minimum_age() -> None: + store = FakeSweepStore([]) + + sweep_stuck_runs( + store, + call_state=lambda _: CallState.ENDED, + now=NOW, + minimum_age=timedelta(minutes=15), + ) + + assert store.cutoffs == [NOW - timedelta(minutes=15)] * 2 + + +def test_report_runs_are_swept_before_simulation_runs_are_listed() -> None: + # A simulation run is listed only once no report run 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-expired:report_run", + "list-expired:simulation_run", + "list:report_run", + "fail:report_run", + "list:simulation_run", + "fail:simulation_run", + ] + + +def test_a_call_running_several_runs_is_asked_once() -> None: + report_run = _run("fc-coordinator", RunKind.REPORT_RUN) + other = _run("fc-coordinator", RunKind.REPORT_RUN) + asked = [] + + def call_state(call_id): + asked.append(call_id) + return CallState.ENDED + + sweep_stuck_runs( + FakeSweepStore([report_run, other]), call_state=call_state, now=NOW + ) + + assert asked == ["fc-coordinator"] + + +def test_a_failed_status_lookup_never_fails_the_run() -> None: + run = _run("fc-error") + store = FakeSweepStore([run]) + + def call_state(_call_id): + msg = "Modal is unreachable" + raise ConnectionError(msg) + + summary = sweep_stuck_runs(store, call_state=call_state, now=NOW) + + assert store.failed == [] + assert summary == SweepSummary( + checked=1, failed=0, running=0, unknown=1, left=0, failed_for_age=0 + ) + + +def test_a_run_another_sweep_failed_first_is_not_counted_again() -> None: + run = _run("fc-ended") + + class AlreadyFailedStore(FakeSweepStore): + def fail_stuck_run(self, run, *, error_message, failed_at): + super().fail_stuck_run( + run, error_message=error_message, failed_at=failed_at + ) + return RunWrite(changed=False, status=RunStatus.FAILED) + + summary = sweep_stuck_runs( + AlreadyFailedStore([run]), call_state=lambda _: CallState.ENDED, now=NOW + ) + + assert summary.failed == 0 + # It finished, or moved to another call, after it was listed. + assert summary.left == 1 + + +def _node(call_id, status, children=()): + return SimpleNamespace( + function_call_id=call_id, status=status, children=list(children) + ) + + +def _fake_modal(statuses=None, error=None, *, graph=None): + """Return a modal stand-in whose call graph holds call fc-1's inputs. + + By default fc-1 is a top-level call; ``graph`` gives the whole graph a + spawned call is part of instead. + """ + + def from_id(call_id): + if error is not None: + raise error + assert call_id == "fc-1" + roots = ( + graph + if graph is not None + else [_node("fc-1", status) for status in (statuses or [])] + ) + return SimpleNamespace(get_call_graph=lambda: roots) + + return SimpleNamespace(FunctionCall=SimpleNamespace(from_id=from_id)) + + +@pytest.mark.parametrize( + ("statuses", "expected"), + [ + ([FakeInputStatus.PENDING], CallState.RUNNING), + ([FakeInputStatus.SUCCESS], CallState.ENDED), + ([FakeInputStatus.FAILURE], CallState.ENDED), + ([FakeInputStatus.TERMINATED], CallState.ENDED), + ([FakeInputStatus.TIMEOUT], CallState.ENDED), + ([FakeInputStatus.INIT_FAILURE], CallState.ENDED), + ([FakeInputStatus.FAILURE, FakeInputStatus.PENDING], CallState.RUNNING), + ([], CallState.UNKNOWN), + ], +) +def test_modal_call_state_reads_the_call_graph(statuses, expected) -> None: + assert modal_call_state("fc-1", modal_module=_fake_modal(statuses)) is expected + + +def test_modal_call_state_is_unknown_when_modal_cannot_answer() -> None: + state = modal_call_state( + "fc-1", modal_module=_fake_modal(error=ConnectionError("unreachable")) + ) + + assert state is CallState.UNKNOWN + + +@pytest.mark.parametrize( + ("coordinator_status", "worker_status", "expected"), + [ + # A coordinator that ran out of memory left its worker running. + (FakeInputStatus.FAILURE, FakeInputStatus.PENDING, CallState.RUNNING), + (FakeInputStatus.PENDING, FakeInputStatus.TERMINATED, CallState.ENDED), + ], +) +def test_modal_call_state_reads_the_recorded_calls_own_inputs( + coordinator_status, worker_status, expected +) -> None: + # A worker is spawned inside its coordinator's input, so Modal nests its + # inputs under the coordinator's in the call graph. + graph = [ + _node( + "fc-coordinator", + coordinator_status, + children=[ + _node("fc-1", worker_status, children=[_node("fc-segment", "x")]), + _node("fc-other-worker", FakeInputStatus.SUCCESS), + ], + ) + ] + + assert modal_call_state("fc-1", modal_module=_fake_modal(graph=graph)) is expected + + +def test_a_call_missing_from_its_call_graph_is_unknown() -> None: + graph = [_node("fc-coordinator", FakeInputStatus.SUCCESS)] + + assert ( + modal_call_state("fc-1", modal_module=_fake_modal(graph=graph)) + is CallState.UNKNOWN + ) + + +def test_a_run_past_the_maximum_age_is_failed_whatever_modal_reports() -> None: + # Modal reports a call it won't retry after an internal failure as + # pending, and no longer describes old calls; neither may keep a run + # running for good, and nor may a run with no recorded call. + reported_running = _run("fc-reported-running", RunKind.REPORT_RUN) + 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 + # Failed for its age first, it is no longer listed for a lookup. + assert summary.running == 0 + + +def test_an_unrecognised_input_status_leaves_the_call_unknown() -> None: + unrecognised = SimpleNamespace(name="SOMETHING_NEW") + graph = [_node("fc-1", unrecognised)] + + assert ( + modal_call_state("fc-1", modal_module=_fake_modal(graph=graph)) + is CallState.UNKNOWN + ) + + +def test_the_age_limit_is_checked_first_with_its_own_cutoff() -> None: + # Before any Modal lookup, so lookups that stall can't keep it from + # running within the scheduled function's time limit. + store = FakeSweepStore([_run("fc-1", RunKind.REPORT_RUN)]) + asked = [] + + def call_state(call_id): + asked.append((call_id, list(store.events))) + return CallState.RUNNING + + sweep_stuck_runs( + store, + call_state=call_state, + now=NOW, + minimum_age=timedelta(minutes=10), + maximum_age=timedelta(hours=24), + ) + + assert store.expired_cutoffs == [NOW - timedelta(hours=24)] * 2 + assert store.cutoffs == [NOW - timedelta(minutes=10)] * 2 + # Both kinds' age checks come before any lookup. + assert asked == [ + ( + "fc-1", + [ + "list-expired:report_run", + "list-expired:simulation_run", + "list:report_run", + ], + ) + ] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py new file mode 100644 index 000000000..f3d47ef13 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py @@ -0,0 +1,82 @@ +"""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" + # The app module imports src.modal at module level, in the container too. + assert calls["add_local_python_source"][0] == ( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + ) + + +def test_the_scheduled_function_sweeps_with_the_report_run_store(monkeypatch) -> None: + module = _load(monkeypatch) + store = object() + swept = [] + monkeypatch.setattr(report_run_access_module, "report_run_store", lambda: store) + + def recording_sweep(received_store): + swept.append(received_store) + return run_sweep.SweepSummary( + checked=3, failed=1, running=1, unknown=1, left=0, failed_for_age=2 + ) + + monkeypatch.setattr(run_sweep, "sweep_stuck_runs", recording_sweep) + + assert module.sweep_stuck_runs() == { + "checked": 3, + "failed": 1, + "running": 1, + "unknown": 1, + "left": 0, + "failed_for_age": 2, + } + assert swept == [store]