diff --git a/changelog.d/v2-budget-window-capacity.changed.md b/changelog.d/v2-budget-window-capacity.changed.md new file mode 100644 index 000000000..02738e725 --- /dev/null +++ b/changelog.d/v2-budget-window-capacity.changed.md @@ -0,0 +1,3 @@ +Raised the v2 functions' Modal container limits so budget windows at their +cap don't wait behind them: the coordinator to 25, each simulation worker +to 50, and the US segment function to 1,000. diff --git a/changelog.d/v2-budget-windows.added.md b/changelog.d/v2-budget-windows.added.md new file mode 100644 index 000000000..8af2aa090 --- /dev/null +++ b/changelog.d/v2-budget-windows.added.md @@ -0,0 +1,10 @@ +Calculated budget windows on the detailed-save path. `POST /reports` takes +`report_type` `budget_window` with `budget_window_size` (1 to 75 years, +ending by 2200) and creates a report run per year, answering 201 with the +report, the first year's run, the window run and every year's runs. At +most five years run at once: the service starts the first ones, and when a +year ends its coordinator starts the window's next waiting years. The +window is marked succeeded when every year has succeeded, or failed as +soon as a year fails. The stuck-run sweep leaves a window's waiting years +to the window, fails a window when it fails one of its years, and settles +a window that no year advances any more. diff --git a/docs/engineering/skills/database-access.md b/docs/engineering/skills/database-access.md index c559f59c7..f6325820e 100644 --- a/docs/engineering/skills/database-access.md +++ b/docs/engineering/skills/database-access.md @@ -46,8 +46,8 @@ and this repository has no migration chain of its own Their read-only models in that library's `input_tables.py` name only the columns the service reads, apart from the models of the tables it writes. - the run tables the v2 Modal coordinator, workers, household functions and - stuck-run sweep write as the Modal worker role: `report_runs`, `simulation_runs`, the - result tables and `household_jobs`. Their mappings in + stuck-run sweep write as the Modal worker role: `report_runs`, `simulation_runs`, + `budget_window_runs`, the result tables and `household_jobs`. Their mappings in `policyengine_stage12_persistence.report_run_tables` and `report_run_results`, separate from the simulation service's models, declare only the columns the coordinator's, workers', household functions' diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index ccb02b6f8..c03186693 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -26,10 +26,13 @@ simulation submission or polling contracts. and reform are separate and start before either is awaited, an eligible report starts 40 region-group calculation calls. These are invocations; Modal may reuse a warm container for more than one invocation. -- Country-specific single-simulation functions and the report coordinator have - a Modal `max_containers` value of 10. The separate US region-group function - has a value of 300 so calculation work cannot wait behind the parent calls - that are collecting it. These are per-function limits, not one shared quota. +- The report coordinator has a Modal `max_containers` value of 25: three + budget windows at five years each, and ten other reports. Each + country-specific single-simulation function has 50, both simulations of + every coordinator call, and the separate US region-group function has + 1,000, every segment of every US simulation call, so calculation work + cannot wait behind the parent calls that are collecting it. These are + per-function limits, not one shared quota. - Stage 12 owns its partition, request and result models, orchestration, merge, and cache implementation. It does not import the corresponding v1 modules. - The v2 functions write canonical private artifacts and use a shared @@ -365,6 +368,56 @@ functions and the household functions, not the coordinator or the economy workers, so it does not catch such a mismatch; a staging economy report run on each staging deployment would. +### Budget windows + +A budget-window report compares the same baseline and reform in each year of a +window of 1 to 75 consecutive years. It is an economy report whose +`budget_window_size` is the window's length. Its calculation is one +`budget_window_runs` row, and each year is a report run naming that row and +the year, with its own simulation runs of that year, so a year can reuse a +single-year report's runs, and the reverse. + +- **Creating and starting.** The simulation service creates every year's runs + in one transaction, all pending with no call, then starts the earliest + years within one ten-second wait for Modal, counted from the first start; + a later year not confirmed within it is deferred (`YearStartDeferred`) + and waits for a started year's coordinator. The contract's + `year_calculation` derives each year's coordinator + input and context from the first year's: only the year, its three runs, the + artifact prefix (`report_run_artifact_prefix`) and the call's creation time + differ. +- **At most five at once.** `BudgetWindowStore.advance_window` locks the + window run's row, then starts the earliest waiting years (pending with no + call) until five are under way (running, or started with a call recorded), + recording each call while the year is still pending with no call. The lock + makes advances that run at once start each year once; a year a duplicate + call reaches is claimed by only one call. +- **Each year's coordinator.** `coordinate_report` runs every call through + `coordinate_window_year`. It looks the year's window up once, marks the + window running before the coordinator runs, and advances the window once the + coordinator returns, fails or is cancelled. The advance starts the next + years in the same application, by `ModalYearStarter`, or finishes the + window. A pre-emption advances nothing; the restarted call advances when it + ends, even when it finds its year already finished. A report run outside any + window makes no window writes, and no window write ever changes the year's + own outcome. +- **The window's outcome.** A failed year, or a start Modal refuses, fails the + window, its waiting years and their unstarted simulation runs, and no + further year starts; years already started run to their end. A start whose + year its own call claimed meanwhile counts as started, and a deferred + start fails nothing. Once every + year has succeeded the window succeeds. +- **The stuck-run sweep.** It leaves a waiting year of an unfinished window to + the window, counts a year started later from when its call was recorded, + fails the window when it fails one of its years, and settles a window that + nothing has changed for ten minutes with no year under way: from its years + when they have all ended, or as stalled (`budget_window_stalled`) when years + are still waiting. + +The coordinator may run 25 calls at once (three windows at five years, and ten +other reports), so each simulation worker may run 50 and the US segment +function 1,000. + ## Report output planning Before starting either child simulation, the Modal report coordinator resolves @@ -926,9 +979,9 @@ This application registers no periodic cleanup function and does not delete the temporary PostgreSQL rows. Those rows remain restricted operational history until Stage 14 removes the temporary schema. -Household computation, budget-window computation, public response fields, -public identifiers, and production persistence are outside this companion -change. +Household computation, public response fields, public identifiers, and +production persistence are outside this companion change; budget windows on +the detailed-save path are described under "Budget windows" above. ## Numerical qualification diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py index 6006c63ea..982db91b3 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/__init__.py @@ -28,6 +28,11 @@ StorageUri, StrictContractModel, ) +from policyengine_simulation_contract.stage12_execution.budget_window_years import ( + report_run_artifact_prefix, + require_artifact_environment, + year_calculation, +) from policyengine_simulation_contract.stage12_execution.comparison_runs import ( MAX_COMPARISON_RUN_ARTIFACT_RETENTION, ComparisonReportPersistenceResult, @@ -112,5 +117,8 @@ "StrictContractModel", "UKLocalAuthorityBoundaryVersion", "UKLocalAuthorityMetadata", + "report_run_artifact_prefix", + "require_artifact_environment", "stage12_output_plan_sha256", + "year_calculation", ] diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/budget_window_years.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/budget_window_years.py new file mode 100644 index 000000000..1c9cd3c38 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution/budget_window_years.py @@ -0,0 +1,126 @@ +"""The coordinator input of each year of a budget window, derived from the first. + +A budget window's years compare the same policies on the same dataset, +region, options and bundle; only the year and the runs each year records on +differ. The simulation service builds the first year's coordinator input, +and the input of every other year, whether the service or an earlier year's +coordinator starts that year, is derived from it here, so all years are +calculated alike. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.stage12_execution.comparison_runs import ( + Stage12InvocationContext, +) +from policyengine_simulation_contract.stage12_execution.report_inputs import ( + ReportExecutionInput, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_contract.stage12_execution.base import ( + SimulationExecutionInput, + ) + from policyengine_simulation_contract.stage12_execution.comparison_runs import ( + ReportRunIds, + ) + + +def require_artifact_environment(environment: str) -> None: + """Raise ValueError for an environment that could leave its artifact prefix.""" + if not environment or "/" in environment or ".." in environment: + msg = "invalid Stage 12 artifact environment" + raise ValueError(msg) + + +def report_run_artifact_prefix( + environment: str, created_at: datetime, run_id: UUID +) -> str: + """Return where a run's Stage 12 artifacts go, by environment, month and run. + + The prefix is ``stage-12-runs////``. Raises + ValueError for an environment that could leave the prefix, or a + ``created_at`` without a time zone. + """ + require_artifact_environment(environment) + if created_at.tzinfo is None: + msg = "Stage 12 artifact timestamp must include a timezone" + raise ValueError(msg) + return f"stage-12-runs/{environment}/{created_at:%Y}/{created_at:%m}/{run_id}" + + +def year_calculation( + report: ReportExecutionInput, + context: Stage12InvocationContext, + *, + year: int, + run_ids: ReportRunIds, + created_at: datetime, +) -> tuple[ReportExecutionInput, Stage12InvocationContext]: + """Return the coordinator input and context of ``year``, recorded on ``run_ids``. + + ``report`` and ``context`` are another year's: the same window's first + year, or the year whose coordinator starts this one. The year, the three + runs, the artifact prefix and the call's creation time change; the + retention period stays as long as the other year's. Every other field is + kept, and the result is validated as any coordinator input is, so a + year outside the contract's range raises ValueError. A context naming no + report run raises ValueError: only a report run's coordinator records a + budget window's year. + """ + if context.report_run is None: + msg = "a call that names no report run starts no budget-window year" + raise ValueError(msg) + report_run_id = run_ids.report_run_id + year_report = ReportExecutionInput.model_validate( + { + **report.model_dump(mode="json"), + "evaluation_id": str(report_run_id), + "baseline": _year_side( + report.baseline, + year=year, + report_run_id=report_run_id, + simulation_run_id=run_ids.baseline_simulation_run_id, + ), + "reform": _year_side( + report.reform, + year=year, + report_run_id=report_run_id, + simulation_run_id=run_ids.reform_simulation_run_id, + ), + } + ) + retention = context.retention_expires_at - context.created_at + year_context = Stage12InvocationContext.model_validate( + { + **context.model_dump(mode="json"), + "report_run": run_ids.model_dump(mode="json"), + "artifact_prefix": report_run_artifact_prefix( + context.environment, created_at, report_run_id + ), + "created_at": created_at.isoformat(), + "retention_expires_at": (created_at + retention).isoformat(), + } + ) + return year_report, year_context + + +def _year_side( + side: SimulationExecutionInput, + *, + year: int, + report_run_id: UUID, + simulation_run_id: UUID, +) -> dict[str, object]: + """Return one side's input for the year, recorded on its simulation run.""" + return { + **side.model_dump(mode="json"), + "year": year, + "evaluation_id": str(report_run_id), + "simulation_execution_id": str(simulation_run_id), + } diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_budget_window_years.py b/libs/policyengine-simulation-contract/tests/test_stage12_budget_window_years.py new file mode 100644 index 000000000..3797fabbd --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_stage12_budget_window_years.py @@ -0,0 +1,228 @@ +"""Each year of a budget window is calculated as its first year is, in its own year. + +A budget window's years differ only in their year and the runs they record +on: the same policies, dataset, region, options and bundle. The simulation +service builds the first year's coordinator input, and every other year's, +whether the service or a coordinator starts it, is derived from it. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from uuid import UUID + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + ReportAggregate, + ReportExecutionInput, + ReportRunIds, + SimulationExecutionInput, + SimulationRole, + Stage12InvocationContext, + report_run_artifact_prefix, + year_calculation, +) + +FIRST_CREATED_AT = datetime(2026, 11, 12, 9, tzinfo=UTC) +LATER_CREATED_AT = datetime(2026, 12, 1, 3, tzinfo=UTC) +DATASET_URI = "hf://policyengine/data/enhanced_cps_2024.h5@revision" +FIRST_YEAR_RUNS = ReportRunIds( + report_run_id=UUID("00000000-0000-0000-0000-0000000000a1"), + baseline_simulation_run_id=UUID("00000000-0000-0000-0000-0000000000b1"), + reform_simulation_run_id=UUID("00000000-0000-0000-0000-0000000000c1"), +) +LATER_YEAR_RUNS = ReportRunIds( + report_run_id=UUID("00000000-0000-0000-0000-0000000000a4"), + baseline_simulation_run_id=UUID("00000000-0000-0000-0000-0000000000b4"), + reform_simulation_run_id=UUID("00000000-0000-0000-0000-0000000000c4"), +) + + +def _side(role: SimulationRole, simulation_run_id: UUID) -> SimulationExecutionInput: + return SimulationExecutionInput( + evaluation_id=FIRST_YEAR_RUNS.report_run_id, + simulation_execution_id=simulation_run_id, + role=role, + policy=( + {} + if role is SimulationRole.BASELINE + else {"gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01": 3000}} + ), + population=DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=DATASET_URI, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256="d" * 64, + ) + ), + year=2026, + geography=GeographySelection(country="us", region="us"), + options={"include_cliffs": True}, + bundle=BundleProvenance( + policyengine_version="6.17.1", + country_package_name="policyengine-us", + country_package_version="1.500.0", + dataset=DatasetProvenance( + identity="enhanced_cps_2024", + uri=DATASET_URI, + artifact_revision="revision", + data_package_name="policyengine-us-data", + data_package_version="1.2.0", + ), + bundle_manifest_sha256="b" * 64, + ), + ) + + +def _first_year() -> tuple[ReportExecutionInput, Stage12InvocationContext]: + report = ReportExecutionInput( + evaluation_id=FIRST_YEAR_RUNS.report_run_id, + baseline=_side( + SimulationRole.BASELINE, FIRST_YEAR_RUNS.baseline_simulation_run_id + ), + reform=_side(SimulationRole.REFORM, FIRST_YEAR_RUNS.reform_simulation_run_id), + requested_aggregates=tuple(ReportAggregate), + ) + context = Stage12InvocationContext( + request_id="request-1", + environment="staging", + modal_environment="staging", + worker_version="6.17.1", + modal_application="policyengine-simulation-v2-py6-17-1", + simulation_callable="run_single_simulation_us", + version_manifest_sha256="a" * 64, + bundle_manifest_sha256="b" * 64, + artifact_prefix=report_run_artifact_prefix( + "staging", FIRST_CREATED_AT, FIRST_YEAR_RUNS.report_run_id + ), + created_at=FIRST_CREATED_AT, + retention_expires_at=FIRST_CREATED_AT + timedelta(days=30), + report_run=FIRST_YEAR_RUNS, + ) + return report, context + + +def test_the_artifact_prefix_names_the_environment_month_and_run() -> None: + assert report_run_artifact_prefix( + "production", LATER_CREATED_AT, LATER_YEAR_RUNS.report_run_id + ) == (f"stage-12-runs/production/2026/12/{LATER_YEAR_RUNS.report_run_id}") + + +@pytest.mark.parametrize("environment", ["", "staging/other", "../production"]) +def test_an_environment_that_could_leave_the_prefix_is_refused(environment) -> None: + with pytest.raises(ValueError, match="invalid Stage 12 artifact environment"): + report_run_artifact_prefix( + environment, LATER_CREATED_AT, LATER_YEAR_RUNS.report_run_id + ) + + +def test_a_creation_time_without_a_time_zone_is_refused() -> None: + with pytest.raises(ValueError, match="timezone"): + report_run_artifact_prefix( + "staging", + LATER_CREATED_AT.replace(tzinfo=None), + LATER_YEAR_RUNS.report_run_id, + ) + + +def test_a_later_year_is_calculated_on_its_own_year_and_runs() -> None: + report, context = _first_year() + + later_report, later_context = year_calculation( + report, + context, + year=2029, + run_ids=LATER_YEAR_RUNS, + created_at=LATER_CREATED_AT, + ) + + assert later_report.evaluation_id == LATER_YEAR_RUNS.report_run_id + for side, simulation_run_id in ( + (later_report.baseline, LATER_YEAR_RUNS.baseline_simulation_run_id), + (later_report.reform, LATER_YEAR_RUNS.reform_simulation_run_id), + ): + assert side.year == 2029 + assert side.evaluation_id == LATER_YEAR_RUNS.report_run_id + assert side.simulation_execution_id == simulation_run_id + assert later_context.report_run == LATER_YEAR_RUNS + assert later_context.artifact_prefix == report_run_artifact_prefix( + "staging", LATER_CREATED_AT, LATER_YEAR_RUNS.report_run_id + ) + assert later_context.created_at == LATER_CREATED_AT + # The first year's retention period, counted from the later year's start. + assert later_context.retention_expires_at == LATER_CREATED_AT + timedelta(days=30) + + +def test_nothing_else_differs_between_the_years() -> None: + report, context = _first_year() + + later_report, later_context = year_calculation( + report, + context, + year=2029, + run_ids=LATER_YEAR_RUNS, + created_at=LATER_CREATED_AT, + ) + + # Put back the year and the runs: what remains is the first year's input. + restored_report = later_report.model_copy( + update={ + "evaluation_id": report.evaluation_id, + "baseline": later_report.baseline.model_copy( + update={ + "year": 2026, + "evaluation_id": report.evaluation_id, + "simulation_execution_id": report.baseline.simulation_execution_id, + } + ), + "reform": later_report.reform.model_copy( + update={ + "year": 2026, + "evaluation_id": report.evaluation_id, + "simulation_execution_id": report.reform.simulation_execution_id, + } + ), + } + ) + restored_context = later_context.model_copy( + update={ + "report_run": context.report_run, + "artifact_prefix": context.artifact_prefix, + "created_at": context.created_at, + "retention_expires_at": context.retention_expires_at, + } + ) + assert restored_report == report + assert restored_context == context + + +def test_a_call_naming_no_report_run_starts_no_year() -> None: + report, context = _first_year() + + with pytest.raises(ValueError, match="names no report run"): + year_calculation( + report, + context.model_copy(update={"report_run": None}), + year=2029, + run_ids=LATER_YEAR_RUNS, + created_at=LATER_CREATED_AT, + ) + + +def test_the_derived_year_is_checked_as_any_input_is() -> None: + report, context = _first_year() + + with pytest.raises(ValueError, match="year"): + year_calculation( + report, + context, + year=2201, + run_ids=LATER_YEAR_RUNS, + created_at=LATER_CREATED_AT, + ) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 4fea4fd37..d6ac6ee63 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -227,9 +227,12 @@ dynamic, household and catalog tables a run reads. `SERVICE_UPDATABLE_COLUMNS` lists the columns it may `UPDATE` on `simulation_runs`, `report_runs` and `household_jobs`: `status`, `modal_call_id`, `error_message`, `completed_at` and `updated_at`, to record a Modal call's ID once Modal accepts the call, or -to fail the run when Modal refuses it. The role holds `UPDATE` on those -columns only, never on a whole table, so the database refuses any other change -to a run. The Public API's privilege manifest grants them all. +to fail the run when Modal refuses it. On `budget_window_runs` it lists +`status`, `error_message`, `completed_at` and `updated_at`, to fail a window +when Modal refuses one of the first years the service starts. The role holds +`UPDATE` on those columns only, never on a whole table, so the database +refuses any other change to a run. The Public API's privilege manifest grants +them all. `missing_table_privileges` returns each one the connected role lacks, ordered by table, as `"INSERT on reports"` or `"UPDATE (modal_call_id) on diff --git a/libs/policyengine-simulation-persistence/public-api-commit b/libs/policyengine-simulation-persistence/public-api-commit index 8e64693d5..90f10c4b5 100644 --- a/libs/policyengine-simulation-persistence/public-api-commit +++ b/libs/policyengine-simulation-persistence/public-api-commit @@ -1 +1 @@ -8819e66c485bcc696c46d40894d8e720dabe4459 +429d348aeff84d1d345799fc93ee7c02f7e2014a 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..bea874cc0 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -6,6 +6,11 @@ missing_table_privileges, read_missing_table_privileges, ) +from policyengine_simulation_persistence.budget_window_records import ( + BudgetWindowYear, + NewBudgetWindowReport, + insert_budget_window_report, +) from policyengine_simulation_persistence.economy_run_reads import ( EconomyReportState, EconomySimulationState, @@ -38,6 +43,7 @@ record_simulation_run_call, ) from policyengine_simulation_persistence.economy_run_store import ( + BudgetWindowInputs, EconomyCatalogRows, EconomyRecordOutcomeUnknownError, EconomyReportInputs, @@ -87,6 +93,8 @@ read_policy_reform, ) from policyengine_simulation_persistence.tables import ( + LATEST_YEAR, + MAX_BUDGET_WINDOW_YEARS, HouseholdJobStatus, ReportRunStatus, ReportRunTrigger, @@ -103,8 +111,12 @@ ) __all__ = [ + "LATEST_YEAR", + "MAX_BUDGET_WINDOW_YEARS", "SERVICE_TABLE_PRIVILEGES", "SERVICE_UPDATABLE_COLUMNS", + "BudgetWindowInputs", + "BudgetWindowYear", "CatalogDataset", "CatalogRegion", "EconomyCatalogRows", @@ -123,6 +135,7 @@ "HouseholdJobStatus", "HouseholdJobStore", "MissingPolicyParametersError", + "NewBudgetWindowReport", "NewEconomyReport", "NewEconomySimulation", "NewHouseholdJob", @@ -160,6 +173,7 @@ "find_reusable_run", "household_jobs", "input_metadata", + "insert_budget_window_report", "insert_economy_report", "insert_economy_simulation", "insert_household_job", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/access.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/access.py index 6a67cd22a..bafb7799e 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/access.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/access.py @@ -81,6 +81,11 @@ _CALL_START: Final = frozenset( {"status", "modal_call_id", "error_message", "completed_at", "updated_at"} ) +# The state the service records for a budget-window run, which has no Modal +# call of its own: the failure when Modal refuses one of its first years. +_WINDOW_FAILURE: Final = frozenset( + {"status", "error_message", "completed_at", "updated_at"} +) # The columns the service's role may UPDATE, by table. It holds UPDATE on # these columns alone, never on a whole table, so the database refuses any # other change to a run. @@ -89,6 +94,7 @@ "simulation_runs": _CALL_START, "report_runs": _CALL_START, "household_jobs": _CALL_START, + "budget_window_runs": _WINDOW_FAILURE, } ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py new file mode 100644 index 000000000..111a025ff --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py @@ -0,0 +1,182 @@ +"""The rows of a budget-window report: one report run per year of its window. + +A budget-window report compares the same baseline and reform in each year +of a window of 1 to 75 consecutive years. It is stored as an economy report +whose ``budget_window_size`` is the window's length, with one +``budget_window_runs`` row for its calculation and, for each year, a report +run naming that row and the year. Each year has its own baseline and reform +simulations, of that year, unless it reuses a run that already succeeded, +so a year's run can be reused by a single-year report, and the reverse. The +report names its first year's simulations. + +The simulation service creates every row pending with no Modal call, then +starts the first years; each year's coordinator starts the next ones (see +``policyengine_stage12_persistence.BudgetWindowStore``). +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any + +from sqlalchemy import insert + +from policyengine_simulation_persistence.economy_run_records import ( + ECONOMY_REPORT_TYPE, + EconomyPopulation, + ReportSimulation, + RunVersions, + insert_simulation_with_run, + references_checked, +) +from policyengine_simulation_persistence.tables import ( + LATEST_YEAR, + MAX_BUDGET_WINDOW_YEARS, + ReportRunStatus, + ReportRunTrigger, + budget_window_runs, + report_runs, + reports, +) + +if TYPE_CHECKING: + from uuid import UUID + + from sqlalchemy import Connection + + +@dataclass(frozen=True) +class BudgetWindowYear: + """One year of a window: its report run and its two sides' runs.""" + + year: int + report_run_id: UUID + baseline: ReportSimulation + reform: ReportSimulation + + +@dataclass(frozen=True) +class NewBudgetWindowReport: + """A budget-window report to create, with its first calculation's runs. + + ``population`` is the first year's; each year's simulations use the same + catalog version, region and dataset in their own year. ``inputs`` holds + the report's options as the Public API reads them. ``years`` runs from + the first year, one year at a time, earliest first. + """ + + id: UUID + window_run_id: UUID + country: str + label: str + description: str | None + tax_benefit_model_id: UUID + population: EconomyPopulation + versions: RunVersions + inputs: dict[str, Any] + years: tuple[BudgetWindowYear, ...] + + def __post_init__(self) -> None: + """Refuse years that aren't 1 to 75 consecutive ones from the first.""" + first_year = self.population.year + expected = tuple(range(first_year, first_year + len(self.years))) + if ( + not self.years + or len(self.years) > MAX_BUDGET_WINDOW_YEARS + or tuple(year.year for year in self.years) != expected + or expected[-1] > LATEST_YEAR + ): + msg = ( + "a budget window's years must be 1 to " + f"{MAX_BUDGET_WINDOW_YEARS} consecutive years from its first, " + f"ending by {LATEST_YEAR}" + ) + raise ValueError(msg) + + @property + def first_year(self) -> BudgetWindowYear: + """The window's first year, whose simulations the report names.""" + return self.years[0] + + def population_in(self, year: int) -> EconomyPopulation: + """Return what the window's simulations of ``year`` calculate on.""" + return replace(self.population, year=year) + + +def insert_budget_window_report( + connection: Connection, report: NewBudgetWindowReport +) -> None: + """Create the report, its budget-window run, and each year's runs, all pending. + + Each year's sides that reuse no run get a new simulation and run of that + year. Raise EconomyRecordReferenceError when the database refuses a row + it names. The refusal ends the caller's transaction, or its savepoint. + """ + with references_checked(): + for window_year in report.years: + population = report.population_in(window_year.year) + for side in (window_year.baseline, window_year.reform): + if not side.reused: + insert_simulation_with_run( + connection, + simulation_id=side.simulation_id, + run_id=side.simulation_run_id, + choice=side.choice, + population=population, + versions=report.versions, + ) + first_year = report.first_year + population = report.population + connection.execute( + insert(reports).values( + id=report.id, + label=report.label, + description=report.description, + country=report.country, + type=ECONOMY_REPORT_TYPE, + user_id=None, + tax_benefit_model_id=report.tax_benefit_model_id, + policy_id=first_year.reform.choice.policy_id, + baseline_simulation_id=first_year.baseline.simulation_id, + reform_simulation_id=first_year.reform.simulation_id, + household_id=None, + dataset_id=population.dataset_id, + region_id=population.region_id, + year=population.year, + budget_window_size=len(report.years), + inputs=report.inputs, + ) + ) + connection.execute( + insert(budget_window_runs).values( + id=report.window_run_id, + report_id=report.id, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + trigger=ReportRunTrigger.INITIAL, + recalculate=False, + status=ReportRunStatus.PENDING, + ) + ) + connection.execute( + insert(report_runs), + [_year_run_row(report, window_year) for window_year in report.years], + ) + + +def _year_run_row( + report: NewBudgetWindowReport, window_year: BudgetWindowYear +) -> dict[str, Any]: + """Return a year's report run, pending with no call, as inserted.""" + return { + "id": window_year.report_run_id, + "report_id": report.id, + "baseline_simulation_run_id": window_year.baseline.simulation_run_id, + "reform_simulation_run_id": window_year.reform.simulation_run_id, + "budget_window_run_id": report.window_run_id, + "year": window_year.year, + "country_package_version": report.versions.country_package_version, + "policyengine_version": report.versions.policyengine_version, + "status": ReportRunStatus.PENDING, + "trigger": ReportRunTrigger.INITIAL, + "modal_call_id": None, + } diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py index dd90ddabd..10c7ea28a 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py @@ -60,6 +60,7 @@ ReportRunTrigger, SimulationStatus, SimulationType, + budget_window_runs, report_runs, reports, simulation_runs, @@ -86,7 +87,13 @@ # the first column it constrains. _KEY_COLUMNS: Final = { str(constraint.name): (table.name, constraint.column_keys[0]) - for table in (simulations, simulation_runs, reports, report_runs) + for table in ( + simulations, + simulation_runs, + reports, + report_runs, + budget_window_runs, + ) for constraint in table.foreign_key_constraints } @@ -374,10 +381,10 @@ def insert_economy_report(connection: Connection, report: NewEconomyReport) -> N Raise EconomyRecordReferenceError when the database refuses a row it names. The refusal ends the caller's transaction, or its savepoint. """ - with _references_checked(): + with references_checked(): for side in (report.baseline, report.reform): if not side.reused: - _insert_simulation( + insert_simulation_with_run( connection, simulation_id=side.simulation_id, run_id=side.simulation_run_id, @@ -431,8 +438,8 @@ def insert_economy_simulation( Raise EconomyRecordReferenceError when the database refuses a row it names. The refusal ends the caller's transaction, or its savepoint. """ - with _references_checked(): - _insert_simulation( + with references_checked(): + insert_simulation_with_run( connection, simulation_id=simulation.id, run_id=simulation.run_id, @@ -442,7 +449,7 @@ def insert_economy_simulation( ) -def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts +def insert_simulation_with_run( # noqa: PLR0913 # the simulation's and its run's parts connection: Connection, *, simulation_id: UUID, @@ -482,7 +489,7 @@ def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts @contextmanager -def _references_checked() -> Iterator[None]: +def references_checked() -> Iterator[None]: """Turn a foreign-key refusal into EconomyRecordReferenceError. Any other integrity error propagates unchanged. 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..2ba26499c 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 @@ -14,6 +14,9 @@ from sqlalchemy.exc import SQLAlchemyError +from policyengine_simulation_persistence.budget_window_records import ( + insert_budget_window_report, +) from policyengine_simulation_persistence.economy_run_reads import ( read_economy_report, read_economy_simulation, @@ -61,6 +64,9 @@ from sqlalchemy import Connection, Engine + from policyengine_simulation_persistence.budget_window_records import ( + NewBudgetWindowReport, + ) from policyengine_simulation_persistence.economy_run_reads import ( EconomyReportState, EconomySimulationState, @@ -199,6 +205,21 @@ class EconomyReportInputs: reform: SimulationInputs +@dataclass(frozen=True) +class BudgetWindowInputs: + """What a budget-window report is read from, and each year's reusable runs. + + ``baseline`` and ``reform`` hold each side's values, the same in every + year. ``reusable_runs`` gives, by year, the baseline's and the reform's + reusable run, each None when there is none. + """ + + catalog: EconomyCatalogRows + baseline: SimulationInputs + reform: SimulationInputs + reusable_runs: dict[int, tuple[ReusableRun | None, ReusableRun | None]] + + class EconomyRunStore: """Read and write economy simulations and reports, each in its own transaction.""" @@ -287,6 +308,71 @@ def read_report_inputs( sides.append(inputs) return EconomyReportInputs(catalog=catalog, baseline=sides[0], reform=sides[1]) + def read_budget_window_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + years: Sequence[int], + ) -> BudgetWindowInputs: + """Read a window's catalog rows and both sides' values, then each year's runs. + + ``request.year`` is the window's first year. With ``reuse``, each + year's sides get the runs ``find_reusable_run`` finds for that year, + and the reform never gets the baseline's. Raise the errors + ``read_report_inputs`` raises. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + baseline_inputs = _side_inputs( + connection, request, catalog, "baseline", baseline + ) + reform_inputs = _side_inputs(connection, request, catalog, "reform", reform) + reusable_runs: dict[int, tuple[ReusableRun | None, ReusableRun | None]] = {} + for year in years: + baseline_run = reform_run = None + if reuse is not None: + population = catalog.population(year) + baseline_run = find_reusable_run( + connection, + choice=baseline, + population=population, + versions=reuse.versions, + inputs=reuse.inputs, + output_uri_prefix=reuse.output_uri_prefix, + ) + reform_run = find_reusable_run( + connection, + choice=reform, + population=population, + versions=reuse.versions, + inputs=reuse.inputs, + output_uri_prefix=reuse.output_uri_prefix, + excluded_run_id=( + None + if baseline_run is None + else baseline_run.simulation_run_id + ), + ) + reusable_runs[year] = (baseline_run, reform_run) + return BudgetWindowInputs( + catalog=catalog, + baseline=baseline_inputs, + reform=reform_inputs, + reusable_runs=reusable_runs, + ) + + def create_budget_window_report(self, report: NewBudgetWindowReport) -> None: + """Create the window's report, its budget-window run and each year's runs. + + Raise ``insert_budget_window_report``'s errors, and + EconomyRecordOutcomeUnknownError when the commit fails after the + rows were inserted. + """ + self._create(lambda connection: insert_budget_window_report(connection, report)) + def create_report(self, report: NewEconomyReport) -> None: """Create the report, its new simulations and its runs. @@ -465,3 +551,19 @@ def _simulation_inputs( return SimulationInputs( choice=choice, reform=combine_reforms(policy_reform, dynamic_reform) ) + + +def _side_inputs( + connection: Connection, + request: EconomyRequest, + catalog: EconomyCatalogRows, + side: ReportSide, + choice: SimulationChoice, +) -> SimulationInputs: + """Read one report side's values; raise ReportSideError naming a refused side.""" + try: + return _simulation_inputs( + connection, request.country, catalog.catalog_version_id, choice + ) + except _SIDE_ERRORS as error: + raise ReportSideError(side, error) from error diff --git a/libs/policyengine-simulation-persistence/tests/test_access.py b/libs/policyengine-simulation-persistence/tests/test_access.py index 69852af37..a1be6d2ca 100644 --- a/libs/policyengine-simulation-persistence/tests/test_access.py +++ b/libs/policyengine-simulation-persistence/tests/test_access.py @@ -48,6 +48,11 @@ CALL_START_COLUMNS = frozenset( {"status", "modal_call_id", "error_message", "completed_at", "updated_at"} ) +# A budget-window run has no Modal call of its own: the service fails it when +# Modal refuses one of the first years it starts. +WINDOW_FAILURE_COLUMNS = frozenset( + {"status", "error_message", "completed_at", "updated_at"} +) def test_the_service_may_create_and_read_rows_in_every_table_it_writes() -> None: @@ -62,11 +67,15 @@ def test_the_service_only_reads_the_rows_a_run_needs() -> None: def test_the_service_updates_only_the_call_state_of_each_run() -> None: - """It records a Modal call's ID, or fails a run Modal refuses, and no more.""" + """It records a Modal call's ID, or fails a run Modal refuses, and no more. + + It fails a budget window when Modal refuses one of its first years. + """ assert dict(SERVICE_UPDATABLE_COLUMNS) == { "simulation_runs": CALL_START_COLUMNS, "report_runs": CALL_START_COLUMNS, "household_jobs": CALL_START_COLUMNS, + "budget_window_runs": WINDOW_FAILURE_COLUMNS, } diff --git a/libs/policyengine-simulation-persistence/tests/test_budget_window_records.py b/libs/policyengine-simulation-persistence/tests/test_budget_window_records.py new file mode 100644 index 000000000..a08d4766f --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_budget_window_records.py @@ -0,0 +1,361 @@ +"""A budget-window report's inputs and rows: one report run per year. + +A budget window compares the same baseline and reform in each year of its +window. Each year has its own simulations, with that year, so a year can +reuse a single-year report's run, and the reverse. These run against a +database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), acting as the simulation service's +role in a transaction the test rolls back. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import replace +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + POLICYENGINE_VERSION, + VERSIONS, + RunShape, + insert_economy_catalog, + insert_report_run_naming, + insert_simulation_run, +) +from policyengine_simulation_persistence import ( + BudgetWindowYear, + EconomyPopulation, + EconomyRecordReferenceError, + EconomyRequest, + EconomyRunStore, + NewBudgetWindowReport, + ReportSimulation, + ReusableRun, + RunReuse, + SimulationChoice, + SimulationStatus, +) +from public_api_rows import StoredValue, insert_policy +from sqlalchemy import text + +if TYPE_CHECKING: + from collections.abc import Iterator + + from economy_rows import EconomyCatalog + 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" +REQUEST = EconomyRequest( + country="us", + policyengine_version=POLICYENGINE_VERSION, + region_code="us", + dataset_name=None, + year=2026, +) +REUSE = RunReuse( + versions=VERSIONS, inputs=OPTIONS, output_uri_prefix=OUTPUT_BUCKET_PREFIX +) +WINDOW_YEARS = (2026, 2027, 2028) + + +def _store_on(connection: Connection) -> EconomyRunStore: + """Return a store whose transactions are savepoints of ``connection``.""" + + @contextmanager + def savepoint() -> Iterator[Connection]: + with connection.begin_nested(): + yield connection + + return EconomyRunStore(savepoint) + + +def _reusable_run_in(connection: Connection, economy: EconomyCatalog, year: int): + """Insert a succeeded current-law run of ``year`` that a report run used.""" + population = economy.population(year=year) + reused = insert_simulation_run(connection, population) + other_side = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming(connection, economy, (reused, other_side)) + return reused + + +def test_a_window_reads_each_year_s_reusable_runs(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(CTC_AMOUNT,)) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + reused_in_2027 = _reusable_run_in(connection, economy, 2027) + act_as_the_service(connection) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + + inputs = _store_on(connection).read_budget_window_inputs( + REQUEST, + baseline=CURRENT_LAW, + reform=reform_choice, + reuse=REUSE, + years=WINDOW_YEARS, + ) + + assert inputs.catalog.catalog_version_id == economy.catalog.version_id + assert (inputs.baseline.choice, inputs.baseline.reform) == (CURRENT_LAW, {}) + assert inputs.reform.choice == reform_choice + assert inputs.reusable_runs == { + 2026: (None, None), + 2027: ( + ReusableRun( + simulation_id=reused_in_2027.simulation_id, + simulation_run_id=reused_in_2027.simulation_run_id, + ), + None, + ), + 2028: (None, None), + } + + +def test_a_window_reads_no_reusable_run_without_reuse(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + _reusable_run_in(connection, economy, 2027) + act_as_the_service(connection) + + inputs = _store_on(connection).read_budget_window_inputs( + REQUEST, + baseline=CURRENT_LAW, + reform=CURRENT_LAW, + reuse=None, + years=WINDOW_YEARS, + ) + + assert set(inputs.reusable_runs.values()) == {(None, None)} + + +def _new_side(choice: SimulationChoice = CURRENT_LAW) -> ReportSimulation: + return ReportSimulation( + simulation_id=uuid4(), simulation_run_id=uuid4(), choice=choice, reused=False + ) + + +def _window(economy: EconomyCatalog, reused_in_2027) -> NewBudgetWindowReport: + years = [] + for year in WINDOW_YEARS: + baseline = ( + ReportSimulation( + simulation_id=reused_in_2027.simulation_id, + simulation_run_id=reused_in_2027.simulation_run_id, + choice=CURRENT_LAW, + reused=True, + ) + if year == 2027 + else _new_side() + ) + years.append( + BudgetWindowYear( + year=year, + report_run_id=uuid4(), + baseline=baseline, + reform=_new_side(), + ) + ) + return NewBudgetWindowReport( + id=uuid4(), + window_run_id=uuid4(), + country="us", + label="Ten-year budget window", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(year=2026), + versions=VERSIONS, + inputs=OPTIONS, + years=tuple(years), + ) + + +def _rows(connection: Connection, table_name: str, **filters): + rows = connection.execute( + text( + f"SELECT * FROM {table_name} WHERE " # noqa: S608 # test-only names + + " AND ".join(f"{name} = :{name}" for name in filters) + ), + filters, + ) + return [dict(row._mapping) for row in rows] + + +def test_a_window_is_created_with_a_report_run_per_year( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused_in_2027 = _reusable_run_in(connection, economy, 2027) + window = _window(economy, reused_in_2027) + act_as_the_service(connection) + + _store_on(connection).create_budget_window_report(window) + + connection.execute(text("RESET ROLE")) + (report,) = _rows(connection, "reports", id=window.id) + (window_run,) = _rows(connection, "budget_window_runs", id=window.window_run_id) + year_runs = { + row["year"]: row + for row in _rows( + connection, "report_runs", budget_window_run_id=window.window_run_id + ) + } + simulations = { + run_id: ( + _rows(connection, "simulation_runs", id=run_id)[0], + _rows( + connection, + "simulations", + id=_rows(connection, "simulation_runs", id=run_id)[0][ + "simulation_id" + ], + )[0], + ) + for year in window.years + for run_id in ( + year.baseline.simulation_run_id, + year.reform.simulation_run_id, + ) + } + + first_year = window.years[0] + assert ( + report["type"], + report["year"], + report["budget_window_size"], + report["inputs"], + report["baseline_simulation_id"], + report["reform_simulation_id"], + report["household_id"], + ) == ( + "economy", + 2026, + 3, + OPTIONS, + first_year.baseline.simulation_id, + first_year.reform.simulation_id, + None, + ) + assert ( + window_run["report_id"], + window_run["status"], + window_run["trigger"], + window_run["recalculate"], + window_run["tax_benefit_model_version_id"], + ) == (window.id, "pending", "initial", False, economy.catalog.version_id) + for year in window.years: + year_run = year_runs[year.year] + assert ( + year_run["id"], + year_run["report_id"], + year_run["status"], + year_run["trigger"], + year_run["modal_call_id"], + year_run["baseline_simulation_run_id"], + year_run["reform_simulation_run_id"], + year_run["country_package_version"], + year_run["policyengine_version"], + ) == ( + year.report_run_id, + window.id, + "pending", + "initial", + None, + year.baseline.simulation_run_id, + year.reform.simulation_run_id, + VERSIONS.country_package_version, + VERSIONS.policyengine_version, + ) + for side in (year.baseline, year.reform): + simulation_run, simulation = simulations[side.simulation_run_id] + # Each year has its own simulations, of that year. + assert (simulation["year"], simulation_run["year"]) == ( + year.year, + year.year, + ) + assert simulation_run["status"] == ( + "succeeded" if side.reused else "pending" + ) + + +def test_a_window_naming_a_missing_policy_is_refused_and_writes_nothing( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused_in_2027 = _reusable_run_in(connection, economy, 2027) + window = _window(economy, reused_in_2027) + missing_policy = SimulationChoice(policy_id=uuid4(), dynamic_id=None) + last_year = window.years[-1] + window = replace( + window, + years=( + *window.years[:-1], + replace( + last_year, reform=replace(last_year.reform, choice=missing_policy) + ), + ), + ) + act_as_the_service(connection) + + with pytest.raises(EconomyRecordReferenceError) as refused: + _store_on(connection).create_budget_window_report(window) + + connection.execute(text("RESET ROLE")) + written = connection.execute( + text("SELECT count(*) FROM reports WHERE id = :id"), {"id": window.id} + ).scalar_one() + + assert (refused.value.table_name, refused.value.column_name) == ( + "simulations", + "policy_id", + ) + assert written == 0 + + +@pytest.mark.parametrize( + "years", + [(), (2026, 2028), tuple(range(2026, 2102))], + ids=["no-years", "a-gap", "76-years"], +) +def test_a_window_s_years_must_be_consecutive_from_its_first(years) -> None: + population = EconomyPopulation( + tax_benefit_model_version_id=uuid4(), + region_id=uuid4(), + dataset_id=uuid4(), + year=years[0] if years else 2026, + ) + with pytest.raises(ValueError, match="budget window"): + NewBudgetWindowReport( + id=uuid4(), + window_run_id=uuid4(), + country="us", + label="Window", + description=None, + tax_benefit_model_id=uuid4(), + population=population, + versions=VERSIONS, + inputs=OPTIONS, + years=tuple( + BudgetWindowYear( + year=year, + report_run_id=uuid4(), + baseline=_new_side(), + reform=_new_side(), + ) + for year in years + ), + ) diff --git a/libs/policyengine-stage12-persistence/README.md b/libs/policyengine-stage12-persistence/README.md index c0d683f29..686c9d944 100644 --- a/libs/policyengine-stage12-persistence/README.md +++ b/libs/policyengine-stage12-persistence/README.md @@ -30,12 +30,25 @@ revisions and live schema. This package never changes a schema. back to pending with no call (`release_simulation_run`). A report run's completion sets its own result columns and inserts its result rows in the same transaction. +- **Budget windows** (`BudgetWindowStore`): `budget_window_runs`, one row per + calculation of a budget-window report, whose years are report runs naming + it and their year. The simulation service creates every year pending and + starts the first ones; each year's coordinator, when its year ends, + advances the window. An advance locks the window run's row, then fails the + window and its waiting years when a year failed, marks it succeeded when + every year succeeded, or starts the earliest waiting years until + `BUDGET_WINDOW_MAX_RUNNING_YEARS` (five) are under way, recording each + call while the year is still pending with no call. A year's coordinator + call marks its window running. The stuck-run sweep leaves a waiting year + to its window, fails the window of a year it fails, and settles a window + no year advances any more. The Modal worker role and the simulation + service's role may both advance a window. ## Tests -The report run tests need a disposable Postgres migrated by the Public API's -v2 chain, with its declared privileges applied. They skip unless both URLs -are set: +The report run and budget-window tests need a disposable Postgres migrated +by the Public API's v2 chain, with its declared privileges applied. They skip +unless both URLs are set: - `REPORT_RUN_TEST_DATABASE_URL`: logs in as `policyengine_v2_modal_worker`; - `REPORT_RUN_TEST_SEED_DATABASE_URL`: logs in as a role that may insert the diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index c0bd88fa5..3f1bca46a 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -1,5 +1,13 @@ -"""Direct DML for temporary Stage 12 comparison state, runs and household jobs.""" +"""Direct DML for Stage 12 comparison state, runs, budget windows and household jobs.""" +from policyengine_stage12_persistence.budget_window_store import ( + BUDGET_WINDOW_MAX_RUNNING_YEARS, + BudgetWindowStore, + WindowAdvance, + YearStartDeferred, + YearStarter, + YearToStart, +) from policyengine_stage12_persistence.report_run_results import ( DistrictCodes, ReportRunResult, @@ -25,7 +33,9 @@ ) __all__ = [ + "BUDGET_WINDOW_MAX_RUNNING_YEARS", "ActiveRun", + "BudgetWindowStore", "ConnectionOpener", "DistrictCodes", "ReportRunResult", @@ -36,6 +46,10 @@ "SimulationRunState", "Stage12PersistenceStore", "Stage12RuntimeAccessError", + "WindowAdvance", + "YearStartDeferred", + "YearStarter", + "YearToStart", "create_stage12_engine", "normalized_database_url", "report_run_result", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/budget_window_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/budget_window_store.py new file mode 100644 index 000000000..4a025babc --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/budget_window_store.py @@ -0,0 +1,601 @@ +"""Starting a budget window's years and recording the window's outcome. + +A budget-window run is one calculation of a budget-window report; its years +are report runs naming it and their year. The simulation service creates +every year pending with no Modal call, then starts the first ones. When a +year's coordinator ends its year, it starts the earliest years still +waiting, so that at most ``BUDGET_WINDOW_MAX_RUNNING_YEARS`` are under way at +once. A failed year fails the window and every year still waiting, with no +further years started, and once every year has succeeded the window is +marked succeeded. + +Each advance locks the budget-window run's row while it starts years and +records their calls, so coordinators ending at once start each year once. +The lock is held across those Modal calls; claims and the window's start +are separate one-statement transactions that never wait on each other, so +other years' coordinators wait only for the lock to be released. +The Modal worker role and the simulation service role may both advance a +window: each may update the columns written here. Only a year's coordinator +call marks a window running, before it claims its year, because only the +Modal worker role may set the window's start time. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from sqlalchemy import and_, exists, func, or_, select, update + +from policyengine_stage12_persistence.report_run_store import ( + ACTIVE_STATUSES, + RunStatus, + no_active_report_run_names, +) +from policyengine_stage12_persistence.report_run_tables import ( + budget_window_runs, + report_runs, + simulation_runs, +) +from policyengine_stage12_persistence.store import ( + ConnectionOpener, + create_stage12_engine, + open_connection_once, +) + +if TYPE_CHECKING: + from collections.abc import Iterator + from datetime import datetime + from uuid import UUID + + from sqlalchemy import ColumnElement, Engine + from sqlalchemy.engine import Connection + +logger = logging.getLogger(__name__) + +# The most years of one budget window under way at once: running, or started +# with a Modal call recorded and waiting for Modal to run it. +BUDGET_WINDOW_MAX_RUNNING_YEARS = 5 + +# Each failure records a fixed code, then a summary that holds no +# calculation detail; a failed year's own message stays on that year's run. +YEAR_FAILED_MESSAGE = ( + "budget_window_year_failed: The budget window's {year} report run failed." +) +WINDOW_FAILED_MESSAGE = ( + "budget_window_failed: Another year of the budget window failed before " + "this year started." +) +CALL_REFUSED_MESSAGE = ( + "modal_call_refused: The call that would have started this year could not " + "be started ({error_type})." +) +WINDOW_STALLED_MESSAGE = ( + "budget_window_stalled: No year of the budget window was under way to " + "start the years still waiting." +) + + +@dataclass(frozen=True) +class YearToStart: + """A waiting year of a window: its year and the three runs it records on.""" + + year: int + run_ids: ReportRunIds + + +class YearStartDeferred(Exception): # noqa: N818 # it names an outcome, not a fault + """A starter left this year, and the later ones, for a later advance. + + Raising it starts nothing more in this advance and fails nothing: the + years stay waiting, and the coordinator of a year already under way + starts them when its year ends. + """ + + +class YearStarter(Protocol): + """Starts a year's coordinator call and returns the call's Modal ID. + + It raises when Modal refuses the call, or ``YearStartDeferred`` to leave + the year waiting. A starter defers only once a year of the window is + under way, whose coordinator starts the waiting years when its year + ends; otherwise nothing starts them until the stuck-run sweep fails the + window as stalled. + """ + + def __call__(self, year_to_start: YearToStart, /) -> str: + """Start the year's coordinator call and return its Modal call ID.""" + ... + + +@dataclass(frozen=True) +class WindowAdvance: + """What one advance of a budget-window run did. + + ``status`` is the window's status afterwards, ``changed`` whether this + advance finished it, and ``started_years`` the years it started, earliest + first. + """ + + status: RunStatus + changed: bool + started_years: tuple[int, ...] = () + + +@dataclass(frozen=True) +class _WindowYear: + """One year of a window as its report run records it.""" + + run_ids: ReportRunIds + year: int + status: RunStatus + modal_call_id: str | None + updated_at: datetime + + @property + def report_run_id(self) -> UUID: + """The year's report run.""" + return self.run_ids.report_run_id + + @property + def waiting(self) -> bool: + """Whether the year is pending with no call: nothing has started it.""" + return self.status is RunStatus.PENDING and self.modal_call_id is None + + @property + def under_way(self) -> bool: + """Whether the year is running, or started and waiting for Modal to run it.""" + return self.status is RunStatus.RUNNING or ( + self.status is RunStatus.PENDING and self.modal_call_id is not None + ) + + +@dataclass(frozen=True) +class _LockedWindow: + """A budget-window run's status and last change, read under its row lock.""" + + status: RunStatus + updated_at: datetime + + @property + def finished(self) -> bool: + """Whether the window has succeeded or failed.""" + return self.status.value not in ACTIVE_STATUSES + + +class BudgetWindowStore: + """Budget-window runs' starts and outcomes in the Public API's run tables.""" + + def __init__( + self, + database_url: str | None = None, + *, + engine: Engine | None = None, + open_connection: ConnectionOpener = open_connection_once, + ) -> None: + """Use ``engine`` if given, else build one from ``database_url``. + + Raises ValueError when neither is given or the URL is invalid. + """ + if engine is None: + if database_url is None: + msg = "budget window database URL is required" + raise ValueError(msg) + engine = create_stage12_engine(database_url) + self.engine = engine + self._open_connection = open_connection + + @contextmanager + def _transaction(self, operation_name: str) -> Iterator[Connection]: + # The opener only opens the connection; every statement and the + # COMMIT run after it returns. + connection = self._open_connection( + self.engine.connect, + operation_name=operation_name, + ) + with connection, connection.begin(): + yield connection + + def window_of_year(self, report_run_id: UUID) -> UUID | None: + """Return the budget-window run a report run is a year of. + + None when the report run is no window's year, or does not exist. + """ + statement = select(report_runs.c.budget_window_run_id).where( + report_runs.c.id == report_run_id + ) + with self._transaction("window_of_year") as connection: + return connection.execute(statement).scalar_one_or_none() + + def start_window(self, report_run_id: UUID, *, started_at: datetime) -> bool: + """Mark the window a claimed year belongs to running; return whether it changed. + + Only a pending window changes, so the first year claimed sets its + start time. A report run that is no budget window's year changes + nothing. + """ + window_of_run = ( + select(report_runs.c.budget_window_run_id) + .where(report_runs.c.id == report_run_id) + .scalar_subquery() + ) + statement = ( + update(budget_window_runs) + .where( + budget_window_runs.c.id == window_of_run, + budget_window_runs.c.status == RunStatus.PENDING.value, + ) + .values( + status=RunStatus.RUNNING.value, + started_at=func.coalesce(budget_window_runs.c.started_at, started_at), + updated_at=started_at, + ) + .returning(budget_window_runs.c.id) + ) + with self._transaction("start_window") as connection: + return connection.execute(statement).scalar_one_or_none() is not None + + def advance_window( + self, + window_run_id: UUID, + *, + now: datetime, + start_year: YearStarter | None = None, + max_running_years: int = BUDGET_WINDOW_MAX_RUNNING_YEARS, + ) -> WindowAdvance: + """Finish the window if its years have ended, else start its waiting years. + + A finished window is left as it is. A failed year fails the window + and its waiting years, and every year succeeded marks it succeeded. + Otherwise, with ``start_year``, the earliest waiting years are + started until ``max_running_years`` are under way, and each started + year's call is recorded while the year is still pending with no call; + a call its coordinator already claimed under keeps that claim. A + start that raises fails that year as refused by Modal, and so the + window, unless the year's own call has claimed it meanwhile; a start + deferred with ``YearStartDeferred`` leaves that year and the later + ones waiting. Raises LookupError when the window run does not exist. + """ + with self._transaction("advance_window") as connection: + window = _lock_window(connection, window_run_id) + if window.finished: + return WindowAdvance(status=window.status, changed=False) + years = _window_years(connection, window_run_id) + finished = _finish_ended_window(connection, window_run_id, years, now=now) + if finished is not None: + return finished + if start_year is None: + return WindowAdvance(status=window.status, changed=False) + return _start_waiting_years( + connection, + window_run_id, + years, + status=window.status, + start_year=start_year, + free_places=max_running_years - sum(year.under_way for year in years), + now=now, + ) + + def list_idle_windows(self, *, idle_since_before: datetime) -> tuple[UUID, ...]: + """Return the unfinished windows no year is advancing, oldest first. + + A window is idle when none of its years is under way and neither it + nor any of its years has changed since ``idle_since_before``: its + last year ended, or the simulation service created it, without + anything recording what follows. + """ + year_of_window = report_runs.c.budget_window_run_id == budget_window_runs.c.id + last_year_change = ( + select(func.max(report_runs.c.updated_at)) + .where(year_of_window) + .scalar_subquery() + ) + statement = ( + select(budget_window_runs.c.id) + .where( + budget_window_runs.c.status.in_(ACTIVE_STATUSES), + ~exists().where(year_of_window, _under_way()), + func.greatest( + budget_window_runs.c.updated_at, + func.coalesce(last_year_change, budget_window_runs.c.updated_at), + ) + < idle_since_before, + ) + .order_by(budget_window_runs.c.created_at, budget_window_runs.c.id) + ) + with self._transaction("list_idle_windows") as connection: + return tuple(connection.execute(statement).scalars()) + + def settle_idle_window( + self, + window_run_id: UUID, + *, + idle_since_before: datetime, + now: datetime, + ) -> WindowAdvance: + """Finish a window that is still idle, as ``list_idle_windows`` judged it. + + A window whose years have all ended is finished as an advance would + finish it; one with years still waiting fails as stalled, with those + years. A window that finished, has a year under way, or changed since + ``idle_since_before`` is left as it is. Raises LookupError when the + window run does not exist. + """ + with self._transaction("settle_idle_window") as connection: + window = _lock_window(connection, window_run_id) + if window.finished: + return WindowAdvance(status=window.status, changed=False) + years = _window_years(connection, window_run_id) + last_change = max([window.updated_at, *(year.updated_at for year in years)]) + if any(year.under_way for year in years) or ( + last_change >= idle_since_before + ): + return WindowAdvance(status=window.status, changed=False) + finished = _finish_ended_window(connection, window_run_id, years, now=now) + if finished is not None: + return finished + _fail_window( + connection, + window_run_id, + WINDOW_STALLED_MESSAGE, + now=now, + waiting_years_message=WINDOW_STALLED_MESSAGE, + ) + return WindowAdvance(status=RunStatus.FAILED, changed=True) + + +def _lock_window(connection: Connection, window_run_id: UUID) -> _LockedWindow: + """Lock the window run's row until the transaction ends, and read it.""" + locked = connection.execute( + select(budget_window_runs.c.status, budget_window_runs.c.updated_at) + .where(budget_window_runs.c.id == window_run_id) + .with_for_update() + ).one_or_none() + if locked is None: + msg = f"budget-window run {window_run_id} does not exist" + raise LookupError(msg) + return _LockedWindow(status=RunStatus(locked.status), updated_at=locked.updated_at) + + +def _window_years(connection: Connection, window_run_id: UUID) -> list[_WindowYear]: + """Return the window's years, earliest first.""" + statement = ( + select( + report_runs.c.id, + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + report_runs.c.year, + report_runs.c.status, + report_runs.c.modal_call_id, + report_runs.c.updated_at, + ) + .where(report_runs.c.budget_window_run_id == window_run_id) + .order_by(report_runs.c.year) + ) + return [ + _WindowYear( + run_ids=ReportRunIds( + report_run_id=row.id, + baseline_simulation_run_id=row.baseline_simulation_run_id, + reform_simulation_run_id=row.reform_simulation_run_id, + ), + year=row.year, + status=RunStatus(row.status), + modal_call_id=row.modal_call_id, + updated_at=row.updated_at, + ) + for row in connection.execute(statement) + ] + + +def _finish_ended_window( + connection: Connection, + window_run_id: UUID, + years: list[_WindowYear], + *, + now: datetime, +) -> WindowAdvance | None: + """Fail the window for its earliest failed year, or mark it succeeded. + + Returns None while no year has failed and not every year has succeeded. + """ + failed_year = next( + (year.year for year in years if year.status is RunStatus.FAILED), None + ) + if failed_year is not None: + _fail_window( + connection, + window_run_id, + YEAR_FAILED_MESSAGE.format(year=failed_year), + now=now, + ) + return WindowAdvance(status=RunStatus.FAILED, changed=True) + if years and all(year.status is RunStatus.SUCCEEDED for year in years): + connection.execute( + update(budget_window_runs) + .where( + budget_window_runs.c.id == window_run_id, + budget_window_runs.c.status.in_(ACTIVE_STATUSES), + ) + .values( + status=RunStatus.SUCCEEDED.value, + error_message=None, + completed_at=now, + updated_at=now, + ) + ) + return WindowAdvance(status=RunStatus.SUCCEEDED, changed=True) + return None + + +def _start_waiting_years( # noqa: PLR0913 # the window's locked state, the starter and the clock + connection: Connection, + window_run_id: UUID, + years: list[_WindowYear], + *, + status: RunStatus, + start_year: YearStarter, + free_places: int, + now: datetime, +) -> WindowAdvance: + """Start the earliest waiting years in the free places and record their calls. + + A start that raises fails that year and the window, and starts nothing + more, unless the year's own call claimed it meanwhile (Modal accepted + the call after the starter stopped waiting): that year is under way, and + the next ones start. A start deferred with ``YearStartDeferred`` leaves + that year and the later ones waiting. + """ + started: list[int] = [] + for year in [year for year in years if year.waiting][: max(free_places, 0)]: + try: + modal_call_id = start_year( + YearToStart(year=year.year, run_ids=year.run_ids) + ) + except YearStartDeferred: + break + except Exception as error: # noqa: BLE001 # any refusal to start the call fails the year and so the window + # The error's text may hold connection details; its type is logged. + logger.warning( + "Could not start budget-window year %s, report run %s (%s)", + year.year, + year.report_run_id, + type(error).__name__, + ) + refused_years = _fail_waiting_years( + connection, + window_run_id, + CALL_REFUSED_MESSAGE.format(error_type=type(error).__name__), + now=now, + only=year.report_run_id, + ) + if not refused_years: + started.append(year.year) + continue + _fail_window( + connection, + window_run_id, + YEAR_FAILED_MESSAGE.format(year=year.year), + now=now, + ) + return WindowAdvance( + status=RunStatus.FAILED, changed=True, started_years=tuple(started) + ) + connection.execute( + update(report_runs) + .where(report_runs.c.id == year.report_run_id, _waiting()) + .values(modal_call_id=modal_call_id, updated_at=now) + ) + started.append(year.year) + return WindowAdvance(status=status, changed=False, started_years=tuple(started)) + + +def _fail_window( + connection: Connection, + window_run_id: UUID, + error_message: str, + *, + now: datetime, + waiting_years_message: str = WINDOW_FAILED_MESSAGE, +) -> None: + """Fail the window run, then its waiting years and their own simulation runs. + + The window records ``error_message``, and its waiting years + ``waiting_years_message``: by default, that another year failed. + """ + connection.execute( + update(budget_window_runs) + .where( + budget_window_runs.c.id == window_run_id, + budget_window_runs.c.status.in_(ACTIVE_STATUSES), + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=now, + updated_at=now, + ) + ) + _fail_waiting_years(connection, window_run_id, waiting_years_message, now=now) + + +def _fail_waiting_years( + connection: Connection, + window_run_id: UUID, + error_message: str, + *, + now: datetime, + only: UUID | None = None, +) -> int: + """Fail the window's waiting years (or only one) and their unstarted simulations. + + Return how many years were failed. + + Each failed year's simulation runs get the same message. A simulation run + is failed only while it is pending with no call and no pending or running + report run names it, so a run another report shares is left to that + report. + """ + conditions = [report_runs.c.budget_window_run_id == window_run_id, _waiting()] + if only is not None: + conditions.append(report_runs.c.id == only) + failed_years = connection.execute( + update(report_runs) + .where(*conditions) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=now, + updated_at=now, + ) + .returning( + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + ) + ).all() + simulation_run_ids = { + run_id + for failed_year in failed_years + for run_id in failed_year + if run_id is not None + } + if not simulation_run_ids: + return len(failed_years) + connection.execute( + update(simulation_runs) + .where( + simulation_runs.c.id.in_(simulation_run_ids), + simulation_runs.c.status == RunStatus.PENDING.value, + simulation_runs.c.modal_call_id.is_(None), + no_active_report_run_names(simulation_runs), + ) + .values( + status=RunStatus.FAILED.value, + error_message=error_message, + completed_at=now, + updated_at=now, + ) + ) + return len(failed_years) + + +def _waiting() -> ColumnElement[bool]: + """Return the condition that a report run is pending with no call recorded.""" + return and_( + report_runs.c.status == RunStatus.PENDING.value, + report_runs.c.modal_call_id.is_(None), + ) + + +def _under_way() -> ColumnElement[bool]: + """Return the condition that a report run is running, or started and pending.""" + return or_( + report_runs.c.status == RunStatus.RUNNING.value, + and_( + report_runs.c.status == RunStatus.PENDING.value, + report_runs.c.modal_call_id.is_not(None), + ), + ) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index 5dc07d2b4..5161ddda9 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -16,10 +16,11 @@ from enum import StrEnum from typing import TYPE_CHECKING, Any -from sqlalchemy import and_, exists, func, insert, or_, select, update +from sqlalchemy import and_, case, exists, func, insert, or_, select, update from sqlalchemy.dialects.postgresql import insert as postgresql_insert from policyengine_stage12_persistence.report_run_tables import ( + budget_window_runs, household_jobs, report_runs, simulation_runs, @@ -80,12 +81,15 @@ class ActiveRun: """A pending or running run or household job, and its recorded Modal call. ``modal_call_id`` is None for a run with no call recorded. + ``budget_window_run_id`` is the budget-window run a report run is a year + of, and None for any other run. """ kind: RunKind run_id: UUID status: RunStatus modal_call_id: str | None + budget_window_run_id: UUID | None = None @dataclass(frozen=True) @@ -563,24 +567,29 @@ def list_active_runs( no pending or running report run names it as its baseline or reform: that report run's coordinator starts its children's calls, records their outcomes, and stops and starts them again when Modal restarts - it, so they are its to judge until it has finished. + it, so they are its to judge until it has finished. A year of a + pending or running budget-window run that is still pending with no + call recorded is waiting for an earlier year to end, and is listed + only once its window has finished: the window is judged as a whole. + A window's year counts, before its first claim, from its last change + rather than its creation: an earlier year's coordinator may record + its call long after the service created it, and nothing else changes + a year that is pending with a call. """ table = RUN_TABLES[kind] - active_since = ( - table.c.created_at - if kind is RunKind.HOUSEHOLD_JOB - else func.coalesce(table.c.started_at, table.c.created_at) - ) + active_since = _active_since(kind) conditions = [ table.c.status.in_(ACTIVE_STATUSES), active_since < active_since_before, ] + columns = [table.c.id, table.c.status, table.c.modal_call_id] if kind is RunKind.SIMULATION_RUN: - conditions.append(_no_active_report_run_names(table)) + conditions.append(no_active_report_run_names(table)) + if kind is RunKind.REPORT_RUN: + conditions.append(~_waiting_window_year()) + columns.append(table.c.budget_window_run_id) statement = ( - select(table.c.id, table.c.status, table.c.modal_call_id) - .where(*conditions) - .order_by(active_since, table.c.id) + select(*columns).where(*conditions).order_by(active_since, table.c.id) ) with self._transaction("list_active_runs") as connection: return tuple( @@ -589,6 +598,9 @@ def list_active_runs( run_id=row.id, status=RunStatus(row.status), modal_call_id=row.modal_call_id, + budget_window_run_id=( + row.budget_window_run_id if kind is RunKind.REPORT_RUN else None + ), ) for row in connection.execute(statement) ) @@ -625,7 +637,7 @@ def fail_stuck_run( table.c.status == run.status.value, same_call, *( - (_no_active_report_run_names(table),) + (no_active_report_run_names(table),) if run.kind is RunKind.SIMULATION_RUN else () ), @@ -705,7 +717,44 @@ def _running_under_call(table: Table, modal_call_id: str) -> ColumnElement[bool] ) -def _no_active_report_run_names(table: Table) -> ColumnElement[bool]: +def _active_since(kind: RunKind) -> ColumnElement[datetime]: + """Return when a run of ``kind`` became active, for the sweep's ages. + + A household job counts from its creation. A run counts from its first + claim or, before it, from its creation, except a budget window's year, + which counts from its last change: when its call was recorded. + """ + table = RUN_TABLES[kind] + if kind is RunKind.HOUSEHOLD_JOB: + return table.c.created_at + unclaimed_since = ( + case( + (table.c.budget_window_run_id.is_not(None), table.c.updated_at), + else_=table.c.created_at, + ) + if kind is RunKind.REPORT_RUN + else table.c.created_at + ) + return func.coalesce(table.c.started_at, unclaimed_since) + + +def _waiting_window_year() -> ColumnElement[bool]: + """Return the condition that a report run is a year waiting in its window. + + It is pending with no call recorded, in a pending or running budget-window + run, so an earlier year's coordinator will start it. + """ + return and_( + report_runs.c.status == RunStatus.PENDING.value, + report_runs.c.modal_call_id.is_(None), + exists().where( + budget_window_runs.c.id == report_runs.c.budget_window_run_id, + budget_window_runs.c.status.in_(ACTIVE_STATUSES), + ), + ) + + +def no_active_report_run_names(table: Table) -> ColumnElement[bool]: """Return the condition that no pending or running report run names the run. A report run names its simulation runs as its baseline and its reform. diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py index bd8b44a6f..9bcc57cb5 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -2,8 +2,9 @@ ``policyengine-api`` owns these tables and every schema migration. The mappings declare only the columns the coordinator, the workers, the household -functions and the stuck-run sweep read or write as the Modal worker role, so -they can issue typed DML; runtime code never emits DDL. +functions, the budget-window scheduling and the stuck-run sweep read or write +as the Modal worker role, so they can issue typed DML; runtime code never +emits DDL. """ from __future__ import annotations @@ -33,7 +34,8 @@ # A report run compares one pair of simulation runs; a report with only a # baseline has no reform run. The cliff impact, SPM receipts and data version -# are the report run's own results. +# are the report run's own results. A budget window's years are report runs +# naming its budget-window run and their year. report_runs = Table( "report_runs", metadata, @@ -41,6 +43,8 @@ Column("created_at", DateTime(timezone=True), nullable=False), Column("baseline_simulation_run_id", UUID(as_uuid=True), nullable=False), Column("reform_simulation_run_id", UUID(as_uuid=True)), + Column("budget_window_run_id", UUID(as_uuid=True)), + Column("year", Integer), Column("status", report_run_status_type, nullable=False), Column("modal_call_id", String(255)), Column("error_message", String), @@ -55,6 +59,21 @@ Column("data_version", String(128)), ) +# One calculation of a budget-window report: its years are the report runs +# naming it. It has no Modal call of its own; the coordinators of its years +# mark it running, succeeded or failed. +budget_window_runs = Table( + "budget_window_runs", + metadata, + Column("id", UUID(as_uuid=True), primary_key=True), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("status", report_run_status_type, nullable=False), + Column("error_message", String), + Column("started_at", DateTime(timezone=True)), + Column("completed_at", DateTime(timezone=True)), + Column("updated_at", DateTime(timezone=True), nullable=False), +) + # A simulation run's year, versions, dataset and region say what it # calculated; the coordinator checks them on a succeeded run that a report # run reuses. diff --git a/libs/policyengine-stage12-persistence/tests/budget_window_database.py b/libs/policyengine-stage12-persistence/tests/budget_window_database.py new file mode 100644 index 000000000..b8548feb6 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/budget_window_database.py @@ -0,0 +1,223 @@ +"""Budget-window runs in the disposable Postgres, as the simulation service makes them. + +A budget-window report compares its baseline and reform in each year of its +window. The simulation service creates one ``budget_window_runs`` row per +calculation, and for each year two simulation runs and a report run naming +that window run and year, all pending with no Modal call. The seed role +inserts them; tests then set each year's status as a scenario needs. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from report_run_database import NOW, VERSIONS, ReportRunSeeder +from sqlalchemy import delete, insert, select + +if TYPE_CHECKING: + from collections.abc import Iterable + +# Every seeded row was created and last changed at NOW, so a test's cutoffs +# don't depend on the clock. +SEEDED_AT = {"created_at": NOW, "updated_at": NOW} + + +@dataclass(frozen=True) +class SeededYear: + """One year of a seeded window: its report run and its two simulation runs.""" + + year: int + report_run_id: UUID + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID + + +@dataclass(frozen=True) +class SeededWindow: + """A seeded budget-window run, its report and its years, earliest first.""" + + window_run_id: UUID + report_id: UUID + years: tuple[SeededYear, ...] + simulation_ids: tuple[UUID, UUID] + model_id: UUID + + def year(self, year: int) -> SeededYear: + """Return the seeded year ``year``.""" + return next(seeded for seeded in self.years if seeded.year == year) + + +class BudgetWindowSeeder: + """Insert, change and remove budget-window runs as the seed role.""" + + def __init__(self, seeder: ReportRunSeeder) -> None: + """Use ``seeder``'s reflected tables, which include ``budget_window_runs``.""" + self.seeder = seeder + self.engine = seeder.engine + + def seed_window(self, first_year: int = 2026, size: int = 7) -> SeededWindow: + """Insert a pending window run with ``size`` pending years from a year.""" + table = self.seeder.table + model_id, version_id, dataset_id = uuid4(), uuid4(), uuid4() + simulation_ids = (uuid4(), uuid4()) + report_id, window_run_id = uuid4(), uuid4() + years = tuple( + SeededYear( + year=year, + report_run_id=uuid4(), + baseline_simulation_run_id=uuid4(), + reform_simulation_run_id=uuid4(), + ) + for year in range(first_year, first_year + size) + ) + with self.engine.begin() as connection: + self.seeder.insert_catalog(connection, model_id, version_id, dataset_id) + for simulation_id in simulation_ids: + connection.execute( + insert(table("simulations")).values( + id=simulation_id, + simulation_type="economy", + dataset_id=dataset_id, + tax_benefit_model_version_id=version_id, + year=first_year, + ) + ) + connection.execute( + insert(table("reports")).values( + id=report_id, + label="Budget window test", + country="us", + tax_benefit_model_id=model_id, + baseline_simulation_id=simulation_ids[0], + reform_simulation_id=simulation_ids[1], + year=first_year, + budget_window_size=size, + inputs={}, + ) + ) + connection.execute( + insert(table("budget_window_runs")).values( + id=window_run_id, + report_id=report_id, + tax_benefit_model_version_id=version_id, + trigger="initial", + recalculate=False, + status="pending", + **SEEDED_AT, + ) + ) + for seeded in years: + for run_id, simulation_id in ( + (seeded.baseline_simulation_run_id, simulation_ids[0]), + (seeded.reform_simulation_run_id, simulation_ids[1]), + ): + connection.execute( + insert(table("simulation_runs")).values( + id=run_id, + simulation_id=simulation_id, + year=seeded.year, + tax_benefit_model_version_id=version_id, + dataset_id=dataset_id, + status="pending", + **SEEDED_AT, + **VERSIONS, + ) + ) + connection.execute( + insert(table("report_runs")).values( + id=seeded.report_run_id, + report_id=report_id, + baseline_simulation_run_id=seeded.baseline_simulation_run_id, + reform_simulation_run_id=seeded.reform_simulation_run_id, + budget_window_run_id=window_run_id, + year=seeded.year, + status="pending", + trigger="initial", + **SEEDED_AT, + **VERSIONS, + ) + ) + return SeededWindow( + window_run_id=window_run_id, + report_id=report_id, + years=years, + simulation_ids=simulation_ids, + model_id=model_id, + ) + + def set_years( + self, window: SeededWindow, years: Iterable[int], **values: Any + ) -> None: # column values of any type + """Set ``values`` on the report runs of ``years`` in ``window``.""" + report_runs = self.seeder.table("report_runs") + run_ids = [window.year(year).report_run_id for year in years] + with self.engine.begin() as connection: + connection.execute( + report_runs.update() + .where(report_runs.c.id.in_(run_ids)) + .values(**values) + ) + + def set_window( + self, window: SeededWindow, **values: Any + ) -> None: # column values of any type + """Set ``values`` on ``window``'s budget-window run.""" + windows = self.seeder.table("budget_window_runs") + with self.engine.begin() as connection: + connection.execute( + windows.update() + .where(windows.c.id == window.window_run_id) + .values(**values) + ) + + def window_row(self, window: SeededWindow) -> dict[str, Any]: + """Return ``window``'s budget-window run as a dictionary.""" + return self.seeder.row("budget_window_runs", window.window_run_id) + + def year_rows(self, window: SeededWindow) -> dict[int, dict[str, Any]]: + """Return each year's report run by year.""" + return { + row["year"]: row + for row in self.seeder.rows( + "report_runs", budget_window_run_id=window.window_run_id + ) + } + + def simulation_run_statuses(self, seeded: SeededYear) -> tuple[str, str]: + """Return the year's baseline and reform simulation runs' statuses.""" + simulation_runs = self.seeder.table("simulation_runs") + statement = select(simulation_runs.c.id, simulation_runs.c.status).where( + simulation_runs.c.id.in_( + (seeded.baseline_simulation_run_id, seeded.reform_simulation_run_id) + ) + ) + with self.engine.connect() as connection: + statuses = {row.id: row.status for row in connection.execute(statement)} + return ( + statuses[seeded.baseline_simulation_run_id], + statuses[seeded.reform_simulation_run_id], + ) + + def remove(self, window: SeededWindow) -> None: + """Delete the seeded rows; the window run, runs and results cascade.""" + table = self.seeder.table + reports, simulations = table("reports"), table("simulations") + datasets, models = table("datasets"), table("tax_benefit_models") + versions = table("tax_benefit_model_versions") + version_ids = select(versions.c.id).where( + versions.c.model_id == window.model_id + ) + with self.engine.begin() as connection: + connection.execute(delete(reports).where(reports.c.id == window.report_id)) + connection.execute( + delete(simulations).where(simulations.c.id.in_(window.simulation_ids)) + ) + connection.execute( + delete(datasets).where( + datasets.c.tax_benefit_model_version_id.in_(version_ids) + ) + ) + # Deleting the model removes its versions. + connection.execute(delete(models).where(models.c.id == window.model_id)) diff --git a/libs/policyengine-stage12-persistence/tests/report_run_database.py b/libs/policyengine-stage12-persistence/tests/report_run_database.py index e11f82d39..280186a54 100644 --- a/libs/policyengine-stage12-persistence/tests/report_run_database.py +++ b/libs/policyengine-stage12-persistence/tests/report_run_database.py @@ -31,6 +31,7 @@ "datasets", "simulations", "reports", + "budget_window_runs", "report_runs", "simulation_runs", "labor_supply_responses", @@ -101,7 +102,7 @@ def seed_report_run(self) -> SeededReportRun: reform_simulation_run_id=uuid4(), ) with self.engine.begin() as connection: - self._insert_catalog(connection, model_id, version_id, dataset_id) + self.insert_catalog(connection, model_id, version_id, dataset_id) for simulation_id in simulation_ids: connection.execute( insert(self.table("simulations")).values( @@ -159,7 +160,7 @@ def seed_household_job(self) -> SeededHouseholdJob: model_id, version_id, dataset_id = uuid4(), uuid4(), uuid4() household_id, job_id = uuid4(), uuid4() with self.engine.begin() as connection: - self._insert_catalog(connection, model_id, version_id, dataset_id) + self.insert_catalog(connection, model_id, version_id, dataset_id) connection.execute( insert(self.table("households")).values( id=household_id, @@ -210,9 +211,10 @@ def remove_household_job(self, seeded: SeededHouseholdJob) -> None: # Deleting the model removes its versions. connection.execute(delete(models).where(models.c.id == seeded.model_id)) - def _insert_catalog( + def insert_catalog( self, connection: Connection, model_id: UUID, version_id: UUID, dataset_id: UUID ) -> None: + """Insert a tax-benefit model, one version of it and one of its datasets.""" connection.execute( insert(self.table("tax_benefit_models")).values( id=model_id, name=f"test-{model_id.hex[:24]}" diff --git a/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py b/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py new file mode 100644 index 000000000..d7b6b8bb7 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py @@ -0,0 +1,643 @@ +"""Starting, finishing and settling budget-window runs, as the Modal worker role. + +A budget window's years are report runs naming its budget-window run. The +simulation service creates them all pending and starts the first ones; each +year's coordinator, when its year ends, starts the earliest years still +waiting, keeping at most five under way, and the last year to end marks the +window. These run against a disposable Postgres that the Public API's v2 +migrations created (see ``report_run_database``). +""" + +from __future__ import annotations + +import logging +import os +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest +from budget_window_database import BudgetWindowSeeder, SeededWindow +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_stage12_persistence import ( + BUDGET_WINDOW_MAX_RUNNING_YEARS, + ActiveRun, + BudgetWindowStore, + ReportRunStore, + RunKind, + RunStatus, + WindowAdvance, + YearStartDeferred, + YearToStart, + budget_window_store, +) +from report_run_database import NOW, SEED_URL_ENV, worker_database_url +from sqlalchemy import Engine, create_engine, event + +# The simulation service's role, which the Public API's manifest grants. +SERVICE_ROLE = "policyengine_v2_simulation_service" + +if TYPE_CHECKING: + from collections.abc import Iterator + from uuid import UUID + + from report_run_database import ReportRunSeeder, SeededReportRun + +LATER = NOW + timedelta(minutes=5) +ALL_YEARS = range(2026, 2033) + + +class RecordingStarter: + """Starts a year by recording it and returning a new call ID. + + A year whose report run is in ``refused`` raises, as Modal does when it + refuses a call. ``during_start`` runs before the call ID is returned. + """ + + def __init__(self, *, refused=(), delay_seconds=0.0, during_start=None): + self.started: list[UUID] = [] + self.years: list[YearToStart] = [] + self.refused = set(refused) + self.delay_seconds = delay_seconds + self.during_start = during_start + self._lock = threading.Lock() + + def __call__(self, year_to_start: YearToStart) -> str: + report_run_id = year_to_start.run_ids.report_run_id + if report_run_id in self.refused: + msg = "Modal refused the call" + raise RuntimeError(msg) + time.sleep(self.delay_seconds) + if self.during_start is not None: + self.during_start(report_run_id) + with self._lock: + self.started.append(report_run_id) + self.years.append(year_to_start) + return f"fc-started-{len(self.started)}" + + +@pytest.fixture +def store() -> Iterator[BudgetWindowStore]: + window_store = BudgetWindowStore(worker_database_url()) + try: + yield window_store + finally: + window_store.engine.dispose() + + +@pytest.fixture +def windows(seeder: ReportRunSeeder) -> BudgetWindowSeeder: + return BudgetWindowSeeder(seeder) + + +@pytest.fixture +def window(windows: BudgetWindowSeeder) -> Iterator[SeededWindow]: + """Yield a pending window run with seven pending years, 2026 to 2032.""" + seeded = windows.seed_window(first_year=2026, size=7) + try: + yield seeded + finally: + windows.remove(seeded) + + +def _run_ids(window: SeededWindow, years) -> list[UUID]: + return [window.year(year).report_run_id for year in years] + + +def _succeeded(windows, window, years) -> None: + windows.set_years( + window, + years, + status="succeeded", + modal_call_id="fc-finished", + started_at=NOW, + completed_at=NOW, + ) + + +def test_the_window_runs_at_most_five_years_at_once() -> None: + assert BUDGET_WINDOW_MAX_RUNNING_YEARS == 5 + + +def test_the_first_advance_starts_the_five_earliest_years( + store, windows, window +) -> None: + starter = RecordingStarter() + + advance = store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + assert starter.started == _run_ids(window, range(2026, 2031)) + assert advance == WindowAdvance( + status=RunStatus.PENDING, + changed=False, + started_years=(2026, 2027, 2028, 2029, 2030), + ) + rows = windows.year_rows(window) + assert [rows[year]["modal_call_id"] for year in ALL_YEARS] == [ + "fc-started-1", + "fc-started-2", + "fc-started-3", + "fc-started-4", + "fc-started-5", + None, + None, + ] + assert {rows[year]["status"] for year in ALL_YEARS} == {"pending"} + assert [rows[year]["updated_at"] for year in range(2026, 2031)] == [LATER] * 5 + assert windows.window_row(window)["status"] == "pending" + + +def test_the_starter_is_given_each_year_and_its_runs(store, window) -> None: + starter = RecordingStarter() + + store.advance_window( + window.window_run_id, now=LATER, start_year=starter, max_running_years=2 + ) + + assert starter.years == [ + YearToStart( + year=year, + run_ids=ReportRunIds( + report_run_id=window.year(year).report_run_id, + baseline_simulation_run_id=window.year(year).baseline_simulation_run_id, + reform_simulation_run_id=window.year(year).reform_simulation_run_id, + ), + ) + for year in (2026, 2027) + ] + + +def test_an_advance_fills_only_the_free_places(store, windows, window) -> None: + _succeeded(windows, window, [2026]) + windows.set_years(window, [2027, 2028], status="running", modal_call_id="fc-run") + # Started, but its coordinator has not claimed it yet. + windows.set_years(window, [2029], modal_call_id="fc-queued") + starter = RecordingStarter() + + advance = store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + assert starter.started == _run_ids(window, [2030, 2031]) + assert advance.started_years == (2030, 2031) + assert windows.year_rows(window)[2032]["modal_call_id"] is None + + +def test_an_advance_without_a_starter_starts_nothing(store, windows, window) -> None: + advance = store.advance_window(window.window_run_id, now=LATER) + + assert advance == WindowAdvance( + status=RunStatus.PENDING, changed=False, started_years=() + ) + assert {row["modal_call_id"] for row in windows.year_rows(window).values()} == { + None + } + + +def test_a_failed_year_fails_the_window_its_waiting_years_and_their_simulations( + store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + windows.set_years( + window, + [2026], + status="failed", + modal_call_id="fc-failed", + error_message="worker_failed: out of memory", + completed_at=NOW, + ) + windows.set_years(window, [2027], status="running", modal_call_id="fc-running") + starter = RecordingStarter() + + advance = store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + assert starter.started == [] + assert advance == WindowAdvance( + status=RunStatus.FAILED, changed=True, started_years=() + ) + window_row = windows.window_row(window) + assert window_row["status"] == "failed" + assert window_row["error_message"].startswith("budget_window_year_failed: ") + assert "2026" in window_row["error_message"] + assert (window_row["completed_at"], window_row["updated_at"]) == (LATER, LATER) + rows = windows.year_rows(window) + # The running year is its coordinator's to finish. + assert (rows[2027]["status"], rows[2027]["completed_at"]) == ("running", None) + for year in range(2028, 2033): + assert rows[year]["status"] == "failed" + assert rows[year]["error_message"].startswith("budget_window_failed: ") + assert rows[year]["completed_at"] == LATER + assert windows.simulation_run_statuses(window.year(year)) == ( + "failed", + "failed", + ) + assert windows.simulation_run_statuses(window.year(2027)) == ( + "pending", + "pending", + ) + + +def test_the_last_year_to_succeed_marks_the_window_succeeded( + store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + _succeeded(windows, window, ALL_YEARS) + + advance = store.advance_window( + window.window_run_id, now=LATER, start_year=RecordingStarter() + ) + + assert advance == WindowAdvance( + status=RunStatus.SUCCEEDED, changed=True, started_years=() + ) + window_row = windows.window_row(window) + assert (window_row["status"], window_row["error_message"]) == ("succeeded", None) + assert (window_row["completed_at"], window_row["updated_at"]) == (LATER, LATER) + + +def test_a_window_with_a_year_still_running_stays_as_it_is( + store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + _succeeded(windows, window, range(2026, 2032)) + windows.set_years(window, [2032], status="running", modal_call_id="fc-last") + + advance = store.advance_window( + window.window_run_id, now=LATER, start_year=RecordingStarter() + ) + + assert advance == WindowAdvance( + status=RunStatus.RUNNING, changed=False, started_years=() + ) + assert windows.window_row(window)["completed_at"] is None + + +def test_a_refused_start_fails_that_year_and_the_window(store, windows, window) -> None: + starter = RecordingStarter(refused=_run_ids(window, [2028])) + + advance = store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + assert starter.started == _run_ids(window, [2026, 2027]) + assert advance == WindowAdvance( + status=RunStatus.FAILED, changed=True, started_years=(2026, 2027) + ) + rows = windows.year_rows(window) + assert [rows[year]["modal_call_id"] for year in (2026, 2027)] == [ + "fc-started-1", + "fc-started-2", + ] + assert rows[2028]["status"] == "failed" + assert rows[2028]["error_message"].startswith("modal_call_refused: ") + # The cause's type, never its text, names why the start failed. + assert "RuntimeError" in rows[2028]["error_message"] + assert "Modal refused the call" not in rows[2028]["error_message"] + assert {rows[year]["status"] for year in range(2029, 2033)} == {"failed"} + window_row = windows.window_row(window) + assert window_row["status"] == "failed" + assert "2028" in window_row["error_message"] + + +def test_a_refused_start_names_its_cause_on_its_own_simulations( + store, windows, window, caplog +) -> None: + starter = RecordingStarter(refused=_run_ids(window, [2026])) + + with caplog.at_level(logging.WARNING, logger=budget_window_store.__name__): + store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + for seeded_year, message_start in ( + (window.year(2026), "modal_call_refused: "), + (window.year(2027), "budget_window_failed: "), + ): + for simulation_run_id in ( + seeded_year.baseline_simulation_run_id, + seeded_year.reform_simulation_run_id, + ): + row = windows.seeder.row("simulation_runs", simulation_run_id) + assert row["status"] == "failed" + assert row["error_message"].startswith(message_start) + # The cause's type is logged; its text may hold connection details. + assert "RuntimeError" in caplog.text + assert str(window.year(2026).report_run_id) in caplog.text + assert "Modal refused the call" not in caplog.text + + +def test_the_service_role_starts_and_fails_a_window(windows, window) -> None: + # The simulation service starts a window's first years with the same + # store, as its own role, and fails the window when Modal refuses one. + service_store = BudgetWindowStore(engine=_engine_as_the_service(windows)) + try: + starter = RecordingStarter(refused=_run_ids(window, [2028])) + + advance = service_store.advance_window( + window.window_run_id, now=LATER, start_year=starter + ) + finally: + service_store.engine.dispose() + + assert advance.started_years == (2026, 2027) + assert windows.window_row(window)["status"] == "failed" + + +def _engine_as_the_service(windows: BudgetWindowSeeder) -> Engine: + """Return an engine whose every connection acts as the service's role.""" + engine = create_engine(os.environ[SEED_URL_ENV]) + + @event.listens_for(engine, "connect") + def act_as_the_service(dbapi_connection, _connection_record) -> None: + with dbapi_connection.cursor() as cursor: + cursor.execute(f"SET ROLE {SERVICE_ROLE}") + + return engine + + +def test_a_year_started_late_counts_from_its_start(run_store, windows, window) -> None: + # The service created every year at submit; an earlier year's coordinator + # recorded this year's call a day later, and Modal hasn't run it yet. + started_late_at = NOW + timedelta(hours=30) + windows.set_window(window, status="running", started_at=NOW) + windows.set_years( + window, [2031], modal_call_id="fc-queued", updated_at=started_late_at + ) + + listed = { + run.run_id + for run in run_store.list_active_runs( + RunKind.REPORT_RUN, active_since_before=NOW + timedelta(hours=25) + ) + } + + assert window.year(2031).report_run_id not in listed + + +def test_a_start_that_fails_after_its_call_claimed_the_year_keeps_the_year( + store, windows, window +) -> None: + # The service waits a limited time for Modal to accept a call; Modal may + # accept it later, and its coordinator claim the year, before the wait's + # end is recorded. + first_year_run_id = window.year(2026).report_run_id + + def claim_then_time_out(year_to_start: YearToStart) -> str: + if year_to_start.run_ids.report_run_id == first_year_run_id: + windows.set_years(window, [2026], status="running", modal_call_id="fc-late") + raise TimeoutError + return f"fc-{year_to_start.year}" + + advance = store.advance_window( + window.window_run_id, now=LATER, start_year=claim_then_time_out + ) + + assert advance == WindowAdvance( + status=RunStatus.PENDING, + changed=False, + started_years=(2026, 2027, 2028, 2029, 2030), + ) + rows = windows.year_rows(window) + assert (rows[2026]["status"], rows[2026]["modal_call_id"]) == ("running", "fc-late") + assert windows.window_row(window)["status"] == "pending" + + +def test_a_deferred_start_leaves_that_year_and_the_later_ones_waiting( + store, windows, window +) -> None: + def defer_from_2028(year_to_start: YearToStart) -> str: + if year_to_start.year >= 2028: + raise YearStartDeferred + return f"fc-{year_to_start.year}" + + advance = store.advance_window( + window.window_run_id, now=LATER, start_year=defer_from_2028 + ) + + assert advance == WindowAdvance( + status=RunStatus.PENDING, changed=False, started_years=(2026, 2027) + ) + rows = windows.year_rows(window) + assert {rows[year]["status"] for year in ALL_YEARS} == {"pending"} + assert [rows[year]["modal_call_id"] for year in ALL_YEARS] == [ + "fc-2026", + "fc-2027", + None, + None, + None, + None, + None, + ] + assert windows.window_row(window)["status"] == "pending" + + +def test_a_finished_window_is_left_as_it_is(store, windows, window) -> None: + windows.set_window( + window, status="failed", error_message="earlier", completed_at=NOW + ) + starter = RecordingStarter() + + advance = store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + assert starter.started == [] + assert advance == WindowAdvance( + status=RunStatus.FAILED, changed=False, started_years=() + ) + assert windows.window_row(window)["error_message"] == "earlier" + + +def test_a_year_its_coordinator_claimed_during_the_start_keeps_its_claim( + store, windows, window +) -> None: + def claim_while_starting(report_run_id: UUID) -> None: + # The started call claims its year before the advance records it. + if report_run_id == window.year(2026).report_run_id: + windows.set_years( + window, [2026], status="running", modal_call_id="fc-claimed" + ) + + starter = RecordingStarter(during_start=claim_while_starting) + + store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + row = windows.year_rows(window)[2026] + assert (row["status"], row["modal_call_id"]) == ("running", "fc-claimed") + + +def test_two_advances_at_once_start_each_year_once(store, windows, window) -> None: + starter = RecordingStarter(delay_seconds=0.2) + barrier = Barrier(2) + + def advance() -> WindowAdvance: + barrier.wait() + return store.advance_window(window.window_run_id, now=LATER, start_year=starter) + + with ThreadPoolExecutor(max_workers=2) as pool: + advances = list(pool.map(lambda _: advance(), range(2))) + + assert sorted(starter.started) == sorted(_run_ids(window, range(2026, 2031))) + assert sorted(year for finished in advances for year in finished.started_years) == [ + 2026, + 2027, + 2028, + 2029, + 2030, + ] + + +def test_an_unknown_window_raises_lookup_error(store) -> None: + with pytest.raises(LookupError, match="budget-window run"): + store.advance_window(uuid4(), now=LATER) + + +def test_a_claimed_year_marks_its_window_running_once(store, windows, window) -> None: + assert store.start_window(window.year(2026).report_run_id, started_at=LATER) + assert not store.start_window( + window.year(2027).report_run_id, started_at=LATER + timedelta(hours=1) + ) + + window_row = windows.window_row(window) + assert window_row["status"] == "running" + assert (window_row["started_at"], window_row["updated_at"]) == (LATER, LATER) + + +def test_a_report_run_outside_any_window_marks_nothing( + store, seeded_run: SeededReportRun +) -> None: + assert not store.start_window(seeded_run.run_ids.report_run_id, started_at=LATER) + + +def test_a_year_names_its_window(store, window, seeded_run) -> None: + assert store.window_of_year(window.year(2029).report_run_id) == ( + window.window_run_id + ) + assert store.window_of_year(seeded_run.run_ids.report_run_id) is None + assert store.window_of_year(uuid4()) is None + + +def test_an_idle_window_with_waiting_years_fails_as_stalled( + store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + _succeeded(windows, window, range(2026, 2031)) + cutoff = NOW + timedelta(minutes=10) + + assert window.window_run_id in store.list_idle_windows(idle_since_before=cutoff) + settled = store.settle_idle_window( + window.window_run_id, idle_since_before=cutoff, now=LATER + ) + + assert settled == WindowAdvance( + status=RunStatus.FAILED, changed=True, started_years=() + ) + window_row = windows.window_row(window) + assert window_row["error_message"].startswith("budget_window_stalled: ") + rows = windows.year_rows(window) + for year in (2031, 2032): + assert rows[year]["status"] == "failed" + # The years record the stall, not another year's failure. + assert rows[year]["error_message"].startswith("budget_window_stalled: ") + for simulation_run_id in ( + window.year(year).baseline_simulation_run_id, + window.year(year).reform_simulation_run_id, + ): + assert windows.seeder.row("simulation_runs", simulation_run_id)[ + "error_message" + ].startswith("budget_window_stalled: ") + assert window.window_run_id not in store.list_idle_windows(idle_since_before=cutoff) + + +def test_an_idle_window_whose_years_all_succeeded_is_marked_succeeded( + store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + _succeeded(windows, window, ALL_YEARS) + cutoff = NOW + timedelta(minutes=10) + + assert window.window_run_id in store.list_idle_windows(idle_since_before=cutoff) + settled = store.settle_idle_window( + window.window_run_id, idle_since_before=cutoff, now=LATER + ) + + assert settled.status is RunStatus.SUCCEEDED + assert windows.window_row(window)["status"] == "succeeded" + + +def test_a_window_whose_service_started_no_year_is_idle(store, window) -> None: + # The simulation service stopped after creating the window's rows. + cutoff = NOW + timedelta(minutes=10) + + assert window.window_run_id in store.list_idle_windows(idle_since_before=cutoff) + assert ( + store.settle_idle_window( + window.window_run_id, idle_since_before=cutoff, now=LATER + ).status + is RunStatus.FAILED + ) + + +def test_a_window_with_a_year_under_way_is_not_idle(store, windows, window) -> None: + windows.set_years(window, [2026], modal_call_id="fc-queued") + cutoff = NOW + timedelta(minutes=10) + + assert window.window_run_id not in store.list_idle_windows(idle_since_before=cutoff) + assert store.settle_idle_window( + window.window_run_id, idle_since_before=cutoff, now=LATER + ) == WindowAdvance(status=RunStatus.PENDING, changed=False, started_years=()) + + +def test_a_window_whose_year_changed_recently_is_not_idle( + store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + _succeeded(windows, window, [2026]) + windows.set_years(window, [2026], updated_at=NOW + timedelta(minutes=20)) + cutoff = NOW + timedelta(minutes=10) + + assert window.window_run_id not in store.list_idle_windows(idle_since_before=cutoff) + assert not store.settle_idle_window( + window.window_run_id, idle_since_before=cutoff, now=LATER + ).changed + + +@pytest.fixture +def run_store() -> Iterator[ReportRunStore]: + report_run_store = ReportRunStore(worker_database_url()) + try: + yield report_run_store + finally: + report_run_store.engine.dispose() + + +def _listed_report_runs(run_store: ReportRunStore) -> dict[UUID, ActiveRun]: + return { + run.run_id: run + for run in run_store.list_active_runs( + RunKind.REPORT_RUN, active_since_before=NOW + timedelta(minutes=30) + ) + } + + +def test_the_sweep_lists_a_window_s_started_years_but_not_its_waiting_ones( + run_store, windows, window +) -> None: + windows.set_window(window, status="running", started_at=NOW) + windows.set_years( + window, [2026], status="running", modal_call_id="fc-running", started_at=NOW + ) + windows.set_years(window, [2027], modal_call_id="fc-queued") + + listed = _listed_report_runs(run_store) + + for year in (2026, 2027): + run = listed[window.year(year).report_run_id] + assert run.budget_window_run_id == window.window_run_id + assert not set(_run_ids(window, range(2028, 2033))) & set(listed) + + +def test_the_sweep_lists_the_waiting_years_of_a_finished_window( + run_store, windows, window +) -> None: + windows.set_window(window, status="failed", error_message="x", completed_at=NOW) + + listed = _listed_report_runs(run_store) + + assert set(_run_ids(window, ALL_YEARS)) <= set(listed) diff --git a/projects/policyengine-apis-integ/tests/test_budget_window_generated_client.py b/projects/policyengine-apis-integ/tests/test_budget_window_generated_client.py new file mode 100644 index 000000000..86b6d3eb4 --- /dev/null +++ b/projects/policyengine-apis-integ/tests/test_budget_window_generated_client.py @@ -0,0 +1,93 @@ +"""Hermetic wire checks of a budget-window report through the generated client. + +The Public API sends a budget window to ``POST /reports`` with +``report_type`` ``budget_window``, and reads the 201 naming the report, the +window run and each year's runs. +""" + +import httpx +from policyengine_api_simulation_client import AuthenticatedClient +from policyengine_api_simulation_client.api.default import create_report +from policyengine_api_simulation_client.models import ( + BudgetWindowReportCreated, + BudgetWindowReportRequest, + EconomyError, + ReportRunStatus, +) + +REPORT_ID = "00000000-0000-4000-8000-0000000000e1" +WINDOW_RUN_ID = "00000000-0000-4000-8000-0000000000e2" +POLICY_ID = "00000000-0000-4000-8000-0000000000c5" +WINDOW_BODY = { + "country": "us", + "report_type": "budget_window", + "label": "Ten-year budget window", + "description": None, + "year": 2026, + "budget_window_size": 2, + "region_code": "us", + "dataset_name": None, + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": POLICY_ID, "dynamic_id": None}, +} +YEARS = [ + { + "year": 2026 + index, + "report_run_id": f"00000000-0000-4000-8000-00000000010{index}", + "baseline_simulation_id": f"00000000-0000-4000-8000-00000000020{index}", + "baseline_simulation_run_id": f"00000000-0000-4000-8000-00000000030{index}", + "baseline_reused": index == 1, + "reform_simulation_id": f"00000000-0000-4000-8000-00000000040{index}", + "reform_simulation_run_id": f"00000000-0000-4000-8000-00000000050{index}", + "reform_reused": False, + } + for index in range(2) +] +CREATED_WINDOW = { + "report_id": REPORT_ID, + "report_run_id": YEARS[0]["report_run_id"], + "budget_window_run_id": WINDOW_RUN_ID, + "status": "pending", + "years": YEARS, +} + + +def _create_window(status_code: int, body: dict): + transport = httpx.MockTransport( + lambda request: httpx.Response(status_code, json=body) + ) + client = AuthenticatedClient(base_url="https://simulation.example", token="test") + client.set_httpx_client( + httpx.Client(base_url="https://simulation.example", transport=transport) + ) + return create_report.sync_detailed( + client=client, body=BudgetWindowReportRequest.from_dict(WINDOW_BODY) + ) + + +def test_a_budget_window_request_is_sent_as_given(): + sent = BudgetWindowReportRequest.from_dict(WINDOW_BODY).to_dict() + + assert sent == WINDOW_BODY + + +def test_a_created_budget_window_is_read_with_its_years(): + response = _create_window(201, CREATED_WINDOW) + + assert isinstance(response.parsed, BudgetWindowReportCreated) + assert response.parsed.status is ReportRunStatus.PENDING + assert response.parsed.to_dict() == CREATED_WINDOW + + +def test_a_refused_first_year_is_read_with_its_report(): + response = _create_window( + 503, + { + "code": "modal_call_refused", + "detail": "Modal did not accept the report's call.", + "report_id": REPORT_ID, + }, + ) + + assert isinstance(response.parsed, EconomyError) + assert str(response.parsed.report_id) == REPORT_ID diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index bc32cc1f6..a0b188b5c 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -151,6 +151,38 @@ a run claimed meanwhile is answered with 201. Every 500 is the JSON error with code `internal_error`, and may follow a write. No `detail` holds an exception's text. +### Budget windows + +`POST /reports` with `report_type` `budget_window` creates a budget-window +report: the same baseline and reform in each year of a window that starts +at `year` and lasts `budget_window_size` years (1 to 75, ending by 2200). +It names the same fields as an economy report, except `include_cliffs`, +which it refuses; it stores the options of a report without cliffs, so its +years and single-year reports reuse each other's runs. In one transaction +the service creates the report (an economy report whose +`budget_window_size` is set), one `budget_window_runs` row, and for each +year a report run naming it and the year, with that year's own simulations +unless a side reuses a run of the same year. It then starts up to the +earliest five years, as its own role, under the window run's row lock +(`BudgetWindowStore.advance_window`): the first year's coordinator input is +built as an economy report's, and each later year's is derived from it. +Each year's coordinator starts the next years when its year ends. The 201 +names the report, the first year's report run, the window run, its status +and every year's runs. A year Modal refuses fails the window, its waiting +years and their new simulation runs, and so does the first year when Modal +doesn't confirm it within ten seconds; the 503 (`modal_call_refused`) names +the report. The window run's +status needs the update grant on `budget_window_runs` that the Public API +gives the service (PolicyEngine/policyengine-api#4267); `/ready` checks it. +The first years share that one ten-second wait for Modal, counted from the +first start: a later year not confirmed within it waits, with the years +after it, for a started year's coordinator to start it. A body with no +`report_type` is refused for that alone, before its other fields are +checked. +`GET /reports/{report_id}` answers a window with one of its years' runs; +the window's own status is in `budget_window_runs`, which the Public API +reads. + `GET /reports/{report_id}` and `GET /simulations/{simulation_id}` return the record and its newest run (created last, ties to the larger ID): its status, versions, error message and times, and a report run's simulation runs. The diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 0c99f05e4..a91636568 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -42,6 +42,7 @@ init_simulation_observability, ) from policyengine_simulation_persistence import EconomyRunStore, HouseholdJobStore +from policyengine_stage12_persistence import BudgetWindowStore from starlette.datastructures import MutableHeaders from policyengine_simulation_entry.auth import ( @@ -51,13 +52,15 @@ CallerAuthorizer, ) from policyengine_simulation_entry.backend import OldGatewayBackend, SimulationBackend +from policyengine_simulation_entry.budget_window_submissions import ( + ReportSubmissionService, +) from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.economy_routes import ( EconomySubmissions, add_economy_routes, ) from policyengine_simulation_entry.economy_submissions import ( - EconomySubmissionService, ServiceSettings, ) from policyengine_simulation_entry.economy_workers import ( @@ -562,7 +565,9 @@ def economy_submissions_for( database, SimulationDatabase ): return None - return EconomySubmissionService( + return ReportSubmissionService( + # The service starts a budget window's first years as its own role. + windows=BudgetWindowStore(engine=database.engine), store=EconomyRunStore.from_engine( database.engine, statement_timeout_milliseconds=ECONOMY_STATEMENT_TIMEOUT_MILLISECONDS, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py new file mode 100644 index 000000000..009b9e537 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py @@ -0,0 +1,149 @@ +"""Request and answer models of a budget-window report on ``POST /reports``. + +A budget-window report compares one region's baseline and reform in each +year of a window of 1 to 75 consecutive years, ending by 2200. The request +names the first year and the window's length, and the same region, dataset, +policies and SPM choice an economy report names (``EconomyReportRequest``); +it sends no cliff option, which a window doesn't calculate. The answers and +their error codes are an economy report's. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal +from uuid import UUID + +from policyengine_simulation_contract.spm import SPMSelection +from policyengine_simulation_contract.stage12_bundle import CountryId +from policyengine_simulation_persistence import ( + LATEST_YEAR, + MAX_BUDGET_WINDOW_YEARS, + ReportRunStatus, +) +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic.json_schema import SkipJsonSchema + +from policyengine_simulation_entry.economy_api_models import ( + CatalogKey, + EconomyReportRequest, + PolicyChoice, + ReportDescription, + ReportLabel, + SimulationYear, + omit_null_default, +) + +BudgetWindowSize = Annotated[ + int, + Field( + ge=1, + le=MAX_BUDGET_WINDOW_YEARS, + description="How many years the window compares, from its first", + ), +] + + +class BudgetWindowReportRequest(BaseModel): + """A budget-window report: a baseline and a reform in each year of a window. + + ``year`` is the window's first year, and ``budget_window_size`` its + length. Every year uses the request's region and dataset, as an economy + report does. ``spm`` selects the SPM forecast scenario and geography; + omit it, rather than sending null, to use the bundle's certified + defaults. + """ + + model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False) + + country: CountryId + report_type: Literal["budget_window"] + label: ReportLabel + description: ReportDescription | None = None + year: SimulationYear + budget_window_size: BudgetWindowSize + region_code: CatalogKey + dataset_name: CatalogKey | None = None + baseline: PolicyChoice + reform: PolicyChoice + spm: SPMSelection | SkipJsonSchema[None] = Field( + default=None, json_schema_extra=omit_null_default + ) + + @field_validator("spm", mode="before") + @classmethod + def refuse_null_spm(cls, value: object) -> object: + """Refuse an explicit null; omit ``spm`` to send no settings.""" + if value is None: + msg = "spm must be an object when given; omit it to send no settings" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def refuse_years_past_the_last(self) -> BudgetWindowReportRequest: + """Refuse a window whose last year is after 2200.""" + if self.year + self.budget_window_size - 1 > LATEST_YEAR: + msg = f"a budget window must end by {LATEST_YEAR}" + raise ValueError(msg) + return self + + @property + def years(self) -> tuple[int, ...]: + """Return the window's years, earliest first.""" + return tuple(range(self.year, self.year + self.budget_window_size)) + + def options(self) -> dict[str, Any]: + """Return the report's options as it stores them (``reports.inputs``). + + They are a single-year economy report's without cliffs, so the + window's years and such reports reuse each other's runs. ``spm`` + holds only the settings the request gave. + """ + options: dict[str, Any] = {"include_cliffs": False} + if self.spm is not None: + options["spm"] = self.spm.model_dump(mode="json") + return options + + +# The body of ``POST /reports``: an economy report or a budget window, told +# apart by ``report_type``. +ReportRequest = Annotated[ + EconomyReportRequest | BudgetWindowReportRequest, + Field(discriminator="report_type"), +] + + +class _Answer(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class BudgetWindowYearCreated(_Answer): + """One year of a created window: its report run and each side's run. + + ``…_reused`` is true for a side that names a run that already succeeded + in that year, rather than a new one. + """ + + year: int + report_run_id: UUID + baseline_simulation_id: UUID + baseline_simulation_run_id: UUID + baseline_reused: bool + reform_simulation_id: UUID + reform_simulation_run_id: UUID + reform_reused: bool + + +class BudgetWindowReportCreated(_Answer): + """The budget-window report a request created, and its years' runs. + + ``report_run_id`` is the first year's report run, and ``status`` the + window's calculation's: ``pending`` once its first years' calls were + started, or when the start couldn't be recorded. ``years`` lists every + year, earliest first. + """ + + report_id: UUID + report_run_id: UUID + budget_window_run_id: UUID + status: ReportRunStatus + years: list[BudgetWindowYearCreated] diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py new file mode 100644 index 000000000..748905f3c --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py @@ -0,0 +1,419 @@ +"""Creating budget-window reports and starting their first years. + +A budget-window report is created as an economy report with one report run +per year of its window (``NewBudgetWindowReport``), all pending with no +Modal call, in one transaction. The service then advances the window as its +own role (``BudgetWindowStore.advance_window``): under the window run's row +lock, up to the earliest five years' coordinators are started and their +calls recorded. Each year's coordinator starts the next years when its year ends. + +The first year's coordinator input is built as a single-year report's is; +every later year's is derived from it (``year_calculation``), so all years +are calculated alike. A year Modal refuses fails the window, and so does +the first year when Modal doesn't confirm it in time; the request is then +answered 503 naming the report, as an economy report whose run Modal +refused is. The first years share one wait for Modal, and a later year not +confirmed within it is left waiting for a started year's coordinator. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import replace +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID, uuid4 + +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + ReportRunIds, + Stage12InvocationContext, + year_calculation, +) +from policyengine_simulation_persistence import ( + BudgetWindowYear, + NewBudgetWindowReport, + ReportRunStatus, + RunReuse, + SimulationChoice, +) +from policyengine_stage12_persistence import RunStatus, YearStartDeferred +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportCreated, + BudgetWindowYearCreated, +) +from policyengine_simulation_entry.economy_calculation_inputs import ( + report_run_calculation, +) +from policyengine_simulation_entry.economy_submissions import ( + CALL_START_TIMEOUT_SECONDS, + EconomyRuns, + EconomySubmissionService, + ServiceSettings, + calculation_scope, + call_refused, + economy_request_for, + simulation_choice, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from policyengine_simulation_persistence import ( + BudgetWindowInputs, + EconomyRequest, + ReusableRun, + SimulationInputs, + ) + from policyengine_stage12_persistence import ( + WindowAdvance, + YearStarter, + YearToStart, + ) + + from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportRequest, + ) + from policyengine_simulation_entry.economy_calculation_inputs import ( + RunCalculation, + ) + from policyengine_simulation_entry.economy_workers import ( + EconomyRunCaller, + EconomyWorker, + EconomyWorkerSource, + ) + +logger = logging.getLogger(__name__) + + +class BudgetWindowRuns(EconomyRuns, Protocol): + """The economy store's reads and writes, with budget windows'.""" + + def read_budget_window_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + years: Sequence[int], + ) -> BudgetWindowInputs: + """Read the window's catalog rows, both sides' values and each year's runs.""" + ... + + def create_budget_window_report(self, report: NewBudgetWindowReport) -> None: + """Create the window's report, its budget-window run and each year's runs.""" + ... + + +class BudgetWindowAdvancer(Protocol): + """Starts a window's waiting years (``BudgetWindowStore`` as the service).""" + + def advance_window( + self, + window_run_id: UUID, + *, + now: datetime, + start_year: YearStarter | None = None, + ) -> WindowAdvance: + """Start the earliest waiting years, at most five under way.""" + ... + + +class ReportSubmissionService(EconomySubmissionService): + """Create economy and budget-window reports and simulations, and start them.""" + + def __init__( # noqa: PLR0913 # the window seams, then the economy service's + self, + *, + windows: BudgetWindowAdvancer, + store: BudgetWindowRuns, + workers: EconomyWorkerSource, + caller: EconomyRunCaller, + settings: ServiceSettings, + call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, + new_id: Callable[[], UUID] = uuid4, + clock: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + """Use these seams; tests replace each of them.""" + super().__init__( + store=store, + workers=workers, + caller=caller, + settings=settings, + call_start_timeout_seconds=call_start_timeout_seconds, + new_id=new_id, + clock=clock, + ) + self._windows = windows + self._window_store = store + + async def submit_budget_window( + self, + request: BudgetWindowReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> BudgetWindowReportCreated: + """Create the window's report and its years' runs, and start its first years. + + Raise the errors ``submit_report`` raises; the 503 after a write + names the report, whose window was failed. + """ + worker = await self._default_worker(request.country) + options = request.options() + reuse_prefix = self._settings.reuse_prefix + reuse = ( + None + if reuse_prefix is None + else RunReuse( + versions=worker.versions, + inputs=options, + output_uri_prefix=reuse_prefix, + ) + ) + economy_request = economy_request_for(request, worker) + inputs = await self._read_inputs( + lambda: self._window_store.read_budget_window_inputs( + economy_request, + baseline=simulation_choice(request.baseline), + reform=simulation_choice(request.reform), + reuse=reuse, + years=request.years, + ), + request=economy_request, + ) + report = NewBudgetWindowReport( + id=self._new_id(), + window_run_id=self._new_id(), + country=request.country, + label=request.label, + description=request.description, + tax_benefit_model_id=inputs.catalog.model_id, + population=inputs.catalog.population(request.year), + versions=worker.versions, + inputs=options, + years=tuple(self._window_year(inputs, year) for year in request.years), + ) + first_year = report.first_year + first_calculation = report_run_calculation( + worker, + calculation_scope(request, worker, inputs.catalog, spm=options.get("spm")), + self._origin(request_id), + run_ids=_run_ids(first_year), + baseline_policy=inputs.baseline.reform, + reform_policy=inputs.reform.reform, + ) + await self._create( + "report", lambda: self._window_store.create_budget_window_report(report) + ) + status = await self._start_first_years( + worker, + report, + first_calculation, + observability_context=observability_context, + ) + return BudgetWindowReportCreated( + report_id=report.id, + report_run_id=first_year.report_run_id, + budget_window_run_id=report.window_run_id, + status=status, + years=[_year_created(window_year) for window_year in report.years], + ) + + def _window_year(self, inputs: BudgetWindowInputs, year: int) -> BudgetWindowYear: + """Return a year naming its own year's reusable runs, or new ones.""" + baseline_run, reform_run = inputs.reusable_runs[year] + return BudgetWindowYear( + year=year, + report_run_id=self._new_id(), + baseline=self._side(_with_reusable(inputs.baseline, baseline_run)), + reform=self._side(_with_reusable(inputs.reform, reform_run)), + ) + + async def _start_first_years( + self, + worker: EconomyWorker, + report: NewBudgetWindowReport, + first_calculation: RunCalculation, + *, + observability_context: dict[str, Any] | None, + ) -> ReportRunStatus: + """Advance the new window, starting its first years; return its status. + + Raise EconomySubmissionUnavailableError naming the report when the + window failed because Modal refused a year, or didn't confirm the + first in time; a later year not confirmed within the shared wait is + left waiting. An advance that can't be written leaves the window pending: + the report exists, so it is answered as created, and a year Modal + started claims itself, while the stuck-run sweep fails a window no + year advances. + """ + start_year = _YearStarts( + worker=worker, + caller=self._caller, + first_year=report.first_year.year, + first_calculation=first_calculation, + observability_context=observability_context, + loop=asyncio.get_running_loop(), + timeout_seconds=self._call_start_timeout_seconds, + clock=self._clock, + ) + try: + advance = await asyncio.to_thread( + self._windows.advance_window, + report.window_run_id, + now=self._clock(), + start_year=start_year, + ) + except SQLAlchemyError as error: + logger.warning( + "budget_window_not_started", + extra={ + "record": "report", + "record_id": str(report.id), + "error_type": type(error).__name__, + }, + ) + return ReportRunStatus.PENDING + if advance.status is RunStatus.FAILED: + # The budget-window store logs the refused call's error type. + logger.error( + "budget_window_call_unconfirmed", + extra={ + "record": "report", + "record_id": str(report.id), + "started_years": list(advance.started_years), + }, + ) + refused = call_refused("report", report.id) + raise refused from None + return ReportRunStatus(advance.status.value) + + +class _YearStarts: + """Start a window year's coordinator from the advance's thread; return its call. + + The first year's call takes the input built for it; each later year's is + derived from the first's. The call starts on the service's event loop, + which waits for the advance meanwhile. All the years share one wait for + Modal, as long as one call's, so the request never waits longer than an + economy report's would. Any failure of the first year's call, or Modal + not confirming it in time, raises, so the advance fails the year and the + window. A later year not confirmed within the shared wait is deferred + (``YearStartDeferred``): it waits, with the years after it, for a started + year's coordinator to start it. + """ + + def __init__( # noqa: PLR0913 # what each year's call is started with + self, + *, + worker: EconomyWorker, + caller: EconomyRunCaller, + first_year: int, + first_calculation: RunCalculation, + observability_context: dict[str, Any] | None, + loop: asyncio.AbstractEventLoop, + timeout_seconds: float, + clock: Callable[[], datetime], + ) -> None: + """Keep what every year's call needs.""" + self._worker = worker + self._caller = caller + self._first_year = first_year + self._first_calculation = first_calculation + self._observability_context = observability_context + self._loop = loop + self._clock = clock + self._timeout_seconds = timeout_seconds + # The shared wait starts with the first start, after the advance's + # lock and reads, so only Modal's time counts against it. + self._wait_ends_at: float | None = None + + def __call__(self, year_to_start: YearToStart, /) -> str: + """Start the year's coordinator and return its Modal call ID.""" + if self._wait_ends_at is None: + self._wait_ends_at = time.monotonic() + self._timeout_seconds + later_year = year_to_start.year != self._first_year + remaining_seconds = self._wait_ends_at - time.monotonic() + if later_year and remaining_seconds <= 0: + logger.info("budget_window_year_deferred", extra=_deferral(year_to_start)) + raise YearStartDeferred + payload, context = self._calculation(year_to_start) + started = asyncio.run_coroutine_threadsafe( + self._caller.start_report_run( + self._worker, + payload, + context, + observability_context=self._observability_context, + ), + self._loop, + ) + try: + return started.result(timeout=max(remaining_seconds, 0)) + except TimeoutError: + started.cancel() + if later_year: + logger.info( + "budget_window_year_deferred", extra=_deferral(year_to_start) + ) + raise YearStartDeferred from None + raise + except BaseException: + started.cancel() + raise + + def _calculation( + self, year_to_start: YearToStart + ) -> tuple[dict[str, Any], dict[str, Any]]: + first = self._first_calculation + if year_to_start.year == self._first_year: + return first.payload, first.context + year_report, year_context = year_calculation( + ReportExecutionInput.model_validate(first.payload), + Stage12InvocationContext.model_validate(first.context), + year=year_to_start.year, + run_ids=year_to_start.run_ids, + created_at=self._clock(), + ) + return year_report.model_dump(mode="json"), year_context.model_dump(mode="json") + + +def _with_reusable( + inputs: SimulationInputs, reusable_run: ReusableRun | None +) -> SimulationInputs: + """Return a side's values naming the year's reusable run, if any.""" + return replace(inputs, reusable_run=reusable_run) + + +def _deferral(year_to_start: YearToStart) -> dict[str, Any]: + """Return a deferred year's log fields: its year and its report run.""" + return { + "year": year_to_start.year, + "record_id": str(year_to_start.run_ids.report_run_id), + } + + +def _run_ids(window_year: BudgetWindowYear) -> ReportRunIds: + return ReportRunIds( + report_run_id=window_year.report_run_id, + baseline_simulation_run_id=window_year.baseline.simulation_run_id, + reform_simulation_run_id=window_year.reform.simulation_run_id, + ) + + +def _year_created(window_year: BudgetWindowYear) -> BudgetWindowYearCreated: + return BudgetWindowYearCreated( + year=window_year.year, + report_run_id=window_year.report_run_id, + baseline_simulation_id=window_year.baseline.simulation_id, + baseline_simulation_run_id=window_year.baseline.simulation_run_id, + baseline_reused=window_year.baseline.reused, + reform_simulation_id=window_year.reform.simulation_id, + reform_simulation_run_id=window_year.reform.simulation_run_id, + reform_reused=window_year.reform.reused, + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py index b161c91e7..84caea1b3 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -84,7 +84,7 @@ ] -def _omit_null_default(schema: dict[str, Any]) -> None: +def omit_null_default(schema: dict[str, Any]) -> None: """Leave out a field's null default: the field is omitted rather than null.""" schema.pop("default", None) @@ -137,7 +137,7 @@ class EconomyReportRequest(_Request): reform: PolicyChoice include_cliffs: StrictBool = False spm: SPMSelection | SkipJsonSchema[None] = Field( - default=None, json_schema_extra=_omit_null_default + default=None, json_schema_extra=omit_null_default ) @field_validator("spm", mode="before") @@ -291,5 +291,5 @@ class EconomyRequestUnprocessable(_Answer): detail: list[EconomyRequestReason] code: EconomyUnprocessableCode | SkipJsonSchema[None] = Field( - default=None, json_schema_extra=_omit_null_default + default=None, json_schema_extra=omit_null_default ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py index f04378af8..879cda5e8 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py @@ -35,6 +35,7 @@ SimulationExecutionInput, SimulationRole, Stage12InvocationContext, + report_run_artifact_prefix, ) if TYPE_CHECKING: @@ -185,9 +186,8 @@ def _context( simulation_callable=worker.simulation_callable, version_manifest_sha256=worker.version_manifest_sha256, bundle_manifest_sha256=worker.bundle_manifest_sha256, - artifact_prefix=( - f"stage-12-runs/{origin.environment}/{created_at:%Y}/{created_at:%m}/" - f"{run_id}" + artifact_prefix=report_run_artifact_prefix( + origin.environment, created_at, run_id ), created_at=created_at, retention_expires_at=created_at + ARTIFACT_RETENTION, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py index 9dd6e5285..1a04277cf 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py @@ -3,7 +3,10 @@ - ``POST /reports`` (permission ``simulation:submit``) creates an economy report, its simulations and its first run, and starts the report run's coordinator, answering 201 with the report's ID, its run's ID and status, - and each side's simulation and simulation run. + and each side's simulation and simulation run. With ``report_type`` + ``budget_window`` it creates a budget-window report with a report run per + year and starts up to its first five years, answering 201 with the report, the + first year's run, the window run, its status and every year's runs. - ``POST /simulations`` (``simulation:submit``) creates a standalone economy simulation and its first run, and starts the run's simulation function, answering 201 with the simulation's ID, its run's ID and status. @@ -38,6 +41,11 @@ from fastapi import APIRouter, Depends, Request, status from fastapi.responses import JSONResponse +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportCreated, + BudgetWindowReportRequest, + ReportRequest, +) from policyengine_simulation_entry.economy_api_models import ( EconomyError, EconomyReportCreated, @@ -56,6 +64,7 @@ EconomySubmissionUnavailableError, EconomySubmissionUnprocessableError, ) +from policyengine_simulation_entry.report_request_errors import ReportRequestRoute from policyengine_simulation_entry.route_support import ( log_unexpected_failure, observability_context, @@ -160,7 +169,7 @@ def _unavailable_description(record: RecordName) -> dict[str, Any]: class EconomySubmissions(Protocol): - """Create economy reports and simulations (``EconomySubmissionService``).""" + """Create reports and simulations (``ReportSubmissionService``).""" async def submit_report( self, @@ -172,6 +181,16 @@ async def submit_report( """Create the report and its runs and start its coordinator.""" ... + async def submit_budget_window( + self, + request: BudgetWindowReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> BudgetWindowReportCreated: + """Create the window's report and its years' runs, and start its first years.""" + ... + async def submit_simulation( self, request: EconomySimulationRequest, @@ -392,21 +411,27 @@ def _submission_routes( ) -> APIRouter: """Return the router of ``POST /reports`` and ``POST /simulations``.""" router = APIRouter() + # A refused report body is located as an economy report's was. + report_router = APIRouter(route_class=ReportRequestRoute) - @router.post( + @report_router.post( REPORTS_PATH, - summary="Create an economy report", + summary="Create an economy or budget-window report", description=( "Create an economy report comparing a baseline and a reform " "simulation of one region's data for one year, on the v2 version " "manifest's default worker, and start its first run. A side whose " - "simulation already succeeded on the same versions reuses that run." + "simulation already succeeded on the same versions reuses that run. " + "A budget-window report (report_type budget_window) compares them " + "in each year of a window of 1 to 75 years, with a report run per " + "year, at most five running at once; each year reuses its own " + "year's runs." ), operation_id="create_report", status_code=status.HTTP_201_CREATED, - response_model=EconomyReportCreated, + response_model=EconomyReportCreated | BudgetWindowReportCreated, responses={ - 201: {"description": "The report and its first run were created"}, + 201: {"description": "The report and its first runs were created"}, 400: _REFUSED, 403: _FORBIDDEN, 422: _UNPROCESSABLE, @@ -416,17 +441,24 @@ def _submission_routes( dependencies=[Depends(submitter)], ) async def create_report( - body: EconomyReportRequest, request: Request + body: ReportRequest, + request: Request, ) -> JSONResponse: if submissions is None: return _not_configured() observability_id = entry.start_workflow_observability(request) try: - created = await submissions.submit_report( - body, - request_id=call_request_id_for(request), - observability_context=observability_context(entry, observability_id), - ) + call_request_id = call_request_id_for(request) + context = observability_context(entry, observability_id) + created: EconomyReportCreated | BudgetWindowReportCreated + if isinstance(body, BudgetWindowReportRequest): + created = await submissions.submit_budget_window( + body, request_id=call_request_id, observability_context=context + ) + else: + created = await submissions.submit_report( + body, request_id=call_request_id, observability_context=context + ) except _SUBMISSION_FAILURES as error: return _submission_failure_response(error, "report") # Every other failure is still answered as the JSON 500 declared. @@ -438,6 +470,8 @@ async def create_report( created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED ) + router.include_router(report_router) + @router.post( SIMULATIONS_PATH, summary="Create an economy simulation", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py index 3247e5ecb..a6011bae7 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py @@ -95,6 +95,9 @@ SimulationInputs, ) + from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportRequest, + ) from policyengine_simulation_entry.economy_api_models import ( EconomyRefusalCode, EconomyReportRequest, @@ -206,7 +209,8 @@ def _database_unavailable() -> EconomySubmissionUnavailableError: ) -def _choice(policy: PolicyChoice) -> SimulationChoice: +def simulation_choice(policy: PolicyChoice) -> SimulationChoice: + """Return a request side's stored policy and dynamic.""" return SimulationChoice(policy_id=policy.policy_id, dynamic_id=policy.dynamic_id) @@ -266,17 +270,19 @@ async def submit_report( output_uri_prefix=reuse_prefix, ) ) - economy_request = _economy_request(request, worker) + economy_request = economy_request_for(request, worker) inputs = await self._read_inputs( lambda: self._store.read_report_inputs( economy_request, - baseline=_choice(request.baseline), - reform=_choice(request.reform), + baseline=simulation_choice(request.baseline), + reform=simulation_choice(request.reform), reuse=reuse, ), request=economy_request, ) - scope = _scope(request, worker, inputs.catalog, spm=options.get("spm")) + scope = calculation_scope( + request, worker, inputs.catalog, spm=options.get("spm") + ) report = NewEconomyReport( id=self._new_id(), run_id=self._new_id(), @@ -331,7 +337,7 @@ async def submit_simulation( its run was failed. """ worker = await self._default_worker(request.country) - economy_request = _economy_request(request, worker) + economy_request = economy_request_for(request, worker) choice = SimulationChoice( policy_id=request.policy_id, dynamic_id=request.dynamic_id ) @@ -339,7 +345,7 @@ async def submit_simulation( lambda: self._store.read_simulation_inputs(economy_request, choice), request=economy_request, ) - scope = _scope(request, worker, inputs.catalog, spm=None) + scope = calculation_scope(request, worker, inputs.catalog, spm=None) simulation = NewEconomySimulation( id=self._new_id(), run_id=self._new_id(), @@ -513,7 +519,7 @@ async def _start_report_run( ), ) if failed: - refused = _call_refused("report", report.id) + refused = call_refused("report", report.id) raise refused from None return ReportRunStatus.PENDING status = await self._record_call( @@ -552,7 +558,7 @@ async def _start_simulation_run( ), ) if failed: - refused = _call_refused("simulation", simulation.id) + refused = call_refused("simulation", simulation.id) raise refused from None return SimulationStatus.PENDING status = await self._record_call( @@ -629,9 +635,10 @@ async def _record_call( return None if status is None else str(status) -def _call_refused( +def call_refused( record: RecordName, record_id: UUID ) -> EconomySubmissionUnavailableError: + """Return the 503 naming a record whose run's call Modal did not accept.""" return EconomySubmissionUnavailableError( code="modal_call_refused", detail=( @@ -642,9 +649,13 @@ def _call_refused( ) -def _economy_request( - request: EconomyReportRequest | EconomySimulationRequest, worker: EconomyWorker +def economy_request_for( + request: EconomyReportRequest + | EconomySimulationRequest + | BudgetWindowReportRequest, + worker: EconomyWorker, ) -> EconomyRequest: + """Return what the request names apart from policies, for the default worker.""" return EconomyRequest( country=request.country, policyengine_version=worker.policyengine_version, @@ -654,8 +665,10 @@ def _economy_request( ) -def _scope( - request: EconomyReportRequest | EconomySimulationRequest, +def calculation_scope( + request: EconomyReportRequest + | EconomySimulationRequest + | BudgetWindowReportRequest, worker: EconomyWorker, catalog: EconomyCatalogRows, *, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py new file mode 100644 index 000000000..bd574462f --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py @@ -0,0 +1,97 @@ +"""Refused ``POST /reports`` bodies, located as before the budget-window type. + +``POST /reports`` takes an economy report or a budget window, told apart by +``report_type``. Pydantic puts the matched type into every refused field's +location (``["body", "economy", "year"]``) and reports a missing or unknown +``report_type`` at the body itself. The Public API passes the locations to +its callers, so the report route answers them as it did when it took only +economy reports: without the type, and with ``report_type`` named. A body +with no ``report_type`` is refused for that alone, before its other fields +are checked, since its type decides which fields it must have. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Final, get_args + +from fastapi.exceptions import RequestValidationError +from fastapi.routing import APIRoute + +from policyengine_simulation_entry.budget_window_api_models import ReportRequest + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine, Sequence + + from fastapi import Request, Response + +# The values of ``report_type``, which Pydantic puts into the locations: one +# per request model the route's body takes. +REPORT_TYPES: Final = tuple( + report_type + for model in get_args(get_args(ReportRequest)[0]) + for report_type in get_args(model.model_fields["report_type"].annotation) +) + + +def _expected(report_types: Sequence[str]) -> str: + """Return the types as Pydantic lists a literal's values: 'a', 'b' or 'c'.""" + quoted = [f"'{report_type}'" for report_type in report_types] + return " or ".join( + [", ".join(quoted[:-1]), quoted[-1]] if len(quoted) > 1 else quoted + ) + + +_EXPECTED_TYPES: Final = _expected(REPORT_TYPES) + + +def report_request_errors(errors: Sequence[Any]) -> list[dict[str, Any]]: + """Return the refusals of a report body with the request type left out. + + A field's location loses the type Pydantic matched; a body with no + ``report_type`` is refused as missing it, and one with an unknown type as + a value outside the types. + """ + return [_located(dict(error)) for error in errors] + + +def _located(error: dict[str, Any]) -> dict[str, Any]: + location = tuple(error.get("loc", ())) + kind = error.get("type") + if location == ("body",) and kind == "union_tag_not_found": + return { + "type": "missing", + "loc": ("body", "report_type"), + "msg": "Field required", + "input": error.get("input"), + } + if location == ("body",) and kind == "union_tag_invalid": + body = error.get("input") + return { + "type": "literal_error", + "loc": ("body", "report_type"), + "msg": f"Input should be {_EXPECTED_TYPES}", + # The value sent, as JSON; the error's own tag is its text. + "input": body.get("report_type") if isinstance(body, dict) else None, + "ctx": {"expected": _EXPECTED_TYPES}, + } + if len(location) >= 2 and location[0] == "body" and location[1] in REPORT_TYPES: # noqa: PLR2004 # the body, then the type + return {**error, "loc": (location[0], *location[2:])} + return error + + +class ReportRequestRoute(APIRoute): + """The report route, whose refused bodies are located as an economy report's.""" + + def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]: + """Return the route's handler, relocating a refused body's errors.""" + handle = super().get_route_handler() + + async def handle_report(request: Request) -> Response: + try: + return await handle(request) + except RequestValidationError as error: + raise RequestValidationError( + report_request_errors(error.errors()), body=error.body + ) from None + + return handle_report diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_requests.py b/projects/policyengine-simulation-entry/tests/test_budget_window_requests.py new file mode 100644 index 000000000..acf3edfbf --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_requests.py @@ -0,0 +1,71 @@ +"""A budget-window report request: one region's baseline and reform over many years.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from economy_service_fakes import REPORT_BODY +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportRequest, +) +from pydantic import ValidationError + +WINDOW_BODY: dict[str, Any] = { + **{name: value for name, value in REPORT_BODY.items() if name != "include_cliffs"}, + "report_type": "budget_window", + "budget_window_size": 10, +} + + +def _request(**changes: Any) -> BudgetWindowReportRequest: + return BudgetWindowReportRequest.model_validate({**WINDOW_BODY, **changes}) + + +def test_a_window_names_its_years_from_its_first() -> None: + request = _request() + + assert request.years == tuple(range(2026, 2036)) + + +def test_a_window_stores_the_options_a_report_without_cliffs_stores() -> None: + # Reuse compares a report's options, so a window's years and single-year + # reports without cliffs reuse each other's runs. + assert _request().options() == { + "include_cliffs": False, + "spm": {"geography_kind": "national"}, + } + without_spm = {name: value for name, value in WINDOW_BODY.items() if name != "spm"} + assert BudgetWindowReportRequest.model_validate(without_spm).options() == { + "include_cliffs": False + } + + +@pytest.mark.parametrize( + "changes", + [ + {"include_cliffs": True}, + {"include_cliffs": False}, + {"budget_window_size": 0}, + {"budget_window_size": 76}, + {"year": 2195, "budget_window_size": 10}, + {"report_type": "economy"}, + {"spm": None}, + ], + ids=[ + "cliffs", + "cliffs-false", + "no-years", + "76-years", + "past-2200", + "economy-type", + "null-spm", + ], +) +def test_a_window_request_is_refused_when(changes) -> None: + with pytest.raises(ValidationError): + _request(**changes) + + +def test_a_window_of_one_year_ending_in_2200_is_accepted() -> None: + assert _request(year=2200, budget_window_size=1).years == (2200,) diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py b/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py new file mode 100644 index 000000000..b9a8f043a --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py @@ -0,0 +1,219 @@ +"""POST /reports with a budget window: the route sends it to the window service.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any +from uuid import UUID + +import pytest +from conftest import FakeBackend, make_settings +from economy_service_fakes import REPORT_BODY +from fastapi.testclient import TestClient +from policyengine_simulation_entry import app as app_module +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportCreated, + BudgetWindowYearCreated, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionUnavailableError, +) +from policyengine_simulation_persistence import ReportRunStatus +from test_budget_window_requests import WINDOW_BODY + +REPORT_ID = UUID("00000000-0000-4000-8000-0000000000e1") +WINDOW_RUN_ID = UUID("00000000-0000-4000-8000-0000000000e2") +YEAR_RUN_IDS = [UUID(f"00000000-0000-4000-8000-0000000001{digit}0") for digit in "12"] +CREATED_WINDOW = BudgetWindowReportCreated( + report_id=REPORT_ID, + report_run_id=YEAR_RUN_IDS[0], + budget_window_run_id=WINDOW_RUN_ID, + status=ReportRunStatus.PENDING, + years=[ + BudgetWindowYearCreated( + year=2026 + index, + report_run_id=run_id, + baseline_simulation_id=UUID(int=10 + index), + baseline_simulation_run_id=UUID(int=20 + index), + baseline_reused=False, + reform_simulation_id=UUID(int=30 + index), + reform_simulation_run_id=UUID(int=40 + index), + reform_reused=False, + ) + for index, run_id in enumerate(YEAR_RUN_IDS) + ], +) + + +@dataclass +class FakeSubmissions: + """Records which service method each report reached.""" + + failure: Exception | None = None + windows: list[Any] = field(default_factory=list) + reports: list[Any] = field(default_factory=list) + + async def submit_budget_window(self, request, *, request_id, observability_context): + self.windows.append(request) + if self.failure is not None: + raise self.failure + return CREATED_WINDOW + + async def submit_report(self, request, *, request_id, observability_context): + self.reports.append(request) + msg = "an economy report reached the test" + raise AssertionError(msg) + + +@pytest.fixture +def client_for(monkeypatch): + def build(submissions: FakeSubmissions) -> TestClient: + monkeypatch.setattr( + app_module, + "economy_submissions_for", + lambda settings, database: submissions, + ) + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + ) + ) + + return build + + +def test_a_budget_window_answers_201_with_its_years(client_for) -> None: + submissions = FakeSubmissions() + + with client_for(submissions) as client: + result = client.post("/reports", json=WINDOW_BODY) + + assert result.status_code == 201 + assert result.json() == CREATED_WINDOW.model_dump(mode="json") + ((request,),) = [submissions.windows] + assert (request.year, request.budget_window_size) == (2026, 10) + assert submissions.reports == [] + + +def test_a_budget_window_sending_cliffs_is_refused_before_the_service( + client_for, +) -> None: + submissions = FakeSubmissions() + + with client_for(submissions) as client: + result = client.post("/reports", json={**WINDOW_BODY, "include_cliffs": True}) + + assert result.status_code == 422 + assert submissions.windows == [] + + +def test_a_refused_first_year_answers_503_naming_the_report(client_for) -> None: + submissions = FakeSubmissions( + failure=EconomySubmissionUnavailableError( + code="modal_call_refused", + detail="Modal did not accept the report's call.", + record_id=REPORT_ID, + ) + ) + + with client_for(submissions) as client: + result = client.post("/reports", json=WINDOW_BODY) + + assert result.status_code == 503 + assert result.json()["report_id"] == str(REPORT_ID) + assert result.json()["code"] == "modal_call_refused" + + +def test_an_economy_report_body_still_reaches_the_report_service(client_for) -> None: + submissions = FakeSubmissions() + + with client_for(submissions) as client: + client.post("/reports", json=REPORT_BODY) + + assert len(submissions.reports) == 1 + assert submissions.windows == [] + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + ( + {**REPORT_BODY, "year": 1800}, + {"loc": ["body", "year"], "type": "greater_than_equal"}, + ), + ( + { + name: value + for name, value in REPORT_BODY.items() + if name != "report_type" + }, + {"loc": ["body", "report_type"], "type": "missing"}, + ), + ( + {**REPORT_BODY, "report_type": "household"}, + {"loc": ["body", "report_type"], "type": "literal_error"}, + ), + ( + {**WINDOW_BODY, "budget_window_size": 0}, + {"loc": ["body", "budget_window_size"], "type": "greater_than_equal"}, + ), + ( + {**WINDOW_BODY, "year": 2195}, + {"loc": ["body"], "type": "value_error"}, + ), + ], + ids=["economy-field", "no-type", "unknown-type", "window-field", "window-years"], +) +def test_a_refused_report_names_its_fields_as_an_economy_report_did( + client_for, body, expected +) -> None: + # Telling the two request types apart adds nothing to a refused field's + # location: the Public API passes the locations to its callers. + with client_for(FakeSubmissions()) as client: + result = client.post("/reports", json=body) + + assert result.status_code == 422 + (error,) = result.json()["detail"] + assert {"loc": error["loc"], "type": error["type"]} == expected + + +@pytest.mark.parametrize("report_type", [None, 5, ["economy"], "household"]) +def test_an_unknown_report_type_is_refused_with_the_value_sent( + client_for, report_type +) -> None: + with client_for(FakeSubmissions()) as client: + result = client.post( + "/reports", json={**REPORT_BODY, "report_type": report_type} + ) + + assert result.status_code == 422 + assert result.json()["detail"] == [ + { + "type": "literal_error", + "loc": ["body", "report_type"], + "msg": "Input should be 'economy' or 'budget_window'", + "input": report_type, + "ctx": {"expected": "'economy' or 'budget_window'"}, + } + ] + + +def test_a_report_with_no_type_is_refused_as_an_economy_report_was( + client_for, +) -> None: + body = {name: value for name, value in REPORT_BODY.items() if name != "report_type"} + + with client_for(FakeSubmissions()) as client: + result = client.post("/reports", json=body) + + assert result.json()["detail"] == [ + { + "type": "missing", + "loc": ["body", "report_type"], + "msg": "Field required", + "input": body, + } + ] diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_submissions.py b/projects/policyengine-simulation-entry/tests/test_budget_window_submissions.py new file mode 100644 index 000000000..f935d2c1c --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_submissions.py @@ -0,0 +1,417 @@ +"""Creating a budget-window report and starting its first years. + +The service creates every year's runs in one transaction, then advances the +window, which starts the earliest years (at most five) by calling the +coordinator with each year's input. The first year's input is built as a +single-year report's is, and every later year's is derived from it. The +store and the window advance are faked here; both are tested against +PostgreSQL in the persistence libraries. +""" + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Any +from uuid import UUID + +import pytest +from economy_service_fakes import ( + CATALOG, + CREATED_AT, + OUTPUT_BUCKET, + REFORM, + FakeCaller, + FakeStore, + FakeWorkers, + numbered_ids, +) +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + ReportRunIds, + Stage12InvocationContext, + year_calculation, +) +from policyengine_simulation_entry import budget_window_submissions +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportRequest, +) +from policyengine_simulation_entry.budget_window_submissions import ( + ReportSubmissionService, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionUnavailableError, +) +from policyengine_simulation_entry.economy_submissions import ServiceSettings +from policyengine_simulation_persistence import ( + BudgetWindowInputs, + ReportRunStatus, + ReusableRun, + RunReuse, + SimulationInputs, +) +from policyengine_stage12_persistence import ( + BUDGET_WINDOW_MAX_RUNNING_YEARS, + RunStatus, + WindowAdvance, + YearStartDeferred, + YearToStart, +) +from sqlalchemy.exc import OperationalError +from test_budget_window_requests import WINDOW_BODY + +OBSERVED = {"observability_id": "observed"} +REUSED_IN_2027 = ReusableRun( + simulation_id=UUID("00000000-0000-4000-8000-0000000000d1"), + simulation_run_id=UUID("00000000-0000-4000-8000-0000000000d2"), +) + + +@dataclass +class FakeWindowStore(FakeStore): + """Also reads and creates budget-window reports.""" + + reusable_by_year: dict[int, tuple[ReusableRun | None, ReusableRun | None]] = field( + default_factory=dict + ) + window_reads: list[tuple[Any, ...]] = field(default_factory=list) + windows: list[Any] = field(default_factory=list) + + def read_budget_window_inputs(self, request, *, baseline, reform, reuse, years): + self.window_reads.append((request, baseline, reform, reuse, tuple(years))) + if self.read_failure is not None: + raise self.read_failure + return BudgetWindowInputs( + catalog=self.catalog, + baseline=SimulationInputs(choice=baseline, reform={}), + reform=SimulationInputs(choice=reform, reform=REFORM), + reusable_runs={ + year: self.reusable_by_year.get(year, (None, None)) for year in years + }, + ) + + def create_budget_window_report(self, report) -> None: + if self.create_failure is not None: + raise self.create_failure + self.windows.append(report) + + +@dataclass +class FakeAdvancer: + """Starts the window's earliest years, as the budget-window store does. + + A start that raises fails the window, as Modal refusing a year does. + """ + + store: FakeWindowStore + failure: Exception | None = None + # How long the advance takes, as the database's lock and reads may, before + # it starts its first year, and between one start and the next. + delay_seconds: float = 0.0 + delay_between_starts_seconds: float = 0.0 + advances: list[UUID] = field(default_factory=list) + started: list[tuple[YearToStart, str]] = field(default_factory=list) + deferred: list[int] = field(default_factory=list) + + def advance_window(self, window_run_id, *, now, start_year): + self.advances.append(window_run_id) + time.sleep(self.delay_seconds) + if self.failure is not None: + raise self.failure + (report,) = self.store.windows + started_years = [] + for window_year in report.years[:BUDGET_WINDOW_MAX_RUNNING_YEARS]: + year_to_start = YearToStart( + year=window_year.year, + run_ids=ReportRunIds( + report_run_id=window_year.report_run_id, + baseline_simulation_run_id=window_year.baseline.simulation_run_id, + reform_simulation_run_id=window_year.reform.simulation_run_id, + ), + ) + if started_years: + time.sleep(self.delay_between_starts_seconds) + try: + call_id = start_year(year_to_start) + except YearStartDeferred: + # As the store, the year and the later ones wait. + self.deferred.append(window_year.year) + break + except Exception: # noqa: BLE001 # as the store, any failed start fails the window + return WindowAdvance( + status=RunStatus.FAILED, + changed=True, + started_years=tuple(started_years), + ) + self.started.append((year_to_start, call_id)) + started_years.append(window_year.year) + return WindowAdvance( + status=RunStatus.PENDING, changed=False, started_years=tuple(started_years) + ) + + +def _service(store, *, caller=None, advancer=None): + return ReportSubmissionService( + windows=advancer or FakeAdvancer(store), + store=store, + workers=FakeWorkers(), + caller=caller or FakeCaller(), + settings=ServiceSettings( + environment="staging", + modal_environment="staging", + output_bucket=OUTPUT_BUCKET, + ), + call_start_timeout_seconds=0.05, + new_id=numbered_ids(), + clock=lambda: CREATED_AT, + ) + + +def _submit(submissions, **changes): + request = BudgetWindowReportRequest.model_validate({**WINDOW_BODY, **changes}) + return asyncio.run( + submissions.submit_budget_window( + request, request_id="observed", observability_context=OBSERVED + ) + ) + + +def test_a_window_is_created_with_a_run_per_year_and_its_first_five_started() -> None: + store = FakeWindowStore() + caller = FakeCaller() + advancer = FakeAdvancer(store) + + created = _submit( + _service(store, caller=caller, advancer=advancer), budget_window_size=7 + ) + + ((request, baseline, reform, reuse, years),) = store.window_reads + assert (baseline.policy_id, str(reform.policy_id)) == ( + None, + WINDOW_BODY["reform"]["policy_id"], + ) + assert (request.year, years) == (2026, tuple(range(2026, 2033))) + assert reuse == RunReuse( + versions=reuse.versions, + inputs={"include_cliffs": False, "spm": {"geography_kind": "national"}}, + output_uri_prefix=f"gs://{OUTPUT_BUCKET}/", + ) + (report,) = store.windows + assert [window_year.year for window_year in report.years] == list(range(2026, 2033)) + assert report.population == CATALOG.population(2026) + assert report.inputs == reuse.inputs + assert advancer.advances == [report.window_run_id] + # The first five years were started, each with its own year and runs. + assert len(caller.report_runs) == 5 + for (year_to_start, call_id), (_, payload, context, observed) in zip( + advancer.started, caller.report_runs, strict=True + ): + report_input = ReportExecutionInput.model_validate(payload) + year_context = Stage12InvocationContext.model_validate(context) + assert ( + report_input.baseline.year == report_input.reform.year == year_to_start.year + ) + assert year_context.report_run == year_to_start.run_ids + assert report_input.reform.policy == REFORM + assert (call_id, observed) == ("fc-coordinator", OBSERVED) + first_year = report.years[0] + assert created.model_dump() == { + "report_id": report.id, + "report_run_id": first_year.report_run_id, + "budget_window_run_id": report.window_run_id, + "status": ReportRunStatus.PENDING, + "years": [ + { + "year": window_year.year, + "report_run_id": window_year.report_run_id, + "baseline_simulation_id": window_year.baseline.simulation_id, + "baseline_simulation_run_id": window_year.baseline.simulation_run_id, + "baseline_reused": False, + "reform_simulation_id": window_year.reform.simulation_id, + "reform_simulation_run_id": window_year.reform.simulation_run_id, + "reform_reused": False, + } + for window_year in report.years + ], + } + + +def test_each_later_year_is_derived_from_the_first() -> None: + store = FakeWindowStore() + caller = FakeCaller() + advancer = FakeAdvancer(store) + + _submit(_service(store, caller=caller, advancer=advancer), budget_window_size=3) + + (_, first_payload, first_context, _) = caller.report_runs[0] + for (year_to_start, _), (_, payload, context, _) in zip( + advancer.started[1:], caller.report_runs[1:], strict=True + ): + derived_report, derived_context = year_calculation( + ReportExecutionInput.model_validate(first_payload), + Stage12InvocationContext.model_validate(first_context), + year=year_to_start.year, + run_ids=year_to_start.run_ids, + created_at=CREATED_AT, + ) + assert payload == derived_report.model_dump(mode="json") + assert context == derived_context.model_dump(mode="json") + + +def test_a_year_reuses_its_own_year_s_run() -> None: + store = FakeWindowStore(reusable_by_year={2027: (REUSED_IN_2027, None)}) + + created = _submit(_service(store), budget_window_size=3) + + years = {window_year.year: window_year for window_year in store.windows[0].years} + assert years[2027].baseline.simulation_run_id == REUSED_IN_2027.simulation_run_id + assert years[2027].baseline.reused + assert not years[2026].baseline.reused + assert [year.baseline_reused for year in created.years] == [False, True, False] + + +def test_a_first_year_modal_refuses_answers_503_naming_the_report() -> None: + store = FakeWindowStore() + caller = FakeCaller(failure=RuntimeError("Modal is unavailable")) + + with pytest.raises(EconomySubmissionUnavailableError) as refused: + _submit(_service(store, caller=caller), budget_window_size=3) + + assert refused.value.code == "modal_call_refused" + assert refused.value.record_id == store.windows[0].id + + +def test_a_start_modal_does_not_confirm_in_time_is_refused() -> None: + store = FakeWindowStore() + + with pytest.raises(EconomySubmissionUnavailableError) as refused: + _submit(_service(store, caller=FakeCaller(hang=True)), budget_window_size=2) + + assert refused.value.record_id == store.windows[0].id + + +def test_a_window_whose_advance_cannot_be_written_is_answered_created() -> None: + # The rows exist; a year Modal started claims itself, and the stuck-run + # sweep fails a window no year advances. + store = FakeWindowStore() + advancer = FakeAdvancer( + store, failure=OperationalError("SELECT 1", {}, Exception("closed")) + ) + + created = _submit(_service(store, advancer=advancer), budget_window_size=2) + + assert created.status is ReportRunStatus.PENDING + assert created.report_id == store.windows[0].id + + +@dataclass +class SlowAfterFirstCaller(FakeCaller): + """Starts the first call at once, then waits past any limit for the rest.""" + + def __post_init__(self) -> None: + self.calls = 0 + + async def start_report_run(self, worker, report_payload, context_payload, **kwargs): + self.calls += 1 + if self.calls > 1: + await asyncio.sleep(60) + return await super().start_report_run( + worker, report_payload, context_payload, **kwargs + ) + + +def test_the_first_years_share_one_wait_for_modal() -> None: + # A request starting five years one after another never waits more than + # one call's limit: a later year not started in time is left for the + # started year's coordinator to start. + store = FakeWindowStore() + caller = SlowAfterFirstCaller() + advancer = FakeAdvancer(store) + started_at = time.monotonic() + + created = _submit( + _service(store, caller=caller, advancer=advancer), budget_window_size=5 + ) + + assert time.monotonic() - started_at < 0.5 + assert created.status is ReportRunStatus.PENDING + assert [year_to_start.year for year_to_start, _ in advancer.started] == [2026] + assert advancer.deferred == [2027] + # Once the wait has run out, no later year's call is even started. + assert caller.calls == 2 + + +def test_the_wait_for_modal_starts_when_the_first_year_starts() -> None: + # The advance's lock and reads take longer than the wait itself; the + # first year's call is still given the whole wait. + store = FakeWindowStore() + caller = FakeCaller() + advancer = FakeAdvancer(store, delay_seconds=0.1) + + created = _submit( + _service(store, caller=caller, advancer=advancer), budget_window_size=2 + ) + + assert created.status is ReportRunStatus.PENDING + assert [year_to_start.year for year_to_start, _ in advancer.started] == [ + 2026, + 2027, + ] + + +@dataclass +class RefusingAfterFirstCaller(FakeCaller): + """Starts the first call, then is refused by Modal for the next.""" + + def __post_init__(self) -> None: + self.calls = 0 + + async def start_report_run(self, worker, report_payload, context_payload, **kwargs): + self.calls += 1 + if self.calls > 1: + message = "Modal refused the call" + raise RuntimeError(message) + return await super().start_report_run( + worker, report_payload, context_payload, **kwargs + ) + + +def test_a_later_year_modal_refuses_fails_the_window() -> None: + store = FakeWindowStore() + + with pytest.raises(EconomySubmissionUnavailableError) as refused: + _submit( + _service(store, caller=RefusingAfterFirstCaller()), budget_window_size=3 + ) + + assert refused.value.record_id == store.windows[0].id + + +def test_a_later_year_is_deferred_without_a_call_once_the_wait_has_run_out( + caplog, +) -> None: + store = FakeWindowStore() + caller = FakeCaller() + advancer = FakeAdvancer(store, delay_between_starts_seconds=0.1) + + with caplog.at_level(logging.INFO, logger=budget_window_submissions.__name__): + created = _submit( + _service(store, caller=caller, advancer=advancer), budget_window_size=3 + ) + + assert created.status is ReportRunStatus.PENDING + # The first year's call used the whole wait; the second is never called. + assert len(caller.report_runs) == 1 + assert advancer.deferred == [2027] + (deferral,) = [ + record + for record in caplog.records + if record.getMessage() == "budget_window_year_deferred" + ] + second_year = store.windows[0].years[1] + assert (deferral.year, deferral.record_id) == ( + 2027, + str(second_year.report_run_id), + ) diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 30f8022af..0bf4a8427 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -90,6 +90,9 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: "/simulations/{simulation_id}", ) ECONOMY_SCHEMAS = ( + "BudgetWindowReportCreated", + "BudgetWindowReportRequest", + "BudgetWindowYearCreated", "EconomyError", "EconomyReportCreated", "EconomyReportRequest", @@ -361,14 +364,15 @@ def test_economy_routes_declare_every_answer(): unprocessable = {"$ref": "#/components/schemas/EconomyRequestUnprocessable"} for path, created in ( - ("/reports", "EconomyReportCreated"), - ("/simulations", "EconomySimulationCreated"), + ("/reports", ("EconomyReportCreated", "BudgetWindowReportCreated")), + ("/simulations", ("EconomySimulationCreated",)), ): responses = _responses(spec, "post", path) assert set(responses) == {"201", "400", "403", "422", "500", "503"} - assert _response_schema(responses["201"]) == { - "$ref": f"#/components/schemas/{created}" - } + created_schema = _response_schema(responses["201"]) + assert created_schema.get("anyOf", [created_schema]) == [ + {"$ref": f"#/components/schemas/{name}"} for name in created + ] for status_code in ("400", "500", "503"): assert _response_schema(responses[status_code]) == error assert _response_schema(responses["422"]) == unprocessable @@ -437,6 +441,39 @@ def test_economy_requests_name_no_version_and_no_record_id(): assert "spm" not in schemas["EconomyReportRequest"].get("required", []) +def test_a_report_request_is_an_economy_report_or_a_budget_window(): + spec = create_app().openapi() + body = spec["paths"]["/reports"]["post"]["requestBody"]["content"][ + "application/json" + ]["schema"] + schemas = spec["components"]["schemas"] + + assert body["discriminator"] == { + "propertyName": "report_type", + "mapping": { + "economy": "#/components/schemas/EconomyReportRequest", + "budget_window": "#/components/schemas/BudgetWindowReportRequest", + }, + } + # A window sends its length and no cliff option, and names no version + # or record ID either. + assert set(schemas["BudgetWindowReportRequest"]["properties"]) == { + "country", + "report_type", + "label", + "description", + "year", + "budget_window_size", + "region_code", + "dataset_name", + "baseline", + "reform", + "spm", + } + size = schemas["BudgetWindowReportRequest"]["properties"]["budget_window_size"] + assert (size["minimum"], size["maximum"]) == (1, 75) + + def test_economy_reads_declare_every_answer(): spec = create_app().openapi() error = {"$ref": "#/components/schemas/EconomyError"} 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..bb566ea51 100644 --- a/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py +++ b/projects/policyengine-simulation-executor/src/modal/run_sweep_app.py @@ -2,13 +2,17 @@ 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 +counted from a run's first claim or, while it is pending, from its creation +(a budget window's year from when its call was recorded), and from a +household job's creation. It asks Modal for the status of each recorded call and marks a run or job failed once its call has ended, and fails a pending one that has no call recorded, because nothing will run it. It also fails any still unfinished 24 hours after it was created or started, -longer than any Modal call of it can run, whatever Modal reports. Its result -counts runs and household jobs together. It connects as the Modal worker +longer than any Modal call of it can run, whatever Modal reports. A failed +year of a budget window fails its window, and a window no year has +advanced for ten minutes is finished from its years, or failed as stalled. +Its result counts runs and household jobs together, and the windows it +finished. It connects as the Modal worker role, from the same Modal secret the v2 coordinator, workers and household functions use. This application is not versioned with the .py bundle: one copy runs per Modal environment. @@ -78,11 +82,14 @@ secrets=[modal_worker_database_secret], ) def sweep_stuck_runs() -> dict: - """Fail the runs and household jobs nothing will finish; count them and why.""" + """Fail the runs, household jobs and windows nothing will finish; count them.""" from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the container, not at deploy time report_run_access, run_sweep, ) - summary = run_sweep.sweep_stuck_runs(report_run_access.report_run_store()) + summary = run_sweep.sweep_stuck_runs( + report_run_access.report_run_store(), + windows=report_run_access.budget_window_store(), + ) return asdict(summary) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 2ef7dabc2..801f12cb9 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -36,6 +36,7 @@ ) from policyengine_simulation_executor.stage12_cache.models import CacheManifest from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore +from policyengine_simulation_executor.stage12_runtime import report_run_access from policyengine_simulation_executor.stage12_runtime.report_run_access import ( MODAL_WORKER_DATABASE_SECRET_NAME, MODAL_WORKER_DATABASE_URL_ENV, @@ -316,7 +317,9 @@ def validate_worker_uk() -> dict: memory=32768, timeout=3600, retries=0, - max_containers=300, + # Every segment of every US simulation call that may run at once: a + # segment queued behind this limit counts against its caller's wait. + max_containers=1000, secrets=worker_secrets, ) def run_single_simulation_segment_us( @@ -356,7 +359,8 @@ def run_single_simulation_segment_us( memory=32768, timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, - max_containers=10, + # Both simulations of every coordinator call that may run at once. + max_containers=50, secrets=report_run_writer_secrets, ) def run_single_simulation_us( @@ -419,7 +423,8 @@ def segmented_calculator( memory=32768, timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, - max_containers=10, + # Both simulations of every coordinator call that may run at once. + max_containers=50, secrets=report_run_writer_secrets, ) def run_single_simulation_uk( @@ -471,7 +476,8 @@ def run_single_simulation_uk( # the production-result wait, and time to finish and record the report. timeout=COORDINATOR_TIMEOUT_SECONDS, retries=0, - max_containers=10, + # Three budget windows at five years each, and ten other reports. + max_containers=25, secrets=report_run_writer_secrets, ) def coordinate_report( @@ -522,15 +528,25 @@ def coordinate_report( }, remote_context=propagated, ): - return stage12_runtime.coordinate_report( - payload, - context, - parent, - invocation=stage12_runtime.CoordinatorInvocation( - application_name=APP_NAME, - invocation_id=coordinator_invocation_id, + # A call naming a budget window's year marks the window running + # first and, when the year ends, starts the window's next years. + return stage12_runtime.coordinate_window_year( + lambda: stage12_runtime.coordinate_report( + payload, + context, + parent, + invocation=stage12_runtime.CoordinatorInvocation( + application_name=APP_NAME, + invocation_id=coordinator_invocation_id, + ), + dependencies=stage12_runtime.CoordinatorDependencies( + runtime=runtime + ), ), - dependencies=stage12_runtime.CoordinatorDependencies(runtime=runtime), + payload=payload, + context=context, + windows=lambda: report_run_access.budget_window_store(runtime), + observability_context=runtime.capture_context, ) finally: runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py index 0044be8e0..64d685a35 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py @@ -22,6 +22,8 @@ RowIdentity, SimulationArtifactDescriptor, Stage12InvocationContext, + report_run_artifact_prefix, + require_artifact_environment, stage12_output_plan_sha256, ) from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata @@ -52,13 +54,6 @@ def canonical_json_bytes(value: object) -> bytes: ).encode("utf-8") -def _require_environment_segment(environment: str) -> None: - """Raise ``ValueError`` for an environment that could escape its key prefix.""" - if not environment or "/" in environment or ".." in environment: - msg = "invalid Stage 12 artifact environment" - raise ValueError(msg) - - def comparison_run_prefix( *, environment: str, @@ -68,15 +63,10 @@ def comparison_run_prefix( """Return ``stage-12-runs////`` for one run. Raises ``ValueError`` for an environment that could escape the prefix or a - naive ``created_at``. + naive ``created_at``. The contract builds the prefix, so the simulation + service and the coordinator name a run's artifacts alike. """ - _require_environment_segment(environment) - if created_at.tzinfo is None: - msg = "Stage 12 artifact timestamp must include a timezone" - raise ValueError(msg) - return ( - f"stage-12-runs/{environment}/{created_at:%Y}/{created_at:%m}/{evaluation_id}" - ) + return report_run_artifact_prefix(environment, created_at, evaluation_id) def _content_named(name: str, content_sha256: str | None) -> str: @@ -112,7 +102,7 @@ def simulation_run_output_path( attempt's. Raises ``ValueError`` for an environment that could escape the prefix. """ - _require_environment_segment(environment) + require_artifact_environment(environment) return ( f"simulation-runs/{environment}/{country}/{simulation_run_id}/" f"output-{content_sha256}.parquet" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index 174071e91..b122f1ad9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -5,6 +5,7 @@ build_aggregate_report, build_spm_result, ) +from .budget_window_scheduling import coordinate_window_year from .coordination import ( CoordinatorDependencies, CoordinatorInvocation, @@ -68,6 +69,7 @@ "calculate_segment", "calculate_simulation_frames", "coordinate_report", + "coordinate_window_year", "merge_segment_results", "redacted_failure_detail", "run_household_job", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/budget_window_scheduling.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/budget_window_scheduling.py new file mode 100644 index 000000000..f4d3b66e2 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/budget_window_scheduling.py @@ -0,0 +1,296 @@ +"""A budget window's years, started and finished by their coordinators. + +A budget window's years are report runs naming its budget-window run. The +simulation service creates every year and starts the first ones. A year's +coordinator call marks the window running before it claims its year and, +once the year has ended, advances the window: it starts the earliest waiting years +in its own Modal application, so that at most five are under way, or marks +the window failed or succeeded (see ``BudgetWindowStore``). + +Neither step changes the year's own outcome. A failed write is logged, +without the error's text, and the stuck-run sweep finishes a window that +no year advances any more. A pre-emption or cancellation propagates, as +everywhere in the coordinator. +""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import UTC, datetime +from importlib import import_module +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID + +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + Stage12InvocationContext, + year_calculation, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_stage12_persistence import ( + WindowAdvance, + YearStarter, + YearToStart, + ) + +logger = logging.getLogger(__name__) + +# The v2 application's coordinator function. +COORDINATOR_FUNCTION_NAME = "coordinate_report" + + +class YearWindowWriter(Protocol): + """A coordinator's budget-window writes; ``BudgetWindowStore`` implements them.""" + + def window_of_year(self, report_run_id: UUID) -> UUID | None: + """Return the budget-window run a report run is a year of, if any.""" + ... + + def start_window(self, report_run_id: UUID, *, started_at: datetime) -> bool: + """Mark the window of a claimed year running; return whether it changed.""" + ... + + def advance_window( + self, + window_run_id: UUID, + *, + now: datetime, + start_year: YearStarter | None = None, + ) -> WindowAdvance: + """Finish the window from its years, else start its waiting years.""" + ... + + +@dataclass(frozen=True) +class ModalYearStarter: + """Starts a year's coordinator call in this coordinator's Modal application. + + Every year of a window runs on the same .py version, so the next year's + coordinator is this application's ``coordinate_report``. Its input is + derived from this coordinator's own ``payload`` and ``context``: only the + year, its runs, its artifact prefix and its call's creation time differ. + They are checked only when a year is started, so a malformed input fails + that start, as Modal refusing the call would. + """ + + payload: Mapping[str, Any] + context: Mapping[str, Any] + observability_context: dict[str, Any] | None = None + clock: Callable[[], datetime] = lambda: datetime.now(UTC) + + def __call__(self, year_to_start: YearToStart, /) -> str: + """Spawn the year's coordinator call and return its Modal call ID. + + Raises ValueError when the coordinator's own input is malformed. + """ + context = Stage12InvocationContext.model_validate(self.context) + year_report, year_context = year_calculation( + ReportExecutionInput.model_validate(self.payload), + context, + year=year_to_start.year, + run_ids=year_to_start.run_ids, + created_at=self.clock(), + ) + modal: Any = import_module("modal") + function = modal.Function.from_name( + context.modal_application, + COORDINATOR_FUNCTION_NAME, + environment_name=context.modal_environment, + ) + call = function.spawn( + year_report.model_dump(mode="json"), + year_context.model_dump(mode="json"), + observability_context=self.observability_context, + ) + return call.object_id + + +def coordinate_window_year[CoordinatorResult]( + coordinate: Callable[[], CoordinatorResult], + *, + payload: Mapping[str, Any], + context: Mapping[str, Any], + windows: Callable[[], YearWindowWriter], + observability_context: Callable[[], dict[str, Any] | None], +) -> CoordinatorResult: + """Run one coordinator call, starting and advancing its year's window. + + A call naming a report run looks its window up once. A window's year + marks the window running before ``coordinate`` runs, and advances it + once ``coordinate`` returns, raises or is cancelled, starting waiting + years in this application. A pre-emption advances nothing: Modal runs + the call again, and the restarted call advances when it ends, even when + it finds its year already finished. Advancing more than once is safe, + because an advance starts only waiting years and never more than five + under way. A report run outside any window, a call naming no report run, + or a store that cannot be built makes no window writes. When the lookup + itself fails, nothing is marked and the window is looked up again once + the year ends. + """ + report_run_id = _report_run_named_by(context) + window_store = None if report_run_id is None else _built(windows) + if report_run_id is None or window_store is None: + return coordinate() + lookup = _window_lookup(window_store, report_run_id) + if lookup.read and lookup.window_run_id is None: + return coordinate() + if lookup.read: + start_window_of_year(window_store, report_run_id) + try: + result = coordinate() + except KeyboardInterrupt: + raise + except BaseException: + _advance_after_year( + window_store, report_run_id, lookup, payload, context, observability_context + ) + raise + _advance_after_year( + window_store, report_run_id, lookup, payload, context, observability_context + ) + return result + + +@dataclass(frozen=True) +class _WindowLookup: + """Whether a year's window was read, and the window it names, if any.""" + + read: bool + window_run_id: UUID | None = None + + +def _window_lookup(windows: YearWindowWriter, report_run_id: UUID) -> _WindowLookup: + """Look the report run's window up; log a failure and say it wasn't read.""" + try: + return _WindowLookup( + read=True, window_run_id=windows.window_of_year(report_run_id) + ) + except Exception as error: # noqa: BLE001 # the year runs whether or not its window can be read + _log_failed_write("read the budget window of", report_run_id, error) + return _WindowLookup(read=False) + + +def _advance_after_year( # noqa: PLR0913 # the store, the year, its window and the call's inputs + windows: YearWindowWriter, + report_run_id: UUID, + lookup: _WindowLookup, + payload: Mapping[str, Any], + context: Mapping[str, Any], + observability_context: Callable[[], dict[str, Any] | None], +) -> None: + """Advance the window of a year whose coordinator call has ended.""" + advance_window_of_year( + windows, + report_run_id, + start_year=ModalYearStarter( + payload=payload, + context=context, + observability_context=_captured(observability_context), + ), + window_run_id=lookup.window_run_id, + ) + + +def _captured( + observability_context: Callable[[], dict[str, Any] | None], +) -> dict[str, Any] | None: + """Capture the call's observability context for the next years, or None.""" + try: + return observability_context() + except Exception as error: # noqa: BLE001 # the next years start without a trace link + logger.warning( + "Could not capture the observability context for a budget window's " + "next years (%s)", + type(error).__name__, + ) + return None + + +def _report_run_named_by(context: Mapping[str, Any]) -> UUID | None: + """Return the report run a coordinator call's context names, if any.""" + report_run = context.get("report_run") + if not isinstance(report_run, Mapping): + return None + try: + return UUID(str(report_run["report_run_id"])) + except (KeyError, ValueError): + # The coordinator refuses the malformed context itself. + return None + + +def _built(windows: Callable[[], YearWindowWriter]) -> YearWindowWriter | None: + """Build the window store, or log why it can't be built and return None.""" + try: + return windows() + except Exception as error: # noqa: BLE001 # the year runs whether or not its window can be recorded + logger.warning( + "Could not open the budget-window store (%s)", type(error).__name__ + ) + return None + + +def start_window_of_year( + windows: YearWindowWriter, + report_run_id: UUID, + *, + now: datetime | None = None, +) -> None: + """Mark the window of a year this call is about to claim running; never raise. + + A report run that is no window's year changes nothing. + """ + try: + windows.start_window(report_run_id, started_at=now or datetime.now(UTC)) + except Exception as error: # noqa: BLE001 # the year runs whether or not its window records the start + _log_failed_write("mark the budget window of", report_run_id, error) + + +def advance_window_of_year( + windows: YearWindowWriter, + report_run_id: UUID, + *, + start_year: YearStarter, + now: datetime | None = None, + window_run_id: UUID | None = None, +) -> WindowAdvance | None: + """Advance the window of a year that has ended; never raise. + + ``window_run_id`` is the year's window when the caller has read it; + otherwise it is looked up. Returns what the advance did, or None when + the report run is no window's year or the advance failed. + """ + try: + if window_run_id is None: + window_run_id = windows.window_of_year(report_run_id) + if window_run_id is None: + return None + advance = windows.advance_window( + window_run_id, now=now or datetime.now(UTC), start_year=start_year + ) + except Exception as error: # noqa: BLE001 # the year's outcome stands whether or not its window advances + _log_failed_write("advance the budget window of", report_run_id, error) + return None + if advance.started_years or advance.changed: + logger.info( + "Budget-window run %s is %s after report run %s; started years %s", + window_run_id, + advance.status.value, + report_run_id, + list(advance.started_years), + ) + return advance + + +def _log_failed_write(action: str, report_run_id: UUID, error: Exception) -> None: + # The error's text may hold connection details, so only its type is logged. + logger.warning( + "Could not %s report run %s (%s)", + action, + report_run_id, + type(error).__name__, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py index 2acf0a594..9b4bab134 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/report_run_access.py @@ -4,9 +4,11 @@ application attaches it to the coordinator, which records the report run its call names, to the simulation workers, each of which claims and records the simulation run its call names, and to the household functions, which -record the household job their call names. The secret must exist in a Modal -environment before the v2 application deploys there: ``modal deploy`` -checks its required key and fails without it. +record the household job their call names. The coordinator also starts a +budget window's waiting years, and it and the stuck-run sweep record the +window's outcome. The secret must exist in a Modal environment before the +v2 application deploys there: ``modal deploy`` checks its required key and +fails without it. """ from __future__ import annotations @@ -15,7 +17,7 @@ from functools import partial from typing import TYPE_CHECKING, Protocol -from policyengine_stage12_persistence import ReportRunStore +from policyengine_stage12_persistence import BudgetWindowStore, ReportRunStore from .database_retry import open_connection_with_retry @@ -173,3 +175,17 @@ def report_run_store(runtime: ObservabilityRuntime | None = None) -> ReportRunSt os.environ.get(MODAL_WORKER_DATABASE_URL_ENV, ""), open_connection=partial(open_connection_with_retry, runtime=runtime), ) + + +def budget_window_store( + runtime: ObservabilityRuntime | None = None, +) -> BudgetWindowStore: + """Build the budget-window store from ``MODAL_WORKER_DATABASE_URL``. + + Raises ``ValueError`` when the variable is unset or not a complete + ``postgresql://`` URL. + """ + return BudgetWindowStore( + os.environ.get(MODAL_WORKER_DATABASE_URL_ENV, ""), + open_connection=partial(open_connection_with_retry, runtime=runtime), + ) 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..26fc6dad1 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 @@ -26,6 +26,18 @@ call it starts on a pending job, and a household function claims the job under that call. A household job counts from its creation, while pending or running, for both ages. + +A budget window's years are report runs naming its budget-window run. A +year still pending with no call, in a window that hasn't finished, waits +for an earlier year's coordinator to start it, so the sweep leaves it to the +window. A year started later counts, until its claim, from when its call +was recorded rather than from its creation at submit. When the sweep fails +a year, it fails the window too, with the window's waiting years. After +judging the runs, it settles each unfinished window that nothing has +changed for the minimum age with no year under way: its last year ended, +or the simulation service created it, without anything recording what +follows. Such a window is finished from its years, or failed as stalled +with the years still waiting. """ from __future__ import annotations @@ -41,8 +53,9 @@ if TYPE_CHECKING: from collections.abc import Callable + from uuid import UUID - from policyengine_stage12_persistence import ActiveRun, RunWrite + from policyengine_stage12_persistence import ActiveRun, RunWrite, WindowAdvance logger = logging.getLogger(__name__) @@ -109,6 +122,24 @@ def fail_stuck_run( ... +class BudgetWindowSettler(Protocol): + """The sweep's budget-window writes; ``BudgetWindowStore`` implements them.""" + + def advance_window(self, window_run_id: UUID, *, now: datetime) -> WindowAdvance: + """Finish the window from its years' outcomes; start no year.""" + ... + + def list_idle_windows(self, *, idle_since_before: datetime) -> tuple[UUID, ...]: + """Return the unfinished windows with no year under way, unchanged since.""" + ... + + def settle_idle_window( + self, window_run_id: UUID, *, idle_since_before: datetime, now: datetime + ) -> WindowAdvance: + """Finish a window still idle: from its years, or as stalled.""" + ... + + # The order the sweep judges each kind in: report runs before simulation # runs, which are listed only once no unfinished report run names them, then # household jobs, which no run names. @@ -127,6 +158,8 @@ class SweepSummary: counts listed runs whose conditional failure changed nothing, because they were claimed, finished, moved to another call, had a call recorded, or a pending or running report run named them after they were listed. + ``windows_failed`` and ``windows_succeeded`` count the budget windows the + sweep finished, after failing one of their years or finding them idle. """ checked: int @@ -136,6 +169,8 @@ class SweepSummary: left: int = 0 failed_for_age: int = 0 failed_without_call: int = 0 + windows_failed: int = 0 + windows_succeeded: int = 0 def modal_call_state(modal_call_id: str, *, modal_module: Any = None) -> CallState: # noqa: ANN401 # the modal module, or a stand-in for it @@ -187,11 +222,83 @@ class _Tally: left: int = 0 failed_for_age: int = 0 failed_without_call: int = 0 + windows_failed: int = 0 + windows_succeeded: int = 0 + + def count_window(self, advance: WindowAdvance | None) -> None: + """Count a window the sweep finished; ``None`` means it was missing.""" + if advance is None or not advance.changed: + return + if advance.status is RunStatus.FAILED: + self.windows_failed += 1 + elif advance.status is RunStatus.SUCCEEDED: + self.windows_succeeded += 1 + + +@dataclass(frozen=True) +class _SweepPass: + """One sweep's shared state: its stores, its clock and its counts.""" + store: StuckRunStore + windows: BudgetWindowSettler | None + now: datetime + tally: _Tally -def sweep_stuck_runs( # the store, then keyword-only seams and limits + def fail(self, run: ActiveRun, error_message: str) -> bool: + """Fail ``run`` as listed, then its window if it is a window's year. + + Returns whether the run changed. + """ + changed = self.store.fail_stuck_run( + run, error_message=error_message, failed_at=self.now + ).changed + if changed and run.budget_window_run_id is not None and self.windows: + windows = self.windows + self.tally.count_window( + _unless_missing( + run.budget_window_run_id, + lambda window_run_id: windows.advance_window( + window_run_id, now=self.now + ), + ) + ) + return changed + + def settle_idle_windows(self, idle_since_before: datetime) -> None: + """Finish each window no year has advanced since ``idle_since_before``.""" + if self.windows is None: + return + windows = self.windows + for window_run_id in windows.list_idle_windows( + idle_since_before=idle_since_before + ): + self.tally.count_window( + _unless_missing( + window_run_id, + lambda idle_window_run_id: windows.settle_idle_window( + idle_window_run_id, + idle_since_before=idle_since_before, + now=self.now, + ), + ) + ) + + +def _unless_missing( + window_run_id: UUID, finish: Callable[[UUID], WindowAdvance] +) -> WindowAdvance | None: + """Finish the window, or log and return None when it no longer exists.""" + try: + return finish(window_run_id) + except LookupError: + logger.info("Budget-window run %s no longer exists", window_run_id) + return None + + +def sweep_stuck_runs( # noqa: PLR0913 # the stores, then keyword-only seams and limits store: StuckRunStore, *, + windows: BudgetWindowSettler | None = None, call_state: Callable[[str], CallState] = modal_call_state, now: datetime | None = None, minimum_age: timedelta = DEFAULT_MINIMUM_AGE, @@ -209,48 +316,46 @@ def sweep_stuck_runs( # the store, then keyword-only seams and limits or whose state is unknown, is left as it is, as is a running run with no call recorded, until the maximum age; so is one whose conditional failure changes nothing, which makes two sweeps at once safe. + + With ``windows``, a failed year of a budget window fails its window, and + the windows no year has advanced for ``minimum_age`` are then settled. """ now = now or datetime.now(UTC) states: dict[str, CallState] = {} - tally = _Tally() + sweep = _SweepPass(store=store, windows=windows, now=now, tally=_Tally()) + tally = sweep.tally for kind in SWEPT_KINDS: for run in store.list_active_runs(kind, active_since_before=now - maximum_age): - if store.fail_stuck_run( - run, error_message=EXCEEDED_AGE_MESSAGE, failed_at=now - ).changed: + if sweep.fail(run, EXCEEDED_AGE_MESSAGE): tally.failed_for_age += 1 else: tally.left += 1 for kind in SWEPT_KINDS: for run in store.list_active_runs(kind, active_since_before=now - minimum_age): - _sweep_run( - store, run, states=states, call_state=call_state, now=now, tally=tally - ) + _sweep_run(sweep, run, states=states, call_state=call_state) + sweep.settle_idle_windows(now - minimum_age) return SweepSummary(**vars(tally)) -def _sweep_run( # noqa: PLR0913 # one run, the store, and the sweep's shared state - store: StuckRunStore, +def _sweep_run( + sweep: _SweepPass, run: ActiveRun, *, states: dict[str, CallState], call_state: Callable[[str], CallState], - now: datetime, - tally: _Tally, ) -> None: """Fail ``run`` if its call has ended or was never started. Each call is asked about once. """ + tally = sweep.tally tally.checked += 1 modal_call_id = run.modal_call_id if modal_call_id is None: if run.status is not RunStatus.PENDING: # Nothing to ask Modal about; the maximum age decides it. tally.unknown += 1 - elif store.fail_stuck_run( - run, error_message=NOT_STARTED_MESSAGE, failed_at=now - ).changed: + elif sweep.fail(run, NOT_STARTED_MESSAGE): tally.failed_without_call += 1 else: tally.left += 1 @@ -266,9 +371,7 @@ def _sweep_run( # noqa: PLR0913 # one run, the store, and the sweep's shared s tally.running += 1 elif state is CallState.UNKNOWN: tally.unknown += 1 - elif store.fail_stuck_run( - run, error_message=ENDED_CALL_MESSAGE, failed_at=now - ).changed: + elif sweep.fail(run, ENDED_CALL_MESSAGE): tally.failed += 1 else: tally.left += 1 diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_capacity.py b/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_capacity.py new file mode 100644 index 000000000..a035518bc --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_capacity.py @@ -0,0 +1,60 @@ +"""The v2 functions have room for budget windows running at their cap. + +A budget window runs at most ``BUDGET_WINDOW_MAX_RUNNING_YEARS`` years at +once, each year one coordinator call that waits for up to two simulation +calls, and a US national simulation call waits for one segment call per +region group. A call queued behind a function's container limit counts +against its caller's wait, so each function has room for every call its +callers may make at once. +""" + +from __future__ import annotations + +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) +from policyengine_stage12_persistence import BUDGET_WINDOW_MAX_RUNNING_YEARS +from test_stage12_modal_app import _load + +# Budget windows the coordinator has room for at once, each at its cap, +# beside the coordinator calls it had for other reports. +WINDOWS_AT_ONCE = 3 +OTHER_REPORTS = 10 +SIMULATIONS_PER_REPORT = 2 + + +def test_the_coordinator_holds_three_windows_at_their_cap_and_other_reports( + monkeypatch, +) -> None: + functions = dict(_load(monkeypatch).app.function_calls) + + assert ( + functions["coordinate_report"]["max_containers"] + == WINDOWS_AT_ONCE * BUDGET_WINDOW_MAX_RUNNING_YEARS + OTHER_REPORTS + == 25 + ) + + +def test_each_simulation_worker_holds_both_simulations_of_every_coordinator( + monkeypatch, +) -> None: + functions = dict(_load(monkeypatch).app.function_calls) + coordinators = functions["coordinate_report"]["max_containers"] + + for worker in ("run_single_simulation_us", "run_single_simulation_uk"): + assert ( + functions[worker]["max_containers"] + == SIMULATIONS_PER_REPORT * coordinators + == 50 + ) + + +def test_every_us_simulation_at_once_has_room_for_its_segments(monkeypatch) -> None: + functions = dict(_load(monkeypatch).app.function_calls) + + assert ( + functions["run_single_simulation_segment_us"]["max_containers"] + == functions["run_single_simulation_us"]["max_containers"] + * len(US_REGION_GROUPS) + == 1000 + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_scheduling.py b/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_scheduling.py new file mode 100644 index 000000000..d4f912226 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_scheduling.py @@ -0,0 +1,469 @@ +"""A budget window's years, started and finished by their coordinators. + +When a year's coordinator claims its report run, the window is marked +running; when the year ends, the coordinator advances the window, which +starts the earliest waiting years in this coordinator's own Modal +application, or finishes the window. Neither step ever changes the year's +own outcome: a failure is logged. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from datetime import UTC, datetime +from types import SimpleNamespace +from uuid import uuid4 + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ReportRunIds, + year_calculation, +) +from policyengine_simulation_executor.stage12_runtime import ( + budget_window_scheduling, +) +from policyengine_simulation_executor.stage12_runtime.budget_window_scheduling import ( + ModalYearStarter, + advance_window_of_year, + coordinate_window_year, + start_window_of_year, +) +from policyengine_stage12_persistence import RunStatus, WindowAdvance, YearToStart +from stage12_fake_report_runs import report_run_context +from stage12_runtime_support import report_execution_input + +NOW = datetime(2026, 11, 12, 12, tzinfo=UTC) + + +class FakeWindowStore: + """Holds which window each year belongs to and records each write.""" + + def __init__(self, windows_of_years=None, *, failing=None): + self.windows_of_years = windows_of_years or {} + self.failing = failing + self.started = [] + self.advanced = [] + + def window_of_year(self, report_run_id): + if self.failing is not None: + raise self.failing + return self.windows_of_years.get(report_run_id) + + def start_window(self, report_run_id, *, started_at): + if self.failing is not None: + raise self.failing + self.started.append((report_run_id, started_at)) + return True + + def advance_window(self, window_run_id, *, now, start_year=None): + self.advanced.append((window_run_id, now, start_year)) + return WindowAdvance( + status=RunStatus.RUNNING, changed=False, started_years=(2031,) + ) + + +def test_the_year_starter_starts_this_application_s_coordinator_for_the_year( + monkeypatch, +) -> None: + calls = [] + + def from_name(*args, **kwargs): + calls.append(("from_name", args, kwargs)) + + def spawn(*spawn_args, **spawn_kwargs): + calls.append(("spawn", spawn_args, spawn_kwargs)) + return SimpleNamespace(object_id="fc-next-year") + + return SimpleNamespace(spawn=spawn) + + fake_modal = SimpleNamespace(Function=SimpleNamespace(from_name=from_name)) + + def import_modal(module_name): + assert module_name == "modal" + return fake_modal + + monkeypatch.setattr(budget_window_scheduling, "import_module", import_modal) + report, context = report_execution_input(), report_run_context() + next_year_runs = ReportRunIds( + report_run_id=uuid4(), + baseline_simulation_run_id=uuid4(), + reform_simulation_run_id=uuid4(), + ) + starter = ModalYearStarter( + payload=report.model_dump(mode="json"), + context=context.model_dump(mode="json"), + observability_context={"traceparent": "trace"}, + clock=lambda: NOW, + ) + + assert starter(YearToStart(year=2029, run_ids=next_year_runs)) == "fc-next-year" + expected_report, expected_context = year_calculation( + report, context, year=2029, run_ids=next_year_runs, created_at=NOW + ) + assert calls == [ + ( + "from_name", + (context.modal_application, "coordinate_report"), + {"environment_name": context.modal_environment}, + ), + ( + "spawn", + ( + expected_report.model_dump(mode="json"), + expected_context.model_dump(mode="json"), + ), + {"observability_context": {"traceparent": "trace"}}, + ), + ] + + +def test_an_ended_year_advances_its_window_with_the_starter() -> None: + year_run_id, window_run_id = uuid4(), uuid4() + windows = FakeWindowStore({year_run_id: window_run_id}) + + def start_year(report_run_id): + return f"fc-{report_run_id}" + + advance = advance_window_of_year( + windows, year_run_id, start_year=start_year, now=NOW + ) + + assert windows.advanced == [(window_run_id, NOW, start_year)] + assert advance == WindowAdvance( + status=RunStatus.RUNNING, changed=False, started_years=(2031,) + ) + + +def test_a_report_run_outside_any_window_advances_nothing() -> None: + windows = FakeWindowStore() + + assert ( + advance_window_of_year( + windows, uuid4(), start_year=lambda report_run_id: "fc", now=NOW + ) + is None + ) + assert windows.advanced == [] + + +def test_a_failed_advance_is_logged_and_never_raised(caplog) -> None: + windows = FakeWindowStore(failing=RuntimeError("database unavailable")) + year_run_id = uuid4() + + with caplog.at_level(logging.WARNING, logger=budget_window_scheduling.__name__): + advance = advance_window_of_year( + windows, year_run_id, start_year=lambda report_run_id: "fc", now=NOW + ) + + assert advance is None + assert str(year_run_id) in caplog.text + assert "RuntimeError" in caplog.text + # The error's own text may hold connection details; it is never logged. + assert "database unavailable" not in caplog.text + + +def test_a_claimed_year_marks_its_window_running() -> None: + year_run_id = uuid4() + windows = FakeWindowStore() + + start_window_of_year(windows, year_run_id, now=NOW) + + assert windows.started == [(year_run_id, NOW)] + + +def test_a_failed_window_start_is_logged_and_never_raised(caplog) -> None: + windows = FakeWindowStore(failing=RuntimeError("database unavailable")) + + with caplog.at_level(logging.WARNING, logger=budget_window_scheduling.__name__): + start_window_of_year(windows, uuid4(), now=NOW) + + assert "RuntimeError" in caplog.text + assert "database unavailable" not in caplog.text + + +@pytest.mark.parametrize("interruption", [KeyboardInterrupt, SystemExit]) +def test_an_interruption_while_advancing_propagates(interruption) -> None: + # Modal pre-empts or cancels the call: the restarted call advances again. + windows = FakeWindowStore(failing=interruption()) + + with pytest.raises(interruption): + advance_window_of_year( + windows, uuid4(), start_year=lambda report_run_id: "fc", now=NOW + ) + + +def test_a_malformed_input_fails_only_the_start() -> None: + # The coordinator refuses the same input; the start raises, so the year + # fails as refused instead of the starter failing to be built. + starter = ModalYearStarter(payload={"report": 1}, context={"context": 1}) + + with pytest.raises(ValueError, match="validation error"): + starter( + YearToStart( + year=2029, + run_ids=ReportRunIds( + report_run_id=uuid4(), + baseline_simulation_run_id=uuid4(), + reform_simulation_run_id=uuid4(), + ), + ) + ) + + +class RecordingWindows(FakeWindowStore): + """Records the order of the window steps and the coordinator call.""" + + def __init__(self, events, *, window_run_id=None): + super().__init__() + self.events = events + self.window_run_id = window_run_id or WINDOW_RUN_ID + + def start_window(self, report_run_id, *, started_at): + self.events.append(("start", report_run_id)) + return True + + def window_of_year(self, report_run_id): + self.events.append(("window_of_year", report_run_id)) + return self.window_run_id + + def advance_window(self, window_run_id, *, now, start_year=None): + self.events.append(("advance", window_run_id, start_year)) + return WindowAdvance(status=RunStatus.RUNNING, changed=False) + + +WINDOW_RUN_ID = uuid4() + + +def _year_call(events, coordinate, *, windows=None, context=None): + context = context or report_run_context().model_dump(mode="json") + return coordinate_window_year( + coordinate, + payload=report_execution_input().model_dump(mode="json"), + context=context, + windows=(lambda: windows or RecordingWindows(events)), + observability_context=lambda: {"traceparent": "trace"}, + ) + + +def test_a_year_marks_its_window_running_then_advances_it_when_it_ends() -> None: + events = [] + + def coordinate(): + events.append(("coordinate",)) + return {"status": "succeeded"} + + assert _year_call(events, coordinate) == {"status": "succeeded"} + + report_run_id = report_run_context().report_run.report_run_id + # The window is looked up once, before the coordinator runs. + assert [event[0] for event in events] == [ + "window_of_year", + "start", + "coordinate", + "advance", + ] + assert events[1] == ("start", report_run_id) + assert events[3][1] == WINDOW_RUN_ID + starter = events[3][2] + assert isinstance(starter, ModalYearStarter) + assert starter.payload == report_execution_input().model_dump(mode="json") + assert starter.context == report_run_context().model_dump(mode="json") + assert starter.observability_context == {"traceparent": "trace"} + + +@pytest.mark.parametrize( + "ending", + [RuntimeError("Stage 12 report coordination failed"), BaseException("cancelled")], + ids=["failure", "cancellation"], +) +def test_a_failed_or_cancelled_year_still_advances_its_window(ending) -> None: + events = [] + + def coordinate(): + raise ending + + with pytest.raises(type(ending)): + _year_call(events, coordinate) + + assert [event[0] for event in events] == ["window_of_year", "start", "advance"] + + +def test_a_preempted_year_leaves_its_window_to_the_restarted_call() -> None: + events = [] + + def coordinate(): + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + _year_call(events, coordinate) + + assert [event[0] for event in events] == ["window_of_year", "start"] + + +def test_a_call_naming_no_report_run_touches_no_window() -> None: + events = [] + comparison_context = ( + report_run_context() + .model_copy(update={"report_run": None}) + .model_dump(mode="json") + ) + + def windows(): + message = "a comparison call has no window" + raise AssertionError(message) + + result = coordinate_window_year( + lambda: {"status": "succeeded"}, + payload={}, + context=comparison_context, + windows=windows, + observability_context=lambda: None, + ) + + assert result == {"status": "succeeded"} + assert events == [] + + +def test_a_window_store_that_cannot_be_built_leaves_the_year_to_run(caplog) -> None: + def windows(): + msg = "MODAL_WORKER_DATABASE_URL is not a complete postgresql:// URL" + raise ValueError(msg) + + with caplog.at_level(logging.WARNING, logger=budget_window_scheduling.__name__): + result = coordinate_window_year( + lambda: {"status": "succeeded"}, + payload=report_execution_input().model_dump(mode="json"), + context=report_run_context().model_dump(mode="json"), + windows=windows, + observability_context=lambda: None, + ) + + assert result == {"status": "succeeded"} + assert "ValueError" in caplog.text + assert "postgresql://" not in caplog.text + + +def test_the_v2_coordinator_runs_each_call_as_a_window_year(monkeypatch) -> None: + from policyengine_simulation_executor import stage12_runtime + from test_stage12_modal_app import RecordingRuntime, _load + + module = _load(monkeypatch) + monkeypatch.setattr( + module.modal, "current_function_call_id", lambda: "fc-1", raising=False + ) + coordinator_trace = {"traceparent": "00-coordinator-trace-01"} + + class TracedRuntime(RecordingRuntime): + """Records whether an operation is open when the trace is captured. + + The real runtime captures the trace of the operation that is open. + """ + + operation_open = False + + @contextmanager + def operation(self, name, **kwargs): + with super().operation(name, **kwargs): + self.operation_open = True + try: + yield + finally: + self.operation_open = False + + def capture_context(self) -> dict: + return {**coordinator_trace, "operation_open": self.operation_open} + + monkeypatch.setattr( + module, "init_process_observability", lambda **_: TracedRuntime() + ) + received = [] + + def window_year(coordinate, *, payload, context, windows, observability_context): + result = coordinate() + # As the wrapper does, the next years' trace is captured once the + # year's coordinator returns, while its operation is still open. + received.append((payload, context, observability_context())) + return result + + monkeypatch.setattr(stage12_runtime, "coordinate_window_year", window_year) + monkeypatch.setattr( + stage12_runtime, + "coordinate_report", + lambda payload, context, parent, **_kwargs: {"status": "succeeded"}, + ) + + result = module.coordinate_report( + {"report": 1}, {"context": 1}, {"parent": 1}, observability_context=None + ) + + assert result == {"status": "succeeded"} + assert received == [ + ({"report": 1}, {"context": 1}, {**coordinator_trace, "operation_open": True}) + ] + + +class NoWindow(RecordingWindows): + """A report run that is no budget window's year.""" + + def window_of_year(self, report_run_id): + self.events.append(("window_of_year", report_run_id)) + + +def test_a_report_run_outside_any_window_makes_no_window_writes() -> None: + events = [] + + result = _year_call( + events, lambda: {"status": "succeeded"}, windows=NoWindow(events) + ) + + assert result == {"status": "succeeded"} + assert [event[0] for event in events] == ["window_of_year"] + + +class UnreadableWindow(RecordingWindows): + """A window lookup that fails before the coordinator runs.""" + + def __init__(self, events): + super().__init__(events) + self.lookups = 0 + + def window_of_year(self, report_run_id): + self.lookups += 1 + if self.lookups == 1: + message = "database unavailable" + raise RuntimeError(message) + return super().window_of_year(report_run_id) + + +def test_a_failed_lookup_still_advances_the_window_when_the_year_ends() -> None: + events = [] + + _year_call( + events, lambda: {"status": "succeeded"}, windows=UnreadableWindow(events) + ) + + # Nothing is marked before the coordinator runs; the window is looked up + # again, and advanced, once the year has ended. + assert [event[0] for event in events] == ["window_of_year", "advance"] + + +def test_an_unreadable_observability_context_still_advances_the_window() -> None: + events = [] + + def observability_context(): + message = "no trace" + raise RuntimeError(message) + + result = coordinate_window_year( + lambda: {"status": "succeeded"}, + payload=report_execution_input().model_dump(mode="json"), + context=report_run_context().model_dump(mode="json"), + windows=lambda: RecordingWindows(events), + observability_context=observability_context, + ) + + assert result == {"status": "succeeded"} + starter = events[-1][2] + assert events[-1][0] == "advance" + assert starter.observability_context is None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_year_start.py b/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_year_start.py new file mode 100644 index 000000000..b44ec1d25 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_budget_window_year_start.py @@ -0,0 +1,113 @@ +"""A year a coordinator starts runs as any report run does. + +The year starter spawns this application's ``coordinate_report`` with an +input and context derived from the starting year's. The call must bind to +the coordinator function's signature, and the coordinator must claim the +started year's report run and calculate it, as it does a report run the +simulation service starts. +""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace +from uuid import UUID + +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_simulation_executor.stage12_runtime import budget_window_scheduling +from policyengine_simulation_executor.stage12_runtime.budget_window_scheduling import ( + ModalYearStarter, +) +from policyengine_stage12_persistence import RunStatus, YearToStart +from stage12_fake_report_runs import ( + ComparisonStoreNotUsed, + FakeReportRuns, + FakeRun, + economy_aggregate, + report_run_context, + report_run_input, +) +from stage12_fake_stores import FakeArtifacts +from stage12_run_worker_support import RunWorkerInvoker +from stage12_runtime_support import coordinate_report +from test_stage12_modal_app import _load + +NEXT_YEAR_RUNS = ReportRunIds( + report_run_id=UUID("00000000-0000-0000-0000-0000000002a1"), + baseline_simulation_run_id=UUID("00000000-0000-0000-0000-0000000002b1"), + reform_simulation_run_id=UUID("00000000-0000-0000-0000-0000000002c1"), +) +NEXT_YEAR_CALL_ID = "fc-next-year" + + +def _started_call(monkeypatch) -> tuple[tuple, dict]: + """Start the next year from the shared report run; return the spawn's arguments.""" + spawned = [] + + def from_name(*_args, **_kwargs): + def spawn(*args, **kwargs): + spawned.append((args, kwargs)) + return SimpleNamespace(object_id=NEXT_YEAR_CALL_ID) + + return SimpleNamespace(spawn=spawn) + + fake_modal = SimpleNamespace(Function=SimpleNamespace(from_name=from_name)) + monkeypatch.setattr(budget_window_scheduling, "import_module", lambda _: fake_modal) + starter = ModalYearStarter( + payload=report_run_input().model_dump(mode="json"), + context=report_run_context().model_dump(mode="json"), + ) + starter(YearToStart(year=2027, run_ids=NEXT_YEAR_RUNS)) + (call,) = spawned + return call + + +def test_the_started_call_binds_to_the_coordinator_function(monkeypatch) -> None: + args, kwargs = _started_call(monkeypatch) + module = _load(monkeypatch) + + bound = inspect.signature(module.coordinate_report).bind(*args, **kwargs) + + # A report run's call passes no evaluation record. + assert "parent" not in bound.arguments + assert bound.arguments["context"]["report_run"] == NEXT_YEAR_RUNS.model_dump( + mode="json" + ) + + +def test_the_coordinator_claims_and_calculates_the_started_year(monkeypatch) -> None: + (payload, context), _ = _started_call(monkeypatch) + report_runs = FakeReportRuns() + report_runs.report_runs[NEXT_YEAR_RUNS.report_run_id] = FakeRun() + for simulation_run_id in ( + NEXT_YEAR_RUNS.baseline_simulation_run_id, + NEXT_YEAR_RUNS.reform_simulation_run_id, + ): + report_runs.simulation_runs[simulation_run_id] = FakeRun(year=2027) + report_runs.simulation_run_pairs[NEXT_YEAR_RUNS.report_run_id] = ( + NEXT_YEAR_RUNS.baseline_simulation_run_id, + NEXT_YEAR_RUNS.reform_simulation_run_id, + ) + artifacts = FakeArtifacts() + + result = coordinate_report( + payload, + context, + application_name=report_run_context().modal_application, + coordinator_invocation_id=NEXT_YEAR_CALL_ID, + store=ComparisonStoreNotUsed(), + artifacts=artifacts, + invoker=RunWorkerInvoker(report_runs, artifacts), + aggregator=economy_aggregate, + report_runs=report_runs, + ) + + assert result == { + "report_run_id": str(NEXT_YEAR_RUNS.report_run_id), + "status": "succeeded", + } + year_run = report_runs.report_runs[NEXT_YEAR_RUNS.report_run_id] + assert (year_run.status, year_run.modal_call_id) == ( + RunStatus.SUCCEEDED, + NEXT_YEAR_CALL_ID, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 112050860..f4be528b8 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -28,6 +28,10 @@ def __init__(self) -> None: def set_context(self, **attributes) -> None: self.contexts.append(attributes) + def capture_context(self) -> dict: + # The coordinator passes it on to start a budget window's next years. + return {} + @contextmanager def operation(self, name, **kwargs): self.operations.append({"name": name, **kwargs}) @@ -345,14 +349,14 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( functions["run_single_simulation_segment_us"]["image"] is module.us_worker_image ) assert functions["run_single_simulation_segment_us"]["timeout"] == 3600 - assert functions["run_single_simulation_segment_us"]["max_containers"] == 300 + assert functions["run_single_simulation_segment_us"]["max_containers"] == 1000 assert functions["run_single_simulation_us"]["timeout"] == 3900 assert functions["run_single_simulation_uk"]["timeout"] == 3000 - assert functions["run_single_simulation_us"]["max_containers"] == 10 - assert functions["run_single_simulation_uk"]["max_containers"] == 10 + assert functions["run_single_simulation_us"]["max_containers"] == 50 + assert functions["run_single_simulation_uk"]["max_containers"] == 50 assert functions["coordinate_report"]["image"] is module.coordinator_image assert functions["coordinate_report"]["timeout"] == 5400 - assert functions["coordinate_report"]["max_containers"] == 10 + assert functions["coordinate_report"]["max_containers"] == 25 assert "asgi_app" not in vars(module) 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..db32d645a 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 @@ -122,12 +122,15 @@ def test_the_sweep_image_installs_the_locked_group_and_the_runtime_sources( def test_the_scheduled_function_sweeps_with_the_report_run_store(monkeypatch) -> None: module = _load(monkeypatch) - store = object() + store, window_store = object(), object() swept = [] monkeypatch.setattr(report_run_access_module, "report_run_store", lambda: store) + monkeypatch.setattr( + report_run_access_module, "budget_window_store", lambda: window_store + ) - def recording_sweep(received_store): - swept.append(received_store) + def recording_sweep(received_store, *, windows): + swept.append((received_store, windows)) return run_sweep.SweepSummary( checked=3, failed=1, @@ -136,6 +139,8 @@ def recording_sweep(received_store): left=0, failed_for_age=2, failed_without_call=4, + windows_failed=5, + windows_succeeded=6, ) monkeypatch.setattr(run_sweep, "sweep_stuck_runs", recording_sweep) @@ -148,5 +153,7 @@ def recording_sweep(received_store): "left": 0, "failed_for_age": 2, "failed_without_call": 4, + "windows_failed": 5, + "windows_succeeded": 6, } - assert swept == [store] + assert swept == [(store, window_store)] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_windows.py b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_windows.py new file mode 100644 index 000000000..e97d33b42 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_run_sweep_windows.py @@ -0,0 +1,201 @@ +"""The stuck-run sweep finishes budget windows that their years no longer advance. + +A budget window's years are report runs naming its budget-window run. Each +year's coordinator starts the next waiting years and marks the window when +its year ends; when the sweep fails a year instead, it fails the window +too, and it settles a window that no year is advancing any more. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from uuid import uuid4 + +from policyengine_simulation_executor.stage12_runtime.run_sweep import ( + DEFAULT_MINIMUM_AGE, + CallState, + SweepSummary, + sweep_stuck_runs, +) +from policyengine_stage12_persistence import ( + ActiveRun, + RunKind, + RunStatus, + RunWrite, + WindowAdvance, +) + +NOW = datetime(2026, 11, 12, 12, tzinfo=UTC) + + +class FakeRunStore: + """Lists ``runs`` for the minimum age and nothing for the maximum age.""" + + def __init__(self, runs, *, unchanged=()): + self.runs = list(runs) + self.unchanged = set(unchanged) + self.failed = [] + + def list_active_runs(self, kind, *, active_since_before): + if active_since_before != NOW - DEFAULT_MINIMUM_AGE: + return () + failed = {run for run, _ in self.failed} + return tuple(run for run in self.runs if run.kind is kind and run not in failed) + + def fail_stuck_run(self, run, *, error_message, failed_at): + if run.run_id in self.unchanged: + return RunWrite(changed=False, status=RunStatus.SUCCEEDED) + self.failed.append((run, error_message)) + return RunWrite(changed=True, status=RunStatus.FAILED) + + +class FakeWindowStore: + """Records each window the sweep advances or settles.""" + + def __init__(self, *, idle=(), settled=None, missing=()): + self.idle = tuple(idle) + self.settled_outcomes = settled or {} + self.missing = set(missing) + self.advanced = [] + self.idle_cutoffs = [] + self.settled = [] + + def advance_window(self, window_run_id, *, now, start_year=None): + self.advanced.append((window_run_id, now, start_year)) + if window_run_id in self.missing: + msg = f"budget-window run {window_run_id} does not exist" + raise LookupError(msg) + return WindowAdvance(status=RunStatus.FAILED, changed=True) + + def list_idle_windows(self, *, idle_since_before): + self.idle_cutoffs.append(idle_since_before) + return self.idle + + def settle_idle_window(self, window_run_id, *, idle_since_before, now): + self.settled.append((window_run_id, idle_since_before, now)) + if window_run_id in self.missing: + msg = f"budget-window run {window_run_id} does not exist" + raise LookupError(msg) + return self.settled_outcomes.get( + window_run_id, WindowAdvance(status=RunStatus.RUNNING, changed=False) + ) + + +def _year_run(call_id, window_run_id, status=RunStatus.RUNNING): + return ActiveRun( + kind=RunKind.REPORT_RUN, + run_id=uuid4(), + status=status, + modal_call_id=call_id, + budget_window_run_id=window_run_id, + ) + + +def test_failing_a_year_fails_its_window_and_starts_nothing() -> None: + window_run_id = uuid4() + ended_year = _year_run("fc-ended", window_run_id) + other_report = ActiveRun( + kind=RunKind.REPORT_RUN, + run_id=uuid4(), + status=RunStatus.RUNNING, + modal_call_id="fc-other", + ) + runs = FakeRunStore([ended_year, other_report]) + windows = FakeWindowStore() + + summary = sweep_stuck_runs( + runs, + windows=windows, + call_state=lambda call_id: CallState.ENDED, + now=NOW, + ) + + assert windows.advanced == [(window_run_id, NOW, None)] + assert summary.failed == 2 + assert summary.windows_failed == 1 + + +def test_a_year_the_sweep_did_not_change_leaves_its_window_alone() -> None: + year = _year_run("fc-ended", uuid4()) + runs = FakeRunStore([year], unchanged={year.run_id}) + windows = FakeWindowStore() + + summary = sweep_stuck_runs( + runs, + windows=windows, + call_state=lambda call_id: CallState.ENDED, + now=NOW, + ) + + assert windows.advanced == [] + assert (summary.left, summary.windows_failed) == (1, 0) + + +def test_idle_windows_are_settled_after_the_runs() -> None: + stalled, finished, busy = uuid4(), uuid4(), uuid4() + windows = FakeWindowStore( + idle=(stalled, finished, busy), + settled={ + stalled: WindowAdvance(status=RunStatus.FAILED, changed=True), + finished: WindowAdvance(status=RunStatus.SUCCEEDED, changed=True), + }, + ) + + summary = sweep_stuck_runs( + FakeRunStore([]), windows=windows, call_state=lambda _: CallState.ENDED, now=NOW + ) + + cutoff = NOW - DEFAULT_MINIMUM_AGE + assert windows.idle_cutoffs == [cutoff] + assert windows.settled == [ + (stalled, cutoff, NOW), + (finished, cutoff, NOW), + (busy, cutoff, NOW), + ] + assert (summary.windows_failed, summary.windows_succeeded) == (1, 1) + + +def test_a_window_deleted_meanwhile_is_skipped() -> None: + gone = uuid4() + year = _year_run("fc-ended", gone) + windows = FakeWindowStore(idle=(gone,), missing={gone}) + + summary = sweep_stuck_runs( + FakeRunStore([year]), + windows=windows, + call_state=lambda _: CallState.ENDED, + now=NOW, + ) + + assert summary.failed == 1 + assert (summary.windows_failed, summary.windows_succeeded) == (0, 0) + + +def test_without_a_window_store_the_sweep_judges_runs_only() -> None: + year = _year_run(None, uuid4(), status=RunStatus.PENDING) + + summary = sweep_stuck_runs( + FakeRunStore([year]), call_state=lambda _: CallState.ENDED, now=NOW + ) + + assert summary == SweepSummary( + checked=1, + failed=0, + running=0, + unknown=0, + failed_without_call=1, + ) + + +def test_the_window_cutoff_is_the_minimum_age() -> None: + windows = FakeWindowStore() + + sweep_stuck_runs( + FakeRunStore([]), + windows=windows, + call_state=lambda _: CallState.ENDED, + now=NOW, + minimum_age=timedelta(minutes=25), + ) + + assert windows.idle_cutoffs == [NOW - timedelta(minutes=25)] diff --git a/rules/fixtures/sim.persistence.missing-privileges-reported/bad_grants.txt b/rules/fixtures/sim.persistence.missing-privileges-reported/bad_grants.txt index 3d1b95edd..ef31953e1 100644 --- a/rules/fixtures/sim.persistence.missing-privileges-reported/bad_grants.txt +++ b/rules/fixtures/sim.persistence.missing-privileges-reported/bad_grants.txt @@ -30,6 +30,10 @@ UPDATE report_runs completed_at UPDATE report_runs error_message UPDATE report_runs status UPDATE report_runs updated_at +UPDATE budget_window_runs completed_at +UPDATE budget_window_runs error_message +UPDATE budget_window_runs status +UPDATE budget_window_runs updated_at UPDATE household_jobs completed_at UPDATE household_jobs error_message UPDATE household_jobs modal_call_id diff --git a/rules/fixtures/sim.persistence.missing-privileges-reported/good_grants.txt b/rules/fixtures/sim.persistence.missing-privileges-reported/good_grants.txt index 38db81875..bba121dfb 100644 --- a/rules/fixtures/sim.persistence.missing-privileges-reported/good_grants.txt +++ b/rules/fixtures/sim.persistence.missing-privileges-reported/good_grants.txt @@ -30,6 +30,10 @@ UPDATE report_runs error_message UPDATE report_runs modal_call_id UPDATE report_runs status UPDATE report_runs updated_at +UPDATE budget_window_runs completed_at +UPDATE budget_window_runs error_message +UPDATE budget_window_runs status +UPDATE budget_window_runs updated_at UPDATE household_jobs completed_at UPDATE household_jobs error_message UPDATE household_jobs modal_call_id