diff --git a/changelog.d/c1-15-household-reports.added.md b/changelog.d/c1-15-household-reports.added.md new file mode 100644 index 000000000..64974d3c7 --- /dev/null +++ b/changelog.d/c1-15-household-reports.added.md @@ -0,0 +1 @@ +The Simulation Entrypoint creates household reports through `POST /reports` with `report_type` `household`: one stored household, for its default year, under a baseline and an optional reform, each a new household simulation run. One household function call calculates the report run, writes each simulation run's `household_result` and finishes the report run. diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index c03186693..ae572dd91 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -719,11 +719,12 @@ computed because no file existed. ## v2 household functions The versioned v2 application also defines a household function for each -country, `calculate_household_us` and `calculate_household_uk`. Each call -calculates one household, in the Public API's situation shape and with +country, `calculate_household_us` and `calculate_household_uk`. A household +calculation is one household, in the Public API's situation shape and with optional axes, for one year under one policy, as a household job or a -household simulation run is one household, one year and one policy; a -baseline and its reform are two calls. It takes a `HouseholdExecutionRequest` +household simulation run is one household, one year and one policy. A call +for a household job, or for deployment validation, makes one calculation: +it takes a `HouseholdExecutionRequest` and returns a `HouseholdExecutionResult`, which holds the calculated household with its warnings and SPM receipt, or a typed `HouseholdExecutionFailure`, from @@ -739,8 +740,8 @@ year. Otherwise the calculation follows the Public API's household calculation, including its input errors, warnings and SPM receipt. A function refuses a request for another country, and returns a `bundle_mismatch` failure for a request that -names another bundle. Nothing calls these functions until the Simulation -Entrypoint has household routes. +names another bundle. The Simulation Entrypoint calls them for +`/household-calculations` and for household reports on `/reports`. A call can name a household job, a `household_jobs` row the Simulation Entrypoint inserts pending, with the `household_job_id` keyword argument. As @@ -759,9 +760,40 @@ only columns the role may update: `status`, `modal_call_id`, `error_message`, `result`, `started_at`, `completed_at` and `updated_at`. A call that names no job, as deployment validation makes, touches no database. +A call can instead name a household report run, with the +`household_report_run_id` keyword argument, and then takes a +`HouseholdReportRunCalculation` +(`policyengine_simulation_contract.household_report_execution`): the +baseline's `HouseholdExecutionRequest` and, optionally, the reform's, each +request's calculation ID naming its simulation run. A call names a job or a +report run, never both. The function then: + +1. claims the report run under its own Modal call ID, as a report run's + coordinator does, and reads which baseline and reform simulation runs it + names; +2. validates the payload and checks that its runs are those; +3. for each side, claims the side's simulation run, calculates it, and + writes its `HouseholdExecutionResult`, or its `HouseholdExecutionFailure` + and one-line summary, to the run's `household_result` and + `error_message`; +4. marks the report run succeeded only if every side's run did, or fails it + with the first such side's reason. + +Both sides are always calculated, so each run has its own outcome. A report +run another call holds, or one already finished, returns `deduplicated`, the +report run's ID and its status, without calculating. A restarted call keeps a +side it already recorded. An error, a refused payload included, or Modal's +cancellation fails the report run and each of its simulation runs left +unfinished, with a summary naming only the error's type, then propagates; a +pre-emption records nothing. An error in the report run's claim, or in the +read of its simulation runs, fails the report run if the claim applied +before the error, and leaves its simulation runs to the stuck-run sweep. + The stuck-run sweep (`src/modal/run_sweep_app.py`) fails household jobs nothing will finish by the rule it applies to report runs and simulation -runs, counting a job's age from its creation: a pending job with no call +runs (a household report run and its simulation runs are judged as any +report run and its simulation runs are), counting a job's age from its +creation: a pending job with no call recorded ten minutes after its creation, a pending or running job once Modal reports its recorded call ended, and any job still unfinished 24 hours after its creation, whatever Modal reports. @@ -771,7 +803,7 @@ after its creation, whatever Modal reports. cache. Building it needs no secret. - Secrets: only `policyengine-modal-worker-database`, whose `MODAL_WORKER_DATABASE_URL` connects as the Modal worker role, to record - the household job a call names. Like the coordinator and the simulation + the household job or household report run a call names. Like the coordinator and the simulation workers, the application cannot deploy to a Modal environment without it. A household calculation reads no private data or Cloud Storage. Telemetry authenticates with the identity token Modal gives every container diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_report_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_report_execution.py new file mode 100644 index 000000000..2f710697b --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_report_execution.py @@ -0,0 +1,72 @@ +"""The calculations of one household report run, as one household function takes them. + +A household report calculates one stored household under a baseline policy +and, optionally, a reform. The simulation service creates a household +simulation run for each side, and a report run naming them, then starts one +household function call with the report run's ID and this document. Each +side is a ``HouseholdExecutionRequest`` whose ``calculation_id`` is the +simulation run it records; both sides calculate the same household, for the +same year, the same way, on the same bundle, and differ only in their +policy. +""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import model_validator + +from policyengine_simulation_contract.household_execution import ( + HOUSEHOLD_CONTRACT_SCHEMA_VERSION, + HouseholdExecutionRequest, + StrictHouseholdModel, +) + +# Which side of a household report a calculation is. +HouseholdReportSide = Literal["baseline", "reform"] +# The request fields both sides share: everything but the policy and the +# simulation run each side records. +_SHARED_FIELDS: set[str] = set(HouseholdExecutionRequest.model_fields) - { + "calculation_id", + "policy", +} + + +class HouseholdReportRunCalculation(StrictHouseholdModel): + """A household report run's baseline and, optionally, reform calculation. + + ``reform`` is None for a report of the baseline alone. + """ + + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + baseline: HouseholdExecutionRequest + reform: HouseholdExecutionRequest | None = None + + @model_validator(mode="after") + def validate_one_household(self) -> HouseholdReportRunCalculation: + """Require two sides to record different runs of one household calculation.""" + if self.reform is None: + return self + if self.reform.calculation_id == self.baseline.calculation_id: + msg = "the baseline and the reform must record different simulation runs" + raise ValueError(msg) + baseline = self.baseline.model_dump(include=_SHARED_FIELDS) + reform = self.reform.model_dump(include=_SHARED_FIELDS) + differing = sorted( + name for name in _SHARED_FIELDS if baseline.get(name) != reform.get(name) + ) + if differing: + msg = ( + "the baseline and the reform must calculate the same household " + f"the same way; they differ in {', '.join(differing)}" + ) + raise ValueError(msg) + return self + + def sides( + self, + ) -> tuple[tuple[HouseholdReportSide, HouseholdExecutionRequest], ...]: + """Return each side's request, the baseline's first.""" + if self.reform is None: + return (("baseline", self.baseline),) + return (("baseline", self.baseline), ("reform", self.reform)) diff --git a/libs/policyengine-simulation-contract/tests/test_household_report_execution.py b/libs/policyengine-simulation-contract/tests/test_household_report_execution.py new file mode 100644 index 000000000..c0236beea --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_household_report_execution.py @@ -0,0 +1,125 @@ +"""Tests for the calculations of one household report run, as one call takes them.""" + +from __future__ import annotations + +from uuid import UUID + +import pytest +from policyengine_simulation_contract.household_execution import ( + HouseholdBundleProvenance, + HouseholdExecutionRequest, +) +from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportRunCalculation, +) +from pydantic import ValidationError + +BASELINE_RUN_ID = UUID("00000000-0000-4000-8000-0000000000b1") +REFORM_RUN_ID = UUID("00000000-0000-4000-8000-0000000000b2") +HOUSEHOLD = { + "people": {"you": {"age": {"2026": 40}, "employment_income": {"2026": 50_000}}}, + "households": { + "your household": { + "members": ["you"], + "household_net_income": {"2026": None}, + } + }, +} +BUNDLE = HouseholdBundleProvenance( + policyengine_version="6.17.1", + core_package_version="3.33.0", + country_package_name="policyengine-us", + country_package_version="1.764.6", + bundle_manifest_sha256="a" * 64, +) +CTC_REFORM = { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2200-12-31": 3_000} +} + + +def _request(calculation_id: UUID, **changes: object) -> dict[str, object]: + """Return one side's household function request as JSON.""" + return { + "calculation_id": str(calculation_id), + "country": "us", + "year": "2026", + "household": HOUSEHOLD, + "policy": {"parameters": {}}, + "add_missing_variables": True, + "bundle": BUNDLE.model_dump(mode="json"), + **changes, + } + + +def _calculation( + *, reform: dict[str, object] | None = None +) -> HouseholdReportRunCalculation: + return HouseholdReportRunCalculation.model_validate( + {"baseline": _request(BASELINE_RUN_ID), "reform": reform} + ) + + +def test_a_report_run_calculates_its_baseline_then_its_reform() -> None: + calculation = _calculation( + reform=_request(REFORM_RUN_ID, policy={"parameters": CTC_REFORM}) + ) + + sides = calculation.sides() + + assert [side for side, _ in sides] == ["baseline", "reform"] + assert [request.calculation_id for _, request in sides] == [ + BASELINE_RUN_ID, + REFORM_RUN_ID, + ] + assert all(isinstance(request, HouseholdExecutionRequest) for _, request in sides) + assert sides[1][1].policy.parameters == CTC_REFORM + + +def test_a_report_run_of_the_baseline_alone_calculates_one_side() -> None: + calculation = _calculation() + + assert calculation.reform is None + assert [side for side, _ in calculation.sides()] == ["baseline"] + + +def test_a_calculation_round_trips_through_json() -> None: + calculation = _calculation(reform=_request(REFORM_RUN_ID)) + + payload = calculation.model_dump(mode="json") + + assert payload["schema_version"] == 1 + assert HouseholdReportRunCalculation.model_validate(payload) == calculation + + +def test_each_side_is_its_own_calculation() -> None: + with pytest.raises(ValidationError, match="different simulation runs"): + _calculation(reform=_request(BASELINE_RUN_ID)) + + +@pytest.mark.parametrize( + "reform_change", + [ + {"year": "2027"}, + {"household": {"people": {"you": {"age": {"2026": 41}}}}}, + {"add_missing_variables": False}, + { + "bundle": { + **BUNDLE.model_dump(mode="json"), + "bundle_manifest_sha256": "b" * 64, + } + }, + ], + ids=["year", "household", "add_missing_variables", "bundle"], +) +def test_both_sides_calculate_one_household_the_same_way( + reform_change: dict[str, object], +) -> None: + with pytest.raises(ValidationError, match="the same household"): + _calculation(reform=_request(REFORM_RUN_ID, **reform_change)) + + +def test_unknown_fields_are_refused() -> None: + with pytest.raises(ValidationError, match="extra"): + HouseholdReportRunCalculation.model_validate( + {"baseline": _request(BASELINE_RUN_ID), "report_run_id": "x"} + ) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index d6ac6ee63..15b5b51d7 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -125,6 +125,29 @@ commit fails after its rows were inserted raises `EconomyRecordOutcomeUnknownError`: PostgreSQL may have applied the commit, so the rows may exist. +## Household reports + +`household_report_records.py` holds the simulation service's writes for a +household report: one stored household, for one year, under a baseline and, +optionally, a reform. `insert_household_report` creates, in the caller's +transaction, a household simulation (no dataset or region) and its first +run for each side, never reusing one; the report, of type `household`, +naming the household, both simulations and the reform's policy, with its +options in `inputs`; and its first report run, trigger `initial`, naming no +reform run for a report of the baseline alone. Every run is `pending` with +no Modal call ID, and a row naming one that does not exist is refused as +`EconomyRecordReferenceError`. + +`HouseholdReportStore` runs the steps in transactions of their own: +`read_inputs` reads the catalog version for the .py version, the country's +model, the stored household and each side's policy values together, +raising `UnknownCatalogVersionError`, `UnknownHouseholdError` or +`ReportSideError` for the side whose policy is missing or sets parameters +the catalog version lacks; `create_report` raises +`EconomyRecordOutcomeUnknownError` when the commit fails after the rows were +inserted; and `record_report_call` and `fail_unstarted_report` use the +economy report run writes. + ## Reading a policy `read_policy_reform` (`policy_reform.py`) reads a stored policy as the reform 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 bea874cc0..461bda49c 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -81,6 +81,16 @@ UnknownCatalogVersionError, UnknownHouseholdError, ) +from policyengine_simulation_persistence.household_report_records import ( + HOUSEHOLD_REPORT_TYPE, + HouseholdReportSimulation, + NewHouseholdReport, + insert_household_report, +) +from policyengine_simulation_persistence.household_report_store import ( + HouseholdReportInputs, + HouseholdReportStore, +) from policyengine_simulation_persistence.input_tables import input_metadata from policyengine_simulation_persistence.policy_reform import ( MissingPolicyParametersError, @@ -111,6 +121,7 @@ ) __all__ = [ + "HOUSEHOLD_REPORT_TYPE", "LATEST_YEAR", "MAX_BUDGET_WINDOW_YEARS", "SERVICE_TABLE_PRIVILEGES", @@ -134,11 +145,15 @@ "HouseholdJobReferenceError", "HouseholdJobStatus", "HouseholdJobStore", + "HouseholdReportInputs", + "HouseholdReportSimulation", + "HouseholdReportStore", "MissingPolicyParametersError", "NewBudgetWindowReport", "NewEconomyReport", "NewEconomySimulation", "NewHouseholdJob", + "NewHouseholdReport", "OverlappingReformError", "PolicyReform", "ReportRunState", @@ -177,6 +192,7 @@ "insert_economy_report", "insert_economy_simulation", "insert_household_job", + "insert_household_report", "missing_table_privileges", "normalized_database_url", "read_catalog_version_id", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/household_report_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/household_report_records.py new file mode 100644 index 000000000..0ed192df0 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/household_report_records.py @@ -0,0 +1,181 @@ +"""The simulation service's writes for household reports. + +A household report calculates one stored household, for one year, under a +baseline policy and, optionally, a reform. The service creates in one +transaction: + +- for each side, a new household ``simulations`` row naming the household + and the side's policy, and its first ``simulation_runs`` row, ``pending`` + with no Modal call ID. Household runs are never reused: each report gets + its own; +- the ``reports`` row, of type ``household``, naming the household, 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, or the baseline's + alone. + +The service then starts one household function call for the report run and +records the call's ID on it, or fails it and its simulation runs at once, +with ``economy_run_records``' report run writes. The household function +writes the rest: each run's ``household_result``, and every status after +``pending``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +from sqlalchemy import insert + +from policyengine_simulation_persistence.economy_run_records import ( + references_checked, +) +from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + ReportRunTrigger, + SimulationStatus, + SimulationType, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from uuid import UUID + + from sqlalchemy import Connection + + from policyengine_simulation_persistence.economy_run_records import RunVersions + +# A report's type, as the Public API reads it. +HOUSEHOLD_REPORT_TYPE: Final = "household" + + +@dataclass(frozen=True) +class HouseholdReportSimulation: + """The new household simulation and run one side of a report creates. + + ``policy_id`` is the side's stored policy, or None for current law. + """ + + simulation_id: UUID + simulation_run_id: UUID + policy_id: UUID | None + + +@dataclass(frozen=True) +class NewHouseholdReport: + """A household report to create, with its first run. + + ``reform`` is None for a report of the baseline alone. ``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 + tax_benefit_model_version_id: UUID + household_id: UUID + year: int + versions: RunVersions + inputs: dict[str, Any] + baseline: HouseholdReportSimulation + reform: HouseholdReportSimulation | None + + @property + def sides(self) -> tuple[HouseholdReportSimulation, ...]: + """Return each side the report calculates, the baseline's first.""" + return (self.baseline,) if self.reform is None else (self.baseline, self.reform) + + @property + def 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.sides) + + +def insert_household_report(connection: Connection, report: NewHouseholdReport) -> None: + """Create the report, its first run and each side's 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.sides: + _insert_household_simulation(connection, report, side) + reform = report.reform + connection.execute( + insert(reports).values( + id=report.id, + label=report.label, + description=report.description, + country=report.country, + type=HOUSEHOLD_REPORT_TYPE, + user_id=None, + tax_benefit_model_id=report.tax_benefit_model_id, + policy_id=None if reform is None else reform.policy_id, + baseline_simulation_id=report.baseline.simulation_id, + reform_simulation_id=None if reform is None else reform.simulation_id, + household_id=report.household_id, + dataset_id=None, + region_id=None, + year=report.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=( + None if reform is None else 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_household_simulation( + connection: Connection, report: NewHouseholdReport, side: HouseholdReportSimulation +) -> None: + """Create one side's household simulation and its first run, pending.""" + connection.execute( + insert(simulations).values( + id=side.simulation_id, + simulation_type=SimulationType.HOUSEHOLD, + dataset_id=None, + household_id=report.household_id, + policy_id=side.policy_id, + dynamic_id=None, + tax_benefit_model_version_id=report.tax_benefit_model_version_id, + region_id=None, + year=report.year, + ) + ) + connection.execute( + insert(simulation_runs).values( + id=side.simulation_run_id, + simulation_id=side.simulation_id, + year=report.year, + tax_benefit_model_version_id=report.tax_benefit_model_version_id, + country_package_version=report.versions.country_package_version, + policyengine_version=report.versions.policyengine_version, + dataset_id=None, + region_id=None, + status=SimulationStatus.PENDING, + modal_call_id=None, + ) + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/household_report_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/household_report_store.py new file mode 100644 index 000000000..9fa51b68a --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/household_report_store.py @@ -0,0 +1,210 @@ +"""The household report store: each read and write in a transaction of its own. + +The Simulation Entrypoint creates a household report in three steps, each +its own transaction, with the household function's call started between +them: it reads the report's inputs, creates the rows, then records the call +on the report run or fails the run. ``HouseholdReportStore`` runs the reads +and writes of ``household_report_records``, ``household_job_records``, +``policy_reform`` and ``economy_run_records``' report run writes that way. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_persistence.economy_run_records import ( + fail_unstarted_report_run, + read_model_id, + record_report_run_call, +) +from policyengine_simulation_persistence.economy_run_store import ( + EconomyRecordOutcomeUnknownError, + ReportSideError, +) +from policyengine_simulation_persistence.engine import bounded_transaction +from policyengine_simulation_persistence.household_job_records import ( + read_catalog_version_id, + read_stored_household, +) +from policyengine_simulation_persistence.household_job_store import ( + UnknownCatalogVersionError, + UnknownHouseholdError, +) +from policyengine_simulation_persistence.household_report_records import ( + insert_household_report, +) +from policyengine_simulation_persistence.policy_reform import ( + MissingPolicyParametersError, + UnknownPolicyError, + read_policy_reform, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping, Sequence + from contextlib import AbstractContextManager + from uuid import UUID + + from sqlalchemy import Connection, Engine + + from policyengine_simulation_persistence.economy_run_store import ReportSide + from policyengine_simulation_persistence.household_job_records import ( + StoredHousehold, + ) + from policyengine_simulation_persistence.household_report_records import ( + NewHouseholdReport, + ) + from policyengine_simulation_persistence.policy_reform import PolicyReform + from policyengine_simulation_persistence.tables import ReportRunStatus + +type TransactionFactory = Callable[[], AbstractContextManager[Connection]] + + +@dataclass(frozen=True) +class HouseholdReportInputs: + """What a household report is read from. + + ``catalog_version_id`` is the country's catalog version for the .py + version the report runs on, and ``tax_benefit_model_id`` its model. + ``reforms`` holds each requested side's policy values by parameter name + and period, empty for current law. + """ + + catalog_version_id: UUID + tax_benefit_model_id: UUID + household: StoredHousehold + reforms: dict[ReportSide, PolicyReform] + + +def _side_reform( + connection: Connection, + side: ReportSide, + *, + country: str, + policy_id: UUID | None, + catalog_version_id: UUID, +) -> PolicyReform: + """Return one side's policy values; raise ReportSideError naming the side.""" + if policy_id is None: + return {} + try: + return read_policy_reform( + connection, + country=country, + policy_id=policy_id, + catalog_version_id=catalog_version_id, + ) + except (UnknownPolicyError, MissingPolicyParametersError) as error: + raise ReportSideError(side, error) from error + + +class HouseholdReportStore: + """Read and write household reports, each operation 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 + ) -> HouseholdReportStore: + """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_inputs( + self, + *, + country: str, + household_id: UUID, + policy_ids: Mapping[ReportSide, UUID | None], + policyengine_version: str, + ) -> HouseholdReportInputs: + """Read the catalog version, the household and each side's policy together. + + ``policy_ids`` names each side's stored policy, None for current law. + Raise UnknownCatalogVersionError or UnknownHouseholdError, or + ReportSideError naming the side whose policy is missing + (UnknownPolicyError) or sets parameters the catalog version lacks + (MissingPolicyParametersError). + """ + with self._transaction() as connection: + catalog_version_id = read_catalog_version_id( + connection, country=country, policyengine_version=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 {policyengine_version}" + raise UnknownCatalogVersionError(msg) + household = read_stored_household( + connection, country=country, household_id=household_id + ) + if household is None: + msg = f"{country} has no household {household_id}" + raise UnknownHouseholdError(msg) + reforms: dict[ReportSide, PolicyReform] = { + side: _side_reform( + connection, + side, + country=country, + policy_id=policy_id, + catalog_version_id=catalog_version_id, + ) + for side, policy_id in policy_ids.items() + } + return HouseholdReportInputs( + catalog_version_id=catalog_version_id, + tax_benefit_model_id=model_id, + household=household, + reforms=reforms, + ) + + def create_report(self, report: NewHouseholdReport) -> None: + """Create the report, its simulations and its runs. + + Raise ``insert_household_report``'s errors, and + EconomyRecordOutcomeUnknownError when the commit fails after the + rows were inserted: they may exist. + """ + inserted = False + try: + with self._transaction() as connection: + insert_household_report(connection, report) + 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 + ) -> 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 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, + ) diff --git a/libs/policyengine-simulation-persistence/tests/test_household_report_store.py b/libs/policyengine-simulation-persistence/tests/test_household_report_store.py new file mode 100644 index 000000000..b7a138465 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_household_report_store.py @@ -0,0 +1,452 @@ +"""The household report store: each read and write in a transaction of its own. + +A household report calculates one stored household under a baseline policy +and, optionally, a reform. These run against a database the Public API's own +v2 migration chain built (``PUBLIC_API_SCHEMA_DATABASE_URL``); each store +transaction 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 policyengine_simulation_persistence import ( + EconomyRecordOutcomeUnknownError, + EconomyRecordReferenceError, + HouseholdReportSimulation, + HouseholdReportStore, + MissingPolicyParametersError, + NewHouseholdReport, + ReportRunStatus, + ReportSideError, + RunVersions, + UnknownCatalogVersionError, + UnknownHouseholdError, + UnknownPolicyError, +) +from public_api_rows import ( + StoredValue, + insert_catalog_version, + insert_household, + insert_model, + insert_policy, +) +from sqlalchemy import text +from sqlalchemy.exc import OperationalError + +if TYPE_CHECKING: + from collections.abc import Iterator + from uuid import UUID + + from public_api_rows import CatalogVersion + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +POLICYENGINE_VERSION = "6.17.1" +VERSIONS = RunVersions( + country_package_version="1.764.6", policyengine_version=POLICYENGINE_VERSION +) +START = datetime(2026, 1, 1, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +DOCUMENT = { + "people": [ + { + "id": "you", + "values": {"age": 40}, + "memberships": {"household": "your household"}, + } + ], + "household": [{"id": "your household", "values": {}}], +} +# A household report has no options: it always adds missing variables. +OPTIONS: dict[str, object] = {} + + +def _store_on(connection: Connection) -> HouseholdReportStore: + """Return a store whose transactions are savepoints of ``connection``.""" + + @contextmanager + def savepoint() -> Iterator[Connection]: + with connection.begin_nested(): + yield connection + + return HouseholdReportStore(savepoint) + + +def _catalog(connection: Connection) -> CatalogVersion: + return insert_catalog_version( + connection, + country="us", + model_id=insert_model(connection, "us"), + policyengine_version=POLICYENGINE_VERSION, + parameter_names=(CTC_AMOUNT,), + ) + + +def _household(connection: Connection, *, default_year: int | None = 2026) -> UUID: + return insert_household( + connection, country="us", default_year=default_year, document=DOCUMENT + ) + + +def _new_report( + catalog: CatalogVersion, + household_id: UUID, + *, + reform_policy_id: UUID | None = None, + with_reform: bool = True, +) -> NewHouseholdReport: + return NewHouseholdReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label="A larger child tax credit", + description="For one family", + tax_benefit_model_id=catalog.model_id, + tax_benefit_model_version_id=catalog.version_id, + household_id=household_id, + year=2026, + versions=VERSIONS, + inputs=OPTIONS, + baseline=HouseholdReportSimulation( + simulation_id=uuid4(), simulation_run_id=uuid4(), policy_id=None + ), + reform=( + HouseholdReportSimulation( + simulation_id=uuid4(), + simulation_run_id=uuid4(), + policy_id=reform_policy_id, + ) + if with_reform + else None + ), + ) + + +def _row(connection: Connection, table_name: str, row_id: UUID) -> dict: + return dict( + connection.execute( + text(f"SELECT * FROM public.{table_name} WHERE id = :id"), # noqa: S608 # fixed table names + {"id": row_id}, + ) + .mappings() + .one() + ) + + +def test_a_household_reports_inputs_are_read_together(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + household_id = _household(connection) + policy_id = insert_policy( + connection, catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + act_as_the_service(connection) + + inputs = _store_on(connection).read_inputs( + country="us", + household_id=household_id, + policy_ids={"baseline": None, "reform": policy_id}, + policyengine_version=POLICYENGINE_VERSION, + ) + + assert inputs.catalog_version_id == catalog.version_id + assert inputs.tax_benefit_model_id == catalog.model_id + assert inputs.household.id == household_id + assert inputs.household.default_year == 2026 + assert inputs.household.document == DOCUMENT + assert inputs.reforms == { + "baseline": {}, + "reform": {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}}, + } + + +def test_each_missing_input_is_refused(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + household_id = _household(connection) + act_as_the_service(connection) + store = _store_on(connection) + missing_policy_id = uuid4() + + with pytest.raises(UnknownCatalogVersionError): + store.read_inputs( + country="us", + household_id=household_id, + policy_ids={"baseline": None}, + policyengine_version="0.0.1", + ) + with pytest.raises(UnknownHouseholdError): + store.read_inputs( + country="us", + household_id=uuid4(), + policy_ids={"baseline": None}, + policyengine_version=POLICYENGINE_VERSION, + ) + with pytest.raises(ReportSideError) as unknown_policy: + store.read_inputs( + country="us", + household_id=household_id, + policy_ids={"baseline": None, "reform": missing_policy_id}, + policyengine_version=POLICYENGINE_VERSION, + ) + + assert catalog.policyengine_version == POLICYENGINE_VERSION + assert unknown_policy.value.side == "reform" + assert isinstance(unknown_policy.value.error, UnknownPolicyError) + + +def test_a_policy_setting_a_parameter_the_version_lacks_is_refused_on_its_side( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + household_id = _household(connection) + # A policy saved on an older catalog version, naming a parameter the + # running version lacks. + older = insert_catalog_version( + connection, + country="us", + model_id=catalog.model_id, + policyengine_version="6.0.0", + parameter_names=("gov.removed.parameter",), + ) + policy_id = insert_policy( + connection, older, [StoredValue("gov.removed.parameter", 1, START)] + ) + act_as_the_service(connection) + + with pytest.raises(ReportSideError) as refused: + _store_on(connection).read_inputs( + country="us", + household_id=household_id, + policy_ids={"baseline": policy_id, "reform": None}, + policyengine_version=POLICYENGINE_VERSION, + ) + + assert refused.value.side == "baseline" + assert isinstance(refused.value.error, MissingPolicyParametersError) + assert refused.value.error.parameter_names == ("gov.removed.parameter",) + + +def test_a_household_report_is_created_with_its_simulations_and_pending_runs( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + household_id = _household(connection) + policy_id = insert_policy(connection, catalog) + act_as_the_service(connection) + report = _new_report(catalog, household_id, reform_policy_id=policy_id) + assert report.reform is not None + + _store_on(connection).create_report(report) + stored_report = _row(connection, "reports", report.id) + stored_run = _row(connection, "report_runs", report.run_id) + simulations = [ + _row(connection, "simulations", side.simulation_id) + for side in (report.baseline, report.reform) + ] + simulation_runs = [ + _row(connection, "simulation_runs", side.simulation_run_id) + for side in (report.baseline, report.reform) + ] + + assert report.simulation_run_ids == ( + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ) + assert { + key: stored_report[key] + for key in ( + "label", + "description", + "country", + "type", + "tax_benefit_model_id", + "policy_id", + "baseline_simulation_id", + "reform_simulation_id", + "household_id", + "dataset_id", + "region_id", + "year", + "budget_window_size", + "inputs", + ) + } == { + "label": "A larger child tax credit", + "description": "For one family", + "country": "us", + "type": "household", + "tax_benefit_model_id": catalog.model_id, + "policy_id": policy_id, + "baseline_simulation_id": report.baseline.simulation_id, + "reform_simulation_id": report.reform.simulation_id, + "household_id": household_id, + "dataset_id": None, + "region_id": None, + "year": 2026, + "budget_window_size": None, + "inputs": OPTIONS, + } + assert (stored_run["status"], stored_run["trigger"]) == ("pending", "initial") + assert stored_run["modal_call_id"] is None + assert ( + stored_run["baseline_simulation_run_id"], + stored_run["reform_simulation_run_id"], + ) == report.simulation_run_ids + assert ( + stored_run["country_package_version"], + stored_run["policyengine_version"], + ) == ("1.764.6", POLICYENGINE_VERSION) + assert [simulation["simulation_type"] for simulation in simulations] == [ + "household", + "household", + ] + assert [simulation["policy_id"] for simulation in simulations] == [ + None, + policy_id, + ] + for simulation in simulations: + assert simulation["household_id"] == household_id + assert simulation["tax_benefit_model_version_id"] == catalog.version_id + assert (simulation["dataset_id"], simulation["region_id"]) == (None, None) + assert (simulation["dynamic_id"], simulation["year"]) == (None, 2026) + for run in simulation_runs: + assert (run["status"], run["modal_call_id"]) == ("pending", None) + assert (run["dataset_id"], run["region_id"], run["year"]) == (None, None, 2026) + assert run["household_result"] is None + assert run["policyengine_version"] == POLICYENGINE_VERSION + + +def test_a_report_of_the_baseline_alone_names_no_reform( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + household_id = _household(connection) + act_as_the_service(connection) + report = _new_report(catalog, household_id, with_reform=False) + + _store_on(connection).create_report(report) + stored_report = _row(connection, "reports", report.id) + stored_run = _row(connection, "report_runs", report.run_id) + + assert report.simulation_run_ids == (report.baseline.simulation_run_id,) + assert stored_report["reform_simulation_id"] is None + assert stored_report["policy_id"] is None + assert stored_run["reform_simulation_run_id"] is None + + +def test_a_report_run_records_its_call_or_is_failed_with_its_runs( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + household_id = _household(connection) + act_as_the_service(connection) + store = _store_on(connection) + recorded, refused = (_new_report(catalog, household_id) for _ in range(2)) + store.create_report(recorded) + store.create_report(refused) + + status = store.record_report_call(recorded.run_id, "fc-household") + recorded_not_failed = store.fail_unstarted_report( + recorded.run_id, recorded.simulation_run_ids, "modal_call_refused: late" + ) + failed = store.fail_unstarted_report( + refused.run_id, refused.simulation_run_ids, "modal_call_refused: no" + ) + recorded_run = _row(connection, "report_runs", recorded.run_id) + refused_run = _row(connection, "report_runs", refused.run_id) + refused_simulation_runs = [ + _row(connection, "simulation_runs", run_id) + for run_id in refused.simulation_run_ids + ] + + assert (status, recorded_not_failed, failed) == ( + ReportRunStatus.PENDING, + False, + True, + ) + assert (recorded_run["status"], recorded_run["modal_call_id"]) == ( + "pending", + "fc-household", + ) + assert (refused_run["status"], refused_run["error_message"]) == ( + "failed", + "modal_call_refused: no", + ) + assert [run["status"] for run in refused_simulation_runs] == ["failed", "failed"] + + +def test_a_report_naming_a_missing_household_is_refused_by_its_key( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = _catalog(connection) + act_as_the_service(connection) + + with pytest.raises(EconomyRecordReferenceError) as refused: + _store_on(connection).create_report(_new_report(catalog, uuid4())) + + assert refused.value.column_name == "household_id" + + +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 test_a_create_whose_commit_fails_may_have_been_written() -> None: + connection = _RecordingConnection() + store = HouseholdReportStore(lambda: _CommitFails(connection)) + catalog_rows = {"model_id": uuid4(), "version_id": uuid4()} + + with pytest.raises(EconomyRecordOutcomeUnknownError) as raised: + store.create_report( + NewHouseholdReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label="Lost commit", + description=None, + tax_benefit_model_id=catalog_rows["model_id"], + tax_benefit_model_version_id=catalog_rows["version_id"], + household_id=uuid4(), + year=2026, + versions=VERSIONS, + inputs=OPTIONS, + baseline=HouseholdReportSimulation(uuid4(), uuid4(), None), + reform=None, + ) + ) + + assert isinstance(raised.value.__cause__, OperationalError) + assert connection.statements diff --git a/libs/policyengine-stage12-persistence/README.md b/libs/policyengine-stage12-persistence/README.md index 686c9d944..9dd737d62 100644 --- a/libs/policyengine-stage12-persistence/README.md +++ b/libs/policyengine-stage12-persistence/README.md @@ -43,6 +43,13 @@ revisions and live schema. This package never changes a schema. to its window, fails the window of a year it fails, and settles a window no year advances any more. The Modal worker role and the simulation service's role may both advance a window. +- **Household report runs** (`HouseholdReportRunStore`, a `ReportRunStore`): + one household function call claims a household report run and each of + its simulation runs, the baseline's and any reform's, records each run's + household result or failure document in `simulation_runs.household_result`, + and then marks the report run succeeded or failed. It reads first which + simulation runs the report run names, to check them against its payload + and to fail them if it stops. ## Tests diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index 3f1bca46a..7e86cb336 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -8,6 +8,9 @@ YearStarter, YearToStart, ) +from policyengine_stage12_persistence.household_report_run_store import ( + HouseholdReportRunStore, +) from policyengine_stage12_persistence.report_run_results import ( DistrictCodes, ReportRunResult, @@ -38,6 +41,7 @@ "BudgetWindowStore", "ConnectionOpener", "DistrictCodes", + "HouseholdReportRunStore", "ReportRunResult", "ReportRunStore", "RunKind", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/household_report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/household_report_run_store.py new file mode 100644 index 000000000..c3b343b85 --- /dev/null +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/household_report_run_store.py @@ -0,0 +1,162 @@ +"""Conditional writes to a household report run and its simulation runs. + +A household report calculates one stored household under a baseline and, +optionally, a reform: each side is a household simulation run, and the +report run names them. One household function call calculates the whole +report run, as the Modal worker role. It claims the report run and each +simulation run with ``ReportRunStore``'s claims, under its own Modal call +ID, reads which simulation runs the report run names, then records each +run's ``HouseholdExecutionResult`` or ``HouseholdExecutionFailure`` +document in ``household_result``, and marks the report run succeeded or +failed. Each write here applies only while the call +still holds its row running. A call that stops fails the runs it left +unfinished with ``ReportRunStore.fail_simulation_run``, which, given no +call, applies while the run is pending or running. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from sqlalchemy import select, update + +from policyengine_stage12_persistence.report_run_store import ( + ReportRunStore, + RunStatus, + running_under_call, +) +from policyengine_stage12_persistence.report_run_tables import ( + report_runs, + simulation_runs, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + from datetime import datetime + from uuid import UUID + + from policyengine_stage12_persistence.report_run_store import RunWrite + + +class HouseholdReportRunStore(ReportRunStore): + """A report run store that also records household simulation runs' results.""" + + def household_report_simulation_runs( + self, report_run_id: UUID + ) -> tuple[UUID, UUID | None]: + """Return the baseline and reform simulation runs the report run names. + + The reform's is None for a report of the baseline alone. Raises + LookupError when the report run does not exist. + """ + statement = select( + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + ).where(report_runs.c.id == report_run_id) + with self._transaction("household_report_simulation_runs") as connection: + named = connection.execute(statement).one_or_none() + if named is None: + msg = f"report run {report_run_id} does not exist" + raise LookupError(msg) + return named.baseline_simulation_run_id, named.reform_simulation_run_id + + def complete_household_simulation_run( + self, + simulation_run_id: UUID, + *, + modal_call_id: str, + household_result: Mapping[str, Any], + completed_at: datetime, + ) -> RunWrite: + """Mark the run succeeded with its household while this call holds it. + + ``household_result`` is the calculation's ``HouseholdExecutionResult`` + as JSON. + """ + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + running_under_call(simulation_runs, modal_call_id), + ) + .values( + status=RunStatus.SUCCEEDED.value, + household_result=dict(household_result), + error_message=None, + completed_at=completed_at, + updated_at=completed_at, + ) + ) + return self._write( + simulation_runs, + simulation_run_id, + statement, + "complete_household_simulation_run", + ) + + def fail_household_simulation_run( + self, + simulation_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + household_result: Mapping[str, Any] | None = None, + ) -> RunWrite: + """Mark the run failed while this call holds it. + + ``household_result`` is the calculation's + ``HouseholdExecutionFailure`` as JSON, or None, stored as SQL NULL, + for a failure outside the calculation. + """ + statement = ( + update(simulation_runs) + .where( + simulation_runs.c.id == simulation_run_id, + running_under_call(simulation_runs, modal_call_id), + ) + .values( + status=RunStatus.FAILED.value, + household_result=( + None if household_result is None else dict(household_result) + ), + error_message=error_message, + completed_at=failed_at, + updated_at=failed_at, + ) + ) + return self._write( + simulation_runs, + simulation_run_id, + statement, + "fail_household_simulation_run", + ) + + def complete_household_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + completed_at: datetime, + ) -> RunWrite: + """Mark the household report run succeeded while this call holds it. + + A household report run has no results of its own: each simulation + run holds its calculated household. + """ + statement = ( + update(report_runs) + .where( + report_runs.c.id == report_run_id, + running_under_call(report_runs, modal_call_id), + ) + .values( + status=RunStatus.SUCCEEDED.value, + error_message=None, + completed_at=completed_at, + updated_at=completed_at, + ) + ) + return self._write( + report_runs, report_run_id, statement, "complete_household_report_run" + ) diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py index 5161ddda9..7cb2d5e48 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_store.py @@ -275,7 +275,7 @@ def complete_report_run( update(report_runs) .where( report_runs.c.id == run_ids.report_run_id, - _running_under_call(report_runs, modal_call_id), + running_under_call(report_runs, modal_call_id), ) .values( status=RunStatus.SUCCEEDED.value, @@ -306,7 +306,7 @@ def fail_report_run( update(report_runs) .where( report_runs.c.id == report_run_id, - _running_under_call(report_runs, modal_call_id), + running_under_call(report_runs, modal_call_id), ) .values( status=RunStatus.FAILED.value, @@ -372,7 +372,7 @@ def complete_simulation_run( update(simulation_runs) .where( simulation_runs.c.id == simulation_run_id, - _running_under_call(simulation_runs, modal_call_id), + running_under_call(simulation_runs, modal_call_id), ) .values( status=RunStatus.SUCCEEDED.value, @@ -405,7 +405,7 @@ def fail_simulation_run( held = ( simulation_runs.c.status.in_(ACTIVE_STATUSES) if modal_call_id is None - else _running_under_call(simulation_runs, modal_call_id) + else running_under_call(simulation_runs, modal_call_id) ) statement = ( update(simulation_runs) @@ -505,7 +505,7 @@ def complete_household_job( update(household_jobs) .where( household_jobs.c.id == household_job_id, - _running_under_call(household_jobs, modal_call_id), + running_under_call(household_jobs, modal_call_id), ) .values( status=RunStatus.SUCCEEDED.value, @@ -538,7 +538,7 @@ def fail_household_job( update(household_jobs) .where( household_jobs.c.id == household_job_id, - _running_under_call(household_jobs, modal_call_id), + running_under_call(household_jobs, modal_call_id), ) .values( status=RunStatus.FAILED.value, @@ -705,7 +705,7 @@ def _uniform_rows(rows: tuple[dict[str, Any], ...]) -> list[dict[str, Any]]: return [dict.fromkeys(columns) | row for row in rows] -def _running_under_call(table: Table, modal_call_id: str) -> ColumnElement[bool]: +def running_under_call(table: Table, modal_call_id: str) -> ColumnElement[bool]: """Return the condition that a row is running under ``modal_call_id``. A run or household job is completed or failed by a call only while that diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py index 9bcc57cb5..68c574b40 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/report_run_tables.py @@ -95,6 +95,9 @@ Column("updated_at", DateTime(timezone=True), nullable=False), Column("output_uri", String(2048)), Column("output_sha256", String(64)), + # A household simulation run's ``HouseholdExecutionResult`` or + # ``HouseholdExecutionFailure`` document, set by the household function. + Column("household_result", JSON(none_as_null=True)), ) # A household job is one household calculation, which a household function diff --git a/libs/policyengine-stage12-persistence/tests/household_report_database.py b/libs/policyengine-stage12-persistence/tests/household_report_database.py new file mode 100644 index 000000000..837b314e0 --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/household_report_database.py @@ -0,0 +1,182 @@ +"""Seed household report runs as the simulation service creates them, for tests. + +A household report calculates one stored household under a baseline policy +and, optionally, a reform. The simulation service creates a household +simulation and a pending simulation run for each side, the report, and its +pending report run naming those runs, with no Modal call recorded on any of +them. The seeder inserts the same rows as the seed role, through +``ReportRunSeeder``'s reflected tables. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +from report_run_database import VERSIONS +from sqlalchemy import delete, insert + +if TYPE_CHECKING: + from report_run_database import ReportRunSeeder + from sqlalchemy import Connection + +HOUSEHOLD_YEAR = 2026 + + +@dataclass(frozen=True) +class SeededHouseholdReportRun: + """A pending household report run and the simulation runs it names. + + ``reform_simulation_run_id`` is None for a report of the baseline alone. + """ + + report_run_id: UUID + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + report_id: UUID + household_id: UUID + model_id: UUID + + +def seed_household_report_run( + seeder: ReportRunSeeder, *, with_reform: bool = True +) -> SeededHouseholdReportRun: + """Insert a household report whose report run and simulation runs are pending.""" + model_id, version_id = uuid4(), uuid4() + household_id, report_id, report_run_id = uuid4(), uuid4(), uuid4() + sides = 2 if with_reform else 1 + simulation_ids = [uuid4() for _ in range(sides)] + simulation_run_ids = [uuid4() for _ in range(sides)] + with seeder.engine.begin() as connection: + _insert_catalog_and_household( + seeder, connection, model_id, version_id, household_id + ) + for simulation_id, run_id in zip( + simulation_ids, simulation_run_ids, strict=True + ): + connection.execute( + insert(seeder.table("simulations")).values( + id=simulation_id, + simulation_type="household", + household_id=household_id, + tax_benefit_model_version_id=version_id, + year=HOUSEHOLD_YEAR, + ) + ) + connection.execute( + insert(seeder.table("simulation_runs")).values( + id=run_id, + simulation_id=simulation_id, + year=HOUSEHOLD_YEAR, + tax_benefit_model_version_id=version_id, + status="pending", + **VERSIONS, + ) + ) + reform_simulation_id = simulation_ids[1] if with_reform else None + reform_run_id = simulation_run_ids[1] if with_reform else None + connection.execute( + insert(seeder.table("reports")).values( + id=report_id, + label="Household report run test", + country="us", + type="household", + tax_benefit_model_id=model_id, + household_id=household_id, + baseline_simulation_id=simulation_ids[0], + reform_simulation_id=reform_simulation_id, + year=HOUSEHOLD_YEAR, + # A household report has no options. + inputs={}, + ) + ) + connection.execute( + insert(seeder.table("report_runs")).values( + id=report_run_id, + report_id=report_id, + baseline_simulation_run_id=simulation_run_ids[0], + reform_simulation_run_id=reform_run_id, + status="pending", + trigger="initial", + **VERSIONS, + ) + ) + return SeededHouseholdReportRun( + report_run_id=report_run_id, + baseline_simulation_run_id=simulation_run_ids[0], + reform_simulation_run_id=reform_run_id, + report_id=report_id, + household_id=household_id, + model_id=model_id, + ) + + +def remove_household_report_run( + seeder: ReportRunSeeder, seeded: SeededHouseholdReportRun +) -> None: + """Delete the seeded report, its runs, simulations, household and catalog rows.""" + report_runs = seeder.table("report_runs") + simulation_runs = seeder.table("simulation_runs") + simulations = seeder.table("simulations") + reports = seeder.table("reports") + households = seeder.table("households") + models = seeder.table("tax_benefit_models") + simulation_run_ids = [ + run_id + for run_id in ( + seeded.baseline_simulation_run_id, + seeded.reform_simulation_run_id, + ) + if run_id is not None + ] + with seeder.engine.begin() as connection: + connection.execute( + delete(report_runs).where(report_runs.c.id == seeded.report_run_id) + ) + connection.execute( + delete(simulation_runs).where(simulation_runs.c.id.in_(simulation_run_ids)) + ) + connection.execute(delete(reports).where(reports.c.id == seeded.report_id)) + connection.execute( + delete(simulations).where(simulations.c.household_id == seeded.household_id) + ) + connection.execute( + delete(households).where(households.c.id == seeded.household_id) + ) + # Deleting the model removes its versions. + connection.execute(delete(models).where(models.c.id == seeded.model_id)) + + +def _insert_catalog_and_household( + seeder: ReportRunSeeder, + connection: Connection, + model_id: UUID, + version_id: UUID, + household_id: UUID, +) -> None: + connection.execute( + insert(seeder.table("tax_benefit_models")).values( + id=model_id, name=f"test-{model_id.hex[:24]}" + ) + ) + connection.execute( + insert(seeder.table("tax_benefit_model_versions")).values( + id=version_id, + model_id=model_id, + version="1.0.0", + current_law_id=1, + metadata_time_periods={}, + ) + ) + connection.execute( + insert(seeder.table("households")).values( + id=household_id, + household_data={"people": {"you": {"age": {"2026": 40}}}}, + country_id="us", + default_year=HOUSEHOLD_YEAR, + canonicalization_version=1, + # Unique per household: households are deduplicated by it. + content_hash=household_id.hex * 2, + ) + ) diff --git a/libs/policyengine-stage12-persistence/tests/test_household_report_run_store.py b/libs/policyengine-stage12-persistence/tests/test_household_report_run_store.py new file mode 100644 index 000000000..62dd39d1d --- /dev/null +++ b/libs/policyengine-stage12-persistence/tests/test_household_report_run_store.py @@ -0,0 +1,299 @@ +"""The Modal worker role's writes to a household report run and its simulation runs. + +A household function given a household report run claims the report run and +each of its simulation runs under its Modal call ID, records each simulation +run's calculated household, or its failure, in ``household_result``, then +marks the report run succeeded or failed. These run against a disposable +Postgres that the Public API's v2 migrations created, connected as the Modal +worker role (see ``report_run_database``). +""" + +from __future__ import annotations + +from datetime import timedelta +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +import pytest +from household_report_database import ( + remove_household_report_run, + seed_household_report_run, +) +from policyengine_simulation_contract.household_execution import ( + CalculatedHousehold, + HouseholdBundleProvenance, + HouseholdCalculationError, + HouseholdCalculationOutput, + HouseholdExecutionFailure, + HouseholdExecutionResult, + HouseholdWorkerIdentity, + parse_household_execution_outcome, +) +from policyengine_stage12_persistence import ( + HouseholdReportRunStore, + RunStatus, + RunWrite, +) +from report_run_database import NOW, worker_database_url +from sqlalchemy import select + +if TYPE_CHECKING: + from collections.abc import Iterator + + from household_report_database import SeededHouseholdReportRun + from report_run_database import ReportRunSeeder + +LATER = NOW + timedelta(minutes=5) +CALL = "fc-household-report" +OTHER_CALL = "fc-other" +WORKER = HouseholdWorkerIdentity( + application_name="policyengine-simulation-v2-py6-12-0", + function_name="calculate_household_us", +) +CALCULATION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +# The documents a household function stores, built from the contract's models +# so they stay valid household function outcomes. +SUCCEEDED_DOCUMENT = HouseholdExecutionResult( + calculation_id=CALCULATION_ID, + output=HouseholdCalculationOutput( + country="us", + bundle=HouseholdBundleProvenance( + policyengine_version="6.12.0", + core_package_version="3.33.0", + country_package_name="policyengine-us", + country_package_version="1.0.0", + bundle_manifest_sha256="a" * 64, + ), + household=CalculatedHousehold.model_validate( + { + "people": {"you": {"age": {"2026": 40}}}, + "households": { + "your household": { + "members": ["you"], + "household_net_income": {"2026": 25_000.0}, + } + }, + } + ), + ), + worker=WORKER, +).model_dump(mode="json") +FAILED_DOCUMENT = HouseholdExecutionFailure( + calculation_id=CALCULATION_ID, + error=HouseholdCalculationError( + code="invalid_inputs", + message="Unrecognized calculate input(s): Unrecognized policy parameter", + ), + worker=WORKER, +).model_dump(mode="json") +FAILURE_SUMMARY = ( + "invalid_inputs: Unrecognized calculate input(s): Unrecognized policy parameter" +) + + +@pytest.fixture +def store() -> Iterator[HouseholdReportRunStore]: + run_store = HouseholdReportRunStore(worker_database_url()) + try: + yield run_store + finally: + run_store.engine.dispose() + + +@pytest.fixture +def seeded_report(seeder: ReportRunSeeder) -> Iterator[SeededHouseholdReportRun]: + seeded = seed_household_report_run(seeder) + try: + yield seeded + finally: + remove_household_report_run(seeder, seeded) + + +@pytest.fixture +def seeded_baseline_only( + seeder: ReportRunSeeder, +) -> Iterator[SeededHouseholdReportRun]: + seeded = seed_household_report_run(seeder, with_reform=False) + try: + yield seeded + finally: + remove_household_report_run(seeder, seeded) + + +def _result_is_null(seeder: ReportRunSeeder, simulation_run_id: UUID) -> bool: + """Whether the run's household_result is SQL NULL, not a stored JSON null.""" + runs = seeder.table("simulation_runs") + with seeder.engine.connect() as connection: + return ( + connection.execute( + select(runs.c.id).where( + runs.c.id == simulation_run_id, runs.c.household_result.is_(None) + ) + ).one_or_none() + is not None + ) + + +def test_a_household_report_runs_simulation_runs_are_read( + store: HouseholdReportRunStore, + seeded_report: SeededHouseholdReportRun, + seeded_baseline_only: SeededHouseholdReportRun, +) -> None: + both = store.household_report_simulation_runs(seeded_report.report_run_id) + baseline_only = store.household_report_simulation_runs( + seeded_baseline_only.report_run_id + ) + + assert both == ( + seeded_report.baseline_simulation_run_id, + seeded_report.reform_simulation_run_id, + ) + assert baseline_only == (seeded_baseline_only.baseline_simulation_run_id, None) + missing_run_id = uuid4() + with pytest.raises(LookupError, match=f"report run {missing_run_id}"): + store.household_report_simulation_runs(missing_run_id) + + +def test_a_household_simulation_run_records_its_result_only_while_its_call_holds_it( + store: HouseholdReportRunStore, + seeder: ReportRunSeeder, + seeded_report: SeededHouseholdReportRun, +) -> None: + run_id = seeded_report.baseline_simulation_run_id + store.claim_simulation_run(run_id, modal_call_id=CALL, claimed_at=NOW) + + by_other_call = store.complete_household_simulation_run( + run_id, + modal_call_id=OTHER_CALL, + household_result=SUCCEEDED_DOCUMENT, + completed_at=LATER, + ) + completed = store.complete_household_simulation_run( + run_id, + modal_call_id=CALL, + household_result=SUCCEEDED_DOCUMENT, + completed_at=LATER, + ) + again = store.complete_household_simulation_run( + run_id, + modal_call_id=CALL, + household_result=FAILED_DOCUMENT, + completed_at=LATER, + ) + + assert by_other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert completed == RunWrite(changed=True, status=RunStatus.SUCCEEDED) + assert again == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + row = seeder.row("simulation_runs", run_id) + assert row["status"] == "succeeded" + assert row["household_result"] == SUCCEEDED_DOCUMENT + assert parse_household_execution_outcome(row["household_result"]).status == ( + "succeeded" + ) + assert row["completed_at"] == LATER + assert row["error_message"] is None + assert (row["output_uri"], row["output_sha256"]) == (None, None) + + +def test_a_household_simulation_run_fails_with_its_failure_document( + store: HouseholdReportRunStore, + seeder: ReportRunSeeder, + seeded_report: SeededHouseholdReportRun, +) -> None: + run_id = seeded_report.reform_simulation_run_id + assert run_id is not None + store.claim_simulation_run(run_id, modal_call_id=CALL, claimed_at=NOW) + + by_other_call = store.fail_household_simulation_run( + run_id, + modal_call_id=OTHER_CALL, + error_message=FAILURE_SUMMARY, + failed_at=LATER, + household_result=FAILED_DOCUMENT, + ) + failed = store.fail_household_simulation_run( + run_id, + modal_call_id=CALL, + error_message=FAILURE_SUMMARY, + failed_at=LATER, + household_result=FAILED_DOCUMENT, + ) + completed_after = store.complete_household_simulation_run( + run_id, + modal_call_id=CALL, + household_result=SUCCEEDED_DOCUMENT, + completed_at=LATER, + ) + + assert by_other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + assert completed_after == RunWrite(changed=False, status=RunStatus.FAILED) + row = seeder.row("simulation_runs", run_id) + assert row["household_result"] == FAILED_DOCUMENT + assert row["error_message"] == FAILURE_SUMMARY + assert row["completed_at"] == LATER + + +def test_a_household_simulation_run_failed_without_a_document_has_no_result( + store: HouseholdReportRunStore, + seeder: ReportRunSeeder, + seeded_report: SeededHouseholdReportRun, +) -> None: + run_id = seeded_report.baseline_simulation_run_id + store.claim_simulation_run(run_id, modal_call_id=CALL, claimed_at=NOW) + + failed = store.fail_household_simulation_run( + run_id, + modal_call_id=CALL, + error_message="household_function_error: The household function raised X.", + failed_at=LATER, + ) + + assert failed == RunWrite(changed=True, status=RunStatus.FAILED) + assert _result_is_null(seeder, run_id) + + +def test_a_household_report_run_completes_only_while_its_call_holds_it( + store: HouseholdReportRunStore, + seeder: ReportRunSeeder, + seeded_baseline_only: SeededHouseholdReportRun, +) -> None: + report_run_id = seeded_baseline_only.report_run_id + claimed = store.claim_report_run(report_run_id, modal_call_id=CALL, claimed_at=NOW) + + by_other_call = store.complete_household_report_run( + report_run_id, modal_call_id=OTHER_CALL, completed_at=LATER + ) + completed = store.complete_household_report_run( + report_run_id, modal_call_id=CALL, completed_at=LATER + ) + failed_after = store.fail_report_run( + report_run_id, + modal_call_id=CALL, + error_message="household_calculation_failed: late", + failed_at=LATER, + ) + + assert claimed == RunWrite(changed=True, status=RunStatus.RUNNING) + assert by_other_call == RunWrite(changed=False, status=RunStatus.RUNNING) + assert completed == RunWrite(changed=True, status=RunStatus.SUCCEEDED) + assert failed_after == RunWrite(changed=False, status=RunStatus.SUCCEEDED) + row = seeder.row("report_runs", report_run_id) + assert row["completed_at"] == LATER + assert row["error_message"] is None + assert row["baseline_cliff_gap"] is None + assert row["spm_provenance"] is None + + +def test_completing_a_missing_simulation_run_raises_lookup_error( + store: HouseholdReportRunStore, +) -> None: + missing_run_id = uuid4() + + with pytest.raises(LookupError, match=f"simulation run {missing_run_id}"): + store.complete_household_simulation_run( + missing_run_id, + modal_call_id=CALL, + household_result=SUCCEEDED_DOCUMENT, + completed_at=LATER, + ) diff --git a/projects/policyengine-apis-integ/tests/test_household_report_generated_client.py b/projects/policyengine-apis-integ/tests/test_household_report_generated_client.py new file mode 100644 index 000000000..f341d7c8b --- /dev/null +++ b/projects/policyengine-apis-integ/tests/test_household_report_generated_client.py @@ -0,0 +1,158 @@ +"""Hermetic wire checks of a household report through the generated client. + +The Public API sends a household report to ``POST /reports`` with +``report_type`` ``household``, and reads the 201 naming the report, its run, +its year and each side's simulation and simulation run, or the error +answers economy reports share. +""" + +import httpx +from policyengine_api_simulation_client import AuthenticatedClient +from policyengine_api_simulation_client.api.default import create_report +from policyengine_api_simulation_client.models import ( + EconomyError, + EconomyReportCreated, + EconomyRequestUnprocessable, + HouseholdReportCreated, + HouseholdReportRequest, + 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"] +HOUSEHOLD_ID = "00000000-0000-4000-8000-0000000000bb" +POLICY_ID = "00000000-0000-4000-8000-0000000000c5" +HOUSEHOLD_REPORT_BODY = { + "country": "us", + "report_type": "household", + "label": "A larger child tax credit for one family", + "description": None, + "household_id": HOUSEHOLD_ID, + "baseline": {"policy_id": None}, + "reform": {"policy_id": POLICY_ID}, +} +CREATED_HOUSEHOLD_REPORT = { + "report_id": REPORT_ID, + "report_run_id": RUN_ID, + "status": "pending", + "year": 2026, + "baseline_simulation_id": SIDE_IDS[0], + "baseline_simulation_run_id": SIDE_IDS[1], + "reform_simulation_id": SIDE_IDS[2], + "reform_simulation_run_id": SIDE_IDS[3], +} +# An economy report's answer, which the same 201 may carry. +CREATED_ECONOMY_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": False, + "reform_simulation_id": SIDE_IDS[2], + "reform_simulation_run_id": SIDE_IDS[3], + "reform_reused": False, +} + + +def _create(status_code: int, body: dict, request: dict = HOUSEHOLD_REPORT_BODY): + sent: list[dict] = [] + + def answer(http_request: httpx.Request) -> httpx.Response: + sent.append(httpx.Response(200, content=http_request.content).json()) + return 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=httpx.MockTransport(answer), + ) + ) + response = create_report.sync_detailed( + client=client, body=HouseholdReportRequest.from_dict(request) + ) + return response, sent + + +def test_a_household_report_request_is_sent_as_given(): + response, sent = _create(201, CREATED_HOUSEHOLD_REPORT) + + assert sent == [HOUSEHOLD_REPORT_BODY] + assert response.status_code == 201 + + +def test_a_report_of_the_baseline_alone_is_sent_without_a_reform(): + request = {**HOUSEHOLD_REPORT_BODY, "reform": None} + + _, sent = _create(201, CREATED_HOUSEHOLD_REPORT, request) + + assert sent[0]["reform"] is None + + +def test_a_created_household_report_is_read_with_its_year_and_sides(): + response, _ = _create(201, CREATED_HOUSEHOLD_REPORT) + + assert isinstance(response.parsed, HouseholdReportCreated) + assert response.parsed.status is ReportRunStatus.PENDING + assert response.parsed.year == 2026 + assert response.parsed.to_dict() == CREATED_HOUSEHOLD_REPORT + + +def test_a_created_report_of_the_baseline_alone_has_no_reform(): + created = { + **CREATED_HOUSEHOLD_REPORT, + "reform_simulation_id": None, + "reform_simulation_run_id": None, + } + + response, _ = _create(201, created) + + assert isinstance(response.parsed, HouseholdReportCreated) + assert response.parsed.reform_simulation_id is None + assert response.parsed.reform_simulation_run_id is None + + +def test_a_created_economy_report_is_still_read_as_one(): + response, _ = _create(201, CREATED_ECONOMY_REPORT) + + assert isinstance(response.parsed, EconomyReportCreated) + + +def test_the_household_refusals_are_read_with_their_codes(): + refused, _ = _create( + 400, + { + "code": "unknown_household", + "detail": "us has no stored household the request names.", + }, + ) + unprocessable, _ = _create( + 422, + { + "detail": [ + { + "loc": ["body", "household_id"], + "msg": "the household has no default year to calculate", + "type": "household_without_default_year", + } + ], + "code": "household_not_calculable", + }, + ) + call_refused, _ = _create( + 503, + { + "code": "modal_call_refused", + "detail": "Modal did not accept the report's call.", + "report_id": REPORT_ID, + }, + ) + + assert isinstance(refused.parsed, EconomyError) + assert refused.parsed.code == "unknown_household" + assert isinstance(unprocessable.parsed, EconomyRequestUnprocessable) + assert unprocessable.parsed.code == "household_not_calculable" + assert isinstance(call_refused.parsed, EconomyError) + assert str(call_refused.parsed.report_id) == REPORT_ID diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index a0b188b5c..0d1408c33 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -188,3 +188,40 @@ 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. +## Household reports + +`POST /reports` with `report_type` `household` creates a household report: +one household the Public API stores, for its `default_year`, under a +baseline and, optionally, a reform. It names the `country`, `label`, +`description`, `household_id`, `baseline` and `reform` (each a stored +`policy_id`, null for current law; a null or omitted `reform` makes a report +of the baseline alone). It names no year, region, dataset, dynamic, .py +version or record ID, and has no options: each side adds every variable the +household lacks, as the Public API's stored-household route does, so the +report holds every value, and `reports.inputs` is `{}`. A household's saved +SPM selection is applied to both sides. + +The service reads the catalog version for the default worker's .py version, +the household and each side's policy, then creates, in one transaction, a +household simulation and a pending simulation run for each side (household +runs are never reused), the report (type `household`) and its pending +report run. It starts one call of the country's household function with the +report run's ID (`household_report_run_id`) and each side's household +contract request, whose calculation ID is the side's simulation run, and +records the call's ID. The household function claims the report run and each +simulation run under its own call, writes each run's result or failure +document to `household_result`, and marks the report run succeeded only if +every side's run did; otherwise it fails with the first such side's reason. +The answer is 201 with the report's and its run's IDs, the run's status, +the `year` and each side's simulation and simulation run. + +Refusals and outages answer as an economy report's do, with three more +codes: 400 `unknown_household` and `no_household_function` (the default +worker has no household function for the country), and 422 +`household_not_calculable` when the household has no default year or +cannot be read, located at `household_id`, or when the household contract +refuses what the report would calculate, located at the field it traces to: +`household_id` or a side's `policy_id`. +A policy setting parameters the catalog version lacks is 422 +`missing_parameters`, located at that side's `policy_id`. + 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 a91636568..c14a96dc5 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -58,6 +58,7 @@ from policyengine_simulation_entry.config import Settings from policyengine_simulation_entry.economy_routes import ( EconomySubmissions, + HouseholdReportSubmissions, add_economy_routes, ) from policyengine_simulation_entry.economy_submissions import ( @@ -83,6 +84,9 @@ from policyengine_simulation_entry.household_calculations import ( HouseholdCalculationService, ) +from policyengine_simulation_entry.household_report_submissions import ( + household_reports_for, +) from policyengine_simulation_entry.household_routes import ( HouseholdCalculations, add_household_calculation_routes, @@ -540,16 +544,24 @@ def household_calculations_for( ) -def _register_economy_routes( +def _register_economy_routes( # noqa: PLR0913 # the app, its services and permissions app: FastAPI, entry: EntryRuntime, submissions: EconomySubmissions | None, + household_reports: HouseholdReportSubmissions | None, *, submitter: AuthorizationDependency, reader: AuthorizationDependency, ) -> None: - """Add the routes that create economy reports and simulations and read them.""" - add_economy_routes(app, entry, submissions, submitter=submitter, reader=reader) + """Add the routes that create reports and economy simulations and read them.""" + add_economy_routes( + app, + entry, + submissions, + submitter=submitter, + reader=reader, + household_reports=household_reports, + ) def economy_submissions_for( @@ -744,6 +756,7 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: app, entry, economy_submissions_for(runtime_settings, runtime_database), + household_reports_for(runtime_settings, runtime_database), submitter=simulation_submitter, reader=simulation_reader, ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py index 009b9e537..467678e71 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/budget_window_api_models.py @@ -25,7 +25,6 @@ from policyengine_simulation_entry.economy_api_models import ( CatalogKey, - EconomyReportRequest, PolicyChoice, ReportDescription, ReportLabel, @@ -104,14 +103,6 @@ def options(self) -> dict[str, Any]: return options -# The body of ``POST /reports``: an economy report or a budget window, told -# apart by ``report_type``. -ReportRequest = Annotated[ - EconomyReportRequest | BudgetWindowReportRequest, - Field(discriminator="report_type"), -] - - class _Answer(BaseModel): model_config = ConfigDict(extra="forbid", frozen=True) 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 84caea1b3..cd0e71b0a 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -59,10 +59,21 @@ "dataset_not_in_bundle", # The database refused a row: a row it names was deleted after the reads. "missing_reference", + # A household report's household is not one of the country's. + "unknown_household", + # The default worker has no household function for the country. + "no_household_function", ] EconomyRefusalCode = Literal[CallerRowRefusalCode, ServiceStateRefusalCode] # Why stored values cannot be calculated (422). -EconomyUnprocessableCode = Literal["missing_parameters", "overlapping_parameters"] +EconomyUnprocessableCode = Literal[ + "missing_parameters", + "overlapping_parameters", + # A household report's stored household has no default year or cannot + # be read, or the household functions' contract refuses what the report + # would calculate: located at the household or a side's policy. + "household_not_calculable", +] # Why a route cannot answer now (503); retry after ``Retry-After`` seconds. EconomyUnavailableCode = Literal[ # The service has no database or v2 Modal environment configured. 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 1a04277cf..129ef37d1 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 @@ -6,7 +6,12 @@ and each side's simulation and simulation run. With ``report_type`` ``budget_window`` it creates a budget-window report with a report run per year and starts up to its first five years, answering 201 with the report, the - first year's run, the window run, its status and every year's runs. + first year's run, the window run, its status and every year's runs. With + ``household`` it creates a household report, a household simulation and + run for its baseline and any reform, and its first run, and starts one + household function call for the report run, answering 201 with the + report, its run, its status, its year 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. @@ -44,7 +49,6 @@ from policyengine_simulation_entry.budget_window_api_models import ( BudgetWindowReportCreated, BudgetWindowReportRequest, - ReportRequest, ) from policyengine_simulation_entry.economy_api_models import ( EconomyError, @@ -64,6 +68,11 @@ EconomySubmissionUnavailableError, EconomySubmissionUnprocessableError, ) +from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportCreated, + HouseholdReportRequest, +) +from policyengine_simulation_entry.report_api_models import ReportRequest from policyengine_simulation_entry.report_request_errors import ReportRequestRoute from policyengine_simulation_entry.route_support import ( log_unexpected_failure, @@ -112,7 +121,9 @@ "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 " + "the caller's own region, dataset or dynamic; a household report's " + "unknown_household and no_household_function name its household and " + "the default worker's household function; any other code is the " "service's state, such as a catalog version missing while a release " "deploys" ), @@ -122,8 +133,9 @@ "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" + "be calculated carry a code: missing_parameters, " + "overlapping_parameters, or household_not_calculable for a household " + "report's household" ), "model": EconomyRequestUnprocessable, } @@ -210,6 +222,19 @@ async def read_simulation(self, simulation_id: UUID) -> StoredSimulation | None: ... +class HouseholdReportSubmissions(Protocol): + """Create household reports (``HouseholdReportSubmissionService``).""" + + async def submit_report( + self, + request: HouseholdReportRequest, + *, + observability_context: dict[str, Any] | None, + ) -> HouseholdReportCreated: + """Create the report and its runs and start its household function call.""" + ... + + def _error_response(status_code: int, error: EconomyError) -> JSONResponse: headers = ( {"Retry-After": str(RETRY_AFTER_SECONDS)} @@ -223,10 +248,12 @@ def _error_response(status_code: int, error: EconomyError) -> JSONResponse: ) -def _not_configured() -> JSONResponse: +def _not_configured( + detail: str = "Economy runs are not configured.", +) -> JSONResponse: return _error_response( status.HTTP_503_SERVICE_UNAVAILABLE, - EconomyError(code="not_configured", detail="Economy runs are not configured."), + EconomyError(code="not_configured", detail=detail), ) @@ -382,19 +409,28 @@ def _unknown(record: RecordName, record_id: UUID) -> JSONResponse: ) -def add_economy_routes( +def add_economy_routes( # noqa: PLR0913 # the app, its services and permissions app: FastAPI, entry: EntryRuntime, submissions: EconomySubmissions | None, *, submitter: AuthorizationDependency, reader: AuthorizationDependency, + household_reports: HouseholdReportSubmissions | None = None, ) -> None: - """Add the economy routes; without a service they answer 503.""" + """Add the economy routes; without a service they answer 503. + + ``POST /reports`` also takes household reports, which answer 503 + without ``household_reports``. + """ 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 + entry, + submissions, + household_reports, + submitter=submitter, + log_string_limit=log_string_limit, ) ) app.include_router( @@ -402,9 +438,75 @@ def add_economy_routes( ) +async def _economy_report_answer( + entry: EntryRuntime, + submissions: EconomySubmissions | None, + body: EconomyReportRequest | BudgetWindowReportRequest, + request: Request, + *, + log_string_limit: int, +) -> JSONResponse: + """Create an economy or budget-window report; answer as ``POST /reports`` does.""" + if submissions is None: + return _not_configured() + observability_id = entry.start_workflow_observability(request) + try: + call_request_id = call_request_id_for(request) + context = observability_context(entry, observability_id) + created: EconomyReportCreated | BudgetWindowReportCreated + if isinstance(body, BudgetWindowReportRequest): + created = await submissions.submit_budget_window( + body, request_id=call_request_id, observability_context=context + ) + else: + created = await submissions.submit_report( + body, request_id=call_request_id, observability_context=context + ) + except _SUBMISSION_FAILURES as error: + return _submission_failure_response(error, "report") + # Every other failure is still answered as the JSON 500 declared. + 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 + ) + + +async def _household_report_answer( + entry: EntryRuntime, + household_reports: HouseholdReportSubmissions | None, + body: HouseholdReportRequest, + request: Request, + *, + log_string_limit: int, +) -> JSONResponse: + """Create a household report and answer as ``POST /reports`` does.""" + if household_reports is None: + return _not_configured("Household reports are not configured.") + observability_id = entry.start_workflow_observability(request) + try: + created = await household_reports.submit_report( + body, observability_context=observability_context(entry, observability_id) + ) + except _SUBMISSION_FAILURES as error: + return _submission_failure_response(error, "report") + # Every other failure, a request the service built wrong included, 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 + ) + + def _submission_routes( entry: EntryRuntime, submissions: EconomySubmissions | None, + household_reports: HouseholdReportSubmissions | None, *, submitter: AuthorizationDependency, log_string_limit: int, @@ -416,7 +518,7 @@ def _submission_routes( @report_router.post( REPORTS_PATH, - summary="Create an economy or budget-window report", + summary="Create an economy, budget-window or household 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 " @@ -425,11 +527,17 @@ def _submission_routes( "A budget-window report (report_type budget_window) compares them " "in each year of a window of 1 to 75 years, with a report run per " "year, at most five running at once; each year reuses its own " - "year's runs." + "year's runs. " + "A household report (report_type household) calculates one stored " + "household for its default year under a baseline and, optionally, " + "a reform, each a new household simulation run, all in one " + "household function call." ), operation_id="create_report", status_code=status.HTTP_201_CREATED, - response_model=EconomyReportCreated | BudgetWindowReportCreated, + response_model=( + EconomyReportCreated | BudgetWindowReportCreated | HouseholdReportCreated + ), responses={ 201: {"description": "The report and its first runs were created"}, 400: _REFUSED, @@ -440,34 +548,17 @@ def _submission_routes( }, dependencies=[Depends(submitter)], ) - async def create_report( - body: ReportRequest, - request: Request, - ) -> JSONResponse: - if submissions is None: - return _not_configured() - observability_id = entry.start_workflow_observability(request) - try: - call_request_id = call_request_id_for(request) - context = observability_context(entry, observability_id) - created: EconomyReportCreated | BudgetWindowReportCreated - if isinstance(body, BudgetWindowReportRequest): - created = await submissions.submit_budget_window( - body, request_id=call_request_id, observability_context=context - ) - else: - created = await submissions.submit_report( - body, request_id=call_request_id, observability_context=context - ) - except _SUBMISSION_FAILURES as error: - return _submission_failure_response(error, "report") - # Every other failure is still answered as the JSON 500 declared. - except Exception as error: # noqa: BLE001 # answered as internal_error - return _unexpected_failure_response( - REPORTS_PATH, error, max_characters=log_string_limit + async def create_report(body: ReportRequest, request: Request) -> JSONResponse: + if isinstance(body, HouseholdReportRequest): + return await _household_report_answer( + entry, + household_reports, + body, + request, + log_string_limit=log_string_limit, ) - return JSONResponse( - created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED + return await _economy_report_answer( + entry, submissions, body, request, log_string_limit=log_string_limit ) router.include_router(report_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 index be55aed4b..a284af30e 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 @@ -49,6 +49,7 @@ "dataset_id": ("unknown_dataset", "dataset"), "dynamic_id": ("unknown_dynamic", "dynamic"), "policy_id": ("unknown_policy", "policy"), + "household_id": ("unknown_household", "household"), "tax_benefit_model_version_id": ("unknown_catalog_version", "catalog version"), "tax_benefit_model_id": ("unknown_catalog_version", "tax-benefit model"), } diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_report_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_report_api_models.py new file mode 100644 index 000000000..11e613c9c --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_report_api_models.py @@ -0,0 +1,88 @@ +"""Request and answer models of a household report on ``POST /reports``. + +A household report calculates one household the Public API stores, for the +household's ``default_year``, under a baseline policy and, optionally, a +reform, each a stored policy or current law. The Simulation Entrypoint reads +the household and the policies from the Public API's tables itself and +calculates on the v2 version manifest's default worker; a request names no +.py version, no year and no record ID. Its answers and their error codes +are an economy report's (``economy_api_models``), with the household codes +``unknown_household``, ``no_household_function`` and +``household_not_calculable``. +""" + +from __future__ import annotations + +from typing import Any, Literal +from uuid import UUID + +from policyengine_simulation_contract.stage12_bundle import CountryId +from policyengine_simulation_persistence import ReportRunStatus +from pydantic import BaseModel, ConfigDict + +from policyengine_simulation_entry.economy_api_models import ( + ReportDescription, + ReportLabel, +) + + +class _Request(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False) + + +class HouseholdPolicyChoice(_Request): + """A side's stored policy; null or omitted means current law.""" + + policy_id: UUID | None = None + + +class HouseholdReportRequest(_Request): + """A household report: one stored household under a baseline and a reform. + + ``reform`` is null, or omitted, for a report of the baseline alone. The + report calculates the household's ``default_year``. Each side's + calculation adds every yearly, monthly and eternal variable the + household lacks, as the Public API's stored-household route does, so the + report holds every value; a household report has no option for it. A + household's saved SPM selection is applied, or the bundle's certified + defaults without one. + """ + + country: CountryId + report_type: Literal["household"] + label: ReportLabel + description: ReportDescription | None = None + household_id: UUID + baseline: HouseholdPolicyChoice + reform: HouseholdPolicyChoice | None = None + + def options(self) -> dict[str, Any]: + """Return the report's options as it stores them (``reports.inputs``). + + A household report has none: the column must hold an object. + """ + return {} + + +class HouseholdReportCreated(BaseModel): + """The household 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 + whatever the household 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``. ``year`` is the + household's default year, which the report calculates. Each side names + its new simulation and simulation run; the reform's are null for a + report of the baseline alone. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + report_id: UUID + report_run_id: UUID + status: ReportRunStatus + year: int + baseline_simulation_id: UUID + baseline_simulation_run_id: UUID + reform_simulation_id: UUID | None + reform_simulation_run_id: UUID | None diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_report_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_report_submissions.py new file mode 100644 index 000000000..987be2b7d --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_report_submissions.py @@ -0,0 +1,693 @@ +"""Create household reports, and start the household function call that calculates them. + +A household report submission names a stored household, the baseline's +stored policy and, optionally, the reform's. The service: + +1. finds the v2 version manifest's default worker's household function for + the country, and its bundle; +2. reads, in one transaction, the country's catalog version for that + worker's .py version, its model, the stored household and each side's + policy values, and builds each side's household contract request for + the household's ``default_year``, with the side's new simulation run ID + as its calculation ID; +3. creates, in one transaction, each side's household simulation and its + pending simulation run, the report, which has no options, and its first + report run, pending with no Modal call ID; +4. starts one household function call with ``spawn``, passing the report + run's ID and both sides' requests, then records the call's ID while the + report run is still pending with none (the function may have claimed it + already). When Modal refuses the call, or does not confirm it in time, it + fails the report run and its simulation runs at once instead, under the + same condition; a run the function claimed meanwhile, because Modal + accepted the call after all, is left as it is and answered ``pending``. + +Whatever is missing or cannot be calculated is refused in steps 1 and 2, +before anything is written. The household function then claims the report +run and each simulation run, writes each run's household document, and +finishes the report run; the stuck-run sweep fails a pending run no call +was recorded for. The errors are the economy routes' +(``economy_submission_errors``), so ``POST /reports`` answers both kinds +alike. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID, uuid4 + +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, +) +from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportRunCalculation, +) +from policyengine_simulation_persistence import ( + EconomyRecordOutcomeUnknownError, + EconomyRecordReferenceError, + HouseholdReportSimulation, + HouseholdReportStore, + MissingPolicyParametersError, + NewHouseholdReport, + ReportRunStatus, + ReportSideError, + RunVersions, + UnknownCatalogVersionError, + UnknownHouseholdError, + UnknownPolicyError, +) +from pydantic import ValidationError +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, + UnprocessableReason, + reference_refusal, +) +from policyengine_simulation_entry.household_contract_fields import ( + job_policy_fields, + request_field_of, +) +from policyengine_simulation_entry.household_conversion import ( + UnreadableHouseholdError, + household_situation, + saved_spm_selection, +) +from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportCreated, +) +from policyengine_simulation_entry.household_workers import ( + HouseholdFunctionUnavailableError, + ManifestHouseholdWorkers, + ModalHouseholdFunctionCaller, + VersionManifestUnavailableError, +) +from policyengine_simulation_entry.simulation_database import SimulationDatabase + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping, Sequence + + from policyengine_simulation_persistence import ( + HouseholdReportInputs, + PolicyReform, + ReportSide, + StoredHousehold, + ) + + from policyengine_simulation_entry.config import Settings + from policyengine_simulation_entry.economy_submission_errors import RecordName + from policyengine_simulation_entry.entry_runtime import ( + SimulationDatabaseConnection, + ) + from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportRequest, + ) + from policyengine_simulation_entry.household_workers import ( + HouseholdReportRunCaller, + HouseholdWorker, + HouseholdWorkerSource, + ) + +logger = logging.getLogger(__name__) + +# How long a submission waits for Modal to accept the household function's +# call before failing the report run. +CALL_START_TIMEOUT_SECONDS = 10.0 +# How long one of the service's statements may run, as a household job's. +STATEMENT_TIMEOUT_MILLISECONDS = 5_000 + + +class HouseholdReports(Protocol): + """The household report reads and writes (``HouseholdReportStore``).""" + + def read_inputs( + self, + *, + country: str, + household_id: UUID, + policy_ids: Mapping[ReportSide, UUID | None], + policyengine_version: str, + ) -> HouseholdReportInputs: + """Read the catalog version, household and each side's policy together.""" + ... + + def create_report(self, report: NewHouseholdReport) -> None: + """Create the report, its simulations and its runs.""" + ... + + 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 simulation runs, while unstarted.""" + ... + + +class HouseholdReportBuildError(Exception): + """The household function's contract refused a field the service built. + + The message names each refused field, never a value; the route answers + it as an internal error. + """ + + +def _database_unavailable() -> EconomySubmissionUnavailableError: + return EconomySubmissionUnavailableError( + code="database_unavailable", detail="The simulation database could not be read." + ) + + +def _not_calculable( + *reasons: UnprocessableReason, +) -> EconomySubmissionUnprocessableError: + return EconomySubmissionUnprocessableError( + code="household_not_calculable", reasons=reasons + ) + + +def _policy_ids(request: HouseholdReportRequest) -> dict[ReportSide, UUID | None]: + """Return each side the report calculates and its stored policy, if any.""" + policy_ids: dict[ReportSide, UUID | None] = {"baseline": request.baseline.policy_id} + if request.reform is not None: + policy_ids["reform"] = request.reform.policy_id + return policy_ids + + +def _read_refusal( + error: Exception, request: HouseholdReportRequest, policyengine_version: str +) -> EconomySubmissionRefusedError | EconomySubmissionUnprocessableError | None: + """Return the refusal a store read's error means, or None for another error. + + Every detail is written from the request, never from the error's text. + """ + if isinstance(error, UnknownCatalogVersionError): + return EconomySubmissionRefusedError( + code="unknown_catalog_version", + detail=( + f"{request.country} has no catalog version for .py " + f"{policyengine_version}, the default v2 worker's version." + ), + ) + if isinstance(error, UnknownHouseholdError): + return EconomySubmissionRefusedError( + code="unknown_household", + detail=f"{request.country} has no stored household the request names.", + ) + if not isinstance(error, ReportSideError): + return None + side, refused = error.side, error.error + if isinstance(refused, UnknownPolicyError): + return EconomySubmissionRefusedError( + code="unknown_policy", + detail=( + f"{request.country} has no stored policy the request's {side} names." + ), + ) + if isinstance(refused, MissingPolicyParametersError): + return EconomySubmissionUnprocessableError( + code="missing_parameters", + reasons=tuple( + UnprocessableReason( + loc=("body", side, "policy_id", name), + msg=( + f"the catalog version for .py {policyengine_version} has " + "no such parameter" + ), + type="missing_parameter", + ) + for name in refused.parameter_names + ), + ) + return None + + +def _spm_fields(document: Mapping[str, Any]) -> dict[str, object]: + """Return the contract's SPM fields: the household's saved selection, if any.""" + saved = saved_spm_selection(document) + if saved is None: + return {} + return {"spm": saved, "spm_selection_source": "saved_household"} + + +def _contract_reasons( + error: ValidationError, side: ReportSide +) -> tuple[UnprocessableReason, ...]: + """Name each field the contract refuses by the request field it traces to. + + Raise HouseholdReportBuildError when the contract refuses a field the + service builds itself, which is no fault of the caller's. + """ + reasons = [] + built_wrong = [] + for item in error.errors(): + location = tuple(item["loc"]) + request_field = request_field_of(location, spm_from_request=False) + # A household report sets add_missing_variables itself. + if request_field is None or request_field == "add_missing_variables": + built_wrong.append(".".join(str(part) for part in location) or "request") + continue + field_location = ( + (side, request_field) if request_field == "policy_id" else (request_field,) + ) + reasons.append( + UnprocessableReason( + loc=("body", *field_location, *location), + msg=item["msg"], + type=item["type"], + ) + ) + if built_wrong: + msg = ( + "the household function's request has fields the service built " + f"wrong: {', '.join(built_wrong)}" + ) + raise HouseholdReportBuildError(msg) + return tuple(reasons) + + +@dataclass(frozen=True) +class HouseholdReportSideInputs: + """One side of a household report run: the simulation run it records. + + ``policy_values`` holds the side's policy values by parameter name and + period, empty for current law, as ``HouseholdReportStore.read_inputs`` + reads them into ``reforms``. + """ + + simulation_run_id: UUID + policy_values: PolicyReform + + +def household_report_calculation( + *, + household: StoredHousehold, + year: int, + baseline: HouseholdReportSideInputs, + reform: HouseholdReportSideInputs | None = None, + worker: HouseholdWorker, +) -> HouseholdReportRunCalculation: + """Build a household report run's household function payload from its inputs. + + Each side's request calculates the stored ``household``, in its + country, for ``year`` with the household's saved SPM selection, on + ``worker``'s bundle, and its calculation ID is the side's simulation + run; ``reform`` is None for a report of the baseline alone. Every side + adds the household's missing variables, so a household report holds + every value. Creating a report passes the household's default year. A + restart or re-run passes the stored report's year, and the household + and policy values as they are now read. + + Raise EconomySubmissionUnprocessableError (``household_not_calculable``) + when the stored household cannot be read, or the household contract + refuses what would be calculated, the baseline's checked first. Its + reasons are located at a report request's fields (``household_id``, a + side's ``policy_id``); a caller rebuilding a stored run reads them as + its own. Raise HouseholdReportBuildError when the contract refuses a + field the service builds, two sides recording one simulation run + included. + """ + country = household.country + try: + situation = household_situation(country, household.document) + except UnreadableHouseholdError: + raise _not_calculable( + UnprocessableReason( + loc=("body", "household_id"), + msg="the stored household's people or groups cannot be read", + type="household_unreadable", + ) + ) from None + sides: dict[ReportSide, HouseholdReportSideInputs] = {"baseline": baseline} + if reform is not None: + sides["reform"] = reform + requests = { + side: _side_request( + side, + side_inputs, + year=year, + situation=situation, + household=household, + worker=worker, + ) + for side, side_inputs in sides.items() + } + try: + return HouseholdReportRunCalculation( + baseline=requests["baseline"], reform=requests.get("reform") + ) + except ValidationError as error: + msg = "the household report run's calculation was built wrong: " + "; ".join( + str(item["msg"]) for item in error.errors() + ) + raise HouseholdReportBuildError(msg) from None + + +def _side_request( # noqa: PLR0913 # the side and what it calculates + side: ReportSide, + side_inputs: HouseholdReportSideInputs, + *, + year: int, + situation: dict[str, Any], + household: StoredHousehold, + worker: HouseholdWorker, +) -> HouseholdExecutionRequest: + """Build one side's household contract request; raise its refusal.""" + try: + return HouseholdExecutionRequest.model_validate( + { + "calculation_id": side_inputs.simulation_run_id, + "country": household.country, + # The report's year, as a situation's period key. + "year": str(year), + "household": situation, + **job_policy_fields(reform=side_inputs.policy_values), + # A household report holds every value, as the Public API's + # stored-household route gives it. + "add_missing_variables": True, + "bundle": worker.bundle, + **_spm_fields(household.document), + } + ) + except ValidationError as error: + raise _not_calculable(*_contract_reasons(error, side)) from None + + +def _report_year(household: StoredHousehold) -> int: + """Return the year a new report calculates: the household's default year.""" + if household.default_year is None: + raise _not_calculable( + UnprocessableReason( + loc=("body", "household_id"), + msg="the household has no default year to calculate", + type="household_without_default_year", + ) + ) + return household.default_year + + +class HouseholdReportSubmissionService: + """Create household reports and start the household function call for each. + + The database and the manifest are read on threads, so the event loop + never waits on them. + """ + + def __init__( + self, + *, + store: HouseholdReports, + workers: HouseholdWorkerSource, + caller: HouseholdReportRunCaller, + call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, + new_id: Callable[[], UUID] = uuid4, + ) -> None: + """Use these seams; tests replace each of them.""" + self._store = store + self._workers = workers + self._caller = caller + self._call_start_timeout_seconds = call_start_timeout_seconds + self._new_id = new_id + + async def submit_report( + self, + request: HouseholdReportRequest, + *, + observability_context: dict[str, Any] | None, + ) -> HouseholdReportCreated: + """Create the report and its runs, and start its household function call. + + Raise EconomySubmissionRefusedError (400) or + EconomySubmissionUnprocessableError (422) before writing anything; + EconomySubmissionUnavailableError (503) when the manifest or the + database cannot be read, or, naming the report, when Modal doesn't + confirm the call and the report run was failed; + EconomySubmissionOutcomeUnknownError (500) when the create's commit + failed after its rows were inserted; and HouseholdReportBuildError + (500) when the service built a request the contract refuses. + """ + worker = await self._default_worker(request.country) + inputs = await self._read_inputs(request, worker) + report, calculation = self._prepare(request, worker, inputs) + await self._create(report) + status = await self._start_report_run( + worker, report, calculation, observability_context=observability_context + ) + return HouseholdReportCreated( + report_id=report.id, + report_run_id=report.run_id, + status=status, + year=report.year, + baseline_simulation_id=report.baseline.simulation_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_id=( + None if report.reform is None else report.reform.simulation_id + ), + reform_simulation_run_id=( + None if report.reform is None else report.reform.simulation_run_id + ), + ) + + async def _default_worker(self, country: str) -> HouseholdWorker: + try: + return await asyncio.to_thread(self._workers.default_worker, country) + except HouseholdFunctionUnavailableError: + raise EconomySubmissionRefusedError( + code="no_household_function", + detail=( + f"The default v2 worker has no household function for {country}." + ), + ) from None + except VersionManifestUnavailableError as error: + logger.warning( + "household_report_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( + self, request: HouseholdReportRequest, worker: HouseholdWorker + ) -> HouseholdReportInputs: + policyengine_version = worker.bundle.policyengine_version + try: + return await asyncio.to_thread( + lambda: self._store.read_inputs( + country=request.country, + household_id=request.household_id, + policy_ids=_policy_ids(request), + policyengine_version=policyengine_version, + ) + ) + except SQLAlchemyError as error: + logger.warning( + "household_report_inputs_unreadable", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + except Exception as error: + refusal = _read_refusal(error, request, policyengine_version) + if refusal is None: + raise + raise refusal from None + + def _prepare( + self, + request: HouseholdReportRequest, + worker: HouseholdWorker, + inputs: HouseholdReportInputs, + ) -> tuple[NewHouseholdReport, HouseholdReportRunCalculation]: + """Return the report to create and the calculation its call is given.""" + household = inputs.household + year = _report_year(household) + report_id, run_id = self._new_id(), self._new_id() + sides: dict[ReportSide, HouseholdReportSimulation] = { + side: HouseholdReportSimulation( + simulation_id=self._new_id(), + simulation_run_id=self._new_id(), + policy_id=policy_id, + ) + for side, policy_id in _policy_ids(request).items() + } + side_inputs = { + side: HouseholdReportSideInputs( + simulation_run_id=created.simulation_run_id, + policy_values=inputs.reforms[side], + ) + for side, created in sides.items() + } + calculation = household_report_calculation( + household=household, + year=year, + baseline=side_inputs["baseline"], + reform=side_inputs.get("reform"), + worker=worker, + ) + report = NewHouseholdReport( + id=report_id, + run_id=run_id, + country=request.country, + label=request.label, + description=request.description, + tax_benefit_model_id=inputs.tax_benefit_model_id, + tax_benefit_model_version_id=inputs.catalog_version_id, + household_id=request.household_id, + year=year, + versions=RunVersions( + country_package_version=worker.bundle.country_package_version, + policyengine_version=worker.bundle.policyengine_version, + ), + inputs=request.options(), + baseline=sides["baseline"], + reform=sides.get("reform"), + ) + return report, calculation + + async def _create(self, report: NewHouseholdReport) -> 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(self._store.create_report, report) + except EconomyRecordReferenceError as error: + raise reference_refusal(error) from None + except EconomyRecordOutcomeUnknownError as error: + # The route logs it, with the commit error's type. + record: RecordName = "report" + raise EconomySubmissionOutcomeUnknownError( + record, error_type=type(error.__cause__).__name__ + ) from None + except SQLAlchemyError as error: + logger.warning( + "household_report_not_created", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + + async def _start_report_run( + self, + worker: HouseholdWorker, + report: NewHouseholdReport, + calculation: HouseholdReportRunCalculation, + *, + observability_context: dict[str, Any] | None, + ) -> ReportRunStatus: + try: + call_id = await asyncio.wait_for( + self._caller.start_report_run( + worker, + calculation.model_dump(mode="json"), + report_run_id=report.run_id, + 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 + if await self._fail_unconfirmed(report, error): + raise EconomySubmissionUnavailableError( + code="modal_call_refused", + detail=( + "Modal did not accept the report's call; the report was " + "created and its run failed." + ), + record_id=report.id, + ) from None + return ReportRunStatus.PENDING + try: + status = await asyncio.to_thread( + self._store.record_report_call, report.run_id, call_id + ) + except SQLAlchemyError as error: + # The household function claims a pending run with no call ID + # itself, so the calculation goes ahead either way. + logger.warning( + "household_report_call_not_recorded", + extra={ + "report_id": str(report.id), + "error_type": type(error).__name__, + }, + ) + return ReportRunStatus.PENDING + return ReportRunStatus.PENDING if status is None else status + + async def _fail_unconfirmed( + self, report: NewHouseholdReport, error: BaseException + ) -> bool: + """Fail the report run of a call Modal didn't confirm, unless it 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 report 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( + "household_report_call_unconfirmed", + extra={"report_id": str(report.id), "error_type": type(error).__name__}, + ) + message = ( + "modal_call_refused: Modal did not accept the report's call " + f"({type(error).__name__})." + ) + try: + return await asyncio.to_thread( + self._store.fail_unstarted_report, + report.run_id, + report.simulation_run_ids, + message, + ) + except SQLAlchemyError as write_error: + logger.warning( + "household_report_run_not_failed", + extra={ + "report_id": str(report.id), + "error_type": type(write_error).__name__, + }, + ) + return False + + +def household_reports_for( + settings: Settings, database: SimulationDatabaseConnection | None +) -> HouseholdReportSubmissionService | None: + """Build the household report service as the household calculation one is. + + It needs the simulation database and the v2 Modal environment; without + them household reports answer 503. Nothing is read from Modal until the + first submission. + """ + if not settings.household_calculations_configured or not isinstance( + database, SimulationDatabase + ): + return None + return HouseholdReportSubmissionService( + store=HouseholdReportStore.from_engine( + database.engine, + statement_timeout_milliseconds=STATEMENT_TIMEOUT_MILLISECONDS, + ), + workers=ManifestHouseholdWorkers.from_modal( + manifest_name=settings.v2_manifest_name, + environment=settings.v2_modal_environment, + ), + caller=ModalHouseholdFunctionCaller(settings.v2_modal_environment), + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py index bf08e7668..0ec81bcd4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_workers.py @@ -72,6 +72,21 @@ async def spawn( ... +class HouseholdReportRunCaller(Protocol): + """Start one household report run's call (``ModalHouseholdFunctionCaller``).""" + + async def start_report_run( + self, + worker: HouseholdWorker, + payload: dict[str, Any], + *, + report_run_id: UUID, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the call without waiting for it and return its Modal call ID.""" + ... + + def default_household_worker( manifest: V2VersionManifest, country: str ) -> HouseholdWorker: @@ -156,16 +171,43 @@ async def spawn( Raise RuntimeError when Modal returns no call ID. """ + return await self._spawn( + worker, + payload, + household_job_id=str(household_job_id), + observability_context=observability_context, + ) + + async def start_report_run( + self, + worker: HouseholdWorker, + payload: dict[str, Any], + *, + report_run_id: UUID, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the household function on a household report run's calculation. + + ``payload`` is the report run's ``HouseholdReportRunCalculation``. + Return the call's Modal call ID; raise RuntimeError when Modal + returns none. + """ + return await self._spawn( + worker, + payload, + household_report_run_id=str(report_run_id), + observability_context=observability_context, + ) + + async def _spawn( + self, worker: HouseholdWorker, payload: dict[str, Any], **keywords: object + ) -> str: function = modal.Function.from_name( worker.application_name, worker.function_name, environment_name=self._environment, ) - call = await function.spawn.aio( - payload, - household_job_id=str(household_job_id), - observability_context=observability_context, - ) + call = await function.spawn.aio(payload, **keywords) call_id = getattr(call, "object_id", None) if not isinstance(call_id, str) or not call_id: msg = "Modal returned no call ID for the household function's call" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_api_models.py new file mode 100644 index 000000000..025be1ec2 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_api_models.py @@ -0,0 +1,24 @@ +"""The body of ``POST /reports``: an economy, budget-window or household report. + +The three are told apart by ``report_type``; each model's own module +describes its fields and its answer. +""" + +from __future__ import annotations + +from typing import Annotated + +from pydantic import Field + +from policyengine_simulation_entry.budget_window_api_models import ( + BudgetWindowReportRequest, +) +from policyengine_simulation_entry.economy_api_models import EconomyReportRequest +from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportRequest, +) + +ReportRequest = Annotated[ + EconomyReportRequest | BudgetWindowReportRequest | HouseholdReportRequest, + Field(discriminator="report_type"), +] diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py index bd574462f..53bd52190 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/report_request_errors.py @@ -1,7 +1,7 @@ -"""Refused ``POST /reports`` bodies, located as before the budget-window type. +"""Refused ``POST /reports`` bodies, located as before the report types. -``POST /reports`` takes an economy report or a budget window, told apart by -``report_type``. Pydantic puts the matched type into every refused field's +``POST /reports`` takes an economy, budget-window or household report, told +apart by ``report_type``. Pydantic puts the matched type into every refused field's location (``["body", "economy", "year"]``) and reports a missing or unknown ``report_type`` at the body itself. The Public API passes the locations to its callers, so the report route answers them as it did when it took only @@ -17,7 +17,7 @@ from fastapi.exceptions import RequestValidationError from fastapi.routing import APIRoute -from policyengine_simulation_entry.budget_window_api_models import ReportRequest +from policyengine_simulation_entry.report_api_models import ReportRequest if TYPE_CHECKING: from collections.abc import Callable, Coroutine, Sequence diff --git a/projects/policyengine-simulation-entry/tests/household_report_fakes.py b/projects/policyengine-simulation-entry/tests/household_report_fakes.py new file mode 100644 index 000000000..851afa294 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/household_report_fakes.py @@ -0,0 +1,240 @@ +"""Fakes of the household report service's store, manifest and Modal call. + +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 itertools import count +from typing import TYPE_CHECKING, Any +from uuid import UUID + +from household_fixtures import household_worker_version +from policyengine_simulation_contract.stage12_manifest import V2VersionManifest +from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportRequest, +) +from policyengine_simulation_entry.household_report_submissions import ( + HouseholdReportSubmissionService, +) +from policyengine_simulation_entry.household_workers import ( + HouseholdWorker, + default_household_worker, +) +from policyengine_simulation_persistence import ( + HouseholdReportInputs, + ReportRunStatus, + StoredHousehold, +) + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from policyengine_simulation_persistence import NewHouseholdReport + +HOUSEHOLD_ID = UUID("00000000-0000-4000-8000-0000000000bb") +POLICY_ID = UUID("00000000-0000-4000-8000-0000000000cc") +CATALOG_VERSION_ID = UUID("00000000-0000-4000-8000-0000000000dd") +MODEL_ID = UUID("00000000-0000-4000-8000-0000000000de") +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +REFORM = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} +STORED_DOCUMENT: dict[str, Any] = { + "people": [ + { + "id": "you", + "values": {"age": 40, "employment_income": {"2026": 30_000}}, + "memberships": { + "household": "your household", + "family": "your family", + "tax_unit": "your tax unit", + "spm_unit": "your spm unit", + "marital_unit": "your marital unit", + }, + } + ], + "household": [{"id": "your household", "values": {"household_net_income": None}}], + "family": [{"id": "your family", "values": {}}], + "tax_unit": [{"id": "your tax unit", "values": {}}], + "spm_unit": [{"id": "your spm unit", "values": {}}], + "marital_unit": [{"id": "your marital unit", "values": {}}], +} +# The stored household above, in the situation shape a household function takes. +SITUATION = { + "people": {"you": {"age": 40, "employment_income": {"2026": 30_000}}}, + "households": { + "your household": {"members": ["you"], "household_net_income": None} + }, + "families": {"your family": {"members": ["you"]}}, + "tax_units": {"your tax unit": {"members": ["you"]}}, + "spm_units": {"your spm unit": {"members": ["you"]}}, + "marital_units": {"your marital unit": {"members": ["you"]}}, +} +HOUSEHOLD_REPORT_BODY: dict[str, Any] = { + "country": "us", + "report_type": "household", + "label": "A larger child tax credit for one family", + "description": "Raises the credit to $3,000", + "household_id": str(HOUSEHOLD_ID), + "baseline": {"policy_id": None}, + "reform": {"policy_id": str(POLICY_ID)}, +} + + +def household_report_request(**changes: Any) -> HouseholdReportRequest: + return HouseholdReportRequest.model_validate({**HOUSEHOLD_REPORT_BODY, **changes}) + + +def manifest_household_worker(country: str) -> HouseholdWorker: + """Return the country's household function in a manifest of one worker.""" + version = household_worker_version() + manifest = V2VersionManifest( + generation="test", + default_version=version.bundle.policyengine_version, + versions={version.bundle.policyengine_version: version}, + ) + return default_household_worker(manifest, country) + + +class FakeHouseholdWorkers: + """The default worker's household functions, or the lookup's failure.""" + + def __init__(self, failure: Exception | None = None) -> None: + self.failure = failure + + def default_worker(self, country: str) -> HouseholdWorker: + if self.failure is not None: + raise self.failure + return manifest_household_worker(country) + + +@dataclass +class FakeHouseholdReportStore: + """Records what the service asks of the database, and answers as told. + + ``claimed`` marks the report run as claimed by its household function: a + failure write then changes nothing. + """ + + default_year: int | None = 2026 + document: dict[str, Any] = field(default_factory=lambda: STORED_DOCUMENT) + read_failure: Exception | None = None + create_failure: Exception | None = None + record_failure: Exception | None = None + fail_failure: Exception | None = None + claimed: bool = False + reads: list[dict[str, Any]] = field(default_factory=list) + reports: list[NewHouseholdReport] = 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_inputs( + self, + *, + country: str, + household_id: UUID, + policy_ids: Mapping[str, UUID | None], + policyengine_version: str, + ) -> HouseholdReportInputs: + self.reads.append( + { + "country": country, + "household_id": household_id, + "policy_ids": dict(policy_ids), + "policyengine_version": policyengine_version, + } + ) + if self.read_failure is not None: + raise self.read_failure + return HouseholdReportInputs( + catalog_version_id=CATALOG_VERSION_ID, + tax_benefit_model_id=MODEL_ID, + household=StoredHousehold( + id=household_id, + country=country, + default_year=self.default_year, + document=self.document, + ), + reforms={ + side: ({} if policy_id is None else REFORM) + for side, policy_id in policy_ids.items() + }, + ) + + def create_report(self, report: NewHouseholdReport) -> None: + if self.create_failure is not None: + raise self.create_failure + self.reports.append(report) + + 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 + + +@dataclass +class FakeHouseholdReportCaller: + """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 + calls: list[tuple[HouseholdWorker, dict, UUID, Any]] = field(default_factory=list) + + async def start_report_run( + self, + worker: HouseholdWorker, + payload: dict[str, Any], + *, + report_run_id: UUID, + observability_context: dict[str, Any] | None, + ) -> str: + if self.hang: + await asyncio.sleep(60) + if self.failure is not None: + raise self.failure + self.calls.append((worker, payload, report_run_id, observability_context)) + return "fc-household-report" + + +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 household_report_service( + store: FakeHouseholdReportStore, + *, + workers: FakeHouseholdWorkers | None = None, + caller: FakeHouseholdReportCaller | None = None, +) -> HouseholdReportSubmissionService: + return HouseholdReportSubmissionService( + store=store, + workers=workers or FakeHouseholdWorkers(), + caller=caller or FakeHouseholdReportCaller(), + call_start_timeout_seconds=0.05, + new_id=numbered_ids(), + ) diff --git a/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py b/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py index b9a8f043a..5a84f1567 100644 --- a/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py +++ b/projects/policyengine-simulation-entry/tests/test_budget_window_routes.py @@ -153,7 +153,7 @@ def test_an_economy_report_body_still_reaches_the_report_service(client_for) -> {"loc": ["body", "report_type"], "type": "missing"}, ), ( - {**REPORT_BODY, "report_type": "household"}, + {**REPORT_BODY, "report_type": "not_a_report_type"}, {"loc": ["body", "report_type"], "type": "literal_error"}, ), ( @@ -180,7 +180,7 @@ def test_a_refused_report_names_its_fields_as_an_economy_report_did( assert {"loc": error["loc"], "type": error["type"]} == expected -@pytest.mark.parametrize("report_type", [None, 5, ["economy"], "household"]) +@pytest.mark.parametrize("report_type", [None, 5, ["economy"], "not_a_report_type"]) def test_an_unknown_report_type_is_refused_with_the_value_sent( client_for, report_type ) -> None: @@ -194,9 +194,9 @@ def test_an_unknown_report_type_is_refused_with_the_value_sent( { "type": "literal_error", "loc": ["body", "report_type"], - "msg": "Input should be 'economy' or 'budget_window'", + "msg": "Input should be 'economy', 'budget_window' or 'household'", "input": report_type, - "ctx": {"expected": "'economy' or 'budget_window'"}, + "ctx": {"expected": "'economy', 'budget_window' or 'household'"}, } ] diff --git a/projects/policyengine-simulation-entry/tests/test_household_report_routes.py b/projects/policyengine-simulation-entry/tests/test_household_report_routes.py new file mode 100644 index 000000000..3f250723d --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_household_report_routes.py @@ -0,0 +1,321 @@ +"""``POST /reports`` for household reports, with the submission services faked.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, NoReturn +from uuid import UUID + +import pytest +from conftest import FakeBackend, make_settings +from economy_service_fakes import REPORT_BODY +from fastapi.testclient import TestClient +from household_report_fakes import HOUSEHOLD_REPORT_BODY, POLICY_ID +from policyengine_simulation_entry import app as app_module +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, + UnprocessableReason, +) +from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportCreated, + HouseholdReportRequest, +) +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_simulation_persistence import ReportRunStatus +from test_economy_routes import CREATED_REPORT, FakeEconomySubmissions + +OBSERVABILITY_ID = "6e1f3d4c-0a5b-4c39-9d2e-7f8a9b0c1d2e" +REPORT_ID = UUID("00000000-0000-4000-8000-0000000000e1") +CREATED_HOUSEHOLD_REPORT = HouseholdReportCreated( + report_id=REPORT_ID, + report_run_id=UUID("00000000-0000-4000-8000-0000000000e2"), + status=ReportRunStatus.PENDING, + year=2026, + baseline_simulation_id=UUID("00000000-0000-4000-8000-0000000000f1"), + baseline_simulation_run_id=UUID("00000000-0000-4000-8000-0000000000f2"), + reform_simulation_id=UUID("00000000-0000-4000-8000-0000000000f3"), + reform_simulation_run_id=UUID("00000000-0000-4000-8000-0000000000f4"), +) + + +def _raise(failure: Exception) -> NoReturn: + raise failure + + +@dataclass +class FakeHouseholdReportSubmissions: + """Records each household report submission, and answers as told.""" + + failure: Exception | None = None + reports: list[tuple[HouseholdReportRequest, Any]] = field(default_factory=list) + + async def submit_report( + self, + request: HouseholdReportRequest, + *, + observability_context: dict[str, Any] | None, + ) -> HouseholdReportCreated: + self.reports.append((request, observability_context)) + if self.failure is not None: + _raise(self.failure) + return CREATED_HOUSEHOLD_REPORT + + +@pytest.fixture +def client_for(monkeypatch): + """Return a function building a client whose app uses these services. + + A household service of None leaves household reports unconfigured. + """ + + def build( + household_reports: FakeHouseholdReportSubmissions | None, + economy: FakeEconomySubmissions | None = None, + ) -> TestClient: + monkeypatch.setattr( + app_module, + "economy_submissions_for", + lambda settings, database: economy or FakeEconomySubmissions(), + ) + monkeypatch.setattr( + app_module, + "household_reports_for", + lambda settings, database: household_reports, + ) + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + ) + ) + + return build + + +def test_a_household_report_answers_201_with_its_records_and_runs(client_for) -> None: + household_reports = FakeHouseholdReportSubmissions() + economy = FakeEconomySubmissions() + + with client_for(household_reports, economy) as client: + result = client.post( + "/reports", + json=HOUSEHOLD_REPORT_BODY, + headers={OBSERVABILITY_ID_HEADER: OBSERVABILITY_ID}, + ) + + assert result.status_code == 201 + assert result.json() == CREATED_HOUSEHOLD_REPORT.model_dump(mode="json") + ((request, context),) = household_reports.reports + assert request.household_id == UUID(HOUSEHOLD_REPORT_BODY["household_id"]) + assert request.reform is not None + assert request.reform.policy_id == POLICY_ID + # A household report has no options: it always adds missing variables. + assert request.options() == {} + assert context["observability_id"] == OBSERVABILITY_ID + assert economy.reports == [] + + +def test_an_economy_report_still_reaches_the_economy_service(client_for) -> None: + household_reports = FakeHouseholdReportSubmissions() + economy = FakeEconomySubmissions() + + with client_for(household_reports, economy) as client: + result = client.post("/reports", json=REPORT_BODY) + + assert result.status_code == 201 + assert result.json() == CREATED_REPORT.model_dump(mode="json") + assert len(economy.reports) == 1 + assert household_reports.reports == [] + + +def test_a_report_of_the_baseline_alone_may_leave_out_its_reform(client_for) -> None: + household_reports = FakeHouseholdReportSubmissions() + body = { + key: value for key, value in HOUSEHOLD_REPORT_BODY.items() if key != "reform" + } + + with client_for(household_reports) as client: + result = client.post("/reports", json=body) + + assert result.status_code == 201 + ((request, _),) = household_reports.reports + assert request.reform is None + + +def test_household_reports_answer_503_until_they_are_configured(client_for) -> None: + with client_for(None) as client: + result = client.post("/reports", json=HOUSEHOLD_REPORT_BODY) + + assert result.status_code == 503 + assert result.headers["Retry-After"] == "10" + assert result.json() == { + "code": "not_configured", + "detail": "Household reports are not configured.", + } + + +@pytest.mark.parametrize( + ("body", "location", "error_type"), + [ + ( + {k: v for k, v in HOUSEHOLD_REPORT_BODY.items() if k != "household_id"}, + ["body", "household_id"], + "missing", + ), + ( + { + **HOUSEHOLD_REPORT_BODY, + "baseline": {"policy_id": None, "dynamic_id": None}, + }, + ["body", "baseline", "dynamic_id"], + "extra_forbidden", + ), + ( + {**HOUSEHOLD_REPORT_BODY, "reform": {"policy_id": "not-a-uuid"}}, + ["body", "reform", "policy_id"], + "uuid_parsing", + ), + ( + {**HOUSEHOLD_REPORT_BODY, "add_missing_variables": True}, + ["body", "add_missing_variables"], + "extra_forbidden", + ), + ( + {**HOUSEHOLD_REPORT_BODY, "year": 2026}, + ["body", "year"], + "extra_forbidden", + ), + ( + {k: v for k, v in HOUSEHOLD_REPORT_BODY.items() if k != "report_type"}, + ["body", "report_type"], + "missing", + ), + ], + ids=[ + "no-household", + "dynamic", + "malformed-policy", + "missing-variables-option", + "year", + "no-report-type", + ], +) +def test_a_refused_household_report_names_its_fields_without_its_type( + client_for, body, location, error_type +) -> None: + household_reports = FakeHouseholdReportSubmissions() + + with client_for(household_reports) as client: + result = client.post("/reports", json=body) + + assert result.status_code == 422 + reasons = result.json()["detail"] + assert [reason["loc"] for reason in reasons] == [location] + assert [reason["type"] for reason in reasons] == [error_type] + assert household_reports.reports == [] + + +def test_an_unknown_report_type_is_refused_with_the_types_there_are( + client_for, +) -> None: + with client_for(FakeHouseholdReportSubmissions()) as client: + result = client.post( + "/reports", + json={**HOUSEHOLD_REPORT_BODY, "report_type": "not_a_report_type"}, + ) + + assert result.status_code == 422 + ((reason),) = result.json()["detail"] + assert reason["loc"] == ["body", "report_type"] + assert reason["type"] == "literal_error" + assert reason["msg"] == "Input should be 'economy', 'budget_window' or 'household'" + assert reason["input"] == "not_a_report_type" + + +@pytest.mark.parametrize( + ("failure", "status_code", "answer"), + [ + ( + EconomySubmissionRefusedError( + code="unknown_household", + detail="us has no stored household the request names.", + ), + 400, + { + "code": "unknown_household", + "detail": "us has no stored household the request names.", + }, + ), + ( + EconomySubmissionUnprocessableError( + code="household_not_calculable", + reasons=( + UnprocessableReason( + loc=("body", "household_id"), + msg="the household has no default year to calculate", + type="household_without_default_year", + ), + ), + ), + 422, + { + "detail": [ + { + "loc": ["body", "household_id"], + "msg": "the household has no default year to calculate", + "type": "household_without_default_year", + } + ], + "code": "household_not_calculable", + }, + ), + ( + EconomySubmissionUnavailableError( + code="modal_call_refused", + detail="Modal did not accept the report's call.", + record_id=REPORT_ID, + ), + 503, + { + "code": "modal_call_refused", + "detail": "Modal did not accept the report's call.", + "report_id": str(REPORT_ID), + }, + ), + ( + EconomySubmissionOutcomeUnknownError( + "report", error_type="OperationalError" + ), + 500, + { + "code": "internal_error", + "detail": ( + "The service may have created the report, but could not " + "confirm it was written." + ), + }, + ), + ( + RuntimeError("the household function's request was built wrong"), + 500, + { + "code": "internal_error", + "detail": "The service could not answer the request.", + }, + ), + ], + ids=["refused", "unprocessable", "call-refused", "outcome-unknown", "unexpected"], +) +def test_a_household_report_the_service_refuses_answers_with_its_code( + client_for, failure, status_code, answer +) -> None: + with client_for(FakeHouseholdReportSubmissions(failure=failure)) as client: + result = client.post("/reports", json=HOUSEHOLD_REPORT_BODY) + + assert result.status_code == status_code + assert result.json() == answer diff --git a/projects/policyengine-simulation-entry/tests/test_household_report_submissions.py b/projects/policyengine-simulation-entry/tests/test_household_report_submissions.py new file mode 100644 index 000000000..48f9b1714 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_household_report_submissions.py @@ -0,0 +1,532 @@ +"""The household report service: what it reads, creates and starts, and refuses. + +The store, the manifest and the Modal call are faked +(``household_report_fakes``); the store's own SQL is tested in the +persistence library. +""" + +from __future__ import annotations + +import asyncio +from uuid import UUID + +import pytest +from household_report_fakes import ( + CATALOG_VERSION_ID, + HOUSEHOLD_ID, + MODEL_ID, + POLICY_ID, + REFORM, + SITUATION, + STORED_DOCUMENT, + FakeHouseholdReportCaller, + FakeHouseholdReportStore, + FakeHouseholdWorkers, + household_report_request, + household_report_service, + manifest_household_worker, +) +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, +) +from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportRunCalculation, +) +from policyengine_simulation_entry import household_report_submissions +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, +) +from policyengine_simulation_entry.household_report_api_models import ( + HouseholdReportCreated, +) +from policyengine_simulation_entry.household_report_submissions import ( + HouseholdReportBuildError, + HouseholdReportSideInputs, + household_report_calculation, +) +from policyengine_simulation_entry.household_workers import ( + HouseholdFunctionUnavailableError, + VersionManifestUnavailableError, +) +from policyengine_simulation_persistence import ( + EconomyRecordOutcomeUnknownError, + EconomyRecordReferenceError, + MissingPolicyParametersError, + ReportRunStatus, + ReportSideError, + RunVersions, + StoredHousehold, + UnknownCatalogVersionError, + UnknownHouseholdError, + UnknownPolicyError, +) +from pydantic import ValidationError +from sqlalchemy.exc import OperationalError + +# The service's IDs, in the order it takes them: the report, its run, then +# each side's simulation and simulation run. +REPORT_ID, RUN_ID, BASELINE_ID, BASELINE_RUN_ID, REFORM_ID, REFORM_RUN_ID = ( + UUID(int=number) for number in range(1, 7) +) +CONTEXT = {"observability_id": "6e1f3d4c-0a5b-4c39-9d2e-7f8a9b0c1d2e"} +DATABASE_ERROR = OperationalError("SELECT 1", {}, Exception("server closed")) + + +def _submit(service, request=None): + return asyncio.run( + service.submit_report( + request or household_report_request(), observability_context=CONTEXT + ) + ) + + +def _side_request(calculation_id: UUID, parameters: dict) -> dict: + worker = manifest_household_worker("us") + return { + "schema_version": 1, + "calculation_id": str(calculation_id), + "country": "us", + "year": "2026", + "household": SITUATION, + "policy": {"parameters": parameters}, + "add_missing_variables": True, + "spm_selection_source": "request", + "bundle": worker.bundle.model_dump(mode="json"), + } + + +def test_a_household_report_is_created_and_its_household_function_started(): + store = FakeHouseholdReportStore() + caller = FakeHouseholdReportCaller() + service = household_report_service(store, caller=caller) + + created = _submit(service) + + assert created == HouseholdReportCreated( + report_id=REPORT_ID, + report_run_id=RUN_ID, + status=ReportRunStatus.PENDING, + year=2026, + baseline_simulation_id=BASELINE_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_id=REFORM_ID, + reform_simulation_run_id=REFORM_RUN_ID, + ) + worker = manifest_household_worker("us") + assert store.reads == [ + { + "country": "us", + "household_id": HOUSEHOLD_ID, + "policy_ids": {"baseline": None, "reform": POLICY_ID}, + "policyengine_version": worker.bundle.policyengine_version, + } + ] + (report,) = store.reports + assert (report.id, report.run_id, report.country) == (REPORT_ID, RUN_ID, "us") + assert report.label == "A larger child tax credit for one family" + assert report.description == "Raises the credit to $3,000" + assert (report.tax_benefit_model_id, report.tax_benefit_model_version_id) == ( + MODEL_ID, + CATALOG_VERSION_ID, + ) + assert (report.household_id, report.year) == (HOUSEHOLD_ID, 2026) + assert report.inputs == {} + assert report.versions == RunVersions( + country_package_version=worker.bundle.country_package_version, + policyengine_version=worker.bundle.policyengine_version, + ) + assert (report.baseline.simulation_id, report.baseline.policy_id) == ( + BASELINE_ID, + None, + ) + assert report.reform is not None + assert (report.reform.simulation_run_id, report.reform.policy_id) == ( + REFORM_RUN_ID, + POLICY_ID, + ) + ((called_worker, payload, report_run_id, context),) = caller.calls + assert called_worker == worker + assert (report_run_id, context) == (RUN_ID, CONTEXT) + assert payload == { + "schema_version": 1, + "baseline": _side_request(BASELINE_RUN_ID, {}), + "reform": _side_request(REFORM_RUN_ID, REFORM), + } + HouseholdReportRunCalculation.model_validate(payload) + assert store.recorded_calls == [(RUN_ID, "fc-household-report")] + + +def test_a_report_of_the_baseline_alone_creates_and_calculates_one_side(): + store = FakeHouseholdReportStore() + caller = FakeHouseholdReportCaller() + service = household_report_service(store, caller=caller) + + created = _submit( + service, + household_report_request(baseline={"policy_id": str(POLICY_ID)}, reform=None), + ) + + assert (created.reform_simulation_id, created.reform_simulation_run_id) == ( + None, + None, + ) + assert store.reads[0]["policy_ids"] == {"baseline": POLICY_ID} + (report,) = store.reports + assert report.reform is None + assert report.simulation_run_ids == (BASELINE_RUN_ID,) + ((_, payload, _, _),) = caller.calls + assert payload["reform"] is None + assert payload["baseline"]["policy"] == {"parameters": REFORM} + + +def test_a_saved_spm_selection_is_applied_to_both_sides(): + saved = {"scenario": "ce_trend", "geography_kind": "national"} + store = FakeHouseholdReportStore(document={**STORED_DOCUMENT, "spm": saved}) + caller = FakeHouseholdReportCaller() + + _submit(household_report_service(store, caller=caller)) + + ((_, payload, _, _),) = caller.calls + for side in ("baseline", "reform"): + assert payload[side]["spm"] == saved + assert payload[side]["spm_selection_source"] == "saved_household" + + +@pytest.mark.parametrize( + ("failure", "expected_type", "code"), + [ + ( + HouseholdFunctionUnavailableError("no household function"), + EconomySubmissionRefusedError, + "no_household_function", + ), + ( + VersionManifestUnavailableError("unreadable"), + EconomySubmissionUnavailableError, + "manifest_unavailable", + ), + ], +) +def test_the_default_workers_household_function_is_needed(failure, expected_type, code): + store = FakeHouseholdReportStore() + service = household_report_service( + store, workers=FakeHouseholdWorkers(failure=failure) + ) + + with pytest.raises(expected_type) as raised: + _submit(service) + + assert raised.value.code == code + assert (store.reads, store.reports) == ([], []) + + +@pytest.mark.parametrize( + ("failure", "expected_type", "code", "detail"), + [ + ( + UnknownCatalogVersionError("no catalog version"), + EconomySubmissionRefusedError, + "unknown_catalog_version", + "us has no catalog version for .py 6.17.1, the default v2 worker's " + "version.", + ), + ( + UnknownHouseholdError("us has no household"), + EconomySubmissionRefusedError, + "unknown_household", + "us has no stored household the request names.", + ), + ( + ReportSideError("reform", UnknownPolicyError("no policy")), + EconomySubmissionRefusedError, + "unknown_policy", + "us has no stored policy the request's reform names.", + ), + ( + DATABASE_ERROR, + EconomySubmissionUnavailableError, + "database_unavailable", + "The simulation database could not be read.", + ), + ], + ids=["catalog-version", "household", "policy", "database"], +) +def test_each_missing_input_is_refused_before_anything_is_written( + failure, expected_type, code, detail +): + store = FakeHouseholdReportStore(read_failure=failure) + caller = FakeHouseholdReportCaller() + + with pytest.raises(expected_type) as raised: + _submit(household_report_service(store, caller=caller)) + + assert (raised.value.code, raised.value.detail) == (code, detail) + assert (store.reports, caller.calls) == ([], []) + + +def test_a_policy_setting_missing_parameters_is_refused_on_its_side(): + store = FakeHouseholdReportStore( + read_failure=ReportSideError( + "baseline", MissingPolicyParametersError(["gov.a", "gov.b"]) + ) + ) + + with pytest.raises(EconomySubmissionUnprocessableError) as raised: + _submit(household_report_service(store)) + + assert raised.value.code == "missing_parameters" + assert [reason.loc for reason in raised.value.reasons] == [ + ("body", "baseline", "policy_id", "gov.a"), + ("body", "baseline", "policy_id", "gov.b"), + ] + assert {reason.type for reason in raised.value.reasons} == {"missing_parameter"} + assert store.reports == [] + + +@pytest.mark.parametrize( + ("store", "reason_type"), + [ + (FakeHouseholdReportStore(default_year=None), "household_without_default_year"), + ( + FakeHouseholdReportStore(document={"people": "not a list of people"}), + "household_unreadable", + ), + ], + ids=["no-default-year", "unreadable"], +) +def test_a_household_that_cannot_be_calculated_is_refused(store, reason_type): + caller = FakeHouseholdReportCaller() + + with pytest.raises(EconomySubmissionUnprocessableError) as raised: + _submit(household_report_service(store, caller=caller)) + + assert raised.value.code == "household_not_calculable" + ((reason),) = raised.value.reasons + assert (reason.loc, reason.type) == (("body", "household_id"), reason_type) + # The detail is the service's own words, never the document's contents. + assert "not a list of people" not in reason.msg + assert (store.reports, caller.calls) == ([], []) + + +def test_a_household_the_contract_refuses_is_located_on_the_household(): + # A group no person belongs to has no members, which the contract refuses. + document = { + **STORED_DOCUMENT, + "family": [*STORED_DOCUMENT["family"], {"id": "an empty family", "values": {}}], + } + store = FakeHouseholdReportStore(document=document) + + with pytest.raises(EconomySubmissionUnprocessableError) as raised: + _submit(household_report_service(store)) + + assert raised.value.code == "household_not_calculable" + assert raised.value.reasons + assert all( + reason.loc[:3] == ("body", "household_id", "household") + for reason in raised.value.reasons + ) + assert store.reports == [] + + +@pytest.mark.parametrize( + ("column", "code"), + [ + ("household_id", "unknown_household"), + ("policy_id", "unknown_policy"), + ("tax_benefit_model_version_id", "unknown_catalog_version"), + (None, "missing_reference"), + ], +) +def test_a_row_deleted_after_the_reads_is_refused_on_create(column, code): + store = FakeHouseholdReportStore( + create_failure=EconomyRecordReferenceError("simulations", column) + ) + caller = FakeHouseholdReportCaller() + + with pytest.raises(EconomySubmissionRefusedError) as raised: + _submit(household_report_service(store, caller=caller)) + + assert raised.value.code == code + assert caller.calls == [] + + +def test_a_create_whose_commit_failed_may_exist(): + store = FakeHouseholdReportStore( + create_failure=EconomyRecordOutcomeUnknownError("commit failed") + ) + caller = FakeHouseholdReportCaller() + + with pytest.raises(EconomySubmissionOutcomeUnknownError) as raised: + _submit(household_report_service(store, caller=caller)) + + assert raised.value.record == "report" + assert caller.calls == [] + + +def test_a_create_the_database_cannot_take_is_unavailable(): + store = FakeHouseholdReportStore(create_failure=DATABASE_ERROR) + + with pytest.raises(EconomySubmissionUnavailableError) as raised: + _submit(household_report_service(store)) + + assert raised.value.code == "database_unavailable" + assert raised.value.record_id is None + + +@pytest.mark.parametrize( + ("caller", "error_type"), + [ + (FakeHouseholdReportCaller(failure=RuntimeError("refused")), "RuntimeError"), + (FakeHouseholdReportCaller(hang=True), "TimeoutError"), + ], + ids=["refused", "not-confirmed-in-time"], +) +def test_a_call_modal_does_not_confirm_fails_the_report_run_and_its_runs( + caller, error_type +): + store = FakeHouseholdReportStore() + + with pytest.raises(EconomySubmissionUnavailableError) as raised: + _submit(household_report_service(store, caller=caller)) + + assert raised.value.code == "modal_call_refused" + assert raised.value.record_id == REPORT_ID + assert store.failed_runs == [ + ( + RUN_ID, + (BASELINE_RUN_ID, REFORM_RUN_ID), + "modal_call_refused: Modal did not accept the report's call " + f"({error_type}).", + ) + ] + assert store.recorded_calls == [] + + +def test_a_call_modal_accepted_late_is_answered_pending(): + # The household function claimed the run, so the failure write changes nothing. + store = FakeHouseholdReportStore(claimed=True) + caller = FakeHouseholdReportCaller(failure=RuntimeError("late")) + + created = _submit(household_report_service(store, caller=caller)) + + assert created.status is ReportRunStatus.PENDING + assert store.failed_runs == [] + + +def test_a_failure_write_that_fails_answers_pending(): + store = FakeHouseholdReportStore(fail_failure=DATABASE_ERROR) + caller = FakeHouseholdReportCaller(failure=RuntimeError("refused")) + + created = _submit(household_report_service(store, caller=caller)) + + assert created.status is ReportRunStatus.PENDING + + +@pytest.mark.parametrize( + ("store", "status"), + [ + (FakeHouseholdReportStore(record_failure=DATABASE_ERROR), "pending"), + (FakeHouseholdReportStore(claimed=True), "running"), + ], + ids=["record-fails", "claimed-first"], +) +def test_the_answer_is_the_report_runs_status_once_its_call_is_recorded(store, status): + created = _submit(household_report_service(store)) + + assert created.status.value == status + + +def _stored_household(document: dict | None = None) -> StoredHousehold: + return StoredHousehold( + id=HOUSEHOLD_ID, + country="us", + default_year=2026, + document=STORED_DOCUMENT if document is None else document, + ) + + +def test_a_report_runs_calculation_is_built_from_its_stored_inputs_alone(): + # A restart rebuilds the call from the stored report: its own year, not + # the household's default year now. + calculation = household_report_calculation( + household=_stored_household(), + year=2027, + baseline=HouseholdReportSideInputs( + simulation_run_id=BASELINE_RUN_ID, policy_values={} + ), + reform=HouseholdReportSideInputs( + simulation_run_id=REFORM_RUN_ID, policy_values=REFORM + ), + worker=manifest_household_worker("us"), + ) + + payload = calculation.model_dump(mode="json") + for side, run_id, parameters in ( + ("baseline", BASELINE_RUN_ID, {}), + ("reform", REFORM_RUN_ID, REFORM), + ): + # Every side adds the household's missing variables. + expected = {**_side_request(run_id, parameters), "year": "2027"} + assert payload[side] == expected + + +def test_a_calculation_of_the_baseline_alone_has_no_reform(): + calculation = household_report_calculation( + household=_stored_household(), + year=2026, + baseline=HouseholdReportSideInputs( + simulation_run_id=BASELINE_RUN_ID, policy_values=REFORM + ), + worker=manifest_household_worker("us"), + ) + + assert calculation.reform is None + assert calculation.baseline.calculation_id == BASELINE_RUN_ID + assert calculation.baseline.policy.parameters == REFORM + + +def test_a_stored_household_the_builder_cannot_read_is_not_calculable(): + with pytest.raises(EconomySubmissionUnprocessableError) as raised: + household_report_calculation( + household=_stored_household({"people": "not a list of people"}), + year=2026, + baseline=HouseholdReportSideInputs( + simulation_run_id=BASELINE_RUN_ID, policy_values={} + ), + worker=manifest_household_worker("us"), + ) + + assert raised.value.code == "household_not_calculable" + assert [reason.type for reason in raised.value.reasons] == ["household_unreadable"] + + +def test_two_sides_recording_one_simulation_run_are_built_wrong(): + side = HouseholdReportSideInputs( + simulation_run_id=BASELINE_RUN_ID, policy_values={} + ) + + with pytest.raises(HouseholdReportBuildError, match="different simulation runs"): + household_report_calculation( + household=_stored_household(), + year=2026, + baseline=side, + reform=side, + worker=manifest_household_worker("us"), + ) + + +def test_a_refusal_of_the_missing_variables_option_is_the_services_fault(): + # A household report sets the option itself, so a refusal of it names no + # field the caller sent. + refused_request = { + **_side_request(BASELINE_RUN_ID, {}), + "add_missing_variables": [1], + } + with pytest.raises(ValidationError) as refused: + HouseholdExecutionRequest.model_validate(refused_request) + + with pytest.raises(HouseholdReportBuildError, match="add_missing_variables"): + household_report_submissions._contract_reasons(refused.value, "baseline") diff --git a/projects/policyengine-simulation-entry/tests/test_household_workers.py b/projects/policyengine-simulation-entry/tests/test_household_workers.py index 661308082..b87d536dc 100644 --- a/projects/policyengine-simulation-entry/tests/test_household_workers.py +++ b/projects/policyengine-simulation-entry/tests/test_household_workers.py @@ -181,3 +181,60 @@ async def test_a_call_modal_gives_no_id_is_refused(monkeypatch): await ModalHouseholdFunctionCaller("staging").spawn( worker, {}, household_job_id=JOB_ID, observability_context=None ) + + +REPORT_RUN_ID = UUID("00000000-0000-4000-8000-0000000000ab") + + +@pytest.mark.asyncio +async def test_a_household_report_run_is_started_by_name_with_its_run(monkeypatch): + function = FakeModalFunction("fc-household-report-1") + looked_up: list[tuple] = [] + + def from_name(application_name, function_name, *, environment_name): + looked_up.append((application_name, function_name, environment_name)) + return function + + monkeypatch.setattr(household_workers.modal.Function, "from_name", from_name) + worker = ManifestHouseholdWorkers( + manifest_store(household_worker_version()) + ).default_worker("uk") + payload = {"baseline": {"calculation_id": str(JOB_ID)}, "reform": None} + + call_id = await ModalHouseholdFunctionCaller("staging").start_report_run( + worker, + payload, + report_run_id=REPORT_RUN_ID, + observability_context={"observability_id": "o-1"}, + ) + + assert call_id == "fc-household-report-1" + assert looked_up == [ + (v2_application_name(POLICYENGINE_VERSION), "calculate_household_uk", "staging") + ] + assert function.spawned == [ + ( + (payload,), + { + "household_report_run_id": str(REPORT_RUN_ID), + "observability_context": {"observability_id": "o-1"}, + }, + ) + ] + + +@pytest.mark.asyncio +async def test_a_report_run_call_modal_gives_no_id_is_refused(monkeypatch): + monkeypatch.setattr( + household_workers.modal.Function, + "from_name", + lambda *names, environment_name: FakeModalFunction(None), + ) + worker = ManifestHouseholdWorkers( + manifest_store(household_worker_version()) + ).default_worker("us") + + with pytest.raises(RuntimeError, match="no call ID"): + await ModalHouseholdFunctionCaller("staging").start_report_run( + worker, {}, report_run_id=REPORT_RUN_ID, observability_context=None + ) diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 0bf4a8427..de8226b4d 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -102,6 +102,9 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: "EconomySimulationCreated", "EconomySimulationRequest", "EconomySimulationState", + "HouseholdPolicyChoice", + "HouseholdReportCreated", + "HouseholdReportRequest", "PolicyChoice", "ReportRunRead", "ReportRunStatus", @@ -364,7 +367,14 @@ def test_economy_routes_declare_every_answer(): unprocessable = {"$ref": "#/components/schemas/EconomyRequestUnprocessable"} for path, created in ( - ("/reports", ("EconomyReportCreated", "BudgetWindowReportCreated")), + ( + "/reports", + ( + "EconomyReportCreated", + "BudgetWindowReportCreated", + "HouseholdReportCreated", + ), + ), ("/simulations", ("EconomySimulationCreated",)), ): responses = _responses(spec, "post", path) @@ -387,9 +397,15 @@ def test_economy_routes_declare_every_answer(): "unknown_dataset", "unknown_dynamic", "unknown_catalog_version", + "unknown_household", + "no_household_function", "modal_call_refused", "internal_error", } <= codes + unprocessable_codes = spec["components"]["schemas"]["EconomyRequestUnprocessable"][ + "properties" + ]["code"]["enum"] + assert "household_not_calculable" in unprocessable_codes def test_every_economy_503_declares_its_retry_after_header(): @@ -441,7 +457,7 @@ def test_economy_requests_name_no_version_and_no_record_id(): assert "spm" not in schemas["EconomyReportRequest"].get("required", []) -def test_a_report_request_is_an_economy_report_or_a_budget_window(): +def test_a_report_request_is_an_economy_report_a_budget_window_or_a_household(): spec = create_app().openapi() body = spec["paths"]["/reports"]["post"]["requestBody"]["content"][ "application/json" @@ -453,6 +469,7 @@ def test_a_report_request_is_an_economy_report_or_a_budget_window(): "mapping": { "economy": "#/components/schemas/EconomyReportRequest", "budget_window": "#/components/schemas/BudgetWindowReportRequest", + "household": "#/components/schemas/HouseholdReportRequest", }, } # A window sends its length and no cliff option, and names no version @@ -474,6 +491,41 @@ def test_a_report_request_is_an_economy_report_or_a_budget_window(): assert (size["minimum"], size["maximum"]) == (1, 75) +def test_a_household_report_names_its_household_and_each_sides_policy(): + schemas = create_app().openapi()["components"]["schemas"] + + # No year, region, dataset, dynamic, version or record ID, and no option: + # a household report always adds the household's missing variables. + household_report = schemas["HouseholdReportRequest"] + assert set(household_report["properties"]) == { + "country", + "report_type", + "label", + "description", + "household_id", + "baseline", + "reform", + } + assert set(household_report["required"]) == { + "country", + "report_type", + "label", + "household_id", + "baseline", + } + assert set(schemas["HouseholdPolicyChoice"]["properties"]) == {"policy_id"} + assert set(schemas["HouseholdReportCreated"]["required"]) == { + "report_id", + "report_run_id", + "status", + "year", + "baseline_simulation_id", + "baseline_simulation_run_id", + "reform_simulation_id", + "reform_simulation_run_id", + } + + def test_economy_reads_declare_every_answer(): spec = create_app().openapi() error = {"$ref": "#/components/schemas/EconomyError"} diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 801f12cb9..ab5d7f646 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -3,10 +3,13 @@ Besides the Stage 12 report coordinator and simulation workers, it defines a household function for each country, ``calculate_household_us`` and -``calculate_household_uk``, which calculates one household under one policy -in an image without datasets. A call that names a household job -records it with the Modal worker role's database secret, the functions' only -secret. +``calculate_household_uk``, in an image without datasets. A call calculates +one household under one policy, and records it on the household job it +names; or it calculates a household report run's baseline and, if it has +one, its reform, and records them on the report run and its simulation +runs. Both record with the Modal worker role's database secret, the +functions' only secret. A call that names neither, as deployment +validation makes, records nothing. This module defines additional resources. It does not rename or modify the existing ``policyengine-simulation-py{version}`` applications. @@ -553,10 +556,14 @@ def coordinate_report( # Household function sizing. A call holds one country model and a simulation -# of the household, under its one policy, at every axis point. Each container +# of the household, under one policy, at every axis point; a household report +# run's call calculates its baseline and then its reform, which the +# simulation service builds without axes. Each container # reserves 4 GiB and may use up to 16 GiB before Modal stops it, so a large # household calculated with every variable at many axis points still fits. -# The timeout is three times the Public API's 300-second request timeout. +# The timeout is three times the Public API's 300-second request timeout; a +# household report run's two calculations, as the simulation service builds +# them without axes, take a small part of it. HOUSEHOLD_FUNCTION_CPU = 2.0 HOUSEHOLD_FUNCTION_MEMORY_MIB = (4096, 16384) HOUSEHOLD_FUNCTION_TIMEOUT_SECONDS = 900 @@ -569,6 +576,7 @@ def _run_household_calculation( *, observability_context: dict | None, household_job_id: str | None, + household_report_run_id: str | None = None, ) -> dict: """Calculate one household request for ``country`` and return its outcome. @@ -577,8 +585,12 @@ def _run_household_calculation( With ``household_job_id``, the call claims that household job under its Modal call ID and records the outcome on it, or returns the job's status without calculating when another call holds it or it has finished - (``stage12_runtime.run_household_job``); without it, the call touches no - database. + (``stage12_runtime.run_household_job``). With ``household_report_run_id``, + ``payload`` is a ``HouseholdReportRunCalculation``: the call claims the + report run and each side's simulation run, calculates and records each + side, and returns the runs' statuses + (``stage12_runtime.run_household_report``). A call names at most one of + the two; without either, the call touches no database. """ from policyengine_simulation_contract.household_execution import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time HouseholdExecutionRequest, @@ -595,9 +607,11 @@ def _run_household_calculation( environment=os.getenv("MODAL_ENVIRONMENT", "local"), ) - def calculate() -> HouseholdExecutionResult | HouseholdExecutionFailure: + def calculate_request( + request: HouseholdExecutionRequest, + ) -> HouseholdExecutionResult | HouseholdExecutionFailure: return household_calculation.calculate_household( - HouseholdExecutionRequest.model_validate(payload), + request, worker=HouseholdWorkerIdentity( application_name=APP_NAME, function_name=household_callable_name(country), @@ -608,7 +622,13 @@ def calculate() -> HouseholdExecutionResult | HouseholdExecutionFailure: ), ) + def calculate() -> HouseholdExecutionResult | HouseholdExecutionFailure: + return calculate_request(HouseholdExecutionRequest.model_validate(payload)) + try: + if household_job_id is not None and household_report_run_id is not None: + msg = "a household call names either a household job or a report run" + raise ValueError(msg) propagated = normalize_observability_context(observability_context) result = None with runtime.operation( @@ -616,7 +636,17 @@ def calculate() -> HouseholdExecutionResult | HouseholdExecutionFailure: attributes={"country": country}, remote_context=propagated, ): - if household_job_id is None: + if household_report_run_id is not None: + # The recording code validates the payload once it holds + # the report run, so a refused payload fails the run. + result = stage12_runtime.run_household_report( + UUID(household_report_run_id), + payload, + calculate_request, + modal_call_id=modal.current_function_call_id(), + runtime=runtime, + ) + elif household_job_id is None: result = calculate().model_dump(mode="json") else: result = stage12_runtime.run_household_job( @@ -647,16 +677,20 @@ def calculate_household_us( *, observability_context: dict | None = None, household_job_id: str | None = None, + household_report_run_id: str | None = None, ) -> dict: """Calculate one US household under the request's policy. With ``household_job_id``, record the calculation on that household job. + With ``household_report_run_id``, ``payload`` is that household report + run's calculation: calculate and record each side, then the report run. """ return _run_household_calculation( "us", payload, observability_context=observability_context, household_job_id=household_job_id, + household_report_run_id=household_report_run_id, ) @@ -674,14 +708,18 @@ def calculate_household_uk( *, observability_context: dict | None = None, household_job_id: str | None = None, + household_report_run_id: str | None = None, ) -> dict: """Calculate one UK household under the request's policy. With ``household_job_id``, record the calculation on that household job. + With ``household_report_run_id``, ``payload`` is that household report + run's calculation: calculate and record each side, then the report run. """ return _run_household_calculation( "uk", payload, observability_context=observability_context, household_job_id=household_job_id, + household_report_run_id=household_report_run_id, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index b122f1ad9..e946f69a5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -25,6 +25,7 @@ run_segmented_simulation, ) from .household_job_recording import run_household_job +from .household_report_recording import run_household_report from .partition import ( should_segment_simulation, stage12_region_groups, @@ -73,6 +74,7 @@ "merge_segment_results", "redacted_failure_detail", "run_household_job", + "run_household_report", "run_segmented_simulation", "run_single_simulation", "should_segment_simulation", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_job_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_job_recording.py index 2825b089a..942db535c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_job_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_job_recording.py @@ -61,7 +61,7 @@ def household_failure_summary(failure: HouseholdExecutionFailure) -> str: return summary[:kept] + _SHORTENED_SUFFIX -def _unexpected_failure_summary(error: BaseException) -> str: +def unexpected_failure_summary(error: BaseException) -> str: """Describe an error outside the calculation, or Modal's cancellation. The summary names only the error's type: its message may repeat the @@ -170,7 +170,7 @@ def run_household_job( # Pre-emption: Modal runs the call again, and it resumes this claim. raise except BaseException as error: - recorder.fail_quietly(_unexpected_failure_summary(error), runtime=runtime) + recorder.fail_quietly(unexpected_failure_summary(error), runtime=runtime) raise if not recorded.changed: logger.warning( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_report_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_report_recording.py new file mode 100644 index 000000000..833a0d73d --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/household_report_recording.py @@ -0,0 +1,445 @@ +"""Recording the household report run a v2 household function's call names. + +The simulation service creates a household report's simulation runs, one +for its baseline and one for any reform, and its report run naming them, +all pending with no Modal call; it starts one household function call with +the report run's ID, then records the call's ID on the report run. The +function claims the report run under its own Modal call ID, then each side: +it claims the side's simulation run, calculates it, and records the +``HouseholdExecutionResult``, or the ``HouseholdExecutionFailure`` with a +one-line summary, in the run's ``household_result``. Both sides are always +calculated, so each run has its own outcome. The report run then succeeds +only if every side's run did; otherwise it fails with the first such side's +reason: its calculation's failure summary, or that another call holds its +run, or that another writer finished the run first. Everything is written +as the Modal worker role. + +A report run another call holds, or one already finished, is not calculated +again. A call Modal restarts after a pre-emption runs under the same call +ID, so it resumes its own claims, and keeps a side it had already recorded. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass +from datetime import UTC, datetime +from functools import partial +from typing import TYPE_CHECKING, Any, Protocol + +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionFailure, +) +from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportRunCalculation, +) +from policyengine_stage12_persistence import HouseholdReportRunStore, RunStatus + +from .database_retry import open_connection_with_retry +from .failure_recording import persist_failure_record +from .household_job_recording import ( + household_failure_summary, + unexpected_failure_summary, +) +from .report_run_access import MODAL_WORKER_DATABASE_URL_ENV + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + from uuid import UUID + + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, + HouseholdExecutionResult, + ) + from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportSide, + ) + from policyengine_stage12_persistence import RunWrite, SimulationRunState + +logger = logging.getLogger(__name__) + +# A report run whose side failed records this code, then the side's summary. +SIDE_FAILED_CODE = "household_calculation_failed" +# A report run whose side another call holds records this code. +SIDE_UNAVAILABLE_CODE = "household_simulation_run_unavailable" + +type HouseholdCalculation = Callable[ + [HouseholdExecutionRequest], + HouseholdExecutionResult | HouseholdExecutionFailure, +] + + +class HouseholdReportRunWriter(Protocol): + """A household report run's writes; ``HouseholdReportRunStore`` implements it. + + Each write is conditional on its row's status and, but for + ``fail_simulation_run`` given no call, on its recorded Modal call; each + returns whether it changed the row, and the row's status afterwards. + """ + + def claim_report_run( + self, report_run_id: UUID, *, modal_call_id: str, claimed_at: datetime + ) -> RunWrite: + """Mark a pending or running report run that records no other call running.""" + ... + + def household_report_simulation_runs( + self, report_run_id: UUID + ) -> tuple[UUID, UUID | None]: + """Return the baseline and reform simulation runs the report run names.""" + ... + + def claim_simulation_run( + self, simulation_run_id: UUID, *, modal_call_id: str, claimed_at: datetime + ) -> RunWrite: + """Mark a pending or running run that records no other call running.""" + ... + + def get_simulation_run(self, simulation_run_id: UUID) -> SimulationRunState: + """Return a simulation run's state, with its recorded call and failure.""" + ... + + def complete_household_simulation_run( + self, + simulation_run_id: UUID, + *, + modal_call_id: str, + household_result: Mapping[str, Any], + completed_at: datetime, + ) -> RunWrite: + """Mark the run succeeded with its household while this call holds it.""" + ... + + def fail_household_simulation_run( + self, + simulation_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + household_result: Mapping[str, Any] | None = None, + ) -> RunWrite: + """Mark the run failed while this call holds it.""" + ... + + def fail_simulation_run( + self, + simulation_run_id: UUID, + *, + error_message: str, + failed_at: datetime, + modal_call_id: str | None = None, + ) -> RunWrite: + """Mark the run failed: while ``modal_call_id`` holds it, or while active.""" + ... + + def complete_household_report_run( + self, report_run_id: UUID, *, modal_call_id: str, completed_at: datetime + ) -> RunWrite: + """Mark the report run succeeded while this call holds it.""" + ... + + def fail_report_run( + self, + report_run_id: UUID, + *, + modal_call_id: str, + error_message: str, + failed_at: datetime, + ) -> RunWrite: + """Mark the report run failed while this call holds it.""" + ... + + +def household_report_run_store( + runtime: ObservabilityRuntime | None = None, +) -> HouseholdReportRunStore: + """Build the household report run store from ``MODAL_WORKER_DATABASE_URL``. + + Raises ``ValueError`` when the variable is unset or not a complete + ``postgresql://`` URL. + """ + return HouseholdReportRunStore( + os.environ.get(MODAL_WORKER_DATABASE_URL_ENV, ""), + open_connection=partial(open_connection_with_retry, runtime=runtime), + ) + + +@dataclass(frozen=True) +class SideOutcome: + """What one side's simulation run ended as, and why it didn't succeed. + + ``reason`` is the report run's failure for a side whose run did not + succeed, and empty for one that did. + """ + + side: HouseholdReportSide + simulation_run_id: UUID + status: RunStatus + reason: str = "" + + def as_json(self) -> dict[str, str]: + """Return the side's outcome as the call answers it.""" + return { + "side": self.side, + "simulation_run_id": str(self.simulation_run_id), + "status": self.status.value, + } + + +@dataclass +class _ReportRunCall: + """One call's writes to the report run and its simulation runs.""" + + writer: HouseholdReportRunWriter + report_run_id: UUID + modal_call_id: str + runtime: ObservabilityRuntime | None + + def side( + self, + side: HouseholdReportSide, + request: HouseholdExecutionRequest, + calculate: HouseholdCalculation, + unfinished: list[UUID], + ) -> SideOutcome: + """Claim, calculate and record one side; return what its run ended as.""" + run_id = request.calculation_id + claimed = self.writer.claim_simulation_run( + run_id, modal_call_id=self.modal_call_id, claimed_at=datetime.now(UTC) + ) + if not claimed.changed: + unfinished.remove(run_id) + return self._side_already_held(side, run_id) + outcome = calculate(request) + document = outcome.model_dump(mode="json") + if isinstance(outcome, HouseholdExecutionFailure): + summary = household_failure_summary(outcome) + recorded = self.writer.fail_household_simulation_run( + run_id, + modal_call_id=self.modal_call_id, + error_message=summary, + failed_at=datetime.now(UTC), + household_result=document, + ) + else: + summary = None + recorded = self.writer.complete_household_simulation_run( + run_id, + modal_call_id=self.modal_call_id, + household_result=document, + completed_at=datetime.now(UTC), + ) + unfinished.remove(run_id) + if recorded.changed: + reason = "" if summary is None else _side_failure(side, summary) + else: + # Another writer finished the run this call held, so the call's + # outcome isn't the run's. + logger.warning( + "Simulation run %s was %s before its calculation was recorded", + run_id, + recorded.status.value, + ) + reason = ( + f"{SIDE_UNAVAILABLE_CODE}: The {side} simulation run was " + f"{recorded.status.value} before its calculation was recorded." + ) + return SideOutcome( + side=side, simulation_run_id=run_id, status=recorded.status, reason=reason + ) + + def _side_already_held( + self, side: HouseholdReportSide, run_id: UUID + ) -> SideOutcome: + """Return a side's run this call could not claim, as it stands. + + A restarted call finds a side it finished: a succeeded run is kept, + and a failed one keeps its recorded failure. A run still pending or + running holds another call, which this call leaves alone. + """ + state = self.writer.get_simulation_run(run_id) + if state.status is RunStatus.SUCCEEDED: + return SideOutcome(side=side, simulation_run_id=run_id, status=state.status) + if state.status is RunStatus.FAILED: + reason = _side_failure( + side, state.error_message or "no failure was recorded" + ) + else: + reason = ( + f"{SIDE_UNAVAILABLE_CODE}: The {side} simulation run is " + f"{state.status.value} under another call." + ) + return SideOutcome( + side=side, simulation_run_id=run_id, status=state.status, reason=reason + ) + + def finish(self, outcomes: list[SideOutcome]) -> RunStatus: + """Mark the report run succeeded, or failed with its first side that wasn't. + + A side succeeds only when its run did: one with any other status fails + the report run, with the side's reason. + """ + unsucceeded = next( + ( + outcome + for outcome in outcomes + if outcome.status is not RunStatus.SUCCEEDED + ), + None, + ) + if unsucceeded is None: + recorded = self.writer.complete_household_report_run( + self.report_run_id, + modal_call_id=self.modal_call_id, + completed_at=datetime.now(UTC), + ) + status = RunStatus.SUCCEEDED + else: + recorded = self.writer.fail_report_run( + self.report_run_id, + modal_call_id=self.modal_call_id, + error_message=unsucceeded.reason, + failed_at=datetime.now(UTC), + ) + status = RunStatus.FAILED + if not recorded.changed: + logger.warning( + "Report run %s was %s before its outcome was recorded", + self.report_run_id, + recorded.status.value, + ) + # Another writer finished it first: answer the run as it stands. + return recorded.status + return status + + def fail_unfinished(self, error: BaseException, *, unfinished: list[UUID]) -> None: + """Fail every run this call left unfinished, logging rather than raising. + + ``unfinished`` holds the report run's simulation runs this call has + neither finished nor found held by another call (``side`` takes each + out of it). Each is failed while still active, whether this call + holds it or its claim rolled back. The write checks no call: the + simulation service creates new simulation runs for each household + report, so only this call works on the runs left in the list. Then + the report run is failed while this call holds it. + """ + summary = unexpected_failure_summary(error) + for run_id in unfinished: + persist_failure_record( + partial( + self.writer.fail_simulation_run, + run_id, + error_message=summary, + failed_at=datetime.now(UTC), + ), + runtime=self.runtime, + failure_scope="household_simulation_run", + ) + persist_failure_record( + partial( + self.writer.fail_report_run, + self.report_run_id, + modal_call_id=self.modal_call_id, + error_message=summary, + failed_at=datetime.now(UTC), + ), + runtime=self.runtime, + failure_scope="report_run", + ) + + +def _require_named_runs( + report_run_id: UUID, + named: tuple[UUID, UUID | None], + calculation: HouseholdReportRunCalculation, +) -> None: + """Raise ValueError unless the payload's runs are the ones the report run names.""" + payload_runs = ( + calculation.baseline.calculation_id, + None if calculation.reform is None else calculation.reform.calculation_id, + ) + if payload_runs != named: + baseline, reform = named + msg = ( + f"report run {report_run_id} names other simulation runs than its " + f"call: baseline {baseline} and reform {reform}" + ) + raise ValueError(msg) + + +def _side_failure(side: HouseholdReportSide, summary: str) -> str: + """Return the report run's failure for a side that failed with ``summary``.""" + return f"{SIDE_FAILED_CODE}: The {side} simulation run failed: {summary}" + + +def run_household_report( # noqa: PLR0913 # the run, its payload and the call's seams + report_run_id: UUID, + payload: Mapping[str, Any], + calculate: HouseholdCalculation, + *, + modal_call_id: str | None, + runtime: ObservabilityRuntime | None = None, + writer: HouseholdReportRunWriter | None = None, +) -> dict[str, Any]: + """Claim the report run, calculate and record each side, then finish the run. + + ``payload`` is the call's ``HouseholdReportRunCalculation`` as sent; it + is validated once the call holds the report run, so a payload the + contract refuses fails the run at once. Returns the report run's and + each side's status, or, for a report run another call holds or one + already finished, a summary of the report run without calculating. + ``writer`` defaults to the store the Modal worker role's database URL + names. An error or Modal's cancellation fails the report run, and each + of its simulation runs the call left unfinished, with a generic summary, + then propagates; a pre-emption doesn't, because Modal runs the call + again under the same call ID. The simulation runs are known once the + call has read them, so an error in the report run's claim, or in that + read, fails the report run if the claim applied before the error, and + leaves its simulation runs to the stuck-run sweep. + Raises RuntimeError, before any write, when ``modal_call_id`` is None. + """ + if modal_call_id is None: + msg = "Modal household function call identifier is unavailable" + raise RuntimeError(msg) + call = _ReportRunCall( + writer=writer or household_report_run_store(runtime), + report_run_id=report_run_id, + modal_call_id=modal_call_id, + runtime=runtime, + ) + # The report run's simulation runs this call may still have to fail if + # it stops, as the database names them. + unfinished: list[UUID] = [] + try: + claimed = call.writer.claim_report_run( + report_run_id, modal_call_id=modal_call_id, claimed_at=datetime.now(UTC) + ) + if not claimed.changed: + return { + "deduplicated": True, + "report_run_id": str(report_run_id), + "status": claimed.status.value, + } + named = call.writer.household_report_simulation_runs(report_run_id) + unfinished.extend(run_id for run_id in named if run_id is not None) + calculation = HouseholdReportRunCalculation.model_validate(payload) + sides = calculation.sides() + _require_named_runs(report_run_id, named, calculation) + outcomes = [] + for side, request in sides: + outcomes.append(call.side(side, request, calculate, unfinished)) + status = call.finish(outcomes) + except KeyboardInterrupt: + # Pre-emption: Modal runs the call again, and it resumes its claims. + raise + except BaseException as error: + call.fail_unfinished(error, unfinished=unfinished) + raise + return { + "report_run_id": str(report_run_id), + "status": status.value, + "simulation_runs": [outcome.as_json() for outcome in outcomes], + } diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_household_reports.py b/projects/policyengine-simulation-executor/tests/test_stage12_household_reports.py new file mode 100644 index 000000000..d41275ee1 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_household_reports.py @@ -0,0 +1,694 @@ +"""A household function records the household report run it is given. + +It claims the report run under its Modal call ID, then each of the report +run's simulation runs, the baseline's and then the reform's: it calculates +each, records each run's household document, and then marks the report run +succeeded only if every side's run did, else failed. A report run another +call holds, or one already finished, is not calculated again, and a side a +restarted call already finished keeps its outcome. Modal pre-empts a call +with SIGINT, which a synchronous function receives as ``KeyboardInterrupt``, +and runs it again under the same call ID, so a pre-emption records nothing. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any +from uuid import UUID + +import pytest +from policyengine_simulation_contract.household_execution import ( + HouseholdCalculationError, + HouseholdExecutionFailure, + HouseholdExecutionRequest, + HouseholdExecutionResult, + HouseholdWorkerIdentity, +) +from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportRunCalculation, +) +from policyengine_simulation_executor.stage12_runtime import household_report_recording +from policyengine_stage12_persistence import RunStatus, RunWrite, SimulationRunState +from pydantic import ValidationError + +REPORT_RUN_ID = UUID("00000000-0000-4000-8000-00000000a001") +BASELINE_RUN_ID = UUID("00000000-0000-4000-8000-0000000000b1") +REFORM_RUN_ID = UUID("00000000-0000-4000-8000-0000000000b2") +CALL_ID = "fc-household-report" +WORKER = HouseholdWorkerIdentity( + application_name="policyengine-simulation-v2-py6-17-1", + function_name="calculate_household_us", +) +BUNDLE = { + "policyengine_version": "6.17.1", + "core_package_version": "3.33.0", + "country_package_name": "policyengine-us", + "country_package_version": "1.764.6", + "bundle_manifest_sha256": "a" * 64, +} +FAILURE_MESSAGE = "Unrecognized calculate input(s): x" +CHANGED = { + "running": RunWrite(changed=True, status=RunStatus.RUNNING), + "succeeded": RunWrite(changed=True, status=RunStatus.SUCCEEDED), + "failed": RunWrite(changed=True, status=RunStatus.FAILED), +} + + +def _request(calculation_id: UUID) -> dict[str, Any]: + return { + "calculation_id": str(calculation_id), + "country": "us", + "year": "2026", + "household": {"people": {"you": {"age": {"2026": 40}}}}, + "policy": {"parameters": {}}, + "bundle": BUNDLE, + } + + +def _calculation(*, with_reform: bool = True) -> HouseholdReportRunCalculation: + return HouseholdReportRunCalculation.model_validate( + { + "baseline": _request(BASELINE_RUN_ID), + "reform": _request(REFORM_RUN_ID) if with_reform else None, + } + ) + + +def _succeeded(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + return HouseholdExecutionResult.model_validate( + { + "calculation_id": str(request.calculation_id), + "output": { + "country": "us", + "bundle": BUNDLE, + "household": {"people": {"you": {"age": {"2026": 40}}}}, + }, + "worker": WORKER.model_dump(), + } + ) + + +def _failed(request: HouseholdExecutionRequest) -> HouseholdExecutionFailure: + return HouseholdExecutionFailure( + calculation_id=request.calculation_id, + error=HouseholdCalculationError(code="invalid_inputs", message=FAILURE_MESSAGE), + worker=WORKER, + ) + + +def _state( + status: RunStatus, *, error_message: str | None = None +) -> SimulationRunState: + return SimulationRunState( + status=status, + year=2026, + country_package_version="1.764.6", + policyengine_version="6.17.1", + dataset_id=None, + region_id=None, + output_uri=None, + output_sha256=None, + modal_call_id="fc-earlier", + error_message=error_message, + ) + + +@dataclass +class FakeHouseholdReportRuns: + """Records each write; a test sets what a claim returns or a read finds. + + ``simulation_claims`` and ``simulation_states`` are by simulation run ID; + a run with no entry is claimed, and read back running under this call. + """ + + report_claim: RunWrite = CHANGED["running"] + claim_error: Exception | None = None + # The baseline and reform simulation runs the report run names. + named_runs: tuple[UUID, UUID | None] = (BASELINE_RUN_ID, REFORM_RUN_ID) + read_error: Exception | None = None + report_record_result: RunWrite | None = None + simulation_claims: dict[UUID, RunWrite] = field(default_factory=dict) + simulation_states: dict[UUID, SimulationRunState] = field(default_factory=dict) + # Simulation runs whose claim raises, by run ID. + simulation_claim_errors: dict[UUID, Exception] = field(default_factory=dict) + record_result: RunWrite | None = None + writes: list[dict[str, Any]] = field(default_factory=list) + + def claim_report_run(self, report_run_id, *, modal_call_id, claimed_at): + self.writes.append( + {"write": "claim_report", "run": report_run_id, "call": modal_call_id} + ) + if self.claim_error is not None: + raise self.claim_error + return self.report_claim + + def household_report_simulation_runs(self, report_run_id): + self.writes.append({"write": "read_runs", "run": report_run_id}) + if self.read_error is not None: + raise self.read_error + return self.named_runs + + def claim_simulation_run(self, simulation_run_id, *, modal_call_id, claimed_at): + self.writes.append( + {"write": "claim", "run": simulation_run_id, "call": modal_call_id} + ) + if simulation_run_id in self.simulation_claim_errors: + raise self.simulation_claim_errors[simulation_run_id] + return self.simulation_claims.get(simulation_run_id, CHANGED["running"]) + + def get_simulation_run(self, simulation_run_id): + return self.simulation_states[simulation_run_id] + + def complete_household_simulation_run( + self, simulation_run_id, *, modal_call_id, household_result, completed_at + ): + self.writes.append( + { + "write": "complete", + "run": simulation_run_id, + "call": modal_call_id, + "household_result": household_result, + } + ) + return self.record_result or CHANGED["succeeded"] + + def fail_household_simulation_run( + self, + simulation_run_id, + *, + modal_call_id, + error_message, + failed_at, + household_result=None, + ): + self.writes.append( + { + "write": "fail", + "run": simulation_run_id, + "call": modal_call_id, + "error_message": error_message, + "household_result": household_result, + } + ) + return self.record_result or CHANGED["failed"] + + def fail_simulation_run( + self, simulation_run_id, *, error_message, failed_at, modal_call_id=None + ): + self.writes.append( + { + "write": "fail_unfinished", + "run": simulation_run_id, + "call": modal_call_id, + "error_message": error_message, + } + ) + return CHANGED["failed"] + + def complete_household_report_run( + self, report_run_id, *, modal_call_id, completed_at + ): + self.writes.append( + {"write": "complete_report", "run": report_run_id, "call": modal_call_id} + ) + return self.report_record_result or CHANGED["succeeded"] + + def fail_report_run( + self, report_run_id, *, modal_call_id, error_message, failed_at + ): + self.writes.append( + { + "write": "fail_report", + "run": report_run_id, + "call": modal_call_id, + "error_message": error_message, + } + ) + return CHANGED["failed"] + + def of(self, kind: str) -> list[dict[str, Any]]: + return [write for write in self.writes if write["write"] == kind] + + @property + def kinds(self) -> list[str]: + return [write["write"] for write in self.writes] + + +class InputCancelledForTest(BaseException): + """Stands in for Modal's InputCancellation, a BaseException but not an Exception.""" + + +def _run(runs, calculate, *, with_reform: bool = True, modal_call_id=CALL_ID): + return household_report_recording.run_household_report( + REPORT_RUN_ID, + _calculation(with_reform=with_reform).model_dump(mode="json"), + calculate, + modal_call_id=modal_call_id, + writer=runs, + ) + + +def test_both_sides_are_calculated_and_recorded_then_the_report_run_succeeds() -> None: + runs = FakeHouseholdReportRuns() + calculated: list[UUID] = [] + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + calculated.append(request.calculation_id) + return _succeeded(request) + + answer = _run(runs, calculate) + + assert calculated == [BASELINE_RUN_ID, REFORM_RUN_ID] + assert runs.kinds == [ + "claim_report", + "read_runs", + "claim", + "complete", + "claim", + "complete", + "complete_report", + ] + assert all(write["call"] == CALL_ID for write in runs.writes if "call" in write) + completed = runs.of("complete") + assert [write["run"] for write in completed] == [BASELINE_RUN_ID, REFORM_RUN_ID] + assert completed[0]["household_result"] == _succeeded( + _calculation().baseline + ).model_dump(mode="json") + assert answer == { + "report_run_id": str(REPORT_RUN_ID), + "status": "succeeded", + "simulation_runs": [ + { + "side": "baseline", + "simulation_run_id": str(BASELINE_RUN_ID), + "status": "succeeded", + }, + { + "side": "reform", + "simulation_run_id": str(REFORM_RUN_ID), + "status": "succeeded", + }, + ], + } + + +def test_a_report_of_the_baseline_alone_calculates_one_side() -> None: + runs = FakeHouseholdReportRuns(named_runs=(BASELINE_RUN_ID, None)) + + answer = _run(runs, _succeeded, with_reform=False) + + assert runs.kinds == [ + "claim_report", + "read_runs", + "claim", + "complete", + "complete_report", + ] + assert answer["status"] == "succeeded" + assert len(answer["simulation_runs"]) == 1 + + +def test_a_failed_side_is_recorded_and_fails_the_report_run_after_both_sides() -> None: + runs = FakeHouseholdReportRuns() + + def calculate(request: HouseholdExecutionRequest): + if request.calculation_id == BASELINE_RUN_ID: + return _failed(request) + return _succeeded(request) + + answer = _run(runs, calculate) + + (failed,) = runs.of("fail") + assert failed["run"] == BASELINE_RUN_ID + assert failed["error_message"] == f"invalid_inputs: {FAILURE_MESSAGE}" + assert failed["household_result"] == _failed(_calculation().baseline).model_dump( + mode="json" + ) + # The reform is still calculated, so each run has its own outcome. + assert [write["run"] for write in runs.of("complete")] == [REFORM_RUN_ID] + (failed_report,) = runs.of("fail_report") + assert failed_report["error_message"] == ( + "household_calculation_failed: The baseline simulation run failed: " + f"invalid_inputs: {FAILURE_MESSAGE}" + ) + assert "complete_report" not in runs.kinds + assert answer["status"] == "failed" + assert [run["status"] for run in answer["simulation_runs"]] == [ + "failed", + "succeeded", + ] + + +@pytest.mark.parametrize("status", [RunStatus.RUNNING, RunStatus.SUCCEEDED]) +def test_a_report_run_another_call_holds_or_that_finished_is_not_calculated( + status: RunStatus, +) -> None: + runs = FakeHouseholdReportRuns(report_claim=RunWrite(changed=False, status=status)) + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + pytest.fail("a report run this call does not hold is not calculated") + + answer = _run(runs, calculate) + + assert runs.kinds == ["claim_report"] + assert answer == { + "deduplicated": True, + "report_run_id": str(REPORT_RUN_ID), + "status": status.value, + } + + +def test_a_restarted_call_keeps_a_side_it_already_recorded() -> None: + runs = FakeHouseholdReportRuns( + simulation_claims={ + BASELINE_RUN_ID: RunWrite(changed=False, status=RunStatus.SUCCEEDED) + }, + simulation_states={BASELINE_RUN_ID: _state(RunStatus.SUCCEEDED)}, + ) + calculated: list[UUID] = [] + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + calculated.append(request.calculation_id) + return _succeeded(request) + + answer = _run(runs, calculate) + + assert calculated == [REFORM_RUN_ID] + assert [write["run"] for write in runs.of("complete")] == [REFORM_RUN_ID] + assert runs.kinds[-1] == "complete_report" + assert answer["status"] == "succeeded" + + +def test_a_side_that_already_failed_fails_the_report_run_with_its_failure() -> None: + recorded = f"invalid_inputs: {FAILURE_MESSAGE}" + runs = FakeHouseholdReportRuns( + simulation_claims={ + REFORM_RUN_ID: RunWrite(changed=False, status=RunStatus.FAILED) + }, + simulation_states={ + REFORM_RUN_ID: _state(RunStatus.FAILED, error_message=recorded) + }, + ) + + answer = _run(runs, _succeeded) + + assert [write["run"] for write in runs.of("complete")] == [BASELINE_RUN_ID] + assert runs.of("fail") == [] + (failed_report,) = runs.of("fail_report") + assert failed_report["error_message"] == ( + "household_calculation_failed: The reform simulation run failed: " + recorded + ) + assert answer["status"] == "failed" + + +def test_a_side_another_call_holds_is_not_calculated_and_fails_the_report_run() -> None: + runs = FakeHouseholdReportRuns( + simulation_claims={ + REFORM_RUN_ID: RunWrite(changed=False, status=RunStatus.RUNNING) + }, + simulation_states={REFORM_RUN_ID: _state(RunStatus.RUNNING)}, + ) + calculated: list[UUID] = [] + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + calculated.append(request.calculation_id) + return _succeeded(request) + + answer = _run(runs, calculate) + + assert calculated == [BASELINE_RUN_ID] + (failed_report,) = runs.of("fail_report") + assert failed_report["error_message"] == ( + "household_simulation_run_unavailable: The reform simulation run is " + "running under another call." + ) + assert answer["status"] == "failed" + + +@pytest.mark.parametrize( + ("error", "summary"), + [ + ( + RuntimeError("the model could not load"), + "household_function_error: The household function raised RuntimeError.", + ), + ( + InputCancelledForTest(), + "household_function_cancelled: Modal cancelled the household " + "function's call.", + ), + ], + ids=["error", "cancellation"], +) +def test_an_unexpected_error_fails_every_unfinished_run_then_propagates( + error: BaseException, summary: str +) -> None: + runs = FakeHouseholdReportRuns() + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + raise error + + with pytest.raises(type(error)): + _run(runs, calculate) + + # The error names only its type: its text may repeat the household. Each + # unfinished run is failed while still active, held by this call or not. + assert runs.of("fail") == [] + unfinished = runs.of("fail_unfinished") + assert [write["run"] for write in unfinished] == [BASELINE_RUN_ID, REFORM_RUN_ID] + assert {write["error_message"] for write in unfinished} == {summary} + assert {write["call"] for write in unfinished} == {None} + (failed_report,) = runs.of("fail_report") + assert (failed_report["call"], failed_report["error_message"]) == (CALL_ID, summary) + + +def test_a_side_whose_claim_raises_is_failed_though_this_call_may_not_hold_it() -> None: + # The claim may have rolled back, leaving the run pending with no call. + runs = FakeHouseholdReportRuns( + simulation_claim_errors={BASELINE_RUN_ID: ConnectionError("lost")} + ) + + with pytest.raises(ConnectionError): + _run(runs, _succeeded) + + assert runs.kinds == [ + "claim_report", + "read_runs", + "claim", + "fail_unfinished", + "fail_unfinished", + "fail_report", + ] + assert [write["run"] for write in runs.of("fail_unfinished")] == [ + BASELINE_RUN_ID, + REFORM_RUN_ID, + ] + assert {write["call"] for write in runs.of("fail_unfinished")} == {None} + + +def _reform_raises(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + if request.calculation_id == REFORM_RUN_ID: + msg = "the model could not load" + raise RuntimeError(msg) + return _succeeded(request) + + +@pytest.mark.parametrize( + "baseline", + ["recorded", "held-by-another-call"], +) +def test_an_error_leaves_a_side_this_call_finished_or_does_not_hold_alone( + baseline: str, +) -> None: + # Only the reform is left to fail: the failure write checks no call, so + # a side out of the list must never reach it. + runs = ( + FakeHouseholdReportRuns() + if baseline == "recorded" + else FakeHouseholdReportRuns( + simulation_claims={ + BASELINE_RUN_ID: RunWrite(changed=False, status=RunStatus.RUNNING) + }, + simulation_states={BASELINE_RUN_ID: _state(RunStatus.RUNNING)}, + ) + ) + + with pytest.raises(RuntimeError): + _run(runs, _reform_raises) + + assert [write["run"] for write in runs.of("fail_unfinished")] == [REFORM_RUN_ID] + assert len(runs.of("fail_report")) == 1 + + +@pytest.mark.parametrize( + ("named_runs", "with_reform"), + [((BASELINE_RUN_ID, None), True), ((BASELINE_RUN_ID, REFORM_RUN_ID), False)], + ids=["payload-adds-a-reform", "payload-leaves-out-the-reform"], +) +def test_a_payload_whose_reform_differs_from_the_report_runs_is_refused( + named_runs: tuple[UUID, UUID | None], *, with_reform: bool +) -> None: + runs = FakeHouseholdReportRuns(named_runs=named_runs) + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + pytest.fail("a payload naming other runs is not calculated") + + with pytest.raises(ValueError, match="names other simulation runs"): + _run(runs, calculate, with_reform=with_reform) + + assert [write["run"] for write in runs.of("fail_unfinished")] == [ + run_id for run_id in named_runs if run_id is not None + ] + assert len(runs.of("fail_report")) == 1 + + +def test_a_preemption_records_nothing_and_propagates() -> None: + runs = FakeHouseholdReportRuns() + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + _run(runs, calculate) + + assert runs.kinds == ["claim_report", "read_runs", "claim"] + + +def test_a_report_run_naming_other_simulation_runs_fails_with_its_own_runs() -> None: + own_runs = ( + UUID("00000000-0000-4000-8000-0000000000d1"), + UUID("00000000-0000-4000-8000-0000000000d2"), + ) + runs = FakeHouseholdReportRuns(named_runs=own_runs) + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + pytest.fail("a report run naming other simulation runs is not calculated") + + with pytest.raises(ValueError, match="names other simulation runs"): + _run(runs, calculate) + + summary = "household_function_error: The household function raised ValueError." + assert runs.kinds == [ + "claim_report", + "read_runs", + "fail_unfinished", + "fail_unfinished", + "fail_report", + ] + # The report run's own runs are failed, never the payload's. + assert [write["run"] for write in runs.of("fail_unfinished")] == list(own_runs) + assert {write["error_message"] for write in runs.of("fail_unfinished")} == {summary} + assert runs.of("fail_report")[0]["error_message"] == summary + + +def test_a_report_run_whose_runs_cannot_be_read_is_failed_alone() -> None: + runs = FakeHouseholdReportRuns(read_error=ConnectionError("connection lost")) + + with pytest.raises(ConnectionError): + _run(runs, _succeeded) + + assert runs.kinds == ["claim_report", "read_runs", "fail_report"] + + +def test_a_report_run_another_writer_finished_is_answered_as_it_stands( + caplog: pytest.LogCaptureFixture, +) -> None: + runs = FakeHouseholdReportRuns( + report_record_result=RunWrite(changed=False, status=RunStatus.FAILED) + ) + + answer = _run(runs, _succeeded) + + assert answer["status"] == "failed" + assert f"Report run {REPORT_RUN_ID} was failed before" in caplog.text + + +def test_a_call_without_a_modal_call_id_writes_nothing() -> None: + runs = FakeHouseholdReportRuns() + + with pytest.raises(RuntimeError, match="call identifier is unavailable"): + _run(runs, _succeeded, modal_call_id=None) + + assert runs.writes == [] + + +def test_a_side_whose_record_changes_nothing_fails_the_report_run( + caplog: pytest.LogCaptureFixture, +) -> None: + # Another writer finished the run this call held before its record. + runs = FakeHouseholdReportRuns( + named_runs=(BASELINE_RUN_ID, None), + record_result=RunWrite(changed=False, status=RunStatus.FAILED), + ) + + answer = _run(runs, _succeeded, with_reform=False) + + assert answer["simulation_runs"][0]["status"] == "failed" + assert f"Simulation run {BASELINE_RUN_ID} was failed before" in caplog.text + assert "complete_report" not in runs.kinds + (failed_report,) = runs.of("fail_report") + assert failed_report["error_message"] == ( + "household_simulation_run_unavailable: The baseline simulation run was " + "failed before its calculation was recorded." + ) + assert answer["status"] == "failed" + + +def test_a_side_found_failed_with_no_recorded_failure_fails_the_report_run() -> None: + runs = FakeHouseholdReportRuns( + simulation_claims={ + REFORM_RUN_ID: RunWrite(changed=False, status=RunStatus.FAILED) + }, + simulation_states={REFORM_RUN_ID: _state(RunStatus.FAILED)}, + ) + + answer = _run(runs, _succeeded) + + (failed_report,) = runs.of("fail_report") + assert failed_report["error_message"] == ( + "household_calculation_failed: The reform simulation run failed: " + "no failure was recorded" + ) + assert answer["status"] == "failed" + + +def test_a_payload_the_contract_refuses_fails_the_report_run_it_claimed() -> None: + runs = FakeHouseholdReportRuns() + + def calculate(request: HouseholdExecutionRequest) -> HouseholdExecutionResult: + pytest.fail("a refused payload is not calculated") + + with pytest.raises(ValidationError): + household_report_recording.run_household_report( + REPORT_RUN_ID, + {"baseline": {"calculation_id": "not a request"}}, + calculate, + modal_call_id=CALL_ID, + writer=runs, + ) + + summary = "household_function_error: The household function raised ValidationError." + # The report run's own simulation runs fail with it, read from the database. + assert runs.kinds == [ + "claim_report", + "read_runs", + "fail_unfinished", + "fail_unfinished", + "fail_report", + ] + assert [write["run"] for write in runs.of("fail_unfinished")] == [ + BASELINE_RUN_ID, + REFORM_RUN_ID, + ] + assert runs.of("fail_report")[0]["error_message"] == summary + + +def test_a_claim_that_raises_fails_the_report_run_it_may_hold_then_propagates() -> None: + # The claim's commit may have applied before its connection failed. + runs = FakeHouseholdReportRuns(claim_error=ConnectionError("connection lost")) + + with pytest.raises(ConnectionError): + _run(runs, _succeeded) + + assert runs.kinds == ["claim_report", "fail_report"] + assert runs.of("fail_report")[0]["call"] == CALL_ID diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index f4be528b8..e3f6ef998 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -731,6 +731,111 @@ def run_household_job(household_job_id, calculate_job, **keywords): assert runtime.shutdown_calls == 1 +def _household_report_payload(country: str) -> dict: + """Return a household report run's calculation: a baseline and a reform.""" + reform = { + **_household_payload(country), + "calculation_id": "00000000-0000-4000-8000-0000000000c2", + } + return {"baseline": _household_payload(country), "reform": reform} + + +@pytest.mark.parametrize("country", ["us", "uk"]) +def test_household_functions_record_the_report_run_their_call_names( + monkeypatch, country +) -> None: + from policyengine_simulation_contract.household_report_execution import ( + HouseholdReportRunCalculation, + ) + from policyengine_simulation_executor import ( + household_calculation, + stage12_runtime, + ) + + module = _load(monkeypatch) + runtime = RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "fc-household-report", + raising=False, + ) + _refuse_household_jobs(monkeypatch) + calculated: list[dict] = [] + recorded: list[dict] = [] + + def calculate(request, **keywords): + calculated.append({"request": request, **keywords}) + return request.calculation_id + + def run_household_report(report_run_id, payload, calculate_side, **keywords): + recorded.append( + {"report_run_id": report_run_id, "payload": payload, **keywords} + ) + calculation = HouseholdReportRunCalculation.model_validate(payload) + outcomes = [calculate_side(request) for _, request in calculation.sides()] + return {"status": "succeeded", "outcomes": [str(item) for item in outcomes]} + + monkeypatch.setattr(household_calculation, "calculate_household", calculate) + monkeypatch.setattr(stage12_runtime, "run_household_report", run_household_report) + + result = getattr(module, f"calculate_household_{country}")( + _household_report_payload(country), + observability_context=_observability_context(), + household_report_run_id="00000000-0000-4000-8000-00000000a002", + ) + + assert result == { + "status": "succeeded", + "outcomes": [ + "00000000-0000-4000-8000-0000000000c1", + "00000000-0000-4000-8000-0000000000c2", + ], + } + (report,) = recorded + assert report["report_run_id"] == UUID("00000000-0000-4000-8000-00000000a002") + # The payload reaches the recording code as sent, which validates it once + # the call holds the report run. + assert report["payload"] == _household_report_payload(country) + assert report["modal_call_id"] == "fc-household-report" + assert report["runtime"] is runtime + assert [call["required_country"] for call in calculated] == [country, country] + assert all( + call["worker"].function_name == f"calculate_household_{country}" + for call in calculated + ) + # The claims, both calculations and the records are one operation. + (operation,) = runtime.operations + assert operation["name"] == "v2_household_execution" + assert runtime.shutdown_calls == 1 + + +def test_a_household_function_refuses_a_job_and_a_report_run_together( + monkeypatch, +) -> None: + from policyengine_simulation_executor import stage12_runtime + + module = _load(monkeypatch) + runtime = RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + _refuse_household_jobs(monkeypatch) + + def refuse(*_args, **_keywords): + msg = "a call naming a job and a report run recorded one" + raise AssertionError(msg) + + monkeypatch.setattr(stage12_runtime, "run_household_report", refuse) + + with pytest.raises(ValueError, match="either a household job or a report run"): + module.calculate_household_us( + _household_report_payload("us"), + household_job_id="00000000-0000-4000-8000-00000000a001", + household_report_run_id="00000000-0000-4000-8000-00000000a002", + ) + assert runtime.shutdown_calls == 1 + + def test_household_functions_given_no_job_record_none(monkeypatch) -> None: # Release validation calls them with validation households, and no job. from policyengine_simulation_executor import household_calculation diff --git a/ruff.toml b/ruff.toml index 38c26d845..325e559de 100644 --- a/ruff.toml +++ b/ruff.toml @@ -110,6 +110,7 @@ runtime-evaluated-base-classes = [ "policyengine_fastapi.observability.contracts.CorrelatedRunFields", "policyengine_fastapi.observability.contracts.ObservabilityModel", "policyengine_simulation_contract.gateway_models.GatewayRequestBase", + "policyengine_simulation_contract.household_execution.StrictHouseholdModel", "policyengine_simulation_contract.macro_output.MacroOutputModel", "policyengine_simulation_contract.macro_output.MacroRootModel", "policyengine_simulation_contract.macro_output.SingleYearMacroOutput",