From 0015c12e06c19429df3ff0b3f711d71ea4a0b9e9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:42:59 +0400 Subject: [PATCH 01/18] Test re-running a stored economy report or simulation on a later catalog version (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/economy_rows.py | 106 ++- .../tests/test_economy_reruns.py | 722 ++++++++++++++++++ 2 files changed, 822 insertions(+), 6 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/tests/test_economy_reruns.py diff --git a/libs/policyengine-simulation-persistence/tests/economy_rows.py b/libs/policyengine-simulation-persistence/tests/economy_rows.py index 4cb6a9014..b203b6aea 100644 --- a/libs/policyengine-simulation-persistence/tests/economy_rows.py +++ b/libs/policyengine-simulation-persistence/tests/economy_rows.py @@ -44,6 +44,12 @@ country_package_version=COUNTRY_PACKAGE_VERSION, policyengine_version=POLICYENGINE_VERSION, ) +# The versions of a release after ``VERSIONS``, with its own catalog version. +LATER_POLICYENGINE_VERSION = "6.18.0" +LATER_VERSIONS = RunVersions( + country_package_version="1.770.0", + policyengine_version=LATER_POLICYENGINE_VERSION, +) OUTPUT_BUCKET_PREFIX = "gs://policyengine-v2-simulation-outputs-staging/" CURRENT_LAW = SimulationChoice(policy_id=None, dynamic_id=None) OPTIONS: dict[str, Any] = {"include_cliffs": False} @@ -75,13 +81,52 @@ def insert_economy_catalog( parameter_names: tuple[str, ...] = (), ) -> EconomyCatalog: """Insert the country's model, a catalog version, a dataset and a region.""" - catalog = insert_catalog_version( + return _insert_catalog_rows( connection, country=country, model_id=insert_model(connection, country), policyengine_version=POLICYENGINE_VERSION, parameter_names=parameter_names, ) + + +def insert_later_catalog( + connection: Connection, + economy: EconomyCatalog, + *, + policyengine_version: str = LATER_POLICYENGINE_VERSION, + parameter_names: tuple[str, ...] = (), +) -> EconomyCatalog: + """Insert a later catalog version of ``economy``'s model. + + It has a dataset and a region with the same name and code as + ``economy``'s, as a release publishing the next .py version's catalog + would. + """ + return _insert_catalog_rows( + connection, + country=economy.catalog.country, + model_id=economy.catalog.model_id, + policyengine_version=policyengine_version, + parameter_names=parameter_names, + ) + + +def _insert_catalog_rows( + connection: Connection, + *, + country: str, + model_id: UUID, + policyengine_version: str, + parameter_names: tuple[str, ...], +) -> EconomyCatalog: + catalog = insert_catalog_version( + connection, + country=country, + model_id=model_id, + policyengine_version=policyengine_version, + parameter_names=parameter_names, + ) dataset_id = insert_dataset(connection, catalog, f"populace_{country}_2024") region_id = insert_region( connection, catalog, country, default_dataset_id=dataset_id @@ -120,9 +165,6 @@ def insert_simulation_run( """Insert an economy simulation and one run of it, as the superuser.""" shape = shape or RunShape() simulation_id = uuid4() - run_id = uuid4() - finished = shape.status in {SimulationStatus.SUCCEEDED, SimulationStatus.FAILED} - succeeded = shape.status is SimulationStatus.SUCCEEDED connection.execute( insert(simulations).values( id=simulation_id, @@ -135,6 +177,23 @@ def insert_simulation_run( year=population.year, ) ) + return insert_run_of(connection, simulation_id, population, shape) + + +def insert_run_of( + connection: Connection, + simulation_id: UUID, + population: EconomyPopulation, + shape: RunShape | None = None, +) -> StoredRun: + """Insert a run of a stored simulation on ``population``, as the superuser. + + ``shape.choice`` is ignored: the run keeps its simulation's policy. + """ + shape = shape or RunShape() + run_id = uuid4() + finished = shape.status in {SimulationStatus.SUCCEEDED, SimulationStatus.FAILED} + succeeded = shape.status is SimulationStatus.SUCCEEDED connection.execute( insert(simulation_runs).values( id=run_id, @@ -179,9 +238,26 @@ def insert_report_run_naming( report run's ID. """ outcome = outcome or ReportRunOutcome() + report_id = insert_report_of(connection, economy, compared, inputs=outcome.inputs) + return insert_report_run_of(connection, report_id, compared, outcome) + + +def insert_report_of( # noqa: PLR0913 # the report's sides and what else it may set + connection: Connection, + economy: EconomyCatalog, + compared: tuple[StoredRun, StoredRun], + *, + inputs: dict[str, Any] | None = None, + budget_window_size: int | None = None, + household_id: UUID | None = None, +) -> UUID: + """Insert an economy report naming the two runs' simulations, as the superuser. + + ``compared`` holds the baseline's run, then the reform's. The report is + on ``economy``'s region and dataset in 2026. Return its ID. + """ baseline, reform = compared report_id = uuid4() - run_id = uuid4() connection.execute( insert(reports).values( id=report_id, @@ -191,12 +267,30 @@ def insert_report_run_naming( tax_benefit_model_id=economy.catalog.model_id, baseline_simulation_id=baseline.simulation_id, reform_simulation_id=reform.simulation_id, + household_id=household_id, dataset_id=economy.dataset_id, region_id=economy.region_id, year=2026, - inputs=OPTIONS if outcome.inputs is None else outcome.inputs, + inputs=OPTIONS if inputs is None else inputs, + budget_window_size=budget_window_size, ) ) + return report_id + + +def insert_report_run_of( + connection: Connection, + report_id: UUID, + compared: tuple[StoredRun, StoredRun], + outcome: ReportRunOutcome | None = None, +) -> UUID: + """Insert a run of a stored report comparing the two runs, as the superuser. + + Return the report run's ID; ``outcome.inputs`` is ignored. + """ + outcome = outcome or ReportRunOutcome() + baseline, reform = compared + run_id = uuid4() finished = outcome.status in {ReportRunStatus.SUCCEEDED, ReportRunStatus.FAILED} connection.execute( insert(report_runs).values( diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py new file mode 100644 index 000000000..9c38e25f3 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py @@ -0,0 +1,722 @@ +"""Re-running a stored economy report or simulation on a later catalog version. + +A re-run reads what the stored report or simulation calculates, creates a new +run of it on the current versions, and keeps every earlier run. A report's +side may reuse a run of its own simulation that succeeded on those versions; +nothing is created while the record has a run pending or running. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the simulation service's role, each +in a transaction rolled back afterwards. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, replace +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + LATER_POLICYENGINE_VERSION, + LATER_VERSIONS, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + EconomyCatalog, + ReportRunOutcome, + RunShape, + StoredRun, + insert_economy_catalog, + insert_later_catalog, + insert_report_of, + insert_report_run_of, + insert_run_of, + insert_simulation_run, +) +from policyengine_simulation_persistence import ( + EconomyRequest, + EconomyRunStore, + MissingPolicyParametersError, + NewReportRerun, + NewSimulationRerun, + ReportDefinition, + ReportRunStatus, + ReportRunTrigger, + ReportSideError, + RerunSide, + ReusableRun, + RunInProgressError, + RunReuse, + SideDefinition, + SimulationChoice, + SimulationDefinition, + SimulationStatus, + UnknownDatasetError, + UnknownRegionError, + find_reusable_run_of, + insert_report_rerun, + insert_simulation_rerun, + lock_reruns_of, + read_report_definition, + read_simulation_definition, +) +from public_api_rows import ( + StoredValue, + insert_dataset, + insert_household, + insert_policy, + insert_region, +) +from sqlalchemy import text +from sqlalchemy.exc import OperationalError + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +START = datetime(2026, 1, 1, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +# What a re-run asks the current catalog version for: the stored report's +# region and dataset, on the later .py version. +LATER_REQUEST = EconomyRequest( + country="us", + policyengine_version=LATER_POLICYENGINE_VERSION, + region_code="us", + dataset_name="populace_us_2024", + year=2026, +) +LATER_REUSE = RunReuse( + versions=LATER_VERSIONS, inputs=OPTIONS, output_uri_prefix=OUTPUT_BUCKET_PREFIX +) + + +@dataclass(frozen=True) +class StoredReport: + """A stored report, the simulation runs its first run compared, and that run.""" + + report_id: UUID + baseline: StoredRun + reform: StoredRun + report_run_id: UUID + + +def _stored_report( + connection: Connection, + economy: EconomyCatalog, + *, + reform_choice: SimulationChoice = CURRENT_LAW, + run_status: ReportRunStatus = ReportRunStatus.SUCCEEDED, +) -> StoredReport: + population = economy.population() + baseline = insert_simulation_run(connection, population) + reform = insert_simulation_run( + connection, population, RunShape(choice=reform_choice) + ) + report_id = insert_report_of(connection, economy, (baseline, reform)) + report_run_id = insert_report_run_of( + connection, report_id, (baseline, reform), ReportRunOutcome(status=run_status) + ) + return StoredReport( + report_id=report_id, + baseline=baseline, + reform=reform, + report_run_id=report_run_id, + ) + + +def _row(connection: Connection, table_name: str, row_id: UUID) -> dict: + row = connection.execute( + text(f"SELECT * FROM public.{table_name} WHERE id = :id"), # noqa: S608 # fixed table names + {"id": row_id}, + ).mappings() + return dict(row.one()) + + +def _count_runs_of(connection: Connection, simulation_id: UUID) -> int: + return connection.execute( + text("SELECT count(*) FROM public.simulation_runs WHERE simulation_id = :id"), + {"id": simulation_id}, + ).scalar_one() + + +def _store_on(connection: Connection) -> EconomyRunStore: + """Return a store whose transactions are savepoints of ``connection``.""" + + @contextmanager + def savepoint() -> Iterator[Connection]: + with connection.begin_nested(): + yield connection + + return EconomyRunStore(savepoint) + + +def test_a_stored_economy_reports_definition_is_read(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + policy_id = insert_policy(connection, economy.catalog) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + stored = _stored_report(connection, economy, reform_choice=reform_choice) + act_as_the_service(connection) + + definition = read_report_definition(connection, stored.report_id) + + assert definition == ReportDefinition( + id=stored.report_id, + country="us", + kind="economy", + year=2026, + region_code="us", + dataset_name="populace_us_2024", + inputs=OPTIONS, + baseline=SideDefinition( + simulation_id=stored.baseline.simulation_id, choice=CURRENT_LAW + ), + reform=SideDefinition( + simulation_id=stored.reform.simulation_id, choice=reform_choice + ), + ) + + +def test_a_budget_window_or_household_report_is_read_as_its_kind( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + population = economy.population() + sides = ( + insert_simulation_run(connection, population), + insert_simulation_run(connection, population), + ) + window_id = insert_report_of(connection, economy, sides, budget_window_size=3) + household_id = insert_household( + connection, country="us", default_year=2026, document={"people": {}} + ) + household_report_id = insert_report_of( + connection, economy, sides, household_id=household_id + ) + act_as_the_service(connection) + + window = read_report_definition(connection, window_id) + household = read_report_definition(connection, household_report_id) + + assert window is not None + assert household is not None + assert (window.kind, household.kind) == ("budget_window", "household") + + +def test_an_unknown_report_or_simulation_has_no_definition( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + act_as_the_service(connection) + + assert read_report_definition(connection, uuid4()) is None + assert read_simulation_definition(connection, uuid4()) is None + + +def test_a_stored_economy_simulations_definition_is_read( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, country="uk") + policy_id = insert_policy(connection, economy.catalog) + choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + stored = insert_simulation_run( + connection, economy.population(year=2027), RunShape(choice=choice) + ) + act_as_the_service(connection) + + definition = read_simulation_definition(connection, stored.simulation_id) + + assert definition == SimulationDefinition( + id=stored.simulation_id, + country="uk", + kind="economy", + year=2027, + region_code="uk", + dataset_name="populace_uk_2024", + choice=choice, + ) + + +def _find_later_run_of( + connection: Connection, + simulation_id: UUID, + later: EconomyCatalog, + *, + excluded_run_id: UUID | None = None, +) -> ReusableRun | None: + return find_reusable_run_of( + connection, + simulation_id=simulation_id, + population=later.population(), + versions=LATER_VERSIONS, + inputs=OPTIONS, + output_uri_prefix=OUTPUT_BUCKET_PREFIX, + excluded_run_id=excluded_run_id, + ) + + +def _used_on_later_catalog( + connection: Connection, stored: StoredReport, later: EconomyCatalog +) -> StoredRun: + """Insert a succeeded later run of the baseline that a succeeded re-run used.""" + later_shape = RunShape(versions=LATER_VERSIONS) + later_baseline = insert_run_of( + connection, stored.baseline.simulation_id, later.population(), later_shape + ) + later_reform = insert_run_of( + connection, + stored.reform.simulation_id, + later.population(), + replace(later_shape, status=SimulationStatus.FAILED), + ) + insert_report_run_of(connection, stored.report_id, (later_baseline, later_reform)) + return later_baseline + + +def test_a_simulations_own_run_on_a_later_catalog_version_is_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy) + later_baseline = _used_on_later_catalog(connection, stored, later) + act_as_the_service(connection) + + found = _find_later_run_of(connection, stored.baseline.simulation_id, later) + excluded = _find_later_run_of( + connection, + stored.baseline.simulation_id, + later, + excluded_run_id=later_baseline.simulation_run_id, + ) + # The run on the earlier catalog version is not on the current versions. + earlier = find_reusable_run_of( + connection, + simulation_id=stored.baseline.simulation_id, + population=economy.population(), + versions=LATER_VERSIONS, + inputs=OPTIONS, + output_uri_prefix=OUTPUT_BUCKET_PREFIX, + ) + + assert found == ReusableRun( + simulation_id=stored.baseline.simulation_id, + simulation_run_id=later_baseline.simulation_run_id, + ) + assert excluded is None + assert earlier is None + + +def test_another_simulations_run_is_not_reused_for_a_simulations_rerun( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy) + _used_on_later_catalog(connection, stored, later) + other = _stored_report(connection, economy) + act_as_the_service(connection) + + found = _find_later_run_of(connection, other.baseline.simulation_id, later) + + assert found is None + + +def _rerun_of( + stored: StoredReport, + later: EconomyCatalog, + *, + baseline_reused: StoredRun | None = None, +) -> NewReportRerun: + baseline = ( + RerunSide( + simulation_id=stored.baseline.simulation_id, + simulation_run_id=uuid4(), + reused=False, + ) + if baseline_reused is None + else RerunSide( + simulation_id=baseline_reused.simulation_id, + simulation_run_id=baseline_reused.simulation_run_id, + reused=True, + ) + ) + return NewReportRerun( + report_id=stored.report_id, + run_id=uuid4(), + population=later.population(), + versions=LATER_VERSIONS, + baseline=baseline, + reform=RerunSide( + simulation_id=stored.reform.simulation_id, + simulation_run_id=uuid4(), + reused=False, + ), + ) + + +def test_a_report_rerun_creates_runs_of_the_reports_own_simulations( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy) + reused = _used_on_later_catalog(connection, stored, later) + rerun = _rerun_of(stored, later, baseline_reused=reused) + act_as_the_service(connection) + + insert_report_rerun(connection, rerun) + report_run = _row(connection, "report_runs", rerun.run_id) + reform_run = _row(connection, "simulation_runs", rerun.reform.simulation_run_id) + earlier_run = _row(connection, "report_runs", stored.report_run_id) + baseline_runs = _count_runs_of(connection, stored.baseline.simulation_id) + + assert rerun.new_simulation_run_ids == (rerun.reform.simulation_run_id,) + assert { + key: report_run[key] + for key in ( + "report_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "country_package_version", + "policyengine_version", + "status", + "trigger", + "modal_call_id", + "budget_window_run_id", + "year", + ) + } == { + "report_id": stored.report_id, + "baseline_simulation_run_id": reused.simulation_run_id, + "reform_simulation_run_id": rerun.reform.simulation_run_id, + "country_package_version": LATER_VERSIONS.country_package_version, + "policyengine_version": LATER_VERSIONS.policyengine_version, + "status": ReportRunStatus.PENDING, + "trigger": ReportRunTrigger.MANUAL, + "modal_call_id": None, + "budget_window_run_id": None, + "year": None, + } + assert { + key: reform_run[key] + for key in ( + "simulation_id", + "year", + "tax_benefit_model_version_id", + "dataset_id", + "region_id", + "country_package_version", + "policyengine_version", + "status", + "modal_call_id", + ) + } == { + "simulation_id": stored.reform.simulation_id, + "year": 2026, + "tax_benefit_model_version_id": later.catalog.version_id, + "dataset_id": later.dataset_id, + "region_id": later.region_id, + "country_package_version": LATER_VERSIONS.country_package_version, + "policyengine_version": LATER_VERSIONS.policyengine_version, + "status": SimulationStatus.PENDING, + "modal_call_id": None, + } + # The earlier run is kept, and the reused side gets no new run. + assert earlier_run["status"] == ReportRunStatus.SUCCEEDED + assert baseline_runs == 2 + + +@pytest.mark.parametrize("status", [ReportRunStatus.PENDING, ReportRunStatus.RUNNING]) +def test_a_report_rerun_is_refused_while_a_run_is_pending_or_running( + public_api_database_url, status: ReportRunStatus +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy, run_status=status) + rerun = _rerun_of(stored, later) + act_as_the_service(connection) + + with pytest.raises(RunInProgressError) as refused: + insert_report_rerun(connection, rerun) + created = connection.execute( + text("SELECT count(*) FROM public.report_runs WHERE id = :id"), + {"id": rerun.run_id}, + ).scalar_one() + + assert refused.value.run_id == stored.report_run_id + assert created == 0 + + +def test_a_report_whose_runs_have_ended_is_rerun_again(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy, run_status=ReportRunStatus.FAILED) + first = _rerun_of(stored, later) + act_as_the_service(connection) + insert_report_rerun(connection, first) + # The first re-run's run fails before a second re-run. + connection.execute(text("RESET ROLE")) + connection.execute( + text( + "UPDATE public.report_runs SET status = 'failed', " + "completed_at = now() WHERE id = :id" + ), + {"id": first.run_id}, + ) + second = _rerun_of(stored, later) + act_as_the_service(connection) + + insert_report_rerun(connection, second) + + assert _row(connection, "report_runs", second.run_id)["status"] == ( + ReportRunStatus.PENDING + ) + + +def test_a_simulation_rerun_creates_a_run_and_is_refused_while_one_is_under_way( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = insert_simulation_run(connection, economy.population()) + rerun = NewSimulationRerun( + simulation_id=stored.simulation_id, + run_id=uuid4(), + population=later.population(), + versions=LATER_VERSIONS, + ) + act_as_the_service(connection) + + insert_simulation_rerun(connection, rerun) + created = _row(connection, "simulation_runs", rerun.run_id) + with pytest.raises(RunInProgressError) as refused: + insert_simulation_rerun(connection, replace(rerun, run_id=uuid4())) + earlier = _row(connection, "simulation_runs", stored.simulation_run_id) + + assert ( + created["simulation_id"], + created["tax_benefit_model_version_id"], + created["policyengine_version"], + created["status"], + ) == ( + stored.simulation_id, + later.catalog.version_id, + LATER_POLICYENGINE_VERSION, + SimulationStatus.PENDING, + ) + assert refused.value.run_id == rerun.run_id + assert earlier["status"] == SimulationStatus.SUCCEEDED + + +def test_a_rerun_holds_its_records_lock_until_its_transaction_ends( + public_api_database_url, +) -> None: + """A second re-run of the record waits; one of another record does not.""" + with rolled_back_transaction(public_api_database_url) as holder: + economy = insert_economy_catalog(holder) + later = insert_later_catalog(holder, economy) + stored = _stored_report(holder, economy, run_status=ReportRunStatus.FAILED) + act_as_the_service(holder) + insert_report_rerun(holder, _rerun_of(stored, later)) + + with rolled_back_transaction(public_api_database_url) as waiter: + act_as_the_service(waiter) + waiter.execute(text("SET LOCAL lock_timeout = '200ms'")) + lock_reruns_of(waiter, uuid4()) + with pytest.raises(OperationalError, match="lock timeout"): + lock_reruns_of(waiter, stored.report_id) + + +def test_a_report_reruns_inputs_are_read_on_the_current_catalog_version( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(CTC_AMOUNT,)) + later = insert_later_catalog(connection, economy, parameter_names=(CTC_AMOUNT,)) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + stored = _stored_report(connection, economy, reform_choice=reform_choice) + reused = _used_on_later_catalog(connection, stored, later) + act_as_the_service(connection) + store = _store_on(connection) + sides = { + "baseline": SideDefinition(stored.baseline.simulation_id, CURRENT_LAW), + "reform": SideDefinition(stored.reform.simulation_id, reform_choice), + } + + inputs = store.read_report_rerun_inputs( + LATER_REQUEST, **sides, reuse=LATER_REUSE + ) + forced = store.read_report_rerun_inputs(LATER_REQUEST, **sides, reuse=None) + + assert (inputs.catalog.catalog_version_id, inputs.catalog.region.id) == ( + later.catalog.version_id, + later.region_id, + ) + assert inputs.catalog.dataset.id == later.dataset_id + assert inputs.baseline.reusable_run == ReusableRun( + simulation_id=stored.baseline.simulation_id, + simulation_run_id=reused.simulation_run_id, + ) + assert inputs.reform.reusable_run is None + assert inputs.reform.reform == {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} + assert forced.baseline.reusable_run is None + + +def test_both_sides_of_one_simulation_never_reuse_the_same_run( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy) + reused = _used_on_later_catalog(connection, stored, later) + act_as_the_service(connection) + side = SideDefinition(stored.baseline.simulation_id, CURRENT_LAW) + + inputs = _store_on(connection).read_report_rerun_inputs( + LATER_REQUEST, baseline=side, reform=side, reuse=LATER_REUSE + ) + + assert inputs.baseline.reusable_run == ReusableRun( + simulation_id=stored.baseline.simulation_id, + simulation_run_id=reused.simulation_run_id, + ) + assert inputs.reform.reusable_run is None + + +def test_a_rerun_names_the_side_whose_policy_the_current_catalog_cannot_set( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(CTC_AMOUNT,)) + insert_later_catalog(connection, economy) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + stored = _stored_report(connection, economy, reform_choice=reform_choice) + act_as_the_service(connection) + + with pytest.raises(ReportSideError) as refused: + _store_on(connection).read_report_rerun_inputs( + LATER_REQUEST, + baseline=SideDefinition(stored.baseline.simulation_id, CURRENT_LAW), + reform=SideDefinition(stored.reform.simulation_id, reform_choice), + reuse=None, + ) + + assert refused.value.side == "reform" + assert isinstance(refused.value.error, MissingPolicyParametersError) + assert refused.value.error.parameter_names == (CTC_AMOUNT,) + + +@pytest.mark.parametrize( + ("missing", "expected"), + [("region", UnknownRegionError), ("dataset", UnknownDatasetError)], +) +def test_a_rerun_is_refused_when_the_current_catalog_lacks_its_region_or_dataset( + public_api_database_url, missing: str, expected: type[Exception] +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + other_dataset_id = insert_dataset(connection, later.catalog, "acs_2024") + insert_region( + connection, + later.catalog, + "state/ca", + default_dataset_id=other_dataset_id, + region_type="state", + ) + stored = _stored_report(connection, economy) + act_as_the_service(connection) + request = ( + replace(LATER_REQUEST, region_code="state/tx") + if missing == "region" + else replace(LATER_REQUEST, dataset_name="populace_us_2023") + ) + side = SideDefinition(stored.baseline.simulation_id, CURRENT_LAW) + + with pytest.raises(expected): + _store_on(connection).read_simulation_rerun_inputs(request, side) + + +def test_a_simulation_reruns_inputs_are_read_on_the_current_catalog_version( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(CTC_AMOUNT,)) + later = insert_later_catalog(connection, economy, parameter_names=(CTC_AMOUNT,)) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) + stored = insert_simulation_run( + connection, economy.population(), RunShape(choice=choice) + ) + act_as_the_service(connection) + + inputs = _store_on(connection).read_simulation_rerun_inputs( + LATER_REQUEST, SideDefinition(stored.simulation_id, choice) + ) + + assert inputs.catalog.catalog_version_id == later.catalog.version_id + assert inputs.simulation.choice == choice + assert inputs.simulation.reform == {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} + + +def test_a_rerun_and_its_runs_are_created_through_the_store( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy) + rerun = _rerun_of(stored, later) + simulation_rerun = NewSimulationRerun( + simulation_id=insert_simulation_run( + connection, economy.population() + ).simulation_id, + run_id=uuid4(), + population=later.population(), + versions=LATER_VERSIONS, + ) + act_as_the_service(connection) + store = _store_on(connection) + + store.create_report_rerun(rerun) + store.create_simulation_rerun(simulation_rerun) + definition = store.read_report_definition(stored.report_id) + simulation = store.read_simulation_definition(simulation_rerun.simulation_id) + with pytest.raises(RunInProgressError): + store.create_report_rerun(_rerun_of(stored, later)) + + assert _row(connection, "report_runs", rerun.run_id)["trigger"] == ( + ReportRunTrigger.MANUAL + ) + assert _row(connection, "simulation_runs", simulation_rerun.run_id)[ + "status" + ] == (SimulationStatus.PENDING) + assert definition is not None + assert definition.id == stored.report_id + assert simulation is not None + assert simulation.id == simulation_rerun.simulation_id From f7575245ba5eeb9dffac885c72c0d52e3690ae53 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:47:26 +0400 Subject: [PATCH 02/18] Read a stored report or simulation, and create its re-run on the current catalog version A re-run keeps the record and every earlier run. It reads the record's region code, dataset name and each simulation's policy and dynamic, so the current catalog version's rows can be found; a report side may reuse a run of its own simulation that succeeded on the current versions, compared on the run's own catalog version, dataset and region. Nothing is created while the record has a run pending or running, checked under a transaction-level advisory lock on the record's ID. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 30 ++ .../economy_rerun_records.py | 409 ++++++++++++++++++ .../economy_run_records.py | 131 +++++- .../economy_run_store.py | 224 ++++++++-- .../tests/test_economy_reruns.py | 16 +- 5 files changed, 749 insertions(+), 61 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 1741cf7b0..d8ea4ef03 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -6,6 +6,21 @@ missing_table_privileges, read_missing_table_privileges, ) +from policyengine_simulation_persistence.economy_rerun_records import ( + NewReportRerun, + NewSimulationRerun, + RecordKind, + ReportDefinition, + RerunSide, + RunInProgressError, + SideDefinition, + SimulationDefinition, + insert_report_rerun, + insert_simulation_rerun, + lock_reruns_of, + read_report_definition, + read_simulation_definition, +) from policyengine_simulation_persistence.economy_run_reads import ( EconomyReportState, EconomySimulationState, @@ -28,6 +43,7 @@ fail_unstarted_report_run, fail_unstarted_simulation_run, find_reusable_run, + find_reusable_run_of, insert_economy_report, insert_economy_simulation, read_dataset, @@ -126,19 +142,27 @@ "NewEconomyReport", "NewEconomySimulation", "NewHouseholdJob", + "NewReportRerun", + "NewSimulationRerun", "OverlappingReformError", "PolicyReform", + "RecordKind", + "ReportDefinition", "ReportRunState", "ReportRunStatus", "ReportRunTrigger", "ReportSide", "ReportSideError", "ReportSimulation", + "RerunSide", "ReusableRun", + "RunInProgressError", "RunReuse", "RunVersions", + "SideDefinition", "SideRefusal", "SimulationChoice", + "SimulationDefinition", "SimulationInputs", "SimulationRunState", "SimulationStatus", @@ -158,11 +182,15 @@ "fail_unstarted_report_run", "fail_unstarted_simulation_run", "find_reusable_run", + "find_reusable_run_of", "household_jobs", "input_metadata", "insert_economy_report", "insert_economy_simulation", "insert_household_job", + "insert_report_rerun", + "insert_simulation_rerun", + "lock_reruns_of", "missing_table_privileges", "normalized_database_url", "read_catalog_version_id", @@ -176,6 +204,8 @@ "read_model_id", "read_policy_reform", "read_region", + "read_report_definition", + "read_simulation_definition", "read_stored_household", "record_household_call", "record_report_run_call", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py new file mode 100644 index 000000000..2e99aba0b --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -0,0 +1,409 @@ +"""The simulation service's reads and writes for re-running a report or simulation. + +A re-run calculates a stored report or simulation again on the versions the +service runs now, as a new run, and keeps every earlier run. The service: + +- reads what the record calculates (``read_report_definition``, + ``read_simulation_definition``): its country, year, region code, dataset + name and each simulation's policy and dynamic, so the current catalog + version's rows with the same code and name can be found; +- creates, in one transaction, the new run and, for each report side not + reusing a run of its own simulation, a new ``simulation_runs`` row of that + simulation, all ``pending`` with no Modal call ID. A report's new run has + trigger ``manual``. + +The record itself is never changed: the service may not update ``reports`` +or ``simulations``, so a report keeps naming the simulations it was created +with, and each new run records the catalog version, dataset and region it +calculates on. Nothing is created while the record has a run that is +``pending`` or ``running``; a transaction-level advisory lock on the +record's ID makes concurrent re-runs of one record wait for each other, as +a row lock on ``reports`` would need ``UPDATE`` on it. + +The service then starts and records the new run's call as it does a new +record's (``economy_run_records``). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Literal + +from sqlalchemy import func, insert, literal, select + +from policyengine_simulation_persistence.economy_run_records import ( + SimulationChoice, + insert_simulation_run, + references_checked, +) +from policyengine_simulation_persistence.input_tables import ( + MODEL_NAME_PREFIX, + datasets, + regions, + tax_benefit_model_versions, + tax_benefit_models, +) +from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + ReportRunTrigger, + SimulationStatus, + SimulationType, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, RowMapping, Table + + from policyengine_simulation_persistence.economy_run_records import ( + EconomyPopulation, + RunVersions, + ) + +# What a stored report or simulation calculates, as a re-run tells them apart: +# a single-year economy report or simulation, a budget window, a household +# report or simulation, or anything else, such as a simulation of a filtered +# population. +type RecordKind = Literal["economy", "budget_window", "household", "unsupported"] +# A report's type, as the Public API and the service write it. +_ECONOMY_TYPE: Final = "economy" +_HOUSEHOLD_TYPE: Final = "household" +# Both run tables' statuses use these values for a run under way. +_UNDER_WAY: Final = (SimulationStatus.PENDING.value, SimulationStatus.RUNNING.value) + + +@dataclass(frozen=True) +class SideDefinition: + """A report side's stored simulation, and its policy and dynamic.""" + + simulation_id: UUID + choice: SimulationChoice + + +@dataclass(frozen=True) +class ReportDefinition: + """What a stored report calculates, as a re-run reads it. + + ``region_code`` and ``dataset_name`` are the code and name of the + report's region and dataset, or None when the report names none (or a + row it named was deleted). ``baseline`` and ``reform`` are None when the + report names no such simulation. + """ + + id: UUID + country: str + kind: RecordKind + year: int | None + region_code: str | None + dataset_name: str | None + inputs: dict[str, Any] + baseline: SideDefinition | None + reform: SideDefinition | None + + +@dataclass(frozen=True) +class SimulationDefinition: + """What a stored simulation calculates, as a re-run reads it. + + ``country`` comes from the simulation's catalog version's model; it is + None for a model the service does not name ``policyengine-``. + """ + + id: UUID + country: str | None + kind: RecordKind + year: int | None + region_code: str | None + dataset_name: str | None + choice: SimulationChoice + + +@dataclass(frozen=True) +class RerunSide: + """The simulation run one side of a report's re-run names. + + The simulation is always the report's own. With ``reused``, the run is + one of its runs that succeeded on the current versions, and nothing is + created for the side; otherwise the run is created, pending. + """ + + simulation_id: UUID + simulation_run_id: UUID + reused: bool + + +@dataclass(frozen=True) +class NewReportRerun: + """A stored report's new run, on the current versions and catalog rows.""" + + report_id: UUID + run_id: UUID + population: EconomyPopulation + versions: RunVersions + baseline: RerunSide + reform: RerunSide + + @property + def new_simulation_run_ids(self) -> tuple[UUID, ...]: + """Return the IDs of the simulation runs created for the re-run.""" + return tuple( + side.simulation_run_id + for side in (self.baseline, self.reform) + if not side.reused + ) + + +@dataclass(frozen=True) +class NewSimulationRerun: + """A stored simulation's new run, on the current versions and catalog rows.""" + + simulation_id: UUID + run_id: UUID + population: EconomyPopulation + versions: RunVersions + + +class RunInProgressError(Exception): + """The record has a run that is pending or running, so no re-run is created. + + ``run_id`` names that run: the newest such one. + """ + + def __init__(self, run_id: UUID) -> None: + """Keep the run under way.""" + super().__init__(f"run {run_id} is pending or running") + self.run_id = run_id + + +def read_report_definition( + connection: Connection, report_id: UUID +) -> ReportDefinition | None: + """Return what the stored report calculates, or None when there is none.""" + region = regions.alias("report_region") + dataset = datasets.alias("report_dataset") + baseline = simulations.alias("baseline_simulation") + reform = simulations.alias("reform_simulation") + row = ( + connection.execute( + select( + reports.c.id, + reports.c.country, + reports.c.type, + reports.c.household_id, + reports.c.budget_window_size, + reports.c.year, + reports.c.inputs, + region.c.code.label("region_code"), + dataset.c.name.label("dataset_name"), + baseline.c.id.label("baseline_id"), + baseline.c.policy_id.label("baseline_policy_id"), + baseline.c.dynamic_id.label("baseline_dynamic_id"), + reform.c.id.label("reform_id"), + reform.c.policy_id.label("reform_policy_id"), + reform.c.dynamic_id.label("reform_dynamic_id"), + ) + .select_from(reports) + .outerjoin(region, region.c.id == reports.c.region_id) + .outerjoin(dataset, dataset.c.id == reports.c.dataset_id) + .outerjoin(baseline, baseline.c.id == reports.c.baseline_simulation_id) + .outerjoin(reform, reform.c.id == reports.c.reform_simulation_id) + .where(reports.c.id == report_id) + ) + .mappings() + .one_or_none() + ) + if row is None: + return None + return ReportDefinition( + id=row["id"], + country=row["country"], + kind=_report_kind(row["type"], row["household_id"], row["budget_window_size"]), + year=row["year"], + region_code=row["region_code"], + dataset_name=row["dataset_name"], + inputs=row["inputs"], + baseline=_side(row, "baseline"), + reform=_side(row, "reform"), + ) + + +def _report_kind( + report_type: str | None, household_id: UUID | None, budget_window_size: int | None +) -> RecordKind: + if household_id is not None or report_type == _HOUSEHOLD_TYPE: + return "household" + if report_type != _ECONOMY_TYPE: + return "unsupported" + return "economy" if budget_window_size is None else "budget_window" + + +def _side( + row: RowMapping, side: Literal["baseline", "reform"] +) -> SideDefinition | None: + simulation_id = row[f"{side}_id"] + if simulation_id is None: + return None + return SideDefinition( + simulation_id=simulation_id, + choice=SimulationChoice( + policy_id=row[f"{side}_policy_id"], dynamic_id=row[f"{side}_dynamic_id"] + ), + ) + + +def read_simulation_definition( + connection: Connection, simulation_id: UUID +) -> SimulationDefinition | None: + """Return what the stored simulation calculates, or None when there is none.""" + row = ( + connection.execute( + select( + simulations.c.id, + simulations.c.simulation_type, + simulations.c.filter_field, + simulations.c.year, + simulations.c.policy_id, + simulations.c.dynamic_id, + tax_benefit_models.c.name.label("model_name"), + regions.c.code.label("region_code"), + datasets.c.name.label("dataset_name"), + ) + .select_from(simulations) + .join( + tax_benefit_model_versions, + tax_benefit_model_versions.c.id + == simulations.c.tax_benefit_model_version_id, + ) + .join( + tax_benefit_models, + tax_benefit_models.c.id == tax_benefit_model_versions.c.model_id, + ) + .outerjoin(regions, regions.c.id == simulations.c.region_id) + .outerjoin(datasets, datasets.c.id == simulations.c.dataset_id) + .where(simulations.c.id == simulation_id) + ) + .mappings() + .one_or_none() + ) + if row is None: + return None + model_name: str = row["model_name"] + simulation_type = SimulationType(row["simulation_type"]) + kind: RecordKind = ( + "household" + if simulation_type is SimulationType.HOUSEHOLD + else "economy" + if row["filter_field"] is None + else "unsupported" + ) + return SimulationDefinition( + id=row["id"], + country=model_name.removeprefix(MODEL_NAME_PREFIX) + if model_name.startswith(MODEL_NAME_PREFIX) + else None, + kind=kind, + year=row["year"], + region_code=row["region_code"], + dataset_name=row["dataset_name"], + choice=SimulationChoice( + policy_id=row["policy_id"], dynamic_id=row["dynamic_id"] + ), + ) + + +def lock_reruns_of(connection: Connection, record_id: UUID) -> None: + """Wait for, then hold until the transaction ends, the record's re-run lock. + + It is a transaction-level advisory lock keyed on the record's ID, which + any role may take: a second re-run of the same report or simulation + waits here until the first one's transaction ends, then sees its run. + """ + connection.execute( + select( + func.pg_advisory_xact_lock( + func.hashtextextended(literal(f"rerun:{record_id}"), 0) + ) + ) + ) + + +def _run_under_way( + connection: Connection, table: Table, condition: ColumnElement[bool] +) -> UUID | None: + """Return the newest pending or running run meeting ``condition``, or None.""" + return connection.execute( + select(table.c.id) + .where(condition, table.c.status.in_(_UNDER_WAY)) + .order_by(table.c.created_at.desc(), table.c.id.desc()) + .limit(1) + ).scalar_one_or_none() + + +def insert_report_rerun(connection: Connection, rerun: NewReportRerun) -> None: + """Create the report's new run and its sides' new simulation runs, pending. + + Take the report's re-run lock first. Raise RunInProgressError, creating + nothing, when the report has a run pending or running, and + EconomyRecordReferenceError when the database refuses a row the new + runs name; that refusal ends the caller's transaction, or its savepoint. + """ + lock_reruns_of(connection, rerun.report_id) + under_way = _run_under_way( + connection, report_runs, report_runs.c.report_id == rerun.report_id + ) + if under_way is not None: + raise RunInProgressError(under_way) + with references_checked(): + for side in (rerun.baseline, rerun.reform): + if not side.reused: + insert_simulation_run( + connection, + run_id=side.simulation_run_id, + simulation_id=side.simulation_id, + population=rerun.population, + versions=rerun.versions, + ) + connection.execute( + insert(report_runs).values( + id=rerun.run_id, + report_id=rerun.report_id, + baseline_simulation_run_id=rerun.baseline.simulation_run_id, + reform_simulation_run_id=rerun.reform.simulation_run_id, + budget_window_run_id=None, + year=None, + country_package_version=rerun.versions.country_package_version, + policyengine_version=rerun.versions.policyengine_version, + status=ReportRunStatus.PENDING, + trigger=ReportRunTrigger.MANUAL, + modal_call_id=None, + ) + ) + + +def insert_simulation_rerun(connection: Connection, rerun: NewSimulationRerun) -> None: + """Create the simulation's new run, pending with no call. + + Take the simulation's re-run lock first. Raise RunInProgressError, + creating nothing, when the simulation has a run pending or running, and + EconomyRecordReferenceError as ``insert_report_rerun`` does. + """ + lock_reruns_of(connection, rerun.simulation_id) + under_way = _run_under_way( + connection, + simulation_runs, + simulation_runs.c.simulation_id == rerun.simulation_id, + ) + if under_way is not None: + raise RunInProgressError(under_way) + with references_checked(): + insert_simulation_run( + connection, + run_id=rerun.run_id, + simulation_id=rerun.simulation_id, + population=rerun.population, + versions=rerun.versions, + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py index dd90ddabd..7800f9462 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py @@ -286,7 +286,7 @@ def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match - of an economy simulation with the same policy, dynamic, catalog version, dataset, region and year; - that succeeded on the same country package and .py versions, for that - year, dataset and region; + catalog version, year, dataset and region; - whose output file is under ``output_uri_prefix``, the output bucket, which never deletes a file; - that a succeeded report run used, in a report with the same ``inputs`` @@ -302,6 +302,77 @@ def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match responses: the refusal fails that report run, and the run is not offered again. """ + return _newest_reusable_run( + connection, + ( + simulations.c.simulation_type == SimulationType.ECONOMY, + simulations.c.policy_id.is_not_distinct_from(choice.policy_id), + simulations.c.dynamic_id.is_not_distinct_from(choice.dynamic_id), + simulations.c.tax_benefit_model_version_id + == population.tax_benefit_model_version_id, + simulations.c.dataset_id == population.dataset_id, + simulations.c.region_id == population.region_id, + simulations.c.year == population.year, + simulations.c.filter_field.is_(None), + ), + _WantedRun( + population=population, + versions=versions, + inputs=inputs, + output_uri_prefix=output_uri_prefix, + excluded_run_id=excluded_run_id, + ), + ) + + +def find_reusable_run_of( # noqa: PLR0913 # the simulation, then each part of the match + connection: Connection, + *, + simulation_id: UUID, + population: EconomyPopulation, + versions: RunVersions, + inputs: dict[str, Any], + output_uri_prefix: str, + excluded_run_id: UUID | None = None, +) -> ReusableRun | None: + """Return the succeeded run of one simulation a re-run of its report may name. + + The run is the newest, by completion, of the simulation's runs that meets + every condition ``find_reusable_run`` sets on a run, or None. Its catalog + version, dataset and region are compared on the run itself: a re-run on + a later .py version calculates the simulation on that version's catalog + rows, while the simulation keeps the ones it was created on. + """ + return _newest_reusable_run( + connection, + (simulation_runs.c.simulation_id == simulation_id,), + _WantedRun( + population=population, + versions=versions, + inputs=inputs, + output_uri_prefix=output_uri_prefix, + excluded_run_id=excluded_run_id, + ), + ) + + +@dataclass(frozen=True) +class _WantedRun: + """What a reusable run must have calculated, and the run it must not be.""" + + population: EconomyPopulation + versions: RunVersions + inputs: dict[str, Any] + output_uri_prefix: str + excluded_run_id: UUID | None + + +def _newest_reusable_run( + connection: Connection, + simulation_conditions: tuple[ColumnElement[bool], ...], + wanted: _WantedRun, +) -> ReusableRun | None: + """Return the newest run meeting the conditions and every run condition.""" used_by = report_runs.alias("used_by_report_run") refused_by = report_runs.alias("refused_by_report_run") used_with_same_options = ( @@ -310,7 +381,7 @@ def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match .where( used_by.c.status == ReportRunStatus.SUCCEEDED, _names_run(used_by), - cast(reports.c.inputs, JSONB) == literal(inputs, JSONB), + cast(reports.c.inputs, JSONB) == literal(wanted.inputs, JSONB), ) .exists() ) @@ -325,29 +396,27 @@ def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match ) .exists() ) + population = wanted.population + versions = wanted.versions conditions = [ - simulations.c.simulation_type == SimulationType.ECONOMY, - simulations.c.policy_id.is_not_distinct_from(choice.policy_id), - simulations.c.dynamic_id.is_not_distinct_from(choice.dynamic_id), - simulations.c.tax_benefit_model_version_id - == population.tax_benefit_model_version_id, - simulations.c.dataset_id == population.dataset_id, - simulations.c.region_id == population.region_id, - simulations.c.year == population.year, - simulations.c.filter_field.is_(None), + *simulation_conditions, simulation_runs.c.status == SimulationStatus.SUCCEEDED, simulation_runs.c.year == population.year, + simulation_runs.c.tax_benefit_model_version_id + == population.tax_benefit_model_version_id, simulation_runs.c.country_package_version == versions.country_package_version, simulation_runs.c.policyengine_version == versions.policyengine_version, simulation_runs.c.dataset_id == population.dataset_id, simulation_runs.c.region_id == population.region_id, - simulation_runs.c.output_uri.startswith(output_uri_prefix, autoescape=True), + simulation_runs.c.output_uri.startswith( + wanted.output_uri_prefix, autoescape=True + ), simulation_runs.c.output_sha256.is_not(None), used_with_same_options, ~refused, ] - if excluded_run_id is not None: - conditions.append(simulation_runs.c.id != excluded_run_id) + if wanted.excluded_run_id is not None: + conditions.append(simulation_runs.c.id != wanted.excluded_run_id) row = connection.execute( select(simulation_runs.c.id, simulation_runs.c.simulation_id) .join(simulations, simulations.c.id == simulation_runs.c.simulation_id) @@ -374,10 +443,10 @@ def insert_economy_report(connection: Connection, report: NewEconomyReport) -> N Raise EconomyRecordReferenceError when the database refuses a row it names. The refusal ends the caller's transaction, or its savepoint. """ - with _references_checked(): + with references_checked(): for side in (report.baseline, report.reform): if not side.reused: - _insert_simulation( + insert_simulation_with_run( connection, simulation_id=side.simulation_id, run_id=side.simulation_run_id, @@ -431,8 +500,8 @@ def insert_economy_simulation( Raise EconomyRecordReferenceError when the database refuses a row it names. The refusal ends the caller's transaction, or its savepoint. """ - with _references_checked(): - _insert_simulation( + with references_checked(): + insert_simulation_with_run( connection, simulation_id=simulation.id, run_id=simulation.run_id, @@ -442,7 +511,7 @@ def insert_economy_simulation( ) -def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts +def insert_simulation_with_run( # noqa: PLR0913 # the simulation's and its run's parts connection: Connection, *, simulation_id: UUID, @@ -465,6 +534,28 @@ def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts year=population.year, ) ) + insert_simulation_run( + connection, + run_id=run_id, + simulation_id=simulation_id, + population=population, + versions=versions, + ) + + +def insert_simulation_run( + connection: Connection, + *, + run_id: UUID, + simulation_id: UUID, + population: EconomyPopulation, + versions: RunVersions, +) -> None: + """Create a run of a stored simulation on ``population``, pending with no call. + + The run records the catalog version, dataset and region it calculates + on, which a re-run's may differ from its simulation's. + """ connection.execute( insert(simulation_runs).values( id=run_id, @@ -482,7 +573,7 @@ def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts @contextmanager -def _references_checked() -> Iterator[None]: +def references_checked() -> Iterator[None]: """Turn a foreign-key refusal into EconomyRecordReferenceError. Any other integrity error propagates unchanged. diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py index 6e66b2a50..3227917e4 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py @@ -9,11 +9,17 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Literal, get_args from sqlalchemy.exc import SQLAlchemyError +from policyengine_simulation_persistence.economy_rerun_records import ( + insert_report_rerun, + insert_simulation_rerun, + read_report_definition, + read_simulation_definition, +) from policyengine_simulation_persistence.economy_run_reads import ( read_economy_report, read_economy_simulation, @@ -28,6 +34,7 @@ fail_unstarted_report_run, fail_unstarted_simulation_run, find_reusable_run, + find_reusable_run_of, insert_economy_report, insert_economy_simulation, read_dataset, @@ -61,6 +68,13 @@ from sqlalchemy import Connection, Engine + from policyengine_simulation_persistence.economy_rerun_records import ( + NewReportRerun, + NewSimulationRerun, + ReportDefinition, + SideDefinition, + SimulationDefinition, + ) from policyengine_simulation_persistence.economy_run_reads import ( EconomyReportState, EconomySimulationState, @@ -78,6 +92,11 @@ type TransactionFactory = Callable[[], AbstractContextManager[Connection]] # Which of a report's simulations a value belongs to. type ReportSide = Literal["baseline", "reform"] +# Finds a report side's reusable run, given the side, its policy and dynamic, +# and the run the other side reuses, which it must not. +type ReusableRunSearch = Callable[ + [ReportSide, SimulationChoice, UUID | None], ReusableRun | None +] # The errors a simulation's own policy or dynamic can raise. type SideRefusal = ( UnknownPolicyError @@ -253,39 +272,88 @@ def read_report_inputs( """ with self._transaction() as connection: catalog = _catalog_rows(connection, request) - population = catalog.population(request.year) - sides = [] - excluded_run_id = None - sides_requested: tuple[tuple[ReportSide, SimulationChoice], ...] = ( - ("baseline", baseline), - ("reform", reform), + search = ( + None + if reuse is None + else _search_by_definition( + connection, catalog.population(request.year), reuse + ) + ) + return _report_inputs( + connection, + request, + catalog, + baseline=baseline, + reform=reform, + search=search, + ) + + def read_report_definition(self, report_id: UUID) -> ReportDefinition | None: + """Return what the stored report calculates, or None when there is none.""" + with self._transaction() as connection: + return read_report_definition(connection, report_id) + + def read_simulation_definition( + self, simulation_id: UUID + ) -> SimulationDefinition | None: + """Return what the stored simulation calculates, or None when there is none.""" + with self._transaction() as connection: + return read_simulation_definition(connection, simulation_id) + + def read_report_rerun_inputs( + self, + request: EconomyRequest, + *, + baseline: SideDefinition, + reform: SideDefinition, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + """Read a report re-run's catalog rows, both sides' values and reusable runs. + + ``request`` names the stored report's region code and dataset name, + found in the catalog version of the current .py version. With + ``reuse``, each side gets the run ``find_reusable_run_of`` finds among + its own simulation's runs, and the reform never gets the baseline's. + Raise the errors ``read_report_inputs`` raises. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + search = ( + None + if reuse is None + else _search_own_runs( + connection, + catalog.population(request.year), + reuse, + { + "baseline": baseline.simulation_id, + "reform": reform.simulation_id, + }, + ) + ) + return _report_inputs( + connection, + request, + catalog, + baseline=baseline.choice, + reform=reform.choice, + search=search, ) - for side, choice in sides_requested: - try: - inputs = _simulation_inputs( - connection, request.country, catalog.catalog_version_id, choice - ) - except _SIDE_ERRORS as error: - raise ReportSideError(side, error) from error - if reuse is not None: - reusable_run = find_reusable_run( - connection, - choice=choice, - population=population, - versions=reuse.versions, - inputs=reuse.inputs, - output_uri_prefix=reuse.output_uri_prefix, - excluded_run_id=excluded_run_id, - ) - if reusable_run is not None: - excluded_run_id = reusable_run.simulation_run_id - inputs = SimulationInputs( - choice=inputs.choice, - reform=inputs.reform, - reusable_run=reusable_run, - ) - sides.append(inputs) - return EconomyReportInputs(catalog=catalog, baseline=sides[0], reform=sides[1]) + + def create_report_rerun(self, rerun: NewReportRerun) -> None: + """Create the report's new run and its sides' new simulation runs. + + Raise ``insert_report_rerun``'s errors, and + EconomyRecordOutcomeUnknownError as ``create_report`` does. + """ + self._create(lambda connection: insert_report_rerun(connection, rerun)) + + def create_simulation_rerun(self, rerun: NewSimulationRerun) -> None: + """Create the simulation's new run; see ``insert_simulation_rerun``. + + Raise EconomyRecordOutcomeUnknownError as ``create_report`` does. + """ + self._create(lambda connection: insert_simulation_rerun(connection, rerun)) def create_report(self, report: NewEconomyReport) -> None: """Create the report, its new simulations and its runs. @@ -465,3 +533,93 @@ def _simulation_inputs( return SimulationInputs( choice=choice, reform=combine_reforms(policy_reform, dynamic_reform) ) + + +def _side_inputs( + connection: Connection, + request: EconomyRequest, + catalog: EconomyCatalogRows, + side: ReportSide, + choice: SimulationChoice, +) -> SimulationInputs: + """Read one report side's values; raise ReportSideError naming a refused side.""" + try: + return _simulation_inputs( + connection, request.country, catalog.catalog_version_id, choice + ) + except _SIDE_ERRORS as error: + raise ReportSideError(side, error) from error + + +def _report_inputs( # noqa: PLR0913 # the request's rows, both sides and the search + connection: Connection, + request: EconomyRequest, + catalog: EconomyCatalogRows, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + search: ReusableRunSearch | None, +) -> EconomyReportInputs: + """Read both sides' values and, with ``search``, each side's reusable run. + + The reform never gets the run the baseline reuses. + """ + sides: list[SimulationInputs] = [] + excluded_run_id: UUID | None = None + requested: tuple[tuple[ReportSide, SimulationChoice], ...] = ( + ("baseline", baseline), + ("reform", reform), + ) + for side, choice in requested: + inputs = _side_inputs(connection, request, catalog, side, choice) + reusable_run = None if search is None else search(side, choice, excluded_run_id) + if reusable_run is not None: + excluded_run_id = reusable_run.simulation_run_id + inputs = replace(inputs, reusable_run=reusable_run) + sides.append(inputs) + return EconomyReportInputs(catalog=catalog, baseline=sides[0], reform=sides[1]) + + +def _search_by_definition( + connection: Connection, population: EconomyPopulation, reuse: RunReuse +) -> ReusableRunSearch: + """Search any simulation with the side's definition (a new report's sides).""" + + def search( + _side: ReportSide, choice: SimulationChoice, excluded_run_id: UUID | None + ) -> ReusableRun | None: + return find_reusable_run( + connection, + choice=choice, + population=population, + versions=reuse.versions, + inputs=reuse.inputs, + output_uri_prefix=reuse.output_uri_prefix, + excluded_run_id=excluded_run_id, + ) + + return search + + +def _search_own_runs( + connection: Connection, + population: EconomyPopulation, + reuse: RunReuse, + simulation_ids: dict[ReportSide, UUID], +) -> ReusableRunSearch: + """Search only each side's own simulation's runs (a stored report's re-run).""" + + def search( + side: ReportSide, _choice: SimulationChoice, excluded_run_id: UUID | None + ) -> ReusableRun | None: + return find_reusable_run_of( + connection, + simulation_id=simulation_ids[side], + population=population, + versions=reuse.versions, + inputs=reuse.inputs, + output_uri_prefix=reuse.output_uri_prefix, + excluded_run_id=excluded_run_id, + ) + + return search diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py index 9c38e25f3..068b20245 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py @@ -654,10 +654,14 @@ def test_a_rerun_is_refused_when_the_current_catalog_lacks_its_region_or_dataset if missing == "region" else replace(LATER_REQUEST, dataset_name="populace_us_2023") ) - side = SideDefinition(stored.baseline.simulation_id, CURRENT_LAW) with pytest.raises(expected): - _store_on(connection).read_simulation_rerun_inputs(request, side) + _store_on(connection).read_report_rerun_inputs( + request, + baseline=SideDefinition(stored.baseline.simulation_id, CURRENT_LAW), + reform=SideDefinition(stored.reform.simulation_id, CURRENT_LAW), + reuse=None, + ) def test_a_simulation_reruns_inputs_are_read_on_the_current_catalog_version( @@ -670,14 +674,10 @@ def test_a_simulation_reruns_inputs_are_read_on_the_current_catalog_version( connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] ) choice = SimulationChoice(policy_id=policy_id, dynamic_id=None) - stored = insert_simulation_run( - connection, economy.population(), RunShape(choice=choice) - ) act_as_the_service(connection) - inputs = _store_on(connection).read_simulation_rerun_inputs( - LATER_REQUEST, SideDefinition(stored.simulation_id, choice) - ) + # A simulation's re-run reads the current catalog version as a new one's. + inputs = _store_on(connection).read_simulation_inputs(LATER_REQUEST, choice) assert inputs.catalog.catalog_version_id == later.catalog.version_id assert inputs.simulation.choice == choice From aa3d01a7e941768a0f109e11bbc4d190f5eb6206 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:52:24 +0400 Subject: [PATCH 03/18] Test re-running economy reports and simulations through POST /reports/{id}/runs and POST /simulations/{id}/runs (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_generated_client.py | 79 ++++ .../tests/economy_rerun_fakes.py | 171 +++++++ .../tests/test_economy_rerun_routes.py | 358 ++++++++++++++ .../tests/test_economy_reruns.py | 437 ++++++++++++++++++ .../tests/test_openapi.py | 51 +- 5 files changed, 1095 insertions(+), 1 deletion(-) create mode 100644 projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_reruns.py diff --git a/projects/policyengine-apis-integ/tests/test_economy_generated_client.py b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py index 5f32692db..636bbf6af 100644 --- a/projects/policyengine-apis-integ/tests/test_economy_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py @@ -14,6 +14,8 @@ create_simulation, get_report, get_simulation, + rerun_report, + rerun_simulation, ) from policyengine_api_simulation_client.models import ( EconomyError, @@ -24,6 +26,7 @@ EconomySimulationCreated, EconomySimulationRequest, EconomySimulationState, + ReportRerunRequest, ReportRunStatus, ) @@ -194,3 +197,79 @@ def test_a_report_and_a_simulation_are_read_with_their_newest_runs(): assert run.status is ReportRunStatus.SUCCEEDED assert isinstance(simulation_response.parsed, EconomySimulationState) assert simulation_response.parsed.to_dict() == simulation + + +def test_a_report_rerun_is_sent_with_force_and_read_with_its_new_run(): + sent: list[httpx.Request] = [] + + def answer(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(201, json=CREATED_REPORT) + + client = AuthenticatedClient(base_url="https://simulation.example", token="test") + client.set_httpx_client( + httpx.Client( + base_url="https://simulation.example", + transport=httpx.MockTransport(answer), + ) + ) + + response = rerun_report.sync_detailed( + UUID(REPORT_ID), client=client, body=ReportRerunRequest(force=True) + ) + + (request,) = sent + assert request.url.path == f"/reports/{REPORT_ID}/runs" + assert request.read() == b'{"force":true}' + assert isinstance(response.parsed, EconomyReportCreated) + assert response.parsed.to_dict() == CREATED_REPORT + + +def test_a_simulation_rerun_is_read_with_its_new_run(): + created = { + "simulation_id": REPORT_ID, + "simulation_run_id": RUN_ID, + "status": "pending", + } + + response = rerun_simulation.sync_detailed( + UUID(REPORT_ID), client=_client_answering(201, created) + ) + + assert isinstance(response.parsed, EconomySimulationCreated) + assert response.parsed.to_dict() == created + + +@pytest.mark.parametrize( + ("status_code", "body"), + [ + (404, {"code": "unknown_report", "detail": f"There is no report {REPORT_ID}."}), + ( + 409, + { + "code": "run_in_progress", + "detail": "The report has a run pending or running.", + "run_id": RUN_ID, + }, + ), + (409, {"code": "rerun_not_supported", "detail": "Not yet."}), + ( + 503, + { + "code": "modal_call_refused", + "detail": "Modal did not accept the call of the report's new run.", + "report_id": REPORT_ID, + "run_id": RUN_ID, + }, + ), + ], +) +def test_a_refused_rerun_is_read_with_its_code_and_any_run(status_code, body): + response = rerun_report.sync_detailed( + UUID(REPORT_ID), + client=_client_answering(status_code, body), + body=ReportRerunRequest(), + ) + + assert isinstance(response.parsed, EconomyError) + assert response.parsed.to_dict() == body diff --git a/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py b/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py new file mode 100644 index 000000000..34ea5696d --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py @@ -0,0 +1,171 @@ +"""A fake of the economy store with re-runs, and stored records to re-run. + +The store's own reads and writes are tested against PostgreSQL in the +persistence library; here a fake records what the re-run service asks of it. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any +from uuid import UUID + +from economy_service_fakes import ( + CREATED_AT, + DYNAMIC_ID, + OUTPUT_BUCKET, + POLICY_ID, + REFORM, + FakeCaller, + FakeStore, + FakeWorkers, + numbered_ids, +) +from policyengine_simulation_entry.economy_reruns import EconomyRerunService +from policyengine_simulation_entry.economy_submissions import ServiceSettings +from policyengine_simulation_persistence import ( + EconomyReportInputs, + EconomySimulationInputs, + ReportDefinition, + SideDefinition, + SimulationChoice, + SimulationDefinition, + SimulationInputs, +) + +if TYPE_CHECKING: + from policyengine_simulation_persistence import ( + EconomyRequest, + NewReportRerun, + NewSimulationRerun, + RunReuse, + ) + +REPORT_ID = UUID("00000000-0000-4000-8000-0000000000a1") +BASELINE_SIMULATION_ID = UUID("00000000-0000-4000-8000-0000000000a2") +REFORM_SIMULATION_ID = UUID("00000000-0000-4000-8000-0000000000a3") +SIMULATION_ID = UUID("00000000-0000-4000-8000-0000000000a4") +RUN_UNDER_WAY_ID = UUID("00000000-0000-4000-8000-0000000000a5") +CURRENT_LAW = SimulationChoice(policy_id=None, dynamic_id=None) +REFORM_CHOICE = SimulationChoice(policy_id=POLICY_ID, dynamic_id=DYNAMIC_ID) +STORED_OPTIONS: dict[str, Any] = { + "include_cliffs": True, + "spm": {"geography_kind": "national"}, +} +STORED_REPORT = ReportDefinition( + id=REPORT_ID, + country="us", + kind="economy", + year=2026, + region_code="us", + dataset_name="populace_us_2024", + inputs=STORED_OPTIONS, + baseline=SideDefinition(simulation_id=BASELINE_SIMULATION_ID, choice=CURRENT_LAW), + reform=SideDefinition(simulation_id=REFORM_SIMULATION_ID, choice=REFORM_CHOICE), +) +STORED_SIMULATION = SimulationDefinition( + id=SIMULATION_ID, + country="us", + kind="economy", + year=2027, + region_code="us", + dataset_name="populace_us_2024", + choice=SimulationChoice(policy_id=POLICY_ID, dynamic_id=None), +) + + +@dataclass +class FakeRerunStore(FakeStore): + """Records the re-run reads and writes, and answers as told. + + ``inputs_failure`` is raised by a re-run's read of its inputs, and + ``rerun_failure`` by its create instead of recording it. + """ + + report_definition: ReportDefinition | None = STORED_REPORT + simulation_definition: SimulationDefinition | None = STORED_SIMULATION + inputs_failure: Exception | None = None + rerun_failure: Exception | None = None + rerun_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) + report_reruns: list[NewReportRerun] = field(default_factory=list) + simulation_reruns: list[NewSimulationRerun] = field(default_factory=list) + + def read_report_definition(self, report_id: UUID) -> ReportDefinition | None: + if self.read_failure is not None: + raise self.read_failure + stored = self.report_definition + return stored if stored is not None and stored.id == report_id else None + + def read_simulation_definition( + self, simulation_id: UUID + ) -> SimulationDefinition | None: + if self.read_failure is not None: + raise self.read_failure + stored = self.simulation_definition + return stored if stored is not None and stored.id == simulation_id else None + + def read_report_rerun_inputs( + self, + request: EconomyRequest, + *, + baseline: SideDefinition, + reform: SideDefinition, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + self.rerun_reads.append((request, (baseline, reform, reuse))) + if self.inputs_failure is not None: + raise self.inputs_failure + baseline_reusable, reform_reusable = self.reusable + return EconomyReportInputs( + catalog=self.catalog, + baseline=SimulationInputs( + choice=baseline.choice, + reform={}, + reusable_run=None if reuse is None else baseline_reusable, + ), + reform=SimulationInputs( + choice=reform.choice, + reform=REFORM, + reusable_run=None if reuse is None else reform_reusable, + ), + ) + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + if self.inputs_failure is not None: + self.simulation_reads.append((request, choice)) + raise self.inputs_failure + return super().read_simulation_inputs(request, choice) + + def create_report_rerun(self, rerun: NewReportRerun) -> None: + if self.rerun_failure is not None: + raise self.rerun_failure + self.report_reruns.append(rerun) + + def create_simulation_rerun(self, rerun: NewSimulationRerun) -> None: + if self.rerun_failure is not None: + raise self.rerun_failure + self.simulation_reruns.append(rerun) + + +def rerun_service( + store: FakeRerunStore, + *, + workers: FakeWorkers | None = None, + caller: FakeCaller | None = None, + output_bucket: str | None = OUTPUT_BUCKET, +) -> EconomyRerunService: + return EconomyRerunService( + store=store, + workers=workers or FakeWorkers(), + caller=caller or FakeCaller(), + settings=ServiceSettings( + environment="staging", + modal_environment="staging", + output_bucket=output_bucket, + ), + call_start_timeout_seconds=0.05, + new_id=numbered_ids(), + clock=lambda: CREATED_AT, + ) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py new file mode 100644 index 000000000..a47be302c --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py @@ -0,0 +1,358 @@ +"""The re-run routes' answers, with the re-run service faked.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, NoReturn +from uuid import UUID + +import pytest +from conftest import FakeBackend, make_settings +from fastapi.testclient import TestClient +from policyengine_simulation_entry import app as app_module +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportCreated, + EconomySimulationCreated, + ReportRerunRequest, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomyRecordNotFoundError, + EconomyRerunConflictError, + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, + UnprocessableReason, +) +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_simulation_persistence import ReportRunStatus, SimulationStatus + +OBSERVABILITY_ID = "6e1f3d4c-0a5b-4c39-9d2e-7f8a9b0c1d2e" +RECORD_ID = UUID("00000000-0000-4000-8000-0000000000e1") +RUN_ID = UUID("00000000-0000-4000-8000-0000000000e2") +SIDE_IDS = [UUID(f"00000000-0000-4000-8000-0000000000f{digit}") for digit in "1234"] +REPORT_RUNS = f"/reports/{RECORD_ID}/runs" +SIMULATION_RUNS = f"/simulations/{RECORD_ID}/runs" +RERUN_REPORT = EconomyReportCreated( + report_id=RECORD_ID, + report_run_id=RUN_ID, + status=ReportRunStatus.PENDING, + baseline_simulation_id=SIDE_IDS[0], + baseline_simulation_run_id=SIDE_IDS[1], + baseline_reused=True, + reform_simulation_id=SIDE_IDS[2], + reform_simulation_run_id=SIDE_IDS[3], + reform_reused=False, +) +RERUN_SIMULATION = EconomySimulationCreated( + simulation_id=RECORD_ID, simulation_run_id=RUN_ID, status=SimulationStatus.PENDING +) + + +def _raise(failure: Exception) -> NoReturn: + raise failure + + +@dataclass +class FakeEconomyReruns: + """Records each re-run, and answers as told.""" + + failure: Exception | None = None + report_reruns: list[tuple[UUID, ReportRerunRequest, str, Any]] = field( + default_factory=list + ) + simulation_reruns: list[tuple[UUID, str, Any]] = field(default_factory=list) + + async def rerun_report( + self, + report_id: UUID, + request: ReportRerunRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + self.report_reruns.append( + (report_id, request, request_id, observability_context) + ) + if self.failure is not None: + _raise(self.failure) + return RERUN_REPORT + + async def rerun_simulation( + self, + simulation_id: UUID, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + self.simulation_reruns.append( + (simulation_id, request_id, observability_context) + ) + if self.failure is not None: + _raise(self.failure) + return RERUN_SIMULATION + + +@pytest.fixture +def client_for(monkeypatch): + """Return a function building a client whose app uses ``reruns``. + + None leaves the app as an unconfigured service builds it. + """ + + def build(reruns: FakeEconomyReruns | None) -> TestClient: + if reruns is not None: + monkeypatch.setattr( + app_module, "economy_submissions_for", lambda settings, database: reruns + ) + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + ) + ) + + return build + + +# Each re-run route, with a body it accepts. +ROUTES = [(REPORT_RUNS, {}), (SIMULATION_RUNS, None)] + + +def test_a_report_rerun_answers_201_with_its_new_run(client_for) -> None: + reruns = FakeEconomyReruns() + + with client_for(reruns) as client: + result = client.post( + REPORT_RUNS, + json={"force": True}, + headers={OBSERVABILITY_ID_HEADER: OBSERVABILITY_ID}, + ) + + assert result.status_code == 201 + assert result.json() == RERUN_REPORT.model_dump(mode="json") + ((report_id, request, request_id, context),) = reruns.report_reruns + assert (report_id, request) == (RECORD_ID, ReportRerunRequest(force=True)) + assert request_id != OBSERVABILITY_ID + assert context["observability_id"] == OBSERVABILITY_ID + + +def test_a_report_rerun_without_force_reuses_runs(client_for) -> None: + reruns = FakeEconomyReruns() + + with client_for(reruns) as client: + result = client.post(REPORT_RUNS, json={}) + + assert result.status_code == 201 + assert reruns.report_reruns[0][1] == ReportRerunRequest(force=False) + + +def test_a_simulation_rerun_answers_201_with_its_new_run(client_for) -> None: + reruns = FakeEconomyReruns() + + with client_for(reruns) as client: + result = client.post(SIMULATION_RUNS) + + assert result.status_code == 201 + assert result.json() == { + "simulation_id": str(RECORD_ID), + "simulation_run_id": str(RUN_ID), + "status": "pending", + } + ((simulation_id, _, context),) = reruns.simulation_reruns + assert simulation_id == RECORD_ID + assert context is not None + + +@pytest.mark.parametrize( + ("path", "body"), + [ + (REPORT_RUNS, {"force": "yes"}), + (REPORT_RUNS, {"force": None}), + (REPORT_RUNS, {"policyengine_version": "6.17.1"}), + ("/reports/not-a-uuid/runs", {}), + ("/simulations/not-a-uuid/runs", None), + ], +) +def test_a_malformed_rerun_is_refused_with_its_reasons(client_for, path, body): + reruns = FakeEconomyReruns() + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 422 + assert all({"loc", "msg"} <= set(reason) for reason in result.json()["detail"]) + assert (reruns.report_reruns, reruns.simulation_reruns) == ([], []) + + +@pytest.mark.parametrize( + ("path", "body", "record"), + [(REPORT_RUNS, {}, "report"), (SIMULATION_RUNS, None, "simulation")], +) +def test_an_unknown_record_answers_404(client_for, path, body, record) -> None: + reruns = FakeEconomyReruns(failure=EconomyRecordNotFoundError(record)) + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 404 + assert result.json() == { + "code": f"unknown_{record}", + "detail": f"There is no {record} {RECORD_ID}.", + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_a_run_under_way_answers_409_naming_it(client_for, path, body) -> None: + reruns = FakeEconomyReruns( + failure=EconomyRerunConflictError( + code="run_in_progress", detail="a run is under way", run_id=RUN_ID + ) + ) + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 409 + assert result.json() == { + "code": "run_in_progress", + "detail": "a run is under way", + "run_id": str(RUN_ID), + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_a_record_the_service_cannot_rerun_answers_409(client_for, path, body) -> None: + reruns = FakeEconomyReruns( + failure=EconomyRerunConflictError( + code="rerun_not_supported", detail="not budget windows yet" + ) + ) + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 409 + assert result.json() == { + "code": "rerun_not_supported", + "detail": "not budget windows yet", + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_a_refused_rerun_answers_400_or_422(client_for, path, body) -> None: + refused = FakeEconomyReruns( + failure=EconomySubmissionRefusedError( + code="unknown_region", detail="no region us, which the report names" + ) + ) + unprocessable = FakeEconomyReruns( + failure=EconomySubmissionUnprocessableError( + code="missing_parameters", + reasons=( + UnprocessableReason( + loc=("report", "reform", "policy_id", "gov.a"), + msg="no such parameter", + type="missing_parameter", + ), + ), + ) + ) + + with client_for(refused) as client: + refusal = client.post(path, json=body) + with client_for(unprocessable) as client: + unprocessable_answer = client.post(path, json=body) + + assert (refusal.status_code, refusal.json()["code"]) == (400, "unknown_region") + assert unprocessable_answer.status_code == 422 + assert unprocessable_answer.json() == { + "detail": [ + { + "loc": ["report", "reform", "policy_id", "gov.a"], + "msg": "no such parameter", + "type": "missing_parameter", + } + ], + "code": "missing_parameters", + } + + +@pytest.mark.parametrize( + ("path", "body", "record"), + [(REPORT_RUNS, {}, "report"), (SIMULATION_RUNS, None, "simulation")], +) +def test_a_rerun_modal_refused_answers_503_naming_the_record_and_its_run( + client_for, path, body, record +) -> None: + reruns = FakeEconomyReruns( + failure=EconomySubmissionUnavailableError( + code="modal_call_refused", + detail="Modal refused it", + record_id=RECORD_ID, + run_id=RUN_ID, + ) + ) + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 503 + assert result.headers["Retry-After"] == "10" + assert result.json() == { + "code": "modal_call_refused", + "detail": "Modal refused it", + f"{record}_id": str(RECORD_ID), + "run_id": str(RUN_ID), + } + + +@pytest.mark.parametrize( + ("path", "body", "record"), + [(REPORT_RUNS, {}, "report"), (SIMULATION_RUNS, None, "simulation")], +) +def test_a_rerun_whose_commit_is_unconfirmed_answers_500( + client_for, path, body, record +) -> None: + reruns = FakeEconomyReruns( + failure=EconomySubmissionOutcomeUnknownError( + record, error_type="OperationalError", rerun=True + ) + ) + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 500 + assert result.json() == { + "code": "internal_error", + "detail": ( + f"The service may have created the {record}'s new run, but could not " + "confirm it was written." + ), + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unexpected_failure_answers_the_json_500(client_for, path, body) -> None: + reruns = FakeEconomyReruns(failure=RuntimeError("secret detail")) + + with client_for(reruns) as client: + result = client.post(path, json=body) + + assert result.status_code == 500 + assert result.json() == { + "code": "internal_error", + "detail": "The service could not answer the request.", + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unconfigured_service_answers_503(client_for, path, body) -> None: + with client_for(None) as client: + result = client.post(path, json=body) + + assert result.status_code == 503 + assert result.json()["code"] == "not_configured" diff --git a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py new file mode 100644 index 000000000..42c33d5d0 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py @@ -0,0 +1,437 @@ +"""Re-running stored economy reports and simulations; the store and Modal are faked.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from uuid import UUID + +import pytest +from economy_rerun_fakes import ( + BASELINE_SIMULATION_ID, + REFORM_SIMULATION_ID, + REPORT_ID, + RUN_UNDER_WAY_ID, + SIMULATION_ID, + STORED_OPTIONS, + STORED_REPORT, + STORED_SIMULATION, + FakeRerunStore, + rerun_service, +) +from economy_service_fakes import ( + CATALOG, + OUTPUT_BUCKET, + REFORM, + FakeCaller, +) +from household_fixtures import POLICYENGINE_VERSION +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + ReportRunIds, + Stage12InvocationContext, +) +from policyengine_simulation_entry.economy_api_models import ReportRerunRequest +from policyengine_simulation_entry.economy_submission_errors import ( + EconomyRecordNotFoundError, + EconomyRerunConflictError, + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, +) +from policyengine_simulation_persistence import ( + CatalogDataset, + EconomyRecordOutcomeUnknownError, + MissingPolicyParametersError, + NewReportRerun, + NewSimulationRerun, + ReportRunStatus, + ReportSideError, + RerunSide, + ReusableRun, + RunInProgressError, + RunReuse, + RunVersions, + SimulationStatus, + UnknownRegionError, +) +from sqlalchemy.exc import OperationalError + +DATABASE_ERROR = OperationalError("SELECT 1", {}, Exception("server closed")) +VERSIONS = RunVersions( + country_package_version="2.2.1", policyengine_version=POLICYENGINE_VERSION +) +OBSERVED = {"observability_id": "observed"} +REUSED_BASELINE = ReusableRun( + simulation_id=BASELINE_SIMULATION_ID, + simulation_run_id=UUID("00000000-0000-4000-8000-0000000000b1"), +) + + +def _rerun_report(service, request: ReportRerunRequest | None = None): + return asyncio.run( + service.rerun_report( + REPORT_ID, + request or ReportRerunRequest(), + request_id="observed", + observability_context=OBSERVED, + ) + ) + + +def _rerun_simulation(service, simulation_id: UUID = SIMULATION_ID): + return asyncio.run( + service.rerun_simulation( + simulation_id, request_id="observed", observability_context=OBSERVED + ) + ) + + +def test_a_report_is_rerun_on_the_default_workers_versions_as_a_new_run() -> None: + store = FakeRerunStore(reusable=(REUSED_BASELINE, None)) + caller = FakeCaller() + + rerun = _rerun_report(rerun_service(store, caller=caller)) + + ((request, (baseline, reform, reuse)),) = store.rerun_reads + assert (request.country, request.policyengine_version, request.year) == ( + "us", + POLICYENGINE_VERSION, + 2026, + ) + assert (request.region_code, request.dataset_name) == ("us", "populace_us_2024") + assert (baseline, reform) == (STORED_REPORT.baseline, STORED_REPORT.reform) + assert reuse == RunReuse( + versions=VERSIONS, + inputs=STORED_OPTIONS, + output_uri_prefix=f"gs://{OUTPUT_BUCKET}/", + ) + assert store.report_reruns == [ + NewReportRerun( + report_id=REPORT_ID, + run_id=UUID(int=1), + population=CATALOG.population(2026), + versions=VERSIONS, + baseline=RerunSide( + simulation_id=BASELINE_SIMULATION_ID, + simulation_run_id=REUSED_BASELINE.simulation_run_id, + reused=True, + ), + reform=RerunSide( + simulation_id=REFORM_SIMULATION_ID, + simulation_run_id=UUID(int=2), + reused=False, + ), + ) + ] + ((_, payload, context, observed),) = caller.report_runs + assert observed == OBSERVED + report = ReportExecutionInput.model_validate(payload) + assert report.evaluation_id == UUID(int=1) + assert report.baseline.simulation_execution_id == REUSED_BASELINE.simulation_run_id + assert (report.reform.simulation_execution_id, report.reform.policy) == ( + UUID(int=2), + REFORM, + ) + assert report.reform.options == STORED_OPTIONS + assert Stage12InvocationContext.model_validate(context).report_run == ReportRunIds( + report_run_id=UUID(int=1), + baseline_simulation_run_id=REUSED_BASELINE.simulation_run_id, + reform_simulation_run_id=UUID(int=2), + ) + assert store.recorded_calls == [(UUID(int=1), "fc-coordinator")] + assert rerun.model_dump() == { + "report_id": REPORT_ID, + "report_run_id": UUID(int=1), + "status": ReportRunStatus.PENDING, + "baseline_simulation_id": BASELINE_SIMULATION_ID, + "baseline_simulation_run_id": REUSED_BASELINE.simulation_run_id, + "baseline_reused": True, + "reform_simulation_id": REFORM_SIMULATION_ID, + "reform_simulation_run_id": UUID(int=2), + "reform_reused": False, + } + + +def test_a_forced_rerun_recalculates_both_simulations() -> None: + store = FakeRerunStore(reusable=(REUSED_BASELINE, None)) + + rerun = _rerun_report(rerun_service(store), ReportRerunRequest(force=True)) + + assert store.rerun_reads[0][1][2] is None + (created,) = store.report_reruns + assert (created.baseline.reused, created.reform.reused) == (False, False) + assert created.new_simulation_run_ids == (UUID(int=2), UUID(int=3)) + assert created.baseline.simulation_id == BASELINE_SIMULATION_ID + assert (rerun.baseline_reused, rerun.reform_reused) == (False, False) + + +def test_without_the_output_bucket_a_rerun_looks_for_no_run_to_reuse() -> None: + store = FakeRerunStore(reusable=(REUSED_BASELINE, None)) + + _rerun_report(rerun_service(store, output_bucket=None)) + + assert store.rerun_reads[0][1][2] is None + + +def test_an_unknown_report_or_simulation_is_not_found() -> None: + store = FakeRerunStore(report_definition=None, simulation_definition=None) + service = rerun_service(store) + + with pytest.raises(EconomyRecordNotFoundError) as report: + _rerun_report(service) + with pytest.raises(EconomyRecordNotFoundError) as simulation: + _rerun_simulation(service) + + assert (report.value.record, simulation.value.record) == ("report", "simulation") + + +@pytest.mark.parametrize( + ("kind", "named"), + [ + ("budget_window", "budget-window"), + ("household", "household"), + ("unsupported", "this type of"), + ], +) +def test_a_report_kind_the_service_does_not_rerun_yet_is_refused(kind, named) -> None: + store = FakeRerunStore(report_definition=replace(STORED_REPORT, kind=kind)) + caller = FakeCaller() + + with pytest.raises(EconomyRerunConflictError) as refused: + _rerun_report(rerun_service(store, caller=caller)) + + assert refused.value.code == "rerun_not_supported" + assert named in refused.value.detail + assert (store.rerun_reads, caller.report_runs) == ([], []) + + +@pytest.mark.parametrize( + ("missing", "named"), + [ + ({"baseline": None}, "baseline simulation"), + ({"reform": None}, "reform simulation"), + ({"region_code": None}, "region"), + ({"dataset_name": None}, "dataset"), + ({"year": None}, "year"), + ], +) +def test_a_report_missing_a_row_it_calculates_on_is_not_rerun(missing, named) -> None: + store = FakeRerunStore(report_definition=replace(STORED_REPORT, **missing)) + + with pytest.raises(EconomyRerunConflictError) as refused: + _rerun_report(rerun_service(store)) + + assert refused.value.code == "rerun_not_supported" + assert named in refused.value.detail + assert store.rerun_reads == [] + + +def test_a_rerun_is_refused_while_a_run_is_under_way() -> None: + store = FakeRerunStore(rerun_failure=RunInProgressError(RUN_UNDER_WAY_ID)) + caller = FakeCaller() + service = rerun_service(store, caller=caller) + + with pytest.raises(EconomyRerunConflictError) as report: + _rerun_report(service) + with pytest.raises(EconomyRerunConflictError) as simulation: + _rerun_simulation(service) + + for refused, record in ((report, "report"), (simulation, "simulation")): + assert refused.value.code == "run_in_progress" + assert refused.value.run_id == RUN_UNDER_WAY_ID + assert f"The {record} has a run" in refused.value.detail + assert (caller.report_runs, caller.simulation_runs) == ([], []) + + +def test_a_rerun_modal_refuses_fails_only_its_new_runs_and_names_them() -> None: + store = FakeRerunStore(reusable=(REUSED_BASELINE, None)) + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _rerun_report( + rerun_service(store, caller=FakeCaller(failure=ConnectionError("no"))) + ) + + error = unavailable.value + assert (error.code, error.record_id, error.run_id) == ( + "modal_call_refused", + REPORT_ID, + UUID(int=1), + ) + assert "the report's new run" in error.detail + ((run_id, simulation_run_ids, message),) = store.failed_runs + assert (run_id, simulation_run_ids) == (UUID(int=1), (UUID(int=2),)) + assert message.startswith("modal_call_refused: Modal did not accept the report") + + +def test_a_rerun_its_coordinator_claimed_meanwhile_is_answered_created() -> None: + store = FakeRerunStore(claimed=True) + + rerun = _rerun_report( + rerun_service(store, caller=FakeCaller(failure=TimeoutError("late"))) + ) + + assert store.failed_runs == [] + assert rerun.report_run_id == UUID(int=1) + + +def test_a_stored_policy_the_current_catalog_cannot_set_is_located_on_its_side() -> ( + None +): + store = FakeRerunStore( + inputs_failure=ReportSideError( + "reform", MissingPolicyParametersError(("gov.old.parameter",)) + ) + ) + + with pytest.raises(EconomySubmissionUnprocessableError) as refused: + _rerun_report(rerun_service(store)) + + assert refused.value.code == "missing_parameters" + ((reason,),) = (refused.value.reasons,) + assert reason.loc == ("report", "reform", "policy_id", "gov.old.parameter") + assert POLICYENGINE_VERSION in reason.msg + assert store.report_reruns == [] + + +def test_a_stored_region_the_current_catalog_lacks_is_refused_naming_the_report() -> ( + None +): + store = FakeRerunStore(inputs_failure=UnknownRegionError("the region is gone")) + + with pytest.raises(EconomySubmissionRefusedError) as refused: + _rerun_report(rerun_service(store)) + + assert refused.value.code == "unknown_region" + assert refused.value.detail == ( + f"us's catalog version for .py {POLICYENGINE_VERSION} has no region us, " + "which the report names." + ) + + +def test_a_stored_dataset_the_bundle_lacks_is_refused() -> None: + store = FakeRerunStore( + catalog=replace( + CATALOG, dataset=CatalogDataset(id=CATALOG.dataset.id, name="cps_2019") + ) + ) + + with pytest.raises(EconomySubmissionRefusedError) as refused: + _rerun_report(rerun_service(store)) + + assert refused.value.code == "dataset_not_in_bundle" + assert "cps_2019" in refused.value.detail + assert store.report_reruns == [] + + +def test_a_database_that_cannot_be_read_is_unavailable_before_any_call() -> None: + store = FakeRerunStore(read_failure=DATABASE_ERROR) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _rerun_report(rerun_service(store, caller=caller)) + + assert unavailable.value.code == "database_unavailable" + assert caller.report_runs == [] + + +def test_a_rerun_whose_commit_is_unconfirmed_may_have_written_its_run() -> None: + store = FakeRerunStore( + rerun_failure=EconomyRecordOutcomeUnknownError("commit failed") + ) + store.rerun_failure.__cause__ = DATABASE_ERROR + + with pytest.raises(EconomySubmissionOutcomeUnknownError) as unknown: + _rerun_report(rerun_service(store)) + + assert (unknown.value.record, unknown.value.rerun) == ("report", True) + assert unknown.value.error_type == "OperationalError" + + +def test_a_simulation_is_rerun_as_a_new_run_of_it() -> None: + store = FakeRerunStore() + caller = FakeCaller() + + rerun = _rerun_simulation(rerun_service(store, caller=caller)) + + ((request, choice),) = store.simulation_reads + assert (request.country, request.region_code, request.dataset_name) == ( + "us", + "us", + "populace_us_2024", + ) + assert (request.policyengine_version, request.year) == (POLICYENGINE_VERSION, 2027) + assert choice == STORED_SIMULATION.choice + assert store.simulation_reruns == [ + NewSimulationRerun( + simulation_id=SIMULATION_ID, + run_id=UUID(int=1), + population=CATALOG.population(2027), + versions=VERSIONS, + ) + ] + ((_, payload, context, observed),) = caller.simulation_runs + assert payload["simulation_execution_id"] == str(UUID(int=1)) + assert (payload["policy"], payload["year"]) == (REFORM, 2027) + assert context["simulation_run_id"] == str(UUID(int=1)) + assert observed == OBSERVED + assert store.recorded_calls == [(UUID(int=1), "fc-worker")] + assert rerun.model_dump() == { + "simulation_id": SIMULATION_ID, + "simulation_run_id": UUID(int=1), + "status": SimulationStatus.PENDING, + } + + +@pytest.mark.parametrize( + ("changes", "named"), + [ + ({"kind": "household"}, "household"), + ({"kind": "unsupported"}, "this type of"), + ({"country": None}, "country"), + ({"region_code": None}, "region"), + ({"dataset_name": None}, "dataset"), + ({"year": None}, "year"), + ], +) +def test_a_simulation_the_service_cannot_rerun_is_refused(changes, named) -> None: + store = FakeRerunStore(simulation_definition=replace(STORED_SIMULATION, **changes)) + + with pytest.raises(EconomyRerunConflictError) as refused: + _rerun_simulation(rerun_service(store)) + + assert refused.value.code == "rerun_not_supported" + assert named in refused.value.detail + assert store.simulation_reads == [] + + +def test_a_simulation_rerun_modal_refuses_is_failed_and_named() -> None: + store = FakeRerunStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _rerun_simulation( + rerun_service(store, caller=FakeCaller(failure=ConnectionError("no"))) + ) + + error = unavailable.value + assert (error.code, error.record_id, error.run_id) == ( + "modal_call_refused", + SIMULATION_ID, + UUID(int=1), + ) + assert store.failed_runs[0][0] == UUID(int=1) + + +def test_a_simulations_stored_policy_is_located_on_the_simulation() -> None: + store = FakeRerunStore( + inputs_failure=MissingPolicyParametersError( + ("gov.old.elasticity",), source="dynamic" + ) + ) + + with pytest.raises(EconomySubmissionUnprocessableError) as refused: + _rerun_simulation(rerun_service(store)) + + ((reason,),) = (refused.value.reasons,) + assert reason.loc == ("simulation", "dynamic_id", "gov.old.elasticity") diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 30f8022af..a44bfb5b0 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -40,9 +40,11 @@ class OpenAPIDocument(TypedDict): ("POST", "/household-calculations"), ("POST", "/ping"), ("POST", "/reports"), + ("POST", "/reports/{report_id}/runs"), ("POST", "/simulate/economy/budget-window"), ("POST", "/simulate/economy/comparison"), ("POST", "/simulations"), + ("POST", "/simulations/{simulation_id}/runs"), } @@ -86,8 +88,10 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: ECONOMY_PATHS = ( "/reports", "/reports/{report_id}", + "/reports/{report_id}/runs", "/simulations", "/simulations/{simulation_id}", + "/simulations/{simulation_id}/runs", ) ECONOMY_SCHEMAS = ( "EconomyError", @@ -100,6 +104,7 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: "EconomySimulationRequest", "EconomySimulationState", "PolicyChoice", + "ReportRerunRequest", "ReportRunRead", "ReportRunStatus", "SimulationRunRead", @@ -376,7 +381,9 @@ def test_economy_routes_declare_every_answer(): assert f"{record}_id" in responses["503"]["description"] assert "nothing was written" in responses["400"]["description"] error_properties = spec["components"]["schemas"]["EconomyError"]["properties"] - assert {"code", "detail", "report_id", "simulation_id"} == set(error_properties) + assert {"code", "detail", "report_id", "simulation_id", "run_id"} == set( + error_properties + ) codes = set(error_properties["code"]["enum"]) assert { "unknown_region", @@ -397,6 +404,48 @@ def test_every_economy_503_declares_its_retry_after_header(): assert headers["Retry-After"]["schema"] == {"type": "integer"} +def test_rerun_routes_declare_every_answer(): + spec = create_app().openapi() + error = {"$ref": "#/components/schemas/EconomyError"} + unprocessable = {"$ref": "#/components/schemas/EconomyRequestUnprocessable"} + + for path, rerun in ( + ("/reports/{report_id}/runs", "EconomyReportCreated"), + ("/simulations/{simulation_id}/runs", "EconomySimulationCreated"), + ): + responses = _responses(spec, "post", path) + assert set(responses) == { + "201", + "400", + "403", + "404", + "409", + "422", + "500", + "503", + } + assert _response_schema(responses["201"]) == { + "$ref": f"#/components/schemas/{rerun}" + } + for status_code in ("400", "404", "409", "500", "503"): + assert _response_schema(responses[status_code]) == error + assert _response_schema(responses["422"]) == unprocessable + assert "run_id" in responses["409"]["description"] + assert "run_id" in responses["503"]["description"] + codes = set( + spec["components"]["schemas"]["EconomyError"]["properties"]["code"]["enum"] + ) + assert {"run_in_progress", "rerun_not_supported"} <= codes + report_rerun = spec["components"]["schemas"]["ReportRerunRequest"] + assert set(report_rerun["properties"]) == {"force"} + assert ( + "requestBody" not in spec["paths"]["/simulations/{simulation_id}/runs"]["post"] + ) + names = operation_ids(cast("OpenAPIDocument", spec)) + assert names[("POST", "/reports/{report_id}/runs")] == "rerun_report" + assert names[("POST", "/simulations/{simulation_id}/runs")] == "rerun_simulation" + + def test_economy_routes_name_their_operations(): names = operation_ids(cast("OpenAPIDocument", create_app().openapi())) From 593e89415ad61f364149027c2d1451418a9e169d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:57:54 +0400 Subject: [PATCH 04/18] Re-run economy reports and simulations through POST /reports/{id}/runs and POST /simulations/{id}/runs A re-run creates the stored record's new run on the v2 version manifest's default worker and starts it by run ID, as a new record's first run is started; every earlier run is kept. A report's re-run reuses a run of each side's own simulation that succeeded on the current versions unless force is set, and its run's trigger is manual. A simulation's re-run is always calculated. A record with a run pending or running, a budget-window or household record, or one missing a row it calculates on is answered 409; an unknown record 404. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/app.py | 22 +- .../economy_api_models.py | 43 +- .../economy_rerun_routes.py | 278 ++++++++++ .../economy_reruns.py | 513 ++++++++++++++++++ .../economy_routes.py | 104 ++-- .../economy_submission_errors.py | 104 +++- .../economy_submissions.py | 96 +++- .../tests/test_economy_reruns.py | 2 + 8 files changed, 1058 insertions(+), 104 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.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 0c99f05e4..fce6b1b7a 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -52,14 +52,13 @@ ) from policyengine_simulation_entry.backend import OldGatewayBackend, SimulationBackend from policyengine_simulation_entry.config import Settings -from policyengine_simulation_entry.economy_routes import ( - EconomySubmissions, - add_economy_routes, -) -from policyengine_simulation_entry.economy_submissions import ( - EconomySubmissionService, - ServiceSettings, +from policyengine_simulation_entry.economy_rerun_routes import ( + EconomyService, + add_economy_rerun_routes, ) +from policyengine_simulation_entry.economy_reruns import EconomyRerunService +from policyengine_simulation_entry.economy_routes import add_economy_routes +from policyengine_simulation_entry.economy_submissions import ServiceSettings from policyengine_simulation_entry.economy_workers import ( ManifestEconomyWorkers, ModalEconomyRunCaller, @@ -540,18 +539,19 @@ def household_calculations_for( def _register_economy_routes( app: FastAPI, entry: EntryRuntime, - submissions: EconomySubmissions | None, + submissions: EconomyService | None, *, submitter: AuthorizationDependency, reader: AuthorizationDependency, ) -> None: - """Add the routes that create economy reports and simulations and read them.""" + """Add the routes that create, read and re-run economy reports and simulations.""" add_economy_routes(app, entry, submissions, submitter=submitter, reader=reader) + add_economy_rerun_routes(app, entry, submissions, submitter=submitter) def economy_submissions_for( settings: Settings, database: SimulationDatabaseConnection | None -) -> EconomySubmissions | None: +) -> EconomyService | None: """Build the economy submission service when the service is configured. It needs the simulation database and the v2 Modal environment; without @@ -562,7 +562,7 @@ def economy_submissions_for( database, SimulationDatabase ): return None - return EconomySubmissionService( + return EconomyRerunService( store=EconomyRunStore.from_engine( database.engine, statement_timeout_milliseconds=ECONOMY_STATEMENT_TIMEOUT_MILLISECONDS, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py index b161c91e7..8ecf0cc1b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -9,16 +9,24 @@ Answers the Public API reads: -- 201 with the created record's ID and its first run's ID; +- 201 with the created record's ID and its first run's ID, or, for a + re-run, the record's ID and its new run's ID; - 400 with ``code`` and ``detail``: one of ``CallerRowRefusalCode`` when the region, dataset or a dynamic the request names does not exist, or another code for the service's own state, such as a catalog version missing while a release deploys. Nothing is written; +- 404 ``unknown_report`` or ``unknown_simulation`` when a re-run names no + stored record; +- 409 with ``code`` and ``detail`` when a stored record is not re-run now: + ``run_in_progress``, naming the run pending or running (``run_id``), or + ``rerun_not_supported`` for a record this service does not re-run. + Nothing is written; - 422 when the request's fields, or the stored values it names, cannot be calculated. Nothing is written; - 503 with ``code`` and ``detail`` when the service cannot answer now. Only ``modal_call_refused`` follows a write, and it names the record created - (``report_id`` or ``simulation_id``), whose run was failed; + or re-run (``report_id`` or ``simulation_id``), whose new run was failed, + and for a re-run that run (``run_id``); - 500 with ``code`` ``internal_error``, possibly after a write: an unexpected failure, or a commit that failed after the rows were inserted. @@ -63,6 +71,14 @@ EconomyRefusalCode = Literal[CallerRowRefusalCode, ServiceStateRefusalCode] # Why stored values cannot be calculated (422). EconomyUnprocessableCode = Literal["missing_parameters", "overlapping_parameters"] +# Why a stored record is not re-run now (409). +EconomyConflictCode = Literal[ + # The record has a run pending or running; ``run_id`` names it. + "run_in_progress", + # The service does not re-run this kind of record, or the record lacks a + # row it calculates on. + "rerun_not_supported", +] # Why a route cannot answer now (503); retry after ``Retry-After`` seconds. EconomyUnavailableCode = Literal[ # The service has no database or v2 Modal environment configured. @@ -77,6 +93,7 @@ EconomyErrorCode = Literal[ EconomyRefusalCode, EconomyUnprocessableCode, + EconomyConflictCode, EconomyUnavailableCode, "internal_error", "unknown_report", @@ -161,12 +178,23 @@ def options(self) -> dict[str, Any]: return options +class ReportRerunRequest(_Request): + """A stored report's re-run on the default worker's versions. + + Without ``force``, a side reuses a run of its own simulation that + succeeded on those versions, as a new report's side does; with it, both + simulations are calculated again. + """ + + force: StrictBool = False + + class _Answer(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) class EconomySimulationCreated(_Answer): - """The simulation a request created, its first run and that run's status. + """The simulation a request created or re-ran, its new run and its status. The status is the run's when its call's ID was recorded: ``pending``, or whatever the simulation function had made it by then if it claimed the @@ -180,7 +208,7 @@ class EconomySimulationCreated(_Answer): class EconomyReportCreated(_Answer): - """The report a request created, its first run and that run's status. + """The report a request created or re-ran, its new run and its status. The status is the run's when its call's ID was recorded: ``pending``, or whatever the coordinator had made it by then if it claimed the run @@ -256,14 +284,17 @@ class EconomyError(_Answer): ``report_id`` (on ``/reports``) or ``simulation_id`` (on ``/simulations``) is set only with ``modal_call_refused``: it names the - record the service created, whose run it failed because Modal refused - the call or didn't accept it in time. + record the service created or re-ran, whose new run it failed because + Modal refused the call or didn't accept it in time. ``run_id`` names + that new run on a re-run, and with ``run_in_progress`` the record's run + that is pending or running. """ code: EconomyErrorCode detail: str report_id: UUID | None = None simulation_id: UUID | None = None + run_id: UUID | None = None class EconomyRequestReason(BaseModel): diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py new file mode 100644 index 000000000..98f8f50cf --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -0,0 +1,278 @@ +"""The re-run routes: calculate a stored report or simulation again as a new run. + +- ``POST /reports/{report_id}/runs`` (permission ``simulation:submit``) + creates the stored report's new run on the v2 version manifest's default + worker and starts its coordinator, answering 201 with the report's ID, the + new run's ID and status, and each side's simulation and simulation run. + ``{"force": true}`` calculates both simulations again; otherwise a side + reuses a run of its own simulation that succeeded on the same versions. +- ``POST /simulations/{simulation_id}/runs`` (``simulation:submit``) creates + the stored simulation's new run and starts its simulation function, + answering 201 with the simulation's ID, the new run's ID and its status. + +Every earlier run is kept. A re-run answers 404 for a record that does not +exist, and 409 when the record has a run pending or running +(``run_in_progress``, naming it in ``run_id``) or is one this service does +not re-run (``rerun_not_supported``); then nothing is written. The other +answers are the creates': 400 or 422 when the stored record cannot be +calculated on the current catalog version, 503 with ``Retry-After`` before +any write or when Modal refused the new run's call (then the run was failed, +and the answer names the record and the run), and the JSON 500. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID + +from fastapi import APIRouter, Depends, Request, status +from fastapi.responses import JSONResponse + +from policyengine_simulation_entry.economy_api_models import ( + EconomyError, + EconomyReportCreated, + EconomyRequestUnprocessable, + EconomySimulationCreated, + ReportRerunRequest, +) +from policyengine_simulation_entry.economy_routes import ( + ERROR_ANSWER, + FORBIDDEN_ANSWER, + RETRY_AFTER_HEADER, + SUBMISSION_FAILURES, + EconomySubmissions, + call_request_id_for, + error_response, + not_configured_response, + submission_failure_response, + unexpected_failure_response, + unknown_record_response, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomyRecordNotFoundError, + EconomyRerunConflictError, +) +from policyengine_simulation_entry.route_support import observability_context + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from fastapi import FastAPI + from pydantic import BaseModel + + from policyengine_simulation_entry.auth import AuthorizationDependency + from policyengine_simulation_entry.economy_submission_errors import RecordName + from policyengine_simulation_entry.entry_runtime import EntryRuntime + +REPORT_RUNS_PATH = "/reports/{report_id}/runs" +SIMULATION_RUNS_PATH = "/simulations/{simulation_id}/runs" + + +class EconomyReruns(Protocol): + """Re-run stored economy reports and simulations (``EconomyRerunService``).""" + + async def rerun_report( + self, + report_id: UUID, + request: ReportRerunRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + """Create the report's new run and start its coordinator.""" + ... + + async def rerun_simulation( + self, + simulation_id: UUID, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + """Create the simulation's new run and start its function.""" + ... + + +class EconomyService(EconomySubmissions, EconomyReruns, Protocol): + """Create, read and re-run economy reports and simulations.""" + + +def _responses(record: RecordName) -> dict[int | str, dict[str, Any]]: + """Return every answer a re-run of ``record`` declares.""" + return { + 201: {"description": f"The {record}'s new run was created and started"}, + 400: { + "description": ( + f"The stored {record}'s region, dataset, policy or dynamic is " + "missing from the catalog version of the default worker's .py " + "version, or a function it needs is missing; nothing was " + "written. Each code is the stored record's or the service's " + "state, never the caller's" + ), + **ERROR_ANSWER, + }, + 403: FORBIDDEN_ANSWER, + 404: { + "description": f"There is no such {record} (unknown_{record})", + **ERROR_ANSWER, + }, + 409: { + "description": ( + f"The {record} is not re-run now, and nothing was written: it " + "has a run pending or running (run_in_progress, named by " + "run_id), or the service does not re-run it " + "(rerun_not_supported)" + ), + **ERROR_ANSWER, + }, + 422: { + "description": ( + "The ID or the body is malformed, or the stored policies and " + "dynamics cannot be calculated on the current catalog version " + "(missing_parameters or overlapping_parameters, located under " + f"{record}); nothing was written" + ), + "model": EconomyRequestUnprocessable, + }, + 500: { + "description": ( + "The service failed unexpectedly, or could not confirm the new " + "run was written (internal_error), possibly after creating it" + ), + **ERROR_ANSWER, + }, + 503: { + "description": ( + "Economy runs are not configured, or the v2 version manifest or " + "the database cannot be read, and nothing was written; or Modal " + "refused the new run's call, or didn't accept it in time " + f"(modal_call_refused): then the {record} named by {record}_id " + "has a new run, named by run_id, which was failed. Retry after " + "Retry-After seconds" + ), + "headers": RETRY_AFTER_HEADER, + **ERROR_ANSWER, + }, + } + + +def _rerun_failure_response( + error: Exception, record: RecordName, record_id: UUID +) -> JSONResponse: + """Answer a re-run the service refused or could not serve.""" + if isinstance(error, EconomyRecordNotFoundError): + return unknown_record_response(record, record_id) + if isinstance(error, EconomyRerunConflictError): + return error_response( + status.HTTP_409_CONFLICT, + EconomyError(code=error.code, detail=error.detail, run_id=error.run_id), + ) + return submission_failure_response(error, record) + + +_RERUN_FAILURES = ( + *SUBMISSION_FAILURES, + EconomyRecordNotFoundError, + EconomyRerunConflictError, +) + + +def add_economy_rerun_routes( + app: FastAPI, + entry: EntryRuntime, + reruns: EconomyReruns | None, + *, + submitter: AuthorizationDependency, +) -> None: + """Add the re-run routes; without a service they answer 503.""" + log_string_limit = entry.runtime.config.limits.max_string_length + router = APIRouter() + + async def answer( # the route, its record and how to re-run it + request: Request, + rerun: Callable[ + [EconomyReruns, str, dict[str, Any] | None], Awaitable[BaseModel] + ], + *, + route: str, + record: RecordName, + record_id: UUID, + ) -> JSONResponse: + if reruns is None: + return not_configured_response() + observability_id = entry.start_workflow_observability(request) + try: + created = await rerun( + reruns, + call_request_id_for(request), + observability_context(entry, observability_id), + ) + except _RERUN_FAILURES as error: + return _rerun_failure_response(error, record, record_id) + # Every other failure is still answered as the JSON 500 declared. + except Exception as error: # noqa: BLE001 # answered as internal_error + return unexpected_failure_response( + route, error, max_characters=log_string_limit + ) + return JSONResponse( + created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED + ) + + @router.post( + REPORT_RUNS_PATH, + summary="Re-run an economy report", + description=( + "Create the stored report's new run on the v2 version manifest's " + "default worker, and start it; every earlier run is kept. Without " + "force, a side reuses a run of its own simulation that succeeded " + "on the same versions; with force, both simulations are " + "calculated again." + ), + operation_id="rerun_report", + status_code=status.HTTP_201_CREATED, + response_model=EconomyReportCreated, + responses=_responses("report"), + dependencies=[Depends(submitter)], + ) + async def rerun_report( + report_id: UUID, body: ReportRerunRequest, request: Request + ) -> JSONResponse: + return await answer( + request, + lambda service, request_id, context: service.rerun_report( + report_id, + body, + request_id=request_id, + observability_context=context, + ), + route=REPORT_RUNS_PATH, + record="report", + record_id=report_id, + ) + + @router.post( + SIMULATION_RUNS_PATH, + summary="Re-run an economy simulation", + description=( + "Create the stored simulation's new run on the v2 version " + "manifest's default worker, and start it; every earlier run is " + "kept, and the new run is always calculated." + ), + operation_id="rerun_simulation", + status_code=status.HTTP_201_CREATED, + response_model=EconomySimulationCreated, + responses=_responses("simulation"), + dependencies=[Depends(submitter)], + ) + async def rerun_simulation(simulation_id: UUID, request: Request) -> JSONResponse: + return await answer( + request, + lambda service, request_id, context: service.rerun_simulation( + simulation_id, request_id=request_id, observability_context=context + ), + route=SIMULATION_RUNS_PATH, + record="simulation", + record_id=simulation_id, + ) + + app.include_router(router) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py new file mode 100644 index 000000000..5d695407d --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -0,0 +1,513 @@ +"""Re-run a stored economy report or simulation on the current .py version. + +A re-run calculates a stored record again on the v2 version manifest's +default worker, as a new run, and keeps every earlier run. The service: + +1. reads what the record calculates: its country, year, region code, + dataset name and each simulation's policy and dynamic. A record this + service does not re-run, or one missing a row it calculates on, is + refused (409 ``rerun_not_supported``), as is one that does not exist + (404); +2. finds the default worker's economy functions for the country, and reads, + in one transaction, the country's catalog version for that worker's .py + version, its region and dataset with the record's code and name, and each + simulation's policy and dynamic values on that catalog version. For a + report, unless the re-run is forced, each side's reusable run is a run of + that side's own simulation that succeeded on the current versions, as a + new report's would be (``find_reusable_run_of``); +3. creates, in one transaction, the new run (a report's with trigger + ``manual``) and each report side's new simulation run of its own + simulation, all pending with no Modal call ID, unless the record has a + run pending or running (409 ``run_in_progress``); +4. starts and records the new run's call as a new record's: the report run's + coordinator, or the simulation function, by run ID (C1 ยท 8's + coordination), failing the new runs when Modal refuses the call. + +Whatever is missing or cannot be calculated is refused in steps 1 and 2, +before anything is written. A budget-window report and a household report +or simulation are refused for now; each gets its re-run with its own +submission. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, Final, Protocol, get_args +from uuid import UUID, uuid4 + +from policyengine_simulation_contract.stage12_bundle import CountryId +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_simulation_persistence import ( + EconomyRequest, + NewReportRerun, + NewSimulationRerun, + RerunSide, + RunInProgressError, + RunReuse, +) + +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportCreated, + EconomySimulationCreated, +) +from policyengine_simulation_entry.economy_calculation_inputs import ( + CalculationScope, + DatasetNotInBundleError, + bundle_dataset, + report_run_calculation, + simulation_options, + simulation_run_calculation, +) +from policyengine_simulation_entry.economy_submission_errors import ( + STORED_REPORT, + STORED_SIMULATION, + EconomyRecordNotFoundError, + EconomyRerunConflictError, + EconomySubmissionRefusedError, +) +from policyengine_simulation_entry.economy_submissions import ( + CALL_START_TIMEOUT_SECONDS, + EconomyRuns, + EconomySubmissionService, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + from policyengine_simulation_persistence import ( + EconomyCatalogRows, + EconomyReportInputs, + RecordKind, + ReportDefinition, + SideDefinition, + SimulationDefinition, + SimulationInputs, + ) + + from policyengine_simulation_entry.economy_api_models import ( + ReportRerunRequest, + ) + from policyengine_simulation_entry.economy_submission_errors import RecordName + from policyengine_simulation_entry.economy_submissions import ServiceSettings + from policyengine_simulation_entry.economy_workers import ( + EconomyRunCaller, + EconomyWorker, + EconomyWorkerSource, + ) + +# How a refusal names each kind of stored record this service does not re-run +# yet. +_NOT_RERUN: Final[dict[RecordKind, str]] = { + "budget_window": "budget-window", + "household": "household", +} +# The countries this service calculates, as the contract names them. +_COUNTRIES: Final[tuple[CountryId, ...]] = get_args(CountryId) +_REPORT: Final[RecordName] = "report" +_SIMULATION: Final[RecordName] = "simulation" + + +class EconomyRerunRuns(EconomyRuns, Protocol): + """The economy store's reads and writes, with re-runs' (``EconomyRunStore``).""" + + def read_report_definition(self, report_id: UUID) -> ReportDefinition | None: + """Return what the stored report calculates, or None.""" + ... + + def read_simulation_definition( + self, simulation_id: UUID + ) -> SimulationDefinition | None: + """Return what the stored simulation calculates, or None.""" + ... + + def read_report_rerun_inputs( + self, + request: EconomyRequest, + *, + baseline: SideDefinition, + reform: SideDefinition, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + """Read a re-run's catalog rows, both sides' values and reusable runs.""" + ... + + def create_report_rerun(self, rerun: NewReportRerun) -> None: + """Create the report's new run and its sides' new simulation runs.""" + ... + + def create_simulation_rerun(self, rerun: NewSimulationRerun) -> None: + """Create the simulation's new run.""" + ... + + +class EconomyRerunService(EconomySubmissionService): + """Create and re-run economy reports and simulations, and start their calls.""" + + def __init__( # noqa: PLR0913 # each seam tests replace + self, + *, + store: EconomyRerunRuns, + workers: EconomyWorkerSource, + caller: EconomyRunCaller, + settings: ServiceSettings, + call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, + new_id: Callable[[], UUID] = uuid4, + clock: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + """Use these seams, as ``EconomySubmissionService`` does.""" + super().__init__( + store=store, + workers=workers, + caller=caller, + settings=settings, + call_start_timeout_seconds=call_start_timeout_seconds, + new_id=new_id, + clock=clock, + ) + self._reruns = store + + async def rerun_report( + self, + report_id: UUID, + request: ReportRerunRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + """Create the stored report's new run on the current versions, and start it. + + Raise EconomyRecordNotFoundError (404) or EconomyRerunConflictError + (409) as the module describes, and otherwise the errors + ``submit_report`` raises, a 503 for Modal's refusal naming the new + run as well as the report. + """ + definition = await self._read_state( + lambda: self._reruns.read_report_definition(report_id) + ) + if definition is None: + raise EconomyRecordNotFoundError(_REPORT) + stored = _rerunnable_report(definition) + worker = await self._default_worker(stored.country) + economy_request = EconomyRequest( + country=stored.country, + policyengine_version=worker.policyengine_version, + region_code=stored.region_code, + dataset_name=stored.dataset_name, + year=stored.year, + ) + reuse = ( + None + if request.force or self._settings.reuse_prefix is None + else RunReuse( + versions=worker.versions, + inputs=definition.inputs, + output_uri_prefix=self._settings.reuse_prefix, + ) + ) + inputs = await self._read_inputs( + lambda: self._reruns.read_report_rerun_inputs( + economy_request, + baseline=stored.baseline, + reform=stored.reform, + reuse=reuse, + ), + request=economy_request, + source=STORED_REPORT, + ) + rerun = NewReportRerun( + report_id=report_id, + run_id=self._new_id(), + population=inputs.catalog.population(stored.year), + versions=worker.versions, + baseline=self._rerun_side(stored.baseline, inputs.baseline), + reform=self._rerun_side(stored.reform, inputs.reform), + ) + calculation = report_run_calculation( + worker, + _rerun_scope( + stored.country, + economy_request, + worker, + inputs.catalog, + options=definition.inputs, + ), + self._origin(request_id), + run_ids=ReportRunIds( + report_run_id=rerun.run_id, + baseline_simulation_run_id=rerun.baseline.simulation_run_id, + reform_simulation_run_id=rerun.reform.simulation_run_id, + ), + baseline_policy=inputs.baseline.reform, + reform_policy=inputs.reform.reform, + ) + await self._create_rerun( + _REPORT, lambda: self._reruns.create_report_rerun(rerun) + ) + status = await self._start_report_run( + worker, + calculation, + report_id=report_id, + run_id=rerun.run_id, + new_simulation_run_ids=rerun.new_simulation_run_ids, + observability_context=observability_context, + rerun=True, + ) + return EconomyReportCreated( + report_id=report_id, + report_run_id=rerun.run_id, + status=status, + baseline_simulation_id=rerun.baseline.simulation_id, + baseline_simulation_run_id=rerun.baseline.simulation_run_id, + baseline_reused=rerun.baseline.reused, + reform_simulation_id=rerun.reform.simulation_id, + reform_simulation_run_id=rerun.reform.simulation_run_id, + reform_reused=rerun.reform.reused, + ) + + async def rerun_simulation( + self, + simulation_id: UUID, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + """Create the stored simulation's new run on the current versions, and start it. + + The run is always calculated: a simulation run is the calculation, + with nothing of it to reuse. Raise the errors ``rerun_report`` raises, + naming the simulation. + """ + definition = await self._read_state( + lambda: self._reruns.read_simulation_definition(simulation_id) + ) + if definition is None: + raise EconomyRecordNotFoundError(_SIMULATION) + stored = _rerunnable_simulation(definition) + worker = await self._default_worker(stored.country) + economy_request = EconomyRequest( + country=stored.country, + policyengine_version=worker.policyengine_version, + region_code=stored.region_code, + dataset_name=stored.dataset_name, + year=stored.year, + ) + inputs = await self._read_inputs( + lambda: self._store.read_simulation_inputs( + economy_request, definition.choice + ), + request=economy_request, + source=STORED_SIMULATION, + ) + rerun = NewSimulationRerun( + simulation_id=simulation_id, + run_id=self._new_id(), + population=inputs.catalog.population(stored.year), + versions=worker.versions, + ) + calculation = simulation_run_calculation( + worker, + _rerun_scope( + stored.country, economy_request, worker, inputs.catalog, options={} + ), + self._origin(request_id), + simulation_run_id=rerun.run_id, + policy=inputs.simulation.reform, + ) + await self._create_rerun( + _SIMULATION, lambda: self._reruns.create_simulation_rerun(rerun) + ) + status = await self._start_simulation_run( + worker, + calculation, + simulation_id=simulation_id, + run_id=rerun.run_id, + observability_context=observability_context, + rerun=True, + ) + return EconomySimulationCreated( + simulation_id=simulation_id, + simulation_run_id=rerun.run_id, + status=status, + ) + + def _rerun_side( + self, stored: SideDefinition, inputs: SimulationInputs + ) -> RerunSide: + """Return a side naming its own simulation's reusable run, or a new run.""" + if inputs.reusable_run is not None: + return RerunSide( + simulation_id=stored.simulation_id, + simulation_run_id=inputs.reusable_run.simulation_run_id, + reused=True, + ) + return RerunSide( + simulation_id=stored.simulation_id, + simulation_run_id=self._new_id(), + reused=False, + ) + + async def _create_rerun( + self, record: RecordName, create: Callable[[], None] + ) -> None: + """Create the re-run's rows; refuse it while the record has a run under way.""" + try: + await self._create(record, create, rerun=True) + except RunInProgressError as error: + raise EconomyRerunConflictError( + code="run_in_progress", + detail=( + f"The {record} has a run pending or running; re-run it once " + "that run has ended." + ), + run_id=error.run_id, + ) from None + + +@dataclass(frozen=True) +class _StoredReportRows: + """The country and rows a stored report calculates on, each present.""" + + country: CountryId + year: int + region_code: str + dataset_name: str + baseline: SideDefinition + reform: SideDefinition + + +@dataclass(frozen=True) +class _StoredSimulationRows: + """The country and rows a stored simulation calculates on, each present.""" + + country: CountryId + year: int + region_code: str + dataset_name: str + + +def _calculated_country(country: str | None) -> CountryId | None: + """Return the country if this service calculates it, or None.""" + return next((known for known in _COUNTRIES if known == country), None) + + +def _refusal( + record: RecordName, kind: RecordKind, missing: tuple[tuple[str, object], ...] +) -> EconomyRerunConflictError: + """Return why a stored record is not re-run: its kind, or the first row it lacks. + + ``missing`` pairs each row's description with the value the record holds. + """ + if kind == "unsupported": + detail = f"The service does not re-run this type of {record}." + elif kind != "economy": + detail = f"The service does not re-run {_NOT_RERUN[kind]} {record}s yet." + else: + lacking = next(name for name, value in missing if value is None) + detail = f"The {record} names no {lacking}, so it cannot be re-run." + return EconomyRerunConflictError(code="rerun_not_supported", detail=detail) + + +def _rerunnable_report(definition: ReportDefinition) -> _StoredReportRows: + """Return the stored report's rows, or refuse one this service can't re-run.""" + baseline, reform = definition.baseline, definition.reform + region_code, dataset_name = definition.region_code, definition.dataset_name + country, year = _calculated_country(definition.country), definition.year + if ( + definition.kind != "economy" + or country is None + or baseline is None + or reform is None + or region_code is None + or dataset_name is None + or year is None + ): + raise _refusal( + _REPORT, + definition.kind, + ( + ("country this service calculates", country), + ("baseline simulation", baseline), + ("reform simulation", reform), + ("region", region_code), + ("dataset", dataset_name), + ("year", year), + ), + ) + return _StoredReportRows( + country=country, + year=year, + region_code=region_code, + dataset_name=dataset_name, + baseline=baseline, + reform=reform, + ) + + +def _rerunnable_simulation(definition: SimulationDefinition) -> _StoredSimulationRows: + """Return the stored simulation's country and rows, or refuse it. + + Its country is None when its catalog version's model is of no country + this service calculates. + """ + country, year = _calculated_country(definition.country), definition.year + region_code, dataset_name = definition.region_code, definition.dataset_name + if ( + definition.kind != "economy" + or country is None + or region_code is None + or dataset_name is None + or year is None + ): + raise _refusal( + _SIMULATION, + definition.kind, + ( + ("country this service calculates", country), + ("region", region_code), + ("dataset", dataset_name), + ("year", year), + ), + ) + return _StoredSimulationRows( + country=country, + year=year, + region_code=region_code, + dataset_name=dataset_name, + ) + + +def _rerun_scope( + country: CountryId, + request: EconomyRequest, + worker: EconomyWorker, + catalog: EconomyCatalogRows, + *, + options: dict[str, Any], +) -> CalculationScope: + """Return what a re-run calculates on, with the bundle's dataset. + + ``options`` are the stored report's (``reports.inputs``). Raise + EconomySubmissionRefusedError (``dataset_not_in_bundle``) when the + default worker's bundle lacks the stored dataset. + """ + try: + dataset = bundle_dataset(worker, catalog.dataset.name) + except DatasetNotInBundleError: + raise EconomySubmissionRefusedError( + code="dataset_not_in_bundle", + detail=( + f"The default v2 worker, .py {worker.version}, has no dataset " + f"{catalog.dataset.name}, which the stored record names." + ), + ) from None + return CalculationScope( + country=country, + region_code=catalog.region.code, + dataset=dataset, + year=request.year, + options=simulation_options( + include_cliffs=bool(options.get("include_cliffs", False)), + spm=options.get("spm"), + ), + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py index 9dd6e5285..b142b3b29 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 @@ -89,14 +89,14 @@ SIMULATION_PATH = "/simulations/{simulation_id}" RETRY_AFTER_SECONDS = 10 -_ERROR = {"model": EconomyError} -_RETRY_AFTER = { +ERROR_ANSWER = {"model": EconomyError} +RETRY_AFTER_HEADER = { "Retry-After": { "description": "Seconds to wait before asking again", "schema": {"type": "integer"}, } } -_FORBIDDEN = { +FORBIDDEN_ANSWER = { "description": "The caller's token is missing, invalid or lacks the permission" } _REFUSED = { @@ -107,7 +107,7 @@ "service's state, such as a catalog version missing while a release " "deploys" ), - **_ERROR, + **ERROR_ANSWER, } _UNPROCESSABLE = { "description": ( @@ -123,7 +123,7 @@ "The service failed unexpectedly, or could not confirm its rows were " "written (internal_error), possibly after creating the record" ), - **_ERROR, + **ERROR_ANSWER, } @@ -136,12 +136,12 @@ "Economy runs are not configured, or the database cannot be read; " "retry after Retry-After seconds" ), - "headers": _RETRY_AFTER, - **_ERROR, + "headers": RETRY_AFTER_HEADER, + **ERROR_ANSWER, } _READ_INTERNAL_ERROR = { "description": "The service failed unexpectedly (internal_error)", - **_ERROR, + **ERROR_ANSWER, } @@ -154,8 +154,8 @@ def _unavailable_description(record: RecordName) -> dict[str, Any]: f"(modal_call_refused): then the {record} named by {record}_id was " "created and its run failed. Retry after Retry-After seconds" ), - "headers": _RETRY_AFTER, - **_ERROR, + "headers": RETRY_AFTER_HEADER, + **ERROR_ANSWER, } @@ -191,7 +191,8 @@ async def read_simulation(self, simulation_id: UUID) -> StoredSimulation | None: ... -def _error_response(status_code: int, error: EconomyError) -> JSONResponse: +def error_response(status_code: int, error: EconomyError) -> JSONResponse: + """Answer ``error`` with its status, and ``Retry-After`` on a 503.""" headers = ( {"Retry-After": str(RETRY_AFTER_SECONDS)} if status_code == status.HTTP_503_SERVICE_UNAVAILABLE @@ -204,14 +205,16 @@ def _error_response(status_code: int, error: EconomyError) -> JSONResponse: ) -def _not_configured() -> JSONResponse: - return _error_response( +def not_configured_response() -> JSONResponse: + """Answer 503 ``not_configured``: the service has no economy runs set up.""" + return error_response( status.HTTP_503_SERVICE_UNAVAILABLE, EconomyError(code="not_configured", detail="Economy runs are not configured."), ) -_SUBMISSION_FAILURES = ( +# The submission service's own failures, each answered with its status. +SUBMISSION_FAILURES = ( EconomySubmissionRefusedError, EconomySubmissionUnprocessableError, EconomySubmissionUnavailableError, @@ -219,10 +222,10 @@ def _not_configured() -> JSONResponse: ) -def _submission_failure_response(error: Exception, record: RecordName) -> JSONResponse: +def submission_failure_response(error: Exception, record: RecordName) -> JSONResponse: """Answer a submission the service refused or could not serve.""" if isinstance(error, EconomySubmissionRefusedError): - return _error_response( + return error_response( status.HTTP_400_BAD_REQUEST, EconomyError(code=error.code, detail=error.detail), ) @@ -239,11 +242,13 @@ def _submission_failure_response(error: Exception, record: RecordName) -> JSONRe ) if isinstance(error, EconomySubmissionOutcomeUnknownError): return _outcome_unknown_response(error) - assert isinstance(error, EconomySubmissionUnavailableError) # noqa: S101 # the last of _SUBMISSION_FAILURES + assert isinstance(error, EconomySubmissionUnavailableError) # noqa: S101 # the last of SUBMISSION_FAILURES record_ids = {} if error.record_id is None else {f"{record}_id": error.record_id} - return _error_response( + return error_response( status.HTTP_503_SERVICE_UNAVAILABLE, - EconomyError(code=error.code, detail=error.detail, **record_ids), + EconomyError( + code=error.code, detail=error.detail, run_id=error.run_id, **record_ids + ), ) @@ -270,24 +275,29 @@ def _outcome_unknown_response( it nothing was created. The log is the fixed event ``economy_record_outcome_unknown``. """ - fields = {"record": error.record, "error_type": error.error_type} + fields = { + "record": error.record, + "error_type": error.error_type, + "rerun": error.rerun, + } + written = f"the {error.record}'s new run" if error.rerun else f"the {error.record}" logger.error( "economy_record_outcome_unknown", extra={**fields, "policyengine_attributes": fields}, ) - return _error_response( + return error_response( status.HTTP_500_INTERNAL_SERVER_ERROR, EconomyError( code="internal_error", detail=( - f"The service may have created the {error.record}, but could not " - "confirm it was written." + f"The service may have created {written}, but could not confirm " + "it was written." ), ), ) -def _unexpected_failure_response( +def unexpected_failure_response( route: str, error: Exception, *, max_characters: int ) -> JSONResponse: """Answer an unexpected failure with the JSON 500 the routes declare. @@ -302,7 +312,7 @@ def _unexpected_failure_response( error=error, max_characters=max_characters, ) - return _error_response( + return error_response( status.HTTP_500_INTERNAL_SERVER_ERROR, EconomyError( code="internal_error", detail="The service could not answer the request." @@ -353,8 +363,9 @@ def _simulation_state(stored: StoredSimulation) -> EconomySimulationState: ) -def _unknown(record: RecordName, record_id: UUID) -> JSONResponse: - return _error_response( +def unknown_record_response(record: RecordName, record_id: UUID) -> JSONResponse: + """Answer 404: there is no such report or simulation.""" + return error_response( status.HTTP_404_NOT_FOUND, EconomyError( code="unknown_report" if record == "report" else "unknown_simulation", @@ -408,7 +419,7 @@ def _submission_routes( responses={ 201: {"description": "The report and its first run were created"}, 400: _REFUSED, - 403: _FORBIDDEN, + 403: FORBIDDEN_ANSWER, 422: _UNPROCESSABLE, 500: _INTERNAL_ERROR, 503: _unavailable_description("report"), @@ -419,7 +430,7 @@ async def create_report( body: EconomyReportRequest, request: Request ) -> JSONResponse: if submissions is None: - return _not_configured() + return not_configured_response() observability_id = entry.start_workflow_observability(request) try: created = await submissions.submit_report( @@ -427,11 +438,11 @@ async def create_report( request_id=call_request_id_for(request), observability_context=observability_context(entry, observability_id), ) - except _SUBMISSION_FAILURES as error: - return _submission_failure_response(error, "report") + except SUBMISSION_FAILURES as error: + return submission_failure_response(error, "report") # Every other failure is still answered as the JSON 500 declared. except Exception as error: # noqa: BLE001 # answered as internal_error - return _unexpected_failure_response( + return unexpected_failure_response( REPORTS_PATH, error, max_characters=log_string_limit ) return JSONResponse( @@ -452,7 +463,7 @@ async def create_report( responses={ 201: {"description": "The simulation and its first run were created"}, 400: _REFUSED, - 403: _FORBIDDEN, + 403: FORBIDDEN_ANSWER, 422: _UNPROCESSABLE, 500: _INTERNAL_ERROR, 503: _unavailable_description("simulation"), @@ -463,7 +474,7 @@ async def create_simulation( body: EconomySimulationRequest, request: Request ) -> JSONResponse: if submissions is None: - return _not_configured() + return not_configured_response() observability_id = entry.start_workflow_observability(request) try: created = await submissions.submit_simulation( @@ -471,11 +482,11 @@ async def create_simulation( request_id=call_request_id_for(request), observability_context=observability_context(entry, observability_id), ) - except _SUBMISSION_FAILURES as error: - return _submission_failure_response(error, "simulation") + except SUBMISSION_FAILURES as error: + return submission_failure_response(error, "simulation") # Every other failure is still answered as the JSON 500 declared. except Exception as error: # noqa: BLE001 # answered as internal_error - return _unexpected_failure_response( + return unexpected_failure_response( SIMULATIONS_PATH, error, max_characters=log_string_limit ) return JSONResponse( @@ -496,20 +507,20 @@ async def _read_answer[Stored]( # noqa: PLR0913 # the read, its record and how ) -> JSONResponse: """Answer a read: the record as it stands, 404, 503 or the JSON 500.""" if submissions is None: - return _not_configured() + return not_configured_response() route = REPORT_PATH if record == "report" else SIMULATION_PATH try: stored = await read(submissions) state = None if stored is None else answer(stored) except EconomySubmissionUnavailableError as error: - return _submission_failure_response(error, record) + return submission_failure_response(error, record) # Every other failure is still answered as the JSON 500 declared. except Exception as error: # noqa: BLE001 # answered as internal_error - return _unexpected_failure_response( + return unexpected_failure_response( route, error, max_characters=log_string_limit ) if state is None: - return _unknown(record, record_id) + return unknown_record_response(record, record_id) return JSONResponse(state.model_dump(mode="json")) @@ -533,8 +544,11 @@ def _read_routes( response_model=EconomyReportState, responses={ 200: {"description": "The report as it stands"}, - 403: _FORBIDDEN, - 404: {"description": "There is no such report (unknown_report)", **_ERROR}, + 403: FORBIDDEN_ANSWER, + 404: { + "description": "There is no such report (unknown_report)", + **ERROR_ANSWER, + }, 422: _READ_UNPROCESSABLE, 500: _READ_INTERNAL_ERROR, 503: _READ_UNAVAILABLE, @@ -562,10 +576,10 @@ async def get_report(report_id: UUID) -> JSONResponse: response_model=EconomySimulationState, responses={ 200: {"description": "The simulation as it stands"}, - 403: _FORBIDDEN, + 403: FORBIDDEN_ANSWER, 404: { "description": "There is no such simulation (unknown_simulation)", - **_ERROR, + **ERROR_ANSWER, }, 422: _READ_UNPROCESSABLE, 500: _READ_INTERNAL_ERROR, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py index be55aed4b..601f12d4f 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py @@ -3,7 +3,8 @@ The submission service raises these, and the routes answer each with its status: EconomySubmissionRefusedError 400, EconomySubmissionUnprocessableError 422, EconomySubmissionUnavailableError 503 and -EconomySubmissionOutcomeUnknownError 500. ``read_refusal`` and +EconomySubmissionOutcomeUnknownError 500; a re-run also raises +EconomyRecordNotFoundError 404 and EconomyRerunConflictError 409. ``read_refusal`` and ``reference_refusal`` turn the store's errors into them, writing each detail from the request rather than from an error's text, since the Public API passes a refusal's detail on to its callers. @@ -35,6 +36,7 @@ ) from policyengine_simulation_entry.economy_api_models import ( + EconomyConflictCode, EconomyRefusalCode, EconomyUnavailableCode, EconomyUnprocessableCode, @@ -52,10 +54,28 @@ "tax_benefit_model_version_id": ("unknown_catalog_version", "catalog version"), "tax_benefit_model_id": ("unknown_catalog_version", "tax-benefit model"), } -# What a submission creates. +# What a submission creates, or a re-run calculates again. RecordName = Literal["report", "simulation"] +@dataclass(frozen=True) +class RefusalSource: + """Where refused values come from, as a refusal names and locates them. + + ``named_by`` says what names the values in a refusal's detail, and + ``location`` starts each unprocessable reason's ``loc``. + """ + + named_by: str + location: str + + +# A new record's request body, or a stored record a re-run reads. +REQUEST_BODY = RefusalSource(named_by="the request", location="body") +STORED_REPORT = RefusalSource(named_by="the report", location="report") +STORED_SIMULATION = RefusalSource(named_by="the simulation", location="simulation") + + class EconomySubmissionRefusedError(Exception): """A row the submission names, or a function it needs, is missing (400).""" @@ -98,21 +118,52 @@ class EconomySubmissionOutcomeUnknownError(Exception): The commit failed after the rows were inserted. The route answers 500, which the Public API reads as a record that may exist, never a 503, - which it reads as nothing created. + which it reads as nothing created. ``rerun`` says the rows are a stored + record's new run, not a new record. """ - def __init__(self, record: RecordName, *, error_type: str) -> None: + def __init__( + self, record: RecordName, *, error_type: str, rerun: bool = False + ) -> None: """Keep what the submission creates, and the commit error's type.""" super().__init__(f"the {record} may have been created") self.record: RecordName = record self.error_type = error_type + self.rerun = rerun + + +class EconomyRecordNotFoundError(LookupError): + """A re-run names no stored report or simulation (404).""" + + def __init__(self, record: RecordName) -> None: + """Keep which kind of record is missing.""" + super().__init__(f"no such {record}") + self.record: RecordName = record + + +class EconomyRerunConflictError(Exception): + """A stored record is not re-run now (409); nothing was written. + + ``run_id`` names the record's run pending or running, with + ``run_in_progress``. + """ + + def __init__( + self, *, code: EconomyConflictCode, detail: str, run_id: UUID | None = None + ) -> None: + """Keep the reason's code, its description and any run under way.""" + super().__init__(detail) + self.code: EconomyConflictCode = code + self.detail = detail + self.run_id = run_id class EconomySubmissionUnavailableError(Exception): """A dependency could not be read or refused the call; ask again later. - ``record_id`` names the report or simulation that was created and whose - run was failed, when one was. + ``record_id`` names the report or simulation that was created, or + re-run, and whose new run was failed, when one was; ``run_id`` names a + re-run's failed run. """ def __init__( @@ -121,41 +172,52 @@ def __init__( code: EconomyUnavailableCode, detail: str, record_id: UUID | None = None, + run_id: UUID | None = None, ) -> None: """Keep the reason's code, its description and any created record.""" super().__init__(detail) self.code: EconomyUnavailableCode = code self.detail = detail self.record_id = record_id + self.run_id = run_id def read_refusal( - error: Exception, request: EconomyRequest + error: Exception, + request: EconomyRequest, + source: RefusalSource = REQUEST_BODY, ) -> EconomySubmissionRefusedError | EconomySubmissionUnprocessableError | None: """Return the refusal a store read's error means, or None for any other error. A missing catalog version, region, dataset, policy or dynamic is refused (400); a policy or dynamic setting parameters the catalog version lacks, or both setting one, cannot be calculated (422). A report side's refusal - (ReportSideError) is located on that side's fields. Every detail is + (ReportSideError) is located on that side's fields, under ``source``: + the request's body, or the stored record a re-run reads. Every detail is written from ``request``, never from the error's text. """ side: ReportSide | None = None if isinstance(error, ReportSideError): side, error = error.side, error.error - return _missing_row_refusal(error, request, side) or _unprocessable( - error, request, side + return _missing_row_refusal(error, request, side, source) or _unprocessable( + error, request, side, source ) def _missing_row_refusal( - error: Exception, request: EconomyRequest, side: ReportSide | None + error: Exception, + request: EconomyRequest, + side: ReportSide | None, + source: RefusalSource, ) -> EconomySubmissionRefusedError | None: """Return the 400 a missing row means, or None for another error.""" catalog = ( f"{request.country}'s catalog version for .py {request.policyengine_version}" ) - named_by = "the request" if side is None else f"the request's {side}" + named_by = source.named_by if side is None else f"{source.named_by}'s {side}" + # A stored record's region and dataset are named by the record, not asked + # for; a request's are the caller's own. + stored_by = "" if source is REQUEST_BODY else f", which {source.named_by} names" refusal: tuple[EconomyRefusalCode, str] | None = None if isinstance(error, UnknownCatalogVersionError): refusal = ( @@ -164,14 +226,17 @@ def _missing_row_refusal( f"{request.policyengine_version}, the default v2 worker's version.", ) elif isinstance(error, UnknownRegionError): - refusal = ("unknown_region", f"{catalog} has no region {request.region_code}.") + refusal = ( + "unknown_region", + f"{catalog} has no region {request.region_code}{stored_by}.", + ) elif isinstance(error, UnknownDatasetError): dataset = ( f"default dataset for region {request.region_code}" if request.dataset_name is None else f"dataset {request.dataset_name}" ) - refusal = ("unknown_dataset", f"{catalog} has no {dataset}.") + refusal = ("unknown_dataset", f"{catalog} has no {dataset}{stored_by}.") elif isinstance(error, UnknownPolicyError): refusal = ( "unknown_policy", @@ -189,14 +254,19 @@ def _missing_row_refusal( def _unprocessable( - error: Exception, request: EconomyRequest, side: ReportSide | None + error: Exception, + request: EconomyRequest, + side: ReportSide | None, + source: RefusalSource, ) -> EconomySubmissionUnprocessableError | None: """Return the 422 refused values mean, or None for another error. Each reason is located at the field that names the values: the request's - own, or the report side's. + own or the stored record's, or the report side's. """ - location: tuple[str, ...] = ("body",) if side is None else ("body", side) + location: tuple[str, ...] = ( + (source.location,) if side is None else (source.location, side) + ) if isinstance(error, MissingPolicyParametersError): field_name = "policy_id" if error.source == "policy" else "dynamic_id" return EconomySubmissionUnprocessableError( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py index 3247e5ecb..eb42e41a1 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py @@ -69,10 +69,12 @@ simulation_run_calculation, ) from policyengine_simulation_entry.economy_submission_errors import ( + REQUEST_BODY, EconomySubmissionOutcomeUnknownError, EconomySubmissionRefusedError, EconomySubmissionUnavailableError, RecordName, + RefusalSource, read_refusal, reference_refusal, ) @@ -304,7 +306,12 @@ async def submit_report( ) await self._create("report", lambda: self._store.create_report(report)) status = await self._start_report_run( - worker, report, calculation, observability_context=observability_context + worker, + calculation, + report_id=report.id, + run_id=report.run_id, + new_simulation_run_ids=report.new_simulation_run_ids, + observability_context=observability_context, ) return EconomyReportCreated( report_id=report.id, @@ -359,8 +366,9 @@ async def submit_simulation( ) status = await self._start_simulation_run( worker, - simulation, calculation, + simulation_id=simulation.id, + run_id=simulation.run_id, observability_context=observability_context, ) return EconomySimulationCreated( @@ -443,9 +451,13 @@ async def _default_worker(self, country: str) -> EconomyWorker: ) from None async def _read_inputs[Inputs]( - self, read: Callable[[], Inputs], *, request: EconomyRequest + self, + read: Callable[[], Inputs], + *, + request: EconomyRequest, + source: RefusalSource = REQUEST_BODY, ) -> Inputs: - """Run ``read`` on a thread; answer each refusal as the request's error.""" + """Run ``read`` on a thread; answer each refusal as ``source``'s error.""" try: return await asyncio.to_thread(read) except SQLAlchemyError as error: @@ -455,17 +467,20 @@ async def _read_inputs[Inputs]( ) raise _database_unavailable() from None except Exception as error: - refusal = read_refusal(error, request) + refusal = read_refusal(error, request, source) if refusal is None: raise raise refusal from None - async def _create(self, record: RecordName, create: Callable[[], None]) -> None: + async def _create( + self, record: RecordName, create: Callable[[], None], *, rerun: bool = False + ) -> None: """Create the rows; answer a refusal, an unknown outcome or an outage. A commit that failed after the rows were inserted may have written them, so it is never answered as unavailable, which would tell the - caller nothing was created. + caller nothing was created. ``rerun`` says the rows are a stored + record's new run. """ try: await asyncio.to_thread(create) @@ -474,7 +489,7 @@ async def _create(self, record: RecordName, create: Callable[[], None]) -> None: except EconomyRecordOutcomeUnknownError as error: # The route logs it, with the commit error's type. raise EconomySubmissionOutcomeUnknownError( - record, error_type=type(error.__cause__).__name__ + record, error_type=type(error.__cause__).__name__, rerun=rerun ) from None except SQLAlchemyError as error: logger.warning( @@ -483,14 +498,23 @@ async def _create(self, record: RecordName, create: Callable[[], None]) -> None: ) raise _database_unavailable() from None - async def _start_report_run( + async def _start_report_run( # noqa: PLR0913 # the call, then the runs it calculates self, worker: EconomyWorker, - report: NewEconomyReport, calculation: RunCalculation, *, + report_id: UUID, + run_id: UUID, + new_simulation_run_ids: Sequence[UUID], observability_context: dict[str, Any] | None, + rerun: bool = False, ) -> ReportRunStatus: + """Start the report run's coordinator and record its call, or fail the run. + + ``new_simulation_run_ids`` are the simulation runs created with the + report run, which fail with it; ``rerun`` says the report existed + before its new run. + """ try: call_id = await asyncio.wait_for( self._caller.start_report_run( @@ -506,31 +530,39 @@ async def _start_report_run( except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike failed = await self._fail_unconfirmed( "report", - report.id, + report_id, error, lambda message: self._store.fail_unstarted_report( - report.run_id, report.new_simulation_run_ids, message + run_id, new_simulation_run_ids, message ), ) if failed: - refused = _call_refused("report", report.id) + refused = _call_refused( + "report", report_id, rerun_id=run_id if rerun else None + ) raise refused from None return ReportRunStatus.PENDING status = await self._record_call( "report", - report.id, - lambda: self._store.record_report_call(report.run_id, call_id), + report_id, + lambda: self._store.record_report_call(run_id, call_id), ) return ReportRunStatus.PENDING if status is None else ReportRunStatus(status) - async def _start_simulation_run( + async def _start_simulation_run( # noqa: PLR0913 # the call, then the run it calculates self, worker: EconomyWorker, - simulation: NewEconomySimulation, calculation: RunCalculation, *, + simulation_id: UUID, + run_id: UUID, observability_context: dict[str, Any] | None, + rerun: bool = False, ) -> SimulationStatus: + """Start the simulation run's function and record its call, or fail the run. + + ``rerun`` says the simulation existed before its new run. + """ try: call_id = await asyncio.wait_for( self._caller.start_simulation_run( @@ -545,20 +577,20 @@ async def _start_simulation_run( except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike failed = await self._fail_unconfirmed( "simulation", - simulation.id, + simulation_id, error, - lambda message: self._store.fail_unstarted_simulation( - simulation.run_id, message - ), + lambda message: self._store.fail_unstarted_simulation(run_id, message), ) if failed: - refused = _call_refused("simulation", simulation.id) + refused = _call_refused( + "simulation", simulation_id, rerun_id=run_id if rerun else None + ) raise refused from None return SimulationStatus.PENDING status = await self._record_call( "simulation", - simulation.id, - lambda: self._store.record_simulation_call(simulation.run_id, call_id), + simulation_id, + lambda: self._store.record_simulation_call(run_id, call_id), ) return SimulationStatus.PENDING if status is None else SimulationStatus(status) @@ -630,8 +662,22 @@ async def _record_call( def _call_refused( - record: RecordName, record_id: UUID + record: RecordName, record_id: UUID, *, rerun_id: UUID | None = None ) -> EconomySubmissionUnavailableError: + """Return the 503 naming a record whose new run's call Modal did not accept. + + ``rerun_id`` names a stored record's new run, which was failed. + """ + if rerun_id is not None: + return EconomySubmissionUnavailableError( + code="modal_call_refused", + detail=( + f"Modal did not accept the call of the {record}'s new run; " + "that run was failed." + ), + record_id=record_id, + run_id=rerun_id, + ) return EconomySubmissionUnavailableError( code="modal_call_refused", detail=( diff --git a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py index 42c33d5d0..412d31e99 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py @@ -210,6 +210,7 @@ def test_a_report_kind_the_service_does_not_rerun_yet_is_refused(kind, named) -> @pytest.mark.parametrize( ("missing", "named"), [ + ({"country": "fr"}, "country"), ({"baseline": None}, "baseline simulation"), ({"reform": None}, "reform simulation"), ({"region_code": None}, "region"), @@ -390,6 +391,7 @@ def test_a_simulation_is_rerun_as_a_new_run_of_it() -> None: ({"kind": "household"}, "household"), ({"kind": "unsupported"}, "this type of"), ({"country": None}, "country"), + ({"country": "fr"}, "country"), ({"region_code": None}, "region"), ({"dataset_name": None}, "dataset"), ({"year": None}, "year"), From 668c1b2f67cc7d9db7e4a831eefa32a5df560000 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 00:59:09 +0400 Subject: [PATCH 05/18] Document re-running economy reports and simulations, and add a changelog fragment Co-Authored-By: Claude Opus 5.5 --- changelog.d/c1-17-reruns.added.md | 1 + .../stage-12-modal-worker-foundation.md | 10 ++++++ .../README.md | 28 ++++++++++++++++ .../policyengine-simulation-entry/README.md | 32 +++++++++++++++++++ 4 files changed, 71 insertions(+) create mode 100644 changelog.d/c1-17-reruns.added.md diff --git a/changelog.d/c1-17-reruns.added.md b/changelog.d/c1-17-reruns.added.md new file mode 100644 index 000000000..9519707f1 --- /dev/null +++ b/changelog.d/c1-17-reruns.added.md @@ -0,0 +1 @@ +The Simulation Entrypoint re-runs a stored economy report or simulation on the current .py version through `POST /reports/{report_id}/runs` and `POST /simulations/{simulation_id}/runs`, as a new run started by run ID, keeping every earlier run. A report's re-run reuses a run of each side's own simulation that succeeded on the same versions unless `force` is set; a record with a run pending or running is refused with 409 `run_in_progress`. diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index ccb02b6f8..d4a0abc9a 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -331,6 +331,16 @@ as for a report comparing it with itself, so a report whose plan needs more columns, for example a reform that turns on labour supply responses, can't reuse it and fails with `reused_simulation_output_incompatible`. +**A re-run.** The service re-runs a stored report or simulation with the same +two calls, by run ID. It creates the record's new run (a report run with +trigger `manual`) and, for a report, a new `simulation_runs` row of each side's +own simulation, or names that simulation's run that already succeeded on the +same versions, unless the re-run is forced. The new runs record the catalog +version, dataset and region of the current .py version's catalog; the +`simulations` and `reports` rows are unchanged. The coordinator and workers +claim and calculate them as they would a new record's: nothing on the Modal +side tells a re-run apart. + **Each worker** claims its simulation run under its own Modal call ID: a pending or running run that records no call or its own. It records its output file, or its failure, only while it holds the run running. It returns the output's diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 4fea4fd37..03fbd6a12 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -125,6 +125,34 @@ commit fails after its rows were inserted raises `EconomyRecordOutcomeUnknownError`: PostgreSQL may have applied the commit, so the rows may exist. +## Re-running a report or simulation + +`economy_rerun_records.py` holds a re-run's reads and writes. A re-run +calculates a stored record again on the current versions, as a new run, and +changes no `reports` or `simulations` row, which the service's role may not +update: a report keeps its simulations, and each new run records the catalog +version, dataset and region it calculates on. + +- `read_report_definition` and `read_simulation_definition` read what the + record calculates: its kind (`economy`, `budget_window`, `household` or + `unsupported`), country, year, its region's code and dataset's name, and + each simulation's policy and dynamic. +- `find_reusable_run_of` is `find_reusable_run` limited to one simulation's + runs, with the catalog version, dataset and region compared on the run. +- `insert_report_rerun` creates a report's new run, trigger `manual`, and a + new run of each side's simulation that reuses none; + `insert_simulation_rerun` creates a simulation's new run. Each refuses with + `RunInProgressError`, naming the run, while the record has a run `pending` + or `running`, and first takes `lock_reruns_of`, a transaction-level + advisory lock on the record's ID, so concurrent re-runs of one record + wait for each other. + +`EconomyRunStore.read_report_rerun_inputs` reads a re-run's catalog rows by +the stored region code and dataset name, and each side's values and +reusable run of its own simulation, never giving both sides the same run. +`create_report_rerun` and `create_simulation_rerun` run the inserts in a +transaction of their own. + ## Reading a policy `read_policy_reform` (`policy_reform.py`) reads a stored policy as the reform diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index bc32cc1f6..91c76fa30 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -44,8 +44,10 @@ Auth0 permission: | `GET /household-calculations/{household_job_id}` | `simulation:read` | | `POST /reports` | `simulation:submit` | | `GET /reports/{report_id}` | `simulation:read` | +| `POST /reports/{report_id}/runs` | `simulation:submit` | | `POST /simulations` | `simulation:submit` | | `GET /simulations/{simulation_id}` | `simulation:read` | +| `POST /simulations/{simulation_id}/runs` | `simulation:submit` | `GET /health`, `GET /ready`, `GET /versions`, `GET /versions/{kind}` and `POST /ping` stay public. The entrypoint answers `GET /health` and `POST /ping` @@ -156,3 +158,33 @@ record and its newest run (created last, ties to the larger ID): its status, versions, error message and times, and a report run's simulation runs. The Public API reads results from the tables itself. +## Re-running a report or simulation + +`POST /reports/{report_id}/runs` calculates a stored economy report again, +as a new run on the default worker's versions, and keeps every earlier run. +Its body is `{"force": false}` by default. The report keeps its simulations: +each side gets a new run of its own simulation, or, without `force`, reuses +that simulation's newest run on the same versions, catalog version, dataset, +region and year that a new report's side could reuse; `force` calculates +both again. The new report run's trigger is `manual`. The service reads the +report's region code and dataset name, finds the rows with that code and +name in the catalog version of the default worker's .py version, and reads +each simulation's stored policy and dynamic on it. +`POST /simulations/{simulation_id}/runs`, with no body, creates and starts +a stored economy simulation's new run, which is always calculated. Each +answers 201 with the same fields as the create, naming the new run. + +A re-run is refused before anything is written: 404 (`unknown_report`, +`unknown_simulation`) for a record that doesn't exist; 409 +`run_in_progress`, naming the run in `run_id`, while the record has a run +pending or running; 409 `rerun_not_supported` for a budget-window or +household record, which this service does not re-run yet, or one that +names no region, dataset, year or simulation; and the create's 400 and 422 +codes when the stored rows can't be calculated on the current catalog +version, worded for the stored record (a 422 reason's `loc` starts with +`report` or `simulation`). The check for a run under way and the inserts +run under a transaction-level advisory lock on the record's ID, so two +re-runs of one record never both create a run. When Modal refuses the new +run's call, the run is failed and the 503 `modal_call_refused` names the +record and the run (`run_id`). + From ca33a901890b4d6cc9b98b42f63c0cdc2af26e71 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:19:02 +0400 Subject: [PATCH 06/18] Test review round 1's findings: rows deleted during a re-run, and the definition reads (failing) The entry tests fail: a record deleted after a re-run's reads should be answered 404, and a row deleted meanwhile should be refused naming the report. The persistence tests cover existing behaviour and pass at once: the reuse search's catalog-version condition on the run row, and the household, filtered, other-model and other-type definition reads. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_reruns.py | 137 +++++++++++++++++- .../tests/test_economy_reruns.py | 44 ++++++ 2 files changed, 180 insertions(+), 1 deletion(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py index 068b20245..7de45c479 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py @@ -26,6 +26,7 @@ LATER_VERSIONS, OPTIONS, OUTPUT_BUCKET_PREFIX, + VERSIONS, EconomyCatalog, ReportRunOutcome, RunShape, @@ -55,23 +56,27 @@ SimulationChoice, SimulationDefinition, SimulationStatus, + SimulationType, UnknownDatasetError, UnknownRegionError, + find_reusable_run, find_reusable_run_of, insert_report_rerun, insert_simulation_rerun, lock_reruns_of, read_report_definition, read_simulation_definition, + simulations, ) from public_api_rows import ( StoredValue, + insert_catalog_version, insert_dataset, insert_household, insert_policy, insert_region, ) -from sqlalchemy import text +from sqlalchemy import insert, text from sqlalchemy.exc import OperationalError if TYPE_CHECKING: @@ -720,3 +725,133 @@ def test_a_rerun_and_its_runs_are_created_through_the_store( assert definition.id == stored.report_id assert simulation is not None assert simulation.id == simulation_rerun.simulation_id + + +def test_a_run_on_another_catalog_version_is_not_reused_by_a_new_report( + public_api_database_url, +) -> None: + """Only the run's own catalog version differs from the request's. + + It is also the newest run, so it would be found without that condition. + """ + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy) + connection.execute( + text( + "UPDATE public.simulation_runs " + "SET tax_benefit_model_version_id = :later, " + "completed_at = completed_at + interval '1 day' WHERE id = :id" + ), + { + "later": later.catalog.version_id, + "id": stored.baseline.simulation_run_id, + }, + ) + act_as_the_service(connection) + + found = find_reusable_run( + connection, + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + inputs=OPTIONS, + output_uri_prefix=OUTPUT_BUCKET_PREFIX, + ) + + # The reform's run, the same simulation definition on the right catalog + # version, is the only one left to reuse. + assert found == ReusableRun( + simulation_id=stored.reform.simulation_id, + simulation_run_id=stored.reform.simulation_run_id, + ) + + +def test_a_household_or_filtered_simulation_and_another_model_are_read_as_such( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + household_id = insert_household( + connection, country="us", default_year=2026, document={"people": {}} + ) + household_simulation_id = uuid4() + filtered_simulation_id = uuid4() + connection.execute( + insert(simulations).values( + id=household_simulation_id, + simulation_type=SimulationType.HOUSEHOLD, + household_id=household_id, + tax_benefit_model_version_id=economy.catalog.version_id, + year=2026, + ) + ) + connection.execute( + insert(simulations).values( + id=filtered_simulation_id, + simulation_type=SimulationType.ECONOMY, + dataset_id=economy.dataset_id, + region_id=economy.region_id, + tax_benefit_model_version_id=economy.catalog.version_id, + year=2026, + filter_field="state_code", + filter_value="CA", + ) + ) + other_model_id = uuid4() + connection.execute( + text( + "INSERT INTO public.tax_benefit_models (id, name) VALUES (:id, :name)" + ), + {"id": other_model_id, "name": "another-model"}, + ) + other = _insert_catalog_of(connection, other_model_id) + other_simulation = insert_simulation_run(connection, other.population()) + act_as_the_service(connection) + + household = read_simulation_definition(connection, household_simulation_id) + filtered = read_simulation_definition(connection, filtered_simulation_id) + of_another_model = read_simulation_definition( + connection, other_simulation.simulation_id + ) + + assert household is not None + assert (household.kind, household.country) == ("household", "us") + assert (household.region_code, household.dataset_name) == (None, None) + assert filtered is not None + assert filtered.kind == "unsupported" + assert of_another_model is not None + assert (of_another_model.kind, of_another_model.country) == ("economy", None) + + +def _insert_catalog_of(connection: Connection, model_id: UUID) -> EconomyCatalog: + catalog = insert_catalog_version( + connection, + country="xx", + model_id=model_id, + policyengine_version=LATER_POLICYENGINE_VERSION, + ) + dataset_id = insert_dataset(connection, catalog, "another_dataset") + region_id = insert_region( + connection, catalog, "everywhere", default_dataset_id=dataset_id + ) + return EconomyCatalog(catalog=catalog, dataset_id=dataset_id, region_id=region_id) + + +def test_a_report_of_another_type_is_read_as_unsupported( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + stored = _stored_report(connection, economy) + connection.execute( + text("UPDATE public.reports SET type = 'custom_total' WHERE id = :id"), + {"id": stored.report_id}, + ) + act_as_the_service(connection) + + definition = read_report_definition(connection, stored.report_id) + + assert definition is not None + assert definition.kind == "unsupported" diff --git a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py index 412d31e99..19c62698c 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py @@ -43,6 +43,7 @@ from policyengine_simulation_persistence import ( CatalogDataset, EconomyRecordOutcomeUnknownError, + EconomyRecordReferenceError, MissingPolicyParametersError, NewReportRerun, NewSimulationRerun, @@ -437,3 +438,46 @@ def test_a_simulations_stored_policy_is_located_on_the_simulation() -> None: ((reason,),) = (refused.value.reasons,) assert reason.loc == ("simulation", "dynamic_id", "gov.old.elasticity") + + +@pytest.mark.parametrize( + ("record", "deleted"), + [ + ("report", EconomyRecordReferenceError("report_runs", "report_id")), + ("simulation", EconomyRecordReferenceError("simulation_runs", "simulation_id")), + ], +) +def test_a_record_deleted_after_its_reads_is_not_found(record, deleted) -> None: + service = rerun_service(FakeRerunStore(rerun_failure=deleted)) + rerun = _rerun_report if record == "report" else _rerun_simulation + + with pytest.raises(EconomyRecordNotFoundError) as missing: + rerun(service) + + assert missing.value.record == record + + +@pytest.mark.parametrize( + ("deleted", "code", "detail"), + [ + ( + EconomyRecordReferenceError("simulation_runs", "dataset_id"), + "unknown_dataset", + "a dataset the report names no longer exists", + ), + ( + EconomyRecordReferenceError("simulation_runs", "simulation_id"), + "missing_reference", + "a simulation the report names no longer exists", + ), + ], +) +def test_a_row_deleted_after_a_reruns_reads_is_refused_naming_the_report( + deleted, code, detail +) -> None: + store = FakeRerunStore(rerun_failure=deleted) + + with pytest.raises(EconomySubmissionRefusedError) as refused: + _rerun_report(rerun_service(store)) + + assert (refused.value.code, refused.value.detail) == (code, detail) From 93f5d65712da653a0d4c01cf5598ac65d7b5641c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:20:49 +0400 Subject: [PATCH 07/18] Fix review round 1's findings: name a re-run's refused rows as the stored record's A record deleted after a re-run read it is answered 404, and any other row the database refuses meanwhile is named as the report's or simulation's. The economy submission service is generic over its store, so the re-run service needs no constructor of its own; the re-run records share the economy report type; the docs say how the caller-row codes apply to a re-run, that the report body is required, and how the in-progress check differs between a report and a simulation. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 10 ++- .../economy_rerun_records.py | 6 +- .../policyengine-simulation-entry/README.md | 13 ++- .../economy_api_models.py | 11 ++- .../economy_rerun_routes.py | 5 +- .../economy_reruns.py | 83 ++++++++++--------- .../economy_submission_errors.py | 14 ++-- .../economy_submissions.py | 18 ++-- 8 files changed, 95 insertions(+), 65 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 03fbd6a12..4c1e37591 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -137,13 +137,17 @@ version, dataset and region it calculates on. record calculates: its kind (`economy`, `budget_window`, `household` or `unsupported`), country, year, its region's code and dataset's name, and each simulation's policy and dynamic. -- `find_reusable_run_of` is `find_reusable_run` limited to one simulation's - runs, with the catalog version, dataset and region compared on the run. +- `find_reusable_run_of` searches one simulation's runs, by its ID, under + `find_reusable_run`'s conditions on a run (versions, the run's own catalog + version, dataset, region and year, the output file, a succeeded report run + with the same options, no refused reuse); it doesn't compare the + simulation row, which keeps the catalog rows it was created on. - `insert_report_rerun` creates a report's new run, trigger `manual`, and a new run of each side's simulation that reuses none; `insert_simulation_rerun` creates a simulation's new run. Each refuses with `RunInProgressError`, naming the run, while the record has a run `pending` - or `running`, and first takes `lock_reruns_of`, a transaction-level + or `running` (for a report, a report run; for a simulation, any of its + runs), and first takes `lock_reruns_of`, a transaction-level advisory lock on the record's ID, so concurrent re-runs of one record wait for each other. diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py index 2e99aba0b..ca2487380 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -32,6 +32,7 @@ from sqlalchemy import func, insert, literal, select from policyengine_simulation_persistence.economy_run_records import ( + ECONOMY_REPORT_TYPE, SimulationChoice, insert_simulation_run, references_checked, @@ -69,8 +70,7 @@ # report or simulation, or anything else, such as a simulation of a filtered # population. type RecordKind = Literal["economy", "budget_window", "household", "unsupported"] -# A report's type, as the Public API and the service write it. -_ECONOMY_TYPE: Final = "economy" +# A household report's type, as the household reports' service writes it. _HOUSEHOLD_TYPE: Final = "household" # Both run tables' statuses use these values for a run under way. _UNDER_WAY: Final = (SimulationStatus.PENDING.value, SimulationStatus.RUNNING.value) @@ -236,7 +236,7 @@ def _report_kind( ) -> RecordKind: if household_id is not None or report_type == _HOUSEHOLD_TYPE: return "household" - if report_type != _ECONOMY_TYPE: + if report_type != ECONOMY_REPORT_TYPE: return "unsupported" return "economy" if budget_window_size is None else "budget_window" diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 91c76fa30..a79c01017 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -162,7 +162,8 @@ Public API reads results from the tables itself. `POST /reports/{report_id}/runs` calculates a stored economy report again, as a new run on the default worker's versions, and keeps every earlier run. -Its body is `{"force": false}` by default. The report keeps its simulations: +Its body is required: `{}` or `{"force": false}` re-runs the report, and +`{"force": true}` forces a recalculation. The report keeps its simulations: each side gets a new run of its own simulation, or, without `force`, reuses that simulation's newest run on the same versions, catalog version, dataset, region and year that a new report's side could reuse; `force` calculates @@ -179,12 +180,18 @@ A re-run is refused before anything is written: 404 (`unknown_report`, `run_in_progress`, naming the run in `run_id`, while the record has a run pending or running; 409 `rerun_not_supported` for a budget-window or household record, which this service does not re-run yet, or one that -names no region, dataset, year or simulation; and the create's 400 and 422 +names no region, dataset, year, simulation or country the service +calculates; and the create's 400 and 422 codes when the stored rows can't be calculated on the current catalog version, worded for the stored record (a 422 reason's `loc` starts with `report` or `simulation`). The check for a run under way and the inserts run under a transaction-level advisory lock on the record's ID, so two -re-runs of one record never both create a run. When Modal refuses the new +re-runs of one record never both create a run. A report's re-run checks only +the report's own runs: a side whose simulation has a run under way for +another report, or for the simulation's own re-run, still gets a new run, +calculated separately. A simulation's re-run checks every run of the +simulation, including those reports created. A record deleted after the +re-run read it answers 404. When Modal refuses the new run's call, the run is failed and the 503 `modal_call_refused` names the record and the run (`run_id`). diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py index 8ecf0cc1b..d4e8a1e7a 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -12,9 +12,10 @@ - 201 with the created record's ID and its first run's ID, or, for a re-run, the record's ID and its new run's ID; - 400 with ``code`` and ``detail``: one of ``CallerRowRefusalCode`` when - the region, dataset or a dynamic the request names does not exist, or - another code for the service's own state, such as a catalog version - missing while a release deploys. Nothing is written; + the region, dataset or a dynamic the request names (or, for a re-run, the + stored record names) does not exist, or another code for the service's own + state, such as a catalog version missing while a release deploys. Nothing + is written; - 404 ``unknown_report`` or ``unknown_simulation`` when a re-run names no stored record; - 409 with ``code`` and ``detail`` when a stored record is not re-run now: @@ -53,7 +54,9 @@ ReportDescription = Annotated[str, Field(max_length=10_000)] # Why a request is refused (400), naming a row the caller chose that doesn't -# exist: the Public API answers these as its caller's error. +# exist: the Public API answers these as its caller's error. On a re-run they +# name the stored record's own region, dataset or dynamic, which the current +# catalog version lacks or which was deleted. CallerRowRefusalCode = Literal["unknown_region", "unknown_dataset", "unknown_dynamic"] # Why a request is refused (400) for the service's own state, or a row the # Public API checks before asking. diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py index 98f8f50cf..0dc5b1d20 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -106,8 +106,9 @@ def _responses(record: RecordName) -> dict[int | str, dict[str, Any]]: f"The stored {record}'s region, dataset, policy or dynamic is " "missing from the catalog version of the default worker's .py " "version, or a function it needs is missing; nothing was " - "written. Each code is the stored record's or the service's " - "state, never the caller's" + "written. unknown_region, unknown_dataset and unknown_dynamic " + f"name the stored {record}'s own rows; any other code is the " + "service's state" ), **ERROR_ANSWER, }, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py index 5d695407d..8cb244787 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -32,13 +32,12 @@ from __future__ import annotations from dataclasses import dataclass -from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, Final, Protocol, get_args -from uuid import UUID, uuid4 from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ReportRunIds from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, EconomyRequest, NewReportRerun, NewSimulationRerun, @@ -67,13 +66,13 @@ EconomySubmissionRefusedError, ) from policyengine_simulation_entry.economy_submissions import ( - CALL_START_TIMEOUT_SECONDS, EconomyRuns, EconomySubmissionService, ) if TYPE_CHECKING: from collections.abc import Callable + from uuid import UUID from policyengine_simulation_persistence import ( EconomyCatalogRows, @@ -88,12 +87,12 @@ from policyengine_simulation_entry.economy_api_models import ( ReportRerunRequest, ) - from policyengine_simulation_entry.economy_submission_errors import RecordName - from policyengine_simulation_entry.economy_submissions import ServiceSettings + from policyengine_simulation_entry.economy_submission_errors import ( + RecordName, + RefusalSource, + ) from policyengine_simulation_entry.economy_workers import ( - EconomyRunCaller, EconomyWorker, - EconomyWorkerSource, ) # How a refusal names each kind of stored record this service does not re-run @@ -106,6 +105,11 @@ _COUNTRIES: Final[tuple[CountryId, ...]] = get_args(CountryId) _REPORT: Final[RecordName] = "report" _SIMULATION: Final[RecordName] = "simulation" +# The table and column of a re-run's new run whose key names the record itself. +_OWN_KEYS: Final[dict[RecordName, tuple[str, str]]] = { + "report": ("report_runs", "report_id"), + "simulation": ("simulation_runs", "simulation_id"), +} class EconomyRerunRuns(EconomyRuns, Protocol): @@ -141,32 +145,9 @@ def create_simulation_rerun(self, rerun: NewSimulationRerun) -> None: ... -class EconomyRerunService(EconomySubmissionService): +class EconomyRerunService(EconomySubmissionService[EconomyRerunRuns]): """Create and re-run economy reports and simulations, and start their calls.""" - def __init__( # noqa: PLR0913 # each seam tests replace - self, - *, - store: EconomyRerunRuns, - workers: EconomyWorkerSource, - caller: EconomyRunCaller, - settings: ServiceSettings, - call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, - new_id: Callable[[], UUID] = uuid4, - clock: Callable[[], datetime] = lambda: datetime.now(UTC), - ) -> None: - """Use these seams, as ``EconomySubmissionService`` does.""" - super().__init__( - store=store, - workers=workers, - caller=caller, - settings=settings, - call_start_timeout_seconds=call_start_timeout_seconds, - new_id=new_id, - clock=clock, - ) - self._reruns = store - async def rerun_report( self, report_id: UUID, @@ -183,7 +164,7 @@ async def rerun_report( run as well as the report. """ definition = await self._read_state( - lambda: self._reruns.read_report_definition(report_id) + lambda: self._store.read_report_definition(report_id) ) if definition is None: raise EconomyRecordNotFoundError(_REPORT) @@ -206,7 +187,7 @@ async def rerun_report( ) ) inputs = await self._read_inputs( - lambda: self._reruns.read_report_rerun_inputs( + lambda: self._store.read_report_rerun_inputs( economy_request, baseline=stored.baseline, reform=stored.reform, @@ -242,7 +223,9 @@ async def rerun_report( reform_policy=inputs.reform.reform, ) await self._create_rerun( - _REPORT, lambda: self._reruns.create_report_rerun(rerun) + _REPORT, + lambda: self._store.create_report_rerun(rerun), + source=STORED_REPORT, ) status = await self._start_report_run( worker, @@ -279,7 +262,7 @@ async def rerun_simulation( naming the simulation. """ definition = await self._read_state( - lambda: self._reruns.read_simulation_definition(simulation_id) + lambda: self._store.read_simulation_definition(simulation_id) ) if definition is None: raise EconomyRecordNotFoundError(_SIMULATION) @@ -315,7 +298,9 @@ async def rerun_simulation( policy=inputs.simulation.reform, ) await self._create_rerun( - _SIMULATION, lambda: self._reruns.create_simulation_rerun(rerun) + _SIMULATION, + lambda: self._store.create_simulation_rerun(rerun), + source=STORED_SIMULATION, ) status = await self._start_simulation_run( worker, @@ -348,11 +333,31 @@ def _rerun_side( ) async def _create_rerun( - self, record: RecordName, create: Callable[[], None] + self, + record: RecordName, + create: Callable[[], None], + *, + source: RefusalSource, ) -> None: - """Create the re-run's rows; refuse it while the record has a run under way.""" + """Create the re-run's rows; refuse it while the record has a run under way. + + A record deleted after its reads is answered as not found: the + database refuses the new run's key to it (``report_runs.report_id`` + or ``simulation_runs.simulation_id``). Any other row it refuses is + named as ``source``'s. + """ + own_key = _OWN_KEYS[record] + + def create_unless_deleted() -> None: + try: + create() + except EconomyRecordReferenceError as error: + if (error.table_name, error.column_name) == own_key: + raise EconomyRecordNotFoundError(record) from None + raise + try: - await self._create(record, create, rerun=True) + await self._create(record, create_unless_deleted, rerun=True, source=source) except RunInProgressError as error: raise EconomyRerunConflictError( code="run_in_progress", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py index 601f12d4f..a842876f7 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py @@ -44,9 +44,12 @@ # Each column whose foreign key can refuse a submission's rows: the code the # refusal answers with and the row it names. A region, dataset or dynamic is -# the caller's own choice; any other row is one the Public API checked first, -# or one the service chose, so its refusal is the service's state. +# the caller's own choice, or on a re-run the stored record's; any other row +# is one the Public API checked first, or one the service chose, so its +# refusal is the service's state. A simulation is refused only on a re-run, +# whose new runs name the report's stored simulations. REFERENCE_REFUSALS: dict[str, tuple[EconomyRefusalCode, str]] = { + "simulation_id": ("missing_reference", "simulation"), "region_id": ("unknown_region", "region"), "dataset_id": ("unknown_dataset", "dataset"), "dynamic_id": ("unknown_dynamic", "dynamic"), @@ -299,15 +302,16 @@ def _unprocessable( def reference_refusal( - error: EconomyRecordReferenceError, + error: EconomyRecordReferenceError, source: RefusalSource = REQUEST_BODY ) -> EconomySubmissionRefusedError: """Return the refusal of rows the database refused for naming a missing row. - It names the kind of row, never its ID, which the caller sent. + It names the kind of row, never its ID, which the caller sent, and what + named it: the request, or the stored record a re-run reads. """ code, row = REFERENCE_REFUSALS.get( error.column_name or "", ("missing_reference", "row") ) return EconomySubmissionRefusedError( - code=code, detail=f"a {row} the request names no longer exists" + code=code, detail=f"a {row} {source.named_by} names no longer exists" ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py index eb42e41a1..bb61f7610 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py @@ -212,17 +212,18 @@ def _choice(policy: PolicyChoice) -> SimulationChoice: return SimulationChoice(policy_id=policy.policy_id, dynamic_id=policy.dynamic_id) -class EconomySubmissionService: +class EconomySubmissionService[Store: EconomyRuns = EconomyRuns]: """Create economy reports and simulations, and start their calls. The database and the manifest are read on threads, so the event loop - never waits on them. + never waits on them. ``Store`` is the store's type, which a subclass + reading more of the database widens. """ def __init__( # noqa: PLR0913 # each seam tests replace self, *, - store: EconomyRuns, + store: Store, workers: EconomyWorkerSource, caller: EconomyRunCaller, settings: ServiceSettings, @@ -473,19 +474,24 @@ async def _read_inputs[Inputs]( raise refusal from None async def _create( - self, record: RecordName, create: Callable[[], None], *, rerun: bool = False + self, + record: RecordName, + create: Callable[[], None], + *, + rerun: bool = False, + source: RefusalSource = REQUEST_BODY, ) -> None: """Create the rows; answer a refusal, an unknown outcome or an outage. A commit that failed after the rows were inserted may have written them, so it is never answered as unavailable, which would tell the caller nothing was created. ``rerun`` says the rows are a stored - record's new run. + record's new run, and ``source`` what names a refused row. """ try: await asyncio.to_thread(create) except EconomyRecordReferenceError as error: - raise reference_refusal(error) from None + raise reference_refusal(error, source) from None except EconomyRecordOutcomeUnknownError as error: # The route logs it, with the commit error's type. raise EconomySubmissionOutcomeUnknownError( From b2fcda52b9bd8029a0eb19a66bb3a711f2006047 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:35:21 +0400 Subject: [PATCH 08/18] Test review round 2's finding: a re-run of a deleted record raises its own error (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_reruns.py | 45 +++++++++++++++++++ .../tests/test_economy_reruns.py | 13 ++---- 2 files changed, 49 insertions(+), 9 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py index 7de45c479..16c1f9a48 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py @@ -39,6 +39,7 @@ insert_simulation_run, ) from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, EconomyRequest, EconomyRunStore, MissingPolicyParametersError, @@ -48,6 +49,7 @@ ReportRunStatus, ReportRunTrigger, ReportSideError, + RerunRecordDeletedError, RerunSide, ReusableRun, RunInProgressError, @@ -855,3 +857,46 @@ def test_a_report_of_another_type_is_read_as_unsupported( assert definition is not None assert definition.kind == "unsupported" + + +def test_a_rerun_of_a_deleted_report_or_simulation_says_the_record_is_gone( + public_api_database_url, +) -> None: + """The record's own key refuses the new run: the record was deleted.""" + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_report(connection, economy, run_status=ReportRunStatus.FAILED) + act_as_the_service(connection) + + with connection.begin_nested(), pytest.raises(RerunRecordDeletedError): + insert_report_rerun( + connection, replace(_rerun_of(stored, later), report_id=uuid4()) + ) + with connection.begin_nested(), pytest.raises(RerunRecordDeletedError): + insert_simulation_rerun( + connection, + NewSimulationRerun( + simulation_id=uuid4(), + run_id=uuid4(), + population=later.population(), + versions=LATER_VERSIONS, + ), + ) + # A report whose side simulation is gone names that simulation instead. + gone_side = replace( + _rerun_of(stored, later), + reform=RerunSide( + simulation_id=uuid4(), simulation_run_id=uuid4(), reused=False + ), + ) + with ( + connection.begin_nested(), + pytest.raises(EconomyRecordReferenceError) as refused, + ): + insert_report_rerun(connection, gone_side) + + assert (refused.value.table_name, refused.value.column_name) == ( + "simulation_runs", + "simulation_id", + ) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py index 19c62698c..7978f3169 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py @@ -49,6 +49,7 @@ NewSimulationRerun, ReportRunStatus, ReportSideError, + RerunRecordDeletedError, RerunSide, ReusableRun, RunInProgressError, @@ -440,15 +441,9 @@ def test_a_simulations_stored_policy_is_located_on_the_simulation() -> None: assert reason.loc == ("simulation", "dynamic_id", "gov.old.elasticity") -@pytest.mark.parametrize( - ("record", "deleted"), - [ - ("report", EconomyRecordReferenceError("report_runs", "report_id")), - ("simulation", EconomyRecordReferenceError("simulation_runs", "simulation_id")), - ], -) -def test_a_record_deleted_after_its_reads_is_not_found(record, deleted) -> None: - service = rerun_service(FakeRerunStore(rerun_failure=deleted)) +@pytest.mark.parametrize("record", ["report", "simulation"]) +def test_a_record_deleted_after_its_reads_is_not_found(record) -> None: + service = rerun_service(FakeRerunStore(rerun_failure=RerunRecordDeletedError())) rerun = _rerun_report if record == "report" else _rerun_simulation with pytest.raises(EconomyRecordNotFoundError) as missing: From 77aea9a020dc71eaf21cd09c699ed689b5a598cb Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:36:49 +0400 Subject: [PATCH 09/18] Fix review round 2's findings: the persistence library tells a deleted re-run record apart insert_report_rerun and insert_simulation_rerun raise RerunRecordDeletedError when the database refuses the new run's key to the record itself, so the re-run service answers 404 without naming tables or columns. The docs say a report deleted with its simulations answers 400 missing_reference, and the re-run routes' 400 names every code that refers to the stored record's rows. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 4 +- .../__init__.py | 2 + .../economy_rerun_records.py | 43 ++++++++++++++++--- .../tests/test_economy_reruns.py | 6 +-- .../policyengine-simulation-entry/README.md | 6 ++- .../economy_rerun_routes.py | 6 +-- .../economy_reruns.py | 28 +++--------- .../economy_submissions.py | 4 +- 8 files changed, 60 insertions(+), 39 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 4c1e37591..745202be9 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -149,7 +149,9 @@ version, dataset and region it calculates on. or `running` (for a report, a report run; for a simulation, any of its runs), and first takes `lock_reruns_of`, a transaction-level advisory lock on the record's ID, so concurrent re-runs of one record - wait for each other. + wait for each other. When the database refuses the new run's key to the + record itself, the record was deleted after it was read: + `RerunRecordDeletedError`. `EconomyRunStore.read_report_rerun_inputs` reads a re-run's catalog rows by the stored region code and dataset name, and each side's values and 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 d8ea4ef03..9fb44a2c6 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -11,6 +11,7 @@ NewSimulationRerun, RecordKind, ReportDefinition, + RerunRecordDeletedError, RerunSide, RunInProgressError, SideDefinition, @@ -154,6 +155,7 @@ "ReportSide", "ReportSideError", "ReportSimulation", + "RerunRecordDeletedError", "RerunSide", "ReusableRun", "RunInProgressError", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py index ca2487380..ead89ef05 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -26,6 +26,7 @@ from __future__ import annotations +from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal @@ -33,6 +34,7 @@ from policyengine_simulation_persistence.economy_run_records import ( ECONOMY_REPORT_TYPE, + EconomyRecordReferenceError, SimulationChoice, insert_simulation_run, references_checked, @@ -56,9 +58,10 @@ ) if TYPE_CHECKING: + from collections.abc import Iterator from uuid import UUID - from sqlalchemy import ColumnElement, Connection, RowMapping, Table + from sqlalchemy import Column, ColumnElement, Connection, RowMapping, Table from policyengine_simulation_persistence.economy_run_records import ( EconomyPopulation, @@ -167,6 +170,14 @@ class NewSimulationRerun: versions: RunVersions +class RerunRecordDeletedError(LookupError): + """The re-run's report or simulation was deleted after the re-run read it. + + The database refused the new run's key to the record itself; nothing was + created. + """ + + class RunInProgressError(Exception): """The record has a run that is pending or running, so no re-run is created. @@ -347,9 +358,11 @@ def insert_report_rerun(connection: Connection, rerun: NewReportRerun) -> None: """Create the report's new run and its sides' new simulation runs, pending. Take the report's re-run lock first. Raise RunInProgressError, creating - nothing, when the report has a run pending or running, and - EconomyRecordReferenceError when the database refuses a row the new - runs name; that refusal ends the caller's transaction, or its savepoint. + nothing, when the report has a run pending or running; + RerunRecordDeletedError when the report no longer exists; and + EconomyRecordReferenceError when the database refuses another row the + new runs name, such as a side's simulation. A refusal ends the caller's + transaction, or its savepoint. """ lock_reruns_of(connection, rerun.report_id) under_way = _run_under_way( @@ -357,7 +370,7 @@ def insert_report_rerun(connection: Connection, rerun: NewReportRerun) -> None: ) if under_way is not None: raise RunInProgressError(under_way) - with references_checked(): + with _record_checked(report_runs.c.report_id), references_checked(): for side in (rerun.baseline, rerun.reform): if not side.reused: insert_simulation_run( @@ -389,7 +402,8 @@ def insert_simulation_rerun(connection: Connection, rerun: NewSimulationRerun) - Take the simulation's re-run lock first. Raise RunInProgressError, creating nothing, when the simulation has a run pending or running, and - EconomyRecordReferenceError as ``insert_report_rerun`` does. + RerunRecordDeletedError or EconomyRecordReferenceError as + ``insert_report_rerun`` does. """ lock_reruns_of(connection, rerun.simulation_id) under_way = _run_under_way( @@ -399,7 +413,7 @@ def insert_simulation_rerun(connection: Connection, rerun: NewSimulationRerun) - ) if under_way is not None: raise RunInProgressError(under_way) - with references_checked(): + with _record_checked(simulation_runs.c.simulation_id), references_checked(): insert_simulation_run( connection, run_id=rerun.run_id, @@ -407,3 +421,18 @@ def insert_simulation_rerun(connection: Connection, rerun: NewSimulationRerun) - population=rerun.population, versions=rerun.versions, ) + + +@contextmanager +def _record_checked(own_key: Column[Any]) -> Iterator[None]: + """Turn a refusal of the new run's key to the record into RerunRecordDeletedError. + + ``own_key`` is that key's column; a refusal of any other row propagates + unchanged. + """ + try: + yield + except EconomyRecordReferenceError as error: + if (error.table_name, error.column_name) == (own_key.table.name, own_key.key): + raise RerunRecordDeletedError from None + raise diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py index 16c1f9a48..53babd5aa 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_reruns.py @@ -869,11 +869,11 @@ def test_a_rerun_of_a_deleted_report_or_simulation_says_the_record_is_gone( stored = _stored_report(connection, economy, run_status=ReportRunStatus.FAILED) act_as_the_service(connection) - with connection.begin_nested(), pytest.raises(RerunRecordDeletedError): + with pytest.raises(RerunRecordDeletedError), connection.begin_nested(): insert_report_rerun( connection, replace(_rerun_of(stored, later), report_id=uuid4()) ) - with connection.begin_nested(), pytest.raises(RerunRecordDeletedError): + with pytest.raises(RerunRecordDeletedError), connection.begin_nested(): insert_simulation_rerun( connection, NewSimulationRerun( @@ -891,8 +891,8 @@ def test_a_rerun_of_a_deleted_report_or_simulation_says_the_record_is_gone( ), ) with ( - connection.begin_nested(), pytest.raises(EconomyRecordReferenceError) as refused, + connection.begin_nested(), ): insert_report_rerun(connection, gone_side) diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index a79c01017..9a8c5ffe3 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -190,8 +190,10 @@ re-runs of one record never both create a run. A report's re-run checks only the report's own runs: a side whose simulation has a run under way for another report, or for the simulation's own re-run, still gets a new run, calculated separately. A simulation's re-run checks every run of the -simulation, including those reports created. A record deleted after the -re-run read it answers 404. When Modal refuses the new +simulation, including those reports created. A report or simulation +deleted after the re-run read it answers 404; a report deleted with its +simulations answers 400 `missing_reference`, naming a simulation. When Modal +refuses the new run's call, the run is failed and the 503 `modal_call_refused` names the record and the run (`run_id`). diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py index 0dc5b1d20..3d0dda954 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -106,9 +106,9 @@ def _responses(record: RecordName) -> dict[int | str, dict[str, Any]]: f"The stored {record}'s region, dataset, policy or dynamic is " "missing from the catalog version of the default worker's .py " "version, or a function it needs is missing; nothing was " - "written. unknown_region, unknown_dataset and unknown_dynamic " - f"name the stored {record}'s own rows; any other code is the " - "service's state" + "written. unknown_region, unknown_dataset, unknown_dynamic, " + f"unknown_policy and missing_reference name the stored {record}'s " + "own rows; any other code is the service's state" ), **ERROR_ANSWER, }, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py index 8cb244787..5c7e35099 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -37,10 +37,10 @@ from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ReportRunIds from policyengine_simulation_persistence import ( - EconomyRecordReferenceError, EconomyRequest, NewReportRerun, NewSimulationRerun, + RerunRecordDeletedError, RerunSide, RunInProgressError, RunReuse, @@ -105,11 +105,6 @@ _COUNTRIES: Final[tuple[CountryId, ...]] = get_args(CountryId) _REPORT: Final[RecordName] = "report" _SIMULATION: Final[RecordName] = "simulation" -# The table and column of a re-run's new run whose key names the record itself. -_OWN_KEYS: Final[dict[RecordName, tuple[str, str]]] = { - "report": ("report_runs", "report_id"), - "simulation": ("simulation_runs", "simulation_id"), -} class EconomyRerunRuns(EconomyRuns, Protocol): @@ -341,23 +336,14 @@ async def _create_rerun( ) -> None: """Create the re-run's rows; refuse it while the record has a run under way. - A record deleted after its reads is answered as not found: the - database refuses the new run's key to it (``report_runs.report_id`` - or ``simulation_runs.simulation_id``). Any other row it refuses is - named as ``source``'s. + A record deleted after its reads is answered as not found. Any other + row the database refuses is named as ``source``'s; for a report, + that includes a side's simulation deleted with it. """ - own_key = _OWN_KEYS[record] - - def create_unless_deleted() -> None: - try: - create() - except EconomyRecordReferenceError as error: - if (error.table_name, error.column_name) == own_key: - raise EconomyRecordNotFoundError(record) from None - raise - try: - await self._create(record, create_unless_deleted, rerun=True, source=source) + await self._create(record, create, rerun=True, source=source) + except RerunRecordDeletedError: + raise EconomyRecordNotFoundError(record) from None except RunInProgressError as error: raise EconomyRerunConflictError( code="run_in_progress", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py index bb61f7610..644b1ed62 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py @@ -216,8 +216,8 @@ class EconomySubmissionService[Store: EconomyRuns = EconomyRuns]: """Create economy reports and simulations, and start their calls. The database and the manifest are read on threads, so the event loop - never waits on them. ``Store`` is the store's type, which a subclass - reading more of the database widens. + never waits on them. ``Store`` is the store's type: a subclass that + reads more of the database names a protocol with more methods. """ def __init__( # noqa: PLR0913 # each seam tests replace From ec1e7e837e4e6893a0158f629fc7e96f58f3d7bd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 01:45:29 +0400 Subject: [PATCH 10/18] Fix review round 3's wording: which codes name a re-run's stored rows, and when a deleted report answers 404 Co-Authored-By: Claude Opus 5.5 --- .../economy_rerun_records.py | 10 +++++----- projects/policyengine-simulation-entry/README.md | 5 +++-- .../economy_api_models.py | 3 ++- .../economy_rerun_routes.py | 16 +++++++++++++--- 4 files changed, 23 insertions(+), 11 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py index ead89ef05..73e6cae18 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -358,11 +358,11 @@ def insert_report_rerun(connection: Connection, rerun: NewReportRerun) -> None: """Create the report's new run and its sides' new simulation runs, pending. Take the report's re-run lock first. Raise RunInProgressError, creating - nothing, when the report has a run pending or running; - RerunRecordDeletedError when the report no longer exists; and - EconomyRecordReferenceError when the database refuses another row the - new runs name, such as a side's simulation. A refusal ends the caller's - transaction, or its savepoint. + nothing, when the report has a run pending or running. Raise + RerunRecordDeletedError when the database refuses the new run's key to + the report, which no longer exists, and EconomyRecordReferenceError when + it refuses another row the new runs name, such as a side's simulation; + either refusal ends the caller's transaction, or its savepoint. """ lock_reruns_of(connection, rerun.report_id) under_way = _run_under_way( diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 9a8c5ffe3..2f0150739 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -191,8 +191,9 @@ the report's own runs: a side whose simulation has a run under way for another report, or for the simulation's own re-run, still gets a new run, calculated separately. A simulation's re-run checks every run of the simulation, including those reports created. A report or simulation -deleted after the re-run read it answers 404; a report deleted with its -simulations answers 400 `missing_reference`, naming a simulation. When Modal +deleted after the re-run read it answers 404. A report deleted with its +simulations answers 400 `missing_reference`, naming a simulation, when a +side gets a new run, and 404 when both sides reuse a run. When Modal refuses the new run's call, the run is failed and the 503 `modal_call_refused` names the record and the run (`run_id`). diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py index d4e8a1e7a..fc2a4ae17 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -59,7 +59,8 @@ # catalog version lacks or which was deleted. CallerRowRefusalCode = Literal["unknown_region", "unknown_dataset", "unknown_dynamic"] # Why a request is refused (400) for the service's own state, or a row the -# Public API checks before asking. +# Public API checks before asking. On a report's re-run, ``unknown_policy`` and +# ``missing_reference`` name the stored report's own policy or simulation. ServiceStateRefusalCode = Literal[ # The country has no catalog version for the default worker's .py version. "unknown_catalog_version", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py index 3d0dda954..25d250379 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -97,6 +97,17 @@ class EconomyService(EconomySubmissions, EconomyReruns, Protocol): """Create, read and re-run economy reports and simulations.""" +# The 400 codes that name a stored record's own rows on its re-run: a +# report's re-run also names its stored simulations (missing_reference). +_STORED_ROW_CODES: dict[RecordName, str] = { + "report": ( + "unknown_region, unknown_dataset, unknown_dynamic, unknown_policy and " + "missing_reference" + ), + "simulation": "unknown_region, unknown_dataset, unknown_dynamic and unknown_policy", +} + + def _responses(record: RecordName) -> dict[int | str, dict[str, Any]]: """Return every answer a re-run of ``record`` declares.""" return { @@ -106,9 +117,8 @@ def _responses(record: RecordName) -> dict[int | str, dict[str, Any]]: f"The stored {record}'s region, dataset, policy or dynamic is " "missing from the catalog version of the default worker's .py " "version, or a function it needs is missing; nothing was " - "written. unknown_region, unknown_dataset, unknown_dynamic, " - f"unknown_policy and missing_reference name the stored {record}'s " - "own rows; any other code is the service's state" + f"written. {_STORED_ROW_CODES[record]} name the stored " + f"{record}'s own rows; any other code is the service's state" ), **ERROR_ANSWER, }, From 206c32f2e966ff24b610bfcc7ef0880d57b8e38e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:56:58 +0400 Subject: [PATCH 11/18] Test re-running a stored budget-window report on a later catalog version (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_window_reruns.py | 484 ++++++++++++++++++ 1 file changed, 484 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py new file mode 100644 index 000000000..a8dfb5e1f --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py @@ -0,0 +1,484 @@ +"""Re-running a stored budget-window report on a later catalog version. + +A budget window's re-run is a new budget-window run of the report, with a +report run per year. Each year's sides are the simulations that year's +report run of the report's newest window run names, which may belong to +another report when that year reused a run. A side may reuse a run of its +own simulation that succeeded on the current versions. These run against a +database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the simulation service's role, in +a transaction rolled back afterwards. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + LATER_POLICYENGINE_VERSION, + LATER_VERSIONS, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + VERSIONS, + EconomyCatalog, + RunShape, + insert_economy_catalog, + insert_later_catalog, + insert_report_run_of, + insert_run_of, +) +from policyengine_simulation_persistence import ( + BudgetWindowYear, + EconomyRequest, + EconomyRunStore, + NewBudgetWindowReport, + NewReportRerun, + NewWindowRerun, + ReportRunStatus, + ReportRunTrigger, + ReportSimulation, + RerunSide, + ReusableRun, + RunInProgressError, + RunReuse, + SimulationStatus, + WindowRerunYear, + WindowYearDefinition, + insert_budget_window_report, + insert_report_rerun, + insert_window_rerun, + read_report_definition, +) +from sqlalchemy import text + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +WINDOW_YEARS = (2026, 2027, 2028) +LATER_REQUEST = EconomyRequest( + country="us", + policyengine_version=LATER_POLICYENGINE_VERSION, + region_code="us", + dataset_name="populace_us_2024", + year=2026, +) +LATER_REUSE = RunReuse( + versions=LATER_VERSIONS, inputs=OPTIONS, output_uri_prefix=OUTPUT_BUCKET_PREFIX +) + + +@dataclass(frozen=True) +class StoredWindow: + """A stored budget-window report, its window run and each year's rows.""" + + report_id: UUID + window_run_id: UUID + years: tuple[BudgetWindowYear, ...] + + +def _new_side() -> ReportSimulation: + return ReportSimulation( + simulation_id=uuid4(), + simulation_run_id=uuid4(), + choice=CURRENT_LAW, + reused=False, + ) + + +def _stored_window( + connection: Connection, + economy: EconomyCatalog, + *, + window_status: ReportRunStatus = ReportRunStatus.SUCCEEDED, +) -> StoredWindow: + """Insert a budget-window report as the service would, then end its runs.""" + years = tuple( + BudgetWindowYear( + year=year, + report_run_id=uuid4(), + baseline=_new_side(), + reform=_new_side(), + ) + for year in WINDOW_YEARS + ) + report = NewBudgetWindowReport( + id=uuid4(), + window_run_id=uuid4(), + country="us", + label="Ten-year cost", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(year=WINDOW_YEARS[0]), + versions=VERSIONS, + inputs=OPTIONS, + years=years, + ) + insert_budget_window_report(connection, report) + # Earlier than anything else the test's transaction creates, whose rows + # all get the transaction's time. + for table, condition in ( + ("budget_window_runs", "id = :window_run_id"), + ("report_runs", "budget_window_run_id = :window_run_id"), + ): + connection.execute( + text( + f"UPDATE public.{table} " # noqa: S608 # fixed table names + f"SET created_at = created_at - interval '1 hour' WHERE {condition}" + ), + {"window_run_id": report.window_run_id}, + ) + if window_status is not ReportRunStatus.PENDING: + ended = window_status in {ReportRunStatus.SUCCEEDED, ReportRunStatus.FAILED} + for table, condition in ( + ("budget_window_runs", "id = :window_run_id"), + ("report_runs", "budget_window_run_id = :window_run_id"), + ( + "simulation_runs", + "id IN (SELECT baseline_simulation_run_id FROM public.report_runs " + "WHERE budget_window_run_id = :window_run_id UNION " + "SELECT reform_simulation_run_id FROM public.report_runs " + "WHERE budget_window_run_id = :window_run_id)", + ), + ): + connection.execute( + text( + f"UPDATE public.{table} SET status = :status, " # noqa: S608 # fixed table names + "completed_at = CASE WHEN :ended THEN now() END " + f"WHERE {condition}" + ), + { + "status": window_status.value, + "ended": ended, + "window_run_id": report.window_run_id, + }, + ) + return StoredWindow( + report_id=report.id, window_run_id=report.window_run_id, years=years + ) + + +def _store_on(connection: Connection) -> EconomyRunStore: + """Return a store whose transactions are savepoints of ``connection``.""" + + @contextmanager + def savepoint() -> Iterator[Connection]: + with connection.begin_nested(): + yield connection + + return EconomyRunStore(savepoint) + + +def _row(connection: Connection, table_name: str, row_id: UUID) -> dict: + row = connection.execute( + text(f"SELECT * FROM public.{table_name} WHERE id = :id"), # noqa: S608 # fixed table names + {"id": row_id}, + ).mappings() + return dict(row.one()) + + +def _window_rerun_of( + stored: StoredWindow, later: EconomyCatalog, *, recalculate: bool = False +) -> NewWindowRerun: + return NewWindowRerun( + report_id=stored.report_id, + window_run_id=uuid4(), + population=later.population(year=WINDOW_YEARS[0]), + versions=LATER_VERSIONS, + recalculate=recalculate, + years=tuple( + WindowRerunYear( + year=window_year.year, + report_run_id=uuid4(), + baseline=RerunSide( + simulation_id=window_year.baseline.simulation_id, + simulation_run_id=uuid4(), + reused=False, + ), + reform=RerunSide( + simulation_id=window_year.reform.simulation_id, + simulation_run_id=uuid4(), + reused=False, + ), + ) + for window_year in stored.years + ), + ) + + +def test_a_budget_windows_definition_names_each_years_simulations( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + stored = _stored_window(connection, economy) + act_as_the_service(connection) + + definition = read_report_definition(connection, stored.report_id) + + assert definition is not None + assert (definition.kind, definition.year, definition.budget_window_size) == ( + "budget_window", + 2026, + 3, + ) + assert definition.window_years == tuple( + WindowYearDefinition( + year=window_year.year, + baseline_simulation_id=window_year.baseline.simulation_id, + reform_simulation_id=window_year.reform.simulation_id, + ) + for window_year in stored.years + ) + + +def test_a_definitions_years_come_from_the_newest_window_run( + public_api_database_url, +) -> None: + """A year that reused another report's run names that report's simulation.""" + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window(connection, economy) + other_2027 = _stored_window(connection, economy).years[1] + rerun = _window_rerun_of(stored, later) + years = list(rerun.years) + years[1] = replace( + years[1], + baseline=RerunSide( + simulation_id=other_2027.baseline.simulation_id, + simulation_run_id=other_2027.baseline.simulation_run_id, + reused=True, + ), + ) + act_as_the_service(connection) + insert_window_rerun(connection, replace(rerun, years=tuple(years))) + + definition = read_report_definition(connection, stored.report_id) + + assert definition is not None + assert [year.baseline_simulation_id for year in definition.window_years] == [ + stored.years[0].baseline.simulation_id, + other_2027.baseline.simulation_id, + stored.years[2].baseline.simulation_id, + ] + + +def test_a_window_rerun_creates_a_window_run_and_each_years_runs( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window(connection, economy) + reused = insert_run_of( + connection, + stored.years[1].baseline.simulation_id, + later.population(year=2027), + RunShape(versions=LATER_VERSIONS), + ) + rerun = _window_rerun_of(stored, later, recalculate=True) + years = list(rerun.years) + years[1] = replace( + years[1], + baseline=RerunSide( + simulation_id=reused.simulation_id, + simulation_run_id=reused.simulation_run_id, + reused=True, + ), + ) + rerun = replace(rerun, years=tuple(years)) + act_as_the_service(connection) + + insert_window_rerun(connection, rerun) + window_run = _row(connection, "budget_window_runs", rerun.window_run_id) + year_runs = [ + _row(connection, "report_runs", year.report_run_id) for year in rerun.years + ] + reform_run_2028 = _row( + connection, "simulation_runs", rerun.years[2].reform.simulation_run_id + ) + earlier_window = _row(connection, "budget_window_runs", stored.window_run_id) + + assert { + key: window_run[key] + for key in ( + "report_id", + "tax_benefit_model_version_id", + "trigger", + "recalculate", + "status", + ) + } == { + "report_id": stored.report_id, + "tax_benefit_model_version_id": later.catalog.version_id, + "trigger": ReportRunTrigger.MANUAL, + "recalculate": True, + "status": ReportRunStatus.PENDING, + } + assert [ + ( + run["year"], + run["budget_window_run_id"], + run["trigger"], + run["status"], + run["policyengine_version"], + ) + for run in year_runs + ] == [ + ( + year, + rerun.window_run_id, + ReportRunTrigger.MANUAL, + ReportRunStatus.PENDING, + LATER_POLICYENGINE_VERSION, + ) + for year in WINDOW_YEARS + ] + assert year_runs[1]["baseline_simulation_run_id"] == reused.simulation_run_id + assert ( + reform_run_2028["simulation_id"], + reform_run_2028["year"], + reform_run_2028["tax_benefit_model_version_id"], + reform_run_2028["status"], + ) == ( + stored.years[2].reform.simulation_id, + 2028, + later.catalog.version_id, + SimulationStatus.PENDING, + ) + assert earlier_window["status"] == ReportRunStatus.SUCCEEDED + + +@pytest.mark.parametrize("status", [ReportRunStatus.PENDING, ReportRunStatus.RUNNING]) +def test_a_window_rerun_is_refused_while_a_window_run_is_under_way( + public_api_database_url, status: ReportRunStatus +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window(connection, economy, window_status=status) + act_as_the_service(connection) + + with pytest.raises(RunInProgressError) as window_refused: + insert_window_rerun(connection, _window_rerun_of(stored, later)) + # A single-year re-run of the report checks the window runs too. + with pytest.raises(RunInProgressError) as report_refused: + insert_report_rerun( + connection, + NewReportRerun( + report_id=stored.report_id, + run_id=uuid4(), + population=later.population(), + versions=LATER_VERSIONS, + baseline=RerunSide( + simulation_id=stored.years[0].baseline.simulation_id, + simulation_run_id=uuid4(), + reused=False, + ), + reform=RerunSide( + simulation_id=stored.years[0].reform.simulation_id, + simulation_run_id=uuid4(), + reused=False, + ), + ), + ) + + assert window_refused.value.run_id == stored.window_run_id + assert report_refused.value.run_id == stored.window_run_id + + +def test_a_window_reruns_inputs_name_each_years_own_reusable_runs( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window(connection, economy) + later_shape = RunShape(versions=LATER_VERSIONS) + reused_2027 = insert_run_of( + connection, + stored.years[1].baseline.simulation_id, + later.population(year=2027), + later_shape, + ) + failed_reform_2027 = insert_run_of( + connection, + stored.years[1].reform.simulation_id, + later.population(year=2027), + replace(later_shape, status=SimulationStatus.FAILED), + ) + insert_report_run_of( + connection, stored.report_id, (reused_2027, failed_reform_2027) + ) + act_as_the_service(connection) + years = _definition_years(stored) + + inputs = _store_on(connection).read_window_rerun_inputs( + LATER_REQUEST, + baseline=CURRENT_LAW, + reform=CURRENT_LAW, + years=years, + reuse=LATER_REUSE, + ) + forced = _store_on(connection).read_window_rerun_inputs( + LATER_REQUEST, + baseline=CURRENT_LAW, + reform=CURRENT_LAW, + years=years, + reuse=None, + ) + + assert inputs.catalog.catalog_version_id == later.catalog.version_id + assert inputs.reusable_runs == { + 2026: (None, None), + 2027: ( + ReusableRun( + simulation_id=stored.years[1].baseline.simulation_id, + simulation_run_id=reused_2027.simulation_run_id, + ), + None, + ), + 2028: (None, None), + } + assert forced.reusable_runs == dict.fromkeys(WINDOW_YEARS, (None, None)) + + +def _definition_years(stored: StoredWindow) -> tuple[WindowYearDefinition, ...]: + """Return the stored window's years as its definition names them.""" + return tuple( + WindowYearDefinition( + year=window_year.year, + baseline_simulation_id=window_year.baseline.simulation_id, + reform_simulation_id=window_year.reform.simulation_id, + ) + for window_year in stored.years + ) + + +def test_a_window_rerun_is_created_through_the_store(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window(connection, economy) + rerun = _window_rerun_of(stored, later) + act_as_the_service(connection) + store = _store_on(connection) + + store.create_window_rerun(rerun) + with pytest.raises(RunInProgressError) as refused: + store.create_window_rerun(_window_rerun_of(stored, later)) + + assert _row(connection, "budget_window_runs", rerun.window_run_id)[ + "trigger" + ] == (ReportRunTrigger.MANUAL) + assert refused.value.run_id == rerun.window_run_id From a61f7994061683d85f98e4e34ea172d8d03b667e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 02:59:02 +0400 Subject: [PATCH 12/18] Test re-running a budget-window report through POST /reports/{id}/runs (failing) Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_generated_client.py | 31 ++++ .../tests/economy_rerun_fakes.py | 78 ++++++++- .../tests/test_economy_rerun_routes.py | 37 ++++- .../tests/test_economy_reruns.py | 157 +++++++++++++++++- .../tests/test_openapi.py | 14 +- 5 files changed, 301 insertions(+), 16 deletions(-) diff --git a/projects/policyengine-apis-integ/tests/test_economy_generated_client.py b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py index 636bbf6af..0f58c439f 100644 --- a/projects/policyengine-apis-integ/tests/test_economy_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py @@ -18,6 +18,7 @@ rerun_simulation, ) from policyengine_api_simulation_client.models import ( + BudgetWindowReportCreated, EconomyError, EconomyReportCreated, EconomyReportRequest, @@ -273,3 +274,33 @@ def test_a_refused_rerun_is_read_with_its_code_and_any_run(status_code, body): assert isinstance(response.parsed, EconomyError) assert response.parsed.to_dict() == body + + +def test_a_budget_window_rerun_is_read_with_its_new_window_run(): + window = { + "report_id": REPORT_ID, + "report_run_id": RUN_ID, + "budget_window_run_id": SIDE_IDS[0], + "status": "pending", + "years": [ + { + "year": 2026, + "report_run_id": RUN_ID, + "baseline_simulation_id": SIDE_IDS[1], + "baseline_simulation_run_id": SIDE_IDS[2], + "baseline_reused": False, + "reform_simulation_id": SIDE_IDS[3], + "reform_simulation_run_id": SIDE_IDS[0], + "reform_reused": True, + } + ], + } + + response = rerun_report.sync_detailed( + UUID(REPORT_ID), + client=_client_answering(201, window), + body=ReportRerunRequest(), + ) + + assert isinstance(response.parsed, BudgetWindowReportCreated) + assert response.parsed.to_dict() == window diff --git a/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py b/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py index e496bca2b..64a4956f2 100644 --- a/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py +++ b/projects/policyengine-simulation-entry/tests/economy_rerun_fakes.py @@ -17,13 +17,13 @@ POLICY_ID, REFORM, FakeCaller, - FakeStore, FakeWorkers, numbered_ids, ) from policyengine_simulation_entry.economy_reruns import EconomyRerunService from policyengine_simulation_entry.economy_submissions import ServiceSettings from policyengine_simulation_persistence import ( + BudgetWindowInputs, EconomyReportInputs, EconomySimulationInputs, ReportDefinition, @@ -31,13 +31,18 @@ SimulationChoice, SimulationDefinition, SimulationInputs, + WindowYearDefinition, ) +from test_budget_window_submissions import FakeAdvancer, FakeWindowStore if TYPE_CHECKING: + from collections.abc import Sequence + from policyengine_simulation_persistence import ( EconomyRequest, NewReportRerun, NewSimulationRerun, + NewWindowRerun, RunReuse, ) @@ -63,6 +68,38 @@ baseline=SideDefinition(simulation_id=BASELINE_SIMULATION_ID, choice=CURRENT_LAW), reform=SideDefinition(simulation_id=REFORM_SIMULATION_ID, choice=REFORM_CHOICE), ) +# A stored budget window of three years, each with its own simulations. +WINDOW_SIMULATION_IDS = { + year: ( + UUID(f"00000000-0000-4000-8000-00000000{year}"), + UUID(f"00000000-0000-4000-8000-00000001{year}"), + ) + for year in (2026, 2027, 2028) +} +STORED_WINDOW = ReportDefinition( + id=REPORT_ID, + country="us", + kind="budget_window", + year=2026, + region_code="us", + dataset_name="populace_us_2024", + inputs=STORED_OPTIONS, + baseline=SideDefinition( + simulation_id=WINDOW_SIMULATION_IDS[2026][0], choice=CURRENT_LAW + ), + reform=SideDefinition( + simulation_id=WINDOW_SIMULATION_IDS[2026][1], choice=REFORM_CHOICE + ), + budget_window_size=3, + window_years=tuple( + WindowYearDefinition( + year=year, + baseline_simulation_id=baseline_id, + reform_simulation_id=reform_id, + ) + for year, (baseline_id, reform_id) in WINDOW_SIMULATION_IDS.items() + ), +) STORED_SIMULATION = SimulationDefinition( id=SIMULATION_ID, country="us", @@ -75,9 +112,12 @@ @dataclass -class FakeRerunStore(FakeStore): +class FakeRerunStore(FakeWindowStore): """Records the re-run reads and writes, and answers as told. + A budget window's re-run is read with ``reusable_by_year``, as a new + window is, and recorded in ``windows``, where the fake advancer finds it. + ``inputs_failure`` is raised by a re-run's read of its inputs, and ``rerun_failure`` by its create instead of recording it. """ @@ -138,6 +178,35 @@ def read_simulation_inputs( raise self.inputs_failure return super().read_simulation_inputs(request, choice) + def read_window_rerun_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + years: Sequence[WindowYearDefinition], + reuse: RunReuse | None, + ) -> BudgetWindowInputs: + self.rerun_reads.append((request, (baseline, reform, tuple(years), reuse))) + if self.inputs_failure is not None: + raise self.inputs_failure + return BudgetWindowInputs( + catalog=self.catalog, + baseline=SimulationInputs(choice=baseline, reform={}), + reform=SimulationInputs(choice=reform, reform=REFORM), + reusable_runs={ + year.year: (None, None) + if reuse is None + else self.reusable_by_year.get(year.year, (None, None)) + for year in years + }, + ) + + def create_window_rerun(self, rerun: NewWindowRerun) -> None: + if self.rerun_failure is not None: + raise self.rerun_failure + self.windows.append(rerun) + def create_report_rerun(self, rerun: NewReportRerun) -> None: if self.rerun_failure is not None: raise self.rerun_failure @@ -150,7 +219,7 @@ def create_simulation_rerun(self, rerun: NewSimulationRerun) -> None: class UnusedAdvancer: - """A window advancer a single-year re-run never asks to start a year.""" + """A window advancer a test of a single-year re-run never expects to be used.""" def advance_window(self, window_run_id: UUID, **_: Any) -> NoReturn: msg = f"a single-year re-run advanced window {window_run_id}" @@ -163,9 +232,10 @@ def rerun_service( workers: FakeWorkers | None = None, caller: FakeCaller | None = None, output_bucket: str | None = OUTPUT_BUCKET, + advancer: FakeAdvancer | None = None, ) -> EconomyRerunService: return EconomyRerunService( - windows=UnusedAdvancer(), + windows=advancer or UnusedAdvancer(), store=store, workers=workers or FakeWorkers(), caller=caller or FakeCaller(), diff --git a/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py index a47be302c..eb5011242 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_rerun_routes.py @@ -11,6 +11,10 @@ from fastapi.testclient import TestClient from policyengine_simulation_entry import app as app_module from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportCreated, + BudgetWindowYearCreated, +) from policyengine_simulation_entry.economy_api_models import ( EconomyReportCreated, EconomySimulationCreated, @@ -59,6 +63,7 @@ class FakeEconomyReruns: """Records each re-run, and answers as told.""" failure: Exception | None = None + report_answer: EconomyReportCreated | BudgetWindowReportCreated = RERUN_REPORT report_reruns: list[tuple[UUID, ReportRerunRequest, str, Any]] = field( default_factory=list ) @@ -71,13 +76,13 @@ async def rerun_report( *, request_id: str, observability_context: dict[str, Any] | None, - ) -> EconomyReportCreated: + ) -> EconomyReportCreated | BudgetWindowReportCreated: self.report_reruns.append( (report_id, request, request_id, observability_context) ) if self.failure is not None: _raise(self.failure) - return RERUN_REPORT + return self.report_answer async def rerun_simulation( self, @@ -139,6 +144,34 @@ def test_a_report_rerun_answers_201_with_its_new_run(client_for) -> None: assert context["observability_id"] == OBSERVABILITY_ID +def test_a_budget_window_rerun_answers_201_with_its_new_window_run(client_for) -> None: + window = BudgetWindowReportCreated( + report_id=RECORD_ID, + report_run_id=RUN_ID, + budget_window_run_id=SIDE_IDS[0], + status=ReportRunStatus.PENDING, + years=[ + BudgetWindowYearCreated( + year=2026, + report_run_id=RUN_ID, + baseline_simulation_id=SIDE_IDS[1], + baseline_simulation_run_id=SIDE_IDS[2], + baseline_reused=False, + reform_simulation_id=SIDE_IDS[3], + reform_simulation_run_id=SIDE_IDS[0], + reform_reused=False, + ) + ], + ) + reruns = FakeEconomyReruns(report_answer=window) + + with client_for(reruns) as client: + result = client.post(REPORT_RUNS, json={}) + + assert result.status_code == 201 + assert result.json() == window.model_dump(mode="json") + + def test_a_report_rerun_without_force_reuses_runs(client_for) -> None: reruns = FakeEconomyReruns() diff --git a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py index 7978f3169..7a6c95b7c 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py @@ -16,6 +16,8 @@ STORED_OPTIONS, STORED_REPORT, STORED_SIMULATION, + STORED_WINDOW, + WINDOW_SIMULATION_IDS, FakeRerunStore, rerun_service, ) @@ -59,12 +61,17 @@ UnknownRegionError, ) from sqlalchemy.exc import OperationalError +from test_budget_window_submissions import FakeAdvancer DATABASE_ERROR = OperationalError("SELECT 1", {}, Exception("server closed")) VERSIONS = RunVersions( country_package_version="2.2.1", policyengine_version=POLICYENGINE_VERSION ) OBSERVED = {"observability_id": "observed"} +REUSED_2027 = ReusableRun( + simulation_id=UUID("00000000-0000-4000-8000-000000002027"), + simulation_run_id=UUID("00000000-0000-4000-8000-0000000000b2"), +) REUSED_BASELINE = ReusableRun( simulation_id=BASELINE_SIMULATION_ID, simulation_run_id=UUID("00000000-0000-4000-8000-0000000000b1"), @@ -191,11 +198,7 @@ def test_an_unknown_report_or_simulation_is_not_found() -> None: @pytest.mark.parametrize( ("kind", "named"), - [ - ("budget_window", "budget-window"), - ("household", "household"), - ("unsupported", "this type of"), - ], + [("household", "household"), ("unsupported", "this type of")], ) def test_a_report_kind_the_service_does_not_rerun_yet_is_refused(kind, named) -> None: store = FakeRerunStore(report_definition=replace(STORED_REPORT, kind=kind)) @@ -476,3 +479,147 @@ def test_a_row_deleted_after_a_reruns_reads_is_refused_naming_the_report( _rerun_report(rerun_service(store)) assert (refused.value.code, refused.value.detail) == (code, detail) + + +def test_a_budget_window_is_rerun_as_a_new_window_run_and_its_first_years_started(): + store = FakeRerunStore( + report_definition=STORED_WINDOW, reusable_by_year={2027: (REUSED_2027, None)} + ) + caller = FakeCaller() + advancer = FakeAdvancer(store) + + rerun = _rerun_report(rerun_service(store, caller=caller, advancer=advancer)) + + ((request, (baseline, reform, years, reuse)),) = store.rerun_reads + assert (request.policyengine_version, request.year) == (POLICYENGINE_VERSION, 2026) + assert (request.region_code, request.dataset_name) == ("us", "populace_us_2024") + assert (baseline, reform, years) == ( + STORED_WINDOW.baseline.choice, + STORED_WINDOW.reform.choice, + STORED_WINDOW.window_years, + ) + assert reuse == RunReuse( + versions=VERSIONS, + inputs=STORED_OPTIONS, + output_uri_prefix=f"gs://{OUTPUT_BUCKET}/", + ) + (window,) = store.windows + assert (window.report_id, window.window_run_id, window.recalculate) == ( + REPORT_ID, + UUID(int=1), + False, + ) + assert (window.population, window.versions) == (CATALOG.population(2026), VERSIONS) + assert [window_year.year for window_year in window.years] == [2026, 2027, 2028] + for window_year in window.years: + baseline_id, reform_id = WINDOW_SIMULATION_IDS[window_year.year] + assert window_year.baseline.simulation_id == baseline_id + assert window_year.reform.simulation_id == reform_id + year_2027 = window.years[1] + assert (year_2027.baseline.simulation_run_id, year_2027.baseline.reused) == ( + REUSED_2027.simulation_run_id, + True, + ) + assert advancer.advances == [window.window_run_id] + assert len(caller.report_runs) == 3 + for (year_to_start, _), (_, payload, context, observed) in zip( + advancer.started, caller.report_runs, strict=True + ): + report = ReportExecutionInput.model_validate(payload) + assert report.baseline.year == report.reform.year == year_to_start.year + assert report.reform.policy == REFORM + assert report.reform.options == STORED_OPTIONS + assert Stage12InvocationContext.model_validate(context).report_run == ( + year_to_start.run_ids + ) + assert observed == OBSERVED + assert rerun.model_dump() == { + "report_id": REPORT_ID, + "report_run_id": window.years[0].report_run_id, + "budget_window_run_id": window.window_run_id, + "status": ReportRunStatus.PENDING, + "years": [ + { + "year": window_year.year, + "report_run_id": window_year.report_run_id, + "baseline_simulation_id": window_year.baseline.simulation_id, + "baseline_simulation_run_id": window_year.baseline.simulation_run_id, + "baseline_reused": window_year.baseline.reused, + "reform_simulation_id": window_year.reform.simulation_id, + "reform_simulation_run_id": window_year.reform.simulation_run_id, + "reform_reused": False, + } + for window_year in window.years + ], + } + + +def test_a_forced_window_rerun_recalculates_every_year() -> None: + store = FakeRerunStore( + report_definition=STORED_WINDOW, reusable_by_year={2027: (REUSED_2027, None)} + ) + + rerun = _rerun_report( + rerun_service(store, advancer=FakeAdvancer(store)), + ReportRerunRequest(force=True), + ) + + assert store.rerun_reads[0][1][3] is None + (window,) = store.windows + assert window.recalculate is True + assert not any(window_year.baseline.reused for window_year in window.years) + assert not any(year.baseline_reused for year in rerun.years) + + +def test_a_window_rerun_whose_first_year_modal_refuses_names_the_window_run() -> None: + store = FakeRerunStore(report_definition=STORED_WINDOW) + caller = FakeCaller(failure=RuntimeError("Modal is unavailable")) + + with pytest.raises(EconomySubmissionUnavailableError) as refused: + _rerun_report(rerun_service(store, caller=caller, advancer=FakeAdvancer(store))) + + error = refused.value + (window,) = store.windows + assert (error.code, error.record_id, error.run_id) == ( + "modal_call_refused", + REPORT_ID, + window.window_run_id, + ) + assert "the report's new run" in error.detail + + +@pytest.mark.parametrize( + ("changes", "named"), + [ + ({"window_years": ()}, "budget-window calculation"), + ({"budget_window_size": 4}, "budget-window calculation"), + ( + { + "window_years": ( + STORED_WINDOW.window_years[0], + STORED_WINDOW.window_years[2], + STORED_WINDOW.window_years[1], + ) + }, + "budget-window calculation", + ), + ( + { + "window_years": ( + *STORED_WINDOW.window_years[:2], + replace(STORED_WINDOW.window_years[2], reform_simulation_id=None), + ) + }, + "budget-window calculation", + ), + ], +) +def test_a_window_whose_years_cannot_be_rerun_is_refused(changes, named) -> None: + store = FakeRerunStore(report_definition=replace(STORED_WINDOW, **changes)) + + with pytest.raises(EconomyRerunConflictError) as refused: + _rerun_report(rerun_service(store, advancer=FakeAdvancer(store))) + + assert refused.value.code == "rerun_not_supported" + assert named in refused.value.detail + assert (store.rerun_reads, store.windows) == ([], []) diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 88ec5b19c..c6f51dbe1 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -414,8 +414,11 @@ def test_rerun_routes_declare_every_answer(): unprocessable = {"$ref": "#/components/schemas/EconomyRequestUnprocessable"} for path, rerun in ( - ("/reports/{report_id}/runs", "EconomyReportCreated"), - ("/simulations/{simulation_id}/runs", "EconomySimulationCreated"), + ( + "/reports/{report_id}/runs", + ("EconomyReportCreated", "BudgetWindowReportCreated"), + ), + ("/simulations/{simulation_id}/runs", ("EconomySimulationCreated",)), ): responses = _responses(spec, "post", path) assert set(responses) == { @@ -428,9 +431,10 @@ def test_rerun_routes_declare_every_answer(): "500", "503", } - assert _response_schema(responses["201"]) == { - "$ref": f"#/components/schemas/{rerun}" - } + rerun_schema = _response_schema(responses["201"]) + assert rerun_schema.get("anyOf", [rerun_schema]) == [ + {"$ref": f"#/components/schemas/{name}"} for name in rerun + ] for status_code in ("400", "404", "409", "500", "503"): assert _response_schema(responses[status_code]) == error assert _response_schema(responses["422"]) == unprocessable From f92e1f4953484fc6d0ea96d68bdcf9ea202154c8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:04:08 +0400 Subject: [PATCH 13/18] Re-run a budget-window report as a new budget-window run A budget window's re-run creates a budget-window run (trigger manual, recalculate set by force) with a report run per year, each year on the simulations that year's report run of the report's newest window run names, reusing a run of them that succeeded on the current versions unless forced. Its first years start as a new window's do: the budget-window service's start of the first years now takes the window's IDs. A report re-run, of either kind, is refused while a budget-window run or a report run of the report is pending or running, and a window whose newest run lacks a year answers 409 rerun_not_supported. The re-run answers with the budget-window create's shape. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 8 + .../economy_rerun_records.py | 193 ++++++++++++- .../economy_run_store.py | 61 +++++ .../budget_window_submissions.py | 28 +- .../economy_rerun_routes.py | 15 +- .../economy_reruns.py | 257 ++++++++++++++++-- 6 files changed, 514 insertions(+), 48 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 c4c0591d3..c8d7198bc 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -14,6 +14,7 @@ from policyengine_simulation_persistence.economy_rerun_records import ( NewReportRerun, NewSimulationRerun, + NewWindowRerun, RecordKind, ReportDefinition, RerunRecordDeletedError, @@ -21,8 +22,11 @@ RunInProgressError, SideDefinition, SimulationDefinition, + WindowRerunYear, + WindowYearDefinition, insert_report_rerun, insert_simulation_rerun, + insert_window_rerun, lock_reruns_of, read_report_definition, read_simulation_definition, @@ -158,6 +162,7 @@ "NewHouseholdJob", "NewReportRerun", "NewSimulationRerun", + "NewWindowRerun", "OverlappingReformError", "PolicyReform", "RecordKind", @@ -189,6 +194,8 @@ "UnknownHouseholdError", "UnknownPolicyError", "UnknownRegionError", + "WindowRerunYear", + "WindowYearDefinition", "bounded_transaction", "budget_window_runs", "combine_reforms", @@ -206,6 +213,7 @@ "insert_household_job", "insert_report_rerun", "insert_simulation_rerun", + "insert_window_rerun", "lock_reruns_of", "missing_table_privileges", "normalized_database_url", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py index 73e6cae18..091584d5c 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -27,7 +27,7 @@ from __future__ import annotations from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, Final, Literal from sqlalchemy import func, insert, literal, select @@ -51,6 +51,7 @@ ReportRunTrigger, SimulationStatus, SimulationType, + budget_window_runs, report_runs, reports, simulation_runs, @@ -87,6 +88,20 @@ class SideDefinition: choice: SimulationChoice +@dataclass(frozen=True) +class WindowYearDefinition: + """One year of a budget window's newest calculation, and its sides' simulations. + + They are the simulations of the runs that year's report run names: a + year that reused a run names that run's simulation, which may be another + report's. ``reform_simulation_id`` is None for a year naming no reform. + """ + + year: int + baseline_simulation_id: UUID + reform_simulation_id: UUID | None + + @dataclass(frozen=True) class ReportDefinition: """What a stored report calculates, as a re-run reads it. @@ -94,7 +109,9 @@ class ReportDefinition: ``region_code`` and ``dataset_name`` are the code and name of the report's region and dataset, or None when the report names none (or a row it named was deleted). ``baseline`` and ``reform`` are None when the - report names no such simulation. + report names no such simulation; a budget window's are its first + year's. A budget window's ``window_years`` are the years of its newest + budget-window run, earliest first, empty when it has none. """ id: UUID @@ -106,6 +123,8 @@ class ReportDefinition: inputs: dict[str, Any] baseline: SideDefinition | None reform: SideDefinition | None + budget_window_size: int | None = None + window_years: tuple[WindowYearDefinition, ...] = () @dataclass(frozen=True) @@ -160,6 +179,37 @@ def new_simulation_run_ids(self) -> tuple[UUID, ...]: ) +@dataclass(frozen=True) +class WindowRerunYear: + """One year of a budget window's re-run: its report run and each side's run.""" + + year: int + report_run_id: UUID + baseline: RerunSide + reform: RerunSide + + +@dataclass(frozen=True) +class NewWindowRerun: + """A stored budget window's new calculation, a report run per year. + + ``population`` is the first year's, on the current catalog version; + each year calculates on the same catalog version, region and dataset in + its own year. ``recalculate`` records whether the re-run was forced. + """ + + report_id: UUID + window_run_id: UUID + population: EconomyPopulation + versions: RunVersions + recalculate: bool + years: tuple[WindowRerunYear, ...] + + def population_in(self, year: int) -> EconomyPopulation: + """Return what the window's simulations of ``year`` calculate on.""" + return replace(self.population, year=year) + + @dataclass(frozen=True) class NewSimulationRerun: """A stored simulation's new run, on the current versions and catalog rows.""" @@ -229,16 +279,61 @@ def read_report_definition( ) if row is None: return None + kind = _report_kind(row["type"], row["household_id"], row["budget_window_size"]) return ReportDefinition( id=row["id"], country=row["country"], - kind=_report_kind(row["type"], row["household_id"], row["budget_window_size"]), + kind=kind, year=row["year"], region_code=row["region_code"], dataset_name=row["dataset_name"], inputs=row["inputs"], baseline=_side(row, "baseline"), reform=_side(row, "reform"), + budget_window_size=row["budget_window_size"], + window_years=( + _window_years(connection, report_id) if kind == "budget_window" else () + ), + ) + + +def _window_years( + connection: Connection, report_id: UUID +) -> tuple[WindowYearDefinition, ...]: + """Return the years of the report's newest budget-window run, earliest first.""" + newest_window_run = ( + select(budget_window_runs.c.id) + .where(budget_window_runs.c.report_id == report_id) + .order_by( + budget_window_runs.c.created_at.desc(), budget_window_runs.c.id.desc() + ) + .limit(1) + .scalar_subquery() + ) + baseline_run = simulation_runs.alias("baseline_run") + reform_run = simulation_runs.alias("reform_run") + rows = connection.execute( + select( + report_runs.c.year, + baseline_run.c.simulation_id.label("baseline_simulation_id"), + reform_run.c.simulation_id.label("reform_simulation_id"), + ) + .join( + baseline_run, baseline_run.c.id == report_runs.c.baseline_simulation_run_id + ) + .outerjoin( + reform_run, reform_run.c.id == report_runs.c.reform_simulation_run_id + ) + .where(report_runs.c.budget_window_run_id == newest_window_run) + .order_by(report_runs.c.year) + ) + return tuple( + WindowYearDefinition( + year=row.year, + baseline_simulation_id=row.baseline_simulation_id, + reform_simulation_id=row.reform_simulation_id, + ) + for row in rows ) @@ -365,11 +460,7 @@ def insert_report_rerun(connection: Connection, rerun: NewReportRerun) -> None: either refusal ends the caller's transaction, or its savepoint. """ lock_reruns_of(connection, rerun.report_id) - under_way = _run_under_way( - connection, report_runs, report_runs.c.report_id == rerun.report_id - ) - if under_way is not None: - raise RunInProgressError(under_way) + _refuse_while_report_runs(connection, rerun.report_id) with _record_checked(report_runs.c.report_id), references_checked(): for side in (rerun.baseline, rerun.reform): if not side.reused: @@ -397,6 +488,81 @@ def insert_report_rerun(connection: Connection, rerun: NewReportRerun) -> None: ) +def insert_window_rerun(connection: Connection, rerun: NewWindowRerun) -> None: + """Create the budget window's new run, each year's report run and new runs. + + The budget-window run has trigger ``manual`` and records ``recalculate``; + each year's report run names it and its year, and each side reusing no + run gets a new run of its year's simulation. All are pending with no + Modal call; the service then starts the first years. Raise as + ``insert_report_rerun`` does. + """ + lock_reruns_of(connection, rerun.report_id) + _refuse_while_report_runs(connection, rerun.report_id) + population = rerun.population + with ( + _record_checked(budget_window_runs.c.report_id, report_runs.c.report_id), + references_checked(), + ): + # First, so a report deleted after its reads is told apart from a + # year's simulation deleted with it. + connection.execute( + insert(budget_window_runs).values( + id=rerun.window_run_id, + report_id=rerun.report_id, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + trigger=ReportRunTrigger.MANUAL, + recalculate=rerun.recalculate, + status=ReportRunStatus.PENDING, + ) + ) + for window_year in rerun.years: + for side in (window_year.baseline, window_year.reform): + if not side.reused: + insert_simulation_run( + connection, + run_id=side.simulation_run_id, + simulation_id=side.simulation_id, + population=rerun.population_in(window_year.year), + versions=rerun.versions, + ) + connection.execute( + insert(report_runs), + [_year_run_row(rerun, window_year) for window_year in rerun.years], + ) + + +def _year_run_row( + rerun: NewWindowRerun, window_year: WindowRerunYear +) -> dict[str, Any]: + """Return a year's new report run, pending with no call, as inserted.""" + return { + "id": window_year.report_run_id, + "report_id": rerun.report_id, + "baseline_simulation_run_id": window_year.baseline.simulation_run_id, + "reform_simulation_run_id": window_year.reform.simulation_run_id, + "budget_window_run_id": rerun.window_run_id, + "year": window_year.year, + "country_package_version": rerun.versions.country_package_version, + "policyengine_version": rerun.versions.policyengine_version, + "status": ReportRunStatus.PENDING, + "trigger": ReportRunTrigger.MANUAL, + "modal_call_id": None, + } + + +def _refuse_while_report_runs(connection: Connection, report_id: UUID) -> None: + """Raise RunInProgressError while the report has a run pending or running. + + A budget-window run under way is named first, then a report run. + """ + under_way = _run_under_way( + connection, budget_window_runs, budget_window_runs.c.report_id == report_id + ) or _run_under_way(connection, report_runs, report_runs.c.report_id == report_id) + if under_way is not None: + raise RunInProgressError(under_way) + + def insert_simulation_rerun(connection: Connection, rerun: NewSimulationRerun) -> None: """Create the simulation's new run, pending with no call. @@ -424,15 +590,16 @@ def insert_simulation_rerun(connection: Connection, rerun: NewSimulationRerun) - @contextmanager -def _record_checked(own_key: Column[Any]) -> Iterator[None]: - """Turn a refusal of the new run's key to the record into RerunRecordDeletedError. +def _record_checked(*own_keys: Column[Any]) -> Iterator[None]: + """Turn a refusal of a new run's key to the record into RerunRecordDeletedError. - ``own_key`` is that key's column; a refusal of any other row propagates - unchanged. + ``own_keys`` are those keys' columns; a refusal of any other row + propagates unchanged. """ try: yield except EconomyRecordReferenceError as error: - if (error.table_name, error.column_name) == (own_key.table.name, own_key.key): + refused = (error.table_name, error.column_name) + if refused in {(key.table.name, key.key) for key in own_keys}: raise RerunRecordDeletedError from None raise diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py index 446a269af..6b803a89d 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py @@ -20,6 +20,7 @@ from policyengine_simulation_persistence.economy_rerun_records import ( insert_report_rerun, insert_simulation_rerun, + insert_window_rerun, read_report_definition, read_simulation_definition, ) @@ -77,9 +78,11 @@ from policyengine_simulation_persistence.economy_rerun_records import ( NewReportRerun, NewSimulationRerun, + NewWindowRerun, ReportDefinition, SideDefinition, SimulationDefinition, + WindowYearDefinition, ) from policyengine_simulation_persistence.economy_run_reads import ( EconomyReportState, @@ -361,6 +364,64 @@ def read_report_rerun_inputs( search=search, ) + def read_window_rerun_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + years: Sequence[WindowYearDefinition], + reuse: RunReuse | None, + ) -> BudgetWindowInputs: + """Read a window re-run's catalog rows, both sides' values and each year's runs. + + ``request.year`` is the window's first year. With ``reuse``, each + year's sides get the runs ``find_reusable_run_of`` finds among that + year's own simulations' runs, and the reform never gets the + baseline's. Raise the errors ``read_report_inputs`` raises. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + baseline_inputs = _side_inputs( + connection, request, catalog, "baseline", baseline + ) + reform_inputs = _side_inputs(connection, request, catalog, "reform", reform) + reusable_runs: dict[int, tuple[ReusableRun | None, ReusableRun | None]] = {} + for window_year in years: + baseline_run = reform_run = None + if reuse is not None and window_year.reform_simulation_id is not None: + search = _search_own_runs( + connection, + catalog.population(window_year.year), + reuse, + { + "baseline": window_year.baseline_simulation_id, + "reform": window_year.reform_simulation_id, + }, + ) + baseline_run = search("baseline", baseline, None) + reform_run = search( + "reform", + reform, + None + if baseline_run is None + else baseline_run.simulation_run_id, + ) + reusable_runs[window_year.year] = (baseline_run, reform_run) + return BudgetWindowInputs( + catalog=catalog, + baseline=baseline_inputs, + reform=reform_inputs, + reusable_runs=reusable_runs, + ) + + def create_window_rerun(self, rerun: NewWindowRerun) -> None: + """Create the window's new run and its years' runs; see ``insert_window_rerun``. + + Raise EconomyRecordOutcomeUnknownError as ``create_report`` does. + """ + self._create(lambda connection: insert_window_rerun(connection, rerun)) + def create_report_rerun(self, rerun: NewReportRerun) -> None: """Create the report's new run and its sides' new simulation runs. diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py index 83d6213ba..c4cf4d59e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py @@ -218,8 +218,10 @@ async def submit_budget_window( ) status = await self._start_first_years( worker, - report, first_calculation, + report_id=report.id, + window_run_id=report.window_run_id, + first_year=first_year.year, observability_context=observability_context, ) return BudgetWindowReportCreated( @@ -240,15 +242,18 @@ def _window_year(self, inputs: BudgetWindowInputs, year: int) -> BudgetWindowYea reform=self._side(_with_reusable(inputs.reform, reform_run)), ) - async def _start_first_years( + async def _start_first_years( # noqa: PLR0913 # the call, then the window it starts self, worker: EconomyWorker, - report: NewBudgetWindowReport, first_calculation: RunCalculation, *, + report_id: UUID, + window_run_id: UUID, + first_year: int, observability_context: dict[str, Any] | None, + rerun: bool = False, ) -> ReportRunStatus: - """Advance the new window, starting its first years; return its status. + """Advance the new window run, starting its first years; return its status. Raise EconomySubmissionUnavailableError naming the report when the window failed because Modal refused a year, or didn't confirm the @@ -256,12 +261,13 @@ async def _start_first_years( left waiting. An advance that can't be written leaves the window pending: the report exists, so it is answered as created, and a year Modal started claims itself, while the stuck-run sweep fails a window no - year advances. + year advances. ``rerun`` says the report existed before the window + run, which the 503 then names. """ start_year = _YearStarts( worker=worker, caller=self._caller, - first_year=report.first_year.year, + first_year=first_year, first_calculation=first_calculation, observability_context=observability_context, loop=asyncio.get_running_loop(), @@ -271,7 +277,7 @@ async def _start_first_years( try: advance = await asyncio.to_thread( self._windows.advance_window, - report.window_run_id, + window_run_id, now=self._clock(), start_year=start_year, ) @@ -280,7 +286,7 @@ async def _start_first_years( "budget_window_not_started", extra={ "record": "report", - "record_id": str(report.id), + "record_id": str(report_id), "error_type": type(error).__name__, }, ) @@ -291,11 +297,13 @@ async def _start_first_years( "budget_window_call_unconfirmed", extra={ "record": "report", - "record_id": str(report.id), + "record_id": str(report_id), "started_years": list(advance.started_years), }, ) - refused = call_refused("report", report.id) + refused = call_refused( + "report", report_id, rerun_id=window_run_id if rerun else None + ) raise refused from None return ReportRunStatus(advance.status.value) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py index 25d250379..235dee607 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -28,6 +28,9 @@ from fastapi import APIRouter, Depends, Request, status from fastapi.responses import JSONResponse +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportCreated, +) from policyengine_simulation_entry.economy_api_models import ( EconomyError, EconomyReportCreated, @@ -78,8 +81,8 @@ async def rerun_report( *, request_id: str, observability_context: dict[str, Any] | None, - ) -> EconomyReportCreated: - """Create the report's new run and start its coordinator.""" + ) -> EconomyReportCreated | BudgetWindowReportCreated: + """Create the report's new run and start it, or a window's first years.""" ... async def rerun_simulation( @@ -231,17 +234,19 @@ async def answer( # the route, its record and how to re-run it @router.post( REPORT_RUNS_PATH, - summary="Re-run an economy report", + summary="Re-run an economy or budget-window report", description=( "Create the stored report's new run on the v2 version manifest's " "default worker, and start it; every earlier run is kept. Without " "force, a side reuses a run of its own simulation that succeeded " "on the same versions; with force, both simulations are " - "calculated again." + "calculated again. A budget-window report's new run is a " + "budget-window run with a report run per year, up to five of them " + "started at once, each year on its own year's simulations." ), operation_id="rerun_report", status_code=status.HTTP_201_CREATED, - response_model=EconomyReportCreated, + response_model=EconomyReportCreated | BudgetWindowReportCreated, responses=_responses("report"), dependencies=[Depends(submitter)], ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py index 5c295701e..c704c4780 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -40,12 +40,18 @@ EconomyRequest, NewReportRerun, NewSimulationRerun, + NewWindowRerun, RerunRecordDeletedError, RerunSide, RunInProgressError, RunReuse, + WindowRerunYear, ) +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportCreated, + BudgetWindowYearCreated, +) from policyengine_simulation_entry.budget_window_submissions import ( BudgetWindowRuns, ReportSubmissionService, @@ -71,17 +77,20 @@ ) if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Sequence from uuid import UUID from policyengine_simulation_persistence import ( + BudgetWindowInputs, EconomyCatalogRows, EconomyReportInputs, RecordKind, ReportDefinition, + ReusableRun, SideDefinition, + SimulationChoice, SimulationDefinition, - SimulationInputs, + WindowYearDefinition, ) from policyengine_simulation_entry.economy_api_models import ( @@ -97,10 +106,11 @@ # How a refusal names each kind of stored record this service does not re-run # yet. -_NOT_RERUN: Final[dict[RecordKind, str]] = { - "budget_window": "budget-window", - "household": "household", -} +_NOT_RERUN: Final[dict[RecordKind, str]] = {"household": "household"} +# The kinds of stored report this service re-runs. +_RERUN_REPORT_KINDS: Final[frozenset[RecordKind]] = frozenset( + {"economy", "budget_window"} +) # The countries this service calculates, as the contract names them. _COUNTRIES: Final[tuple[CountryId, ...]] = get_args(CountryId) _REPORT: Final[RecordName] = "report" @@ -134,6 +144,22 @@ def read_report_rerun_inputs( """Read a re-run's catalog rows, both sides' values and reusable runs.""" ... + def read_window_rerun_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + years: Sequence[WindowYearDefinition], + reuse: RunReuse | None, + ) -> BudgetWindowInputs: + """Read a window re-run's catalog rows, both sides' values and years' runs.""" + ... + + def create_window_rerun(self, rerun: NewWindowRerun) -> None: + """Create the window's new run, each year's report run and new runs.""" + ... + def create_report_rerun(self, rerun: NewReportRerun) -> None: """Create the report's new run and its sides' new simulation runs.""" ... @@ -153,13 +179,15 @@ async def rerun_report( *, request_id: str, observability_context: dict[str, Any] | None, - ) -> EconomyReportCreated: + ) -> EconomyReportCreated | BudgetWindowReportCreated: """Create the stored report's new run on the current versions, and start it. - Raise EconomyRecordNotFoundError (404) or EconomyRerunConflictError - (409) as the module describes, and otherwise the errors - ``submit_report`` raises, a 503 for Modal's refusal naming the new - run as well as the report. + A budget window's new run is a budget-window run with a report run + per year, whose first years are started (``_rerun_window``). Raise + EconomyRecordNotFoundError (404) or EconomyRerunConflictError (409) + as the module describes, and otherwise the errors ``submit_report`` + raises, a 503 for Modal's refusal naming the new run as well as the + report. """ definition = await self._read_state( lambda: self._store.read_report_definition(report_id) @@ -167,6 +195,11 @@ async def rerun_report( if definition is None: raise EconomyRecordNotFoundError(_REPORT) stored = _rerunnable_report(definition) + window_years = ( + _complete_window_years(definition) + if definition.kind == "budget_window" + else () + ) worker = await self._default_worker(stored.country) economy_request = EconomyRequest( country=stored.country, @@ -184,6 +217,20 @@ async def rerun_report( output_uri_prefix=self._settings.reuse_prefix, ) ) + if window_years: + return await self._rerun_window( + _WindowRerunRequest( + stored=stored, + options=definition.inputs, + years=window_years, + economy_request=economy_request, + reuse=reuse, + recalculate=request.force, + ), + worker, + request_id=request_id, + observability_context=observability_context, + ) inputs = await self._read_inputs( lambda: self._store.read_report_rerun_inputs( economy_request, @@ -199,8 +246,12 @@ async def rerun_report( run_id=self._new_id(), population=inputs.catalog.population(stored.year), versions=worker.versions, - baseline=self._rerun_side(stored.baseline, inputs.baseline), - reform=self._rerun_side(stored.reform, inputs.reform), + baseline=self._rerun_side( + stored.baseline.simulation_id, inputs.baseline.reusable_run + ), + reform=self._rerun_side( + stored.reform.simulation_id, inputs.reform.reusable_run + ), ) calculation = report_run_calculation( worker, @@ -314,18 +365,107 @@ async def rerun_simulation( status=status, ) + async def _rerun_window( + self, + window: _WindowRerunRequest, + worker: EconomyWorker, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> BudgetWindowReportCreated: + """Create the window's new run and each year's runs; start its first years. + + Each year's sides are that year's own simulations, as its report run + in the report's newest window run names them, and each year reuses a + run of its own simulations unless the re-run is forced. The years + start as a new window's do (``_start_first_years``). + """ + stored = window.stored + inputs = await self._read_inputs( + lambda: self._store.read_window_rerun_inputs( + window.economy_request, + baseline=stored.baseline.choice, + reform=stored.reform.choice, + years=[year.definition for year in window.years], + reuse=window.reuse, + ), + request=window.economy_request, + source=STORED_REPORT, + ) + window_run_id = self._new_id() + rerun = NewWindowRerun( + report_id=stored.report_id, + window_run_id=window_run_id, + population=inputs.catalog.population(stored.year), + versions=worker.versions, + recalculate=window.recalculate, + years=tuple(self._window_rerun_year(year, inputs) for year in window.years), + ) + first_year = rerun.years[0] + first_calculation = report_run_calculation( + worker, + _rerun_scope( + stored.country, + window.economy_request, + worker, + inputs.catalog, + options=window.options, + ), + self._origin(request_id), + run_ids=ReportRunIds( + report_run_id=first_year.report_run_id, + baseline_simulation_run_id=first_year.baseline.simulation_run_id, + reform_simulation_run_id=first_year.reform.simulation_run_id, + ), + baseline_policy=inputs.baseline.reform, + reform_policy=inputs.reform.reform, + ) + await self._create_rerun( + _REPORT, + lambda: self._store.create_window_rerun(rerun), + source=STORED_REPORT, + ) + status = await self._start_first_years( + worker, + first_calculation, + report_id=stored.report_id, + window_run_id=window_run_id, + first_year=first_year.year, + observability_context=observability_context, + rerun=True, + ) + return BudgetWindowReportCreated( + report_id=stored.report_id, + report_run_id=first_year.report_run_id, + budget_window_run_id=window_run_id, + status=status, + years=[_rerun_year_created(window_year) for window_year in rerun.years], + ) + + def _window_rerun_year( + self, year: _CompleteWindowYear, inputs: BudgetWindowInputs + ) -> WindowRerunYear: + """Return a year naming its own simulations' reusable runs, or new runs.""" + baseline_run, reform_run = inputs.reusable_runs[year.year] + return WindowRerunYear( + year=year.year, + report_run_id=self._new_id(), + baseline=self._rerun_side(year.baseline_simulation_id, baseline_run), + reform=self._rerun_side(year.reform_simulation_id, reform_run), + ) + def _rerun_side( - self, stored: SideDefinition, inputs: SimulationInputs + self, simulation_id: UUID, reusable_run: ReusableRun | None ) -> RerunSide: """Return a side naming its own simulation's reusable run, or a new run.""" - if inputs.reusable_run is not None: + if reusable_run is not None: return RerunSide( - simulation_id=stored.simulation_id, - simulation_run_id=inputs.reusable_run.simulation_run_id, + simulation_id=simulation_id, + simulation_run_id=reusable_run.simulation_run_id, reused=True, ) return RerunSide( - simulation_id=stored.simulation_id, + simulation_id=simulation_id, simulation_run_id=self._new_id(), reused=False, ) @@ -362,6 +502,7 @@ async def _create_rerun( class _StoredReportRows: """The country and rows a stored report calculates on, each present.""" + report_id: UUID country: CountryId year: int region_code: str @@ -394,7 +535,7 @@ def _refusal( """ if kind == "unsupported": detail = f"The service does not re-run this type of {record}." - elif kind != "economy": + elif kind in _NOT_RERUN: detail = f"The service does not re-run {_NOT_RERUN[kind]} {record}s yet." else: lacking = next(name for name, value in missing if value is None) @@ -408,7 +549,7 @@ def _rerunnable_report(definition: ReportDefinition) -> _StoredReportRows: region_code, dataset_name = definition.region_code, definition.dataset_name country, year = _calculated_country(definition.country), definition.year if ( - definition.kind != "economy" + definition.kind not in _RERUN_REPORT_KINDS or country is None or baseline is None or reform is None @@ -429,6 +570,7 @@ def _rerunnable_report(definition: ReportDefinition) -> _StoredReportRows: ), ) return _StoredReportRows( + report_id=definition.id, country=country, year=year, region_code=region_code, @@ -438,6 +580,81 @@ def _rerunnable_report(definition: ReportDefinition) -> _StoredReportRows: ) +@dataclass(frozen=True) +class _CompleteWindowYear: + """A year of a stored window's newest calculation, with both sides' simulations.""" + + definition: WindowYearDefinition + baseline_simulation_id: UUID + reform_simulation_id: UUID + + @property + def year(self) -> int: + """The year.""" + return self.definition.year + + +@dataclass(frozen=True) +class _WindowRerunRequest: + """What a budget window's re-run calculates, as the service has read it.""" + + stored: _StoredReportRows + options: dict[str, Any] + years: tuple[_CompleteWindowYear, ...] + economy_request: EconomyRequest + reuse: RunReuse | None + recalculate: bool + + +def _complete_window_years( + definition: ReportDefinition, +) -> tuple[_CompleteWindowYear, ...]: + """Return the stored window's years, or refuse a window that lacks one. + + Its newest calculation must have a report run naming both sides in each + year of its window, from the report's first year. + """ + expected = ( + () + if definition.year is None or definition.budget_window_size is None + else tuple( + range(definition.year, definition.year + definition.budget_window_size) + ) + ) + complete = tuple( + _CompleteWindowYear( + definition=year, + baseline_simulation_id=year.baseline_simulation_id, + reform_simulation_id=year.reform_simulation_id, + ) + for year in definition.window_years + if year.reform_simulation_id is not None + ) + if not expected or tuple(year.year for year in complete) != expected: + raise EconomyRerunConflictError( + code="rerun_not_supported", + detail=( + "The report names no complete budget-window calculation, a " + "report run naming both sides in each year of its window, so " + "it cannot be re-run." + ), + ) + return complete + + +def _rerun_year_created(window_year: WindowRerunYear) -> BudgetWindowYearCreated: + return BudgetWindowYearCreated( + year=window_year.year, + report_run_id=window_year.report_run_id, + baseline_simulation_id=window_year.baseline.simulation_id, + baseline_simulation_run_id=window_year.baseline.simulation_run_id, + baseline_reused=window_year.baseline.reused, + reform_simulation_id=window_year.reform.simulation_id, + reform_simulation_run_id=window_year.reform.simulation_run_id, + reform_reused=window_year.reform.reused, + ) + + def _rerunnable_simulation(definition: SimulationDefinition) -> _StoredSimulationRows: """Return the stored simulation's country and rows, or refuse it. From 2068202a4a81c7b39ea26a646716bb60bb68aee4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:04:42 +0400 Subject: [PATCH 14/18] Document re-running a budget-window report, with a changelog fragment Co-Authored-By: Claude Opus 5.5 --- changelog.d/c1-17-budget-window-reruns.added.md | 1 + .../stage-12-modal-worker-foundation.md | 8 +++++--- .../README.md | 13 +++++++++++-- projects/policyengine-simulation-entry/README.md | 16 ++++++++++++---- .../economy_rerun_routes.py | 3 +++ .../economy_reruns.py | 10 ++++++---- 6 files changed, 38 insertions(+), 13 deletions(-) create mode 100644 changelog.d/c1-17-budget-window-reruns.added.md diff --git a/changelog.d/c1-17-budget-window-reruns.added.md b/changelog.d/c1-17-budget-window-reruns.added.md new file mode 100644 index 000000000..481580859 --- /dev/null +++ b/changelog.d/c1-17-budget-window-reruns.added.md @@ -0,0 +1 @@ +`POST /reports/{report_id}/runs` also re-runs a stored budget-window report, as a new budget-window run with a report run per year on each year's own simulations, its first years started as a new window's are. diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index a01a48d80..80c741de4 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -340,9 +340,11 @@ trigger `manual`) and, for a report, a new `simulation_runs` row of each side's own simulation, or names that simulation's run that already succeeded on the same versions, unless the re-run is forced. The new runs record the catalog version, dataset and region of the current .py version's catalog; the -`simulations` and `reports` rows are unchanged. The coordinator and workers -claim and calculate them as they would a new record's: nothing on the Modal -side tells a re-run apart. +`simulations` and `reports` rows are unchanged. A budget-window report's +re-run is a new `budget_window_runs` row (trigger `manual`) with a report run +per year, each on its year's own simulations, started as a new window's +first years are. The coordinator and workers claim and calculate them as +they would a new record's: nothing on the Modal side tells a re-run apart. **Each worker** claims its simulation run under its own Modal call ID: a pending or running run that records no call or its own. It records its output file, or diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 608d012ba..885b9f668 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -153,11 +153,20 @@ version, dataset and region it calculates on. record itself, the record was deleted after it was read: `RerunRecordDeletedError`. +A budget window's definition also names its newest budget-window run's +years (`WindowYearDefinition`): each year's sides are the simulations of the +runs that year's report run names. `insert_window_rerun` creates the +window's new budget-window run (trigger `manual`, `recalculate` set), a +report run per year and each year's new runs of its own simulations. Every +report re-run is refused while a budget-window run or a report run of the +report is `pending` or `running`. + `EconomyRunStore.read_report_rerun_inputs` reads a re-run's catalog rows by the stored region code and dataset name, and each side's values and reusable run of its own simulation, never giving both sides the same run. -`create_report_rerun` and `create_simulation_rerun` run the inserts in a -transaction of their own. +`read_window_rerun_inputs` does the same for each year of a window. +`create_report_rerun`, `create_window_rerun` and `create_simulation_rerun` +run the inserts in a transaction of their own. ## Reading a policy diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index bc28253ae..11570676a 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -203,6 +203,13 @@ both again. The new report run's trigger is `manual`. The service reads the report's region code and dataset name, finds the rows with that code and name in the catalog version of the default worker's .py version, and reads each simulation's stored policy and dynamic on it. +A budget-window report's re-run is a new budget-window run, trigger +`manual`, with `recalculate` set by `force`, and a report run per year. Each +year's sides are the simulations that year's report run of the report's +newest window run names (a year that reused another report's run names that +report's simulation), and each reuses a run of them as a single-year side +does. The first five years start as a new window's do; it answers 201 with +the budget-window create's fields, naming the window run and every year. `POST /simulations/{simulation_id}/runs`, with no body, creates and starts a stored economy simulation's new run, which is always calculated. Each answers 201 with the same fields as the create, naming the new run. @@ -210,10 +217,11 @@ answers 201 with the same fields as the create, naming the new run. A re-run is refused before anything is written: 404 (`unknown_report`, `unknown_simulation`) for a record that doesn't exist; 409 `run_in_progress`, naming the run in `run_id`, while the record has a run -pending or running; 409 `rerun_not_supported` for a budget-window or -household record, which this service does not re-run yet, or one that -names no region, dataset, year, simulation or country the service -calculates; and the create's 400 and 422 +pending or running (for a report, a report run or a budget-window run); +409 `rerun_not_supported` for a household record, which this service does +not re-run yet, a budget window whose newest window run lacks a year, or a +record that names no region, dataset, year, simulation or country the +service calculates; and the create's 400 and 422 codes when the stored rows can't be calculated on the current catalog version, worded for the stored record (a 422 reason's `loc` starts with `report` or `simulation`). The check for a run under way and the inserts diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py index 235dee607..4923f8b12 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -6,6 +6,9 @@ new run's ID and status, and each side's simulation and simulation run. ``{"force": true}`` calculates both simulations again; otherwise a side reuses a run of its own simulation that succeeded on the same versions. + A budget-window report's new run is a budget-window run with a report run + per year, of which the first five are started; it answers as the + budget-window create does, naming the window run and every year's runs. - ``POST /simulations/{simulation_id}/runs`` (``simulation:submit``) creates the stored simulation's new run and starts its simulation function, answering 201 with the simulation's ID, the new run's ID and its status. diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py index c704c4780..97e5a5878 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -23,10 +23,12 @@ coordinator, or the simulation function, by run ID (C1 ยท 8's coordination), failing the new runs when Modal refuses the call. -Whatever is missing or cannot be calculated is refused in steps 1 and 2, -before anything is written. A budget-window report and a household report -or simulation are refused for now; each gets its re-run with its own -submission. +A budget-window report's re-run is a new budget-window run with a report +run per year, each year on the simulations that year's report run of the +report's newest window run names; its first years start as a new window's +do. Whatever is missing or cannot be calculated is refused in steps 1 and +2, before anything is written. A household report or simulation is refused +for now; it gets its re-run with its own submission. """ from __future__ import annotations From 4356480399972702b377a8b0e0dd03b73e4091ff Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:24:32 +0400 Subject: [PATCH 15/18] Test review round 4's findings for budget-window re-runs (failing) Fail: a window re-run checks its years are its window's, and a year naming no reform is left out of a window's definition. Pass at once, covering existing behaviour: a window re-run of a deleted report or year simulation, one refused while a year still runs after its window ended, and the budget-window store starting only a report's later window run's years. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_window_reruns.py | 117 ++++++++++++++++++ .../tests/budget_window_database.py | 81 ++++++++++++ .../tests/test_budget_window_store.py | 22 ++++ .../tests/test_economy_reruns.py | 38 +++--- 4 files changed, 234 insertions(+), 24 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py index a8dfb5e1f..88eb4c22f 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py @@ -35,6 +35,8 @@ ) from policyengine_simulation_persistence import ( BudgetWindowYear, + EconomyPopulation, + EconomyRecordReferenceError, EconomyRequest, EconomyRunStore, NewBudgetWindowReport, @@ -43,6 +45,7 @@ ReportRunStatus, ReportRunTrigger, ReportSimulation, + RerunRecordDeletedError, RerunSide, ReusableRun, RunInProgressError, @@ -482,3 +485,117 @@ def test_a_window_rerun_is_created_through_the_store(public_api_database_url) -> "trigger" ] == (ReportRunTrigger.MANUAL) assert refused.value.run_id == rerun.window_run_id + + +def test_a_window_rerun_refuses_years_that_are_not_its_window() -> None: + """Its years run one at a time from its first year, at most 75 of them.""" + year = WindowRerunYear( + year=2026, + report_run_id=uuid4(), + baseline=RerunSide( + simulation_id=uuid4(), simulation_run_id=uuid4(), reused=False + ), + reform=RerunSide( + simulation_id=uuid4(), simulation_run_id=uuid4(), reused=False + ), + ) + rerun = NewWindowRerun( + report_id=uuid4(), + window_run_id=uuid4(), + population=EconomyPopulation( + tax_benefit_model_version_id=uuid4(), + region_id=uuid4(), + dataset_id=uuid4(), + year=2026, + ), + versions=LATER_VERSIONS, + recalculate=False, + years=(year,), + ) + + for years in ( + (), + (replace(year, year=2027),), + (year, replace(year, year=2028)), + tuple(replace(year, year=2026 + offset) for offset in range(76)), + ): + with pytest.raises(ValueError, match="consecutive years"): + replace(rerun, years=years) + + +def test_a_year_naming_no_reform_is_left_out_of_a_windows_definition( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + stored = _stored_window(connection, economy) + connection.execute( + text( + "UPDATE public.report_runs SET reform_simulation_run_id = NULL " + "WHERE id = :id" + ), + {"id": stored.years[1].report_run_id}, + ) + act_as_the_service(connection) + + definition = read_report_definition(connection, stored.report_id) + + assert definition is not None + assert [year.year for year in definition.window_years] == [2026, 2028] + + +def test_a_window_rerun_of_a_deleted_report_or_year_simulation_is_refused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window(connection, economy) + act_as_the_service(connection) + rerun = _window_rerun_of(stored, later) + gone_year = replace( + rerun.years[2], + reform=RerunSide( + simulation_id=uuid4(), simulation_run_id=uuid4(), reused=False + ), + ) + + with pytest.raises(RerunRecordDeletedError), connection.begin_nested(): + insert_window_rerun(connection, replace(rerun, report_id=uuid4())) + with ( + pytest.raises(EconomyRecordReferenceError) as refused, + connection.begin_nested(), + ): + insert_window_rerun( + connection, replace(rerun, years=(*rerun.years[:2], gone_year)) + ) + + assert (refused.value.table_name, refused.value.column_name) == ( + "simulation_runs", + "simulation_id", + ) + + +def test_a_window_rerun_is_refused_while_a_year_still_runs_after_its_window_ended( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + later = insert_later_catalog(connection, economy) + stored = _stored_window( + connection, economy, window_status=ReportRunStatus.FAILED + ) + still_running = stored.years[2].report_run_id + connection.execute( + text( + "UPDATE public.report_runs SET status = 'running', " + "started_at = now(), completed_at = NULL WHERE id = :id" + ), + {"id": still_running}, + ) + act_as_the_service(connection) + + with pytest.raises(RunInProgressError) as refused: + insert_window_rerun(connection, _window_rerun_of(stored, later)) + + assert refused.value.run_id == still_running diff --git a/libs/policyengine-stage12-persistence/tests/budget_window_database.py b/libs/policyengine-stage12-persistence/tests/budget_window_database.py index b8548feb6..0313febe7 100644 --- a/libs/policyengine-stage12-persistence/tests/budget_window_database.py +++ b/libs/policyengine-stage12-persistence/tests/budget_window_database.py @@ -10,6 +10,7 @@ from __future__ import annotations from dataclasses import dataclass +from datetime import timedelta from typing import TYPE_CHECKING, Any from uuid import UUID, uuid4 @@ -147,6 +148,86 @@ def seed_window(self, first_year: int = 2026, size: int = 7) -> SeededWindow: model_id=model_id, ) + def seed_window_rerun(self, window: SeededWindow) -> SeededWindow: + """Insert a later pending window run of ``window``'s report, as a re-run. + + Each year gets a new report run and two new runs of the report's + simulations; the earlier window run and its years are left as they are. + """ + table = self.seeder.table + with self.engine.connect() as connection: + first_run = connection.execute( + select(table("simulation_runs")).where( + table("simulation_runs").c.id + == window.years[0].baseline_simulation_run_id + ) + ).one() + window_run_id = uuid4() + years = tuple( + SeededYear( + year=seeded.year, + report_run_id=uuid4(), + baseline_simulation_run_id=uuid4(), + reform_simulation_run_id=uuid4(), + ) + for seeded in window.years + ) + later = {"created_at": NOW + timedelta(minutes=1), "updated_at": NOW} + with self.engine.begin() as connection: + connection.execute( + insert(table("budget_window_runs")).values( + id=window_run_id, + report_id=window.report_id, + tax_benefit_model_version_id=first_run.tax_benefit_model_version_id, + trigger="manual", + recalculate=True, + status="pending", + **later, + ) + ) + for seeded in years: + for run_id, simulation_id in zip( + ( + seeded.baseline_simulation_run_id, + seeded.reform_simulation_run_id, + ), + window.simulation_ids, + strict=True, + ): + connection.execute( + insert(table("simulation_runs")).values( + id=run_id, + simulation_id=simulation_id, + year=seeded.year, + tax_benefit_model_version_id=first_run.tax_benefit_model_version_id, + dataset_id=first_run.dataset_id, + status="pending", + **later, + **VERSIONS, + ) + ) + connection.execute( + insert(table("report_runs")).values( + id=seeded.report_run_id, + report_id=window.report_id, + baseline_simulation_run_id=seeded.baseline_simulation_run_id, + reform_simulation_run_id=seeded.reform_simulation_run_id, + budget_window_run_id=window_run_id, + year=seeded.year, + status="pending", + trigger="manual", + **later, + **VERSIONS, + ) + ) + return SeededWindow( + window_run_id=window_run_id, + report_id=window.report_id, + years=years, + simulation_ids=window.simulation_ids, + model_id=window.model_id, + ) + def set_years( self, window: SeededWindow, years: Iterable[int], **values: Any ) -> None: # column values of any type diff --git a/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py b/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py index d7b6b8bb7..4c3f2cad4 100644 --- a/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_budget_window_store.py @@ -641,3 +641,25 @@ def test_the_sweep_lists_the_waiting_years_of_a_finished_window( listed = _listed_report_runs(run_store) assert set(_run_ids(window, ALL_YEARS)) <= set(listed) + + +def test_a_reports_later_window_run_starts_only_its_own_years( + store, windows, window +) -> None: + """A re-run's window run of the same report and simulations advances alone.""" + _succeeded(windows, window, ALL_YEARS) + windows.set_window(window, status="succeeded", started_at=NOW, completed_at=NOW) + rerun = windows.seed_window_rerun(window) + starter = RecordingStarter() + + advance = store.advance_window(rerun.window_run_id, now=LATER, start_year=starter) + + assert starter.started == _run_ids(rerun, range(2026, 2031)) + assert advance.started_years == (2026, 2027, 2028, 2029, 2030) + assert windows.window_row(window)["status"] == "succeeded" + assert {row["status"] for row in windows.year_rows(window).values()} == { + "succeeded" + } + assert {row["modal_call_id"] for row in windows.year_rows(window).values()} == { + "fc-finished" + } diff --git a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py index 7a6c95b7c..6dd37eb2a 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_reruns.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_reruns.py @@ -589,37 +589,27 @@ def test_a_window_rerun_whose_first_year_modal_refuses_names_the_window_run() -> @pytest.mark.parametrize( - ("changes", "named"), + "changes", [ - ({"window_years": ()}, "budget-window calculation"), - ({"budget_window_size": 4}, "budget-window calculation"), - ( - { - "window_years": ( - STORED_WINDOW.window_years[0], - STORED_WINDOW.window_years[2], - STORED_WINDOW.window_years[1], - ) - }, - "budget-window calculation", - ), - ( - { - "window_years": ( - *STORED_WINDOW.window_years[:2], - replace(STORED_WINDOW.window_years[2], reform_simulation_id=None), - ) - }, - "budget-window calculation", - ), + {"window_years": ()}, + {"budget_window_size": 4}, + # A year missing from the newest window run, as one naming no reform is. + {"window_years": STORED_WINDOW.window_years[::2]}, + { + "window_years": ( + STORED_WINDOW.window_years[0], + STORED_WINDOW.window_years[2], + STORED_WINDOW.window_years[1], + ) + }, ], ) -def test_a_window_whose_years_cannot_be_rerun_is_refused(changes, named) -> None: +def test_a_window_whose_years_cannot_be_rerun_is_refused(changes) -> None: store = FakeRerunStore(report_definition=replace(STORED_WINDOW, **changes)) with pytest.raises(EconomyRerunConflictError) as refused: _rerun_report(rerun_service(store, advancer=FakeAdvancer(store))) assert refused.value.code == "rerun_not_supported" - assert named in refused.value.detail + assert "complete budget-window calculation" in refused.value.detail assert (store.rerun_reads, store.windows) == ([], []) From a7ca243e8cbc90ff2f1449010617d93af09573a3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:28:41 +0400 Subject: [PATCH 16/18] Fix review round 4's findings: share the window code, check a re-run's years A window re-run's dataclass checks its years as a new window's does, with the same check; a window's definition leaves out a year naming no reform, so the service needs no copy of each year. A new window and a re-run share the per-year reuse search, the year's report run row and the year's answer; the budget-window service reads its store by one name. The docs say which run a 409 run_in_progress names, that up to five years start at once, and that a deleted budget-window report always answers 404. Co-Authored-By: Claude Opus 5.5 --- .../budget_window_records.py | 73 ++++++--- .../economy_rerun_records.py | 59 ++++---- .../economy_run_store.py | 140 +++++++++--------- .../policyengine-simulation-entry/README.md | 11 +- .../budget_window_submissions.py | 13 +- .../economy_api_models.py | 10 +- .../economy_rerun_routes.py | 14 +- .../economy_reruns.py | 56 ++----- 8 files changed, 198 insertions(+), 178 deletions(-) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py index 111a025ff..b7ef5293d 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py @@ -40,10 +40,15 @@ ) if TYPE_CHECKING: + from collections.abc import Sequence from uuid import UUID from sqlalchemy import Connection + from policyengine_simulation_persistence.economy_rerun_records import ( + WindowRerunYear, + ) + @dataclass(frozen=True) class BudgetWindowYear: @@ -78,20 +83,7 @@ class NewBudgetWindowReport: def __post_init__(self) -> None: """Refuse years that aren't 1 to 75 consecutive ones from the first.""" - first_year = self.population.year - expected = tuple(range(first_year, first_year + len(self.years))) - if ( - not self.years - or len(self.years) > MAX_BUDGET_WINDOW_YEARS - or tuple(year.year for year in self.years) != expected - or expected[-1] > LATEST_YEAR - ): - msg = ( - "a budget window's years must be 1 to " - f"{MAX_BUDGET_WINDOW_YEARS} consecutive years from its first, " - f"ending by {LATEST_YEAR}" - ) - raise ValueError(msg) + check_window_years(self.population.year, [year.year for year in self.years]) @property def first_year(self) -> BudgetWindowYear: @@ -103,6 +95,23 @@ def population_in(self, year: int) -> EconomyPopulation: return replace(self.population, year=year) +def check_window_years(first_year: int, years: Sequence[int]) -> None: + """Raise ValueError unless ``years`` are 1 to 75 in a row from ``first_year``.""" + expected = tuple(range(first_year, first_year + len(years))) + if ( + not years + or len(years) > MAX_BUDGET_WINDOW_YEARS + or tuple(years) != expected + or expected[-1] > LATEST_YEAR + ): + msg = ( + "a budget window's years must be 1 to " + f"{MAX_BUDGET_WINDOW_YEARS} consecutive years from its first, " + f"ending by {LATEST_YEAR}" + ) + raise ValueError(msg) + + def insert_budget_window_report( connection: Connection, report: NewBudgetWindowReport ) -> None: @@ -159,24 +168,42 @@ def insert_budget_window_report( ) connection.execute( insert(report_runs), - [_year_run_row(report, window_year) for window_year in report.years], + [ + year_run_row( + window_year, + report_id=report.id, + window_run_id=report.window_run_id, + versions=report.versions, + trigger=ReportRunTrigger.INITIAL, + ) + for window_year in report.years + ], ) -def _year_run_row( - report: NewBudgetWindowReport, window_year: BudgetWindowYear +def year_run_row( + window_year: BudgetWindowYear | WindowRerunYear, + *, + report_id: UUID, + window_run_id: UUID, + versions: RunVersions, + trigger: ReportRunTrigger, ) -> dict[str, Any]: - """Return a year's report run, pending with no call, as inserted.""" + """Return a year's report run in a window run, pending with no call, as inserted. + + ``window_year`` is a new window's year or a re-run's; ``trigger`` says + which. + """ return { "id": window_year.report_run_id, - "report_id": report.id, + "report_id": report_id, "baseline_simulation_run_id": window_year.baseline.simulation_run_id, "reform_simulation_run_id": window_year.reform.simulation_run_id, - "budget_window_run_id": report.window_run_id, + "budget_window_run_id": window_run_id, "year": window_year.year, - "country_package_version": report.versions.country_package_version, - "policyengine_version": report.versions.policyengine_version, + "country_package_version": versions.country_package_version, + "policyengine_version": versions.policyengine_version, "status": ReportRunStatus.PENDING, - "trigger": ReportRunTrigger.INITIAL, + "trigger": trigger, "modal_call_id": None, } diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py index 091584d5c..d71f3265f 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -10,13 +10,16 @@ - creates, in one transaction, the new run and, for each report side not reusing a run of its own simulation, a new ``simulation_runs`` row of that simulation, all ``pending`` with no Modal call ID. A report's new run has - trigger ``manual``. + trigger ``manual``. A budget window's new run is a ``budget_window_runs`` + row with a report run per year, each year on the simulations its report + run in the report's newest window run names. The record itself is never changed: the service may not update ``reports`` or ``simulations``, so a report keeps naming the simulations it was created with, and each new run records the catalog version, dataset and region it calculates on. Nothing is created while the record has a run that is -``pending`` or ``running``; a transaction-level advisory lock on the +``pending`` or ``running`` (for a report, a budget-window run or a report +run of it); a transaction-level advisory lock on the record's ID makes concurrent re-runs of one record wait for each other, as a row lock on ``reports`` would need ``UPDATE`` on it. @@ -32,6 +35,10 @@ from sqlalchemy import func, insert, literal, select +from policyengine_simulation_persistence.budget_window_records import ( + check_window_years, + year_run_row, +) from policyengine_simulation_persistence.economy_run_records import ( ECONOMY_REPORT_TYPE, EconomyRecordReferenceError, @@ -76,7 +83,8 @@ type RecordKind = Literal["economy", "budget_window", "household", "unsupported"] # A household report's type, as the household reports' service writes it. _HOUSEHOLD_TYPE: Final = "household" -# Both run tables' statuses use these values for a run under way. +# Every run table's statuses (simulation, report and budget-window runs) use +# these values for a run under way. _UNDER_WAY: Final = (SimulationStatus.PENDING.value, SimulationStatus.RUNNING.value) @@ -94,12 +102,12 @@ class WindowYearDefinition: They are the simulations of the runs that year's report run names: a year that reused a run names that run's simulation, which may be another - report's. ``reform_simulation_id`` is None for a year naming no reform. + report's. """ year: int baseline_simulation_id: UUID - reform_simulation_id: UUID | None + reform_simulation_id: UUID @dataclass(frozen=True) @@ -111,7 +119,8 @@ class ReportDefinition: row it named was deleted). ``baseline`` and ``reform`` are None when the report names no such simulation; a budget window's are its first year's. A budget window's ``window_years`` are the years of its newest - budget-window run, earliest first, empty when it has none. + budget-window run whose report run names both sides, earliest first, + empty when it has none. """ id: UUID @@ -205,6 +214,10 @@ class NewWindowRerun: recalculate: bool years: tuple[WindowRerunYear, ...] + def __post_init__(self) -> None: + """Refuse years that aren't 1 to 75 consecutive ones from the first.""" + check_window_years(self.population.year, [year.year for year in self.years]) + def population_in(self, year: int) -> EconomyPopulation: """Return what the window's simulations of ``year`` calculate on.""" return replace(self.population, year=year) @@ -321,9 +334,7 @@ def _window_years( .join( baseline_run, baseline_run.c.id == report_runs.c.baseline_simulation_run_id ) - .outerjoin( - reform_run, reform_run.c.id == report_runs.c.reform_simulation_run_id - ) + .join(reform_run, reform_run.c.id == report_runs.c.reform_simulation_run_id) .where(report_runs.c.budget_window_run_id == newest_window_run) .order_by(report_runs.c.year) ) @@ -528,29 +539,19 @@ def insert_window_rerun(connection: Connection, rerun: NewWindowRerun) -> None: ) connection.execute( insert(report_runs), - [_year_run_row(rerun, window_year) for window_year in rerun.years], + [ + year_run_row( + window_year, + report_id=rerun.report_id, + window_run_id=rerun.window_run_id, + versions=rerun.versions, + trigger=ReportRunTrigger.MANUAL, + ) + for window_year in rerun.years + ], ) -def _year_run_row( - rerun: NewWindowRerun, window_year: WindowRerunYear -) -> dict[str, Any]: - """Return a year's new report run, pending with no call, as inserted.""" - return { - "id": window_year.report_run_id, - "report_id": rerun.report_id, - "baseline_simulation_run_id": window_year.baseline.simulation_run_id, - "reform_simulation_run_id": window_year.reform.simulation_run_id, - "budget_window_run_id": rerun.window_run_id, - "year": window_year.year, - "country_package_version": rerun.versions.country_package_version, - "policyengine_version": rerun.versions.policyengine_version, - "status": ReportRunStatus.PENDING, - "trigger": ReportRunTrigger.MANUAL, - "modal_call_id": None, - } - - def _refuse_while_report_runs(connection: Connection, report_id: UUID) -> None: """Raise RunInProgressError while the report has a run pending or running. diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py index 6b803a89d..4666ed427 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py @@ -66,7 +66,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Sequence + from collections.abc import Callable, Iterable, Sequence from contextlib import AbstractContextManager from uuid import UUID @@ -382,38 +382,30 @@ def read_window_rerun_inputs( """ with self._transaction() as connection: catalog = _catalog_rows(connection, request) - baseline_inputs = _side_inputs( - connection, request, catalog, "baseline", baseline - ) - reform_inputs = _side_inputs(connection, request, catalog, "reform", reform) - reusable_runs: dict[int, tuple[ReusableRun | None, ReusableRun | None]] = {} - for window_year in years: - baseline_run = reform_run = None - if reuse is not None and window_year.reform_simulation_id is not None: - search = _search_own_runs( - connection, - catalog.population(window_year.year), - reuse, - { - "baseline": window_year.baseline_simulation_id, - "reform": window_year.reform_simulation_id, - }, - ) - baseline_run = search("baseline", baseline, None) - reform_run = search( - "reform", - reform, + return _window_inputs( + connection, + request, + catalog, + baseline=baseline, + reform=reform, + searches=( + ( + window_year.year, None - if baseline_run is None - else baseline_run.simulation_run_id, + if reuse is None + else _search_own_runs( + connection, + catalog.population(window_year.year), + reuse, + { + "baseline": window_year.baseline_simulation_id, + "reform": window_year.reform_simulation_id, + }, + ), ) - reusable_runs[window_year.year] = (baseline_run, reform_run) - return BudgetWindowInputs( - catalog=catalog, - baseline=baseline_inputs, - reform=reform_inputs, - reusable_runs=reusable_runs, - ) + for window_year in years + ), + ) def create_window_rerun(self, rerun: NewWindowRerun) -> None: """Create the window's new run and its years' runs; see ``insert_window_rerun``. @@ -455,43 +447,24 @@ def read_budget_window_inputs( """ with self._transaction() as connection: catalog = _catalog_rows(connection, request) - baseline_inputs = _side_inputs( - connection, request, catalog, "baseline", baseline - ) - reform_inputs = _side_inputs(connection, request, catalog, "reform", reform) - reusable_runs: dict[int, tuple[ReusableRun | None, ReusableRun | None]] = {} - for year in years: - baseline_run = reform_run = None - if reuse is not None: - population = catalog.population(year) - baseline_run = find_reusable_run( - connection, - choice=baseline, - population=population, - versions=reuse.versions, - inputs=reuse.inputs, - output_uri_prefix=reuse.output_uri_prefix, - ) - reform_run = find_reusable_run( - connection, - choice=reform, - population=population, - versions=reuse.versions, - inputs=reuse.inputs, - output_uri_prefix=reuse.output_uri_prefix, - excluded_run_id=( - None - if baseline_run is None - else baseline_run.simulation_run_id + return _window_inputs( + connection, + request, + catalog, + baseline=baseline, + reform=reform, + searches=( + ( + year, + None + if reuse is None + else _search_by_definition( + connection, catalog.population(year), reuse ), ) - reusable_runs[year] = (baseline_run, reform_run) - return BudgetWindowInputs( - catalog=catalog, - baseline=baseline_inputs, - reform=reform_inputs, - reusable_runs=reusable_runs, - ) + for year in years + ), + ) def create_budget_window_report(self, report: NewBudgetWindowReport) -> None: """Create the window's report, its budget-window run and each year's runs. @@ -727,6 +700,41 @@ def _report_inputs( # noqa: PLR0913 # the request's rows, both sides and the s return EconomyReportInputs(catalog=catalog, baseline=sides[0], reform=sides[1]) +def _window_inputs( # noqa: PLR0913 # the request's rows, both sides and each year's search + connection: Connection, + request: EconomyRequest, + catalog: EconomyCatalogRows, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + searches: Iterable[tuple[int, ReusableRunSearch | None]], +) -> BudgetWindowInputs: + """Read both sides' values, and each year's reusable runs with its search. + + A year without a search reuses nothing; in each year the reform never + gets the run the baseline reuses. + """ + baseline_inputs = _side_inputs(connection, request, catalog, "baseline", baseline) + reform_inputs = _side_inputs(connection, request, catalog, "reform", reform) + reusable_runs: dict[int, tuple[ReusableRun | None, ReusableRun | None]] = {} + for year, search in searches: + baseline_run = reform_run = None + if search is not None: + baseline_run = search("baseline", baseline, None) + reform_run = search( + "reform", + reform, + None if baseline_run is None else baseline_run.simulation_run_id, + ) + reusable_runs[year] = (baseline_run, reform_run) + return BudgetWindowInputs( + catalog=catalog, + baseline=baseline_inputs, + reform=reform_inputs, + reusable_runs=reusable_runs, + ) + + def _search_by_definition( connection: Connection, population: EconomyPopulation, reuse: RunReuse ) -> ReusableRunSearch: diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 11570676a..d48273250 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -208,8 +208,10 @@ A budget-window report's re-run is a new budget-window run, trigger year's sides are the simulations that year's report run of the report's newest window run names (a year that reused another report's run names that report's simulation), and each reuses a run of them as a single-year side -does. The first five years start as a new window's do; it answers 201 with -the budget-window create's fields, naming the window run and every year. +does. Up to five years start at once, as a new window's do (a later year +Modal doesn't confirm in time waits for an earlier year's coordinator); it +answers 201 with the budget-window create's fields, naming the window run +and every year. `POST /simulations/{simulation_id}/runs`, with no body, creates and starts a stored economy simulation's new run, which is always calculated. Each answers 201 with the same fields as the create, naming the new run. @@ -233,7 +235,10 @@ calculated separately. A simulation's re-run checks every run of the simulation, including those reports created. A report or simulation deleted after the re-run read it answers 404. A report deleted with its simulations answers 400 `missing_reference`, naming a simulation, when a -side gets a new run, and 404 when both sides reuse a run. When Modal +side gets a new run, and 404 when both sides reuse a run; a deleted +budget-window report always answers 404, as its window run is created +first. A 409 `run_in_progress` names the report's budget-window run under +way, or else its report run, or the simulation's run. When Modal refuses the new run's call, the run is failed and the 503 `modal_call_refused` names the record and the run (`run_id`). diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py index c4cf4d59e..efcd23e30 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py @@ -68,6 +68,7 @@ EconomyRequest, ReusableRun, SimulationInputs, + WindowRerunYear, ) from policyengine_stage12_persistence import ( WindowAdvance, @@ -155,7 +156,6 @@ def __init__( # noqa: PLR0913 # the window seams, then the economy service's clock=clock, ) self._windows = windows - self._window_store = store async def submit_budget_window( self, @@ -183,7 +183,7 @@ async def submit_budget_window( ) economy_request = economy_request_for(request, worker) inputs = await self._read_inputs( - lambda: self._window_store.read_budget_window_inputs( + lambda: self._store.read_budget_window_inputs( economy_request, baseline=simulation_choice(request.baseline), reform=simulation_choice(request.reform), @@ -214,7 +214,7 @@ async def submit_budget_window( reform_policy=inputs.reform.reform, ) await self._create( - "report", lambda: self._window_store.create_budget_window_report(report) + "report", lambda: self._store.create_budget_window_report(report) ) status = await self._start_first_years( worker, @@ -229,7 +229,7 @@ async def submit_budget_window( report_run_id=first_year.report_run_id, budget_window_run_id=report.window_run_id, status=status, - years=[_year_created(window_year) for window_year in report.years], + years=[year_created(window_year) for window_year in report.years], ) def _window_year(self, inputs: BudgetWindowInputs, year: int) -> BudgetWindowYear: @@ -419,7 +419,10 @@ def _run_ids(window_year: BudgetWindowYear) -> ReportRunIds: ) -def _year_created(window_year: BudgetWindowYear) -> BudgetWindowYearCreated: +def year_created( + window_year: BudgetWindowYear | WindowRerunYear, +) -> BudgetWindowYearCreated: + """Return a new window's year, or a re-run's, as the answer lists it.""" return BudgetWindowYearCreated( year=window_year.year, report_run_id=window_year.report_run_id, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py index 8f09eebe5..fdfdce51b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -77,7 +77,9 @@ EconomyUnprocessableCode = Literal["missing_parameters", "overlapping_parameters"] # Why a stored record is not re-run now (409). EconomyConflictCode = Literal[ - # The record has a run pending or running; ``run_id`` names it. + # The record has a run pending or running; ``run_id`` names it: for a + # report, its budget-window run, or else its report run; for a + # simulation, its simulation run. "run_in_progress", # The service does not re-run this kind of record, or the record lacks a # row it calculates on. @@ -290,8 +292,10 @@ class EconomyError(_Answer): ``/simulations``) is set only with ``modal_call_refused``: it names the record the service created or re-ran, whose new run it failed because Modal refused the call or didn't accept it in time. ``run_id`` names - that new run on a re-run, and with ``run_in_progress`` the record's run - that is pending or running. + that new run on a re-run (a budget window's: its budget-window run), + and with ``run_in_progress`` the record's run that is pending or + running: a report's budget-window run, or else its report run, or a + simulation's run. """ code: EconomyErrorCode diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py index 4923f8b12..36c43f5d2 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_rerun_routes.py @@ -7,8 +7,9 @@ ``{"force": true}`` calculates both simulations again; otherwise a side reuses a run of its own simulation that succeeded on the same versions. A budget-window report's new run is a budget-window run with a report run - per year, of which the first five are started; it answers as the - budget-window create does, naming the window run and every year's runs. + per year, up to five of them started at once as a new window's are; it + answers as the budget-window create does, naming the window run and every + year's runs. - ``POST /simulations/{simulation_id}/runs`` (``simulation:submit``) creates the stored simulation's new run and starts its simulation function, answering 201 with the simulation's ID, the new run's ID and its status. @@ -103,6 +104,11 @@ class EconomyService(EconomySubmissions, EconomyReruns, Protocol): """Create, read and re-run economy reports and simulations.""" +# The run a 409 run_in_progress names, by the record re-run. +_RUN_UNDER_WAY: dict[RecordName, str] = { + "report": "the report's budget-window run, or else its report run", + "simulation": "the simulation's run", +} # The 400 codes that name a stored record's own rows on its re-run: a # report's re-run also names its stored simulations (missing_reference). _STORED_ROW_CODES: dict[RecordName, str] = { @@ -136,8 +142,8 @@ def _responses(record: RecordName) -> dict[int | str, dict[str, Any]]: 409: { "description": ( f"The {record} is not re-run now, and nothing was written: it " - "has a run pending or running (run_in_progress, named by " - "run_id), or the service does not re-run it " + "has a run pending or running (run_in_progress; run_id names " + f"{_RUN_UNDER_WAY[record]}), or the service does not re-run it " "(rerun_not_supported)" ), **ERROR_ANSWER, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py index 97e5a5878..bc6d386ba 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -52,11 +52,11 @@ from policyengine_simulation_entry.budget_window_api_models import ( BudgetWindowReportCreated, - BudgetWindowYearCreated, ) from policyengine_simulation_entry.budget_window_submissions import ( BudgetWindowRuns, ReportSubmissionService, + year_created, ) from policyengine_simulation_entry.economy_api_models import ( EconomyReportCreated, @@ -388,7 +388,7 @@ async def _rerun_window( window.economy_request, baseline=stored.baseline.choice, reform=stored.reform.choice, - years=[year.definition for year in window.years], + years=window.years, reuse=window.reuse, ), request=window.economy_request, @@ -441,11 +441,11 @@ async def _rerun_window( report_run_id=first_year.report_run_id, budget_window_run_id=window_run_id, status=status, - years=[_rerun_year_created(window_year) for window_year in rerun.years], + years=[year_created(window_year) for window_year in rerun.years], ) def _window_rerun_year( - self, year: _CompleteWindowYear, inputs: BudgetWindowInputs + self, year: WindowYearDefinition, inputs: BudgetWindowInputs ) -> WindowRerunYear: """Return a year naming its own simulations' reusable runs, or new runs.""" baseline_run, reform_run = inputs.reusable_runs[year.year] @@ -582,27 +582,13 @@ def _rerunnable_report(definition: ReportDefinition) -> _StoredReportRows: ) -@dataclass(frozen=True) -class _CompleteWindowYear: - """A year of a stored window's newest calculation, with both sides' simulations.""" - - definition: WindowYearDefinition - baseline_simulation_id: UUID - reform_simulation_id: UUID - - @property - def year(self) -> int: - """The year.""" - return self.definition.year - - @dataclass(frozen=True) class _WindowRerunRequest: """What a budget window's re-run calculates, as the service has read it.""" stored: _StoredReportRows options: dict[str, Any] - years: tuple[_CompleteWindowYear, ...] + years: tuple[WindowYearDefinition, ...] economy_request: EconomyRequest reuse: RunReuse | None recalculate: bool @@ -610,11 +596,12 @@ class _WindowRerunRequest: def _complete_window_years( definition: ReportDefinition, -) -> tuple[_CompleteWindowYear, ...]: +) -> tuple[WindowYearDefinition, ...]: """Return the stored window's years, or refuse a window that lacks one. Its newest calculation must have a report run naming both sides in each - year of its window, from the report's first year. + year of its window, from the report's first year; the definition leaves + out a year naming no reform. """ expected = ( () @@ -623,16 +610,8 @@ def _complete_window_years( range(definition.year, definition.year + definition.budget_window_size) ) ) - complete = tuple( - _CompleteWindowYear( - definition=year, - baseline_simulation_id=year.baseline_simulation_id, - reform_simulation_id=year.reform_simulation_id, - ) - for year in definition.window_years - if year.reform_simulation_id is not None - ) - if not expected or tuple(year.year for year in complete) != expected: + years = definition.window_years + if not expected or tuple(year.year for year in years) != expected: raise EconomyRerunConflictError( code="rerun_not_supported", detail=( @@ -641,20 +620,7 @@ def _complete_window_years( "it cannot be re-run." ), ) - return complete - - -def _rerun_year_created(window_year: WindowRerunYear) -> BudgetWindowYearCreated: - return BudgetWindowYearCreated( - year=window_year.year, - report_run_id=window_year.report_run_id, - baseline_simulation_id=window_year.baseline.simulation_id, - baseline_simulation_run_id=window_year.baseline.simulation_run_id, - baseline_reused=window_year.baseline.reused, - reform_simulation_id=window_year.reform.simulation_id, - reform_simulation_run_id=window_year.reform.simulation_run_id, - reform_reused=window_year.reform.reused, - ) + return years def _rerunnable_simulation(definition: SimulationDefinition) -> _StoredSimulationRows: From 8803608930eadc50475c7735828d2a580d027ed1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:38:42 +0400 Subject: [PATCH 17/18] Fix review round 5's findings: type a window's years by the fields read, and test the insert order The shared year row and year answer take a protocol of the fields they read, so the new-window modules no longer name the re-run's year type. A new test, which passes at once, shows a window re-run of a report deleted with a year's simulation is refused as a deleted report, since its window run is inserted first. The years check, the window years read and the README say what they do in full. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 4 ++ .../budget_window_records.py | 58 +++++++++++++++++-- .../economy_rerun_records.py | 5 +- .../tests/test_economy_window_reruns.py | 7 +++ .../policyengine-simulation-entry/README.md | 6 +- .../budget_window_submissions.py | 6 +- 6 files changed, 72 insertions(+), 14 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 c8d7198bc..3fda391e1 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -9,6 +9,8 @@ from policyengine_simulation_persistence.budget_window_records import ( BudgetWindowYear, NewBudgetWindowReport, + WindowSide, + WindowYearRuns, insert_budget_window_report, ) from policyengine_simulation_persistence.economy_rerun_records import ( @@ -195,7 +197,9 @@ "UnknownPolicyError", "UnknownRegionError", "WindowRerunYear", + "WindowSide", "WindowYearDefinition", + "WindowYearRuns", "bounded_transaction", "budget_window_runs", "combine_reforms", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py index b7ef5293d..9f40b6c40 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/budget_window_records.py @@ -17,7 +17,7 @@ from __future__ import annotations from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol from sqlalchemy import insert @@ -45,9 +45,52 @@ from sqlalchemy import Connection - from policyengine_simulation_persistence.economy_rerun_records import ( - WindowRerunYear, - ) + +class WindowSide(Protocol): + """A window year's side: its simulation and the run the year names.""" + + @property + def simulation_id(self) -> UUID: + """The side's simulation.""" + ... + + @property + def simulation_run_id(self) -> UUID: + """The run of it the year names.""" + ... + + @property + def reused(self) -> bool: + """Whether the run already succeeded, rather than being created.""" + ... + + +class WindowYearRuns(Protocol): + """A year of a window run: its year, its report run and each side. + + A new window's ``BudgetWindowYear`` and a re-run's ``WindowRerunYear`` + are both one. + """ + + @property + def year(self) -> int: + """The year.""" + ... + + @property + def report_run_id(self) -> UUID: + """The year's report run.""" + ... + + @property + def baseline(self) -> WindowSide: + """The baseline side.""" + ... + + @property + def reform(self) -> WindowSide: + """The reform side.""" + ... @dataclass(frozen=True) @@ -96,7 +139,10 @@ def population_in(self, year: int) -> EconomyPopulation: def check_window_years(first_year: int, years: Sequence[int]) -> None: - """Raise ValueError unless ``years`` are 1 to 75 in a row from ``first_year``.""" + """Raise ValueError unless ``years`` are 1 to 75 in a row from ``first_year``. + + The last of them must also be no later than the tables' latest year. + """ expected = tuple(range(first_year, first_year + len(years))) if ( not years @@ -182,7 +228,7 @@ def insert_budget_window_report( def year_run_row( - window_year: BudgetWindowYear | WindowRerunYear, + window_year: WindowYearRuns, *, report_id: UUID, window_run_id: UUID, diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py index d71f3265f..311cbaae2 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_rerun_records.py @@ -313,7 +313,10 @@ def read_report_definition( def _window_years( connection: Connection, report_id: UUID ) -> tuple[WindowYearDefinition, ...]: - """Return the years of the report's newest budget-window run, earliest first.""" + """Return the years of the report's newest budget-window run, earliest first. + + A year whose report run names no reform is left out. + """ newest_window_run = ( select(budget_window_runs.c.id) .where(budget_window_runs.c.report_id == report_id) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py b/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py index 88eb4c22f..6d7619f6b 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_window_reruns.py @@ -562,6 +562,13 @@ def test_a_window_rerun_of_a_deleted_report_or_year_simulation_is_refused( with pytest.raises(RerunRecordDeletedError), connection.begin_nested(): insert_window_rerun(connection, replace(rerun, report_id=uuid4())) + # A report deleted with a year's simulation: the window run is + # inserted first, so the report's own key refuses it. + with pytest.raises(RerunRecordDeletedError), connection.begin_nested(): + insert_window_rerun( + connection, + replace(rerun, report_id=uuid4(), years=(*rerun.years[:2], gone_year)), + ) with ( pytest.raises(EconomyRecordReferenceError) as refused, connection.begin_nested(), diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index d48273250..6b7d8db06 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -219,7 +219,8 @@ answers 201 with the same fields as the create, naming the new run. A re-run is refused before anything is written: 404 (`unknown_report`, `unknown_simulation`) for a record that doesn't exist; 409 `run_in_progress`, naming the run in `run_id`, while the record has a run -pending or running (for a report, a report run or a budget-window run); +pending or running (for a report, its budget-window run, or else its report +run); 409 `rerun_not_supported` for a household record, which this service does not re-run yet, a budget window whose newest window run lacks a year, or a record that names no region, dataset, year, simulation or country the @@ -237,8 +238,7 @@ deleted after the re-run read it answers 404. A report deleted with its simulations answers 400 `missing_reference`, naming a simulation, when a side gets a new run, and 404 when both sides reuse a run; a deleted budget-window report always answers 404, as its window run is created -first. A 409 `run_in_progress` names the report's budget-window run under -way, or else its report run, or the simulation's run. When Modal +first. When Modal refuses the new run's call, the run is failed and the 503 `modal_call_refused` names the record and the run (`run_id`). diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py index efcd23e30..a421ff9f4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py @@ -68,7 +68,7 @@ EconomyRequest, ReusableRun, SimulationInputs, - WindowRerunYear, + WindowYearRuns, ) from policyengine_stage12_persistence import ( WindowAdvance, @@ -419,9 +419,7 @@ def _run_ids(window_year: BudgetWindowYear) -> ReportRunIds: ) -def year_created( - window_year: BudgetWindowYear | WindowRerunYear, -) -> BudgetWindowYearCreated: +def year_created(window_year: WindowYearRuns) -> BudgetWindowYearCreated: """Return a new window's year, or a re-run's, as the answer lists it.""" return BudgetWindowYearCreated( year=window_year.year, From fe8756e350d4aa885b292f908c29021e05ac5849 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 13 Nov 2026 03:44:01 +0400 Subject: [PATCH 18/18] Fix review round 6's finding: build a window year's run IDs in one place Co-Authored-By: Claude Opus 5.5 --- .../budget_window_submissions.py | 5 +++-- .../src/policyengine_simulation_entry/economy_reruns.py | 7 ++----- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py index a421ff9f4..62aa85692 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_submissions.py @@ -209,7 +209,7 @@ async def submit_budget_window( worker, calculation_scope(request, worker, inputs.catalog, spm=options.get("spm")), self._origin(request_id), - run_ids=_run_ids(first_year), + run_ids=year_run_ids(first_year), baseline_policy=inputs.baseline.reform, reform_policy=inputs.reform.reform, ) @@ -411,7 +411,8 @@ def _deferral(year_to_start: YearToStart) -> dict[str, Any]: } -def _run_ids(window_year: BudgetWindowYear) -> ReportRunIds: +def year_run_ids(window_year: WindowYearRuns) -> ReportRunIds: + """Return a window year's report run and its sides' runs, as its call names them.""" return ReportRunIds( report_run_id=window_year.report_run_id, baseline_simulation_run_id=window_year.baseline.simulation_run_id, diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py index bc6d386ba..e9accda41 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_reruns.py @@ -57,6 +57,7 @@ BudgetWindowRuns, ReportSubmissionService, year_created, + year_run_ids, ) from policyengine_simulation_entry.economy_api_models import ( EconomyReportCreated, @@ -414,11 +415,7 @@ async def _rerun_window( options=window.options, ), self._origin(request_id), - run_ids=ReportRunIds( - report_run_id=first_year.report_run_id, - baseline_simulation_run_id=first_year.baseline.simulation_run_id, - reform_simulation_run_id=first_year.reform.simulation_run_id, - ), + run_ids=year_run_ids(first_year), baseline_policy=inputs.baseline.reform, reform_policy=inputs.reform.reform, )