From 257e9be11df8a20496c3024e843a83a5c1a722c1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:17:27 +0400 Subject: [PATCH 01/19] Test the restart check's reads and writes for unstarted runs (failing) List the pending runs and household jobs with no Modal call recorded, mark one restart of each, and read a report run's or standalone simulation run's stored inputs back, as the Modal worker role. Co-Authored-By: Claude Opus 5.5 --- .../tests/conftest.py | 12 + .../tests/test_run_restart_records.py | 540 ++++++++++++++++++ 2 files changed, 552 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/tests/test_run_restart_records.py 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..fddb135a8 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py @@ -0,0 +1,540 @@ +"""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, + RunRestartStore, + SimulationChoice, + SimulationType, + StoredRunUnreadableError, + UnstartedRun, + budget_window_runs, + insert_economy_report, + insert_economy_simulation, + insert_household_job, + list_unstarted_runs, + mark_restart, + 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=12), + 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=12), + 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) + + +@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]) + + assert set(by_kind) == {RestartKind.REPORT_RUN, RestartKind.SIMULATION_RUN} + 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" + ) From 8ad78f6358ae7add849abe39a9975e3f5a99ad32 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:19:33 +0400 Subject: [PATCH 02/19] List, mark and read back runs whose Modal call never started The restart check lists the household jobs, report runs and standalone simulation runs still pending with no Modal call that nothing has changed since their insert, marks each one before restarting it so it is restarted at most once, and reads back the stored rows its call is built from, against the run's own catalog version. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 28 ++ .../run_restart_records.py | 419 ++++++++++++++++++ .../run_restart_store.py | 85 ++++ 3 files changed, 532 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_records.py create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_store.py 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 1741cf7b0..e1c2048a0 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -86,6 +86,21 @@ read_dynamic_reform, read_policy_reform, ) +from policyengine_simulation_persistence.run_restart_records import ( + RESTART_TABLES, + EconomyScope, + RestartKind, + StoredEconomySide, + StoredReportRun, + StoredRunUnreadableError, + StoredSimulationRun, + UnstartedRun, + list_unstarted_runs, + mark_restart, + read_unstarted_report_run, + read_unstarted_simulation_run, +) +from policyengine_simulation_persistence.run_restart_store import RunRestartStore from policyengine_simulation_persistence.tables import ( HouseholdJobStatus, ReportRunStatus, @@ -103,6 +118,7 @@ ) __all__ = [ + "RESTART_TABLES", "SERVICE_TABLE_PRIVILEGES", "SERVICE_UPDATABLE_COLUMNS", "CatalogDataset", @@ -115,6 +131,7 @@ "EconomyReportState", "EconomyRequest", "EconomyRunStore", + "EconomyScope", "EconomySimulationInputs", "EconomySimulationState", "HouseholdCalculationInputs", @@ -134,7 +151,9 @@ "ReportSide", "ReportSideError", "ReportSimulation", + "RestartKind", "ReusableRun", + "RunRestartStore", "RunReuse", "RunVersions", "SideRefusal", @@ -143,13 +162,18 @@ "SimulationRunState", "SimulationStatus", "SimulationType", + "StoredEconomySide", "StoredHousehold", + "StoredReportRun", + "StoredRunUnreadableError", + "StoredSimulationRun", "UnknownCatalogVersionError", "UnknownDatasetError", "UnknownDynamicError", "UnknownHouseholdError", "UnknownPolicyError", "UnknownRegionError", + "UnstartedRun", "bounded_transaction", "budget_window_runs", "combine_reforms", @@ -163,6 +187,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", @@ -177,6 +203,8 @@ "read_policy_reform", "read_region", "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/run_restart_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_records.py new file mode 100644 index 000000000..98f5f36bb --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_records.py @@ -0,0 +1,419 @@ +"""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: their own tasks restart them. +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.input_tables import ( + MODEL_NAME_PREFIX, + datasets, + regions, + tax_benefit_model_versions, + tax_benefit_models, +) +from policyengine_simulation_persistence.policy_reform import ( + combine_reforms, + read_dynamic_reform, + read_policy_reform, +) +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 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 _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) + policy_reform = ( + {} + if choice.policy_id is None + else read_policy_reform( + connection, + country=country, + policy_id=choice.policy_id, + catalog_version_id=scope.catalog_version_id, + ) + ) + dynamic_reform = ( + {} + if choice.dynamic_id is None + else read_dynamic_reform( + connection, + dynamic_id=choice.dynamic_id, + catalog_version_id=scope.catalog_version_id, + ) + ) + return StoredEconomySide( + simulation_run_id=side.id, + choice=choice, + reform=combine_reforms(policy_reform, dynamic_reform), + ) + + +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..f7005c662 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_restart_store.py @@ -0,0 +1,85 @@ +"""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_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 ( + 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) From a95c14b8ff72369544dac84b69d3649065a2edda Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:24:13 +0400 Subject: [PATCH 03/19] Test restarting runs and household jobs whose call never started (failing) The restart rebuilds each run's call from its stored rows and starts the same call its submission started, on the worker of the run's own .py version; a run that can't be rebuilt, or whose restart Modal refuses, is failed with run_not_started, and each listed run is one telemetry scope with its outcome. Co-Authored-By: Claude Opus 5.5 --- .../tests/economy_service_fakes.py | 3 +- .../tests/test_economy_workers.py | 30 + .../tests/test_household_workers.py | 29 + .../tests/test_run_restarts.py | 783 ++++++++++++++++++ 4 files changed, 844 insertions(+), 1 deletion(-) create mode 100644 projects/policyengine-simulation-entry/tests/test_run_restarts.py 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_restarts.py b/projects/policyengine-simulation-entry/tests/test_run_restarts.py new file mode 100644 index 000000000..726194210 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_run_restarts.py @@ -0,0 +1,783 @@ +"""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 ( + RESTART_AFTER, + RESTART_UNTIL, + RESTARTS_PER_CHECK, + RestartCallers, + RestartOutcome, + RestartSettings, + RestartStores, + RestartSummary, + RestartWorkers, + RunRestartService, +) +from policyengine_simulation_persistence import ( + EconomyScope, + HouseholdCalculationInputs, + HouseholdJobRecord, + HouseholdJobStatus, + RestartKind, + SimulationChoice, + StoredEconomySide, + StoredHousehold, + StoredReportRun, + StoredRunUnreadableError, + StoredSimulationRun, + UnstartedRun, +) +from sqlalchemy.exc import OperationalError + +if TYPE_CHECKING: + from collections.abc import 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) + read_failure: Exception | None = None + marked_elsewhere: bool = False + windows: list[tuple[datetime, datetime, int]] = field(default_factory=list) + marked: list[UnstartedRun] = field(default_factory=list) + + 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 + 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: dict[str, object] | 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 + 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, + 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_twelve_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=12), + ) == (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") From 7ede4eb542344da2fe762ce5ef55252b11e285ef Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:28:04 +0400 Subject: [PATCH 04/19] Start a run or household job again when its Modal call never started The restart service lists the rows created 3 to 12 minutes ago that are still pending with no call, rebuilds each call from its stored rows on the worker of the row's own .py version, marks the row, starts the call and records its ID. A row that can't be rebuilt, or whose restart Modal refuses, is failed with run_not_started; one whose rows or manifest can't be read now is left for the next check. Co-Authored-By: Claude Opus 5.5 --- .../economy_workers.py | 60 +- .../household_workers.py | 59 +- .../run_restarts.py | 789 ++++++++++++++++++ 3 files changed, 896 insertions(+), 12 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py 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_restarts.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py new file mode 100644 index 000000000..14b46f204 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -0,0 +1,789 @@ +"""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 12 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 is left unmarked, for the next check. 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 uuid +from collections import Counter +from contextlib import nullcontext +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from enum import StrEnum +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, Sequence + from contextlib import AbstractContextManager + from uuid import UUID + + from policyengine_simulation_persistence import ( + EconomyScope, + HouseholdCalculationInputs, + HouseholdJobRecord, + HouseholdJobStatus, + NewHouseholdJob, + ReportRunStatus, + RunVersions, + SimulationStatus, + StoredReportRun, + StoredSimulationRun, + UnstartedRun, + ) + + from policyengine_simulation_entry.economy_workers import ( + EconomyRunCaller, + EconomyWorker, + ) + from policyengine_simulation_entry.household_workers import ( + HouseholdFunctionCaller, + HouseholdWorker, + ) + +# 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, so a restart has time to record its call before the sweep judges it. +RESTART_UNTIL: Final = timedelta(minutes=12) +# At most this many rows are restarted by one check; the next takes the rest. +RESTARTS_PER_CHECK: Final = 50 +# How long the check waits for Modal to accept a restarted call. +CALL_START_TIMEOUT_SECONDS: Final = 10.0 +# 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"} +) +_NOT_STARTED: Final = ( + "run_not_started: The run stayed pending with no Modal call recorded for it" +) +_PENDING: Final = "pending" + + +class RestartOutcome(StrEnum): + """How the check ended for one listed row. + + ``restarted``: its new call started and the call's ID was recorded. + ``call_not_recorded``: its new call started, but the ID couldn't be + recorded; the call claims the row itself. ``started_meanwhile``: another + call claimed the row after the check started its own, which that call + then can't claim. ``left``: the row was gone, or changed before the check + started a call; also a row whose failure couldn't be written, which the + sweep fails later. ``deferred``: its rows or the manifest couldn't be + read, so it is left unmarked for the next check. ``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 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) + + +@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 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.""" + ... + + +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 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.""" + ... + + +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 + + +@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 _DeferredError(Exception): + """The row's rows or the manifest couldn't be read now.""" + + +class _CannotRebuildError(Exception): + """The row's call can't be built again; ``reason`` says why, in a sentence part. + + ``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: + 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) + + +class RunRestartService: + """Start the calls of the runs and household jobs whose call never started. + + The store reads and writes are synchronous: the check is a batch job of + its own, so nothing else waits on its event loop. + """ + + 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), + 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._workers = workers + self._callers = callers + self._settings = settings + self._telemetry = telemetry or NoRestartTelemetry() + self._clock = clock + self._call_start_timeout_seconds = call_start_timeout_seconds + self._new_request_id = new_request_id + + async def restart_unstarted_runs(self) -> RestartSummary: + """Start once more each row created 3 to 12 minutes ago with no call. + + Raise the store's error when the rows can't be listed. + """ + 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 run in runs: + 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()}, + ) + + 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 = self._rebuilt_call(run) + except _DeferredError as deferred: + return RestartOutcome.DEFERRED, str(deferred) + except _CannotRebuildError 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 + if rebuilt is None: + return RestartOutcome.LEFT, None + start_call, simulation_run_ids = rebuilt + return await self._start_marked(run, start_call, 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.""" + 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 another call claimed the + # row; one that couldn't be written leaves the row to the sweep. + outcome = { + True: RestartOutcome.FAILED, + False: RestartOutcome.STARTED_MEANWHILE, + None: RestartOutcome.LEFT, + }[failed] + return outcome, type(error).__name__ + return self._recorded(run, call_id) + + def _recorded( + self, run: UnstartedRun, call_id: str + ) -> tuple[RestartOutcome, str | None]: + """Record the started call's ID; return how the restart ended.""" + try: + status = self._record(run, call_id) + except SQLAlchemyError as error: + return RestartOutcome.CALL_NOT_RECORDED, type(error).__name__ + if status is None: + return RestartOutcome.LEFT, None + if status == _PENDING: + return RestartOutcome.RESTARTED, None + return RestartOutcome.STARTED_MEANWHILE, None + + def _rebuilt_call( + self, run: UnstartedRun + ) -> tuple[CallStart, tuple[UUID, ...]] | None: + """Return what starts the row's call again and its report's simulation runs. + + Return None when the row is gone. Raise _DeferredError when its rows or + the manifest can't be read now, and _CannotRebuildError when its call + can't be built again. + """ + simulation_run_ids: tuple[UUID, ...] = () + try: + if run.kind is RestartKind.REPORT_RUN: + stored_report = self._stores.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 self._report_call(stored_report), simulation_run_ids + if run.kind is RestartKind.SIMULATION_RUN: + stored_simulation = self._stores.restarts.read_simulation_run( + run.run_id + ) + if stored_simulation is None: + return None + return self._simulation_call(stored_simulation), () + job = self._stores.households.read(run.run_id) + if job is None: + return None + return self._household_call(job), () + except (SQLAlchemyError, VersionManifestUnavailableError) as error: + raise _DeferredError(type(error).__name__) from error + except _REBUILD_ERRORS as error: + raise _CannotRebuildError( + _rebuild_reason(error), + error_type=type(error).__name__, + simulation_run_ids=simulation_run_ids, + ) from error + + 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._stores.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 + + 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) -> str | None: + """Record the call's ID while the row is unstarted; return its status after.""" + if run.kind is RestartKind.REPORT_RUN: + status = self._stores.economy.record_report_call(run.run_id, call_id) + elif run.kind is RestartKind.SIMULATION_RUN: + status = self._stores.economy.record_simulation_call(run.run_id, call_id) + else: + status = self._stores.households.record_call(run.run_id, call_id) + return None if status is None else str(status) From fa2236a188471c190947b90d38c7321e2840eab8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:32:54 +0400 Subject: [PATCH 05/19] Build a restarted call in a module of its own The restart service keeps the check's loop and its writes; the call builder, its stores' protocols and the errors that defer or fail a restart move to run_restart_calls. Co-Authored-By: Claude Opus 5.5 --- .../run_restart_calls.py | 495 ++++++++++++++++++ .../run_restarts.py | 489 ++--------------- 2 files changed, 538 insertions(+), 446 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_calls.py 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..8c7d12b9d --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_calls.py @@ -0,0 +1,495 @@ +"""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, + 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.""" + ... + + +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: + raise CallNotRebuiltError( + _rebuild_reason(error), + error_type=type(error).__name__, + simulation_run_ids=simulation_run_ids, + ) from error + + 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_restarts.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py index 14b46f204..0dfe483d7 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -38,76 +38,51 @@ from dataclasses import dataclass from datetime import UTC, datetime, timedelta from enum import StrEnum -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 typing import TYPE_CHECKING, Final, Protocol + +from policyengine_simulation_persistence import RestartKind 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, +from policyengine_simulation_entry.run_restart_calls import ( + CallNotRebuiltError, + DeferredRestartError, + HouseholdJobRestarts, + RestartCallBuilder, + RestartCallers, + RestartSettings, + RestartWorkers, + UnstartedRunSource, ) if TYPE_CHECKING: - from collections.abc import Awaitable, Callable, Sequence + from collections.abc import Callable, Sequence from contextlib import AbstractContextManager from uuid import UUID from policyengine_simulation_persistence import ( - EconomyScope, - HouseholdCalculationInputs, - HouseholdJobRecord, - HouseholdJobStatus, - NewHouseholdJob, ReportRunStatus, - RunVersions, SimulationStatus, - StoredReportRun, - StoredSimulationRun, UnstartedRun, ) - from policyengine_simulation_entry.economy_workers import ( - EconomyRunCaller, - EconomyWorker, - ) - from policyengine_simulation_entry.household_workers import ( - HouseholdFunctionCaller, - HouseholdWorker, - ) + from policyengine_simulation_entry.run_restart_calls import CallStart + +__all__ = [ + "CALL_START_TIMEOUT_SECONDS", + "RESTARTS_PER_CHECK", + "RESTART_AFTER", + "RESTART_UNTIL", + "EconomyRunWrites", + "NoRestartTelemetry", + "RestartCallers", + "RestartOutcome", + "RestartSettings", + "RestartStores", + "RestartSummary", + "RestartTelemetry", + "RestartWorkers", + "RunRestartService", +] # 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. @@ -119,12 +94,6 @@ RESTARTS_PER_CHECK: Final = 50 # How long the check waits for Modal to accept a restarted call. CALL_START_TIMEOUT_SECONDS: Final = 10.0 -# 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"} -) _NOT_STARTED: Final = ( "run_not_started: The run stayed pending with no Modal call recorded for it" ) @@ -153,33 +122,6 @@ class RestartOutcome(StrEnum): FAILED = "failed" -@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) - - @dataclass(frozen=True) class RestartSummary: """How many rows one check listed, and how each ended (``RestartOutcome``).""" @@ -193,30 +135,6 @@ class RestartSummary: failed: int = 0 -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.""" - ... - - class EconomyRunWrites(Protocol): """The submission's own call writes on economy runs (``EconomyRunStore``).""" @@ -248,55 +166,6 @@ def fail_unstarted_simulation( ... -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.""" - ... - - class RestartTelemetry(Protocol): """Where the check reports each listed row (``ObservabilityRestartTelemetry``).""" @@ -356,120 +225,6 @@ class RestartStores: households: HouseholdJobRestarts -@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 _DeferredError(Exception): - """The row's rows or the manifest couldn't be read now.""" - - -class _CannotRebuildError(Exception): - """The row's call can't be built again; ``reason`` says why, in a sentence part. - - ``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: - 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) - - class RunRestartService: """Start the calls of the runs and household jobs whose call never started. @@ -491,13 +246,17 @@ def __init__( # noqa: PLR0913 # each seam tests replace ) -> None: """Use these seams; tests replace each of them.""" self._stores = stores - self._workers = workers - self._callers = callers - self._settings = settings + 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._call_start_timeout_seconds = call_start_timeout_seconds - self._new_request_id = new_request_id async def restart_unstarted_runs(self) -> RestartSummary: """Start once more each row created 3 to 12 minutes ago with no call. @@ -525,10 +284,10 @@ async def restart_unstarted_runs(self) -> RestartSummary: 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 = self._rebuilt_call(run) - except _DeferredError as deferred: + rebuilt = self._calls.build(run) + except DeferredRestartError as deferred: return RestartOutcome.DEFERRED, str(deferred) - except _CannotRebuildError as cannot: + except CallNotRebuiltError as cannot: message = ( f"{_NOT_STARTED}, and its call could not be built again: " f"{cannot.reason}." @@ -543,8 +302,7 @@ async def _restart(self, run: UnstartedRun) -> tuple[RestartOutcome, str | None] return outcome, cannot.error_type if rebuilt is None: return RestartOutcome.LEFT, None - start_call, simulation_run_ids = rebuilt - return await self._start_marked(run, start_call, simulation_run_ids) + return await self._start_marked(run, rebuilt.start, rebuilt.simulation_run_ids) async def _start_marked( self, @@ -595,167 +353,6 @@ def _recorded( return RestartOutcome.RESTARTED, None return RestartOutcome.STARTED_MEANWHILE, None - def _rebuilt_call( - self, run: UnstartedRun - ) -> tuple[CallStart, tuple[UUID, ...]] | None: - """Return what starts the row's call again and its report's simulation runs. - - Return None when the row is gone. Raise _DeferredError when its rows or - the manifest can't be read now, and _CannotRebuildError when its call - can't be built again. - """ - simulation_run_ids: tuple[UUID, ...] = () - try: - if run.kind is RestartKind.REPORT_RUN: - stored_report = self._stores.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 self._report_call(stored_report), simulation_run_ids - if run.kind is RestartKind.SIMULATION_RUN: - stored_simulation = self._stores.restarts.read_simulation_run( - run.run_id - ) - if stored_simulation is None: - return None - return self._simulation_call(stored_simulation), () - job = self._stores.households.read(run.run_id) - if job is None: - return None - return self._household_call(job), () - except (SQLAlchemyError, VersionManifestUnavailableError) as error: - raise _DeferredError(type(error).__name__) from error - except _REBUILD_ERRORS as error: - raise _CannotRebuildError( - _rebuild_reason(error), - error_type=type(error).__name__, - simulation_run_ids=simulation_run_ids, - ) from error - - 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._stores.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 - def _fail( self, run: UnstartedRun, message: str, simulation_run_ids: tuple[UUID, ...] ) -> bool | None: From 116780d76b1b4c5e9acfc55ad1f7e2a093975851 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:36:39 +0400 Subject: [PATCH 06/19] Test the scheduled restart check, its telemetry and the sweep's age (failing) The restart check runs every two minutes in the sweep's Modal app on the Modal worker's database secret, with the Modal environment and the v2 workers' telemetry settings baked into its image; each check and each restart is an operation named by a registered stage, with the restart's outcome logged as an event. The sweep now waits 15 minutes before failing a pending run with no call, after the restart window. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stages.py | 11 + .../tests/test_run_restart_check.py | 221 ++++++++++++++++++ .../tests/test_stage12_run_sweep.py | 16 ++ .../tests/test_stage12_run_sweep_app.py | 127 +++++++++- 4 files changed, 370 insertions(+), 5 deletions(-) create mode 100644 projects/policyengine-simulation-entry/tests/test_run_restart_check.py 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/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-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py index f72b380b8..0ad0be027 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,10 @@ from uuid import uuid4 import pytest +from policyengine_simulation_entry.run_restarts import RESTART_UNTIL from policyengine_simulation_executor.stage12_runtime.run_sweep import ( DEFAULT_MAXIMUM_AGE, + DEFAULT_MINIMUM_AGE, CallState, SweepSummary, modal_call_state, @@ -144,6 +146,20 @@ 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 12 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 + assert DEFAULT_MINIMUM_AGE > RESTART_UNTIL + + 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 ffccee8ab..3266f683d 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,12 @@ import yaml from fixtures.fake_modal import install_fake_modal +from policyengine_simulation_entry import run_restart_check 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 +29,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 +84,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 +111,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 +135,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", ) @@ -150,3 +178,92 @@ def recording_sweep(received_store): "failed_without_call": 4, } assert swept == [store] + + +def test_the_restart_check_runs_every_two_minutes_with_the_worker_database_secret( + monkeypatch, +) -> None: + """It restarts runs 3 to 12 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["timeout"] <= 600 + assert options["image"] is module.sweep_image + + +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 From ace6c6427914805ff3d71ace0776713fd809d616 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:40:30 +0400 Subject: [PATCH 07/19] Run the restart check every two minutes beside the stuck-run sweep The sweep's Modal app gains a scheduled function that runs the simulation service's restart check as the Modal worker role, with the Modal environment and the v2 workers' telemetry settings baked into its image at deployment. Each check and each restarted run is an operation named by a registered stage, and each restart's outcome is logged as an event. The sweep now fails a pending run with no call at 15 minutes, after the restart window; the entry package loads its web app lazily so the restart code imports in the sweep's image. Co-Authored-By: Claude Opus 5.5 --- .../workflows/simulation-deploy.reusable.yml | 28 +- .../stages.py | 10 + .../policyengine_simulation_entry/__init__.py | 23 +- .../run_restart_check.py | 242 ++++++++++++++++++ .../pyproject.toml | 6 + .../src/modal/run_sweep_app.py | 92 ++++++- .../stage12_runtime/run_sweep.py | 24 +- .../policyengine-simulation-executor/uv.lock | 67 +++++ 8 files changed, 458 insertions(+), 34 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_check.py 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/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/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/run_restart_check.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_check.py new file mode 100644 index 000000000..0030b49ac --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restart_check.py @@ -0,0 +1,242 @@ +"""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 lists is one ``v2_run_restart`` operation, with the row's kind, + ID and age, so the operation count metric counts the runs whose call never + started; +- 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. + +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-executor/pyproject.toml b/projects/policyengine-simulation-executor/pyproject.toml index 2b4f42ac6..f566c5f93 100644 --- a/projects/policyengine-simulation-executor/pyproject.toml +++ b/projects/policyengine-simulation-executor/pyproject.toml @@ -19,6 +19,10 @@ dependencies = [ "policyengine-simulation-observability", "policyengine-simulation-contract[modal]", "policyengine-stage12-persistence", + # The restart check in the stuck-run sweep's Modal app runs the + # simulation service's own restart code on its persistence library. + "policyengine-simulation-entry", + "policyengine-simulation-persistence", # PolicyEngine.py's models extra is the only selector for core, country # models, and the SPM calculator. Do not duplicate those versions here. "policyengine[models]==6.17.1", @@ -80,6 +84,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 722c6d8b1..1054a9d12 100644 --- a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -1,17 +1,29 @@ -"""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, -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. Its result -counts runs and household jobs together. 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. +"""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, 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. Its result counts runs and household jobs together. + +The restart check runs every two minutes. It starts once more the call of +each run or household job created 3 to 12 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 @@ -23,6 +35,7 @@ from __future__ import annotations +import os from dataclasses import asdict from pathlib import Path @@ -30,11 +43,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) @@ -48,6 +67,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, @@ -62,6 +87,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, ) ) @@ -86,3 +114,39 @@ def sweep_stuck_runs() -> dict: summary = run_sweep.sweep_stuck_runs(report_run_access.report_run_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 14841c4bd..708118ece 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 12 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 @@ -49,8 +51,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 12 +# 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/uv.lock b/projects/policyengine-simulation-executor/uv.lock index b46aa9520..8ddd1c0d5 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" @@ -1905,7 +1948,9 @@ dependencies = [ { name = "policyengine-fastapi" }, { name = "policyengine-observability", extra = ["fastapi", "google", "otlp-grpc"] }, { name = "policyengine-simulation-contract", extra = ["modal"] }, + { name = "policyengine-simulation-entry" }, { name = "policyengine-simulation-observability" }, + { name = "policyengine-simulation-persistence" }, { name = "policyengine-stage12-persistence" }, { name = "pyarrow" }, { name = "pydantic-settings" }, @@ -1957,7 +2002,9 @@ requires-dist = [ { name = "policyengine-fastapi", editable = "../../libs/policyengine-fastapi" }, { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.7,<4" }, { name = "policyengine-simulation-contract", extras = ["modal"], editable = "../../libs/policyengine-simulation-contract" }, + { name = "policyengine-simulation-entry", editable = "../policyengine-simulation-entry" }, { 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 = "pyarrow", specifier = ">=20,<24" }, { name = "pydantic-settings", specifier = ">=2.7.1,<3.0.0" }, @@ -2049,6 +2096,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" From 83d0b80989ba2bf30efc8f3f34e514a9f224aa1d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:41:46 +0400 Subject: [PATCH 08/19] Describe the restart check and the sweep's 15 minutes Co-Authored-By: Claude Opus 5.5 --- .../c1-19-restart-unstarted-runs.added.md | 1 + .../stage-12-modal-worker-foundation.md | 30 ++++++++++++++-- .../README.md | 26 ++++++++++++++ .../policyengine-simulation-entry/README.md | 36 +++++++++++++++++-- .../README.md | 5 +-- 5 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 changelog.d/c1-19-restart-unstarted-runs.added.md 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..f81d27024 --- /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 12 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 ccb02b6f8..2865fb8b5 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -356,7 +356,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 @@ -709,10 +709,36 @@ 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. +**Restarting a run whose call never started.** The same Modal application +runs a restart check every two minutes, as the Modal worker role. It lists +the household jobs, report runs and standalone simulation runs created 3 to +12 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, it reads +back what the entrypoint built the call from, builds the same call with the +entrypoint's own builders on the worker of the .py version the row records, +sets the row's `updated_at` so it 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`; one whose rows or +manifest it can't read now is left for the next check. The sweep's 15 +minutes come after that 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 the report runs of a budget window and the runs of +household simulations, whose own calls start them. Each check is a +`v2_run_restart_check` operation and each listed row 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, so the operation count +metric and the events say how often runs need a restart and how each ended. + - Image: the CA store, the frozen `modal-simulation-image` dependency group and the source packages, without the bundle's datasets or the Stage 12 cache. Building it needs no secret. diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 4fea4fd37..183845032 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -125,6 +125,32 @@ 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. + +`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/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index bc32cc1f6..9e00f34be 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -100,12 +100,44 @@ 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 12 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 12 minutes ago that are still `pending` + with no call and that nothing has changed since their insert, at most 50 + a check, oldest first. A budget window's report runs and household + simulations' runs are left out. +- 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. +- 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 row whose rows or manifest can't be read now is left + unmarked for the next check. +- 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-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 From 6aacd1af79f2e0b9927489c5d50c2ba73a1789e0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:42:44 +0400 Subject: [PATCH 09/19] Test that a restart check waiting on every call ends within its timeout (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_run_sweep_app.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) 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 3266f683d..553d2d214 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 @@ -9,6 +9,10 @@ import yaml from fixtures.fake_modal import install_fake_modal from policyengine_simulation_entry import run_restart_check +from policyengine_simulation_entry.run_restarts import ( + CALL_START_TIMEOUT_SECONDS, + RESTARTS_PER_CHECK, +) from policyengine_simulation_executor.stage12_runtime import ( report_run_access as report_run_access_module, ) @@ -197,10 +201,21 @@ def test_the_restart_check_runs_every_two_minutes_with_the_worker_database_secre # 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["timeout"] <= 600 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.""" + module = _load(monkeypatch) + options = _functions(module)["restart_unstarted_runs"] + longest_call_waits = RESTARTS_PER_CHECK * CALL_START_TIMEOUT_SECONDS + + # A minute for the database and the manifest. + assert longest_call_waits + 60 <= options["timeout"] + + def test_the_image_bakes_the_modal_environment_and_telemetry_settings_in( monkeypatch, ) -> None: From 5a8f7e9f162493611a7b65268ff906016973ca98 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:43:01 +0400 Subject: [PATCH 10/19] Restart at most 20 runs a check, so a check ends within its timeout Co-Authored-By: Claude Opus 5.5 --- projects/policyengine-simulation-entry/README.md | 2 +- .../src/policyengine_simulation_entry/run_restarts.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 9e00f34be..240b68e8f 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -118,7 +118,7 @@ this service: the stuck-run sweep's Modal application runs it every two minutes as the Modal worker role. - It takes the rows created 3 to 12 minutes ago that are still `pending` - with no call and that nothing has changed since their insert, at most 50 + 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. - It rebuilds each call with this service's own builders 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 index 0dfe483d7..631de8d18 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -91,7 +91,10 @@ # minutes, so a restart has time to record its call before the sweep judges it. RESTART_UNTIL: Final = timedelta(minutes=12) # At most this many rows are restarted by one check; the next takes the rest. -RESTARTS_PER_CHECK: Final = 50 +# With every call waited on for its full time, a check still ends within the +# scheduled function's five-minute timeout, so no row it marks is left +# without its call. +RESTARTS_PER_CHECK: Final = 20 # How long the check waits for Modal to accept a restarted call. CALL_START_TIMEOUT_SECONDS: Final = 10.0 _NOT_STARTED: Final = ( From d5ee476e41b93cf292cdf974fe277188658a947c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:10:48 +0400 Subject: [PATCH 11/19] Test review round 1's restart findings (failing) The executor's image builds only from the directories its Dockerfile copies; an unreadable report run is failed with its simulation runs; a restart is counted by whose call holds the row afterwards; a check marks no row once its time limit has passed; and a call is rebuilt off the event loop. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_run_restart_records.py | 60 ++++++ .../tests/test_run_restarts.py | 175 +++++++++++++++++- .../tests/test_executor_image_dependencies.py | 51 +++++ .../tests/test_stage12_run_sweep_app.py | 22 ++- 4 files changed, 301 insertions(+), 7 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_executor_image_dependencies.py diff --git a/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py index fddb135a8..0ceac82e5 100644 --- a/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py +++ b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py @@ -31,6 +31,7 @@ ReportRunTrigger, ReportSimulation, RestartKind, + RunCall, RunRestartStore, SimulationChoice, SimulationType, @@ -42,6 +43,8 @@ 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, @@ -477,6 +480,56 @@ def test_a_report_run_whose_sides_differ_cannot_be_read(connection) -> None: 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``.""" @@ -509,8 +562,15 @@ def test_the_store_lists_marks_and_reads_in_transactions_of_its_own( 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 diff --git a/projects/policyengine-simulation-entry/tests/test_run_restarts.py b/projects/policyengine-simulation-entry/tests/test_run_restarts.py index 726194210..42c8cce0c 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_restarts.py +++ b/projects/policyengine-simulation-entry/tests/test_run_restarts.py @@ -43,6 +43,7 @@ ManifestHouseholdWorkers, ) from policyengine_simulation_entry.run_restarts import ( + CHECK_TIME_LIMIT, RESTART_AFTER, RESTART_UNTIL, RESTARTS_PER_CHECK, @@ -60,6 +61,7 @@ HouseholdJobRecord, HouseholdJobStatus, RestartKind, + RunCall, SimulationChoice, StoredEconomySide, StoredHousehold, @@ -71,7 +73,7 @@ from sqlalchemy.exc import OperationalError if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator from datetime import datetime from policyengine_simulation_entry.household_workers import HouseholdWorker @@ -118,11 +120,28 @@ class FakeRestartStore: 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 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, ...]: @@ -288,7 +307,8 @@ def _rig( # noqa: PLR0913 # each fake a test may replace households: FakeHouseholdJobs | None = None, economy_caller: FakeCaller | None = None, household_caller: FakeHouseholdCaller | None = None, - manifest: dict[str, object] | None = None, + manifest: object | None = None, + timer: Callable[[], float] | None = None, ) -> RestartRig: restarts = restarts or FakeRestartStore() economy = economy or FakeStore() @@ -297,6 +317,16 @@ def _rig( # noqa: PLR0913 # each fake a test may replace 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( @@ -307,6 +337,7 @@ def _rig( # noqa: PLR0913 # each fake a test may replace 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", ) @@ -781,3 +812,143 @@ def test_the_deployment_label_follows_the_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) 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_stage12_run_sweep_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_app.py index 553d2d214..2410e7f8b 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 @@ -9,9 +9,12 @@ 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, - RESTARTS_PER_CHECK, + CHECK_TIME_LIMIT, ) from policyengine_simulation_executor.stage12_runtime import ( report_run_access as report_run_access_module, @@ -207,13 +210,22 @@ def test_the_restart_check_runs_every_two_minutes_with_the_worker_database_secre 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.""" + """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"] - longest_call_waits = RESTARTS_PER_CHECK * CALL_START_TIMEOUT_SECONDS + statement_seconds = RESTART_STATEMENT_TIMEOUT_MILLISECONDS / 1000 + longest_after_the_limit = CALL_START_TIMEOUT_SECONDS + 3 * statement_seconds - # A minute for the database and the manifest. - assert longest_call_waits + 60 <= options["timeout"] + # 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( From 5e158d399ac6538366973fe6f5ba289902674351 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:14:35 +0400 Subject: [PATCH 12/19] Fix review round 1's restart findings The executor's Docker image builds again: the Simulation Entrypoint and its persistence library move to the executor's dev group, which the deploy job installs for modal deploy and the image does not. A report run whose rows can't be rebuilt is failed with its simulation runs. A restart is counted by whose call the row records afterwards, so a restarted call that claims its run before the record counts as restarted. A check marks no row after three minutes, rebuilds each call on a thread so Modal's blocking reads don't hold its event loop, and a restart reads its values with the submission's own reader. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 55 ++++---- .../README.md | 7 +- .../__init__.py | 8 ++ .../economy_run_store.py | 27 +++- .../run_restart_records.py | 79 ++++++++---- .../run_restart_store.py | 15 +++ .../policyengine-simulation-entry/README.md | 10 +- .../run_restart_calls.py | 26 ++++ .../run_restarts.py | 119 +++++++++++------- .../pyproject.toml | 11 +- .../policyengine-simulation-executor/uv.lock | 12 +- 11 files changed, 258 insertions(+), 111 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 2865fb8b5..47c930f1a 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -713,32 +713,6 @@ 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. -**Restarting a run whose call never started.** The same Modal application -runs a restart check every two minutes, as the Modal worker role. It lists -the household jobs, report runs and standalone simulation runs created 3 to -12 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, it reads -back what the entrypoint built the call from, builds the same call with the -entrypoint's own builders on the worker of the .py version the row records, -sets the row's `updated_at` so it 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`; one whose rows or -manifest it can't read now is left for the next check. The sweep's 15 -minutes come after that 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 the report runs of a budget window and the runs of -household simulations, whose own calls start them. Each check is a -`v2_run_restart_check` operation and each listed row 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, so the operation count -metric and the events say how often runs need a restart and how each ended. - - Image: the CA store, the frozen `modal-simulation-image` dependency group and the source packages, without the bundle's datasets or the Stage 12 cache. Building it needs no secret. @@ -765,6 +739,35 @@ 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 12 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; one whose rows or manifest it can't read now is left +for the next check. 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 listed row 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, so the operation count +metric and the events say how often runs need a restart and how each ended. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 183845032..4c4b8cd9b 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -145,7 +145,12 @@ the run tables and their inputs, and `UPDATE` of `updated_at`, 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. + 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 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 e1c2048a0..f7236c1f0 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -51,6 +51,7 @@ SimulationInputs, UnknownDatasetError, UnknownRegionError, + read_simulation_reform, ) from policyengine_simulation_persistence.engine import ( bounded_transaction, @@ -90,6 +91,7 @@ RESTART_TABLES, EconomyScope, RestartKind, + RunCall, StoredEconomySide, StoredReportRun, StoredRunUnreadableError, @@ -97,6 +99,8 @@ UnstartedRun, list_unstarted_runs, mark_restart, + read_report_run_simulation_run_ids, + read_run_call, read_unstarted_report_run, read_unstarted_simulation_run, ) @@ -153,6 +157,7 @@ "ReportSimulation", "RestartKind", "ReusableRun", + "RunCall", "RunRestartStore", "RunReuse", "RunVersions", @@ -202,6 +207,9 @@ "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", 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 6e66b2a50..ca05c3b3a 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 @@ -443,6 +443,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 @@ -462,6 +485,4 @@ 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) 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 index 98f5f36bb..0613d1f5c 100644 --- 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 @@ -11,7 +11,8 @@ (``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: their own tasks restart them. + 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 @@ -40,6 +41,9 @@ RunVersions, SimulationChoice, ) +from policyengine_simulation_persistence.economy_run_store import ( + read_simulation_reform, +) from policyengine_simulation_persistence.input_tables import ( MODEL_NAME_PREFIX, datasets, @@ -47,11 +51,6 @@ tax_benefit_model_versions, tax_benefit_models, ) -from policyengine_simulation_persistence.policy_reform import ( - combine_reforms, - read_dynamic_reform, - read_policy_reform, -) from policyengine_simulation_persistence.tables import ( SimulationStatus, SimulationType, @@ -102,6 +101,14 @@ 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. @@ -256,6 +263,40 @@ def mark_restart(connection: Connection, run: UnstartedRun) -> bool: 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( @@ -312,29 +353,15 @@ def _side( Raise the policy reader's errors, or OverlappingReformError. """ choice = SimulationChoice(policy_id=side.policy_id, dynamic_id=side.dynamic_id) - policy_reform = ( - {} - if choice.policy_id is None - else read_policy_reform( - connection, - country=country, - policy_id=choice.policy_id, - catalog_version_id=scope.catalog_version_id, - ) - ) - dynamic_reform = ( - {} - if choice.dynamic_id is None - else read_dynamic_reform( - connection, - dynamic_id=choice.dynamic_id, - catalog_version_id=scope.catalog_version_id, - ) - ) return StoredEconomySide( simulation_run_id=side.id, choice=choice, - reform=combine_reforms(policy_reform, dynamic_reform), + reform=read_simulation_reform( + connection, + country=country, + catalog_version_id=scope.catalog_version_id, + choice=choice, + ), ) 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 index f7005c662..ce1733723 100644 --- 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 @@ -14,6 +14,8 @@ 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, ) @@ -27,6 +29,7 @@ from sqlalchemy import Connection, Engine from policyengine_simulation_persistence.run_restart_records import ( + RunCall, StoredReportRun, StoredSimulationRun, UnstartedRun, @@ -83,3 +86,15 @@ def read_simulation_run( """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/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 240b68e8f..fee081036 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -120,7 +120,8 @@ runs it every two minutes as the Modal worker role. - It takes the rows created 3 to 12 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. + 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 @@ -129,11 +130,14 @@ runs it every two minutes as the Modal worker role. (`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. + 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 row whose rows or manifest can't be read now is left + it, and …`; a report run is failed with its simulation runs still pending + with no call. A row whose rows or manifest can't be read now is left unmarked for the next check. - Each check, each row and each outcome are telemetry; see `run_restart_check`'s docstring for the operations and event. 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 index 8c7d12b9d..40d089da9 100644 --- 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 @@ -68,6 +68,7 @@ HouseholdJobRecord, HouseholdJobStatus, NewHouseholdJob, + RunCall, RunVersions, StoredReportRun, StoredSimulationRun, @@ -141,6 +142,16 @@ def read_simulation_run( """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``).""" @@ -367,12 +378,27 @@ def build(self, run: UnstartedRun) -> RebuiltCall | None: 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( 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 index 631de8d18..95a7ba665 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -32,6 +32,7 @@ from __future__ import annotations import asyncio +import time import uuid from collections import Counter from contextlib import nullcontext @@ -69,6 +70,7 @@ __all__ = [ "CALL_START_TIMEOUT_SECONDS", + "CHECK_TIME_LIMIT", "RESTARTS_PER_CHECK", "RESTART_AFTER", "RESTART_UNTIL", @@ -90,31 +92,35 @@ # A row older than this is left to the stuck-run sweep, which fails it at 15 # minutes, so a restart has time to record its call before the sweep judges it. RESTART_UNTIL: Final = timedelta(minutes=12) -# At most this many rows are restarted by one check; the next takes the rest. -# With every call waited on for its full time, a check still ends within the -# scheduled function's five-minute timeout, so no row it marks is left -# without its call. +# 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" ) -_PENDING: Final = "pending" class RestartOutcome(StrEnum): """How the check ended for one listed row. - ``restarted``: its new call started and the call's ID was recorded. - ``call_not_recorded``: its new call started, but the ID couldn't be - recorded; the call claims the row itself. ``started_meanwhile``: another - call claimed the row after the check started its own, which that call - then can't claim. ``left``: the row was gone, or changed before the check - started a call; also a row whose failure couldn't be written, which the + ``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 row whose failure couldn't be written, which the sweep fails later. ``deferred``: its rows or the manifest couldn't be - read, so it is left unmarked for the next check. ``failed``: it couldn't - be started again and was failed with ``run_not_started``. + read, or the check's time ran out, so it is left unmarked for the next + check. ``failed``: it couldn't be started again and was failed with + ``run_not_started``. """ RESTARTED = "restarted" @@ -231,8 +237,9 @@ class RestartStores: class RunRestartService: """Start the calls of the runs and household jobs whose call never started. - The store reads and writes are synchronous: the check is a batch job of - its own, so nothing else waits on its event loop. + 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 @@ -244,6 +251,7 @@ def __init__( # noqa: PLR0913 # each seam tests replace 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: @@ -259,13 +267,18 @@ def __init__( # noqa: PLR0913 # each seam tests replace ) 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 12 minutes ago with no call. - Raise the store's error when the rows can't be listed. + 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, @@ -273,7 +286,10 @@ async def restart_unstarted_runs(self) -> RestartSummary: limit=RESTARTS_PER_CHECK, ) outcomes: Counter[RestartOutcome] = Counter() - for run in runs: + 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) @@ -284,10 +300,13 @@ async def restart_unstarted_runs(self) -> RestartSummary: **{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 = self._calls.build(run) + rebuilt = await asyncio.to_thread(self._calls.build, run) except DeferredRestartError as deferred: return RestartOutcome.DEFERRED, str(deferred) except CallNotRebuiltError as cannot: @@ -314,6 +333,8 @@ async def _start_marked( 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 @@ -332,29 +353,42 @@ async def _start_marked( f"again ({type(error).__name__})." ) failed = self._fail(run, message, simulation_run_ids) - # A failure that changed nothing means another call claimed the - # row; one that couldn't be written leaves the row to the sweep. - outcome = { - True: RestartOutcome.FAILED, - False: RestartOutcome.STARTED_MEANWHILE, - None: RestartOutcome.LEFT, - }[failed] + # 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__ - return self._recorded(run, call_id) - - def _recorded( - self, run: UnstartedRun, call_id: str - ) -> tuple[RestartOutcome, str | None]: - """Record the started call's ID; return how the restart ended.""" try: - status = self._record(run, call_id) + self._record(run, call_id) except SQLAlchemyError as error: return RestartOutcome.CALL_NOT_RECORDED, type(error).__name__ - if status is None: - return RestartOutcome.LEFT, None - if status == _PENDING: - return RestartOutcome.RESTARTED, None - return RestartOutcome.STARTED_MEANWHILE, None + 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, ...] @@ -378,12 +412,11 @@ def _fail( return None return failed - def _record(self, run: UnstartedRun, call_id: str) -> str | None: - """Record the call's ID while the row is unstarted; return its status after.""" + 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: - status = self._stores.economy.record_report_call(run.run_id, call_id) + self._stores.economy.record_report_call(run.run_id, call_id) elif run.kind is RestartKind.SIMULATION_RUN: - status = self._stores.economy.record_simulation_call(run.run_id, call_id) + self._stores.economy.record_simulation_call(run.run_id, call_id) else: - status = self._stores.households.record_call(run.run_id, call_id) - return None if status is None else str(status) + self._stores.households.record_call(run.run_id, call_id) diff --git a/projects/policyengine-simulation-executor/pyproject.toml b/projects/policyengine-simulation-executor/pyproject.toml index f566c5f93..cdb6f32e3 100644 --- a/projects/policyengine-simulation-executor/pyproject.toml +++ b/projects/policyengine-simulation-executor/pyproject.toml @@ -19,10 +19,6 @@ dependencies = [ "policyengine-simulation-observability", "policyengine-simulation-contract[modal]", "policyengine-stage12-persistence", - # The restart check in the stuck-run sweep's Modal app runs the - # simulation service's own restart code on its persistence library. - "policyengine-simulation-entry", - "policyengine-simulation-persistence", # PolicyEngine.py's models extra is the only selector for core, country # models, and the SPM calculator. Do not duplicate those versions here. "policyengine[models]==6.17.1", @@ -75,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] diff --git a/projects/policyengine-simulation-executor/uv.lock b/projects/policyengine-simulation-executor/uv.lock index 8ddd1c0d5..2fab1f614 100644 --- a/projects/policyengine-simulation-executor/uv.lock +++ b/projects/policyengine-simulation-executor/uv.lock @@ -1948,9 +1948,7 @@ dependencies = [ { name = "policyengine-fastapi" }, { name = "policyengine-observability", extra = ["fastapi", "google", "otlp-grpc"] }, { name = "policyengine-simulation-contract", extra = ["modal"] }, - { name = "policyengine-simulation-entry" }, { name = "policyengine-simulation-observability" }, - { name = "policyengine-simulation-persistence" }, { name = "policyengine-stage12-persistence" }, { name = "pyarrow" }, { name = "pydantic-settings" }, @@ -1973,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" }, @@ -2002,9 +2002,7 @@ requires-dist = [ { name = "policyengine-fastapi", editable = "../../libs/policyengine-fastapi" }, { name = "policyengine-observability", extras = ["fastapi", "google", "otlp-grpc"], specifier = ">=3.0.7,<4" }, { name = "policyengine-simulation-contract", extras = ["modal"], editable = "../../libs/policyengine-simulation-contract" }, - { name = "policyengine-simulation-entry", editable = "../policyengine-simulation-entry" }, { 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 = "pyarrow", specifier = ">=20,<24" }, { name = "pydantic-settings", specifier = ">=2.7.1,<3.0.0" }, @@ -2018,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" }, From 1768f26f4f91d752e3cbe594b1b0034e2b216d6f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:41:38 +0400 Subject: [PATCH 13/19] Test review round 2's restart findings (failing) A restart is recorded with a minute to spare before the sweep judges its run, an unexpected rebuild error defers only its own run, and the restart check imports with only the sweep image's packages (this last test passed when first run: it guards what already holds). Co-Authored-By: Claude Opus 5.5 --- .../tests/test_run_restarts.py | 35 ++++- .../tests/test_run_restart_image.py | 132 ++++++++++++++++++ .../tests/test_stage12_run_sweep.py | 16 ++- 3 files changed, 178 insertions(+), 5 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_run_restart_image.py diff --git a/projects/policyengine-simulation-entry/tests/test_run_restarts.py b/projects/policyengine-simulation-entry/tests/test_run_restarts.py index 42c8cce0c..68ecd2eda 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_restarts.py +++ b/projects/policyengine-simulation-entry/tests/test_run_restarts.py @@ -124,6 +124,8 @@ class FakeRestartStore: 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``). @@ -157,6 +159,8 @@ def mark_restart(self, run: UnstartedRun) -> bool: 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( @@ -452,7 +456,7 @@ 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_twelve_minutes_ago() -> None: +def test_the_check_lists_runs_created_three_to_ten_minutes_ago() -> None: rig = _rig() summary = rig.restart() @@ -462,7 +466,7 @@ def test_the_check_lists_runs_created_three_to_twelve_minutes_ago() -> None: ] assert ( timedelta(minutes=3), - timedelta(minutes=12), + timedelta(minutes=10), ) == (RESTART_AFTER, RESTART_UNTIL) assert summary == RestartSummary(listed=0) @@ -952,3 +956,30 @@ def test_a_runs_call_is_rebuilt_off_the_event_loop() -> None: 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() -> None: + """One row the check can't rebuild for an unforeseen reason holds up no other.""" + 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) 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..6f36bd130 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py @@ -0,0 +1,132 @@ +"""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 +every other installed distribution's modules. +""" + +from __future__ import annotations + +import re +import subprocess +import sys +import tomllib +from importlib import metadata +from pathlib import Path + +PROJECT = Path(__file__).resolve().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", + } +) +# 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 _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" + script.write_text( + "import importlib, sys\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", + 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-executor/tests/test_stage12_run_sweep.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep.py index 0ad0be027..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,7 +8,11 @@ from uuid import uuid4 import pytest -from policyengine_simulation_entry.run_restarts import RESTART_UNTIL +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, @@ -147,7 +151,7 @@ def test_the_sweep_lists_runs_older_than_its_minimum_age() -> None: 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 12 minutes old. + """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. @@ -157,7 +161,13 @@ def test_by_default_unstarted_runs_wait_until_the_restart_check_has_passed() -> sweep_stuck_runs(store, call_state=lambda _: CallState.ENDED, now=NOW) assert store.cutoffs == [NOW - timedelta(minutes=15)] * 3 - assert DEFAULT_MINIMUM_AGE > RESTART_UNTIL + # 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: From b0db539efedd1bc426e9a90569f647447bd9f59b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:42:55 +0400 Subject: [PATCH 14/19] Fix review round 2's restart findings The restart window is now 3 to 10 minutes after creation, so a row marked at the end of a check's time limit has its call recorded minutes before the sweep judges it; an unexpected error rebuilding one row's call defers that row and the check goes on; and the telemetry docs say how to count runs from restart attempts. Co-Authored-By: Claude Opus 5.5 --- changelog.d/c1-19-restart-unstarted-runs.added.md | 2 +- docs/migration/stage-12-modal-worker-foundation.md | 12 +++++++----- projects/policyengine-simulation-entry/README.md | 4 ++-- .../run_restart_check.py | 8 +++++--- .../policyengine_simulation_entry/run_restarts.py | 14 ++++++++++---- .../src/modal/run_sweep_app.py | 2 +- .../stage12_runtime/run_sweep.py | 4 ++-- .../tests/test_stage12_run_sweep_app.py | 2 +- 8 files changed, 29 insertions(+), 19 deletions(-) diff --git a/changelog.d/c1-19-restart-unstarted-runs.added.md b/changelog.d/c1-19-restart-unstarted-runs.added.md index f81d27024..2ad18f24b 100644 --- a/changelog.d/c1-19-restart-unstarted-runs.added.md +++ b/changelog.d/c1-19-restart-unstarted-runs.added.md @@ -1 +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 12 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. +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 47c930f1a..84311afae 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -742,7 +742,7 @@ 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 12 minutes ago that are still pending +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 @@ -763,10 +763,12 @@ 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 listed row 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, so the operation count -metric and the events say how often runs need a restart and how each ended. +`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 runs that needed a restart +are the distinct run IDs, or the outcomes other than `deferred`; the events +say how each ended. ## Temporary direct runner endpoint diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index fee081036..42dd96279 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -101,7 +101,7 @@ 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 restart check starts -its call once more 3 to 12 minutes after its creation, and the stuck-run +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 @@ -117,7 +117,7 @@ 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 12 minutes ago that are still `pending` +- 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 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 index 0030b49ac..dbdae0e00 100644 --- 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 @@ -10,9 +10,11 @@ - 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 lists is one ``v2_run_restart`` operation, with the row's kind, - ID and age, so the operation count metric counts the runs whose call never - started; +- 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; the runs + whose call never started are the distinct ``run_id`` values, or the + operations with any outcome but ``deferred``; - 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 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 index 95a7ba665..be7b230b4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -8,7 +8,7 @@ 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 12 minutes ago that are still pending with +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; @@ -90,8 +90,10 @@ # 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, so a restart has time to record its call before the sweep judges it. -RESTART_UNTIL: Final = timedelta(minutes=12) +# 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. @@ -272,7 +274,7 @@ def __init__( # noqa: PLR0913 # each seam tests replace self._deadline = 0.0 async def restart_unstarted_runs(self) -> RestartSummary: - """Start once more each row created 3 to 12 minutes ago with no call. + """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 @@ -322,6 +324,10 @@ async def _restart(self, run: UnstartedRun) -> tuple[RestartOutcome, str | None] 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. + except Exception as error: # noqa: BLE001 # one row's unforeseen error holds up no other + 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) 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 1054a9d12..1ed512e20 100644 --- a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -11,7 +11,7 @@ whatever Modal reports. Its result counts runs and household jobs together. The restart check runs every two minutes. It starts once more the call of -each run or household job created 3 to 12 minutes ago that is still pending +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 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 708118ece..4e6722e04 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 @@ -9,7 +9,7 @@ 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 12 minutes after its creation, could +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 @@ -51,7 +51,7 @@ # 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. It is longer than the restart check's window (3 to 12 +# 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) 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 2410e7f8b..b482069a0 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 @@ -190,7 +190,7 @@ def recording_sweep(received_store): def test_the_restart_check_runs_every_two_minutes_with_the_worker_database_secret( monkeypatch, ) -> None: - """It restarts runs 3 to 12 minutes old, so each is seen several times.""" + """It restarts runs 3 to 10 minutes old, so each is seen several times.""" module = _load(monkeypatch) options = _functions(module)["restart_unstarted_runs"] From a9fb34e66556b144f79f5c395d47903398d4be2f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:59:06 +0400 Subject: [PATCH 15/19] Test review round 3's restart findings (failing) An unexpected rebuild error is logged with its traceback, the image test also requires every module the restart check loads to come from the standard library, a local source or an image distribution, and the persistence window test uses the 3-to-10-minute window. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_run_restart_records.py | 4 +- .../tests/test_run_restarts.py | 17 ++++- .../tests/test_run_restart_image.py | 76 ++++++++++++++++++- 3 files changed, 90 insertions(+), 7 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py index 0ceac82e5..7573c4f69 100644 --- a/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py +++ b/libs/policyengine-simulation-persistence/tests/test_run_restart_records.py @@ -200,14 +200,14 @@ def test_a_run_created_outside_the_window_is_not_listed(connection) -> None: too_new = list_unstarted_runs( connection, - created_after=created_at - timedelta(minutes=12), + 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=12), + created_before=created_at + timedelta(minutes=10), limit=LISTING_LIMIT, ) diff --git a/projects/policyengine-simulation-entry/tests/test_run_restarts.py b/projects/policyengine-simulation-entry/tests/test_run_restarts.py index 68ecd2eda..bf78b53b8 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_restarts.py +++ b/projects/policyengine-simulation-entry/tests/test_run_restarts.py @@ -958,8 +958,13 @@ def test_a_runs_call_is_rebuilt_off_the_event_loop() -> None: assert not any(manifest.read_on_event_loop) -def test_an_unexpected_rebuild_error_defers_its_run_and_the_check_goes_on() -> None: - """One row the check can't rebuild for an unforeseen reason holds up no other.""" +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( @@ -983,3 +988,11 @@ def test_an_unexpected_rebuild_error_defers_its_run_and_the_check_goes_on() -> N ) 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) + assert (logged.run_kind, logged.run_id) == ("report_run", str(broken_run_id)) diff --git a/projects/policyengine-simulation-executor/tests/test_run_restart_image.py b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py index 6f36bd130..b078d2b3f 100644 --- a/projects/policyengine-simulation-executor/tests/test_run_restart_image.py +++ b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py @@ -4,20 +4,27 @@ ``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 -every other installed distribution's modules. +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" @@ -32,6 +39,25 @@ "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 modules the test's own process loads that are not the restart check's: +# its script (also named ``__mp_main__`` once multiprocessing is imported) and +# the virtual environment's start-up module. +NOT_THE_RESTART_CHECK = ("__main__", "__mp_main__", "_virtualenv") # What the scheduled function's body imports. RESTART_MODULES = ( "policyengine_simulation_entry.run_restart_check", @@ -88,6 +114,32 @@ def _top_level_modules(distribution: metadata.Distribution) -> set[str]: 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 @@ -109,6 +161,7 @@ def test_the_restart_check_imports_with_only_the_sweep_images_packages( # 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" f"refused = {sorted(refused)!r}\n" @@ -118,7 +171,12 @@ def test_the_restart_check_imports_with_only_the_sweep_images_packages( " 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", + " 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" + f" if name not in {list(NOT_THE_RESTART_CHECK)!r}}} - {{''}})\n" + f"open({str(loaded_list)!r}, 'w').write(json.dumps(loaded))\n", encoding="utf-8", ) @@ -130,3 +188,15 @@ def test_the_restart_check_imports_with_only_the_sweep_images_packages( ) 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 == [] From 80a2e56d03f0a7984455861dcb63eb3a77abd3c3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:59:39 +0400 Subject: [PATCH 16/19] Fix review round 3's restart findings An unexpected error rebuilding a row's call is logged with its traceback, the descriptions of a deferred row name that case, and the telemetry docs count the runs that needed a restart by distinct run ID. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 8 +++++--- .../policyengine-simulation-entry/README.md | 3 ++- .../run_restart_check.py | 11 +++++++---- .../run_restarts.py | 19 ++++++++++++++----- 4 files changed, 28 insertions(+), 13 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 84311afae..80441db51 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -752,7 +752,8 @@ worker of the .py version the row records. It then sets the row's 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; one whose rows or manifest it can't read now is left +pending simulation runs; one whose rows or manifest it can't read now, or +whose rebuild raises an unexpected error (logged with its traceback), is left for the next check. 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 @@ -767,8 +768,9 @@ runs: it can't rebuild their calls yet. Each check is 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 runs that needed a restart -are the distinct run IDs, or the outcomes other than `deferred`; the events -say how each ended. +are the distinct run IDs (the outcomes other than `deferred` leave out a row +deferred in every check, which the sweep then fails); the events say how +each ended. ## Temporary direct runner endpoint diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 42dd96279..1beaae9f3 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -137,7 +137,8 @@ runs it every two minutes as the Modal worker role. 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 whose rows or manifest can't be read now is left + with no call. A row whose rows or manifest can't be read now, or whose + rebuild raises an unexpected error (logged with its traceback), is left unmarked for the next check. - Each check, each row and each outcome are telemetry; see `run_restart_check`'s docstring for the operations and event. 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 index dbdae0e00..93c2b40a6 100644 --- 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 @@ -12,13 +12,16 @@ 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; the runs - whose call never started are the distinct ``run_id`` values, or the - operations with any outcome but ``deferred``; + check, so the operation count metric counts restart attempts, and the + runs whose call never started are the distinct ``run_id`` values. (A + count of the outcomes other than ``deferred`` leaves out the rows deferred + in every check, which the sweep then fails.) - 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. + 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. 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 index be7b230b4..9f11e8d44 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -18,7 +18,8 @@ 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 is left unmarked, for the next check. A +rows or manifest can't be read now, or whose rebuild raises an unexpected +error, is left unmarked, for the next check. 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 @@ -32,6 +33,7 @@ from __future__ import annotations import asyncio +import logging import time import uuid from collections import Counter @@ -86,6 +88,8 @@ "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) @@ -120,8 +124,9 @@ class RestartOutcome(StrEnum): ``left``: the row was gone, changed before the check started a call, or failed meanwhile; also a row whose failure couldn't be written, which the sweep fails later. ``deferred``: its rows or the manifest couldn't be - read, or the check's time ran out, so it is left unmarked for the next - check. ``failed``: it couldn't be started again and was failed with + read, rebuilding its call raised an unexpected error (logged with its + traceback), or the check's time ran out, so it is left unmarked for the + next check. ``failed``: it couldn't be started again and was failed with ``run_not_started``. """ @@ -325,8 +330,12 @@ async def _restart(self, run: UnstartedRun) -> tuple[RestartOutcome, str | None] }[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. - except Exception as error: # noqa: BLE001 # one row's unforeseen error holds up no other + # the check goes on to the next row; its traceback is logged. + except Exception as error: + logger.exception( + "v2_run_restart_unexpected_error", + extra={"run_kind": run.kind.value, "run_id": str(run.run_id)}, + ) return RestartOutcome.DEFERRED, type(error).__name__ if rebuilt is None: return RestartOutcome.LEFT, None From a78cf822ec98618cab1c2e157408bcb6a18f9712 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:13:37 +0400 Subject: [PATCH 17/19] Test that the unexpected-error log carries its run to the runtime (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/test_run_restarts.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-entry/tests/test_run_restarts.py b/projects/policyengine-simulation-entry/tests/test_run_restarts.py index bf78b53b8..8eaabe323 100644 --- a/projects/policyengine-simulation-entry/tests/test_run_restarts.py +++ b/projects/policyengine-simulation-entry/tests/test_run_restarts.py @@ -995,4 +995,8 @@ def test_an_unexpected_rebuild_error_defers_its_run_and_the_check_goes_on( ] assert logged.exc_info is not None assert isinstance(logged.exc_info[1], RuntimeError) - assert (logged.run_kind, logged.run_id) == ("report_run", str(broken_run_id)) + # The observability runtime's log handler writes only these attributes. + assert logged.policyengine_attributes == { + "run_kind": "report_run", + "run_id": str(broken_run_id), + } From 37980447a4dd26bb9309958aac0f48b0feb94708 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:14:53 +0400 Subject: [PATCH 18/19] Fix review round 4's restart findings The unexpected-error log passes its run kind and ID as policyengine_attributes, which the observability runtime's log handler writes; the image test checks only the modules the import adds; and the descriptions of a deferred row and of the restart counts name every case. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 8 ++++---- .../run_restart_check.py | 8 +++++--- .../run_restarts.py | 20 ++++++++++++------- .../tests/test_run_restart_image.py | 13 +++++++----- 4 files changed, 30 insertions(+), 19 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 80441db51..ae02b2132 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -767,10 +767,10 @@ 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 runs that needed a restart -are the distinct run IDs (the outcomes other than `deferred` leave out a row -deferred in every check, which the sweep then fails); the events say how -each ended. +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 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 index 93c2b40a6..f4c254d42 100644 --- 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 @@ -13,9 +13,11 @@ - 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 - runs whose call never started are the distinct ``run_id`` values. (A - count of the outcomes other than ``deferred`` leaves out the rows deferred - in every check, which the sweep then fails.) + 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 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 index 9f11e8d44..ac7ad6bc1 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -19,7 +19,8 @@ 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. A +error, is left unmarked, for the next check, as is one the check has no +time left for or whose mark or failure can't be written. 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 @@ -122,11 +123,13 @@ class RestartOutcome(StrEnum): 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 row whose failure couldn't be written, which the - sweep fails later. ``deferred``: its rows or the manifest couldn't be - read, rebuilding its call raised an unexpected error (logged with its - traceback), or the check's time ran out, so it is left unmarked for the - next check. ``failed``: it couldn't be started again and was failed with + failed meanwhile; also a marked row whose failure couldn't be written, + which the sweep fails later. ``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``. """ @@ -332,9 +335,12 @@ async def _restart(self, run: UnstartedRun) -> tuple[RestartOutcome, str | None] # 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={"run_kind": run.kind.value, "run_id": str(run.run_id)}, + extra={**fields, "policyengine_attributes": fields}, ) return RestartOutcome.DEFERRED, type(error).__name__ if rebuilt is None: diff --git a/projects/policyengine-simulation-executor/tests/test_run_restart_image.py b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py index b078d2b3f..5273226f3 100644 --- a/projects/policyengine-simulation-executor/tests/test_run_restart_image.py +++ b/projects/policyengine-simulation-executor/tests/test_run_restart_image.py @@ -54,10 +54,9 @@ ) ), ) -# The modules the test's own process loads that are not the restart check's: -# its script (also named ``__mp_main__`` once multiprocessing is imported) and -# the virtual environment's start-up module. -NOT_THE_RESTART_CHECK = ("__main__", "__mp_main__", "_virtualenv") +# 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", @@ -164,6 +163,8 @@ def test_the_restart_check_imports_with_only_the_sweep_images_packages( 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" @@ -175,7 +176,9 @@ def test_the_restart_check_imports_with_only_the_sweep_images_packages( "import json\n" "loaded = sorted({getattr(module, '__file__', None) or ''\n" " for name, module in list(sys.modules.items())\n" - f" if name not in {list(NOT_THE_RESTART_CHECK)!r}}} - {{''}})\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", ) From b00679187db0fce7de2473c8cf22887871f15d6e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:25:38 +0400 Subject: [PATCH 19/19] Describe every way a restart is left or deferred Wording fixes from review round 5; no behaviour changes. Co-Authored-By: Claude Opus 5.5 --- .../migration/stage-12-modal-worker-foundation.md | 11 ++++++----- projects/policyengine-simulation-entry/README.md | 7 ++++--- .../policyengine_simulation_entry/run_restarts.py | 15 +++++++++------ 3 files changed, 19 insertions(+), 14 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index ae02b2132..185e8f44c 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -752,11 +752,12 @@ worker of the .py version the row records. It then sets the row's 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; one whose rows or manifest it can't read now, or -whose rebuild raises an unexpected error (logged with its traceback), is left -for the next check. 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 +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 diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 1beaae9f3..ab0890cc3 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -137,9 +137,10 @@ runs it every two minutes as the Modal worker role. 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 whose rows or manifest can't be read now, or whose - rebuild raises an unexpected error (logged with its traceback), is left - unmarked for the next check. + 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. 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 index ac7ad6bc1..5c7afb78f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_restarts.py @@ -20,7 +20,9 @@ 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 or whose mark or failure can't be written. A +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 @@ -124,11 +126,12 @@ class RestartOutcome(StrEnum): 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. ``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. + 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``. """