Fail report runs and simulation runs left running after their Modal call ended

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.

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",

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,
+    )