Calculate budget windows on the detailed-save path, five years at a time

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.

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.

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'

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
 

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

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/<environment>/<YYYY>/<MM>/<run_id>``. 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),
+    }