diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index b89de5283..5df51c15c 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -1189,22 +1189,34 @@ jobs: run: ../../.github/scripts/stage12-v2-default-verify.sh # The scheduled sweep that fails report runs, simulation runs and - # household jobs nothing will finish. Its app is unversioned and its - # schedule starts when it deploys, so it deploys here, after every - # check, with the release's Stage 12 default; if it fails, the next - # step puts the default back. + # household jobs nothing will finish, and the restart check that starts + # again those whose Modal call never started. Their app is unversioned + # and its schedules start when it deploys, so it deploys here, after + # every check, with the release's Stage 12 default; if it fails, the + # next step puts the default back. # Only from main: a dispatch from another branch, which may be behind - # main, leaves the environment's sweep as it is. With no status + # main, leaves the environment's app as it is. With no status # function, the condition also requires every earlier step to pass. - # Its database URL comes from the Modal secret + # Their database URL comes from the Modal secret # policyengine-modal-worker-database, which the v2 workers' deploy in - # deploy_stage12_v2 already requires. - - name: Deploy the stuck-run sweep + # deploy_stage12_v2 already requires. The image bakes in the Modal + # environment the restart check starts calls in, and the v2 workers' + # telemetry settings, from the repository variables they use. + - name: Deploy the stuck-run sweep and the restart check if: ${{ github.ref == 'refs/heads/main' }} working-directory: projects/policyengine-simulation-executor env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + SIMULATION_V2_MODAL_ENVIRONMENT: ${{ inputs.modal_environment }} + OBSERVABILITY_SERVICE_NAMESPACE: ${{ vars.OBSERVABILITY_SERVICE_NAMESPACE }} + OBSERVABILITY_TRACE_PROJECT_ID: ${{ vars.OBSERVABILITY_TRACE_PROJECT_ID }} + OBSERVABILITY_LOGGING_PROJECT_ID: ${{ vars.OBSERVABILITY_LOGGING_PROJECT_ID }} + OBSERVABILITY_LOG_NAME: ${{ vars.OBSERVABILITY_LOG_NAME }} + OBSERVABILITY_GOOGLE_WORKLOAD_IDENTITY_PROVIDER: ${{ vars.OBSERVABILITY_GOOGLE_WORKLOAD_IDENTITY_PROVIDER }} + OBSERVABILITY_GOOGLE_SERVICE_ACCOUNT_EMAIL: ${{ vars.OBSERVABILITY_GOOGLE_SERVICE_ACCOUNT_EMAIL }} + OTEL_EXPORTER_OTLP_ENDPOINT: ${{ vars.OBSERVABILITY_OTLP_ENDPOINT }} + POLICYENGINE_OTEL_GOOGLE_AUDIENCE: ${{ vars.OBSERVABILITY_OTLP_GOOGLE_AUDIENCE }} run: uv run --frozen modal deploy --env="${{ inputs.modal_environment }}" src/modal/run_sweep_app.py # Writes back the manifest the move replaced, only while the stored diff --git a/changelog.d/c1-19-restart-unstarted-runs.added.md b/changelog.d/c1-19-restart-unstarted-runs.added.md new file mode 100644 index 000000000..2ad18f24b --- /dev/null +++ b/changelog.d/c1-19-restart-unstarted-runs.added.md @@ -0,0 +1 @@ +A restart check in the stuck-run sweep's Modal application starts the Modal call of a household job, economy report run or standalone economy simulation run again when the run stayed pending with no call recorded 3 to 10 minutes after its creation, as when the Simulation Entrypoint stopped before starting its call; its telemetry counts every restart and how it ended. The sweep now fails a run still pending with no call at 15 minutes instead of 10. diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index c03186693..b74e7e47e 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -359,7 +359,7 @@ run's maximum age of 24 hours. anything: their `coordinate_report` requires an evaluation record, and their context model refuses `simulation_run_id`. Modal accepts the call, which fails at once, and the run stays pending with its call recorded until the sweep fails -it (`modal_call_ended`), at least ten minutes after its creation. A call +it (`modal_call_ended`), at least 15 minutes after its creation. A call reaches such an application only while the manifest's default, or `STAGE12_V2_WORKER_VERSION`, names one: during the deployment of this change, if it adds a newer .py version, until it moves the default to it, or while a @@ -410,7 +410,7 @@ single-year report's runs, and the reverse. - **The stuck-run sweep.** It leaves a waiting year of an unfinished window to the window, counts a year started later from when its call was recorded, fails the window when it fails one of its years, and settles a window that - nothing has changed for ten minutes with no year under way: from its years + nothing has changed for 15 minutes with no year under way: from its years when they have all ended, or as stalled (`budget_window_stalled`) when years are still waiting. @@ -762,7 +762,7 @@ call that names no job, as deployment validation makes, touches no database. The stuck-run sweep (`src/modal/run_sweep_app.py`) fails household jobs nothing will finish by the rule it applies to report runs and simulation runs, counting a job's age from its creation: a pending job with no call -recorded ten minutes after its creation, a pending or running job once +recorded 15 minutes after its creation, a pending or running job once Modal reports its recorded call ended, and any job still unfinished 24 hours after its creation, whatever Modal reports. @@ -792,6 +792,40 @@ validation like any other validation failure: the manifest is not published, and, as described under "Release and deployment constraints", the deployment does not promote the new Simulation Entrypoint revision. +**Restarting a run whose call never started.** The stuck-run sweep's Modal +application also runs a restart check every two minutes, as the Modal worker +role. It lists, at most 20 a check, the household jobs, report runs and +standalone simulation runs created 3 to 10 minutes ago that are still pending +with no call recorded and that nothing has changed since their insert, as +when the Simulation Entrypoint stopped between creating a run and starting +its call. For each, on a thread, it reads back what the entrypoint built the +call from and builds the same call with the entrypoint's own builders, on the +worker of the .py version the row records. It then sets the row's +`updated_at`, so the row is never restarted twice, starts the call and +records its ID while the row is still pending with none. A row it can't start +again (its worker left the manifest, its stored rows changed, or Modal +refuses the call) is failed with `run_not_started`, a report run with its +pending simulation runs. A row is left unmarked for the next check when its +rows or manifest can't be read now, its rebuild raises an unexpected error +(logged with its traceback), its mark can't be written, its call can't be +rebuilt and its failure can't be written, or the check's time is up. After +three minutes a check marks no more rows, so it ends within its five-minute +timeout without leaving a marked row with no call. The sweep's 15 minutes come after the restart window, so a second miss +is still failed with `run_not_started`. The restart code is the simulation +service's (`policyengine_simulation_entry.run_restarts`, `run_restart_calls` +and `run_restart_check`, on `policyengine_simulation_persistence`), which the +image adds as source. + +It leaves alone a budget window's report runs and household simulations' +runs: it can't rebuild their calls yet. Each check is a +`v2_run_restart_check` operation and each row it takes up a +`v2_run_restart` operation, with the row's kind, ID, age and outcome, and +each outcome is also logged as a `v2_run_restart_outcome` event. A deferred +row is taken up again by the next check, so the distinct run IDs are the +runs that needed a restart that a check took up; a row no check took up, in +a burst larger than the checks in its window had time for, is failed by the +sweep with `run_not_started` instead. The events say how each ended. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py index d2c72e5f5..5f72afe1a 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -27,6 +27,7 @@ class RunConfiguration(StrEnum): STAGE12_SIMULATION = "stage12_simulation" STAGE12_SEGMENT = "stage12_segment" V2_HOUSEHOLD = "v2_household" + V2_RUN_RESTART = "v2_run_restart" class Stage(StrEnum): @@ -116,6 +117,8 @@ class Stage(StrEnum): V2_HOUSEHOLD_COUNTRY_MODEL_LOAD = "v2_household_country_model_load" V2_HOUSEHOLD_INPUT_NORMALIZATION = "v2_household_input_normalization" V2_HOUSEHOLD_CALCULATION = "v2_household_calculation" + V2_RUN_RESTART_CHECK = "v2_run_restart_check" + V2_RUN_RESTART = "v2_run_restart" @dataclass(frozen=True) @@ -287,6 +290,12 @@ def name(self, stage: Stage) -> str: Stage.V2_HOUSEHOLD_CALCULATION, ), ), + # One scheduled restart check, then each run or household job it + # lists because its Modal call never started. + RunConfiguration.V2_RUN_RESTART: StagePlan( + RunConfiguration.V2_RUN_RESTART, + (Stage.V2_RUN_RESTART_CHECK, Stage.V2_RUN_RESTART), + ), } ) @@ -302,3 +311,4 @@ def name(self, stage: Stage) -> str: STAGE12_SIMULATION_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SIMULATION] STAGE12_SEGMENT_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SEGMENT] V2_HOUSEHOLD_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.V2_HOUSEHOLD] +V2_RUN_RESTART_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.V2_RUN_RESTART] diff --git a/libs/policyengine-simulation-observability/tests/test_stages.py b/libs/policyengine-simulation-observability/tests/test_stages.py index 0861363b9..bfa157a42 100644 --- a/libs/policyengine-simulation-observability/tests/test_stages.py +++ b/libs/policyengine-simulation-observability/tests/test_stages.py @@ -1,6 +1,7 @@ import pytest from policyengine_simulation_observability.stages import ( RUN_STAGE_REGISTRY, + V2_RUN_RESTART_STAGES, RunConfiguration, Stage, ) @@ -66,3 +67,13 @@ def test_v2_household_plan_measures_model_load_inputs_and_each_policy() -> None: Stage.V2_HOUSEHOLD_CALCULATION, ) assert household.name(Stage.V2_HOUSEHOLD_CALCULATION) == "v2_household_calculation" + + +def test_v2_run_restart_plan_measures_each_check_and_each_restart() -> None: + """The scheduled check is one operation; each row it lists is one more.""" + restart = RUN_STAGE_REGISTRY[RunConfiguration.V2_RUN_RESTART] + + assert restart is V2_RUN_RESTART_STAGES + assert restart.stages == (Stage.V2_RUN_RESTART_CHECK, Stage.V2_RUN_RESTART) + assert restart.name(Stage.V2_RUN_RESTART_CHECK) == "v2_run_restart_check" + assert restart.name(Stage.V2_RUN_RESTART) == "v2_run_restart" diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index d6ac6ee63..bdfd3794b 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -125,6 +125,37 @@ commit fails after its rows were inserted raises `EconomyRecordOutcomeUnknownError`: PostgreSQL may have applied the commit, so the rows may exist. +## Restarting runs whose call never started + +`run_restart_records.py` holds the restart check's reads and writes. The +check runs as the Modal worker role, whose grants cover them: `SELECT` on +the run tables and their inputs, and `UPDATE` of `updated_at`, +`modal_call_id`, `status`, `error_message` and `completed_at`. + +- `list_unstarted_runs` returns the oldest rows, up to a limit, created in a + window that are still `pending` with no Modal call and whose `updated_at` + still equals `created_at`: household jobs, report runs outside a budget + window whose baseline is an economy simulation, and runs of economy + simulations that no report run names (standalone runs). +- `mark_restart` sets a listed row's `updated_at` to the statement's clock + while it is still unstarted, so it is listed and marked only once. +- `read_unstarted_report_run` and `read_unstarted_simulation_run` read back + what the Simulation Entrypoint built a call from: each side's simulation + run, policy and dynamic, with the values read against the run's own + catalog version, the region, dataset and year the run records, the + report's options and the run's versions. `StoredRunUnreadableError` says + the rows no longer describe one calculation: a side's region or dataset + is gone, or the two sides differ. The values are read with + `read_simulation_reform`, the reader a submission uses. +- `read_report_run_simulation_run_ids` returns the simulation runs a report + run names, so one that can't be rebuilt is failed with them, and + `read_run_call` returns a row's status and recorded call, which says whose + call holds it after a restart. + +`RunRestartStore` runs each in a transaction of its own. The check records +a call, or fails a row, with `EconomyRunStore`'s and `HouseholdJobStore`'s +own conditional writes. + ## Reading a policy `read_policy_reform` (`policy_reform.py`) reads a stored policy as the reform diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index bea874cc0..3192a61e1 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -57,6 +57,7 @@ SimulationInputs, UnknownDatasetError, UnknownRegionError, + read_simulation_reform, ) from policyengine_simulation_persistence.engine import ( bounded_transaction, @@ -92,6 +93,24 @@ read_dynamic_reform, read_policy_reform, ) +from policyengine_simulation_persistence.run_restart_records import ( + RESTART_TABLES, + EconomyScope, + RestartKind, + RunCall, + StoredEconomySide, + StoredReportRun, + StoredRunUnreadableError, + StoredSimulationRun, + UnstartedRun, + list_unstarted_runs, + mark_restart, + read_report_run_simulation_run_ids, + read_run_call, + read_unstarted_report_run, + read_unstarted_simulation_run, +) +from policyengine_simulation_persistence.run_restart_store import RunRestartStore from policyengine_simulation_persistence.tables import ( LATEST_YEAR, MAX_BUDGET_WINDOW_YEARS, @@ -113,6 +132,7 @@ __all__ = [ "LATEST_YEAR", "MAX_BUDGET_WINDOW_YEARS", + "RESTART_TABLES", "SERVICE_TABLE_PRIVILEGES", "SERVICE_UPDATABLE_COLUMNS", "BudgetWindowInputs", @@ -127,6 +147,7 @@ "EconomyReportState", "EconomyRequest", "EconomyRunStore", + "EconomyScope", "EconomySimulationInputs", "EconomySimulationState", "HouseholdCalculationInputs", @@ -147,7 +168,10 @@ "ReportSide", "ReportSideError", "ReportSimulation", + "RestartKind", "ReusableRun", + "RunCall", + "RunRestartStore", "RunReuse", "RunVersions", "SideRefusal", @@ -156,13 +180,18 @@ "SimulationRunState", "SimulationStatus", "SimulationType", + "StoredEconomySide", "StoredHousehold", + "StoredReportRun", + "StoredRunUnreadableError", + "StoredSimulationRun", "UnknownCatalogVersionError", "UnknownDatasetError", "UnknownDynamicError", "UnknownHouseholdError", "UnknownPolicyError", "UnknownRegionError", + "UnstartedRun", "bounded_transaction", "budget_window_runs", "combine_reforms", @@ -177,6 +206,8 @@ "insert_economy_report", "insert_economy_simulation", "insert_household_job", + "list_unstarted_runs", + "mark_restart", "missing_table_privileges", "normalized_database_url", "read_catalog_version_id", @@ -190,7 +221,12 @@ "read_model_id", "read_policy_reform", "read_region", + "read_report_run_simulation_run_ids", + "read_run_call", + "read_simulation_reform", "read_stored_household", + "read_unstarted_report_run", + "read_unstarted_simulation_run", "record_household_call", "record_report_run_call", "record_simulation_run_call", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py index 2ba26499c..f7644b9ee 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py @@ -529,6 +529,29 @@ def _simulation_inputs( choice: SimulationChoice, ) -> SimulationInputs: """Read one simulation's policy and dynamic values, combined.""" + return SimulationInputs( + choice=choice, + reform=read_simulation_reform( + connection, + country=country, + catalog_version_id=catalog_version_id, + choice=choice, + ), + ) + + +def read_simulation_reform( + connection: Connection, + *, + country: str, + catalog_version_id: UUID, + choice: SimulationChoice, +) -> PolicyReform: + """Read a simulation's policy and dynamic values on a catalog version, combined. + + A submission and a restart of its run read them this way. Raise the + policy reader's errors, or OverlappingReformError. + """ policy_reform = ( {} if choice.policy_id is None @@ -548,9 +571,7 @@ def _simulation_inputs( catalog_version_id=catalog_version_id, ) ) - return SimulationInputs( - choice=choice, reform=combine_reforms(policy_reform, dynamic_reform) - ) + return combine_reforms(policy_reform, dynamic_reform) def _side_inputs( diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_records.py new file mode 100644 index 000000000..0613d1f5c --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_records.py @@ -0,0 +1,446 @@ +"""The restart check's reads and writes, for runs whose Modal call never started. + +The simulation service creates a run or household job ``pending`` with no +Modal call ID, starts its call, then records the call's ID. When the service +stops between the insert and the start, nothing ever records a call, and the +stuck-run sweep would fail the run. The restart check, a scheduled Modal +function, starts such a run's call once more: + +1. ``list_unstarted_runs`` lists the rows still pending with no call that were + created within a window and that nothing has changed since their insert + (``updated_at`` still equals ``created_at``): household jobs, report runs, + and standalone simulation runs, which no report run names. A report run's + own simulation runs are its coordinator's to start. Budget-window years + and household simulations are left out: the check can't rebuild their + calls yet. +2. ``mark_restart`` sets the row's ``updated_at`` before the check starts its + call, so the row is never listed again and is restarted at most once. It + uses the statement's clock rather than the transaction's, so the mark + differs from a creation time in the same transaction too. +3. ``read_unstarted_report_run`` and ``read_unstarted_simulation_run`` read + back what the service built the call from: each side's simulation and + values, read against the run's own catalog version, the region, dataset + and year the run records, and the report's options. + +The check then records its call's ID, or fails the row, with the service's +own conditional writes (``record_report_run_call`` and the rest). Every +function here takes a connection, so the caller chooses the transactions, +and every update sets only columns both the service's and the Modal worker's +roles may update. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING, Any, Final + +from sqlalchemy import exists, func, or_, select, update + +from policyengine_simulation_persistence.economy_run_records import ( + RunVersions, + SimulationChoice, +) +from policyengine_simulation_persistence.economy_run_store import ( + read_simulation_reform, +) +from policyengine_simulation_persistence.input_tables import ( + MODEL_NAME_PREFIX, + datasets, + regions, + tax_benefit_model_versions, + tax_benefit_models, +) +from policyengine_simulation_persistence.tables import ( + SimulationStatus, + SimulationType, + household_jobs, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, Row, Select, Table + + from policyengine_simulation_persistence.policy_reform import PolicyReform + +# Every restartable table's status uses this value for a row not yet claimed. +_PENDING: Final = SimulationStatus.PENDING.value + + +class RestartKind(StrEnum): + """What kind of row the restart check starts again.""" + + HOUSEHOLD_JOB = "household_job" + REPORT_RUN = "report_run" + SIMULATION_RUN = "simulation_run" + + +RESTART_TABLES: Final[dict[RestartKind, Table]] = { + RestartKind.HOUSEHOLD_JOB: household_jobs, + RestartKind.REPORT_RUN: report_runs, + RestartKind.SIMULATION_RUN: simulation_runs, +} + + +@dataclass(frozen=True) +class UnstartedRun: + """A row still pending with no Modal call, and when it was created.""" + + kind: RestartKind + run_id: UUID + created_at: datetime + + +class StoredRunUnreadableError(LookupError): + """A run's stored rows no longer describe one calculation to start.""" + + +@dataclass(frozen=True) +class RunCall: + """A row's status and the Modal call it records, if any.""" + + status: str + modal_call_id: str | None + + +@dataclass(frozen=True) +class EconomyScope: + """What every simulation of an economy run calculates on, as stored. + + ``catalog_version_id`` is the catalog version the run was created on, and + its region and dataset are named by code and name, as the worker reads + them. + """ + + catalog_version_id: UUID + region_code: str + dataset_name: str + year: int + + +@dataclass(frozen=True) +class StoredEconomySide: + """One simulation run of an economy run, with its simulation's values. + + ``reform`` holds the simulation's policy and dynamic values together, by + parameter name and period, read against the run's catalog version. + """ + + simulation_run_id: UUID + choice: SimulationChoice + reform: PolicyReform + + +@dataclass(frozen=True) +class StoredReportRun: + """A report run as its coordinator's call was built from it. + + ``options`` are the report's (``reports.inputs``), and ``versions`` those + the report run records, which choose the worker that calculates it. + """ + + report_run_id: UUID + country: str + versions: RunVersions + options: dict[str, Any] + scope: EconomyScope + baseline: StoredEconomySide + reform: StoredEconomySide + created_at: datetime + + +@dataclass(frozen=True) +class StoredSimulationRun: + """A standalone simulation run as its simulation function's call was built.""" + + country: str + versions: RunVersions + scope: EconomyScope + run: StoredEconomySide + created_at: datetime + + +def _unstarted(table: Table) -> tuple[ColumnElement[bool], ...]: + """Return the conditions that a row is pending, with no call, never changed.""" + return ( + table.c.status == _PENDING, + table.c.modal_call_id.is_(None), + table.c.updated_at == table.c.created_at, + ) + + +def _economy_simulation_run( + simulation_run_id: ColumnElement[Any], +) -> ColumnElement[bool]: + """Return the condition that the simulation run is an economy simulation's.""" + return exists().where( + simulations.c.id == simulation_runs.c.simulation_id, + simulation_runs.c.id == simulation_run_id, + simulations.c.simulation_type == SimulationType.ECONOMY, + ) + + +def _restart_candidates(kind: RestartKind) -> Select[Any]: + """Return the query for one kind's unstarted rows, before the window.""" + table = RESTART_TABLES[kind] + statement = select(table.c.id, table.c.created_at).where(*_unstarted(table)) + if kind is RestartKind.REPORT_RUN: + return statement.where( + report_runs.c.budget_window_run_id.is_(None), + _economy_simulation_run(report_runs.c.baseline_simulation_run_id), + ) + if kind is RestartKind.SIMULATION_RUN: + naming_report_run = report_runs.alias("naming_report_run") + return statement.where( + exists().where( + simulations.c.id == simulation_runs.c.simulation_id, + simulations.c.simulation_type == SimulationType.ECONOMY, + ), + ~exists().where( + or_( + naming_report_run.c.baseline_simulation_run_id + == simulation_runs.c.id, + naming_report_run.c.reform_simulation_run_id + == simulation_runs.c.id, + ) + ), + ) + return statement + + +def list_unstarted_runs( + connection: Connection, + *, + created_after: datetime, + created_before: datetime, + limit: int, +) -> tuple[UnstartedRun, ...]: + """Return the oldest restartable rows created in the window, at most ``limit``. + + A row is restartable while it is pending with no Modal call recorded and + nothing has changed it since its insert. Report runs and standalone + simulation runs are listed only for economy simulations, and report runs + only outside a budget window. + """ + listed: list[UnstartedRun] = [] + for kind, table in RESTART_TABLES.items(): + statement = ( + _restart_candidates(kind) + .where( + table.c.created_at > created_after, + table.c.created_at <= created_before, + ) + .order_by(table.c.created_at, table.c.id) + .limit(limit) + ) + listed.extend( + UnstartedRun(kind=kind, run_id=row.id, created_at=row.created_at) + for row in connection.execute(statement) + ) + listed.sort(key=lambda run: (run.created_at, run.run_id)) + return tuple(listed[:limit]) + + +def mark_restart(connection: Connection, run: UnstartedRun) -> bool: + """Mark the row as restarted while it is still unstarted; return whether it was. + + A row given a call, claimed, finished or already marked since it was + listed is left as it is, so two checks at once restart it only once. + """ + table = RESTART_TABLES[run.kind] + marked = connection.execute( + update(table) + .where(table.c.id == run.run_id, *_unstarted(table)) + .values(updated_at=func.clock_timestamp()) + .returning(table.c.id) + ).scalar_one_or_none() + return marked is not None + + +def read_run_call(connection: Connection, run: UnstartedRun) -> RunCall | None: + """Return the row's status and recorded call, or None when it is gone.""" + table = RESTART_TABLES[run.kind] + row = connection.execute( + select(table.c.status, table.c.modal_call_id).where(table.c.id == run.run_id) + ).one_or_none() + if row is None: + return None + return RunCall(status=str(row.status), modal_call_id=row.modal_call_id) + + +def read_report_run_simulation_run_ids( + connection: Connection, report_run_id: UUID +) -> tuple[UUID, ...]: + """Return the simulation runs the report run names, baseline first. + + Return no IDs when there is no such run. A report run the check can't + rebuild is failed with them, as its submission would fail it. + """ + row = connection.execute( + select( + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + ).where(report_runs.c.id == report_run_id) + ).one_or_none() + if row is None: + return () + return tuple( + run_id + for run_id in (row.baseline_simulation_run_id, row.reform_simulation_run_id) + if run_id is not None + ) + + +def _side_rows(connection: Connection, simulation_run_id: UUID) -> Row[Any] | None: + """Return a simulation run's scope and its simulation's policy and dynamic.""" + return connection.execute( + select( + simulation_runs.c.id, + simulation_runs.c.created_at, + simulation_runs.c.year, + simulation_runs.c.tax_benefit_model_version_id, + simulation_runs.c.country_package_version, + simulation_runs.c.policyengine_version, + simulations.c.policy_id, + simulations.c.dynamic_id, + regions.c.code.label("region_code"), + datasets.c.name.label("dataset_name"), + tax_benefit_models.c.name.label("model_name"), + ) + .join(simulations, simulations.c.id == simulation_runs.c.simulation_id) + .join( + tax_benefit_model_versions, + tax_benefit_model_versions.c.id + == simulation_runs.c.tax_benefit_model_version_id, + ) + .join( + tax_benefit_models, + tax_benefit_models.c.id == tax_benefit_model_versions.c.model_id, + ) + .outerjoin(regions, regions.c.id == simulation_runs.c.region_id) + .outerjoin(datasets, datasets.c.id == simulation_runs.c.dataset_id) + .where(simulation_runs.c.id == simulation_run_id) + ).one_or_none() + + +def _scope(side: Row[Any]) -> EconomyScope: + """Return what the side calculates on; raise when its region or dataset is gone.""" + if side.region_code is None: + msg = f"simulation run {side.id} names no region" + raise StoredRunUnreadableError(msg) + if side.dataset_name is None: + msg = f"simulation run {side.id} names no dataset" + raise StoredRunUnreadableError(msg) + return EconomyScope( + catalog_version_id=side.tax_benefit_model_version_id, + region_code=side.region_code, + dataset_name=side.dataset_name, + year=side.year, + ) + + +def _side( + connection: Connection, country: str, scope: EconomyScope, side: Row[Any] +) -> StoredEconomySide: + """Read the side's policy and dynamic values against the scope's catalog version. + + Raise the policy reader's errors, or OverlappingReformError. + """ + choice = SimulationChoice(policy_id=side.policy_id, dynamic_id=side.dynamic_id) + return StoredEconomySide( + simulation_run_id=side.id, + choice=choice, + reform=read_simulation_reform( + connection, + country=country, + catalog_version_id=scope.catalog_version_id, + choice=choice, + ), + ) + + +def read_unstarted_report_run( + connection: Connection, report_run_id: UUID +) -> StoredReportRun | None: + """Return what the report run's coordinator call is built from, or None. + + Raise StoredRunUnreadableError when the report run names no reform, a + side's run is missing or names no region or dataset, or its sides differ + in catalog version, region, dataset or year; and the policy reader's + errors when a side's values can't be read. + """ + run = connection.execute( + select( + report_runs.c.created_at, + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + report_runs.c.country_package_version, + report_runs.c.policyengine_version, + reports.c.country, + reports.c.inputs, + ) + .join(reports, reports.c.id == report_runs.c.report_id) + .where(report_runs.c.id == report_run_id) + ).one_or_none() + if run is None: + return None + if run.reform_simulation_run_id is None: + msg = f"report run {report_run_id} names no reform simulation run" + raise StoredRunUnreadableError(msg) + baseline_rows = _side_rows(connection, run.baseline_simulation_run_id) + reform_rows = _side_rows(connection, run.reform_simulation_run_id) + if baseline_rows is None or reform_rows is None: + msg = f"report run {report_run_id} names a missing simulation run" + raise StoredRunUnreadableError(msg) + scope = _scope(baseline_rows) + if _scope(reform_rows) != scope: + msg = ( + f"report run {report_run_id}'s sides calculate on different catalog " + "versions, regions, datasets or years" + ) + raise StoredRunUnreadableError(msg) + return StoredReportRun( + report_run_id=report_run_id, + country=run.country, + versions=RunVersions( + country_package_version=run.country_package_version, + policyengine_version=run.policyengine_version, + ), + options=dict(run.inputs), + scope=scope, + baseline=_side(connection, run.country, scope, baseline_rows), + reform=_side(connection, run.country, scope, reform_rows), + created_at=run.created_at, + ) + + +def read_unstarted_simulation_run( + connection: Connection, simulation_run_id: UUID +) -> StoredSimulationRun | None: + """Return what a standalone run's simulation function call is built from. + + Return None when there is no such run. Raise StoredRunUnreadableError + when it names no region or dataset, and the policy reader's errors when + its values can't be read. + """ + side_rows = _side_rows(connection, simulation_run_id) + if side_rows is None: + return None + country = str(side_rows.model_name).removeprefix(MODEL_NAME_PREFIX) + scope = _scope(side_rows) + return StoredSimulationRun( + country=country, + versions=RunVersions( + country_package_version=side_rows.country_package_version, + policyengine_version=side_rows.policyengine_version, + ), + scope=scope, + run=_side(connection, country, scope, side_rows), + created_at=side_rows.created_at, + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_store.py new file mode 100644 index 000000000..ce1733723 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_store.py @@ -0,0 +1,100 @@ +"""The restart check's store: each read and write in a transaction of its own. + +The restart check lists the runs and household jobs whose Modal call never +started, marks each one it restarts, and reads back what its call is built +from, each in its own transaction, with Modal calls started between them. +``RunRestartStore`` runs ``run_restart_records``' reads and writes that way. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_persistence.engine import bounded_transaction +from policyengine_simulation_persistence.run_restart_records import ( + list_unstarted_runs, + mark_restart, + read_report_run_simulation_run_ids, + read_run_call, + read_unstarted_report_run, + read_unstarted_simulation_run, +) + +if TYPE_CHECKING: + from collections.abc import Callable + from contextlib import AbstractContextManager + from datetime import datetime + from uuid import UUID + + from sqlalchemy import Connection, Engine + + from policyengine_simulation_persistence.run_restart_records import ( + RunCall, + StoredReportRun, + StoredSimulationRun, + UnstartedRun, + ) + +type TransactionFactory = Callable[[], AbstractContextManager[Connection]] + + +class RunRestartStore: + """List, mark and read unstarted runs, each operation in its own transaction.""" + + def __init__(self, transaction: TransactionFactory) -> None: + """Open each transaction with ``transaction``, which commits on success.""" + self._transaction = transaction + + @classmethod + def from_engine( + cls, engine: Engine, *, statement_timeout_milliseconds: int + ) -> RunRestartStore: + """Use ``engine``, bounding each statement to the given time.""" + + def transaction() -> AbstractContextManager[Connection]: + return bounded_transaction( + engine, statement_timeout_milliseconds=statement_timeout_milliseconds + ) + + return cls(transaction) + + def list_unstarted( + self, *, created_after: datetime, created_before: datetime, limit: int + ) -> tuple[UnstartedRun, ...]: + """Return the oldest restartable rows created in the window; see the reader.""" + with self._transaction() as connection: + return list_unstarted_runs( + connection, + created_after=created_after, + created_before=created_before, + limit=limit, + ) + + def mark_restart(self, run: UnstartedRun) -> bool: + """Mark the row restarted while it is still unstarted; return whether it was.""" + with self._transaction() as connection: + return mark_restart(connection, run) + + def read_report_run(self, report_run_id: UUID) -> StoredReportRun | None: + """Return what the report run's call is built from, or None.""" + with self._transaction() as connection: + return read_unstarted_report_run(connection, report_run_id) + + def read_simulation_run( + self, simulation_run_id: UUID + ) -> StoredSimulationRun | None: + """Return what the standalone simulation run's call is built from, or None.""" + with self._transaction() as connection: + return read_unstarted_simulation_run(connection, simulation_run_id) + + def read_call(self, run: UnstartedRun) -> RunCall | None: + """Return the row's status and recorded call, or None when it is gone.""" + with self._transaction() as connection: + return read_run_call(connection, run) + + def read_report_run_simulation_run_ids( + self, report_run_id: UUID + ) -> tuple[UUID, ...]: + """Return the simulation runs the report run names, baseline first.""" + with self._transaction() as connection: + return read_report_run_simulation_run_ids(connection, report_run_id) diff --git a/libs/policyengine-simulation-persistence/tests/conftest.py b/libs/policyengine-simulation-persistence/tests/conftest.py index 328c731b3..da9a52c38 100644 --- a/libs/policyengine-simulation-persistence/tests/conftest.py +++ b/libs/policyengine-simulation-persistence/tests/conftest.py @@ -43,6 +43,9 @@ # The role the Public API's privilege manifest manages for the simulation # service; the deployed SIMULATION_SERVICE_DATABASE_URL logs in as it. SERVICE_ROLE = "policyengine_v2_simulation_service" +# The role the Public API's privilege manifest manages for the v2 Modal +# functions; the Modal secret policyengine-modal-worker-database logs in as it. +MODAL_WORKER_ROLE = "policyengine_v2_modal_worker" @dataclass(frozen=True) @@ -140,6 +143,15 @@ def act_as_the_service(connection: Connection) -> None: connection.execute(text(f"SET LOCAL ROLE {SERVICE_ROLE}")) +def act_as_the_modal_worker(connection: Connection) -> None: + """Take the Modal worker's role for the rest of the connection's transaction. + + The restart check connects as this role, from the Modal secret the v2 + functions and the stuck-run sweep use. + """ + connection.execute(text(f"SET LOCAL ROLE {MODAL_WORKER_ROLE}")) + + @contextmanager def connection_as_the_service(database_url: str) -> Iterator[Connection]: """Yield a connection acting as the service's role, rolled back afterwards.""" diff --git a/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py new file mode 100644 index 000000000..7573c4f69 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py @@ -0,0 +1,600 @@ +"""The restart check's reads and writes for runs whose Modal call never started. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the Modal worker role the restart +check connects as, each in a transaction rolled back afterwards. The rows are +inserted first, as the superuser, with the simulation service's own writes. +Every row a test inserts shares the transaction's start time as its +``created_at``, so each listing names its window around that time. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_modal_worker, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + VERSIONS, + EconomyCatalog, + insert_economy_catalog, + insert_reusable_run, +) +from policyengine_simulation_persistence import ( + NewEconomyReport, + NewEconomySimulation, + NewHouseholdJob, + ReportRunTrigger, + ReportSimulation, + RestartKind, + RunCall, + RunRestartStore, + SimulationChoice, + SimulationType, + StoredRunUnreadableError, + UnstartedRun, + budget_window_runs, + insert_economy_report, + insert_economy_simulation, + insert_household_job, + list_unstarted_runs, + mark_restart, + read_report_run_simulation_run_ids, + read_run_call, + read_unstarted_report_run, + read_unstarted_simulation_run, + record_household_call, + record_report_run_call, + report_runs, + simulation_runs, + simulations, +) +from public_api_rows import ( + StoredValue, + insert_dynamic, + insert_household, + insert_policy, +) +from sqlalchemy import func, insert, select, text, update + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +START = datetime(2026, 1, 1, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +ELASTICITY = "gov.simulation.labor_supply_responses.elasticities.income" +REPORT_OPTIONS = {"include_cliffs": True, "spm": {"geography_kind": "national"}} +HOUSEHOLD_DOCUMENT = { + "people": [ + { + "id": "you", + "values": {"age": 40}, + "memberships": {"household": "your household"}, + } + ], + "household": [{"id": "your household", "values": {}}], +} +LISTING_LIMIT = 100 + + +def _transaction_start(connection: Connection) -> datetime: + """Return the time every row the test inserts records as its creation.""" + return connection.execute(select(func.now())).scalar_one() + + +def _listed( + connection: Connection, *, around: datetime, limit: int = LISTING_LIMIT +) -> tuple[UnstartedRun, ...]: + """List the unstarted runs created within a second of ``around``.""" + return list_unstarted_runs( + connection, + created_after=around - timedelta(seconds=1), + created_before=around + timedelta(seconds=1), + limit=limit, + ) + + +def _new_side(choice: SimulationChoice = CURRENT_LAW) -> ReportSimulation: + return ReportSimulation( + simulation_id=uuid4(), simulation_run_id=uuid4(), choice=choice, reused=False + ) + + +def _new_report( + economy: EconomyCatalog, + *, + baseline: ReportSimulation | None = None, + reform: ReportSimulation | None = None, +) -> NewEconomyReport: + return NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country=economy.catalog.country, + label="Child tax credit", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=REPORT_OPTIONS, + baseline=baseline or _new_side(), + reform=reform or _new_side(), + ) + + +def _new_simulation( + economy: EconomyCatalog, choice: SimulationChoice = CURRENT_LAW +) -> NewEconomySimulation: + return NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=choice, + population=economy.population(), + versions=VERSIONS, + ) + + +def _insert_household_job(connection: Connection, economy: EconomyCatalog) -> UUID: + household_id = insert_household( + connection, country="us", default_year=2026, document=HOUSEHOLD_DOCUMENT + ) + job_id = uuid4() + insert_household_job( + connection, + NewHouseholdJob( + id=job_id, + country="us", + household_id=household_id, + policy_id=None, + year=2026, + request_data={"country": "us", "household_id": str(household_id)}, + tax_benefit_model_version_id=economy.catalog.version_id, + country_package_version=VERSIONS.country_package_version, + policyengine_version=VERSIONS.policyengine_version, + ), + ) + return job_id + + +@pytest.fixture +def connection(public_api_database_url) -> Iterator[Connection]: + """Yield the superuser's connection in a transaction rolled back afterwards.""" + with rolled_back_transaction(public_api_database_url) as superuser_connection: + yield superuser_connection + + +def test_an_unstarted_run_of_each_kind_is_listed_and_a_reports_sides_are_not( + connection, +) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + simulation = _new_simulation(economy) + insert_economy_simulation(connection, simulation) + job_id = _insert_household_job(connection, economy) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + listed = _listed(connection, around=created_at) + + assert {(run.kind, run.run_id) for run in listed} == { + (RestartKind.REPORT_RUN, report.run_id), + (RestartKind.SIMULATION_RUN, simulation.run_id), + (RestartKind.HOUSEHOLD_JOB, job_id), + } + assert {run.created_at for run in listed} == {created_at} + + +def test_a_run_created_outside_the_window_is_not_listed(connection) -> None: + economy = insert_economy_catalog(connection) + insert_economy_simulation(connection, _new_simulation(economy)) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + too_new = list_unstarted_runs( + connection, + created_after=created_at - timedelta(minutes=10), + created_before=created_at - timedelta(minutes=3), + limit=LISTING_LIMIT, + ) + too_old = list_unstarted_runs( + connection, + created_after=created_at + timedelta(minutes=3), + created_before=created_at + timedelta(minutes=10), + limit=LISTING_LIMIT, + ) + + assert too_new == () + assert too_old == () + + +def test_the_listing_takes_the_oldest_runs_up_to_its_limit(connection) -> None: + economy = insert_economy_catalog(connection) + for _ in range(3): + insert_economy_simulation(connection, _new_simulation(economy)) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + assert len(_listed(connection, around=created_at, limit=2)) == 2 + + +def test_a_run_with_a_call_or_no_longer_pending_is_not_listed(connection) -> None: + economy = insert_economy_catalog(connection) + called = _new_report(economy) + insert_economy_report(connection, called) + record_report_run_call( + connection, report_run_id=called.run_id, modal_call_id="fc-original" + ) + claimed = _new_simulation(economy) + insert_economy_simulation(connection, claimed) + connection.execute( + update(simulation_runs) + .where(simulation_runs.c.id == claimed.run_id) + .values(status="running", started_at=func.now()) + ) + job_id = _insert_household_job(connection, economy) + record_household_call(connection, job_id=job_id, modal_call_id="fc-household") + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + listed = _listed(connection, around=created_at) + + # The called report run's sides stay named by a pending report run. + assert listed == () + + +def test_a_marked_run_is_marked_once_and_never_listed_again(connection) -> None: + economy = insert_economy_catalog(connection) + simulation = _new_simulation(economy) + insert_economy_simulation(connection, simulation) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + (unstarted,) = _listed(connection, around=created_at) + + first = mark_restart(connection, unstarted) + second = mark_restart(connection, unstarted) + + assert (first, second) == (True, False) + assert _listed(connection, around=created_at) == () + row = connection.execute( + select(simulation_runs.c.status, simulation_runs.c.modal_call_id).where( + simulation_runs.c.id == simulation.run_id + ) + ).one() + assert (row.status, row.modal_call_id) == ("pending", None) + + +def test_a_run_given_a_call_after_it_was_listed_is_not_marked(connection) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + (unstarted,) = _listed(connection, around=created_at) + record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-late" + ) + + assert mark_restart(connection, unstarted) is False + + +def test_each_kind_is_marked_in_its_own_table(connection) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + _insert_household_job(connection, economy) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + marked = { + run.kind: mark_restart(connection, run) + for run in _listed(connection, around=created_at) + } + + assert marked == {RestartKind.REPORT_RUN: True, RestartKind.HOUSEHOLD_JOB: True} + assert _listed(connection, around=created_at) == () + + +def test_a_budget_window_year_is_not_listed(connection) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + window_run_id = uuid4() + connection.execute( + insert(budget_window_runs).values( + id=window_run_id, + report_id=report.id, + tax_benefit_model_version_id=economy.catalog.version_id, + trigger=ReportRunTrigger.INITIAL, + recalculate=False, + status="pending", + ) + ) + connection.execute( + update(report_runs) + .where(report_runs.c.id == report.run_id) + .values(budget_window_run_id=window_run_id, year=2026) + ) + # The update above changes the year's updated_at; put it back to its + # creation, as a year the service creates with its window stays. + connection.execute( + update(report_runs) + .where(report_runs.c.id == report.run_id) + .values(updated_at=report_runs.c.created_at) + ) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + assert _listed(connection, around=created_at) == () + + +def test_a_household_simulations_run_and_report_run_are_not_listed( + connection, +) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + simulation = _new_simulation(economy) + insert_economy_simulation(connection, simulation) + household_id = insert_household( + connection, country="us", default_year=2026, document=HOUSEHOLD_DOCUMENT + ) + # Make every simulation a household simulation, as C1 ยท 15's household + # reports and household simulations will be. + for simulation_id in ( + report.baseline.simulation_id, + report.reform.simulation_id, + simulation.id, + ): + connection.execute( + update(simulations) + .where(simulations.c.id == simulation_id) + .values( + simulation_type=SimulationType.HOUSEHOLD, + dataset_id=None, + household_id=household_id, + ) + ) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + assert _listed(connection, around=created_at) == () + + +def test_a_report_runs_stored_inputs_are_read_back(connection) -> None: + economy = insert_economy_catalog( + connection, parameter_names=(CTC_AMOUNT, ELASTICITY) + ) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=dynamic_id) + reused = insert_reusable_run(connection, economy) + report = _new_report( + economy, + baseline=ReportSimulation( + simulation_id=reused.simulation_id, + simulation_run_id=reused.simulation_run_id, + choice=CURRENT_LAW, + reused=True, + ), + reform=_new_side(reform_choice), + ) + insert_economy_report(connection, report) + act_as_the_modal_worker(connection) + + stored = read_unstarted_report_run(connection, report.run_id) + + assert stored is not None + assert stored.report_run_id == report.run_id + assert stored.country == "us" + assert stored.versions == VERSIONS + assert stored.options == REPORT_OPTIONS + assert stored.scope.catalog_version_id == economy.catalog.version_id + assert stored.scope.region_code == "us" + assert stored.scope.dataset_name == "populace_us_2024" + assert stored.scope.year == 2026 + assert stored.baseline.simulation_run_id == reused.simulation_run_id + assert stored.baseline.choice == CURRENT_LAW + assert stored.baseline.reform == {} + assert stored.reform.simulation_run_id == report.reform.simulation_run_id + assert stored.reform.choice == reform_choice + assert stored.reform.reform == { + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + ELASTICITY: {"2026-01-01.2200-12-31": -0.05}, + } + assert stored.created_at == _transaction_start(connection) + + +def test_a_standalone_runs_stored_inputs_are_read_back(connection) -> None: + economy = insert_economy_catalog(connection, parameter_names=(CTC_AMOUNT,)) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 2_500, START)] + ) + choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + simulation = _new_simulation(economy, choice) + insert_economy_simulation(connection, simulation) + act_as_the_modal_worker(connection) + + stored = read_unstarted_simulation_run(connection, simulation.run_id) + + assert stored is not None + assert stored.country == "us" + assert stored.versions == VERSIONS + assert stored.scope.catalog_version_id == economy.catalog.version_id + assert (stored.scope.region_code, stored.scope.dataset_name) == ( + "us", + "populace_us_2024", + ) + assert stored.scope.year == 2026 + assert stored.run.simulation_run_id == simulation.run_id + assert stored.run.choice == choice + assert stored.run.reform == {CTC_AMOUNT: {"2026-01-01.2200-12-31": 2_500}} + assert stored.created_at == _transaction_start(connection) + + +def test_a_missing_run_reads_as_none(connection) -> None: + act_as_the_modal_worker(connection) + + assert read_unstarted_report_run(connection, uuid4()) is None + assert read_unstarted_simulation_run(connection, uuid4()) is None + + +def test_a_run_whose_region_is_gone_cannot_be_read(connection) -> None: + economy = insert_economy_catalog(connection) + simulation = _new_simulation(economy) + insert_economy_simulation(connection, simulation) + connection.execute( + update(simulation_runs) + .where(simulation_runs.c.id == simulation.run_id) + .values(region_id=None) + ) + act_as_the_modal_worker(connection) + + with pytest.raises(StoredRunUnreadableError, match="region"): + read_unstarted_simulation_run(connection, simulation.run_id) + + +def test_a_report_run_whose_sides_differ_cannot_be_read(connection) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + connection.execute( + update(simulation_runs) + .where(simulation_runs.c.id == report.reform.simulation_run_id) + .values(year=2027) + ) + act_as_the_modal_worker(connection) + + with pytest.raises(StoredRunUnreadableError, match="sides"): + read_unstarted_report_run(connection, report.run_id) + + +def test_a_report_runs_simulation_runs_are_read_even_when_its_rows_are_not( + connection, +) -> None: + """A report run that can't be rebuilt is failed with its new simulation runs.""" + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + connection.execute( + update(simulation_runs) + .where(simulation_runs.c.id == report.reform.simulation_run_id) + .values(region_id=None) + ) + act_as_the_modal_worker(connection) + + with pytest.raises(StoredRunUnreadableError): + read_unstarted_report_run(connection, report.run_id) + assert read_report_run_simulation_run_ids(connection, report.run_id) == ( + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ) + assert read_report_run_simulation_run_ids(connection, uuid4()) == () + + +def test_a_rows_status_and_recorded_call_are_read_back(connection) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-restarted" + ) + job_id = _insert_household_job(connection, economy) + act_as_the_modal_worker(connection) + + report_run = UnstartedRun( + kind=RestartKind.REPORT_RUN, run_id=report.run_id, created_at=START + ) + job = UnstartedRun(kind=RestartKind.HOUSEHOLD_JOB, run_id=job_id, created_at=START) + missing = UnstartedRun( + kind=RestartKind.SIMULATION_RUN, run_id=uuid4(), created_at=START + ) + + assert read_run_call(connection, report_run) == RunCall( + status="pending", modal_call_id="fc-restarted" + ) + assert read_run_call(connection, job) == RunCall( + status="pending", modal_call_id=None + ) + assert read_run_call(connection, missing) is None + + +@pytest.fixture +def savepoint_store(connection) -> RunRestartStore: + """Return a store whose every transaction is a savepoint of ``connection``.""" + + @contextmanager + def transaction() -> Iterator[Connection]: + with connection.begin_nested(): + yield connection + + return RunRestartStore(transaction) + + +def test_the_store_lists_marks_and_reads_in_transactions_of_its_own( + connection, savepoint_store +) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + simulation = _new_simulation(economy) + insert_economy_simulation(connection, simulation) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + + listed = savepoint_store.list_unstarted( + created_after=created_at - timedelta(seconds=1), + created_before=created_at + timedelta(seconds=1), + limit=LISTING_LIMIT, + ) + by_kind = {run.kind: run for run in listed} + report_run = savepoint_store.read_report_run(report.run_id) + simulation_run = savepoint_store.read_simulation_run(simulation.run_id) + marked = savepoint_store.mark_restart(by_kind[RestartKind.REPORT_RUN]) + call = savepoint_store.read_call(by_kind[RestartKind.SIMULATION_RUN]) + sides = savepoint_store.read_report_run_simulation_run_ids(report.run_id) + + assert set(by_kind) == {RestartKind.REPORT_RUN, RestartKind.SIMULATION_RUN} + assert call == RunCall(status="pending", modal_call_id=None) + assert sides == ( + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ) + assert report_run is not None + assert report_run.report_run_id == report.run_id + assert simulation_run is not None + assert simulation_run.run.simulation_run_id == simulation.run_id + assert marked is True + + +def test_the_restart_writes_need_only_the_modal_worker_roles_privileges( + connection, +) -> None: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + insert_economy_report(connection, report) + created_at = _transaction_start(connection) + act_as_the_modal_worker(connection) + (unstarted,) = _listed(connection, around=created_at) + + assert mark_restart(connection, unstarted) is True + assert ( + record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-restarted" + ) + == "pending" + ) + assert connection.execute(text("SELECT current_user")).scalar_one() == ( + "policyengine_v2_modal_worker" + ) diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index a0b188b5c..7426e83f7 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -100,12 +100,50 @@ can't be read. When Modal refuses the call, or doesn't confirm it within ten seconds, the job is failed at once and the 503 (`modal_call_refused`) names it; a job the household function claimed meanwhile is answered with 202 as it stands, and a job whose state couldn't then be recorded or read with 202 -`pending`, to be read rather than submitted again (the stuck-run sweep fails -it if no call ever starts). Every 500 of either route, from an unexpected +`pending`, to be read rather than submitted again (the restart check starts +its call once more 3 to 10 minutes after its creation, and the stuck-run +sweep fails it if no call ever starts). Every 500 of either route, from an unexpected failure or a request the service built and the household contract refuses, is the JSON error with code `internal_error`. `DEPLOYMENT.md` describes the deployment settings and the staging check. +## Restarting a run whose call never started + +When the service stops between creating a household job, a report run or a +standalone simulation run and starting its Modal call, nothing records a +call for it. The restart check (`run_restarts`, with `run_restart_calls` and +`run_restart_check`) starts such a row's call once more. It doesn't run in +this service: the stuck-run sweep's Modal application +(`projects/policyengine-simulation-executor/src/modal/run_sweep_app.py`) +runs it every two minutes as the Modal worker role. + +- It takes the rows created 3 to 10 minutes ago that are still `pending` + with no call and that nothing has changed since their insert, at most 20 + a check, oldest first. A budget window's report runs and household + simulations' runs are left out. After three minutes a check marks no more + rows and leaves the rest for the next. +- It rebuilds each call with this service's own builders + (`report_run_calculation`, `simulation_run_calculation`, + `prepare_household_job`) from the stored rows, on the manifest's worker + for the .py version the row records, even when that worker is no longer + the default. The context's request ID is the check's own + (`run-restart-`), and its creation time is the row's. +- It sets the row's `updated_at` before starting the call, so a row is + restarted at most once, then records the call's ID with the submission's + own conditional write. The restart counts as done when the row then + records its call; a row recording another call was claimed by that call + first. +- A row whose call can't be rebuilt, or whose restart Modal refuses or + doesn't confirm within ten seconds, is failed with + `run_not_started: The run stayed pending with no Modal call recorded for + it, and โ€ฆ`; a report run is failed with its simulation runs still pending + with no call. A row is left unmarked for the next check when its rows or + manifest can't be read now, its rebuild raises an unexpected error + (logged with its traceback), its mark can't be written, or its call can't + be rebuilt and its failure can't be written. +- Each check, each row and each outcome are telemetry; see + `run_restart_check`'s docstring for the operations and event. + ## Economy reports and simulations `POST /reports` creates an economy report: a baseline and a reform diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/__init__.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/__init__.py index 722f82b54..b34ddc052 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/__init__.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/__init__.py @@ -1,5 +1,24 @@ -"""PolicyEngine Cloud Run Simulation Entrypoint.""" +"""PolicyEngine Cloud Run Simulation Entrypoint. -from policyengine_simulation_entry.app import create_app +``create_app`` is loaded on first use, so importing one of the package's +modules, such as the restart check that runs in the stuck-run sweep's Modal +image, never imports the web service and the packages only it needs. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from policyengine_simulation_entry.app import create_app __all__ = ["create_app"] + + +def __getattr__(name: str) -> Any: # noqa: ANN401 # the attribute the app module defines + """Load ``create_app`` from the app module when it is first asked for.""" + if name == "create_app": + return import_module("policyengine_simulation_entry.app").create_app + msg = f"module {__name__!r} has no attribute {name!r}" + raise AttributeError(msg) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py index b07a6e693..8e123beb4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py @@ -75,6 +75,43 @@ def default_economy_worker(manifest: V2VersionManifest, country: str) -> Economy country's economy runs. """ version = manifest.default_version + return _economy_worker( + manifest, + country, + version, + described_as=f"the default v2 worker, .py {version},", + ) + + +def economy_worker_for_version( + manifest: V2VersionManifest, country: str, policyengine_version: str +) -> EconomyWorker: + """Return the economy functions of the manifest's worker for a .py version. + + A run is calculated by the worker of the .py version it records, which + may no longer be the default. Raise EconomyWorkerUnavailableError when the + manifest no longer holds that worker, or it cannot calculate the + country's economy runs. + """ + if policyengine_version not in manifest.versions: + msg = f"the v2 version manifest has no worker for .py {policyengine_version}" + raise EconomyWorkerUnavailableError(msg) + return _economy_worker( + manifest, + country, + policyengine_version, + described_as=f"the v2 worker for .py {policyengine_version}", + ) + + +def _economy_worker( + manifest: V2VersionManifest, country: str, version: str, *, described_as: str +) -> EconomyWorker: + """Return the economy functions of the manifest's worker ``version``. + + ``described_as`` names the worker in the error raised when it cannot + calculate the country's economy runs. + """ worker = manifest.versions[version] country_worker = worker.economy_worker(country) country_bundle = next( @@ -85,10 +122,7 @@ def default_economy_worker(manifest: V2VersionManifest, country: str) -> Economy or country_worker.single_simulation_callable is None or country_bundle is None ): - msg = ( - f"the default v2 worker, .py {version}, cannot calculate economy " - f"runs for {country}" - ) + msg = f"{described_as} cannot calculate economy runs for {country}" raise EconomyWorkerUnavailableError(msg) return EconomyWorker( version=version, @@ -139,13 +173,27 @@ def default_worker(self, country: str) -> EconomyWorker: Raise EconomyWorkerUnavailableError when it has none, and VersionManifestUnavailableError when the manifest cannot be read. """ + return default_economy_worker(self._manifest(), country) + + def worker_for_version( + self, country: str, policyengine_version: str + ) -> EconomyWorker: + """Return the economy functions of the worker for a .py version. + + Raise EconomyWorkerUnavailableError when the manifest has none, and + VersionManifestUnavailableError when it cannot be read. + """ + return economy_worker_for_version( + self._manifest(), country, policyengine_version + ) + + def _manifest(self) -> V2VersionManifest: try: - manifest = self._loader.load() + return self._loader.load() # Modal raises its own errors, and any of them means the same here. except Exception as error: msg = "the v2 version manifest could not be read" raise VersionManifestUnavailableError(msg) from error - return default_economy_worker(manifest, country) class EconomyRunCaller(Protocol): diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py index bf08e7668..dd8a0646e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py @@ -80,13 +80,46 @@ def default_household_worker( Raise HouseholdFunctionUnavailableError when it has none. """ version = manifest.default_version + return _household_worker( + manifest, + country, + version, + described_as=f"the default v2 worker, .py {version},", + ) + + +def household_worker_for_version( + manifest: V2VersionManifest, country: str, policyengine_version: str +) -> HouseholdWorker: + """Return the household function of the manifest's worker for a .py version. + + A household job is calculated by the worker of the .py version it + records, which may no longer be the default. Raise + HouseholdFunctionUnavailableError when the manifest no longer holds that + worker, or it has no household function for the country. + """ + if policyengine_version not in manifest.versions: + msg = f"the v2 version manifest has no worker for .py {policyengine_version}" + raise HouseholdFunctionUnavailableError(msg) + return _household_worker( + manifest, + country, + policyengine_version, + described_as=f"the v2 worker for .py {policyengine_version}", + ) + + +def _household_worker( + manifest: V2VersionManifest, country: str, version: str, *, described_as: str +) -> HouseholdWorker: + """Return the household function of the manifest's worker ``version``. + + ``described_as`` names the worker in the error raised when it has none. + """ worker = manifest.versions[version] country_worker = worker.household_worker(country) if country_worker is None or country_worker.household_callable is None: - msg = ( - f"the default v2 worker, .py {version}, has no household function " - f"for {country}" - ) + msg = f"{described_as} has no household function for {country}" raise HouseholdFunctionUnavailableError(msg) return HouseholdWorker( application_name=worker.application_name, @@ -128,13 +161,27 @@ def default_worker(self, country: str) -> HouseholdWorker: Raise HouseholdFunctionUnavailableError when it has none, and VersionManifestUnavailableError when the manifest cannot be read. """ + return default_household_worker(self._manifest(), country) + + def worker_for_version( + self, country: str, policyengine_version: str + ) -> HouseholdWorker: + """Return the household function of the worker for a .py version. + + Raise HouseholdFunctionUnavailableError when the manifest has none, + and VersionManifestUnavailableError when it cannot be read. + """ + return household_worker_for_version( + self._manifest(), country, policyengine_version + ) + + def _manifest(self) -> V2VersionManifest: try: - manifest = self._loader.load() + return self._loader.load() # Modal raises its own errors, and any of them means the same here. except Exception as error: msg = "the v2 version manifest could not be read" raise VersionManifestUnavailableError(msg) from error - return default_household_worker(manifest, country) class ModalHouseholdFunctionCaller: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_calls.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_calls.py new file mode 100644 index 000000000..40d089da9 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_calls.py @@ -0,0 +1,521 @@ +"""Rebuild the Modal call of a run or household job whose call never started. + +The restart check (``run_restarts``) starts such a row's call once more. The +call must be the one the row's submission started: ``RestartCallBuilder`` +reads back what the submission built it from and builds it again with the +submission's own builders (``report_run_calculation``, +``simulation_run_calculation`` and ``prepare_household_job``), on the worker +of the .py version the row records. A household job whose rebuilt job would +record anything else (another year, catalog version or policy) can't be +rebuilt, so the restart never calculates something its row doesn't record. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, Protocol, get_args + +from policyengine_simulation_contract.stage12_bundle import CountryId +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_simulation_persistence import ( + MissingPolicyParametersError, + OverlappingReformError, + RestartKind, + StoredRunUnreadableError, + UnknownCatalogVersionError, + UnknownDynamicError, + UnknownHouseholdError, + UnknownPolicyError, +) +from pydantic import ValidationError +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_entry.economy_calculation_inputs import ( + CalculationScope, + CallOrigin, + DatasetNotInBundleError, + bundle_dataset, + report_run_calculation, + simulation_options, + simulation_run_calculation, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, +) +from policyengine_simulation_entry.household_api_models import ( + HouseholdCalculationRequest, +) +from policyengine_simulation_entry.household_calculations import ( + HouseholdCalculationInternalError, + HouseholdCalculationUnprocessableError, + prepare_household_job, +) +from policyengine_simulation_entry.household_workers import ( + HouseholdFunctionUnavailableError, + VersionManifestUnavailableError, +) + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_persistence import ( + EconomyScope, + HouseholdCalculationInputs, + HouseholdJobRecord, + HouseholdJobStatus, + NewHouseholdJob, + RunCall, + RunVersions, + StoredReportRun, + StoredSimulationRun, + UnstartedRun, + ) + + from policyengine_simulation_entry.economy_workers import ( + EconomyRunCaller, + EconomyWorker, + ) + from policyengine_simulation_entry.household_workers import ( + HouseholdFunctionCaller, + HouseholdWorker, + ) + +# The deployment each Modal environment holding the v2 workers belongs to, as +# the Simulation Entrypoint's deploy pairs them (APP_ENVIRONMENT). A call's +# context and artifact prefix name the deployment. +DEPLOYMENT_ENVIRONMENTS: Final = MappingProxyType( + {"staging": "staging", "main": "production"} +) + + +@dataclass(frozen=True) +class RestartSettings: + """The deployment the check restarts runs for. + + ``environment`` labels each call's context and artifacts, and + ``modal_environment`` holds the v2 workers. + """ + + environment: str + modal_environment: str + + @classmethod + def for_modal_environment(cls, modal_environment: str) -> RestartSettings: + """Return the settings of the deployment whose workers are in the environment. + + Raise ValueError for an environment no deployment uses. + """ + environment = DEPLOYMENT_ENVIRONMENTS.get(modal_environment) + if environment is None: + msg = ( + "no deployment runs the v2 workers in Modal environment " + f"{modal_environment!r}" + ) + raise ValueError(msg) + return cls(environment=environment, modal_environment=modal_environment) + + +class UnstartedRunSource(Protocol): + """The restart check's own reads and writes (``RunRestartStore``).""" + + def list_unstarted( + self, *, created_after: datetime, created_before: datetime, limit: int + ) -> tuple[UnstartedRun, ...]: + """Return the oldest restartable rows created in the window.""" + ... + + def mark_restart(self, run: UnstartedRun) -> bool: + """Mark the row restarted while it is still unstarted.""" + ... + + def read_report_run(self, report_run_id: UUID) -> StoredReportRun | None: + """Return what the report run's call is built from.""" + ... + + def read_simulation_run( + self, simulation_run_id: UUID + ) -> StoredSimulationRun | None: + """Return what the standalone simulation run's call is built from.""" + ... + + def read_report_run_simulation_run_ids( + self, report_run_id: UUID + ) -> tuple[UUID, ...]: + """Return the simulation runs the report run names, baseline first.""" + ... + + def read_call(self, run: UnstartedRun) -> RunCall | None: + """Return the row's status and recorded call, or None when it is gone.""" + ... + + +class HouseholdJobRestarts(Protocol): + """The household job reads and writes a restart uses (``HouseholdJobStore``).""" + + def read(self, job_id: UUID) -> HouseholdJobRecord | None: + """Return the job as it stands.""" + ... + + def read_inputs( + self, + *, + country: str, + household_id: UUID, + policy_id: UUID | None, + policyengine_version: str, + ) -> HouseholdCalculationInputs: + """Read the catalog version, household and policy together.""" + ... + + def record_call( + self, job_id: UUID, modal_call_id: str + ) -> HouseholdJobStatus | None: + """Record the call ID while the job is pending with none.""" + ... + + def fail_unstarted(self, job_id: UUID, error_message: str) -> bool: + """Fail the job while it is pending with no call ID.""" + ... + + +class VersionedEconomyWorkers(Protocol): + """Find a .py version's economy functions (``ManifestEconomyWorkers``).""" + + def worker_for_version( + self, country: str, policyengine_version: str + ) -> EconomyWorker: + """Return the economy functions of the worker for the .py version.""" + ... + + +class VersionedHouseholdWorkers(Protocol): + """Find a .py version's household function (``ManifestHouseholdWorkers``).""" + + def worker_for_version( + self, country: str, policyengine_version: str + ) -> HouseholdWorker: + """Return the household function of the worker for the .py version.""" + ... + + +@dataclass(frozen=True) +class RestartWorkers: + """Where a restart finds the worker of a row's .py version.""" + + economy: VersionedEconomyWorkers + households: VersionedHouseholdWorkers + + +@dataclass(frozen=True) +class RestartCallers: + """What starts a restarted call on Modal.""" + + economy: EconomyRunCaller + households: HouseholdFunctionCaller + + +type CallStart = Callable[[dict[str, Any] | None], Awaitable[str]] + + +class DeferredRestartError(Exception): + """The row's rows or the manifest couldn't be read now; the message is its type.""" + + +class CallNotRebuiltError(Exception): + """The row's call can't be built again; ``reason`` says why, as a clause. + + ``simulation_run_ids`` are a report run's simulation runs, failed with it. + """ + + def __init__( + self, + reason: str, + *, + error_type: str, + simulation_run_ids: tuple[UUID, ...] = (), + ) -> None: + """Keep why, the error type that said so, and the report's simulation runs.""" + super().__init__(reason) + self.reason = reason + self.error_type = error_type + self.simulation_run_ids = simulation_run_ids + + +class _ChangedRunError(LookupError): + """What the call would be built from now differs from what the row records.""" + + +# The errors that mean a row's call can't be built again, and why, in words a +# failed run's error message can hold. A worker's absence is described by the +# error itself, which names only the .py version and the country. +_REBUILD_REASONS: Final[dict[type[Exception], str | None]] = { + EconomyWorkerUnavailableError: None, + HouseholdFunctionUnavailableError: None, + _ChangedRunError: None, + StoredRunUnreadableError: "its stored rows no longer describe one calculation", + UnknownCatalogVersionError: "its catalog version is gone", + UnknownHouseholdError: "its household is gone", + UnknownPolicyError: "its policy is gone", + UnknownDynamicError: "its dynamic is gone", + MissingPolicyParametersError: "its policy names parameters its catalog lacks", + OverlappingReformError: "its policy and dynamic set the same parameters", + DatasetNotInBundleError: "its worker's bundle has no dataset it names", + HouseholdCalculationUnprocessableError: "its household can no longer be read", + HouseholdCalculationInternalError: "its household request no longer validates", + ValidationError: "its stored request no longer validates", +} +_REBUILD_ERRORS: Final = tuple(_REBUILD_REASONS) + + +def _rebuild_reason(error: Exception) -> str: + reason = next( + ( + reason + for error_class, reason in _REBUILD_REASONS.items() + if isinstance(error, error_class) + ), + None, + ) + return str(error) if reason is None else reason + + +def _same_job(rebuilt: NewHouseholdJob, stored: HouseholdJobRecord) -> bool: + """Return whether the rebuilt job records what the stored job records.""" + return ( + rebuilt.id, + rebuilt.country, + rebuilt.household_id, + rebuilt.policy_id, + rebuilt.year, + rebuilt.request_data, + rebuilt.tax_benefit_model_version_id, + rebuilt.country_package_version, + rebuilt.policyengine_version, + ) == ( + stored.id, + stored.country, + stored.household_id, + stored.policy_id, + stored.year, + stored.request_data, + stored.tax_benefit_model_version_id, + stored.country_package_version, + stored.policyengine_version, + ) + + +def _country(country: str) -> CountryId: + """Return the stored country as the contract names it.""" + for known in get_args(CountryId): + if country == known: + return known + msg = f"no v2 worker calculates economy runs for country {country}" + raise StoredRunUnreadableError(msg) + + +@dataclass(frozen=True) +class RebuiltCall: + """What starts a row's call again, and a report run's simulation runs. + + ``start`` takes the observability context the call continues from. + """ + + start: CallStart + simulation_run_ids: tuple[UUID, ...] = () + + +class RestartCallBuilder: + """Build a row's call again from its stored rows.""" + + def __init__( # noqa: PLR0913 # each seam tests replace + self, + *, + restarts: UnstartedRunSource, + households: HouseholdJobRestarts, + workers: RestartWorkers, + callers: RestartCallers, + settings: RestartSettings, + new_request_id: Callable[[], str] = lambda: f"run-restart-{uuid.uuid4()}", + ) -> None: + """Read rows from the two stores; start calls with ``callers``.""" + self._restarts = restarts + self._households = households + self._workers = workers + self._callers = callers + self._settings = settings + self._new_request_id = new_request_id + + def build(self, run: UnstartedRun) -> RebuiltCall | None: + """Return what starts the row's call again, or None when the row is gone. + + Raise DeferredRestartError when its rows or the manifest can't be read + now, and CallNotRebuiltError when its call can't be built again. + """ + simulation_run_ids: tuple[UUID, ...] = () + try: + if run.kind is RestartKind.REPORT_RUN: + stored_report = self._restarts.read_report_run(run.run_id) + if stored_report is None: + return None + simulation_run_ids = ( + stored_report.baseline.simulation_run_id, + stored_report.reform.simulation_run_id, + ) + return RebuiltCall(self._report_call(stored_report), simulation_run_ids) + if run.kind is RestartKind.SIMULATION_RUN: + stored_simulation = self._restarts.read_simulation_run(run.run_id) + if stored_simulation is None: + return None + return RebuiltCall(self._simulation_call(stored_simulation)) + job = self._households.read(run.run_id) + if job is None: + return None + return RebuiltCall(self._household_call(job)) + except (SQLAlchemyError, VersionManifestUnavailableError) as error: + raise DeferredRestartError(type(error).__name__) from error + except _REBUILD_ERRORS as error: + if run.kind is RestartKind.REPORT_RUN and not simulation_run_ids: + simulation_run_ids = self._report_simulation_run_ids(run.run_id) + raise CallNotRebuiltError( + _rebuild_reason(error), + error_type=type(error).__name__, + simulation_run_ids=simulation_run_ids, + ) from error + + def _report_simulation_run_ids(self, report_run_id: UUID) -> tuple[UUID, ...]: + """Return the simulation runs a report run its rows can't rebuild names. + + It is failed with them, as its submission fails it with the simulation + runs created for it; only those still pending with no call are failed. + When even they can't be read, it is failed alone, and the sweep fails + them later. + """ + try: + return self._restarts.read_report_run_simulation_run_ids(report_run_id) + except SQLAlchemyError: + return () + + def _economy_worker(self, country: str, versions: RunVersions) -> EconomyWorker: + """Return the worker of the run's .py version, which must match its versions.""" + worker = self._workers.economy.worker_for_version( + country, versions.policyengine_version + ) + if worker.versions != versions: + msg = ( + f"the v2 worker for .py {versions.policyengine_version} calculates " + "on other versions than the run records" + ) + raise _ChangedRunError(msg) + return worker + + def _economy_scope( + self, + worker: EconomyWorker, + country: str, + scope: EconomyScope, + options: dict[str, Any], + ) -> CalculationScope: + return CalculationScope( + country=_country(country), + region_code=scope.region_code, + dataset=bundle_dataset(worker, scope.dataset_name), + year=scope.year, + options=simulation_options( + include_cliffs=bool(options.get("include_cliffs", False)), + spm=options.get("spm"), + ), + ) + + def _origin(self, created_at: datetime) -> CallOrigin: + """Return where the restarted call comes from, dated as its row was created. + + The row's creation time places the call's artifacts where its first + call's would have been. + """ + return CallOrigin( + environment=self._settings.environment, + modal_environment=self._settings.modal_environment, + request_id=self._new_request_id(), + created_at=created_at, + ) + + def _report_call(self, stored: StoredReportRun) -> CallStart: + worker = self._economy_worker(stored.country, stored.versions) + calculation = report_run_calculation( + worker, + self._economy_scope(worker, stored.country, stored.scope, stored.options), + self._origin(stored.created_at), + run_ids=ReportRunIds( + report_run_id=stored.report_run_id, + baseline_simulation_run_id=stored.baseline.simulation_run_id, + reform_simulation_run_id=stored.reform.simulation_run_id, + ), + baseline_policy=stored.baseline.reform, + reform_policy=stored.reform.reform, + ) + caller = self._callers.economy + + def start(observability_context: dict[str, Any] | None) -> Awaitable[str]: + return caller.start_report_run( + worker, + calculation.payload, + calculation.context, + observability_context=observability_context, + ) + + return start + + def _simulation_call(self, stored: StoredSimulationRun) -> CallStart: + worker = self._economy_worker(stored.country, stored.versions) + calculation = simulation_run_calculation( + worker, + # A standalone simulation takes no report options. + self._economy_scope(worker, stored.country, stored.scope, {}), + self._origin(stored.created_at), + simulation_run_id=stored.run.simulation_run_id, + policy=stored.run.reform, + ) + caller = self._callers.economy + + def start(observability_context: dict[str, Any] | None) -> Awaitable[str]: + return caller.start_simulation_run( + worker, + calculation.payload, + calculation.context, + observability_context=observability_context, + ) + + return start + + def _household_call(self, job: HouseholdJobRecord) -> CallStart: + request = HouseholdCalculationRequest.model_validate(job.request_data) + worker = self._workers.households.worker_for_version( + job.country, job.policyengine_version + ) + inputs = self._households.read_inputs( + country=request.country, + household_id=request.household_id, + policy_id=request.policy_id, + policyengine_version=job.policyengine_version, + ) + prepared = prepare_household_job(request, worker, inputs, job_id=job.id) + if not _same_job(prepared.job, job): + msg = ( + "its household, policy or catalog version changed since it was created" + ) + raise _ChangedRunError(msg) + caller = self._callers.households + + def start(observability_context: dict[str, Any] | None) -> Awaitable[str]: + return caller.spawn( + worker, + prepared.payload, + household_job_id=job.id, + observability_context=observability_context, + ) + + return start diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_check.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_check.py new file mode 100644 index 000000000..f4c254d42 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_check.py @@ -0,0 +1,249 @@ +"""The scheduled restart check: its Modal wiring and its telemetry. + +The stuck-run sweep's Modal application runs ``run_restart_check`` every two +minutes, as the Modal worker role, from the sweep's Modal secret. It builds +the restart service on the run tables (``scheduled_restart_service``) and +runs one check (``run_scheduled_check``) in its own telemetry operation. + +How often runs need a restart is the check's telemetry: + +- each check is one ``v2_run_restart_check`` operation, with the number of + rows it listed and how each ended (``restart_listed``, ``restart_failed`` + and the rest) set on it; +- each row it takes up is one ``v2_run_restart`` operation, with the row's + kind, ID and age. A row left ``deferred`` is taken up again by the next + check, so the operation count metric counts restart attempts, and the + distinct ``run_id`` values are the runs whose call never started that a + check took up. A row no check took up (more rows at once than the checks + in its window had time for) is failed by the sweep with + ``run_not_started`` instead; a count of the outcomes other than + ``deferred`` also leaves out the rows deferred in every check. +- each row's outcome is set on its operation (``restart_outcome``, with + ``restart_error_type`` when an error ended it) and logged as a + ``v2_run_restart_outcome`` event, at WARNING when the row could not be + started again or its rows could not be read; an unexpected error + rebuilding a row's call is also logged with its traceback + (``v2_run_restart_unexpected_error``). + +This module imports none of the web service, so the sweep's image, which +has none of its packages, can run it. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Final, Protocol + +from policyengine_simulation_contract.stage12_manifest import V2_VERSION_MANIFEST_NAME +from policyengine_simulation_observability.stages import ( + V2_RUN_RESTART_STAGES, + Stage, +) +from policyengine_simulation_persistence import ( + EconomyRunStore, + HouseholdJobStore, + RunRestartStore, + create_simulation_engine, +) + +from policyengine_simulation_entry.economy_workers import ( + ManifestEconomyWorkers, + ModalEconomyRunCaller, +) +from policyengine_simulation_entry.household_workers import ( + ManifestHouseholdWorkers, + ModalHouseholdFunctionCaller, +) +from policyengine_simulation_entry.run_restarts import ( + RestartCallers, + RestartOutcome, + RestartSettings, + RestartStores, + RestartWorkers, + RunRestartService, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + from contextlib import AbstractContextManager + + from policyengine_simulation_persistence import UnstartedRun + from sqlalchemy import Engine + + from policyengine_simulation_entry.run_restarts import RestartSummary + +# The longest any one statement of a check may take. +RESTART_STATEMENT_TIMEOUT_MILLISECONDS: Final = 5_000 +OUTCOME_EVENT: Final = "v2_run_restart_outcome" +# The outcomes worth a warning: the row could not be started again, or the +# check could not tell, so a person may need to look. +WARNING_OUTCOMES: Final = frozenset( + { + RestartOutcome.FAILED, + RestartOutcome.DEFERRED, + RestartOutcome.CALL_NOT_RECORDED, + } +) + + +class RestartRuntime(Protocol): + """What the check's telemetry uses of ``ObservabilityRuntime``.""" + + def operation( + self, name: str, *, attributes: Mapping[str, Any] | None = None + ) -> AbstractContextManager[Any]: + """Return the scope of one operation.""" + ... + + def set_context(self, **attributes: Any) -> None: # noqa: ANN401 # the runtime's own signature + """Set attributes on the innermost open operation.""" + ... + + def event( + self, + name: str, + *, + severity: str = "INFO", + attributes: Mapping[str, Any] | None = None, + ) -> None: + """Log one structured event.""" + ... + + def capture_context(self) -> dict[str, str]: + """Return the context a call started now continues from.""" + ... + + +class ObservabilityRestartTelemetry: + """Report each listed row as an operation, and its outcome as an event.""" + + def __init__(self, runtime: RestartRuntime) -> None: + """Report through ``runtime``, which never raises.""" + self._runtime = runtime + + def restart_scope( + self, run: UnstartedRun, *, age_seconds: float + ) -> AbstractContextManager[Any]: + """Return the operation one row's restart runs in.""" + return self._runtime.operation( + V2_RUN_RESTART_STAGES.name(Stage.V2_RUN_RESTART), + attributes={ + "run_kind": run.kind.value, + "run_id": str(run.run_id), + "run_age_seconds": age_seconds, + }, + ) + + def record_outcome( + self, + run: UnstartedRun, + outcome: RestartOutcome, + *, + error_type: str | None = None, + ) -> None: + """Set the outcome on the row's operation and log it as an event.""" + outcome_attributes: dict[str, str] = {"restart_outcome": outcome.value} + if error_type is not None: + outcome_attributes["restart_error_type"] = error_type + self._runtime.set_context(**outcome_attributes) + self._runtime.event( + OUTCOME_EVENT, + severity="WARNING" if outcome in WARNING_OUTCOMES else "INFO", + attributes={ + "run_kind": run.kind.value, + "run_id": str(run.run_id), + **outcome_attributes, + }, + ) + + def capture_context(self) -> dict[str, str] | None: + """Return the context a restarted call continues the check's trace from.""" + return self._runtime.capture_context() + + +class RestartCheck(Protocol): + """One restart check (``RunRestartService``).""" + + async def restart_unstarted_runs(self) -> RestartSummary: + """Restart the rows whose call never started; count how each ended.""" + ... + + +def run_scheduled_check( + service: RestartCheck, runtime: RestartRuntime +) -> dict[str, int]: + """Run one check in its own operation; return how many rows ended each way.""" + with runtime.operation(V2_RUN_RESTART_STAGES.name(Stage.V2_RUN_RESTART_CHECK)): + summary = asyncio.run(service.restart_unstarted_runs()) + counts = asdict(summary) + runtime.set_context( + **{f"restart_{name}": count for name, count in counts.items()} + ) + return counts + + +@dataclass(frozen=True) +class ScheduledRestart: + """The restart service on the run tables, and the engine it connects through.""" + + service: RunRestartService + settings: RestartSettings + engine: Engine + + +def scheduled_restart_service( + *, database_url: str, modal_environment: str, runtime: RestartRuntime +) -> ScheduledRestart: + """Build the restart service for the Modal environment's deployment. + + ``database_url`` is the Modal worker role's ``postgresql://`` URL. Nothing + connects to the database or reads Modal until the check runs. Raise + ValueError for a Modal environment no deployment uses, or a malformed URL. + """ + settings = RestartSettings.for_modal_environment(modal_environment) + engine = create_simulation_engine(database_url) + timeout = RESTART_STATEMENT_TIMEOUT_MILLISECONDS + service = RunRestartService( + stores=RestartStores( + restarts=RunRestartStore.from_engine( + engine, statement_timeout_milliseconds=timeout + ), + economy=EconomyRunStore.from_engine( + engine, statement_timeout_milliseconds=timeout + ), + households=HouseholdJobStore.from_engine( + engine, statement_timeout_milliseconds=timeout + ), + ), + workers=RestartWorkers( + economy=ManifestEconomyWorkers.from_modal( + manifest_name=V2_VERSION_MANIFEST_NAME, environment=modal_environment + ), + households=ManifestHouseholdWorkers.from_modal( + manifest_name=V2_VERSION_MANIFEST_NAME, environment=modal_environment + ), + ), + callers=RestartCallers( + economy=ModalEconomyRunCaller(modal_environment), + households=ModalHouseholdFunctionCaller(modal_environment), + ), + settings=settings, + telemetry=ObservabilityRestartTelemetry(runtime), + ) + return ScheduledRestart(service=service, settings=settings, engine=engine) + + +def run_restart_check( + *, database_url: str, modal_environment: str, runtime: RestartRuntime +) -> dict[str, int]: + """Run one restart check on the run tables; return how each row ended.""" + scheduled = scheduled_restart_service( + database_url=database_url, + modal_environment=modal_environment, + runtime=runtime, + ) + try: + return run_scheduled_check(scheduled.service, runtime) + finally: + scheduled.engine.dispose() diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py new file mode 100644 index 000000000..5c7afb78f --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -0,0 +1,446 @@ +"""Start a run or household job again when its Modal call never started. + +The simulation service creates a household job, a report run or a standalone +simulation run ``pending`` with no Modal call ID, starts the call that +calculates it, then records the call's ID. When the service stops between +the insert and the start (a Cloud Run instance shut down, for example), +nothing records a call, and the stuck-run sweep would fail the run 15 +minutes after its creation. The restart check, a scheduled Modal function +beside the sweep, starts such a row's call once more: + +1. it lists the rows created 3 to 10 minutes ago that are still pending with + no call and that nothing has changed since their insert; +2. for each, it reads back what the service built the call from and builds + the same call, on the worker of the .py version the row records; +3. it marks the row, so it is never restarted twice, starts the call, and + records the call's ID while the row is still pending with none. + +A row whose call can't be built again (its worker left the version manifest, +or its stored rows changed), or whose new call Modal refuses, is failed at +once with ``run_not_started``, as the sweep would fail it later. A row whose +rows or manifest can't be read now, or whose rebuild raises an unexpected +error, is left unmarked, for the next check, as is one the check has no +time left for, whose mark can't be written, or whose call can't be rebuilt +and whose failure can't be written. A marked row whose failure can't be +written is left to the sweep. A +call the service started after all claims the row only while the row records +no call or that call's own ID, so once the restart records its call, the +earlier call's claim is refused; if the earlier call claimed it first, the +restart's call is refused instead. + +Each listed row is one telemetry scope, ``restart_scope``, with its age, and +its outcome is reported through ``record_outcome``, so how often runs need a +restart, and how each restart ended, can be counted. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +import uuid +from collections import Counter +from contextlib import nullcontext +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from enum import StrEnum +from typing import TYPE_CHECKING, Final, Protocol + +from policyengine_simulation_persistence import RestartKind +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_entry.run_restart_calls import ( + CallNotRebuiltError, + DeferredRestartError, + HouseholdJobRestarts, + RestartCallBuilder, + RestartCallers, + RestartSettings, + RestartWorkers, + UnstartedRunSource, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from contextlib import AbstractContextManager + from uuid import UUID + + from policyengine_simulation_persistence import ( + ReportRunStatus, + SimulationStatus, + UnstartedRun, + ) + + from policyengine_simulation_entry.run_restart_calls import CallStart + +__all__ = [ + "CALL_START_TIMEOUT_SECONDS", + "CHECK_TIME_LIMIT", + "RESTARTS_PER_CHECK", + "RESTART_AFTER", + "RESTART_UNTIL", + "EconomyRunWrites", + "NoRestartTelemetry", + "RestartCallers", + "RestartOutcome", + "RestartSettings", + "RestartStores", + "RestartSummary", + "RestartTelemetry", + "RestartWorkers", + "RunRestartService", +] + +logger = logging.getLogger(__name__) + +# A row is restarted once it has stayed pending with no call this long, which +# is far longer than a submission takes to start and record its call. +RESTART_AFTER: Final = timedelta(minutes=3) +# A row older than this is left to the stuck-run sweep, which fails it at 15 +# minutes. A row listed just under this age may be marked as late as the +# check's time limit allows, and its call still starts and is recorded +# minutes before the sweep judges it. +RESTART_UNTIL: Final = timedelta(minutes=10) +# At most this many rows are listed by one check; the next takes the rest. +RESTARTS_PER_CHECK: Final = 20 +# How long the check waits for Modal to accept a restarted call. +CALL_START_TIMEOUT_SECONDS: Final = 10.0 +# No row is marked once a check has run this long: the rest are left unmarked +# for the next check. A marked row then needs only its call's start and a few +# bounded statements, so the check ends within the scheduled function's +# timeout and no row it marks is left without its call. +CHECK_TIME_LIMIT: Final = timedelta(minutes=3) +_NOT_STARTED: Final = ( + "run_not_started: The run stayed pending with no Modal call recorded for it" +) + + +class RestartOutcome(StrEnum): + """How the check ended for one listed row. + + ``restarted``: its new call started, and the row records that call. + ``call_not_recorded``: its new call started, but the check couldn't + record or read the call's ID; the call claims the row itself. + ``started_meanwhile``: the row records another call, which claimed it + after the check marked it (or the check's own call, accepted after the + check stopped waiting for it); the other call then can't claim it. + ``left``: the row was gone, changed before the check started a call, or + failed meanwhile; also a marked row whose failure couldn't be written, + which the sweep fails later, and a row whose call couldn't be read + after Modal didn't confirm the check's call. ``deferred``: the row is + left unmarked for the next check, because its rows or the manifest + couldn't be read, rebuilding its call raised an unexpected error (logged + with its traceback), the failure of a row that can't be rebuilt couldn't + be written, the mark couldn't be written, or the check's time ran out. + ``failed``: it couldn't be started again and was failed with + ``run_not_started``. + """ + + RESTARTED = "restarted" + CALL_NOT_RECORDED = "call_not_recorded" + STARTED_MEANWHILE = "started_meanwhile" + LEFT = "left" + DEFERRED = "deferred" + FAILED = "failed" + + +@dataclass(frozen=True) +class RestartSummary: + """How many rows one check listed, and how each ended (``RestartOutcome``).""" + + listed: int + restarted: int = 0 + call_not_recorded: int = 0 + started_meanwhile: int = 0 + left: int = 0 + deferred: int = 0 + failed: int = 0 + + +class EconomyRunWrites(Protocol): + """The submission's own call writes on economy runs (``EconomyRunStore``).""" + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + """Record the report run's call ID while it is pending with none.""" + ... + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + """Fail the report run, and its simulation runs, while unstarted.""" + ... + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + """Record the simulation run's call ID while it is pending with none.""" + ... + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + """Fail the simulation run while it is pending with no call ID.""" + ... + + +class RestartTelemetry(Protocol): + """Where the check reports each listed row (``ObservabilityRestartTelemetry``).""" + + def restart_scope( + self, run: UnstartedRun, *, age_seconds: float + ) -> AbstractContextManager[None]: + """Return the scope one row's restart runs in.""" + ... + + def record_outcome( + self, + run: UnstartedRun, + outcome: RestartOutcome, + *, + error_type: str | None = None, + ) -> None: + """Report how the row's restart ended, and the error type that ended it.""" + ... + + def capture_context(self) -> dict[str, str] | None: + """Return the context a restarted call continues the check's trace from.""" + ... + + +class NoRestartTelemetry: + """Report nothing.""" + + def restart_scope( + self, + run: UnstartedRun, # noqa: ARG002 # the protocol's argument + *, + age_seconds: float, # noqa: ARG002 # the protocol's argument + ) -> AbstractContextManager[None]: + """Return a scope that does nothing.""" + return nullcontext() + + def record_outcome( + self, + run: UnstartedRun, + outcome: RestartOutcome, + *, + error_type: str | None = None, + ) -> None: + """Report nothing.""" + + def capture_context(self) -> dict[str, str] | None: + """Return no context.""" + return None + + +@dataclass(frozen=True) +class RestartStores: + """The stores a restart reads and writes.""" + + restarts: UnstartedRunSource + economy: EconomyRunWrites + households: HouseholdJobRestarts + + +class RunRestartService: + """Start the calls of the runs and household jobs whose call never started. + + A call is rebuilt on a thread: the reads, Modal's manifest read included, + are synchronous. The check's other database writes are short and run on + its own event loop, which only the check uses. + """ + + def __init__( # noqa: PLR0913 # each seam tests replace + self, + *, + stores: RestartStores, + workers: RestartWorkers, + callers: RestartCallers, + settings: RestartSettings, + telemetry: RestartTelemetry | None = None, + clock: Callable[[], datetime] = lambda: datetime.now(UTC), + timer: Callable[[], float] = time.monotonic, + call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, + new_request_id: Callable[[], str] = lambda: f"run-restart-{uuid.uuid4()}", + ) -> None: + """Use these seams; tests replace each of them.""" + self._stores = stores + self._calls = RestartCallBuilder( + restarts=stores.restarts, + households=stores.households, + workers=workers, + callers=callers, + settings=settings, + new_request_id=new_request_id, + ) + self._telemetry = telemetry or NoRestartTelemetry() + self._clock = clock + self._timer = timer + self._call_start_timeout_seconds = call_start_timeout_seconds + self._deadline = 0.0 + + async def restart_unstarted_runs(self) -> RestartSummary: + """Start once more each row created 3 to 10 minutes ago with no call. + + Rows the check has no time left for are counted ``deferred``, with no + telemetry scope, since the next check lists them again. Raise the + store's error when the rows can't be listed. + """ + self._deadline = self._timer() + CHECK_TIME_LIMIT.total_seconds() + now = self._clock() + runs = self._stores.restarts.list_unstarted( + created_after=now - RESTART_UNTIL, + created_before=now - RESTART_AFTER, + limit=RESTARTS_PER_CHECK, + ) + outcomes: Counter[RestartOutcome] = Counter() + for position, run in enumerate(runs): + if self._out_of_time(): + outcomes[RestartOutcome.DEFERRED] += len(runs) - position + break + age_seconds = (now - run.created_at).total_seconds() + with self._telemetry.restart_scope(run, age_seconds=age_seconds): + outcome, error_type = await self._restart(run) + self._telemetry.record_outcome(run, outcome, error_type=error_type) + outcomes[outcome] += 1 + return RestartSummary( + listed=len(runs), + **{outcome.value: count for outcome, count in outcomes.items()}, + ) + + def _out_of_time(self) -> bool: + return self._timer() >= self._deadline + + async def _restart(self, run: UnstartedRun) -> tuple[RestartOutcome, str | None]: + """Start the row's call once more; return how it ended and any error type.""" + try: + rebuilt = await asyncio.to_thread(self._calls.build, run) + except DeferredRestartError as deferred: + return RestartOutcome.DEFERRED, str(deferred) + except CallNotRebuiltError as cannot: + message = ( + f"{_NOT_STARTED}, and its call could not be built again: " + f"{cannot.reason}." + ) + failed = self._fail(run, message, cannot.simulation_run_ids) + # An unmarked row whose failure couldn't be written is listed again. + outcome = { + True: RestartOutcome.FAILED, + False: RestartOutcome.LEFT, + None: RestartOutcome.DEFERRED, + }[failed] + return outcome, cannot.error_type + # Any other error leaves this row unmarked for the next check, and + # the check goes on to the next row; its traceback is logged. + except Exception as error: + fields = {"run_kind": run.kind.value, "run_id": str(run.run_id)} + # The observability runtime's log handler writes only + # ``policyengine_attributes``, as route_support's own log does. + logger.exception( + "v2_run_restart_unexpected_error", + extra={**fields, "policyengine_attributes": fields}, + ) + return RestartOutcome.DEFERRED, type(error).__name__ + if rebuilt is None: + return RestartOutcome.LEFT, None + return await self._start_marked(run, rebuilt.start, rebuilt.simulation_run_ids) + + async def _start_marked( + self, + run: UnstartedRun, + start_call: CallStart, + simulation_run_ids: tuple[UUID, ...], + ) -> tuple[RestartOutcome, str | None]: + """Mark the row, then start its call and record the call's ID.""" + if self._out_of_time(): + return RestartOutcome.DEFERRED, None + try: + if not self._stores.restarts.mark_restart(run): + return RestartOutcome.LEFT, None + except SQLAlchemyError as error: + return RestartOutcome.DEFERRED, type(error).__name__ + try: + call_id = await asyncio.wait_for( + start_call(self._telemetry.capture_context()), + timeout=self._call_start_timeout_seconds, + ) + # Modal raises its own errors; any of them, or the timeout, means the + # call was not confirmed, and the row is failed unless it has started. + except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike + message = ( + f"{_NOT_STARTED}, and Modal did not accept the call that started it " + f"again ({type(error).__name__})." + ) + failed = self._fail(run, message, simulation_run_ids) + # A failure that changed nothing means a call holds the row now, + # or something failed it; one that couldn't be written leaves the + # row to the sweep. + if failed is None: + outcome = RestartOutcome.LEFT + elif failed: + outcome = RestartOutcome.FAILED + else: + outcome = self._held_since(run, call_id=None) + return outcome, type(error).__name__ + try: + self._record(run, call_id) + except SQLAlchemyError as error: + return RestartOutcome.CALL_NOT_RECORDED, type(error).__name__ + return self._held_since(run, call_id=call_id), None + + def _held_since(self, run: UnstartedRun, *, call_id: str | None) -> RestartOutcome: + """Return how the restart ended, from the call the row records now. + + ``call_id`` is the check's own call, or None when Modal didn't + confirm one. The record write keeps a call recorded or claimed + before it, so whose call the row records says whose start counts. + """ + try: + held = self._stores.restarts.read_call(run) + except SQLAlchemyError: + return ( + RestartOutcome.LEFT + if call_id is None + else RestartOutcome.CALL_NOT_RECORDED + ) + if held is None or held.modal_call_id is None: + return RestartOutcome.LEFT + if held.modal_call_id == call_id: + return RestartOutcome.RESTARTED + return RestartOutcome.STARTED_MEANWHILE + + def _fail( + self, run: UnstartedRun, message: str, simulation_run_ids: tuple[UUID, ...] + ) -> bool | None: + """Fail the row while it is unstarted; return whether it was failed. + + Return None when the failure couldn't be written. + """ + try: + if run.kind is RestartKind.REPORT_RUN: + failed = self._stores.economy.fail_unstarted_report( + run.run_id, simulation_run_ids, message + ) + elif run.kind is RestartKind.SIMULATION_RUN: + failed = self._stores.economy.fail_unstarted_simulation( + run.run_id, message + ) + else: + failed = self._stores.households.fail_unstarted(run.run_id, message) + except SQLAlchemyError: + return None + return failed + + def _record(self, run: UnstartedRun, call_id: str) -> None: + """Record the call's ID while the row is pending with no call.""" + if run.kind is RestartKind.REPORT_RUN: + self._stores.economy.record_report_call(run.run_id, call_id) + elif run.kind is RestartKind.SIMULATION_RUN: + self._stores.economy.record_simulation_call(run.run_id, call_id) + else: + self._stores.households.record_call(run.run_id, call_id) diff --git a/projects/policyengine-simulation-entry/tests/economy_service_fakes.py b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py index 535873f0f..8becd4e9e 100644 --- a/projects/policyengine-simulation-entry/tests/economy_service_fakes.py +++ b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: from collections.abc import Sequence + from policyengine_simulation_entry.economy_workers import EconomyWorkerSource from policyengine_simulation_persistence import ( EconomyRequest, NewEconomyReport, @@ -296,7 +297,7 @@ def numbered_ids(): def service( store: FakeStore, *, - workers: FakeWorkers | None = None, + workers: EconomyWorkerSource | None = None, caller: FakeCaller | None = None, output_bucket: str | None = OUTPUT_BUCKET, ) -> EconomySubmissionService: diff --git a/projects/policyengine-simulation-entry/tests/test_economy_workers.py b/projects/policyengine-simulation-entry/tests/test_economy_workers.py index 9192e4f27..38632d558 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_workers.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_workers.py @@ -194,3 +194,33 @@ def test_a_call_modal_gives_no_id_is_an_error(modal_functions, object_id) -> Non worker, {}, {}, observability_context=None ) ) + + +def test_the_worker_for_a_runs_version_is_found_when_it_is_not_the_default() -> None: + """A run keeps the .py version it was created on after the default moves.""" + workers = ManifestEconomyWorkers( + manifest_store(household_worker_version(), household_worker_version("6.18.0")) + ) + + worker = workers.worker_for_version("us", "6.18.0") + + assert worker.version == "6.18.0" + assert worker.policyengine_version == "6.18.0" + assert worker.application_name == v2_application_name("6.18.0") + assert worker.simulation_callable == "run_single_simulation_us" + + +def test_a_version_the_manifest_no_longer_holds_has_no_economy_worker() -> None: + workers = ManifestEconomyWorkers(manifest_store(household_worker_version())) + + with pytest.raises(EconomyWorkerUnavailableError) as raised: + workers.worker_for_version("us", "6.0.0") + + assert str(raised.value) == "the v2 version manifest has no worker for .py 6.0.0" + + +def test_a_workers_version_lookup_reads_an_unreadable_manifest_as_unavailable() -> None: + with pytest.raises(VersionManifestUnavailableError): + ManifestEconomyWorkers(UnreadableStore()).worker_for_version( + "us", POLICYENGINE_VERSION + ) diff --git a/projects/policyengine-simulation-entry/tests/test_household_workers.py b/projects/policyengine-simulation-entry/tests/test_household_workers.py index 661308082..cbc2786eb 100644 --- a/projects/policyengine-simulation-entry/tests/test_household_workers.py +++ b/projects/policyengine-simulation-entry/tests/test_household_workers.py @@ -181,3 +181,32 @@ async def test_a_call_modal_gives_no_id_is_refused(monkeypatch): await ModalHouseholdFunctionCaller("staging").spawn( worker, {}, household_job_id=JOB_ID, observability_context=None ) + + +def test_the_household_function_for_a_jobs_version_is_found_when_not_the_default() -> ( + None +): + """A job keeps the .py version it was created on after the default moves.""" + workers = ManifestHouseholdWorkers( + manifest_store(household_worker_version(), household_worker_version("6.18.0")) + ) + + worker = workers.worker_for_version("us", "6.18.0") + + assert worker.bundle.policyengine_version == "6.18.0" + assert worker.application_name == v2_application_name("6.18.0") + assert worker.function_name == "calculate_household_us" + + +def test_a_version_the_manifest_no_longer_holds_has_no_household_function() -> None: + workers = ManifestHouseholdWorkers(manifest_store(household_worker_version())) + + with pytest.raises(HouseholdFunctionUnavailableError) as raised: + workers.worker_for_version("us", "6.0.0") + + assert str(raised.value) == "the v2 version manifest has no worker for .py 6.0.0" + + +def test_a_household_version_lookup_reads_an_unreadable_store_as_unavailable() -> None: + with pytest.raises(VersionManifestUnavailableError): + ManifestHouseholdWorkers({}).worker_for_version("us", POLICYENGINE_VERSION) diff --git a/projects/policyengine-simulation-entry/tests/test_run_restart_check.py b/projects/policyengine-simulation-entry/tests/test_run_restart_check.py new file mode 100644 index 000000000..4b778a3e1 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_restart_check.py @@ -0,0 +1,221 @@ +"""The scheduled restart check's telemetry, and the service it runs.""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any +from uuid import UUID + +import pytest +from policyengine_simulation_entry.run_restart_check import ( + ObservabilityRestartTelemetry, + run_scheduled_check, + scheduled_restart_service, +) +from policyengine_simulation_entry.run_restarts import ( + RestartOutcome, + RestartSettings, + RestartSummary, + RunRestartService, +) +from policyengine_simulation_persistence import RestartKind, UnstartedRun + +if TYPE_CHECKING: + from collections.abc import Iterator + +RUN = UnstartedRun( + kind=RestartKind.REPORT_RUN, + run_id=UUID("00000000-0000-4000-8000-0000000000e1"), + created_at=datetime(2026, 11, 13, 9, tzinfo=UTC), +) +DATABASE_URL = "postgresql://worker:secret@database.example/postgres" + + +@dataclass +class RecordingRuntime: + """Records what the telemetry asks of the observability runtime.""" + + operations: list[tuple[str, dict[str, Any]]] = field(default_factory=list) + contexts: list[dict[str, Any]] = field(default_factory=list) + events: list[tuple[str, str, dict[str, Any]]] = field(default_factory=list) + open_operations: list[str] = field(default_factory=list) + + @contextmanager + def operation( + self, name: str, *, attributes: dict[str, Any] | None = None + ) -> Iterator[None]: + self.operations.append((name, dict(attributes or {}))) + self.open_operations.append(name) + try: + yield + finally: + self.open_operations.pop() + + def set_context(self, **attributes: Any) -> None: + self.contexts.append({"within": list(self.open_operations), **attributes}) + + def event( + self, + name: str, + *, + severity: str = "INFO", + attributes: dict[str, Any] | None = None, + ) -> None: + self.events.append((name, severity, dict(attributes or {}))) + + def capture_context(self) -> dict[str, str]: + return {"traceparent": "00-restart-check"} + + +def test_each_restart_is_an_operation_named_by_its_registered_stage() -> None: + runtime = RecordingRuntime() + telemetry = ObservabilityRestartTelemetry(runtime) + + with telemetry.restart_scope(RUN, age_seconds=210.5): + within = list(runtime.open_operations) + + assert within == ["v2_run_restart"] + assert runtime.operations == [ + ( + "v2_run_restart", + { + "run_kind": "report_run", + "run_id": str(RUN.run_id), + "run_age_seconds": 210.5, + }, + ) + ] + + +@pytest.mark.parametrize( + ("outcome", "error_type", "severity"), + [ + (RestartOutcome.RESTARTED, None, "INFO"), + (RestartOutcome.STARTED_MEANWHILE, None, "INFO"), + (RestartOutcome.LEFT, None, "INFO"), + (RestartOutcome.DEFERRED, "OperationalError", "WARNING"), + (RestartOutcome.CALL_NOT_RECORDED, "OperationalError", "WARNING"), + (RestartOutcome.FAILED, "ConnectionError", "WARNING"), + ], +) +def test_an_outcome_is_set_on_the_restart_and_logged_as_an_event( + outcome, error_type, severity +) -> None: + runtime = RecordingRuntime() + telemetry = ObservabilityRestartTelemetry(runtime) + + with telemetry.restart_scope(RUN, age_seconds=200.0): + telemetry.record_outcome(RUN, outcome, error_type=error_type) + + expected = {"restart_outcome": outcome.value} + if error_type is not None: + expected["restart_error_type"] = error_type + assert runtime.contexts == [{"within": ["v2_run_restart"], **expected}] + assert runtime.events == [ + ( + "v2_run_restart_outcome", + severity, + {"run_kind": "report_run", "run_id": str(RUN.run_id), **expected}, + ) + ] + + +def test_a_restarted_call_continues_the_checks_trace() -> None: + telemetry = ObservabilityRestartTelemetry(RecordingRuntime()) + + assert telemetry.capture_context() == {"traceparent": "00-restart-check"} + + +@dataclass +class FakeRestartService: + """Answers one check with a summary, noting the operation it ran in.""" + + runtime: RecordingRuntime + summary: RestartSummary + ran_within: list[list[str]] = field(default_factory=list) + + async def restart_unstarted_runs(self) -> RestartSummary: + self.ran_within.append(list(self.runtime.open_operations)) + return self.summary + + +def test_a_check_runs_in_its_own_operation_and_answers_its_counts() -> None: + runtime = RecordingRuntime() + summary = RestartSummary(listed=3, restarted=1, failed=1, left=1) + restart_service = FakeRestartService(runtime, summary) + + counts = run_scheduled_check(restart_service, runtime) + + assert counts == { + "listed": 3, + "restarted": 1, + "call_not_recorded": 0, + "started_meanwhile": 0, + "left": 1, + "deferred": 0, + "failed": 1, + } + assert restart_service.ran_within == [["v2_run_restart_check"]] + assert runtime.operations == [("v2_run_restart_check", {})] + assert runtime.contexts == [ + {"within": ["v2_run_restart_check"]} + | {f"restart_{name}": count for name, count in counts.items()} + ] + + +def test_the_scheduled_service_restarts_in_its_modal_environments_deployment() -> None: + scheduled = scheduled_restart_service( + database_url=DATABASE_URL, + modal_environment="main", + runtime=RecordingRuntime(), + ) + try: + assert isinstance(scheduled.service, RunRestartService) + assert scheduled.settings == RestartSettings( + environment="production", modal_environment="main" + ) + finally: + scheduled.engine.dispose() + + +def test_the_scheduled_service_refuses_a_modal_environment_no_deployment_uses() -> None: + with pytest.raises(ValueError, match="local"): + scheduled_restart_service( + database_url=DATABASE_URL, + modal_environment="local", + runtime=RecordingRuntime(), + ) + + +def test_the_restart_module_imports_without_the_web_service(tmp_path) -> None: + """The restart check runs in the sweep's image, which has no web packages. + + Importing it must not import the FastAPI app, whose authentication needs + PyJWT and cryptography, which that image does not install. + """ + import subprocess # this test alone starts a process + import sys # this test alone starts a process + + script = tmp_path / "import_restart_check.py" + script.write_text( + "import sys\n" + "class Refuse:\n" + " def find_spec(self, name, path=None, target=None):\n" + " if name.split('.')[0] in {'jwt', 'cryptography', 'uvicorn'}:\n" + " raise ImportError(name)\n" + "sys.meta_path.insert(0, Refuse())\n" + "import policyengine_simulation_entry.run_restart_check\n" + "assert 'policyengine_simulation_entry.app' not in sys.modules\n", + encoding="utf-8", + ) + + completed = subprocess.run( # the test's own interpreter and script + [sys.executable, str(script)], + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr diff --git a/projects/policyengine-simulation-entry/tests/test_run_restarts.py b/projects/policyengine-simulation-entry/tests/test_run_restarts.py new file mode 100644 index 000000000..8eaabe323 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_restarts.py @@ -0,0 +1,1002 @@ +"""Starting a run or household job again when its Modal call never started. + +The restart check rebuilds a run's call from its stored rows. Each test of a +restart first submits the run through the service that creates it, captures +the call that submission started, and requires the restart to start the same +call: the same payload, and the same context apart from the request ID. The +store's reads are tested against PostgreSQL in the persistence library; here +fakes answer as the stored rows would. +""" + +from __future__ import annotations + +import asyncio +from contextlib import contextmanager +from dataclasses import dataclass, field, replace +from datetime import timedelta +from typing import TYPE_CHECKING, Any +from uuid import UUID + +import pytest +from economy_fixtures import economy_worker_version +from economy_service_fakes import ( + CATALOG_VERSION_ID, + CREATED_AT, + DYNAMIC_ID, + POLICY_ID, + REFORM, + FakeCaller, + FakeStore, + report_request, + service, + simulation_request, +) +from household_fixtures import manifest_store +from policyengine_simulation_entry.economy_workers import ManifestEconomyWorkers +from policyengine_simulation_entry.household_api_models import ( + HouseholdCalculationRequest, +) +from policyengine_simulation_entry.household_calculations import ( + HouseholdCalculationService, +) +from policyengine_simulation_entry.household_workers import ( + ManifestHouseholdWorkers, +) +from policyengine_simulation_entry.run_restarts import ( + CHECK_TIME_LIMIT, + RESTART_AFTER, + RESTART_UNTIL, + RESTARTS_PER_CHECK, + RestartCallers, + RestartOutcome, + RestartSettings, + RestartStores, + RestartSummary, + RestartWorkers, + RunRestartService, +) +from policyengine_simulation_persistence import ( + EconomyScope, + HouseholdCalculationInputs, + HouseholdJobRecord, + HouseholdJobStatus, + RestartKind, + RunCall, + SimulationChoice, + StoredEconomySide, + StoredHousehold, + StoredReportRun, + StoredRunUnreadableError, + StoredSimulationRun, + UnstartedRun, +) +from sqlalchemy.exc import OperationalError + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from datetime import datetime + + from policyengine_simulation_entry.household_workers import HouseholdWorker + from policyengine_simulation_persistence import NewHouseholdJob + +NOW = CREATED_AT + timedelta(minutes=4) +DATABASE_ERROR = OperationalError("SELECT 1", {}, Exception("server closed")) +HOUSEHOLD_ID = UUID("00000000-0000-4000-8000-0000000000bb") +HOUSEHOLD_JOB_ID = UUID("00000000-0000-4000-8000-0000000000aa") +HOUSEHOLD_DOCUMENT = { + "people": [ + { + "id": "you", + "values": {"age": 40, "employment_income": 30_000}, + "memberships": { + "household": "your household", + "family": "your family", + "tax_unit": "your tax unit", + "spm_unit": "your spm unit", + "marital_unit": "your marital unit", + }, + } + ], + "household": [{"id": "your household", "values": {"state_name": "TX"}}], + "family": [{"id": "your family", "values": {}}], + "tax_unit": [{"id": "your tax unit", "values": {}}], + "spm_unit": [{"id": "your spm unit", "values": {}}], + "marital_unit": [{"id": "your marital unit", "values": {}}], +} +SETTINGS = RestartSettings(environment="staging", modal_environment="staging") +# The manifest both a submission and its restart read: the call's context +# records the digest of the manifest its worker was read from. +MANIFEST_STORE = manifest_store(economy_worker_version()) + + +@dataclass +class FakeRestartStore: + """Lists the unstarted runs it is given and reads back their stored rows. + + ``marked_elsewhere`` makes every mark change nothing, as when another + check marked the run first. + """ + + unstarted: tuple[UnstartedRun, ...] = () + report_runs: dict[UUID, StoredReportRun] = field(default_factory=dict) + simulation_runs: dict[UUID, StoredSimulationRun] = field(default_factory=dict) + report_simulation_run_ids: dict[UUID, tuple[UUID, ...]] = field( + default_factory=dict + ) + read_failure: Exception | None = None + # A read of one run that raises, by run ID. + read_failures: dict[UUID, Exception] = field(default_factory=dict) + marked_elsewhere: bool = False + # Each row's status and call as stored, by run ID; a row not named here is + # as the rig's other fakes left it (``holder``). + held: dict[UUID, RunCall | None] = field(default_factory=dict) + holder: Callable[[UnstartedRun], RunCall | None] | None = None + windows: list[tuple[datetime, datetime, int]] = field(default_factory=list) + marked: list[UnstartedRun] = field(default_factory=list) + + def read_call(self, run: UnstartedRun) -> RunCall | None: + if run.run_id in self.held: + return self.held[run.run_id] + return None if self.holder is None else self.holder(run) + + def read_report_run_simulation_run_ids( + self, report_run_id: UUID + ) -> tuple[UUID, ...]: + return self.report_simulation_run_ids.get(report_run_id, ()) + + def list_unstarted( + self, *, created_after: datetime, created_before: datetime, limit: int + ) -> tuple[UnstartedRun, ...]: + self.windows.append((created_after, created_before, limit)) + return self.unstarted + + def mark_restart(self, run: UnstartedRun) -> bool: + if self.marked_elsewhere: + return False + self.marked.append(run) + return True + + def read_report_run(self, report_run_id: UUID) -> StoredReportRun | None: + if self.read_failure is not None: + raise self.read_failure + if report_run_id in self.read_failures: + raise self.read_failures[report_run_id] + return self.report_runs.get(report_run_id) + + def read_simulation_run( + self, simulation_run_id: UUID + ) -> StoredSimulationRun | None: + if self.read_failure is not None: + raise self.read_failure + return self.simulation_runs.get(simulation_run_id) + + +@dataclass +class FakeHouseholdJobs: + """Holds the household jobs created, and answers the restart's reads of them.""" + + default_year: int | None = 2026 + reform: dict[str, dict[str, object]] = field(default_factory=dict) + claimed: bool = False + created: list[NewHouseholdJob] = field(default_factory=list) + recorded: list[tuple[UUID, str]] = field(default_factory=list) + failed: list[tuple[UUID, str]] = field(default_factory=list) + + def read_inputs(self, **arguments: Any) -> HouseholdCalculationInputs: + return HouseholdCalculationInputs( + catalog_version_id=CATALOG_VERSION_ID, + household=StoredHousehold( + id=arguments["household_id"], + country=arguments["country"], + default_year=self.default_year, + document=HOUSEHOLD_DOCUMENT, + ), + reform=self.reform, + ) + + def create_job(self, job: NewHouseholdJob) -> None: + self.created.append(job) + + def record_call(self, job_id: UUID, modal_call_id: str) -> HouseholdJobStatus: + self.recorded.append((job_id, modal_call_id)) + return ( + HouseholdJobStatus.RUNNING if self.claimed else HouseholdJobStatus.PENDING + ) + + def fail_unstarted(self, job_id: UUID, error_message: str) -> bool: + if self.claimed: + return False + self.failed.append((job_id, error_message)) + return True + + def read(self, job_id: UUID) -> HouseholdJobRecord | None: + job = next((job for job in self.created if job.id == job_id), None) + if job is None: + return None + return HouseholdJobRecord( + id=job.id, + status=HouseholdJobStatus.PENDING, + country=job.country, + household_id=job.household_id, + policy_id=job.policy_id, + year=job.year, + request_data=job.request_data, + tax_benefit_model_version_id=job.tax_benefit_model_version_id, + country_package_version=job.country_package_version, + policyengine_version=job.policyengine_version, + modal_call_id=None, + result=None, + error_message=None, + created_at=CREATED_AT, + started_at=None, + completed_at=None, + ) + + +@dataclass +class FakeHouseholdCaller: + """Records each household function call; ``failure`` is raised instead.""" + + failure: Exception | None = None + calls: list[dict[str, Any]] = field(default_factory=list) + + async def spawn( + self, + worker: HouseholdWorker, + payload: dict[str, Any], + *, + household_job_id: UUID, + observability_context: dict[str, Any] | None, + ) -> str: + if self.failure is not None: + raise self.failure + self.calls.append( + { + "worker": worker, + "payload": payload, + "household_job_id": household_job_id, + "observability_context": observability_context, + } + ) + return "fc-household-restarted" + + +@dataclass +class RecordingTelemetry: + """Records each restart scope and outcome the check reports.""" + + scopes: list[tuple[RestartKind, UUID, float]] = field(default_factory=list) + outcomes: list[tuple[RestartKind, UUID, RestartOutcome, str | None]] = field( + default_factory=list + ) + + @contextmanager + def restart_scope(self, run: UnstartedRun, *, age_seconds: float) -> Iterator[None]: + self.scopes.append((run.kind, run.run_id, age_seconds)) + yield + + def record_outcome( + self, + run: UnstartedRun, + outcome: RestartOutcome, + *, + error_type: str | None = None, + ) -> None: + self.outcomes.append((run.kind, run.run_id, outcome, error_type)) + + def capture_context(self) -> dict[str, str]: + return {"traceparent": "00-restart-check"} + + +@dataclass +class RestartRig: + """The restart service and every fake it was built with.""" + + restarts: FakeRestartStore + economy: FakeStore + households: FakeHouseholdJobs + economy_caller: FakeCaller + household_caller: FakeHouseholdCaller + telemetry: RecordingTelemetry + service: RunRestartService + + def restart(self) -> RestartSummary: + return asyncio.run(self.service.restart_unstarted_runs()) + + +def _rig( # noqa: PLR0913 # each fake a test may replace + *, + restarts: FakeRestartStore | None = None, + economy: FakeStore | None = None, + households: FakeHouseholdJobs | None = None, + economy_caller: FakeCaller | None = None, + household_caller: FakeHouseholdCaller | None = None, + manifest: object | None = None, + timer: Callable[[], float] | None = None, +) -> RestartRig: + restarts = restarts or FakeRestartStore() + economy = economy or FakeStore() + households = households or FakeHouseholdJobs() + economy_caller = economy_caller or FakeCaller() + household_caller = household_caller or FakeHouseholdCaller() + telemetry = RecordingTelemetry() + store = MANIFEST_STORE if manifest is None else manifest + + def holder(run: UnstartedRun) -> RunCall: + """Return the row as the fakes left it: claimed, given a call, or not.""" + if economy.claimed or households.claimed: + return RunCall(status="running", modal_call_id="fc-original") + recorded = dict([*economy.recorded_calls, *households.recorded]) + return RunCall(status="pending", modal_call_id=recorded.get(run.run_id)) + + if restarts.holder is None: + restarts.holder = holder + restart_service = RunRestartService( + stores=RestartStores(restarts=restarts, economy=economy, households=households), + workers=RestartWorkers( + economy=ManifestEconomyWorkers(store), + households=ManifestHouseholdWorkers(store), + ), + callers=RestartCallers(economy=economy_caller, households=household_caller), + settings=SETTINGS, + telemetry=telemetry, + clock=lambda: NOW, + timer=timer or (lambda: 0.0), + call_start_timeout_seconds=0.05, + new_request_id=lambda: "restart-request", + ) + return RestartRig( + restarts=restarts, + economy=economy, + households=households, + economy_caller=economy_caller, + household_caller=household_caller, + telemetry=telemetry, + service=restart_service, + ) + + +def _submitted_report() -> tuple[StoredReportRun, tuple[dict, dict]]: + """Submit a report as the service does; return its stored run and its call.""" + caller = FakeCaller() + created = asyncio.run( + service( + FakeStore(), workers=ManifestEconomyWorkers(MANIFEST_STORE), caller=caller + ).submit_report( + report_request(), request_id="original-request", observability_context=None + ) + ) + (worker, payload, context, _observability) = caller.report_runs[0] + stored = StoredReportRun( + report_run_id=created.report_run_id, + country="us", + versions=worker.versions, + options=report_request().options(), + scope=EconomyScope( + catalog_version_id=CATALOG_VERSION_ID, + region_code="us", + dataset_name="populace_us_2024", + year=2026, + ), + baseline=StoredEconomySide( + simulation_run_id=created.baseline_simulation_run_id, + choice=SimulationChoice(policy_id=None, dynamic_id=None), + reform={}, + ), + reform=StoredEconomySide( + simulation_run_id=created.reform_simulation_run_id, + choice=SimulationChoice(policy_id=POLICY_ID, dynamic_id=DYNAMIC_ID), + reform=REFORM, + ), + created_at=CREATED_AT, + ) + return stored, (payload, context) + + +def _submitted_simulation() -> tuple[StoredSimulationRun, tuple[dict, dict]]: + """Submit a standalone simulation; return its stored run and its call.""" + caller = FakeCaller() + created = asyncio.run( + service( + FakeStore(), workers=ManifestEconomyWorkers(MANIFEST_STORE), caller=caller + ).submit_simulation( + simulation_request(), + request_id="original-request", + observability_context=None, + ) + ) + (worker, payload, context, _observability) = caller.simulation_runs[0] + stored = StoredSimulationRun( + country="us", + versions=worker.versions, + scope=EconomyScope( + catalog_version_id=CATALOG_VERSION_ID, + region_code="us", + dataset_name="populace_us_2024", + year=2026, + ), + run=StoredEconomySide( + simulation_run_id=created.simulation_run_id, + choice=SimulationChoice(policy_id=POLICY_ID, dynamic_id=None), + reform=REFORM, + ), + created_at=CREATED_AT, + ) + return stored, (payload, context) + + +def _submitted_household_job( + households: FakeHouseholdJobs, **request_fields: Any +) -> dict[str, Any]: + """Submit a household calculation into ``households``; return its call's payload.""" + caller = FakeHouseholdCaller() + household_service = HouseholdCalculationService( + store=households, + workers=ManifestHouseholdWorkers(MANIFEST_STORE), + caller=caller, + new_job_id=lambda: HOUSEHOLD_JOB_ID, + ) + asyncio.run( + household_service.submit( + HouseholdCalculationRequest.model_validate( + {"country": "us", "household_id": str(HOUSEHOLD_ID), **request_fields} + ), + observability_context=None, + ) + ) + households.recorded.clear() + return caller.calls[0]["payload"] + + +def _unstarted(kind: RestartKind, run_id: UUID) -> UnstartedRun: + return UnstartedRun(kind=kind, run_id=run_id, created_at=CREATED_AT) + + +def _without_request_id(context: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in context.items() if key != "request_id"} + + +def test_the_check_lists_runs_created_three_to_ten_minutes_ago() -> None: + rig = _rig() + + summary = rig.restart() + + assert rig.restarts.windows == [ + (NOW - RESTART_UNTIL, NOW - RESTART_AFTER, RESTARTS_PER_CHECK) + ] + assert ( + timedelta(minutes=3), + timedelta(minutes=10), + ) == (RESTART_AFTER, RESTART_UNTIL) + assert summary == RestartSummary(listed=0) + + +def test_a_report_run_is_started_again_with_the_call_its_submission_started() -> None: + stored, (payload, context) = _submitted_report() + run = _unstarted(RestartKind.REPORT_RUN, stored.report_run_id) + rig = _rig( + restarts=FakeRestartStore( + unstarted=(run,), report_runs={stored.report_run_id: stored} + ) + ) + + summary = rig.restart() + + ((_worker, restarted_payload, restarted_context, observability),) = ( + rig.economy_caller.report_runs + ) + assert restarted_payload == payload + assert _without_request_id(restarted_context) == _without_request_id(context) + assert restarted_context["request_id"] == "restart-request" + assert observability == {"traceparent": "00-restart-check"} + assert rig.restarts.marked == [run] + assert rig.economy.recorded_calls == [(stored.report_run_id, "fc-coordinator")] + assert summary == RestartSummary(listed=1, restarted=1) + + +def test_a_standalone_simulation_run_is_started_again_with_its_submissions_call() -> ( + None +): + stored, (payload, context) = _submitted_simulation() + run_id = stored.run.simulation_run_id + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.SIMULATION_RUN, run_id),), + simulation_runs={run_id: stored}, + ) + ) + + summary = rig.restart() + + ((_worker, restarted_payload, restarted_context, _observability),) = ( + rig.economy_caller.simulation_runs + ) + assert restarted_payload == payload + assert _without_request_id(restarted_context) == _without_request_id(context) + assert rig.economy.recorded_calls == [(run_id, "fc-worker")] + assert summary == RestartSummary(listed=1, restarted=1) + + +def test_a_household_job_is_started_again_with_its_submissions_call() -> None: + households = FakeHouseholdJobs(reform=REFORM) + payload = _submitted_household_job( + households, policy_id=str(POLICY_ID), add_missing_variables=True + ) + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.HOUSEHOLD_JOB, HOUSEHOLD_JOB_ID),) + ), + households=households, + ) + + summary = rig.restart() + + (call,) = rig.household_caller.calls + assert call["payload"] == payload + assert call["household_job_id"] == HOUSEHOLD_JOB_ID + assert call["observability_context"] == {"traceparent": "00-restart-check"} + assert households.recorded == [(HOUSEHOLD_JOB_ID, "fc-household-restarted")] + assert summary == RestartSummary(listed=1, restarted=1) + + +def test_a_run_another_check_marked_first_is_left_without_a_call() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + marked_elsewhere=True, + ) + ) + + summary = rig.restart() + + assert rig.economy_caller.report_runs == [] + assert rig.economy.failed_runs == [] + assert summary == RestartSummary(listed=1, left=1) + + +def test_a_run_gone_before_it_is_read_is_left_alone() -> None: + rig = _rig( + restarts=FakeRestartStore( + unstarted=( + _unstarted(RestartKind.REPORT_RUN, UUID(int=1)), + _unstarted(RestartKind.SIMULATION_RUN, UUID(int=2)), + _unstarted(RestartKind.HOUSEHOLD_JOB, UUID(int=3)), + ) + ) + ) + + summary = rig.restart() + + assert rig.restarts.marked == [] + assert summary == RestartSummary(listed=3, left=3) + + +def test_a_run_whose_rows_cannot_be_read_is_left_for_the_next_check() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + read_failure=DATABASE_ERROR, + ) + ) + + summary = rig.restart() + + assert rig.restarts.marked == [] + assert rig.economy.failed_runs == [] + assert summary == RestartSummary(listed=1, deferred=1) + + +def test_an_unreadable_manifest_leaves_the_run_for_the_next_check() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + ), + manifest={}, + ) + + summary = rig.restart() + + assert rig.restarts.marked == [] + assert summary == RestartSummary(listed=1, deferred=1) + + +def test_a_run_whose_worker_left_the_manifest_is_failed_as_not_started() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + ), + manifest=manifest_store(economy_worker_version("6.18.0")), + ) + + summary = rig.restart() + + ((run_id, simulation_run_ids, message),) = rig.economy.failed_runs + assert run_id == stored.report_run_id + assert simulation_run_ids == ( + stored.baseline.simulation_run_id, + stored.reform.simulation_run_id, + ) + assert message.startswith("run_not_started: ") + assert f".py {stored.versions.policyengine_version}" in message + assert rig.economy_caller.report_runs == [] + assert summary == RestartSummary(listed=1, failed=1) + + +def test_a_run_whose_stored_rows_no_longer_describe_it_is_failed() -> None: + run_id = UUID(int=7) + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.SIMULATION_RUN, run_id),), + read_failure=StoredRunUnreadableError("simulation run 7 names no region"), + ) + ) + + summary = rig.restart() + + ((failed_run_id, _sides, message),) = rig.economy.failed_runs + assert failed_run_id == run_id + assert message.startswith("run_not_started: ") + assert "stored rows" in message + assert summary == RestartSummary(listed=1, failed=1) + + +def test_a_household_job_whose_household_changed_is_failed_as_not_started() -> None: + households = FakeHouseholdJobs() + _submitted_household_job(households) + households.default_year = 2027 + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.HOUSEHOLD_JOB, HOUSEHOLD_JOB_ID),) + ), + households=households, + ) + + summary = rig.restart() + + ((job_id, message),) = households.failed + assert job_id == HOUSEHOLD_JOB_ID + assert message.startswith("run_not_started: ") + assert rig.household_caller.calls == [] + assert summary == RestartSummary(listed=1, failed=1) + + +def test_a_restart_modal_refuses_fails_the_run_as_not_started() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + ), + economy_caller=FakeCaller(failure=ConnectionError("Modal is down")), + ) + + summary = rig.restart() + + ((run_id, simulation_run_ids, message),) = rig.economy.failed_runs + assert run_id == stored.report_run_id + assert len(simulation_run_ids) == 2 + assert message == ( + "run_not_started: The run stayed pending with no Modal call recorded for " + "it, and Modal did not accept the call that started it again " + "(ConnectionError)." + ) + assert summary == RestartSummary(listed=1, failed=1) + assert rig.telemetry.outcomes == [ + ( + RestartKind.REPORT_RUN, + stored.report_run_id, + RestartOutcome.FAILED, + "ConnectionError", + ) + ] + + +def test_a_restart_modal_does_not_confirm_in_time_fails_the_run() -> None: + stored, _call = _submitted_simulation() + run_id = stored.run.simulation_run_id + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.SIMULATION_RUN, run_id),), + simulation_runs={run_id: stored}, + ), + economy_caller=FakeCaller(hang=True), + ) + + summary = rig.restart() + + ((failed_run_id, _sides, message),) = rig.economy.failed_runs + assert failed_run_id == run_id + assert "(TimeoutError)" in message + assert summary == RestartSummary(listed=1, failed=1) + + +def test_a_refused_restart_of_a_run_another_call_claimed_leaves_it() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + ), + economy=FakeStore(claimed=True), + economy_caller=FakeCaller(failure=ConnectionError("Modal is down")), + ) + + summary = rig.restart() + + assert rig.economy.failed_runs == [] + assert summary == RestartSummary(listed=1, started_meanwhile=1) + + +def test_a_restarted_call_another_call_claimed_first_counts_as_started_meanwhile() -> ( + None +): + households = FakeHouseholdJobs() + _submitted_household_job(households) + households.claimed = True + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.HOUSEHOLD_JOB, HOUSEHOLD_JOB_ID),) + ), + households=households, + ) + + summary = rig.restart() + + assert len(rig.household_caller.calls) == 1 + assert summary == RestartSummary(listed=1, started_meanwhile=1) + + +def test_a_restart_whose_call_cannot_be_recorded_still_counts_its_call() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + ), + economy=FakeStore(record_failure=DATABASE_ERROR), + ) + + summary = rig.restart() + + assert len(rig.economy_caller.report_runs) == 1 + assert rig.economy.failed_runs == [] + assert summary == RestartSummary(listed=1, call_not_recorded=1) + + +def test_each_listed_run_is_one_telemetry_scope_with_its_age_and_outcome() -> None: + report, _report_call = _submitted_report() + simulation, _simulation_call = _submitted_simulation() + simulation_run_id = simulation.run.simulation_run_id + report_run = replace( + _unstarted(RestartKind.REPORT_RUN, report.report_run_id), + created_at=CREATED_AT + timedelta(seconds=30), + ) + rig = _rig( + restarts=FakeRestartStore( + unstarted=( + report_run, + _unstarted(RestartKind.SIMULATION_RUN, simulation_run_id), + ), + report_runs={report.report_run_id: report}, + ) + ) + + summary = rig.restart() + + assert rig.telemetry.scopes == [ + (RestartKind.REPORT_RUN, report.report_run_id, 210.0), + (RestartKind.SIMULATION_RUN, simulation_run_id, 240.0), + ] + assert rig.telemetry.outcomes == [ + (RestartKind.REPORT_RUN, report.report_run_id, RestartOutcome.RESTARTED, None), + (RestartKind.SIMULATION_RUN, simulation_run_id, RestartOutcome.LEFT, None), + ] + assert summary == RestartSummary(listed=2, restarted=1, left=1) + + +@pytest.mark.parametrize( + ("modal_environment", "environment"), + [("staging", "staging"), ("main", "production")], +) +def test_the_deployment_label_follows_the_modal_environment( + modal_environment, environment +) -> None: + assert RestartSettings.for_modal_environment(modal_environment) == ( + RestartSettings(environment=environment, modal_environment=modal_environment) + ) + + +def test_a_modal_environment_no_deployment_uses_is_refused() -> None: + with pytest.raises(ValueError, match="dev"): + RestartSettings.for_modal_environment("dev") + + +def test_an_unreadable_report_run_is_failed_with_its_simulation_runs() -> None: + """As a submission fails a report run, with the simulation runs made for it.""" + run_id = UUID(int=8) + sides = (UUID(int=81), UUID(int=82)) + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, run_id),), + read_failure=StoredRunUnreadableError("report run 8 names no region"), + report_simulation_run_ids={run_id: sides}, + ) + ) + + summary = rig.restart() + + ((failed_run_id, simulation_run_ids, message),) = rig.economy.failed_runs + assert (failed_run_id, simulation_run_ids) == (run_id, sides) + assert message.startswith("run_not_started: ") + assert summary == RestartSummary(listed=1, failed=1) + + +def test_a_restarted_call_claiming_its_run_before_the_record_is_restarted() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + held={ + stored.report_run_id: RunCall( + status="running", modal_call_id="fc-coordinator" + ) + }, + ), + economy=FakeStore(claimed=True), + ) + + assert rig.restart() == RestartSummary(listed=1, restarted=1) + + +def test_a_run_failed_before_its_restart_was_recorded_is_left() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + held={stored.report_run_id: RunCall(status="failed", modal_call_id=None)}, + ), + economy=FakeStore(claimed=True), + ) + + assert rig.restart() == RestartSummary(listed=1, left=1) + + +def test_a_refused_restart_of_a_run_the_sweep_failed_meanwhile_is_left() -> None: + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + held={stored.report_run_id: RunCall(status="failed", modal_call_id=None)}, + ), + economy=FakeStore(claimed=True), + economy_caller=FakeCaller(failure=ConnectionError("Modal is down")), + ) + + assert rig.restart() == RestartSummary(listed=1, left=1) + + +@dataclass +class AdvancingCaller(FakeCaller): + """Starts calls as ``FakeCaller`` does, moving a timer on by each start.""" + + elapsed: list[float] = field(default_factory=lambda: [0.0]) + step_seconds: float = 0.0 + + async def start_report_run(self, *arguments: Any, **keywords: Any) -> str: + self.elapsed[0] += self.step_seconds + return await super().start_report_run(*arguments, **keywords) + + +def test_a_check_past_its_time_limit_leaves_the_other_runs_for_the_next() -> None: + """No row is marked once the check's time is up, so none waits for the sweep.""" + first, _first_call = _submitted_report() + second, _second_call = _submitted_report() + caller = AdvancingCaller(step_seconds=CHECK_TIME_LIMIT.total_seconds()) + rig = _rig( + restarts=FakeRestartStore( + unstarted=( + _unstarted(RestartKind.REPORT_RUN, first.report_run_id), + _unstarted(RestartKind.REPORT_RUN, second.report_run_id), + ), + report_runs={first.report_run_id: first, second.report_run_id: second}, + ), + economy_caller=caller, + timer=lambda: caller.elapsed[0], + ) + + summary = rig.restart() + + assert [run.run_id for run in rig.restarts.marked] == [first.report_run_id] + assert [scope[1] for scope in rig.telemetry.scopes] == [first.report_run_id] + assert summary == RestartSummary(listed=2, restarted=1, deferred=1) + + +class LoopWatchingManifestStore: + """A manifest store that notes whether each read ran on an event loop.""" + + def __init__(self, store: dict[str, object]) -> None: + self.store = store + self.read_on_event_loop: list[bool] = [] + + def get(self, key: str, default: object = None) -> object: + try: + asyncio.get_running_loop() + except RuntimeError: + self.read_on_event_loop.append(False) + else: + self.read_on_event_loop.append(True) + return self.store.get(key, default) + + def __setitem__(self, key: str, value: object) -> None: + raise AssertionError(key) + + +def test_a_runs_call_is_rebuilt_off_the_event_loop() -> None: + """Modal's blocking reads, such as the manifest's, never hold up the loop.""" + stored, _call = _submitted_report() + manifest = LoopWatchingManifestStore(MANIFEST_STORE) + rig = _rig( + restarts=FakeRestartStore( + unstarted=(_unstarted(RestartKind.REPORT_RUN, stored.report_run_id),), + report_runs={stored.report_run_id: stored}, + ), + manifest=manifest, + ) + + assert rig.restart() == RestartSummary(listed=1, restarted=1) + assert manifest.read_on_event_loop + assert not any(manifest.read_on_event_loop) + + +def test_an_unexpected_rebuild_error_defers_its_run_and_the_check_goes_on( + caplog, +) -> None: + """One row the check can't rebuild for an unforeseen reason holds up no other. + + The error is logged with its traceback, so it can be diagnosed. + """ + broken_run_id = UUID(int=9) + stored, _call = _submitted_report() + rig = _rig( + restarts=FakeRestartStore( + unstarted=( + _unstarted(RestartKind.REPORT_RUN, broken_run_id), + _unstarted(RestartKind.REPORT_RUN, stored.report_run_id), + ), + report_runs={stored.report_run_id: stored}, + read_failures={broken_run_id: RuntimeError("unforeseen")}, + ) + ) + + summary = rig.restart() + + assert rig.telemetry.outcomes[0] == ( + RestartKind.REPORT_RUN, + broken_run_id, + RestartOutcome.DEFERRED, + "RuntimeError", + ) + assert [run.run_id for run in rig.restarts.marked] == [stored.report_run_id] + assert summary == RestartSummary(listed=2, restarted=1, deferred=1) + (logged,) = [ + record + for record in caplog.records + if record.getMessage() == "v2_run_restart_unexpected_error" + ] + assert logged.exc_info is not None + assert isinstance(logged.exc_info[1], RuntimeError) + # The observability runtime's log handler writes only these attributes. + assert logged.policyengine_attributes == { + "run_kind": "report_run", + "run_id": str(broken_run_id), + } diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index acf16e000..fe41ed136 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -154,8 +154,9 @@ with its own mechanism, separate from the v1 routing jobs above: while the stored manifest is still what it wrote (`--restore-plan`). In a deployment from `main`, the same job then deploys the unversioned stuck-run sweep (`src/modal/run_sweep_app.py`), which fails the report - runs, simulation runs and household jobs nothing will finish, and whose - schedule starts as soon as it deploys; a failed sweep deploy, including a failed build of its + runs, simulation runs and household jobs nothing will finish, with the + restart check, which starts again those whose Modal call never started, + and whose schedules start as soon as it deploys; a failed sweep deploy, including a failed build of its image, fails the job and so puts the previous default back. A deployment from another branch leaves the sweep as it is. `deployment_ready` requires this job to succeed whenever Stage 12 was diff --git a/projects/policyengine-simulation-executor/pyproject.toml b/projects/policyengine-simulation-executor/pyproject.toml index 2b4f42ac6..cdb6f32e3 100644 --- a/projects/policyengine-simulation-executor/pyproject.toml +++ b/projects/policyengine-simulation-executor/pyproject.toml @@ -71,8 +71,15 @@ modal-simulation-image = [ ] # The executor's budget-window scheduler tests exercise the gateway app # end-to-end; the gateway project rides in as a dev-only path dependency. +# The restart check in the stuck-run sweep's Modal app runs the simulation +# service's own restart code: ``modal deploy`` adds the Simulation Entrypoint +# and its persistence library to that app's image as local sources, so the +# deploy job's ``uv sync`` installs them here. Nothing the executor's own +# Docker image runs imports them, and it is built without this group. dev = [ "policyengine-simulation-gateway", + "policyengine-simulation-entry", + "policyengine-simulation-persistence", ] [tool.uv.sources] @@ -80,6 +87,8 @@ policyengine-fastapi = { path = "../../libs/policyengine-fastapi", editable = tr policyengine-simulation-observability = { path = "../../libs/policyengine-simulation-observability", editable = true } policyengine-simulation-contract = { path = "../../libs/policyengine-simulation-contract", editable = true } policyengine-stage12-persistence = { path = "../../libs/policyengine-stage12-persistence", editable = true } +policyengine-simulation-entry = { path = "../policyengine-simulation-entry", editable = true } +policyengine-simulation-persistence = { path = "../../libs/policyengine-simulation-persistence", editable = true } policyengine-simulation-gateway = { path = "../policyengine-simulation-gateway", editable = true } [project.optional-dependencies] 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 bb566ea51..b0e346f63 100644 --- a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -1,21 +1,33 @@ -"""A scheduled Modal function that fails runs and household jobs nothing will finish. - -Every ten minutes it lists the Public API's report runs, simulation runs and -household jobs that have been pending or running for at least ten minutes, -counted from a run's first claim or, while it is pending, from its creation -(a budget window's year from when its call was recorded), and from a -household job's creation. It asks Modal for the status of each -recorded call and marks a run or job failed once its call has ended, and -fails a pending one that has no call recorded, because nothing will run it. -It also fails any still unfinished 24 hours after it was created or started, -longer than any Modal call of it can run, whatever Modal reports. A failed -year of a budget window fails its window, and a window no year has -advanced for ten minutes is finished from its years, or failed as stalled. +"""Scheduled Modal functions that fail or restart runs and household jobs. + +The stuck-run sweep runs every ten minutes. It lists the Public API's report +runs, simulation runs and household jobs that have been pending or running +for at least 15 minutes, counted from a run's first claim or, while it is +pending, from its creation (a budget window's year from when its call was +recorded), and from a household job's creation. It asks Modal for the status +of each recorded call and marks a run or job failed once its call has ended, +and fails a pending one that has no call recorded, because nothing will run +it. It also fails any still unfinished 24 hours after it was created or +started, longer than any Modal call of it can run, whatever Modal reports. A +failed year of a budget window fails its window, and a window no year has +advanced for 15 minutes is finished from its years, or failed as stalled. Its result counts runs and household jobs together, and the windows it -finished. It connects as the Modal worker -role, from the same Modal secret the v2 coordinator, workers and household -functions use. This application is not versioned with the .py bundle: one -copy runs per Modal environment. +finished. + +The restart check runs every two minutes. It starts once more the call of +each run or household job created 3 to 10 minutes ago that is still pending +with no call recorded, as when the simulation service stopped between +creating it and starting its call; the simulation service's own restart code +(``policyengine_simulation_entry.run_restart_check``) rebuilds the call from +the stored rows and records it. A run it can't start again is failed with +``run_not_started``, as the sweep would fail it at 15 minutes. Its telemetry +counts each restart and how it ended. + +Both connect as the Modal worker role, from the same Modal secret the v2 +coordinator, workers and household functions use. This application is not +versioned with the .py bundle: one copy runs per Modal environment. Its +image bakes in, at deployment, the Modal environment the restart check starts +calls in and the v2 workers' telemetry settings. Each deployment from ``main`` deploys it to its Modal environment after every check, in ``publish_stage12_v2_default``, once that job has moved the @@ -27,6 +39,7 @@ from __future__ import annotations +import os from dataclasses import asdict from pathlib import Path @@ -34,11 +47,17 @@ MODAL_WORKER_DATABASE_SECRET_NAME, MODAL_WORKER_DATABASE_URL_ENV, ) +from policyengine_simulation_observability.observability import ( + modal_image_environment, +) from src.modal.ca_certificates import GRPC_CA_BUNDLE_ENV, OPERATING_SYSTEM_CA_PACKAGE import modal APP_NAME = "policyengine-simulation-run-sweep" +# The Modal environment the restart check reads the v2 manifest from and +# starts calls in, baked into the image from the deploy step's environment. +V2_MODAL_ENVIRONMENT_ENV = "SIMULATION_V2_MODAL_ENVIRONMENT" _UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." app = modal.App(APP_NAME) @@ -52,6 +71,12 @@ # bundle's data, which the sweep never reads. .apt_install(OPERATING_SYSTEM_CA_PACKAGE) .env(GRPC_CA_BUNDLE_ENV) + .env( + { + **modal_image_environment(), + V2_MODAL_ENVIRONMENT_ENV: os.getenv(V2_MODAL_ENVIRONMENT_ENV, ""), + } + ) .uv_sync( uv_project_dir=_UV_PROJECT_DIR, frozen=True, @@ -66,6 +91,9 @@ "policyengine_simulation_observability", "policyengine_simulation_contract", "policyengine_stage12_persistence", + # The restart check runs the simulation service's restart code. + "policyengine_simulation_persistence", + "policyengine_simulation_entry", copy=True, ) ) @@ -93,3 +121,39 @@ def sweep_stuck_runs() -> dict: windows=report_run_access.budget_window_store(), ) return asdict(summary) + + +@app.function( + image=sweep_image, + schedule=modal.Cron("*/2 * * * *"), + cpu=0.25, + memory=512, + timeout=300, + retries=0, + max_containers=1, + secrets=[modal_worker_database_secret], +) +def restart_unstarted_runs() -> dict: + """Start again the runs and household jobs whose call never started; count them.""" + from policyengine_simulation_entry import ( # noqa: PLC0415 # Modal function body: loaded in the container, not at deploy time + run_restart_check, + ) + from policyengine_simulation_observability import ( # noqa: PLC0415 # Modal function body: loaded in the container, not at deploy time + observability, + ) + + modal_environment = os.environ[V2_MODAL_ENVIRONMENT_ENV] + runtime = observability.init_process_observability( + service_name="policyengine-simulation-run-restart", + service_role="v2_run_restart_check", + platform="modal", + environment=modal_environment, + ) + try: + return run_restart_check.run_restart_check( + database_url=os.environ[MODAL_WORKER_DATABASE_URL_ENV], + modal_environment=modal_environment, + runtime=runtime, + ) + finally: + runtime.shutdown() 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 26fc6dad1..5da02e9c6 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 @@ -6,14 +6,16 @@ call ends without recording the run's outcome, for example because its container ran out of memory, the run stays pending or running. The sweep asks Modal for the status of each recorded call and marks the run failed -once the call has ended. A pending run with no call recorded soon after its -creation is failed without asking Modal: usually the simulation service -stopped before starting its call. If it started the call but failed to -record it, that call's claim is refused once the run has failed, so it -finishes nothing. The sweep also fails a run still unfinished long after it -was created or started, longer than any Modal call of it can run, whatever -Modal reports and with or without a recorded call, so a run Modal can't -decide doesn't stay unfinished for good. +once the call has ended. A pending run with no call recorded 15 minutes +after its creation is failed without asking Modal: usually the simulation +service stopped before starting its call, and the restart check, which +starts such a run's call once more 3 to 10 minutes after its creation, could +not start it either. If a call was started but not recorded, its claim is +refused once the run has failed, so it finishes nothing. The sweep also +fails a run still unfinished long after it was created or started, longer +than any Modal call of it can run, whatever Modal reports and with or +without a recorded call, so a run Modal can't decide doesn't stay +unfinished for good. Modal restarts a pre-empted call under the same call ID and reports it pending meanwhile, so a pre-empted coordinator's or worker's run is never @@ -62,8 +64,10 @@ # A run must have been unfinished at least this long, counted from its first # claim or, while it is pending, from its creation (a household job's from its # creation), before the sweep judges it; a pending run with no call recorded -# by then is failed. -DEFAULT_MINIMUM_AGE = timedelta(minutes=10) +# by then is failed. It is longer than the restart check's window (3 to 10 +# minutes after creation), so a run that check started again is judged by +# its call. +DEFAULT_MINIMUM_AGE = timedelta(minutes=15) ENDED_CALL_MESSAGE = ( "modal_call_ended: The Modal call recorded for this run ended without " "recording the run's outcome." diff --git a/projects/policyengine-simulation-executor/tests/test_executor_image_dependencies.py b/projects/policyengine-simulation-executor/tests/test_executor_image_dependencies.py new file mode 100644 index 000000000..64bbc3e18 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_executor_image_dependencies.py @@ -0,0 +1,51 @@ +"""The executor's Docker image can install every dependency it declares. + +The image's build context is the repository root, and its Dockerfile copies +only some directories before ``uv sync --frozen --no-dev``. A main dependency +installed from a local path outside those directories can't be found, and +the image's build fails. A development group's local packages, such as the +gateway or the Simulation Entrypoint, are not installed there. +""" + +from __future__ import annotations + +import re +import tomllib +from pathlib import Path + +PROJECT = Path(__file__).resolve().parents[1] +REPOSITORY = PROJECT.parents[1] +DEPENDENCY_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*") + + +def _copied_directories() -> list[Path]: + """Return each repository directory the Dockerfile copies into the image.""" + copied = [] + for dockerfile_line in (PROJECT / "Dockerfile").read_text().splitlines(): + words = dockerfile_line.split() + if len(words) == 3 and words[0] == "COPY" and not words[1].startswith("--"): + copied.append((REPOSITORY / words[1]).resolve()) + return copied + + +def test_every_main_dependency_from_a_local_path_is_in_the_images_build() -> None: + pyproject = tomllib.loads((PROJECT / "pyproject.toml").read_text()) + sources = pyproject["tool"]["uv"]["sources"] + main_dependencies = { + DEPENDENCY_NAME.match(requirement).group(0) + for requirement in pyproject["project"]["dependencies"] + } + local_main_dependencies = { + name: (PROJECT / source["path"]).resolve() + for name, source in sources.items() + if "path" in source and name in main_dependencies + } + copied = _copied_directories() + + assert local_main_dependencies + outside_the_build = sorted( + name + for name, path in local_main_dependencies.items() + if not any(path.is_relative_to(directory) for directory in copied) + ) + assert outside_the_build == [] diff --git a/projects/policyengine-simulation-executor/tests/test_run_restart_image.py b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py new file mode 100644 index 000000000..5273226f3 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py @@ -0,0 +1,205 @@ +"""The restart check imports with only the packages the sweep's image installs. + +The sweep's Modal image installs the executor's locked +``modal-simulation-image`` group alone, and Modal adds its own client; the +restart check's code comes in as local sources. The scheduled function +imports that code only when it runs, so ``modal deploy`` would not notice a +package the image lacks. Here the import runs in a process that refuses the +top-level modules only other installed distributions provide, then every +module it loaded must come from the standard library, a local source, or a +distribution the image installs: a module under a namespace that an image +distribution shares, or from an editable install, is caught there. +""" + +from __future__ import annotations + +import json +import re +import site +import subprocess +import sys +import sysconfig +import tomllib +from importlib import metadata +from pathlib import Path + +PROJECT = Path(__file__).resolve().parents[1] +REPOSITORY = PROJECT.parents[1] +IMAGE_GROUP = "modal-simulation-image" +# Modal installs its client into every image it builds. +MODAL_CLIENT = "modal" +# The local sources the sweep's image adds, which the lock doesn't install. +LOCAL_SOURCES = frozenset( + { + "policyengine-simulation-executor", + "policyengine-simulation-observability", + "policyengine-simulation-contract", + "policyengine-stage12-persistence", + "policyengine-simulation-persistence", + "policyengine-simulation-entry", + } +) +# Where each local source's package lies in the repository. +LOCAL_SOURCE_DIRECTORIES = ( + PROJECT / "src" / "policyengine_simulation_executor", + REPOSITORY + / "projects/policyengine-simulation-entry/src/policyengine_simulation_entry", + *( + REPOSITORY / "libs" / library / "src" / library.replace("-", "_") + for library in ( + "policyengine-simulation-observability", + "policyengine-simulation-contract", + "policyengine-stage12-persistence", + "policyengine-simulation-persistence", + ) + ), +) +# The script's own module, which multiprocessing, once imported, also names +# ``__mp_main__``; it is not the restart check's. +NOT_THE_RESTART_CHECK = ("__mp_main__",) +# What the scheduled function's body imports. +RESTART_MODULES = ( + "policyengine_simulation_entry.run_restart_check", + "policyengine_simulation_observability.observability", +) + + +def _normalized(name: str) -> str: + return re.sub(r"[-_.]+", "-", name).lower() + + +def _image_distributions() -> set[str]: + """Return the locked image group's distributions, with their dependencies.""" + lock = tomllib.loads((PROJECT / "uv.lock").read_text()) + packages = {_normalized(package["name"]): package for package in lock["package"]} + project = packages["policyengine-simulation-executor"] + wanted = [ + (_normalized(dependency["name"]), tuple(dependency.get("extra", ()))) + for dependency in project["dev-dependencies"][IMAGE_GROUP] + ] + wanted.append((MODAL_CLIENT, ())) + installed: set[str] = set() + seen: set[tuple[str, tuple[str, ...]]] = set() + while wanted: + name, extras = wanted.pop() + if (name, extras) in seen or name not in packages: + continue + seen.add((name, extras)) + installed.add(name) + package = packages[name] + dependencies = list(package.get("dependencies", ())) + for extra in extras: + dependencies.extend(package.get("optional-dependencies", {}).get(extra, ())) + wanted.extend( + (_normalized(dependency["name"]), tuple(dependency.get("extra", ()))) + for dependency in dependencies + ) + return installed + + +def _top_level_modules(distribution: metadata.Distribution) -> set[str]: + """Return the importable top-level names a distribution installs.""" + listed = distribution.read_text("top_level.txt") + if listed: + return {name.strip() for name in listed.splitlines() if name.strip()} + modules = set() + for installed_file in distribution.files or (): + first = installed_file.parts[0] + if first.endswith((".dist-info", ".data")) or first in {"..", "__pycache__"}: + continue + if len(installed_file.parts) == 1 and installed_file.suffix != ".py": + continue + modules.add(first.removesuffix(".py")) + return modules + + +def _image_files() -> set[Path]: + """Return every file the image's distributions install here.""" + allowed = _image_distributions() + return { + Path(str(distribution.locate_file(installed_file))).resolve() + for distribution in metadata.distributions() + if _normalized(distribution.metadata["Name"]) in allowed + for installed_file in distribution.files or () + } + + +def _from_the_image(module_file: Path, image_files: set[Path]) -> bool: + """Return whether a loaded module's file is one the sweep's image has.""" + site_packages = [Path(path).resolve() for path in site.getsitepackages()] + standard_library = [ + Path(sysconfig.get_paths()[name]).resolve() for name in ("stdlib", "platstdlib") + ] + if module_file in image_files: + return True + if any(module_file.is_relative_to(source) for source in LOCAL_SOURCE_DIRECTORIES): + return True + return any( + module_file.is_relative_to(library) for library in standard_library + ) and not any(module_file.is_relative_to(packages) for packages in site_packages) + + +def _refused_modules() -> set[str]: + """Return the modules only distributions outside the image install.""" + allowed = _image_distributions() | LOCAL_SOURCES + allowed_modules: set[str] = set() + other_modules: set[str] = set() + for distribution in metadata.distributions(): + modules = _top_level_modules(distribution) + if _normalized(distribution.metadata["Name"]) in allowed: + allowed_modules |= modules + else: + other_modules |= modules + return other_modules - allowed_modules + + +def test_the_restart_check_imports_with_only_the_sweep_images_packages( + tmp_path, +) -> None: + refused = _refused_modules() + # The executor's own environment holds packages the image doesn't. + assert {"jwt", "uvicorn"} <= refused + script = tmp_path / "import_restart_check.py" + loaded_list = tmp_path / "loaded_module_files.json" + script.write_text( + "import importlib, sys\n" + # Only what the import adds is checked, not what started the process. + "loaded_before = set(sys.modules)\n" + f"refused = {sorted(refused)!r}\n" + "class Refuse:\n" + " def find_spec(self, name, path=None, target=None):\n" + " if name.split('.')[0] in refused:\n" + " raise ImportError(f'not in the sweep image: {name}')\n" + "sys.meta_path.insert(0, Refuse())\n" + f"for module in {list(RESTART_MODULES)!r}:\n" + " importlib.import_module(module)\n" + "import json\n" + "loaded = sorted({getattr(module, '__file__', None) or ''\n" + " for name, module in list(sys.modules.items())\n" + " if name not in loaded_before\n" + f" and name not in {list(NOT_THE_RESTART_CHECK)!r}}}\n" + " - {''})\n" + f"open({str(loaded_list)!r}, 'w').write(json.dumps(loaded))\n", + encoding="utf-8", + ) + + completed = subprocess.run( # the test's own interpreter and script + [sys.executable, str(script)], + capture_output=True, + text=True, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + image_files = _image_files() + loaded = [Path(name).resolve() for name in json.loads(loaded_list.read_text())] + assert any( + module_file.is_relative_to(LOCAL_SOURCE_DIRECTORIES[1]) + for module_file in loaded + ) + outside_the_image = [ + str(module_file) + for module_file in loaded + if not _from_the_image(module_file, image_files) + ] + assert outside_the_image == [] 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 f72b380b8..8ec46da1f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py @@ -8,8 +8,14 @@ from uuid import uuid4 import pytest +from policyengine_simulation_entry.run_restarts import ( + CALL_START_TIMEOUT_SECONDS, + CHECK_TIME_LIMIT, + RESTART_UNTIL, +) from policyengine_simulation_executor.stage12_runtime.run_sweep import ( DEFAULT_MAXIMUM_AGE, + DEFAULT_MINIMUM_AGE, CallState, SweepSummary, modal_call_state, @@ -144,6 +150,26 @@ def test_the_sweep_lists_runs_older_than_its_minimum_age() -> None: assert store.cutoffs == [NOW - timedelta(minutes=15)] * 3 +def test_by_default_unstarted_runs_wait_until_the_restart_check_has_passed() -> None: + """The restart check starts a pending run with no call 3 to 10 minutes old. + + The sweep fails such a run only once it is 15 minutes old, so a run whose + restart started is judged by its call, and only a second miss is failed. + """ + store = FakeSweepStore([]) + + sweep_stuck_runs(store, call_state=lambda _: CallState.ENDED, now=NOW) + + assert store.cutoffs == [NOW - timedelta(minutes=15)] * 3 + # A row listed at the end of the window may be marked at the end of the + # check's time limit; its call starts and is recorded with a minute to + # spare before the sweep judges it. + latest_restart = ( + RESTART_UNTIL + CHECK_TIME_LIMIT + timedelta(seconds=CALL_START_TIMEOUT_SECONDS) + ) + assert latest_restart + timedelta(minutes=1) <= DEFAULT_MINIMUM_AGE + + def test_report_runs_are_swept_before_simulation_runs_are_listed() -> None: # A simulation run is listed only once no report run naming it is # running, so failing a dead coordinator's report run first lets the same 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 db32d645a..3580ab8bc 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 @@ -8,10 +8,19 @@ import yaml from fixtures.fake_modal import install_fake_modal +from policyengine_simulation_entry import run_restart_check +from policyengine_simulation_entry.run_restart_check import ( + RESTART_STATEMENT_TIMEOUT_MILLISECONDS, +) +from policyengine_simulation_entry.run_restarts import ( + CALL_START_TIMEOUT_SECONDS, + CHECK_TIME_LIMIT, +) from policyengine_simulation_executor.stage12_runtime import ( report_run_access as report_run_access_module, ) from policyengine_simulation_executor.stage12_runtime import run_sweep +from policyengine_simulation_observability import observability REUSABLE_DEPLOY_WORKFLOW = ( Path(__file__).resolve().parents[3] @@ -27,6 +36,23 @@ def _load(monkeypatch): return importlib.import_module("src.modal.run_sweep_app") +def _deploy_steps(job: dict, app_file: str) -> list[dict]: + return [step for step in job.get("steps", []) if app_file in step.get("run", "")] + + +def _telemetry_settings(environment: dict[str, str]) -> dict[str, str]: + """Return the deploy-time settings ``modal_image_environment`` bakes in.""" + return { + name: value + for name, value in environment.items() + if name.startswith(("OBSERVABILITY_", "OTEL_", "POLICYENGINE_OTEL_")) + } + + +def _functions(module) -> dict[str, dict]: + return dict(module.app.function_calls) + + def _sweep_deploy_steps(job: dict) -> list[dict]: return [ step @@ -65,11 +91,18 @@ def test_the_sweep_deploys_with_the_stage12_default_after_every_check() -> None: assert ( sweep_step["working-directory"] == "projects/policyengine-simulation-executor" ) - # Only Modal's tokens: the database URL reaches the scheduled function from - # its Modal secret, which must exist before this step can deploy it. + # Modal's tokens, the Modal environment the restart check starts calls in, + # and the v2 workers' own telemetry settings, which the image bakes in. + # The database URL reaches the scheduled functions from their Modal + # secret, which must exist before this step can deploy them. + (v2_workers_step,) = _deploy_steps(jobs["deploy_stage12_v2"], "src/modal/v2_app.py") + telemetry = _telemetry_settings(v2_workers_step["env"]) + assert len(telemetry) == 8 assert sweep_step["env"] == { "MODAL_TOKEN_ID": "${{ secrets.MODAL_TOKEN_ID }}", "MODAL_TOKEN_SECRET": "${{ secrets.MODAL_TOKEN_SECRET }}", + "SIMULATION_V2_MODAL_ENVIRONMENT": "${{ inputs.modal_environment }}", + **telemetry, } # Only from main. With no status function in it, the condition also # requires every earlier step to have passed (an implicit success()). @@ -85,10 +118,9 @@ def test_the_sweep_runs_every_ten_minutes_with_the_worker_database_secret( monkeypatch, ) -> None: module = _load(monkeypatch) - ((name, options),) = module.app.function_calls + options = _functions(module)["sweep_stuck_runs"] assert module.app.name == "policyengine-simulation-run-sweep" - assert name == "sweep_stuck_runs" assert options["schedule"] == {"schedule": "*/10 * * * *"} assert options["secrets"] == [ { @@ -110,13 +142,16 @@ 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. + # The app module imports src.modal at module level, in the container too; + # the restart check runs the simulation service's own restart code. assert calls["add_local_python_source"][0] == ( "src.modal", "policyengine_simulation_executor", "policyengine_simulation_observability", "policyengine_simulation_contract", "policyengine_stage12_persistence", + "policyengine_simulation_persistence", + "policyengine_simulation_entry", ) @@ -157,3 +192,112 @@ def recording_sweep(received_store, *, windows): "windows_succeeded": 6, } assert swept == [(store, window_store)] + + +def test_the_restart_check_runs_every_two_minutes_with_the_worker_database_secret( + monkeypatch, +) -> None: + """It restarts runs 3 to 10 minutes old, so each is seen several times.""" + module = _load(monkeypatch) + options = _functions(module)["restart_unstarted_runs"] + + assert options["schedule"] == {"schedule": "*/2 * * * *"} + assert options["secrets"] == [ + { + "args": ("policyengine-modal-worker-database",), + "kwargs": {"required_keys": ["MODAL_WORKER_DATABASE_URL"]}, + } + ] + # One check at a time; two at once would still restart each run once. + assert options["max_containers"] == 1 + assert options["retries"] == 0 + assert options["image"] is module.sweep_image + + +def test_a_check_that_waits_on_every_call_still_ends_within_its_timeout( + monkeypatch, +) -> None: + """A row marked just before the timeout stops the check is never restarted. + + The check marks no row once its time limit has passed; after a row is + marked, it waits at most for the call's start and for the mark, record or + failure statements, each bounded by the statement timeout. + """ + module = _load(monkeypatch) + options = _functions(module)["restart_unstarted_runs"] + statement_seconds = RESTART_STATEMENT_TIMEOUT_MILLISECONDS / 1000 + longest_after_the_limit = CALL_START_TIMEOUT_SECONDS + 3 * statement_seconds + + # Half a minute more for connecting and Modal's own start. + assert ( + CHECK_TIME_LIMIT.total_seconds() + longest_after_the_limit + 30 + <= options["timeout"] + ) + + +def test_the_image_bakes_the_modal_environment_and_telemetry_settings_in( + monkeypatch, +) -> None: + monkeypatch.setenv("SIMULATION_V2_MODAL_ENVIRONMENT", "staging") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "https://collector.example") + monkeypatch.setenv("OBSERVABILITY_LOG_NAME", "policyengine-api-v1-modal") + module = _load(monkeypatch) + baked: dict[str, str] = {} + for call in module.sweep_image.calls: + if call[0] == "env": + baked.update(call[1]) + + assert baked["SIMULATION_V2_MODAL_ENVIRONMENT"] == "staging" + assert baked["OTEL_EXPORTER_OTLP_ENDPOINT"] == "https://collector.example" + assert baked["OBSERVABILITY_LOG_NAME"] == "policyengine-api-v1-modal" + assert baked["OTEL_METRICS_EXPORTER"] == "otlp" + + +class RecordingRuntime: + """Stands in for the observability runtime; records its shutdown.""" + + def __init__(self) -> None: + self.shut_down = False + + def shutdown(self) -> None: + self.shut_down = True + + +def test_the_scheduled_restart_runs_the_check_on_the_worker_database( + monkeypatch, +) -> None: + module = _load(monkeypatch) + monkeypatch.setenv("MODAL_WORKER_DATABASE_URL", "postgresql://worker@db/postgres") + monkeypatch.setenv("SIMULATION_V2_MODAL_ENVIRONMENT", "main") + runtime = RecordingRuntime() + built = [] + checked = [] + + def recording_runtime(**identity): + built.append(identity) + return runtime + + def recording_check(**arguments): + checked.append(arguments) + return {"listed": 2, "restarted": 2} + + monkeypatch.setattr(observability, "init_process_observability", recording_runtime) + monkeypatch.setattr(run_restart_check, "run_restart_check", recording_check) + + assert module.restart_unstarted_runs() == {"listed": 2, "restarted": 2} + assert checked == [ + { + "database_url": "postgresql://worker@db/postgres", + "modal_environment": "main", + "runtime": runtime, + } + ] + assert built == [ + { + "service_name": "policyengine-simulation-run-restart", + "service_role": "v2_run_restart_check", + "platform": "modal", + "environment": "main", + } + ] + assert runtime.shut_down diff --git a/projects/policyengine-simulation-executor/uv.lock b/projects/policyengine-simulation-executor/uv.lock index b46aa9520..2fab1f614 100644 --- a/projects/policyengine-simulation-executor/uv.lock +++ b/projects/policyengine-simulation-executor/uv.lock @@ -1855,6 +1855,9 @@ google = [ { name = "google-auth" }, { name = "google-cloud-logging" }, ] +httpx = [ + { name = "httpx" }, +] otlp-grpc = [ { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp-proto-grpc" }, @@ -1891,6 +1894,46 @@ requires-dist = [ ] provides-extras = ["modal", "test", "build"] +[[package]] +name = "policyengine-simulation-entry" +version = "0.1.0" +source = { editable = "../policyengine-simulation-entry" } +dependencies = [ + { name = "cryptography" }, + { name = "fastapi" }, + { name = "httpx" }, + { name = "modal" }, + { name = "policyengine-observability", extra = ["fastapi", "google", "httpx", "otlp-grpc"] }, + { name = "policyengine-simulation-contract" }, + { name = "policyengine-simulation-observability" }, + { name = "policyengine-simulation-persistence" }, + { name = "policyengine-stage12-persistence" }, + { name = "pyjwt" }, + { name = "uvicorn" }, +] + +[package.metadata] +requires-dist = [ + { name = "cryptography", specifier = ">=41.0.0" }, + { name = "fastapi", specifier = ">=0.115.0,<1" }, + { name = "httpx", specifier = ">=0.28.0,<1" }, + { name = "modal", specifier = ">=1.4,<2" }, + { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, + { name = "policyengine-observability", extras = ["fastapi", "google", "httpx", "otlp-grpc"], specifier = ">=3.0.7,<4" }, + { name = "policyengine-simulation-contract", editable = "../../libs/policyengine-simulation-contract" }, + { name = "policyengine-simulation-observability", editable = "../../libs/policyengine-simulation-observability" }, + { name = "policyengine-simulation-persistence", editable = "../../libs/policyengine-simulation-persistence" }, + { name = "policyengine-stage12-persistence", editable = "../../libs/policyengine-stage12-persistence" }, + { name = "pyjwt", specifier = ">=2.10.1,<3" }, + { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, + { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, + { name = "pytest-asyncio", marker = "extra == 'test'", specifier = ">=0.25.3" }, + { name = "pytest-cov", marker = "extra == 'test'", specifier = ">=6.1.1" }, + { name = "pyyaml", marker = "extra == 'test'", specifier = ">=6.0.2" }, + { name = "uvicorn", specifier = ">=0.35.0,<1" }, +] +provides-extras = ["test", "build"] + [[package]] name = "policyengine-simulation-executor" version = "0.1.0" @@ -1928,7 +1971,9 @@ test = [ [package.dev-dependencies] dev = [ + { name = "policyengine-simulation-entry" }, { name = "policyengine-simulation-gateway" }, + { name = "policyengine-simulation-persistence" }, ] modal-simulation-image = [ { name = "fastapi" }, @@ -1971,7 +2016,11 @@ requires-dist = [ provides-extras = ["test", "build"] [package.metadata.requires-dev] -dev = [{ name = "policyengine-simulation-gateway", editable = "../policyengine-simulation-gateway" }] +dev = [ + { name = "policyengine-simulation-entry", editable = "../policyengine-simulation-entry" }, + { name = "policyengine-simulation-gateway", editable = "../policyengine-simulation-gateway" }, + { name = "policyengine-simulation-persistence", editable = "../../libs/policyengine-simulation-persistence" }, +] modal-simulation-image = [ { name = "fastapi", specifier = ">=0.115.0" }, { name = "google-cloud-storage", specifier = ">=2" }, @@ -2049,6 +2098,26 @@ requires-dist = [ ] provides-extras = ["test", "build"] +[[package]] +name = "policyengine-simulation-persistence" +version = "0.1.0" +source = { editable = "../../libs/policyengine-simulation-persistence" } +dependencies = [ + { name = "psycopg", extra = ["binary"] }, + { name = "sqlalchemy" }, +] + +[package.metadata] +requires-dist = [ + { name = "psycopg", extras = ["binary"], specifier = ">=3.2,<4" }, + { name = "pyright", marker = "extra == 'build'", specifier = ">=1.1.401" }, + { name = "pytest", marker = "extra == 'test'", specifier = ">=8.3.4" }, + { name = "pytest-cov", marker = "extra == 'test'", specifier = ">=6.1.1" }, + { name = "pyyaml", marker = "extra == 'test'", specifier = ">=6.0.2" }, + { name = "sqlalchemy", specifier = ">=2,<2.1" }, +] +provides-extras = ["test", "build"] + [[package]] name = "policyengine-stage12-persistence" version = "0.1.0"