From eac7dd28fd862be439e61f093e2a07ac56a3d1ab Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:01:27 +0400 Subject: [PATCH 01/20] Record budget-window reports and one report run per year of each window A report gains an optional budget window size: an economy report with one compares its two simulations in each year of the window. Each calculation of such a report is a budget-window run (new table simulation.budget_window_runs) on one catalog version, which records its trigger, whether its years calculate their simulations again, how many years it calculates at once, its status, when the safety check next looks at it, and the lease of the caller starting its next years. Each year is a report run that names its budget-window run, of the same report and catalog version, with at most one per year. BudgetWindowRunStore creates, claims, releases and finishes runs, and reads each year's runs and budget rows; nothing stores a total. Revision 0005 adds the column, the table and the link; its downgrade removes them. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 20 + .../budget_window_runs.py | 222 +++++++ .../budget_window_store.py | 353 +++++++++++ .../reports.py | 38 ++ .../run_store.py | 30 + .../runs.py | 22 +- .../schema.py | 2 +- .../tests/budget_window_support.py | 117 ++++ .../tests/test_budget_window_runs.py | 551 ++++++++++++++++++ .../tests/test_schema.py | 1 + .../versions/0005_budget_window_runs.py | 193 ++++++ 11 files changed, 1547 insertions(+), 2 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py create mode 100644 libs/policyengine-simulation-persistence/tests/budget_window_support.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py create mode 100644 migrations/versions/0005_budget_window_runs.py diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 1b49c5d6a..c18029fd6 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -4,6 +4,17 @@ the migration environment compares with the database. """ +from policyengine_simulation_persistence.budget_window_runs import ( + BudgetWindowRunRecord, + BudgetWindowStatus, + BudgetWindowYear, + NewBudgetWindowRun, + StartedBudgetWindowRun, + YearBudget, +) +from policyengine_simulation_persistence.budget_window_store import ( + BudgetWindowRunStore, +) from policyengine_simulation_persistence.catalog_reads import ( CATALOG_READ_TABLES, CatalogError, @@ -41,6 +52,7 @@ record_output_files, ) from policyengine_simulation_persistence.reports import ( + MAX_BUDGET_WINDOW_YEARS, NewReport, ReportKind, ReportRecord, @@ -85,6 +97,7 @@ "ALEMBIC_VERSION_TABLE", "CATALOG_READ_TABLES", "HEAD_REVISION", + "MAX_BUDGET_WINDOW_YEARS", "OPTIONAL_SECTIONS", "OUTPUT_FILE_RETENTION", "RESULT_FIELD_STORAGE", @@ -92,6 +105,10 @@ "SIMULATION_SCHEMA", "SUCCEEDED_RUN_REUSE_WINDOW", "AbandonedDispatches", + "BudgetWindowRunRecord", + "BudgetWindowRunStore", + "BudgetWindowStatus", + "BudgetWindowYear", "CatalogError", "CatalogReader", "CatalogVersion", @@ -100,6 +117,7 @@ "EconomyResultLabels", "GeographyLabel", "HouseholdOutput", + "NewBudgetWindowRun", "NewReport", "NewReportRun", "NewSimulation", @@ -123,8 +141,10 @@ "SimulationRecord", "SimulationRunRecord", "SimulationStore", + "StartedBudgetWindowRun", "StartedReportRun", "StartedSimulationRun", + "YearBudget", "available_output_files", "check_storable", "create_simulation_engine", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py new file mode 100644 index 000000000..c8f3f985a --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py @@ -0,0 +1,222 @@ +"""``budget_window_runs``: calculations of budget-window reports. + +A budget-window report (``reports.budget_window_size``) compares its two +simulations in each year of its window. One budget-window run is one +calculation of it on a catalog version: it records what started it (its +trigger), whether its years calculate their simulations again (``force``), how +many years it calculates at once, its status and error, and when the safety +check next looks at it. Each year it calculates is one of its report runs +(``report_runs.budget_window_run_id``), at most one per year, of its report and +catalog version. The run stores no figures of its own: its totals are the sums +of its years' budget rows, added up when it is read. + +A run is ``running`` until every year of its window has succeeded, when it is +``succeeded``, or until one year has failed, when it is ``failed`` with that +year's error. ``scheduling_claimed_until`` is the lease of the one caller +starting its next years at a time. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from enum import StrEnum +from typing import TYPE_CHECKING, Self + +from sqlalchemy import ( + Boolean, + CheckConstraint, + Column, + DateTime, + ForeignKeyConstraint, + Index, + Integer, + Table, + Text, + UniqueConstraint, + Uuid, +) + +from policyengine_simulation_persistence import public_api +from policyengine_simulation_persistence.reports import reports +from policyengine_simulation_persistence.runs import ( + ReportRunRecord, + RunTrigger, + SimulationRunRecord, +) +from policyengine_simulation_persistence.schema import simulation_metadata + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from sqlalchemy import RowMapping + + +class BudgetWindowStatus(StrEnum): + """Where a budget-window run stands.""" + + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +FINISHED_BUDGET_WINDOW_STATUSES = ( + BudgetWindowStatus.SUCCEEDED, + BudgetWindowStatus.FAILED, +) + + +def _values(enum: type[StrEnum]) -> str: + return ", ".join(f"'{member.value}'" for member in enum) + + +budget_window_runs = Table( + "budget_window_runs", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("report_id", Uuid(as_uuid=True), nullable=False), + Column("tax_benefit_model_version_id", Uuid(as_uuid=True), nullable=False), + Column("trigger", Text, nullable=False), + Column("force", Boolean, nullable=False), + Column("max_concurrent_years", Integer, nullable=False), + Column("status", Text, nullable=False), + Column("error_code", Text, nullable=True), + Column("error_message", Text, nullable=True), + Column("next_check_at", DateTime(timezone=True), nullable=True), + Column("scheduling_claimed_until", DateTime(timezone=True), nullable=True), + Column("completed_at", DateTime(timezone=True), nullable=True), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"), + CheckConstraint(f"status IN ({_values(BudgetWindowStatus)})", name="status"), + CheckConstraint("max_concurrent_years >= 1", name="max_concurrent_years"), + # Only a running run is checked, and every running run is due a check. + CheckConstraint( + "(status = 'running') = (next_check_at IS NOT NULL)", name="next_check_at" + ), + CheckConstraint( + "(status IN ('succeeded', 'failed')) = (completed_at IS NOT NULL)", + name="completed_at", + ), + CheckConstraint("(status = 'failed') = (error_code IS NOT NULL)", name="error"), + ForeignKeyConstraint( + ["report_id"], + [reports.c.id], + name="fk_budget_window_runs_report", + ondelete="RESTRICT", + ), + ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + [public_api.tax_benefit_model_versions.c.id], + name="fk_budget_window_runs_catalog_version", + ondelete="RESTRICT", + ), + # Each year's report run refers to the run with its report and catalog + # version, so it is a run of the same report on the same version. + UniqueConstraint( + "id", + "report_id", + "tax_benefit_model_version_id", + name="uq_budget_window_runs_id_report_catalog_version", + ), +) +# A report's newest budget-window run is its current one. +Index( + "ix_budget_window_runs_report_id_created_at", + budget_window_runs.c.report_id, + budget_window_runs.c.created_at, +) +# The safety check reads only running runs that are due a check. +Index( + "ix_budget_window_runs_due_checks", + budget_window_runs.c.next_check_at, + postgresql_where=budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value, +) +# A Public API delete of a catalog version checks this for references. +Index( + "ix_budget_window_runs_tax_benefit_model_version_id", + budget_window_runs.c.tax_benefit_model_version_id, +) + + +@dataclass(frozen=True) +class NewBudgetWindowRun: + """A calculation of a budget-window report to create. + + ``force`` calculates every year's simulations again instead of reusing + runs that exist; ``max_concurrent_years`` caps how many of its years are + calculated at once. + """ + + report_id: UUID + tax_benefit_model_version_id: UUID + trigger: RunTrigger + max_concurrent_years: int + force: bool = False + + def __post_init__(self) -> None: + """Reject a cap below one year without touching the database.""" + if self.max_concurrent_years < 1: + msg = "max_concurrent_years must be at least 1" + raise ValueError(msg) + + +@dataclass(frozen=True) +class BudgetWindowRunRecord: + """One ``simulation.budget_window_runs`` row.""" + + id: UUID + report_id: UUID + tax_benefit_model_version_id: UUID + trigger: RunTrigger + force: bool + max_concurrent_years: int + status: BudgetWindowStatus + error_code: str | None + error_message: str | None + next_check_at: datetime | None + scheduling_claimed_until: datetime | None + completed_at: datetime | None + created_at: datetime + updated_at: datetime + + @property + def finished(self) -> bool: + """Whether the run has succeeded or failed.""" + return self.status in FINISHED_BUDGET_WINDOW_STATUSES + + @classmethod + def from_row(cls, row: RowMapping) -> Self: + """Build the record from a selected ``budget_window_runs`` row.""" + values = {field.name: row[field.name] for field in fields(cls)} + values["trigger"] = RunTrigger(values["trigger"]) + values["status"] = BudgetWindowStatus(values["status"]) + return cls(**values) + + +@dataclass(frozen=True) +class StartedBudgetWindowRun: + """The run a first start resolved to, and whether this call created it.""" + + run: BudgetWindowRunRecord + created: bool + + +@dataclass(frozen=True) +class BudgetWindowYear: + """One year of a budget-window run: its report run and both simulation runs.""" + + report_run: ReportRunRecord + baseline_run: SimulationRunRecord + reform_run: SimulationRunRecord + + +@dataclass(frozen=True) +class YearBudget: + """A year's budget figures, from its comparison's ``budget_impacts`` row.""" + + year: int + budgetary_impact: float + tax_revenue_impact: float + state_tax_revenue_impact: float + benefit_spending_impact: float diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py new file mode 100644 index 000000000..318d45fc1 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py @@ -0,0 +1,353 @@ +"""Create, read, schedule and finish budget-window runs, and read their years. + +A budget-window run starts ``running``. Whoever starts its next years first +claims it (``claim_scheduling``): a claim is a lease, so a caller that stopped +loses it once the lease ends, and another caller can carry on. Releasing the +claim sets when the safety check next looks at the run. The run is finished +once, with a conditional write, as ``succeeded`` or ``failed``. + +A run's years are its report runs, read with both simulation runs of each. +Its budget figures are each year's ``budget_impacts`` row, read in year order, +so a reader sums them; nothing stores a total. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from uuid import uuid4 + +from sqlalchemy import func, or_, select, update +from sqlalchemy.dialects.postgresql import insert as postgres_insert + +from policyengine_simulation_persistence.budget_window_runs import ( + FINISHED_BUDGET_WINDOW_STATUSES, + BudgetWindowRunRecord, + BudgetWindowStatus, + BudgetWindowYear, + StartedBudgetWindowRun, + YearBudget, + budget_window_runs, +) +from policyengine_simulation_persistence.reports import reports +from policyengine_simulation_persistence.result_tables import budget_impacts +from policyengine_simulation_persistence.run_store import utc_now +from policyengine_simulation_persistence.runs import ( + ReportRunRecord, + SimulationRunRecord, + bounded_error_message, + report_runs, + simulation_runs, +) + +if TYPE_CHECKING: + from collections.abc import Callable + from datetime import datetime, timedelta + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, Engine + + from policyengine_simulation_persistence.budget_window_runs import ( + NewBudgetWindowRun, + ) + from policyengine_simulation_persistence.run_transitions import RunError + +# The first of the two keys of the transaction-level advisory locks that +# serialise a report's first budget-window runs; the second is the hash of +# the report ID. ``RunStore`` uses 2_020_101 for first report runs. +FIRST_BUDGET_WINDOW_RUN_LOCK_SPACE = 2_030_101 + + +class BudgetWindowRunStore: + """Typed SQLAlchemy Core reads and writes for budget-window runs.""" + + def __init__( + self, + engine: Engine, + *, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Use ``engine`` (the runtime role) and ``clock`` for every timestamp.""" + self.engine = engine + self._clock = clock + + # --- Creating runs ----------------------------------------------------- + + def create_run( + self, new_run: NewBudgetWindowRun, *, first_check_at: datetime + ) -> BudgetWindowRunRecord: + """Add a running budget-window run, due its first check at a time. + + Raise LookupError when the report does not exist, and ValueError + when it is not a budget-window report. + """ + with self.engine.begin() as connection: + return self._insert(connection, new_run, first_check_at) + + def create_first_run( + self, new_run: NewBudgetWindowRun, *, first_check_at: datetime + ) -> StartedBudgetWindowRun: + """Add the report's first budget-window run, or return its newest one. + + Callers starting a report's first run at the same time wait for each + other on a transaction-level advisory lock on the report, so only the + first adds a run. Raise as ``create_run`` does. + """ + with self.engine.begin() as connection: + connection.execute( + select( + func.pg_advisory_xact_lock( + FIRST_BUDGET_WINDOW_RUN_LOCK_SPACE, + func.hashtext(str(new_run.report_id)), + ) + ) + ) + newest = self._newest( + connection, budget_window_runs.c.report_id == new_run.report_id + ) + if newest is not None: + return StartedBudgetWindowRun(run=newest, created=False) + return StartedBudgetWindowRun( + run=self._insert(connection, new_run, first_check_at), created=True + ) + + def _insert( + self, + connection: Connection, + new_run: NewBudgetWindowRun, + first_check_at: datetime, + ) -> BudgetWindowRunRecord: + window_size = connection.execute( + select(reports.c.id, reports.c.budget_window_size).where( + reports.c.id == new_run.report_id + ) + ).first() + if window_size is None: + msg = f"no report {new_run.report_id}" + raise LookupError(msg) + if window_size.budget_window_size is None: + msg = f"report {new_run.report_id} is not a budget-window report" + raise ValueError(msg) + now = self._clock() + row = ( + connection.execute( + postgres_insert(budget_window_runs) + .values( + id=uuid4(), + report_id=new_run.report_id, + tax_benefit_model_version_id=new_run.tax_benefit_model_version_id, + trigger=new_run.trigger.value, + force=new_run.force, + max_concurrent_years=new_run.max_concurrent_years, + status=BudgetWindowStatus.RUNNING.value, + next_check_at=first_check_at, + created_at=now, + updated_at=now, + ) + .returning(*budget_window_runs.c) + ) + .mappings() + .one() + ) + return BudgetWindowRunRecord.from_row(row) + + # --- Reading runs ------------------------------------------------------ + + def get_run(self, run_id: UUID) -> BudgetWindowRunRecord | None: + """Return the budget-window run with this ID, or None.""" + with self.engine.connect() as connection: + return self._newest(connection, budget_window_runs.c.id == run_id) + + def latest_run(self, report_id: UUID) -> BudgetWindowRunRecord | None: + """Return the report's newest budget-window run: the report's status.""" + with self.engine.connect() as connection: + return self._newest(connection, budget_window_runs.c.report_id == report_id) + + def runs(self, report_id: UUID) -> list[BudgetWindowRunRecord]: + """Return the report's budget-window runs, newest first.""" + with self.engine.connect() as connection: + rows = connection.execute( + select(budget_window_runs) + .where(budget_window_runs.c.report_id == report_id) + .order_by( + budget_window_runs.c.created_at.desc(), + budget_window_runs.c.id.desc(), + ) + ).mappings() + return [BudgetWindowRunRecord.from_row(row) for row in rows] + + @staticmethod + def _newest( + connection: Connection, condition: ColumnElement[bool] + ) -> BudgetWindowRunRecord | None: + row = ( + connection.execute( + select(budget_window_runs) + .where(condition) + .order_by( + budget_window_runs.c.created_at.desc(), + budget_window_runs.c.id.desc(), + ) + .limit(1) + ) + .mappings() + .first() + ) + return None if row is None else BudgetWindowRunRecord.from_row(row) + + def due_runs(self, *, now: datetime, limit: int) -> list[BudgetWindowRunRecord]: + """Return running runs whose next check is due, the longest-due first.""" + with self.engine.connect() as connection: + rows = connection.execute( + select(budget_window_runs) + .where( + budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value, + budget_window_runs.c.next_check_at <= now, + ) + .order_by( + budget_window_runs.c.next_check_at, + budget_window_runs.c.created_at, + budget_window_runs.c.id, + ) + .limit(limit) + ).mappings() + return [BudgetWindowRunRecord.from_row(row) for row in rows] + + # --- A run's years ----------------------------------------------------- + + def years(self, run_id: UUID) -> list[BudgetWindowYear]: + """Return the years the run has started, in year order, with their runs.""" + with self.engine.connect() as connection: + year_runs = [ + ReportRunRecord.from_row(row) + for row in connection.execute( + select(report_runs) + .where(report_runs.c.budget_window_run_id == run_id) + .order_by(report_runs.c.year) + ).mappings() + ] + simulation_run_ids = { + run_id + for year_run in year_runs + for run_id in ( + year_run.baseline_simulation_run_id, + year_run.reform_simulation_run_id, + ) + if run_id is not None + } + by_id = { + row["id"]: SimulationRunRecord.from_row(row) + for row in connection.execute( + select(simulation_runs).where( + simulation_runs.c.id.in_(simulation_run_ids) + ) + ).mappings() + } + return [ + BudgetWindowYear( + report_run=year_run, + baseline_run=by_id[year_run.baseline_simulation_run_id], + reform_run=by_id[ + year_run.reform_simulation_run_id + or year_run.baseline_simulation_run_id + ], + ) + for year_run in year_runs + ] + + def year_budgets(self, run_id: UUID) -> list[YearBudget]: + """Return the budget figures of each year whose comparison stored them.""" + with self.engine.connect() as connection: + rows = connection.execute( + select( + report_runs.c.year, + budget_impacts.c.budgetary_impact, + budget_impacts.c.tax_revenue_impact, + budget_impacts.c.state_tax_revenue_impact, + budget_impacts.c.benefit_spending_impact, + ) + .join( + budget_impacts, budget_impacts.c.report_run_id == report_runs.c.id + ) + .where(report_runs.c.budget_window_run_id == run_id) + .order_by(report_runs.c.year) + ).mappings() + return [YearBudget(**row) for row in rows] + + # --- Scheduling and finishing ------------------------------------------ + + def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: + """Claim a running run to start its next years; return whether this did. + + A claim held by another caller is refused until its lease ends. + """ + now = self._clock() + with self.engine.begin() as connection: + claimed = connection.execute( + update(budget_window_runs) + .where( + budget_window_runs.c.id == run_id, + budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value, + or_( + budget_window_runs.c.scheduling_claimed_until.is_(None), + budget_window_runs.c.scheduling_claimed_until <= now, + ), + ) + .values(scheduling_claimed_until=now + lease, updated_at=now) + ) + return claimed.rowcount == 1 + + def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: + """Release a running run's claim and set when it is next checked.""" + now = self._clock() + with self.engine.begin() as connection: + connection.execute( + update(budget_window_runs) + .where( + budget_window_runs.c.id == run_id, + budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value, + ) + .values( + scheduling_claimed_until=None, + next_check_at=next_check_at, + updated_at=now, + ) + ) + + def finish_run( + self, + run_id: UUID, + status: BudgetWindowStatus, + error: RunError | None = None, + ) -> bool: + """Mark a running run succeeded, or failed with its error; return if it was. + + Raise ValueError for a status that does not finish a run, a failure + without its error, or a success with one. + """ + if status not in FINISHED_BUDGET_WINDOW_STATUSES: + msg = "a budget-window run is finished as succeeded or failed" + raise ValueError(msg) + if (status is BudgetWindowStatus.FAILED) != (error is not None): + msg = "a failed budget-window run, and only a failed one, has an error" + raise ValueError(msg) + now = self._clock() + with self.engine.begin() as connection: + finished = connection.execute( + update(budget_window_runs) + .where( + budget_window_runs.c.id == run_id, + budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value, + ) + .values( + status=status.value, + error_code=None if error is None else error.code.value, + error_message=( + None if error is None else bounded_error_message(error.message) + ), + completed_at=now, + next_check_at=None, + scheduling_claimed_until=None, + updated_at=now, + ) + ) + return finished.rowcount == 1 diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py index ec911e10f..48bd9c6a2 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/reports.py @@ -9,6 +9,11 @@ A report's calculations are its report runs (item 1.06); a report with none has not been calculated yet. A report migrated from the Public API's v1 database keeps its v1 report ID in ``legacy_report_id``. + +A budget-window report is an economy report with a ``budget_window_size``: it +compares its two simulations in each of that many years, starting with its +``year``. Each of its calculations is a budget-window run +(``budget_window_runs``), which holds one report run per year. """ from __future__ import annotations @@ -54,6 +59,8 @@ FIRST_YEAR = 1900 LAST_YEAR = 2200 +# The most years one budget-window report compares. +MAX_BUDGET_WINDOW_YEARS = 75 class ReportKind(StrEnum): @@ -83,6 +90,7 @@ class ReportKind(StrEnum): nullable=False, server_default=func.now(), ), + Column("budget_window_size", Integer, nullable=True), CheckConstraint(f"kind IN ({_KIND_VALUES})", name="kind"), CheckConstraint("country ~ '^[a-z]{2}$'", name="country"), CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), @@ -95,6 +103,14 @@ class ReportKind(StrEnum): "legacy_report_id IS NULL OR legacy_report_id > 0", name="legacy_report_id", ), + # Only an economy report has a budget window, and its last year is a year + # a run may have. + CheckConstraint( + "budget_window_size IS NULL OR (kind = 'economy' " + f"AND budget_window_size BETWEEN 1 AND {MAX_BUDGET_WINDOW_YEARS} " + f"AND year + budget_window_size - 1 <= {LAST_YEAR})", + name="budget_window_size", + ), ForeignKeyConstraint( ["baseline_simulation_id", "country", "kind"], _SIMULATION_KEY, @@ -131,6 +147,7 @@ class NewReport: reform_simulation_id: UUID | None = None options: dict[str, Any] = dataclass_field(default_factory=dict) legacy_report_id: int | None = None + budget_window_size: int | None = None def __post_init__(self) -> None: """Reject malformed fields without touching the database.""" @@ -144,6 +161,21 @@ def __post_init__(self) -> None: msg = "options must be a JSON object" raise TypeError(msg) check_legacy_id("legacy_report_id", self.legacy_report_id) + self._check_budget_window() + + def _check_budget_window(self) -> None: + size = self.budget_window_size + if size is None: + return + if self.kind is not ReportKind.ECONOMY: + msg = "only an economy report has a budget window" + raise ValueError(msg) + if not 1 <= size <= MAX_BUDGET_WINDOW_YEARS: + msg = f"budget_window_size must be between 1 and {MAX_BUDGET_WINDOW_YEARS}" + raise ValueError(msg) + if self.year + size - 1 > LAST_YEAR: + msg = f"a budget window must end in {LAST_YEAR} or earlier" + raise ValueError(msg) @dataclass(frozen=True) @@ -159,6 +191,12 @@ class ReportRecord: options: dict[str, Any] legacy_report_id: int | None created_at: datetime + budget_window_size: int | None = None + + @property + def window_years(self) -> range: + """Return the years the report compares: its year, or its budget window.""" + return range(self.year, self.year + (self.budget_window_size or 1)) @classmethod def from_row(cls, row: RowMapping) -> ReportRecord: diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py index 5180d2b38..b1f7d175a 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py @@ -107,6 +107,32 @@ def _check_report_run_simulations( return baseline +def _check_budget_window_year( + report: RowMapping, year: int, new_report_run: NewReportRun +) -> None: + """Require a budget-window report's run to be one year of its window. + + A report run of any other report belongs to no budget-window run. The + database checks that the budget-window run is one of the report's, on + the report run's catalog version. + """ + window_size = report["budget_window_size"] + if window_size is None: + if new_report_run.budget_window_run_id is not None: + msg = "only a budget-window report's report runs are years of a window" + raise ValueError(msg) + return + if new_report_run.budget_window_run_id is None: + msg = "a budget-window report's report run is a year of a budget-window run" + raise ValueError(msg) + if not report["year"] <= year < report["year"] + window_size: + msg = ( + f"{year} is not a year of the report's window, which starts in " + f"{report['year']} and has {window_size} years" + ) + raise ValueError(msg) + + class RunStore: """Typed SQLAlchemy Core reads and inserts for runs.""" @@ -328,6 +354,8 @@ def _insert_report_run( reports.c.kind, reports.c.baseline_simulation_id, reports.c.reform_simulation_id, + reports.c.year, + reports.c.budget_window_size, ).where(reports.c.id == new_report_run.report_id) ) .mappings() @@ -345,6 +373,7 @@ def _insert_report_run( baseline = _check_report_run_simulations( report, baseline, reform, new_report_run ) + _check_budget_window_year(report, baseline["year"], new_report_run) economy = report["kind"] == ReportKind.ECONOMY.value row = ( connection.execute( @@ -361,6 +390,7 @@ def _insert_report_run( new_report_run.baseline_simulation_run_id ), reform_simulation_run_id=new_report_run.reform_simulation_run_id, + budget_window_run_id=new_report_run.budget_window_run_id, trigger=new_report_run.trigger.value, comparison_status=(RunStatus.PENDING.value if economy else None), created_at=now, diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py index 843277f85..4748b8627 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/runs.py @@ -12,7 +12,9 @@ single-simulation household report), plus, for an economy report, the comparison step's own status and Modal call. A household report has no comparison step: one Modal call computes both households, so its simulation -runs share that call. +runs share that call. Each report run of a budget-window report is one year of +a budget-window run (``budget_window_run_id``), which has at most one report +run per year. Each status moves ``pending`` -> ``dispatching`` -> ``running`` -> ``succeeded`` or ``failed``. The region and dataset keys include the catalog version, through @@ -289,6 +291,7 @@ def _execution_columns(prefix: str, *, required: bool) -> list[Column[Any]]: Column("spm_provenance", JSONB(none_as_null=True), nullable=True), Column("result_model_version", Text, nullable=True), Column("result_data_version", Text, nullable=True), + Column("budget_window_run_id", Uuid(as_uuid=True), nullable=True), CheckConstraint(f"year BETWEEN {FIRST_YEAR} AND {LAST_YEAR}", name="year"), CheckConstraint(f"trigger IN ({_values(RunTrigger)})", name="trigger"), CheckConstraint( @@ -332,6 +335,19 @@ def _execution_columns(prefix: str, *, required: bool) -> list[Column[Any]]: name="fk_report_runs_reform_simulation_run", ondelete="RESTRICT", ), + # A year of a budget-window run is a report run of the same report and + # catalog version (``budget_window_runs``), and each year has one. + ForeignKeyConstraint( + ["budget_window_run_id", "report_id", "tax_benefit_model_version_id"], + [ + "simulation.budget_window_runs.id", + "simulation.budget_window_runs.report_id", + "simulation.budget_window_runs.tax_benefit_model_version_id", + ], + name="fk_report_runs_budget_window_run", + ondelete="RESTRICT", + ), + UniqueConstraint("budget_window_run_id", "year"), ) # A report's status is its newest report run's. Index( @@ -409,6 +425,9 @@ class NewReportRun: baseline_simulation_run_id: UUID trigger: RunTrigger reform_simulation_run_id: UUID | None = None + # The budget-window run this report run is a year of; a budget-window + # report's report runs each have one, and other reports' none. + budget_window_run_id: UUID | None = None def _record_values(cls: type, row: RowMapping) -> dict[str, Any]: @@ -480,6 +499,7 @@ class ReportRunRecord: comparison_completed_at: datetime | None created_at: datetime updated_at: datetime + budget_window_run_id: UUID | None = None @classmethod def from_row(cls, row: RowMapping) -> Self: diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py index 00ec22e10..b1cad0df4 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py @@ -17,7 +17,7 @@ # The newest revision in migrations/versions. Each migration PR updates it, and # tests/test_schema.py fails when it differs from the chain's head. The # readiness check reports it next to the database's revision. -HEAD_REVISION = "0004" +HEAD_REVISION = "0005" NAMING_CONVENTION = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", diff --git a/libs/policyengine-simulation-persistence/tests/budget_window_support.py b/libs/policyengine-simulation-persistence/tests/budget_window_support.py new file mode 100644 index 000000000..de89aae0f --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/budget_window_support.py @@ -0,0 +1,117 @@ +"""Builders for the budget-window run tests.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_persistence import ( + BudgetWindowRunStore, + NewBudgetWindowRun, + NewReportRun, + RunTrigger, + create_simulation_engine, +) +from policyengine_simulation_persistence.result_tables import budget_impacts +from report_support import economy_pair, economy_report, report_store +from run_support import economy_run, run_store + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_persistence import ( + ReportRecord, + ReportRunRecord, + SimulationRecord, + ) + from public_api_support import PublicApiRows + from transition_support import ManualClock + +WINDOW_START = 2026 +WINDOW_SIZE = 3 + + +def budget_window_store( + database, clock: ManualClock | None = None +) -> BudgetWindowRunStore: + """Return a store that connects as the runtime role.""" + engine = create_simulation_engine(database.runtime_url) + if clock is None: + return BudgetWindowRunStore(engine) + return BudgetWindowRunStore(engine, clock=clock) + + +def budget_window_report( + database, rows: PublicApiRows, *, window_size: int = WINDOW_SIZE +) -> tuple[ReportRecord, tuple[SimulationRecord, SimulationRecord]]: + """Create a California budget-window report from 2026, and its simulations.""" + baseline, reform = economy_pair(database, rows) + report = report_store(database).create( + economy_report( + baseline, reform, year=WINDOW_START, budget_window_size=window_size + ) + ) + return report, (baseline, reform) + + +def new_budget_window_run( + report: ReportRecord, rows: PublicApiRows, **overrides: object +) -> NewBudgetWindowRun: + """Return a first calculation of the report on the first catalog version.""" + values: dict[str, object] = { + "report_id": report.id, + "tax_benefit_model_version_id": rows.catalog_version_id, + "trigger": RunTrigger.INITIAL, + "max_concurrent_years": 2, + } + values.update(overrides) + return NewBudgetWindowRun(**values) # type: ignore[arg-type] + + +def year_report_run( + database, + rows: PublicApiRows, + window: tuple[ReportRecord, tuple[SimulationRecord, SimulationRecord]], + *, + year: int, + budget_window_run_id: UUID | None, +) -> ReportRunRecord: + """Start both simulations' runs for ``year`` and join them in a report run. + + ``window`` is a report and its two simulations, as ``budget_window_report`` + returns them. + """ + report, simulations = window + store = run_store(database) + baseline_run, reform_run = ( + store.start_simulation_run(economy_run(simulation, rows, year=year)).run + for simulation in simulations + ) + return store.create_report_run( + NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + budget_window_run_id=budget_window_run_id, + ) + ) + + +def write_year_budget(database, report_run: ReportRunRecord, scale: float) -> None: + """Store a year's budget figures, as a recorded comparison would.""" + engine = create_simulation_engine(database.runtime_url) + try: + with engine.begin() as connection: + connection.execute( + budget_impacts.insert().values( + report_run_id=report_run.id, + budgetary_impact=-1.5 * scale, + tax_revenue_impact=2.25 * scale, + state_tax_revenue_impact=0.25 * scale, + benefit_spending_impact=3.75 * scale, + households=100.0, + baseline_net_income=1000.0, + ) + ) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py new file mode 100644 index 000000000..d71eabb59 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py @@ -0,0 +1,551 @@ +"""budget_window_runs: budget-window calculations, one report run per year.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from datetime import timedelta +from threading import Barrier +from uuid import uuid4 + +import pytest +from alembic import command +from budget_window_support import ( + WINDOW_SIZE, + WINDOW_START, + budget_window_report, + budget_window_store, + new_budget_window_run, + write_year_budget, + year_report_run, +) +from migration_support import alembic_config, migration_cycle_failures +from policyengine_simulation_persistence import ( + BudgetWindowStatus, + NewBudgetWindowRun, + NewReport, + ReportKind, + RunError, + RunErrorCode, + RunStatus, + RunTrigger, + YearBudget, + create_simulation_engine, +) +from public_api_support import superuser_engine +from report_support import economy_pair, economy_report, report_store +from run_support import set_simulation_run_status +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError, ProgrammingError +from transition_support import ManualClock + +pytestmark = pytest.mark.postgres + +FIRST_CHECK = timedelta(minutes=5) + + +def test_the_budget_window_revision_matches_its_tables_and_downgrades( + disposable_database, +) -> None: + config = alembic_config() + + assert migration_cycle_failures(config, check=True) == [] + command.downgrade(config, "0004") + command.upgrade(config, "head") + command.check(config) + + +def test_a_budget_window_report_records_its_window_size( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + + assert report.year == WINDOW_START + assert report.budget_window_size == WINDOW_SIZE + assert report_store(migrated_database).get(report.id) == report + + +def test_a_single_year_report_has_no_window(migrated_database, public_api_rows) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + + report = report_store(migrated_database).create(economy_report(baseline, reform)) + + assert report.budget_window_size is None + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"budget_window_size": 0}, "between 1 and 75"), + ({"budget_window_size": 76}, "between 1 and 75"), + ({"budget_window_size": 3, "year": 2199}, "2200"), + ({"budget_window_size": 3, "kind": ReportKind.HOUSEHOLD}, "economy"), + ], +) +def test_a_window_size_is_checked_before_the_database(overrides, message) -> None: + values: dict[str, object] = { + "kind": ReportKind.ECONOMY, + "country": "us", + "baseline_simulation_id": uuid4(), + "reform_simulation_id": uuid4(), + "year": 2026, + } + values.update(overrides) + + with pytest.raises(ValueError, match=message): + NewReport(**values) # type: ignore[arg-type] + + +@pytest.mark.parametrize("window_size", [0, 76]) +def test_the_database_checks_the_window_size( + migrated_database, public_api_rows, window_size +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + engine = superuser_engine(migrated_database.superuser_url) + try: + with ( + pytest.raises(IntegrityError, match="ck_reports_budget_window_size"), + engine.begin() as connection, + ): + connection.execute( + text( + "INSERT INTO simulation.reports (id, kind, country, " + "baseline_simulation_id, reform_simulation_id, year, " + "budget_window_size) VALUES (:id, 'economy', 'us', " + ":baseline, :reform, 2026, :size)" + ), + { + "id": uuid4(), + "baseline": report.baseline_simulation_id, + "reform": report.reform_simulation_id, + "size": window_size, + }, + ) + finally: + engine.dispose() + + +def test_a_budget_window_run_starts_running_and_due_a_check( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + + run = store.create_run( + new_budget_window_run(report, public_api_rows, force=True), + first_check_at=clock() + FIRST_CHECK, + ) + + assert run.report_id == report.id + assert run.tax_benefit_model_version_id == public_api_rows.catalog_version_id + assert run.status is BudgetWindowStatus.RUNNING + assert run.trigger is RunTrigger.INITIAL + assert run.force is True + assert run.max_concurrent_years == 2 + assert run.next_check_at == clock() + FIRST_CHECK + assert run.completed_at is None + assert store.get_run(run.id) == run + assert store.get_run(uuid4()) is None + + +def test_only_a_budget_window_report_has_budget_window_runs( + migrated_database, public_api_rows +) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + single_year = report_store(migrated_database).create( + economy_report(baseline, reform) + ) + store = budget_window_store(migrated_database) + + with pytest.raises(ValueError, match="budget-window report"): + store.create_run( + new_budget_window_run(single_year, public_api_rows), + first_check_at=ManualClock()(), + ) + unknown = new_budget_window_run(single_year, public_api_rows, report_id=uuid4()) + with pytest.raises(LookupError, match="report"): + store.create_run(unknown, first_check_at=ManualClock()()) + + +def test_a_run_calculates_at_least_one_year_at_once() -> None: + with pytest.raises(ValueError, match="max_concurrent_years"): + NewBudgetWindowRun( + report_id=uuid4(), + tax_benefit_model_version_id=uuid4(), + trigger=RunTrigger.INITIAL, + max_concurrent_years=0, + ) + + +def test_each_year_is_one_report_run_of_the_budget_window_run( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + store = budget_window_store(migrated_database) + run = store.create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + + later, first = ( + year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=year, + budget_window_run_id=run.id, + ) + for year in (WINDOW_START + 1, WINDOW_START) + ) + + years = store.years(run.id) + assert [year.report_run.year for year in years] == [WINDOW_START, WINDOW_START + 1] + assert [year.report_run for year in years] == [first, later] + assert first.budget_window_run_id == run.id + assert {year.baseline_run.simulation_id for year in years} == {simulations[0].id} + assert {year.reform_run.simulation_id for year in years} == {simulations[1].id} + assert {year.baseline_run.year for year in years} == { + WINDOW_START, + WINDOW_START + 1, + } + with pytest.raises(IntegrityError, match="uq_report_runs_budget_window_run_id"): + year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=run.id, + ) + + +def test_a_year_outside_the_window_is_refused( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + run = budget_window_store(migrated_database).create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + + for year in (WINDOW_START - 1, WINDOW_START + WINDOW_SIZE): + with pytest.raises(ValueError, match="window"): + year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=year, + budget_window_run_id=run.id, + ) + + +def test_a_budget_window_reports_runs_belong_to_its_budget_window_runs( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + other_report, other_simulations = budget_window_report( + migrated_database, public_api_rows + ) + other_run = budget_window_store(migrated_database).create_run( + new_budget_window_run(other_report, public_api_rows), + first_check_at=ManualClock()(), + ) + + with pytest.raises(ValueError, match="budget-window run"): + year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=None, + ) + with pytest.raises(IntegrityError, match="fk_report_runs_budget_window_run"): + year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=other_run.id, + ) + del other_simulations + + +def test_a_single_year_report_run_has_no_budget_window_run( + migrated_database, public_api_rows +) -> None: + baseline, reform = economy_pair(migrated_database, public_api_rows) + single_year = report_store(migrated_database).create( + economy_report(baseline, reform) + ) + window_report, _ = budget_window_report(migrated_database, public_api_rows) + window_run = budget_window_store(migrated_database).create_run( + new_budget_window_run(window_report, public_api_rows), + first_check_at=ManualClock()(), + ) + + with pytest.raises(ValueError, match="window"): + year_report_run( + migrated_database, + public_api_rows, + (single_year, (baseline, reform)), + year=2026, + budget_window_run_id=window_run.id, + ) + + +def test_a_budget_window_run_finishes_once(migrated_database, public_api_rows) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + succeeded = store.create_run( + new_budget_window_run(report, public_api_rows), first_check_at=clock() + ) + failed = store.create_run( + new_budget_window_run(report, public_api_rows), first_check_at=clock() + ) + error = RunError(RunErrorCode.WORKER_ERROR, "The 2027 comparison failed.") + + assert store.finish_run(succeeded.id, BudgetWindowStatus.SUCCEEDED) is True + assert store.finish_run(failed.id, BudgetWindowStatus.FAILED, error) is True + assert store.finish_run(failed.id, BudgetWindowStatus.SUCCEEDED) is False + + finished = store.get_run(succeeded.id) + assert finished is not None + assert finished.status is BudgetWindowStatus.SUCCEEDED + assert finished.completed_at == clock() + assert finished.next_check_at is None + failure = store.get_run(failed.id) + assert failure is not None + assert failure.status is BudgetWindowStatus.FAILED + assert (failure.error_code, failure.error_message) == ( + "worker_error", + "The 2027 comparison failed.", + ) + with pytest.raises(ValueError, match="error"): + store.finish_run(succeeded.id, BudgetWindowStatus.FAILED) + with pytest.raises(ValueError, match="finished"): + store.finish_run(succeeded.id, BudgetWindowStatus.RUNNING) + + +def test_one_scheduler_claims_a_run_at_a_time( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + run = store.create_run( + new_budget_window_run(report, public_api_rows), first_check_at=clock() + ) + lease = timedelta(minutes=2) + + assert store.claim_scheduling(run.id, lease=lease) is True + assert store.claim_scheduling(run.id, lease=lease) is False + store.release_scheduling(run.id, next_check_at=clock() + FIRST_CHECK) + assert store.claim_scheduling(run.id, lease=lease) is True + clock.advance(lease) + # A scheduler that stopped without releasing its claim loses it. + assert store.claim_scheduling(run.id, lease=lease) is True + released = store.get_run(run.id) + assert released is not None + store.finish_run(run.id, BudgetWindowStatus.SUCCEEDED) + assert store.claim_scheduling(run.id, lease=lease) is False + + +SIMULTANEOUS_SCHEDULERS = 6 + + +def test_simultaneous_schedulers_claim_a_run_once( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + store = budget_window_store(migrated_database) + run = store.create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + all_ready = Barrier(SIMULTANEOUS_SCHEDULERS) + + def claim(_: int) -> bool: + all_ready.wait() + return store.claim_scheduling(run.id, lease=timedelta(minutes=2)) + + with ThreadPoolExecutor(max_workers=SIMULTANEOUS_SCHEDULERS) as pool: + claims = list(pool.map(claim, range(SIMULTANEOUS_SCHEDULERS))) + + assert sum(claims) == 1 + + +def test_releasing_a_claim_sets_the_next_check( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + run = store.create_run( + new_budget_window_run(report, public_api_rows), first_check_at=clock() + ) + store.claim_scheduling(run.id, lease=timedelta(minutes=2)) + + store.release_scheduling(run.id, next_check_at=clock() + FIRST_CHECK) + + released = store.get_run(run.id) + assert released is not None + assert released.next_check_at == clock() + FIRST_CHECK + + +def test_due_runs_are_the_running_runs_whose_check_is_due( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + later, due, finished = ( + store.create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=clock() + offset, + ) + for offset in (FIRST_CHECK, timedelta(0), timedelta(0)) + ) + store.finish_run(finished.id, BudgetWindowStatus.SUCCEEDED) + + assert [run.id for run in store.due_runs(now=clock(), limit=10)] == [due.id] + clock.advance(FIRST_CHECK) + assert [run.id for run in store.due_runs(now=clock(), limit=10)] == [ + due.id, + later.id, + ] + assert len(store.due_runs(now=clock(), limit=1)) == 1 + + +def test_a_reports_budget_window_runs_are_listed_newest_first( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + first = store.create_run( + new_budget_window_run(report, public_api_rows), first_check_at=clock() + ) + clock.advance(timedelta(minutes=1)) + rerun = store.create_run( + new_budget_window_run(report, public_api_rows, trigger=RunTrigger.MANUAL), + first_check_at=clock(), + ) + + assert store.runs(report.id) == [rerun, first] + assert store.latest_run(report.id) == rerun + assert store.latest_run(uuid4()) is None + + +SIMULTANEOUS_FIRST_OPENS = 6 + + +def test_simultaneous_first_runs_of_a_report_add_one( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + store = budget_window_store(migrated_database) + first_open = new_budget_window_run( + report, public_api_rows, trigger=RunTrigger.SYSTEM + ) + all_ready = Barrier(SIMULTANEOUS_FIRST_OPENS) + + def open_first_time(_: int): + all_ready.wait() + return store.create_first_run(first_open, first_check_at=ManualClock()()) + + with ThreadPoolExecutor(max_workers=SIMULTANEOUS_FIRST_OPENS) as pool: + started = list(pool.map(open_first_time, range(SIMULTANEOUS_FIRST_OPENS))) + + assert sum(start.created for start in started) == 1 + assert {start.run.id for start in started} == {started[0].run.id} + assert len(store.runs(report.id)) == 1 + again = store.create_first_run(first_open, first_check_at=ManualClock()()) + assert again.created is False + + +def test_each_years_budget_figures_are_read_back_in_year_order( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + store = budget_window_store(migrated_database) + run = store.create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + for year in (WINDOW_START + 1, WINDOW_START): + report_run = year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=year, + budget_window_run_id=run.id, + ) + write_year_budget(migrated_database, report_run, scale=year - 2025) + + assert store.year_budgets(run.id) == [ + YearBudget( + year=WINDOW_START, + budgetary_impact=-1.5, + tax_revenue_impact=2.25, + state_tax_revenue_impact=0.25, + benefit_spending_impact=3.75, + ), + YearBudget( + year=WINDOW_START + 1, + budgetary_impact=-3.0, + tax_revenue_impact=4.5, + state_tax_revenue_impact=0.5, + benefit_spending_impact=7.5, + ), + ] + + +def test_a_years_simulation_runs_are_read_with_its_report_run( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + store = budget_window_store(migrated_database) + run = store.create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + report_run = year_report_run( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=run.id, + ) + set_simulation_run_status( + migrated_database, report_run.reform_simulation_run_id, "failed" + ) + + (year,) = store.years(run.id) + + assert year.baseline_run.status is RunStatus.PENDING + assert year.reform_run.status is RunStatus.FAILED + assert year.reform_run.error_code == "test_failure" + + +def test_the_runtime_role_updates_budget_window_runs_but_cannot_delete_them( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + budget_window_store(migrated_database).create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + engine = create_simulation_engine(migrated_database.runtime_url) + try: + with engine.begin() as connection: + updated = connection.execute( + text("UPDATE simulation.budget_window_runs SET updated_at = now()") + ) + assert updated.rowcount == 1 + with ( + pytest.raises(ProgrammingError, match="permission denied"), + engine.begin() as connection, + ): + connection.execute(text("DELETE FROM simulation.budget_window_runs")) + finally: + engine.dispose() diff --git a/libs/policyengine-simulation-persistence/tests/test_schema.py b/libs/policyengine-simulation-persistence/tests/test_schema.py index d116d72e0..655a05c0a 100644 --- a/libs/policyengine-simulation-persistence/tests/test_schema.py +++ b/libs/policyengine-simulation-persistence/tests/test_schema.py @@ -121,6 +121,7 @@ def test_importing_the_package_registers_every_table() -> None: "simulation.reports", "simulation.simulation_runs", "simulation.report_runs", + "simulation.budget_window_runs", *( f"simulation.{name}" for name in ( diff --git a/migrations/versions/0005_budget_window_runs.py b/migrations/versions/0005_budget_window_runs.py new file mode 100644 index 000000000..f10fb2289 --- /dev/null +++ b/migrations/versions/0005_budget_window_runs.py @@ -0,0 +1,193 @@ +"""Add budget-window reports and their runs. + +A report gains ``budget_window_size``: an economy report with one compares its +two simulations in each of that many years, from its year on (at most 75, and +none past the last year a run may have). ``budget_window_runs`` records each +calculation of such a report on a catalog version: its trigger, whether its +years calculate their simulations again, how many years it calculates at +once, its status and error, when the safety check next looks at it, and the +lease of the caller starting its next years. A report run gains +``budget_window_run_id``: the budget-window run it is a year of, of the same +report and catalog version (a composite foreign key), with at most one report +run per year. The default privileges let the runtime role read, insert and +update the new table, and not delete from it. Nothing writes the new columns +or table yet. + +Revision ID: 0005 +Revises: 0004 +Create Date: 2026-10-29 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "0005" +down_revision = "0004" +branch_labels = None +depends_on = None + +SCHEMA = "simulation" +TRIGGERS = "'initial', 'manual', 'system'" +STATUSES = "'running', 'succeeded', 'failed'" +MAX_BUDGET_WINDOW_YEARS = 75 +LAST_YEAR = 2200 + + +def _add_budget_window_size() -> None: + op.add_column( + "reports", + sa.Column("budget_window_size", sa.Integer(), nullable=True), + schema=SCHEMA, + ) + op.create_check_constraint( + op.f("ck_reports_budget_window_size"), + "reports", + "budget_window_size IS NULL OR (kind = 'economy' " + f"AND budget_window_size BETWEEN 1 AND {MAX_BUDGET_WINDOW_YEARS} " + f"AND year + budget_window_size - 1 <= {LAST_YEAR})", + schema=SCHEMA, + ) + + +def _create_budget_window_runs() -> None: + op.create_table( + "budget_window_runs", + sa.Column("id", sa.Uuid(), nullable=False), + sa.Column("report_id", sa.Uuid(), nullable=False), + sa.Column("tax_benefit_model_version_id", sa.Uuid(), nullable=False), + sa.Column("trigger", sa.Text(), nullable=False), + sa.Column("force", sa.Boolean(), nullable=False), + sa.Column("max_concurrent_years", sa.Integer(), nullable=False), + sa.Column("status", sa.Text(), nullable=False), + sa.Column("error_code", sa.Text(), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column("next_check_at", sa.DateTime(timezone=True), nullable=True), + sa.Column( + "scheduling_claimed_until", sa.DateTime(timezone=True), nullable=True + ), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + f"trigger IN ({TRIGGERS})", name=op.f("ck_budget_window_runs_trigger") + ), + sa.CheckConstraint( + f"status IN ({STATUSES})", name=op.f("ck_budget_window_runs_status") + ), + sa.CheckConstraint( + "max_concurrent_years >= 1", + name=op.f("ck_budget_window_runs_max_concurrent_years"), + ), + sa.CheckConstraint( + "(status = 'running') = (next_check_at IS NOT NULL)", + name=op.f("ck_budget_window_runs_next_check_at"), + ), + sa.CheckConstraint( + "(status IN ('succeeded', 'failed')) = (completed_at IS NOT NULL)", + name=op.f("ck_budget_window_runs_completed_at"), + ), + sa.CheckConstraint( + "(status = 'failed') = (error_code IS NOT NULL)", + name=op.f("ck_budget_window_runs_error"), + ), + sa.ForeignKeyConstraint( + ["report_id"], + [f"{SCHEMA}.reports.id"], + name="fk_budget_window_runs_report", + ondelete="RESTRICT", + ), + sa.ForeignKeyConstraint( + ["tax_benefit_model_version_id"], + ["public.tax_benefit_model_versions.id"], + name="fk_budget_window_runs_catalog_version", + ondelete="RESTRICT", + ), + sa.PrimaryKeyConstraint("id", name="pk_budget_window_runs"), + sa.UniqueConstraint( + "id", + "report_id", + "tax_benefit_model_version_id", + name="uq_budget_window_runs_id_report_catalog_version", + ), + schema=SCHEMA, + ) + op.create_index( + "ix_budget_window_runs_report_id_created_at", + "budget_window_runs", + ["report_id", "created_at"], + schema=SCHEMA, + ) + op.create_index( + "ix_budget_window_runs_due_checks", + "budget_window_runs", + ["next_check_at"], + schema=SCHEMA, + postgresql_where=sa.text("status = 'running'"), + ) + op.create_index( + "ix_budget_window_runs_tax_benefit_model_version_id", + "budget_window_runs", + ["tax_benefit_model_version_id"], + schema=SCHEMA, + ) + + +def _add_budget_window_run_id() -> None: + op.add_column( + "report_runs", + sa.Column("budget_window_run_id", sa.Uuid(), nullable=True), + schema=SCHEMA, + ) + op.create_foreign_key( + "fk_report_runs_budget_window_run", + "report_runs", + "budget_window_runs", + ["budget_window_run_id", "report_id", "tax_benefit_model_version_id"], + ["id", "report_id", "tax_benefit_model_version_id"], + source_schema=SCHEMA, + referent_schema=SCHEMA, + ondelete="RESTRICT", + ) + op.create_unique_constraint( + "uq_report_runs_budget_window_run_id_year", + "report_runs", + ["budget_window_run_id", "year"], + schema=SCHEMA, + ) + + +def upgrade() -> None: + """Add the window size, the budget-window runs and each year's link.""" + _add_budget_window_size() + _create_budget_window_runs() + _add_budget_window_run_id() + + +def downgrade() -> None: + """Drop what ``upgrade`` added, with its indexes and keys.""" + op.drop_constraint( + "uq_report_runs_budget_window_run_id_year", + "report_runs", + schema=SCHEMA, + type_="unique", + ) + op.drop_constraint( + "fk_report_runs_budget_window_run", + "report_runs", + schema=SCHEMA, + type_="foreignkey", + ) + op.drop_column("report_runs", "budget_window_run_id", schema=SCHEMA) + for index in ( + "ix_budget_window_runs_tax_benefit_model_version_id", + "ix_budget_window_runs_due_checks", + "ix_budget_window_runs_report_id_created_at", + ): + op.drop_index(index, table_name="budget_window_runs", schema=SCHEMA) + op.drop_table("budget_window_runs", schema=SCHEMA) + op.drop_constraint( + op.f("ck_reports_budget_window_size"), "reports", schema=SCHEMA, type_="check" + ) + op.drop_column("reports", "budget_window_size", schema=SCHEMA) From d63d3715172b4802c44b70fc3d69ae067fd05e62 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:02:01 +0400 Subject: [PATCH 02/20] Run an economy report's two simulations for any year of its window EconomyReports gains three public steps the budget-window path needs: prepare_report creates a report and its simulations without running it, optionally with a budget window; report_setup reads a report's simulations for a catalog version; run_year finds or creates both simulation runs and a report run for one year, recorded as one year of a budget-window run even on the system trigger, and starts what can start, with the same reuse, forcing and labour supply rules as a single-year report. submit and run_report use them unchanged. A budget-window report is not a single-year economy report: run_report and the economy report reads answer that no such economy report exists. Co-Authored-By: Claude Opus 5.5 --- .../economy_reports.py | 120 ++++++++++-- .../economy_routes.py | 8 +- .../tests/budget_window_fakes.py | 112 +++++++++++ .../tests/test_economy_report_years.py | 185 ++++++++++++++++++ 4 files changed, 402 insertions(+), 23 deletions(-) create mode 100644 projects/policyengine-simulation-entry/tests/budget_window_fakes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_report_years.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py index 796a589e2..301a1e889 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py @@ -14,6 +14,11 @@ both runs have succeeded: at once when both were reused, or when a call's report says that its run succeeded (``start_comparisons_after``). +A budget-window report (``budget_window_reports``) runs the same steps once +per year of its window: ``run_year`` records each year's report run as one +year of its budget-window run. The single-year path neither runs it nor +counts its runs as the report's. + A report whose policies change a labour supply elasticity needs both simulations' labour supply response columns, which a run calculates only when asked, so its runs are always calculated again (forced) and asked for @@ -117,6 +122,15 @@ class _PlannedRun: refusal: str | None = None +@dataclass(frozen=True) +class ReportSetup: + """A report and its two simulations, to run on one catalog version.""" + + report: ReportRecord + version: CatalogVersion + simulations: tuple[SimulationRecord, SimulationRecord] + + @dataclass class EconomyReports: """The v2 economy path: report submission, runs, calls and comparisons.""" @@ -136,15 +150,34 @@ async def submit( ) -> tuple[ReportRecord, ReportRunRecord]: """Create a report's simulations, the report and its first run, and start it. - Raise EconomyRequestError, before anything is written, when the - country has no catalog for the policyengine.py version, the version - has no such region or dataset, options are given, or a policy or - dynamic does not exist for the country. + Raise EconomyRequestError, before anything is written, as + ``prepare_report`` does. + """ + setup = await self.prepare_report(submission) + report_run = await self._run( + setup, setup.report.year, RunTrigger.INITIAL, force=False + ) + return setup.report, report_run + + async def prepare_report( + self, + submission: EconomyReportSubmission, + *, + budget_window_size: int | None = None, + ) -> ReportSetup: + """Create a report and its simulations, without running it. + + ``budget_window_size`` makes it a budget-window report of that many + years from the submission's year. Raise EconomyRequestError, before + anything is written, when the country has no catalog for the + policyengine.py version, the version has no such region or dataset, + options are given, or a policy or dynamic does not exist for the + country. """ if submission.options: msg = "economy reports take no options yet" raise EconomyRequestError(msg) - version = await self._version( + version = await self.catalog_version( submission.country, submission.policyengine_version ) await self._population(version, submission.region_code, submission.dataset_name) @@ -177,12 +210,10 @@ async def submit( baseline_simulation_id=baseline.id, reform_simulation_id=reform.id, year=submission.year, + budget_window_size=budget_window_size, ), ) - report_run = await self._run( - report, version, (baseline, reform), RunTrigger.INITIAL, force=False - ) - return report, report_run + return ReportSetup(report, version, (baseline, reform)) async def run_report( self, @@ -209,7 +240,11 @@ async def run_report( if report is None: msg = f"no report {report_id}" raise ReportNotFoundError(msg) - if report.kind is not ReportKind.ECONOMY or report.reform_simulation_id is None: + if ( + report.kind is not ReportKind.ECONOMY + or report.reform_simulation_id is None + or report.budget_window_size is not None + ): msg = f"no economy report {report_id}" raise ReportNotFoundError(msg) if report.options: @@ -219,22 +254,58 @@ async def run_report( newest = await asyncio.to_thread(self.runs.latest_report_run, report_id) if newest is not None: return newest - version = await self._version(report.country, policyengine_version) + version = await self.catalog_version(report.country, policyengine_version) + setup = await self.report_setup(report, version) + return await self._run(setup, report.year, trigger, force=force) + + async def report_setup( + self, report: ReportRecord, version: CatalogVersion + ) -> ReportSetup: + """Return an economy report with its two simulations, to run on a version.""" simulations = [] for simulation_id in ( report.baseline_simulation_id, report.reform_simulation_id, ): - simulation = await asyncio.to_thread(self.simulations.get, simulation_id) + simulation = ( + None + if simulation_id is None + else await asyncio.to_thread(self.simulations.get, simulation_id) + ) if simulation is None: # pragma: no cover - a foreign key keeps it msg = f"no simulation {simulation_id}" raise ReportNotFoundError(msg) simulations.append(simulation) + return ReportSetup(report, version, (simulations[0], simulations[1])) + + async def run_year( + self, + setup: ReportSetup, + *, + year: int, + trigger: RunTrigger, + force: bool, + budget_window_run_id: UUID | None = None, + ) -> ReportRunRecord: + """Run a report's two simulations and their comparison for one year. + + The report run is one year of ``budget_window_run_id``, a + budget-window run of the report, and is recorded as a run of its + own whatever the trigger. Runs are reused, forced and started as + for a single-year report. + """ return await self._run( - report, version, (simulations[0], simulations[1]), trigger, force=force + setup, + year, + trigger, + force=force, + budget_window_run_id=budget_window_run_id, ) - async def _version(self, country: str, policyengine_version: str) -> CatalogVersion: + async def catalog_version( + self, country: str, policyengine_version: str + ) -> CatalogVersion: + """Return the country's catalog version; raise EconomyRequestError if none.""" version = await asyncio.to_thread( self.catalog.catalog_version, country, policyengine_version ) @@ -287,15 +358,18 @@ async def _plan( async def _run( self, - report: ReportRecord, - version: CatalogVersion, - simulations: tuple[SimulationRecord, SimulationRecord], + setup: ReportSetup, + year: int, trigger: RunTrigger, *, force: bool, + budget_window_run_id: UUID | None = None, ) -> ReportRunRecord: """Find or create both runs and the report run, then start what can start.""" - planned = [await self._plan(version, simulation) for simulation in simulations] + report, version = setup.report, setup.version + planned = [ + await self._plan(version, simulation) for simulation in setup.simulations + ] labor_supply_response = any( run.changes is not None and policy_affects_labor_supply(report.country, run.changes) @@ -313,7 +387,7 @@ async def _run( NewSimulationRun( simulation_id=run.simulation.id, tax_benefit_model_version_id=version.id, - year=report.year, + year=year, trigger=trigger, region_code=run.population.region_code, dataset_id=run.population.dataset_id, @@ -328,6 +402,7 @@ async def _run( baseline_simulation_run_id=baseline.id, reform_simulation_run_id=reform.id, trigger=trigger, + budget_window_run_id=budget_window_run_id, ) ) compared = { @@ -370,9 +445,12 @@ async def _record_report_run(self, new_report_run: NewReportRun) -> ReportRunRec Simultaneous first opens of a report without a run then share one report run. A caller that loses still starts the runs it created that the recorded run compares, and leaves the others pending for the - safety check. + safety check. A year of a budget-window run is always its own run. """ - if new_report_run.trigger is not RunTrigger.SYSTEM: + if ( + new_report_run.trigger is not RunTrigger.SYSTEM + or new_report_run.budget_window_run_id is not None + ): return await asyncio.to_thread(self.runs.create_report_run, new_report_run) first = await asyncio.to_thread( self.runs.create_first_report_run, new_report_run 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 e7a8cddc8..c1c6df58f 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 @@ -336,8 +336,12 @@ def simulation(self, simulation_id: UUID) -> SimulationResponse: @staticmethod def _report(record: ReportRecord | None, identifier: object) -> ReportRecord: - """Return an economy report; a household report is not one.""" - if record is None or record.kind is not ReportKind.ECONOMY: + """Return an economy report; a household or budget-window report is not one.""" + if ( + record is None + or record.kind is not ReportKind.ECONOMY + or record.budget_window_size is not None + ): msg = f"no economy report {identifier}" raise ReportNotFoundError(msg) return record diff --git a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py new file mode 100644 index 000000000..4c6a0d9b1 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py @@ -0,0 +1,112 @@ +"""In-memory stores and a path builder for the budget-window tests.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from typing import TYPE_CHECKING, Any +from uuid import uuid4 + +from economy_fakes import ( + CTC_REFORM, + LABOR_SUPPLY_REFORM, + FakeCatalog, + FakeReports, + FakeRuns, + FakeSimulations, +) +from policyengine_simulation_contract.economy_runs import ( + EconomyComparisonCall, + EconomySimulationCall, +) +from policyengine_simulation_entry.economy_report_types import ( + EconomyReportSubmission, + SimulationChoice, +) +from policyengine_simulation_entry.economy_reports import EconomyReports +from policyengine_simulation_entry.run_dispatch import RunDispatchService +from run_fakes import FakeDispatcher + +if TYPE_CHECKING: + from policyengine_simulation_persistence import NewReportRun, ReportRunRecord + +REFORM_POLICY = uuid4() +LABOR_SUPPLY_POLICY = uuid4() + + +class WindowRuns(FakeRuns): + """``FakeRuns`` whose report runs keep the budget-window run they belong to.""" + + def create_report_run(self, new_report_run: NewReportRun) -> ReportRunRecord: + row = replace( + super().create_report_run(new_report_run), + budget_window_run_id=new_report_run.budget_window_run_id, + ) + self.report_runs_by_id[row.id] = row + return row + + +class EconomyPath: + """The economy path over in-memory stores, with every spawned request kept.""" + + def __init__(self) -> None: + self.catalog = FakeCatalog() + self.catalog.changes.update( + {REFORM_POLICY: CTC_REFORM, LABOR_SUPPLY_POLICY: LABOR_SUPPLY_REFORM} + ) + self.simulations = FakeSimulations() + self.simulations.policies_of_country["us"].update( + {REFORM_POLICY, LABOR_SUPPLY_POLICY} + ) + self.reports = FakeReports() + self.runs = WindowRuns() + self.dispatcher = FakeDispatcher() + self.economy = EconomyReports( + simulations=self.simulations, + reports=self.reports, + runs=self.runs, + catalog=self.catalog, + transitions=self.runs, + dispatch=RunDispatchService( + transitions=self.runs, dispatcher=self.dispatcher, clock=self.runs.clock + ), + artifact_environment="staging", + ) + + @staticmethod + def submission(**changes: object) -> EconomyReportSubmission: + """Return a California 2026 submission of the CTC reform.""" + values: dict[str, Any] = { + "country": "us", + "policyengine_version": "6.7.5", + "year": 2026, + "region_code": "state/ca", + "baseline": SimulationChoice(), + "reform": SimulationChoice(policy_id=REFORM_POLICY), + } + values.update(changes) + return EconomyReportSubmission(**values) + + def prepare(self, *, window_size: int | None = 3, **changes: object): + """Create a report, by default a three-year budget window from 2026.""" + return asyncio.run( + self.economy.prepare_report( + self.submission(**changes), budget_window_size=window_size + ) + ) + + def simulation_calls(self) -> list[EconomySimulationCall]: + """Return the simulation calls started so far.""" + return [ + EconomySimulationCall.model_validate(request.arguments[0]) + for request in self.dispatcher.requests + if request.target.function.startswith("run_economy_simulation") + ] + + def comparison_calls(self) -> list[EconomyComparisonCall]: + """Return the comparison calls started so far.""" + return [ + EconomyComparisonCall.model_validate(request.arguments[0]) + for request in self.dispatcher.requests + if request.target.function == "compare_economy_simulations" + ] diff --git a/projects/policyengine-simulation-entry/tests/test_economy_report_years.py b/projects/policyengine-simulation-entry/tests/test_economy_report_years.py new file mode 100644 index 000000000..292224f88 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_report_years.py @@ -0,0 +1,185 @@ +"""An economy report's runs for any year of its window, and budget-window reports. + +A budget-window report compares its simulations in each year of its window, +one report run per year; the single-year path neither runs nor reads it. +""" + +from __future__ import annotations + +import asyncio +from uuid import uuid4 + +import pytest +from budget_window_fakes import LABOR_SUPPLY_POLICY, EconomyPath +from economy_fakes import US_VERSION +from policyengine_simulation_entry.economy_report_types import ( + EconomyRequestError, + ReportNotFoundError, + SimulationChoice, +) +from policyengine_simulation_entry.economy_routes import EconomyReportReads +from policyengine_simulation_persistence import RunTrigger + + +def _run_year(path: EconomyPath, setup, year: int, **changes: object): + values: dict[str, object] = { + "year": year, + "trigger": RunTrigger.INITIAL, + "force": False, + "budget_window_run_id": uuid4(), + } + values.update(changes) + return asyncio.run(path.economy.run_year(setup, **values)) # type: ignore[arg-type] + + +def test_a_budget_window_report_records_its_window_without_running() -> None: + path = EconomyPath() + + setup = path.prepare(window_size=3) + + assert setup.report.year == 2026 + assert setup.report.budget_window_size == 3 + assert setup.version == US_VERSION + assert [simulation.policy_id for simulation in setup.simulations] == [ + None, + setup.simulations[1].policy_id, + ] + assert path.runs.report_runs_by_id == {} + assert path.dispatcher.requests == [] + + +def test_a_budget_window_report_is_refused_before_anything_is_written() -> None: + path = EconomyPath() + + with pytest.raises(EconomyRequestError, match="region"): + path.prepare(region_code="state/zz") + + assert path.simulations.rows == {} + assert path.reports.rows == {} + + +def test_a_year_runs_both_simulations_in_that_year() -> None: + path = EconomyPath() + setup = path.prepare() + window_run_id = uuid4() + + report_run = _run_year(path, setup, 2027, budget_window_run_id=window_run_id) + + assert report_run.year == 2027 + assert report_run.report_id == setup.report.id + assert report_run.budget_window_run_id == window_run_id + assert report_run.tax_benefit_model_version_id == US_VERSION.id + calls = path.simulation_calls() + assert {call.year for call in calls} == {2027} + assert {call.simulation_run_id for call in calls} == { + report_run.baseline_simulation_run_id, + report_run.reform_simulation_run_id, + } + + +def test_each_year_has_its_own_runs_and_a_year_reuses_its_runs() -> None: + path = EconomyPath() + setup = path.prepare() + + first_year = _run_year(path, setup, 2026) + second_year = _run_year(path, setup, 2027) + again = _run_year(path, setup, 2027, trigger=RunTrigger.MANUAL) + forced = _run_year(path, setup, 2027, trigger=RunTrigger.MANUAL, force=True) + + assert second_year.baseline_simulation_run_id != ( + first_year.baseline_simulation_run_id + ) + assert again.baseline_simulation_run_id == second_year.baseline_simulation_run_id + assert forced.baseline_simulation_run_id != second_year.baseline_simulation_run_id + assert path.runs.simulation_runs[forced.baseline_simulation_run_id].attempt == 2 + + +def test_a_system_year_is_a_report_run_of_its_own() -> None: + """A ``system`` year is not the report's newest run, as a first open is.""" + path = EconomyPath() + setup = path.prepare() + first_year = _run_year(path, setup, 2026, trigger=RunTrigger.SYSTEM) + + second_year = _run_year(path, setup, 2027, trigger=RunTrigger.SYSTEM) + + assert second_year.id != first_year.id + assert second_year.year == 2027 + + +def test_a_labour_supply_year_calculates_its_runs_with_their_columns() -> None: + path = EconomyPath() + setup = path.prepare(reform=SimulationChoice(policy_id=LABOR_SUPPLY_POLICY)) + + _run_year(path, setup, 2026) + + assert {call.labor_supply_response for call in path.simulation_calls()} == {True} + + +def test_a_year_whose_runs_succeeded_starts_its_comparison_at_once() -> None: + path = EconomyPath() + setup = path.prepare() + first = _run_year(path, setup, 2027) + for run_id in (first.baseline_simulation_run_id, first.reform_simulation_run_id): + path.runs.succeed(run_id) + + again = _run_year(path, setup, 2027, trigger=RunTrigger.MANUAL) + + assert [call.report_run_id for call in path.comparison_calls()] == [ + first.id, + again.id, + ] + assert {call.year for call in path.comparison_calls()} == {2027} + + +def test_the_single_year_path_does_not_run_a_budget_window_report() -> None: + path = EconomyPath() + setup = path.prepare() + + with pytest.raises(ReportNotFoundError, match="economy report"): + asyncio.run( + path.economy.run_report( + setup.report.id, + policyengine_version="6.7.5", + trigger=RunTrigger.MANUAL, + force=False, + ) + ) + assert path.dispatcher.requests == [] + + +def test_the_single_year_reads_do_not_find_a_budget_window_report() -> None: + path = EconomyPath() + setup = path.prepare() + reads = EconomyReportReads( + simulations=path.simulations, + reports=path.reports, + runs=path.runs, + catalog=path.catalog, + results=None, # type: ignore[arg-type] + ) + + with pytest.raises(ReportNotFoundError, match="economy report"): + reads.report_by_id(setup.report.id) + with pytest.raises(ReportNotFoundError, match="economy report"): + reads.report_runs(setup.report.id) + + +def test_a_single_year_report_is_set_up_as_before() -> None: + path = EconomyPath() + + setup = path.prepare(window_size=None) + report, report_run = asyncio.run(path.economy.submit(path.submission())) + + assert setup.report.budget_window_size is None + assert report.budget_window_size is None + assert report_run.year == 2026 + assert report_run.budget_window_run_id is None + + +def test_a_setup_reads_the_reports_simulations_on_a_version() -> None: + path = EconomyPath() + prepared = path.prepare() + + setup = asyncio.run(path.economy.report_setup(prepared.report, US_VERSION)) + + assert setup == prepared From a2b66aec3c3f10c4c8b2c61eb9a592142291880b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:07:42 +0400 Subject: [PATCH 03/20] Schedule a budget window's years, a few at a time, until one fails or all succeed BudgetWindowReports submits a budget-window report and its first run, and runs one again (manual, or system, which returns the report's newest run). Each run starts its years through EconomyReports.run_year in year order, at most max_concurrent_years (five by default) at once, so a US national window of up to 75 years cannot take every segment container. A caller first claims the run for a two-minute lease, so two callers never start the same year; it releases the claim with the run's next safety-check time, and tries again if a year finished while it held the claim. The next years start when a reported call frees a year's place (start_years_after_call) and on each safety-check pass (resume_waiting_work), so a run whose caller stopped carries on. The first failed year fails the run with that year's error, and a year that cannot start fails it as invalid input; years already running keep their results. Once every year has succeeded the run has succeeded. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_reports.py | 538 ++++++++++++++++++ .../economy_routes.py | 6 +- .../tests/budget_window_fakes.py | 221 ++++++- .../tests/test_budget_window_reports.py | 332 +++++++++++ 4 files changed, 1092 insertions(+), 5 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py create mode 100644 projects/policyengine-simulation-entry/tests/test_budget_window_reports.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py new file mode 100644 index 000000000..7d8160382 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -0,0 +1,538 @@ +"""Submit, re-run and schedule budget-window reports: one report run per year. + +A budget-window report (``simulation.reports.budget_window_size``) compares +its two simulations in each year of its window. Each calculation of it is a +budget-window run on one catalog version (``simulation.budget_window_runs``), +whose years are report runs that ``EconomyReports.run_year`` starts exactly as +it starts a single-year report's: it reuses each year's simulation runs that +exist (the current-law baseline, most often) unless the run is forced, starts +one Modal call per run it created, and starts the year's comparison once both +runs have succeeded. + +At most ``max_concurrent_years`` of a run's years are calculated at once, so +one window cannot take every Modal container (a US national year is two +simulations of 20 segment calls each). The next years start when a year's +call reports its end (``start_years_after_call``) and on each safety-check +pass (``resume_waiting_work``). Whoever starts years first claims the run, a +lease of ``SCHEDULING_LEASE``, so two callers do not start the same year; +the database allows one report run per year of a run in any case. + +The first year that fails fails the run, with that year's error, and no more +years start; years already running finish and keep their results. Once every +year has succeeded, the run has succeeded. The run stores no totals: they are +the sums of its years' budget rows, added up when it is read. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass, field +from datetime import timedelta +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_persistence import ( + BudgetWindowStatus, + NewBudgetWindowRun, + ReportKind, + RunError, + RunErrorCode, + RunTrigger, +) +from policyengine_simulation_persistence.run_store import utc_now +from sqlalchemy.exc import IntegrityError + +from policyengine_simulation_entry.economy_report_types import ( + EconomyRequestError, + ReportNotFoundError, +) +from policyengine_simulation_entry.economy_routes import report_run_progress +from policyengine_simulation_entry.economy_schemas import ReportStatus +from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL + +if TYPE_CHECKING: + from collections.abc import Callable, Collection + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_persistence import ( + BudgetWindowRunRecord, + BudgetWindowYear, + CatalogVersion, + ReportRecord, + ReportRunRecord, + StartedBudgetWindowRun, + YearBudget, + ) + + from policyengine_simulation_entry.economy_report_types import ( + EconomyCatalog, + EconomyReportSubmission, + ReportDefinitions, + ) + from policyengine_simulation_entry.economy_reports import EconomyReports + +logger = logging.getLogger(__name__) + +# How many of a run's years are calculated at once. A US national year is up +# to two simulation calls of 20 segment calls each, so five years use at most +# 200 of the US segment function's containers. +DEFAULT_MAX_CONCURRENT_YEARS = 5 +# A caller starting a run's years holds the run this long; one that stopped +# loses the claim afterwards, and the safety check carries on. +SCHEDULING_LEASE = timedelta(minutes=2) +# A running run is checked by the safety check this long after it was last +# scheduled: its years' calls report their ends, so this only catches a run +# whose next years a stopped caller never started. +BUDGET_WINDOW_CHECK_INTERVAL = SAFETY_CHECK_INTERVAL +RERUN_TRIGGERS = frozenset({RunTrigger.MANUAL, RunTrigger.SYSTEM}) +WAITING_WORK_BATCH_SIZE = 50 +# A caller tries again after releasing its claim, in case a year finished +# while it held the claim (its own read was then too early); at most this +# many times. +SCHEDULING_ATTEMPTS = 3 + + +class BudgetWindowRuns(Protocol): + """``BudgetWindowRunStore``'s writes and reads the scheduler needs.""" + + def create_run( + self, new_run: NewBudgetWindowRun, *, first_check_at: datetime + ) -> BudgetWindowRunRecord: + """Add a running budget-window run.""" + ... + + def create_first_run( + self, new_run: NewBudgetWindowRun, *, first_check_at: datetime + ) -> StartedBudgetWindowRun: + """Add the report's first budget-window run, or return its newest.""" + ... + + def get_run(self, run_id: UUID) -> BudgetWindowRunRecord | None: + """Return the run, or None.""" + ... + + def latest_run(self, report_id: UUID) -> BudgetWindowRunRecord | None: + """Return the report's newest run, or None.""" + ... + + def runs(self, report_id: UUID) -> list[BudgetWindowRunRecord]: + """Return the report's runs, newest first.""" + ... + + def due_runs(self, *, now: datetime, limit: int) -> list[BudgetWindowRunRecord]: + """Return running runs whose next check is due.""" + ... + + def years(self, run_id: UUID) -> list[BudgetWindowYear]: + """Return the years the run has started, with their runs.""" + ... + + def year_budgets(self, run_id: UUID) -> list[YearBudget]: + """Return each stored year's budget figures, in year order.""" + ... + + def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: + """Claim a running run to start its years.""" + ... + + def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: + """Release the claim and set the next check.""" + ... + + def finish_run( + self, + run_id: UUID, + status: BudgetWindowStatus, + error: RunError | None = None, + ) -> bool: + """Mark a running run succeeded or failed.""" + ... + + +class ReportRunLookups(Protocol): + """``RunStore``'s reads of the report runs a reported call finished.""" + + def get_report_run(self, report_run_id: UUID) -> ReportRunRecord | None: + """Return the report run, or None.""" + ... + + def report_runs_using(self, simulation_run_id: UUID) -> list[ReportRunRecord]: + """Return the report runs comparing a simulation run.""" + ... + + +@dataclass(frozen=True) +class BudgetWindowReportSubmission: + """A new budget-window report: an economy report's fields and a window. + + ``report.year`` is the window's first year. + """ + + report: EconomyReportSubmission + window_size: int + + +@dataclass(frozen=True) +class _WindowState: + """What a run's started years say about it.""" + + failed: RunError | None + succeeded: bool + in_flight: int + started_years: frozenset[int] + + +def _run_error_code(code: str) -> RunErrorCode: + try: + return RunErrorCode(code) + except ValueError: # pragma: no cover - sim writes only these codes + return RunErrorCode.WORKER_ERROR + + +def window_state(report: ReportRecord, years: list[BudgetWindowYear]) -> _WindowState: + """Return whether a year failed, whether every year succeeded, and how many run. + + A year is failed, succeeded or running as its report run is + (``report_run_progress``); the first failed year, in year order, is the + run's error. + """ + failed: RunError | None = None + succeeded_years = 0 + in_flight = 0 + for year in years: + status, _, error = report_run_progress( + year.report_run, (year.baseline_run, year.reform_run) + ) + if status is ReportStatus.FAILED and failed is None: + detail = "" if error is None or not error.message else f": {error.message}" + failed = RunError( + _run_error_code(error.code if error is not None else ""), + f"The {year.report_run.year} calculation failed{detail}", + ) + elif status is ReportStatus.SUCCEEDED: + succeeded_years += 1 + elif status is not ReportStatus.FAILED: + in_flight += 1 + return _WindowState( + failed=failed, + succeeded=succeeded_years == len(report.window_years), + in_flight=in_flight, + started_years=frozenset(year.report_run.year for year in years), + ) + + +@dataclass +class BudgetWindowReports: + """Budget-window reports: submission, re-runs and their years' scheduling.""" + + economy: EconomyReports + reports: ReportDefinitions + windows: BudgetWindowRuns + runs: ReportRunLookups + catalog: EconomyCatalog + max_concurrent_years: int = DEFAULT_MAX_CONCURRENT_YEARS + clock: Callable[[], datetime] = field(default=utc_now) + + # --- Reports ----------------------------------------------------------- + + async def submit( + self, submission: BudgetWindowReportSubmission + ) -> tuple[ReportRecord, BudgetWindowRunRecord]: + """Create a budget-window report and its first run, and start its years. + + Raise EconomyRequestError, before anything is written, as + ``EconomyReports.prepare_report`` does. + """ + setup = await self.economy.prepare_report( + submission.report, budget_window_size=submission.window_size + ) + run = await asyncio.to_thread( + lambda: self.windows.create_run( + self._new_run(setup.report, setup.version, RunTrigger.INITIAL), + first_check_at=self._next_check(), + ) + ) + return setup.report, await self._started(run) + + async def run_report( + self, + report_id: UUID, + *, + policyengine_version: str, + trigger: RunTrigger, + force: bool, + ) -> BudgetWindowRunRecord: + """Calculate a budget-window report again on a catalog version. + + ``trigger`` is ``manual`` for a user's re-run and ``system`` for a + run the caller starts itself; a ``system`` run returns the report's + newest run when it has one, and simultaneous ones of a report + without one add a single run. Each year's simulation runs are + reused unless ``force`` is set. Raise ReportNotFoundError for a + report that is not a budget-window report, and EconomyRequestError + for one that cannot run. + """ + if trigger not in RERUN_TRIGGERS: + msg = "a report is run again with the manual or system trigger" + raise EconomyRequestError(msg) + report = await asyncio.to_thread(self.reports.get, report_id) + if ( + report is None + or report.kind is not ReportKind.ECONOMY + or report.budget_window_size is None + ): + msg = f"no budget-window report {report_id}" + raise ReportNotFoundError(msg) + if report.options: + msg = "economy reports take no options yet" + raise EconomyRequestError(msg) + version = await self.economy.catalog_version( + report.country, policyengine_version + ) + new_run = self._new_run(report, version, trigger, force=force) + if trigger is RunTrigger.SYSTEM: + first = await asyncio.to_thread( + lambda: self.windows.create_first_run( + new_run, first_check_at=self._next_check() + ) + ) + if not first.created: + return first.run + return await self._started(first.run) + run = await asyncio.to_thread( + lambda: self.windows.create_run(new_run, first_check_at=self._next_check()) + ) + return await self._started(run) + + def _new_run( + self, + report: ReportRecord, + version: CatalogVersion, + trigger: RunTrigger, + *, + force: bool = False, + ) -> NewBudgetWindowRun: + return NewBudgetWindowRun( + report_id=report.id, + tax_benefit_model_version_id=version.id, + trigger=trigger, + max_concurrent_years=self.max_concurrent_years, + force=force, + ) + + async def _started(self, run: BudgetWindowRunRecord) -> BudgetWindowRunRecord: + """Start the new run's first years; return it as it then stands.""" + await self.start_years(run.id) + current = await asyncio.to_thread(self.windows.get_run, run.id) + return run if current is None else current + + def _next_check(self) -> datetime: + return self.clock() + BUDGET_WINDOW_CHECK_INTERVAL + + # --- Scheduling -------------------------------------------------------- + + async def start_years(self, budget_window_run_id: UUID) -> int: + """Start the run's next years up to its cap; return how many started. + + A run whose year failed is failed, and one whose years all succeeded + is succeeded, instead. A run another caller has claimed is left to + it. The claim is released, with the run's next check, however + starting ends. + """ + started = 0 + for _ in range(SCHEDULING_ATTEMPTS): + claimed = await asyncio.to_thread( + lambda: self.windows.claim_scheduling( + budget_window_run_id, lease=SCHEDULING_LEASE + ) + ) + if not claimed: + break + try: + started += await self._schedule(budget_window_run_id) + finally: + await asyncio.to_thread( + lambda: self.windows.release_scheduling( + budget_window_run_id, next_check_at=self._next_check() + ) + ) + if not await self._has_room(budget_window_run_id): + break + return started + + async def _has_room(self, budget_window_run_id: UUID) -> bool: + """Whether a running run could start a year now, or should finish. + + A year that finished while this caller held the claim, after it read + the run's years, leaves room it did not use. + """ + run = await asyncio.to_thread(self.windows.get_run, budget_window_run_id) + if run is None or run.finished: + return False + report = await asyncio.to_thread(self.reports.get, run.report_id) + if report is None: # pragma: no cover - a foreign key keeps it + return False + state = window_state( + report, await asyncio.to_thread(self.windows.years, run.id) + ) + queued = len(report.window_years) - len(state.started_years) + return ( + state.failed is not None + or state.succeeded + or (queued > 0 and state.in_flight < run.max_concurrent_years) + ) + + async def _schedule(self, budget_window_run_id: UUID) -> int: + """Finish the claimed run, or start its next years; return how many started.""" + run = await asyncio.to_thread(self.windows.get_run, budget_window_run_id) + if run is None or run.finished: + return 0 + report = await asyncio.to_thread(self.reports.get, run.report_id) + if report is None: + msg = f"budget-window run {run.id} has no report" + raise LookupError(msg) + state = window_state( + report, await asyncio.to_thread(self.windows.years, run.id) + ) + if state.failed is not None: + await self._finish(run, BudgetWindowStatus.FAILED, state.failed) + return 0 + if state.succeeded: + await self._finish(run, BudgetWindowStatus.SUCCEEDED, None) + return 0 + room = run.max_concurrent_years - state.in_flight + queued = [ + year for year in report.window_years if year not in state.started_years + ] + if room <= 0 or not queued: + return 0 + return await self._start_queued(run, report, queued[:room]) + + async def _start_queued( + self, run: BudgetWindowRunRecord, report: ReportRecord, years: list[int] + ) -> int: + """Start these years of the run in order; return how many started. + + A year that cannot start fails the run, and one another caller + started ends this pass. + """ + version = await asyncio.to_thread( + self.catalog.catalog_version_by_id, run.tax_benefit_model_version_id + ) + if version is None: # pragma: no cover - a foreign key keeps it + msg = f"budget-window run {run.id} has no catalog version" + raise LookupError(msg) + setup = await self.economy.report_setup(report, version) + started = 0 + for year in years: + try: + await self.economy.run_year( + setup, + year=year, + trigger=run.trigger, + force=run.force, + budget_window_run_id=run.id, + ) + except EconomyRequestError as error: + await self._finish( + run, + BudgetWindowStatus.FAILED, + RunError( + RunErrorCode.INVALID_INPUT, + f"The {year} calculation cannot start: {error}", + ), + ) + return started + # Another caller started this year after its claim's lease ended. + except IntegrityError: + logger.warning( + "budget_window_year_already_started", + extra={"budget_window_run_id": str(run.id), "year": year}, + ) + return started + started += 1 + return started + + async def _finish( + self, + run: BudgetWindowRunRecord, + status: BudgetWindowStatus, + error: RunError | None, + ) -> None: + if await asyncio.to_thread(self.windows.finish_run, run.id, status, error): + logger.info( + "budget_window_run_finished", + extra={ + "budget_window_run_id": str(run.id), + "budget_window_status": status.value, + "error_code": None if error is None else error.code.value, + }, + ) + + # --- After a call ended, and the safety check -------------------------- + + async def start_years_after_call( + self, + *, + simulation_run_ids: Collection[UUID] | None, + report_run_id: UUID | None, + ) -> int: + """Start the years that a reported call's end lets start. + + A finished comparison frees its year's place in its budget-window + run; a failed simulation run fails the years comparing it, and so + their runs. Return the years started. + """ + report_runs: list[ReportRunRecord] = [] + if report_run_id is not None: + report_run = await asyncio.to_thread( + self.runs.get_report_run, report_run_id + ) + if report_run is not None: + report_runs.append(report_run) + for simulation_run_id in simulation_run_ids or (): + report_runs.extend( + await asyncio.to_thread(self.runs.report_runs_using, simulation_run_id) + ) + window_ids = sorted( + { + report_run.budget_window_run_id + for report_run in report_runs + if report_run.budget_window_run_id is not None + }, + key=str, + ) + started = 0 + for window_id in window_ids: + started += await self.start_years(window_id) + return started + + async def resume_waiting_work(self, *, now: datetime) -> int: + """Start the years of running runs whose check is due; return how many. + + A run left waiting by a stopped caller, or by a lost report, carries + on. A run that fails to resume is logged and tried again on the next + pass, so it cannot stop the others. + """ + due = await asyncio.to_thread( + lambda: self.windows.due_runs(now=now, limit=WAITING_WORK_BATCH_SIZE) + ) + started = 0 + for run in due: + started += await self._resume(run.id) + return started + + async def _resume(self, budget_window_run_id: UUID) -> int: + """Start one due run's years; a failure is logged and leaves it due.""" + try: + return await self.start_years(budget_window_run_id) + except Exception as error: # noqa: BLE001 # logged; the next pass tries again + logger.warning( + "budget_window_years_not_resumed", + extra={ + "budget_window_run_id": str(budget_window_run_id), + "error_type": type(error).__name__, + }, + ) + return 0 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 c1c6df58f..7ef4b1906 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 @@ -185,7 +185,7 @@ def _simulation_run(record: SimulationRunRecord) -> SimulationRunResponse: ) -def _progress( +def report_run_progress( report_run: ReportRunRecord, runs: tuple[SimulationRunRecord, SimulationRunRecord], ) -> tuple[ReportStatus, JobStage | None, RunErrorResponse | None]: @@ -242,7 +242,9 @@ def report_run(self, report_run: ReportRunRecord) -> EconomyReportRunResponse: version = self.catalog.catalog_version_by_id( report_run.tax_benefit_model_version_id ) - report_status, stage, error = _progress(report_run, (baseline, reform)) + report_status, stage, error = report_run_progress( + report_run, (baseline, reform) + ) return EconomyReportRunResponse( id=report_run.id, report_id=report_run.report_id, diff --git a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py index 4c6a0d9b1..f022fba15 100644 --- a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py +++ b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py @@ -3,9 +3,10 @@ from __future__ import annotations import asyncio -from dataclasses import replace +from dataclasses import dataclass, field, replace +from datetime import timedelta from typing import TYPE_CHECKING, Any -from uuid import uuid4 +from uuid import UUID, uuid4 from economy_fakes import ( CTC_REFORM, @@ -19,16 +20,35 @@ EconomyComparisonCall, EconomySimulationCall, ) +from policyengine_simulation_entry.budget_window_reports import ( + BudgetWindowReports, + BudgetWindowReportSubmission, +) from policyengine_simulation_entry.economy_report_types import ( EconomyReportSubmission, SimulationChoice, ) from policyengine_simulation_entry.economy_reports import EconomyReports from policyengine_simulation_entry.run_dispatch import RunDispatchService +from policyengine_simulation_persistence import ( + BudgetWindowRunRecord, + BudgetWindowStatus, + BudgetWindowYear, + RunStatus, + StartedBudgetWindowRun, + YearBudget, +) from run_fakes import FakeDispatcher if TYPE_CHECKING: - from policyengine_simulation_persistence import NewReportRun, ReportRunRecord + from datetime import datetime + + from policyengine_simulation_persistence import ( + NewBudgetWindowRun, + NewReportRun, + ReportRunRecord, + RunError, + ) REFORM_POLICY = uuid4() LABOR_SUPPLY_POLICY = uuid4() @@ -46,6 +66,166 @@ def create_report_run(self, new_report_run: NewReportRun) -> ReportRunRecord: return row +@dataclass +class FakeBudgetWindows: + """``BudgetWindowRunStore`` in memory, reading years from ``WindowRuns``.""" + + economy_runs: WindowRuns + rows: dict[UUID, BudgetWindowRunRecord] = field(default_factory=dict) + budgets: dict[UUID, YearBudget] = field(default_factory=dict) + claims_refused: bool = False + + def _insert(self, new_run: NewBudgetWindowRun, first_check_at: datetime): + now = self.economy_runs.clock() + row = BudgetWindowRunRecord( + id=uuid4(), + report_id=new_run.report_id, + tax_benefit_model_version_id=new_run.tax_benefit_model_version_id, + trigger=new_run.trigger, + force=new_run.force, + max_concurrent_years=new_run.max_concurrent_years, + status=BudgetWindowStatus.RUNNING, + error_code=None, + error_message=None, + next_check_at=first_check_at, + scheduling_claimed_until=None, + completed_at=None, + created_at=now + timedelta(microseconds=len(self.rows)), + updated_at=now, + ) + self.rows[row.id] = row + return row + + def create_run(self, new_run: NewBudgetWindowRun, *, first_check_at: datetime): + return self._insert(new_run, first_check_at) + + def create_first_run( + self, new_run: NewBudgetWindowRun, *, first_check_at: datetime + ) -> StartedBudgetWindowRun: + newest = self.latest_run(new_run.report_id) + if newest is not None: + return StartedBudgetWindowRun(run=newest, created=False) + return StartedBudgetWindowRun( + run=self._insert(new_run, first_check_at), created=True + ) + + def get_run(self, run_id: UUID) -> BudgetWindowRunRecord | None: + return self.rows.get(run_id) + + def runs(self, report_id: UUID) -> list[BudgetWindowRunRecord]: + return sorted( + (row for row in self.rows.values() if row.report_id == report_id), + key=lambda row: row.created_at, + reverse=True, + ) + + def latest_run(self, report_id: UUID) -> BudgetWindowRunRecord | None: + newest = self.runs(report_id) + return newest[0] if newest else None + + def due_runs(self, *, now: datetime, limit: int) -> list[BudgetWindowRunRecord]: + due = sorted( + ( + row + for row in self.rows.values() + if row.status is BudgetWindowStatus.RUNNING + and row.next_check_at is not None + and row.next_check_at <= now + ), + key=lambda row: (row.next_check_at, row.created_at), + ) + return due[:limit] + + def years(self, run_id: UUID) -> list[BudgetWindowYear]: + year_runs = sorted( + ( + report_run + for report_run in self.economy_runs.report_runs_by_id.values() + if report_run.budget_window_run_id == run_id + ), + key=lambda report_run: report_run.year, + ) + simulation_runs = self.economy_runs.simulation_runs + return [ + BudgetWindowYear( + report_run=report_run, + baseline_run=simulation_runs[report_run.baseline_simulation_run_id], + reform_run=simulation_runs[ + report_run.reform_simulation_run_id + or report_run.baseline_simulation_run_id + ], + ) + for report_run in year_runs + ] + + def year_budgets(self, run_id: UUID) -> list[YearBudget]: + return [ + self.budgets[year.report_run.id] + for year in self.years(run_id) + if year.report_run.id in self.budgets + ] + + def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: + row = self.rows[run_id] + now = self.economy_runs.clock() + if ( + self.claims_refused + or row.status is not BudgetWindowStatus.RUNNING + or ( + row.scheduling_claimed_until is not None + and row.scheduling_claimed_until > now + ) + ): + return False + self.rows[run_id] = replace(row, scheduling_claimed_until=now + lease) + return True + + def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: + row = self.rows[run_id] + if row.status is BudgetWindowStatus.RUNNING: + self.rows[run_id] = replace( + row, scheduling_claimed_until=None, next_check_at=next_check_at + ) + + def finish_run( + self, + run_id: UUID, + status: BudgetWindowStatus, + error: RunError | None = None, + ) -> bool: + row = self.rows[run_id] + if row.status is not BudgetWindowStatus.RUNNING: + return False + self.rows[run_id] = replace( + row, + status=status, + error_code=None if error is None else error.code.value, + error_message=None if error is None else error.message, + completed_at=self.economy_runs.clock(), + next_check_at=None, + scheduling_claimed_until=None, + ) + return True + + +def finish_year(runs: WindowRuns, report_run: ReportRunRecord) -> None: + """Succeed a year's two simulation runs and its comparison.""" + for run_id in ( + report_run.baseline_simulation_run_id, + report_run.reform_simulation_run_id, + ): + if run_id is not None and ( + runs.simulation_runs[run_id].status is not RunStatus.SUCCEEDED + ): + runs.succeed(run_id) + runs.report_runs_by_id[report_run.id] = replace( + runs.report_runs_by_id[report_run.id], + comparison_status=RunStatus.SUCCEEDED, + comparison_modal_call_id=f"fc-comparison-{report_run.id.hex[:6]}", + comparison_completed_at=runs.clock(), + ) + + class EconomyPath: """The economy path over in-memory stores, with every spawned request kept.""" @@ -110,3 +290,38 @@ def comparison_calls(self) -> list[EconomyComparisonCall]: for request in self.dispatcher.requests if request.target.function == "compare_economy_simulations" ] + + +class WindowPath(EconomyPath): + """The budget-window path over the economy path's in-memory stores.""" + + def __init__(self, *, max_concurrent_years: int = 2) -> None: + super().__init__() + self.windows = FakeBudgetWindows(self.runs) + self.budget_windows = BudgetWindowReports( + economy=self.economy, + reports=self.reports, + windows=self.windows, + runs=self.runs, + catalog=self.catalog, + max_concurrent_years=max_concurrent_years, + clock=self.runs.clock, + ) + + def submit_window(self, *, window_size: int = 4, **changes: object): + """Submit a budget window of the CTC reform, from 2026 by default.""" + return asyncio.run( + self.budget_windows.submit( + BudgetWindowReportSubmission( + report=self.submission(**changes), window_size=window_size + ) + ) + ) + + def year_runs(self, budget_window_run_id: UUID) -> dict[int, ReportRunRecord]: + """Return the run's report runs by year.""" + return { + report_run.year: report_run + for report_run in self.runs.report_runs_by_id.values() + if report_run.budget_window_run_id == budget_window_run_id + } diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py new file mode 100644 index 000000000..16c15d83f --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py @@ -0,0 +1,332 @@ +"""Budget-window reports: one report run per year, at most a few years at once.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from datetime import timedelta +from uuid import uuid4 + +import pytest +from budget_window_fakes import ( + LABOR_SUPPLY_POLICY, + WindowPath, + finish_year, +) +from economy_fakes import NOW, US_VERSION +from policyengine_simulation_entry.budget_window_reports import ( + BUDGET_WINDOW_CHECK_INTERVAL, + SCHEDULING_LEASE, +) +from policyengine_simulation_entry.economy_report_types import ( + EconomyRequestError, + ReportNotFoundError, + SimulationChoice, +) +from policyengine_simulation_persistence import ( + BudgetWindowStatus, + NewBudgetWindowRun, + RunStatus, + RunTrigger, +) + + +def _rerun(path: WindowPath, report_id, **changes: object): + values: dict[str, object] = { + "policyengine_version": "6.7.5", + "trigger": RunTrigger.MANUAL, + "force": False, + } + values.update(changes) + return asyncio.run(path.budget_windows.run_report(report_id, **values)) # type: ignore[arg-type] + + +def _start_years_after(path: WindowPath, **reported: object): + values: dict[str, object] = {"simulation_run_ids": None, "report_run_id": None} + values.update(reported) + return asyncio.run(path.budget_windows.start_years_after_call(**values)) # type: ignore[arg-type] + + +def test_a_submission_starts_its_first_years_up_to_the_cap() -> None: + path = WindowPath(max_concurrent_years=2) + + report, run = path.submit_window(window_size=4) + + assert report.year == 2026 + assert report.budget_window_size == 4 + assert run.report_id == report.id + assert run.trigger is RunTrigger.INITIAL + assert run.tax_benefit_model_version_id == US_VERSION.id + assert run.max_concurrent_years == 2 + assert run.status is BudgetWindowStatus.RUNNING + assert sorted(path.year_runs(run.id)) == [2026, 2027] + assert {call.year for call in path.simulation_calls()} == {2026, 2027} + assert run.next_check_at == NOW + BUDGET_WINDOW_CHECK_INTERVAL + assert run.scheduling_claimed_until is None + + +def test_a_years_recorded_comparison_starts_the_next_year() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + first_year = path.year_runs(run.id)[2026] + finish_year(path.runs, first_year) + + started = _start_years_after(path, report_run_id=first_year.id) + + assert started == 1 + assert sorted(path.year_runs(run.id)) == [2026, 2027, 2028] + assert {call.year for call in path.simulation_calls()} == {2026, 2027, 2028} + + +def test_a_year_whose_comparison_is_still_running_keeps_its_place() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + first_year = path.year_runs(run.id)[2026] + for run_id in ( + first_year.baseline_simulation_run_id, + first_year.reform_simulation_run_id, + ): + path.runs.succeed(run_id) + + started = _start_years_after( + path, + simulation_run_ids=[ + first_year.baseline_simulation_run_id, + first_year.reform_simulation_run_id, + ], + ) + + assert started == 0 + assert sorted(path.year_runs(run.id)) == [2026, 2027] + + +def test_every_year_succeeding_finishes_the_run() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=2) + for year_run in path.year_runs(run.id).values(): + finish_year(path.runs, year_run) + + _start_years_after(path, report_run_id=path.year_runs(run.id)[2027].id) + + finished = path.windows.get_run(run.id) + assert finished is not None + assert finished.status is BudgetWindowStatus.SUCCEEDED + assert finished.completed_at == NOW + + +def test_a_failed_year_fails_the_run_and_starts_no_more_years() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + failed_year = path.year_runs(run.id)[2027] + reform_run_id = failed_year.reform_simulation_run_id + assert reform_run_id is not None + path.runs.simulation_runs[reform_run_id] = replace( + path.runs.simulation_runs[reform_run_id], + status=RunStatus.FAILED, + error_code="worker_error", + error_message="The Modal function raised RuntimeError.", + completed_at=NOW, + ) + finish_year(path.runs, path.year_runs(run.id)[2026]) + + _start_years_after(path, simulation_run_ids=[reform_run_id]) + + failed = path.windows.get_run(run.id) + assert failed is not None + assert failed.status is BudgetWindowStatus.FAILED + assert failed.error_code == "worker_error" + assert failed.error_message == ( + "The 2027 calculation failed: The Modal function raised RuntimeError." + ) + assert sorted(path.year_runs(run.id)) == [2026, 2027] + + +def test_a_year_that_cannot_start_fails_the_run() -> None: + path = WindowPath(max_concurrent_years=1) + _, run = path.submit_window(window_size=3) + finish_year(path.runs, path.year_runs(run.id)[2026]) + del path.catalog.populations["state/ca"] + + _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) + + failed = path.windows.get_run(run.id) + assert failed is not None + assert failed.status is BudgetWindowStatus.FAILED + assert failed.error_code == "invalid_input" + assert failed.error_message is not None + assert failed.error_message.startswith("The 2027 calculation cannot start:") + assert sorted(path.year_runs(run.id)) == [2026] + + +def test_a_run_another_caller_is_scheduling_is_left_to_it() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + finish_year(path.runs, path.year_runs(run.id)[2026]) + path.windows.claims_refused = True + + started = _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) + + assert started == 0 + assert sorted(path.year_runs(run.id)) == [2026, 2027] + + +def test_a_claim_is_released_with_the_next_check_even_when_starting_fails() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + finish_year(path.runs, path.year_runs(run.id)[2026]) + + def refuse_to_start(*_args, **_kwargs): + msg = "the database dropped the connection" + raise ConnectionError(msg) + + path.runs.start_simulation_run = refuse_to_start # type: ignore[method-assign] + + with pytest.raises(ConnectionError): + _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) + + released = path.windows.get_run(run.id) + assert released is not None + assert released.status is BudgetWindowStatus.RUNNING + assert released.scheduling_claimed_until is None + assert released.next_check_at == NOW + BUDGET_WINDOW_CHECK_INTERVAL + + +def test_a_claim_lasts_its_lease() -> None: + assert timedelta(0) < SCHEDULING_LEASE <= BUDGET_WINDOW_CHECK_INTERVAL + + +def test_a_rerun_reuses_each_years_runs_unless_forced() -> None: + path = WindowPath(max_concurrent_years=2) + report, first = path.submit_window(window_size=2) + calls_before = len(path.simulation_calls()) + + rerun = _rerun(path, report.id) + forced = _rerun(path, report.id, force=True) + + assert rerun.id != first.id + assert rerun.trigger is RunTrigger.MANUAL + assert forced.force is True + first_years, rerun_years, forced_years = ( + path.year_runs(run.id) for run in (first, rerun, forced) + ) + assert { + year: run.baseline_simulation_run_id for year, run in rerun_years.items() + } == {year: run.baseline_simulation_run_id for year, run in first_years.items()} + assert all( + forced_years[year].baseline_simulation_run_id + != first_years[year].baseline_simulation_run_id + for year in (2026, 2027) + ) + # Only the forced run's years started calls: two years of two simulations. + assert len(path.simulation_calls()) == calls_before + 4 + assert path.windows.latest_run(report.id) == path.windows.get_run(forced.id) + + +def test_a_system_run_of_a_report_that_has_a_run_returns_its_newest() -> None: + path = WindowPath() + report, first = path.submit_window(window_size=2) + path.dispatcher.requests.clear() + + again = _rerun(path, report.id, trigger=RunTrigger.SYSTEM) + + assert again.id == first.id + assert path.dispatcher.requests == [] + + +def test_a_system_run_of_a_report_without_one_calculates_it() -> None: + path = WindowPath() + setup = path.prepare(window_size=2) + + run = _rerun(path, setup.report.id, trigger=RunTrigger.SYSTEM) + + assert run.trigger is RunTrigger.SYSTEM + assert sorted(path.year_runs(run.id)) == [2026, 2027] + assert {report_run.trigger for report_run in path.year_runs(run.id).values()} == { + RunTrigger.SYSTEM + } + + +def test_a_rerun_needs_a_budget_window_report_and_a_rerun_trigger() -> None: + path = WindowPath() + report, _ = asyncio.run(path.economy.submit(path.submission())) + window_report, _ = path.submit_window(window_size=2) + + with pytest.raises(ReportNotFoundError, match="budget-window report"): + _rerun(path, report.id) + with pytest.raises(ReportNotFoundError, match="budget-window report"): + _rerun(path, uuid4()) + with pytest.raises(EconomyRequestError, match="manual or system"): + _rerun(path, window_report.id, trigger=RunTrigger.INITIAL) + with pytest.raises(EconomyRequestError, match="catalog"): + _rerun(path, window_report.id, policyengine_version="9.9.9") + + +def test_a_submission_that_cannot_run_is_refused_before_anything_starts() -> None: + path = WindowPath() + + with pytest.raises(EconomyRequestError, match="region"): + path.submit_window(region_code="state/zz") + + assert path.reports.rows == {} + assert path.windows.rows == {} + assert path.dispatcher.requests == [] + + +def test_a_labour_supply_window_calculates_every_year_with_its_columns() -> None: + path = WindowPath(max_concurrent_years=3) + + path.submit_window( + window_size=3, reform=SimulationChoice(policy_id=LABOR_SUPPLY_POLICY) + ) + + calls = path.simulation_calls() + assert {call.year for call in calls} == {2026, 2027, 2028} + assert {call.labor_supply_response for call in calls} == {True} + + +def test_the_safety_check_starts_the_years_of_runs_left_waiting() -> None: + path = WindowPath(max_concurrent_years=2) + setup = path.prepare(window_size=3) + stranded = path.windows.create_run( + NewBudgetWindowRun( + report_id=setup.report.id, + tax_benefit_model_version_id=US_VERSION.id, + trigger=RunTrigger.INITIAL, + max_concurrent_years=2, + ), + first_check_at=NOW + BUDGET_WINDOW_CHECK_INTERVAL, + ) + + early = asyncio.run(path.budget_windows.resume_waiting_work(now=NOW)) + started = asyncio.run( + path.budget_windows.resume_waiting_work(now=NOW + BUDGET_WINDOW_CHECK_INTERVAL) + ) + + assert early == 0 + assert started == 2 + assert sorted(path.year_runs(stranded.id)) == [2026, 2027] + + +def test_one_run_that_cannot_resume_does_not_stop_the_others() -> None: + path = WindowPath(max_concurrent_years=1) + broken_setup = path.prepare(window_size=2) + working_setup = path.prepare(window_size=2, region_code="us") + new_runs = [ + path.windows.create_run( + NewBudgetWindowRun( + report_id=setup.report.id, + tax_benefit_model_version_id=US_VERSION.id, + trigger=RunTrigger.INITIAL, + max_concurrent_years=1, + ), + first_check_at=NOW, + ) + for setup in (broken_setup, working_setup) + ] + del path.reports.rows[broken_setup.report.id] + + started = asyncio.run(path.budget_windows.resume_waiting_work(now=NOW)) + + assert started == 1 + assert path.year_runs(new_runs[0].id) == {} + assert sorted(path.year_runs(new_runs[1].id)) == [2026] From f11b4ac402ada1fba8efa3334ab94c519b904ed9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:13:10 +0400 Subject: [PATCH 04/20] Read budget-window reports with each year's progress and totals summed on read A budget-window report reads with its newest run, and a run with every year of its window: a started year has its report run's status, stage and error; a year not started yet is queued, and one a finished run never started was not run. A run has failed as soon as one year has, before the scheduler records it, and has succeeded once every year has; progress is the share of years that succeeded. Once it has succeeded, its result is today's budget-window format, built from each year's stored budget row, with totals summed in decimal arithmetic and returned as floats. A succeeded run missing a year's figures is a stored-result error. The request and response models of the budget-window routes are added; a window has 1 to 75 years and ends in 2099 or earlier. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_reads.py | 299 +++++++++++++++++ .../budget_window_schemas.py | 125 +++++++ .../economy_routes.py | 9 +- .../tests/test_budget_window_reads.py | 304 ++++++++++++++++++ 4 files changed, 733 insertions(+), 4 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py create mode 100644 projects/policyengine-simulation-entry/tests/test_budget_window_reads.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py new file mode 100644 index 000000000..6fb83bc64 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py @@ -0,0 +1,299 @@ +"""Read budget-window reports: each year's progress, and the totals, summed on read. + +A budget-window report reads with its newest run. A run reads with each year +of its window: the years it started have their report runs' status, stage and +error; the others are queued, or never run once the run has finished. A run +has failed as soon as one year has, before the scheduler records it, and has +succeeded once every year has. Once it has succeeded, its result is today's +budget-window format (``BudgetWindowResult``), built from each year's stored +budget row. The totals are summed in decimal arithmetic, so a ten-year sum +of large figures carries no binary rounding error, then returned as floats, +as the old gateway's budget windows are. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from decimal import Decimal +from typing import TYPE_CHECKING, Protocol + +from policyengine_simulation_contract.gateway_models import ( + BudgetWindowAnnualImpact, + BudgetWindowResult, + BudgetWindowTotals, +) +from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_persistence import BudgetWindowStatus, ReportKind + +from policyengine_simulation_entry.budget_window_reports import window_state +from policyengine_simulation_entry.budget_window_schemas import ( + BudgetWindowReportResponse, + BudgetWindowRunResponse, + BudgetWindowRunsResponse, + BudgetWindowYearResponse, +) +from policyengine_simulation_entry.economy_report_types import ReportNotFoundError +from policyengine_simulation_entry.economy_routes import ( + StoredResultError, + report_run_progress, + simulation_response, +) +from policyengine_simulation_entry.economy_schemas import ( + ReportStatus, + RunErrorResponse, +) + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_persistence import ( + BudgetWindowRunRecord, + BudgetWindowYear, + ReportRecord, + YearBudget, + ) + + from policyengine_simulation_entry.economy_report_types import ( + EconomyCatalog, + ReportDefinitions, + SimulationDefinitions, + ) + +_BUDGET_FIELDS = ( + "taxRevenueImpact", + "federalTaxRevenueImpact", + "stateTaxRevenueImpact", + "benefitSpendingImpact", + "budgetaryImpact", +) + + +class BudgetWindowRunReads(Protocol): + """``BudgetWindowRunStore``'s reads.""" + + def latest_run(self, report_id: UUID) -> BudgetWindowRunRecord | None: + """Return the report's newest run, or None.""" + ... + + def runs(self, report_id: UUID) -> list[BudgetWindowRunRecord]: + """Return the report's runs, newest first.""" + ... + + def years(self, run_id: UUID) -> list[BudgetWindowYear]: + """Return the years the run has started, with their runs.""" + ... + + def year_budgets(self, run_id: UUID) -> list[YearBudget]: + """Return each stored year's budget figures, in year order.""" + ... + + +def _annual_impact(budget: YearBudget) -> BudgetWindowAnnualImpact: + return BudgetWindowAnnualImpact( + year=str(budget.year), + taxRevenueImpact=budget.tax_revenue_impact, + federalTaxRevenueImpact=( + budget.tax_revenue_impact - budget.state_tax_revenue_impact + ), + stateTaxRevenueImpact=budget.state_tax_revenue_impact, + benefitSpendingImpact=budget.benefit_spending_impact, + budgetaryImpact=budget.budgetary_impact, + ) + + +def _totals(annual_impacts: list[BudgetWindowAnnualImpact]) -> BudgetWindowTotals: + """Sum the years' figures in decimal arithmetic, from each float's shortest form.""" + sums = { + name: sum( + (Decimal(str(getattr(impact, name))) for impact in annual_impacts), + Decimal(0), + ) + for name in _BUDGET_FIELDS + } + return BudgetWindowTotals( + taxRevenueImpact=float(sums["taxRevenueImpact"]), + federalTaxRevenueImpact=float(sums["federalTaxRevenueImpact"]), + stateTaxRevenueImpact=float(sums["stateTaxRevenueImpact"]), + benefitSpendingImpact=float(sums["benefitSpendingImpact"]), + budgetaryImpact=float(sums["budgetaryImpact"]), + ) + + +def budget_window_result( + report: ReportRecord, budgets: list[YearBudget] +) -> BudgetWindowResult: + """Return the window's result from each year's budget figures. + + Raise LookupError when a year of the window has no stored figures. + """ + years = list(report.window_years) + stored = [budget.year for budget in budgets] + if stored != years: + missing = sorted(set(years) - set(stored)) + msg = f"years {missing} of the window have no stored budget figures" + raise LookupError(msg) + annual_impacts = [_annual_impact(budget) for budget in budgets] + return BudgetWindowResult( + startYear=str(years[0]), + endYear=str(years[-1]), + windowSize=len(years), + annualImpacts=annual_impacts, + totals=_totals(annual_impacts), + ) + + +def _error(code: str | None, message: str | None) -> RunErrorResponse | None: + return None if code is None else RunErrorResponse(code=code, message=message) + + +def _run_status( + run: BudgetWindowRunRecord, + report: ReportRecord, + years: list[BudgetWindowYear], +) -> tuple[ReportStatus, RunErrorResponse | None]: + if run.status is BudgetWindowStatus.FAILED: + return ReportStatus.FAILED, _error(run.error_code, run.error_message) + state = window_state(report, years) + if state.failed is not None: + return ReportStatus.FAILED, _error( + state.failed.code.value, state.failed.message + ) + if run.status is BudgetWindowStatus.SUCCEEDED or state.succeeded: + return ReportStatus.SUCCEEDED, None + calculating = any( + report_run_progress(year.report_run, (year.baseline_run, year.reform_run))[0] + is not ReportStatus.PENDING + for year in years + ) + return (ReportStatus.RUNNING if calculating else ReportStatus.PENDING), None + + +@dataclass +class BudgetWindowReportReads: + """Build budget-window report and run responses from the stored rows.""" + + simulations: SimulationDefinitions + reports: ReportDefinitions + windows: BudgetWindowRunReads + catalog: EconomyCatalog + + def run( + self, run: BudgetWindowRunRecord, report: ReportRecord + ) -> BudgetWindowRunResponse: + """Describe one budget-window run of ``report``, with each of its years.""" + years = self.windows.years(run.id) + status, error = _run_status(run, report, years) + started = {year.report_run.year: year for year in years} + version = self.catalog.catalog_version_by_id(run.tax_benefit_model_version_id) + year_responses = [ + self._year(window_year, started.get(window_year), finished=run.finished) + for window_year in report.window_years + ] + succeeded = sum( + year.status is ReportStatus.SUCCEEDED for year in year_responses + ) + return BudgetWindowRunResponse( + id=run.id, + report_id=run.report_id, + trigger=run.trigger.value, + catalog_version_id=run.tax_benefit_model_version_id, + policyengine_version=None + if version is None + else version.policyengine_version, + force=run.force, + max_concurrent_years=run.max_concurrent_years, + status=status, + error=error, + progress=100 * succeeded // len(year_responses), + years=year_responses, + created_at=run.created_at, + completed_at=run.completed_at, + ) + + @staticmethod + def _year( + window_year: int, year: BudgetWindowYear | None, *, finished: bool + ) -> BudgetWindowYearResponse: + if year is None: + return BudgetWindowYearResponse( + year=window_year, + status=ReportStatus.NOT_RUN if finished else ReportStatus.PENDING, + stage=None if finished else JobStage.QUEUED, + error=None, + report_run_id=None, + ) + status, stage, error = report_run_progress( + year.report_run, (year.baseline_run, year.reform_run) + ) + return BudgetWindowYearResponse( + year=window_year, + status=status, + stage=stage, + error=error, + report_run_id=year.report_run.id, + ) + + def report(self, report: ReportRecord) -> BudgetWindowReportResponse: + """Describe a report with its newest run and, once it succeeded, its result. + + Raise StoredResultError when a succeeded run's figures cannot be read. + """ + baseline, reform = ( + self.simulations.get(simulation_id) + for simulation_id in ( + report.baseline_simulation_id, + report.reform_simulation_id or report.baseline_simulation_id, + ) + ) + if baseline is None or reform is None: # pragma: no cover - foreign keys + msg = f"report {report.id} has no simulation" + raise LookupError(msg) + newest = self.windows.latest_run(report.id) + run = None if newest is None else self.run(newest, report) + result = None + if ( + newest is not None + and run is not None + and run.status is ReportStatus.SUCCEEDED + ): + try: + result = budget_window_result( + report, self.windows.year_budgets(newest.id) + ) + except LookupError as error: + raise StoredResultError(newest.id, error) from error + return BudgetWindowReportResponse( + id=report.id, + country=report.country, + start_year=report.year, + window_size=len(report.window_years), + options=report.options, + created_at=report.created_at, + baseline_simulation=simulation_response(baseline), + reform_simulation=simulation_response(reform), + status=ReportStatus.NOT_RUN if run is None else run.status, + run=run, + result=result, + ) + + def report_by_id(self, report_id: UUID) -> BudgetWindowReportResponse: + """Describe a report; raise ReportNotFoundError for one that is not a window.""" + return self.report(self._report(report_id)) + + def report_runs(self, report_id: UUID) -> BudgetWindowRunsResponse: + """List a budget-window report's runs, newest first.""" + report = self._report(report_id) + return BudgetWindowRunsResponse( + runs=[self.run(run, report) for run in self.windows.runs(report_id)] + ) + + def _report(self, report_id: UUID) -> ReportRecord: + record = self.reports.get(report_id) + if ( + record is None + or record.kind is not ReportKind.ECONOMY + or record.budget_window_size is None + ): + msg = f"no budget-window report {report_id}" + raise ReportNotFoundError(msg) + return record diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py new file mode 100644 index 000000000..755c98f26 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py @@ -0,0 +1,125 @@ +"""Request and response bodies of the v2 budget-window report routes. + +A budget-window report's result is today's budget-window format +(``BudgetWindowResult``): each year's budget figures and their totals, which +are summed when the report is read. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Annotated, Self +from uuid import UUID + +from policyengine_simulation_contract.gateway_models import BudgetWindowResult +from policyengine_simulation_contract.job_progress import JobStage +from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator + +from policyengine_simulation_entry.economy_schemas import ( + CatalogKey, + CountryCode, + PolicyEngineVersion, + ReportStatus, + RunErrorResponse, + SimulationChoiceRequest, + SimulationResponse, +) + +# A window may end no later than this year, as the old gateway's budget +# windows may, and have at most this many years. +LAST_WINDOW_YEAR = 2099 +MAX_WINDOW_YEARS = 75 + + +class BudgetWindowReportSubmissionRequest(BaseModel): + """A new budget-window report and its first run. + + The report compares ``baseline`` and ``reform`` in each of + ``window_size`` years from ``start_year``, ending in 2099 or earlier. + The other fields are an economy report's: ``policyengine_version`` + selects the country's catalog version, ``region_code`` and + ``dataset_name`` name its region and dataset, and ``options`` must be + empty for now. + """ + + model_config = ConfigDict(extra="forbid") + + country: CountryCode + policyengine_version: PolicyEngineVersion + start_year: Annotated[int, Field(ge=1900, le=LAST_WINDOW_YEAR)] + window_size: Annotated[int, Field(ge=1, le=MAX_WINDOW_YEARS)] + region_code: CatalogKey + dataset_name: CatalogKey | None = None + baseline: SimulationChoiceRequest + reform: SimulationChoiceRequest + options: dict[str, JsonValue] = Field(default_factory=dict) + + @model_validator(mode="after") + def _window_ends_in_time(self) -> Self: + if self.start_year + self.window_size - 1 > LAST_WINDOW_YEAR: + msg = f"a budget window must end in {LAST_WINDOW_YEAR} or earlier" + raise ValueError(msg) + return self + + +class BudgetWindowYearResponse(BaseModel): + """One year of a budget-window run. + + A year the run has started has its report run, with that run's status, + its stage while it runs (``queued``, ``simulating`` or ``outputs``) and + its error. A year not started yet is ``pending`` and ``queued``; one a + finished run never started is ``not_run``. + """ + + year: int + status: ReportStatus + stage: JobStage | None + error: RunErrorResponse | None + report_run_id: UUID | None + + +class BudgetWindowRunResponse(BaseModel): + """One calculation of a budget-window report on a catalog version. + + ``status`` is ``failed`` as soon as one year has failed, ``succeeded`` + once every year has, ``running`` while any year is calculated, and + ``pending`` before then. ``progress`` is the share of the window's years + that have succeeded, from 0 to 100. At most ``max_concurrent_years`` are + calculated at once; ``force`` calculates every year's simulations again. + """ + + id: UUID + report_id: UUID + trigger: str + catalog_version_id: UUID + policyengine_version: str | None + force: bool + max_concurrent_years: int + status: ReportStatus + error: RunErrorResponse | None + progress: int + years: list[BudgetWindowYearResponse] + created_at: datetime + completed_at: datetime | None + + +class BudgetWindowRunsResponse(BaseModel): + """A budget-window report's runs, newest first.""" + + runs: list[BudgetWindowRunResponse] + + +class BudgetWindowReportResponse(BaseModel): + """A budget-window report, its newest run and, once it succeeded, its result.""" + + id: UUID + country: str + start_year: int + window_size: int + options: dict[str, JsonValue] + created_at: datetime + baseline_simulation: SimulationResponse + reform_simulation: SimulationResponse + status: ReportStatus + run: BudgetWindowRunResponse | None + result: BudgetWindowResult | None 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 7ef4b1906..7156f7d68 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 @@ -157,7 +157,8 @@ def _error(code: str | None, message: str | None) -> RunErrorResponse | None: return None if code is None else RunErrorResponse(code=code, message=message) -def _simulation(record: SimulationRecord) -> SimulationResponse: +def simulation_response(record: SimulationRecord) -> SimulationResponse: + """Describe a simulation.""" return SimulationResponse( id=record.id, kind=record.kind.value, @@ -301,8 +302,8 @@ def report(self, report: ReportRecord) -> EconomyReportResponse: options=report.options, legacy_report_id=report.legacy_report_id, created_at=report.created_at, - baseline_simulation=_simulation(baseline), - reform_simulation=_simulation(reform), + baseline_simulation=simulation_response(baseline), + reform_simulation=simulation_response(reform), status=ReportStatus.NOT_RUN if latest is None else latest.status, run=latest, result=result, @@ -334,7 +335,7 @@ def simulation(self, simulation_id: UUID) -> SimulationResponse: if record is None: msg = f"no simulation {simulation_id}" raise ReportNotFoundError(msg) - return _simulation(record) + return simulation_response(record) @staticmethod def _report(record: ReportRecord | None, identifier: object) -> ReportRecord: diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py new file mode 100644 index 000000000..e38e29d94 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py @@ -0,0 +1,304 @@ +"""Reading budget-window reports: each year's progress, and totals summed on read.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from uuid import uuid4 + +import pytest +from budget_window_fakes import WindowPath, finish_year +from economy_fakes import NOW, US_VERSION +from policyengine_simulation_entry.budget_window_reads import BudgetWindowReportReads +from policyengine_simulation_entry.budget_window_schemas import ( + BudgetWindowReportSubmissionRequest, +) +from policyengine_simulation_entry.economy_report_types import ReportNotFoundError +from policyengine_simulation_entry.economy_routes import StoredResultError +from policyengine_simulation_entry.economy_schemas import ReportStatus +from policyengine_simulation_persistence import ( + BudgetWindowStatus, + NewBudgetWindowRun, + RunError, + RunErrorCode, + RunStatus, + RunTrigger, + YearBudget, +) +from pydantic import ValidationError + + +def _reads(path: WindowPath) -> BudgetWindowReportReads: + return BudgetWindowReportReads( + simulations=path.simulations, + reports=path.reports, + windows=path.windows, + catalog=path.catalog, + ) + + +def _budget(year: int, budgetary_impact: float, **changes: float) -> YearBudget: + values = { + "year": year, + "budgetary_impact": budgetary_impact, + "tax_revenue_impact": 2.0, + "state_tax_revenue_impact": 0.5, + "benefit_spending_impact": 1.0, + } + values.update(changes) + return YearBudget(**values) # type: ignore[arg-type] + + +def _finish_window(path: WindowPath, run_id, budgets: dict[int, float]) -> None: + """Succeed every year with its budgetary impact, and the run itself.""" + while True: + years = path.year_runs(run_id) + unfinished = [ + report_run + for report_run in years.values() + if report_run.comparison_status is not RunStatus.SUCCEEDED + ] + if not unfinished: + break + for report_run in unfinished: + finish_year(path.runs, report_run) + path.windows.budgets[report_run.id] = _budget( + report_run.year, budgets[report_run.year] + ) + asyncio.run( + path.budget_windows.start_years_after_call( + simulation_run_ids=None, report_run_id=report_run.id + ) + ) + + +def test_a_new_window_reads_its_started_and_queued_years() -> None: + path = WindowPath(max_concurrent_years=2) + report, run = path.submit_window(window_size=3) + + read = _reads(path).report_by_id(report.id) + + assert read.start_year == 2026 + assert read.window_size == 3 + # The first two years' calls have started; the third waits its turn. + assert read.status is ReportStatus.RUNNING + assert read.result is None + assert read.run is not None + assert read.run.id == run.id + assert read.run.catalog_version_id == US_VERSION.id + assert read.run.policyengine_version == "6.7.5" + assert read.run.max_concurrent_years == 2 + assert read.run.progress == 0 + started = path.year_runs(run.id) + assert [ + (year.year, year.status, year.stage, year.report_run_id) + for year in read.run.years + ] == [ + (2026, ReportStatus.RUNNING, "simulating", started[2026].id), + (2027, ReportStatus.RUNNING, "simulating", started[2027].id), + (2028, ReportStatus.PENDING, "queued", None), + ] + + +def test_a_window_whose_years_have_not_started_calling_is_pending() -> None: + path = WindowPath(max_concurrent_years=2) + setup = path.prepare(window_size=2) + # The run a submission creates, before it starts any year. + run = path.windows.create_run( + NewBudgetWindowRun( + report_id=setup.report.id, + tax_benefit_model_version_id=US_VERSION.id, + trigger=RunTrigger.INITIAL, + max_concurrent_years=2, + ), + first_check_at=NOW, + ) + + read = _reads(path).report_by_id(setup.report.id) + + assert read.status is ReportStatus.PENDING + assert read.run is not None + assert read.run.id == run.id + assert {(year.status, year.stage) for year in read.run.years} == { + (ReportStatus.PENDING, "queued") + } + + +def test_a_window_whose_comparison_runs_reads_its_stage() -> None: + path = WindowPath(max_concurrent_years=2) + report, run = path.submit_window(window_size=3) + first_year = path.year_runs(run.id)[2026] + for run_id in ( + first_year.baseline_simulation_run_id, + first_year.reform_simulation_run_id, + ): + path.runs.succeed(run_id) + + read = _reads(path).report_by_id(report.id) + + assert read.status is ReportStatus.RUNNING + assert read.run is not None + assert (read.run.years[0].status, read.run.years[0].stage) == ( + ReportStatus.RUNNING, + "outputs", + ) + + +def test_a_finished_window_sums_its_years_exactly_when_read() -> None: + path = WindowPath(max_concurrent_years=2) + report, run = path.submit_window(window_size=3) + _finish_window(path, run.id, {2026: 0.1, 2027: 0.2, 2028: 0.3}) + + read = _reads(path).report_by_id(report.id) + + assert read.status is ReportStatus.SUCCEEDED + assert read.run is not None + assert read.run.progress == 100 + assert read.run.completed_at == NOW + result = read.result + assert result is not None + assert (result.startYear, result.endYear, result.windowSize) == ("2026", "2028", 3) + assert [impact.year for impact in result.annualImpacts] == ["2026", "2027", "2028"] + assert [impact.budgetaryImpact for impact in result.annualImpacts] == [ + 0.1, + 0.2, + 0.3, + ] + # Binary floats would make 0.6000000000000001. + assert result.totals.budgetaryImpact == 0.6 + assert result.totals.year == "Total" + first = result.annualImpacts[0] + assert first.taxRevenueImpact == 2.0 + assert first.stateTaxRevenueImpact == 0.5 + assert first.federalTaxRevenueImpact == 1.5 + assert first.benefitSpendingImpact == 1.0 + assert result.totals.federalTaxRevenueImpact == 4.5 + + +def test_a_failed_year_fails_the_window_before_the_run_is_finished() -> None: + path = WindowPath(max_concurrent_years=2) + report, run = path.submit_window(window_size=3) + reform_run_id = path.year_runs(run.id)[2027].reform_simulation_run_id + assert reform_run_id is not None + path.runs.simulation_runs[reform_run_id] = replace( + path.runs.simulation_runs[reform_run_id], + status=RunStatus.FAILED, + error_code="worker_timeout", + error_message="The Modal function timed out.", + completed_at=NOW, + ) + + read = _reads(path).report_by_id(report.id) + + assert read.status is ReportStatus.FAILED + assert read.run is not None + assert read.run.error is not None + assert read.run.error.code == "worker_timeout" + assert read.run.error.message == ( + "The 2027 calculation failed: The Modal function timed out." + ) + assert read.run.years[1].status is ReportStatus.FAILED + assert read.result is None + + +def test_a_failed_run_reads_its_own_error_and_its_years_never_started() -> None: + path = WindowPath(max_concurrent_years=1) + report, run = path.submit_window(window_size=3) + path.windows.finish_run( + run.id, + BudgetWindowStatus.FAILED, + RunError(RunErrorCode.INVALID_INPUT, "The 2027 calculation cannot start."), + ) + + read = _reads(path).report_by_id(report.id) + + assert read.status is ReportStatus.FAILED + assert read.run is not None + assert read.run.error is not None + assert (read.run.error.code, read.run.error.message) == ( + "invalid_input", + "The 2027 calculation cannot start.", + ) + assert [year.status for year in read.run.years] == [ + ReportStatus.RUNNING, + ReportStatus.NOT_RUN, + ReportStatus.NOT_RUN, + ] + + +def test_a_succeeded_window_missing_a_years_budget_is_a_stored_result_error() -> None: + path = WindowPath(max_concurrent_years=2) + report, run = path.submit_window(window_size=2) + _finish_window(path, run.id, {2026: 1.0, 2027: 2.0}) + del path.windows.budgets[path.year_runs(run.id)[2027].id] + + with pytest.raises(StoredResultError) as error: + _reads(path).report_by_id(report.id) + + assert error.value.report_run_id == run.id + + +def test_a_reports_runs_are_listed_newest_first() -> None: + path = WindowPath() + report, first = path.submit_window(window_size=2) + rerun = asyncio.run( + path.budget_windows.run_report( + report.id, + policyengine_version="6.7.5", + trigger=RunTrigger.MANUAL, + force=True, + ) + ) + + listed = _reads(path).report_runs(report.id) + + assert [run.id for run in listed.runs] == [rerun.id, first.id] + assert listed.runs[0].force is True + assert listed.runs[0].trigger == "manual" + + +def test_a_report_without_a_run_has_not_run() -> None: + path = WindowPath() + setup = path.prepare(window_size=2) + + read = _reads(path).report_by_id(setup.report.id) + + assert read.status is ReportStatus.NOT_RUN + assert read.run is None + assert read.result is None + + +def test_only_budget_window_reports_are_read() -> None: + path = WindowPath() + single_year, _ = asyncio.run(path.economy.submit(path.submission())) + reads = _reads(path) + + for report_id in (single_year.id, uuid4()): + with pytest.raises(ReportNotFoundError, match="budget-window report"): + reads.report_by_id(report_id) + with pytest.raises(ReportNotFoundError, match="budget-window report"): + reads.report_runs(report_id) + + +def test_a_window_must_end_in_2099_or_earlier() -> None: + values = { + "country": "us", + "policyengine_version": "6.7.5", + "region_code": "us", + "baseline": {}, + "reform": {}, + } + + ends_in_2099 = BudgetWindowReportSubmissionRequest.model_validate( + {**values, "start_year": 2090, "window_size": 10} + ) + assert ends_in_2099.window_size == 10 + with pytest.raises(ValidationError, match="2099 or earlier"): + BudgetWindowReportSubmissionRequest.model_validate( + {**values, "start_year": 2091, "window_size": 10} + ) + for window_size in (0, 76): + with pytest.raises(ValidationError, match="window_size"): + BudgetWindowReportSubmissionRequest.model_validate( + {**values, "start_year": 2026, "window_size": window_size} + ) From 0acc18b41e485c79ad196fa980acf123c570c195 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:17:33 +0400 Subject: [PATCH 05/20] Serve budget-window reports on /v2/budget-window-reports Four routes submit a budget-window report, read it with its newest run, run it again (manual, or system), and list its runs. They need the same permissions as the economy report routes, declare the same error answers for the generated client (404 and 503 as EconomyErrorResponse, 422 as EconomyRefusalResponse, and 500 on the read when a succeeded run's stored figures cannot be read back), and answer 503 without the simulation database or the Modal environment. The app builds the budget-window path from the simulation database on the economy path, and wraps the call recorder so that a reported call's end starts the budget-window years it frees, after the comparisons it lets start. create_app's wiring of the report paths moves to _report_paths. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/app.py | 72 ++++- .../budget_window_reads.py | 22 +- .../budget_window_reports.py | 48 ++- .../budget_window_routes.py | 215 +++++++++++++ .../economy_routes.py | 14 +- .../simulation_database.py | 30 ++ .../tests/test_budget_window_reads.py | 10 +- .../tests/test_budget_window_routes.py | 302 ++++++++++++++++++ .../tests/test_openapi.py | 46 ++- 9 files changed, 721 insertions(+), 38 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_budget_window_routes.py 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 b0273eeb4..a8bdb82c7 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -52,6 +52,11 @@ CallerAuthorizer, ) from policyengine_simulation_entry.backend import OldGatewayBackend, SimulationBackend +from policyengine_simulation_entry.budget_window_reports import YearStartingRecorder +from policyengine_simulation_entry.budget_window_routes import ( + BudgetWindowRouteServices, + register_budget_window_routes, +) from policyengine_simulation_entry.call_report_route import ( CallReportRecorder, ModalCallerVerifier, @@ -123,6 +128,7 @@ class RunServices: call_recorder: CallReportRecorder | None = None modal_identity_verifier: ModalCallerVerifier | None = None economy: EconomyRouteServices | None = None + budget_windows: BudgetWindowRouteServices | None = None def _temporary_submission_payload( @@ -580,6 +586,48 @@ async def ping(body: PingRequest, request: Request) -> Response: ) +def _report_paths( + services: RunServices, + runtime_settings: Settings, + runtime_database: SimulationDatabaseConnection | None, +) -> tuple[ + CallReportRecorder | None, + EconomyRouteServices | None, + BudgetWindowRouteServices | None, +]: + """Return the call recorder and the report paths, building each one not given. + + The economy and budget-window paths need the simulation database and the + Modal environment calls start in. A reported call's end starts the + comparisons and budget-window years it lets start. + """ + call_recorder = services.call_recorder + if call_recorder is None and isinstance(runtime_database, SimulationDatabase): + call_recorder = runtime_database.call_recorder() + economy = services.economy + if ( + economy is None + and isinstance(runtime_database, SimulationDatabase) + and runtime_settings.call_report_modal_environment + ): + economy = runtime_database.economy_routes( + artifact_environment=runtime_settings.environment, + modal_environment=runtime_settings.call_report_modal_environment, + ) + budget_windows = services.budget_windows + if ( + budget_windows is None + and isinstance(runtime_database, SimulationDatabase) + and economy is not None + ): + budget_windows = runtime_database.budget_window_routes(economy) + if call_recorder is not None and economy is not None: + call_recorder = ComparisonStartingRecorder(call_recorder, economy.reports) + if call_recorder is not None and budget_windows is not None: + call_recorder = YearStartingRecorder(call_recorder, budget_windows.reports) + return call_recorder, economy, budget_windows + + def create_app( *, settings: Settings | None = None, @@ -596,21 +644,9 @@ def create_app( runtime_database = SimulationDatabase.from_url( runtime_settings.simulation_database_url ) - call_recorder = services.call_recorder - if call_recorder is None and isinstance(runtime_database, SimulationDatabase): - call_recorder = runtime_database.call_recorder() - economy = services.economy - if ( - economy is None - and isinstance(runtime_database, SimulationDatabase) - and runtime_settings.call_report_modal_environment - ): - economy = runtime_database.economy_routes( - artifact_environment=runtime_settings.environment, - modal_environment=runtime_settings.call_report_modal_environment, - ) - if call_recorder is not None and economy is not None: - call_recorder = ComparisonStartingRecorder(call_recorder, economy.reports) + call_recorder, economy, budget_windows = _report_paths( + services, runtime_settings, runtime_database + ) modal_identity_verifier = services.modal_identity_verifier expected_modal_caller = runtime_settings.expected_modal_caller if modal_identity_verifier is None and expected_modal_caller is not None: @@ -686,6 +722,12 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: submitter=simulation_submitter, reader=simulation_reader, ) + register_budget_window_routes( + app, + services=budget_windows, + submitter=simulation_submitter, + reader=simulation_reader, + ) return app diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py index 6fb83bc64..be78b7b47 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py @@ -34,7 +34,6 @@ ) from policyengine_simulation_entry.economy_report_types import ReportNotFoundError from policyengine_simulation_entry.economy_routes import ( - StoredResultError, report_run_progress, simulation_response, ) @@ -68,6 +67,15 @@ ) +class StoredWindowResultError(Exception): + """A succeeded budget-window run's stored figures cannot be read back.""" + + def __init__(self, budget_window_run_id: UUID, cause: LookupError) -> None: + """Name the run and keep the reader's error as the cause.""" + super().__init__(f"budget-window run {budget_window_run_id}: {cause}") + self.budget_window_run_id = budget_window_run_id + + class BudgetWindowRunReads(Protocol): """``BudgetWindowRunStore``'s reads.""" @@ -236,7 +244,8 @@ def _year( def report(self, report: ReportRecord) -> BudgetWindowReportResponse: """Describe a report with its newest run and, once it succeeded, its result. - Raise StoredResultError when a succeeded run's figures cannot be read. + Raise StoredWindowResultError when a succeeded run's figures cannot + be read. """ baseline, reform = ( self.simulations.get(simulation_id) @@ -261,7 +270,7 @@ def report(self, report: ReportRecord) -> BudgetWindowReportResponse: report, self.windows.year_budgets(newest.id) ) except LookupError as error: - raise StoredResultError(newest.id, error) from error + raise StoredWindowResultError(newest.id, error) from error return BudgetWindowReportResponse( id=report.id, country=report.country, @@ -278,16 +287,17 @@ def report(self, report: ReportRecord) -> BudgetWindowReportResponse: def report_by_id(self, report_id: UUID) -> BudgetWindowReportResponse: """Describe a report; raise ReportNotFoundError for one that is not a window.""" - return self.report(self._report(report_id)) + return self.report(self.report_record(report_id)) def report_runs(self, report_id: UUID) -> BudgetWindowRunsResponse: """List a budget-window report's runs, newest first.""" - report = self._report(report_id) + report = self.report_record(report_id) return BudgetWindowRunsResponse( runs=[self.run(run, report) for run in self.windows.runs(report_id)] ) - def _report(self, report_id: UUID) -> ReportRecord: + def report_record(self, report_id: UUID) -> ReportRecord: + """Return a budget-window report; raise ReportNotFoundError for any other.""" record = self.reports.get(report_id) if ( record is None diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index 7d8160382..f4a9eccf5 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -17,7 +17,9 @@ lease of ``SCHEDULING_LEASE``, so two callers do not start the same year; the database allows one report run per year of a run in any case. -The first year that fails fails the run, with that year's error, and no more +``YearStartingRecorder`` wraps the recording of a reported call, so its end +starts the years it frees. The first year that fails fails the run, with that +year's error, and no more years start; years already running finish and keep their results. Once every year has succeeded, the run has succeeded. The run stores no totals: they are the sums of its years' budget rows, added up when it is read. @@ -31,6 +33,7 @@ from datetime import timedelta from typing import TYPE_CHECKING, Protocol +from policyengine_simulation_contract.call_reports import CallReportStatus from policyengine_simulation_persistence import ( BudgetWindowStatus, NewBudgetWindowRun, @@ -55,6 +58,7 @@ from datetime import datetime from uuid import UUID + from policyengine_simulation_contract.call_reports import CallEndedReport from policyengine_simulation_persistence import ( BudgetWindowRunRecord, BudgetWindowYear, @@ -536,3 +540,45 @@ async def _resume(self, budget_window_run_id: UUID) -> int: }, ) return 0 + + +class ReportedCallRecorder(Protocol): + """Record a reported call's outcome (``RunRecorder``, possibly wrapped).""" + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + """Record the call's outcome if Modal has one.""" + ... + + +_RECORDED = frozenset({CallReportStatus.RECORDED, CallReportStatus.ALREADY_RECORDED}) + + +@dataclass(frozen=True) +class YearStartingRecorder: + """Record a reported call, then start the budget-window years its end frees. + + A failure to start them is logged and changes nothing in the answer: the + safety check starts a run's due years. + """ + + recorder: ReportedCallRecorder + windows: BudgetWindowReports + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + """Record the reported call; start the years it lets start.""" + report_status = await self.recorder.record_reported_call(report) + if report_status in _RECORDED: + try: + await self.windows.start_years_after_call( + simulation_run_ids=report.simulation_run_ids, + report_run_id=report.report_run_id, + ) + except Exception as error: # noqa: BLE001 # logged; the safety check starts them later + logger.warning( + "budget_window_years_not_started", + extra={ + "modal_call_id": report.modal_call_id, + "error_type": type(error).__name__, + }, + ) + return report_status diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py new file mode 100644 index 000000000..8271a49bc --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py @@ -0,0 +1,215 @@ +"""The v2 budget-window report routes. + +- ``POST /v2/budget-window-reports``: create a budget-window report and its + first run, and start its first years (``BudgetWindowReports.submit``); 201. +- ``GET /v2/budget-window-reports/{report_id}``: the report, its newest run + with each year's progress and, once every year has succeeded, its result + in today's budget-window format, with totals summed when read. +- ``POST /v2/budget-window-reports/{report_id}/runs``: calculate it again; + 201. A ``system`` run returns the report's newest run when it has one. +- ``GET /v2/budget-window-reports/{report_id}/runs``: its runs, newest first. + +Submissions need ``simulation:submit`` and reads ``simulation:read``. Without +the simulation database, or the Modal environment sim starts calls in, every +route answers 503. The routes declare their error answers as the economy +report routes do (``economy_routes``), so the generated client parses them; +the report read also answers 500 when a succeeded run's stored figures +cannot be read back. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any +from uuid import UUID # noqa: TC003 # FastAPI evaluates the handlers' path parameters + +from fastapi import Depends, FastAPI, HTTPException, status +from policyengine_simulation_persistence import RunTrigger + +from policyengine_simulation_entry.budget_window_reads import StoredWindowResultError +from policyengine_simulation_entry.budget_window_reports import ( + BudgetWindowReportSubmission, +) +from policyengine_simulation_entry.budget_window_schemas import ( + BudgetWindowReportResponse, + BudgetWindowReportSubmissionRequest, + BudgetWindowRunResponse, + BudgetWindowRunsResponse, +) +from policyengine_simulation_entry.economy_report_types import ( + EconomyReportSubmission, + SimulationChoice, +) +from policyengine_simulation_entry.economy_routes import ( + REFUSED, + RESULT_UNREADABLE, + UNAVAILABLE, + answer_route, +) +from policyengine_simulation_entry.economy_schemas import ( + EconomyErrorResponse, + EconomyReportRunRequest, +) + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from policyengine_simulation_entry.auth import AuthorizationDependency + from policyengine_simulation_entry.budget_window_reads import ( + BudgetWindowReportReads, + ) + from policyengine_simulation_entry.budget_window_reports import ( + BudgetWindowReports, + ) + +logger = logging.getLogger(__name__) + +NOT_FOUND: dict[int | str, dict[str, Any]] = { + status.HTTP_404_NOT_FOUND: { + "model": EconomyErrorResponse, + "description": "No such budget-window report", + }, +} + + +@dataclass(frozen=True) +class BudgetWindowRouteServices: + """What the routes run on: the budget-window path and its status reads.""" + + reports: BudgetWindowReports + reads: BudgetWindowReportReads + + +async def _answer[ResultT](call: Callable[[], Awaitable[ResultT]]) -> ResultT: + """Answer as the economy routes do, and 500 for unreadable stored figures.""" + try: + return await answer_route(call) + except StoredWindowResultError as error: + logger.error( # noqa: TRY400 # result failure logs carry only the run ID and the error type; a traceback could expose stored result data + "budget_window_result_unreadable", + extra={ + "budget_window_run_id": str(error.budget_window_run_id), + "error_type": type(error.__cause__).__name__, + }, + ) + raise HTTPException( + status.HTTP_500_INTERNAL_SERVER_ERROR, + "The report's stored result cannot be read back.", + ) from None + + +def _configured( + services: BudgetWindowRouteServices | None, +) -> BudgetWindowRouteServices: + if services is None: + raise HTTPException( + status.HTTP_503_SERVICE_UNAVAILABLE, + "Budget-window reports are not configured on this deployment.", + ) + return services + + +def _submission( + body: BudgetWindowReportSubmissionRequest, +) -> BudgetWindowReportSubmission: + return BudgetWindowReportSubmission( + report=EconomyReportSubmission( + country=body.country, + policyengine_version=body.policyengine_version, + year=body.start_year, + region_code=body.region_code, + dataset_name=body.dataset_name, + baseline=SimulationChoice(**body.baseline.model_dump()), + reform=SimulationChoice(**body.reform.model_dump()), + options=body.options, + ), + window_size=body.window_size, + ) + + +def register_budget_window_routes( + app: FastAPI, + *, + services: BudgetWindowRouteServices | None, + submitter: AuthorizationDependency, + reader: AuthorizationDependency, +) -> None: + """Add the v2 budget-window report routes.""" + + @app.post( + "/v2/budget-window-reports", + status_code=status.HTTP_201_CREATED, + response_model=BudgetWindowReportResponse, + operation_id="submit_budget_window_report", + responses={**REFUSED, **UNAVAILABLE}, + dependencies=[Depends(submitter)], + ) + async def submit_budget_window_report( + body: BudgetWindowReportSubmissionRequest, + ) -> BudgetWindowReportResponse: + windows = _configured(services) + + async def submit() -> BudgetWindowReportResponse: + report, _ = await windows.reports.submit(_submission(body)) + return await asyncio.to_thread(windows.reads.report, report) + + return await _answer(submit) + + @app.get( + "/v2/budget-window-reports/{report_id}", + response_model=BudgetWindowReportResponse, + operation_id="get_budget_window_report", + responses={**NOT_FOUND, **REFUSED, **RESULT_UNREADABLE, **UNAVAILABLE}, + dependencies=[Depends(reader)], + ) + async def get_budget_window_report(report_id: UUID) -> BudgetWindowReportResponse: + windows = _configured(services) + return await _answer( + lambda: asyncio.to_thread(windows.reads.report_by_id, report_id) + ) + + @app.post( + "/v2/budget-window-reports/{report_id}/runs", + status_code=status.HTTP_201_CREATED, + response_model=BudgetWindowRunResponse, + operation_id="run_budget_window_report", + responses={**NOT_FOUND, **REFUSED, **UNAVAILABLE}, + dependencies=[Depends(submitter)], + ) + async def run_budget_window_report( + report_id: UUID, body: EconomyReportRunRequest + ) -> BudgetWindowRunResponse: + windows = _configured(services) + + async def run() -> BudgetWindowRunResponse: + budget_window_run = await windows.reports.run_report( + report_id, + policyengine_version=body.policyengine_version, + trigger=RunTrigger(body.trigger), + force=body.force, + ) + return await asyncio.to_thread( + lambda: windows.reads.run( + budget_window_run, + windows.reads.report_record(budget_window_run.report_id), + ) + ) + + return await _answer(run) + + @app.get( + "/v2/budget-window-reports/{report_id}/runs", + response_model=BudgetWindowRunsResponse, + operation_id="list_budget_window_report_runs", + responses={**NOT_FOUND, **REFUSED, **UNAVAILABLE}, + dependencies=[Depends(reader)], + ) + async def list_budget_window_report_runs( + report_id: UUID, + ) -> BudgetWindowRunsResponse: + windows = _configured(services) + return await _answer( + lambda: asyncio.to_thread(windows.reads.report_runs, report_id) + ) 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 7156f7d68..370e4b36c 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 @@ -358,7 +358,7 @@ class EconomyRouteServices: reads: EconomyReportReads -async def _answer[ResultT](call: Callable[[], Awaitable[ResultT]]) -> ResultT: +async def answer_route[ResultT](call: Callable[[], Awaitable[ResultT]]) -> ResultT: """Run a route's work, turning sim's refusals into HTTP errors.""" try: return await call() @@ -435,7 +435,7 @@ async def submit() -> EconomyReportResponse: ) return await asyncio.to_thread(economy.reads.report, report) - return await _answer(submit) + return await answer_route(submit) @app.get( "/v2/economy-reports/by-legacy-id/{legacy_report_id}", @@ -448,7 +448,7 @@ async def get_economy_report_by_legacy_id( legacy_report_id: int, ) -> EconomyReportResponse: economy = _configured(services) - return await _answer( + return await answer_route( lambda: asyncio.to_thread( economy.reads.report_by_legacy_id, legacy_report_id ) @@ -463,7 +463,7 @@ async def get_economy_report_by_legacy_id( ) async def get_economy_report(report_id: UUID) -> EconomyReportResponse: economy = _configured(services) - return await _answer( + return await answer_route( lambda: asyncio.to_thread(economy.reads.report_by_id, report_id) ) @@ -489,7 +489,7 @@ async def run() -> EconomyReportRunResponse: ) return await asyncio.to_thread(economy.reads.report_run, report_run) - return await _answer(run) + return await answer_route(run) @app.get( "/v2/economy-reports/{report_id}/runs", @@ -500,7 +500,7 @@ async def run() -> EconomyReportRunResponse: ) async def list_economy_report_runs(report_id: UUID) -> EconomyReportRunsResponse: economy = _configured(services) - return await _answer( + return await answer_route( lambda: asyncio.to_thread(economy.reads.report_runs, report_id) ) @@ -513,6 +513,6 @@ async def list_economy_report_runs(report_id: UUID) -> EconomyReportRunsResponse ) async def get_simulation(simulation_id: UUID) -> SimulationResponse: economy = _configured(services) - return await _answer( + return await answer_route( lambda: asyncio.to_thread(economy.reads.simulation, simulation_id) ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py index 17b8f7f90..014a23426 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_database.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING from policyengine_simulation_persistence import ( + BudgetWindowRunStore, CatalogReader, EconomyResultLabels, ReportStore, @@ -18,6 +19,11 @@ read_schema_revision, ) +from policyengine_simulation_entry.budget_window_reads import BudgetWindowReportReads +from policyengine_simulation_entry.budget_window_reports import BudgetWindowReports +from policyengine_simulation_entry.budget_window_routes import ( + BudgetWindowRouteServices, +) from policyengine_simulation_entry.economy_recorders import economy_result_recorders from policyengine_simulation_entry.economy_reports import EconomyReports from policyengine_simulation_entry.economy_routes import ( @@ -130,6 +136,30 @@ def economy_routes( ), ) + def budget_window_routes( + self, economy: EconomyRouteServices + ) -> BudgetWindowRouteServices: + """Return the budget-window path, running its years on ``economy``'s path.""" + engine = self.engine + reports = ReportStore(engine) + windows = BudgetWindowRunStore(engine) + catalog = CatalogReader(engine) + return BudgetWindowRouteServices( + reports=BudgetWindowReports( + economy=economy.reports, + reports=reports, + windows=windows, + runs=RunStore(engine), + catalog=catalog, + ), + reads=BudgetWindowReportReads( + simulations=SimulationStore(engine), + reports=reports, + windows=windows, + catalog=catalog, + ), + ) + def dispose(self) -> None: """Close every pooled connection.""" self.engine.dispose() diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py index e38e29d94..f9ec743e4 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py @@ -9,12 +9,14 @@ import pytest from budget_window_fakes import WindowPath, finish_year from economy_fakes import NOW, US_VERSION -from policyengine_simulation_entry.budget_window_reads import BudgetWindowReportReads +from policyengine_simulation_entry.budget_window_reads import ( + BudgetWindowReportReads, + StoredWindowResultError, +) from policyengine_simulation_entry.budget_window_schemas import ( BudgetWindowReportSubmissionRequest, ) from policyengine_simulation_entry.economy_report_types import ReportNotFoundError -from policyengine_simulation_entry.economy_routes import StoredResultError from policyengine_simulation_entry.economy_schemas import ReportStatus from policyengine_simulation_persistence import ( BudgetWindowStatus, @@ -232,10 +234,10 @@ def test_a_succeeded_window_missing_a_years_budget_is_a_stored_result_error() -> _finish_window(path, run.id, {2026: 1.0, 2027: 2.0}) del path.windows.budgets[path.year_runs(run.id)[2027].id] - with pytest.raises(StoredResultError) as error: + with pytest.raises(StoredWindowResultError) as error: _reads(path).report_by_id(report.id) - assert error.value.report_run_id == run.id + assert error.value.budget_window_run_id == run.id def test_a_reports_runs_are_listed_newest_first() -> None: 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..c312c8e43 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py @@ -0,0 +1,302 @@ +"""The v2 budget-window report routes, over in-memory stores.""" + +from __future__ import annotations + +import asyncio +from uuid import UUID, uuid4 + +import pytest +from budget_window_fakes import REFORM_POLICY, WindowPath, finish_year +from conftest import FakeBackend, make_settings +from fastapi.testclient import TestClient +from policyengine_simulation_contract.call_reports import ( + CallEndedReport, + CallReportStatus, +) +from policyengine_simulation_entry.app import RunServices, create_app +from policyengine_simulation_entry.budget_window_reads import BudgetWindowReportReads +from policyengine_simulation_entry.budget_window_reports import ( + DEFAULT_MAX_CONCURRENT_YEARS, + YearStartingRecorder, +) +from policyengine_simulation_entry.budget_window_routes import ( + BudgetWindowRouteServices, +) +from policyengine_simulation_entry.economy_routes import ( + EconomyReportReads, + EconomyRouteServices, +) +from policyengine_simulation_entry.economy_schemas import ( + EconomyErrorResponse, + EconomyRefusalResponse, +) +from policyengine_simulation_entry.simulation_database import SimulationDatabase +from policyengine_simulation_persistence import RunStatus, YearBudget + +SUBMISSION = { + "country": "us", + "policyengine_version": "6.7.5", + "start_year": 2026, + "window_size": 3, + "region_code": "state/ca", + "baseline": {}, + "reform": {"policy_id": str(REFORM_POLICY)}, +} + + +class _RoutePath(WindowPath): + def __init__(self, *, configured: bool = True) -> None: + super().__init__(max_concurrent_years=2) + self.services = BudgetWindowRouteServices( + reports=self.budget_windows, + reads=BudgetWindowReportReads( + simulations=self.simulations, + reports=self.reports, + windows=self.windows, + catalog=self.catalog, + ), + ) + economy = EconomyRouteServices( + reports=self.economy, + reads=EconomyReportReads( + simulations=self.simulations, + reports=self.reports, + runs=self.runs, + catalog=self.catalog, + results=None, # type: ignore[arg-type] + ), + ) + self.client = TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + run_services=RunServices( + economy=economy, + budget_windows=self.services if configured else None, + ), + ) + ) + + def submit(self, **changes: object) -> dict: + response = self.client.post( + "/v2/budget-window-reports", json={**SUBMISSION, **changes} + ) + assert response.status_code == 201, response.text + return response.json() + + def finish_started_years(self, run_id: UUID, budgetary_impact: float) -> None: + for report_run in self.year_runs(run_id).values(): + if report_run.comparison_status is RunStatus.SUCCEEDED: + continue + finish_year(self.runs, report_run) + self.windows.budgets[report_run.id] = YearBudget( + year=report_run.year, + budgetary_impact=budgetary_impact, + tax_revenue_impact=1.0, + state_tax_revenue_impact=0.0, + benefit_spending_impact=2.0, + ) + + +@pytest.fixture +def window_path(): + path = _RoutePath() + with path.client: + yield path + + +def test_a_submitted_window_is_calculating_its_first_years(window_path) -> None: + report = window_path.submit() + + assert report["country"] == "us" + assert report["start_year"] == 2026 + assert report["window_size"] == 3 + assert report["status"] == "running" + assert report["result"] is None + assert report["reform_simulation"]["policy_id"] == str(REFORM_POLICY) + run = report["run"] + assert run["trigger"] == "initial" + assert run["policyengine_version"] == "6.7.5" + assert run["max_concurrent_years"] == 2 + assert [(year["year"], year["status"]) for year in run["years"]] == [ + (2026, "running"), + (2027, "running"), + (2028, "pending"), + ] + + +def test_a_window_reads_back_its_result_once_every_year_succeeded( + window_path, +) -> None: + report = window_path.submit() + run_id = UUID(report["run"]["id"]) + window_path.finish_started_years(run_id, -1.25) + asyncio.run( + window_path.budget_windows.start_years_after_call( + simulation_run_ids=None, + report_run_id=window_path.year_runs(run_id)[2026].id, + ) + ) + window_path.finish_started_years(run_id, -1.25) + + body = window_path.client.get(f"/v2/budget-window-reports/{report['id']}").json() + + assert body["status"] == "succeeded" + assert body["run"]["progress"] == 100 + assert body["result"]["kind"] == "budgetWindow" + assert body["result"]["totals"]["budgetaryImpact"] == -3.75 + assert [impact["year"] for impact in body["result"]["annualImpacts"]] == [ + "2026", + "2027", + "2028", + ] + + +def test_a_window_is_run_again_and_its_runs_are_listed(window_path) -> None: + report = window_path.submit() + + rerun = window_path.client.post( + f"/v2/budget-window-reports/{report['id']}/runs", + json={"policyengine_version": "6.7.5", "trigger": "manual", "force": True}, + ) + listed = window_path.client.get(f"/v2/budget-window-reports/{report['id']}/runs") + + assert rerun.status_code == 201, rerun.text + assert rerun.json()["force"] is True + assert rerun.json()["trigger"] == "manual" + assert [run["id"] for run in listed.json()["runs"]] == [ + rerun.json()["id"], + report["run"]["id"], + ] + + +def test_unknown_and_single_year_reports_are_not_found(window_path) -> None: + single_year, _ = asyncio.run(window_path.economy.submit(window_path.submission())) + + for report_id in (uuid4(), single_year.id): + for method, path, body in ( + ("GET", f"/v2/budget-window-reports/{report_id}", None), + ("GET", f"/v2/budget-window-reports/{report_id}/runs", None), + ( + "POST", + f"/v2/budget-window-reports/{report_id}/runs", + {"policyengine_version": "6.7.5", "trigger": "manual"}, + ), + ): + response = window_path.client.request(method, path, json=body) + assert response.status_code == 404, (method, path) + assert "budget-window report" in ( + EconomyErrorResponse.model_validate(response.json()).detail + ) + + +def test_a_window_sim_cannot_run_or_that_is_malformed_is_refused( + window_path, +) -> None: + unknown_region = window_path.client.post( + "/v2/budget-window-reports", json={**SUBMISSION, "region_code": "state/zz"} + ) + too_late = window_path.client.post( + "/v2/budget-window-reports", + json={**SUBMISSION, "start_year": 2095, "window_size": 10}, + ) + + assert unknown_region.status_code == 422 + assert "state/zz" in str( + EconomyRefusalResponse.model_validate(unknown_region.json()).detail + ) + assert too_late.status_code == 422 + assert "2099" in str(EconomyRefusalResponse.model_validate(too_late.json()).detail) + assert window_path.windows.rows == {} + + +def test_without_the_budget_window_path_every_route_answers_503() -> None: + path = _RoutePath(configured=False) + report_id = uuid4() + + with path.client: + responses = [ + path.client.post("/v2/budget-window-reports", json=SUBMISSION), + path.client.get(f"/v2/budget-window-reports/{report_id}"), + path.client.post( + f"/v2/budget-window-reports/{report_id}/runs", + json={"policyengine_version": "6.7.5", "trigger": "manual"}, + ), + path.client.get(f"/v2/budget-window-reports/{report_id}/runs"), + ] + + assert {response.status_code for response in responses} == {503} + + +class _Recorder: + def __init__(self, answer: CallReportStatus) -> None: + self.answer = answer + self.reports: list[CallEndedReport] = [] + + async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: + self.reports.append(report) + return self.answer + + +def test_a_recorded_comparison_starts_its_windows_next_year() -> None: + path = WindowPath(max_concurrent_years=1) + _, run = path.submit_window(window_size=2) + first_year = path.year_runs(run.id)[2026] + finish_year(path.runs, first_year) + inner = _Recorder(CallReportStatus.RECORDED) + recorder = YearStartingRecorder(inner, path.budget_windows) + report = CallEndedReport(modal_call_id="fc-comparison", report_run_id=first_year.id) + + answer = asyncio.run(recorder.record_reported_call(report)) + + assert answer is CallReportStatus.RECORDED + assert inner.reports == [report] + assert sorted(path.year_runs(run.id)) == [2026, 2027] + + +def test_a_report_not_recorded_starts_nothing_and_a_failure_is_only_logged() -> None: + path = WindowPath(max_concurrent_years=1) + _, run = path.submit_window(window_size=2) + first_year = path.year_runs(run.id)[2026] + finish_year(path.runs, first_year) + report = CallEndedReport(modal_call_id="fc-comparison", report_run_id=first_year.id) + + still_running = asyncio.run( + YearStartingRecorder( + _Recorder(CallReportStatus.STILL_RUNNING), path.budget_windows + ).record_reported_call(report) + ) + assert still_running is CallReportStatus.STILL_RUNNING + assert sorted(path.year_runs(run.id)) == [2026] + + def lose_the_database(*_args, **_kwargs): + msg = "the database dropped the connection" + raise ConnectionError(msg) + + path.runs.get_report_run = lose_the_database # type: ignore[method-assign] + recorded = asyncio.run( + YearStartingRecorder( + _Recorder(CallReportStatus.RECORDED), path.budget_windows + ).record_reported_call(report) + ) + assert recorded is CallReportStatus.RECORDED + assert sorted(path.year_runs(run.id)) == [2026] + + +def test_the_database_builds_the_budget_window_path_on_the_economy_path() -> None: + database = SimulationDatabase.from_url( + "postgresql://simulation_runtime:secret@db.example:5432/postgres" + ) + try: + economy = database.economy_routes( + artifact_environment="staging", modal_environment="main" + ) + + windows = database.budget_window_routes(economy) + finally: + database.dispose() + + assert windows.reports.economy is economy.reports + assert windows.reports.max_concurrent_years == DEFAULT_MAX_CONCURRENT_YEARS + assert windows.reads.windows is windows.reports.windows diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 69ebee8ee..1287f2bd3 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -43,10 +43,19 @@ class OpenAPIDocument(TypedDict): ("GET", "/v2/simulations/{simulation_id}"), ("POST", "/v2/economy-reports"), ("POST", "/v2/economy-reports/{report_id}/runs"), + ("GET", "/v2/budget-window-reports/{report_id}"), + ("GET", "/v2/budget-window-reports/{report_id}/runs"), + ("POST", "/v2/budget-window-reports"), + ("POST", "/v2/budget-window-reports/{report_id}/runs"), } -# The v2 economy report routes are the entrypoint's own; the old gateway has -# none of them or their bodies. +# The v2 economy and budget-window report routes are the entrypoint's own; the +# old gateway has none of them or their bodies. V2_SCHEMAS = ( + "BudgetWindowReportResponse", + "BudgetWindowReportSubmissionRequest", + "BudgetWindowRunResponse", + "BudgetWindowRunsResponse", + "BudgetWindowYearResponse", "ComparisonResponse", "EconomyErrorResponse", "EconomyRefusalResponse", @@ -229,10 +238,33 @@ def test_v2_economy_routes_are_public_for_the_generated_client(): assert set(V2_SCHEMAS) <= set(spec["components"]["schemas"]) +def test_v2_budget_window_routes_are_public_for_the_generated_client(): + spec = cast("OpenAPIDocument", create_app().openapi()) + ids = operation_ids(spec) + + assert { + ids[("POST", "/v2/budget-window-reports")], + ids[("GET", "/v2/budget-window-reports/{report_id}")], + ids[("POST", "/v2/budget-window-reports/{report_id}/runs")], + ids[("GET", "/v2/budget-window-reports/{report_id}/runs")], + } == { + "submit_budget_window_report", + "get_budget_window_report", + "run_budget_window_report", + "list_budget_window_report_runs", + } + + V2_ERROR_ANSWERS = { - ("POST", "/v2/economy-reports"): { - "422": "EconomyRefusalResponse", - "503": "EconomyErrorResponse", + **{ + submission: { + "422": "EconomyRefusalResponse", + "503": "EconomyErrorResponse", + } + for submission in ( + ("POST", "/v2/economy-reports"), + ("POST", "/v2/budget-window-reports"), + ) }, **{ operation: { @@ -246,6 +278,9 @@ def test_v2_economy_routes_are_public_for_the_generated_client(): ("POST", "/v2/economy-reports/{report_id}/runs"), ("GET", "/v2/economy-reports/{report_id}/runs"), ("GET", "/v2/simulations/{simulation_id}"), + ("GET", "/v2/budget-window-reports/{report_id}"), + ("POST", "/v2/budget-window-reports/{report_id}/runs"), + ("GET", "/v2/budget-window-reports/{report_id}/runs"), ) }, } @@ -253,6 +288,7 @@ def test_v2_economy_routes_are_public_for_the_generated_client(): for _result_read in ( ("GET", "/v2/economy-reports/{report_id}"), ("GET", "/v2/economy-reports/by-legacy-id/{legacy_report_id}"), + ("GET", "/v2/budget-window-reports/{report_id}"), ): V2_ERROR_ANSWERS[_result_read] = { **V2_ERROR_ANSWERS[_result_read], From cdda15cce405fbbd44988aa905116ef8bec4f549 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:18:50 +0400 Subject: [PATCH 06/20] Start the years of budget windows left waiting on each safety-check pass After it records due calls and starts the economy work left waiting, each safety-check pass starts the next years of budget-window runs whose check is due, so a run whose caller stopped between years, or whose call report never arrived, carries on. The pass counts the years it started. The job builds the budget-window path on the same economy path, in its own Modal environment, and starts none without it. Co-Authored-By: Claude Opus 5.5 --- .../run_recording.py | 2 + .../safety_check.py | 37 ++++++++--- .../tests/test_budget_window_safety_check.py | 62 +++++++++++++++++++ 3 files changed, 91 insertions(+), 10 deletions(-) create mode 100644 projects/policyengine-simulation-entry/tests/test_budget_window_safety_check.py diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py index 30a00fc3f..b8b3a49af 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/run_recording.py @@ -234,6 +234,8 @@ class RecordingPass: # had succeeded, and runs left pending by a submitter that stopped. comparisons_started: int = 0 runs_started: int = 0 + # Budget-window years the pass started for runs a stopped caller left. + budget_window_years_started: int = 0 @property def found_work(self) -> bool: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py index fd8335000..6ebee0971 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/safety_check.py @@ -12,7 +12,9 @@ one interval ago, and records what it reports, with the same conditional writes as a report; - starts the economy work left waiting: comparison steps whose runs both - succeeded, and runs a submitter left pending (``SafetyCheckPass``). + succeeded, and runs a submitter left pending (``SafetyCheckPass``); +- starts the next years of budget-window runs whose check is due, so a run + whose caller stopped between years carries on. The job runs only the image of the revision that serves traffic: the deploy updates it when that revision is promoted, and a rollback moves it back, so @@ -91,21 +93,34 @@ async def resume_waiting_work(self, *, now: datetime) -> ResumedWork: ... +class WaitingBudgetWindows(Protocol): + """Start budget-window years left waiting (``BudgetWindowReports``).""" + + async def resume_waiting_work(self, *, now: datetime) -> int: + """Start the years of runs whose check is due; return how many.""" + ... + + @dataclass(frozen=True) class SafetyCheckPass: """One pass: record the due calls, then start the work left waiting.""" recorder: DueOutcomeRecorder waiting: WaitingWork | None + budget_windows: WaitingBudgetWindows | None = None clock: Callable[[], datetime] = utc_now async def record_due_outcomes(self) -> RecordingPass: - """Record due calls' outcomes, then start waiting comparisons and runs.""" + """Record due calls, then start waiting comparisons, runs and years.""" recording = await self.recorder.record_due_outcomes() if self.waiting is not None: resumed = await self.waiting.resume_waiting_work(now=self.clock()) recording.comparisons_started = resumed.comparisons_started recording.runs_started = resumed.runs_started + if self.budget_windows is not None: + recording.budget_window_years_started = ( + await self.budget_windows.resume_waiting_work(now=self.clock()) + ) return recording @@ -115,15 +130,17 @@ def _database_recorder( """Build the pass; it starts waiting work when the Modal environment is set.""" database = SimulationDatabase.from_url(database_url) modal_environment = environment.get("SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT", "") - waiting = ( - database.economy_routes( - artifact_environment=environment.get("APP_ENVIRONMENT", "local").lower(), - modal_environment=modal_environment, - ).reports - if modal_environment - else None + if not modal_environment: + return SafetyCheckPass(database.call_recorder(), None) + economy = database.economy_routes( + artifact_environment=environment.get("APP_ENVIRONMENT", "local").lower(), + modal_environment=modal_environment, + ) + return SafetyCheckPass( + database.call_recorder(), + economy.reports, + budget_windows=database.budget_window_routes(economy).reports, ) - return SafetyCheckPass(database.call_recorder(), waiting) def main( diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_safety_check.py b/projects/policyengine-simulation-entry/tests/test_budget_window_safety_check.py new file mode 100644 index 000000000..1e2299cbf --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_safety_check.py @@ -0,0 +1,62 @@ +"""The safety check also starts the years of budget windows left waiting.""" + +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime + +from policyengine_simulation_entry import safety_check +from policyengine_simulation_entry.budget_window_reports import BudgetWindowReports +from policyengine_simulation_entry.run_recording import RecordingPass + +DATABASE_URL = "postgresql://runtime:secret@db.example/postgres" +NOW = datetime(2026, 10, 29, 12, tzinfo=UTC) + + +class _OnePass: + async def record_due_outcomes(self) -> RecordingPass: + return RecordingPass(checked=1) + + +class _WaitingWindows: + def __init__(self) -> None: + self.resumed_at: list[datetime] = [] + + async def resume_waiting_work(self, *, now: datetime) -> int: + self.resumed_at.append(now) + return 3 + + +def test_a_pass_starts_the_years_of_budget_windows_left_waiting() -> None: + windows = _WaitingWindows() + recorder = safety_check.SafetyCheckPass( + _OnePass(), None, budget_windows=windows, clock=lambda: NOW + ) + + recording = asyncio.run(recorder.record_due_outcomes()) + + assert windows.resumed_at == [NOW] + assert recording.budget_window_years_started == 3 + assert recording.checked == 1 + + +def test_a_pass_without_budget_windows_starts_no_years() -> None: + recording = asyncio.run( + safety_check.SafetyCheckPass(_OnePass(), None).record_due_outcomes() + ) + + assert recording.budget_window_years_started == 0 + + +def test_the_job_resumes_budget_windows_in_its_own_modal_environment() -> None: + environment = { + "SIMULATION_CALL_REPORT_MODAL_ENVIRONMENT": "staging", + "APP_ENVIRONMENT": "Staging", + } + + configured = safety_check._database_recorder(DATABASE_URL, environment=environment) + unconfigured = safety_check._database_recorder(DATABASE_URL, environment={}) + + assert isinstance(configured.budget_windows, BudgetWindowReports) + assert configured.budget_windows.economy is configured.waiting + assert unconfigured.budget_windows is None From 9cf2a3dd687906b5cb391d487c02d8b3317e9452 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:22:53 +0400 Subject: [PATCH 07/20] Give the v2 economy functions room for budget windows at their cap The Simulation Entrypoint calculates at most five years of a budget window at once, a value the contract now holds (BUDGET_WINDOW_MAX_CONCURRENT_YEARS) so the entrypoint and the Modal app share it. Each year is up to two simulation calls and one comparison, and a US national simulation runs 20 segment calls. The economy simulation functions and the comparison function rise from 10 to 40 containers, room for four windows at the cap, and the US segment function from 300 to 600, room for three US national windows. The time limits are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../economy_runs.py | 3 + .../budget_window_reports.py | 8 ++- .../tests/test_budget_window_reports.py | 8 +++ .../src/modal/v2_app.py | 16 ++++-- .../tests/test_economy_run_capacity.py | 55 +++++++++++++++++++ .../tests/test_stage12_modal_app.py | 2 +- 6 files changed, 85 insertions(+), 7 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/economy_runs.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/economy_runs.py index 5d157d2ab..16420f4d4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/economy_runs.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/economy_runs.py @@ -51,6 +51,9 @@ "uk": "run_economy_simulation_uk", } ECONOMY_COMPARISON_FUNCTION = "compare_economy_simulations" +# How many years of one budget window the Simulation Entrypoint calculates at +# once. The functions above leave room for several windows at this cap. +BUDGET_WINDOW_MAX_CONCURRENT_YEARS = 5 # The parameters whose change makes policyengine.py calculate a labour supply # response, by country: a parameter at or below one of these paths. They are diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index f4a9eccf5..685ccf1bd 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -34,6 +34,9 @@ from typing import TYPE_CHECKING, Protocol from policyengine_simulation_contract.call_reports import CallReportStatus +from policyengine_simulation_contract.economy_runs import ( + BUDGET_WINDOW_MAX_CONCURRENT_YEARS, +) from policyengine_simulation_persistence import ( BudgetWindowStatus, NewBudgetWindowRun, @@ -80,8 +83,9 @@ # How many of a run's years are calculated at once. A US national year is up # to two simulation calls of 20 segment calls each, so five years use at most -# 200 of the US segment function's containers. -DEFAULT_MAX_CONCURRENT_YEARS = 5 +# 200 of the US segment function's containers; the v2 Modal app's functions +# have room for several windows at this cap. +DEFAULT_MAX_CONCURRENT_YEARS = BUDGET_WINDOW_MAX_CONCURRENT_YEARS # A caller starting a run's years holds the run this long; one that stopped # loses the claim afterwards, and the safety check carries on. SCHEDULING_LEASE = timedelta(minutes=2) diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py index 16c15d83f..f1724333c 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py @@ -14,8 +14,12 @@ finish_year, ) from economy_fakes import NOW, US_VERSION +from policyengine_simulation_contract.economy_runs import ( + BUDGET_WINDOW_MAX_CONCURRENT_YEARS, +) from policyengine_simulation_entry.budget_window_reports import ( BUDGET_WINDOW_CHECK_INTERVAL, + DEFAULT_MAX_CONCURRENT_YEARS, SCHEDULING_LEASE, ) from policyengine_simulation_entry.economy_report_types import ( @@ -195,6 +199,10 @@ def test_a_claim_lasts_its_lease() -> None: assert timedelta(0) < SCHEDULING_LEASE <= BUDGET_WINDOW_CHECK_INTERVAL +def test_the_cap_is_the_one_the_modal_functions_have_room_for() -> None: + assert DEFAULT_MAX_CONCURRENT_YEARS == BUDGET_WINDOW_MAX_CONCURRENT_YEARS == 5 + + def test_a_rerun_reuses_each_years_runs_unless_forced() -> None: path = WindowPath(max_concurrent_years=2) report, first = path.submit_window(window_size=2) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 2b08e0aab..4650ab614 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -287,7 +287,9 @@ def validate_worker_uk() -> dict: memory=32768, timeout=3600, retries=0, - max_containers=300, + # Three US national budget windows at the Simulation Entrypoint's cap of + # five years each: two simulations a year, 20 segments each. + max_containers=600, secrets=worker_secrets, ) def run_single_simulation_segment_us( @@ -653,7 +655,9 @@ def segmented_calculator( memory=32768, timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, - max_containers=10, + # Four budget windows at the Simulation Entrypoint's cap of five years, + # two simulations a year. + max_containers=40, secrets=worker_secrets, ) def run_economy_simulation_us( @@ -669,7 +673,9 @@ def run_economy_simulation_us( memory=32768, timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, - max_containers=10, + # Four budget windows at the Simulation Entrypoint's cap of five years, + # two simulations a year. + max_containers=40, secrets=worker_secrets, ) def run_economy_simulation_uk( @@ -685,7 +691,9 @@ def run_economy_simulation_uk( memory=8192, timeout=ECONOMY_COMPARISON_TIMEOUT_SECONDS, retries=0, - max_containers=10, + # Four budget windows at the Simulation Entrypoint's cap of five years, + # one comparison a year, with as much room again for other reports. + max_containers=40, secrets=worker_secrets, ) def compare_economy_simulations( diff --git a/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py b/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py new file mode 100644 index 000000000..73d1cb88b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py @@ -0,0 +1,55 @@ +"""The v2 economy functions have room for budget windows calculated at their cap. + +The Simulation Entrypoint calculates at most +``BUDGET_WINDOW_MAX_CONCURRENT_YEARS`` years of a budget window at once, each +year up to two simulation calls and one comparison call, and a US national +simulation call runs one segment call per region group. +""" + +from __future__ import annotations + +from policyengine_simulation_contract.economy_runs import ( + BUDGET_WINDOW_MAX_CONCURRENT_YEARS, + ECONOMY_COMPARISON_FUNCTION, + ECONOMY_SIMULATION_FUNCTIONS, +) +from policyengine_simulation_executor.national_partition import ( + US_NATIONAL_REGION_GROUPS, +) +from test_stage12_modal_app import _load + +# Budget windows each function has room for at once, every year at the cap. +WINDOWS_AT_ONCE = 4 +US_NATIONAL_WINDOWS_AT_ONCE = 3 +SIMULATIONS_PER_YEAR = 2 + + +def test_a_budget_window_calculates_a_few_years_at_once() -> None: + assert BUDGET_WINDOW_MAX_CONCURRENT_YEARS == 5 + + +def test_the_simulation_and_comparison_functions_hold_four_windows( + monkeypatch, +) -> None: + functions = dict(_load(monkeypatch).app.function_calls) + years = WINDOWS_AT_ONCE * BUDGET_WINDOW_MAX_CONCURRENT_YEARS + + for country in ("us", "uk"): + simulation = functions[ECONOMY_SIMULATION_FUNCTIONS[country]] + assert simulation["max_containers"] == SIMULATIONS_PER_YEAR * years == 40 + assert functions[ECONOMY_COMPARISON_FUNCTION]["max_containers"] == 2 * years + + +def test_the_us_segment_function_holds_three_national_windows(monkeypatch) -> None: + functions = dict(_load(monkeypatch).app.function_calls) + + assert ( + functions["run_single_simulation_segment_us"]["max_containers"] + == ( + US_NATIONAL_WINDOWS_AT_ONCE + * BUDGET_WINDOW_MAX_CONCURRENT_YEARS + * SIMULATIONS_PER_YEAR + * len(US_NATIONAL_REGION_GROUPS) + ) + == 600 + ) 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 db04536e3..d482667cb 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -342,7 +342,7 @@ 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"] == 600 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 From 0e47dd6a945515b6af0e3ac3b61644f2e5ff6b2f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:26:55 +0400 Subject: [PATCH 08/20] Add the staging qualification of budget windows against the old gateway An operator runs authenticated_tests/test_staging_budget_window_qualification.py once against staging after the budget-window routes are deployed; no workflow runs it. It submits the same reform, region and years to the v2 route and to the old gateway's budget-window route, and requires each year's figures and the totals to agree within a relative and an absolute tolerance (0.1% and 1,000 by default). The window has more years than sim calculates at once (seven against five), so its last years start only when earlier years report their ends or the safety check passes. Running the report again unforced must read back the same figures. It needs a reform saved in the staging Public API database, given as its policy ID for the v2 report and as its parameter values for the old gateway, and both windows cost Modal time for every year, so it does not run on every staging deployment. It runs only where SIMULATION_STAGING_BUDGET_WINDOW_QUALIFICATION is 1. Co-Authored-By: Claude Opus 5.5 --- ...est_staging_budget_window_qualification.py | 272 ++++++++++++++++++ .../tests/test_budget_window_qualification.py | 91 ++++++ 2 files changed, 363 insertions(+) create mode 100644 projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py create mode 100644 projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py b/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py new file mode 100644 index 000000000..7297efe54 --- /dev/null +++ b/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py @@ -0,0 +1,272 @@ +"""Qualify v2 budget windows on staging against the old gateway's: an operator step. + +An operator runs this once against staging after the v2 budget-window routes +are deployed there; no workflow runs it. It needs a reform saved in the +staging Public API database, given twice: its policy ID for the v2 report, +which reads the policy itself, and its parameter values for the old gateway, +which takes them in the request. Both windows cost Modal time for every year, +which is why this is not run on every staging deployment. + +It checks: + +- **totals:** the same reform, region and years submitted to the v2 route + (``/v2/budget-window-reports``) and to the old gateway's + (``/simulate/economy/budget-window``) give each year's budget figures and + the totals within the tolerance; +- **recovery of queued years:** the window has more years than sim + calculates at once, so its last years start only when earlier years' + calls report their ends or the safety check passes, without the submitter; +- **recovery of results:** running the report again unforced, so that its + years reuse the simulation runs that exist, reads back the same figures. + +Run it from ``projects/policyengine-simulation-entry`` with +``SIMULATION_STAGING_BUDGET_WINDOW_QUALIFICATION=1`` and: + +- ``SIMULATION_ENTRYPOINT_TEST_BASE_URL``: the staging stable URL; +- ``SIMULATION_ENTRYPOINT_TEST_AUTH_ISSUER``, ``_AUDIENCE``, ``_CLIENT_ID`` + and ``_CLIENT_SECRET``: the test client's, as the staging checks use; +- ``SIMULATION_QUALIFICATION_POLICYENGINE_VERSION``: the deployed + PolicyEngine.py version, which a staging catalog version must exist for; +- ``SIMULATION_QUALIFICATION_MODEL_VERSION``: the country package version + of that release, which the old gateway routes by; +- ``SIMULATION_QUALIFICATION_REFORM_POLICY_ID`` and + ``SIMULATION_QUALIFICATION_REFORM``: the reform's v2 policy ID and its + parameter values as JSON (``{"gov....": {"2026-01-01.2100-12-31": 1}}``); +- optionally ``SIMULATION_QUALIFICATION_COUNTRY`` (``us``), + ``_REGION`` (``state/ut``), ``_START_YEAR`` (2026), ``_WINDOW_SIZE`` (7), + ``_RELATIVE_TOLERANCE`` (0.001), ``_ABSOLUTE_TOLERANCE`` (1000, in the + country's currency) and ``_TIMEOUT_MINUTES`` (120). + +For example: + + uv run --frozen pytest \ + authenticated_tests/test_staging_budget_window_qualification.py -v + +The deployed-auth runs collect this directory too; the test runs only where +``SIMULATION_STAGING_BUDGET_WINDOW_QUALIFICATION`` is 1. +""" + +from __future__ import annotations + +import json +import math +import os +import time +from typing import Any + +import httpx +import pytest + +POLL_INTERVAL_SECONDS = 30 +FINISHED = {"succeeded", "failed"} +OLD_GATEWAY_FINISHED = {"complete", "failed"} +# Sim calculates at most this many years of a window at once +# (BUDGET_WINDOW_MAX_CONCURRENT_YEARS); a longer window needs recovery of +# its queued years. +MAX_CONCURRENT_YEARS = 5 +BUDGET_FIELDS = ( + "taxRevenueImpact", + "federalTaxRevenueImpact", + "stateTaxRevenueImpact", + "benefitSpendingImpact", + "budgetaryImpact", +) + +pytestmark = pytest.mark.skipif( + os.environ.get("SIMULATION_STAGING_BUDGET_WINDOW_QUALIFICATION") != "1", + reason="an operator step: runs only when it is asked for", +) + + +def _required(name: str) -> str: + value = os.environ.get(name, "").strip() + if not value: + msg = f"{name} is required for the budget-window qualification." + raise RuntimeError(msg) + return value + + +def _optional(name: str, default: str) -> str: + return os.environ.get(name, "").strip() or default + + +def figures_differ( + v2_value: float, old_value: float, *, relative: float, absolute: float +) -> bool: + """Whether two budget figures differ by more than both tolerances allow.""" + return not math.isclose(v2_value, old_value, rel_tol=relative, abs_tol=absolute) + + +def window_differences( + v2_result: dict[str, Any], + old_result: dict[str, Any], + *, + relative: float, + absolute: float, +) -> list[str]: + """Return each year's and each total's figure that differs between the paths. + + Both results are in the budget-window format (``BudgetWindowResult``); + a year one path has and the other lacks is a difference too. + """ + v2_years = {impact["year"]: impact for impact in v2_result["annualImpacts"]} + old_years = {impact["year"]: impact for impact in old_result["annualImpacts"]} + differences = [ + f"{year}: only one path has it" + for year in sorted(set(v2_years) ^ set(old_years)) + ] + rows = [ + *( + (year, v2_years[year], old_years[year]) + for year in sorted(v2_years) + if year in old_years + ), + ("Total", v2_result["totals"], old_result["totals"]), + ] + differences.extend( + f"{year} {field}: v2 {v2_row[field]} against the old gateway's {old_row[field]}" + for year, v2_row, old_row in rows + for field in BUDGET_FIELDS + if figures_differ( + v2_row[field], old_row[field], relative=relative, absolute=absolute + ) + ) + return differences + + +@pytest.fixture(scope="module") +def base_url() -> str: + """Return the staging stable URL, without a trailing slash.""" + return _required("SIMULATION_ENTRYPOINT_TEST_BASE_URL").rstrip("/") + + +@pytest.fixture(scope="module") +def headers() -> dict[str, str]: + """Return the test client's bearer token for the entrypoint's audience.""" + issuer = _required("SIMULATION_ENTRYPOINT_TEST_AUTH_ISSUER").rstrip("/") + response = httpx.post( + f"{issuer}/oauth/token", + json={ + "client_id": _required("SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_ID"), + "client_secret": _required("SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET"), + "audience": _required("SIMULATION_ENTRYPOINT_TEST_AUTH_AUDIENCE"), + "grant_type": "client_credentials", + }, + timeout=30, + ) + response.raise_for_status() + token = response.json().get("access_token") + if not isinstance(token, str) or not token: + msg = "Auth0 response omitted access_token." + raise RuntimeError(msg) + return {"Authorization": f"Bearer {token}"} + + +def _wait(read, finished_statuses: set[str], deadline: float) -> dict[str, Any]: + while True: + body = read() + if body["status"] in finished_statuses: + return body + if time.monotonic() > deadline: + pytest.fail(f"still {body['status']} at the deadline") + time.sleep(POLL_INTERVAL_SECONDS) + + +def test_v2_budget_windows_match_the_old_gateways_and_recover( + base_url: str, headers: dict[str, str] +) -> None: + """Submit both windows, compare them, and run the v2 one again.""" + country = _optional("SIMULATION_QUALIFICATION_COUNTRY", "us") + region = _optional("SIMULATION_QUALIFICATION_REGION", "state/ut") + start_year = int(_optional("SIMULATION_QUALIFICATION_START_YEAR", "2026")) + window_size = int(_optional("SIMULATION_QUALIFICATION_WINDOW_SIZE", "7")) + relative = float(_optional("SIMULATION_QUALIFICATION_RELATIVE_TOLERANCE", "0.001")) + absolute = float(_optional("SIMULATION_QUALIFICATION_ABSOLUTE_TOLERANCE", "1000")) + deadline = time.monotonic() + 60 * float( + _optional("SIMULATION_QUALIFICATION_TIMEOUT_MINUTES", "120") + ) + policyengine_version = _required("SIMULATION_QUALIFICATION_POLICYENGINE_VERSION") + assert window_size > MAX_CONCURRENT_YEARS, ( + "the window must have queued years to qualify their recovery" + ) + + submitted = httpx.post( + f"{base_url}/v2/budget-window-reports", + json={ + "country": country, + "policyengine_version": policyengine_version, + "start_year": start_year, + "window_size": window_size, + "region_code": region, + "baseline": {}, + "reform": { + "policy_id": _required("SIMULATION_QUALIFICATION_REFORM_POLICY_ID") + }, + }, + headers=headers, + timeout=300, + ) + assert submitted.status_code == 201, submitted.text + report_id = submitted.json()["id"] + old_submitted = httpx.post( + f"{base_url}/simulate/economy/budget-window", + json={ + "country": country, + "version": _required("SIMULATION_QUALIFICATION_MODEL_VERSION"), + "region": region, + "scope": "macro", + "reform": json.loads(_required("SIMULATION_QUALIFICATION_REFORM")), + "start_year": str(start_year), + "window_size": window_size, + }, + headers=headers, + timeout=300, + ) + assert old_submitted.status_code == 200, old_submitted.text + batch_job_id = old_submitted.json()["batch_job_id"] + + def read_v2() -> dict[str, Any]: + response = httpx.get( + f"{base_url}/v2/budget-window-reports/{report_id}", + headers=headers, + timeout=120, + ) + response.raise_for_status() + return response.json() + + def read_old() -> dict[str, Any]: + response = httpx.get( + f"{base_url}/budget-window-jobs/{batch_job_id}", + headers=headers, + timeout=120, + ) + response.raise_for_status() + return response.json() + + v2_report = _wait(read_v2, FINISHED, deadline) + old_status = _wait(read_old, OLD_GATEWAY_FINISHED, deadline) + + assert v2_report["status"] == "succeeded", v2_report["run"]["error"] + assert old_status["status"] == "complete", old_status.get("error") + assert len(v2_report["run"]["years"]) == window_size + assert ( + window_differences( + v2_report["result"], + old_status["result"], + relative=relative, + absolute=absolute, + ) + == [] + ) + + rerun = httpx.post( + f"{base_url}/v2/budget-window-reports/{report_id}/runs", + json={"policyengine_version": policyengine_version, "trigger": "manual"}, + headers=headers, + timeout=300, + ) + assert rerun.status_code == 201, rerun.text + rerun_report = _wait(read_v2, FINISHED, deadline) + assert rerun_report["run"]["id"] == rerun.json()["id"] + assert rerun_report["status"] == "succeeded", rerun_report["run"]["error"] + assert rerun_report["result"] == v2_report["result"] diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py b/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py new file mode 100644 index 000000000..83aab9a42 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py @@ -0,0 +1,91 @@ +"""The staging budget-window qualification's comparison with the old gateway.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +QUALIFICATION = ( + Path(__file__).resolve().parents[1] + / "authenticated_tests" + / "test_staging_budget_window_qualification.py" +) + + +def _qualification(): + spec = importlib.util.spec_from_file_location( + "budget_window_qualification", QUALIFICATION + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _result(*budgetary_impacts: float) -> dict: + def row(year: str, budgetary_impact: float) -> dict: + return { + "year": year, + "taxRevenueImpact": 100.0, + "federalTaxRevenueImpact": 80.0, + "stateTaxRevenueImpact": 20.0, + "benefitSpendingImpact": 50.0, + "budgetaryImpact": budgetary_impact, + } + + return { + "annualImpacts": [ + row(str(2026 + offset), value) + for offset, value in enumerate(budgetary_impacts) + ], + "totals": row("Total", sum(budgetary_impacts)), + } + + +def test_windows_within_the_tolerance_have_no_differences() -> None: + qualification = _qualification() + + assert ( + qualification.window_differences( + _result(1_000_000.0, 2_000_000.0), + _result(1_000_500.0, 1_999_000.0), + relative=0.001, + absolute=0.0, + ) + == [] + ) + + +def test_a_year_and_its_total_outside_the_tolerance_are_named() -> None: + qualification = _qualification() + + differences = qualification.window_differences( + _result(1_000_000.0, 2_000_000.0), + _result(1_000_000.0, 2_100_000.0), + relative=0.001, + absolute=1000.0, + ) + + assert differences == [ + "2027 budgetaryImpact: v2 2000000.0 against the old gateway's 2100000.0", + "Total budgetaryImpact: v2 3000000.0 against the old gateway's 3100000.0", + ] + + +def test_a_year_only_one_path_has_is_a_difference() -> None: + qualification = _qualification() + + assert qualification.window_differences( + _result(1.0, 2.0), _result(1.0), relative=0.001, absolute=1000.0 + ) == ["2027: only one path has it"] + + +def test_the_qualification_runs_only_when_an_operator_asks(monkeypatch) -> None: + monkeypatch.delenv("SIMULATION_STAGING_BUDGET_WINDOW_QUALIFICATION", raising=False) + qualification = _qualification() + + marker = qualification.pytestmark.mark + assert marker.name == "skipif" + assert marker.args == (True,) + assert qualification.MAX_CONCURRENT_YEARS == 5 From 5eacd8aa7c83a97dc297204de5ccd807bc6566a7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:27:27 +0400 Subject: [PATCH 09/20] Describe budget-window reports and runs in the persistence README Co-Authored-By: Claude Opus 5.5 --- .../README.md | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 5eb232973..ba9b077fd 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -13,9 +13,10 @@ writes them with DML only. | Table | Holds | | --- | --- | | `simulations` | What to run, independent of model versions: kind (`economy` or `household`), country, policy (null for current law), dynamic, and either a household or a region code (`state/ca`) with an optional dataset name (null means the region's default dataset in each catalog version) | -| `reports` | Which two simulations a report compares (an economy report needs a reform; a household report may have only a baseline), both of the report's own country and kind, with its year and options | +| `reports` | Which two simulations a report compares (an economy report needs a reform; a household report may have only a baseline), both of the report's own country and kind, with its year and options, and, for a budget-window report, its window size (1 to 75 years from its year) | | `simulation_runs` | One execution of a simulation for a catalog version and year: the region code and dataset row it resolved in that version (economy runs only), its status, Modal call ID, timings, error, trigger (`initial`, `manual` or `system`) and attempt number | -| `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call | +| `report_runs` | One calculation of a report: a baseline and a reform simulation run of the same catalog version and year (no reform run for a single-simulation household report), the trigger, and, for an economy report, the comparison step's own status and Modal call; a budget-window report's report run names the budget-window run it is a year of | +| `budget_window_runs` | One calculation of a budget-window report on a catalog version: its trigger, whether its years calculate their simulations again, how many years it calculates at once, its status and error, when the safety check next looks at it, and the lease of the caller starting its next years | `simulations` and `reports` are definitions: the runtime role may read and insert them but not update or delete them. Rows migrated from the Public API's v1 database keep their v1 @@ -70,6 +71,20 @@ steps whose simulation run failed) are failed so the work can be submitted again. A comparison step can start only once both of its simulation runs have succeeded; starting it is left to the stage 2 report path. +### Budget-window runs + +A budget-window run groups one report run per year of its report's window, +each of the same report and catalog version (a composite foreign key) and at +most one per year (`uq_report_runs_budget_window_run_id_year`). +`BudgetWindowRunStore` creates runs (`create_first_run` adds one only while +the report has none, under a transaction-level advisory lock, as first report +runs are), reads each year's report run with both of its simulation runs, and +reads each year's `budget_impacts` row; nothing stores a total, so readers sum +the years. A caller starting a run's next years first claims it +(`claim_scheduling`, a lease on `scheduling_claimed_until`), and releasing the +claim sets the run's next safety-check time. A run is finished once, with a +conditional write, as `succeeded` or `failed`. + ### Results Every economy result is stored as rows (`result_tables`), one table per @@ -179,7 +194,8 @@ These rules are registered in `rules/registry.toml` and checked by The initial revision grants the runtime role its privileges, including default privileges on tables the migrator creates later; revision 0002 revokes `UPDATE` on the definition tables, which never change, and revision 0004 on -the result tables, whose rows are written once. Both roles must exist +the result tables, whose rows are written once; revision 0005's +`budget_window_runs` takes the default privileges. Both roles must exist before the first migration runs; an operator creates them. Before revision 0002 is deployed, the owner of the Public API's tables must also grant the migrator `REFERENCES` on `policies`, `households` and `dynamics`, and From 1e80ee09f2eb90ad1b6c4282cdd471ac7367d93f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:48:34 +0400 Subject: [PATCH 10/20] Give v2 budget-window results models of their own The budget-window report's result used the old gateway's models, so a change there would have changed this contract and its generated client, and the document advertised the old gateway's SPM provenance, which v2 budget windows never have. BudgetWindowReportResult, BudgetWindowYearImpact and BudgetWindowTotalImpact describe the same JSON on the v2 side, without it. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_reads.py | 22 ++++---- .../budget_window_schemas.py | 53 ++++++++++++++++--- .../tests/test_budget_window_reads.py | 3 ++ .../tests/test_openapi.py | 15 ++++++ 4 files changed, 75 insertions(+), 18 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py index be78b7b47..c1b746187 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py @@ -5,7 +5,7 @@ error; the others are queued, or never run once the run has finished. A run has failed as soon as one year has, before the scheduler records it, and has succeeded once every year has. Once it has succeeded, its result is today's -budget-window format (``BudgetWindowResult``), built from each year's stored +budget-window format (``BudgetWindowReportResult``), built from each year's stored budget row. The totals are summed in decimal arithmetic, so a ten-year sum of large figures carries no binary rounding error, then returned as floats, as the old gateway's budget windows are. @@ -17,19 +17,17 @@ from decimal import Decimal from typing import TYPE_CHECKING, Protocol -from policyengine_simulation_contract.gateway_models import ( - BudgetWindowAnnualImpact, - BudgetWindowResult, - BudgetWindowTotals, -) from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_persistence import BudgetWindowStatus, ReportKind from policyengine_simulation_entry.budget_window_reports import window_state from policyengine_simulation_entry.budget_window_schemas import ( BudgetWindowReportResponse, + BudgetWindowReportResult, BudgetWindowRunResponse, BudgetWindowRunsResponse, + BudgetWindowTotalImpact, + BudgetWindowYearImpact, BudgetWindowYearResponse, ) from policyengine_simulation_entry.economy_report_types import ReportNotFoundError @@ -96,8 +94,8 @@ def year_budgets(self, run_id: UUID) -> list[YearBudget]: ... -def _annual_impact(budget: YearBudget) -> BudgetWindowAnnualImpact: - return BudgetWindowAnnualImpact( +def _annual_impact(budget: YearBudget) -> BudgetWindowYearImpact: + return BudgetWindowYearImpact( year=str(budget.year), taxRevenueImpact=budget.tax_revenue_impact, federalTaxRevenueImpact=( @@ -109,7 +107,7 @@ def _annual_impact(budget: YearBudget) -> BudgetWindowAnnualImpact: ) -def _totals(annual_impacts: list[BudgetWindowAnnualImpact]) -> BudgetWindowTotals: +def _totals(annual_impacts: list[BudgetWindowYearImpact]) -> BudgetWindowTotalImpact: """Sum the years' figures in decimal arithmetic, from each float's shortest form.""" sums = { name: sum( @@ -118,7 +116,7 @@ def _totals(annual_impacts: list[BudgetWindowAnnualImpact]) -> BudgetWindowTotal ) for name in _BUDGET_FIELDS } - return BudgetWindowTotals( + return BudgetWindowTotalImpact( taxRevenueImpact=float(sums["taxRevenueImpact"]), federalTaxRevenueImpact=float(sums["federalTaxRevenueImpact"]), stateTaxRevenueImpact=float(sums["stateTaxRevenueImpact"]), @@ -129,7 +127,7 @@ def _totals(annual_impacts: list[BudgetWindowAnnualImpact]) -> BudgetWindowTotal def budget_window_result( report: ReportRecord, budgets: list[YearBudget] -) -> BudgetWindowResult: +) -> BudgetWindowReportResult: """Return the window's result from each year's budget figures. Raise LookupError when a year of the window has no stored figures. @@ -141,7 +139,7 @@ def budget_window_result( msg = f"years {missing} of the window have no stored budget figures" raise LookupError(msg) annual_impacts = [_annual_impact(budget) for budget in budgets] - return BudgetWindowResult( + return BudgetWindowReportResult( startYear=str(years[0]), endYear=str(years[-1]), windowSize=len(years), diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py index 755c98f26..60f8a6cfa 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py @@ -1,17 +1,19 @@ """Request and response bodies of the v2 budget-window report routes. -A budget-window report's result is today's budget-window format -(``BudgetWindowResult``): each year's budget figures and their totals, which -are summed when the report is read. +A budget-window report's result is today's budget-window format, the JSON +the old gateway's budget windows answer with: each year's budget figures and +their totals, which are summed when the report is read. Its models are the v2 +path's own (``BudgetWindowReportResult``), so a change to the old gateway's +models never changes this contract; they leave out the old gateway's SPM +provenance, which v2 budget windows never have. """ from __future__ import annotations from datetime import datetime -from typing import Annotated, Self +from typing import Annotated, Literal, Self from uuid import UUID -from policyengine_simulation_contract.gateway_models import BudgetWindowResult from policyengine_simulation_contract.job_progress import JobStage from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator @@ -62,6 +64,45 @@ def _window_ends_in_time(self) -> Self: return self +class BudgetWindowYearImpact(BaseModel): + """One year's budget figures; the federal figure is the total less the state's.""" + + model_config = ConfigDict(extra="forbid") + + year: str + taxRevenueImpact: float # noqa: N815 # today's budget-window JSON field + federalTaxRevenueImpact: float # noqa: N815 # today's budget-window JSON field + stateTaxRevenueImpact: float # noqa: N815 # today's budget-window JSON field + benefitSpendingImpact: float # noqa: N815 # today's budget-window JSON field + budgetaryImpact: float # noqa: N815 # today's budget-window JSON field + + +class BudgetWindowTotalImpact(BaseModel): + """The sums of every year's budget figures.""" + + model_config = ConfigDict(extra="forbid") + + year: Literal["Total"] = "Total" + taxRevenueImpact: float # noqa: N815 # today's budget-window JSON field + federalTaxRevenueImpact: float # noqa: N815 # today's budget-window JSON field + stateTaxRevenueImpact: float # noqa: N815 # today's budget-window JSON field + benefitSpendingImpact: float # noqa: N815 # today's budget-window JSON field + budgetaryImpact: float # noqa: N815 # today's budget-window JSON field + + +class BudgetWindowReportResult(BaseModel): + """A budget window's result: each year's figures, in year order, and totals.""" + + model_config = ConfigDict(extra="forbid") + + kind: Literal["budgetWindow"] = "budgetWindow" + startYear: str # noqa: N815 # today's budget-window JSON field + endYear: str # noqa: N815 # today's budget-window JSON field + windowSize: int # noqa: N815 # today's budget-window JSON field + annualImpacts: list[BudgetWindowYearImpact] # noqa: N815 # today's budget-window JSON field + totals: BudgetWindowTotalImpact + + class BudgetWindowYearResponse(BaseModel): """One year of a budget-window run. @@ -122,4 +163,4 @@ class BudgetWindowReportResponse(BaseModel): reform_simulation: SimulationResponse status: ReportStatus run: BudgetWindowRunResponse | None - result: BudgetWindowResult | None + result: BudgetWindowReportResult | None diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py index f9ec743e4..3bc67156b 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reads.py @@ -14,6 +14,7 @@ StoredWindowResultError, ) from policyengine_simulation_entry.budget_window_schemas import ( + BudgetWindowReportResult, BudgetWindowReportSubmissionRequest, ) from policyengine_simulation_entry.economy_report_types import ReportNotFoundError @@ -169,6 +170,8 @@ def test_a_finished_window_sums_its_years_exactly_when_read() -> None: # Binary floats would make 0.6000000000000001. assert result.totals.budgetaryImpact == 0.6 assert result.totals.year == "Total" + assert isinstance(result, BudgetWindowReportResult) + assert "spm_provenance" not in result.model_dump()["annualImpacts"][0] first = result.annualImpacts[0] assert first.taxRevenueImpact == 2.0 assert first.stateTaxRevenueImpact == 0.5 diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 1287f2bd3..a777a09e9 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -52,9 +52,12 @@ class OpenAPIDocument(TypedDict): # old gateway has none of them or their bodies. V2_SCHEMAS = ( "BudgetWindowReportResponse", + "BudgetWindowReportResult", "BudgetWindowReportSubmissionRequest", "BudgetWindowRunResponse", "BudgetWindowRunsResponse", + "BudgetWindowTotalImpact", + "BudgetWindowYearImpact", "BudgetWindowYearResponse", "ComparisonResponse", "EconomyErrorResponse", @@ -238,6 +241,18 @@ def test_v2_economy_routes_are_public_for_the_generated_client(): assert set(V2_SCHEMAS) <= set(spec["components"]["schemas"]) +def test_v2_budget_window_results_have_models_of_their_own(): + """The old gateway's budget-window models stay the old gateway's.""" + spec = cast("dict[str, Any]", create_app().openapi()) + report = spec["components"]["schemas"]["BudgetWindowReportResponse"] + + assert report["properties"]["result"]["anyOf"][0] == { + "$ref": "#/components/schemas/BudgetWindowReportResult" + } + year = spec["components"]["schemas"]["BudgetWindowYearImpact"] + assert "spm_provenance" not in year["properties"] + + def test_v2_budget_window_routes_are_public_for_the_generated_client(): spec = cast("OpenAPIDocument", create_app().openapi()) ids = operation_ids(spec) From 081b4508132bef47a3ebbcc4c9916fee45880204 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:52:02 +0400 Subject: [PATCH 11/20] Refuse a budget-window re-run that cannot run, before anything is written A re-run on a catalog version that lacks the report's region or dataset was recorded as a run that failed when its first year could not start, so the report's newest run became a failed one. It is refused with 422 now, as a single-year re-run is: EconomyReports.check_setup resolves both simulations' rows before the run is written. A system run returns the report's newest run before the catalog is read, as the single-year path does, so a migrated reader's first opening never fails on a version check. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_reports.py | 6 +++++ .../economy_reports.py | 12 ++++++++++ .../tests/test_budget_window_reports.py | 24 +++++++++++++++++++ 3 files changed, 42 insertions(+) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index 685ccf1bd..faa104a12 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -295,9 +295,15 @@ async def run_report( if report.options: msg = "economy reports take no options yet" raise EconomyRequestError(msg) + if trigger is RunTrigger.SYSTEM: + newest = await asyncio.to_thread(self.windows.latest_run, report_id) + if newest is not None: + return newest version = await self.economy.catalog_version( report.country, policyengine_version ) + # Refuse a run its catalog version cannot resolve before it is written. + await self.economy.check_setup(await self.economy.report_setup(report, version)) new_run = self._new_run(report, version, trigger, force=force) if trigger is RunTrigger.SYSTEM: first = await asyncio.to_thread( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py index 301a1e889..98febf83c 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py @@ -278,6 +278,18 @@ async def report_setup( simulations.append(simulation) return ReportSetup(report, version, (simulations[0], simulations[1])) + async def check_setup(self, setup: ReportSetup) -> None: + """Raise EconomyRequestError unless both simulations' rows resolve. + + Each simulation's region and dataset must exist in the setup's + catalog version, in a form a call can carry, so a caller can refuse + a report before it writes anything. + """ + for simulation in setup.simulations: + await self._population( + setup.version, simulation.region_code, simulation.dataset_name + ) + async def run_year( self, setup: ReportSetup, diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py index f1724333c..5378a1b45 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py @@ -241,6 +241,30 @@ def test_a_system_run_of_a_report_that_has_a_run_returns_its_newest() -> None: assert path.dispatcher.requests == [] +def test_a_rerun_its_version_cannot_run_is_refused_before_anything_is_written() -> None: + path = WindowPath() + report, first = path.submit_window(window_size=2) + del path.catalog.populations["state/ca"] + path.dispatcher.requests.clear() + + with pytest.raises(EconomyRequestError, match="state/ca"): + _rerun(path, report.id) + + assert path.windows.runs(report.id) == [path.windows.get_run(first.id)] + assert path.dispatcher.requests == [] + + +def test_a_system_run_returns_the_newest_run_before_reading_the_catalog() -> None: + path = WindowPath() + report, first = path.submit_window(window_size=2) + + again = _rerun( + path, report.id, trigger=RunTrigger.SYSTEM, policyengine_version="9.9.9" + ) + + assert again.id == first.id + + def test_a_system_run_of_a_report_without_one_calculates_it() -> None: path = WindowPath() setup = path.prepare(window_size=2) From c5606e00a59ef4442657850379174ccd044e65cb Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:53:15 +0400 Subject: [PATCH 12/20] Release only the scheduling claim the caller holds Releasing a budget-window run's scheduling claim cleared whatever claim it held, so a caller that finished after its lease ended could clear a newer caller's claim and let a third caller start years at the same time. claim_scheduling now returns when the claim ends, and release_scheduling releases, and sets the next check, only while the run still holds that claim. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_store.py | 30 +++++++---- .../tests/test_budget_window_runs.py | 54 +++++++++++++++---- .../budget_window_reports.py | 20 ++++--- .../tests/budget_window_fakes.py | 15 ++++-- 4 files changed, 86 insertions(+), 33 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py index 318d45fc1..4e3066df5 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py @@ -2,8 +2,10 @@ A budget-window run starts ``running``. Whoever starts its next years first claims it (``claim_scheduling``): a claim is a lease, so a caller that stopped -loses it once the lease ends, and another caller can carry on. Releasing the -claim sets when the safety check next looks at the run. The run is finished +loses it once the lease ends, and another caller can carry on. A caller +releases only its own claim, named by the time it lasts until, so a caller +that finishes after its lease ended leaves the newer claim alone; releasing +sets when the safety check next looks at the run. The run is finished once, with a conditional write, as ``succeeded`` or ``failed``. A run's years are its report runs, read with both simulation runs of each. @@ -275,12 +277,15 @@ def year_budgets(self, run_id: UUID) -> list[YearBudget]: # --- Scheduling and finishing ------------------------------------------ - def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: - """Claim a running run to start its next years; return whether this did. + def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> datetime | None: + """Claim a running run to start its next years. - A claim held by another caller is refused until its lease ends. + Return when the claim ends, which names it for ``release_scheduling``, + or None when another caller holds a claim whose lease has not ended, + or the run has finished. """ now = self._clock() + claimed_until = now + lease with self.engine.begin() as connection: claimed = connection.execute( update(budget_window_runs) @@ -292,12 +297,18 @@ def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: budget_window_runs.c.scheduling_claimed_until <= now, ), ) - .values(scheduling_claimed_until=now + lease, updated_at=now) + .values(scheduling_claimed_until=claimed_until, updated_at=now) ) - return claimed.rowcount == 1 + return claimed_until if claimed.rowcount == 1 else None - def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: - """Release a running run's claim and set when it is next checked.""" + def release_scheduling( + self, run_id: UUID, *, claimed_until: datetime, next_check_at: datetime + ) -> None: + """Release this caller's claim on a running run and set its next check. + + ``claimed_until`` is what ``claim_scheduling`` returned. A claim + another caller took after this one's lease ended is left in place. + """ now = self._clock() with self.engine.begin() as connection: connection.execute( @@ -305,6 +316,7 @@ def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: .where( budget_window_runs.c.id == run_id, budget_window_runs.c.status == BudgetWindowStatus.RUNNING.value, + budget_window_runs.c.scheduling_claimed_until == claimed_until, ) .values( scheduling_claimed_until=None, diff --git a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py index d71eabb59..618a8dfb8 100644 --- a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py +++ b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py @@ -337,17 +337,45 @@ def test_one_scheduler_claims_a_run_at_a_time( ) lease = timedelta(minutes=2) - assert store.claim_scheduling(run.id, lease=lease) is True - assert store.claim_scheduling(run.id, lease=lease) is False - store.release_scheduling(run.id, next_check_at=clock() + FIRST_CHECK) - assert store.claim_scheduling(run.id, lease=lease) is True + first = store.claim_scheduling(run.id, lease=lease) + assert first == clock() + lease + assert store.claim_scheduling(run.id, lease=lease) is None + store.release_scheduling( + run.id, claimed_until=first, next_check_at=clock() + FIRST_CHECK + ) + assert store.claim_scheduling(run.id, lease=lease) is not None clock.advance(lease) # A scheduler that stopped without releasing its claim loses it. - assert store.claim_scheduling(run.id, lease=lease) is True - released = store.get_run(run.id) - assert released is not None + assert store.claim_scheduling(run.id, lease=lease) == clock() + lease store.finish_run(run.id, BudgetWindowStatus.SUCCEEDED) - assert store.claim_scheduling(run.id, lease=lease) is False + assert store.claim_scheduling(run.id, lease=lease) is None + + +def test_a_late_release_leaves_a_newer_claim_in_place( + migrated_database, public_api_rows +) -> None: + report, _ = budget_window_report(migrated_database, public_api_rows) + clock = ManualClock() + store = budget_window_store(migrated_database, clock) + run = store.create_run( + new_budget_window_run(report, public_api_rows), first_check_at=clock() + ) + lease = timedelta(minutes=2) + expired = store.claim_scheduling(run.id, lease=lease) + clock.advance(lease) + current = store.claim_scheduling(run.id, lease=lease) + assert expired is not None + assert current is not None + + # The first scheduler finishes after its lease ended. + store.release_scheduling( + run.id, claimed_until=expired, next_check_at=clock() + FIRST_CHECK + ) + + held = store.get_run(run.id) + assert held is not None + assert held.scheduling_claimed_until == current + assert store.claim_scheduling(run.id, lease=lease) is None SIMULTANEOUS_SCHEDULERS = 6 @@ -366,7 +394,7 @@ def test_simultaneous_schedulers_claim_a_run_once( def claim(_: int) -> bool: all_ready.wait() - return store.claim_scheduling(run.id, lease=timedelta(minutes=2)) + return store.claim_scheduling(run.id, lease=timedelta(minutes=2)) is not None with ThreadPoolExecutor(max_workers=SIMULTANEOUS_SCHEDULERS) as pool: claims = list(pool.map(claim, range(SIMULTANEOUS_SCHEDULERS))) @@ -383,13 +411,17 @@ def test_releasing_a_claim_sets_the_next_check( run = store.create_run( new_budget_window_run(report, public_api_rows), first_check_at=clock() ) - store.claim_scheduling(run.id, lease=timedelta(minutes=2)) + claimed_until = store.claim_scheduling(run.id, lease=timedelta(minutes=2)) + assert claimed_until is not None - store.release_scheduling(run.id, next_check_at=clock() + FIRST_CHECK) + store.release_scheduling( + run.id, claimed_until=claimed_until, next_check_at=clock() + FIRST_CHECK + ) released = store.get_run(run.id) assert released is not None assert released.next_check_at == clock() + FIRST_CHECK + assert released.scheduling_claimed_until is None def test_due_runs_are_the_running_runs_whose_check_is_due( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index faa104a12..1213b7e47 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -140,12 +140,14 @@ def year_budgets(self, run_id: UUID) -> list[YearBudget]: """Return each stored year's budget figures, in year order.""" ... - def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: - """Claim a running run to start its years.""" + def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> datetime | None: + """Claim a running run to start its years; return when the claim ends.""" ... - def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: - """Release the claim and set the next check.""" + def release_scheduling( + self, run_id: UUID, *, claimed_until: datetime, next_check_at: datetime + ) -> None: + """Release this caller's claim and set the next check.""" ... def finish_run( @@ -356,19 +358,21 @@ async def start_years(self, budget_window_run_id: UUID) -> int: """ started = 0 for _ in range(SCHEDULING_ATTEMPTS): - claimed = await asyncio.to_thread( + claimed_until = await asyncio.to_thread( lambda: self.windows.claim_scheduling( budget_window_run_id, lease=SCHEDULING_LEASE ) ) - if not claimed: + if claimed_until is None: break try: started += await self._schedule(budget_window_run_id) finally: await asyncio.to_thread( - lambda: self.windows.release_scheduling( - budget_window_run_id, next_check_at=self._next_check() + lambda claim=claimed_until: self.windows.release_scheduling( + budget_window_run_id, + claimed_until=claim, + next_check_at=self._next_check(), ) ) if not await self._has_room(budget_window_run_id): diff --git a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py index f022fba15..23d07b78b 100644 --- a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py +++ b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py @@ -165,7 +165,7 @@ def year_budgets(self, run_id: UUID) -> list[YearBudget]: if year.report_run.id in self.budgets ] - def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: + def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> datetime | None: row = self.rows[run_id] now = self.economy_runs.clock() if ( @@ -176,13 +176,18 @@ def claim_scheduling(self, run_id: UUID, *, lease: timedelta) -> bool: and row.scheduling_claimed_until > now ) ): - return False + return None self.rows[run_id] = replace(row, scheduling_claimed_until=now + lease) - return True + return now + lease - def release_scheduling(self, run_id: UUID, *, next_check_at: datetime) -> None: + def release_scheduling( + self, run_id: UUID, *, claimed_until: datetime, next_check_at: datetime + ) -> None: row = self.rows[run_id] - if row.status is BudgetWindowStatus.RUNNING: + if ( + row.status is BudgetWindowStatus.RUNNING + and row.scheduling_claimed_until == claimed_until + ): self.rows[run_id] = replace( row, scheduling_claimed_until=None, next_check_at=next_check_at ) From 5312ea3e367e5cbe6e5e0dd5a84eb6660fe98c1f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:54:50 +0400 Subject: [PATCH 13/20] Treat only a lost year's insert as another caller's year The scheduler took any integrity error while starting a year as another caller having started it, so a different constraint's refusal would have left that year unstarted and the run running indefinitely. It now checks that the error names the one-report-run-per-year constraint (ONE_REPORT_RUN_PER_YEAR) and raises any other, after releasing its claim. Tests cover that branch and the second pass after a year finishes while the claim is held. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 2 + .../budget_window_runs.py | 4 + .../tests/test_budget_window_runs.py | 3 +- .../budget_window_reports.py | 5 +- .../tests/test_budget_window_reports.py | 76 +++++++++++++++++++ 5 files changed, 88 insertions(+), 2 deletions(-) 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 c18029fd6..3a7d63d70 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -5,6 +5,7 @@ """ from policyengine_simulation_persistence.budget_window_runs import ( + ONE_REPORT_RUN_PER_YEAR, BudgetWindowRunRecord, BudgetWindowStatus, BudgetWindowYear, @@ -98,6 +99,7 @@ "CATALOG_READ_TABLES", "HEAD_REVISION", "MAX_BUDGET_WINDOW_YEARS", + "ONE_REPORT_RUN_PER_YEAR", "OPTIONAL_SECTIONS", "OUTPUT_FILE_RETENTION", "RESULT_FIELD_STORAGE", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py index c8f3f985a..3cc6baf94 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_runs.py @@ -60,6 +60,10 @@ class BudgetWindowStatus(StrEnum): FAILED = "failed" +# The unique constraint that allows one report run per year of a +# budget-window run; a caller that loses a year's insert reads its name. +ONE_REPORT_RUN_PER_YEAR = "uq_report_runs_budget_window_run_id_year" + FINISHED_BUDGET_WINDOW_STATUSES = ( BudgetWindowStatus.SUCCEEDED, BudgetWindowStatus.FAILED, diff --git a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py index 618a8dfb8..2dd8ab33a 100644 --- a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py +++ b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py @@ -20,6 +20,7 @@ ) from migration_support import alembic_config, migration_cycle_failures from policyengine_simulation_persistence import ( + ONE_REPORT_RUN_PER_YEAR, BudgetWindowStatus, NewBudgetWindowRun, NewReport, @@ -208,7 +209,7 @@ def test_each_year_is_one_report_run_of_the_budget_window_run( WINDOW_START, WINDOW_START + 1, } - with pytest.raises(IntegrityError, match="uq_report_runs_budget_window_run_id"): + with pytest.raises(IntegrityError, match=ONE_REPORT_RUN_PER_YEAR): year_report_run( migrated_database, public_api_rows, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index 1213b7e47..c1d360b6c 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -38,6 +38,7 @@ BUDGET_WINDOW_MAX_CONCURRENT_YEARS, ) from policyengine_simulation_persistence import ( + ONE_REPORT_RUN_PER_YEAR, BudgetWindowStatus, NewBudgetWindowRun, ReportKind, @@ -463,7 +464,9 @@ async def _start_queued( ) return started # Another caller started this year after its claim's lease ended. - except IntegrityError: + except IntegrityError as error: + if ONE_REPORT_RUN_PER_YEAR not in str(error): + raise logger.warning( "budget_window_year_already_started", extra={"budget_window_run_id": str(run.id), "year": year}, diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py index 5378a1b45..5a7457000 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py @@ -28,11 +28,13 @@ SimulationChoice, ) from policyengine_simulation_persistence import ( + ONE_REPORT_RUN_PER_YEAR, BudgetWindowStatus, NewBudgetWindowRun, RunStatus, RunTrigger, ) +from sqlalchemy.exc import IntegrityError def _rerun(path: WindowPath, report_id, **changes: object): @@ -362,3 +364,77 @@ def test_one_run_that_cannot_resume_does_not_stop_the_others() -> None: assert started == 1 assert path.year_runs(new_runs[0].id) == {} assert sorted(path.year_runs(new_runs[1].id)) == [2026] + + +def _integrity_error(constraint: str) -> IntegrityError: + statement = "INSERT INTO simulation.report_runs" + return IntegrityError( + statement, {}, Exception(f'violates unique constraint "{constraint}"') + ) + + +def test_a_year_another_caller_started_ends_the_pass_without_failing() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + finish_year(path.runs, path.year_runs(run.id)[2026]) + real_create = path.runs.create_report_run + conflicts = [_integrity_error(ONE_REPORT_RUN_PER_YEAR)] + + def lose_the_year(new_report_run): + # Another caller, whose claim's lease had ended, records the year first. + winner = real_create(new_report_run) + if conflicts: + raise conflicts.pop() + return winner + + path.runs.create_report_run = lose_the_year # type: ignore[method-assign] + + started = _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) + + still_running = path.windows.get_run(run.id) + assert still_running is not None + assert still_running.status is BudgetWindowStatus.RUNNING + assert started == 0 + assert sorted(path.year_runs(run.id)) == [2026, 2027, 2028] + + +def test_any_other_integrity_error_is_raised_and_the_claim_released() -> None: + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + finish_year(path.runs, path.year_runs(run.id)[2026]) + + other_constraint = "fk_report_runs_budget_window_run" + + def refuse(_new_report_run): + raise _integrity_error(other_constraint) + + path.runs.create_report_run = refuse # type: ignore[method-assign] + + with pytest.raises(IntegrityError): + _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) + + released = path.windows.get_run(run.id) + assert released is not None + assert released.scheduling_claimed_until is None + + +def test_a_year_that_finished_while_the_claim_was_held_frees_its_place() -> None: + """The caller read its years just before a year's comparison was recorded.""" + path = WindowPath(max_concurrent_years=2) + _, run = path.submit_window(window_size=4) + real_years = path.windows.years + reads: list[int] = [] + + def read_then_finish_a_year(run_id): + years = real_years(run_id) + reads.append(len(years)) + if len(reads) == 1: + finish_year(path.runs, path.year_runs(run.id)[2026]) + return years + + path.windows.years = read_then_finish_a_year # type: ignore[method-assign] + + started = asyncio.run(path.budget_windows.start_years(run.id)) + + assert started == 1 + assert sorted(path.year_runs(run.id)) == [2026, 2027, 2028] From f4428728c6470dbc183b52f404d29167ba35f760 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:55:43 +0400 Subject: [PATCH 14/20] Make the v2 functions' container limits agree with each other The economy simulation functions' 40 containers were all for four budget windows, leaving none for other reports, and the US segment function's 600 could not run the segments of the 40 US simulation calls those functions allow, let alone Stage 12's own. Now the simulation functions keep 40 (three windows at the cap and the 10 they had for other reports), the comparison function has 25 (three windows' comparisons and its 10), and the US segment function 1,000: the 20 segments of every v2 and Stage 12 US simulation call that may run at once, so no segment waits behind the limit, which would count against its caller's time limit. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 22 +++++----- .../tests/test_economy_run_capacity.py | 44 ++++++++++++------- .../tests/test_stage12_modal_app.py | 2 +- 3 files changed, 41 insertions(+), 27 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 4650ab614..5eef4b9d9 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -287,9 +287,11 @@ def validate_worker_uk() -> dict: memory=32768, timeout=3600, retries=0, - # Three US national budget windows at the Simulation Entrypoint's cap of - # five years each: two simulations a year, 20 segments each. - max_containers=600, + # Room for the segments of every US simulation call that may run at once + # (40 v2 economy simulations and 10 Stage 12 ones, 20 segments each), so + # no segment waits behind this limit, which would count against its + # caller's time limit. + max_containers=1000, secrets=worker_secrets, ) def run_single_simulation_segment_us( @@ -655,8 +657,8 @@ def segmented_calculator( memory=32768, timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, - # Four budget windows at the Simulation Entrypoint's cap of five years, - # two simulations a year. + # Three budget windows at the Simulation Entrypoint's cap of five years, + # two simulations a year, and the 10 it had for other reports. max_containers=40, secrets=worker_secrets, ) @@ -673,8 +675,8 @@ def run_economy_simulation_us( memory=32768, timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, - # Four budget windows at the Simulation Entrypoint's cap of five years, - # two simulations a year. + # Three budget windows at the Simulation Entrypoint's cap of five years, + # two simulations a year, and the 10 it had for other reports. max_containers=40, secrets=worker_secrets, ) @@ -691,9 +693,9 @@ def run_economy_simulation_uk( memory=8192, timeout=ECONOMY_COMPARISON_TIMEOUT_SECONDS, retries=0, - # Four budget windows at the Simulation Entrypoint's cap of five years, - # one comparison a year, with as much room again for other reports. - max_containers=40, + # Three budget windows at the Simulation Entrypoint's cap of five years, + # one comparison a year, and the 10 it had for other reports. + max_containers=25, secrets=worker_secrets, ) def compare_economy_simulations( diff --git a/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py b/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py index 73d1cb88b..c7c6990c6 100644 --- a/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py +++ b/projects/policyengine-simulation-executor/tests/test_economy_run_capacity.py @@ -2,8 +2,9 @@ The Simulation Entrypoint calculates at most ``BUDGET_WINDOW_MAX_CONCURRENT_YEARS`` years of a budget window at once, each -year up to two simulation calls and one comparison call, and a US national -simulation call runs one segment call per region group. +year up to two simulation calls and one comparison call. The functions keep +the room they had for other reports, and a US national simulation call runs +one segment call per region group. """ from __future__ import annotations @@ -18,9 +19,10 @@ ) from test_stage12_modal_app import _load -# Budget windows each function has room for at once, every year at the cap. -WINDOWS_AT_ONCE = 4 -US_NATIONAL_WINDOWS_AT_ONCE = 3 +# Budget windows each function has room for at once, every year at the cap, +# beside the containers the functions had for other reports. +WINDOWS_AT_ONCE = 3 +OTHER_REPORTS = 10 SIMULATIONS_PER_YEAR = 2 @@ -28,7 +30,7 @@ def test_a_budget_window_calculates_a_few_years_at_once() -> None: assert BUDGET_WINDOW_MAX_CONCURRENT_YEARS == 5 -def test_the_simulation_and_comparison_functions_hold_four_windows( +def test_the_simulation_and_comparison_functions_hold_three_windows_and_more( monkeypatch, ) -> None: functions = dict(_load(monkeypatch).app.function_calls) @@ -36,20 +38,30 @@ def test_the_simulation_and_comparison_functions_hold_four_windows( for country in ("us", "uk"): simulation = functions[ECONOMY_SIMULATION_FUNCTIONS[country]] - assert simulation["max_containers"] == SIMULATIONS_PER_YEAR * years == 40 - assert functions[ECONOMY_COMPARISON_FUNCTION]["max_containers"] == 2 * years + assert ( + simulation["max_containers"] + == (SIMULATIONS_PER_YEAR * years + OTHER_REPORTS) + == 40 + ) + assert ( + functions[ECONOMY_COMPARISON_FUNCTION]["max_containers"] + == (years + OTHER_REPORTS) + == 25 + ) -def test_the_us_segment_function_holds_three_national_windows(monkeypatch) -> None: +def test_every_us_simulation_call_at_once_has_room_for_its_segments( + monkeypatch, +) -> None: + """A segment queued behind the limit counts against its caller's time limit.""" functions = dict(_load(monkeypatch).app.function_calls) + us_simulation_calls = ( + functions[ECONOMY_SIMULATION_FUNCTIONS["us"]]["max_containers"] + + functions["run_single_simulation_us"]["max_containers"] + ) assert ( functions["run_single_simulation_segment_us"]["max_containers"] - == ( - US_NATIONAL_WINDOWS_AT_ONCE - * BUDGET_WINDOW_MAX_CONCURRENT_YEARS - * SIMULATIONS_PER_YEAR - * len(US_NATIONAL_REGION_GROUPS) - ) - == 600 + == (us_simulation_calls * len(US_NATIONAL_REGION_GROUPS)) + == 1000 ) 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 d482667cb..34084c75c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -342,7 +342,7 @@ 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"] == 600 + 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 From 302dccb6abe40180887b1162f5b4f0c9823c1774 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:57:05 +0400 Subject: [PATCH 15/20] Describe report runs' progress in a module of its own, and fix doc drift report_run_progress and the run error response it builds move from the economy routes to report_progress, which the economy reads, the budget-window scheduler and the budget-window reads all use, so the scheduler no longer reads statuses from a routes module and the error builder is not written twice. window_state returns the public WindowState. Revision 0005's docstring no longer says nothing writes its table, the safety-check resumption's docstring says how a failed run is tried again, and a comprehension in BudgetWindowRunStore.years no longer reuses its parameter's name. No behaviour changes. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_store.py | 6 +- .../versions/0005_budget_window_runs.py | 3 +- .../budget_window_reads.py | 19 +++--- .../budget_window_reports.py | 14 ++-- .../economy_routes.py | 48 ++------------ .../report_progress.py | 66 +++++++++++++++++++ 6 files changed, 94 insertions(+), 62 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_progress.py diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py index 4e3066df5..56345953c 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_store.py @@ -228,13 +228,13 @@ def years(self, run_id: UUID) -> list[BudgetWindowYear]: ).mappings() ] simulation_run_ids = { - run_id + simulation_run_id for year_run in year_runs - for run_id in ( + for simulation_run_id in ( year_run.baseline_simulation_run_id, year_run.reform_simulation_run_id, ) - if run_id is not None + if simulation_run_id is not None } by_id = { row["id"]: SimulationRunRecord.from_row(row) diff --git a/migrations/versions/0005_budget_window_runs.py b/migrations/versions/0005_budget_window_runs.py index f10fb2289..6fab5d27f 100644 --- a/migrations/versions/0005_budget_window_runs.py +++ b/migrations/versions/0005_budget_window_runs.py @@ -10,8 +10,7 @@ ``budget_window_run_id``: the budget-window run it is a year of, of the same report and catalog version (a composite foreign key), with at most one report run per year. The default privileges let the runtime role read, insert and -update the new table, and not delete from it. Nothing writes the new columns -or table yet. +update the new table, and not delete from it. Revision ID: 0005 Revises: 0004 diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py index c1b746187..2430bf392 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py @@ -31,14 +31,15 @@ BudgetWindowYearResponse, ) from policyengine_simulation_entry.economy_report_types import ReportNotFoundError -from policyengine_simulation_entry.economy_routes import ( - report_run_progress, - simulation_response, -) +from policyengine_simulation_entry.economy_routes import simulation_response from policyengine_simulation_entry.economy_schemas import ( ReportStatus, RunErrorResponse, ) +from policyengine_simulation_entry.report_progress import ( + report_run_progress, + run_error_response, +) if TYPE_CHECKING: from uuid import UUID @@ -148,20 +149,18 @@ def budget_window_result( ) -def _error(code: str | None, message: str | None) -> RunErrorResponse | None: - return None if code is None else RunErrorResponse(code=code, message=message) - - def _run_status( run: BudgetWindowRunRecord, report: ReportRecord, years: list[BudgetWindowYear], ) -> tuple[ReportStatus, RunErrorResponse | None]: if run.status is BudgetWindowStatus.FAILED: - return ReportStatus.FAILED, _error(run.error_code, run.error_message) + return ReportStatus.FAILED, run_error_response( + run.error_code, run.error_message + ) state = window_state(report, years) if state.failed is not None: - return ReportStatus.FAILED, _error( + return ReportStatus.FAILED, run_error_response( state.failed.code.value, state.failed.message ) if run.status is BudgetWindowStatus.SUCCEEDED or state.succeeded: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index c1d360b6c..a440e876f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -53,8 +53,8 @@ EconomyRequestError, ReportNotFoundError, ) -from policyengine_simulation_entry.economy_routes import report_run_progress from policyengine_simulation_entry.economy_schemas import ReportStatus +from policyengine_simulation_entry.report_progress import report_run_progress from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL if TYPE_CHECKING: @@ -185,7 +185,7 @@ class BudgetWindowReportSubmission: @dataclass(frozen=True) -class _WindowState: +class WindowState: """What a run's started years say about it.""" failed: RunError | None @@ -201,7 +201,7 @@ def _run_error_code(code: str) -> RunErrorCode: return RunErrorCode.WORKER_ERROR -def window_state(report: ReportRecord, years: list[BudgetWindowYear]) -> _WindowState: +def window_state(report: ReportRecord, years: list[BudgetWindowYear]) -> WindowState: """Return whether a year failed, whether every year succeeded, and how many run. A year is failed, succeeded or running as its report run is @@ -225,7 +225,7 @@ def window_state(report: ReportRecord, years: list[BudgetWindowYear]) -> _Window succeeded_years += 1 elif status is not ReportStatus.FAILED: in_flight += 1 - return _WindowState( + return WindowState( failed=failed, succeeded=succeeded_years == len(report.window_years), in_flight=in_flight, @@ -545,7 +545,11 @@ async def resume_waiting_work(self, *, now: datetime) -> int: return started async def _resume(self, budget_window_run_id: UUID) -> int: - """Start one due run's years; a failure is logged and leaves it due.""" + """Start one due run's years; a failure is logged, and a later pass retries. + + Releasing the claim moves the run's next check one interval on, so the + next pass due after that tries it again. + """ try: return await self.start_years(budget_window_run_id) except Exception as error: # noqa: BLE001 # logged; the next pass tries again 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 370e4b36c..c3f38bc37 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 @@ -36,7 +36,6 @@ from uuid import UUID # noqa: TC003 # FastAPI evaluates the handlers' path parameters from fastapi import Depends, FastAPI, HTTPException, status -from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_persistence import ReportKind, RunStatus, RunTrigger from sqlalchemy.exc import DBAPIError @@ -57,10 +56,13 @@ EconomyReportRunsResponse, EconomyReportSubmissionRequest, ReportStatus, - RunErrorResponse, SimulationResponse, SimulationRunResponse, ) +from policyengine_simulation_entry.report_progress import ( + report_run_progress, + run_error_response, +) if TYPE_CHECKING: from collections.abc import Awaitable, Callable @@ -153,10 +155,6 @@ def report_runs(self, report_id: UUID) -> list[ReportRunRecord]: ... -def _error(code: str | None, message: str | None) -> RunErrorResponse | None: - return None if code is None else RunErrorResponse(code=code, message=message) - - def simulation_response(record: SimulationRecord) -> SimulationResponse: """Describe a simulation.""" return SimulationResponse( @@ -179,47 +177,13 @@ def _simulation_run(record: SimulationRunRecord) -> SimulationRunResponse: status=record.status.value, attempt=record.attempt, trigger=record.trigger.value, - error=_error(record.error_code, record.error_message), + error=run_error_response(record.error_code, record.error_message), dispatched_at=record.dispatched_at, completed_at=record.completed_at, created_at=record.created_at, ) -def report_run_progress( - report_run: ReportRunRecord, - runs: tuple[SimulationRunRecord, SimulationRunRecord], -) -> tuple[ReportStatus, JobStage | None, RunErrorResponse | None]: - """Return a report run's status, its stage while it runs, and its error. - - A failed simulation run fails the report run at once, before the safety - check fails its comparison step. - """ - failed = next((run for run in runs if run.status is RunStatus.FAILED), None) - if failed is not None: - return ( - ReportStatus.FAILED, - None, - _error(failed.error_code, failed.error_message), - ) - comparison = report_run.comparison_status - if comparison is RunStatus.FAILED: - return ( - ReportStatus.FAILED, - None, - _error( - report_run.comparison_error_code, report_run.comparison_error_message - ), - ) - if comparison is RunStatus.SUCCEEDED: - return ReportStatus.SUCCEEDED, None, None - if all(run.status is RunStatus.SUCCEEDED for run in runs): - return ReportStatus.RUNNING, JobStage.OUTPUTS, None - if all(run.status is RunStatus.PENDING for run in runs): - return ReportStatus.PENDING, JobStage.QUEUED, None - return ReportStatus.RUNNING, JobStage.SIMULATING, None - - @dataclass class EconomyReportReads: """Build report and report run responses from the stored rows.""" @@ -262,7 +226,7 @@ def report_run(self, report_run: ReportRunRecord) -> EconomyReportRunResponse: reform_run=_simulation_run(reform), comparison=ComparisonResponse( status=(report_run.comparison_status or RunStatus.PENDING).value, - error=_error( + error=run_error_response( report_run.comparison_error_code, report_run.comparison_error_message, ), diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_progress.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_progress.py new file mode 100644 index 000000000..9d182ace6 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_progress.py @@ -0,0 +1,66 @@ +"""A report run's status, stage and error, from its rows. + +The economy report reads and the budget-window scheduler and reads all +describe a report run the same way: failed as soon as a simulation run or its +comparison failed, succeeded once its comparison succeeded, and otherwise +pending until a call starts and running after. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_persistence import RunStatus + +from policyengine_simulation_entry.economy_schemas import ( + ReportStatus, + RunErrorResponse, +) + +if TYPE_CHECKING: + from policyengine_simulation_persistence import ( + ReportRunRecord, + SimulationRunRecord, + ) + + +def run_error_response( + code: str | None, message: str | None +) -> RunErrorResponse | None: + """Describe a run's error, or None when it has none.""" + return None if code is None else RunErrorResponse(code=code, message=message) + + +def report_run_progress( + report_run: ReportRunRecord, + runs: tuple[SimulationRunRecord, SimulationRunRecord], +) -> tuple[ReportStatus, JobStage | None, RunErrorResponse | None]: + """Return a report run's status, its stage while it runs, and its error. + + A failed simulation run fails the report run at once, before the safety + check fails its comparison step. + """ + failed = next((run for run in runs if run.status is RunStatus.FAILED), None) + if failed is not None: + return ( + ReportStatus.FAILED, + None, + run_error_response(failed.error_code, failed.error_message), + ) + comparison = report_run.comparison_status + if comparison is RunStatus.FAILED: + return ( + ReportStatus.FAILED, + None, + run_error_response( + report_run.comparison_error_code, report_run.comparison_error_message + ), + ) + if comparison is RunStatus.SUCCEEDED: + return ReportStatus.SUCCEEDED, None, None + if all(run.status is RunStatus.SUCCEEDED for run in runs): + return ReportStatus.RUNNING, JobStage.OUTPUTS, None + if all(run.status is RunStatus.PENDING for run in runs): + return ReportStatus.PENDING, JobStage.QUEUED, None + return ReportStatus.RUNNING, JobStage.SIMULATING, None From 0e3cfdd19eccfb58c3087bd987356284275e1805 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:57:50 +0400 Subject: [PATCH 16/20] Read a failed old-gateway batch in the qualification, and describe recovery The old gateway's batch status route answers 500, with the batch's status and error, when a batch failed; the qualification raised on that answer instead of reporting the batch's error. old_gateway_status reads it, and raises any other error answer. The cap comes from the contract instead of a copy, and the module describes the operator's recovery check: promoting or rolling the entrypoint while the window runs, which it must survive. Co-Authored-By: Claude Opus 5.5 --- ...est_staging_budget_window_qualification.py | 47 +++++++++++++++---- .../tests/test_budget_window_qualification.py | 37 ++++++++++++++- 2 files changed, 73 insertions(+), 11 deletions(-) diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py b/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py index 7297efe54..56d624c9e 100644 --- a/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py +++ b/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py @@ -37,6 +37,14 @@ ``_RELATIVE_TOLERANCE`` (0.001), ``_ABSOLUTE_TOLERANCE`` (1000, in the country's currency) and ``_TIMEOUT_MINUTES`` (120). +**Recovery after the caller stops** is an operator action beside it: while +the v2 window runs (its run reads with queued years), promote a new +entrypoint revision, or roll back to the previous one and forward again, +with the deploy and rollback workflows. The window must still finish within +the timeout: its later years are then started by the new revision's call +reports and by the safety check, not by the instance that took the +submission. + For example: uv run --frozen pytest \ @@ -56,14 +64,16 @@ import httpx import pytest +from policyengine_simulation_contract.economy_runs import ( + BUDGET_WINDOW_MAX_CONCURRENT_YEARS, +) POLL_INTERVAL_SECONDS = 30 FINISHED = {"succeeded", "failed"} OLD_GATEWAY_FINISHED = {"complete", "failed"} -# Sim calculates at most this many years of a window at once -# (BUDGET_WINDOW_MAX_CONCURRENT_YEARS); a longer window needs recovery of -# its queued years. -MAX_CONCURRENT_YEARS = 5 +# Sim calculates at most this many years of a window at once; a longer +# window needs recovery of its queued years. +MAX_CONCURRENT_YEARS = BUDGET_WINDOW_MAX_CONCURRENT_YEARS BUDGET_FIELDS = ( "taxRevenueImpact", "federalTaxRevenueImpact", @@ -90,6 +100,23 @@ def _optional(name: str, default: str) -> str: return os.environ.get(name, "").strip() or default +def old_gateway_status(response: httpx.Response) -> dict[str, Any]: + """Return an old-gateway batch's status from its status route's answer. + + The route answers 500 with the batch's status and error when the batch + failed; any other error answer is raised. + """ + if response.status_code == httpx.codes.INTERNAL_SERVER_ERROR: + try: + body = response.json() + except ValueError: + body = None + if isinstance(body, dict) and body.get("status") == "failed": + return body + response.raise_for_status() + return response.json() + + def figures_differ( v2_value: float, old_value: float, *, relative: float, absolute: float ) -> bool: @@ -235,13 +262,13 @@ def read_v2() -> dict[str, Any]: return response.json() def read_old() -> dict[str, Any]: - response = httpx.get( - f"{base_url}/budget-window-jobs/{batch_job_id}", - headers=headers, - timeout=120, + return old_gateway_status( + httpx.get( + f"{base_url}/budget-window-jobs/{batch_job_id}", + headers=headers, + timeout=120, + ) ) - response.raise_for_status() - return response.json() v2_report = _wait(read_v2, FINISHED, deadline) old_status = _wait(read_old, OLD_GATEWAY_FINISHED, deadline) diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py b/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py index 83aab9a42..0e4169bc0 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py @@ -88,4 +88,39 @@ def test_the_qualification_runs_only_when_an_operator_asks(monkeypatch) -> None: marker = qualification.pytestmark.mark assert marker.name == "skipif" assert marker.args == (True,) - assert qualification.MAX_CONCURRENT_YEARS == 5 + + +def test_a_failed_old_gateway_batch_is_read_as_its_status() -> None: + """The old gateway answers 500 with the batch's status when it failed.""" + import httpx + + qualification = _qualification() + request = httpx.Request("GET", "https://staging.example/budget-window-jobs/fc-1") + failed = httpx.Response( + 500, json={"status": "failed", "error": "year 2027 failed"}, request=request + ) + running = httpx.Response(200, json={"status": "running"}, request=request) + + assert qualification.old_gateway_status(failed) == { + "status": "failed", + "error": "year 2027 failed", + } + assert qualification.old_gateway_status(running) == {"status": "running"} + for unreadable in ( + httpx.Response(502, text="Bad Gateway", request=request), + httpx.Response(500, json={"detail": "Internal Server Error"}, request=request), + ): + try: + qualification.old_gateway_status(unreadable) + except httpx.HTTPStatusError: + continue + msg = f"{unreadable.status_code} was read as a status" + raise AssertionError(msg) + + +def test_the_qualification_uses_the_cap_sim_uses() -> None: + from policyengine_simulation_contract.economy_runs import ( + BUDGET_WINDOW_MAX_CONCURRENT_YEARS, + ) + + assert _qualification().MAX_CONCURRENT_YEARS is BUDGET_WINDOW_MAX_CONCURRENT_YEARS From 1801eb897cd9bf91e3272785131ece446de7f872 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:23:12 +0400 Subject: [PATCH 17/20] Start the runs a caller created for a year another caller recorded first When a caller whose scheduling lease had ended lost a year's report-run insert to the caller that took the lease over, the simulation runs it had created stayed pending until the safety check started them, ten minutes or more later: the other caller had reused them and so did not start them. RunStore.create_window_year_run now returns the report run already recorded for the year, and the losing caller starts the runs it created that the year compares, as a first open of a single-year report that loses does. The scheduler no longer needs to catch the conflict. Co-Authored-By: Claude Opus 5.5 --- .../run_store.py | 42 +++++++++ .../tests/test_budget_window_runs.py | 93 ++++++++++++++++++- .../budget_window_reports.py | 22 ++--- .../economy_report_types.py | 4 + .../economy_reports.py | 14 ++- .../tests/budget_window_fakes.py | 21 +++++ .../tests/test_budget_window_reports.py | 16 +++- 7 files changed, 188 insertions(+), 24 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py index b1f7d175a..46da3ef7d 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/run_store.py @@ -16,7 +16,11 @@ from sqlalchemy import func, select from sqlalchemy.dialects.postgresql import insert as postgres_insert +from sqlalchemy.exc import IntegrityError +from policyengine_simulation_persistence.budget_window_runs import ( + ONE_REPORT_RUN_PER_YEAR, +) from policyengine_simulation_persistence.output_files import ( OutputFileRecord, available_output_files, @@ -344,6 +348,44 @@ def create_first_report_run(self, new_report_run: NewReportRun) -> StartedReport run=self._insert_report_run(connection, new_report_run), created=True ) + def create_window_year_run(self, new_report_run: NewReportRun) -> StartedReportRun: + """Add a budget-window year's report run, or return the one already added. + + Callers that start the same year of a budget-window run at once meet + at the unique year (``ONE_REPORT_RUN_PER_YEAR``): the first adds the + report run, and the others return it, so a caller that loses can + still start the simulation runs it created that the year compares. + Raise ValueError when the report run names no budget-window run, and + otherwise as ``create_report_run`` does. + """ + window_run_id = new_report_run.budget_window_run_id + if window_run_id is None: + msg = "a budget-window year's report run names its budget-window run" + raise ValueError(msg) + with self.engine.begin() as connection: + try: + with connection.begin_nested(): + created = self._insert_report_run(connection, new_report_run) + except IntegrityError as error: + if ONE_REPORT_RUN_PER_YEAR not in str(error): + raise + baseline = self._run_key( + connection, new_report_run.baseline_simulation_run_id + ) + existing = ( + None + if baseline is None + else self._newest_report_run_where( + connection, + (report_runs.c.budget_window_run_id == window_run_id) + & (report_runs.c.year == baseline["year"]), + ) + ) + if existing is None: # pragma: no cover - the conflict proves a row + raise + return StartedReportRun(run=existing, created=False) + return StartedReportRun(run=created, created=True) + def _insert_report_run( self, connection: Connection, new_report_run: NewReportRun ) -> ReportRunRecord: diff --git a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py index 2dd8ab33a..64a5817f8 100644 --- a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py +++ b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py @@ -24,6 +24,7 @@ BudgetWindowStatus, NewBudgetWindowRun, NewReport, + NewReportRun, ReportKind, RunError, RunErrorCode, @@ -34,7 +35,7 @@ ) from public_api_support import superuser_engine from report_support import economy_pair, economy_report, report_store -from run_support import set_simulation_run_status +from run_support import economy_run, run_store, set_simulation_run_status from sqlalchemy import text from sqlalchemy.exc import IntegrityError, ProgrammingError from transition_support import ManualClock @@ -582,3 +583,93 @@ def test_the_runtime_role_updates_budget_window_runs_but_cannot_delete_them( connection.execute(text("DELETE FROM simulation.budget_window_runs")) finally: engine.dispose() + + +def _year_report_run_input(database, rows, window, *, year, budget_window_run_id): + report, simulations = window + store = run_store(database) + baseline_run, reform_run = ( + store.start_simulation_run(economy_run(simulation, rows, year=year)).run + for simulation in simulations + ) + return NewReportRun( + report_id=report.id, + baseline_simulation_run_id=baseline_run.id, + reform_simulation_run_id=reform_run.id, + trigger=RunTrigger.INITIAL, + budget_window_run_id=budget_window_run_id, + ) + + +def test_a_year_recorded_twice_returns_the_first_report_run( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + run = budget_window_store(migrated_database).create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + year_input = _year_report_run_input( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=run.id, + ) + store = run_store(migrated_database) + + first = store.create_window_year_run(year_input) + again = store.create_window_year_run(year_input) + + assert first.created is True + assert first.run.budget_window_run_id == run.id + assert again.created is False + assert again.run == first.run + + +SIMULTANEOUS_YEAR_STARTS = 6 + + +def test_simultaneous_starts_of_a_year_add_one_report_run( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + run = budget_window_store(migrated_database).create_run( + new_budget_window_run(report, public_api_rows), + first_check_at=ManualClock()(), + ) + year_input = _year_report_run_input( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=run.id, + ) + store = run_store(migrated_database) + all_ready = Barrier(SIMULTANEOUS_YEAR_STARTS) + + def start_the_year(_: int): + all_ready.wait() + return store.create_window_year_run(year_input) + + with ThreadPoolExecutor(max_workers=SIMULTANEOUS_YEAR_STARTS) as pool: + started = list(pool.map(start_the_year, range(SIMULTANEOUS_YEAR_STARTS))) + + assert sum(start.created for start in started) == 1 + assert {start.run.id for start in started} == {started[0].run.id} + + +def test_only_a_year_of_a_budget_window_run_is_recorded_as_one( + migrated_database, public_api_rows +) -> None: + report, simulations = budget_window_report(migrated_database, public_api_rows) + year_input = _year_report_run_input( + migrated_database, + public_api_rows, + (report, simulations), + year=WINDOW_START, + budget_window_run_id=None, + ) + + with pytest.raises(ValueError, match="budget-window run"): + run_store(migrated_database).create_window_year_run(year_input) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index a440e876f..13ecfa036 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -14,8 +14,10 @@ simulations of 20 segment calls each). The next years start when a year's call reports its end (``start_years_after_call``) and on each safety-check pass (``resume_waiting_work``). Whoever starts years first claims the run, a -lease of ``SCHEDULING_LEASE``, so two callers do not start the same year; -the database allows one report run per year of a run in any case. +lease of ``SCHEDULING_LEASE``, so two callers do not start the same year. +The database allows one report run per year of a run in any case: a caller +whose lease ended while it worked, and that loses a year's insert, shares +the other caller's report run and starts the simulation runs it created. ``YearStartingRecorder`` wraps the recording of a reported call, so its end starts the years it frees. The first year that fails fails the run, with that @@ -38,7 +40,6 @@ BUDGET_WINDOW_MAX_CONCURRENT_YEARS, ) from policyengine_simulation_persistence import ( - ONE_REPORT_RUN_PER_YEAR, BudgetWindowStatus, NewBudgetWindowRun, ReportKind, @@ -47,7 +48,6 @@ RunTrigger, ) from policyengine_simulation_persistence.run_store import utc_now -from sqlalchemy.exc import IntegrityError from policyengine_simulation_entry.economy_report_types import ( EconomyRequestError, @@ -433,8 +433,9 @@ async def _start_queued( ) -> int: """Start these years of the run in order; return how many started. - A year that cannot start fails the run, and one another caller - started ends this pass. + A year that cannot start fails the run. A year another caller + started at the same time is that caller's report run; this caller + starts the simulation runs it created for it. """ version = await asyncio.to_thread( self.catalog.catalog_version_by_id, run.tax_benefit_model_version_id @@ -463,15 +464,6 @@ async def _start_queued( ), ) return started - # Another caller started this year after its claim's lease ended. - except IntegrityError as error: - if ONE_REPORT_RUN_PER_YEAR not in str(error): - raise - logger.warning( - "budget_window_year_already_started", - extra={"budget_window_run_id": str(run.id), "year": year}, - ) - return started started += 1 return started diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_report_types.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_report_types.py index 31db24030..a360401fb 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_report_types.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_report_types.py @@ -126,6 +126,10 @@ def create_first_report_run(self, new_report_run: NewReportRun) -> StartedReport """Insert the report's first report run, or return its newest one.""" ... + def create_window_year_run(self, new_report_run: NewReportRun) -> StartedReportRun: + """Insert a budget-window year's report run, or return the one added.""" + ... + def latest_report_run(self, report_id: UUID) -> ReportRunRecord | None: """Return the report's newest report run, or None.""" ... diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py index 98febf83c..32bfb33e4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reports.py @@ -457,12 +457,16 @@ async def _record_report_run(self, new_report_run: NewReportRun) -> ReportRunRec Simultaneous first opens of a report without a run then share one report run. A caller that loses still starts the runs it created that the recorded run compares, and leaves the others pending for the - safety check. A year of a budget-window run is always its own run. + safety check. A year of a budget-window run is its own run, whatever + the trigger; a caller that starts a year another caller has just + started shares that caller's report run in the same way. """ - if ( - new_report_run.trigger is not RunTrigger.SYSTEM - or new_report_run.budget_window_run_id is not None - ): + if new_report_run.budget_window_run_id is not None: + year = await asyncio.to_thread( + self.runs.create_window_year_run, new_report_run + ) + return year.run + if new_report_run.trigger is not RunTrigger.SYSTEM: return await asyncio.to_thread(self.runs.create_report_run, new_report_run) first = await asyncio.to_thread( self.runs.create_first_report_run, new_report_run diff --git a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py index 23d07b78b..a453a6722 100644 --- a/projects/policyengine-simulation-entry/tests/budget_window_fakes.py +++ b/projects/policyengine-simulation-entry/tests/budget_window_fakes.py @@ -31,14 +31,17 @@ from policyengine_simulation_entry.economy_reports import EconomyReports from policyengine_simulation_entry.run_dispatch import RunDispatchService from policyengine_simulation_persistence import ( + ONE_REPORT_RUN_PER_YEAR, BudgetWindowRunRecord, BudgetWindowStatus, BudgetWindowYear, RunStatus, StartedBudgetWindowRun, + StartedReportRun, YearBudget, ) from run_fakes import FakeDispatcher +from sqlalchemy.exc import IntegrityError if TYPE_CHECKING: from datetime import datetime @@ -65,6 +68,24 @@ def create_report_run(self, new_report_run: NewReportRun) -> ReportRunRecord: self.report_runs_by_id[row.id] = row return row + def create_window_year_run(self, new_report_run: NewReportRun) -> StartedReportRun: + """Add a year's report run, or return the one another caller added.""" + try: + return StartedReportRun( + run=self.create_report_run(new_report_run), created=True + ) + except IntegrityError as error: + if ONE_REPORT_RUN_PER_YEAR not in str(error): + raise + year = self.simulation_runs[new_report_run.baseline_simulation_run_id].year + existing = next( + report_run + for report_run in self.report_runs_by_id.values() + if report_run.budget_window_run_id == new_report_run.budget_window_run_id + and report_run.year == year + ) + return StartedReportRun(run=existing, created=False) + @dataclass class FakeBudgetWindows: diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py index 5a7457000..5be12011e 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_reports.py @@ -373,7 +373,12 @@ def _integrity_error(constraint: str) -> IntegrityError: ) -def test_a_year_another_caller_started_ends_the_pass_without_failing() -> None: +def test_a_caller_that_loses_a_year_still_starts_the_runs_it_created() -> None: + """Another caller records the year first, comparing the runs this one created. + + That caller reused the runs, so only this one starts them, as a first + open that loses does for a single-year report. + """ path = WindowPath(max_concurrent_years=2) _, run = path.submit_window(window_size=4) finish_year(path.runs, path.year_runs(run.id)[2026]) @@ -388,14 +393,19 @@ def lose_the_year(new_report_run): return winner path.runs.create_report_run = lose_the_year # type: ignore[method-assign] + path.dispatcher.requests.clear() - started = _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) + _start_years_after(path, report_run_id=path.year_runs(run.id)[2026].id) still_running = path.windows.get_run(run.id) assert still_running is not None assert still_running.status is BudgetWindowStatus.RUNNING - assert started == 0 assert sorted(path.year_runs(run.id)) == [2026, 2027, 2028] + year_2028 = path.year_runs(run.id)[2028] + assert {call.simulation_run_id for call in path.simulation_calls()} == { + year_2028.baseline_simulation_run_id, + year_2028.reform_simulation_run_id, + } def test_any_other_integrity_error_is_raised_and_the_claim_released() -> None: From 27b66967d9756f68156b4cb3cd3b5f384215a61a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:24:25 +0400 Subject: [PATCH 18/20] Refuse to downgrade revision 0005 while budget-window reports exist Without their window size and their years' link, a budget-window report's years would read as report runs of a single-year report, and the economy report reads would show the newest year's run as the report's. The downgrade now raises while any report has a budget window, and removes the revision's objects only when none does. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_budget_window_runs.py | 10 ++++++++++ .../versions/0005_budget_window_runs.py | 19 +++++++++++++++++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py index 64a5817f8..101103020 100644 --- a/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py +++ b/libs/policyengine-simulation-persistence/tests/test_budget_window_runs.py @@ -45,6 +45,16 @@ FIRST_CHECK = timedelta(minutes=5) +def test_the_revision_is_not_downgraded_while_budget_windows_exist( + migrated_database, public_api_rows +) -> None: + """Their years would otherwise read as runs of single-year reports.""" + budget_window_report(migrated_database, public_api_rows) + + with pytest.raises(ProgrammingError, match="budget-window reports exist"): + command.downgrade(alembic_config(), "0004") + + def test_the_budget_window_revision_matches_its_tables_and_downgrades( disposable_database, ) -> None: diff --git a/migrations/versions/0005_budget_window_runs.py b/migrations/versions/0005_budget_window_runs.py index 6fab5d27f..1e0b8767a 100644 --- a/migrations/versions/0005_budget_window_runs.py +++ b/migrations/versions/0005_budget_window_runs.py @@ -10,7 +10,10 @@ ``budget_window_run_id``: the budget-window run it is a year of, of the same report and catalog version (a composite foreign key), with at most one report run per year. The default privileges let the runtime role read, insert and -update the new table, and not delete from it. +update the new table, and not delete from it. The downgrade refuses to run +while any budget-window report exists: without its window size and its +years' link, each year's report run would read as a run of a single-year +report. Revision ID: 0005 Revises: 0004 @@ -164,8 +167,20 @@ def upgrade() -> None: _add_budget_window_run_id() +def _refuse_while_budget_windows_exist() -> None: + op.execute( + "DO $$ BEGIN " + "IF EXISTS (SELECT 1 FROM simulation.reports " + "WHERE budget_window_size IS NOT NULL) " + "THEN RAISE EXCEPTION 'budget-window reports exist; revision 0005 is not " + "downgraded while they do'; " + "END IF; END $$" + ) + + def downgrade() -> None: - """Drop what ``upgrade`` added, with its indexes and keys.""" + """Drop what ``upgrade`` added, unless a budget-window report exists.""" + _refuse_while_budget_windows_exist() op.drop_constraint( "uq_report_runs_budget_window_run_id_year", "report_runs", From 3fd0489a478612de513b6f08801203a0d00740c3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:25:15 +0400 Subject: [PATCH 19/20] Share the economy path's re-run triggers, batch size and recorder protocol The budget-window path re-declared RERUN_TRIGGERS, WAITING_WORK_BATCH_SIZE, the reported-call recorder protocol and the recorded answers; it now uses the economy path's, so the two paths' refusals cannot drift apart (economy_recorders' recorded answers become the public RECORDED_STATUSES). The scheduler's docstring also says how far recovery reaches: a caller that stops between years is covered, and one that stops mid-dispatch fails its year after ten minutes, as for a single-year report. No behaviour changes. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_reports.py | 39 ++++++++++--------- .../economy_recorders.py | 7 +++- 2 files changed, 26 insertions(+), 20 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py index 13ecfa036..446fad95e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reports.py @@ -20,8 +20,14 @@ the other caller's report run and starts the simulation runs it created. ``YearStartingRecorder`` wraps the recording of a reported call, so its end -starts the years it frees. The first year that fails fails the run, with that -year's error, and no more +starts the years it frees. + +Recovery covers a caller that stops between years. A caller that stops after +claiming a simulation run's dispatch, before recording its Modal call, leaves +the run dispatching; the safety check fails it after ten minutes, which fails +its year and so the run, as it fails a single-year report. + +The first year that fails fails the run, with that year's error, and no more years start; years already running finish and keep their results. Once every year has succeeded, the run has succeeded. The run stores no totals: they are the sums of its years' budget rows, added up when it is read. @@ -35,7 +41,6 @@ from datetime import timedelta from typing import TYPE_CHECKING, Protocol -from policyengine_simulation_contract.call_reports import CallReportStatus from policyengine_simulation_contract.economy_runs import ( BUDGET_WINDOW_MAX_CONCURRENT_YEARS, ) @@ -49,10 +54,15 @@ ) from policyengine_simulation_persistence.run_store import utc_now +from policyengine_simulation_entry.economy_recorders import RECORDED_STATUSES from policyengine_simulation_entry.economy_report_types import ( EconomyRequestError, ReportNotFoundError, ) +from policyengine_simulation_entry.economy_reports import ( + RERUN_TRIGGERS, + WAITING_WORK_BATCH_SIZE, +) from policyengine_simulation_entry.economy_schemas import ReportStatus from policyengine_simulation_entry.report_progress import report_run_progress from policyengine_simulation_entry.run_recording import SAFETY_CHECK_INTERVAL @@ -62,7 +72,10 @@ from datetime import datetime from uuid import UUID - from policyengine_simulation_contract.call_reports import CallEndedReport + from policyengine_simulation_contract.call_reports import ( + CallEndedReport, + CallReportStatus, + ) from policyengine_simulation_persistence import ( BudgetWindowRunRecord, BudgetWindowYear, @@ -73,6 +86,9 @@ YearBudget, ) + from policyengine_simulation_entry.economy_recorders import ( + ReportedCallRecorder, + ) from policyengine_simulation_entry.economy_report_types import ( EconomyCatalog, EconomyReportSubmission, @@ -94,8 +110,6 @@ # scheduled: its years' calls report their ends, so this only catches a run # whose next years a stopped caller never started. BUDGET_WINDOW_CHECK_INTERVAL = SAFETY_CHECK_INTERVAL -RERUN_TRIGGERS = frozenset({RunTrigger.MANUAL, RunTrigger.SYSTEM}) -WAITING_WORK_BATCH_SIZE = 50 # A caller tries again after releasing its claim, in case a year finished # while it held the claim (its own read was then too early); at most this # many times. @@ -555,17 +569,6 @@ async def _resume(self, budget_window_run_id: UUID) -> int: return 0 -class ReportedCallRecorder(Protocol): - """Record a reported call's outcome (``RunRecorder``, possibly wrapped).""" - - async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: - """Record the call's outcome if Modal has one.""" - ... - - -_RECORDED = frozenset({CallReportStatus.RECORDED, CallReportStatus.ALREADY_RECORDED}) - - @dataclass(frozen=True) class YearStartingRecorder: """Record a reported call, then start the budget-window years its end frees. @@ -580,7 +583,7 @@ class YearStartingRecorder: async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: """Record the reported call; start the years it lets start.""" report_status = await self.recorder.record_reported_call(report) - if report_status in _RECORDED: + if report_status in RECORDED_STATUSES: try: await self.windows.start_years_after_call( simulation_run_ids=report.simulation_run_ids, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_recorders.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_recorders.py index 6cc342628..708569b9f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_recorders.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_recorders.py @@ -110,7 +110,10 @@ async def record_reported_call(self, report: CallEndedReport) -> CallReportStatu ... -_RECORDED = frozenset({CallReportStatus.RECORDED, CallReportStatus.ALREADY_RECORDED}) +# The answers that mean the reported call's outcome is recorded. +RECORDED_STATUSES = frozenset( + {CallReportStatus.RECORDED, CallReportStatus.ALREADY_RECORDED} +) @dataclass(frozen=True) @@ -129,7 +132,7 @@ class ComparisonStartingRecorder: async def record_reported_call(self, report: CallEndedReport) -> CallReportStatus: """Record the reported call; start the comparisons it lets start.""" report_status = await self.recorder.record_reported_call(report) - if report.simulation_run_ids and report_status in _RECORDED: + if report.simulation_run_ids and report_status in RECORDED_STATUSES: try: await self.comparisons.start_comparisons_after( report.simulation_run_ids From b02b5897d16ac267650e05cf806864b4f7dfec84 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:46:45 +0400 Subject: [PATCH 20/20] Describe the budget-window routes' own refusals, and tidy the final review's points The budget-window routes published the economy routes' 422 and 503 descriptions, which name economy reports; they now have their own. The window-size limit comes from the persistence library, the totals are summed field by field so the type checker sees each one, and the staging qualification also reads a failed old-gateway batch answered with 400 (a failure carrying the caller's errors) and types its polling helper. Co-Authored-By: Claude Opus 5.5 --- ...est_staging_budget_window_qualification.py | 17 ++++++--- .../budget_window_reads.py | 35 ++++++++----------- .../budget_window_routes.py | 21 +++++++++-- .../budget_window_schemas.py | 5 +-- .../tests/test_budget_window_qualification.py | 15 ++++++++ .../tests/test_openapi.py | 8 +++++ 6 files changed, 73 insertions(+), 28 deletions(-) diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py b/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py index 56d624c9e..f1d44720f 100644 --- a/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py +++ b/projects/policyengine-simulation-entry/authenticated_tests/test_staging_budget_window_qualification.py @@ -60,7 +60,7 @@ import math import os import time -from typing import Any +from typing import TYPE_CHECKING, Any import httpx import pytest @@ -82,6 +82,9 @@ "budgetaryImpact", ) +if TYPE_CHECKING: + from collections.abc import Callable + pytestmark = pytest.mark.skipif( os.environ.get("SIMULATION_STAGING_BUDGET_WINDOW_QUALIFICATION") != "1", reason="an operator step: runs only when it is asked for", @@ -104,9 +107,13 @@ def old_gateway_status(response: httpx.Response) -> dict[str, Any]: """Return an old-gateway batch's status from its status route's answer. The route answers 500 with the batch's status and error when the batch - failed; any other error answer is raised. + failed, or 400 when the failure carries the caller's errors; any other + error answer is raised. """ - if response.status_code == httpx.codes.INTERNAL_SERVER_ERROR: + if response.status_code in { + httpx.codes.BAD_REQUEST, + httpx.codes.INTERNAL_SERVER_ERROR, + }: try: body = response.json() except ValueError: @@ -189,7 +196,9 @@ def headers() -> dict[str, str]: return {"Authorization": f"Bearer {token}"} -def _wait(read, finished_statuses: set[str], deadline: float) -> dict[str, Any]: +def _wait( + read: Callable[[], dict[str, Any]], finished_statuses: set[str], deadline: float +) -> dict[str, Any]: while True: body = read() if body["status"] in finished_statuses: diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py index 2430bf392..8e0bb32f5 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_reads.py @@ -57,14 +57,6 @@ SimulationDefinitions, ) -_BUDGET_FIELDS = ( - "taxRevenueImpact", - "federalTaxRevenueImpact", - "stateTaxRevenueImpact", - "benefitSpendingImpact", - "budgetaryImpact", -) - class StoredWindowResultError(Exception): """A succeeded budget-window run's stored figures cannot be read back.""" @@ -110,19 +102,22 @@ def _annual_impact(budget: YearBudget) -> BudgetWindowYearImpact: def _totals(annual_impacts: list[BudgetWindowYearImpact]) -> BudgetWindowTotalImpact: """Sum the years' figures in decimal arithmetic, from each float's shortest form.""" - sums = { - name: sum( - (Decimal(str(getattr(impact, name))) for impact in annual_impacts), - Decimal(0), - ) - for name in _BUDGET_FIELDS - } + + def total(figures: list[float]) -> float: + return float(sum((Decimal(str(figure)) for figure in figures), Decimal(0))) + return BudgetWindowTotalImpact( - taxRevenueImpact=float(sums["taxRevenueImpact"]), - federalTaxRevenueImpact=float(sums["federalTaxRevenueImpact"]), - stateTaxRevenueImpact=float(sums["stateTaxRevenueImpact"]), - benefitSpendingImpact=float(sums["benefitSpendingImpact"]), - budgetaryImpact=float(sums["budgetaryImpact"]), + taxRevenueImpact=total([year.taxRevenueImpact for year in annual_impacts]), + federalTaxRevenueImpact=total( + [year.federalTaxRevenueImpact for year in annual_impacts] + ), + stateTaxRevenueImpact=total( + [year.stateTaxRevenueImpact for year in annual_impacts] + ), + benefitSpendingImpact=total( + [year.benefitSpendingImpact for year in annual_impacts] + ), + budgetaryImpact=total([year.budgetaryImpact for year in annual_impacts]), ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py index 8271a49bc..fa444cd93 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_routes.py @@ -43,13 +43,12 @@ SimulationChoice, ) from policyengine_simulation_entry.economy_routes import ( - REFUSED, RESULT_UNREADABLE, - UNAVAILABLE, answer_route, ) from policyengine_simulation_entry.economy_schemas import ( EconomyErrorResponse, + EconomyRefusalResponse, EconomyReportRunRequest, ) @@ -72,6 +71,24 @@ "description": "No such budget-window report", }, } +REFUSED: dict[int | str, dict[str, Any]] = { + status.HTTP_422_UNPROCESSABLE_ENTITY: { + "model": EconomyRefusalResponse, + "description": ( + "Refused: sim cannot run the budget window (its reason as text), or " + "a request field is invalid (one entry per field)" + ), + }, +} +UNAVAILABLE: dict[int | str, dict[str, Any]] = { + status.HTTP_503_SERVICE_UNAVAILABLE: { + "model": EconomyErrorResponse, + "description": ( + "The simulation database cannot be reached, or budget-window " + "reports are not configured on this deployment" + ), + }, +} @dataclass(frozen=True) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py index 60f8a6cfa..acb09536a 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_schemas.py @@ -15,6 +15,7 @@ from uuid import UUID from policyengine_simulation_contract.job_progress import JobStage +from policyengine_simulation_persistence import MAX_BUDGET_WINDOW_YEARS from pydantic import BaseModel, ConfigDict, Field, JsonValue, model_validator from policyengine_simulation_entry.economy_schemas import ( @@ -28,9 +29,9 @@ ) # A window may end no later than this year, as the old gateway's budget -# windows may, and have at most this many years. +# windows may, and have at most as many years as a report may. LAST_WINDOW_YEAR = 2099 -MAX_WINDOW_YEARS = 75 +MAX_WINDOW_YEARS = MAX_BUDGET_WINDOW_YEARS class BudgetWindowReportSubmissionRequest(BaseModel): diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py b/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py index 0e4169bc0..2efa42986 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_qualification.py @@ -124,3 +124,18 @@ def test_the_qualification_uses_the_cap_sim_uses() -> None: ) assert _qualification().MAX_CONCURRENT_YEARS is BUDGET_WINDOW_MAX_CONCURRENT_YEARS + + +def test_a_failed_batch_with_caller_errors_is_read_as_its_status() -> None: + """The old gateway answers 400 for a failed batch that carries caller errors.""" + import httpx + + qualification = _qualification() + request = httpx.Request("GET", "https://staging.example/budget-window-jobs/fc-1") + failed = httpx.Response( + 400, + json={"status": "failed", "errors": [{"message": "no SPM for 2030"}]}, + request=request, + ) + + assert qualification.old_gateway_status(failed)["status"] == "failed" diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index a777a09e9..327c32775 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -339,3 +339,11 @@ def test_v2_routes_declare_each_error_answer_with_its_body(): "type": "array", }, ] + + +def test_v2_budget_window_routes_describe_their_own_refusals(): + spec = cast("dict[str, Any]", create_app().openapi()) + submit = spec["paths"]["/v2/budget-window-reports"]["post"]["responses"] + + assert "budget-window reports" in submit["503"]["description"] + assert "economy reports" not in submit["503"]["description"]