From 21624a557da667a195628c38bd52ed474ea97347 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:09:21 +0400 Subject: [PATCH 01/26] Test reading a stored dynamic's values and combining them with a policy's (failing) A simulation names a dynamic as well as a policy. Its parameter values are read as reform values, with the policy reader's periods, and refused when the running catalog version lacks a parameter they set. A parameter both set is refused. Co-Authored-By: Claude Opus 5.5 --- .../tests/public_api_rows.py | 36 ++++ .../tests/test_dynamic_reform.py | 155 ++++++++++++++++++ 2 files changed, 191 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py diff --git a/libs/policyengine-simulation-persistence/tests/public_api_rows.py b/libs/policyengine-simulation-persistence/tests/public_api_rows.py index 7bb74ec47..5c0908a98 100644 --- a/libs/policyengine-simulation-persistence/tests/public_api_rows.py +++ b/libs/policyengine-simulation-persistence/tests/public_api_rows.py @@ -148,6 +148,42 @@ def insert_policy( return policy_id +def insert_dynamic( + connection: Connection, + catalog: CatalogVersion, + values: Iterable[StoredValue] = (), +) -> UUID: + """Insert a dynamic and the values it sets on ``catalog``'s parameters. + + A dynamic belongs to no country; each value names a parameter of that + catalog version by name. + """ + dynamic_id = uuid4() + connection.execute( + text("INSERT INTO public.dynamics (id, name) VALUES (:id, :name)"), + {"id": dynamic_id, "name": f"dynamic {dynamic_id}"}, + ) + parameter_ids = _parameter_ids(connection, catalog.version_id) + insert_value = text( + "INSERT INTO public.parameter_values " + "(id, parameter_id, value_json, start_date, end_date, dynamic_id) " + "VALUES (:id, :parameter_id, :value, :start_date, :end_date, :dynamic_id)" + ).bindparams(bindparam("value", type_=JSONB)) + for stored in values: + connection.execute( + insert_value, + { + "id": uuid4(), + "parameter_id": parameter_ids[stored.parameter_name], + "value": stored.value, + "start_date": stored.start_date, + "end_date": stored.end_date, + "dynamic_id": dynamic_id, + }, + ) + return dynamic_id + + def insert_household( connection: Connection, *, diff --git a/libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py b/libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py new file mode 100644 index 000000000..0bce6c2ea --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py @@ -0,0 +1,155 @@ +"""A stored dynamic's parameter values, and combining them with a policy's. + +The reads 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. Combining needs no database. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from uuid import uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from policyengine_simulation_persistence import ( + MissingPolicyParametersError, + OverlappingReformError, + UnknownDynamicError, + combine_reforms, + read_dynamic_reform, +) +from public_api_rows import ( + StoredValue, + insert_catalog_version, + insert_dynamic, + insert_model, +) + +START = datetime(2026, 1, 1, tzinfo=UTC) +END = datetime(2027, 12, 31, tzinfo=UTC) +ELASTICITY = "gov.simulation.labor_supply_responses.elasticities.income" +SUBSTITUTION = "gov.simulation.labor_supply_responses.elasticities.substitution.all" +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" + + +@pytest.mark.public_api_schema +def test_a_dynamic_is_read_as_its_values_by_parameter_name_and_period( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = insert_catalog_version( + connection, + country="us", + model_id=insert_model(connection, "us"), + policyengine_version="6.17.1", + parameter_names=[ELASTICITY, SUBSTITUTION], + ) + dynamic_id = insert_dynamic( + connection, + catalog, + [ + StoredValue(ELASTICITY, -0.05, START, END), + StoredValue(SUBSTITUTION, 0.25, START), + ], + ) + act_as_the_service(connection) + + reform = read_dynamic_reform( + connection, dynamic_id=dynamic_id, catalog_version_id=catalog.version_id + ) + + assert reform == { + ELASTICITY: {"2026-01-01.2027-12-31": -0.05}, + SUBSTITUTION: {"2026-01-01.2200-12-31": 0.25}, + } + + +@pytest.mark.public_api_schema +def test_an_unknown_dynamic_is_refused(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = insert_catalog_version( + connection, + country="us", + model_id=insert_model(connection, "us"), + policyengine_version="6.17.1", + ) + act_as_the_service(connection) + missing_dynamic_id = uuid4() + + with pytest.raises(UnknownDynamicError, match=str(missing_dynamic_id)): + read_dynamic_reform( + connection, + dynamic_id=missing_dynamic_id, + catalog_version_id=catalog.version_id, + ) + + +@pytest.mark.public_api_schema +def test_a_dynamic_changing_parameters_the_run_lacks_names_each_of_them( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + model_id = insert_model(connection, "us") + saved_on = insert_catalog_version( + connection, + country="us", + model_id=model_id, + policyengine_version="6.16.0", + parameter_names=[ELASTICITY, SUBSTITUTION], + ) + running = insert_catalog_version( + connection, + country="us", + model_id=model_id, + policyengine_version="6.17.1", + parameter_names=[ELASTICITY], + ) + dynamic_id = insert_dynamic( + connection, + saved_on, + [ + StoredValue(ELASTICITY, -0.05, START), + StoredValue(SUBSTITUTION, 0.25, START), + ], + ) + act_as_the_service(connection) + + with pytest.raises(MissingPolicyParametersError) as refusal: + read_dynamic_reform( + connection, + dynamic_id=dynamic_id, + catalog_version_id=running.version_id, + ) + + assert refusal.value.parameter_names == (SUBSTITUTION,) + assert refusal.value.source == "dynamic" + assert str(refusal.value).startswith("the dynamic sets parameters") + + +def test_a_policy_and_a_dynamic_combine_by_parameter() -> None: + policy = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} + dynamic = {ELASTICITY: {"2026-01-01.2200-12-31": -0.05}} + + assert combine_reforms(policy, dynamic) == { + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + ELASTICITY: {"2026-01-01.2200-12-31": -0.05}, + } + assert combine_reforms(policy, {}) == policy + assert combine_reforms({}, dynamic) == dynamic + + +def test_a_parameter_the_policy_and_the_dynamic_both_set_is_refused() -> None: + policy = { + ELASTICITY: {"2026-01-01.2200-12-31": -0.1}, + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + } + dynamic = { + ELASTICITY: {"2030-01-01.2200-12-31": -0.05}, + SUBSTITUTION: {"2026-01-01.2200-12-31": 0.25}, + } + + with pytest.raises(OverlappingReformError) as refusal: + combine_reforms(policy, dynamic) + + assert refusal.value.parameter_names == (ELASTICITY,) From f4370052ac4dd9c30e852dd400f41ca95074e89b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:10:37 +0400 Subject: [PATCH 02/26] Read a stored dynamic's values, and combine them with a policy's read_dynamic_reform reads a dynamic's parameter values with the policy reader's periods, and refuses one setting parameters the running catalog version lacks (MissingPolicyParametersError, now naming its source). combine_reforms joins a simulation's policy and dynamic values, refusing a parameter both set (OverlappingReformError). The read-only models gain the dynamics table and parameter_values.dynamic_id. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 8 ++ .../input_tables.py | 30 +++-- .../policy_reform.py | 103 ++++++++++++++++-- 3 files changed, 122 insertions(+), 19 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 11f4713b4..dd492b76a 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -32,8 +32,12 @@ from policyengine_simulation_persistence.input_tables import input_metadata from policyengine_simulation_persistence.policy_reform import ( MissingPolicyParametersError, + OverlappingReformError, PolicyReform, + UnknownDynamicError, UnknownPolicyError, + combine_reforms, + read_dynamic_reform, read_policy_reform, ) from policyengine_simulation_persistence.tables import ( @@ -62,6 +66,7 @@ "HouseholdJobStore", "MissingPolicyParametersError", "NewHouseholdJob", + "OverlappingReformError", "PolicyReform", "ReportRunStatus", "ReportRunTrigger", @@ -69,10 +74,12 @@ "SimulationType", "StoredHousehold", "UnknownCatalogVersionError", + "UnknownDynamicError", "UnknownHouseholdError", "UnknownPolicyError", "bounded_transaction", "budget_window_runs", + "combine_reforms", "create_simulation_engine", "fail_unstarted_household_job", "household_jobs", @@ -81,6 +88,7 @@ "missing_table_privileges", "normalized_database_url", "read_catalog_version_id", + "read_dynamic_reform", "read_household_job", "read_missing_table_privileges", "read_policy_reform", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py index 1bde68b74..fa4a7d0fd 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py @@ -1,14 +1,14 @@ """Read-only models of the Public API's tables a calculation is read from. -The Public API writes households, policies and their parameter values, and -its model catalog; the simulation service only reads them, with ``SELECT``. -Each table here names only the columns the service reads, so these are not -models of the whole tables: the Public API's other columns, defaults, -constraints and indexes are its own. The tables live in ``input_metadata``, -apart from the models of the tables the service writes (``service_metadata``) -and from the tables those name by ``id`` alone (``referenced_metadata``). -``tests/test_public_api_schema.py`` compares each column with the Public -API's table. +The Public API writes households, policies, dynamics and their parameter +values, and its model catalog; the simulation service only reads them, with +``SELECT``. Each table here names only the columns the service reads, so these +are not models of the whole tables: the Public API's other columns, +defaults, constraints and indexes are its own. The tables live in +``input_metadata``, apart from the models of the tables the service writes +(``service_metadata``) and from the tables those name by ``id`` alone +(``referenced_metadata``). ``tests/test_public_api_schema.py`` compares each +column with the Public API's table. """ from __future__ import annotations @@ -70,7 +70,8 @@ ) # A parameter's value from ``start_date`` until ``end_date``, or with no end -# date indefinitely; a policy's own values name it in ``policy_id``. +# date indefinitely; a policy's own values name it in ``policy_id``, and a +# dynamic's in ``dynamic_id``. parameter_values = Table( "parameter_values", input_metadata, @@ -79,4 +80,13 @@ Column("start_date", DateTime(timezone=True), nullable=False), Column("end_date", DateTime(timezone=True), nullable=True), Column("policy_id", Uuid, nullable=True), + Column("dynamic_id", Uuid, nullable=True), +) + +# Named parameter values a simulation applies beside its policy, such as +# behavioural responses. A dynamic belongs to no country. +dynamics = Table( + "dynamics", + input_metadata, + Column("id", Uuid, nullable=False), ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py index c87ebc985..e1261160b 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py @@ -1,11 +1,13 @@ -"""Read a stored policy as the reform values a calculation applies. +"""Read a stored policy, or dynamic, as the reform values a calculation applies. The Public API stores a policy's parameter values against the parameters of the catalog version it was saved on, each from a start date until an end date, or with no end date indefinitely. A calculation on any catalog version applies them by parameter name, as ``{parameter name: {period: value}}``, the shape of the Public API's policy JSON that the household functions' contract -and the economy requests take. +and the economy requests take. A dynamic's values are stored and read the same +way. A simulation applies its policy's and its dynamic's values together, and +is refused when both set the same parameter (``combine_reforms``). Each period is a range, ``YYYY-MM-DD.YYYY-MM-DD``, with the dates read in UTC. A value with no end date follows .py's own rule @@ -29,11 +31,12 @@ from __future__ import annotations from datetime import UTC, date, timedelta -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from sqlalchemy import select from policyengine_simulation_persistence.input_tables import ( + dynamics, parameter_values, parameters, policies, @@ -45,31 +48,58 @@ from datetime import datetime from uuid import UUID - from sqlalchemy import Connection + from sqlalchemy import ColumnElement, Connection # The last day a value without an end date is applied. OPEN_ENDED_STOP: Final = date(LATEST_YEAR, 12, 31) # A policy's reform values: each parameter's values by period. type PolicyReform = dict[str, dict[str, object]] +# What stores the values a calculation applies. +type ReformSource = Literal["policy", "dynamic"] class UnknownPolicyError(LookupError): """The Public API has no such policy in the country.""" +class UnknownDynamicError(LookupError): + """The Public API has no such dynamic.""" + + class MissingPolicyParametersError(ValueError): - """The policy sets parameters the calculation's catalog version lacks.""" + """The policy, or dynamic, sets parameters the calculation's catalog version lacks. - def __init__(self, parameter_names: Sequence[str]) -> None: + ``source`` says which of the two sets them. + """ + + def __init__( + self, parameter_names: Sequence[str], *, source: ReformSource = "policy" + ) -> None: """Name every parameter the catalog version lacks, in name order.""" self.parameter_names = tuple(parameter_names) + self.source: ReformSource = source super().__init__( - "the policy sets parameters the catalog version does not have: " + f"the {source} sets parameters the catalog version does not have: " + ", ".join(self.parameter_names) ) +class OverlappingReformError(ValueError): + """A simulation's policy and dynamic both set the same parameters. + + Neither is applied over the other, so the simulation cannot be + calculated. + """ + + def __init__(self, parameter_names: Sequence[str]) -> None: + """Name every parameter both set, in name order.""" + self.parameter_names = tuple(parameter_names) + super().__init__( + "the policy and the dynamic both set: " + ", ".join(self.parameter_names) + ) + + def _period( start_date: datetime, end_date: datetime | None, @@ -118,6 +148,61 @@ def read_policy_reform( if found is None: msg = f"{country} has no policy {policy_id}" raise UnknownPolicyError(msg) + return _read_reform( + connection, + owner=parameter_values.c.policy_id == policy_id, + catalog_version_id=catalog_version_id, + source="policy", + ) + + +def read_dynamic_reform( + connection: Connection, *, dynamic_id: UUID, catalog_version_id: UUID +) -> PolicyReform: + """Return the values the dynamic sets, by parameter name and period. + + The periods follow the same rule as a policy's. Raise + UnknownDynamicError when there is no such dynamic, and + MissingPolicyParametersError, naming the dynamic as its source, when + ``catalog_version_id`` lacks a parameter it sets. + """ + found = connection.execute( + select(dynamics.c.id).where(dynamics.c.id == dynamic_id) + ).scalar_one_or_none() + if found is None: + msg = f"there is no dynamic {dynamic_id}" + raise UnknownDynamicError(msg) + return _read_reform( + connection, + owner=parameter_values.c.dynamic_id == dynamic_id, + catalog_version_id=catalog_version_id, + source="dynamic", + ) + + +def combine_reforms(policy: PolicyReform, dynamic: PolicyReform) -> PolicyReform: + """Return a simulation's policy values and its dynamic's, as one reform. + + Raise OverlappingReformError when both set a parameter, for any period. + """ + overlapping = sorted(set(policy) & set(dynamic)) + if overlapping: + raise OverlappingReformError(overlapping) + return {**policy, **dynamic} + + +def _read_reform( + connection: Connection, + *, + owner: ColumnElement[bool], + catalog_version_id: UUID, + source: ReformSource, +) -> PolicyReform: + """Return the values ``owner`` selects, by parameter name and .py's periods. + + Raise MissingPolicyParametersError, naming ``source``, when + ``catalog_version_id`` lacks a parameter they set. + """ values = connection.execute( select( parameters.c.name, @@ -126,7 +211,7 @@ def read_policy_reform( parameter_values.c.end_date, ) .join(parameters, parameters.c.id == parameter_values.c.parameter_id) - .where(parameter_values.c.policy_id == policy_id) + .where(owner) .order_by(parameters.c.name, parameter_values.c.start_date) ).all() names = sorted({row.name for row in values}) @@ -140,7 +225,7 @@ def read_policy_reform( ) missing = [name for name in names if name not in available] if missing: - raise MissingPolicyParametersError(missing) + raise MissingPolicyParametersError(missing, source=source) reform: PolicyReform = {} # The rows are ordered by parameter name, then start date, the order # .py takes them in; each value's periods are .py's. From f86250f100959e5852ae83572e752393a6b0db42 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:14:32 +0400 Subject: [PATCH 03/26] Test the economy simulation and report records, and the choice of a run to reuse (failing) The service reads the country's model, a region and a dataset of the catalog version, and creates a report, its new simulations and runs, and its first report run in one transaction, or a simulation and its first run; a foreign-key refusal names the column. A report side may name a run that already succeeded instead. It records a Modal call's ID, or fails the run and its new simulation runs, only while the run is pending with no call ID. A run is reused only when it is the newest succeeded run of the same simulation on the same versions, its file is in the output bucket, a succeeded report run with the same options used it, and no report run refused to reuse it. Co-Authored-By: Claude Opus 5.5 --- .../tests/economy_rows.py | 237 +++++++ .../tests/public_api_rows.py | 49 ++ .../tests/test_economy_run_records.py | 602 ++++++++++++++++++ .../tests/test_economy_run_reuse.py | 307 +++++++++ 4 files changed, 1195 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/tests/economy_rows.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_economy_run_records.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py diff --git a/libs/policyengine-simulation-persistence/tests/economy_rows.py b/libs/policyengine-simulation-persistence/tests/economy_rows.py new file mode 100644 index 000000000..4cb6a9014 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/economy_rows.py @@ -0,0 +1,237 @@ +"""Rows for the economy simulation and report tests. + +The catalog rows a request names are the Public API's, inserted as the +superuser. A run that already succeeded, and the report run that used it, are +written as the superuser too, through the service tables' models: the +service's role can neither mark a run succeeded nor record its output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from policyengine_simulation_persistence import ( + EconomyPopulation, + ReportRunStatus, + ReportRunTrigger, + RunVersions, + SimulationChoice, + SimulationStatus, + SimulationType, + report_runs, + reports, + simulation_runs, + simulations, +) +from public_api_rows import ( + CatalogVersion, + insert_catalog_version, + insert_dataset, + insert_model, + insert_region, +) +from sqlalchemy import insert + +if TYPE_CHECKING: + from sqlalchemy import Connection + +POLICYENGINE_VERSION = "6.17.1" +COUNTRY_PACKAGE_VERSION = "1.764.6" +VERSIONS = RunVersions( + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version=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} +COMPLETED_AT = datetime(2026, 11, 1, tzinfo=UTC) + + +@dataclass(frozen=True) +class EconomyCatalog: + """A country's catalog version, with one dataset and one region using it.""" + + catalog: CatalogVersion + dataset_id: UUID + region_id: UUID + + def population(self, *, year: int = 2026) -> EconomyPopulation: + """Return the population a request for the region and year uses.""" + return EconomyPopulation( + tax_benefit_model_version_id=self.catalog.version_id, + region_id=self.region_id, + dataset_id=self.dataset_id, + year=year, + ) + + +def insert_economy_catalog( + connection: Connection, + *, + country: str = "us", + parameter_names: tuple[str, ...] = (), +) -> EconomyCatalog: + """Insert the country's model, a catalog version, a dataset and a region.""" + catalog = insert_catalog_version( + connection, + country=country, + model_id=insert_model(connection, country), + 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 + ) + return EconomyCatalog(catalog=catalog, dataset_id=dataset_id, region_id=region_id) + + +@dataclass(frozen=True) +class StoredRun: + """A simulation and one of its runs, as stored.""" + + simulation_id: UUID + simulation_run_id: UUID + + +@dataclass(frozen=True) +class RunShape: + """What a stored economy simulation and its run hold. + + A succeeded run records an output file in the output bucket unless + ``output_uri`` names another. + """ + + choice: SimulationChoice = CURRENT_LAW + status: SimulationStatus = SimulationStatus.SUCCEEDED + versions: RunVersions = VERSIONS + output_uri: str | None = None + completed_at: datetime = COMPLETED_AT + + +def insert_simulation_run( + connection: Connection, + population: EconomyPopulation, + shape: RunShape | None = None, +) -> StoredRun: + """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, + simulation_type=SimulationType.ECONOMY, + dataset_id=population.dataset_id, + policy_id=shape.choice.policy_id, + dynamic_id=shape.choice.dynamic_id, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + region_id=population.region_id, + year=population.year, + ) + ) + connection.execute( + insert(simulation_runs).values( + id=run_id, + simulation_id=simulation_id, + year=population.year, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + country_package_version=shape.versions.country_package_version, + policyengine_version=shape.versions.policyengine_version, + dataset_id=population.dataset_id, + region_id=population.region_id, + status=shape.status, + completed_at=shape.completed_at if finished else None, + output_uri=( + (shape.output_uri or f"{OUTPUT_BUCKET_PREFIX}outputs/{run_id}.parquet") + if succeeded + else None + ), + output_sha256="d" * 64 if succeeded else None, + ) + ) + return StoredRun(simulation_id=simulation_id, simulation_run_id=run_id) + + +@dataclass(frozen=True) +class ReportRunOutcome: + """How a stored report run ended, and the options its report has.""" + + status: ReportRunStatus = ReportRunStatus.SUCCEEDED + inputs: dict[str, Any] | None = None + error_message: str | None = None + + +def insert_report_run_naming( + connection: Connection, + economy: EconomyCatalog, + compared: tuple[StoredRun, StoredRun], + outcome: ReportRunOutcome | None = None, +) -> UUID: + """Insert a report comparing the two runs, and a run of it, as the superuser. + + ``compared`` holds the baseline's run, then the reform's. Return the + report run's ID. + """ + outcome = outcome or ReportRunOutcome() + baseline, reform = compared + report_id = uuid4() + run_id = uuid4() + connection.execute( + insert(reports).values( + id=report_id, + label="An earlier report", + country=economy.catalog.country, + type="economy", + tax_benefit_model_id=economy.catalog.model_id, + baseline_simulation_id=baseline.simulation_id, + reform_simulation_id=reform.simulation_id, + dataset_id=economy.dataset_id, + region_id=economy.region_id, + year=2026, + inputs=OPTIONS if outcome.inputs is None else outcome.inputs, + ) + ) + finished = outcome.status in {ReportRunStatus.SUCCEEDED, ReportRunStatus.FAILED} + connection.execute( + insert(report_runs).values( + id=run_id, + report_id=report_id, + baseline_simulation_run_id=baseline.simulation_run_id, + reform_simulation_run_id=reform.simulation_run_id, + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version=POLICYENGINE_VERSION, + status=outcome.status, + trigger=ReportRunTrigger.INITIAL, + completed_at=COMPLETED_AT if finished else None, + error_message=outcome.error_message, + ) + ) + return run_id + + +def insert_reusable_run( + connection: Connection, + economy: EconomyCatalog, + shape: RunShape | None = None, + *, + inputs: dict[str, Any] | None = None, +) -> StoredRun: + """Insert a run shaped ``shape`` that a succeeded report run used, with ``inputs``. + + The report run's other simulation run failed, so it is never reusable. + """ + population = economy.population() + reused = insert_simulation_run(connection, population, shape) + other_side = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming( + connection, economy, (reused, other_side), ReportRunOutcome(inputs=inputs) + ) + return reused diff --git a/libs/policyengine-simulation-persistence/tests/public_api_rows.py b/libs/policyengine-simulation-persistence/tests/public_api_rows.py index 5c0908a98..37bfaf22b 100644 --- a/libs/policyengine-simulation-persistence/tests/public_api_rows.py +++ b/libs/policyengine-simulation-persistence/tests/public_api_rows.py @@ -184,6 +184,55 @@ def insert_dynamic( return dynamic_id +def insert_dataset( + connection: Connection, catalog: CatalogVersion, name: str, *, year: int = 2024 +) -> UUID: + """Insert a dataset of ``catalog`` with this name.""" + dataset_id = uuid4() + connection.execute( + text( + "INSERT INTO public.datasets (id, name, year, is_output_dataset, " + "tax_benefit_model_version_id) " + "VALUES (:id, :name, :year, false, :version_id)" + ), + { + "id": dataset_id, + "name": name, + "year": year, + "version_id": catalog.version_id, + }, + ) + return dataset_id + + +def insert_region( + connection: Connection, + catalog: CatalogVersion, + code: str, + *, + default_dataset_id: UUID, + region_type: str = "national", +) -> UUID: + """Insert a region of ``catalog`` with this code and default dataset.""" + region_id = uuid4() + connection.execute( + text( + "INSERT INTO public.regions (id, code, label, region_type, " + "requires_filter, default_dataset_id, tax_benefit_model_version_id) " + "VALUES (:id, :code, :code, CAST(:region_type AS v2_region_type), " + "false, :dataset_id, :version_id)" + ), + { + "id": region_id, + "code": code, + "region_type": region_type, + "dataset_id": default_dataset_id, + "version_id": catalog.version_id, + }, + ) + return region_id + + def insert_household( connection: Connection, *, diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_records.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_records.py new file mode 100644 index 000000000..1ddf50262 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_records.py @@ -0,0 +1,602 @@ +"""The simulation service's reads and writes for economy simulations and reports. + +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. The Public API's rows a request names +are inserted first, as the superuser. +""" + +from __future__ import annotations + +from dataclasses import 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, + VERSIONS, + EconomyCatalog, + insert_economy_catalog, + insert_reusable_run, +) +from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, + NewEconomyReport, + NewEconomySimulation, + ReportRunStatus, + ReportSimulation, + SimulationChoice, + SimulationStatus, + fail_unstarted_report_run, + fail_unstarted_simulation_run, + insert_economy_report, + insert_economy_simulation, + read_dataset, + read_dataset_by_id, + read_model_id, + read_region, + record_report_run_call, + record_simulation_run_call, + report_runs, + simulation_runs, +) +from public_api_rows import insert_dataset, insert_policy, insert_region +from sqlalchemy import select, text, update + +if TYPE_CHECKING: + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +REPORT_OPTIONS = {"include_cliffs": True, "spm": {"geography_kind": "national"}} + + +def _new_side(choice: SimulationChoice = CURRENT_LAW) -> ReportSimulation: + return ReportSimulation( + simulation_id=uuid4(), + simulation_run_id=uuid4(), + choice=choice, + reused=False, + ) + + +def _new_report( + economy: EconomyCatalog, + *, + baseline: ReportSimulation | None = None, + reform: ReportSimulation | None = None, +) -> NewEconomyReport: + return NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country=economy.catalog.country, + label="Child tax credit", + description="A larger credit", + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=REPORT_OPTIONS, + baseline=baseline or _new_side(), + reform=reform or _new_side(), + ) + + +def _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(connection: Connection, table_name: str) -> int: + return connection.execute( + text(f"SELECT count(*) FROM public.{table_name}") # noqa: S608 # fixed table names + ).scalar_one() + + +def test_the_countrys_model_is_found_by_its_name(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + + assert read_model_id(connection, country="us") == economy.catalog.model_id + assert read_model_id(connection, country="uk") is None + + +def test_a_region_is_found_by_code_in_its_catalog_version_only( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + other = insert_economy_catalog(connection, country="uk") + state_id = insert_region( + connection, + economy.catalog, + "state/ca", + default_dataset_id=economy.dataset_id, + region_type="state", + ) + act_as_the_service(connection) + + region = read_region( + connection, + catalog_version_id=economy.catalog.version_id, + region_code="state/ca", + ) + assert ( + read_region( + connection, + catalog_version_id=other.catalog.version_id, + region_code="state/ca", + ) + is None + ) + + assert region is not None + assert (region.id, region.code, region.default_dataset_id) == ( + state_id, + "state/ca", + economy.dataset_id, + ) + + +def test_a_dataset_is_found_by_name_or_id_in_its_catalog_version_only( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + other = insert_economy_catalog(connection, country="uk") + local_id = insert_dataset(connection, economy.catalog, "acs_local_2024") + act_as_the_service(connection) + version_id = economy.catalog.version_id + + by_name = read_dataset( + connection, catalog_version_id=version_id, name="acs_local_2024" + ) + by_id = read_dataset_by_id( + connection, catalog_version_id=version_id, dataset_id=economy.dataset_id + ) + other_version = other.catalog.version_id + assert ( + read_dataset( + connection, catalog_version_id=other_version, name="acs_local_2024" + ) + is None + ) + assert ( + read_dataset_by_id( + connection, catalog_version_id=other_version, dataset_id=local_id + ) + is None + ) + + assert by_name is not None + assert (by_name.id, by_name.name) == (local_id, "acs_local_2024") + assert by_id is not None + assert (by_id.id, by_id.name) == (economy.dataset_id, "populace_us_2024") + + +def test_a_report_its_simulations_and_runs_are_created_pending( + 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) + report = _new_report( + economy, reform=_new_side(SimulationChoice(policy_id, dynamic_id=None)) + ) + act_as_the_service(connection) + + insert_economy_report(connection, report) + stored_report = _row(connection, "reports", report.id) + stored_run = _row(connection, "report_runs", report.run_id) + baseline = _row(connection, "simulations", report.baseline.simulation_id) + reform = _row(connection, "simulations", report.reform.simulation_id) + baseline_run = _row( + connection, "simulation_runs", report.baseline.simulation_run_id + ) + reform_run = _row( + connection, "simulation_runs", report.reform.simulation_run_id + ) + + population = economy.population() + assert { + key: stored_report[key] + for key in ( + "label", + "description", + "country", + "type", + "user_id", + "tax_benefit_model_id", + "policy_id", + "baseline_simulation_id", + "reform_simulation_id", + "household_id", + "dataset_id", + "region_id", + "year", + "inputs", + "budget_window_size", + ) + } == { + "label": "Child tax credit", + "description": "A larger credit", + "country": "us", + "type": "economy", + "user_id": None, + "tax_benefit_model_id": economy.catalog.model_id, + "policy_id": policy_id, + "baseline_simulation_id": report.baseline.simulation_id, + "reform_simulation_id": report.reform.simulation_id, + "household_id": None, + "dataset_id": economy.dataset_id, + "region_id": economy.region_id, + "year": 2026, + "inputs": REPORT_OPTIONS, + "budget_window_size": None, + } + assert { + key: stored_run[key] + for key in ( + "report_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "budget_window_run_id", + "year", + "country_package_version", + "policyengine_version", + "status", + "trigger", + "modal_call_id", + "started_at", + "completed_at", + ) + } == { + "report_id": report.id, + "baseline_simulation_run_id": report.baseline.simulation_run_id, + "reform_simulation_run_id": report.reform.simulation_run_id, + "budget_window_run_id": None, + "year": None, + "country_package_version": VERSIONS.country_package_version, + "policyengine_version": VERSIONS.policyengine_version, + "status": "pending", + "trigger": "initial", + "modal_call_id": None, + "started_at": None, + "completed_at": None, + } + for simulation, run, expected_policy in ( + (baseline, baseline_run, None), + (reform, reform_run, policy_id), + ): + assert ( + simulation["simulation_type"], + simulation["policy_id"], + simulation["dynamic_id"], + simulation["household_id"], + simulation["dataset_id"], + simulation["region_id"], + simulation["tax_benefit_model_version_id"], + simulation["year"], + simulation["filter_field"], + ) == ( + "economy", + expected_policy, + None, + None, + population.dataset_id, + population.region_id, + population.tax_benefit_model_version_id, + 2026, + None, + ) + assert ( + run["simulation_id"], + run["year"], + run["tax_benefit_model_version_id"], + run["country_package_version"], + run["policyengine_version"], + run["dataset_id"], + run["region_id"], + run["status"], + run["modal_call_id"], + run["output_uri"], + ) == ( + simulation["id"], + 2026, + population.tax_benefit_model_version_id, + VERSIONS.country_package_version, + VERSIONS.policyengine_version, + population.dataset_id, + population.region_id, + "pending", + None, + None, + ) + + +def test_a_side_reusing_a_succeeded_run_creates_nothing_for_that_side( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused = insert_reusable_run(connection, economy) + report = _new_report( + economy, + baseline=ReportSimulation( + simulation_id=reused.simulation_id, + simulation_run_id=reused.simulation_run_id, + choice=CURRENT_LAW, + reused=True, + ), + ) + simulations_before = _count(connection, "simulations") + runs_before = _count(connection, "simulation_runs") + act_as_the_service(connection) + + insert_economy_report(connection, report) + stored_report = _row(connection, "reports", report.id) + stored_run = _row(connection, "report_runs", report.run_id) + simulations_after = _count(connection, "simulations") + runs_after = _count(connection, "simulation_runs") + + assert stored_report["baseline_simulation_id"] == reused.simulation_id + assert stored_run["baseline_simulation_run_id"] == reused.simulation_run_id + assert stored_run["reform_simulation_run_id"] == report.reform.simulation_run_id + assert (simulations_after, runs_after) == ( + simulations_before + 1, + runs_before + 1, + ) + + +@pytest.mark.parametrize( + ("missing_column", "table_name"), + [("region_id", "simulations"), ("dataset_id", "simulations")], +) +def test_a_report_naming_a_missing_row_names_the_column_that_refused_it( + public_api_database_url, missing_column: str, table_name: str +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + population = replace(report.population, **{missing_column: uuid4()}) + act_as_the_service(connection) + + with pytest.raises(EconomyRecordReferenceError) as refusal: + insert_economy_report(connection, replace(report, population=population)) + + assert (refusal.value.table_name, refusal.value.column_name) == ( + table_name, + missing_column, + ) + + +def test_a_report_naming_a_missing_dynamic_names_its_column( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + report = _new_report( + economy, reform=_new_side(SimulationChoice(None, dynamic_id=uuid4())) + ) + act_as_the_service(connection) + + with pytest.raises(EconomyRecordReferenceError) as refusal: + insert_economy_report(connection, report) + + assert (refusal.value.table_name, refusal.value.column_name) == ( + "simulations", + "dynamic_id", + ) + + +def test_a_simulation_and_its_run_are_created_pending(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) + simulation = NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=SimulationChoice(policy_id, dynamic_id=None), + population=economy.population(year=2027), + versions=VERSIONS, + ) + act_as_the_service(connection) + + insert_economy_simulation(connection, simulation) + stored = _row(connection, "simulations", simulation.id) + run = _row(connection, "simulation_runs", simulation.run_id) + + assert ( + stored["simulation_type"], + stored["policy_id"], + stored["dynamic_id"], + stored["dataset_id"], + stored["region_id"], + stored["year"], + ) == ("economy", policy_id, None, economy.dataset_id, economy.region_id, 2027) + assert ( + run["simulation_id"], + run["year"], + run["status"], + run["modal_call_id"], + run["policyengine_version"], + ) == (simulation.id, 2027, "pending", None, VERSIONS.policyengine_version) + + +def _created_report(connection: Connection) -> NewEconomyReport: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + act_as_the_service(connection) + insert_economy_report(connection, report) + return report + + +def test_a_report_runs_call_is_recorded_once_while_it_is_pending_without_one( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + report = _created_report(connection) + + first = record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-first" + ) + second = record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-second" + ) + unknown = record_report_run_call( + connection, report_run_id=uuid4(), modal_call_id="fc-other" + ) + call_id = _row(connection, "report_runs", report.run_id)["modal_call_id"] + + assert (first, second, unknown) == ( + ReportRunStatus.PENDING, + ReportRunStatus.PENDING, + None, + ) + assert call_id == "fc-first" + + +def test_a_claimed_report_run_keeps_its_claim(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + report = _created_report(connection) + connection.execute(text("RESET ROLE")) + connection.execute( + update(report_runs) + .where(report_runs.c.id == report.run_id) + .values(status=ReportRunStatus.RUNNING, modal_call_id="fc-coordinator") + ) + act_as_the_service(connection) + + status = record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-late" + ) + failed = fail_unstarted_report_run( + connection, + report_run_id=report.run_id, + simulation_run_ids=(report.baseline.simulation_run_id,), + error_message="modal_call_refused: Modal refused the call.", + ) + stored = _row(connection, "report_runs", report.run_id) + baseline = _row( + connection, "simulation_runs", report.baseline.simulation_run_id + ) + + assert status is ReportRunStatus.RUNNING + assert failed is False + assert (stored["status"], stored["modal_call_id"]) == ("running", "fc-coordinator") + assert baseline["status"] == "pending" + + +def test_an_unstarted_report_run_and_its_new_simulation_runs_fail_at_once( + public_api_database_url, +) -> None: + message = "modal_call_refused: Modal refused the coordinator's call." + with rolled_back_transaction(public_api_database_url) as connection: + report = _created_report(connection) + + failed = fail_unstarted_report_run( + connection, + report_run_id=report.run_id, + simulation_run_ids=( + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ), + error_message=message, + ) + again = fail_unstarted_report_run( + connection, + report_run_id=report.run_id, + simulation_run_ids=(), + error_message="a second failure", + ) + stored = _row(connection, "report_runs", report.run_id) + stored_simulation_runs = connection.execute( + select( + simulation_runs.c.status, + simulation_runs.c.error_message, + simulation_runs.c.completed_at, + ).where( + simulation_runs.c.id.in_( + [ + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ] + ) + ) + ).all() + + assert (failed, again) == (True, False) + assert (stored["status"], stored["error_message"]) == ("failed", message) + assert stored["completed_at"] is not None + assert [ + (row.status, row.error_message, row.completed_at is not None) + for row in stored_simulation_runs + ] == [("failed", message, True)] * 2 + + +def test_a_simulation_runs_call_is_recorded_or_the_run_failed_while_unstarted( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + recorded = NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + ) + refused = NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + ) + act_as_the_service(connection) + insert_economy_simulation(connection, recorded) + insert_economy_simulation(connection, refused) + + status = record_simulation_run_call( + connection, simulation_run_id=recorded.run_id, modal_call_id="fc-worker" + ) + not_failed = fail_unstarted_simulation_run( + connection, + simulation_run_id=recorded.run_id, + error_message="modal_call_refused: refused", + ) + failed = fail_unstarted_simulation_run( + connection, + simulation_run_id=refused.run_id, + error_message="modal_call_refused: refused", + ) + not_recorded = record_simulation_run_call( + connection, simulation_run_id=refused.run_id, modal_call_id="fc-late" + ) + unknown = record_simulation_run_call( + connection, simulation_run_id=uuid4(), modal_call_id="fc-other" + ) + recorded_run = _row(connection, "simulation_runs", recorded.run_id) + refused_run = _row(connection, "simulation_runs", refused.run_id) + + assert (status, not_failed, failed, not_recorded, unknown) == ( + SimulationStatus.PENDING, + False, + True, + SimulationStatus.FAILED, + None, + ) + assert (recorded_run["status"], recorded_run["modal_call_id"]) == ( + "pending", + "fc-worker", + ) + assert ( + refused_run["status"], + refused_run["modal_call_id"], + refused_run["error_message"], + ) == ("failed", None, "modal_call_refused: refused") diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py new file mode 100644 index 000000000..ee750e10d --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py @@ -0,0 +1,307 @@ +"""Which succeeded simulation run a new report reuses, if any. + +A report reuses the newest run of a simulation with the same policy, dynamic, +catalog version, dataset, region and year, that succeeded on the same country +package and .py versions, whose output file is in the output bucket, and that +a succeeded report run with the same options used. A run any report run +refused to reuse is never offered again. + +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 dataclasses import dataclass, field, replace +from datetime import timedelta +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + COMPLETED_AT, + CURRENT_LAW, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + VERSIONS, + EconomyCatalog, + ReportRunOutcome, + RunShape, + StoredRun, + insert_economy_catalog, + insert_report_run_naming, + insert_reusable_run, + insert_simulation_run, +) +from policyengine_simulation_persistence import ( + ReportRunStatus, + ReusableRun, + RunVersions, + SimulationChoice, + SimulationStatus, + find_reusable_run, +) +from public_api_rows import insert_dataset, insert_dynamic, insert_policy, insert_region + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + + +@dataclass(frozen=True) +class Request: + """What a new report's simulation asks for, as the reuse search sees it.""" + + choice: SimulationChoice = CURRENT_LAW + year: int = 2026 + versions: RunVersions = VERSIONS + inputs: dict[str, Any] = field(default_factory=lambda: dict(OPTIONS)) + excluded_run_id: UUID | None = None + + +def _find( + connection: Connection, economy: EconomyCatalog, request: Request | None = None +) -> ReusableRun | None: + request = request or Request() + return find_reusable_run( + connection, + choice=request.choice, + population=economy.population(year=request.year), + versions=request.versions, + inputs=request.inputs, + output_uri_prefix=OUTPUT_BUCKET_PREFIX, + excluded_run_id=request.excluded_run_id, + ) + + +def _reusable(run: StoredRun) -> ReusableRun: + return ReusableRun( + simulation_id=run.simulation_id, simulation_run_id=run.simulation_run_id + ) + + +def test_a_succeeded_run_a_report_used_with_the_same_options_is_reused( + 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) + dynamic_id = insert_dynamic(connection, economy.catalog) + choice = SimulationChoice(policy_id=policy_id, dynamic_id=dynamic_id) + reused = insert_reusable_run(connection, economy, RunShape(choice=choice)) + act_as_the_service(connection) + + found = _find(connection, economy, Request(choice=choice)) + + assert found == _reusable(reused) + + +def _another_policy(connection: Connection, economy: EconomyCatalog) -> dict: + return { + "choice": SimulationChoice(insert_policy(connection, economy.catalog), None) + } + + +def _another_dynamic(connection: Connection, economy: EconomyCatalog) -> dict: + return { + "choice": SimulationChoice(None, insert_dynamic(connection, economy.catalog)) + } + + +# Each makes the request differ from the stored run in one way. +DIFFERENCES: dict[str, Callable[[Connection, EconomyCatalog], dict[str, Any]]] = { + "another policy": _another_policy, + "another dynamic": _another_dynamic, + "another year": lambda _connection, _economy: {"year": 2027}, + "another .py version": lambda _connection, _economy: { + "versions": replace(VERSIONS, policyengine_version="6.18.0") + }, + "another country package version": lambda _connection, _economy: { + "versions": replace(VERSIONS, country_package_version="1.765.0") + }, + "other options": lambda _connection, _economy: {"inputs": {"include_cliffs": True}}, +} + + +@pytest.mark.parametrize("difference", sorted(DIFFERENCES)) +def test_a_run_differing_from_the_request_is_not_reused( + public_api_database_url, difference: str +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run(connection, economy) + request = DIFFERENCES[difference](connection, economy) + act_as_the_service(connection) + + assert _find(connection, economy, Request(**request)) is None + + +def test_a_run_of_another_region_or_dataset_is_not_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run(connection, economy) + other_dataset_id = insert_dataset(connection, economy.catalog, "acs_2024") + other_region_id = insert_region( + connection, + economy.catalog, + "state/ca", + default_dataset_id=economy.dataset_id, + region_type="state", + ) + act_as_the_service(connection) + + for other in ( + replace(economy, dataset_id=other_dataset_id), + replace(economy, region_id=other_region_id), + ): + assert _find(connection, other) is None + + +@pytest.mark.parametrize( + "status", + [SimulationStatus.PENDING, SimulationStatus.RUNNING, SimulationStatus.FAILED], +) +def test_an_unfinished_or_failed_run_is_not_reused( + public_api_database_url, status: SimulationStatus +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run(connection, economy, RunShape(status=status)) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +def test_a_run_whose_file_is_outside_the_output_bucket_is_not_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run( + connection, + economy, + RunShape( + output_uri="gs://policyengine-stage12-staging-artifacts/runs/output.parquet" + ), + ) + # A bucket whose name starts with the output bucket's is another bucket. + insert_reusable_run( + connection, + economy, + RunShape( + output_uri=OUTPUT_BUCKET_PREFIX.rstrip("/") + "-old/output.parquet" + ), + ) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +@pytest.mark.parametrize( + "report_run_status", + [ReportRunStatus.PENDING, ReportRunStatus.RUNNING, ReportRunStatus.FAILED], +) +def test_a_run_no_succeeded_report_run_used_is_not_reused( + public_api_database_url, report_run_status: ReportRunStatus +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + population = economy.population() + stored = insert_simulation_run(connection, population) + other_side = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming( + connection, + economy, + (stored, other_side), + ReportRunOutcome(status=report_run_status), + ) + # A run no report run names at all, such as a simulation's own. + insert_simulation_run(connection, population) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +@pytest.mark.parametrize("refused_as", ["baseline", "reform"]) +def test_a_run_a_report_run_refused_to_reuse_is_not_reused_again( + public_api_database_url, refused_as: str +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused = insert_reusable_run(connection, economy) + other_side = insert_simulation_run( + connection, economy.population(), RunShape(status=SimulationStatus.FAILED) + ) + baseline, reform = ( + (reused, other_side) if refused_as == "baseline" else (other_side, reused) + ) + insert_report_run_naming( + connection, + economy, + (baseline, reform), + ReportRunOutcome( + status=ReportRunStatus.FAILED, + error_message=( + "reused_simulation_output_incompatible: a planned column is missing" + ), + ), + ) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +def test_a_run_whose_report_run_failed_for_another_reason_is_still_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused = insert_reusable_run(connection, economy) + other_side = insert_simulation_run( + connection, economy.population(), RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming( + connection, + economy, + (reused, other_side), + ReportRunOutcome( + status=ReportRunStatus.FAILED, + error_message=( + "simulation_failed: the reform's worker ran out of memory" + ), + ), + ) + act_as_the_service(connection) + + assert _find(connection, economy) == _reusable(reused) + + +def test_the_newest_matching_run_is_reused_unless_excluded( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + older = insert_reusable_run(connection, economy) + newer = insert_reusable_run( + connection, economy, RunShape(completed_at=COMPLETED_AT + timedelta(days=1)) + ) + act_as_the_service(connection) + + newest = _find(connection, economy) + without_newest = _find( + connection, economy, Request(excluded_run_id=newer.simulation_run_id) + ) + without_any = _find(connection, economy, Request(excluded_run_id=uuid4())) + + assert newest == _reusable(newer) + assert without_newest == _reusable(older) + assert without_any == _reusable(newer) From b1dd6971b0c1832ba1bb59d1b8f86c779f858b9d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:16:20 +0400 Subject: [PATCH 04/26] Create economy simulations, reports and their runs, and choose a run to reuse economy_run_records reads the country's model and a catalog version's region and dataset, and creates in one transaction a report, its first run and each side's new simulation and pending run, or a standalone simulation and its first run. A report side may name a succeeded run instead (find_reusable_run): the newest of the same simulation on the same versions, with its file in the output bucket, used by a succeeded report run with the same options, and never refused for reuse. A foreign-key refusal is EconomyRecordReferenceError, naming the table and column. A Modal call's ID is recorded, or the run failed (with the report run's new simulation runs), only while the run is pending with no call ID. The read-only models gain the catalog's regions and datasets columns. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 44 ++ .../economy_run_records.py | 600 ++++++++++++++++++ .../input_tables.py | 20 + 3 files changed, 664 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_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 dd492b76a..71a21952e 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,29 @@ missing_table_privileges, read_missing_table_privileges, ) +from policyengine_simulation_persistence.economy_run_records import ( + CatalogDataset, + CatalogRegion, + EconomyPopulation, + EconomyRecordReferenceError, + NewEconomyReport, + NewEconomySimulation, + ReportSimulation, + ReusableRun, + RunVersions, + SimulationChoice, + fail_unstarted_report_run, + fail_unstarted_simulation_run, + find_reusable_run, + insert_economy_report, + insert_economy_simulation, + read_dataset, + read_dataset_by_id, + read_model_id, + read_region, + record_report_run_call, + record_simulation_run_call, +) from policyengine_simulation_persistence.engine import ( bounded_transaction, create_simulation_engine, @@ -59,17 +82,27 @@ __all__ = [ "SERVICE_TABLE_PRIVILEGES", "SERVICE_UPDATABLE_COLUMNS", + "CatalogDataset", + "CatalogRegion", + "EconomyPopulation", + "EconomyRecordReferenceError", "HouseholdCalculationInputs", "HouseholdJobRecord", "HouseholdJobReferenceError", "HouseholdJobStatus", "HouseholdJobStore", "MissingPolicyParametersError", + "NewEconomyReport", + "NewEconomySimulation", "NewHouseholdJob", "OverlappingReformError", "PolicyReform", "ReportRunStatus", "ReportRunTrigger", + "ReportSimulation", + "ReusableRun", + "RunVersions", + "SimulationChoice", "SimulationStatus", "SimulationType", "StoredHousehold", @@ -82,18 +115,29 @@ "combine_reforms", "create_simulation_engine", "fail_unstarted_household_job", + "fail_unstarted_report_run", + "fail_unstarted_simulation_run", + "find_reusable_run", "household_jobs", "input_metadata", + "insert_economy_report", + "insert_economy_simulation", "insert_household_job", "missing_table_privileges", "normalized_database_url", "read_catalog_version_id", + "read_dataset", + "read_dataset_by_id", "read_dynamic_reform", "read_household_job", "read_missing_table_privileges", + "read_model_id", "read_policy_reform", + "read_region", "read_stored_household", "record_household_call", + "record_report_run_call", + "record_simulation_run_call", "referenced_metadata", "report_runs", "reports", 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 new file mode 100644 index 000000000..cb44d76f4 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py @@ -0,0 +1,600 @@ +"""The simulation service's reads and writes for economy simulations and reports. + +An economy report compares a baseline and a reform simulation of one +region's data for one year. The service reads the catalog rows the request +names from the Public API's tables, then creates in one transaction: + +- for each side, a new ``simulations`` row and its first ``simulation_runs`` + row, ``pending`` with no Modal call ID; or, when a run of the same + simulation already succeeded and may be reused (``find_reusable_run``), + nothing: the report names that run and its simulation instead; +- the ``reports`` row, naming both simulations and the reform's policy, with + the report's options in ``inputs``; +- its first ``report_runs`` row, ``pending`` with no Modal call ID and + trigger ``initial``, naming both simulation runs. + +A standalone economy simulation is one ``simulations`` row and its first +run, created the same way. + +The service then starts the Modal call that calculates the run it created, +and records the call's ID on that run, or fails it at once when Modal +refuses the call. Both updates apply only while the run is still +``pending`` with no call ID: the coordinator or worker may already have +claimed it, and the stuck-run sweep fails a pending run no call was +recorded for. Failing a report run also fails the simulation runs created +for it, under the same condition. The coordinator and workers write the +rest: ``running``, the results, and ``succeeded`` or ``failed``. + +Every function takes a connection, so the caller chooses the transactions. +Each update sets only the columns the service's role may update, with +``updated_at``, whose automatic update the Public API's models make on their +own side only. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +from sqlalchemy import ( + cast, + func, + insert, + literal, + or_, + select, + update, +) +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.exc import IntegrityError + +from policyengine_simulation_persistence.input_tables import ( + MODEL_NAME_PREFIX, + datasets, + regions, + tax_benefit_models, +) +from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + ReportRunTrigger, + SimulationStatus, + SimulationType, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, FromClause, Table + +# PostgreSQL's SQLSTATE for a row naming one that does not exist. +FOREIGN_KEY_VIOLATION: Final = "23503" +# A report's type, as the Public API reads it. +ECONOMY_REPORT_TYPE: Final = "economy" +# The start of the error message of a report run whose coordinator refused a +# reused simulation run's output (``reused_simulation_: ``). +REFUSED_REUSE_PREFIX: Final = "reused_simulation_" +# Both run tables' statuses use these values. +_PENDING: Final = SimulationStatus.PENDING.value +_FAILED: Final = SimulationStatus.FAILED.value +# Each foreign key of the tables the service writes, by name: its table and +# the first column it constrains. +_KEY_COLUMNS: Final = { + str(constraint.name): (table.name, constraint.column_keys[0]) + for table in (simulations, simulation_runs, reports, report_runs) + for constraint in table.foreign_key_constraints +} + + +@dataclass(frozen=True) +class CatalogRegion: + """A region of a catalog version, and the dataset it uses by default.""" + + id: UUID + code: str + default_dataset_id: UUID + + +@dataclass(frozen=True) +class CatalogDataset: + """A dataset of a catalog version.""" + + id: UUID + name: str + + +@dataclass(frozen=True) +class RunVersions: + """The country package and .py versions a run calculates on.""" + + country_package_version: str + policyengine_version: str + + +@dataclass(frozen=True) +class EconomyPopulation: + """What every simulation of one request calculates on. + + The catalog version, the region and dataset of that version, and the + year. + """ + + tax_benefit_model_version_id: UUID + region_id: UUID + dataset_id: UUID + year: int + + +@dataclass(frozen=True) +class SimulationChoice: + """A simulation's stored policy and dynamic; neither means current law.""" + + policy_id: UUID | None + dynamic_id: UUID | None + + +@dataclass(frozen=True) +class ReusableRun: + """A succeeded simulation run a new report may name, and its simulation.""" + + simulation_id: UUID + simulation_run_id: UUID + + +@dataclass(frozen=True) +class ReportSimulation: + """The simulation and run one side of a report names. + + With ``reused``, both exist already and the run has succeeded, so nothing + is created for the side. Otherwise both are created, the run pending. + """ + + simulation_id: UUID + simulation_run_id: UUID + choice: SimulationChoice + reused: bool + + +@dataclass(frozen=True) +class NewEconomyReport: + """An economy report to create, with its first run. + + ``inputs`` holds the report's options as the Public API reads them. + """ + + id: UUID + run_id: UUID + country: str + label: str + description: str | None + tax_benefit_model_id: UUID + population: EconomyPopulation + versions: RunVersions + inputs: dict[str, Any] + baseline: ReportSimulation + reform: ReportSimulation + + @property + def new_simulation_run_ids(self) -> tuple[UUID, ...]: + """Return the IDs of the simulation runs created for the report.""" + return tuple( + side.simulation_run_id + for side in (self.baseline, self.reform) + if not side.reused + ) + + +@dataclass(frozen=True) +class NewEconomySimulation: + """A standalone economy simulation to create, with its first run.""" + + id: UUID + run_id: UUID + choice: SimulationChoice + population: EconomyPopulation + versions: RunVersions + + +class EconomyRecordReferenceError(LookupError): + """The database refused a row naming one that does not exist. + + ``table_name`` and ``column_name`` name the column whose foreign key + refused it; both are None when the database named a key the service's + tables don't have. + """ + + def __init__(self, table_name: str | None, column_name: str | None) -> None: + """Keep the table and column whose key refused the row.""" + self.table_name = table_name + self.column_name = column_name + named = ( + "a row" if table_name is None else f"the {table_name} row's {column_name}" + ) + super().__init__(f"{named} names no stored row") + + +def read_model_id(connection: Connection, *, country: str) -> UUID | None: + """Return the id of the country's tax-benefit model, or None.""" + return connection.execute( + select(tax_benefit_models.c.id).where( + tax_benefit_models.c.name == f"{MODEL_NAME_PREFIX}{country}" + ) + ).scalar_one_or_none() + + +def read_region( + connection: Connection, *, catalog_version_id: UUID, region_code: str +) -> CatalogRegion | None: + """Return the catalog version's region with this code, or None.""" + row = connection.execute( + select(regions.c.id, regions.c.code, regions.c.default_dataset_id).where( + regions.c.tax_benefit_model_version_id == catalog_version_id, + regions.c.code == region_code, + ) + ).one_or_none() + if row is None: + return None + return CatalogRegion( + id=row.id, code=row.code, default_dataset_id=row.default_dataset_id + ) + + +def read_dataset( + connection: Connection, *, catalog_version_id: UUID, name: str +) -> CatalogDataset | None: + """Return the catalog version's dataset with this name, or None.""" + return _dataset(connection, catalog_version_id, datasets.c.name == name) + + +def read_dataset_by_id( + connection: Connection, *, catalog_version_id: UUID, dataset_id: UUID +) -> CatalogDataset | None: + """Return the catalog version's dataset with this id, or None.""" + return _dataset(connection, catalog_version_id, datasets.c.id == dataset_id) + + +def _dataset( + connection: Connection, catalog_version_id: UUID, condition: ColumnElement[bool] +) -> CatalogDataset | None: + row = connection.execute( + select(datasets.c.id, datasets.c.name).where( + datasets.c.tax_benefit_model_version_id == catalog_version_id, condition + ) + ).one_or_none() + return None if row is None else CatalogDataset(id=row.id, name=row.name) + + +def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match + connection: Connection, + *, + choice: SimulationChoice, + population: EconomyPopulation, + versions: RunVersions, + inputs: dict[str, Any], + output_uri_prefix: str, + excluded_run_id: UUID | None = None, +) -> ReusableRun | None: + """Return the succeeded run a new report's simulation may name, or None. + + The run is the newest, by completion, of a run: + + - 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; + - 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`` + (options), so its file holds what such a report needs; + - that no report run has refused to reuse (an error message starting + ``reused_simulation_``), so a submission after a refused reuse + calculates the simulation again; and + - that is not ``excluded_run_id``, the other side's run. + + The coordinator may still refuse it, for example when labour supply + responses differ between the two reports. + """ + used_by = report_runs.alias("used_by_report_run") + refused_by = report_runs.alias("refused_by_report_run") + used_with_same_options = ( + select(used_by.c.id) + .join(reports, reports.c.id == used_by.c.report_id) + .where( + used_by.c.status == ReportRunStatus.SUCCEEDED, + _names_run(used_by), + cast(reports.c.inputs, JSONB) == literal(inputs, JSONB), + ) + .exists() + ) + refused = ( + select(refused_by.c.id) + .where( + refused_by.c.status == ReportRunStatus.FAILED, + _names_run(refused_by), + refused_by.c.error_message.startswith( + REFUSED_REUSE_PREFIX, autoescape=True + ), + ) + .exists() + ) + 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_runs.c.status == SimulationStatus.SUCCEEDED, + simulation_runs.c.year == population.year, + 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_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) + row = connection.execute( + select(simulation_runs.c.id, simulation_runs.c.simulation_id) + .join(simulations, simulations.c.id == simulation_runs.c.simulation_id) + .where(*conditions) + .order_by(simulation_runs.c.completed_at.desc(), simulation_runs.c.id.desc()) + .limit(1) + ).one_or_none() + if row is None: + return None + return ReusableRun(simulation_id=row.simulation_id, simulation_run_id=row.id) + + +def _names_run(report_run: FromClause) -> ColumnElement[bool]: + """Return whether ``report_run`` names the outer query's simulation run.""" + return or_( + report_run.c.baseline_simulation_run_id == simulation_runs.c.id, + report_run.c.reform_simulation_run_id == simulation_runs.c.id, + ) + + +def insert_economy_report(connection: Connection, report: NewEconomyReport) -> None: + """Create the report, its first run and each side's new simulation and run. + + Raise EconomyRecordReferenceError when the database refuses a row it + names. The refusal ends the caller's transaction, or its savepoint. + """ + with _references_checked(): + for side in (report.baseline, report.reform): + if not side.reused: + _insert_simulation( + connection, + simulation_id=side.simulation_id, + run_id=side.simulation_run_id, + choice=side.choice, + population=report.population, + versions=report.versions, + ) + population = report.population + connection.execute( + insert(reports).values( + id=report.id, + label=report.label, + description=report.description, + country=report.country, + type=ECONOMY_REPORT_TYPE, + user_id=None, + tax_benefit_model_id=report.tax_benefit_model_id, + policy_id=report.reform.choice.policy_id, + baseline_simulation_id=report.baseline.simulation_id, + reform_simulation_id=report.reform.simulation_id, + household_id=None, + dataset_id=population.dataset_id, + region_id=population.region_id, + year=population.year, + budget_window_size=None, + inputs=report.inputs, + ) + ) + connection.execute( + insert(report_runs).values( + id=report.run_id, + report_id=report.id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_run_id=report.reform.simulation_run_id, + budget_window_run_id=None, + year=None, + country_package_version=report.versions.country_package_version, + policyengine_version=report.versions.policyengine_version, + status=ReportRunStatus.PENDING, + trigger=ReportRunTrigger.INITIAL, + modal_call_id=None, + ) + ) + + +def insert_economy_simulation( + connection: Connection, simulation: NewEconomySimulation +) -> None: + """Create the standalone simulation and its first run, pending. + + 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( + connection, + simulation_id=simulation.id, + run_id=simulation.run_id, + choice=simulation.choice, + population=simulation.population, + versions=simulation.versions, + ) + + +def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts + connection: Connection, + *, + simulation_id: UUID, + run_id: UUID, + choice: SimulationChoice, + population: EconomyPopulation, + versions: RunVersions, +) -> None: + """Create an economy simulation and its first run, pending with no call.""" + connection.execute( + insert(simulations).values( + id=simulation_id, + simulation_type=SimulationType.ECONOMY, + dataset_id=population.dataset_id, + household_id=None, + policy_id=choice.policy_id, + dynamic_id=choice.dynamic_id, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + region_id=population.region_id, + year=population.year, + ) + ) + connection.execute( + insert(simulation_runs).values( + id=run_id, + simulation_id=simulation_id, + year=population.year, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + country_package_version=versions.country_package_version, + policyengine_version=versions.policyengine_version, + dataset_id=population.dataset_id, + region_id=population.region_id, + status=SimulationStatus.PENDING, + modal_call_id=None, + ) + ) + + +@contextmanager +def _references_checked() -> Iterator[None]: + """Turn a foreign-key refusal into EconomyRecordReferenceError. + + Any other integrity error propagates unchanged. + """ + try: + yield + except IntegrityError as error: + driver_error = error.orig + if getattr(driver_error, "sqlstate", None) != FOREIGN_KEY_VIOLATION: + raise + constraint_name = getattr( + getattr(driver_error, "diag", None), "constraint_name", None + ) + table_name, column_name = _KEY_COLUMNS.get(str(constraint_name), (None, None)) + raise EconomyRecordReferenceError(table_name, column_name) from None + + +def _pending_without_call( + table: Table, run_id: UUID +) -> tuple[ColumnElement[bool], ...]: + return ( + table.c.id == run_id, + table.c.status == _PENDING, + table.c.modal_call_id.is_(None), + ) + + +def record_report_run_call( + connection: Connection, *, report_run_id: UUID, modal_call_id: str +) -> ReportRunStatus | None: + """Record the report run's Modal call ID while it is pending with none. + + Return the run's status afterwards, or None when there is no such run. A + run the coordinator has claimed already keeps its claim. + """ + status = _record_call(connection, report_runs, report_run_id, modal_call_id) + return None if status is None else ReportRunStatus(status) + + +def record_simulation_run_call( + connection: Connection, *, simulation_run_id: UUID, modal_call_id: str +) -> SimulationStatus | None: + """Record the simulation run's Modal call ID while it is pending with none. + + Return the run's status afterwards, or None when there is no such run. A + run the worker has claimed already keeps its claim. + """ + status = _record_call(connection, simulation_runs, simulation_run_id, modal_call_id) + return None if status is None else SimulationStatus(status) + + +def _record_call( + connection: Connection, table: Table, run_id: UUID, modal_call_id: str +) -> str | None: + recorded = connection.execute( + update(table) + .where(*_pending_without_call(table, run_id)) + .values(modal_call_id=modal_call_id, updated_at=func.now()) + .returning(table.c.status) + ).scalar_one_or_none() + if recorded is not None: + return str(recorded) + status = connection.execute( + select(table.c.status).where(table.c.id == run_id) + ).scalar_one_or_none() + return None if status is None else str(status) + + +def fail_unstarted_report_run( + connection: Connection, + *, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, +) -> bool: + """Fail the report run at once while it is pending with no Modal call ID. + + When it was failed, also fail each of ``simulation_run_ids``, the + simulation runs created for it, that is still pending with no call ID. + Return whether the report run was failed: one a call has been recorded + for, or one the coordinator has claimed, is left as it is, and so are + its simulation runs. + """ + if not _fail_unstarted(connection, report_runs, report_run_id, error_message): + return False + for simulation_run_id in simulation_run_ids: + _fail_unstarted(connection, simulation_runs, simulation_run_id, error_message) + return True + + +def fail_unstarted_simulation_run( + connection: Connection, *, simulation_run_id: UUID, error_message: str +) -> bool: + """Fail the simulation run at once while it is pending with no Modal call ID. + + Return whether it was failed: a run a call has been recorded for, or one + a worker has claimed, is left as it is. + """ + return _fail_unstarted( + connection, simulation_runs, simulation_run_id, error_message + ) + + +def _fail_unstarted( + connection: Connection, table: Table, run_id: UUID, error_message: str +) -> bool: + failed = connection.execute( + update(table) + .where(*_pending_without_call(table, run_id)) + .values( + status=_FAILED, + error_message=error_message, + completed_at=func.now(), + updated_at=func.now(), + ) + .returning(table.c.id) + ).scalar_one_or_none() + return failed is not None diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py index fa4a7d0fd..ec9797fad 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py @@ -42,6 +42,26 @@ Column("version", String(128), nullable=False), ) +# A region of one catalog version, unique by code within it, and the dataset +# it uses unless a calculation names another. +regions = Table( + "regions", + input_metadata, + Column("id", Uuid, nullable=False), + Column("code", String(255), nullable=False), + Column("tax_benefit_model_version_id", Uuid, nullable=False), + Column("default_dataset_id", Uuid, nullable=False), +) + +# A dataset of one catalog version, unique by name within it. +datasets = Table( + "datasets", + input_metadata, + Column("id", Uuid, nullable=False), + Column("name", String(255), nullable=False), + Column("tax_benefit_model_version_id", Uuid, nullable=False), +) + # ``household_data`` lists the people and each group entity as records, and # may hold a saved SPM selection; ``default_year`` is the year it describes. households = Table( From ebff4f87b23542405b5bcac3766ff8b212f43ac2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:18:20 +0400 Subject: [PATCH 05/26] Test the economy store's reads and writes, each in its own transaction (failing) The store reads a request's catalog version, model, region and dataset, each side's policy and dynamic values, and each side's reusable run in one transaction, never giving both sides the same run; it refuses each missing row with its own error. It creates a report or a simulation, records a call or fails an unstarted run, each in its own transaction. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_run_store.py | 324 ++++++++++++++++++ 1 file changed, 324 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/tests/test_economy_run_store.py diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py new file mode 100644 index 000000000..b9cb84487 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py @@ -0,0 +1,324 @@ +"""The economy store: each read and write in a transaction of its own. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``). The store opens each transaction +through a function it is given; here each is a savepoint of one connection, +acting as the simulation service's role, which the test rolls back. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import UTC, datetime +from typing import TYPE_CHECKING +from uuid import uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + POLICYENGINE_VERSION, + VERSIONS, + insert_economy_catalog, + insert_reusable_run, +) +from policyengine_simulation_persistence import ( + EconomyRequest, + EconomyRunStore, + MissingPolicyParametersError, + NewEconomyReport, + NewEconomySimulation, + OverlappingReformError, + ReportRunStatus, + ReportSimulation, + ReusableRun, + RunReuse, + SimulationChoice, + SimulationStatus, + UnknownCatalogVersionError, + UnknownDatasetError, + UnknownDynamicError, + UnknownPolicyError, + UnknownRegionError, +) +from public_api_rows import StoredValue, insert_dataset, insert_dynamic, insert_policy +from sqlalchemy import text + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +START = datetime(2026, 1, 1, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +ELASTICITY = "gov.simulation.labor_supply_responses.elasticities.income" +REQUEST = EconomyRequest( + country="us", + policyengine_version=POLICYENGINE_VERSION, + region_code="us", + dataset_name=None, + year=2026, +) +REUSE = RunReuse( + versions=VERSIONS, inputs=OPTIONS, output_uri_prefix=OUTPUT_BUCKET_PREFIX +) + + +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_reports_inputs_are_read_together(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog( + connection, parameter_names=(CTC_AMOUNT, ELASTICITY) + ) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + reused = insert_reusable_run(connection, economy) + act_as_the_service(connection) + store = _store_on(connection) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=dynamic_id) + + inputs = store.read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=reform_choice, reuse=REUSE + ) + without_reuse = store.read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=reform_choice, reuse=None + ) + + catalog = inputs.catalog + assert ( + catalog.catalog_version_id, + catalog.model_id, + catalog.region.id, + catalog.dataset.id, + catalog.dataset.name, + ) == ( + economy.catalog.version_id, + economy.catalog.model_id, + economy.region_id, + economy.dataset_id, + "populace_us_2024", + ) + assert catalog.population(2026) == economy.population() + assert (inputs.baseline.choice, inputs.baseline.reform) == (CURRENT_LAW, {}) + assert inputs.reform.choice == reform_choice + assert inputs.reform.reform == { + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + ELASTICITY: {"2026-01-01.2200-12-31": -0.05}, + } + assert inputs.baseline.reusable_run == ReusableRun( + simulation_id=reused.simulation_id, simulation_run_id=reused.simulation_run_id + ) + assert inputs.reform.reusable_run is None + assert without_reuse.baseline.reusable_run is None + + +def test_the_two_sides_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) + only = insert_reusable_run(connection, economy) + act_as_the_service(connection) + + inputs = _store_on(connection).read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=CURRENT_LAW, reuse=REUSE + ) + + assert inputs.baseline.reusable_run == ReusableRun( + simulation_id=only.simulation_id, simulation_run_id=only.simulation_run_id + ) + assert inputs.reform.reusable_run is None + + +def test_a_named_dataset_is_read_instead_of_the_regions_default( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + local_id = insert_dataset(connection, economy.catalog, "acs_local_2024") + act_as_the_service(connection) + + inputs = _store_on(connection).read_simulation_inputs( + EconomyRequest(**{**vars(REQUEST), "dataset_name": "acs_local_2024"}), + CURRENT_LAW, + ) + + assert (inputs.catalog.dataset.id, inputs.catalog.dataset.name) == ( + local_id, + "acs_local_2024", + ) + assert (inputs.simulation.choice, inputs.simulation.reform) == (CURRENT_LAW, {}) + assert inputs.simulation.reusable_run is None + + +@pytest.mark.parametrize( + ("missing", "error_type"), + [ + ("catalog version", UnknownCatalogVersionError), + ("region", UnknownRegionError), + ("dataset", UnknownDatasetError), + ("policy", UnknownPolicyError), + ("dynamic", UnknownDynamicError), + ], +) +def test_each_missing_row_a_request_names_is_refused( + public_api_database_url, missing: str, error_type: type[Exception] +) -> None: + request_fields = { + "catalog version": {"policyengine_version": "6.16.0"}, + "region": {"region_code": "state/zz"}, + "dataset": {"dataset_name": "no_such_dataset"}, + }.get(missing, {}) + choice = { + "policy": SimulationChoice(policy_id=uuid4(), dynamic_id=None), + "dynamic": SimulationChoice(policy_id=None, dynamic_id=uuid4()), + }.get(missing, CURRENT_LAW) + with rolled_back_transaction(public_api_database_url) as connection: + insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + request = EconomyRequest(**{**vars(REQUEST), **request_fields}) + + with pytest.raises(error_type): + store.read_report_inputs( + request, baseline=CURRENT_LAW, reform=choice, reuse=REUSE + ) + with pytest.raises(error_type): + store.read_simulation_inputs(request, choice) + + +def test_values_the_running_version_lacks_or_both_set_are_refused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(ELASTICITY,)) + saved_on = insert_economy_catalog( + connection, country="uk", parameter_names=(CTC_AMOUNT,) + ) + # A policy of the US saved on parameters the running version lacks. + lacking_policy_id = insert_policy( + connection, + saved_on.catalog.__class__(**{**vars(saved_on.catalog), "country": "us"}), + [StoredValue(CTC_AMOUNT, 3_000, START)], + ) + overlapping_policy_id = insert_policy( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.1, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + act_as_the_service(connection) + store = _store_on(connection) + + with pytest.raises(MissingPolicyParametersError): + store.read_simulation_inputs( + REQUEST, SimulationChoice(lacking_policy_id, dynamic_id=None) + ) + with pytest.raises(OverlappingReformError): + store.read_simulation_inputs( + REQUEST, SimulationChoice(overlapping_policy_id, dynamic_id) + ) + + +def _status(connection: Connection, table_name: str, run_id) -> tuple: + return tuple( + connection.execute( + text( + f"SELECT status, modal_call_id FROM public.{table_name} " # noqa: S608 # fixed table names + "WHERE id = :id" + ), + {"id": run_id}, + ).one() + ) + + +def test_a_report_is_created_then_its_call_recorded_or_its_run_failed( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + reports = [ + NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label=f"Report {number}", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=OPTIONS, + baseline=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + reform=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + ) + for number in (1, 2) + ] + for report in reports: + store.create_report(report) + recorded, refused = reports + + status = store.record_report_call(recorded.run_id, "fc-coordinator") + failed = store.fail_unstarted_report( + refused.run_id, refused.new_simulation_run_ids, "modal_call_refused: no" + ) + recorded_state = _status(connection, "report_runs", recorded.run_id) + refused_state = _status(connection, "report_runs", refused.run_id) + refused_baseline = _status( + connection, "simulation_runs", refused.baseline.simulation_run_id + ) + + assert (status, failed) == (ReportRunStatus.PENDING, True) + assert recorded_state == ("pending", "fc-coordinator") + assert refused_state == ("failed", None) + assert refused_baseline == ("failed", None) + + +def test_a_simulation_is_created_then_its_call_recorded_or_its_run_failed( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + simulations = [ + NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + ) + for _ in range(2) + ] + for simulation in simulations: + store.create_simulation(simulation) + recorded, refused = simulations + + status = store.record_simulation_call(recorded.run_id, "fc-worker") + failed = store.fail_unstarted_simulation( + refused.run_id, "modal_call_refused: no" + ) + recorded_state = _status(connection, "simulation_runs", recorded.run_id) + refused_state = _status(connection, "simulation_runs", refused.run_id) + + assert (status, failed) == (SimulationStatus.PENDING, True) + assert recorded_state == ("pending", "fc-worker") + assert refused_state == ("failed", None) From 9668c725137afe4accdcfa5cb40f24c5c705b5f8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:19:37 +0400 Subject: [PATCH 06/26] Read and write economy simulations and reports, each in its own transaction EconomyRunStore reads a request's catalog version, model, region and dataset (the region's default unless one is named), each side's combined policy and dynamic values and, when asked, each side's reusable run, in one transaction; the reform never reuses the baseline's run. A missing region or dataset is UnknownRegionError or UnknownDatasetError. Creating a report or simulation, recording a call and failing an unstarted run each take their own transaction. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 20 + .../economy_run_store.py | 374 ++++++++++++++++++ 2 files changed, 394 insertions(+) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 71a21952e..e8ce4e7a3 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -29,6 +29,17 @@ record_report_run_call, record_simulation_run_call, ) +from policyengine_simulation_persistence.economy_run_store import ( + EconomyCatalogRows, + EconomyReportInputs, + EconomyRequest, + EconomyRunStore, + EconomySimulationInputs, + RunReuse, + SimulationInputs, + UnknownDatasetError, + UnknownRegionError, +) from policyengine_simulation_persistence.engine import ( bounded_transaction, create_simulation_engine, @@ -84,8 +95,13 @@ "SERVICE_UPDATABLE_COLUMNS", "CatalogDataset", "CatalogRegion", + "EconomyCatalogRows", "EconomyPopulation", "EconomyRecordReferenceError", + "EconomyReportInputs", + "EconomyRequest", + "EconomyRunStore", + "EconomySimulationInputs", "HouseholdCalculationInputs", "HouseholdJobRecord", "HouseholdJobReferenceError", @@ -101,15 +117,19 @@ "ReportRunTrigger", "ReportSimulation", "ReusableRun", + "RunReuse", "RunVersions", "SimulationChoice", + "SimulationInputs", "SimulationStatus", "SimulationType", "StoredHousehold", "UnknownCatalogVersionError", + "UnknownDatasetError", "UnknownDynamicError", "UnknownHouseholdError", "UnknownPolicyError", + "UnknownRegionError", "bounded_transaction", "budget_window_runs", "combine_reforms", 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 new file mode 100644 index 000000000..b90921b4b --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py @@ -0,0 +1,374 @@ +"""The economy store: each read and write in a transaction of its own. + +The Simulation Entrypoint creates an economy report, or a standalone economy +simulation, in three steps, each its own transaction, with the Modal call +started between them: it reads the request's inputs, creates the rows, then +records the call or fails the run it started. ``EconomyRunStore`` runs the +reads and writes of ``economy_run_records`` and ``policy_reform`` that way. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_persistence.economy_run_records import ( + CatalogDataset, + CatalogRegion, + EconomyPopulation, + ReusableRun, + RunVersions, + SimulationChoice, + fail_unstarted_report_run, + fail_unstarted_simulation_run, + find_reusable_run, + insert_economy_report, + insert_economy_simulation, + read_dataset, + read_dataset_by_id, + read_model_id, + read_region, + record_report_run_call, + record_simulation_run_call, +) +from policyengine_simulation_persistence.engine import bounded_transaction +from policyengine_simulation_persistence.household_job_records import ( + read_catalog_version_id, +) +from policyengine_simulation_persistence.household_job_store import ( + UnknownCatalogVersionError, +) +from policyengine_simulation_persistence.policy_reform import ( + combine_reforms, + read_dynamic_reform, + read_policy_reform, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from contextlib import AbstractContextManager + from uuid import UUID + + from sqlalchemy import Connection, Engine + + from policyengine_simulation_persistence.economy_run_records import ( + NewEconomyReport, + NewEconomySimulation, + ) + from policyengine_simulation_persistence.policy_reform import PolicyReform + from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + SimulationStatus, + ) + +type TransactionFactory = Callable[[], AbstractContextManager[Connection]] + + +class UnknownRegionError(LookupError): + """The catalog version has no region with the requested code.""" + + +class UnknownDatasetError(LookupError): + """The catalog version has no dataset with the requested name.""" + + +@dataclass(frozen=True) +class EconomyRequest: + """What an economy simulation or report request names, apart from policies. + + ``policyengine_version`` is the .py version the calculation runs on, + which chooses the country's catalog version. ``dataset_name`` None means + the region's default dataset. + """ + + country: str + policyengine_version: str + region_code: str + dataset_name: str | None + year: int + + +@dataclass(frozen=True) +class RunReuse: + """How a report looks for simulation runs to reuse (``find_reusable_run``). + + ``versions`` are those the new runs would calculate on, ``inputs`` the + report's options, and ``output_uri_prefix`` the output bucket's + ``gs:///``. + """ + + versions: RunVersions + inputs: dict[str, Any] + output_uri_prefix: str + + +@dataclass(frozen=True) +class EconomyCatalogRows: + """The catalog rows a request is calculated on.""" + + catalog_version_id: UUID + model_id: UUID + region: CatalogRegion + dataset: CatalogDataset + + def population(self, year: int) -> EconomyPopulation: + """Return what every simulation of the request calculates on.""" + return EconomyPopulation( + tax_benefit_model_version_id=self.catalog_version_id, + region_id=self.region.id, + dataset_id=self.dataset.id, + year=year, + ) + + +@dataclass(frozen=True) +class SimulationInputs: + """One simulation's policy and dynamic, and the values they set together. + + ``reusable_run`` names a succeeded run the simulation may use instead of + a new one, when the request looked for one and one exists. + """ + + choice: SimulationChoice + reform: PolicyReform + reusable_run: ReusableRun | None = None + + +@dataclass(frozen=True) +class EconomySimulationInputs: + """What a standalone economy simulation is read from.""" + + catalog: EconomyCatalogRows + simulation: SimulationInputs + + +@dataclass(frozen=True) +class EconomyReportInputs: + """What an economy report is read from: the catalog rows and both sides.""" + + catalog: EconomyCatalogRows + baseline: SimulationInputs + reform: SimulationInputs + + +class EconomyRunStore: + """Read and write economy simulations and reports, each in its own transaction.""" + + def __init__(self, transaction: TransactionFactory) -> None: + """Open each transaction with ``transaction``, which commits on success.""" + self._transaction = transaction + + @classmethod + def from_engine( + cls, engine: Engine, *, statement_timeout_milliseconds: int + ) -> EconomyRunStore: + """Use ``engine``, bounding each statement to the given time.""" + + def transaction() -> AbstractContextManager[Connection]: + return bounded_transaction( + engine, statement_timeout_milliseconds=statement_timeout_milliseconds + ) + + return cls(transaction) + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + """Read a standalone simulation's catalog rows and values together. + + Raise the errors ``read_report_inputs`` names. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + simulation = _simulation_inputs( + connection, request.country, catalog.catalog_version_id, choice + ) + return EconomySimulationInputs(catalog=catalog, simulation=simulation) + + def read_report_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + """Read a report's catalog rows, both sides' values and reusable runs. + + With ``reuse``, each side gets the run ``find_reusable_run`` finds, + and the reform never gets the baseline's. Raise + UnknownCatalogVersionError, UnknownRegionError, UnknownDatasetError, + the policy reader's UnknownPolicyError, UnknownDynamicError or + MissingPolicyParametersError, or OverlappingReformError. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + population = catalog.population(request.year) + sides = [] + excluded_run_id = None + for choice in (baseline, reform): + inputs = _simulation_inputs( + connection, request.country, catalog.catalog_version_id, choice + ) + 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(self, report: NewEconomyReport) -> None: + """Create the report, its new simulations and its runs.""" + with self._transaction() as connection: + insert_economy_report(connection, report) + + def create_simulation(self, simulation: NewEconomySimulation) -> None: + """Create the simulation and its run; see ``insert_economy_simulation``.""" + with self._transaction() as connection: + insert_economy_simulation(connection, simulation) + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + """Record the report run's call ID while it is pending with none.""" + with self._transaction() as connection: + return record_report_run_call( + connection, report_run_id=report_run_id, modal_call_id=modal_call_id + ) + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + """Fail the report run, and its new simulation runs, while it is unstarted.""" + with self._transaction() as connection: + return fail_unstarted_report_run( + connection, + report_run_id=report_run_id, + simulation_run_ids=simulation_run_ids, + error_message=error_message, + ) + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + """Record the simulation run's call ID while it is pending with none.""" + with self._transaction() as connection: + return record_simulation_run_call( + connection, + simulation_run_id=simulation_run_id, + modal_call_id=modal_call_id, + ) + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + """Fail the simulation run while it is pending with no call ID.""" + with self._transaction() as connection: + return fail_unstarted_simulation_run( + connection, + simulation_run_id=simulation_run_id, + error_message=error_message, + ) + + +def _catalog_rows( + connection: Connection, request: EconomyRequest +) -> EconomyCatalogRows: + """Read the catalog version, model, region and dataset a request names. + + Raise UnknownCatalogVersionError, UnknownRegionError or + UnknownDatasetError. + """ + country = request.country + catalog_version_id = read_catalog_version_id( + connection, + country=country, + policyengine_version=request.policyengine_version, + ) + model_id = read_model_id(connection, country=country) + if catalog_version_id is None or model_id is None: + msg = f"{country} has no catalog version for .py {request.policyengine_version}" + raise UnknownCatalogVersionError(msg) + region = read_region( + connection, + catalog_version_id=catalog_version_id, + region_code=request.region_code, + ) + if region is None: + msg = ( + f"{country}'s catalog version for .py {request.policyengine_version} " + f"has no region {request.region_code}" + ) + raise UnknownRegionError(msg) + dataset = ( + read_dataset_by_id( + connection, + catalog_version_id=catalog_version_id, + dataset_id=region.default_dataset_id, + ) + if request.dataset_name is None + else read_dataset( + connection, + catalog_version_id=catalog_version_id, + name=request.dataset_name, + ) + ) + if dataset is None: + msg = ( + f"{country}'s catalog version for .py {request.policyengine_version} " + f"has no dataset {request.dataset_name}" + ) + raise UnknownDatasetError(msg) + return EconomyCatalogRows( + catalog_version_id=catalog_version_id, + model_id=model_id, + region=region, + dataset=dataset, + ) + + +def _simulation_inputs( + connection: Connection, + country: str, + catalog_version_id: UUID, + choice: SimulationChoice, +) -> SimulationInputs: + """Read one simulation's policy and dynamic values, combined.""" + policy_reform = ( + {} + if choice.policy_id is None + else read_policy_reform( + connection, + country=country, + policy_id=choice.policy_id, + catalog_version_id=catalog_version_id, + ) + ) + dynamic_reform = ( + {} + if choice.dynamic_id is None + else read_dynamic_reform( + connection, + dynamic_id=choice.dynamic_id, + catalog_version_id=catalog_version_id, + ) + ) + return SimulationInputs( + choice=choice, reform=combine_reforms(policy_reform, dynamic_reform) + ) From 40c0691d639e1f8cd8ac22f2fcd4167760b1c779 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:32:32 +0400 Subject: [PATCH 07/26] Test POST /reports and POST /simulations, and the economy calls they start (failing) The service finds the default v2 worker's economy functions, refuses each missing row with its code before writing, creates the records and runs, then starts the run's call: the coordinator, given the report with its run IDs and no evaluation record, or a standalone run's simulation function, given the run's ID. It records the call's ID, or fails the run and answers 503 naming the record. Reports reuse a succeeded run only when the output bucket is configured. The routes answer 201, 400, 422, 503 with Retry-After, and a JSON 500, and the OpenAPI document declares each. Co-Authored-By: Claude Opus 5.5 --- .../tests/economy_fixtures.py | 49 ++ .../tests/economy_service_fakes.py | 299 +++++++++++ .../tests/test_config.py | 23 + .../tests/test_economy_calculation_inputs.py | 152 ++++++ .../tests/test_economy_routes.py | 372 ++++++++++++++ .../tests/test_economy_submissions.py | 469 ++++++++++++++++++ .../tests/test_economy_workers.py | 196 ++++++++ .../tests/test_openapi.py | 104 +++- 8 files changed, 1662 insertions(+), 2 deletions(-) create mode 100644 projects/policyengine-simulation-entry/tests/economy_fixtures.py create mode 100644 projects/policyengine-simulation-entry/tests/economy_service_fakes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_routes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_submissions.py create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_workers.py diff --git a/projects/policyengine-simulation-entry/tests/economy_fixtures.py b/projects/policyengine-simulation-entry/tests/economy_fixtures.py new file mode 100644 index 000000000..d24a831bd --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/economy_fixtures.py @@ -0,0 +1,49 @@ +"""Test values for the economy report and simulation routes. + +They build on the household path's v2 worker, which can calculate economy +runs for the US and the UK, with a second US dataset that state regions use. +""" + +from __future__ import annotations + +from household_fixtures import POLICYENGINE_VERSION, household_worker_version +from policyengine_simulation_contract.stage12_bundle import Stage12Dataset +from policyengine_simulation_contract.stage12_manifest import ( + V2VersionManifest, + V2WorkerVersion, +) + +STATE_DATASET = Stage12Dataset( + identity="populace_us_2024_acs_local", + uri="hf://policyengine/data/populace_us_2024_acs_local.h5@acs-local-revision", + artifact_revision="acs-local-revision", + sha256="c" * 64, + repo_type="dataset", +) + + +def economy_worker_version( + policyengine_version: str = POLICYENGINE_VERSION, +) -> V2WorkerVersion: + """Return a v2 worker whose US bundle also holds the state dataset.""" + base = household_worker_version(policyengine_version) + countries = tuple( + country.model_copy(update={"datasets": (*country.datasets, STATE_DATASET)}) + if country.country == "us" + else country + for country in base.bundle.countries + ) + return base.model_copy( + update={"bundle": base.bundle.model_copy(update={"countries": countries})} + ) + + +def economy_manifest( + default_worker: V2WorkerVersion | None = None, +) -> V2VersionManifest: + """Return a manifest whose only worker, its default, is ``default_worker``.""" + worker = default_worker or economy_worker_version() + version = worker.bundle.policyengine_version + return V2VersionManifest( + generation="test", default_version=version, versions={version: worker} + ) diff --git a/projects/policyengine-simulation-entry/tests/economy_service_fakes.py b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py new file mode 100644 index 000000000..3be3fd4c1 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py @@ -0,0 +1,299 @@ +"""Fakes of the economy submission service's store, manifest and Modal calls. + +The store's own reads and writes are tested against PostgreSQL in the +persistence library; here a fake records what the service asks of it. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from datetime import UTC, datetime +from itertools import count +from typing import TYPE_CHECKING, Any +from uuid import UUID + +from economy_fixtures import economy_manifest +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportRequest, + EconomySimulationRequest, +) +from policyengine_simulation_entry.economy_submissions import ( + EconomySubmissionService, + ServiceSettings, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorker, + default_economy_worker, +) +from policyengine_simulation_persistence import ( + CatalogDataset, + CatalogRegion, + EconomyCatalogRows, + EconomyReportInputs, + EconomySimulationInputs, + ReportRunStatus, + ReusableRun, + SimulationInputs, + SimulationStatus, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from policyengine_simulation_persistence import ( + EconomyRequest, + NewEconomyReport, + NewEconomySimulation, + RunReuse, + SimulationChoice, + ) + +CATALOG_VERSION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +MODEL_ID = UUID("00000000-0000-4000-8000-0000000000c2") +REGION_ID = UUID("00000000-0000-4000-8000-0000000000c3") +DATASET_ID = UUID("00000000-0000-4000-8000-0000000000c4") +POLICY_ID = UUID("00000000-0000-4000-8000-0000000000c5") +DYNAMIC_ID = UUID("00000000-0000-4000-8000-0000000000c6") +REUSED = ReusableRun( + simulation_id=UUID("00000000-0000-4000-8000-0000000000d1"), + simulation_run_id=UUID("00000000-0000-4000-8000-0000000000d2"), +) +CREATED_AT = datetime(2026, 11, 12, 15, tzinfo=UTC) +OUTPUT_BUCKET = "policyengine-v2-simulation-outputs-staging" +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +REFORM = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} +CATALOG = EconomyCatalogRows( + catalog_version_id=CATALOG_VERSION_ID, + model_id=MODEL_ID, + region=CatalogRegion(id=REGION_ID, code="us", default_dataset_id=DATASET_ID), + dataset=CatalogDataset(id=DATASET_ID, name="populace_us_2024"), +) +REPORT_BODY: dict[str, Any] = { + "country": "us", + "report_type": "economy", + "label": "A larger child tax credit", + "description": "Raises the credit to $3,000", + "year": 2026, + "region_code": "us", + "dataset_name": None, + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": str(POLICY_ID), "dynamic_id": str(DYNAMIC_ID)}, + "include_cliffs": True, + "spm": {"geography_kind": "national"}, +} +SIMULATION_BODY: dict[str, Any] = { + "country": "us", + "simulation_type": "economy", + "year": 2026, + "region_code": "us", + "dataset_name": None, + "policy_id": str(POLICY_ID), + "dynamic_id": None, +} + + +def report_request(**changes: Any) -> EconomyReportRequest: + return EconomyReportRequest.model_validate({**REPORT_BODY, **changes}) + + +def simulation_request(**changes: Any) -> EconomySimulationRequest: + return EconomySimulationRequest.model_validate({**SIMULATION_BODY, **changes}) + + +class FakeWorkers: + """The default worker's economy functions, or the lookup's failure.""" + + def __init__(self, failure: Exception | None = None) -> None: + self.failure = failure + + def default_worker(self, country: str) -> EconomyWorker: + if self.failure is not None: + raise self.failure + return default_economy_worker(economy_manifest(), country) + + +@dataclass +class FakeStore: + """Records what the service asks of the database, and answers as told. + + ``claimed`` marks the run as claimed by its coordinator or function: a + failure write then changes nothing. + """ + + catalog: EconomyCatalogRows = CATALOG + reusable: tuple[ReusableRun | None, ReusableRun | None] = (None, None) + read_failure: Exception | None = None + create_failure: Exception | None = None + record_failure: Exception | None = None + fail_failure: Exception | None = None + claimed: bool = False + report_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) + simulation_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) + reports: list[NewEconomyReport] = field(default_factory=list) + simulations: list[NewEconomySimulation] = field(default_factory=list) + recorded_calls: list[tuple[UUID, str]] = field(default_factory=list) + failed_runs: list[tuple[UUID, tuple[UUID, ...], str]] = field(default_factory=list) + + def read_report_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + self.report_reads.append((request, (baseline, reform, reuse))) + if self.read_failure is not None: + raise self.read_failure + baseline_reusable, reform_reusable = self.reusable + return EconomyReportInputs( + catalog=self.catalog, + baseline=SimulationInputs( + choice=baseline, reform={}, reusable_run=baseline_reusable + ), + reform=SimulationInputs( + choice=reform, reform=REFORM, reusable_run=reform_reusable + ), + ) + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + self.simulation_reads.append((request, choice)) + if self.read_failure is not None: + raise self.read_failure + return EconomySimulationInputs( + catalog=self.catalog, + simulation=SimulationInputs(choice=choice, reform=REFORM), + ) + + def create_report(self, report: NewEconomyReport) -> None: + if self.create_failure is not None: + raise self.create_failure + self.reports.append(report) + + def create_simulation(self, simulation: NewEconomySimulation) -> None: + if self.create_failure is not None: + raise self.create_failure + self.simulations.append(simulation) + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + if self.record_failure is not None: + raise self.record_failure + self.recorded_calls.append((report_run_id, modal_call_id)) + return ReportRunStatus.RUNNING if self.claimed else ReportRunStatus.PENDING + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + if self.fail_failure is not None: + raise self.fail_failure + if self.claimed: + return False + self.failed_runs.append( + (report_run_id, tuple(simulation_run_ids), error_message) + ) + return True + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + if self.record_failure is not None: + raise self.record_failure + self.recorded_calls.append((simulation_run_id, modal_call_id)) + return SimulationStatus.RUNNING if self.claimed else SimulationStatus.PENDING + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + if self.fail_failure is not None: + raise self.fail_failure + if self.claimed: + return False + self.failed_runs.append((simulation_run_id, (), error_message)) + return True + + +@dataclass +class FakeCaller: + """Starts calls by recording them; ``failure`` is raised instead when set. + + ``hang`` makes a start wait until the service stops waiting for it. + """ + + failure: Exception | None = None + hang: bool = False + report_runs: list[tuple[EconomyWorker, dict, dict, Any]] = field( + default_factory=list + ) + simulation_runs: list[tuple[EconomyWorker, dict, dict, Any]] = field( + default_factory=list + ) + + async def start_report_run( + self, + worker: EconomyWorker, + report_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + await self._maybe_fail() + self.report_runs.append( + (worker, report_payload, context_payload, observability_context) + ) + return "fc-coordinator" + + async def start_simulation_run( + self, + worker: EconomyWorker, + simulation_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + await self._maybe_fail() + self.simulation_runs.append( + (worker, simulation_payload, context_payload, observability_context) + ) + return "fc-worker" + + async def _maybe_fail(self) -> None: + if self.hang: + await asyncio.sleep(60) + if self.failure is not None: + raise self.failure + + +def numbered_ids(): + """Return a function giving 00…01, 00…02 and so on, in order.""" + numbers = count(1) + return lambda: UUID(int=next(numbers)) + + +def service( + store: FakeStore, + *, + workers: FakeWorkers | None = None, + caller: FakeCaller | None = None, + output_bucket: str | None = OUTPUT_BUCKET, +) -> EconomySubmissionService: + return EconomySubmissionService( + 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_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index 5546a4668..d3b72af5a 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -410,3 +410,26 @@ def test_the_v2_manifest_settings_configure_no_stage12_resource(): settings.validate() assert not settings.stage12_resources_configured + + +def test_the_v2_output_bucket_is_read_from_the_environment(monkeypatch): + from policyengine_simulation_entry.config import Settings + + monkeypatch.setenv("V2_SIMULATION_OUTPUT_BUCKET", "policyengine-v2-outputs-staging") + + assert Settings.from_env().v2_simulation_output_bucket == ( + "policyengine-v2-outputs-staging" + ) + monkeypatch.delenv("V2_SIMULATION_OUTPUT_BUCKET") + assert Settings.from_env().v2_simulation_output_bucket == "" + + +def test_economy_runs_need_the_database_and_the_v2_modal_environment(): + assert not make_settings().economy_runs_configured + assert not make_settings( + simulation_database_url=DATABASE_URL + ).economy_runs_configured + assert not make_settings(v2_modal_environment="staging").economy_runs_configured + assert make_settings( + simulation_database_url=DATABASE_URL, v2_modal_environment="staging" + ).economy_runs_configured diff --git a/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py b/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py new file mode 100644 index 000000000..f14f28fc7 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py @@ -0,0 +1,152 @@ +"""What the coordinator or a simulation function is given to calculate a run.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from uuid import UUID + +import pytest +from economy_fixtures import STATE_DATASET, economy_manifest +from household_fixtures import POLICYENGINE_VERSION +from policyengine_simulation_contract.stage12_execution import ( + ReportAggregate, + ReportExecutionInput, + ReportRunIds, + SimulationExecutionInput, + SimulationRole, + Stage12InvocationContext, +) +from policyengine_simulation_entry.economy_calculation_inputs import ( + SIMULATION_RUN_ID_FIELD, + CalculationScope, + CallOrigin, + DatasetNotInBundleError, + bundle_dataset, + report_run_calculation, + simulation_options, + simulation_run_calculation, +) +from policyengine_simulation_entry.economy_workers import default_economy_worker + +REPORT_RUN_ID = UUID("00000000-0000-4000-8000-000000000001") +BASELINE_RUN_ID = UUID("00000000-0000-4000-8000-000000000002") +REFORM_RUN_ID = UUID("00000000-0000-4000-8000-000000000003") +CREATED_AT = datetime(2026, 11, 12, 15, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +REFORM = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} +ORIGIN = CallOrigin( + environment="staging", + modal_environment="staging", + request_id="observed-request", + created_at=CREATED_AT, +) +RUN_IDS = ReportRunIds( + report_run_id=REPORT_RUN_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_run_id=REFORM_RUN_ID, +) + + +def _worker(): + return default_economy_worker(economy_manifest(), "us") + + +def _scope(**changes) -> CalculationScope: + values = { + "country": "us", + "region_code": "state/ca", + "dataset": STATE_DATASET, + "year": 2026, + "options": {"include_cliffs": True}, + } + return CalculationScope(**{**values, **changes}) + + +def test_the_bundle_dataset_is_the_one_the_catalog_dataset_names() -> None: + worker = _worker() + + assert bundle_dataset(worker, STATE_DATASET.identity) == STATE_DATASET + with pytest.raises(DatasetNotInBundleError, match="enhanced_cps_2024"): + bundle_dataset(worker, "enhanced_cps_2024") + + +def test_the_options_hold_cliffs_only_when_asked_and_the_spm_selection_given() -> None: + assert simulation_options(include_cliffs=False, spm=None) == {} + assert simulation_options( + include_cliffs=True, spm={"geography_kind": "national"} + ) == {"include_cliffs": True, "spm": {"geography_kind": "national"}} + + +def test_a_report_run_is_named_by_its_run_ids_with_no_evaluation_record() -> None: + worker = _worker() + + calculation = report_run_calculation( + worker, + _scope(), + ORIGIN, + run_ids=RUN_IDS, + baseline_policy={}, + reform_policy=REFORM, + ) + + report = ReportExecutionInput.model_validate(calculation.payload) + assert report.evaluation_id == REPORT_RUN_ID + assert report.requested_aggregates == tuple(ReportAggregate) + for side, role, run_id, policy in ( + (report.baseline, SimulationRole.BASELINE, BASELINE_RUN_ID, {}), + (report.reform, SimulationRole.REFORM, REFORM_RUN_ID, REFORM), + ): + assert (side.role, side.simulation_execution_id, side.policy) == ( + role, + run_id, + policy, + ) + assert side.evaluation_id == REPORT_RUN_ID + assert side.year == 2026 + assert (side.geography.country, side.geography.region) == ("us", "state/ca") + assert side.options == {"include_cliffs": True} + assert side.population.artifact.uri == STATE_DATASET.uri + assert side.population.artifact.content_sha256 == STATE_DATASET.sha256 + assert side.bundle.policyengine_version == POLICYENGINE_VERSION + assert side.bundle.country_package_version == "2.2.1" + assert side.bundle.dataset.identity == STATE_DATASET.identity + assert side.bundle.dataset.artifact_revision == "acs-local-revision" + assert side.bundle.bundle_manifest_sha256 == worker.bundle_manifest_sha256 + + context = Stage12InvocationContext.model_validate(calculation.context) + assert context.report_run == RUN_IDS + assert context.production_function_call_id is None + assert (context.environment, context.modal_environment, context.request_id) == ( + "staging", + "staging", + "observed-request", + ) + assert context.worker_version == worker.version + assert context.modal_application == worker.application_name + assert context.simulation_callable == "run_single_simulation_us" + assert context.version_manifest_sha256 == worker.version_manifest_sha256 + assert context.artifact_prefix == f"stage-12-runs/staging/2026/11/{REPORT_RUN_ID}" + assert context.retention_expires_at - context.created_at == timedelta(days=30) + + +def test_a_standalone_run_is_named_in_its_functions_context() -> None: + worker = _worker() + + calculation = simulation_run_calculation( + worker, + _scope(options={}), + ORIGIN, + simulation_run_id=BASELINE_RUN_ID, + policy=REFORM, + ) + + simulation = SimulationExecutionInput.model_validate(calculation.payload) + assert simulation.role is SimulationRole.STANDALONE + assert simulation.simulation_execution_id == BASELINE_RUN_ID + assert simulation.evaluation_id == BASELINE_RUN_ID + assert simulation.policy == REFORM + assert calculation.context[SIMULATION_RUN_ID_FIELD] == str(BASELINE_RUN_ID) + assert calculation.context["report_run"] is None + assert calculation.context["artifact_prefix"] == ( + f"stage-12-runs/staging/2026/11/{BASELINE_RUN_ID}" + ) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py new file mode 100644 index 000000000..35985f7a5 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -0,0 +1,372 @@ +"""The economy routes' answers, with the submission service faked.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from typing import Any, NoReturn +from uuid import UUID + +import pytest +from conftest import FakeBackend, make_settings +from economy_service_fakes import REPORT_BODY, SIMULATION_BODY +from fastapi.testclient import TestClient +from policyengine_simulation_entry import app as app_module +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportCreated, + EconomyReportRequest, + EconomySimulationCreated, + EconomySimulationRequest, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, + UnprocessableReason, +) +from policyengine_simulation_entry.economy_submissions import ( + EconomySubmissionService, +) +from policyengine_simulation_entry.simulation_database import SimulationDatabase +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"] +CREATED_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, +) +CREATED_SIMULATION = EconomySimulationCreated( + simulation_id=RECORD_ID, simulation_run_id=RUN_ID, status=SimulationStatus.PENDING +) +UNEXPECTED_FAILURE_ANSWER = { + "code": "internal_error", + "detail": "The service could not answer the request.", +} + + +def _raise(failure: Exception) -> NoReturn: + raise failure + + +@dataclass +class FakeEconomySubmissions: + """Records each submission, and answers as told.""" + + failure: Exception | None = None + reports: list[tuple[EconomyReportRequest, str, Any]] = field(default_factory=list) + simulations: list[tuple[EconomySimulationRequest, str, Any]] = field( + default_factory=list + ) + + async def submit_report( + self, + request: EconomyReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + self.reports.append((request, request_id, observability_context)) + if self.failure is not None: + _raise(self.failure) + return CREATED_REPORT + + async def submit_simulation( + self, + request: EconomySimulationRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + self.simulations.append((request, request_id, observability_context)) + if self.failure is not None: + _raise(self.failure) + return CREATED_SIMULATION + + +@pytest.fixture +def client_for(monkeypatch): + """Return a function building a client whose app uses ``submissions``. + + None leaves the app as an unconfigured service builds it. + """ + + def build(submissions: FakeEconomySubmissions | None) -> TestClient: + if submissions is not None: + monkeypatch.setattr( + app_module, + "economy_submissions_for", + lambda settings, database: submissions, + ) + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + ) + ) + + return build + + +ROUTES = [("/reports", REPORT_BODY), ("/simulations", SIMULATION_BODY)] + + +def test_a_report_answers_201_with_its_records_and_runs(client_for) -> None: + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post( + "/reports", + json=REPORT_BODY, + headers={OBSERVABILITY_ID_HEADER: OBSERVABILITY_ID}, + ) + + assert result.status_code == 201 + assert result.json() == CREATED_REPORT.model_dump(mode="json") + assert result.headers[OBSERVABILITY_ID_HEADER] == OBSERVABILITY_ID + ((request, request_id, context),) = submissions.reports + assert request.label == REPORT_BODY["label"] + assert request.options() == { + "include_cliffs": True, + "spm": {"geography_kind": "national"}, + } + assert request_id == OBSERVABILITY_ID + assert context["observability_id"] == OBSERVABILITY_ID + + +def test_a_simulation_answers_201_with_its_record_and_run(client_for) -> None: + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post("/simulations", json=SIMULATION_BODY) + + assert result.status_code == 201 + assert result.json() == { + "simulation_id": str(RECORD_ID), + "simulation_run_id": str(RUN_ID), + "status": "pending", + } + ((request, request_id, _),) = submissions.simulations + assert str(request.policy_id) == SIMULATION_BODY["policy_id"] + assert request_id == result.headers[OBSERVABILITY_ID_HEADER] + + +@pytest.mark.parametrize( + ("path", "body"), + [ + ("/reports", {**REPORT_BODY, "report_type": "household"}), + ("/reports", {**REPORT_BODY, "spm": None}), + ("/reports", {**REPORT_BODY, "include_cliffs": "yes"}), + ("/reports", {**REPORT_BODY, "report_id": str(RECORD_ID)}), + ("/reports", {**REPORT_BODY, "label": ""}), + ("/reports", {**REPORT_BODY, "year": 2201}), + ("/reports", {**REPORT_BODY, "policyengine_version": "6.17.1"}), + ("/reports", {**REPORT_BODY, "reform": {"policy_id": "not-a-uuid"}}), + ("/simulations", {**SIMULATION_BODY, "country": "fr"}), + ("/simulations", {**SIMULATION_BODY, "region_code": "state ca"}), + ("/simulations", {**SIMULATION_BODY, "simulation_id": str(RECORD_ID)}), + ("/simulations", {**SIMULATION_BODY, "simulation_type": "household"}), + ], +) +def test_a_malformed_request_is_refused_with_its_reasons(client_for, path, body): + submissions = FakeEconomySubmissions() + + with client_for(submissions) 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 (submissions.reports, submissions.simulations) == ([], []) + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_a_refusal_answers_400_with_its_code(client_for, path, body) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionRefusedError( + code="unknown_region", detail="no region state/zz" + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 400 + assert result.json() == {"code": "unknown_region", "detail": "no region state/zz"} + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_stored_values_that_cannot_be_calculated_answer_422_with_a_code( + client_for, path, body +) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnprocessableError( + code="missing_parameters", + reasons=( + UnprocessableReason( + loc=("body", "policy_id", "gov.a"), + msg="no such parameter", + type="missing_parameter", + ), + ), + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 422 + assert result.json() == { + "detail": [ + { + "loc": ["body", "policy_id", "gov.a"], + "msg": "no such parameter", + "type": "missing_parameter", + } + ], + "code": "missing_parameters", + } + + +@pytest.mark.parametrize( + ("path", "body", "record_field"), + [ + ("/reports", REPORT_BODY, "report_id"), + ("/simulations", SIMULATION_BODY, "simulation_id"), + ], +) +def test_a_refused_call_answers_503_naming_the_record_it_created( + client_for, path, body, record_field +) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnavailableError( + code="modal_call_refused", detail="Modal refused", record_id=RECORD_ID + ) + ) + + with client_for(submissions) 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", + record_field: str(RECORD_ID), + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unavailable_dependency_answers_503_naming_no_record( + client_for, path, body +) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnavailableError( + code="database_unavailable", detail="The database could not be read." + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 503 + assert result.json() == { + "code": "database_unavailable", + "detail": "The database could not be read.", + } + + +@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.headers["Retry-After"] == "10" + assert result.json() == { + "code": "not_configured", + "detail": "Economy runs are not configured.", + } + + +def _emitted_records(printed: str, message: str) -> list[dict[str, Any]]: + """Return the records with ``message`` the observability runtime printed.""" + records = [ + json.loads(printed_line) + for printed_line in printed.splitlines() + if printed_line.startswith("{") + ] + return [record for record in records if record.get("message") == message] + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unexpected_failure_answers_a_json_500_without_its_values( + client_for, capsys, path, body +) -> None: + submissions = FakeEconomySubmissions( + failure=RuntimeError("a value that must not be logged") + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 500 + assert result.json() == UNEXPECTED_FAILURE_ANSWER + printed = capsys.readouterr().out + (record,) = _emitted_records(printed, "economy_route_failed") + attributes = record["attributes"] + assert (attributes["route"], attributes["error_type"]) == (path, "RuntimeError") + assert "test_economy_routes.py" in attributes["stack_frames"].split("; ")[0] + assert "a value that must not be logged" not in printed + + +DATABASE_URL = "postgresql://simulation_service:secret@db.example:5432/postgres" + + +def test_a_configured_service_builds_economy_submissions_without_connecting(): + """Nothing reaches the database or Modal until the first submission.""" + database = SimulationDatabase.from_url(DATABASE_URL) + try: + submissions = app_module.economy_submissions_for( + make_settings( + simulation_database_url=DATABASE_URL, + v2_modal_environment="staging", + v2_simulation_output_bucket="policyengine-v2-outputs-staging", + ), + database, + ) + finally: + database.dispose() + + assert isinstance(submissions, EconomySubmissionService) + + +@pytest.mark.parametrize( + "settings_changes", + [ + {}, + {"simulation_database_url": DATABASE_URL}, + {"v2_modal_environment": "staging"}, + ], +) +def test_without_its_database_and_v2_environment_none_is_built(settings_changes): + database = SimulationDatabase.from_url(DATABASE_URL) + try: + submissions = app_module.economy_submissions_for( + make_settings(**settings_changes), database + ) + finally: + database.dispose() + + assert submissions is None diff --git a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py new file mode 100644 index 000000000..bc05dd61e --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py @@ -0,0 +1,469 @@ +"""Creating economy reports and simulations, with the database and Modal faked.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from uuid import UUID + +import pytest +from economy_fixtures import STATE_DATASET +from economy_service_fakes import ( + CATALOG, + CREATED_AT, + DYNAMIC_ID, + OUTPUT_BUCKET, + POLICY_ID, + REFORM, + REPORT_BODY, + REUSED, + FakeCaller, + FakeStore, + FakeWorkers, + report_request, + service, + simulation_request, +) +from household_fixtures import POLICYENGINE_VERSION +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + ReportRunIds, + Stage12InvocationContext, +) +from policyengine_simulation_entry.economy_api_models import EconomyReportRequest +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, +) +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) +from policyengine_simulation_persistence import ( + CatalogDataset, + EconomyRecordReferenceError, + MissingPolicyParametersError, + OverlappingReformError, + ReportRunStatus, + RunReuse, + RunVersions, + SimulationChoice, + SimulationStatus, + UnknownCatalogVersionError, + UnknownDatasetError, + UnknownDynamicError, + UnknownPolicyError, + 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"} + + +def _submit_report(submissions, request=None): + return asyncio.run( + submissions.submit_report( + request or report_request(), + request_id="observed", + observability_context=OBSERVED, + ) + ) + + +def _submit_simulation(submissions, request=None): + return asyncio.run( + submissions.submit_simulation( + request or simulation_request(), + request_id="observed", + observability_context=OBSERVED, + ) + ) + + +def test_a_report_is_created_with_new_runs_and_its_coordinator_started() -> None: + store = FakeStore() + caller = FakeCaller() + + created = _submit_report(service(store, caller=caller)) + + (request, (baseline, reform, reuse)) = store.report_reads[0] + assert (request.country, request.policyengine_version, request.region_code) == ( + "us", + POLICYENGINE_VERSION, + "us", + ) + assert (request.dataset_name, request.year) == (None, 2026) + assert baseline == SimulationChoice(policy_id=None, dynamic_id=None) + assert reform == SimulationChoice(policy_id=POLICY_ID, dynamic_id=DYNAMIC_ID) + options = {"include_cliffs": True, "spm": {"geography_kind": "national"}} + assert reuse == RunReuse( + versions=VERSIONS, inputs=options, output_uri_prefix=f"gs://{OUTPUT_BUCKET}/" + ) + (report,) = store.reports + assert (report.id, report.run_id) == (UUID(int=1), UUID(int=2)) + assert (report.label, report.description, report.country) == ( + "A larger child tax credit", + "Raises the credit to $3,000", + "us", + ) + assert report.tax_benefit_model_id == CATALOG.model_id + assert report.population == CATALOG.population(2026) + assert (report.versions, report.inputs) == (VERSIONS, options) + assert (report.baseline.reused, report.reform.reused) == (False, False) + assert report.new_simulation_run_ids == (UUID(int=4), UUID(int=6)) + assert report.reform.choice == reform + ((_, payload, context, observed),) = caller.report_runs + assert observed == OBSERVED + assert ReportExecutionInput.model_validate(payload).reform.policy == REFORM + assert Stage12InvocationContext.model_validate(context).report_run == ReportRunIds( + report_run_id=report.run_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_run_id=report.reform.simulation_run_id, + ) + assert store.recorded_calls == [(report.run_id, "fc-coordinator")] + assert created.model_dump() == { + "report_id": report.id, + "report_run_id": report.run_id, + "status": ReportRunStatus.PENDING, + "baseline_simulation_id": report.baseline.simulation_id, + "baseline_simulation_run_id": report.baseline.simulation_run_id, + "baseline_reused": False, + "reform_simulation_id": report.reform.simulation_id, + "reform_simulation_run_id": report.reform.simulation_run_id, + "reform_reused": False, + } + + +def test_a_side_with_a_reusable_run_names_it_instead_of_a_new_one() -> None: + store = FakeStore(reusable=(REUSED, None)) + caller = FakeCaller() + + created = _submit_report(service(store, caller=caller)) + + (report,) = store.reports + assert report.baseline.reused is True + assert (report.baseline.simulation_id, report.baseline.simulation_run_id) == ( + REUSED.simulation_id, + REUSED.simulation_run_id, + ) + assert report.new_simulation_run_ids == (report.reform.simulation_run_id,) + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.baseline.simulation_execution_id == REUSED.simulation_run_id + assert (created.baseline_reused, created.reform_reused) == (True, False) + assert created.baseline_simulation_run_id == REUSED.simulation_run_id + + +def test_without_the_output_bucket_no_run_is_looked_for() -> None: + store = FakeStore() + + _submit_report(service(store, output_bucket=None)) + + assert store.report_reads[0][1][2] is None + + +def test_the_report_stores_only_the_spm_settings_given() -> None: + store = FakeStore() + caller = FakeCaller() + + _submit_report( + service(store, caller=caller), + report_request(include_cliffs=False, spm={"scenario": "zero"}), + ) + + assert store.reports[0].inputs == { + "include_cliffs": False, + "spm": {"scenario": "zero"}, + } + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.baseline.options == {"spm": {"scenario": "zero"}} + + +def test_a_report_without_spm_settings_stores_none() -> None: + store = FakeStore() + caller = FakeCaller() + without_spm = {key: value for key, value in REPORT_BODY.items() if key != "spm"} + + _submit_report( + service(store, caller=caller), + EconomyReportRequest.model_validate(without_spm), + ) + + assert store.reports[0].inputs == {"include_cliffs": True} + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.reform.options == {"include_cliffs": True} + + +@pytest.mark.parametrize( + ("error", "code"), + [ + ( + UnknownCatalogVersionError("us has no catalog version"), + "unknown_catalog_version", + ), + (UnknownRegionError("no region state/zz"), "unknown_region"), + (UnknownDatasetError("no dataset acs"), "unknown_dataset"), + (UnknownPolicyError("us has no policy"), "unknown_policy"), + (UnknownDynamicError("there is no dynamic"), "unknown_dynamic"), + ], +) +def test_a_missing_row_is_refused_before_anything_is_written(error, code) -> None: + store = FakeStore(read_failure=error) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionRefusedError) as report_refusal: + _submit_report(service(store, caller=caller)) + with pytest.raises(EconomySubmissionRefusedError) as simulation_refusal: + _submit_simulation(service(store, caller=caller)) + + assert report_refusal.value.code == simulation_refusal.value.code == code + assert str(error) in report_refusal.value.detail + assert (store.reports, store.simulations, caller.report_runs) == ([], [], []) + + +def test_values_the_catalog_version_lacks_cannot_be_calculated() -> None: + error = MissingPolicyParametersError(["gov.a", "gov.b"], source="dynamic") + store = FakeStore(read_failure=error) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_report(service(store)) + + assert refusal.value.code == "missing_parameters" + assert [reason.loc for reason in refusal.value.reasons] == [ + ("body", "dynamic_id", "gov.a"), + ("body", "dynamic_id", "gov.b"), + ] + assert POLICYENGINE_VERSION in refusal.value.reasons[0].msg + assert store.reports == [] + + +def test_a_parameter_the_policy_and_dynamic_both_set_cannot_be_calculated() -> None: + store = FakeStore(read_failure=OverlappingReformError(["gov.a"])) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_simulation(service(store)) + + assert refusal.value.code == "overlapping_parameters" + assert [reason.type for reason in refusal.value.reasons] == [ + "overlapping_parameter" + ] + + +@pytest.mark.parametrize( + ("dataset_name", "code"), + [(None, "dataset_not_in_bundle"), ("enhanced_cps_2024", "unknown_dataset")], +) +def test_a_catalog_dataset_the_bundle_lacks_is_refused(dataset_name, code) -> None: + catalog = replace( + CATALOG, dataset=CatalogDataset(id=CATALOG.dataset.id, name="enhanced_cps_2024") + ) + store = FakeStore(catalog=catalog) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(store), report_request(dataset_name=dataset_name)) + + assert refusal.value.code == code + assert store.reports == [] + + +def test_a_named_dataset_in_the_bundle_is_the_one_calculated() -> None: + catalog = replace( + CATALOG, + dataset=CatalogDataset(id=CATALOG.dataset.id, name=STATE_DATASET.identity), + ) + store = FakeStore(catalog=catalog) + caller = FakeCaller() + + _submit_report( + service(store, caller=caller), + report_request(dataset_name=STATE_DATASET.identity), + ) + + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.baseline.bundle.dataset.identity == STATE_DATASET.identity + + +def test_a_country_the_default_worker_cannot_calculate_is_refused() -> None: + store = FakeStore() + workers = FakeWorkers(EconomyWorkerUnavailableError("cannot calculate uk")) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(store, workers=workers)) + + assert refusal.value.code == "no_economy_worker" + assert store.report_reads == [] + + +def test_a_manifest_that_cannot_be_read_is_unavailable() -> None: + failure = VersionManifestUnavailableError("unreadable") + failure.__cause__ = ConnectionError("Modal") + store = FakeStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_report(service(store, workers=FakeWorkers(failure))) + + assert (unavailable.value.code, unavailable.value.record_id) == ( + "manifest_unavailable", + None, + ) + + +@pytest.mark.parametrize("failing", ["read", "create"]) +def test_a_database_that_cannot_be_used_is_unavailable_before_any_call( + failing, +) -> None: + store = FakeStore(**{f"{failing}_failure": DATABASE_ERROR}) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_report(service(store, caller=caller)) + + assert (unavailable.value.code, unavailable.value.record_id) == ( + "database_unavailable", + None, + ) + assert caller.report_runs == [] + + +@pytest.mark.parametrize( + ("column", "code"), + [ + ("region_id", "unknown_region"), + ("dataset_id", "unknown_dataset"), + ("dynamic_id", "unknown_dynamic"), + ("policy_id", "unknown_policy"), + ("tax_benefit_model_version_id", "unknown_catalog_version"), + ("baseline_simulation_run_id", "missing_reference"), + (None, "missing_reference"), + ], +) +def test_a_row_deleted_after_the_reads_is_refused_by_its_column(column, code) -> None: + store = FakeStore( + create_failure=EconomyRecordReferenceError( + None if column is None else "simulations", column + ) + ) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(store, caller=caller)) + + assert refusal.value.code == code + assert "no longer exists" in refusal.value.detail + assert caller.report_runs == [] + + +@pytest.mark.parametrize( + "caller", + [FakeCaller(failure=ConnectionError("refused")), FakeCaller(hang=True)], + ids=["refused", "not accepted in time"], +) +def test_a_call_modal_does_not_confirm_fails_the_report_run_and_names_it( + caller, +) -> None: + store = FakeStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_report(service(store, caller=caller)) + + (report,) = store.reports + assert (unavailable.value.code, unavailable.value.record_id) == ( + "modal_call_refused", + report.id, + ) + ((run_id, simulation_run_ids, message),) = store.failed_runs + assert run_id == report.run_id + assert simulation_run_ids == report.new_simulation_run_ids + assert message.startswith("modal_call_refused: Modal did not accept the report") + assert store.recorded_calls == [] + + +def test_a_report_run_its_coordinator_claimed_meanwhile_is_answered_created() -> None: + store = FakeStore(claimed=True) + + created = _submit_report( + service(store, caller=FakeCaller(failure=TimeoutError("late"))) + ) + + assert store.failed_runs == [] + assert created.report_id == store.reports[0].id + assert created.status is ReportRunStatus.PENDING + + +@pytest.mark.parametrize("failing", ["fail", "record"]) +def test_a_run_whose_state_cannot_be_written_is_answered_created_pending( + failing, +) -> None: + store = FakeStore(**{f"{failing}_failure": DATABASE_ERROR}) + caller = FakeCaller( + failure=ConnectionError("refused") if failing == "fail" else None + ) + + created = _submit_report(service(store, caller=caller)) + + assert created.report_id == store.reports[0].id + assert created.status is ReportRunStatus.PENDING + + +def test_a_simulation_is_created_and_its_run_started_alone() -> None: + store = FakeStore() + caller = FakeCaller() + + created = _submit_simulation(service(store, caller=caller)) + + (request, choice) = store.simulation_reads[0] + assert (request.policyengine_version, request.year) == (POLICYENGINE_VERSION, 2026) + assert choice == SimulationChoice(policy_id=POLICY_ID, dynamic_id=None) + (simulation,) = store.simulations + assert (simulation.id, simulation.run_id) == (UUID(int=1), UUID(int=2)) + assert simulation.population == CATALOG.population(2026) + assert simulation.versions == VERSIONS + ((_, payload, context, observed),) = caller.simulation_runs + assert payload["simulation_execution_id"] == str(simulation.run_id) + assert payload["policy"] == REFORM + assert context["simulation_run_id"] == str(simulation.run_id) + assert observed == OBSERVED + assert store.recorded_calls == [(simulation.run_id, "fc-worker")] + assert created.model_dump() == { + "simulation_id": simulation.id, + "simulation_run_id": simulation.run_id, + "status": SimulationStatus.PENDING, + } + + +def test_a_simulation_run_modal_refuses_is_failed_and_the_simulation_named() -> None: + store = FakeStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_simulation( + service(store, caller=FakeCaller(failure=ConnectionError("refused"))) + ) + + (simulation,) = store.simulations + assert (unavailable.value.code, unavailable.value.record_id) == ( + "modal_call_refused", + simulation.id, + ) + assert store.failed_runs[0][0] == simulation.run_id + assert store.failed_runs[0][2].startswith( + "modal_call_refused: Modal did not accept the simulation" + ) + + +def test_the_calls_context_records_when_and_where_it_was_made() -> None: + caller = FakeCaller() + + _submit_report(service(FakeStore(), caller=caller)) + + context = Stage12InvocationContext.model_validate(caller.report_runs[0][2]) + assert (context.created_at, context.request_id, context.environment) == ( + CREATED_AT, + "observed", + "staging", + ) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_workers.py b/projects/policyengine-simulation-entry/tests/test_economy_workers.py new file mode 100644 index 000000000..9192e4f27 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_workers.py @@ -0,0 +1,196 @@ +"""The default v2 worker's economy functions, and starting their calls on Modal.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +import pytest +from economy_fixtures import economy_manifest, economy_worker_version +from household_fixtures import ( + BUNDLE_MANIFEST_SHA256, + POLICYENGINE_VERSION, + household_worker_version, + manifest_store, +) +from policyengine_simulation_contract.stage12_manifest import ( + V2CountryWorker, + manifest_digest, + v2_application_name, +) +from policyengine_simulation_entry import economy_workers +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, + ManifestEconomyWorkers, + ModalEconomyRunCaller, + default_economy_worker, +) +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) +from policyengine_simulation_persistence import RunVersions + + +def test_the_default_workers_economy_functions_and_bundle_are_found() -> None: + manifest = economy_manifest() + + worker = default_economy_worker(manifest, "us") + + assert worker.version == POLICYENGINE_VERSION + assert worker.version_manifest_sha256 == manifest_digest(manifest) + assert worker.application_name == v2_application_name(POLICYENGINE_VERSION) + assert worker.coordinator_callable == "coordinate_report" + assert worker.simulation_callable == "run_single_simulation_us" + assert worker.bundle_manifest_sha256 == BUNDLE_MANIFEST_SHA256 + assert worker.country.country == "us" + assert worker.versions == RunVersions( + country_package_version="2.2.1", policyengine_version=POLICYENGINE_VERSION + ) + + +def test_a_country_the_default_worker_cannot_calculate_is_refused() -> None: + base = economy_worker_version() + household_only = base.model_copy( + update={ + "countries": tuple( + V2CountryWorker( + country=item.country, + economy_capable=False, + single_simulation_callable=None, + household_callable=item.household_callable, + ) + if item.country == "uk" + else item + for item in base.countries + ), + "validation": base.validation.model_copy( + update={"country_validation_invocation_ids": {"us": "validation-us"}} + ), + } + ) + manifest = economy_manifest(household_only) + + with pytest.raises(EconomyWorkerUnavailableError, match="uk"): + default_economy_worker(manifest, "uk") + assert default_economy_worker(manifest, "us").country.country == "us" + + +def test_the_manifest_store_is_read_for_the_default_worker() -> None: + newer = household_worker_version("6.18.0") + workers = ManifestEconomyWorkers(manifest_store(household_worker_version(), newer)) + + worker = workers.default_worker("uk") + + assert worker.version == POLICYENGINE_VERSION + assert worker.simulation_callable == "run_single_simulation_uk" + + +class UnreadableStore: + def get(self, key: str, default: object = None) -> object: + msg = "Modal is unreachable" + raise ConnectionError(msg) + + def __setitem__(self, key: str, value: object) -> None: + raise AssertionError + + +def test_a_manifest_that_cannot_be_read_is_unavailable() -> None: + with pytest.raises(VersionManifestUnavailableError) as raised: + ManifestEconomyWorkers(UnreadableStore()).default_worker("us") + + assert isinstance(raised.value.__cause__, ConnectionError) + + +class FakeFunction: + """A Modal function looked up by name, recording each spawn.""" + + def __init__(self, object_id: object) -> None: + self.object_id = object_id + self.spawns: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.spawn = SimpleNamespace(aio=self._spawn) + + async def _spawn(self, *arguments: Any, **keywords: Any) -> SimpleNamespace: + self.spawns.append((arguments, keywords)) + return SimpleNamespace(object_id=self.object_id) + + +@pytest.fixture +def modal_functions(monkeypatch): + """Replace Modal's lookup by name; return the lookups and the function.""" + looked_up: list[tuple[str, str, str]] = [] + function = FakeFunction("fc-started") + + def from_name(application: str, name: str, *, environment_name: str): + looked_up.append((application, name, environment_name)) + return function + + monkeypatch.setattr( + economy_workers.modal.Function, "from_name", staticmethod(from_name) + ) + return looked_up, function + + +def test_a_report_run_starts_the_coordinator_with_its_report_and_context( + modal_functions, +) -> None: + looked_up, function = modal_functions + worker = default_economy_worker(economy_manifest(), "us") + + call_id = asyncio.run( + ModalEconomyRunCaller("staging").start_report_run( + worker, + {"report": "payload"}, + {"report_run": "context"}, + observability_context={"observability_id": "observed"}, + ) + ) + + assert call_id == "fc-started" + assert looked_up == [(worker.application_name, "coordinate_report", "staging")] + # No evaluation record follows the context. + assert function.spawns == [ + ( + ({"report": "payload"}, {"report_run": "context"}), + {"observability_context": {"observability_id": "observed"}}, + ) + ] + + +def test_a_standalone_run_starts_the_countrys_simulation_function( + modal_functions, +) -> None: + looked_up, function = modal_functions + worker = default_economy_worker(economy_manifest(), "uk") + + call_id = asyncio.run( + ModalEconomyRunCaller("main").start_simulation_run( + worker, + {"simulation": 1}, + {"simulation_run_id": "run"}, + observability_context=None, + ) + ) + + assert call_id == "fc-started" + assert looked_up == [(worker.application_name, "run_single_simulation_uk", "main")] + assert function.spawns == [ + ( + ({"simulation": 1}, {"simulation_run_id": "run"}), + {"observability_context": None}, + ) + ] + + +@pytest.mark.parametrize("object_id", [None, ""]) +def test_a_call_modal_gives_no_id_is_an_error(modal_functions, object_id) -> None: + _, function = modal_functions + function.object_id = object_id + worker = default_economy_worker(economy_manifest(), "us") + + with pytest.raises(RuntimeError, match="no call ID"): + asyncio.run( + ModalEconomyRunCaller("staging").start_report_run( + worker, {}, {}, observability_context=None + ) + ) diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 631db41d6..95beaad89 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -37,8 +37,10 @@ class OpenAPIDocument(TypedDict): ("GET", "/versions/{kind}"), ("POST", "/household-calculations"), ("POST", "/ping"), + ("POST", "/reports"), ("POST", "/simulate/economy/budget-window"), ("POST", "/simulate/economy/comparison"), + ("POST", "/simulations"), } @@ -77,11 +79,28 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: ) +# The economy report and simulation routes, which only the Cloud Run +# entrypoint has. +ECONOMY_PATHS = ("/reports", "/simulations") +ECONOMY_SCHEMAS = ( + "EconomyError", + "EconomyReportCreated", + "EconomyReportRequest", + "EconomyRequestReason", + "EconomyRequestUnprocessable", + "EconomySimulationCreated", + "EconomySimulationRequest", + "PolicyChoice", + "ReportRunStatus", + "SimulationStatus", +) + + def normalized_compatibility_paths(spec: OpenAPIDocument) -> OpenAPIPathMap: """Remove intentional Cloud Run and additive SPM error documentation.""" paths = deepcopy(spec["paths"]) paths.pop("/ready", None) - for path in HOUSEHOLD_PATHS: + for path in (*HOUSEHOLD_PATHS, *ECONOMY_PATHS): paths.pop(path, None) for operations in paths.values(): for method, operation in operations.items(): @@ -212,7 +231,7 @@ def test_normalized_contract_matches_old_gateway(): cloud_run_schemas.pop("SimulationEntrypointReadinessResponse") cloud_run_schemas.pop("SimulationDatabaseReadiness") cloud_run_schemas.pop("SimulationErrorResponse") - for schema_name in HOUSEHOLD_SCHEMAS: + for schema_name in (*HOUSEHOLD_SCHEMAS, *ECONOMY_SCHEMAS): cloud_run_schemas.pop(schema_name) assert cloud_run_schemas == gateway_spec["components"]["schemas"] gateway_operation_ids = operation_ids(gateway_spec) @@ -323,3 +342,84 @@ def test_a_household_calculation_request_names_its_fields_only(): "succeeded", "failed", ] + + +def test_economy_routes_declare_every_answer(): + spec = create_app().openapi() + error = {"$ref": "#/components/schemas/EconomyError"} + unprocessable = {"$ref": "#/components/schemas/EconomyRequestUnprocessable"} + + for path, created in ( + ("/reports", "EconomyReportCreated"), + ("/simulations", "EconomySimulationCreated"), + ): + responses = _responses(spec, "post", path) + assert set(responses) == {"201", "400", "403", "422", "500", "503"} + assert _response_schema(responses["201"]) == { + "$ref": f"#/components/schemas/{created}" + } + for status_code in ("400", "500", "503"): + assert _response_schema(responses[status_code]) == error + assert _response_schema(responses["422"]) == unprocessable + record = path.removeprefix("/").removesuffix("s") + 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) + codes = set(error_properties["code"]["enum"]) + assert { + "unknown_region", + "unknown_dataset", + "unknown_dynamic", + "unknown_catalog_version", + "modal_call_refused", + "internal_error", + } <= codes + + +def test_every_economy_503_declares_its_retry_after_header(): + spec = create_app().openapi() + + for path in ECONOMY_PATHS: + headers = _responses(spec, "post", path)["503"]["headers"] + assert headers["Retry-After"]["schema"] == {"type": "integer"} + + +def test_economy_routes_name_their_operations(): + names = operation_ids(cast("OpenAPIDocument", create_app().openapi())) + + assert names[("POST", "/reports")] == "create_report" + assert names[("POST", "/simulations")] == "create_simulation" + + +def test_economy_requests_name_no_version_and_no_record_id(): + """The service runs its default worker's version and chooses every ID.""" + schemas = create_app().openapi()["components"]["schemas"] + + assert set(schemas["EconomyReportRequest"]["properties"]) == { + "country", + "report_type", + "label", + "description", + "year", + "region_code", + "dataset_name", + "baseline", + "reform", + "include_cliffs", + "spm", + } + assert set(schemas["EconomySimulationRequest"]["properties"]) == { + "country", + "simulation_type", + "year", + "region_code", + "dataset_name", + "policy_id", + "dynamic_id", + } + # Omitted, not null, when the caller gives no SPM settings. + spm = schemas["EconomyReportRequest"]["properties"]["spm"] + assert spm["$ref"] == "#/components/schemas/SPMSelection" + assert "default" not in spm + assert "spm" not in schemas["EconomyReportRequest"].get("required", []) From d570a2ab39432047945d044fbb37a456483f2bf1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:33:20 +0400 Subject: [PATCH 08/26] Create economy reports and simulations through POST /reports and POST /simulations The Simulation Entrypoint's new routes find the v2 version manifest's default worker's economy functions, read the request's catalog rows and stored values (economy_submissions), create the records and runs in one transaction, and start the run's Modal call by name: the report run's coordinator with the report and its run IDs, and no Stage 12 evaluation record, or a standalone simulation run's simulation function with its run's ID (economy_calculation_inputs, economy_workers). They record the call's ID, or fail the run, and its new simulation runs, and answer 503 naming the record. A report side reuses a succeeded run when V2_SIMULATION_OUTPUT_BUCKET names the output bucket. The routes answer 201, 400 with the Public API's caller-row codes or the service's own, 422, 503 with Retry-After and a JSON 500, never with an exception's text. The household routes' failure log helpers move, unchanged, to route_support, which both use. Co-Authored-By: Claude Opus 5.5 --- .../src/policyengine_simulation_entry/app.py | 64 +- .../policyengine_simulation_entry/config.py | 10 + .../economy_api_models.py | 240 +++++++ .../economy_calculation_inputs.py | 259 ++++++++ .../economy_routes.py | 314 +++++++++ .../economy_submission_errors.py | 180 +++++ .../economy_submissions.py | 628 ++++++++++++++++++ .../economy_workers.py | 245 +++++++ .../household_routes.py | 113 +--- .../route_support.py | 124 ++++ 10 files changed, 2075 insertions(+), 102 deletions(-) create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.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 2d8bbb804..2ca53a5c4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -41,7 +41,7 @@ from policyengine_simulation_observability.observability import ( init_simulation_observability, ) -from policyengine_simulation_persistence import HouseholdJobStore +from policyengine_simulation_persistence import EconomyRunStore, HouseholdJobStore from starlette.datastructures import MutableHeaders from policyengine_simulation_entry.auth import ( @@ -52,6 +52,18 @@ ) 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_workers import ( + ManifestEconomyWorkers, + ModalEconomyRunCaller, +) from policyengine_simulation_entry.entry_runtime import ( BACKEND_RESPONSE_HEADER, STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS, @@ -110,6 +122,8 @@ SIMULATION_DATABASE_READINESS_TIMEOUT_SECONDS = 3.0 # The database ends each statement of a household route that runs longer. HOUSEHOLD_STATEMENT_TIMEOUT_MILLISECONDS = 5000 +# The database ends each statement of an economy route that runs longer. +ECONOMY_STATEMENT_TIMEOUT_MILLISECONDS = 5000 def _temporary_submission_payload( @@ -523,6 +537,48 @@ def household_calculations_for( ) +def _register_economy_routes( + app: FastAPI, + entry: EntryRuntime, + submissions: EconomySubmissions | None, + *, + submitter: AuthorizationDependency, +) -> None: + """Add the routes that create economy reports and simulations.""" + add_economy_routes(app, entry, submissions, submitter=submitter) + + +def economy_submissions_for( + settings: Settings, database: SimulationDatabaseConnection | None +) -> EconomySubmissions | None: + """Build the economy submission service when the service is configured. + + It needs the simulation database and the v2 Modal environment; without + them the economy routes answer 503. Nothing is read from Modal until the + first submission. Without the output bucket, no report reuses a run. + """ + if not settings.economy_runs_configured or not isinstance( + database, SimulationDatabase + ): + return None + return EconomySubmissionService( + store=EconomyRunStore.from_engine( + database.engine, + statement_timeout_milliseconds=ECONOMY_STATEMENT_TIMEOUT_MILLISECONDS, + ), + workers=ManifestEconomyWorkers.from_modal( + manifest_name=settings.v2_manifest_name, + environment=settings.v2_modal_environment, + ), + caller=ModalEconomyRunCaller(settings.v2_modal_environment), + settings=ServiceSettings( + environment=settings.environment, + modal_environment=settings.v2_modal_environment, + output_bucket=settings.v2_simulation_output_bucket or None, + ), + ) + + def _register_service_routes(app: FastAPI, entry: EntryRuntime) -> None: """Add the public version, health, readiness and ping routes.""" database_readiness = ( @@ -678,6 +734,12 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: submitter=simulation_submitter, reader=simulation_reader, ) + _register_economy_routes( + app, + entry, + economy_submissions_for(runtime_settings, runtime_database), + submitter=simulation_submitter, + ) _register_service_routes(app, entry) return app diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index 883b3c077..5b6ae75dd 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -102,6 +102,10 @@ class Settings: # manifest, and the manifest's name. Empty outside deployed environments. v2_modal_environment: str = "" v2_manifest_name: str = V2_VERSION_MANIFEST_NAME + # The bucket that keeps v2 simulation output files for good. An economy + # report reuses only a succeeded run whose file is in it; without it, no + # run is reused. Empty outside deployed environments. + v2_simulation_output_bucket: str = "" @classmethod def from_env(cls) -> Settings: @@ -157,6 +161,7 @@ def from_env(cls) -> Settings: v2_manifest_name=os.getenv( "SIMULATION_V2_MANIFEST_NAME", V2_VERSION_MANIFEST_NAME ), + v2_simulation_output_bucket=os.getenv("V2_SIMULATION_OUTPUT_BUCKET", ""), ) @property @@ -180,6 +185,11 @@ def household_calculations_configured(self) -> bool: """Whether the database and the v2 Modal environment are both set.""" return bool(self.simulation_database_url and self.v2_modal_environment) + @property + def economy_runs_configured(self) -> bool: + """Whether economy reports and simulations can run: as household ones.""" + return self.household_calculations_configured + def validate(self) -> None: """Raise ConfigurationError unless the service can start safely. 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 new file mode 100644 index 000000000..7b2134637 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -0,0 +1,240 @@ +"""Request and response models of the economy simulation and report routes. + +The Public API names a region and dataset of the country's catalog by code +and name, and each simulation's stored policy and dynamic by ID; the +Simulation Entrypoint reads them from the Public API's tables itself, and +calculates on the v2 version manifest's default worker. A request names no +.py version and no record ID: the service runs its default worker's version +and chooses every ID. + +Answers the Public API reads: + +- 201 with the created record's ID and its first run's ID; +- 400 with ``code`` and ``detail``: one of ``CALLER_ROW_REFUSAL_CODES`` 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; +- 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; +- 500 with ``code`` ``internal_error``, possibly after a write. + +``detail`` never holds an exception's text, and ``code`` is a fixed +identifier. +""" + +from __future__ import annotations + +from typing import Annotated, Any, Final, Literal +from uuid import UUID + +from policyengine_simulation_contract.spm import SPMSelection +from policyengine_simulation_contract.stage12_bundle import CountryId +from policyengine_simulation_persistence import ReportRunStatus, SimulationStatus +from pydantic import BaseModel, ConfigDict, Field, StrictBool, field_validator +from pydantic.json_schema import SkipJsonSchema + +# The bounds the tables accept. +SimulationYear = Annotated[int, Field(ge=1900, le=2200, description="The year")] +CatalogKey = Annotated[str, Field(min_length=1, max_length=255, pattern=r"^\S+$")] +ReportLabel = Annotated[str, Field(min_length=1, max_length=255)] +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. +CallerRowRefusalCode = Literal["unknown_region", "unknown_dataset", "unknown_dynamic"] +CALLER_ROW_REFUSAL_CODES: Final = frozenset( + {"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. +ServiceStateRefusalCode = Literal[ + # The country has no catalog version for the default worker's .py version. + "unknown_catalog_version", + "unknown_policy", + # The default worker cannot calculate the country's economy runs. + "no_economy_worker", + # The default worker's bundle lacks the region's default dataset. + "dataset_not_in_bundle", + # The database refused a row: a row it names was deleted after the reads. + "missing_reference", +] +EconomyRefusalCode = Literal[CallerRowRefusalCode, ServiceStateRefusalCode] +# Why stored values cannot be calculated (422). +EconomyUnprocessableCode = Literal["missing_parameters", "overlapping_parameters"] +# Why a route cannot answer now (503); retry after ``Retry-After`` seconds. +EconomyUnavailableCode = Literal[ + # The service has no database or v2 Modal environment configured. + "not_configured", + "manifest_unavailable", + "database_unavailable", + # Modal refused the run's call, or did not accept it in time; the record + # named was created and its run failed at once. + "modal_call_refused", +] +# Every code an economy route's error answer carries. +EconomyErrorCode = Literal[ + EconomyRefusalCode, + EconomyUnprocessableCode, + EconomyUnavailableCode, + "internal_error", +] + + +def _omit_null_default(schema: dict[str, Any]) -> None: + """Leave out a field's null default: the field is omitted rather than null.""" + schema.pop("default", None) + + +class _Request(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False) + + +class PolicyChoice(_Request): + """A simulation's stored policy and dynamic; neither means current law.""" + + policy_id: UUID | None = None + dynamic_id: UUID | None = None + + +class EconomySimulationRequest(PolicyChoice): + """An economy simulation: a policy and dynamic applied to one region's data. + + ``region_code`` is a region of the country's catalog version for the + default worker's .py version (``us``, ``state/ca``, ``country/wales``), + and ``dataset_name`` one of that version's datasets; without it, the + region's default dataset is used. + """ + + country: CountryId + simulation_type: Literal["economy"] + year: SimulationYear + region_code: CatalogKey + dataset_name: CatalogKey | None = None + + +class EconomyReportRequest(_Request): + """An economy report: a baseline and a reform simulation of one region's data. + + Both simulations use the report's year, region and dataset, as + ``EconomySimulationRequest`` describes them. ``include_cliffs`` also + calculates the baseline's and reform's cliff impact. ``spm`` selects the + SPM forecast scenario and geography; omit it, rather than sending null, + to use the bundle's certified defaults. + """ + + country: CountryId + report_type: Literal["economy"] + label: ReportLabel + description: ReportDescription | None = None + year: SimulationYear + region_code: CatalogKey + dataset_name: CatalogKey | None = None + baseline: PolicyChoice + reform: PolicyChoice + include_cliffs: StrictBool = False + spm: SPMSelection | SkipJsonSchema[None] = Field( + default=None, json_schema_extra=_omit_null_default + ) + + @field_validator("spm", mode="before") + @classmethod + def refuse_null_spm(cls, value: object) -> object: + """Refuse an explicit null; omit ``spm`` to send no settings.""" + if value is None: + msg = "spm must be an object when given; omit it to send no settings" + raise ValueError(msg) + return value + + def options(self) -> dict[str, Any]: + """Return the report's options as it stores them (``reports.inputs``). + + ``spm`` holds only the settings the request gave, and is left out + when it gave none. + """ + options: dict[str, Any] = {"include_cliffs": self.include_cliffs} + if self.spm is not None: + options["spm"] = self.spm.model_dump(mode="json") + return options + + +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 status is the run's as the service last wrote or read it: ``pending`` + unless the simulation function had claimed the run by then. + """ + + simulation_id: UUID + simulation_run_id: UUID + status: SimulationStatus + + +class EconomyReportCreated(_Answer): + """The report a request created, its first run and that run's status. + + The status is the run's as the service last wrote or read it: ``pending`` + unless the coordinator had claimed the run by then. Each side names its + simulation and simulation run; ``…_reused`` is true for a side that + names a run that already succeeded, rather than a new one. + """ + + report_id: UUID + report_run_id: UUID + status: ReportRunStatus + baseline_simulation_id: UUID + baseline_simulation_run_id: UUID + baseline_reused: bool + reform_simulation_id: UUID + reform_simulation_run_id: UUID + reform_reused: bool + + +class EconomyError(_Answer): + """Why an economy route refused a request or cannot answer now. + + ``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. + """ + + code: EconomyErrorCode + detail: str + report_id: UUID | None = None + simulation_id: UUID | None = None + + +class EconomyRequestReason(BaseModel): + """One reason a request cannot be calculated, as FastAPI writes a field error. + + ``loc`` starts with the request field the reason traces to. FastAPI's own + field errors may also carry ``input`` and ``ctx``. + """ + + model_config = ConfigDict(extra="allow", frozen=True) + + loc: list[str | int] + msg: str + type: str + + +class EconomyRequestUnprocessable(_Answer): + """A request that cannot be calculated (422), one reason per refused field. + + Request fields FastAPI refuses carry no ``code``. Stored values that + cannot be calculated carry one: ``missing_parameters`` for a policy or + dynamic setting parameters the catalog version lacks, and + ``overlapping_parameters`` for a policy and dynamic setting the same one. + """ + + detail: list[EconomyRequestReason] + code: EconomyUnprocessableCode | SkipJsonSchema[None] = Field( + default=None, json_schema_extra=_omit_null_default + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py new file mode 100644 index 000000000..5154e57c1 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py @@ -0,0 +1,259 @@ +"""What the coordinator or a simulation function is given to calculate a run. + +A report run's coordinator takes a ``ReportExecutionInput`` naming the +report run as its ``evaluation_id`` and each side's simulation run as that +side's ``simulation_execution_id``, and a ``Stage12InvocationContext`` whose +``report_run`` names the three runs. It claims the report run and gives +each simulation function the run it claims; no evaluation record is passed. +A standalone simulation run's function takes that run's +``SimulationExecutionInput`` and a context naming the run +(``simulation_run_id``). + +Both sides of a report calculate on the same dataset, region, year, +options and bundle; only their policies differ. The dataset is the bundle's +dataset whose identity is the catalog dataset's name, the region is the +catalog region's code, which .py's region registry reads, and the policies +are the stored policy and dynamic values by parameter name and period. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final + +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + ReportAggregate, + ReportExecutionInput, + ReportRunIds, + SimulationExecutionInput, + SimulationRole, + Stage12InvocationContext, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_contract.stage12_bundle import Stage12Dataset + from policyengine_simulation_contract.stage12_execution import CountryId + + from policyengine_simulation_entry.economy_workers import EconomyWorker + +# Every aggregate a report computes, as the coordinator requires. +REPORT_AGGREGATES: Final = tuple(ReportAggregate) +# How long the coordinator keeps a call's Stage 12 artifacts; the longest the +# contract allows. Simulation output files go to the output bucket instead. +ARTIFACT_RETENTION: Final = timedelta(days=30) +# The context field naming a standalone simulation run's own run. +SIMULATION_RUN_ID_FIELD: Final = "simulation_run_id" + + +class DatasetNotInBundleError(LookupError): + """The default worker's bundle has no dataset with the catalog dataset's name.""" + + +@dataclass(frozen=True) +class CalculationScope: + """What every simulation of a request calculates on, as the worker reads it. + + ``options`` are the simulation options both sides share: ``include_cliffs`` + when true, and the SPM selection when one was given. + """ + + country: CountryId + region_code: str + dataset: Stage12Dataset + year: int + options: dict[str, Any] + + +@dataclass(frozen=True) +class CallOrigin: + """Where a call comes from: the deployment, the request and when it was made.""" + + environment: str + modal_environment: str + request_id: str + created_at: datetime + + +@dataclass(frozen=True) +class RunCalculation: + """The two arguments the run's Modal function is called with.""" + + payload: dict[str, Any] + context: dict[str, Any] + + +def bundle_dataset(worker: EconomyWorker, dataset_name: str) -> Stage12Dataset: + """Return the worker's bundle dataset whose identity is ``dataset_name``. + + Raise DatasetNotInBundleError when it has none. + """ + dataset = next( + (item for item in worker.country.datasets if item.identity == dataset_name), + None, + ) + if dataset is None: + msg = ( + f"the default v2 worker, .py {worker.version}, has no dataset " + f"{dataset_name}" + ) + raise DatasetNotInBundleError(msg) + return dataset + + +def simulation_options( + *, include_cliffs: bool, spm: dict[str, Any] | None +) -> dict[str, Any]: + """Return the options both sides of a calculation share.""" + options: dict[str, Any] = {} + if include_cliffs: + options["include_cliffs"] = True + if spm is not None: + options["spm"] = spm + return options + + +def _simulation_input( # noqa: PLR0913 # the run, its role, its policy and what it calculates + worker: EconomyWorker, + scope: CalculationScope, + *, + evaluation_id: UUID, + simulation_run_id: UUID, + role: SimulationRole, + policy: dict[str, Any], +) -> SimulationExecutionInput: + country = worker.country + dataset = scope.dataset + return SimulationExecutionInput( + evaluation_id=evaluation_id, + simulation_execution_id=simulation_run_id, + role=role, + policy=policy, + population=DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256=dataset.sha256, + ) + ), + year=scope.year, + geography=GeographySelection(country=scope.country, region=scope.region_code), + options=scope.options, + bundle=BundleProvenance( + policyengine_version=worker.policyengine_version, + country_package_name=country.country_package_name, + country_package_version=country.country_package_version, + dataset=DatasetProvenance( + identity=dataset.identity, + uri=dataset.uri, + artifact_revision=dataset.artifact_revision, + data_package_name=country.data_package_name, + data_package_version=country.data_package_version, + ), + bundle_manifest_sha256=worker.bundle_manifest_sha256, + ), + ) + + +def _context( + worker: EconomyWorker, + origin: CallOrigin, + *, + run_id: UUID, + report_run: ReportRunIds | None, +) -> Stage12InvocationContext: + created_at = origin.created_at + return Stage12InvocationContext( + request_id=origin.request_id, + environment=origin.environment, + modal_environment=origin.modal_environment, + worker_version=worker.version, + modal_application=worker.application_name, + simulation_callable=worker.simulation_callable, + version_manifest_sha256=worker.version_manifest_sha256, + bundle_manifest_sha256=worker.bundle_manifest_sha256, + artifact_prefix=( + f"stage-12-runs/{origin.environment}/{created_at:%Y}/{created_at:%m}/" + f"{run_id}" + ), + created_at=created_at, + retention_expires_at=created_at + ARTIFACT_RETENTION, + report_run=report_run, + ) + + +def report_run_calculation( # noqa: PLR0913 # the runs, both policies and what they calculate + worker: EconomyWorker, + scope: CalculationScope, + origin: CallOrigin, + *, + run_ids: ReportRunIds, + baseline_policy: dict[str, Any], + reform_policy: dict[str, Any], +) -> RunCalculation: + """Return the coordinator's report and context for the report run.""" + report_run_id = run_ids.report_run_id + report = ReportExecutionInput( + evaluation_id=report_run_id, + baseline=_simulation_input( + worker, + scope, + evaluation_id=report_run_id, + simulation_run_id=run_ids.baseline_simulation_run_id, + role=SimulationRole.BASELINE, + policy=baseline_policy, + ), + reform=_simulation_input( + worker, + scope, + evaluation_id=report_run_id, + simulation_run_id=run_ids.reform_simulation_run_id, + role=SimulationRole.REFORM, + policy=reform_policy, + ), + requested_aggregates=REPORT_AGGREGATES, + ) + context = _context(worker, origin, run_id=report_run_id, report_run=run_ids) + return RunCalculation( + payload=report.model_dump(mode="json"), + context=context.model_dump(mode="json"), + ) + + +def simulation_run_calculation( + worker: EconomyWorker, + scope: CalculationScope, + origin: CallOrigin, + *, + simulation_run_id: UUID, + policy: dict[str, Any], +) -> RunCalculation: + """Return the simulation function's input and context for a standalone run. + + The context names the run the function claims and records. + """ + simulation = _simulation_input( + worker, + scope, + evaluation_id=simulation_run_id, + simulation_run_id=simulation_run_id, + role=SimulationRole.STANDALONE, + policy=policy, + ) + context = _context(worker, origin, run_id=simulation_run_id, report_run=None) + return RunCalculation( + payload=simulation.model_dump(mode="json"), + context={ + **context.model_dump(mode="json"), + SIMULATION_RUN_ID_FIELD: str(simulation_run_id), + }, + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py new file mode 100644 index 000000000..d0924244a --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py @@ -0,0 +1,314 @@ +"""The economy routes: create a report or a simulation, and start its calculation. + +- ``POST /reports`` (permission ``simulation:submit``) creates an economy + report, its simulations and its first run, and starts the report run's + coordinator, answering 201 with the report's ID, its run's ID and status, + and each side's simulation and simulation run. +- ``POST /simulations`` (``simulation:submit``) creates a standalone + economy simulation and its first run, and starts the run's simulation + function, answering 201 with the simulation's ID, its run's ID and status. + +A refused request answers 400 (a row it names, or a function it needs, is +missing; ``code`` says which) or 422 (its fields, or the stored values it +names, cannot be calculated), and writes nothing. Either route answers 503 +with ``Retry-After`` when it is not configured or a dependency cannot be +read, before writing anything, or when Modal refused the run's call: then +the record was created and its run failed, and the answer names it. Every +500 is the JSON error with code ``internal_error``, and may follow a write. + +A request gets its observability ID from the caller's header, or a new one, +and passes it to the Modal call it starts. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Protocol + +from fastapi import APIRouter, Depends, Request, status +from fastapi.responses import JSONResponse + +from policyengine_simulation_entry.economy_api_models import ( + EconomyError, + EconomyReportCreated, + EconomyReportRequest, + EconomyRequestUnprocessable, + EconomySimulationCreated, + EconomySimulationRequest, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, +) +from policyengine_simulation_entry.route_support import ( + log_unexpected_failure, + observability_context, +) + +if TYPE_CHECKING: + from fastapi import FastAPI + + from policyengine_simulation_entry.auth import AuthorizationDependency + from policyengine_simulation_entry.economy_submission_errors import RecordName + from policyengine_simulation_entry.entry_runtime import EntryRuntime + +logger = logging.getLogger(__name__) + +REPORTS_PATH = "/reports" +SIMULATIONS_PATH = "/simulations" +RETRY_AFTER_SECONDS = 10 + +_ERROR = {"model": EconomyError} +_RETRY_AFTER = { + "Retry-After": { + "description": "Seconds to wait before asking again", + "schema": {"type": "integer"}, + } +} +_FORBIDDEN = { + "description": "The caller's token is missing, invalid or lacks the permission" +} +_REFUSED = { + "description": ( + "A row the request names, or a function it needs, is missing; nothing " + "was written. unknown_region, unknown_dataset and unknown_dynamic name " + "the caller's own region, dataset or dynamic; any other code is the " + "service's state, such as a catalog version missing while a release " + "deploys" + ), + **_ERROR, +} +_UNPROCESSABLE = { + "description": ( + "The request's fields, or the stored policies and dynamics it names, " + "cannot be calculated; nothing was written. Stored values that cannot " + "be calculated carry a code: missing_parameters or " + "overlapping_parameters" + ), + "model": EconomyRequestUnprocessable, +} +_INTERNAL_ERROR = { + "description": ( + "The service failed unexpectedly (internal_error), possibly after " + "creating the record" + ), + **_ERROR, +} + + +def _unavailable_description(record: RecordName) -> dict[str, Any]: + return { + "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 run's call, or didn't accept it in time " + 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, + } + + +class EconomySubmissions(Protocol): + """Create economy reports and simulations (``EconomySubmissionService``).""" + + async def submit_report( + self, + request: EconomyReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + """Create the report and its runs and start its coordinator.""" + ... + + async def submit_simulation( + self, + request: EconomySimulationRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + """Create the simulation and its run and start the run's function.""" + ... + + +def _error_response(status_code: int, error: EconomyError) -> JSONResponse: + headers = ( + {"Retry-After": str(RETRY_AFTER_SECONDS)} + if status_code == status.HTTP_503_SERVICE_UNAVAILABLE + else None + ) + return JSONResponse( + error.model_dump(mode="json", exclude_none=True), + status_code=status_code, + headers=headers, + ) + + +def _not_configured() -> JSONResponse: + return _error_response( + status.HTTP_503_SERVICE_UNAVAILABLE, + EconomyError(code="not_configured", detail="Economy runs are not configured."), + ) + + +_SUBMISSION_FAILURES = ( + EconomySubmissionRefusedError, + EconomySubmissionUnprocessableError, + EconomySubmissionUnavailableError, +) + + +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( + status.HTTP_400_BAD_REQUEST, + EconomyError(code=error.code, detail=error.detail), + ) + if isinstance(error, EconomySubmissionUnprocessableError): + return JSONResponse( + { + "detail": [ + {"loc": list(reason.loc), "msg": reason.msg, "type": reason.type} + for reason in error.reasons + ], + "code": error.code, + }, + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ) + 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( + status.HTTP_503_SERVICE_UNAVAILABLE, + EconomyError(code=error.code, detail=error.detail, **record_ids), + ) + + +def _unexpected_failure_response( + route: str, error: Exception, *, max_characters: int +) -> JSONResponse: + """Answer an unexpected failure with the JSON 500 the routes declare. + + The log is the fixed event ``economy_route_failed``; see + ``log_unexpected_failure``. + """ + log_unexpected_failure( + logger, + "economy_route_failed", + route=route, + error=error, + max_characters=max_characters, + ) + return _error_response( + status.HTTP_500_INTERNAL_SERVER_ERROR, + EconomyError( + code="internal_error", detail="The service could not answer the request." + ), + ) + + +def add_economy_routes( + app: FastAPI, + entry: EntryRuntime, + submissions: EconomySubmissions | None, + *, + submitter: AuthorizationDependency, +) -> None: + """Add the economy routes; without a service they answer 503.""" + router = APIRouter() + log_string_limit = entry.runtime.config.limits.max_string_length + + @router.post( + REPORTS_PATH, + summary="Create an economy report", + description=( + "Create an economy report comparing a baseline and a reform " + "simulation of one region's data for one year, on the v2 version " + "manifest's default worker, and start its first run. A side whose " + "simulation already succeeded on the same versions reuses that run." + ), + operation_id="create_report", + status_code=status.HTTP_201_CREATED, + response_model=EconomyReportCreated, + responses={ + 201: {"description": "The report and its first run were created"}, + 400: _REFUSED, + 403: _FORBIDDEN, + 422: _UNPROCESSABLE, + 500: _INTERNAL_ERROR, + 503: _unavailable_description("report"), + }, + dependencies=[Depends(submitter)], + ) + async def create_report( + body: EconomyReportRequest, request: Request + ) -> JSONResponse: + if submissions is None: + return _not_configured() + observability_id = entry.start_workflow_observability(request) + try: + created = await submissions.submit_report( + body, + request_id=observability_id, + observability_context=observability_context(entry, observability_id), + ) + 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( + REPORTS_PATH, error, max_characters=log_string_limit + ) + return JSONResponse( + created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED + ) + + @router.post( + SIMULATIONS_PATH, + summary="Create an economy simulation", + description=( + "Create a standalone economy simulation of one region's data for " + "one year under one stored policy and dynamic, on the v2 version " + "manifest's default worker, and start its first run." + ), + operation_id="create_simulation", + status_code=status.HTTP_201_CREATED, + response_model=EconomySimulationCreated, + responses={ + 201: {"description": "The simulation and its first run were created"}, + 400: _REFUSED, + 403: _FORBIDDEN, + 422: _UNPROCESSABLE, + 500: _INTERNAL_ERROR, + 503: _unavailable_description("simulation"), + }, + dependencies=[Depends(submitter)], + ) + async def create_simulation( + body: EconomySimulationRequest, request: Request + ) -> JSONResponse: + if submissions is None: + return _not_configured() + observability_id = entry.start_workflow_observability(request) + try: + created = await submissions.submit_simulation( + body, + request_id=observability_id, + observability_context=observability_context(entry, observability_id), + ) + 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( + SIMULATIONS_PATH, error, max_characters=log_string_limit + ) + return JSONResponse( + created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED + ) + + app.include_router(router) 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 new file mode 100644 index 000000000..f4382898c --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py @@ -0,0 +1,180 @@ +"""Why an economy submission is refused, cannot be calculated, or cannot be served. + +The submission service raises these, and the routes answer each with its +status: EconomySubmissionRefusedError 400, EconomySubmissionUnprocessableError +422 and EconomySubmissionUnavailableError 503. ``read_refusal`` and +``reference_refusal`` turn the store's errors into them. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +from policyengine_simulation_persistence import ( + MissingPolicyParametersError, + OverlappingReformError, + UnknownCatalogVersionError, + UnknownDatasetError, + UnknownDynamicError, + UnknownPolicyError, + UnknownRegionError, +) + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_persistence import EconomyRecordReferenceError + + from policyengine_simulation_entry.economy_api_models import ( + EconomyRefusalCode, + EconomyUnavailableCode, + EconomyUnprocessableCode, + ) + +# 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. +REFERENCE_REFUSALS: dict[str, tuple[EconomyRefusalCode, str]] = { + "region_id": ("unknown_region", "region"), + "dataset_id": ("unknown_dataset", "dataset"), + "dynamic_id": ("unknown_dynamic", "dynamic"), + "policy_id": ("unknown_policy", "policy"), + "tax_benefit_model_version_id": ("unknown_catalog_version", "catalog version"), + "tax_benefit_model_id": ("unknown_catalog_version", "tax-benefit model"), +} +# What a submission creates. +RecordName = Literal["report", "simulation"] + + +class EconomySubmissionRefusedError(Exception): + """A row the submission names, or a function it needs, is missing (400).""" + + def __init__(self, *, code: EconomyRefusalCode, detail: str) -> None: + """Keep the reason's code and its description.""" + super().__init__(detail) + self.code: EconomyRefusalCode = code + self.detail = detail + + +@dataclass(frozen=True) +class UnprocessableReason: + """One reason stored values cannot be calculated, shaped as a FastAPI field error. + + ``loc`` starts with ``body`` and the request field the reason traces to. + """ + + loc: tuple[str | int, ...] + msg: str + type: str + + +class EconomySubmissionUnprocessableError(Exception): + """The stored policies or dynamics the submission names cannot be calculated.""" + + def __init__( + self, + *, + code: EconomyUnprocessableCode, + reasons: tuple[UnprocessableReason, ...], + ) -> None: + """Keep the code and every reason, in order.""" + super().__init__("; ".join(reason.msg for reason in reasons)) + self.code: EconomyUnprocessableCode = code + self.reasons = reasons + + +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. + """ + + def __init__( + self, + *, + code: EconomyUnavailableCode, + detail: str, + record_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 + + +# Each store error naming a row the request names that doesn't exist, and +# the code it is refused with. +_MISSING_ROW_CODES: dict[type[LookupError], EconomyRefusalCode] = { + UnknownRegionError: "unknown_region", + UnknownDatasetError: "unknown_dataset", + UnknownDynamicError: "unknown_dynamic", + UnknownPolicyError: "unknown_policy", +} + + +def read_refusal( + error: Exception, *, policyengine_version: str +) -> 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 for + ``policyengine_version`` lacks, or both setting one, cannot be + calculated (422). + """ + if isinstance(error, UnknownCatalogVersionError): + return EconomySubmissionRefusedError( + code="unknown_catalog_version", + detail=f"{error}, the default v2 worker's version", + ) + for error_type, code in _MISSING_ROW_CODES.items(): + if isinstance(error, error_type): + return EconomySubmissionRefusedError(code=code, detail=str(error)) + if isinstance(error, MissingPolicyParametersError): + field_name = "policy_id" if error.source == "policy" else "dynamic_id" + return EconomySubmissionUnprocessableError( + code="missing_parameters", + reasons=tuple( + UnprocessableReason( + loc=("body", field_name, name), + msg=( + f"the catalog version for .py {policyengine_version} has " + "no such parameter" + ), + type="missing_parameter", + ) + for name in error.parameter_names + ), + ) + if isinstance(error, OverlappingReformError): + return EconomySubmissionUnprocessableError( + code="overlapping_parameters", + reasons=tuple( + UnprocessableReason( + loc=("body", "dynamic_id", name), + msg="the policy and the dynamic both set this parameter", + type="overlapping_parameter", + ) + for name in error.parameter_names + ), + ) + return None + + +def reference_refusal( + error: EconomyRecordReferenceError, +) -> 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. + """ + 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" + ) 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 new file mode 100644 index 000000000..bd48345a2 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py @@ -0,0 +1,628 @@ +"""Create economy reports and simulations, and start the calls that calculate them. + +A report submission names a region, a dataset, a year and each side's +stored policy and dynamic. The service: + +1. finds the v2 version manifest's default worker's economy functions for + the country, and its bundle; +2. reads, in one transaction, the country's catalog version for that + worker's .py version, the region and dataset, each side's policy and + dynamic values, and, when the output bucket is configured, each side's + reusable run: a run of the same simulation that succeeded on the same + versions, used by a succeeded report run with the same options, and + never refused for reuse; +3. creates, in one transaction, each side's new simulation and pending run + (none for a reused side), the report with its options, and its first + report run, pending with no Modal call ID; +4. starts the coordinator's call with ``spawn``, passing the report and its + run IDs, then records the call's ID while the report run is still + pending with none (the coordinator may have claimed it already). When + Modal refuses the call, or does not confirm it in time, it fails the + report run and its new simulation runs at once instead, under the same + condition; a run the coordinator claimed meanwhile, because Modal + accepted the call after all, is answered as it stands. + +A standalone simulation is created and started the same way, with one +simulation run calculated by the country's simulation function. + +Whatever is missing or cannot be calculated is refused in steps 1 and 2, +before anything is written. The coordinator and the simulation functions +then claim their runs and record their results or failures; the stuck-run +sweep fails a pending run no call was recorded for. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID, uuid4 + +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, + EconomyRequest, + NewEconomyReport, + NewEconomySimulation, + ReportRunStatus, + ReportSimulation, + RunReuse, + SimulationChoice, + SimulationStatus, +) +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportCreated, + EconomySimulationCreated, +) +from policyengine_simulation_entry.economy_calculation_inputs import ( + CalculationScope, + CallOrigin, + DatasetNotInBundleError, + bundle_dataset, + report_run_calculation, + simulation_options, + simulation_run_calculation, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + RecordName, + read_refusal, + reference_refusal, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, +) +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from policyengine_simulation_persistence import ( + EconomyCatalogRows, + EconomyReportInputs, + EconomySimulationInputs, + SimulationInputs, + ) + + from policyengine_simulation_entry.economy_api_models import ( + EconomyRefusalCode, + EconomyReportRequest, + EconomySimulationRequest, + PolicyChoice, + ) + from policyengine_simulation_entry.economy_calculation_inputs import ( + RunCalculation, + ) + from policyengine_simulation_entry.economy_workers import ( + EconomyRunCaller, + EconomyWorker, + EconomyWorkerSource, + ) + +logger = logging.getLogger(__name__) + +# How long a submission waits for Modal to accept the run's call before +# failing the run. +CALL_START_TIMEOUT_SECONDS = 10.0 + + +class EconomyRuns(Protocol): + """The economy reads and writes (``EconomyRunStore``).""" + + def read_report_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + """Read a report's catalog rows, both sides' values and reusable runs.""" + ... + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + """Read a standalone simulation's catalog rows and values.""" + ... + + def create_report(self, report: NewEconomyReport) -> None: + """Create the report, its new simulations and its runs.""" + ... + + def create_simulation(self, simulation: NewEconomySimulation) -> None: + """Create the simulation and its first run.""" + ... + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + """Record the report run's call ID while it is pending with none.""" + ... + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + """Fail the report run, and its new simulation runs, while unstarted.""" + ... + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + """Record the simulation run's call ID while it is pending with none.""" + ... + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + """Fail the simulation run while it is pending with no call ID.""" + ... + + +@dataclass(frozen=True) +class ServiceSettings: + """The deployment the service runs in. + + ``environment`` labels the calls' artifacts, ``modal_environment`` holds + the v2 workers, and ``output_bucket`` is the bucket that keeps + simulation output files for good. Without it, no report reuses a run. + """ + + environment: str + modal_environment: str + output_bucket: str | None = None + + @property + def reuse_prefix(self) -> str | None: + """Return ``gs:///``, under which reusable output files lie.""" + return None if not self.output_bucket else f"gs://{self.output_bucket}/" + + +def _database_unavailable() -> EconomySubmissionUnavailableError: + return EconomySubmissionUnavailableError( + code="database_unavailable", detail="The simulation database could not be read." + ) + + +def _choice(policy: PolicyChoice) -> SimulationChoice: + return SimulationChoice(policy_id=policy.policy_id, dynamic_id=policy.dynamic_id) + + +class EconomySubmissionService: + """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. + """ + + def __init__( # noqa: PLR0913 # each seam tests replace + self, + *, + store: EconomyRuns, + workers: EconomyWorkerSource, + caller: EconomyRunCaller, + settings: ServiceSettings, + call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, + new_id: Callable[[], UUID] = uuid4, + clock: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + """Use these seams; tests replace each of them.""" + self._store = store + self._workers = workers + self._caller = caller + self._settings = settings + self._call_start_timeout_seconds = call_start_timeout_seconds + self._new_id = new_id + self._clock = clock + + async def submit_report( + self, + request: EconomyReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + """Create the report and its runs and start its coordinator. + + Raise EconomySubmissionRefusedError (400) or + EconomySubmissionUnprocessableError (422) before writing anything, + and EconomySubmissionUnavailableError (503) when the manifest or + database cannot be read, or, naming the report, when Modal doesn't + confirm the call and the report run was failed. + """ + worker = await self._default_worker(request.country) + options = request.options() + reuse_prefix = self._settings.reuse_prefix + reuse = ( + None + if reuse_prefix is None + else RunReuse( + versions=worker.versions, + inputs=options, + output_uri_prefix=reuse_prefix, + ) + ) + economy_request = _economy_request(request, worker) + inputs = await self._read_inputs( + lambda: self._store.read_report_inputs( + economy_request, + baseline=_choice(request.baseline), + reform=_choice(request.reform), + reuse=reuse, + ), + policyengine_version=worker.policyengine_version, + ) + scope = _scope(request, worker, inputs.catalog, spm=options.get("spm")) + report = NewEconomyReport( + id=self._new_id(), + run_id=self._new_id(), + country=request.country, + label=request.label, + description=request.description, + tax_benefit_model_id=inputs.catalog.model_id, + population=inputs.catalog.population(request.year), + versions=worker.versions, + inputs=options, + baseline=self._side(inputs.baseline), + reform=self._side(inputs.reform), + ) + calculation = report_run_calculation( + worker, + scope, + self._origin(request_id), + run_ids=ReportRunIds( + report_run_id=report.run_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_run_id=report.reform.simulation_run_id, + ), + baseline_policy=inputs.baseline.reform, + reform_policy=inputs.reform.reform, + ) + await self._create(lambda: self._store.create_report(report)) + status = await self._start_report_run( + worker, report, calculation, observability_context=observability_context + ) + return EconomyReportCreated( + report_id=report.id, + report_run_id=report.run_id, + status=status, + baseline_simulation_id=report.baseline.simulation_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + baseline_reused=report.baseline.reused, + reform_simulation_id=report.reform.simulation_id, + reform_simulation_run_id=report.reform.simulation_run_id, + reform_reused=report.reform.reused, + ) + + async def submit_simulation( + self, + request: EconomySimulationRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + """Create the simulation and its first run, and start the run's function. + + Raise the errors ``submit_report`` names, naming the simulation when + its run was failed. + """ + worker = await self._default_worker(request.country) + economy_request = _economy_request(request, worker) + choice = SimulationChoice( + policy_id=request.policy_id, dynamic_id=request.dynamic_id + ) + inputs = await self._read_inputs( + lambda: self._store.read_simulation_inputs(economy_request, choice), + policyengine_version=worker.policyengine_version, + ) + scope = _scope(request, worker, inputs.catalog, spm=None) + simulation = NewEconomySimulation( + id=self._new_id(), + run_id=self._new_id(), + choice=inputs.simulation.choice, + population=inputs.catalog.population(request.year), + versions=worker.versions, + ) + calculation = simulation_run_calculation( + worker, + scope, + self._origin(request_id), + simulation_run_id=simulation.run_id, + policy=inputs.simulation.reform, + ) + await self._create(lambda: self._store.create_simulation(simulation)) + status = await self._start_simulation_run( + worker, + simulation, + calculation, + observability_context=observability_context, + ) + return EconomySimulationCreated( + simulation_id=simulation.id, + simulation_run_id=simulation.run_id, + status=status, + ) + + def _side(self, inputs: SimulationInputs) -> ReportSimulation: + """Return a side naming its reusable run, or a new simulation and run.""" + if inputs.reusable_run is not None: + return ReportSimulation( + simulation_id=inputs.reusable_run.simulation_id, + simulation_run_id=inputs.reusable_run.simulation_run_id, + choice=inputs.choice, + reused=True, + ) + return ReportSimulation( + simulation_id=self._new_id(), + simulation_run_id=self._new_id(), + choice=inputs.choice, + reused=False, + ) + + def _origin(self, request_id: str) -> CallOrigin: + return CallOrigin( + environment=self._settings.environment, + modal_environment=self._settings.modal_environment, + request_id=request_id, + created_at=self._clock(), + ) + + async def _default_worker(self, country: str) -> EconomyWorker: + try: + return await asyncio.to_thread(self._workers.default_worker, country) + except EconomyWorkerUnavailableError as error: + raise EconomySubmissionRefusedError( + code="no_economy_worker", detail=str(error) + ) from None + except VersionManifestUnavailableError as error: + logger.warning( + "economy_manifest_unavailable", + extra={"error_type": type(error.__cause__).__name__}, + ) + raise EconomySubmissionUnavailableError( + code="manifest_unavailable", + detail="The v2 version manifest could not be read.", + ) from None + + async def _read_inputs[Inputs]( + self, read: Callable[[], Inputs], *, policyengine_version: str + ) -> Inputs: + """Run ``read`` on a thread; answer each refusal as the request's error.""" + try: + return await asyncio.to_thread(read) + except SQLAlchemyError as error: + logger.warning( + "economy_inputs_unreadable", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + except Exception as error: + refusal = read_refusal(error, policyengine_version=policyengine_version) + if refusal is None: + raise + raise refusal from None + + async def _create(self, create: Callable[[], None]) -> None: + try: + await asyncio.to_thread(create) + except EconomyRecordReferenceError as error: + raise reference_refusal(error) from None + except SQLAlchemyError as error: + logger.warning( + "economy_records_not_created", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + + async def _start_report_run( + self, + worker: EconomyWorker, + report: NewEconomyReport, + calculation: RunCalculation, + *, + observability_context: dict[str, Any] | None, + ) -> ReportRunStatus: + try: + call_id = await asyncio.wait_for( + self._caller.start_report_run( + worker, + calculation.payload, + calculation.context, + observability_context=observability_context, + ), + timeout=self._call_start_timeout_seconds, + ) + # Modal raises its own errors; any of them, or the timeout, means the + # call was not confirmed, and the run is failed unless it has started. + except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike + failed = await self._fail_unconfirmed( + "report", + report.id, + error, + lambda message: self._store.fail_unstarted_report( + report.run_id, report.new_simulation_run_ids, message + ), + ) + if failed: + refused = _call_refused("report", report.id) + 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), + ) + return ReportRunStatus.PENDING if status is None else ReportRunStatus(status) + + async def _start_simulation_run( + self, + worker: EconomyWorker, + simulation: NewEconomySimulation, + calculation: RunCalculation, + *, + observability_context: dict[str, Any] | None, + ) -> SimulationStatus: + try: + call_id = await asyncio.wait_for( + self._caller.start_simulation_run( + worker, + calculation.payload, + calculation.context, + observability_context=observability_context, + ), + timeout=self._call_start_timeout_seconds, + ) + # As for a report run: any error or the timeout leaves it unconfirmed. + except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike + failed = await self._fail_unconfirmed( + "simulation", + simulation.id, + error, + lambda message: self._store.fail_unstarted_simulation( + simulation.run_id, message + ), + ) + if failed: + refused = _call_refused("simulation", simulation.id) + 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), + ) + return SimulationStatus.PENDING if status is None else SimulationStatus(status) + + async def _fail_unconfirmed( + self, + record: RecordName, + record_id: UUID, + error: BaseException, + fail: Callable[[str], bool], + ) -> bool: + """Fail the run of a call Modal didn't confirm, unless it has started. + + Return whether it was failed. Modal may have accepted the call after + the wait ended, and its function claimed the run first: the failure + write then changes nothing. When the write itself fails, the run is + left pending: the record exists, so it is answered as created, and + the stuck-run sweep fails the run if no call ever claims it. + """ + # The log names the error type only; its text could repeat request data. + logger.error( + "economy_call_unconfirmed", + extra={ + "record": record, + "record_id": str(record_id), + "error_type": type(error).__name__, + }, + ) + message = ( + f"modal_call_refused: Modal did not accept the {record}'s call " + f"({type(error).__name__})." + ) + try: + return await asyncio.to_thread(fail, message) + except SQLAlchemyError as write_error: + logger.warning( + "economy_run_not_failed", + extra={ + "record": record, + "record_id": str(record_id), + "error_type": type(write_error).__name__, + }, + ) + return False + + async def _record_call( + self, + record: RecordName, + record_id: UUID, + recording: Callable[[], ReportRunStatus | SimulationStatus | None], + ) -> str | None: + """Record the call's ID; when that fails, answer the run as created. + + The coordinator or function claims a pending run with no call ID + itself, so the calculation goes ahead either way. + """ + try: + status = await asyncio.to_thread(recording) + except SQLAlchemyError as error: + logger.warning( + "economy_call_not_recorded", + extra={ + "record": record, + "record_id": str(record_id), + "error_type": type(error).__name__, + }, + ) + return None + return None if status is None else str(status) + + +def _call_refused( + record: RecordName, record_id: UUID +) -> EconomySubmissionUnavailableError: + return EconomySubmissionUnavailableError( + code="modal_call_refused", + detail=( + f"Modal did not accept the {record}'s call; the {record} was created " + "and its run failed." + ), + record_id=record_id, + ) + + +def _economy_request( + request: EconomyReportRequest | EconomySimulationRequest, worker: EconomyWorker +) -> EconomyRequest: + return EconomyRequest( + country=request.country, + policyengine_version=worker.policyengine_version, + region_code=request.region_code, + dataset_name=request.dataset_name, + year=request.year, + ) + + +def _scope( + request: EconomyReportRequest | EconomySimulationRequest, + worker: EconomyWorker, + catalog: EconomyCatalogRows, + *, + spm: dict[str, Any] | None, +) -> CalculationScope: + """Return what the calculation runs on, with the bundle's dataset. + + Raise EconomySubmissionRefusedError when the bundle lacks the catalog + dataset: ``unknown_dataset`` when the request named it, and + ``dataset_not_in_bundle`` for the region's default. + """ + try: + dataset = bundle_dataset(worker, catalog.dataset.name) + except DatasetNotInBundleError as error: + code: EconomyRefusalCode = ( + "dataset_not_in_bundle" + if request.dataset_name is None + else "unknown_dataset" + ) + raise EconomySubmissionRefusedError(code=code, detail=str(error)) from None + return CalculationScope( + country=request.country, + region_code=catalog.region.code, + dataset=dataset, + year=request.year, + options=simulation_options( + include_cliffs=getattr(request, "include_cliffs", False), spm=spm + ), + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py new file mode 100644 index 000000000..b07a6e693 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py @@ -0,0 +1,245 @@ +"""The default v2 worker's economy functions, and starting their calls. + +The economy routes use only the v2 version manifest's default worker: the +Modal application deployed for one .py version, which names its report +coordinator, each country's simulation function, and the bundle they run. +A report run is calculated by the coordinator, which starts a simulation +function for each simulation run it doesn't reuse; a standalone simulation +run is calculated by the simulation function alone. Each is called by name +with ``spawn``, so a route waits only for Modal to accept the call, and is +given the run IDs it claims and records, never a Stage 12 evaluation +record. + +This is the economy path's own Modal access, apart from the Stage 12 +comparison path's. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +import modal +from policyengine_simulation_contract.stage12_manifest import ( + V2ManifestLoader, + manifest_digest, +) +from policyengine_simulation_persistence import RunVersions + +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import Stage12CountryBundle + from policyengine_simulation_contract.stage12_manifest import ( + V2ManifestStore, + V2VersionManifest, + ) + + +class EconomyWorkerUnavailableError(LookupError): + """The default v2 worker cannot calculate the country's economy runs.""" + + +@dataclass(frozen=True) +class EconomyWorker: + """A country's economy functions in the default worker, and its bundle. + + ``version`` is the worker's .py version, the manifest's default, and + ``version_manifest_sha256`` the digest of the manifest it was read from. + """ + + version: str + version_manifest_sha256: str + application_name: str + coordinator_callable: str + simulation_callable: str + bundle_manifest_sha256: str + policyengine_version: str + country: Stage12CountryBundle + + @property + def versions(self) -> RunVersions: + """Return the versions every run this worker calculates records.""" + return RunVersions( + country_package_version=self.country.country_package_version, + policyengine_version=self.policyengine_version, + ) + + +def default_economy_worker(manifest: V2VersionManifest, country: str) -> EconomyWorker: + """Return the manifest's default worker's economy functions for ``country``. + + Raise EconomyWorkerUnavailableError when it cannot calculate the + country's economy runs. + """ + version = manifest.default_version + worker = manifest.versions[version] + country_worker = worker.economy_worker(country) + country_bundle = next( + (item for item in worker.bundle.countries if item.country == country), None + ) + if ( + country_worker is None + or country_worker.single_simulation_callable is None + or country_bundle is None + ): + msg = ( + f"the default v2 worker, .py {version}, cannot calculate economy " + f"runs for {country}" + ) + raise EconomyWorkerUnavailableError(msg) + return EconomyWorker( + version=version, + version_manifest_sha256=manifest_digest(manifest), + application_name=worker.application_name, + coordinator_callable=worker.report_coordinator_callable, + simulation_callable=country_worker.single_simulation_callable, + bundle_manifest_sha256=worker.bundle_manifest_sha256, + policyengine_version=worker.bundle.policyengine_version, + country=country_bundle, + ) + + +class EconomyWorkerSource(Protocol): + """Find the default worker's economy functions (``ManifestEconomyWorkers``).""" + + def default_worker(self, country: str) -> EconomyWorker: + """Return the country's economy functions, or raise why there are none.""" + ... + + +class ManifestEconomyWorkers: + """Find economy functions in the v2 version manifest's default worker. + + Each lookup reads the manifest, which is small, so a moved default is + used at once. The loader keeps its last valid copy for a stored value + that does not validate; a store that cannot be read is unavailable. + """ + + def __init__(self, store: V2ManifestStore) -> None: + """Read the manifest from ``store``; nothing is read until a lookup.""" + self._loader = V2ManifestLoader(store) + + @classmethod + def from_modal( + cls, *, manifest_name: str, environment: str + ) -> ManifestEconomyWorkers: + """Read the manifest from its Modal Dict, which must already exist.""" + return cls( + modal.Dict.from_name( + manifest_name, environment_name=environment, create_if_missing=False + ) + ) + + def default_worker(self, country: str) -> EconomyWorker: + """Return the default worker's economy functions for ``country``. + + Raise EconomyWorkerUnavailableError when it has none, and + VersionManifestUnavailableError when the manifest cannot be read. + """ + try: + manifest = self._loader.load() + # Modal raises its own errors, and any of them means the same here. + except Exception as error: + msg = "the v2 version manifest could not be read" + raise VersionManifestUnavailableError(msg) from error + return default_economy_worker(manifest, country) + + +class EconomyRunCaller(Protocol): + """Start the Modal call that calculates a run (``ModalEconomyRunCaller``).""" + + async def start_report_run( + self, + worker: EconomyWorker, + report_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the report run's coordinator and return its Modal call ID.""" + ... + + async def start_simulation_run( + self, + worker: EconomyWorker, + simulation_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the standalone simulation run's function; return its call ID.""" + ... + + +class ModalEconomyRunCaller: + """Start economy calls on Modal, by application and function name.""" + + def __init__(self, environment: str) -> None: + """Look up the functions in this Modal environment.""" + self._environment = environment + + async def start_report_run( + self, + worker: EconomyWorker, + report_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the coordinator with the report and its run IDs; return its call ID. + + Raise RuntimeError when Modal returns no call ID. + """ + return await self._spawn( + worker, + worker.coordinator_callable, + report_payload, + context_payload, + observability_context=observability_context, + ) + + async def start_simulation_run( + self, + worker: EconomyWorker, + simulation_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the simulation function with its run's ID; return its call ID. + + Raise RuntimeError when Modal returns no call ID. + """ + return await self._spawn( + worker, + worker.simulation_callable, + simulation_payload, + context_payload, + observability_context=observability_context, + ) + + async def _spawn( + self, + worker: EconomyWorker, + function_name: str, + payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + function = modal.Function.from_name( + worker.application_name, + function_name, + environment_name=self._environment, + ) + call = await function.spawn.aio( + payload, context_payload, observability_context=observability_context + ) + call_id = getattr(call, "object_id", None) + if not isinstance(call_id, str) or not call_id: + msg = f"Modal returned no call ID for {function_name}'s call" + raise RuntimeError(msg) + return call_id diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py index 7709bded3..660d4d13b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py @@ -25,9 +25,7 @@ from __future__ import annotations -import contextlib import logging -import traceback from typing import TYPE_CHECKING, Any, Protocol from uuid import UUID @@ -47,6 +45,10 @@ HouseholdCalculationUnavailableError, HouseholdCalculationUnprocessableError, ) +from policyengine_simulation_entry.route_support import ( + log_unexpected_failure, + observability_context, +) if TYPE_CHECKING: from fastapi import FastAPI @@ -207,96 +209,20 @@ def _submission_failure_response(error: Exception) -> JSONResponse: ) -# Separates the frames in the 500 log's ``stack_frames`` attribute. -STACK_FRAME_SEPARATOR = "; " -# The service's own packages, whose frames the 500 log keeps before others. -SERVICE_PACKAGES = frozenset( - {"policyengine_simulation_entry", "policyengine_simulation_persistence"} -) -# The directories a package's files sit under: an installed package's, then -# a checkout's. A frame's path is shortened to its part after one of them. -_PACKAGE_PARENT_DIRECTORIES = ("/site-packages/", "/src/") - - -def _package_path(file_path: str) -> str: - """Return ``file_path`` from its package's directory on. - - That is its part after the last ``site-packages/``, where the image - installs the service and its libraries, or else after the last ``src/``, - where a checkout keeps the service's packages. Any other path is - returned whole. - """ - for parent_directory in _PACKAGE_PARENT_DIRECTORIES: - _, found, package_path = file_path.rpartition(parent_directory) - if found: - return package_path - return file_path - - -def _stack_frames(error: Exception, max_characters: int) -> tuple[str, int]: - """Return the error's frames that fit ``max_characters``, and how many don't. - - Each frame is ``path:line in function``, with its path from its - package's directory on (``_package_path``). The kept frames are written - innermost first, from where the error was raised out to the route's - handler. When they don't all fit, frames are taken in this order until - the next one would pass ``max_characters``: the frame that raised the - error, then the service's own frames (``SERVICE_PACKAGES``) from the - innermost out, then the others from the innermost out. So other - libraries' outer frames are left out first, and left-out frames may sit - between two kept ones. Staying within ``max_characters`` keeps the - string whole: the observability runtime keeps only the start of a longer - string attribute. - """ - frames: list[str] = [] - from_the_service: list[bool] = [] - for frame in reversed(traceback.extract_tb(error.__traceback__)): - package_path = _package_path(frame.filename) - frames.append(f"{package_path}:{frame.lineno} in {frame.name}") - from_the_service.append(package_path.split("/", 1)[0] in SERVICE_PACKAGES) - positions_by_priority = sorted( - range(len(frames)), - key=lambda position: (position > 0, not from_the_service[position], position), - ) - kept_positions: list[int] = [] - length = 0 - for position in positions_by_priority: - added = len(frames[position]) + ( - len(STACK_FRAME_SEPARATOR) if kept_positions else 0 - ) - if length + added > max_characters: - break - kept_positions.append(position) - length += added - kept = [frames[position] for position in sorted(kept_positions)] - return STACK_FRAME_SEPARATOR.join(kept), len(frames) - len(kept) - - def _unexpected_failure_response( route: str, error: Exception, *, max_characters: int ) -> JSONResponse: """Answer an unexpected failure with the JSON 500 the routes declare. - The log is the fixed event ``household_route_failed`` with the route, the - error's type, its stack frames (path, line and function; see - ``_stack_frames``) as one string of at most ``max_characters``, the - observability runtime's string limit, and the number of frames left out. - It never carries the error's message or any value, which could repeat - request or stored data. The fields go in ``policyengine_attributes`` too, - the only extra fields the observability runtime's log handler keeps, and - each is a string or an integer: the runtime keeps only strings, numbers - and booleans. + The log is the fixed event ``household_route_failed``; see + ``log_unexpected_failure``. """ - stack_frames, omitted_frame_count = _stack_frames(error, max_characters) - fields = { - "route": route, - "error_type": type(error).__name__, - "stack_frames": stack_frames, - "stack_frames_omitted": omitted_frame_count, - } - logger.error( + log_unexpected_failure( + logger, "household_route_failed", - extra={**fields, "policyengine_attributes": fields}, + route=route, + error=error, + max_characters=max_characters, ) return _error_response( status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -306,21 +232,6 @@ def _unexpected_failure_response( ) -def _observability_context( - entry: EntryRuntime, observability_id: str -) -> dict[str, Any]: - """Capture the trace context for the household function, with the ID. - - Telemetry is non-fatal: without a captured context, only the ID is sent. - """ - captured: object = None - with contextlib.suppress(Exception): - captured = entry.runtime.capture_context() - context = dict(captured) if isinstance(captured, dict) else {} - context["observability_id"] = observability_id - return context - - def _state(record: HouseholdJobRecord) -> HouseholdCalculationState: return HouseholdCalculationState( job_id=record.id, @@ -395,7 +306,7 @@ async def submit_household_calculation( try: accepted = await calculations.submit( body, - observability_context=_observability_context(entry, observability_id), + observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: return _submission_failure_response(error) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.py new file mode 100644 index 000000000..457e467cd --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.py @@ -0,0 +1,124 @@ +"""What the v2 calculation routes share: their failure logs and trace context. + +The household and economy routes answer an unexpected failure with a JSON +500 and log it as a fixed event with the error's stack frames +(``log_unexpected_failure``), and pass each submission's trace context and +observability ID to the Modal call they start (``observability_context``). +""" + +from __future__ import annotations + +import contextlib +import traceback +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import logging + + from policyengine_simulation_entry.entry_runtime import EntryRuntime + +# Separates the frames in the 500 log's ``stack_frames`` attribute. +STACK_FRAME_SEPARATOR = "; " +# The service's own packages, whose frames the 500 log keeps before others. +SERVICE_PACKAGES = frozenset( + {"policyengine_simulation_entry", "policyengine_simulation_persistence"} +) +# The directories a package's files sit under: an installed package's, then +# a checkout's. A frame's path is shortened to its part after one of them. +_PACKAGE_PARENT_DIRECTORIES = ("/site-packages/", "/src/") + + +def package_path(file_path: str) -> str: + """Return ``file_path`` from its package's directory on. + + That is its part after the last ``site-packages/``, where the image + installs the service and its libraries, or else after the last ``src/``, + where a checkout keeps the service's packages. Any other path is + returned whole. + """ + for parent_directory in _PACKAGE_PARENT_DIRECTORIES: + _, found, path_after_parent = file_path.rpartition(parent_directory) + if found: + return path_after_parent + return file_path + + +def stack_frames(error: Exception, max_characters: int) -> tuple[str, int]: + """Return the error's frames that fit ``max_characters``, and how many don't. + + Each frame is ``path:line in function``, with its path from its + package's directory on (``package_path``). The kept frames are written + innermost first, from where the error was raised out to the route's + handler. When they don't all fit, frames are taken in this order until + the next one would pass ``max_characters``: the frame that raised the + error, then the service's own frames (``SERVICE_PACKAGES``) from the + innermost out, then the others from the innermost out. So other + libraries' outer frames are left out first, and left-out frames may sit + between two kept ones. Staying within ``max_characters`` keeps the + string whole: the observability runtime keeps only the start of a longer + string attribute. + """ + frames: list[str] = [] + from_the_service: list[bool] = [] + for frame in reversed(traceback.extract_tb(error.__traceback__)): + path_in_package = package_path(frame.filename) + frames.append(f"{path_in_package}:{frame.lineno} in {frame.name}") + from_the_service.append(path_in_package.split("/", 1)[0] in SERVICE_PACKAGES) + positions_by_priority = sorted( + range(len(frames)), + key=lambda position: (position > 0, not from_the_service[position], position), + ) + kept_positions: list[int] = [] + length = 0 + for position in positions_by_priority: + added = len(frames[position]) + ( + len(STACK_FRAME_SEPARATOR) if kept_positions else 0 + ) + if length + added > max_characters: + break + kept_positions.append(position) + length += added + kept = [frames[position] for position in sorted(kept_positions)] + return STACK_FRAME_SEPARATOR.join(kept), len(frames) - len(kept) + + +def log_unexpected_failure( + logger: logging.Logger, + event: str, + *, + route: str, + error: Exception, + max_characters: int, +) -> None: + """Log an unexpected failure of ``route`` as the fixed ``event``. + + The record carries the route, the error's type, its stack frames (path, + line and function; see ``stack_frames``) as one string of at most + ``max_characters``, the observability runtime's string limit, and the + number of frames left out. It never carries the error's message or any + value, which could repeat request or stored data. The fields go in + ``policyengine_attributes`` too, the only extra fields the observability + runtime's log handler keeps, and each is a string or an integer: the + runtime keeps only strings, numbers and booleans. + """ + frames, omitted_frame_count = stack_frames(error, max_characters) + fields = { + "route": route, + "error_type": type(error).__name__, + "stack_frames": frames, + "stack_frames_omitted": omitted_frame_count, + } + logger.error(event, extra={**fields, "policyengine_attributes": fields}) + + +def observability_context(entry: EntryRuntime, observability_id: str) -> dict[str, Any]: + """Capture the trace context for the Modal call a submission starts, with the ID. + + Telemetry is non-fatal: without a captured context, only the ID is sent. + """ + captured: object = None + with contextlib.suppress(Exception): + captured = entry.runtime.capture_context() + context = dict(captured) if isinstance(captured, dict) else {} + context["observability_id"] = observability_id + return context From eddc263f0fb9d7b849abdc99ef48efbafee8e082 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:36:13 +0400 Subject: [PATCH 09/26] Test reading an economy report or simulation with its newest run's status (failing) GET /reports/{report_id} and GET /simulations/{simulation_id} answer the record with its newest run, by creation and then the larger ID: its status, versions, error message and times, and a report run's simulation runs. An unknown ID answers 404 (unknown_report, unknown_simulation), a database that cannot be read 503, and the OpenAPI document declares each answer. The staging economy check polls the report this way. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_run_reads.py | 198 ++++++++++++++++++ .../tests/test_economy_run_store.py | 39 ++++ .../tests/economy_service_fakes.py | 16 ++ .../tests/test_economy_routes.py | 140 ++++++++++++- .../tests/test_economy_submissions.py | 61 ++++++ .../tests/test_openapi.py | 37 +++- 6 files changed, 488 insertions(+), 3 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py new file mode 100644 index 000000000..836b63c23 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py @@ -0,0 +1,198 @@ +"""Reading an economy report or simulation as it stands, with its newest run. + +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. The rows are written as the +superuser, to give runs their own creation times. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + COUNTRY_PACKAGE_VERSION, + POLICYENGINE_VERSION, + EconomyCatalog, + ReportRunOutcome, + RunShape, + insert_economy_catalog, + insert_report_run_naming, + insert_simulation_run, +) +from policyengine_simulation_persistence import ( + ReportRunStatus, + SimulationStatus, + read_economy_report, + read_economy_simulation, + report_runs, + simulation_runs, +) +from sqlalchemy import update + +pytestmark = pytest.mark.public_api_schema + +EARLIER = datetime(2026, 11, 1, tzinfo=UTC) + + +def _report_with_runs(connection, economy: EconomyCatalog) -> tuple[UUID, list[UUID]]: + """Insert a report and three runs of it; return the report and its runs.""" + population = economy.population() + baseline = insert_simulation_run(connection, population) + reform = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + run_ids = [ + insert_report_run_naming( + connection, + economy, + (baseline, reform), + ReportRunOutcome(status=ReportRunStatus.FAILED, error_message="first"), + ) + ] + report_id = ( + connection.execute(report_runs.select().where(report_runs.c.id == run_ids[0])) + .one() + .report_id + ) + for _ in range(2): + run_id = uuid4() + connection.execute( + report_runs.insert().values( + id=run_id, + report_id=report_id, + baseline_simulation_run_id=baseline.simulation_run_id, + reform_simulation_run_id=reform.simulation_run_id, + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version=POLICYENGINE_VERSION, + status=ReportRunStatus.PENDING, + trigger="manual", + ) + ) + run_ids.append(run_id) + return report_id, run_ids + + +def test_a_report_is_read_with_its_newest_run(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + report_id, (first, second, third) = _report_with_runs(connection, economy) + # The first run is the oldest; the other two share a creation time, so + # the larger ID is the newer. + connection.execute( + update(report_runs) + .where(report_runs.c.id == first) + .values(created_at=EARLIER) + ) + shared = EARLIER + timedelta(days=1) + connection.execute( + update(report_runs) + .where(report_runs.c.id.in_([second, third])) + .values(created_at=shared) + ) + act_as_the_service(connection) + + report = read_economy_report(connection, report_id) + missing = read_economy_report(connection, uuid4()) + + assert missing is None + assert report is not None + assert (report.id, report.country) == (report_id, "us") + run = report.run + assert run is not None + assert run.id == max(second, third) + assert (run.status, run.error_message) == (ReportRunStatus.PENDING, None) + assert (run.country_package_version, run.policyengine_version) == ( + COUNTRY_PACKAGE_VERSION, + POLICYENGINE_VERSION, + ) + assert run.created_at == shared + assert (run.started_at, run.completed_at) == (None, None) + + +def test_a_failed_newest_run_is_read_with_its_error(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + population = economy.population() + baseline = insert_simulation_run(connection, population) + reform = insert_simulation_run(connection, population) + run_id = insert_report_run_naming( + connection, + economy, + (baseline, reform), + ReportRunOutcome( + status=ReportRunStatus.FAILED, + error_message="simulation_failed: the reform failed", + ), + ) + report_id = ( + connection.execute(report_runs.select().where(report_runs.c.id == run_id)) + .one() + .report_id + ) + act_as_the_service(connection) + + report = read_economy_report(connection, report_id) + + assert report is not None + assert (report.baseline_simulation_id, report.reform_simulation_id) == ( + baseline.simulation_id, + reform.simulation_id, + ) + assert report.run is not None + assert ( + report.run.baseline_simulation_run_id, + report.run.reform_simulation_run_id, + ) == ( + baseline.simulation_run_id, + reform.simulation_run_id, + ) + assert (report.run.status, report.run.error_message) == ( + ReportRunStatus.FAILED, + "simulation_failed: the reform failed", + ) + assert report.run.completed_at is not None + + +def test_a_simulation_is_read_with_its_newest_run(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + stored = insert_simulation_run(connection, economy.population()) + newer_run_id = uuid4() + connection.execute( + simulation_runs.insert().values( + id=newer_run_id, + simulation_id=stored.simulation_id, + year=2026, + tax_benefit_model_version_id=economy.catalog.version_id, + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version="6.18.0", + dataset_id=economy.dataset_id, + region_id=economy.region_id, + status=SimulationStatus.RUNNING, + started_at=EARLIER + timedelta(days=2), + ) + ) + connection.execute( + update(simulation_runs) + .where(simulation_runs.c.id == stored.simulation_run_id) + .values(created_at=EARLIER) + ) + act_as_the_service(connection) + + simulation = read_economy_simulation(connection, stored.simulation_id) + missing = read_economy_simulation(connection, uuid4()) + + assert missing is None + assert simulation is not None + assert simulation.id == stored.simulation_id + assert simulation.run is not None + assert (simulation.run.id, simulation.run.status) == ( + newer_run_id, + SimulationStatus.RUNNING, + ) + assert simulation.run.policyengine_version == "6.18.0" + assert simulation.run.started_at == EARLIER + timedelta(days=2) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py index b9cb84487..e09fce0c6 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py @@ -322,3 +322,42 @@ def test_a_simulation_is_created_then_its_call_recorded_or_its_run_failed( assert (status, failed) == (SimulationStatus.PENDING, True) assert recorded_state == ("pending", "fc-worker") assert refused_state == ("failed", None) + + +def test_a_report_and_a_simulation_are_read_as_they_stand( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + report = NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label="Read back", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=OPTIONS, + baseline=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + reform=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + ) + store.create_report(report) + store.record_report_call(report.run_id, "fc-coordinator") + + stored_report = store.read_report(report.id) + stored_simulation = store.read_simulation(report.baseline.simulation_id) + missing = (store.read_report(uuid4()), store.read_simulation(uuid4())) + + assert missing == (None, None) + assert stored_report is not None + assert stored_report.run is not None + assert (stored_report.run.id, stored_report.run.status) == ( + report.run_id, + ReportRunStatus.PENDING, + ) + assert stored_simulation is not None + assert stored_simulation.run is not None + assert stored_simulation.run.id == report.baseline.simulation_run_id diff --git a/projects/policyengine-simulation-entry/tests/economy_service_fakes.py b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py index 3be3fd4c1..535873f0f 100644 --- a/projects/policyengine-simulation-entry/tests/economy_service_fakes.py +++ b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py @@ -31,7 +31,9 @@ CatalogRegion, EconomyCatalogRows, EconomyReportInputs, + EconomyReportState, EconomySimulationInputs, + EconomySimulationState, ReportRunStatus, ReusableRun, SimulationInputs, @@ -128,6 +130,8 @@ class FakeStore: record_failure: Exception | None = None fail_failure: Exception | None = None claimed: bool = False + stored_report: EconomyReportState | None = None + stored_simulation: EconomySimulationState | None = None report_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) simulation_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) reports: list[NewEconomyReport] = field(default_factory=list) @@ -168,6 +172,18 @@ def read_simulation_inputs( simulation=SimulationInputs(choice=choice, reform=REFORM), ) + def read_report(self, report_id: UUID) -> EconomyReportState | None: + if self.read_failure is not None: + raise self.read_failure + stored = self.stored_report + return stored if stored is not None and stored.id == report_id else None + + def read_simulation(self, simulation_id: UUID) -> EconomySimulationState | None: + if self.read_failure is not None: + raise self.read_failure + stored = self.stored_simulation + return stored if stored is not None and stored.id == simulation_id else None + def create_report(self, report: NewEconomyReport) -> None: if self.create_failure is not None: raise self.create_failure diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py index 35985f7a5..4b67af99f 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -4,6 +4,7 @@ import json from dataclasses import dataclass, field +from datetime import datetime from typing import Any, NoReturn from uuid import UUID @@ -30,7 +31,13 @@ ) from policyengine_simulation_entry.simulation_database import SimulationDatabase from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER -from policyengine_simulation_persistence import ReportRunStatus, SimulationStatus +from policyengine_simulation_persistence import ( + EconomyReportState, + EconomySimulationState, + ReportRunState, + ReportRunStatus, + SimulationStatus, +) OBSERVABILITY_ID = "6e1f3d4c-0a5b-4c39-9d2e-7f8a9b0c1d2e" RECORD_ID = UUID("00000000-0000-4000-8000-0000000000e1") @@ -65,6 +72,8 @@ class FakeEconomySubmissions: """Records each submission, and answers as told.""" failure: Exception | None = None + stored_report: EconomyReportState | None = None + stored_simulation: EconomySimulationState | None = None reports: list[tuple[EconomyReportRequest, str, Any]] = field(default_factory=list) simulations: list[tuple[EconomySimulationRequest, str, Any]] = field( default_factory=list @@ -94,6 +103,20 @@ async def submit_simulation( _raise(self.failure) return CREATED_SIMULATION + async def read_report(self, report_id: UUID) -> EconomyReportState | None: + if self.failure is not None: + _raise(self.failure) + stored = self.stored_report + return stored if stored is not None and stored.id == report_id else None + + async def read_simulation( + self, simulation_id: UUID + ) -> EconomySimulationState | None: + if self.failure is not None: + _raise(self.failure) + stored = self.stored_simulation + return stored if stored is not None and stored.id == simulation_id else None + @pytest.fixture def client_for(monkeypatch): @@ -370,3 +393,118 @@ def test_without_its_database_and_v2_environment_none_is_built(settings_changes) database.dispose() assert submissions is None + + +READ_AT = "2026-11-12T15:00:00Z" +STORED_REPORT = EconomyReportState( + id=RECORD_ID, + country="uk", + baseline_simulation_id=SIDE_IDS[0], + reform_simulation_id=SIDE_IDS[2], + run=ReportRunState( + id=RUN_ID, + status=ReportRunStatus.FAILED, + country_package_version="2.174.1", + policyengine_version="6.17.1", + error_message="simulation_failed: the reform failed", + created_at=datetime.fromisoformat(READ_AT), + started_at=datetime.fromisoformat(READ_AT), + completed_at=datetime.fromisoformat(READ_AT), + baseline_simulation_run_id=SIDE_IDS[1], + reform_simulation_run_id=SIDE_IDS[3], + ), +) +STORED_SIMULATION = EconomySimulationState(id=RECORD_ID, run=None) + + +def test_a_report_is_read_with_its_newest_runs_status(client_for) -> None: + submissions = FakeEconomySubmissions(stored_report=STORED_REPORT) + + with client_for(submissions) as client: + result = client.get(f"/reports/{RECORD_ID}") + missing = client.get(f"/reports/{RUN_ID}") + + assert result.status_code == 200 + assert result.json() == { + "report_id": str(RECORD_ID), + "country": "uk", + "baseline_simulation_id": str(SIDE_IDS[0]), + "reform_simulation_id": str(SIDE_IDS[2]), + "run": { + "run_id": str(RUN_ID), + "status": "failed", + "country_package_version": "2.174.1", + "policyengine_version": "6.17.1", + "error_message": "simulation_failed: the reform failed", + "created_at": READ_AT, + "started_at": READ_AT, + "completed_at": READ_AT, + "baseline_simulation_run_id": str(SIDE_IDS[1]), + "reform_simulation_run_id": str(SIDE_IDS[3]), + }, + } + assert missing.status_code == 404 + assert missing.json() == { + "code": "unknown_report", + "detail": f"There is no report {RUN_ID}.", + } + + +def test_a_simulation_is_read_with_its_newest_runs_status(client_for) -> None: + submissions = FakeEconomySubmissions(stored_simulation=STORED_SIMULATION) + + with client_for(submissions) as client: + result = client.get(f"/simulations/{RECORD_ID}") + missing = client.get(f"/simulations/{RUN_ID}") + + assert result.status_code == 200 + assert result.json() == {"simulation_id": str(RECORD_ID), "run": None} + assert missing.status_code == 404 + assert missing.json()["code"] == "unknown_simulation" + + +@pytest.mark.parametrize("path", ["/reports/not-a-uuid", "/simulations/not-a-uuid"]) +def test_a_read_of_an_id_that_is_not_a_uuid_is_refused(client_for, path) -> None: + with client_for(FakeEconomySubmissions()) as client: + result = client.get(path) + + assert result.status_code == 422 + + +@pytest.mark.parametrize("record", ["reports", "simulations"]) +def test_a_read_the_database_cannot_answer_is_unavailable(client_for, record) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnavailableError( + code="database_unavailable", detail="The database could not be read." + ) + ) + + with client_for(submissions) as client: + result = client.get(f"/{record}/{RECORD_ID}") + + assert result.status_code == 503 + assert result.headers["Retry-After"] == "10" + assert result.json() == { + "code": "database_unavailable", + "detail": "The database could not be read.", + } + + +@pytest.mark.parametrize("record", ["reports", "simulations"]) +def test_an_unconfigured_service_answers_reads_with_503(client_for, record) -> None: + with client_for(None) as client: + result = client.get(f"/{record}/{RECORD_ID}") + + assert result.status_code == 503 + assert result.json()["code"] == "not_configured" + + +@pytest.mark.parametrize("record", ["reports", "simulations"]) +def test_an_unexpected_read_failure_answers_a_json_500(client_for, record) -> None: + submissions = FakeEconomySubmissions(failure=RuntimeError("must not be logged")) + + with client_for(submissions) as client: + result = client.get(f"/{record}/{RECORD_ID}") + + assert result.status_code == 500 + assert result.json() == UNEXPECTED_FAILURE_ANSWER diff --git a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py index bc05dd61e..6c69bad37 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py @@ -45,12 +45,16 @@ from policyengine_simulation_persistence import ( CatalogDataset, EconomyRecordReferenceError, + EconomyReportState, + EconomySimulationState, MissingPolicyParametersError, OverlappingReformError, + ReportRunState, ReportRunStatus, RunReuse, RunVersions, SimulationChoice, + SimulationRunState, SimulationStatus, UnknownCatalogVersionError, UnknownDatasetError, @@ -467,3 +471,60 @@ def test_the_calls_context_records_when_and_where_it_was_made() -> None: "observed", "staging", ) + + +REPORT_RUN = ReportRunState( + id=UUID(int=11), + status=ReportRunStatus.SUCCEEDED, + country_package_version="2.2.1", + policyengine_version=POLICYENGINE_VERSION, + error_message=None, + created_at=CREATED_AT, + started_at=CREATED_AT, + completed_at=CREATED_AT, + baseline_simulation_run_id=UUID(int=12), + reform_simulation_run_id=UUID(int=13), +) +STORED_REPORT = EconomyReportState( + id=UUID(int=10), + country="us", + baseline_simulation_id=UUID(int=14), + reform_simulation_id=UUID(int=15), + run=REPORT_RUN, +) +STORED_SIMULATION = EconomySimulationState( + id=UUID(int=20), + run=SimulationRunState( + id=UUID(int=21), + status=SimulationStatus.FAILED, + country_package_version="2.2.1", + policyengine_version=POLICYENGINE_VERSION, + error_message="simulation_failed: no", + created_at=CREATED_AT, + started_at=None, + completed_at=CREATED_AT, + ), +) + + +def test_a_report_and_a_simulation_are_read_as_they_stand() -> None: + submissions = service( + FakeStore(stored_report=STORED_REPORT, stored_simulation=STORED_SIMULATION) + ) + + assert asyncio.run(submissions.read_report(STORED_REPORT.id)) == STORED_REPORT + assert ( + asyncio.run(submissions.read_simulation(STORED_SIMULATION.id)) + == STORED_SIMULATION + ) + assert asyncio.run(submissions.read_report(UUID(int=99))) is None + assert asyncio.run(submissions.read_simulation(UUID(int=99))) is None + + +def test_a_read_the_database_cannot_answer_is_unavailable() -> None: + submissions = service(FakeStore(read_failure=DATABASE_ERROR)) + + for read in (submissions.read_report, submissions.read_simulation): + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + asyncio.run(read(UUID(int=1))) + assert unavailable.value.code == "database_unavailable" diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 95beaad89..30f8022af 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -33,6 +33,8 @@ class OpenAPIDocument(TypedDict): ("GET", "/health"), ("GET", "/jobs/{job_id}"), ("GET", "/ready"), + ("GET", "/reports/{report_id}"), + ("GET", "/simulations/{simulation_id}"), ("GET", "/versions"), ("GET", "/versions/{kind}"), ("POST", "/household-calculations"), @@ -81,17 +83,26 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: # The economy report and simulation routes, which only the Cloud Run # entrypoint has. -ECONOMY_PATHS = ("/reports", "/simulations") +ECONOMY_PATHS = ( + "/reports", + "/reports/{report_id}", + "/simulations", + "/simulations/{simulation_id}", +) ECONOMY_SCHEMAS = ( "EconomyError", "EconomyReportCreated", "EconomyReportRequest", + "EconomyReportState", "EconomyRequestReason", "EconomyRequestUnprocessable", "EconomySimulationCreated", "EconomySimulationRequest", + "EconomySimulationState", "PolicyChoice", + "ReportRunRead", "ReportRunStatus", + "SimulationRunRead", "SimulationStatus", ) @@ -381,7 +392,8 @@ def test_every_economy_503_declares_its_retry_after_header(): spec = create_app().openapi() for path in ECONOMY_PATHS: - headers = _responses(spec, "post", path)["503"]["headers"] + method = "get" if path.endswith("}") else "post" + headers = _responses(spec, method, path)["503"]["headers"] assert headers["Retry-After"]["schema"] == {"type": "integer"} @@ -423,3 +435,24 @@ def test_economy_requests_name_no_version_and_no_record_id(): assert spm["$ref"] == "#/components/schemas/SPMSelection" assert "default" not in spm assert "spm" not in schemas["EconomyReportRequest"].get("required", []) + + +def test_economy_reads_declare_every_answer(): + spec = create_app().openapi() + error = {"$ref": "#/components/schemas/EconomyError"} + + for path, state in ( + ("/reports/{report_id}", "EconomyReportState"), + ("/simulations/{simulation_id}", "EconomySimulationState"), + ): + responses = _responses(spec, "get", path) + assert set(responses) == {"200", "403", "404", "422", "500", "503"} + assert _response_schema(responses["200"]) == { + "$ref": f"#/components/schemas/{state}" + } + for status_code in ("404", "500", "503"): + assert _response_schema(responses[status_code]) == error + assert "Modal" not in responses["503"]["description"] + names = operation_ids(cast("OpenAPIDocument", spec)) + assert names[("GET", "/reports/{report_id}")] == "get_report" + assert names[("GET", "/simulations/{simulation_id}")] == "get_simulation" From eb7002e6755b0dfb3cf761890c07cf43cfb8a803 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:39:26 +0400 Subject: [PATCH 10/26] Read an economy report or simulation with its newest run's status economy_run_reads returns a report, its two simulations and its newest run (created last, ties to the larger ID), or a simulation and its newest run, with each run's status, versions, error message and times. The store reads each in its own transaction. GET /reports/{report_id} and GET /simulations/{simulation_id} (permission simulation:read) answer it, 404 for an unknown ID, 503 when the database cannot be read, and the JSON 500 otherwise. Results stay the Public API's to read. Co-Authored-By: Claude Opus 5.5 --- .../__init__.py | 14 ++ .../economy_run_reads.py | 181 ++++++++++++++ .../economy_run_store.py | 18 ++ .../src/policyengine_simulation_entry/app.py | 6 +- .../economy_api_models.py | 53 +++++ .../economy_routes.py | 223 +++++++++++++++++- .../economy_submissions.py | 38 +++ 7 files changed, 529 insertions(+), 4 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.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 e8ce4e7a3..85b0c7407 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,14 @@ missing_table_privileges, read_missing_table_privileges, ) +from policyengine_simulation_persistence.economy_run_reads import ( + EconomyReportState, + EconomySimulationState, + ReportRunState, + SimulationRunState, + read_economy_report, + read_economy_simulation, +) from policyengine_simulation_persistence.economy_run_records import ( CatalogDataset, CatalogRegion, @@ -99,9 +107,11 @@ "EconomyPopulation", "EconomyRecordReferenceError", "EconomyReportInputs", + "EconomyReportState", "EconomyRequest", "EconomyRunStore", "EconomySimulationInputs", + "EconomySimulationState", "HouseholdCalculationInputs", "HouseholdJobRecord", "HouseholdJobReferenceError", @@ -113,6 +123,7 @@ "NewHouseholdJob", "OverlappingReformError", "PolicyReform", + "ReportRunState", "ReportRunStatus", "ReportRunTrigger", "ReportSimulation", @@ -121,6 +132,7 @@ "RunVersions", "SimulationChoice", "SimulationInputs", + "SimulationRunState", "SimulationStatus", "SimulationType", "StoredHousehold", @@ -149,6 +161,8 @@ "read_dataset", "read_dataset_by_id", "read_dynamic_reform", + "read_economy_report", + "read_economy_simulation", "read_household_job", "read_missing_table_privileges", "read_model_id", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.py new file mode 100644 index 000000000..1cde162ad --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.py @@ -0,0 +1,181 @@ +"""Read a report or a simulation as it stands, with its newest run. + +The newest run is the one created last, with ties broken by the larger ID, +whatever its status and versions, as the Public API reads it. The run's +results are the Public API's to read: these reads give only its state. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from sqlalchemy import select + +from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + SimulationStatus, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from sqlalchemy import Connection, RowMapping + + +@dataclass(frozen=True) +class SimulationRunState: + """A simulation run's state: its status, versions, error and times.""" + + id: UUID + status: SimulationStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + + +@dataclass(frozen=True) +class ReportRunState: + """A report run's state, and the simulation runs it compares.""" + + id: UUID + status: ReportRunStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + + +@dataclass(frozen=True) +class EconomyReportState: + """A report, its two simulations, and its newest run (None with no run).""" + + id: UUID + country: str + baseline_simulation_id: UUID | None + reform_simulation_id: UUID | None + run: ReportRunState | None + + +@dataclass(frozen=True) +class EconomySimulationState: + """A simulation and its newest run (None with no run).""" + + id: UUID + run: SimulationRunState | None + + +def read_economy_report( + connection: Connection, report_id: UUID +) -> EconomyReportState | None: + """Return the report with its newest run, or None when there is no such report.""" + report = connection.execute( + select( + reports.c.id, + reports.c.country, + reports.c.baseline_simulation_id, + reports.c.reform_simulation_id, + ).where(reports.c.id == report_id) + ).one_or_none() + if report is None: + return None + run = ( + connection.execute( + select( + report_runs.c.id, + report_runs.c.status, + report_runs.c.country_package_version, + report_runs.c.policyengine_version, + report_runs.c.error_message, + report_runs.c.created_at, + report_runs.c.started_at, + report_runs.c.completed_at, + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + ) + .where(report_runs.c.report_id == report_id) + .order_by(report_runs.c.created_at.desc(), report_runs.c.id.desc()) + .limit(1) + ) + .mappings() + .one_or_none() + ) + return EconomyReportState( + id=report.id, + country=report.country, + baseline_simulation_id=report.baseline_simulation_id, + reform_simulation_id=report.reform_simulation_id, + run=None if run is None else _report_run_state(run), + ) + + +def read_economy_simulation( + connection: Connection, simulation_id: UUID +) -> EconomySimulationState | None: + """Return the simulation with its newest run, or None when there is none.""" + found = connection.execute( + select(simulations.c.id).where(simulations.c.id == simulation_id) + ).scalar_one_or_none() + if found is None: + return None + run = ( + connection.execute( + select( + simulation_runs.c.id, + simulation_runs.c.status, + simulation_runs.c.country_package_version, + simulation_runs.c.policyengine_version, + simulation_runs.c.error_message, + simulation_runs.c.created_at, + simulation_runs.c.started_at, + simulation_runs.c.completed_at, + ) + .where(simulation_runs.c.simulation_id == simulation_id) + .order_by(simulation_runs.c.created_at.desc(), simulation_runs.c.id.desc()) + .limit(1) + ) + .mappings() + .one_or_none() + ) + return EconomySimulationState( + id=simulation_id, + run=None + if run is None + else SimulationRunState( + id=run["id"], + status=SimulationStatus(run["status"]), + country_package_version=run["country_package_version"], + policyengine_version=run["policyengine_version"], + error_message=run["error_message"], + created_at=run["created_at"], + started_at=run["started_at"], + completed_at=run["completed_at"], + ), + ) + + +def _report_run_state(run: RowMapping) -> ReportRunState: + return ReportRunState( + id=run["id"], + status=ReportRunStatus(run["status"]), + country_package_version=run["country_package_version"], + policyengine_version=run["policyengine_version"], + error_message=run["error_message"], + created_at=run["created_at"], + started_at=run["started_at"], + completed_at=run["completed_at"], + baseline_simulation_run_id=run["baseline_simulation_run_id"], + reform_simulation_run_id=run["reform_simulation_run_id"], + ) 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 b90921b4b..ec41c804a 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 @@ -12,6 +12,10 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any +from policyengine_simulation_persistence.economy_run_reads import ( + read_economy_report, + read_economy_simulation, +) from policyengine_simulation_persistence.economy_run_records import ( CatalogDataset, CatalogRegion, @@ -51,6 +55,10 @@ from sqlalchemy import Connection, Engine + from policyengine_simulation_persistence.economy_run_reads import ( + EconomyReportState, + EconomySimulationState, + ) from policyengine_simulation_persistence.economy_run_records import ( NewEconomyReport, NewEconomySimulation, @@ -286,6 +294,16 @@ def fail_unstarted_simulation( error_message=error_message, ) + def read_report(self, report_id: UUID) -> EconomyReportState | None: + """Return the report with its newest run, or None when there is none.""" + with self._transaction() as connection: + return read_economy_report(connection, report_id) + + def read_simulation(self, simulation_id: UUID) -> EconomySimulationState | None: + """Return the simulation with its newest run, or None when there is none.""" + with self._transaction() as connection: + return read_economy_simulation(connection, simulation_id) + def _catalog_rows( connection: Connection, request: EconomyRequest 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 2ca53a5c4..0c99f05e4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -543,9 +543,10 @@ def _register_economy_routes( submissions: EconomySubmissions | None, *, submitter: AuthorizationDependency, + reader: AuthorizationDependency, ) -> None: - """Add the routes that create economy reports and simulations.""" - add_economy_routes(app, entry, submissions, submitter=submitter) + """Add the routes that create economy reports and simulations and read them.""" + add_economy_routes(app, entry, submissions, submitter=submitter, reader=reader) def economy_submissions_for( @@ -739,6 +740,7 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: entry, economy_submissions_for(runtime_settings, runtime_database), submitter=simulation_submitter, + reader=simulation_reader, ) _register_service_routes(app, entry) return app 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 7b2134637..b67454886 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 @@ -27,6 +27,7 @@ from __future__ import annotations +from datetime import datetime from typing import Annotated, Any, Final, Literal from uuid import UUID @@ -80,6 +81,8 @@ EconomyUnprocessableCode, EconomyUnavailableCode, "internal_error", + "unknown_report", + "unknown_simulation", ] @@ -196,6 +199,56 @@ class EconomyReportCreated(_Answer): reform_reused: bool +class SimulationRunRead(_Answer): + """A simulation run's state: its status, versions, error message and times.""" + + run_id: UUID + status: SimulationStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + + +class ReportRunRead(_Answer): + """A report run's state, and the simulation runs it compares.""" + + run_id: UUID + status: ReportRunStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + + +class EconomyReportState(_Answer): + """A report, its two simulations and its newest run, as they stand. + + The newest run is the one created last, ties going to the larger ID; + ``run`` is null for a report with none. Its results are read from the + Public API. + """ + + report_id: UUID + country: str + baseline_simulation_id: UUID | None + reform_simulation_id: UUID | None + run: ReportRunRead | None + + +class EconomySimulationState(_Answer): + """A simulation and its newest run, as they stand; ``run`` null with none.""" + + simulation_id: UUID + run: SimulationRunRead | None + + class EconomyError(_Answer): """Why an economy route refused a request or cannot answer now. 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 d0924244a..e54be88f6 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 @@ -7,6 +7,10 @@ - ``POST /simulations`` (``simulation:submit``) creates a standalone economy simulation and its first run, and starts the run's simulation function, answering 201 with the simulation's ID, its run's ID and status. +- ``GET /reports/{report_id}`` and ``GET /simulations/{simulation_id}`` + (``simulation:read``) return the record with its newest run's status, + versions, error message and times. The Public API reads results from the + tables itself; these reads let a deploy check follow a run. A refused request answers 400 (a row it names, or a function it needs, is missing; ``code`` says which) or 422 (its fields, or the stored values it @@ -24,6 +28,7 @@ import logging from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID from fastapi import APIRouter, Depends, Request, status from fastapi.responses import JSONResponse @@ -32,9 +37,13 @@ EconomyError, EconomyReportCreated, EconomyReportRequest, + EconomyReportState, EconomyRequestUnprocessable, EconomySimulationCreated, EconomySimulationRequest, + EconomySimulationState, + ReportRunRead, + SimulationRunRead, ) from policyengine_simulation_entry.economy_submission_errors import ( EconomySubmissionRefusedError, @@ -47,7 +56,16 @@ ) if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + from fastapi import FastAPI + from policyengine_simulation_persistence import ( + EconomyReportState as StoredReport, + ) + from policyengine_simulation_persistence import ( + EconomySimulationState as StoredSimulation, + ) + from pydantic import BaseModel from policyengine_simulation_entry.auth import AuthorizationDependency from policyengine_simulation_entry.economy_submission_errors import RecordName @@ -56,7 +74,9 @@ logger = logging.getLogger(__name__) REPORTS_PATH = "/reports" +REPORT_PATH = "/reports/{report_id}" SIMULATIONS_PATH = "/simulations" +SIMULATION_PATH = "/simulations/{simulation_id}" RETRY_AFTER_SECONDS = 10 _ERROR = {"model": EconomyError} @@ -97,6 +117,24 @@ } +_READ_UNPROCESSABLE = { + "description": "The ID is not a UUID", + "model": EconomyRequestUnprocessable, +} +_READ_UNAVAILABLE = { + "description": ( + "Economy runs are not configured, or the database cannot be read; " + "retry after Retry-After seconds" + ), + "headers": _RETRY_AFTER, + **_ERROR, +} +_READ_INTERNAL_ERROR = { + "description": "The service failed unexpectedly (internal_error)", + **_ERROR, +} + + def _unavailable_description(record: RecordName) -> dict[str, Any]: return { "description": ( @@ -134,6 +172,14 @@ async def submit_simulation( """Create the simulation and its run and start the run's function.""" ... + async def read_report(self, report_id: UUID) -> StoredReport | None: + """Return the report with its newest run, or None.""" + ... + + async def read_simulation(self, simulation_id: UUID) -> StoredSimulation | None: + """Return the simulation with its newest run, or None.""" + ... + def _error_response(status_code: int, error: EconomyError) -> JSONResponse: headers = ( @@ -211,16 +257,88 @@ def _unexpected_failure_response( ) +def _report_state(stored: StoredReport) -> EconomyReportState: + run = stored.run + return EconomyReportState( + report_id=stored.id, + country=stored.country, + baseline_simulation_id=stored.baseline_simulation_id, + reform_simulation_id=stored.reform_simulation_id, + run=None + if run is None + else ReportRunRead( + run_id=run.id, + status=run.status, + country_package_version=run.country_package_version, + policyengine_version=run.policyengine_version, + error_message=run.error_message, + created_at=run.created_at, + started_at=run.started_at, + completed_at=run.completed_at, + baseline_simulation_run_id=run.baseline_simulation_run_id, + reform_simulation_run_id=run.reform_simulation_run_id, + ), + ) + + +def _simulation_state(stored: StoredSimulation) -> EconomySimulationState: + run = stored.run + return EconomySimulationState( + simulation_id=stored.id, + run=None + if run is None + else SimulationRunRead( + run_id=run.id, + status=run.status, + country_package_version=run.country_package_version, + policyengine_version=run.policyengine_version, + error_message=run.error_message, + created_at=run.created_at, + started_at=run.started_at, + completed_at=run.completed_at, + ), + ) + + +def _unknown(record: RecordName, record_id: UUID) -> JSONResponse: + return _error_response( + status.HTTP_404_NOT_FOUND, + EconomyError( + code="unknown_report" if record == "report" else "unknown_simulation", + detail=f"There is no {record} {record_id}.", + ), + ) + + def add_economy_routes( app: FastAPI, entry: EntryRuntime, submissions: EconomySubmissions | None, *, submitter: AuthorizationDependency, + reader: AuthorizationDependency, ) -> None: """Add the economy routes; without a service they answer 503.""" - router = APIRouter() log_string_limit = entry.runtime.config.limits.max_string_length + app.include_router( + _submission_routes( + entry, submissions, submitter=submitter, log_string_limit=log_string_limit + ) + ) + app.include_router( + _read_routes(submissions, reader=reader, log_string_limit=log_string_limit) + ) + + +def _submission_routes( + entry: EntryRuntime, + submissions: EconomySubmissions | None, + *, + submitter: AuthorizationDependency, + log_string_limit: int, +) -> APIRouter: + """Return the router of ``POST /reports`` and ``POST /simulations``.""" + router = APIRouter() @router.post( REPORTS_PATH, @@ -311,4 +429,105 @@ async def create_simulation( created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED ) - app.include_router(router) + return router + + +async def _read_answer[Stored]( # noqa: PLR0913 # the read, its record and how to answer it + submissions: EconomySubmissions | None, + read: Callable[[EconomySubmissions], Awaitable[Stored | None]], + answer: Callable[[Stored], BaseModel], + *, + record: RecordName, + record_id: UUID, + log_string_limit: int, +) -> JSONResponse: + """Answer a read: the record as it stands, 404, 503 or the JSON 500.""" + if submissions is None: + return _not_configured() + 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) + # 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 + ) + if state is None: + return _unknown(record, record_id) + return JSONResponse(state.model_dump(mode="json")) + + +def _read_routes( + submissions: EconomySubmissions | None, + *, + reader: AuthorizationDependency, + log_string_limit: int, +) -> APIRouter: + """Return the router of the report and simulation reads.""" + router = APIRouter() + + @router.get( + REPORT_PATH, + summary="Get an economy report's newest run", + description=( + "Return the report, its two simulations and its newest run's " + "status, versions, error message and times." + ), + operation_id="get_report", + response_model=EconomyReportState, + responses={ + 200: {"description": "The report as it stands"}, + 403: _FORBIDDEN, + 404: {"description": "There is no such report (unknown_report)", **_ERROR}, + 422: _READ_UNPROCESSABLE, + 500: _READ_INTERNAL_ERROR, + 503: _READ_UNAVAILABLE, + }, + dependencies=[Depends(reader)], + ) + async def get_report(report_id: UUID) -> JSONResponse: + return await _read_answer( + submissions, + lambda service: service.read_report(report_id), + _report_state, + record="report", + record_id=report_id, + log_string_limit=log_string_limit, + ) + + @router.get( + SIMULATION_PATH, + summary="Get an economy simulation's newest run", + description=( + "Return the simulation and its newest run's status, versions, error " + "message and times." + ), + operation_id="get_simulation", + response_model=EconomySimulationState, + responses={ + 200: {"description": "The simulation as it stands"}, + 403: _FORBIDDEN, + 404: { + "description": "There is no such simulation (unknown_simulation)", + **_ERROR, + }, + 422: _READ_UNPROCESSABLE, + 500: _READ_INTERNAL_ERROR, + 503: _READ_UNAVAILABLE, + }, + dependencies=[Depends(reader)], + ) + async def get_simulation(simulation_id: UUID) -> JSONResponse: + return await _read_answer( + submissions, + lambda service: service.read_simulation(simulation_id), + _simulation_state, + record="simulation", + record_id=simulation_id, + log_string_limit=log_string_limit, + ) + + return router 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 bd48345a2..7bcd0fad6 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 @@ -87,7 +87,9 @@ from policyengine_simulation_persistence import ( EconomyCatalogRows, EconomyReportInputs, + EconomyReportState, EconomySimulationInputs, + EconomySimulationState, SimulationInputs, ) @@ -168,6 +170,14 @@ def fail_unstarted_simulation( """Fail the simulation run while it is pending with no call ID.""" ... + def read_report(self, report_id: UUID) -> EconomyReportState | None: + """Return the report with its newest run.""" + ... + + def read_simulation(self, simulation_id: UUID) -> EconomySimulationState | None: + """Return the simulation with its newest run.""" + ... + @dataclass(frozen=True) class ServiceSettings: @@ -353,6 +363,34 @@ async def submit_simulation( status=status, ) + async def read_report(self, report_id: UUID) -> EconomyReportState | None: + """Return the report with its newest run, or None when there is none. + + Raise EconomySubmissionUnavailableError when it cannot be read. + """ + return await self._read_state(lambda: self._store.read_report(report_id)) + + async def read_simulation( + self, simulation_id: UUID + ) -> EconomySimulationState | None: + """Return the simulation with its newest run, or None when there is none. + + Raise EconomySubmissionUnavailableError when it cannot be read. + """ + return await self._read_state( + lambda: self._store.read_simulation(simulation_id) + ) + + async def _read_state[State](self, read: Callable[[], State]) -> State: + try: + return await asyncio.to_thread(read) + except SQLAlchemyError as error: + logger.warning( + "economy_state_unreadable", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + def _side(self, inputs: SimulationInputs) -> ReportSimulation: """Return a side naming its reusable run, or a new simulation and run.""" if inputs.reusable_run is not None: From daa1f5c8e0673981314e1e4c19358ce48f582ccf Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:42:23 +0400 Subject: [PATCH 11/26] Test the output bucket on every deploy and the staging economy check (failing) Every deploy names V2_SIMULATION_OUTPUT_BUCKET, from the environment's GitHub variable, and stops without a valid bucket name. Every staging deploy creates a small report, current law against current law for Wales in 2026, through the candidate's POST /reports, follows its run through GET /reports/{report_id} and requires it to succeed; deployment_ready requires the check, so a failure stops the production deploy. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_deployment_assets_support.py | 1 + .../tests/test_economy_deployment.py | 84 ++++ .../tests/test_staging_economy_check.py | 399 ++++++++++++++++++ 3 files changed, 484 insertions(+) create mode 100644 projects/policyengine-simulation-entry/tests/test_economy_deployment.py create mode 100644 projects/policyengine-simulation-entry/tests/test_staging_economy_check.py diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py index a0c8ada56..5d8dc1eb0 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py @@ -644,4 +644,5 @@ def cloud_run_deploy_environment(tmp_path: Path) -> dict[str, str]: "SIMULATION_SERVICE_DATABASE_URL_SECRET_NAME": ( "simulation-staging-service-database-url" ), + "V2_SIMULATION_OUTPUT_BUCKET_VALUE": "policyengine-v2-outputs-staging", } diff --git a/projects/policyengine-simulation-entry/tests/test_economy_deployment.py b/projects/policyengine-simulation-entry/tests/test_economy_deployment.py new file mode 100644 index 000000000..b500d1be7 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_deployment.py @@ -0,0 +1,84 @@ +"""Deployment wiring for the economy report and simulation routes. + +Every deploy gives the service the bucket that keeps v2 simulation output +files for good (the environment variable ``V2_SIMULATION_OUTPUT_BUCKET``), +so a report reuses only a succeeded run whose file is in it. The service +compares output URIs with it and never reads the bucket. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets_support import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + cloud_run_deploy_environment, + workflow_jobs, +) + + +def _run_deploy(tmp_path, **overrides): + environment = cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + return result, environment + + +@pytest.mark.parametrize("deploy_stage12", ["true", "false"]) +def test_every_deploy_names_the_v2_output_bucket(tmp_path, deploy_stage12): + result, environment = _run_deploy( + tmp_path, + DEPLOY_STAGE12_V2=deploy_stage12, + V2_SIMULATION_OUTPUT_BUCKET_VALUE="policyengine-v2-outputs-staging", + ) + + assert result.returncode == 0, result.stderr + runtime_environment = json.loads( + Path(environment["GCLOUD_ENV_FILE_CAPTURE"]).read_text(encoding="utf-8") + ) + assert runtime_environment["V2_SIMULATION_OUTPUT_BUCKET"] == ( + "policyengine-v2-outputs-staging" + ) + + +@pytest.mark.parametrize( + ("bucket", "message"), + [ + ("", "V2_SIMULATION_OUTPUT_BUCKET_VALUE must be set"), + ( + "gs://policyengine-v2-outputs", + "V2_SIMULATION_OUTPUT_BUCKET_VALUE is invalid", + ), + ("Policyengine", "V2_SIMULATION_OUTPUT_BUCKET_VALUE is invalid"), + ], +) +def test_a_deploy_without_a_valid_output_bucket_stops(tmp_path, bucket, message): + result, environment = _run_deploy( + tmp_path, V2_SIMULATION_OUTPUT_BUCKET_VALUE=bucket + ) + + assert result.returncode != 0 + assert message in result.stderr + assert not Path(environment["GCLOUD_ARGUMENTS_CAPTURE"]).exists() + + +def test_the_deploy_step_passes_the_environments_output_bucket(): + deploy_entrypoint = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + deploy_step = next( + step + for step in deploy_entrypoint["steps"] + if step.get("name") == "Deploy Cloud Run entrypoint candidate" + ) + + assert deploy_step["env"]["V2_SIMULATION_OUTPUT_BUCKET_VALUE"] == ( + "${{ vars.V2_SIMULATION_OUTPUT_BUCKET }}" + ) diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py new file mode 100644 index 000000000..9ff95623e --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -0,0 +1,399 @@ +"""The staging economy check, run on every staging deploy. + +``simulation-staging-economy-check.sh`` creates a small economy report, +current law against current law for Wales in 2026, through the deployed +candidate's ``POST /reports``, then follows its run through +``GET /reports/{report_id}`` and requires it to succeed. It runs after the +v2 manifest's default has moved to the release's worker, so the report runs +on the release's own coordinator. Once a .py version's first check has +succeeded, later checks on it reuse those simulation runs, and check the +coordinator's use of their stored output. +""" + +from __future__ import annotations + +import json + +import pytest +from test_deployment_assets_support import ( + REUSABLE_DEPLOY_WORKFLOW, + evaluate_github_expression, + routing_context, + workflow_jobs, +) +from workflow_script_support import FakeCommands, run_script + +CHECK_SCRIPT = "simulation-staging-economy-check.sh" +READY_SCRIPT = "simulation-staging-economy-ready.sh" +CHECK_JOB = "staging_economy_check" +BASE_URL = "https://policyengine-simulation-entry-staging-s-123.a.run.app" +REPORT_ID = "00000000-0000-4000-8000-0000000000e1" +REPORT_RUN_ID = "00000000-0000-4000-8000-0000000000e2" +OLDER_RUN_ID = "00000000-0000-4000-8000-0000000000e3" +ACCESS_TOKEN = "test-access-token" + +# The fake curl answers the token request, the submission and each poll in +# turn, and keeps each request body it was sent on standard input. +FAKE_CURL = r""" +output="" +url="" +method="GET" +data="" +while (($#)); do + case "$1" in + --output) output="$2"; shift 2 ;; + --request) method="$2"; shift 2 ;; + --data) data="$2"; shift 2 ;; + --header | --write-out) shift 2 ;; + *) url="$1"; shift ;; + esac +done +calls="$(($(cat "${FAKE_STATE}/calls" 2>/dev/null || echo 0) + 1))" +echo "${calls}" > "${FAKE_STATE}/calls" +if [[ "${data}" == "@-" ]]; then + cat > "${FAKE_STATE}/body-${calls}.json" +fi +case "${method} ${url}" in + "POST https://issuer.example/oauth/token") + printf '{"access_token": "%s"}' "${FAKE_ACCESS_TOKEN}" > "${output}" + printf '200' + ;; + "POST ${FAKE_BASE_URL}/reports") + printf '%s' "${FAKE_SUBMISSION_BODY}" > "${output}" + printf '%s' "${FAKE_SUBMISSION_STATUS}" + ;; + "GET ${FAKE_BASE_URL}/reports/"*) + echo "${url}" > "${FAKE_STATE}/read-url" + polls="$(($(cat "${FAKE_STATE}/polls" 2>/dev/null || echo 0) + 1))" + echo "${polls}" > "${FAKE_STATE}/polls" + answer="$(sed -n "${polls}p" "${FAKE_STATE}/answers")" + printf '%s' "${answer#* }" > "${output}" + printf '%s' "${answer%% *}" + ;; + *) + printf 'unexpected request: %s %s\n' "${method}" "${url}" >&2 + exit 7 + ;; +esac +""" +CREATED = { + "report_id": REPORT_ID, + "report_run_id": REPORT_RUN_ID, + "status": "pending", + "baseline_simulation_id": "00000000-0000-4000-8000-0000000000f1", + "baseline_simulation_run_id": "00000000-0000-4000-8000-0000000000f2", + "baseline_reused": True, + "reform_simulation_id": "00000000-0000-4000-8000-0000000000f3", + "reform_simulation_run_id": "00000000-0000-4000-8000-0000000000f4", + "reform_reused": False, +} + + +def _report(status: str, *, run_id: str = REPORT_RUN_ID, **run_fields) -> str: + return json.dumps( + { + "report_id": REPORT_ID, + "country": "uk", + "baseline_simulation_id": CREATED["baseline_simulation_id"], + "reform_simulation_id": CREATED["reform_simulation_id"], + "run": { + "run_id": run_id, + "status": status, + "country_package_version": "2.174.1", + "policyengine_version": "6.17.1", + "error_message": None, + "created_at": "2026-11-12T15:00:00Z", + "started_at": None, + "completed_at": None, + "baseline_simulation_run_id": CREATED["baseline_simulation_run_id"], + "reform_simulation_run_id": CREATED["reform_simulation_run_id"], + **run_fields, + }, + } + ) + + +@pytest.fixture +def fakes(tmp_path) -> FakeCommands: + commands = FakeCommands.create(tmp_path) + commands.add("curl", FAKE_CURL) + commands.add("sleep") + return commands + + +def _run(fakes, tmp_path, answers: list[str], **overrides): + """Run the check; ``answers`` are the polls' " " in order.""" + state = tmp_path / "state" + state.mkdir() + (state / "answers").write_text("\n".join(answers) + "\n", encoding="utf-8") + environment = { + "SIMULATION_TEST_AUTH_ISSUER": "https://issuer.example/", + "SIMULATION_TEST_AUTH_AUDIENCE": "simulation-entry", + "SIMULATION_TEST_AUTH_CLIENT_ID": "staging-test-client", + "SIMULATION_TEST_AUTH_CLIENT_SECRET": "staging-test-secret", + "FAKE_STATE": str(state), + "FAKE_BASE_URL": BASE_URL, + "FAKE_ACCESS_TOKEN": ACCESS_TOKEN, + "FAKE_SUBMISSION_STATUS": "201", + "FAKE_SUBMISSION_BODY": json.dumps(CREATED), + **overrides, + } + result = run_script(CHECK_SCRIPT, BASE_URL + "/", env=fakes.env(**environment)) + return result, state + + +def test_the_check_creates_a_small_report_and_waits_for_its_run_to_succeed( + fakes, tmp_path +): + result, state = _run( + fakes, + tmp_path, + [ + f"200 {_report('pending')}", + f"200 {_report('running')}", + f"200 {_report('succeeded', completed_at='2026-11-12T15:05:00Z')}", + ], + ) + + assert result.returncode == 0, result.stderr + assert REPORT_ID in result.stdout + assert "reused" in result.stdout + token_request = json.loads((state / "body-1.json").read_text(encoding="utf-8")) + assert token_request["grant_type"] == "client_credentials" + submission = json.loads((state / "body-2.json").read_text(encoding="utf-8")) + assert submission == { + "country": "uk", + "report_type": "economy", + "label": "Staging economy check", + "year": 2026, + "region_code": "country/wales", + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": None, "dynamic_id": None}, + "include_cliffs": False, + } + assert (state / "read-url").read_text(encoding="utf-8").strip() == ( + f"{BASE_URL}/reports/{REPORT_ID}" + ) + assert [call[1:] for call in fakes.calls("sleep")] == [["15"], ["15"]] + + +def test_no_secret_reaches_a_command_line(fakes, tmp_path): + """The token is sent from a file, and the client secret on standard input.""" + result, _ = _run(fakes, tmp_path, [f"200 {_report('succeeded')}"]) + + assert result.returncode == 0, result.stderr + arguments = [argument for call in fakes.calls("curl") for argument in call] + assert not any(ACCESS_TOKEN in argument for argument in arguments) + assert not any("staging-test-secret" in argument for argument in arguments) + + +def test_a_failed_run_fails_the_check_with_its_error_message(fakes, tmp_path): + failed = _report( + "failed", error_message="simulation_failed: the reform ran out of memory" + ) + + result, _ = _run(fakes, tmp_path, [f"200 {failed}"]) + + assert result.returncode == 1 + assert "simulation_failed: the reform ran out of memory" in result.stderr + assert fakes.calls("sleep") == [] + + +def test_a_newest_run_other_than_the_created_one_fails_the_check(fakes, tmp_path): + result, _ = _run( + fakes, tmp_path, [f"200 {_report('succeeded', run_id=OLDER_RUN_ID)}"] + ) + + assert result.returncode == 1 + assert OLDER_RUN_ID in result.stderr + + +@pytest.mark.parametrize( + ("status", "body"), + [ + ("400", {"code": "unknown_region", "detail": "no region country/wales"}), + ("503", {"code": "modal_call_refused", "detail": "x", "report_id": REPORT_ID}), + ], +) +def test_a_refused_submission_fails_the_check_with_its_answer( + fakes, tmp_path, status, body +): + result, state = _run( + fakes, + tmp_path, + [], + FAKE_SUBMISSION_STATUS=status, + FAKE_SUBMISSION_BODY=json.dumps(body), + ) + + assert result.returncode == 1 + assert status in result.stderr + assert body["code"] in result.stderr + assert not (state / "polls").exists() + + +def test_an_unavailable_read_is_asked_again(fakes, tmp_path): + unavailable = json.dumps({"code": "database_unavailable", "detail": "x"}) + + result, _ = _run( + fakes, tmp_path, [f"503 {unavailable}", f"200 {_report('succeeded')}"] + ) + + assert result.returncode == 0, result.stderr + + +def test_any_other_read_answer_fails_the_check(fakes, tmp_path): + result, _ = _run(fakes, tmp_path, ['404 {"code": "unknown_report"}']) + + assert result.returncode == 1 + assert "404" in result.stderr + + +def test_a_run_that_does_not_finish_in_time_fails_the_check(fakes, tmp_path): + result, _ = _run( + fakes, + tmp_path, + [f"200 {_report('pending')}"], + ECONOMY_CHECK_DEADLINE_SECONDS="0", + ) + + assert result.returncode == 1 + assert "did not finish within 0 seconds" in result.stderr + + +def test_the_check_needs_the_test_client_credentials(fakes, tmp_path): + result, _ = _run(fakes, tmp_path, [], SIMULATION_TEST_AUTH_CLIENT_ID="") + + assert result.returncode == 1 + assert "SIMULATION_TEST_AUTH_CLIENT_ID" in result.stderr + assert fakes.calls("curl") == [] + + +# --- The workflow job -------------------------------------------------------- + + +def _jobs() -> dict: + return workflow_jobs(REUSABLE_DEPLOY_WORKFLOW) + + +def _context(**overrides): + context = routing_context() + context["inputs"]["deployment_environment"] = "staging" + context["inputs"]["deploy_stage12_v2"] = True + context["needs"]["deploy_entrypoint"] = {"result": "success"} + context["needs"]["publish_stage12_v2_default"] = {"result": "success"} + for path, value in overrides.items(): + target = context + *parents, last = path.split("__") + for part in parents: + target = target[part] + target[last] = value + return context + + +def test_the_check_job_creates_the_report_through_the_candidate(): + job = _jobs()[CHECK_JOB] + (step,) = [step for step in job["steps"] if "run" in step] + + assert set(job["needs"]) == { + "deploy_entrypoint", + "authenticated_test", + "publish_stage12_v2_default", + } + assert job["environment"] == "${{ inputs.release_environment }}" + assert step["run"] == ( + f".github/scripts/{CHECK_SCRIPT} " + '"${{ needs.deploy_entrypoint.outputs.candidate_url }}"' + ) + assert step["env"] == { + "SIMULATION_TEST_AUTH_ISSUER": ( + "${{ secrets.SIMULATION_ENTRYPOINT_AUTH_ISSUER }}" + ), + "SIMULATION_TEST_AUTH_AUDIENCE": ( + "${{ secrets.SIMULATION_ENTRYPOINT_AUTH_AUDIENCE }}" + ), + "SIMULATION_TEST_AUTH_CLIENT_ID": ( + "${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_ID }}" + ), + "SIMULATION_TEST_AUTH_CLIENT_SECRET": ( + "${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET }}" + ), + } + + +@pytest.mark.parametrize( + ("overrides", "runs"), + [ + ({}, True), + ({"inputs__deployment_environment": "production"}, False), + ({"needs__deploy_entrypoint__result": "failure"}, False), + ({"needs__authenticated_test__result": "failure"}, False), + ({"needs__publish_stage12_v2_default__result": "failure"}, False), + ({"needs__publish_stage12_v2_default__result": "skipped"}, False), + ( + { + "inputs__deploy_stage12_v2": False, + "needs__publish_stage12_v2_default__result": "skipped", + }, + True, + ), + ({"cancelled": True}, False), + ], + ids=[ + "staging", + "production", + "entrypoint-failed", + "authenticated-tests-failed", + "default-move-failed", + "default-move-skipped", + "no-stage12", + "cancelled", + ], +) +def test_the_check_runs_on_every_staging_deploy_after_the_default_moves( + overrides, runs +): + job = _jobs()[CHECK_JOB] + + assert ( + bool(evaluate_github_expression(job["if"], _context(**overrides), job_if=True)) + is runs + ) + + +def test_deployment_ready_requires_the_staging_economy_check(): + ready = _jobs()["deployment_ready"] + (step,) = [ + step + for step in ready["steps"] + if step.get("name") == "Check the staging economy check result" + ] + + assert CHECK_JOB in ready["needs"] + assert step["run"] == f".github/scripts/{READY_SCRIPT}" + assert step["env"] == { + "REQUESTED": "${{ inputs.deployment_environment == 'staging' }}", + "CHECK_RESULT": "${{ needs.staging_economy_check.result }}", + } + + +@pytest.mark.parametrize( + ("requested", "check_result", "passes"), + [ + ("true", "success", True), + ("true", "failure", False), + ("true", "skipped", False), + ("true", "cancelled", False), + ("false", "skipped", True), + ("false", "success", False), + ], +) +def test_the_ready_check_requires_the_check_exactly_in_staging( + fakes, requested, check_result, passes +): + result = run_script( + READY_SCRIPT, + env=fakes.env(REQUESTED=requested, CHECK_RESULT=check_result), + ) + + assert result.returncode == (0 if passes else 1) From 876328cb5f6abf9a6da8e2d033f05572a58b6a01 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:44:15 +0400 Subject: [PATCH 12/26] Name the output bucket on every deploy, and check a staging economy report The entrypoint's deploy passes V2_SIMULATION_OUTPUT_BUCKET, from the environment's GitHub variable of that name, to Cloud Run on every deploy, and stops without a valid bucket name. The new staging_economy_check job creates a small report on every staging deploy, after the v2 default has moved, through the candidate's POST /reports, follows its run through GET /reports/{report_id} and requires it to succeed within 30 minutes (simulation-staging-economy-check.sh). deployment_ready requires its result (simulation-staging-economy-ready.sh), so a failure stops the production deploy. The release environment count now includes the new job. Co-Authored-By: Claude Opus 5.5 --- .../deploy-cloud-run-simulation-entry.sh | 7 ++ .../simulation-staging-economy-check.sh | 119 ++++++++++++++++++ .../simulation-staging-economy-ready.sh | 13 ++ .../workflows/simulation-deploy.reusable.yml | 44 ++++++- .../tests/test_deployment_assets.py | 7 +- 5 files changed, 186 insertions(+), 4 deletions(-) create mode 100755 .github/scripts/simulation-staging-economy-check.sh create mode 100755 .github/scripts/simulation-staging-economy-ready.sh diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index cbe60559e..6f0108438 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -102,6 +102,7 @@ required_variables=( SIMULATION_SERVICE_DATABASE_URL_SECRET_NAME MODAL_TOKEN_ID_SECRET_NAME MODAL_TOKEN_SECRET_SECRET_NAME + V2_SIMULATION_OUTPUT_BUCKET_VALUE OBSERVABILITY_SERVICE_NAMESPACE OBSERVABILITY_TRACE_PROJECT_ID OTEL_EXPORTER_OTLP_ENDPOINT @@ -128,6 +129,11 @@ esac || fail "TAG is invalid" [[ "${ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ || fail "ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT is invalid" +# The bucket that keeps v2 simulation output files for good: a report +# reuses only a succeeded run whose file is in it. The service compares +# output URIs with it and never reads it. +[[ "${V2_SIMULATION_OUTPUT_BUCKET_VALUE}" =~ ^[a-z0-9][a-z0-9._-]{1,61}[a-z0-9]$ ]] \ + || fail "V2_SIMULATION_OUTPUT_BUCKET_VALUE is invalid" [[ "${ENTRYPOINT_MIN_INSTANCES}" =~ ^[0-9]+$ ]] \ || fail "ENTRYPOINT_MIN_INSTANCES must be a non-negative integer" [[ "${ENTRYPOINT_MAX_INSTANCES}" =~ ^[0-9]+$ ]] \ @@ -195,6 +201,7 @@ jq -n ' POLICYENGINE_OTEL_GOOGLE_AUDIENCE: env.POLICYENGINE_OTEL_GOOGLE_AUDIENCE, SIMULATION_V2_MODAL_ENVIRONMENT: env.MODAL_ENVIRONMENT, SIMULATION_V2_MANIFEST_NAME: "simulation-api-v2-version-manifest", + V2_SIMULATION_OUTPUT_BUCKET: env.V2_SIMULATION_OUTPUT_BUCKET_VALUE, STAGE12_ENABLED: ( if env.DEPLOY_STAGE12_V2 == "true" then env.STAGE12_ENABLED_VALUE else "0" end ) diff --git a/.github/scripts/simulation-staging-economy-check.sh b/.github/scripts/simulation-staging-economy-check.sh new file mode 100755 index 000000000..fdae8766f --- /dev/null +++ b/.github/scripts/simulation-staging-economy-check.sh @@ -0,0 +1,119 @@ +#!/bin/bash +# Create a small economy report through a deployed Simulation Entrypoint, and +# require its run to succeed: current law against current law for Wales in +# 2026, which needs no stored policy. +# Usage: simulation-staging-economy-check.sh +# Required env vars: the test client's SIMULATION_TEST_AUTH_ISSUER, +# SIMULATION_TEST_AUTH_AUDIENCE, SIMULATION_TEST_AUTH_CLIENT_ID and +# SIMULATION_TEST_AUTH_CLIENT_SECRET. +# Optional: ECONOMY_CHECK_DEADLINE_SECONDS (1800) and +# ECONOMY_CHECK_POLL_SECONDS (15). +# The access token is sent from a file and the client secret on standard +# input, so neither appears on a command line. + +set -euo pipefail + +fail() { + printf 'Staging economy check failed: %s\n' "$1" >&2 + exit 1 +} + +base_url="${1:?Simulation Entrypoint base URL is required}" +base_url="${base_url%/}" +deadline_seconds="${ECONOMY_CHECK_DEADLINE_SECONDS:-1800}" +poll_seconds="${ECONOMY_CHECK_POLL_SECONDS:-15}" +uuid_pattern='^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$' + +for variable_name in SIMULATION_TEST_AUTH_ISSUER SIMULATION_TEST_AUTH_AUDIENCE \ + SIMULATION_TEST_AUTH_CLIENT_ID SIMULATION_TEST_AUTH_CLIENT_SECRET; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} must be set" +done +[[ "${deadline_seconds}" =~ ^[0-9]+$ ]] \ + || fail "ECONOMY_CHECK_DEADLINE_SECONDS must be a whole number of seconds" +[[ "${poll_seconds}" =~ ^[0-9]+$ ]] \ + || fail "ECONOMY_CHECK_POLL_SECONDS must be a whole number of seconds" + +work_directory="$(mktemp -d)" +trap 'rm -rf "${work_directory}"' EXIT +chmod 700 "${work_directory}" + +# Send one request; write its body to .json and print its HTTP status. +send() { + local name="$1" + shift + curl --silent --show-error --output "${work_directory}/${name}.json" \ + --write-out '%{http_code}' "$@" +} + +token_status="$( + jq -n \ + --arg client_id "${SIMULATION_TEST_AUTH_CLIENT_ID}" \ + --arg client_secret "${SIMULATION_TEST_AUTH_CLIENT_SECRET}" \ + --arg audience "${SIMULATION_TEST_AUTH_AUDIENCE}" \ + '{client_id: $client_id, client_secret: $client_secret, audience: $audience, grant_type: "client_credentials"}' \ + | send token --request POST --header 'Content-Type: application/json' \ + --data @- "${SIMULATION_TEST_AUTH_ISSUER%/}/oauth/token" +)" +[[ "${token_status}" == "200" ]] || fail "the token request answered ${token_status}" +access_token="$(jq -r '.access_token // empty' "${work_directory}/token.json")" +[[ -n "${access_token}" ]] || fail "the token response has no access_token" +authorization_header="${work_directory}/authorization-header" +printf 'Authorization: Bearer %s\n' "${access_token}" > "${authorization_header}" + +submission_status="$( + jq -n '{ + country: "uk", + report_type: "economy", + label: "Staging economy check", + year: 2026, + region_code: "country/wales", + baseline: {policy_id: null, dynamic_id: null}, + reform: {policy_id: null, dynamic_id: null}, + include_cliffs: false + }' \ + | send submission --request POST --header "@${authorization_header}" \ + --header 'Content-Type: application/json' --data @- \ + "${base_url}/reports" +)" +[[ "${submission_status}" == "201" ]] \ + || fail "POST /reports answered ${submission_status}: $(cat "${work_directory}/submission.json")" +submission="${work_directory}/submission.json" +report_id="$(jq -r '.report_id // empty' "${submission}")" +report_run_id="$(jq -r '.report_run_id // empty' "${submission}")" +[[ "${report_id}" =~ ${uuid_pattern} ]] || fail "the submission named no report" +[[ "${report_run_id}" =~ ${uuid_pattern} ]] || fail "the submission named no report run" +printf 'Created report %s, run %s (baseline reused: %s, reform reused: %s)\n' \ + "${report_id}" "${report_run_id}" \ + "$(jq -r '.baseline_reused' "${submission}")" "$(jq -r '.reform_reused' "${submission}")" + +started_at="${SECONDS}" +while true; do + read_status="$( + send report --header "@${authorization_header}" "${base_url}/reports/${report_id}" + )" + report_body="${work_directory}/report.json" + if [[ "${read_status}" == "200" ]]; then + newest_run_id="$(jq -r '.run.run_id // empty' "${report_body}")" + [[ "${newest_run_id}" == "${report_run_id}" ]] \ + || fail "report ${report_id}'s newest run is ${newest_run_id:-missing}, not the created run ${report_run_id}" + run_status="$(jq -r '.run.status' "${report_body}")" + case "${run_status}" in + succeeded) + printf 'Report run %s succeeded on .py %s with %s\n' "${report_run_id}" \ + "$(jq -r '.run.policyengine_version' "${report_body}")" \ + "$(jq -r '.run.country_package_version' "${report_body}")" + exit 0 + ;; + failed) + fail "report run ${report_run_id} failed: $(jq -r '.run.error_message // "no error message"' "${report_body}")" + ;; + pending | running) ;; + *) fail "report run ${report_run_id} has an unknown status: ${run_status}" ;; + esac + elif [[ "${read_status}" != "503" ]]; then + fail "GET /reports/${report_id} answered ${read_status}: $(cat "${report_body}")" + fi + ((SECONDS - started_at < deadline_seconds)) \ + || fail "report run ${report_run_id} did not finish within ${deadline_seconds} seconds" + sleep "${poll_seconds}" +done diff --git a/.github/scripts/simulation-staging-economy-ready.sh b/.github/scripts/simulation-staging-economy-ready.sh new file mode 100755 index 000000000..a10c4f66a --- /dev/null +++ b/.github/scripts/simulation-staging-economy-ready.sh @@ -0,0 +1,13 @@ +#!/bin/bash +# Require the staging economy check to have succeeded in a staging +# deployment, and to have been skipped in any other. +# Required env vars: REQUESTED ("true" in a staging deployment) and +# CHECK_RESULT (the staging_economy_check job result). + +set -euo pipefail + +if [[ "${REQUESTED}" == "true" ]]; then + test "${CHECK_RESULT}" = "success" +else + test "${CHECK_RESULT}" = "skipped" +fi diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 1775e6ddd..aee05a9b7 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -247,6 +247,7 @@ jobs: STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} MODAL_TOKEN_ID_SECRET_NAME: ${{ vars.MODAL_TOKEN_ID_SECRET_NAME }} MODAL_TOKEN_SECRET_SECRET_NAME: ${{ vars.MODAL_TOKEN_SECRET_SECRET_NAME }} + V2_SIMULATION_OUTPUT_BUCKET_VALUE: ${{ vars.V2_SIMULATION_OUTPUT_BUCKET }} run: .github/scripts/deploy-cloud-run-simulation-entry.sh - name: Resolve exact candidate metadata @@ -1252,6 +1253,41 @@ jobs: SIMULATION_TEST_AUTH_CLIENT_SECRET: ${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET }} run: .github/scripts/simulation-staging-household-check.sh "${{ needs.deploy_entrypoint.outputs.candidate_url }}" + staging_economy_check: + name: Create a staging economy report through the candidate + # Every staging deployment creates a small economy report, current law + # against current law for Wales in 2026, through the candidate's + # POST /reports, and requires its run to succeed. The routes use only + # the v2 manifest's default worker, so the check runs after the default + # has moved to the release's worker (or, in a deployment without Stage + # 12, with the default as it is). It calls the tagged candidate URL, + # which serves this release's revision whether or not it was promoted. + # The report's rows are written in the staging database, and its + # coordinator records the run. Once a .py version's first check has + # succeeded, later checks on it reuse those simulation runs. + if: >- + ${{ + !cancelled() && + inputs.deployment_environment == 'staging' && + needs.deploy_entrypoint.result == 'success' && + needs.authenticated_test.result == 'success' && + ((inputs.deploy_stage12_v2 && needs.publish_stage12_v2_default.result == 'success') || + (!inputs.deploy_stage12_v2 && needs.publish_stage12_v2_default.result == 'skipped')) + }} + needs: [deploy_entrypoint, authenticated_test, publish_stage12_v2_default] + runs-on: ubuntu-latest + environment: ${{ inputs.release_environment }} + steps: + - uses: actions/checkout@v6 + + - name: Create the staging economy report and require its run to succeed + env: + SIMULATION_TEST_AUTH_ISSUER: ${{ secrets.SIMULATION_ENTRYPOINT_AUTH_ISSUER }} + SIMULATION_TEST_AUTH_AUDIENCE: ${{ secrets.SIMULATION_ENTRYPOINT_AUTH_AUDIENCE }} + SIMULATION_TEST_AUTH_CLIENT_ID: ${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_ID }} + SIMULATION_TEST_AUTH_CLIENT_SECRET: ${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET }} + run: .github/scripts/simulation-staging-economy-check.sh "${{ needs.deploy_entrypoint.outputs.candidate_url }}" + finalize_routing: name: Finalize the Modal routing publication # Finish the durable journal after every gate, including failed uploads @@ -1310,7 +1346,7 @@ jobs: # The reporter only reads job results. Journal completion happens before # it, so a failure or cancellation here cannot strand a publication. if: ${{ !cancelled() }} - needs: [deploy_entrypoint, update_routing, integration, authenticated_test, promote_entrypoint, publish_latest, finalize_routing, publish_stage12_v2_default, staging_household_check] + needs: [deploy_entrypoint, update_routing, integration, authenticated_test, promote_entrypoint, publish_latest, finalize_routing, publish_stage12_v2_default, staging_household_check, staging_economy_check] runs-on: ubuntu-latest defaults: run: @@ -1345,6 +1381,12 @@ jobs: CHECK_RESULT: ${{ needs.staging_household_check.result }} run: .github/scripts/simulation-staging-household-ready.sh + - name: Check the staging economy check result + env: + REQUESTED: ${{ inputs.deployment_environment == 'staging' }} + CHECK_RESULT: ${{ needs.staging_economy_check.result }} + run: .github/scripts/simulation-staging-economy-ready.sh + cleanup_entrypoint_revisions: name: Remove superseded Cloud Run revisions if: ${{ !cancelled() && inputs.promote_entrypoint && needs.deployment_ready.result == 'success' }} diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index b46faf65d..24e2d1a3e 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -215,10 +215,11 @@ def _assert_deployment_environment_settings(reusable_workflow: str) -> None: assert ( # Every job but stage12_v2_ready declares the release environment. # The job that moves latest after the checks reads the prod - # environment's custom-hostname variable, and the staging household + # environment's custom-hostname variable, the staging household # check reads the beta variable SIMULATION_STAGING_HOUSEHOLD_ID and - # the environment's auth secrets. - reusable_workflow.count("environment: ${{ inputs.release_environment }}") == 18 + # the environment's auth secrets, and the staging economy check + # reads those auth secrets too. + reusable_workflow.count("environment: ${{ inputs.release_environment }}") == 19 ) assert "APP_ENVIRONMENT: ${{ inputs.deployment_environment }}" in reusable_workflow assert "STAGE12_ENABLED_VALUE: ${{ vars.STAGE12_ENABLED }}" in reusable_workflow From 9788790aec9b11fb194a84a040fb1ca730c53a23 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:46:05 +0400 Subject: [PATCH 13/26] Check the economy routes through the regenerated simulation client The generated client sends the Public API's report and simulation bodies, keeps omitted SPM settings omitted, and reads every answer the routes declare: the created records, each error with its code and any record, the 422 with its code, and both reads. Run against a client generated from main, the test fails to import create_report; against this branch's, it passes. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_generated_client.py | 194 ++++++++++++++++++ 1 file changed, 194 insertions(+) create mode 100644 projects/policyengine-apis-integ/tests/test_economy_generated_client.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 new file mode 100644 index 000000000..186929556 --- /dev/null +++ b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py @@ -0,0 +1,194 @@ +"""Hermetic wire checks of the economy routes' generated client models. + +The Public API sends its report and simulation requests through these +models, and reads each answer the routes declare. +""" + +from uuid import UUID + +import httpx +import pytest +from policyengine_api_simulation_client import AuthenticatedClient +from policyengine_api_simulation_client.api.default import ( + create_report, + create_simulation, + get_report, + get_simulation, +) +from policyengine_api_simulation_client.models import ( + EconomyError, + EconomyReportCreated, + EconomyReportRequest, + EconomyReportState, + EconomyRequestUnprocessable, + EconomySimulationCreated, + EconomySimulationRequest, + EconomySimulationState, + ReportRunStatus, +) + +REPORT_ID = "00000000-0000-4000-8000-0000000000e1" +RUN_ID = "00000000-0000-4000-8000-0000000000e2" +SIDE_IDS = [f"00000000-0000-4000-8000-0000000000f{digit}" for digit in "1234"] +POLICY_ID = "00000000-0000-4000-8000-0000000000c5" +# The body the Public API sends (``simulation_reports.report_request``). +REPORT_BODY = { + "country": "us", + "report_type": "economy", + "label": "A larger child tax credit", + "description": None, + "year": 2026, + "region_code": "us", + "dataset_name": None, + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": POLICY_ID, "dynamic_id": None}, + "include_cliffs": False, +} +SIMULATION_BODY = { + "country": "uk", + "simulation_type": "economy", + "year": 2026, + "region_code": "country/wales", + "dataset_name": None, + "policy_id": None, + "dynamic_id": None, +} +CREATED_REPORT = { + "report_id": REPORT_ID, + "report_run_id": RUN_ID, + "status": "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, +} + + +def test_a_report_request_keeps_omitted_spm_settings_omitted(): + """Omitting spm uses the bundle's defaults; null is not a selection.""" + omitted = EconomyReportRequest.from_dict(REPORT_BODY).to_dict() + given = EconomyReportRequest.from_dict( + {**REPORT_BODY, "spm": {"geography_kind": "national"}} + ).to_dict() + + assert "spm" not in omitted + assert omitted["reform"] == {"policy_id": POLICY_ID, "dynamic_id": None} + assert given["spm"] == {"geography_kind": "national"} + + +def _client_answering(status_code: int, body: dict) -> AuthenticatedClient: + transport = httpx.MockTransport( + lambda request: httpx.Response(status_code, json=body) + ) + client = AuthenticatedClient(base_url="https://simulation.example", token="test") + client.set_httpx_client( + httpx.Client(base_url="https://simulation.example", transport=transport) + ) + return client + + +def _create_report(status_code: int, body: dict): + return create_report.sync_detailed( + client=_client_answering(status_code, body), + body=EconomyReportRequest.from_dict(REPORT_BODY), + ) + + +def test_a_created_report_is_read_with_its_runs(): + response = _create_report(201, CREATED_REPORT) + + assert isinstance(response.parsed, EconomyReportCreated) + assert response.parsed.status is ReportRunStatus.PENDING + assert response.parsed.to_dict() == CREATED_REPORT + + +def test_a_created_simulation_is_read_with_its_run(): + created = { + "simulation_id": REPORT_ID, + "simulation_run_id": RUN_ID, + "status": "pending", + } + + response = create_simulation.sync_detailed( + client=_client_answering(201, created), + body=EconomySimulationRequest.from_dict(SIMULATION_BODY), + ) + + assert isinstance(response.parsed, EconomySimulationCreated) + assert response.parsed.to_dict() == created + + +@pytest.mark.parametrize( + ("status_code", "body"), + [ + (400, {"code": "unknown_region", "detail": "no region us"}), + ( + 503, + { + "code": "modal_call_refused", + "detail": "Modal did not accept the report's call.", + "report_id": REPORT_ID, + }, + ), + (500, {"code": "internal_error", "detail": "The service could not answer."}), + ], +) +def test_an_error_is_read_with_its_code_and_any_record(status_code, body): + response = _create_report(status_code, body) + + assert isinstance(response.parsed, EconomyError) + assert response.parsed.to_dict() == body + + +def test_stored_values_that_cannot_be_calculated_are_read_with_their_code(): + body = { + "detail": [ + { + "loc": ["body", "policy_id", "gov.a"], + "msg": "no such parameter", + "type": "missing_parameter", + } + ], + "code": "missing_parameters", + } + + response = _create_report(422, body) + + assert isinstance(response.parsed, EconomyRequestUnprocessable) + assert response.parsed.to_dict() == body + + +def test_a_report_and_a_simulation_are_read_with_their_newest_runs(): + report = { + "report_id": REPORT_ID, + "country": "us", + "baseline_simulation_id": SIDE_IDS[0], + "reform_simulation_id": SIDE_IDS[2], + "run": { + "run_id": RUN_ID, + "status": "succeeded", + "country_package_version": "2.2.1", + "policyengine_version": "6.17.1", + "error_message": None, + "created_at": "2026-11-12T15:00:00Z", + "started_at": "2026-11-12T15:00:05Z", + "completed_at": "2026-11-12T15:04:00Z", + "baseline_simulation_run_id": SIDE_IDS[1], + "reform_simulation_run_id": SIDE_IDS[3], + }, + } + simulation = {"simulation_id": REPORT_ID, "run": None} + + report_response = get_report.sync_detailed( + UUID(REPORT_ID), client=_client_answering(200, report) + ) + simulation_response = get_simulation.sync_detailed( + UUID(REPORT_ID), client=_client_answering(200, simulation) + ) + + assert isinstance(report_response.parsed, EconomyReportState) + assert report_response.parsed.run.status is ReportRunStatus.SUCCEEDED + assert isinstance(simulation_response.parsed, EconomySimulationState) + assert simulation_response.parsed.to_dict() == simulation From db02539a5b2b330005fd85b672d866ed37af0d39 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 19:47:20 +0400 Subject: [PATCH 14/26] Document the economy routes, their deployment and their records The entrypoint's README describes POST /reports, POST /simulations and their reads, the reuse rule and every answer; DEPLOYMENT.md the output bucket setting and the staging economy check; the persistence README the economy records, the reuse search, the store and the dynamic reader. The changelog fragment announces the routes. Co-Authored-By: Claude Opus 5.5 --- changelog.d/c1-9-economy-reports.added.md | 1 + .../README.md | 49 +++++++++++++++-- .../DEPLOYMENT.md | 26 +++++++++ .../policyengine-simulation-entry/README.md | 54 +++++++++++++++++++ 4 files changed, 127 insertions(+), 3 deletions(-) create mode 100644 changelog.d/c1-9-economy-reports.added.md diff --git a/changelog.d/c1-9-economy-reports.added.md b/changelog.d/c1-9-economy-reports.added.md new file mode 100644 index 000000000..4d3d782c4 --- /dev/null +++ b/changelog.d/c1-9-economy-reports.added.md @@ -0,0 +1 @@ +The Simulation Entrypoint creates economy reports and simulations through `POST /reports` and `POST /simulations`, in the Public API's run tables, and starts their calculation on the v2 coordinator and workers by run ID; a report reuses a matching simulation run that already succeeded. `GET /reports/{report_id}` and `GET /simulations/{simulation_id}` return each record's newest run, and every staging deploy checks a small economy report. diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index a049a0575..4051d4862 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -50,9 +50,9 @@ year's report run names that row and its year, and belongs to the same report. ## Input tables `input_tables.py` holds read-only models of the Public API's tables a -calculation is read from: `households`, `policies`, `parameters`, -`parameter_values`, `tax_benefit_models` and `tax_benefit_model_versions`, in -`input_metadata`. The service only reads them, with `SELECT`. Each names only +calculation is read from: `households`, `policies`, `dynamics`, `parameters`, +`parameter_values`, `regions`, `datasets`, `tax_benefit_models` and +`tax_benefit_model_versions`, in `input_metadata`. The service only reads them, with `SELECT`. Each names only the columns the service reads, so they are not models of the whole tables, and they are kept apart from the models of the tables the service writes (`service_metadata`) and from the tables those name by `id` alone @@ -81,6 +81,46 @@ which commits when its block ends and sets `statement_timeout` with `set_config`, local to the transaction, so the database ends a statement that runs too long and the limit never stays on a pooled connection. +## Economy simulations and reports + +`economy_run_records.py` holds the simulation service's reads and writes for +an economy report, or a standalone economy simulation. It reads the +country's model (`read_model_id`) and a catalog version's region by code +(`read_region`) and dataset by name or ID (`read_dataset`, +`read_dataset_by_id`). `insert_economy_report` creates, in the caller's +transaction, each side's new simulation and first run, the report naming +both simulations and the reform's policy with its options in `inputs`, and +its first report run, trigger `initial`; every run is `pending` with no +Modal call ID. A side may instead name a run that already succeeded +(`ReportSimulation.reused`), for which nothing is created. +`insert_economy_simulation` creates a standalone simulation and its first +run. A row naming one that does not exist is refused by its foreign key with +SQLSTATE `23503`; `EconomyRecordReferenceError` names the table and column. + +`find_reusable_run` returns the run a new report's side may reuse: the +newest, by completion, of the same economy simulation (policy, dynamic, +catalog version, dataset, region and year) that succeeded on the same +country package and .py versions, whose output file is under the output +bucket's `gs:///`, that a succeeded report run with the same options +used, and that no report run refused to reuse (an error message starting +`reused_simulation_`). So a submission after a refused reuse calculates the +simulation again. + +`record_report_run_call` and `record_simulation_run_call` record a Modal +call's ID, and `fail_unstarted_report_run` and +`fail_unstarted_simulation_run` fail a run at once, each only while it is +`pending` with no call ID; failing a report run also fails the simulation +runs created for it. `economy_run_reads.py` reads a report, or a simulation, +with its newest run's state. + +`EconomyRunStore` runs them as the Simulation Entrypoint does, each step in +a transaction of its own: `read_report_inputs` reads the catalog version, +model, region and dataset, each side's policy and dynamic values, and each +side's reusable run together, never giving both sides the same run; +`read_simulation_inputs` does the same for one simulation; the writes and +reads each take their own. A missing region or dataset raises +`UnknownRegionError` or `UnknownDatasetError`. + ## Reading a policy `read_policy_reform` (`policy_reform.py`) reads a stored policy as the reform @@ -91,6 +131,9 @@ the parameters of the catalog version it was saved on, so a calculation on any catalog version applies them by name; a parameter the calculation's version lacks raises `MissingPolicyParametersError`, naming each one, and a policy the country doesn't have raises `UnknownPolicyError`. +`read_dynamic_reform` reads a stored dynamic the same way (`UnknownDynamicError` +for none), and `combine_reforms` joins a simulation's policy and dynamic +values, refusing a parameter both set (`OverlappingReformError`). Each period is a range, `YYYY-MM-DD.YYYY-MM-DD`, with the stored dates read in UTC. A value with no end date follows .py's own rule diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index e467e87b8..5972daaff 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -116,6 +116,32 @@ be a US household stored in the staging Public API with a default year; each run adds one `household_jobs` row for it. An empty or malformed variable fails the check. +## Economy reports and simulations + +`POST /reports` and `POST /simulations`, and their reads, use the same +database, Modal token and v2 manifest settings as the household routes, and +the deploy's `V2_SIMULATION_OUTPUT_BUCKET`, the `beta` or `prod` variable of +that name: the bucket that keeps v2 simulation output files for good. A +report reuses only a succeeded simulation run whose file is in it; the +service compares output URIs with it and never reads the bucket, so it needs +no storage permission. A missing or invalid bucket name stops the deploy. +The coordinator and workers get the same bucket from the Modal secret the +Stage 12 secret sync writes. + +### The staging economy check + +Every staging deployment runs `staging_economy_check` when the household +check runs: `.github/scripts/simulation-staging-economy-check.sh` creates a +small report, current law against current law for `country/wales` in 2026, +through the candidate's tagged URL with the authenticated tests' client +credentials, follows it through `GET /reports/{report_id}`, and requires its +run to succeed within 30 minutes. `deployment_ready` requires the check to +pass in staging, so a failure stops the production deployment. Each run adds +a report and its runs to the staging database. Once a .py version's first +check has succeeded, later checks on that version reuse its simulation runs, +so they check the coordinator's use of stored outputs rather than the +workers', which the v2 deployment's own validation calls. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 2e3a3f6ba..bc32cc1f6 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -42,6 +42,10 @@ Auth0 permission: | `GET /internal/stage12/reports/{evaluation_id}` | `stage12:operate` | | `POST /household-calculations` | `simulation:submit` | | `GET /household-calculations/{household_job_id}` | `simulation:read` | +| `POST /reports` | `simulation:submit` | +| `GET /reports/{report_id}` | `simulation:read` | +| `POST /simulations` | `simulation:submit` | +| `GET /simulations/{simulation_id}` | `simulation:read` | `GET /health`, `GET /ready`, `GET /versions`, `GET /versions/{kind}` and `POST /ping` stay public. The entrypoint answers `GET /health` and `POST /ping` @@ -102,3 +106,53 @@ failure or a request the service built and the household contract refuses, is the JSON error with code `internal_error`. `DEPLOYMENT.md` describes the deployment settings and the staging check. +## Economy reports and simulations + +`POST /reports` creates an economy report: a baseline and a reform +simulation of one region's data for one year. It names the `country`, +`report_type` (`economy`), `label`, `description`, `year`, `region_code` +and `dataset_name` (a region and dataset of the country's catalog version +for the default worker's .py version; null for the region's default +dataset), each side's stored `policy_id` and `dynamic_id` (null for current +law), `include_cliffs`, and `spm` (SPM settings; omit it, since null is +refused). `POST /simulations` creates a standalone economy simulation from +the same fields for one policy and dynamic. Neither names a .py version or a +record ID: the service runs the v2 version manifest's default worker and +chooses every ID. + +The service reads the rows itself and, in one transaction, creates the +records and their first runs, `pending` with no Modal call ID. A report +side reuses a simulation run that already succeeded instead of creating +one, when it is the newest run of the same policy, dynamic, catalog +version, dataset, region and year on the same country package and .py +versions, its output file is in the output bucket +(`V2_SIMULATION_OUTPUT_BUCKET`), a succeeded report run with the same +options used it, and no report run has refused to reuse it; the two sides +never name the same run. The service then starts the run's Modal call by +name: the report run's coordinator, given the report and its run IDs and no +Stage 12 evaluation record, or a standalone run's simulation function, +given its run's ID. It records the call's ID, and answers 201 with the +record's ID, its run's ID and status, and, for a report, each side's +simulation and run and whether it was reused. + +A request is refused before anything is written: with 400 and a `code` when +a row it names is missing (`unknown_region`, `unknown_dataset` and +`unknown_dynamic` name the caller's own choice; `unknown_catalog_version`, +`unknown_policy`, `no_economy_worker`, `dataset_not_in_bundle` and +`missing_reference` are the service's state, such as a catalog version +missing while a release deploys), and with 422 and a `code` when a policy or +dynamic sets parameters the catalog version lacks (`missing_parameters`) or +both set one (`overlapping_parameters`). The routes answer 503 with +`Retry-After` when the manifest or the database can't be read. When Modal +refuses the call, or doesn't confirm it within ten seconds, the run (and a +report run's new simulation runs) is failed at once and the 503 +(`modal_call_refused`) names the record, as `report_id` or `simulation_id`; +a run claimed meanwhile is answered with 201. Every 500 is the JSON error +with code `internal_error`, and may follow a write. No `detail` holds an +exception's text. + +`GET /reports/{report_id}` and `GET /simulations/{simulation_id}` return the +record and its newest run (created last, ties to the larger ID): its status, +versions, error message and times, and a report run's simulation runs. The +Public API reads results from the tables itself. + From 40a84c2799db95e008ed838955f1bba452c1e866 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 20:14:23 +0400 Subject: [PATCH 15/26] Test review round 1's fixes and refusal details built from the request (failing) - A report's refused policy or dynamic names its side: the store raises ReportSideError, and a report's 422 locates each reason at baseline. or reform.. - A create whose commit can't be confirmed raises EconomyRecordOutcomeUnknownError, and the route answers 500, which tells the Public API the record may exist, never a 503 saying it doesn't. - Every refusal's detail is built from the request, never from a store error's text. - The Modal call's request ID is the request's own X-Request-ID. - The staging check asks again after a read that got no answer, gives each request a time limit, and its job one longer than its deadline. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_run_store.py | 141 +++++++++++++++- .../tests/test_economy_routes.py | 56 +++++++ .../tests/test_economy_submissions.py | 154 +++++++++++++++++- .../tests/test_staging_economy_check.py | 26 ++- 4 files changed, 372 insertions(+), 5 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py index e09fce0c6..1369aae7a 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py @@ -9,6 +9,7 @@ from __future__ import annotations from contextlib import contextmanager +from dataclasses import replace from datetime import UTC, datetime from typing import TYPE_CHECKING from uuid import uuid4 @@ -25,6 +26,8 @@ insert_reusable_run, ) from policyengine_simulation_persistence import ( + EconomyPopulation, + EconomyRecordOutcomeUnknownError, EconomyRequest, EconomyRunStore, MissingPolicyParametersError, @@ -32,6 +35,7 @@ NewEconomySimulation, OverlappingReformError, ReportRunStatus, + ReportSideError, ReportSimulation, ReusableRun, RunReuse, @@ -45,6 +49,7 @@ ) from public_api_rows import StoredValue, insert_dataset, insert_dynamic, insert_policy from sqlalchemy import text +from sqlalchemy.exc import OperationalError if TYPE_CHECKING: from collections.abc import Iterator @@ -214,7 +219,7 @@ def test_values_the_running_version_lacks_or_both_set_are_refused( # A policy of the US saved on parameters the running version lacks. lacking_policy_id = insert_policy( connection, - saved_on.catalog.__class__(**{**vars(saved_on.catalog), "country": "us"}), + replace(saved_on.catalog, country="us"), [StoredValue(CTC_AMOUNT, 3_000, START)], ) overlapping_policy_id = insert_policy( @@ -361,3 +366,137 @@ def test_a_report_and_a_simulation_are_read_as_they_stand( assert stored_simulation is not None assert stored_simulation.run is not None assert stored_simulation.run.id == report.baseline.simulation_run_id + + +@pytest.mark.parametrize( + ("side", "error_type"), + [ + ("baseline", UnknownDynamicError), + ("reform", UnknownPolicyError), + ("reform", MissingPolicyParametersError), + ("baseline", OverlappingReformError), + ], +) +def test_a_reports_refused_values_name_the_side_they_belong_to( + public_api_database_url, side: str, error_type: type[Exception] +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(ELASTICITY,)) + saved_on = insert_economy_catalog( + connection, country="uk", parameter_names=(CTC_AMOUNT,) + ) + lacking_policy_id = insert_policy( + connection, + replace(saved_on.catalog, country="us"), + [StoredValue(CTC_AMOUNT, 3_000, START)], + ) + overlapping_policy_id = insert_policy( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.1, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + act_as_the_service(connection) + refused = { + UnknownDynamicError: SimulationChoice(None, uuid4()), + UnknownPolicyError: SimulationChoice(uuid4(), None), + MissingPolicyParametersError: SimulationChoice(lacking_policy_id, None), + OverlappingReformError: SimulationChoice(overlapping_policy_id, dynamic_id), + }[error_type] + choices = {"baseline": CURRENT_LAW, "reform": CURRENT_LAW, side: refused} + + with pytest.raises(ReportSideError) as raised: + _store_on(connection).read_report_inputs( + REQUEST, + baseline=choices["baseline"], + reform=choices["reform"], + reuse=None, + ) + + assert raised.value.side == side + assert isinstance(raised.value.error, error_type) + + +class _CommitFails: + """A transaction whose block runs, and whose commit then fails.""" + + def __init__(self, connection) -> None: + self.connection = connection + + def __enter__(self): + return self.connection + + def __exit__(self, error_type, error, traceback) -> None: + if error is None: + statement = "COMMIT" + raise OperationalError(statement, {}, Exception("connection lost")) + + +class _RecordingConnection: + def __init__(self) -> None: + self.statements: list[object] = [] + + def execute(self, statement, *arguments): + self.statements.append(statement) + + +def _new_report() -> NewEconomyReport: + return NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label="Lost commit", + description=None, + tax_benefit_model_id=uuid4(), + population=EconomyPopulation( + tax_benefit_model_version_id=uuid4(), + region_id=uuid4(), + dataset_id=uuid4(), + year=2026, + ), + versions=VERSIONS, + inputs=OPTIONS, + baseline=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + reform=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + ) + + +def test_a_create_whose_commit_fails_may_have_been_written() -> None: + connection = _RecordingConnection() + store = EconomyRunStore(lambda: _CommitFails(connection)) + report = _new_report() + + with pytest.raises(EconomyRecordOutcomeUnknownError) as raised: + store.create_report(report) + with pytest.raises(EconomyRecordOutcomeUnknownError): + store.create_simulation( + NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=report.population, + versions=VERSIONS, + ) + ) + + assert isinstance(raised.value.__cause__, OperationalError) + assert connection.statements + + +class _StatementFails: + def __enter__(self): + return self + + def __exit__(self, error_type, error, traceback) -> None: + return None + + def execute(self, statement, *arguments): + failed_statement = "INSERT" + raise OperationalError(failed_statement, {}, Exception("connection lost")) + + +def test_a_create_whose_statement_fails_is_not_written() -> None: + store = EconomyRunStore(_StatementFails) + + with pytest.raises(OperationalError): + store.create_report(_new_report()) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py index 4b67af99f..1d725e501 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -12,6 +12,7 @@ from conftest import FakeBackend, make_settings from economy_service_fakes import REPORT_BODY, SIMULATION_BODY from fastapi.testclient import TestClient +from policyengine_observability import REQUEST_ID_HEADER 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 ( @@ -21,6 +22,7 @@ EconomySimulationRequest, ) from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, EconomySubmissionRefusedError, EconomySubmissionUnavailableError, EconomySubmissionUnprocessableError, @@ -508,3 +510,57 @@ def test_an_unexpected_read_failure_answers_a_json_500(client_for, record) -> No assert result.status_code == 500 assert result.json() == UNEXPECTED_FAILURE_ANSWER + + +def test_the_calls_request_id_is_the_requests_own(client_for) -> None: + """The observability ID travels in the observability context instead.""" + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + client.post( + "/reports", + json=REPORT_BODY, + headers={ + REQUEST_ID_HEADER: "request-1", + OBSERVABILITY_ID_HEADER: OBSERVABILITY_ID, + }, + ) + client.post("/simulations", json=SIMULATION_BODY) + + ((_, report_request_id, context),) = submissions.reports + ((_, simulation_request_id, _),) = submissions.simulations + assert report_request_id == "request-1" + assert context["observability_id"] == OBSERVABILITY_ID + assert simulation_request_id not in {"", OBSERVABILITY_ID} + + +@pytest.mark.parametrize( + ("path", "body", "record"), + [ + ("/reports", REPORT_BODY, "report"), + ("/simulations", SIMULATION_BODY, "simulation"), + ], +) +def test_a_create_whose_commit_is_unconfirmed_answers_a_500( + client_for, capsys, path, body, record +) -> None: + """A 500 tells the Public API the record may exist; no 503 may say it doesn't.""" + submissions = FakeEconomySubmissions( + failure=EconomySubmissionOutcomeUnknownError(record) + ) + + with client_for(submissions) 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}, but could not confirm " + "it was written." + ), + } + (logged,) = _emitted_records( + capsys.readouterr().out, "economy_record_outcome_unknown" + ) + assert logged["attributes"]["record"] == record diff --git a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py index 6c69bad37..f2435134d 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py @@ -32,6 +32,7 @@ ) from policyengine_simulation_entry.economy_api_models import EconomyReportRequest from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, EconomySubmissionRefusedError, EconomySubmissionUnavailableError, EconomySubmissionUnprocessableError, @@ -44,6 +45,7 @@ ) from policyengine_simulation_persistence import ( CatalogDataset, + EconomyRecordOutcomeUnknownError, EconomyRecordReferenceError, EconomyReportState, EconomySimulationState, @@ -51,6 +53,7 @@ OverlappingReformError, ReportRunState, ReportRunStatus, + ReportSideError, RunReuse, RunVersions, SimulationChoice, @@ -227,7 +230,7 @@ def test_a_missing_row_is_refused_before_anything_is_written(error, code) -> Non _submit_simulation(service(store, caller=caller)) assert report_refusal.value.code == simulation_refusal.value.code == code - assert str(error) in report_refusal.value.detail + assert str(error) not in report_refusal.value.detail assert (store.reports, store.simulations, caller.report_runs) == ([], [], []) @@ -236,7 +239,7 @@ def test_values_the_catalog_version_lacks_cannot_be_calculated() -> None: store = FakeStore(read_failure=error) with pytest.raises(EconomySubmissionUnprocessableError) as refusal: - _submit_report(service(store)) + _submit_simulation(service(store)) assert refusal.value.code == "missing_parameters" assert [reason.loc for reason in refusal.value.reasons] == [ @@ -244,7 +247,7 @@ def test_values_the_catalog_version_lacks_cannot_be_calculated() -> None: ("body", "dynamic_id", "gov.b"), ] assert POLICYENGINE_VERSION in refusal.value.reasons[0].msg - assert store.reports == [] + assert store.simulations == [] def test_a_parameter_the_policy_and_dynamic_both_set_cannot_be_calculated() -> None: @@ -528,3 +531,148 @@ def test_a_read_the_database_cannot_answer_is_unavailable() -> None: with pytest.raises(EconomySubmissionUnavailableError) as unavailable: asyncio.run(read(UUID(int=1))) assert unavailable.value.code == "database_unavailable" + + +@pytest.mark.parametrize( + ("side", "source", "field"), + [("reform", "policy", "policy_id"), ("baseline", "dynamic", "dynamic_id")], +) +def test_a_reports_missing_parameters_are_located_on_their_side( + side, source, field +) -> None: + error = ReportSideError( + side, MissingPolicyParametersError(["gov.a"], source=source) + ) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_report(service(FakeStore(read_failure=error))) + + assert refusal.value.code == "missing_parameters" + assert [reason.loc for reason in refusal.value.reasons] == [ + ("body", side, field, "gov.a") + ] + + +def test_a_reports_overlapping_parameters_are_located_on_their_side() -> None: + error = ReportSideError("baseline", OverlappingReformError(["gov.a"])) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_report(service(FakeStore(read_failure=error))) + + assert [reason.loc for reason in refusal.value.reasons] == [ + ("body", "baseline", "dynamic_id", "gov.a") + ] + + +# Text a store error carries, which no answer may pass on. +ERROR_TEXT = "internal text that must not be passed on" + + +@pytest.mark.parametrize( + ("error", "request_fields", "code", "detail"), + [ + ( + UnknownCatalogVersionError(ERROR_TEXT), + {}, + "unknown_catalog_version", + f"us has no catalog version for .py {POLICYENGINE_VERSION}, the " + "default v2 worker's version.", + ), + ( + UnknownRegionError(ERROR_TEXT), + {"region_code": "state/zz"}, + "unknown_region", + f"us's catalog version for .py {POLICYENGINE_VERSION} has no region " + "state/zz.", + ), + ( + UnknownDatasetError(ERROR_TEXT), + {"dataset_name": "acs_2099"}, + "unknown_dataset", + f"us's catalog version for .py {POLICYENGINE_VERSION} has no dataset " + "acs_2099.", + ), + ( + UnknownDatasetError(ERROR_TEXT), + {}, + "unknown_dataset", + f"us's catalog version for .py {POLICYENGINE_VERSION} has no default " + "dataset for region us.", + ), + ( + ReportSideError("reform", UnknownPolicyError(ERROR_TEXT)), + {}, + "unknown_policy", + "us has no stored policy the request's reform names.", + ), + ( + ReportSideError("baseline", UnknownDynamicError(ERROR_TEXT)), + {}, + "unknown_dynamic", + "No stored dynamic matches the one the request's baseline names.", + ), + ], +) +def test_a_refusals_detail_is_built_from_the_request_not_the_error( + error, request_fields, code, detail +) -> None: + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report( + service(FakeStore(read_failure=error)), report_request(**request_fields) + ) + + assert (refusal.value.code, refusal.value.detail) == (code, detail) + assert ERROR_TEXT not in refusal.value.detail + + +def test_a_standalone_simulations_unknown_rows_name_its_own_fields() -> None: + for error, detail in ( + (UnknownPolicyError(ERROR_TEXT), "us has no stored policy the request names."), + ( + UnknownDynamicError(ERROR_TEXT), + "No stored dynamic matches the one the request names.", + ), + ): + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_simulation(service(FakeStore(read_failure=error))) + assert refusal.value.detail == detail + + +def test_a_country_without_an_economy_worker_is_refused_in_the_services_words() -> None: + workers = FakeWorkers(EconomyWorkerUnavailableError(ERROR_TEXT)) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(FakeStore(), workers=workers)) + + assert refusal.value.detail == ( + "The default v2 worker cannot calculate economy runs for us." + ) + + +def test_a_dataset_the_bundle_lacks_is_named_in_the_services_words() -> None: + catalog = replace( + CATALOG, dataset=CatalogDataset(id=CATALOG.dataset.id, name="enhanced_cps_2024") + ) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(FakeStore(catalog=catalog))) + + assert refusal.value.detail == ( + f"The default v2 worker, .py {POLICYENGINE_VERSION}, has no dataset " + "enhanced_cps_2024." + ) + + +@pytest.mark.parametrize("record", ["report", "simulation"]) +def test_a_create_whose_commit_is_unconfirmed_may_have_written_the_record( + record, +) -> None: + store = FakeStore(create_failure=EconomyRecordOutcomeUnknownError("lost")) + caller = FakeCaller() + submit = _submit_report if record == "report" else _submit_simulation + + with pytest.raises(EconomySubmissionOutcomeUnknownError) as unknown: + submit(service(store, caller=caller)) + + assert unknown.value.record == record + assert (caller.report_runs, caller.simulation_runs) == ([], []) diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py index 9ff95623e..5f62407da 100644 --- a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -44,7 +44,7 @@ --output) output="$2"; shift 2 ;; --request) method="$2"; shift 2 ;; --data) data="$2"; shift 2 ;; - --header | --write-out) shift 2 ;; + --header | --write-out | --max-time) shift 2 ;; *) url="$1"; shift ;; esac done @@ -67,6 +67,10 @@ polls="$(($(cat "${FAKE_STATE}/polls" 2>/dev/null || echo 0) + 1))" echo "${polls}" > "${FAKE_STATE}/polls" answer="$(sed -n "${polls}p" "${FAKE_STATE}/answers")" + if [[ "${answer}" == "TRANSPORT" ]]; then + # curl's exit status when the request times out. + exit 28 + fi printf '%s' "${answer#* }" > "${output}" printf '%s' "${answer%% *}" ;; @@ -242,6 +246,24 @@ def test_an_unavailable_read_is_asked_again(fakes, tmp_path): assert result.returncode == 0, result.stderr +def test_a_read_that_gets_no_answer_is_asked_again(fakes, tmp_path): + result, _ = _run( + fakes, tmp_path, ["TRANSPORT", "TRANSPORT", f"200 {_report('succeeded')}"] + ) + + assert result.returncode == 0, result.stderr + assert "no answer" in result.stderr + + +def test_every_request_has_a_time_limit(fakes, tmp_path): + result, _ = _run(fakes, tmp_path, [f"200 {_report('succeeded')}"]) + + assert result.returncode == 0, result.stderr + calls = fakes.calls("curl") + assert calls + assert all("--max-time" in call for call in calls) + + def test_any_other_read_answer_fails_the_check(fakes, tmp_path): result, _ = _run(fakes, tmp_path, ['404 {"code": "unknown_report"}']) @@ -301,6 +323,8 @@ def test_the_check_job_creates_the_report_through_the_candidate(): "publish_stage12_v2_default", } assert job["environment"] == "${{ inputs.release_environment }}" + # Longer than the script's own 30-minute deadline, so the script fails first. + assert job["timeout-minutes"] == 40 assert step["run"] == ( f".github/scripts/{CHECK_SCRIPT} " '"${{ needs.deploy_entrypoint.outputs.candidate_url }}"' From 1eca0d43ef54f92a153524b2c3eda9748e05718c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 20:20:16 +0400 Subject: [PATCH 16/26] Fix review round 1: locate a report side's refusals, and never answer 503 after a possible write - A report side's refused policy or dynamic is raised as ReportSideError, and a report's 422 locates each reason at baseline. or reform. rather than at a field the body doesn't have. - A create whose commit fails after its rows were inserted raises EconomyRecordOutcomeUnknownError: the commit may have been applied. The route answers the JSON 500 (internal_error), which the Public API reads as a record that may exist, instead of a 503 saying nothing was created, and logs economy_record_outcome_unknown. - Every refusal's detail is written from the request and the worker, never from an error's text, since the Public API passes 4xx details on to its callers. - The Modal call's request_id is the request's X-Request-ID, not the observability ID, which travels in the observability context. - The staging check gives each request 30 seconds, asks again after a read that got no answer, and its job stops after 40 minutes. - Two tests of this branch follow the changes: a report read's policy and dynamic refusals now arrive with their side, and refusal details no longer hold the store's text. Docstrings now say a run claimed meanwhile is answered pending, and that a reused run may still be refused; the unused CALLER_ROW_REFUSAL_CODES is gone. Co-Authored-By: Claude Opus 5.5 --- .../simulation-staging-economy-check.sh | 23 ++-- .../workflows/simulation-deploy.reusable.yml | 2 + .../__init__.py | 4 + .../economy_run_records.py | 9 +- .../economy_run_store.py | 103 +++++++++++++--- .../tests/test_economy_run_store.py | 8 +- .../economy_api_models.py | 20 +-- .../economy_routes.py | 45 +++++-- .../economy_submission_errors.py | 115 ++++++++++++++---- .../economy_submissions.py | 51 ++++++-- .../tests/test_economy_routes.py | 5 +- .../tests/test_economy_submissions.py | 10 +- 12 files changed, 305 insertions(+), 90 deletions(-) diff --git a/.github/scripts/simulation-staging-economy-check.sh b/.github/scripts/simulation-staging-economy-check.sh index fdae8766f..23d092b05 100755 --- a/.github/scripts/simulation-staging-economy-check.sh +++ b/.github/scripts/simulation-staging-economy-check.sh @@ -9,7 +9,9 @@ # Optional: ECONOMY_CHECK_DEADLINE_SECONDS (1800) and # ECONOMY_CHECK_POLL_SECONDS (15). # The access token is sent from a file and the client secret on standard -# input, so neither appears on a command line. +# input, so neither appears on a command line. Each request is given +# REQUEST_TIMEOUT_SECONDS; a read that gets no answer is asked again, as a +# 503 is, until the deadline. set -euo pipefail @@ -22,6 +24,7 @@ base_url="${1:?Simulation Entrypoint base URL is required}" base_url="${base_url%/}" deadline_seconds="${ECONOMY_CHECK_DEADLINE_SECONDS:-1800}" poll_seconds="${ECONOMY_CHECK_POLL_SECONDS:-15}" +request_timeout_seconds=30 uuid_pattern='^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$' for variable_name in SIMULATION_TEST_AUTH_ISSUER SIMULATION_TEST_AUTH_AUDIENCE \ @@ -38,11 +41,12 @@ trap 'rm -rf "${work_directory}"' EXIT chmod 700 "${work_directory}" # Send one request; write its body to .json and print its HTTP status. +# It fails, printing nothing, when no answer arrives. send() { local name="$1" shift - curl --silent --show-error --output "${work_directory}/${name}.json" \ - --write-out '%{http_code}' "$@" + curl --silent --show-error --max-time "${request_timeout_seconds}" \ + --output "${work_directory}/${name}.json" --write-out '%{http_code}' "$@" } token_status="$( @@ -53,7 +57,7 @@ token_status="$( '{client_id: $client_id, client_secret: $client_secret, audience: $audience, grant_type: "client_credentials"}' \ | send token --request POST --header 'Content-Type: application/json' \ --data @- "${SIMULATION_TEST_AUTH_ISSUER%/}/oauth/token" -)" +)" || fail "the token request got no answer" [[ "${token_status}" == "200" ]] || fail "the token request answered ${token_status}" access_token="$(jq -r '.access_token // empty' "${work_directory}/token.json")" [[ -n "${access_token}" ]] || fail "the token response has no access_token" @@ -74,7 +78,7 @@ submission_status="$( | send submission --request POST --header "@${authorization_header}" \ --header 'Content-Type: application/json' --data @- \ "${base_url}/reports" -)" +)" || fail "POST /reports got no answer; it may have created a report" [[ "${submission_status}" == "201" ]] \ || fail "POST /reports answered ${submission_status}: $(cat "${work_directory}/submission.json")" submission="${work_directory}/submission.json" @@ -88,11 +92,12 @@ printf 'Created report %s, run %s (baseline reused: %s, reform reused: %s)\n' \ started_at="${SECONDS}" while true; do - read_status="$( - send report --header "@${authorization_header}" "${base_url}/reports/${report_id}" - )" report_body="${work_directory}/report.json" - if [[ "${read_status}" == "200" ]]; then + if ! read_status="$( + send report --header "@${authorization_header}" "${base_url}/reports/${report_id}" + )"; then + printf 'GET /reports/%s got no answer; asking again\n' "${report_id}" >&2 + elif [[ "${read_status}" == "200" ]]; then newest_run_id="$(jq -r '.run.run_id // empty' "${report_body}")" [[ "${newest_run_id}" == "${report_run_id}" ]] \ || fail "report ${report_id}'s newest run is ${newest_run_id:-missing}, not the created run ${report_run_id}" diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index aee05a9b7..b89de5283 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -1276,6 +1276,8 @@ jobs: }} needs: [deploy_entrypoint, authenticated_test, publish_stage12_v2_default] runs-on: ubuntu-latest + # Longer than the script's own 30-minute deadline, so the script fails first. + timeout-minutes: 40 environment: ${{ inputs.release_environment }} steps: - uses: actions/checkout@v6 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 85b0c7407..1183bedf9 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -39,10 +39,12 @@ ) from policyengine_simulation_persistence.economy_run_store import ( EconomyCatalogRows, + EconomyRecordOutcomeUnknownError, EconomyReportInputs, EconomyRequest, EconomyRunStore, EconomySimulationInputs, + ReportSideError, RunReuse, SimulationInputs, UnknownDatasetError, @@ -105,6 +107,7 @@ "CatalogRegion", "EconomyCatalogRows", "EconomyPopulation", + "EconomyRecordOutcomeUnknownError", "EconomyRecordReferenceError", "EconomyReportInputs", "EconomyReportState", @@ -126,6 +129,7 @@ "ReportRunState", "ReportRunStatus", "ReportRunTrigger", + "ReportSideError", "ReportSimulation", "ReusableRun", "RunReuse", 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 cb44d76f4..dd90ddabd 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 @@ -290,14 +290,17 @@ def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match - 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`` - (options), so its file holds what such a report needs; + (options), so its file holds the outputs a report with those options + planned; - that no report run has refused to reuse (an error message starting ``reused_simulation_``), so a submission after a refused reuse calculates the simulation again; and - that is not ``excluded_run_id``, the other side's run. - The coordinator may still refuse it, for example when labour supply - responses differ between the two reports. + The coordinator may still refuse it, when the new report plans outputs + that report did not, as for a reform that turns on labour supply + responses: the refusal fails that report run, and the run is not offered + again. """ used_by = report_runs.alias("used_by_report_run") refused_by = report_runs.alias("refused_by_report_run") diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py index ec41c804a..26874604d 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 @@ -10,7 +10,9 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal + +from sqlalchemy.exc import SQLAlchemyError from policyengine_simulation_persistence.economy_run_reads import ( read_economy_report, @@ -43,6 +45,10 @@ UnknownCatalogVersionError, ) from policyengine_simulation_persistence.policy_reform import ( + MissingPolicyParametersError, + OverlappingReformError, + UnknownDynamicError, + UnknownPolicyError, combine_reforms, read_dynamic_reform, read_policy_reform, @@ -70,6 +76,15 @@ ) type TransactionFactory = Callable[[], AbstractContextManager[Connection]] +# Which of a report's simulations a value belongs to. +type ReportSide = Literal["baseline", "reform"] +# The errors a simulation's own policy or dynamic can raise. +_SIDE_ERRORS = ( + UnknownPolicyError, + UnknownDynamicError, + MissingPolicyParametersError, + OverlappingReformError, +) class UnknownRegionError(LookupError): @@ -80,6 +95,29 @@ class UnknownDatasetError(LookupError): """The catalog version has no dataset with the requested name.""" +class ReportSideError(Exception): + """A report side's policy or dynamic is refused. + + ``side`` names the side, and ``error`` is the policy reader's refusal: + UnknownPolicyError, UnknownDynamicError, MissingPolicyParametersError or + OverlappingReformError. + """ + + def __init__(self, side: ReportSide, error: Exception) -> None: + """Keep the side and the refusal.""" + super().__init__(f"the report's {side}: {error}") + self.side: ReportSide = side + self.error = error + + +class EconomyRecordOutcomeUnknownError(RuntimeError): + """A create's statements ran, but its commit failed: the rows may exist. + + PostgreSQL may have applied the commit before the connection failed, so + the caller must not answer that nothing was created. + """ + + @dataclass(frozen=True) class EconomyRequest: """What an economy simulation or report request names, apart from policies. @@ -184,7 +222,10 @@ def read_simulation_inputs( ) -> EconomySimulationInputs: """Read a standalone simulation's catalog rows and values together. - Raise the errors ``read_report_inputs`` names. + Raise UnknownCatalogVersionError, UnknownRegionError or + UnknownDatasetError, or the policy reader's UnknownPolicyError, + UnknownDynamicError or MissingPolicyParametersError, or + OverlappingReformError. """ with self._transaction() as connection: catalog = _catalog_rows(connection, request) @@ -205,19 +246,25 @@ def read_report_inputs( With ``reuse``, each side gets the run ``find_reusable_run`` finds, and the reform never gets the baseline's. Raise - UnknownCatalogVersionError, UnknownRegionError, UnknownDatasetError, - the policy reader's UnknownPolicyError, UnknownDynamicError or - MissingPolicyParametersError, or OverlappingReformError. + UnknownCatalogVersionError, UnknownRegionError or UnknownDatasetError, + or ReportSideError naming the side whose policy or dynamic is refused. """ with self._transaction() as connection: catalog = _catalog_rows(connection, request) population = catalog.population(request.year) sides = [] excluded_run_id = None - for choice in (baseline, reform): - inputs = _simulation_inputs( - connection, request.country, catalog.catalog_version_id, choice - ) + sides_requested: tuple[tuple[ReportSide, SimulationChoice], ...] = ( + ("baseline", baseline), + ("reform", reform), + ) + 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, @@ -239,14 +286,40 @@ def read_report_inputs( return EconomyReportInputs(catalog=catalog, baseline=sides[0], reform=sides[1]) def create_report(self, report: NewEconomyReport) -> None: - """Create the report, its new simulations and its runs.""" - with self._transaction() as connection: - insert_economy_report(connection, report) + """Create the report, its new simulations and its runs. + + Raise ``insert_economy_report``'s errors, and + EconomyRecordOutcomeUnknownError when the commit fails after the + rows were inserted. + """ + self._create(lambda connection: insert_economy_report(connection, report)) def create_simulation(self, simulation: NewEconomySimulation) -> None: - """Create the simulation and its run; see ``insert_economy_simulation``.""" - with self._transaction() as connection: - insert_economy_simulation(connection, simulation) + """Create the simulation and its run; see ``insert_economy_simulation``. + + Raise EconomyRecordOutcomeUnknownError as ``create_report`` does. + """ + self._create( + lambda connection: insert_economy_simulation(connection, simulation) + ) + + def _create(self, insert: Callable[[Connection], None]) -> None: + """Run ``insert`` in a transaction, telling a failed commit apart. + + An error from ``insert`` leaves nothing written and propagates. An + error once it has run is the commit's, after which the rows may + exist: EconomyRecordOutcomeUnknownError. + """ + inserted = False + try: + with self._transaction() as connection: + insert(connection) + inserted = True + except SQLAlchemyError as error: + if not inserted: + raise + msg = "the create's commit failed after its rows were inserted" + raise EconomyRecordOutcomeUnknownError(msg) from error def record_report_call( self, report_run_id: UUID, modal_call_id: str diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py index 1369aae7a..2a40fe4da 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py @@ -200,13 +200,19 @@ def test_each_missing_row_a_request_names_is_refused( store = _store_on(connection) request = EconomyRequest(**{**vars(REQUEST), **request_fields}) - with pytest.raises(error_type): + # A report names the side whose policy or dynamic is missing. + side_refused = missing in {"policy", "dynamic"} + with pytest.raises(ReportSideError if side_refused else error_type) as raised: store.read_report_inputs( request, baseline=CURRENT_LAW, reform=choice, reuse=REUSE ) with pytest.raises(error_type): store.read_simulation_inputs(request, choice) + if side_refused: + assert raised.value.side == "reform" + assert isinstance(raised.value.error, error_type) + def test_values_the_running_version_lacks_or_both_set_are_refused( public_api_database_url, 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 b67454886..a4df869bd 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 @@ -19,7 +19,8 @@ - 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; -- 500 with ``code`` ``internal_error``, possibly after a write. +- 500 with ``code`` ``internal_error``, possibly after a write: an + unexpected failure, or a commit that failed after the rows were inserted. ``detail`` never holds an exception's text, and ``code`` is a fixed identifier. @@ -28,7 +29,7 @@ from __future__ import annotations from datetime import datetime -from typing import Annotated, Any, Final, Literal +from typing import Annotated, Any, Literal from uuid import UUID from policyengine_simulation_contract.spm import SPMSelection @@ -46,9 +47,6 @@ # 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. CallerRowRefusalCode = Literal["unknown_region", "unknown_dataset", "unknown_dynamic"] -CALLER_ROW_REFUSAL_CODES: Final = frozenset( - {"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. ServiceStateRefusalCode = Literal[ @@ -170,8 +168,10 @@ class _Answer(BaseModel): class EconomySimulationCreated(_Answer): """The simulation a request created, its first run and that run's status. - The status is the run's as the service last wrote or read it: ``pending`` - unless the simulation function had claimed the run by then. + The status is the run's when its call's ID was recorded: ``pending``, or + ``running`` when the simulation function had claimed it first. A run + whose call Modal didn't confirm, or whose call's ID couldn't be + recorded, is answered ``pending``. """ simulation_id: UUID @@ -182,8 +182,10 @@ class EconomySimulationCreated(_Answer): class EconomyReportCreated(_Answer): """The report a request created, its first run and that run's status. - The status is the run's as the service last wrote or read it: ``pending`` - unless the coordinator had claimed the run by then. Each side names its + The status is the run's when its call's ID was recorded: ``pending``, or + ``running`` when the coordinator had claimed it first. A run whose call + Modal didn't confirm, or whose call's ID couldn't be recorded, is + answered ``pending``. Each side names its simulation and simulation run; ``…_reused`` is true for a side that names a run that already succeeded, rather than a new one. """ 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 e54be88f6..b8d41ec89 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 @@ -21,7 +21,8 @@ 500 is the JSON error with code ``internal_error``, and may follow a write. A request gets its observability ID from the caller's header, or a new one, -and passes it to the Modal call it starts. +and passes it to the Modal call it starts in the observability context; the +call's ``request_id`` is the request's own ``X-Request-ID``. """ from __future__ import annotations @@ -46,6 +47,7 @@ SimulationRunRead, ) from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, EconomySubmissionRefusedError, EconomySubmissionUnavailableError, EconomySubmissionUnprocessableError, @@ -110,8 +112,8 @@ } _INTERNAL_ERROR = { "description": ( - "The service failed unexpectedly (internal_error), possibly after " - "creating the record" + "The service failed unexpectedly, or could not confirm its rows were " + "written (internal_error), possibly after creating the record" ), **_ERROR, } @@ -205,6 +207,7 @@ def _not_configured() -> JSONResponse: EconomySubmissionRefusedError, EconomySubmissionUnprocessableError, EconomySubmissionUnavailableError, + EconomySubmissionOutcomeUnknownError, ) @@ -226,6 +229,8 @@ def _submission_failure_response(error: Exception, record: RecordName) -> JSONRe }, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, ) + if isinstance(error, EconomySubmissionOutcomeUnknownError): + return _outcome_unknown_response(error) 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( @@ -234,6 +239,32 @@ def _submission_failure_response(error: Exception, record: RecordName) -> JSONRe ) +def _outcome_unknown_response( + error: EconomySubmissionOutcomeUnknownError, +) -> JSONResponse: + """Answer a create whose commit can't be confirmed with the JSON 500. + + The Public API reads a 500 as a record that may exist; a 503 would tell + it nothing was created. The log is the fixed event + ``economy_record_outcome_unknown``. + """ + fields = {"record": error.record} + logger.error( + "economy_record_outcome_unknown", + extra={**fields, "policyengine_attributes": fields}, + ) + 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." + ), + ), + ) + + def _unexpected_failure_response( route: str, error: Exception, *, max_characters: int ) -> JSONResponse: @@ -371,7 +402,7 @@ async def create_report( try: created = await submissions.submit_report( body, - request_id=observability_id, + request_id=request.state.request_id, observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: @@ -415,7 +446,7 @@ async def create_simulation( try: created = await submissions.submit_simulation( body, - request_id=observability_id, + request_id=request.state.request_id, observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: @@ -471,7 +502,7 @@ def _read_routes( @router.get( REPORT_PATH, - summary="Get an economy report's newest run", + summary="Get a report's newest run", description=( "Return the report, its two simulations and its newest run's " "status, versions, error message and times." @@ -500,7 +531,7 @@ async def get_report(report_id: UUID) -> JSONResponse: @router.get( SIMULATION_PATH, - summary="Get an economy simulation's newest run", + summary="Get a simulation's newest run", description=( "Return the simulation and its newest run's status, versions, error " "message and times." 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 f4382898c..0a1bda43d 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 @@ -2,8 +2,11 @@ The submission service raises these, and the routes answer each with its status: EconomySubmissionRefusedError 400, EconomySubmissionUnprocessableError -422 and EconomySubmissionUnavailableError 503. ``read_refusal`` and -``reference_refusal`` turn the store's errors into them. +422, EconomySubmissionUnavailableError 503 and +EconomySubmissionOutcomeUnknownError 500. ``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. """ from __future__ import annotations @@ -14,6 +17,7 @@ from policyengine_simulation_persistence import ( MissingPolicyParametersError, OverlappingReformError, + ReportSideError, UnknownCatalogVersionError, UnknownDatasetError, UnknownDynamicError, @@ -24,7 +28,10 @@ if TYPE_CHECKING: from uuid import UUID - from policyengine_simulation_persistence import EconomyRecordReferenceError + from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, + EconomyRequest, + ) from policyengine_simulation_entry.economy_api_models import ( EconomyRefusalCode, @@ -85,6 +92,20 @@ def __init__( self.reasons = reasons +class EconomySubmissionOutcomeUnknownError(Exception): + """The submission's rows may have been written, but that can't be confirmed. + + 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. + """ + + def __init__(self, record: RecordName) -> None: + """Keep what the submission creates.""" + super().__init__(f"the {record} may have been created") + self.record: RecordName = record + + class EconomySubmissionUnavailableError(Exception): """A dependency could not be read or refused the call; ask again later. @@ -106,44 +127,84 @@ def __init__( self.record_id = record_id -# Each store error naming a row the request names that doesn't exist, and -# the code it is refused with. -_MISSING_ROW_CODES: dict[type[LookupError], EconomyRefusalCode] = { - UnknownRegionError: "unknown_region", - UnknownDatasetError: "unknown_dataset", - UnknownDynamicError: "unknown_dynamic", - UnknownPolicyError: "unknown_policy", -} - - def read_refusal( - error: Exception, *, policyengine_version: str + error: Exception, request: EconomyRequest ) -> 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 for - ``policyengine_version`` lacks, or both setting one, cannot be - calculated (422). + (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 + written from ``request``, never from the error's text. """ + side = None + if isinstance(error, ReportSideError): + side, error = error.side, error.error + return _missing_row_refusal(error, request, side) or _unprocessable( + error, request, side + ) + + +def _missing_row_refusal( + error: Exception, request: EconomyRequest, side: str | None +) -> 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}" + refusal: tuple[EconomyRefusalCode, str] | None = None if isinstance(error, UnknownCatalogVersionError): - return EconomySubmissionRefusedError( - code="unknown_catalog_version", - detail=f"{error}, the default v2 worker's version", + refusal = ( + "unknown_catalog_version", + f"{request.country} has no catalog version for .py " + 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}.") + 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}.") + elif isinstance(error, UnknownPolicyError): + refusal = ( + "unknown_policy", + f"{request.country} has no stored policy {named_by} names.", ) - for error_type, code in _MISSING_ROW_CODES.items(): - if isinstance(error, error_type): - return EconomySubmissionRefusedError(code=code, detail=str(error)) + elif isinstance(error, UnknownDynamicError): + refusal = ( + "unknown_dynamic", + f"No stored dynamic matches the one {named_by} names.", + ) + if refusal is None: + return None + code, detail = refusal + return EconomySubmissionRefusedError(code=code, detail=detail) + + +def _unprocessable( + error: Exception, request: EconomyRequest, side: str | None +) -> 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. + """ + location: tuple[str, ...] = ("body",) if side is None else ("body", side) if isinstance(error, MissingPolicyParametersError): field_name = "policy_id" if error.source == "policy" else "dynamic_id" return EconomySubmissionUnprocessableError( code="missing_parameters", reasons=tuple( UnprocessableReason( - loc=("body", field_name, name), + loc=(*location, field_name, name), msg=( - f"the catalog version for .py {policyengine_version} has " - "no such parameter" + f"the catalog version for .py {request.policyengine_version} " + "has no such parameter" ), type="missing_parameter", ) @@ -155,7 +216,7 @@ def read_refusal( code="overlapping_parameters", reasons=tuple( UnprocessableReason( - loc=("body", "dynamic_id", name), + loc=(*location, "dynamic_id", name), msg="the policy and the dynamic both set this parameter", type="overlapping_parameter", ) 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 7bcd0fad6..6a1c549e5 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 @@ -20,7 +20,7 @@ Modal refuses the call, or does not confirm it in time, it fails the report run and its new simulation runs at once instead, under the same condition; a run the coordinator claimed meanwhile, because Modal - accepted the call after all, is answered as it stands. + accepted the call after all, is left as it is and answered ``pending``. A standalone simulation is created and started the same way, with one simulation run calculated by the country's simulation function. @@ -42,6 +42,7 @@ from policyengine_simulation_contract.stage12_execution import ReportRunIds from policyengine_simulation_persistence import ( + EconomyRecordOutcomeUnknownError, EconomyRecordReferenceError, EconomyRequest, NewEconomyReport, @@ -68,6 +69,7 @@ simulation_run_calculation, ) from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, EconomySubmissionRefusedError, EconomySubmissionUnavailableError, RecordName, @@ -270,7 +272,7 @@ async def submit_report( reform=_choice(request.reform), reuse=reuse, ), - policyengine_version=worker.policyengine_version, + request=economy_request, ) scope = _scope(request, worker, inputs.catalog, spm=options.get("spm")) report = NewEconomyReport( @@ -298,7 +300,7 @@ async def submit_report( baseline_policy=inputs.baseline.reform, reform_policy=inputs.reform.reform, ) - await self._create(lambda: self._store.create_report(report)) + await self._create("report", lambda: self._store.create_report(report)) status = await self._start_report_run( worker, report, calculation, observability_context=observability_context ) @@ -333,7 +335,7 @@ async def submit_simulation( ) inputs = await self._read_inputs( lambda: self._store.read_simulation_inputs(economy_request, choice), - policyengine_version=worker.policyengine_version, + request=economy_request, ) scope = _scope(request, worker, inputs.catalog, spm=None) simulation = NewEconomySimulation( @@ -350,7 +352,9 @@ async def submit_simulation( simulation_run_id=simulation.run_id, policy=inputs.simulation.reform, ) - await self._create(lambda: self._store.create_simulation(simulation)) + await self._create( + "simulation", lambda: self._store.create_simulation(simulation) + ) status = await self._start_simulation_run( worker, simulation, @@ -418,9 +422,13 @@ def _origin(self, request_id: str) -> CallOrigin: async def _default_worker(self, country: str) -> EconomyWorker: try: return await asyncio.to_thread(self._workers.default_worker, country) - except EconomyWorkerUnavailableError as error: + except EconomyWorkerUnavailableError: raise EconomySubmissionRefusedError( - code="no_economy_worker", detail=str(error) + code="no_economy_worker", + detail=( + "The default v2 worker cannot calculate economy runs for " + f"{country}." + ), ) from None except VersionManifestUnavailableError as error: logger.warning( @@ -433,7 +441,7 @@ async def _default_worker(self, country: str) -> EconomyWorker: ) from None async def _read_inputs[Inputs]( - self, read: Callable[[], Inputs], *, policyengine_version: str + self, read: Callable[[], Inputs], *, request: EconomyRequest ) -> Inputs: """Run ``read`` on a thread; answer each refusal as the request's error.""" try: @@ -445,16 +453,31 @@ async def _read_inputs[Inputs]( ) raise _database_unavailable() from None except Exception as error: - refusal = read_refusal(error, policyengine_version=policyengine_version) + refusal = read_refusal(error, request) if refusal is None: raise raise refusal from None - async def _create(self, create: Callable[[], None]) -> None: + async def _create(self, record: RecordName, create: Callable[[], None]) -> 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. + """ try: await asyncio.to_thread(create) except EconomyRecordReferenceError as error: raise reference_refusal(error) from None + except EconomyRecordOutcomeUnknownError as error: + logger.warning( + "economy_records_outcome_unknown", + extra={ + "record": record, + "error_type": type(error.__cause__).__name__, + }, + ) + raise EconomySubmissionOutcomeUnknownError(record) from None except SQLAlchemyError as error: logger.warning( "economy_records_not_created", @@ -648,13 +671,17 @@ def _scope( """ try: dataset = bundle_dataset(worker, catalog.dataset.name) - except DatasetNotInBundleError as error: + except DatasetNotInBundleError: code: EconomyRefusalCode = ( "dataset_not_in_bundle" if request.dataset_name is None else "unknown_dataset" ) - raise EconomySubmissionRefusedError(code=code, detail=str(error)) from None + detail = ( + f"The default v2 worker, .py {worker.version}, has no dataset " + f"{catalog.dataset.name}." + ) + raise EconomySubmissionRefusedError(code=code, detail=detail) from None return CalculationScope( country=request.country, region_code=catalog.region.code, diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py index 1d725e501..82c859dd4 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -167,7 +167,8 @@ def test_a_report_answers_201_with_its_records_and_runs(client_for) -> None: "include_cliffs": True, "spm": {"geography_kind": "national"}, } - assert request_id == OBSERVABILITY_ID + # The call's request ID is the request's own (see the request ID test). + assert request_id != OBSERVABILITY_ID assert context["observability_id"] == OBSERVABILITY_ID @@ -185,7 +186,7 @@ def test_a_simulation_answers_201_with_its_record_and_run(client_for) -> None: } ((request, request_id, _),) = submissions.simulations assert str(request.policy_id) == SIMULATION_BODY["policy_id"] - assert request_id == result.headers[OBSERVABILITY_ID_HEADER] + assert request_id != result.headers[OBSERVABILITY_ID_HEADER] @pytest.mark.parametrize( diff --git a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py index f2435134d..aff3b7786 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py @@ -211,13 +211,13 @@ def test_a_report_without_spm_settings_stores_none() -> None: ("error", "code"), [ ( - UnknownCatalogVersionError("us has no catalog version"), + UnknownCatalogVersionError("store text: no catalog version"), "unknown_catalog_version", ), - (UnknownRegionError("no region state/zz"), "unknown_region"), - (UnknownDatasetError("no dataset acs"), "unknown_dataset"), - (UnknownPolicyError("us has no policy"), "unknown_policy"), - (UnknownDynamicError("there is no dynamic"), "unknown_dynamic"), + (UnknownRegionError("store text: no region"), "unknown_region"), + (UnknownDatasetError("store text: no dataset"), "unknown_dataset"), + (UnknownPolicyError("store text: no policy"), "unknown_policy"), + (UnknownDynamicError("store text: no dynamic"), "unknown_dynamic"), ], ) def test_a_missing_row_is_refused_before_anything_is_written(error, code) -> None: From fd9b624027c499910c31847666ef5409b9a94757 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 20:37:11 +0400 Subject: [PATCH 17/26] Test review round 2's fix: a request ID the call can't carry is replaced (failing) A request whose X-Request-ID is longer than the call's context takes, or holds characters outside the observability runtime's, gets a new UUID as the call's request ID instead of failing with 500. An unknown commit is logged once, with its error type, and the staging check's submission gets 60 seconds, its other requests 30. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_routes.py | 31 ++++++++++++++++--- .../tests/test_economy_submissions.py | 9 ++++-- .../tests/test_staging_economy_check.py | 12 +++++-- 3 files changed, 43 insertions(+), 9 deletions(-) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py index 82c859dd4..bc4564f7f 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -547,7 +547,9 @@ def test_a_create_whose_commit_is_unconfirmed_answers_a_500( ) -> None: """A 500 tells the Public API the record may exist; no 503 may say it doesn't.""" submissions = FakeEconomySubmissions( - failure=EconomySubmissionOutcomeUnknownError(record) + failure=EconomySubmissionOutcomeUnknownError( + record, error_type="OperationalError" + ) ) with client_for(submissions) as client: @@ -561,7 +563,28 @@ def test_a_create_whose_commit_is_unconfirmed_answers_a_500( "it was written." ), } - (logged,) = _emitted_records( - capsys.readouterr().out, "economy_record_outcome_unknown" - ) + printed = capsys.readouterr().out + (logged,) = _emitted_records(printed, "economy_record_outcome_unknown") assert logged["attributes"]["record"] == record + assert logged["attributes"]["error_type"] == "OperationalError" + # One event for one outcome. + assert _emitted_records(printed, "economy_records_outcome_unknown") == [] + + +@pytest.mark.parametrize( + "request_id", + ["r" * 256, "has spaces in it", ""], + ids=["too-long", "spaces", "empty"], +) +def test_a_request_id_the_call_cannot_carry_is_replaced(client_for, request_id) -> None: + """The call's context takes 1 to 128 of the observability runtime's characters.""" + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post( + "/reports", json=REPORT_BODY, headers={REQUEST_ID_HEADER: request_id} + ) + + assert result.status_code == 201 + ((_, call_request_id, _),) = submissions.reports + assert str(UUID(call_request_id)) == call_request_id diff --git a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py index aff3b7786..062fa9b7c 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py @@ -667,12 +667,17 @@ def test_a_dataset_the_bundle_lacks_is_named_in_the_services_words() -> None: def test_a_create_whose_commit_is_unconfirmed_may_have_written_the_record( record, ) -> None: - store = FakeStore(create_failure=EconomyRecordOutcomeUnknownError("lost")) + lost_commit = EconomyRecordOutcomeUnknownError("lost") + lost_commit.__cause__ = DATABASE_ERROR + store = FakeStore(create_failure=lost_commit) caller = FakeCaller() submit = _submit_report if record == "report" else _submit_simulation with pytest.raises(EconomySubmissionOutcomeUnknownError) as unknown: submit(service(store, caller=caller)) - assert unknown.value.record == record + assert (unknown.value.record, unknown.value.error_type) == ( + record, + "OperationalError", + ) assert (caller.report_runs, caller.simulation_runs) == ([], []) diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py index 5f62407da..e3cedaae2 100644 --- a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -256,12 +256,18 @@ def test_a_read_that_gets_no_answer_is_asked_again(fakes, tmp_path): def test_every_request_has_a_time_limit(fakes, tmp_path): + """The submission, which isn't asked again, may take Cloud Run's full minute.""" result, _ = _run(fakes, tmp_path, [f"200 {_report('succeeded')}"]) assert result.returncode == 0, result.stderr - calls = fakes.calls("curl") - assert calls - assert all("--max-time" in call for call in calls) + limits = { + call[-1]: call[call.index("--max-time") + 1] for call in fakes.calls("curl") + } + assert limits == { + "https://issuer.example/oauth/token": "30", + f"{BASE_URL}/reports": "60", + f"{BASE_URL}/reports/{REPORT_ID}": "30", + } def test_any_other_read_answer_fails_the_check(fakes, tmp_path): From edff811cb2b5adfab3fe80c37dd76dc0e06c1006 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 20:38:31 +0400 Subject: [PATCH 18/26] Fix review round 2: give the Modal call a request ID its context can carry The request middleware keeps the caller's X-Request-ID unchecked, and a Modal call's context takes at most 255 characters, so a longer header made every economy submission answer 500. The routes now pass the request's ID only when it fits the observability runtime's rule (1 to 128 of [A-Za-z0-9._:-]), and otherwise a new UUID. Also from round 2's findings: - An unknown commit is logged once, by the route, with its error type in the attributes the observability runtime keeps. - The staging check's submission, which isn't asked again, gets 60 seconds, Cloud Run's own limit; its other requests keep 30. - ReportSide and SideRefusal type a report side and its refusal, and the refusal helpers use them. - Docstrings: the created answers' status, submit_report's errors, the module docstring's caller-row codes, the persistence README's new errors and the staging script's comments. Co-Authored-By: Claude Opus 5.5 --- .../simulation-staging-economy-check.sh | 24 ++++++++------- .../README.md | 6 +++- .../__init__.py | 4 +++ .../economy_run_store.py | 8 ++++- .../economy_api_models.py | 14 ++++----- .../economy_routes.py | 29 +++++++++++++++---- .../economy_submission_errors.py | 12 ++++---- .../economy_submissions.py | 20 ++++++------- 8 files changed, 76 insertions(+), 41 deletions(-) diff --git a/.github/scripts/simulation-staging-economy-check.sh b/.github/scripts/simulation-staging-economy-check.sh index 23d092b05..ef3c2efb8 100755 --- a/.github/scripts/simulation-staging-economy-check.sh +++ b/.github/scripts/simulation-staging-economy-check.sh @@ -9,9 +9,10 @@ # Optional: ECONOMY_CHECK_DEADLINE_SECONDS (1800) and # ECONOMY_CHECK_POLL_SECONDS (15). # The access token is sent from a file and the client secret on standard -# input, so neither appears on a command line. Each request is given -# REQUEST_TIMEOUT_SECONDS; a read that gets no answer is asked again, as a -# 503 is, until the deadline. +# input, so neither appears on a command line. The submission, which isn't +# asked again, may take 60 seconds, Cloud Run's own limit; every other +# request 30. A read that gets no answer is asked again, as a 503 is, until +# the deadline. set -euo pipefail @@ -25,6 +26,7 @@ base_url="${base_url%/}" deadline_seconds="${ECONOMY_CHECK_DEADLINE_SECONDS:-1800}" poll_seconds="${ECONOMY_CHECK_POLL_SECONDS:-15}" request_timeout_seconds=30 +submission_timeout_seconds=60 uuid_pattern='^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$' for variable_name in SIMULATION_TEST_AUTH_ISSUER SIMULATION_TEST_AUTH_AUDIENCE \ @@ -40,12 +42,13 @@ work_directory="$(mktemp -d)" trap 'rm -rf "${work_directory}"' EXIT chmod 700 "${work_directory}" -# Send one request; write its body to .json and print its HTTP status. -# It fails, printing nothing, when no answer arrives. +# Send one request, allowing it ; write its body to .json and +# print its HTTP status. It fails when no answer arrives. send() { local name="$1" - shift - curl --silent --show-error --max-time "${request_timeout_seconds}" \ + local seconds="$2" + shift 2 + curl --silent --show-error --max-time "${seconds}" \ --output "${work_directory}/${name}.json" --write-out '%{http_code}' "$@" } @@ -55,7 +58,7 @@ token_status="$( --arg client_secret "${SIMULATION_TEST_AUTH_CLIENT_SECRET}" \ --arg audience "${SIMULATION_TEST_AUTH_AUDIENCE}" \ '{client_id: $client_id, client_secret: $client_secret, audience: $audience, grant_type: "client_credentials"}' \ - | send token --request POST --header 'Content-Type: application/json' \ + | send token "${request_timeout_seconds}" --request POST --header 'Content-Type: application/json' \ --data @- "${SIMULATION_TEST_AUTH_ISSUER%/}/oauth/token" )" || fail "the token request got no answer" [[ "${token_status}" == "200" ]] || fail "the token request answered ${token_status}" @@ -75,7 +78,7 @@ submission_status="$( reform: {policy_id: null, dynamic_id: null}, include_cliffs: false }' \ - | send submission --request POST --header "@${authorization_header}" \ + | send submission "${submission_timeout_seconds}" --request POST --header "@${authorization_header}" \ --header 'Content-Type: application/json' --data @- \ "${base_url}/reports" )" || fail "POST /reports got no answer; it may have created a report" @@ -94,7 +97,8 @@ started_at="${SECONDS}" while true; do report_body="${work_directory}/report.json" if ! read_status="$( - send report --header "@${authorization_header}" "${base_url}/reports/${report_id}" + send report "${request_timeout_seconds}" --header "@${authorization_header}" \ + "${base_url}/reports/${report_id}" )"; then printf 'GET /reports/%s got no answer; asking again\n' "${report_id}" >&2 elif [[ "${read_status}" == "200" ]]; then diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 4051d4862..4fea4fd37 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -119,7 +119,11 @@ model, region and dataset, each side's policy and dynamic values, and each side's reusable run together, never giving both sides the same run; `read_simulation_inputs` does the same for one simulation; the writes and reads each take their own. A missing region or dataset raises -`UnknownRegionError` or `UnknownDatasetError`. +`UnknownRegionError` or `UnknownDatasetError`, and a report side's refused +policy or dynamic raises `ReportSideError`, naming the side. A create whose +commit fails after its rows were inserted raises +`EconomyRecordOutcomeUnknownError`: PostgreSQL may have applied the commit, +so the rows may exist. ## Reading a policy 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 1183bedf9..1741cf7b0 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -44,8 +44,10 @@ EconomyRequest, EconomyRunStore, EconomySimulationInputs, + ReportSide, ReportSideError, RunReuse, + SideRefusal, SimulationInputs, UnknownDatasetError, UnknownRegionError, @@ -129,11 +131,13 @@ "ReportRunState", "ReportRunStatus", "ReportRunTrigger", + "ReportSide", "ReportSideError", "ReportSimulation", "ReusableRun", "RunReuse", "RunVersions", + "SideRefusal", "SimulationChoice", "SimulationInputs", "SimulationRunState", 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 26874604d..ee20f6b51 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 @@ -79,6 +79,12 @@ # Which of a report's simulations a value belongs to. type ReportSide = Literal["baseline", "reform"] # The errors a simulation's own policy or dynamic can raise. +type SideRefusal = ( + UnknownPolicyError + | UnknownDynamicError + | MissingPolicyParametersError + | OverlappingReformError +) _SIDE_ERRORS = ( UnknownPolicyError, UnknownDynamicError, @@ -103,7 +109,7 @@ class ReportSideError(Exception): OverlappingReformError. """ - def __init__(self, side: ReportSide, error: Exception) -> None: + def __init__(self, side: ReportSide, error: SideRefusal) -> None: """Keep the side and the refusal.""" super().__init__(f"the report's {side}: {error}") self.side: ReportSide = side 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 a4df869bd..b161c91e7 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 @@ -10,7 +10,7 @@ Answers the Public API reads: - 201 with the created record's ID and its first run's ID; -- 400 with ``code`` and ``detail``: one of ``CALLER_ROW_REFUSAL_CODES`` when +- 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; @@ -169,9 +169,9 @@ class EconomySimulationCreated(_Answer): """The simulation a request created, its first run and that run's status. The status is the run's when its call's ID was recorded: ``pending``, or - ``running`` when the simulation function had claimed it first. A run - whose call Modal didn't confirm, or whose call's ID couldn't be - recorded, is answered ``pending``. + whatever the simulation function had made it by then if it claimed the + run first. A run whose call Modal didn't confirm, or whose call's ID + couldn't be recorded, is answered ``pending``. """ simulation_id: UUID @@ -183,9 +183,9 @@ class EconomyReportCreated(_Answer): """The report a request created, its first run and that run's status. The status is the run's when its call's ID was recorded: ``pending``, or - ``running`` when the coordinator had claimed it first. A run whose call - Modal didn't confirm, or whose call's ID couldn't be recorded, is - answered ``pending``. Each side names its + whatever the coordinator had made it by then if it claimed the run + first. A run whose call Modal didn't confirm, or whose call's ID couldn't + be recorded, is answered ``pending``. Each side names its simulation and simulation run; ``…_reused`` is true for a side that names a run that already succeeded, rather than a new one. """ 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 b8d41ec89..235b25a71 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 @@ -22,14 +22,16 @@ A request gets its observability ID from the caller's header, or a new one, and passes it to the Modal call it starts in the observability context; the -call's ``request_id`` is the request's own ``X-Request-ID``. +call's ``request_id`` is the request's own ``X-Request-ID``, or a new UUID +when that doesn't fit the call's context. """ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Protocol -from uuid import UUID +import re +from typing import TYPE_CHECKING, Any, Final, Protocol +from uuid import UUID, uuid4 from fastapi import APIRouter, Depends, Request, status from fastapi.responses import JSONResponse @@ -75,6 +77,10 @@ logger = logging.getLogger(__name__) +# The request IDs the observability runtime accepts, which a Modal call's +# context also carries (``Stage12InvocationContext.request_id`` takes at most +# 255 characters). Any other request ID is replaced by a new one for the call. +CALL_REQUEST_ID: Final = re.compile(r"[A-Za-z0-9._:-]{1,128}") REPORTS_PATH = "/reports" REPORT_PATH = "/reports/{report_id}" SIMULATIONS_PATH = "/simulations" @@ -239,6 +245,17 @@ def _submission_failure_response(error: Exception, record: RecordName) -> JSONRe ) +def _call_request_id(request: Request) -> str: + """Return the request's ID for the Modal call, or a new one that fits it. + + The request middleware keeps the caller's ``X-Request-ID`` unchecked. + """ + request_id = getattr(request.state, "request_id", "") + if isinstance(request_id, str) and CALL_REQUEST_ID.fullmatch(request_id): + return request_id + return str(uuid4()) + + def _outcome_unknown_response( error: EconomySubmissionOutcomeUnknownError, ) -> JSONResponse: @@ -248,7 +265,7 @@ def _outcome_unknown_response( it nothing was created. The log is the fixed event ``economy_record_outcome_unknown``. """ - fields = {"record": error.record} + fields = {"record": error.record, "error_type": error.error_type} logger.error( "economy_record_outcome_unknown", extra={**fields, "policyengine_attributes": fields}, @@ -402,7 +419,7 @@ async def create_report( try: created = await submissions.submit_report( body, - request_id=request.state.request_id, + request_id=_call_request_id(request), observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: @@ -446,7 +463,7 @@ async def create_simulation( try: created = await submissions.submit_simulation( body, - request_id=request.state.request_id, + request_id=_call_request_id(request), observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as 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 0a1bda43d..be55aed4b 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 @@ -31,6 +31,7 @@ from policyengine_simulation_persistence import ( EconomyRecordReferenceError, EconomyRequest, + ReportSide, ) from policyengine_simulation_entry.economy_api_models import ( @@ -100,10 +101,11 @@ class EconomySubmissionOutcomeUnknownError(Exception): which it reads as nothing created. """ - def __init__(self, record: RecordName) -> None: - """Keep what the submission creates.""" + def __init__(self, record: RecordName, *, error_type: str) -> 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 class EconomySubmissionUnavailableError(Exception): @@ -138,7 +140,7 @@ def read_refusal( (ReportSideError) is located on that side's fields. Every detail is written from ``request``, never from the error's text. """ - side = None + side: ReportSide | None = None if isinstance(error, ReportSideError): side, error = error.side, error.error return _missing_row_refusal(error, request, side) or _unprocessable( @@ -147,7 +149,7 @@ def read_refusal( def _missing_row_refusal( - error: Exception, request: EconomyRequest, side: str | None + error: Exception, request: EconomyRequest, side: ReportSide | None ) -> EconomySubmissionRefusedError | None: """Return the 400 a missing row means, or None for another error.""" catalog = ( @@ -187,7 +189,7 @@ def _missing_row_refusal( def _unprocessable( - error: Exception, request: EconomyRequest, side: str | None + error: Exception, request: EconomyRequest, side: ReportSide | None ) -> EconomySubmissionUnprocessableError | None: """Return the 422 refused values mean, or None for another error. 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 6a1c549e5..3247e5ecb 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 @@ -247,10 +247,12 @@ async def submit_report( """Create the report and its runs and start its coordinator. Raise EconomySubmissionRefusedError (400) or - EconomySubmissionUnprocessableError (422) before writing anything, - and EconomySubmissionUnavailableError (503) when the manifest or + EconomySubmissionUnprocessableError (422) before writing anything; + EconomySubmissionUnavailableError (503) when the manifest or database cannot be read, or, naming the report, when Modal doesn't - confirm the call and the report run was failed. + confirm the call and the report run was failed; and + EconomySubmissionOutcomeUnknownError (500) when the create's commit + failed after its rows were inserted. """ worker = await self._default_worker(request.country) options = request.options() @@ -470,14 +472,10 @@ async def _create(self, record: RecordName, create: Callable[[], None]) -> None: except EconomyRecordReferenceError as error: raise reference_refusal(error) from None except EconomyRecordOutcomeUnknownError as error: - logger.warning( - "economy_records_outcome_unknown", - extra={ - "record": record, - "error_type": type(error.__cause__).__name__, - }, - ) - raise EconomySubmissionOutcomeUnknownError(record) from None + # The route logs it, with the commit error's type. + raise EconomySubmissionOutcomeUnknownError( + record, error_type=type(error.__cause__).__name__ + ) from None except SQLAlchemyError as error: logger.warning( "economy_records_not_created", From dbfa8ed8104481dc9b1f1969227dcd095e5c3907 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 20:51:38 +0400 Subject: [PATCH 19/26] Test review round 3's quality findings (failing) - The call keeps a request ID of 128 characters and replaces one of 129, on both POST routes, and a request whose state holds no usable ID gets a new one (call_request_id_for). - Every error SideRefusal names is raised from a report read with its side. This passes before the change too: the change derives the caught errors from that type, so the two can't drift apart. - The staging check asks again after a 429 or any 5xx Cloud Run answers in the service's place, not only a 503. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_run_store.py | 42 ++++++++++++++++++- .../tests/test_economy_routes.py | 36 ++++++++++++---- .../tests/test_staging_economy_check.py | 10 +++++ 3 files changed, 79 insertions(+), 9 deletions(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py index 2a40fe4da..7ba3e04c8 100644 --- a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py @@ -11,7 +11,7 @@ from contextlib import contextmanager from dataclasses import replace from datetime import UTC, datetime -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, get_args from uuid import uuid4 import pytest @@ -26,6 +26,9 @@ insert_reusable_run, ) from policyengine_simulation_persistence import ( + CatalogDataset, + CatalogRegion, + EconomyCatalogRows, EconomyPopulation, EconomyRecordOutcomeUnknownError, EconomyRequest, @@ -39,6 +42,7 @@ ReportSimulation, ReusableRun, RunReuse, + SideRefusal, SimulationChoice, SimulationStatus, UnknownCatalogVersionError, @@ -46,6 +50,7 @@ UnknownDynamicError, UnknownPolicyError, UnknownRegionError, + economy_run_store, ) from public_api_rows import StoredValue, insert_dataset, insert_dynamic, insert_policy from sqlalchemy import text @@ -506,3 +511,38 @@ def test_a_create_whose_statement_fails_is_not_written() -> None: with pytest.raises(OperationalError): store.create_report(_new_report()) + + +def test_every_side_refusal_is_raised_with_its_side(monkeypatch) -> None: + """Every error the side refusal type names is wrapped; none slips through.""" + catalog = EconomyCatalogRows( + catalog_version_id=uuid4(), + model_id=uuid4(), + region=CatalogRegion(id=uuid4(), code="us", default_dataset_id=uuid4()), + dataset=CatalogDataset(id=uuid4(), name="populace_us_2024"), + ) + monkeypatch.setattr(economy_run_store, "_catalog_rows", lambda *_: catalog) + refusals = { + UnknownPolicyError: UnknownPolicyError("no policy"), + UnknownDynamicError: UnknownDynamicError("no dynamic"), + MissingPolicyParametersError: MissingPolicyParametersError(["gov.a"]), + OverlappingReformError: OverlappingReformError(["gov.a"]), + } + assert set(get_args(SideRefusal.__value__)) == set(refusals) + + @contextmanager + def no_database(): + yield None + + store = EconomyRunStore(no_database) + for refusal in refusals.values(): + + def refuse(*_, refusal=refusal): + raise refusal + + monkeypatch.setattr(economy_run_store, "_simulation_inputs", refuse) + with pytest.raises(ReportSideError) as raised: + store.read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=CURRENT_LAW, reuse=None + ) + assert (raised.value.side, raised.value.error) == ("baseline", refusal) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py index bc4564f7f..90ca6e108 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -5,6 +5,7 @@ import json from dataclasses import dataclass, field from datetime import datetime +from types import SimpleNamespace from typing import Any, NoReturn from uuid import UUID @@ -21,6 +22,7 @@ EconomySimulationCreated, EconomySimulationRequest, ) +from policyengine_simulation_entry.economy_routes import call_request_id_for from policyengine_simulation_entry.economy_submission_errors import ( EconomySubmissionOutcomeUnknownError, EconomySubmissionRefusedError, @@ -571,20 +573,38 @@ def test_a_create_whose_commit_is_unconfirmed_answers_a_500( assert _emitted_records(printed, "economy_records_outcome_unknown") == [] +@pytest.mark.parametrize(("path", "body"), ROUTES) @pytest.mark.parametrize( - "request_id", - ["r" * 256, "has spaces in it", ""], - ids=["too-long", "spaces", "empty"], + ("request_id", "kept"), + [ + ("r" * 128, True), + ("r" * 129, False), + ("r" * 256, False), + ("has spaces in it", False), + ], + ids=["128-kept", "129-replaced", "256-replaced", "spaces-replaced"], ) -def test_a_request_id_the_call_cannot_carry_is_replaced(client_for, request_id) -> None: +def test_a_request_id_the_call_cannot_carry_is_replaced( + client_for, path, body, request_id, kept +) -> None: """The call's context takes 1 to 128 of the observability runtime's characters.""" submissions = FakeEconomySubmissions() with client_for(submissions) as client: - result = client.post( - "/reports", json=REPORT_BODY, headers={REQUEST_ID_HEADER: request_id} - ) + result = client.post(path, json=body, headers={REQUEST_ID_HEADER: request_id}) assert result.status_code == 201 - ((_, call_request_id, _),) = submissions.reports + ((_, call_request_id, _),) = submissions.reports or submissions.simulations + if kept: + assert call_request_id == request_id + else: + assert str(UUID(call_request_id)) == call_request_id + + +@pytest.mark.parametrize("state", [{"request_id": ""}, {"request_id": None}, {}]) +def test_a_request_without_a_usable_id_gets_a_new_one(state) -> None: + request = SimpleNamespace(state=SimpleNamespace(**state)) + + call_request_id = call_request_id_for(request) + assert str(UUID(call_request_id)) == call_request_id diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py index e3cedaae2..c2e3bda45 100644 --- a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -236,6 +236,16 @@ def test_a_refused_submission_fails_the_check_with_its_answer( assert not (state / "polls").exists() +@pytest.mark.parametrize("status", ["429", "500", "502", "504"]) +def test_a_read_cloud_run_answers_in_the_services_place_is_asked_again( + fakes, tmp_path, status +): + result, _ = _run(fakes, tmp_path, [f"{status} {{}}", f"200 {_report('succeeded')}"]) + + assert result.returncode == 0, result.stderr + assert status in result.stderr + + def test_an_unavailable_read_is_asked_again(fakes, tmp_path): unavailable = json.dumps({"code": "database_unavailable", "detail": "x"}) From 70203a9886832925b88da4cd96f12a5f09e9574b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 20:52:33 +0400 Subject: [PATCH 20/26] Fix review round 3's quality findings - call_request_id_for, public now, is tested at the 128-character limit on both POST routes and with no usable state value. Its docstrings, and the routes' module docstring, name the real header, X-PolicyEngine-Request-Id, say a replaced ID matches none of the request's log records (the observability ID still joins them), and that a failure before a handler runs gets the middleware's plain-text 500. - The errors a report read wraps with their side are taken from SideRefusal itself, so the type and the caught errors can't name different errors. - The staging check asks again after a 429 or any 5xx, which Cloud Run's front end may answer in the service's place, as it does after a 503. Co-Authored-By: Claude Opus 5.5 --- .../simulation-staging-economy-check.sh | 9 ++++++--- .../economy_run_store.py | 10 +++------- .../economy_routes.py | 19 ++++++++++++------- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/.github/scripts/simulation-staging-economy-check.sh b/.github/scripts/simulation-staging-economy-check.sh index ef3c2efb8..d146457af 100755 --- a/.github/scripts/simulation-staging-economy-check.sh +++ b/.github/scripts/simulation-staging-economy-check.sh @@ -11,8 +11,8 @@ # The access token is sent from a file and the client secret on standard # input, so neither appears on a command line. The submission, which isn't # asked again, may take 60 seconds, Cloud Run's own limit; every other -# request 30. A read that gets no answer is asked again, as a 503 is, until -# the deadline. +# request 30. A read that gets no answer, a 429 or any 5xx is asked again +# until the deadline. set -euo pipefail @@ -119,7 +119,10 @@ while true; do pending | running) ;; *) fail "report run ${report_run_id} has an unknown status: ${run_status}" ;; esac - elif [[ "${read_status}" != "503" ]]; then + elif [[ "${read_status}" == "429" || "${read_status}" == 5?? ]]; then + # The service, or Cloud Run's front end in its place, can't answer now. + printf 'GET /reports/%s answered %s; asking again\n' "${report_id}" "${read_status}" >&2 + else fail "GET /reports/${report_id} answered ${read_status}: $(cat "${report_body}")" fi ((SECONDS - started_at < deadline_seconds)) \ 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 ee20f6b51..6e66b2a50 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 @@ -10,7 +10,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, get_args from sqlalchemy.exc import SQLAlchemyError @@ -85,12 +85,8 @@ | MissingPolicyParametersError | OverlappingReformError ) -_SIDE_ERRORS = ( - UnknownPolicyError, - UnknownDynamicError, - MissingPolicyParametersError, - OverlappingReformError, -) +# Caught by the type itself, so the two never name different errors. +_SIDE_ERRORS: tuple[type[SideRefusal], ...] = get_args(SideRefusal.__value__) class UnknownRegionError(LookupError): 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 235b25a71..b93333bf2 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 @@ -18,12 +18,14 @@ with ``Retry-After`` when it is not configured or a dependency cannot be read, before writing anything, or when Modal refused the run's call: then the record was created and its run failed, and the answer names it. Every -500 is the JSON error with code ``internal_error``, and may follow a write. +500 the routes answer is the JSON error with code ``internal_error``, and may +follow a write; a failure before a route's handler runs gets the request +middleware's own plain-text 500, which the Public API reads alike. A request gets its observability ID from the caller's header, or a new one, and passes it to the Modal call it starts in the observability context; the -call's ``request_id`` is the request's own ``X-Request-ID``, or a new UUID -when that doesn't fit the call's context. +call's ``request_id`` is the request's own ``X-PolicyEngine-Request-Id``, +or a new UUID when that doesn't fit the call's context. """ from __future__ import annotations @@ -245,10 +247,13 @@ def _submission_failure_response(error: Exception, record: RecordName) -> JSONRe ) -def _call_request_id(request: Request) -> str: +def call_request_id_for(request: Request) -> str: """Return the request's ID for the Modal call, or a new one that fits it. - The request middleware keeps the caller's ``X-Request-ID`` unchecked. + The request middleware keeps the caller's ``X-PolicyEngine-Request-Id`` + unchecked. A replaced ID matches none of the request's own log records; + the observability ID, which the call's observability context carries, + still joins them. """ request_id = getattr(request.state, "request_id", "") if isinstance(request_id, str) and CALL_REQUEST_ID.fullmatch(request_id): @@ -419,7 +424,7 @@ async def create_report( try: created = await submissions.submit_report( body, - request_id=_call_request_id(request), + request_id=call_request_id_for(request), observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: @@ -463,7 +468,7 @@ async def create_simulation( try: created = await submissions.submit_simulation( body, - request_id=_call_request_id(request), + request_id=call_request_id_for(request), observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: From 92778f161400c7a10fe0945b96ff688485091ec9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:02:36 +0400 Subject: [PATCH 21/26] Test the standalone context through the contract's field, and the deadline message (failing) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit C1 · 8's contract now has Stage12InvocationContext.simulation_run_id: a standalone run's context validates with it set, and a report run's with it unset. These pass before the change too, which only builds the context through the model instead of adding the key to its JSON. The staging check's deadline failure says what the last read found: the run's status, the HTTP status, or no answer. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_economy_calculation_inputs.py | 8 +++++--- .../tests/test_staging_economy_check.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py b/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py index f14f28fc7..722ee115c 100644 --- a/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py +++ b/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py @@ -17,7 +17,6 @@ Stage12InvocationContext, ) from policyengine_simulation_entry.economy_calculation_inputs import ( - SIMULATION_RUN_ID_FIELD, CalculationScope, CallOrigin, DatasetNotInBundleError, @@ -115,6 +114,7 @@ def test_a_report_run_is_named_by_its_run_ids_with_no_evaluation_record() -> Non context = Stage12InvocationContext.model_validate(calculation.context) assert context.report_run == RUN_IDS + assert context.simulation_run_id is None assert context.production_function_call_id is None assert (context.environment, context.modal_environment, context.request_id) == ( "staging", @@ -145,8 +145,10 @@ def test_a_standalone_run_is_named_in_its_functions_context() -> None: assert simulation.simulation_execution_id == BASELINE_RUN_ID assert simulation.evaluation_id == BASELINE_RUN_ID assert simulation.policy == REFORM - assert calculation.context[SIMULATION_RUN_ID_FIELD] == str(BASELINE_RUN_ID) - assert calculation.context["report_run"] is None + context = Stage12InvocationContext.model_validate(calculation.context) + assert context.simulation_run_id == BASELINE_RUN_ID + assert context.report_run is None + assert context.records_runs assert calculation.context["artifact_prefix"] == ( f"stage-12-runs/staging/2026/11/{BASELINE_RUN_ID}" ) diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py index c2e3bda45..3ebc0dd6a 100644 --- a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -287,6 +287,23 @@ def test_any_other_read_answer_fails_the_check(fakes, tmp_path): assert "404" in result.stderr +@pytest.mark.parametrize( + ("answer", "last_read"), + [ + (f"200 {_report('running')}", "the run was running"), + ('500 {"code": "internal_error", "detail": "x"}', "it answered 500"), + ("TRANSPORT", "it got no answer"), + ], +) +def test_the_deadline_failure_says_what_the_last_read_found( + fakes, tmp_path, answer, last_read +): + result, _ = _run(fakes, tmp_path, [answer], ECONOMY_CHECK_DEADLINE_SECONDS="0") + + assert result.returncode == 1 + assert f"last read: {last_read}" in result.stderr + + def test_a_run_that_does_not_finish_in_time_fails_the_check(fakes, tmp_path): result, _ = _run( fakes, From 628cd3c1ef994ff18f8f2851d2e2cf009cd33f65 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:03:28 +0400 Subject: [PATCH 22/26] =?UTF-8?q?Build=20the=20standalone=20context=20thro?= =?UTF-8?q?ugh=20C1=20=C2=B7=208's=20field,=20and=20name=20the=20last=20re?= =?UTF-8?q?ad=20at=20the=20deadline?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - With C1 · 8 merged, a standalone simulation run's context sets Stage12InvocationContext.simulation_run_id through the model, so the contract validates it (a call names a report run or a simulation run, not both), instead of adding the key to the dumped JSON. - The staging check's deadline failure ends with what the last read found: the run's status, the status a read answered, or no answer. - The routes' module docstring says a failure outside a handler's try gets the middleware's plain-text 500 (review round 4). Co-Authored-By: Claude Opus 5.5 --- .../scripts/simulation-staging-economy-check.sh | 7 ++++++- .../economy_calculation_inputs.py | 17 ++++++++++------- .../economy_routes.py | 2 +- 3 files changed, 17 insertions(+), 9 deletions(-) diff --git a/.github/scripts/simulation-staging-economy-check.sh b/.github/scripts/simulation-staging-economy-check.sh index d146457af..ae24e8673 100755 --- a/.github/scripts/simulation-staging-economy-check.sh +++ b/.github/scripts/simulation-staging-economy-check.sh @@ -94,18 +94,22 @@ printf 'Created report %s, run %s (baseline reused: %s, reform reused: %s)\n' \ "$(jq -r '.baseline_reused' "${submission}")" "$(jq -r '.reform_reused' "${submission}")" started_at="${SECONDS}" +# What the last read found, for the deadline's failure message. +last_read="nothing yet" while true; do report_body="${work_directory}/report.json" if ! read_status="$( send report "${request_timeout_seconds}" --header "@${authorization_header}" \ "${base_url}/reports/${report_id}" )"; then + last_read="it got no answer" printf 'GET /reports/%s got no answer; asking again\n' "${report_id}" >&2 elif [[ "${read_status}" == "200" ]]; then newest_run_id="$(jq -r '.run.run_id // empty' "${report_body}")" [[ "${newest_run_id}" == "${report_run_id}" ]] \ || fail "report ${report_id}'s newest run is ${newest_run_id:-missing}, not the created run ${report_run_id}" run_status="$(jq -r '.run.status' "${report_body}")" + last_read="the run was ${run_status}" case "${run_status}" in succeeded) printf 'Report run %s succeeded on .py %s with %s\n' "${report_run_id}" \ @@ -121,11 +125,12 @@ while true; do esac elif [[ "${read_status}" == "429" || "${read_status}" == 5?? ]]; then # The service, or Cloud Run's front end in its place, can't answer now. + last_read="it answered ${read_status}" printf 'GET /reports/%s answered %s; asking again\n' "${report_id}" "${read_status}" >&2 else fail "GET /reports/${report_id} answered ${read_status}: $(cat "${report_body}")" fi ((SECONDS - started_at < deadline_seconds)) \ - || fail "report run ${report_run_id} did not finish within ${deadline_seconds} seconds" + || fail "report run ${report_run_id} did not finish within ${deadline_seconds} seconds (last read: ${last_read})" sleep "${poll_seconds}" done diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py index 5154e57c1..bc9075388 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py @@ -51,8 +51,6 @@ # How long the coordinator keeps a call's Stage 12 artifacts; the longest the # contract allows. Simulation output files go to the output bucket instead. ARTIFACT_RETENTION: Final = timedelta(days=30) -# The context field naming a standalone simulation run's own run. -SIMULATION_RUN_ID_FIELD: Final = "simulation_run_id" class DatasetNotInBundleError(LookupError): @@ -170,6 +168,7 @@ def _context( *, run_id: UUID, report_run: ReportRunIds | None, + simulation_run_id: UUID | None = None, ) -> Stage12InvocationContext: created_at = origin.created_at return Stage12InvocationContext( @@ -188,6 +187,7 @@ def _context( created_at=created_at, retention_expires_at=created_at + ARTIFACT_RETENTION, report_run=report_run, + simulation_run_id=simulation_run_id, ) @@ -249,11 +249,14 @@ def simulation_run_calculation( role=SimulationRole.STANDALONE, policy=policy, ) - context = _context(worker, origin, run_id=simulation_run_id, report_run=None) + context = _context( + worker, + origin, + run_id=simulation_run_id, + report_run=None, + simulation_run_id=simulation_run_id, + ) return RunCalculation( payload=simulation.model_dump(mode="json"), - context={ - **context.model_dump(mode="json"), - SIMULATION_RUN_ID_FIELD: str(simulation_run_id), - }, + context=context.model_dump(mode="json"), ) 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 b93333bf2..9dd6e5285 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 @@ -19,7 +19,7 @@ read, before writing anything, or when Modal refused the run's call: then the record was created and its run failed, and the answer names it. Every 500 the routes answer is the JSON error with code ``internal_error``, and may -follow a write; a failure before a route's handler runs gets the request +follow a write; a failure outside a handler's ``try`` gets the request middleware's own plain-text 500, which the Public API reads alike. A request gets its observability ID from the caller's header, or a new one, From 31c4c7f8cb30a8c65e9d2b7b515479d895bc0e41 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:12:32 +0400 Subject: [PATCH 23/26] Fix CI: list the new input tables, and check a read run is present before its status - test_tables' list of the read-only input tables gains dynamics, regions and datasets, which this branch models. - The generated-client test asserts the report's run is present before reading its status, as Pyright requires for an optional field. Co-Authored-By: Claude Opus 5.5 --- libs/policyengine-simulation-persistence/tests/test_tables.py | 3 +++ .../tests/test_economy_generated_client.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/policyengine-simulation-persistence/tests/test_tables.py b/libs/policyengine-simulation-persistence/tests/test_tables.py index b36958a86..2d7ade37f 100644 --- a/libs/policyengine-simulation-persistence/tests/test_tables.py +++ b/libs/policyengine-simulation-persistence/tests/test_tables.py @@ -193,8 +193,11 @@ def test_the_input_models_are_tables_the_service_reads() -> None: assert {table.name for table in input_metadata.tables.values()} == { "households", "policies", + "dynamics", "parameters", "parameter_values", + "regions", + "datasets", "tax_benefit_models", "tax_benefit_model_versions", } 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 186929556..5f32692db 100644 --- a/projects/policyengine-apis-integ/tests/test_economy_generated_client.py +++ b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py @@ -189,6 +189,8 @@ def test_a_report_and_a_simulation_are_read_with_their_newest_runs(): ) assert isinstance(report_response.parsed, EconomyReportState) - assert report_response.parsed.run.status is ReportRunStatus.SUCCEEDED + run = report_response.parsed.run + assert run is not None + assert run.status is ReportRunStatus.SUCCEEDED assert isinstance(simulation_response.parsed, EconomySimulationState) assert simulation_response.parsed.to_dict() == simulation From 0dc29a2605340280657487e7c0b0fa4dce4736a6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:14:44 +0400 Subject: [PATCH 24/26] Test that the deadline failure names the latest read, not the first A running run, then a front end answering 503 until a 2-second deadline: the failure names the 503. It passes before and after; review round 5 found nothing tested that a later read replaces an earlier one in the message. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_staging_economy_check.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py index 3ebc0dd6a..adf951855 100644 --- a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -304,6 +304,21 @@ def test_the_deadline_failure_says_what_the_last_read_found( assert f"last read: {last_read}" in result.stderr +def test_the_deadline_failure_names_the_latest_read_not_the_first(fakes, tmp_path): + """A running run, then a front end answering 503 until the deadline. + + Bash counts the deadline in whole seconds, so a 2-second deadline can't be + reached on the first read, whenever its second starts. + """ + unavailable = '503 {"code": "database_unavailable", "detail": "x"}' + answers = [f"200 {_report('running')}", *([unavailable] * 5000)] + + result, _ = _run(fakes, tmp_path, answers, ECONOMY_CHECK_DEADLINE_SECONDS="2") + + assert result.returncode == 1 + assert "(last read: it answered 503)" in result.stderr + + def test_a_run_that_does_not_finish_in_time_fails_the_check(fakes, tmp_path): result, _ = _run( fakes, From 695b88feab9deaa4a6d9f8fc45720238ac1370b3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:14:44 +0400 Subject: [PATCH 25/26] Name a call's runs once: the context's prefix follows the runs it names economy_calculation_inputs._context takes one runs argument, a report run's IDs or a standalone simulation run's ID, and builds both the context's run fields and its artifact prefix from it, so a prefix naming another run can't be built (review round 5). The existing prefix tests cover both calls, before and after. Co-Authored-By: Claude Opus 5.5 --- .../economy_calculation_inputs.py | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py index bc9075388..05155cf18 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py @@ -163,13 +163,17 @@ def _simulation_input( # noqa: PLR0913 # the run, its role, its policy and wha def _context( - worker: EconomyWorker, - origin: CallOrigin, - *, - run_id: UUID, - report_run: ReportRunIds | None, - simulation_run_id: UUID | None = None, + worker: EconomyWorker, origin: CallOrigin, *, runs: ReportRunIds | UUID ) -> Stage12InvocationContext: + """Return the call's context, naming ``runs``: a report run's, or one run. + + The artifact prefix names the report run, or the standalone simulation + run, that the context names, so the two can't differ. + """ + if isinstance(runs, ReportRunIds): + report_run, simulation_run_id, run_id = runs, None, runs.report_run_id + else: + report_run, simulation_run_id, run_id = None, runs, runs created_at = origin.created_at return Stage12InvocationContext( request_id=origin.request_id, @@ -222,7 +226,7 @@ def report_run_calculation( # noqa: PLR0913 # the runs, both policies and what ), requested_aggregates=REPORT_AGGREGATES, ) - context = _context(worker, origin, run_id=report_run_id, report_run=run_ids) + context = _context(worker, origin, runs=run_ids) return RunCalculation( payload=report.model_dump(mode="json"), context=context.model_dump(mode="json"), @@ -249,13 +253,7 @@ def simulation_run_calculation( role=SimulationRole.STANDALONE, policy=policy, ) - context = _context( - worker, - origin, - run_id=simulation_run_id, - report_run=None, - simulation_run_id=simulation_run_id, - ) + context = _context(worker, origin, runs=simulation_run_id) return RunCalculation( payload=simulation.model_dump(mode="json"), context=context.model_dump(mode="json"), From 4826fa20738453c09379c92e887840cdb394280a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 12 Nov 2026 21:20:10 +0400 Subject: [PATCH 26/26] Describe the deadline test's 503s and the context's runs argument exactly Review round 6's wording notes: the queued 503s are the service's own answer, and _context's runs is a report run's IDs or a standalone simulation run's ID. Co-Authored-By: Claude Opus 5.5 --- .../economy_calculation_inputs.py | 7 ++++--- .../tests/test_staging_economy_check.py | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py index 05155cf18..f04378af8 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py @@ -165,10 +165,11 @@ def _simulation_input( # noqa: PLR0913 # the run, its role, its policy and wha def _context( worker: EconomyWorker, origin: CallOrigin, *, runs: ReportRunIds | UUID ) -> Stage12InvocationContext: - """Return the call's context, naming ``runs``: a report run's, or one run. + """Return the call's context, naming ``runs``. - The artifact prefix names the report run, or the standalone simulation - run, that the context names, so the two can't differ. + ``runs`` is a report run's IDs, for its coordinator, or a standalone + simulation run's ID, for its simulation function. The artifact prefix + names that report run or simulation run, so the two can't differ. """ if isinstance(runs, ReportRunIds): report_run, simulation_run_id, run_id = runs, None, runs.report_run_id diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py index adf951855..097ab1974 100644 --- a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -305,7 +305,7 @@ def test_the_deadline_failure_says_what_the_last_read_found( def test_the_deadline_failure_names_the_latest_read_not_the_first(fakes, tmp_path): - """A running run, then a front end answering 503 until the deadline. + """A running run, then the service answering 503 until the deadline. Bash counts the deadline in whole seconds, so a 2-second deadline can't be reached on the first read, whenever its second starts.