diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index cbe60559e..6f0108438 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -102,6 +102,7 @@ required_variables=( SIMULATION_SERVICE_DATABASE_URL_SECRET_NAME MODAL_TOKEN_ID_SECRET_NAME MODAL_TOKEN_SECRET_SECRET_NAME + V2_SIMULATION_OUTPUT_BUCKET_VALUE OBSERVABILITY_SERVICE_NAMESPACE OBSERVABILITY_TRACE_PROJECT_ID OTEL_EXPORTER_OTLP_ENDPOINT @@ -128,6 +129,11 @@ esac || fail "TAG is invalid" [[ "${ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT}" =~ ^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.gserviceaccount\.com$ ]] \ || fail "ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT is invalid" +# The bucket that keeps v2 simulation output files for good: a report +# reuses only a succeeded run whose file is in it. The service compares +# output URIs with it and never reads it. +[[ "${V2_SIMULATION_OUTPUT_BUCKET_VALUE}" =~ ^[a-z0-9][a-z0-9._-]{1,61}[a-z0-9]$ ]] \ + || fail "V2_SIMULATION_OUTPUT_BUCKET_VALUE is invalid" [[ "${ENTRYPOINT_MIN_INSTANCES}" =~ ^[0-9]+$ ]] \ || fail "ENTRYPOINT_MIN_INSTANCES must be a non-negative integer" [[ "${ENTRYPOINT_MAX_INSTANCES}" =~ ^[0-9]+$ ]] \ @@ -195,6 +201,7 @@ jq -n ' POLICYENGINE_OTEL_GOOGLE_AUDIENCE: env.POLICYENGINE_OTEL_GOOGLE_AUDIENCE, SIMULATION_V2_MODAL_ENVIRONMENT: env.MODAL_ENVIRONMENT, SIMULATION_V2_MANIFEST_NAME: "simulation-api-v2-version-manifest", + V2_SIMULATION_OUTPUT_BUCKET: env.V2_SIMULATION_OUTPUT_BUCKET_VALUE, STAGE12_ENABLED: ( if env.DEPLOY_STAGE12_V2 == "true" then env.STAGE12_ENABLED_VALUE else "0" end ) diff --git a/.github/scripts/simulation-staging-economy-check.sh b/.github/scripts/simulation-staging-economy-check.sh new file mode 100755 index 000000000..ae24e8673 --- /dev/null +++ b/.github/scripts/simulation-staging-economy-check.sh @@ -0,0 +1,136 @@ +#!/bin/bash +# Create a small economy report through a deployed Simulation Entrypoint, and +# require its run to succeed: current law against current law for Wales in +# 2026, which needs no stored policy. +# Usage: simulation-staging-economy-check.sh +# Required env vars: the test client's SIMULATION_TEST_AUTH_ISSUER, +# SIMULATION_TEST_AUTH_AUDIENCE, SIMULATION_TEST_AUTH_CLIENT_ID and +# SIMULATION_TEST_AUTH_CLIENT_SECRET. +# Optional: ECONOMY_CHECK_DEADLINE_SECONDS (1800) and +# ECONOMY_CHECK_POLL_SECONDS (15). +# The access token is sent from a file and the client secret on standard +# input, so neither appears on a command line. The submission, which isn't +# asked again, may take 60 seconds, Cloud Run's own limit; every other +# request 30. A read that gets no answer, a 429 or any 5xx is asked again +# until the deadline. + +set -euo pipefail + +fail() { + printf 'Staging economy check failed: %s\n' "$1" >&2 + exit 1 +} + +base_url="${1:?Simulation Entrypoint base URL is required}" +base_url="${base_url%/}" +deadline_seconds="${ECONOMY_CHECK_DEADLINE_SECONDS:-1800}" +poll_seconds="${ECONOMY_CHECK_POLL_SECONDS:-15}" +request_timeout_seconds=30 +submission_timeout_seconds=60 +uuid_pattern='^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$' + +for variable_name in SIMULATION_TEST_AUTH_ISSUER SIMULATION_TEST_AUTH_AUDIENCE \ + SIMULATION_TEST_AUTH_CLIENT_ID SIMULATION_TEST_AUTH_CLIENT_SECRET; do + [[ -n "${!variable_name:-}" ]] || fail "${variable_name} must be set" +done +[[ "${deadline_seconds}" =~ ^[0-9]+$ ]] \ + || fail "ECONOMY_CHECK_DEADLINE_SECONDS must be a whole number of seconds" +[[ "${poll_seconds}" =~ ^[0-9]+$ ]] \ + || fail "ECONOMY_CHECK_POLL_SECONDS must be a whole number of seconds" + +work_directory="$(mktemp -d)" +trap 'rm -rf "${work_directory}"' EXIT +chmod 700 "${work_directory}" + +# Send one request, allowing it ; write its body to .json and +# print its HTTP status. It fails when no answer arrives. +send() { + local name="$1" + local seconds="$2" + shift 2 + curl --silent --show-error --max-time "${seconds}" \ + --output "${work_directory}/${name}.json" --write-out '%{http_code}' "$@" +} + +token_status="$( + jq -n \ + --arg client_id "${SIMULATION_TEST_AUTH_CLIENT_ID}" \ + --arg client_secret "${SIMULATION_TEST_AUTH_CLIENT_SECRET}" \ + --arg audience "${SIMULATION_TEST_AUTH_AUDIENCE}" \ + '{client_id: $client_id, client_secret: $client_secret, audience: $audience, grant_type: "client_credentials"}' \ + | send token "${request_timeout_seconds}" --request POST --header 'Content-Type: application/json' \ + --data @- "${SIMULATION_TEST_AUTH_ISSUER%/}/oauth/token" +)" || fail "the token request got no answer" +[[ "${token_status}" == "200" ]] || fail "the token request answered ${token_status}" +access_token="$(jq -r '.access_token // empty' "${work_directory}/token.json")" +[[ -n "${access_token}" ]] || fail "the token response has no access_token" +authorization_header="${work_directory}/authorization-header" +printf 'Authorization: Bearer %s\n' "${access_token}" > "${authorization_header}" + +submission_status="$( + jq -n '{ + country: "uk", + report_type: "economy", + label: "Staging economy check", + year: 2026, + region_code: "country/wales", + baseline: {policy_id: null, dynamic_id: null}, + reform: {policy_id: null, dynamic_id: null}, + include_cliffs: false + }' \ + | send submission "${submission_timeout_seconds}" --request POST --header "@${authorization_header}" \ + --header 'Content-Type: application/json' --data @- \ + "${base_url}/reports" +)" || fail "POST /reports got no answer; it may have created a report" +[[ "${submission_status}" == "201" ]] \ + || fail "POST /reports answered ${submission_status}: $(cat "${work_directory}/submission.json")" +submission="${work_directory}/submission.json" +report_id="$(jq -r '.report_id // empty' "${submission}")" +report_run_id="$(jq -r '.report_run_id // empty' "${submission}")" +[[ "${report_id}" =~ ${uuid_pattern} ]] || fail "the submission named no report" +[[ "${report_run_id}" =~ ${uuid_pattern} ]] || fail "the submission named no report run" +printf 'Created report %s, run %s (baseline reused: %s, reform reused: %s)\n' \ + "${report_id}" "${report_run_id}" \ + "$(jq -r '.baseline_reused' "${submission}")" "$(jq -r '.reform_reused' "${submission}")" + +started_at="${SECONDS}" +# What the last read found, for the deadline's failure message. +last_read="nothing yet" +while true; do + report_body="${work_directory}/report.json" + if ! read_status="$( + send report "${request_timeout_seconds}" --header "@${authorization_header}" \ + "${base_url}/reports/${report_id}" + )"; then + last_read="it got no answer" + printf 'GET /reports/%s got no answer; asking again\n' "${report_id}" >&2 + elif [[ "${read_status}" == "200" ]]; then + newest_run_id="$(jq -r '.run.run_id // empty' "${report_body}")" + [[ "${newest_run_id}" == "${report_run_id}" ]] \ + || fail "report ${report_id}'s newest run is ${newest_run_id:-missing}, not the created run ${report_run_id}" + run_status="$(jq -r '.run.status' "${report_body}")" + last_read="the run was ${run_status}" + case "${run_status}" in + succeeded) + printf 'Report run %s succeeded on .py %s with %s\n' "${report_run_id}" \ + "$(jq -r '.run.policyengine_version' "${report_body}")" \ + "$(jq -r '.run.country_package_version' "${report_body}")" + exit 0 + ;; + failed) + fail "report run ${report_run_id} failed: $(jq -r '.run.error_message // "no error message"' "${report_body}")" + ;; + pending | running) ;; + *) fail "report run ${report_run_id} has an unknown status: ${run_status}" ;; + esac + elif [[ "${read_status}" == "429" || "${read_status}" == 5?? ]]; then + # The service, or Cloud Run's front end in its place, can't answer now. + last_read="it answered ${read_status}" + printf 'GET /reports/%s answered %s; asking again\n' "${report_id}" "${read_status}" >&2 + else + fail "GET /reports/${report_id} answered ${read_status}: $(cat "${report_body}")" + fi + ((SECONDS - started_at < deadline_seconds)) \ + || fail "report run ${report_run_id} did not finish within ${deadline_seconds} seconds (last read: ${last_read})" + sleep "${poll_seconds}" +done diff --git a/.github/scripts/simulation-staging-economy-ready.sh b/.github/scripts/simulation-staging-economy-ready.sh new file mode 100755 index 000000000..a10c4f66a --- /dev/null +++ b/.github/scripts/simulation-staging-economy-ready.sh @@ -0,0 +1,13 @@ +#!/bin/bash +# Require the staging economy check to have succeeded in a staging +# deployment, and to have been skipped in any other. +# Required env vars: REQUESTED ("true" in a staging deployment) and +# CHECK_RESULT (the staging_economy_check job result). + +set -euo pipefail + +if [[ "${REQUESTED}" == "true" ]]; then + test "${CHECK_RESULT}" = "success" +else + test "${CHECK_RESULT}" = "skipped" +fi diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 1775e6ddd..b89de5283 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -247,6 +247,7 @@ jobs: STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} MODAL_TOKEN_ID_SECRET_NAME: ${{ vars.MODAL_TOKEN_ID_SECRET_NAME }} MODAL_TOKEN_SECRET_SECRET_NAME: ${{ vars.MODAL_TOKEN_SECRET_SECRET_NAME }} + V2_SIMULATION_OUTPUT_BUCKET_VALUE: ${{ vars.V2_SIMULATION_OUTPUT_BUCKET }} run: .github/scripts/deploy-cloud-run-simulation-entry.sh - name: Resolve exact candidate metadata @@ -1252,6 +1253,43 @@ jobs: SIMULATION_TEST_AUTH_CLIENT_SECRET: ${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET }} run: .github/scripts/simulation-staging-household-check.sh "${{ needs.deploy_entrypoint.outputs.candidate_url }}" + staging_economy_check: + name: Create a staging economy report through the candidate + # Every staging deployment creates a small economy report, current law + # against current law for Wales in 2026, through the candidate's + # POST /reports, and requires its run to succeed. The routes use only + # the v2 manifest's default worker, so the check runs after the default + # has moved to the release's worker (or, in a deployment without Stage + # 12, with the default as it is). It calls the tagged candidate URL, + # which serves this release's revision whether or not it was promoted. + # The report's rows are written in the staging database, and its + # coordinator records the run. Once a .py version's first check has + # succeeded, later checks on it reuse those simulation runs. + if: >- + ${{ + !cancelled() && + inputs.deployment_environment == 'staging' && + needs.deploy_entrypoint.result == 'success' && + needs.authenticated_test.result == 'success' && + ((inputs.deploy_stage12_v2 && needs.publish_stage12_v2_default.result == 'success') || + (!inputs.deploy_stage12_v2 && needs.publish_stage12_v2_default.result == 'skipped')) + }} + needs: [deploy_entrypoint, authenticated_test, publish_stage12_v2_default] + runs-on: ubuntu-latest + # Longer than the script's own 30-minute deadline, so the script fails first. + timeout-minutes: 40 + environment: ${{ inputs.release_environment }} + steps: + - uses: actions/checkout@v6 + + - name: Create the staging economy report and require its run to succeed + env: + SIMULATION_TEST_AUTH_ISSUER: ${{ secrets.SIMULATION_ENTRYPOINT_AUTH_ISSUER }} + SIMULATION_TEST_AUTH_AUDIENCE: ${{ secrets.SIMULATION_ENTRYPOINT_AUTH_AUDIENCE }} + SIMULATION_TEST_AUTH_CLIENT_ID: ${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_ID }} + SIMULATION_TEST_AUTH_CLIENT_SECRET: ${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET }} + run: .github/scripts/simulation-staging-economy-check.sh "${{ needs.deploy_entrypoint.outputs.candidate_url }}" + finalize_routing: name: Finalize the Modal routing publication # Finish the durable journal after every gate, including failed uploads @@ -1310,7 +1348,7 @@ jobs: # The reporter only reads job results. Journal completion happens before # it, so a failure or cancellation here cannot strand a publication. if: ${{ !cancelled() }} - needs: [deploy_entrypoint, update_routing, integration, authenticated_test, promote_entrypoint, publish_latest, finalize_routing, publish_stage12_v2_default, staging_household_check] + needs: [deploy_entrypoint, update_routing, integration, authenticated_test, promote_entrypoint, publish_latest, finalize_routing, publish_stage12_v2_default, staging_household_check, staging_economy_check] runs-on: ubuntu-latest defaults: run: @@ -1345,6 +1383,12 @@ jobs: CHECK_RESULT: ${{ needs.staging_household_check.result }} run: .github/scripts/simulation-staging-household-ready.sh + - name: Check the staging economy check result + env: + REQUESTED: ${{ inputs.deployment_environment == 'staging' }} + CHECK_RESULT: ${{ needs.staging_economy_check.result }} + run: .github/scripts/simulation-staging-economy-ready.sh + cleanup_entrypoint_revisions: name: Remove superseded Cloud Run revisions if: ${{ !cancelled() && inputs.promote_entrypoint && needs.deployment_ready.result == 'success' }} diff --git a/changelog.d/c1-9-economy-reports.added.md b/changelog.d/c1-9-economy-reports.added.md new file mode 100644 index 000000000..4d3d782c4 --- /dev/null +++ b/changelog.d/c1-9-economy-reports.added.md @@ -0,0 +1 @@ +The Simulation Entrypoint creates economy reports and simulations through `POST /reports` and `POST /simulations`, in the Public API's run tables, and starts their calculation on the v2 coordinator and workers by run ID; a report reuses a matching simulation run that already succeeded. `GET /reports/{report_id}` and `GET /simulations/{simulation_id}` return each record's newest run, and every staging deploy checks a small economy report. diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index a049a0575..4fea4fd37 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -50,9 +50,9 @@ year's report run names that row and its year, and belongs to the same report. ## Input tables `input_tables.py` holds read-only models of the Public API's tables a -calculation is read from: `households`, `policies`, `parameters`, -`parameter_values`, `tax_benefit_models` and `tax_benefit_model_versions`, in -`input_metadata`. The service only reads them, with `SELECT`. Each names only +calculation is read from: `households`, `policies`, `dynamics`, `parameters`, +`parameter_values`, `regions`, `datasets`, `tax_benefit_models` and +`tax_benefit_model_versions`, in `input_metadata`. The service only reads them, with `SELECT`. Each names only the columns the service reads, so they are not models of the whole tables, and they are kept apart from the models of the tables the service writes (`service_metadata`) and from the tables those name by `id` alone @@ -81,6 +81,50 @@ which commits when its block ends and sets `statement_timeout` with `set_config`, local to the transaction, so the database ends a statement that runs too long and the limit never stays on a pooled connection. +## Economy simulations and reports + +`economy_run_records.py` holds the simulation service's reads and writes for +an economy report, or a standalone economy simulation. It reads the +country's model (`read_model_id`) and a catalog version's region by code +(`read_region`) and dataset by name or ID (`read_dataset`, +`read_dataset_by_id`). `insert_economy_report` creates, in the caller's +transaction, each side's new simulation and first run, the report naming +both simulations and the reform's policy with its options in `inputs`, and +its first report run, trigger `initial`; every run is `pending` with no +Modal call ID. A side may instead name a run that already succeeded +(`ReportSimulation.reused`), for which nothing is created. +`insert_economy_simulation` creates a standalone simulation and its first +run. A row naming one that does not exist is refused by its foreign key with +SQLSTATE `23503`; `EconomyRecordReferenceError` names the table and column. + +`find_reusable_run` returns the run a new report's side may reuse: the +newest, by completion, of the same economy simulation (policy, dynamic, +catalog version, dataset, region and year) that succeeded on the same +country package and .py versions, whose output file is under the output +bucket's `gs:///`, that a succeeded report run with the same options +used, and that no report run refused to reuse (an error message starting +`reused_simulation_`). So a submission after a refused reuse calculates the +simulation again. + +`record_report_run_call` and `record_simulation_run_call` record a Modal +call's ID, and `fail_unstarted_report_run` and +`fail_unstarted_simulation_run` fail a run at once, each only while it is +`pending` with no call ID; failing a report run also fails the simulation +runs created for it. `economy_run_reads.py` reads a report, or a simulation, +with its newest run's state. + +`EconomyRunStore` runs them as the Simulation Entrypoint does, each step in +a transaction of its own: `read_report_inputs` reads the catalog version, +model, region and dataset, each side's policy and dynamic values, and each +side's reusable run together, never giving both sides the same run; +`read_simulation_inputs` does the same for one simulation; the writes and +reads each take their own. A missing region or dataset raises +`UnknownRegionError` or `UnknownDatasetError`, and a report side's refused +policy or dynamic raises `ReportSideError`, naming the side. A create whose +commit fails after its rows were inserted raises +`EconomyRecordOutcomeUnknownError`: PostgreSQL may have applied the commit, +so the rows may exist. + ## Reading a policy `read_policy_reform` (`policy_reform.py`) reads a stored policy as the reform @@ -91,6 +135,9 @@ the parameters of the catalog version it was saved on, so a calculation on any catalog version applies them by name; a parameter the calculation's version lacks raises `MissingPolicyParametersError`, naming each one, and a policy the country doesn't have raises `UnknownPolicyError`. +`read_dynamic_reform` reads a stored dynamic the same way (`UnknownDynamicError` +for none), and `combine_reforms` joins a simulation's policy and dynamic +values, refusing a parameter both set (`OverlappingReformError`). Each period is a range, `YYYY-MM-DD.YYYY-MM-DD`, with the stored dates read in UTC. A value with no end date follows .py's own rule diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py index 11f4713b4..1741cf7b0 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/__init__.py @@ -6,6 +6,52 @@ missing_table_privileges, read_missing_table_privileges, ) +from policyengine_simulation_persistence.economy_run_reads import ( + EconomyReportState, + EconomySimulationState, + ReportRunState, + SimulationRunState, + read_economy_report, + read_economy_simulation, +) +from policyengine_simulation_persistence.economy_run_records import ( + CatalogDataset, + CatalogRegion, + EconomyPopulation, + EconomyRecordReferenceError, + NewEconomyReport, + NewEconomySimulation, + ReportSimulation, + ReusableRun, + RunVersions, + SimulationChoice, + fail_unstarted_report_run, + fail_unstarted_simulation_run, + find_reusable_run, + insert_economy_report, + insert_economy_simulation, + read_dataset, + read_dataset_by_id, + read_model_id, + read_region, + record_report_run_call, + record_simulation_run_call, +) +from policyengine_simulation_persistence.economy_run_store import ( + EconomyCatalogRows, + EconomyRecordOutcomeUnknownError, + EconomyReportInputs, + EconomyRequest, + EconomyRunStore, + EconomySimulationInputs, + ReportSide, + ReportSideError, + RunReuse, + SideRefusal, + SimulationInputs, + UnknownDatasetError, + UnknownRegionError, +) from policyengine_simulation_persistence.engine import ( bounded_transaction, create_simulation_engine, @@ -32,8 +78,12 @@ from policyengine_simulation_persistence.input_tables import input_metadata from policyengine_simulation_persistence.policy_reform import ( MissingPolicyParametersError, + OverlappingReformError, PolicyReform, + UnknownDynamicError, UnknownPolicyError, + combine_reforms, + read_dynamic_reform, read_policy_reform, ) from policyengine_simulation_persistence.tables import ( @@ -55,37 +105,81 @@ __all__ = [ "SERVICE_TABLE_PRIVILEGES", "SERVICE_UPDATABLE_COLUMNS", + "CatalogDataset", + "CatalogRegion", + "EconomyCatalogRows", + "EconomyPopulation", + "EconomyRecordOutcomeUnknownError", + "EconomyRecordReferenceError", + "EconomyReportInputs", + "EconomyReportState", + "EconomyRequest", + "EconomyRunStore", + "EconomySimulationInputs", + "EconomySimulationState", "HouseholdCalculationInputs", "HouseholdJobRecord", "HouseholdJobReferenceError", "HouseholdJobStatus", "HouseholdJobStore", "MissingPolicyParametersError", + "NewEconomyReport", + "NewEconomySimulation", "NewHouseholdJob", + "OverlappingReformError", "PolicyReform", + "ReportRunState", "ReportRunStatus", "ReportRunTrigger", + "ReportSide", + "ReportSideError", + "ReportSimulation", + "ReusableRun", + "RunReuse", + "RunVersions", + "SideRefusal", + "SimulationChoice", + "SimulationInputs", + "SimulationRunState", "SimulationStatus", "SimulationType", "StoredHousehold", "UnknownCatalogVersionError", + "UnknownDatasetError", + "UnknownDynamicError", "UnknownHouseholdError", "UnknownPolicyError", + "UnknownRegionError", "bounded_transaction", "budget_window_runs", + "combine_reforms", "create_simulation_engine", "fail_unstarted_household_job", + "fail_unstarted_report_run", + "fail_unstarted_simulation_run", + "find_reusable_run", "household_jobs", "input_metadata", + "insert_economy_report", + "insert_economy_simulation", "insert_household_job", "missing_table_privileges", "normalized_database_url", "read_catalog_version_id", + "read_dataset", + "read_dataset_by_id", + "read_dynamic_reform", + "read_economy_report", + "read_economy_simulation", "read_household_job", "read_missing_table_privileges", + "read_model_id", "read_policy_reform", + "read_region", "read_stored_household", "record_household_call", + "record_report_run_call", + "record_simulation_run_call", "referenced_metadata", "report_runs", "reports", diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.py new file mode 100644 index 000000000..1cde162ad --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_reads.py @@ -0,0 +1,181 @@ +"""Read a report or a simulation as it stands, with its newest run. + +The newest run is the one created last, with ties broken by the larger ID, +whatever its status and versions, as the Public API reads it. The run's +results are the Public API's to read: these reads give only its state. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from sqlalchemy import select + +from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + SimulationStatus, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from sqlalchemy import Connection, RowMapping + + +@dataclass(frozen=True) +class SimulationRunState: + """A simulation run's state: its status, versions, error and times.""" + + id: UUID + status: SimulationStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + + +@dataclass(frozen=True) +class ReportRunState: + """A report run's state, and the simulation runs it compares.""" + + id: UUID + status: ReportRunStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + + +@dataclass(frozen=True) +class EconomyReportState: + """A report, its two simulations, and its newest run (None with no run).""" + + id: UUID + country: str + baseline_simulation_id: UUID | None + reform_simulation_id: UUID | None + run: ReportRunState | None + + +@dataclass(frozen=True) +class EconomySimulationState: + """A simulation and its newest run (None with no run).""" + + id: UUID + run: SimulationRunState | None + + +def read_economy_report( + connection: Connection, report_id: UUID +) -> EconomyReportState | None: + """Return the report with its newest run, or None when there is no such report.""" + report = connection.execute( + select( + reports.c.id, + reports.c.country, + reports.c.baseline_simulation_id, + reports.c.reform_simulation_id, + ).where(reports.c.id == report_id) + ).one_or_none() + if report is None: + return None + run = ( + connection.execute( + select( + report_runs.c.id, + report_runs.c.status, + report_runs.c.country_package_version, + report_runs.c.policyengine_version, + report_runs.c.error_message, + report_runs.c.created_at, + report_runs.c.started_at, + report_runs.c.completed_at, + report_runs.c.baseline_simulation_run_id, + report_runs.c.reform_simulation_run_id, + ) + .where(report_runs.c.report_id == report_id) + .order_by(report_runs.c.created_at.desc(), report_runs.c.id.desc()) + .limit(1) + ) + .mappings() + .one_or_none() + ) + return EconomyReportState( + id=report.id, + country=report.country, + baseline_simulation_id=report.baseline_simulation_id, + reform_simulation_id=report.reform_simulation_id, + run=None if run is None else _report_run_state(run), + ) + + +def read_economy_simulation( + connection: Connection, simulation_id: UUID +) -> EconomySimulationState | None: + """Return the simulation with its newest run, or None when there is none.""" + found = connection.execute( + select(simulations.c.id).where(simulations.c.id == simulation_id) + ).scalar_one_or_none() + if found is None: + return None + run = ( + connection.execute( + select( + simulation_runs.c.id, + simulation_runs.c.status, + simulation_runs.c.country_package_version, + simulation_runs.c.policyengine_version, + simulation_runs.c.error_message, + simulation_runs.c.created_at, + simulation_runs.c.started_at, + simulation_runs.c.completed_at, + ) + .where(simulation_runs.c.simulation_id == simulation_id) + .order_by(simulation_runs.c.created_at.desc(), simulation_runs.c.id.desc()) + .limit(1) + ) + .mappings() + .one_or_none() + ) + return EconomySimulationState( + id=simulation_id, + run=None + if run is None + else SimulationRunState( + id=run["id"], + status=SimulationStatus(run["status"]), + country_package_version=run["country_package_version"], + policyengine_version=run["policyengine_version"], + error_message=run["error_message"], + created_at=run["created_at"], + started_at=run["started_at"], + completed_at=run["completed_at"], + ), + ) + + +def _report_run_state(run: RowMapping) -> ReportRunState: + return ReportRunState( + id=run["id"], + status=ReportRunStatus(run["status"]), + country_package_version=run["country_package_version"], + policyengine_version=run["policyengine_version"], + error_message=run["error_message"], + created_at=run["created_at"], + started_at=run["started_at"], + completed_at=run["completed_at"], + baseline_simulation_run_id=run["baseline_simulation_run_id"], + reform_simulation_run_id=run["reform_simulation_run_id"], + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py new file mode 100644 index 000000000..dd90ddabd --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_records.py @@ -0,0 +1,603 @@ +"""The simulation service's reads and writes for economy simulations and reports. + +An economy report compares a baseline and a reform simulation of one +region's data for one year. The service reads the catalog rows the request +names from the Public API's tables, then creates in one transaction: + +- for each side, a new ``simulations`` row and its first ``simulation_runs`` + row, ``pending`` with no Modal call ID; or, when a run of the same + simulation already succeeded and may be reused (``find_reusable_run``), + nothing: the report names that run and its simulation instead; +- the ``reports`` row, naming both simulations and the reform's policy, with + the report's options in ``inputs``; +- its first ``report_runs`` row, ``pending`` with no Modal call ID and + trigger ``initial``, naming both simulation runs. + +A standalone economy simulation is one ``simulations`` row and its first +run, created the same way. + +The service then starts the Modal call that calculates the run it created, +and records the call's ID on that run, or fails it at once when Modal +refuses the call. Both updates apply only while the run is still +``pending`` with no call ID: the coordinator or worker may already have +claimed it, and the stuck-run sweep fails a pending run no call was +recorded for. Failing a report run also fails the simulation runs created +for it, under the same condition. The coordinator and workers write the +rest: ``running``, the results, and ``succeeded`` or ``failed``. + +Every function takes a connection, so the caller chooses the transactions. +Each update sets only the columns the service's role may update, with +``updated_at``, whose automatic update the Public API's models make on their +own side only. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final + +from sqlalchemy import ( + cast, + func, + insert, + literal, + or_, + select, + update, +) +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.exc import IntegrityError + +from policyengine_simulation_persistence.input_tables import ( + MODEL_NAME_PREFIX, + datasets, + regions, + tax_benefit_models, +) +from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + ReportRunTrigger, + SimulationStatus, + SimulationType, + report_runs, + reports, + simulation_runs, + simulations, +) + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + from uuid import UUID + + from sqlalchemy import ColumnElement, Connection, FromClause, Table + +# PostgreSQL's SQLSTATE for a row naming one that does not exist. +FOREIGN_KEY_VIOLATION: Final = "23503" +# A report's type, as the Public API reads it. +ECONOMY_REPORT_TYPE: Final = "economy" +# The start of the error message of a report run whose coordinator refused a +# reused simulation run's output (``reused_simulation_: ``). +REFUSED_REUSE_PREFIX: Final = "reused_simulation_" +# Both run tables' statuses use these values. +_PENDING: Final = SimulationStatus.PENDING.value +_FAILED: Final = SimulationStatus.FAILED.value +# Each foreign key of the tables the service writes, by name: its table and +# the first column it constrains. +_KEY_COLUMNS: Final = { + str(constraint.name): (table.name, constraint.column_keys[0]) + for table in (simulations, simulation_runs, reports, report_runs) + for constraint in table.foreign_key_constraints +} + + +@dataclass(frozen=True) +class CatalogRegion: + """A region of a catalog version, and the dataset it uses by default.""" + + id: UUID + code: str + default_dataset_id: UUID + + +@dataclass(frozen=True) +class CatalogDataset: + """A dataset of a catalog version.""" + + id: UUID + name: str + + +@dataclass(frozen=True) +class RunVersions: + """The country package and .py versions a run calculates on.""" + + country_package_version: str + policyengine_version: str + + +@dataclass(frozen=True) +class EconomyPopulation: + """What every simulation of one request calculates on. + + The catalog version, the region and dataset of that version, and the + year. + """ + + tax_benefit_model_version_id: UUID + region_id: UUID + dataset_id: UUID + year: int + + +@dataclass(frozen=True) +class SimulationChoice: + """A simulation's stored policy and dynamic; neither means current law.""" + + policy_id: UUID | None + dynamic_id: UUID | None + + +@dataclass(frozen=True) +class ReusableRun: + """A succeeded simulation run a new report may name, and its simulation.""" + + simulation_id: UUID + simulation_run_id: UUID + + +@dataclass(frozen=True) +class ReportSimulation: + """The simulation and run one side of a report names. + + With ``reused``, both exist already and the run has succeeded, so nothing + is created for the side. Otherwise both are created, the run pending. + """ + + simulation_id: UUID + simulation_run_id: UUID + choice: SimulationChoice + reused: bool + + +@dataclass(frozen=True) +class NewEconomyReport: + """An economy report to create, with its first run. + + ``inputs`` holds the report's options as the Public API reads them. + """ + + id: UUID + run_id: UUID + country: str + label: str + description: str | None + tax_benefit_model_id: UUID + population: EconomyPopulation + versions: RunVersions + inputs: dict[str, Any] + baseline: ReportSimulation + reform: ReportSimulation + + @property + def new_simulation_run_ids(self) -> tuple[UUID, ...]: + """Return the IDs of the simulation runs created for the report.""" + return tuple( + side.simulation_run_id + for side in (self.baseline, self.reform) + if not side.reused + ) + + +@dataclass(frozen=True) +class NewEconomySimulation: + """A standalone economy simulation to create, with its first run.""" + + id: UUID + run_id: UUID + choice: SimulationChoice + population: EconomyPopulation + versions: RunVersions + + +class EconomyRecordReferenceError(LookupError): + """The database refused a row naming one that does not exist. + + ``table_name`` and ``column_name`` name the column whose foreign key + refused it; both are None when the database named a key the service's + tables don't have. + """ + + def __init__(self, table_name: str | None, column_name: str | None) -> None: + """Keep the table and column whose key refused the row.""" + self.table_name = table_name + self.column_name = column_name + named = ( + "a row" if table_name is None else f"the {table_name} row's {column_name}" + ) + super().__init__(f"{named} names no stored row") + + +def read_model_id(connection: Connection, *, country: str) -> UUID | None: + """Return the id of the country's tax-benefit model, or None.""" + return connection.execute( + select(tax_benefit_models.c.id).where( + tax_benefit_models.c.name == f"{MODEL_NAME_PREFIX}{country}" + ) + ).scalar_one_or_none() + + +def read_region( + connection: Connection, *, catalog_version_id: UUID, region_code: str +) -> CatalogRegion | None: + """Return the catalog version's region with this code, or None.""" + row = connection.execute( + select(regions.c.id, regions.c.code, regions.c.default_dataset_id).where( + regions.c.tax_benefit_model_version_id == catalog_version_id, + regions.c.code == region_code, + ) + ).one_or_none() + if row is None: + return None + return CatalogRegion( + id=row.id, code=row.code, default_dataset_id=row.default_dataset_id + ) + + +def read_dataset( + connection: Connection, *, catalog_version_id: UUID, name: str +) -> CatalogDataset | None: + """Return the catalog version's dataset with this name, or None.""" + return _dataset(connection, catalog_version_id, datasets.c.name == name) + + +def read_dataset_by_id( + connection: Connection, *, catalog_version_id: UUID, dataset_id: UUID +) -> CatalogDataset | None: + """Return the catalog version's dataset with this id, or None.""" + return _dataset(connection, catalog_version_id, datasets.c.id == dataset_id) + + +def _dataset( + connection: Connection, catalog_version_id: UUID, condition: ColumnElement[bool] +) -> CatalogDataset | None: + row = connection.execute( + select(datasets.c.id, datasets.c.name).where( + datasets.c.tax_benefit_model_version_id == catalog_version_id, condition + ) + ).one_or_none() + return None if row is None else CatalogDataset(id=row.id, name=row.name) + + +def find_reusable_run( # noqa: PLR0913 # each keyword is one part of the match + connection: Connection, + *, + choice: SimulationChoice, + population: EconomyPopulation, + versions: RunVersions, + inputs: dict[str, Any], + output_uri_prefix: str, + excluded_run_id: UUID | None = None, +) -> ReusableRun | None: + """Return the succeeded run a new report's simulation may name, or None. + + The run is the newest, by completion, of a run: + + - of an economy simulation with the same policy, dynamic, catalog + version, dataset, region and year; + - that succeeded on the same country package and .py versions, for that + year, dataset and region; + - whose output file is under ``output_uri_prefix``, the output bucket, + which never deletes a file; + - that a succeeded report run used, in a report with the same ``inputs`` + (options), so its file holds the outputs a report with those options + planned; + - that no report run has refused to reuse (an error message starting + ``reused_simulation_``), so a submission after a refused reuse + calculates the simulation again; and + - that is not ``excluded_run_id``, the other side's run. + + The coordinator may still refuse it, when the new report plans outputs + that report did not, as for a reform that turns on labour supply + responses: the refusal fails that report run, and the run is not offered + again. + """ + used_by = report_runs.alias("used_by_report_run") + refused_by = report_runs.alias("refused_by_report_run") + used_with_same_options = ( + select(used_by.c.id) + .join(reports, reports.c.id == used_by.c.report_id) + .where( + used_by.c.status == ReportRunStatus.SUCCEEDED, + _names_run(used_by), + cast(reports.c.inputs, JSONB) == literal(inputs, JSONB), + ) + .exists() + ) + refused = ( + select(refused_by.c.id) + .where( + refused_by.c.status == ReportRunStatus.FAILED, + _names_run(refused_by), + refused_by.c.error_message.startswith( + REFUSED_REUSE_PREFIX, autoescape=True + ), + ) + .exists() + ) + conditions = [ + simulations.c.simulation_type == SimulationType.ECONOMY, + simulations.c.policy_id.is_not_distinct_from(choice.policy_id), + simulations.c.dynamic_id.is_not_distinct_from(choice.dynamic_id), + simulations.c.tax_benefit_model_version_id + == population.tax_benefit_model_version_id, + simulations.c.dataset_id == population.dataset_id, + simulations.c.region_id == population.region_id, + simulations.c.year == population.year, + simulations.c.filter_field.is_(None), + simulation_runs.c.status == SimulationStatus.SUCCEEDED, + simulation_runs.c.year == population.year, + simulation_runs.c.country_package_version == versions.country_package_version, + simulation_runs.c.policyengine_version == versions.policyengine_version, + simulation_runs.c.dataset_id == population.dataset_id, + simulation_runs.c.region_id == population.region_id, + simulation_runs.c.output_uri.startswith(output_uri_prefix, autoescape=True), + simulation_runs.c.output_sha256.is_not(None), + used_with_same_options, + ~refused, + ] + if excluded_run_id is not None: + conditions.append(simulation_runs.c.id != excluded_run_id) + row = connection.execute( + select(simulation_runs.c.id, simulation_runs.c.simulation_id) + .join(simulations, simulations.c.id == simulation_runs.c.simulation_id) + .where(*conditions) + .order_by(simulation_runs.c.completed_at.desc(), simulation_runs.c.id.desc()) + .limit(1) + ).one_or_none() + if row is None: + return None + return ReusableRun(simulation_id=row.simulation_id, simulation_run_id=row.id) + + +def _names_run(report_run: FromClause) -> ColumnElement[bool]: + """Return whether ``report_run`` names the outer query's simulation run.""" + return or_( + report_run.c.baseline_simulation_run_id == simulation_runs.c.id, + report_run.c.reform_simulation_run_id == simulation_runs.c.id, + ) + + +def insert_economy_report(connection: Connection, report: NewEconomyReport) -> None: + """Create the report, its first run and each side's new simulation and run. + + Raise EconomyRecordReferenceError when the database refuses a row it + names. The refusal ends the caller's transaction, or its savepoint. + """ + with _references_checked(): + for side in (report.baseline, report.reform): + if not side.reused: + _insert_simulation( + connection, + simulation_id=side.simulation_id, + run_id=side.simulation_run_id, + choice=side.choice, + population=report.population, + versions=report.versions, + ) + population = report.population + connection.execute( + insert(reports).values( + id=report.id, + label=report.label, + description=report.description, + country=report.country, + type=ECONOMY_REPORT_TYPE, + user_id=None, + tax_benefit_model_id=report.tax_benefit_model_id, + policy_id=report.reform.choice.policy_id, + baseline_simulation_id=report.baseline.simulation_id, + reform_simulation_id=report.reform.simulation_id, + household_id=None, + dataset_id=population.dataset_id, + region_id=population.region_id, + year=population.year, + budget_window_size=None, + inputs=report.inputs, + ) + ) + connection.execute( + insert(report_runs).values( + id=report.run_id, + report_id=report.id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_run_id=report.reform.simulation_run_id, + budget_window_run_id=None, + year=None, + country_package_version=report.versions.country_package_version, + policyengine_version=report.versions.policyengine_version, + status=ReportRunStatus.PENDING, + trigger=ReportRunTrigger.INITIAL, + modal_call_id=None, + ) + ) + + +def insert_economy_simulation( + connection: Connection, simulation: NewEconomySimulation +) -> None: + """Create the standalone simulation and its first run, pending. + + Raise EconomyRecordReferenceError when the database refuses a row it + names. The refusal ends the caller's transaction, or its savepoint. + """ + with _references_checked(): + _insert_simulation( + connection, + simulation_id=simulation.id, + run_id=simulation.run_id, + choice=simulation.choice, + population=simulation.population, + versions=simulation.versions, + ) + + +def _insert_simulation( # noqa: PLR0913 # the simulation's and its run's parts + connection: Connection, + *, + simulation_id: UUID, + run_id: UUID, + choice: SimulationChoice, + population: EconomyPopulation, + versions: RunVersions, +) -> None: + """Create an economy simulation and its first run, pending with no call.""" + connection.execute( + insert(simulations).values( + id=simulation_id, + simulation_type=SimulationType.ECONOMY, + dataset_id=population.dataset_id, + household_id=None, + policy_id=choice.policy_id, + dynamic_id=choice.dynamic_id, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + region_id=population.region_id, + year=population.year, + ) + ) + connection.execute( + insert(simulation_runs).values( + id=run_id, + simulation_id=simulation_id, + year=population.year, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + country_package_version=versions.country_package_version, + policyengine_version=versions.policyengine_version, + dataset_id=population.dataset_id, + region_id=population.region_id, + status=SimulationStatus.PENDING, + modal_call_id=None, + ) + ) + + +@contextmanager +def _references_checked() -> Iterator[None]: + """Turn a foreign-key refusal into EconomyRecordReferenceError. + + Any other integrity error propagates unchanged. + """ + try: + yield + except IntegrityError as error: + driver_error = error.orig + if getattr(driver_error, "sqlstate", None) != FOREIGN_KEY_VIOLATION: + raise + constraint_name = getattr( + getattr(driver_error, "diag", None), "constraint_name", None + ) + table_name, column_name = _KEY_COLUMNS.get(str(constraint_name), (None, None)) + raise EconomyRecordReferenceError(table_name, column_name) from None + + +def _pending_without_call( + table: Table, run_id: UUID +) -> tuple[ColumnElement[bool], ...]: + return ( + table.c.id == run_id, + table.c.status == _PENDING, + table.c.modal_call_id.is_(None), + ) + + +def record_report_run_call( + connection: Connection, *, report_run_id: UUID, modal_call_id: str +) -> ReportRunStatus | None: + """Record the report run's Modal call ID while it is pending with none. + + Return the run's status afterwards, or None when there is no such run. A + run the coordinator has claimed already keeps its claim. + """ + status = _record_call(connection, report_runs, report_run_id, modal_call_id) + return None if status is None else ReportRunStatus(status) + + +def record_simulation_run_call( + connection: Connection, *, simulation_run_id: UUID, modal_call_id: str +) -> SimulationStatus | None: + """Record the simulation run's Modal call ID while it is pending with none. + + Return the run's status afterwards, or None when there is no such run. A + run the worker has claimed already keeps its claim. + """ + status = _record_call(connection, simulation_runs, simulation_run_id, modal_call_id) + return None if status is None else SimulationStatus(status) + + +def _record_call( + connection: Connection, table: Table, run_id: UUID, modal_call_id: str +) -> str | None: + recorded = connection.execute( + update(table) + .where(*_pending_without_call(table, run_id)) + .values(modal_call_id=modal_call_id, updated_at=func.now()) + .returning(table.c.status) + ).scalar_one_or_none() + if recorded is not None: + return str(recorded) + status = connection.execute( + select(table.c.status).where(table.c.id == run_id) + ).scalar_one_or_none() + return None if status is None else str(status) + + +def fail_unstarted_report_run( + connection: Connection, + *, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, +) -> bool: + """Fail the report run at once while it is pending with no Modal call ID. + + When it was failed, also fail each of ``simulation_run_ids``, the + simulation runs created for it, that is still pending with no call ID. + Return whether the report run was failed: one a call has been recorded + for, or one the coordinator has claimed, is left as it is, and so are + its simulation runs. + """ + if not _fail_unstarted(connection, report_runs, report_run_id, error_message): + return False + for simulation_run_id in simulation_run_ids: + _fail_unstarted(connection, simulation_runs, simulation_run_id, error_message) + return True + + +def fail_unstarted_simulation_run( + connection: Connection, *, simulation_run_id: UUID, error_message: str +) -> bool: + """Fail the simulation run at once while it is pending with no Modal call ID. + + Return whether it was failed: a run a call has been recorded for, or one + a worker has claimed, is left as it is. + """ + return _fail_unstarted( + connection, simulation_runs, simulation_run_id, error_message + ) + + +def _fail_unstarted( + connection: Connection, table: Table, run_id: UUID, error_message: str +) -> bool: + failed = connection.execute( + update(table) + .where(*_pending_without_call(table, run_id)) + .values( + status=_FAILED, + error_message=error_message, + completed_at=func.now(), + updated_at=func.now(), + ) + .returning(table.c.id) + ).scalar_one_or_none() + return failed is not None diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py new file mode 100644 index 000000000..6e66b2a50 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/economy_run_store.py @@ -0,0 +1,467 @@ +"""The economy store: each read and write in a transaction of its own. + +The Simulation Entrypoint creates an economy report, or a standalone economy +simulation, in three steps, each its own transaction, with the Modal call +started between them: it reads the request's inputs, creates the rows, then +records the call or fails the run it started. ``EconomyRunStore`` runs the +reads and writes of ``economy_run_records`` and ``policy_reform`` that way. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, get_args + +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_persistence.economy_run_reads import ( + read_economy_report, + read_economy_simulation, +) +from policyengine_simulation_persistence.economy_run_records import ( + CatalogDataset, + CatalogRegion, + EconomyPopulation, + ReusableRun, + RunVersions, + SimulationChoice, + fail_unstarted_report_run, + fail_unstarted_simulation_run, + find_reusable_run, + insert_economy_report, + insert_economy_simulation, + read_dataset, + read_dataset_by_id, + read_model_id, + read_region, + record_report_run_call, + record_simulation_run_call, +) +from policyengine_simulation_persistence.engine import bounded_transaction +from policyengine_simulation_persistence.household_job_records import ( + read_catalog_version_id, +) +from policyengine_simulation_persistence.household_job_store import ( + UnknownCatalogVersionError, +) +from policyengine_simulation_persistence.policy_reform import ( + MissingPolicyParametersError, + OverlappingReformError, + UnknownDynamicError, + UnknownPolicyError, + combine_reforms, + read_dynamic_reform, + read_policy_reform, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + from contextlib import AbstractContextManager + from uuid import UUID + + from sqlalchemy import Connection, Engine + + from policyengine_simulation_persistence.economy_run_reads import ( + EconomyReportState, + EconomySimulationState, + ) + from policyengine_simulation_persistence.economy_run_records import ( + NewEconomyReport, + NewEconomySimulation, + ) + from policyengine_simulation_persistence.policy_reform import PolicyReform + from policyengine_simulation_persistence.tables import ( + ReportRunStatus, + SimulationStatus, + ) + +type TransactionFactory = Callable[[], AbstractContextManager[Connection]] +# Which of a report's simulations a value belongs to. +type ReportSide = Literal["baseline", "reform"] +# The errors a simulation's own policy or dynamic can raise. +type SideRefusal = ( + UnknownPolicyError + | UnknownDynamicError + | MissingPolicyParametersError + | OverlappingReformError +) +# Caught by the type itself, so the two never name different errors. +_SIDE_ERRORS: tuple[type[SideRefusal], ...] = get_args(SideRefusal.__value__) + + +class UnknownRegionError(LookupError): + """The catalog version has no region with the requested code.""" + + +class UnknownDatasetError(LookupError): + """The catalog version has no dataset with the requested name.""" + + +class ReportSideError(Exception): + """A report side's policy or dynamic is refused. + + ``side`` names the side, and ``error`` is the policy reader's refusal: + UnknownPolicyError, UnknownDynamicError, MissingPolicyParametersError or + OverlappingReformError. + """ + + def __init__(self, side: ReportSide, error: SideRefusal) -> None: + """Keep the side and the refusal.""" + super().__init__(f"the report's {side}: {error}") + self.side: ReportSide = side + self.error = error + + +class EconomyRecordOutcomeUnknownError(RuntimeError): + """A create's statements ran, but its commit failed: the rows may exist. + + PostgreSQL may have applied the commit before the connection failed, so + the caller must not answer that nothing was created. + """ + + +@dataclass(frozen=True) +class EconomyRequest: + """What an economy simulation or report request names, apart from policies. + + ``policyengine_version`` is the .py version the calculation runs on, + which chooses the country's catalog version. ``dataset_name`` None means + the region's default dataset. + """ + + country: str + policyengine_version: str + region_code: str + dataset_name: str | None + year: int + + +@dataclass(frozen=True) +class RunReuse: + """How a report looks for simulation runs to reuse (``find_reusable_run``). + + ``versions`` are those the new runs would calculate on, ``inputs`` the + report's options, and ``output_uri_prefix`` the output bucket's + ``gs:///``. + """ + + versions: RunVersions + inputs: dict[str, Any] + output_uri_prefix: str + + +@dataclass(frozen=True) +class EconomyCatalogRows: + """The catalog rows a request is calculated on.""" + + catalog_version_id: UUID + model_id: UUID + region: CatalogRegion + dataset: CatalogDataset + + def population(self, year: int) -> EconomyPopulation: + """Return what every simulation of the request calculates on.""" + return EconomyPopulation( + tax_benefit_model_version_id=self.catalog_version_id, + region_id=self.region.id, + dataset_id=self.dataset.id, + year=year, + ) + + +@dataclass(frozen=True) +class SimulationInputs: + """One simulation's policy and dynamic, and the values they set together. + + ``reusable_run`` names a succeeded run the simulation may use instead of + a new one, when the request looked for one and one exists. + """ + + choice: SimulationChoice + reform: PolicyReform + reusable_run: ReusableRun | None = None + + +@dataclass(frozen=True) +class EconomySimulationInputs: + """What a standalone economy simulation is read from.""" + + catalog: EconomyCatalogRows + simulation: SimulationInputs + + +@dataclass(frozen=True) +class EconomyReportInputs: + """What an economy report is read from: the catalog rows and both sides.""" + + catalog: EconomyCatalogRows + baseline: SimulationInputs + reform: SimulationInputs + + +class EconomyRunStore: + """Read and write economy simulations and reports, each in its own transaction.""" + + def __init__(self, transaction: TransactionFactory) -> None: + """Open each transaction with ``transaction``, which commits on success.""" + self._transaction = transaction + + @classmethod + def from_engine( + cls, engine: Engine, *, statement_timeout_milliseconds: int + ) -> EconomyRunStore: + """Use ``engine``, bounding each statement to the given time.""" + + def transaction() -> AbstractContextManager[Connection]: + return bounded_transaction( + engine, statement_timeout_milliseconds=statement_timeout_milliseconds + ) + + return cls(transaction) + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + """Read a standalone simulation's catalog rows and values together. + + Raise UnknownCatalogVersionError, UnknownRegionError or + UnknownDatasetError, or the policy reader's UnknownPolicyError, + UnknownDynamicError or MissingPolicyParametersError, or + OverlappingReformError. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + simulation = _simulation_inputs( + connection, request.country, catalog.catalog_version_id, choice + ) + return EconomySimulationInputs(catalog=catalog, simulation=simulation) + + def read_report_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + """Read a report's catalog rows, both sides' values and reusable runs. + + With ``reuse``, each side gets the run ``find_reusable_run`` finds, + and the reform never gets the baseline's. Raise + UnknownCatalogVersionError, UnknownRegionError or UnknownDatasetError, + or ReportSideError naming the side whose policy or dynamic is refused. + """ + with self._transaction() as connection: + catalog = _catalog_rows(connection, request) + population = catalog.population(request.year) + sides = [] + excluded_run_id = None + sides_requested: tuple[tuple[ReportSide, SimulationChoice], ...] = ( + ("baseline", baseline), + ("reform", reform), + ) + for side, choice in sides_requested: + try: + inputs = _simulation_inputs( + connection, request.country, catalog.catalog_version_id, choice + ) + except _SIDE_ERRORS as error: + raise ReportSideError(side, error) from error + if reuse is not None: + reusable_run = find_reusable_run( + connection, + choice=choice, + population=population, + versions=reuse.versions, + inputs=reuse.inputs, + output_uri_prefix=reuse.output_uri_prefix, + excluded_run_id=excluded_run_id, + ) + if reusable_run is not None: + excluded_run_id = reusable_run.simulation_run_id + inputs = SimulationInputs( + choice=inputs.choice, + reform=inputs.reform, + reusable_run=reusable_run, + ) + sides.append(inputs) + return EconomyReportInputs(catalog=catalog, baseline=sides[0], reform=sides[1]) + + def create_report(self, report: NewEconomyReport) -> None: + """Create the report, its new simulations and its runs. + + Raise ``insert_economy_report``'s errors, and + EconomyRecordOutcomeUnknownError when the commit fails after the + rows were inserted. + """ + self._create(lambda connection: insert_economy_report(connection, report)) + + def create_simulation(self, simulation: NewEconomySimulation) -> None: + """Create the simulation and its run; see ``insert_economy_simulation``. + + Raise EconomyRecordOutcomeUnknownError as ``create_report`` does. + """ + self._create( + lambda connection: insert_economy_simulation(connection, simulation) + ) + + def _create(self, insert: Callable[[Connection], None]) -> None: + """Run ``insert`` in a transaction, telling a failed commit apart. + + An error from ``insert`` leaves nothing written and propagates. An + error once it has run is the commit's, after which the rows may + exist: EconomyRecordOutcomeUnknownError. + """ + inserted = False + try: + with self._transaction() as connection: + insert(connection) + inserted = True + except SQLAlchemyError as error: + if not inserted: + raise + msg = "the create's commit failed after its rows were inserted" + raise EconomyRecordOutcomeUnknownError(msg) from error + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + """Record the report run's call ID while it is pending with none.""" + with self._transaction() as connection: + return record_report_run_call( + connection, report_run_id=report_run_id, modal_call_id=modal_call_id + ) + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + """Fail the report run, and its new simulation runs, while it is unstarted.""" + with self._transaction() as connection: + return fail_unstarted_report_run( + connection, + report_run_id=report_run_id, + simulation_run_ids=simulation_run_ids, + error_message=error_message, + ) + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + """Record the simulation run's call ID while it is pending with none.""" + with self._transaction() as connection: + return record_simulation_run_call( + connection, + simulation_run_id=simulation_run_id, + modal_call_id=modal_call_id, + ) + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + """Fail the simulation run while it is pending with no call ID.""" + with self._transaction() as connection: + return fail_unstarted_simulation_run( + connection, + simulation_run_id=simulation_run_id, + error_message=error_message, + ) + + def read_report(self, report_id: UUID) -> EconomyReportState | None: + """Return the report with its newest run, or None when there is none.""" + with self._transaction() as connection: + return read_economy_report(connection, report_id) + + def read_simulation(self, simulation_id: UUID) -> EconomySimulationState | None: + """Return the simulation with its newest run, or None when there is none.""" + with self._transaction() as connection: + return read_economy_simulation(connection, simulation_id) + + +def _catalog_rows( + connection: Connection, request: EconomyRequest +) -> EconomyCatalogRows: + """Read the catalog version, model, region and dataset a request names. + + Raise UnknownCatalogVersionError, UnknownRegionError or + UnknownDatasetError. + """ + country = request.country + catalog_version_id = read_catalog_version_id( + connection, + country=country, + policyengine_version=request.policyengine_version, + ) + model_id = read_model_id(connection, country=country) + if catalog_version_id is None or model_id is None: + msg = f"{country} has no catalog version for .py {request.policyengine_version}" + raise UnknownCatalogVersionError(msg) + region = read_region( + connection, + catalog_version_id=catalog_version_id, + region_code=request.region_code, + ) + if region is None: + msg = ( + f"{country}'s catalog version for .py {request.policyengine_version} " + f"has no region {request.region_code}" + ) + raise UnknownRegionError(msg) + dataset = ( + read_dataset_by_id( + connection, + catalog_version_id=catalog_version_id, + dataset_id=region.default_dataset_id, + ) + if request.dataset_name is None + else read_dataset( + connection, + catalog_version_id=catalog_version_id, + name=request.dataset_name, + ) + ) + if dataset is None: + msg = ( + f"{country}'s catalog version for .py {request.policyengine_version} " + f"has no dataset {request.dataset_name}" + ) + raise UnknownDatasetError(msg) + return EconomyCatalogRows( + catalog_version_id=catalog_version_id, + model_id=model_id, + region=region, + dataset=dataset, + ) + + +def _simulation_inputs( + connection: Connection, + country: str, + catalog_version_id: UUID, + choice: SimulationChoice, +) -> SimulationInputs: + """Read one simulation's policy and dynamic values, combined.""" + policy_reform = ( + {} + if choice.policy_id is None + else read_policy_reform( + connection, + country=country, + policy_id=choice.policy_id, + catalog_version_id=catalog_version_id, + ) + ) + dynamic_reform = ( + {} + if choice.dynamic_id is None + else read_dynamic_reform( + connection, + dynamic_id=choice.dynamic_id, + catalog_version_id=catalog_version_id, + ) + ) + return SimulationInputs( + choice=choice, reform=combine_reforms(policy_reform, dynamic_reform) + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py index 1bde68b74..ec9797fad 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/input_tables.py @@ -1,14 +1,14 @@ """Read-only models of the Public API's tables a calculation is read from. -The Public API writes households, policies and their parameter values, and -its model catalog; the simulation service only reads them, with ``SELECT``. -Each table here names only the columns the service reads, so these are not -models of the whole tables: the Public API's other columns, defaults, -constraints and indexes are its own. The tables live in ``input_metadata``, -apart from the models of the tables the service writes (``service_metadata``) -and from the tables those name by ``id`` alone (``referenced_metadata``). -``tests/test_public_api_schema.py`` compares each column with the Public -API's table. +The Public API writes households, policies, dynamics and their parameter +values, and its model catalog; the simulation service only reads them, with +``SELECT``. Each table here names only the columns the service reads, so these +are not models of the whole tables: the Public API's other columns, +defaults, constraints and indexes are its own. The tables live in +``input_metadata``, apart from the models of the tables the service writes +(``service_metadata``) and from the tables those name by ``id`` alone +(``referenced_metadata``). ``tests/test_public_api_schema.py`` compares each +column with the Public API's table. """ from __future__ import annotations @@ -42,6 +42,26 @@ Column("version", String(128), nullable=False), ) +# A region of one catalog version, unique by code within it, and the dataset +# it uses unless a calculation names another. +regions = Table( + "regions", + input_metadata, + Column("id", Uuid, nullable=False), + Column("code", String(255), nullable=False), + Column("tax_benefit_model_version_id", Uuid, nullable=False), + Column("default_dataset_id", Uuid, nullable=False), +) + +# A dataset of one catalog version, unique by name within it. +datasets = Table( + "datasets", + input_metadata, + Column("id", Uuid, nullable=False), + Column("name", String(255), nullable=False), + Column("tax_benefit_model_version_id", Uuid, nullable=False), +) + # ``household_data`` lists the people and each group entity as records, and # may hold a saved SPM selection; ``default_year`` is the year it describes. households = Table( @@ -70,7 +90,8 @@ ) # A parameter's value from ``start_date`` until ``end_date``, or with no end -# date indefinitely; a policy's own values name it in ``policy_id``. +# date indefinitely; a policy's own values name it in ``policy_id``, and a +# dynamic's in ``dynamic_id``. parameter_values = Table( "parameter_values", input_metadata, @@ -79,4 +100,13 @@ Column("start_date", DateTime(timezone=True), nullable=False), Column("end_date", DateTime(timezone=True), nullable=True), Column("policy_id", Uuid, nullable=True), + Column("dynamic_id", Uuid, nullable=True), +) + +# Named parameter values a simulation applies beside its policy, such as +# behavioural responses. A dynamic belongs to no country. +dynamics = Table( + "dynamics", + input_metadata, + Column("id", Uuid, nullable=False), ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py index c87ebc985..e1261160b 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/policy_reform.py @@ -1,11 +1,13 @@ -"""Read a stored policy as the reform values a calculation applies. +"""Read a stored policy, or dynamic, as the reform values a calculation applies. The Public API stores a policy's parameter values against the parameters of the catalog version it was saved on, each from a start date until an end date, or with no end date indefinitely. A calculation on any catalog version applies them by parameter name, as ``{parameter name: {period: value}}``, the shape of the Public API's policy JSON that the household functions' contract -and the economy requests take. +and the economy requests take. A dynamic's values are stored and read the same +way. A simulation applies its policy's and its dynamic's values together, and +is refused when both set the same parameter (``combine_reforms``). Each period is a range, ``YYYY-MM-DD.YYYY-MM-DD``, with the dates read in UTC. A value with no end date follows .py's own rule @@ -29,11 +31,12 @@ from __future__ import annotations from datetime import UTC, date, timedelta -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from sqlalchemy import select from policyengine_simulation_persistence.input_tables import ( + dynamics, parameter_values, parameters, policies, @@ -45,31 +48,58 @@ from datetime import datetime from uuid import UUID - from sqlalchemy import Connection + from sqlalchemy import ColumnElement, Connection # The last day a value without an end date is applied. OPEN_ENDED_STOP: Final = date(LATEST_YEAR, 12, 31) # A policy's reform values: each parameter's values by period. type PolicyReform = dict[str, dict[str, object]] +# What stores the values a calculation applies. +type ReformSource = Literal["policy", "dynamic"] class UnknownPolicyError(LookupError): """The Public API has no such policy in the country.""" +class UnknownDynamicError(LookupError): + """The Public API has no such dynamic.""" + + class MissingPolicyParametersError(ValueError): - """The policy sets parameters the calculation's catalog version lacks.""" + """The policy, or dynamic, sets parameters the calculation's catalog version lacks. - def __init__(self, parameter_names: Sequence[str]) -> None: + ``source`` says which of the two sets them. + """ + + def __init__( + self, parameter_names: Sequence[str], *, source: ReformSource = "policy" + ) -> None: """Name every parameter the catalog version lacks, in name order.""" self.parameter_names = tuple(parameter_names) + self.source: ReformSource = source super().__init__( - "the policy sets parameters the catalog version does not have: " + f"the {source} sets parameters the catalog version does not have: " + ", ".join(self.parameter_names) ) +class OverlappingReformError(ValueError): + """A simulation's policy and dynamic both set the same parameters. + + Neither is applied over the other, so the simulation cannot be + calculated. + """ + + def __init__(self, parameter_names: Sequence[str]) -> None: + """Name every parameter both set, in name order.""" + self.parameter_names = tuple(parameter_names) + super().__init__( + "the policy and the dynamic both set: " + ", ".join(self.parameter_names) + ) + + def _period( start_date: datetime, end_date: datetime | None, @@ -118,6 +148,61 @@ def read_policy_reform( if found is None: msg = f"{country} has no policy {policy_id}" raise UnknownPolicyError(msg) + return _read_reform( + connection, + owner=parameter_values.c.policy_id == policy_id, + catalog_version_id=catalog_version_id, + source="policy", + ) + + +def read_dynamic_reform( + connection: Connection, *, dynamic_id: UUID, catalog_version_id: UUID +) -> PolicyReform: + """Return the values the dynamic sets, by parameter name and period. + + The periods follow the same rule as a policy's. Raise + UnknownDynamicError when there is no such dynamic, and + MissingPolicyParametersError, naming the dynamic as its source, when + ``catalog_version_id`` lacks a parameter it sets. + """ + found = connection.execute( + select(dynamics.c.id).where(dynamics.c.id == dynamic_id) + ).scalar_one_or_none() + if found is None: + msg = f"there is no dynamic {dynamic_id}" + raise UnknownDynamicError(msg) + return _read_reform( + connection, + owner=parameter_values.c.dynamic_id == dynamic_id, + catalog_version_id=catalog_version_id, + source="dynamic", + ) + + +def combine_reforms(policy: PolicyReform, dynamic: PolicyReform) -> PolicyReform: + """Return a simulation's policy values and its dynamic's, as one reform. + + Raise OverlappingReformError when both set a parameter, for any period. + """ + overlapping = sorted(set(policy) & set(dynamic)) + if overlapping: + raise OverlappingReformError(overlapping) + return {**policy, **dynamic} + + +def _read_reform( + connection: Connection, + *, + owner: ColumnElement[bool], + catalog_version_id: UUID, + source: ReformSource, +) -> PolicyReform: + """Return the values ``owner`` selects, by parameter name and .py's periods. + + Raise MissingPolicyParametersError, naming ``source``, when + ``catalog_version_id`` lacks a parameter they set. + """ values = connection.execute( select( parameters.c.name, @@ -126,7 +211,7 @@ def read_policy_reform( parameter_values.c.end_date, ) .join(parameters, parameters.c.id == parameter_values.c.parameter_id) - .where(parameter_values.c.policy_id == policy_id) + .where(owner) .order_by(parameters.c.name, parameter_values.c.start_date) ).all() names = sorted({row.name for row in values}) @@ -140,7 +225,7 @@ def read_policy_reform( ) missing = [name for name in names if name not in available] if missing: - raise MissingPolicyParametersError(missing) + raise MissingPolicyParametersError(missing, source=source) reform: PolicyReform = {} # The rows are ordered by parameter name, then start date, the order # .py takes them in; each value's periods are .py's. diff --git a/libs/policyengine-simulation-persistence/tests/economy_rows.py b/libs/policyengine-simulation-persistence/tests/economy_rows.py new file mode 100644 index 000000000..4cb6a9014 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/economy_rows.py @@ -0,0 +1,237 @@ +"""Rows for the economy simulation and report tests. + +The catalog rows a request names are the Public API's, inserted as the +superuser. A run that already succeeded, and the report run that used it, are +written as the superuser too, through the service tables' models: the +service's role can neither mark a run succeeded nor record its output. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from policyengine_simulation_persistence import ( + EconomyPopulation, + ReportRunStatus, + ReportRunTrigger, + RunVersions, + SimulationChoice, + SimulationStatus, + SimulationType, + report_runs, + reports, + simulation_runs, + simulations, +) +from public_api_rows import ( + CatalogVersion, + insert_catalog_version, + insert_dataset, + insert_model, + insert_region, +) +from sqlalchemy import insert + +if TYPE_CHECKING: + from sqlalchemy import Connection + +POLICYENGINE_VERSION = "6.17.1" +COUNTRY_PACKAGE_VERSION = "1.764.6" +VERSIONS = RunVersions( + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version=POLICYENGINE_VERSION, +) +OUTPUT_BUCKET_PREFIX = "gs://policyengine-v2-simulation-outputs-staging/" +CURRENT_LAW = SimulationChoice(policy_id=None, dynamic_id=None) +OPTIONS: dict[str, Any] = {"include_cliffs": False} +COMPLETED_AT = datetime(2026, 11, 1, tzinfo=UTC) + + +@dataclass(frozen=True) +class EconomyCatalog: + """A country's catalog version, with one dataset and one region using it.""" + + catalog: CatalogVersion + dataset_id: UUID + region_id: UUID + + def population(self, *, year: int = 2026) -> EconomyPopulation: + """Return the population a request for the region and year uses.""" + return EconomyPopulation( + tax_benefit_model_version_id=self.catalog.version_id, + region_id=self.region_id, + dataset_id=self.dataset_id, + year=year, + ) + + +def insert_economy_catalog( + connection: Connection, + *, + country: str = "us", + parameter_names: tuple[str, ...] = (), +) -> EconomyCatalog: + """Insert the country's model, a catalog version, a dataset and a region.""" + catalog = insert_catalog_version( + connection, + country=country, + model_id=insert_model(connection, country), + policyengine_version=POLICYENGINE_VERSION, + parameter_names=parameter_names, + ) + dataset_id = insert_dataset(connection, catalog, f"populace_{country}_2024") + region_id = insert_region( + connection, catalog, country, default_dataset_id=dataset_id + ) + return EconomyCatalog(catalog=catalog, dataset_id=dataset_id, region_id=region_id) + + +@dataclass(frozen=True) +class StoredRun: + """A simulation and one of its runs, as stored.""" + + simulation_id: UUID + simulation_run_id: UUID + + +@dataclass(frozen=True) +class RunShape: + """What a stored economy simulation and its run hold. + + A succeeded run records an output file in the output bucket unless + ``output_uri`` names another. + """ + + choice: SimulationChoice = CURRENT_LAW + status: SimulationStatus = SimulationStatus.SUCCEEDED + versions: RunVersions = VERSIONS + output_uri: str | None = None + completed_at: datetime = COMPLETED_AT + + +def insert_simulation_run( + connection: Connection, + population: EconomyPopulation, + shape: RunShape | None = None, +) -> StoredRun: + """Insert an economy simulation and one run of it, as the superuser.""" + shape = shape or RunShape() + simulation_id = uuid4() + run_id = uuid4() + finished = shape.status in {SimulationStatus.SUCCEEDED, SimulationStatus.FAILED} + succeeded = shape.status is SimulationStatus.SUCCEEDED + connection.execute( + insert(simulations).values( + id=simulation_id, + simulation_type=SimulationType.ECONOMY, + dataset_id=population.dataset_id, + policy_id=shape.choice.policy_id, + dynamic_id=shape.choice.dynamic_id, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + region_id=population.region_id, + year=population.year, + ) + ) + connection.execute( + insert(simulation_runs).values( + id=run_id, + simulation_id=simulation_id, + year=population.year, + tax_benefit_model_version_id=population.tax_benefit_model_version_id, + country_package_version=shape.versions.country_package_version, + policyengine_version=shape.versions.policyengine_version, + dataset_id=population.dataset_id, + region_id=population.region_id, + status=shape.status, + completed_at=shape.completed_at if finished else None, + output_uri=( + (shape.output_uri or f"{OUTPUT_BUCKET_PREFIX}outputs/{run_id}.parquet") + if succeeded + else None + ), + output_sha256="d" * 64 if succeeded else None, + ) + ) + return StoredRun(simulation_id=simulation_id, simulation_run_id=run_id) + + +@dataclass(frozen=True) +class ReportRunOutcome: + """How a stored report run ended, and the options its report has.""" + + status: ReportRunStatus = ReportRunStatus.SUCCEEDED + inputs: dict[str, Any] | None = None + error_message: str | None = None + + +def insert_report_run_naming( + connection: Connection, + economy: EconomyCatalog, + compared: tuple[StoredRun, StoredRun], + outcome: ReportRunOutcome | None = None, +) -> UUID: + """Insert a report comparing the two runs, and a run of it, as the superuser. + + ``compared`` holds the baseline's run, then the reform's. Return the + report run's ID. + """ + outcome = outcome or ReportRunOutcome() + baseline, reform = compared + report_id = uuid4() + run_id = uuid4() + connection.execute( + insert(reports).values( + id=report_id, + label="An earlier report", + country=economy.catalog.country, + type="economy", + tax_benefit_model_id=economy.catalog.model_id, + baseline_simulation_id=baseline.simulation_id, + reform_simulation_id=reform.simulation_id, + dataset_id=economy.dataset_id, + region_id=economy.region_id, + year=2026, + inputs=OPTIONS if outcome.inputs is None else outcome.inputs, + ) + ) + finished = outcome.status in {ReportRunStatus.SUCCEEDED, ReportRunStatus.FAILED} + connection.execute( + insert(report_runs).values( + id=run_id, + report_id=report_id, + baseline_simulation_run_id=baseline.simulation_run_id, + reform_simulation_run_id=reform.simulation_run_id, + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version=POLICYENGINE_VERSION, + status=outcome.status, + trigger=ReportRunTrigger.INITIAL, + completed_at=COMPLETED_AT if finished else None, + error_message=outcome.error_message, + ) + ) + return run_id + + +def insert_reusable_run( + connection: Connection, + economy: EconomyCatalog, + shape: RunShape | None = None, + *, + inputs: dict[str, Any] | None = None, +) -> StoredRun: + """Insert a run shaped ``shape`` that a succeeded report run used, with ``inputs``. + + The report run's other simulation run failed, so it is never reusable. + """ + population = economy.population() + reused = insert_simulation_run(connection, population, shape) + other_side = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming( + connection, economy, (reused, other_side), ReportRunOutcome(inputs=inputs) + ) + return reused diff --git a/libs/policyengine-simulation-persistence/tests/public_api_rows.py b/libs/policyengine-simulation-persistence/tests/public_api_rows.py index 7bb74ec47..37bfaf22b 100644 --- a/libs/policyengine-simulation-persistence/tests/public_api_rows.py +++ b/libs/policyengine-simulation-persistence/tests/public_api_rows.py @@ -148,6 +148,91 @@ def insert_policy( return policy_id +def insert_dynamic( + connection: Connection, + catalog: CatalogVersion, + values: Iterable[StoredValue] = (), +) -> UUID: + """Insert a dynamic and the values it sets on ``catalog``'s parameters. + + A dynamic belongs to no country; each value names a parameter of that + catalog version by name. + """ + dynamic_id = uuid4() + connection.execute( + text("INSERT INTO public.dynamics (id, name) VALUES (:id, :name)"), + {"id": dynamic_id, "name": f"dynamic {dynamic_id}"}, + ) + parameter_ids = _parameter_ids(connection, catalog.version_id) + insert_value = text( + "INSERT INTO public.parameter_values " + "(id, parameter_id, value_json, start_date, end_date, dynamic_id) " + "VALUES (:id, :parameter_id, :value, :start_date, :end_date, :dynamic_id)" + ).bindparams(bindparam("value", type_=JSONB)) + for stored in values: + connection.execute( + insert_value, + { + "id": uuid4(), + "parameter_id": parameter_ids[stored.parameter_name], + "value": stored.value, + "start_date": stored.start_date, + "end_date": stored.end_date, + "dynamic_id": dynamic_id, + }, + ) + return dynamic_id + + +def insert_dataset( + connection: Connection, catalog: CatalogVersion, name: str, *, year: int = 2024 +) -> UUID: + """Insert a dataset of ``catalog`` with this name.""" + dataset_id = uuid4() + connection.execute( + text( + "INSERT INTO public.datasets (id, name, year, is_output_dataset, " + "tax_benefit_model_version_id) " + "VALUES (:id, :name, :year, false, :version_id)" + ), + { + "id": dataset_id, + "name": name, + "year": year, + "version_id": catalog.version_id, + }, + ) + return dataset_id + + +def insert_region( + connection: Connection, + catalog: CatalogVersion, + code: str, + *, + default_dataset_id: UUID, + region_type: str = "national", +) -> UUID: + """Insert a region of ``catalog`` with this code and default dataset.""" + region_id = uuid4() + connection.execute( + text( + "INSERT INTO public.regions (id, code, label, region_type, " + "requires_filter, default_dataset_id, tax_benefit_model_version_id) " + "VALUES (:id, :code, :code, CAST(:region_type AS v2_region_type), " + "false, :dataset_id, :version_id)" + ), + { + "id": region_id, + "code": code, + "region_type": region_type, + "dataset_id": default_dataset_id, + "version_id": catalog.version_id, + }, + ) + return region_id + + def insert_household( connection: Connection, *, diff --git a/libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py b/libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py new file mode 100644 index 000000000..0bce6c2ea --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_dynamic_reform.py @@ -0,0 +1,155 @@ +"""A stored dynamic's parameter values, and combining them with a policy's. + +The reads run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the simulation service's role, in a +transaction rolled back afterwards. Combining needs no database. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from uuid import uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from policyengine_simulation_persistence import ( + MissingPolicyParametersError, + OverlappingReformError, + UnknownDynamicError, + combine_reforms, + read_dynamic_reform, +) +from public_api_rows import ( + StoredValue, + insert_catalog_version, + insert_dynamic, + insert_model, +) + +START = datetime(2026, 1, 1, tzinfo=UTC) +END = datetime(2027, 12, 31, tzinfo=UTC) +ELASTICITY = "gov.simulation.labor_supply_responses.elasticities.income" +SUBSTITUTION = "gov.simulation.labor_supply_responses.elasticities.substitution.all" +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" + + +@pytest.mark.public_api_schema +def test_a_dynamic_is_read_as_its_values_by_parameter_name_and_period( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = insert_catalog_version( + connection, + country="us", + model_id=insert_model(connection, "us"), + policyengine_version="6.17.1", + parameter_names=[ELASTICITY, SUBSTITUTION], + ) + dynamic_id = insert_dynamic( + connection, + catalog, + [ + StoredValue(ELASTICITY, -0.05, START, END), + StoredValue(SUBSTITUTION, 0.25, START), + ], + ) + act_as_the_service(connection) + + reform = read_dynamic_reform( + connection, dynamic_id=dynamic_id, catalog_version_id=catalog.version_id + ) + + assert reform == { + ELASTICITY: {"2026-01-01.2027-12-31": -0.05}, + SUBSTITUTION: {"2026-01-01.2200-12-31": 0.25}, + } + + +@pytest.mark.public_api_schema +def test_an_unknown_dynamic_is_refused(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + catalog = insert_catalog_version( + connection, + country="us", + model_id=insert_model(connection, "us"), + policyengine_version="6.17.1", + ) + act_as_the_service(connection) + missing_dynamic_id = uuid4() + + with pytest.raises(UnknownDynamicError, match=str(missing_dynamic_id)): + read_dynamic_reform( + connection, + dynamic_id=missing_dynamic_id, + catalog_version_id=catalog.version_id, + ) + + +@pytest.mark.public_api_schema +def test_a_dynamic_changing_parameters_the_run_lacks_names_each_of_them( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + model_id = insert_model(connection, "us") + saved_on = insert_catalog_version( + connection, + country="us", + model_id=model_id, + policyengine_version="6.16.0", + parameter_names=[ELASTICITY, SUBSTITUTION], + ) + running = insert_catalog_version( + connection, + country="us", + model_id=model_id, + policyengine_version="6.17.1", + parameter_names=[ELASTICITY], + ) + dynamic_id = insert_dynamic( + connection, + saved_on, + [ + StoredValue(ELASTICITY, -0.05, START), + StoredValue(SUBSTITUTION, 0.25, START), + ], + ) + act_as_the_service(connection) + + with pytest.raises(MissingPolicyParametersError) as refusal: + read_dynamic_reform( + connection, + dynamic_id=dynamic_id, + catalog_version_id=running.version_id, + ) + + assert refusal.value.parameter_names == (SUBSTITUTION,) + assert refusal.value.source == "dynamic" + assert str(refusal.value).startswith("the dynamic sets parameters") + + +def test_a_policy_and_a_dynamic_combine_by_parameter() -> None: + policy = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} + dynamic = {ELASTICITY: {"2026-01-01.2200-12-31": -0.05}} + + assert combine_reforms(policy, dynamic) == { + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + ELASTICITY: {"2026-01-01.2200-12-31": -0.05}, + } + assert combine_reforms(policy, {}) == policy + assert combine_reforms({}, dynamic) == dynamic + + +def test_a_parameter_the_policy_and_the_dynamic_both_set_is_refused() -> None: + policy = { + ELASTICITY: {"2026-01-01.2200-12-31": -0.1}, + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + } + dynamic = { + ELASTICITY: {"2030-01-01.2200-12-31": -0.05}, + SUBSTITUTION: {"2026-01-01.2200-12-31": 0.25}, + } + + with pytest.raises(OverlappingReformError) as refusal: + combine_reforms(policy, dynamic) + + assert refusal.value.parameter_names == (ELASTICITY,) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py new file mode 100644 index 000000000..836b63c23 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_reads.py @@ -0,0 +1,198 @@ +"""Reading an economy report or simulation as it stands, with its newest run. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the simulation service's role, each +in a transaction rolled back afterwards. The rows are written as the +superuser, to give runs their own creation times. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + COUNTRY_PACKAGE_VERSION, + POLICYENGINE_VERSION, + EconomyCatalog, + ReportRunOutcome, + RunShape, + insert_economy_catalog, + insert_report_run_naming, + insert_simulation_run, +) +from policyengine_simulation_persistence import ( + ReportRunStatus, + SimulationStatus, + read_economy_report, + read_economy_simulation, + report_runs, + simulation_runs, +) +from sqlalchemy import update + +pytestmark = pytest.mark.public_api_schema + +EARLIER = datetime(2026, 11, 1, tzinfo=UTC) + + +def _report_with_runs(connection, economy: EconomyCatalog) -> tuple[UUID, list[UUID]]: + """Insert a report and three runs of it; return the report and its runs.""" + population = economy.population() + baseline = insert_simulation_run(connection, population) + reform = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + run_ids = [ + insert_report_run_naming( + connection, + economy, + (baseline, reform), + ReportRunOutcome(status=ReportRunStatus.FAILED, error_message="first"), + ) + ] + report_id = ( + connection.execute(report_runs.select().where(report_runs.c.id == run_ids[0])) + .one() + .report_id + ) + for _ in range(2): + run_id = uuid4() + connection.execute( + report_runs.insert().values( + id=run_id, + report_id=report_id, + baseline_simulation_run_id=baseline.simulation_run_id, + reform_simulation_run_id=reform.simulation_run_id, + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version=POLICYENGINE_VERSION, + status=ReportRunStatus.PENDING, + trigger="manual", + ) + ) + run_ids.append(run_id) + return report_id, run_ids + + +def test_a_report_is_read_with_its_newest_run(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + report_id, (first, second, third) = _report_with_runs(connection, economy) + # The first run is the oldest; the other two share a creation time, so + # the larger ID is the newer. + connection.execute( + update(report_runs) + .where(report_runs.c.id == first) + .values(created_at=EARLIER) + ) + shared = EARLIER + timedelta(days=1) + connection.execute( + update(report_runs) + .where(report_runs.c.id.in_([second, third])) + .values(created_at=shared) + ) + act_as_the_service(connection) + + report = read_economy_report(connection, report_id) + missing = read_economy_report(connection, uuid4()) + + assert missing is None + assert report is not None + assert (report.id, report.country) == (report_id, "us") + run = report.run + assert run is not None + assert run.id == max(second, third) + assert (run.status, run.error_message) == (ReportRunStatus.PENDING, None) + assert (run.country_package_version, run.policyengine_version) == ( + COUNTRY_PACKAGE_VERSION, + POLICYENGINE_VERSION, + ) + assert run.created_at == shared + assert (run.started_at, run.completed_at) == (None, None) + + +def test_a_failed_newest_run_is_read_with_its_error(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + population = economy.population() + baseline = insert_simulation_run(connection, population) + reform = insert_simulation_run(connection, population) + run_id = insert_report_run_naming( + connection, + economy, + (baseline, reform), + ReportRunOutcome( + status=ReportRunStatus.FAILED, + error_message="simulation_failed: the reform failed", + ), + ) + report_id = ( + connection.execute(report_runs.select().where(report_runs.c.id == run_id)) + .one() + .report_id + ) + act_as_the_service(connection) + + report = read_economy_report(connection, report_id) + + assert report is not None + assert (report.baseline_simulation_id, report.reform_simulation_id) == ( + baseline.simulation_id, + reform.simulation_id, + ) + assert report.run is not None + assert ( + report.run.baseline_simulation_run_id, + report.run.reform_simulation_run_id, + ) == ( + baseline.simulation_run_id, + reform.simulation_run_id, + ) + assert (report.run.status, report.run.error_message) == ( + ReportRunStatus.FAILED, + "simulation_failed: the reform failed", + ) + assert report.run.completed_at is not None + + +def test_a_simulation_is_read_with_its_newest_run(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + stored = insert_simulation_run(connection, economy.population()) + newer_run_id = uuid4() + connection.execute( + simulation_runs.insert().values( + id=newer_run_id, + simulation_id=stored.simulation_id, + year=2026, + tax_benefit_model_version_id=economy.catalog.version_id, + country_package_version=COUNTRY_PACKAGE_VERSION, + policyengine_version="6.18.0", + dataset_id=economy.dataset_id, + region_id=economy.region_id, + status=SimulationStatus.RUNNING, + started_at=EARLIER + timedelta(days=2), + ) + ) + connection.execute( + update(simulation_runs) + .where(simulation_runs.c.id == stored.simulation_run_id) + .values(created_at=EARLIER) + ) + act_as_the_service(connection) + + simulation = read_economy_simulation(connection, stored.simulation_id) + missing = read_economy_simulation(connection, uuid4()) + + assert missing is None + assert simulation is not None + assert simulation.id == stored.simulation_id + assert simulation.run is not None + assert (simulation.run.id, simulation.run.status) == ( + newer_run_id, + SimulationStatus.RUNNING, + ) + assert simulation.run.policyengine_version == "6.18.0" + assert simulation.run.started_at == EARLIER + timedelta(days=2) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_records.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_records.py new file mode 100644 index 000000000..1ddf50262 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_records.py @@ -0,0 +1,602 @@ +"""The simulation service's reads and writes for economy simulations and reports. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the simulation service's role, each +in a transaction rolled back afterwards. The Public API's rows a request names +are inserted first, as the superuser. +""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + VERSIONS, + EconomyCatalog, + insert_economy_catalog, + insert_reusable_run, +) +from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, + NewEconomyReport, + NewEconomySimulation, + ReportRunStatus, + ReportSimulation, + SimulationChoice, + SimulationStatus, + fail_unstarted_report_run, + fail_unstarted_simulation_run, + insert_economy_report, + insert_economy_simulation, + read_dataset, + read_dataset_by_id, + read_model_id, + read_region, + record_report_run_call, + record_simulation_run_call, + report_runs, + simulation_runs, +) +from public_api_rows import insert_dataset, insert_policy, insert_region +from sqlalchemy import select, text, update + +if TYPE_CHECKING: + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +REPORT_OPTIONS = {"include_cliffs": True, "spm": {"geography_kind": "national"}} + + +def _new_side(choice: SimulationChoice = CURRENT_LAW) -> ReportSimulation: + return ReportSimulation( + simulation_id=uuid4(), + simulation_run_id=uuid4(), + choice=choice, + reused=False, + ) + + +def _new_report( + economy: EconomyCatalog, + *, + baseline: ReportSimulation | None = None, + reform: ReportSimulation | None = None, +) -> NewEconomyReport: + return NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country=economy.catalog.country, + label="Child tax credit", + description="A larger credit", + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=REPORT_OPTIONS, + baseline=baseline or _new_side(), + reform=reform or _new_side(), + ) + + +def _row(connection: Connection, table_name: str, row_id: UUID) -> dict: + row = connection.execute( + text(f"SELECT * FROM public.{table_name} WHERE id = :id"), # noqa: S608 # fixed table names + {"id": row_id}, + ).mappings() + return dict(row.one()) + + +def _count(connection: Connection, table_name: str) -> int: + return connection.execute( + text(f"SELECT count(*) FROM public.{table_name}") # noqa: S608 # fixed table names + ).scalar_one() + + +def test_the_countrys_model_is_found_by_its_name(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + + assert read_model_id(connection, country="us") == economy.catalog.model_id + assert read_model_id(connection, country="uk") is None + + +def test_a_region_is_found_by_code_in_its_catalog_version_only( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + other = insert_economy_catalog(connection, country="uk") + state_id = insert_region( + connection, + economy.catalog, + "state/ca", + default_dataset_id=economy.dataset_id, + region_type="state", + ) + act_as_the_service(connection) + + region = read_region( + connection, + catalog_version_id=economy.catalog.version_id, + region_code="state/ca", + ) + assert ( + read_region( + connection, + catalog_version_id=other.catalog.version_id, + region_code="state/ca", + ) + is None + ) + + assert region is not None + assert (region.id, region.code, region.default_dataset_id) == ( + state_id, + "state/ca", + economy.dataset_id, + ) + + +def test_a_dataset_is_found_by_name_or_id_in_its_catalog_version_only( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + other = insert_economy_catalog(connection, country="uk") + local_id = insert_dataset(connection, economy.catalog, "acs_local_2024") + act_as_the_service(connection) + version_id = economy.catalog.version_id + + by_name = read_dataset( + connection, catalog_version_id=version_id, name="acs_local_2024" + ) + by_id = read_dataset_by_id( + connection, catalog_version_id=version_id, dataset_id=economy.dataset_id + ) + other_version = other.catalog.version_id + assert ( + read_dataset( + connection, catalog_version_id=other_version, name="acs_local_2024" + ) + is None + ) + assert ( + read_dataset_by_id( + connection, catalog_version_id=other_version, dataset_id=local_id + ) + is None + ) + + assert by_name is not None + assert (by_name.id, by_name.name) == (local_id, "acs_local_2024") + assert by_id is not None + assert (by_id.id, by_id.name) == (economy.dataset_id, "populace_us_2024") + + +def test_a_report_its_simulations_and_runs_are_created_pending( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + policy_id = insert_policy(connection, economy.catalog) + report = _new_report( + economy, reform=_new_side(SimulationChoice(policy_id, dynamic_id=None)) + ) + act_as_the_service(connection) + + insert_economy_report(connection, report) + stored_report = _row(connection, "reports", report.id) + stored_run = _row(connection, "report_runs", report.run_id) + baseline = _row(connection, "simulations", report.baseline.simulation_id) + reform = _row(connection, "simulations", report.reform.simulation_id) + baseline_run = _row( + connection, "simulation_runs", report.baseline.simulation_run_id + ) + reform_run = _row( + connection, "simulation_runs", report.reform.simulation_run_id + ) + + population = economy.population() + assert { + key: stored_report[key] + for key in ( + "label", + "description", + "country", + "type", + "user_id", + "tax_benefit_model_id", + "policy_id", + "baseline_simulation_id", + "reform_simulation_id", + "household_id", + "dataset_id", + "region_id", + "year", + "inputs", + "budget_window_size", + ) + } == { + "label": "Child tax credit", + "description": "A larger credit", + "country": "us", + "type": "economy", + "user_id": None, + "tax_benefit_model_id": economy.catalog.model_id, + "policy_id": policy_id, + "baseline_simulation_id": report.baseline.simulation_id, + "reform_simulation_id": report.reform.simulation_id, + "household_id": None, + "dataset_id": economy.dataset_id, + "region_id": economy.region_id, + "year": 2026, + "inputs": REPORT_OPTIONS, + "budget_window_size": None, + } + assert { + key: stored_run[key] + for key in ( + "report_id", + "baseline_simulation_run_id", + "reform_simulation_run_id", + "budget_window_run_id", + "year", + "country_package_version", + "policyengine_version", + "status", + "trigger", + "modal_call_id", + "started_at", + "completed_at", + ) + } == { + "report_id": report.id, + "baseline_simulation_run_id": report.baseline.simulation_run_id, + "reform_simulation_run_id": report.reform.simulation_run_id, + "budget_window_run_id": None, + "year": None, + "country_package_version": VERSIONS.country_package_version, + "policyengine_version": VERSIONS.policyengine_version, + "status": "pending", + "trigger": "initial", + "modal_call_id": None, + "started_at": None, + "completed_at": None, + } + for simulation, run, expected_policy in ( + (baseline, baseline_run, None), + (reform, reform_run, policy_id), + ): + assert ( + simulation["simulation_type"], + simulation["policy_id"], + simulation["dynamic_id"], + simulation["household_id"], + simulation["dataset_id"], + simulation["region_id"], + simulation["tax_benefit_model_version_id"], + simulation["year"], + simulation["filter_field"], + ) == ( + "economy", + expected_policy, + None, + None, + population.dataset_id, + population.region_id, + population.tax_benefit_model_version_id, + 2026, + None, + ) + assert ( + run["simulation_id"], + run["year"], + run["tax_benefit_model_version_id"], + run["country_package_version"], + run["policyengine_version"], + run["dataset_id"], + run["region_id"], + run["status"], + run["modal_call_id"], + run["output_uri"], + ) == ( + simulation["id"], + 2026, + population.tax_benefit_model_version_id, + VERSIONS.country_package_version, + VERSIONS.policyengine_version, + population.dataset_id, + population.region_id, + "pending", + None, + None, + ) + + +def test_a_side_reusing_a_succeeded_run_creates_nothing_for_that_side( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused = insert_reusable_run(connection, economy) + report = _new_report( + economy, + baseline=ReportSimulation( + simulation_id=reused.simulation_id, + simulation_run_id=reused.simulation_run_id, + choice=CURRENT_LAW, + reused=True, + ), + ) + simulations_before = _count(connection, "simulations") + runs_before = _count(connection, "simulation_runs") + act_as_the_service(connection) + + insert_economy_report(connection, report) + stored_report = _row(connection, "reports", report.id) + stored_run = _row(connection, "report_runs", report.run_id) + simulations_after = _count(connection, "simulations") + runs_after = _count(connection, "simulation_runs") + + assert stored_report["baseline_simulation_id"] == reused.simulation_id + assert stored_run["baseline_simulation_run_id"] == reused.simulation_run_id + assert stored_run["reform_simulation_run_id"] == report.reform.simulation_run_id + assert (simulations_after, runs_after) == ( + simulations_before + 1, + runs_before + 1, + ) + + +@pytest.mark.parametrize( + ("missing_column", "table_name"), + [("region_id", "simulations"), ("dataset_id", "simulations")], +) +def test_a_report_naming_a_missing_row_names_the_column_that_refused_it( + public_api_database_url, missing_column: str, table_name: str +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + population = replace(report.population, **{missing_column: uuid4()}) + act_as_the_service(connection) + + with pytest.raises(EconomyRecordReferenceError) as refusal: + insert_economy_report(connection, replace(report, population=population)) + + assert (refusal.value.table_name, refusal.value.column_name) == ( + table_name, + missing_column, + ) + + +def test_a_report_naming_a_missing_dynamic_names_its_column( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + report = _new_report( + economy, reform=_new_side(SimulationChoice(None, dynamic_id=uuid4())) + ) + act_as_the_service(connection) + + with pytest.raises(EconomyRecordReferenceError) as refusal: + insert_economy_report(connection, report) + + assert (refusal.value.table_name, refusal.value.column_name) == ( + "simulations", + "dynamic_id", + ) + + +def test_a_simulation_and_its_run_are_created_pending(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + policy_id = insert_policy(connection, economy.catalog) + simulation = NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=SimulationChoice(policy_id, dynamic_id=None), + population=economy.population(year=2027), + versions=VERSIONS, + ) + act_as_the_service(connection) + + insert_economy_simulation(connection, simulation) + stored = _row(connection, "simulations", simulation.id) + run = _row(connection, "simulation_runs", simulation.run_id) + + assert ( + stored["simulation_type"], + stored["policy_id"], + stored["dynamic_id"], + stored["dataset_id"], + stored["region_id"], + stored["year"], + ) == ("economy", policy_id, None, economy.dataset_id, economy.region_id, 2027) + assert ( + run["simulation_id"], + run["year"], + run["status"], + run["modal_call_id"], + run["policyengine_version"], + ) == (simulation.id, 2027, "pending", None, VERSIONS.policyengine_version) + + +def _created_report(connection: Connection) -> NewEconomyReport: + economy = insert_economy_catalog(connection) + report = _new_report(economy) + act_as_the_service(connection) + insert_economy_report(connection, report) + return report + + +def test_a_report_runs_call_is_recorded_once_while_it_is_pending_without_one( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + report = _created_report(connection) + + first = record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-first" + ) + second = record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-second" + ) + unknown = record_report_run_call( + connection, report_run_id=uuid4(), modal_call_id="fc-other" + ) + call_id = _row(connection, "report_runs", report.run_id)["modal_call_id"] + + assert (first, second, unknown) == ( + ReportRunStatus.PENDING, + ReportRunStatus.PENDING, + None, + ) + assert call_id == "fc-first" + + +def test_a_claimed_report_run_keeps_its_claim(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + report = _created_report(connection) + connection.execute(text("RESET ROLE")) + connection.execute( + update(report_runs) + .where(report_runs.c.id == report.run_id) + .values(status=ReportRunStatus.RUNNING, modal_call_id="fc-coordinator") + ) + act_as_the_service(connection) + + status = record_report_run_call( + connection, report_run_id=report.run_id, modal_call_id="fc-late" + ) + failed = fail_unstarted_report_run( + connection, + report_run_id=report.run_id, + simulation_run_ids=(report.baseline.simulation_run_id,), + error_message="modal_call_refused: Modal refused the call.", + ) + stored = _row(connection, "report_runs", report.run_id) + baseline = _row( + connection, "simulation_runs", report.baseline.simulation_run_id + ) + + assert status is ReportRunStatus.RUNNING + assert failed is False + assert (stored["status"], stored["modal_call_id"]) == ("running", "fc-coordinator") + assert baseline["status"] == "pending" + + +def test_an_unstarted_report_run_and_its_new_simulation_runs_fail_at_once( + public_api_database_url, +) -> None: + message = "modal_call_refused: Modal refused the coordinator's call." + with rolled_back_transaction(public_api_database_url) as connection: + report = _created_report(connection) + + failed = fail_unstarted_report_run( + connection, + report_run_id=report.run_id, + simulation_run_ids=( + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ), + error_message=message, + ) + again = fail_unstarted_report_run( + connection, + report_run_id=report.run_id, + simulation_run_ids=(), + error_message="a second failure", + ) + stored = _row(connection, "report_runs", report.run_id) + stored_simulation_runs = connection.execute( + select( + simulation_runs.c.status, + simulation_runs.c.error_message, + simulation_runs.c.completed_at, + ).where( + simulation_runs.c.id.in_( + [ + report.baseline.simulation_run_id, + report.reform.simulation_run_id, + ] + ) + ) + ).all() + + assert (failed, again) == (True, False) + assert (stored["status"], stored["error_message"]) == ("failed", message) + assert stored["completed_at"] is not None + assert [ + (row.status, row.error_message, row.completed_at is not None) + for row in stored_simulation_runs + ] == [("failed", message, True)] * 2 + + +def test_a_simulation_runs_call_is_recorded_or_the_run_failed_while_unstarted( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + recorded = NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + ) + refused = NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + ) + act_as_the_service(connection) + insert_economy_simulation(connection, recorded) + insert_economy_simulation(connection, refused) + + status = record_simulation_run_call( + connection, simulation_run_id=recorded.run_id, modal_call_id="fc-worker" + ) + not_failed = fail_unstarted_simulation_run( + connection, + simulation_run_id=recorded.run_id, + error_message="modal_call_refused: refused", + ) + failed = fail_unstarted_simulation_run( + connection, + simulation_run_id=refused.run_id, + error_message="modal_call_refused: refused", + ) + not_recorded = record_simulation_run_call( + connection, simulation_run_id=refused.run_id, modal_call_id="fc-late" + ) + unknown = record_simulation_run_call( + connection, simulation_run_id=uuid4(), modal_call_id="fc-other" + ) + recorded_run = _row(connection, "simulation_runs", recorded.run_id) + refused_run = _row(connection, "simulation_runs", refused.run_id) + + assert (status, not_failed, failed, not_recorded, unknown) == ( + SimulationStatus.PENDING, + False, + True, + SimulationStatus.FAILED, + None, + ) + assert (recorded_run["status"], recorded_run["modal_call_id"]) == ( + "pending", + "fc-worker", + ) + assert ( + refused_run["status"], + refused_run["modal_call_id"], + refused_run["error_message"], + ) == ("failed", None, "modal_call_refused: refused") diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py new file mode 100644 index 000000000..ee750e10d --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_reuse.py @@ -0,0 +1,307 @@ +"""Which succeeded simulation run a new report reuses, if any. + +A report reuses the newest run of a simulation with the same policy, dynamic, +catalog version, dataset, region and year, that succeeded on the same country +package and .py versions, whose output file is in the output bucket, and that +a succeeded report run with the same options used. A run any report run +refused to reuse is never offered again. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``), as the simulation service's role, each +in a transaction rolled back afterwards. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from datetime import timedelta +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + COMPLETED_AT, + CURRENT_LAW, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + VERSIONS, + EconomyCatalog, + ReportRunOutcome, + RunShape, + StoredRun, + insert_economy_catalog, + insert_report_run_naming, + insert_reusable_run, + insert_simulation_run, +) +from policyengine_simulation_persistence import ( + ReportRunStatus, + ReusableRun, + RunVersions, + SimulationChoice, + SimulationStatus, + find_reusable_run, +) +from public_api_rows import insert_dataset, insert_dynamic, insert_policy, insert_region + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + + +@dataclass(frozen=True) +class Request: + """What a new report's simulation asks for, as the reuse search sees it.""" + + choice: SimulationChoice = CURRENT_LAW + year: int = 2026 + versions: RunVersions = VERSIONS + inputs: dict[str, Any] = field(default_factory=lambda: dict(OPTIONS)) + excluded_run_id: UUID | None = None + + +def _find( + connection: Connection, economy: EconomyCatalog, request: Request | None = None +) -> ReusableRun | None: + request = request or Request() + return find_reusable_run( + connection, + choice=request.choice, + population=economy.population(year=request.year), + versions=request.versions, + inputs=request.inputs, + output_uri_prefix=OUTPUT_BUCKET_PREFIX, + excluded_run_id=request.excluded_run_id, + ) + + +def _reusable(run: StoredRun) -> ReusableRun: + return ReusableRun( + simulation_id=run.simulation_id, simulation_run_id=run.simulation_run_id + ) + + +def test_a_succeeded_run_a_report_used_with_the_same_options_is_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + policy_id = insert_policy(connection, economy.catalog) + dynamic_id = insert_dynamic(connection, economy.catalog) + choice = SimulationChoice(policy_id=policy_id, dynamic_id=dynamic_id) + reused = insert_reusable_run(connection, economy, RunShape(choice=choice)) + act_as_the_service(connection) + + found = _find(connection, economy, Request(choice=choice)) + + assert found == _reusable(reused) + + +def _another_policy(connection: Connection, economy: EconomyCatalog) -> dict: + return { + "choice": SimulationChoice(insert_policy(connection, economy.catalog), None) + } + + +def _another_dynamic(connection: Connection, economy: EconomyCatalog) -> dict: + return { + "choice": SimulationChoice(None, insert_dynamic(connection, economy.catalog)) + } + + +# Each makes the request differ from the stored run in one way. +DIFFERENCES: dict[str, Callable[[Connection, EconomyCatalog], dict[str, Any]]] = { + "another policy": _another_policy, + "another dynamic": _another_dynamic, + "another year": lambda _connection, _economy: {"year": 2027}, + "another .py version": lambda _connection, _economy: { + "versions": replace(VERSIONS, policyengine_version="6.18.0") + }, + "another country package version": lambda _connection, _economy: { + "versions": replace(VERSIONS, country_package_version="1.765.0") + }, + "other options": lambda _connection, _economy: {"inputs": {"include_cliffs": True}}, +} + + +@pytest.mark.parametrize("difference", sorted(DIFFERENCES)) +def test_a_run_differing_from_the_request_is_not_reused( + public_api_database_url, difference: str +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run(connection, economy) + request = DIFFERENCES[difference](connection, economy) + act_as_the_service(connection) + + assert _find(connection, economy, Request(**request)) is None + + +def test_a_run_of_another_region_or_dataset_is_not_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run(connection, economy) + other_dataset_id = insert_dataset(connection, economy.catalog, "acs_2024") + other_region_id = insert_region( + connection, + economy.catalog, + "state/ca", + default_dataset_id=economy.dataset_id, + region_type="state", + ) + act_as_the_service(connection) + + for other in ( + replace(economy, dataset_id=other_dataset_id), + replace(economy, region_id=other_region_id), + ): + assert _find(connection, other) is None + + +@pytest.mark.parametrize( + "status", + [SimulationStatus.PENDING, SimulationStatus.RUNNING, SimulationStatus.FAILED], +) +def test_an_unfinished_or_failed_run_is_not_reused( + public_api_database_url, status: SimulationStatus +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run(connection, economy, RunShape(status=status)) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +def test_a_run_whose_file_is_outside_the_output_bucket_is_not_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + insert_reusable_run( + connection, + economy, + RunShape( + output_uri="gs://policyengine-stage12-staging-artifacts/runs/output.parquet" + ), + ) + # A bucket whose name starts with the output bucket's is another bucket. + insert_reusable_run( + connection, + economy, + RunShape( + output_uri=OUTPUT_BUCKET_PREFIX.rstrip("/") + "-old/output.parquet" + ), + ) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +@pytest.mark.parametrize( + "report_run_status", + [ReportRunStatus.PENDING, ReportRunStatus.RUNNING, ReportRunStatus.FAILED], +) +def test_a_run_no_succeeded_report_run_used_is_not_reused( + public_api_database_url, report_run_status: ReportRunStatus +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + population = economy.population() + stored = insert_simulation_run(connection, population) + other_side = insert_simulation_run( + connection, population, RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming( + connection, + economy, + (stored, other_side), + ReportRunOutcome(status=report_run_status), + ) + # A run no report run names at all, such as a simulation's own. + insert_simulation_run(connection, population) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +@pytest.mark.parametrize("refused_as", ["baseline", "reform"]) +def test_a_run_a_report_run_refused_to_reuse_is_not_reused_again( + public_api_database_url, refused_as: str +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused = insert_reusable_run(connection, economy) + other_side = insert_simulation_run( + connection, economy.population(), RunShape(status=SimulationStatus.FAILED) + ) + baseline, reform = ( + (reused, other_side) if refused_as == "baseline" else (other_side, reused) + ) + insert_report_run_naming( + connection, + economy, + (baseline, reform), + ReportRunOutcome( + status=ReportRunStatus.FAILED, + error_message=( + "reused_simulation_output_incompatible: a planned column is missing" + ), + ), + ) + act_as_the_service(connection) + + assert _find(connection, economy) is None + + +def test_a_run_whose_report_run_failed_for_another_reason_is_still_reused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + reused = insert_reusable_run(connection, economy) + other_side = insert_simulation_run( + connection, economy.population(), RunShape(status=SimulationStatus.FAILED) + ) + insert_report_run_naming( + connection, + economy, + (reused, other_side), + ReportRunOutcome( + status=ReportRunStatus.FAILED, + error_message=( + "simulation_failed: the reform's worker ran out of memory" + ), + ), + ) + act_as_the_service(connection) + + assert _find(connection, economy) == _reusable(reused) + + +def test_the_newest_matching_run_is_reused_unless_excluded( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + older = insert_reusable_run(connection, economy) + newer = insert_reusable_run( + connection, economy, RunShape(completed_at=COMPLETED_AT + timedelta(days=1)) + ) + act_as_the_service(connection) + + newest = _find(connection, economy) + without_newest = _find( + connection, economy, Request(excluded_run_id=newer.simulation_run_id) + ) + without_any = _find(connection, economy, Request(excluded_run_id=uuid4())) + + assert newest == _reusable(newer) + assert without_newest == _reusable(older) + assert without_any == _reusable(newer) diff --git a/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py new file mode 100644 index 000000000..7ba3e04c8 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_economy_run_store.py @@ -0,0 +1,548 @@ +"""The economy store: each read and write in a transaction of its own. + +These run against a database the Public API's own v2 migration chain built +(``PUBLIC_API_SCHEMA_DATABASE_URL``). The store opens each transaction +through a function it is given; here each is a savepoint of one connection, +acting as the simulation service's role, which the test rolls back. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import replace +from datetime import UTC, datetime +from typing import TYPE_CHECKING, get_args +from uuid import uuid4 + +import pytest +from conftest import act_as_the_service, rolled_back_transaction +from economy_rows import ( + CURRENT_LAW, + OPTIONS, + OUTPUT_BUCKET_PREFIX, + POLICYENGINE_VERSION, + VERSIONS, + insert_economy_catalog, + insert_reusable_run, +) +from policyengine_simulation_persistence import ( + CatalogDataset, + CatalogRegion, + EconomyCatalogRows, + EconomyPopulation, + EconomyRecordOutcomeUnknownError, + EconomyRequest, + EconomyRunStore, + MissingPolicyParametersError, + NewEconomyReport, + NewEconomySimulation, + OverlappingReformError, + ReportRunStatus, + ReportSideError, + ReportSimulation, + ReusableRun, + RunReuse, + SideRefusal, + SimulationChoice, + SimulationStatus, + UnknownCatalogVersionError, + UnknownDatasetError, + UnknownDynamicError, + UnknownPolicyError, + UnknownRegionError, + economy_run_store, +) +from public_api_rows import StoredValue, insert_dataset, insert_dynamic, insert_policy +from sqlalchemy import text +from sqlalchemy.exc import OperationalError + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlalchemy import Connection + +pytestmark = pytest.mark.public_api_schema + +START = datetime(2026, 1, 1, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +ELASTICITY = "gov.simulation.labor_supply_responses.elasticities.income" +REQUEST = EconomyRequest( + country="us", + policyengine_version=POLICYENGINE_VERSION, + region_code="us", + dataset_name=None, + year=2026, +) +REUSE = RunReuse( + versions=VERSIONS, inputs=OPTIONS, output_uri_prefix=OUTPUT_BUCKET_PREFIX +) + + +def _store_on(connection: Connection) -> EconomyRunStore: + """Return a store whose transactions are savepoints of ``connection``.""" + + @contextmanager + def savepoint() -> Iterator[Connection]: + with connection.begin_nested(): + yield connection + + return EconomyRunStore(savepoint) + + +def test_a_reports_inputs_are_read_together(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog( + connection, parameter_names=(CTC_AMOUNT, ELASTICITY) + ) + policy_id = insert_policy( + connection, economy.catalog, [StoredValue(CTC_AMOUNT, 3_000, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + reused = insert_reusable_run(connection, economy) + act_as_the_service(connection) + store = _store_on(connection) + reform_choice = SimulationChoice(policy_id=policy_id, dynamic_id=dynamic_id) + + inputs = store.read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=reform_choice, reuse=REUSE + ) + without_reuse = store.read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=reform_choice, reuse=None + ) + + catalog = inputs.catalog + assert ( + catalog.catalog_version_id, + catalog.model_id, + catalog.region.id, + catalog.dataset.id, + catalog.dataset.name, + ) == ( + economy.catalog.version_id, + economy.catalog.model_id, + economy.region_id, + economy.dataset_id, + "populace_us_2024", + ) + assert catalog.population(2026) == economy.population() + assert (inputs.baseline.choice, inputs.baseline.reform) == (CURRENT_LAW, {}) + assert inputs.reform.choice == reform_choice + assert inputs.reform.reform == { + CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}, + ELASTICITY: {"2026-01-01.2200-12-31": -0.05}, + } + assert inputs.baseline.reusable_run == ReusableRun( + simulation_id=reused.simulation_id, simulation_run_id=reused.simulation_run_id + ) + assert inputs.reform.reusable_run is None + assert without_reuse.baseline.reusable_run is None + + +def test_the_two_sides_never_reuse_the_same_run(public_api_database_url) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + only = insert_reusable_run(connection, economy) + act_as_the_service(connection) + + inputs = _store_on(connection).read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=CURRENT_LAW, reuse=REUSE + ) + + assert inputs.baseline.reusable_run == ReusableRun( + simulation_id=only.simulation_id, simulation_run_id=only.simulation_run_id + ) + assert inputs.reform.reusable_run is None + + +def test_a_named_dataset_is_read_instead_of_the_regions_default( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + local_id = insert_dataset(connection, economy.catalog, "acs_local_2024") + act_as_the_service(connection) + + inputs = _store_on(connection).read_simulation_inputs( + EconomyRequest(**{**vars(REQUEST), "dataset_name": "acs_local_2024"}), + CURRENT_LAW, + ) + + assert (inputs.catalog.dataset.id, inputs.catalog.dataset.name) == ( + local_id, + "acs_local_2024", + ) + assert (inputs.simulation.choice, inputs.simulation.reform) == (CURRENT_LAW, {}) + assert inputs.simulation.reusable_run is None + + +@pytest.mark.parametrize( + ("missing", "error_type"), + [ + ("catalog version", UnknownCatalogVersionError), + ("region", UnknownRegionError), + ("dataset", UnknownDatasetError), + ("policy", UnknownPolicyError), + ("dynamic", UnknownDynamicError), + ], +) +def test_each_missing_row_a_request_names_is_refused( + public_api_database_url, missing: str, error_type: type[Exception] +) -> None: + request_fields = { + "catalog version": {"policyengine_version": "6.16.0"}, + "region": {"region_code": "state/zz"}, + "dataset": {"dataset_name": "no_such_dataset"}, + }.get(missing, {}) + choice = { + "policy": SimulationChoice(policy_id=uuid4(), dynamic_id=None), + "dynamic": SimulationChoice(policy_id=None, dynamic_id=uuid4()), + }.get(missing, CURRENT_LAW) + with rolled_back_transaction(public_api_database_url) as connection: + insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + request = EconomyRequest(**{**vars(REQUEST), **request_fields}) + + # A report names the side whose policy or dynamic is missing. + side_refused = missing in {"policy", "dynamic"} + with pytest.raises(ReportSideError if side_refused else error_type) as raised: + store.read_report_inputs( + request, baseline=CURRENT_LAW, reform=choice, reuse=REUSE + ) + with pytest.raises(error_type): + store.read_simulation_inputs(request, choice) + + if side_refused: + assert raised.value.side == "reform" + assert isinstance(raised.value.error, error_type) + + +def test_values_the_running_version_lacks_or_both_set_are_refused( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(ELASTICITY,)) + saved_on = insert_economy_catalog( + connection, country="uk", parameter_names=(CTC_AMOUNT,) + ) + # A policy of the US saved on parameters the running version lacks. + lacking_policy_id = insert_policy( + connection, + replace(saved_on.catalog, country="us"), + [StoredValue(CTC_AMOUNT, 3_000, START)], + ) + overlapping_policy_id = insert_policy( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.1, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + act_as_the_service(connection) + store = _store_on(connection) + + with pytest.raises(MissingPolicyParametersError): + store.read_simulation_inputs( + REQUEST, SimulationChoice(lacking_policy_id, dynamic_id=None) + ) + with pytest.raises(OverlappingReformError): + store.read_simulation_inputs( + REQUEST, SimulationChoice(overlapping_policy_id, dynamic_id) + ) + + +def _status(connection: Connection, table_name: str, run_id) -> tuple: + return tuple( + connection.execute( + text( + f"SELECT status, modal_call_id FROM public.{table_name} " # noqa: S608 # fixed table names + "WHERE id = :id" + ), + {"id": run_id}, + ).one() + ) + + +def test_a_report_is_created_then_its_call_recorded_or_its_run_failed( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + reports = [ + NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label=f"Report {number}", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=OPTIONS, + baseline=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + reform=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + ) + for number in (1, 2) + ] + for report in reports: + store.create_report(report) + recorded, refused = reports + + status = store.record_report_call(recorded.run_id, "fc-coordinator") + failed = store.fail_unstarted_report( + refused.run_id, refused.new_simulation_run_ids, "modal_call_refused: no" + ) + recorded_state = _status(connection, "report_runs", recorded.run_id) + refused_state = _status(connection, "report_runs", refused.run_id) + refused_baseline = _status( + connection, "simulation_runs", refused.baseline.simulation_run_id + ) + + assert (status, failed) == (ReportRunStatus.PENDING, True) + assert recorded_state == ("pending", "fc-coordinator") + assert refused_state == ("failed", None) + assert refused_baseline == ("failed", None) + + +def test_a_simulation_is_created_then_its_call_recorded_or_its_run_failed( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + simulations = [ + NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=economy.population(), + versions=VERSIONS, + ) + for _ in range(2) + ] + for simulation in simulations: + store.create_simulation(simulation) + recorded, refused = simulations + + status = store.record_simulation_call(recorded.run_id, "fc-worker") + failed = store.fail_unstarted_simulation( + refused.run_id, "modal_call_refused: no" + ) + recorded_state = _status(connection, "simulation_runs", recorded.run_id) + refused_state = _status(connection, "simulation_runs", refused.run_id) + + assert (status, failed) == (SimulationStatus.PENDING, True) + assert recorded_state == ("pending", "fc-worker") + assert refused_state == ("failed", None) + + +def test_a_report_and_a_simulation_are_read_as_they_stand( + public_api_database_url, +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection) + act_as_the_service(connection) + store = _store_on(connection) + report = NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label="Read back", + description=None, + tax_benefit_model_id=economy.catalog.model_id, + population=economy.population(), + versions=VERSIONS, + inputs=OPTIONS, + baseline=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + reform=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + ) + store.create_report(report) + store.record_report_call(report.run_id, "fc-coordinator") + + stored_report = store.read_report(report.id) + stored_simulation = store.read_simulation(report.baseline.simulation_id) + missing = (store.read_report(uuid4()), store.read_simulation(uuid4())) + + assert missing == (None, None) + assert stored_report is not None + assert stored_report.run is not None + assert (stored_report.run.id, stored_report.run.status) == ( + report.run_id, + ReportRunStatus.PENDING, + ) + assert stored_simulation is not None + assert stored_simulation.run is not None + assert stored_simulation.run.id == report.baseline.simulation_run_id + + +@pytest.mark.parametrize( + ("side", "error_type"), + [ + ("baseline", UnknownDynamicError), + ("reform", UnknownPolicyError), + ("reform", MissingPolicyParametersError), + ("baseline", OverlappingReformError), + ], +) +def test_a_reports_refused_values_name_the_side_they_belong_to( + public_api_database_url, side: str, error_type: type[Exception] +) -> None: + with rolled_back_transaction(public_api_database_url) as connection: + economy = insert_economy_catalog(connection, parameter_names=(ELASTICITY,)) + saved_on = insert_economy_catalog( + connection, country="uk", parameter_names=(CTC_AMOUNT,) + ) + lacking_policy_id = insert_policy( + connection, + replace(saved_on.catalog, country="us"), + [StoredValue(CTC_AMOUNT, 3_000, START)], + ) + overlapping_policy_id = insert_policy( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.1, START)] + ) + dynamic_id = insert_dynamic( + connection, economy.catalog, [StoredValue(ELASTICITY, -0.05, START)] + ) + act_as_the_service(connection) + refused = { + UnknownDynamicError: SimulationChoice(None, uuid4()), + UnknownPolicyError: SimulationChoice(uuid4(), None), + MissingPolicyParametersError: SimulationChoice(lacking_policy_id, None), + OverlappingReformError: SimulationChoice(overlapping_policy_id, dynamic_id), + }[error_type] + choices = {"baseline": CURRENT_LAW, "reform": CURRENT_LAW, side: refused} + + with pytest.raises(ReportSideError) as raised: + _store_on(connection).read_report_inputs( + REQUEST, + baseline=choices["baseline"], + reform=choices["reform"], + reuse=None, + ) + + assert raised.value.side == side + assert isinstance(raised.value.error, error_type) + + +class _CommitFails: + """A transaction whose block runs, and whose commit then fails.""" + + def __init__(self, connection) -> None: + self.connection = connection + + def __enter__(self): + return self.connection + + def __exit__(self, error_type, error, traceback) -> None: + if error is None: + statement = "COMMIT" + raise OperationalError(statement, {}, Exception("connection lost")) + + +class _RecordingConnection: + def __init__(self) -> None: + self.statements: list[object] = [] + + def execute(self, statement, *arguments): + self.statements.append(statement) + + +def _new_report() -> NewEconomyReport: + return NewEconomyReport( + id=uuid4(), + run_id=uuid4(), + country="us", + label="Lost commit", + description=None, + tax_benefit_model_id=uuid4(), + population=EconomyPopulation( + tax_benefit_model_version_id=uuid4(), + region_id=uuid4(), + dataset_id=uuid4(), + year=2026, + ), + versions=VERSIONS, + inputs=OPTIONS, + baseline=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + reform=ReportSimulation(uuid4(), uuid4(), CURRENT_LAW, reused=False), + ) + + +def test_a_create_whose_commit_fails_may_have_been_written() -> None: + connection = _RecordingConnection() + store = EconomyRunStore(lambda: _CommitFails(connection)) + report = _new_report() + + with pytest.raises(EconomyRecordOutcomeUnknownError) as raised: + store.create_report(report) + with pytest.raises(EconomyRecordOutcomeUnknownError): + store.create_simulation( + NewEconomySimulation( + id=uuid4(), + run_id=uuid4(), + choice=CURRENT_LAW, + population=report.population, + versions=VERSIONS, + ) + ) + + assert isinstance(raised.value.__cause__, OperationalError) + assert connection.statements + + +class _StatementFails: + def __enter__(self): + return self + + def __exit__(self, error_type, error, traceback) -> None: + return None + + def execute(self, statement, *arguments): + failed_statement = "INSERT" + raise OperationalError(failed_statement, {}, Exception("connection lost")) + + +def test_a_create_whose_statement_fails_is_not_written() -> None: + store = EconomyRunStore(_StatementFails) + + with pytest.raises(OperationalError): + store.create_report(_new_report()) + + +def test_every_side_refusal_is_raised_with_its_side(monkeypatch) -> None: + """Every error the side refusal type names is wrapped; none slips through.""" + catalog = EconomyCatalogRows( + catalog_version_id=uuid4(), + model_id=uuid4(), + region=CatalogRegion(id=uuid4(), code="us", default_dataset_id=uuid4()), + dataset=CatalogDataset(id=uuid4(), name="populace_us_2024"), + ) + monkeypatch.setattr(economy_run_store, "_catalog_rows", lambda *_: catalog) + refusals = { + UnknownPolicyError: UnknownPolicyError("no policy"), + UnknownDynamicError: UnknownDynamicError("no dynamic"), + MissingPolicyParametersError: MissingPolicyParametersError(["gov.a"]), + OverlappingReformError: OverlappingReformError(["gov.a"]), + } + assert set(get_args(SideRefusal.__value__)) == set(refusals) + + @contextmanager + def no_database(): + yield None + + store = EconomyRunStore(no_database) + for refusal in refusals.values(): + + def refuse(*_, refusal=refusal): + raise refusal + + monkeypatch.setattr(economy_run_store, "_simulation_inputs", refuse) + with pytest.raises(ReportSideError) as raised: + store.read_report_inputs( + REQUEST, baseline=CURRENT_LAW, reform=CURRENT_LAW, reuse=None + ) + assert (raised.value.side, raised.value.error) == ("baseline", refusal) diff --git a/libs/policyengine-simulation-persistence/tests/test_tables.py b/libs/policyengine-simulation-persistence/tests/test_tables.py index b36958a86..2d7ade37f 100644 --- a/libs/policyengine-simulation-persistence/tests/test_tables.py +++ b/libs/policyengine-simulation-persistence/tests/test_tables.py @@ -193,8 +193,11 @@ def test_the_input_models_are_tables_the_service_reads() -> None: assert {table.name for table in input_metadata.tables.values()} == { "households", "policies", + "dynamics", "parameters", "parameter_values", + "regions", + "datasets", "tax_benefit_models", "tax_benefit_model_versions", } diff --git a/projects/policyengine-apis-integ/tests/test_economy_generated_client.py b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py new file mode 100644 index 000000000..5f32692db --- /dev/null +++ b/projects/policyengine-apis-integ/tests/test_economy_generated_client.py @@ -0,0 +1,196 @@ +"""Hermetic wire checks of the economy routes' generated client models. + +The Public API sends its report and simulation requests through these +models, and reads each answer the routes declare. +""" + +from uuid import UUID + +import httpx +import pytest +from policyengine_api_simulation_client import AuthenticatedClient +from policyengine_api_simulation_client.api.default import ( + create_report, + create_simulation, + get_report, + get_simulation, +) +from policyengine_api_simulation_client.models import ( + EconomyError, + EconomyReportCreated, + EconomyReportRequest, + EconomyReportState, + EconomyRequestUnprocessable, + EconomySimulationCreated, + EconomySimulationRequest, + EconomySimulationState, + ReportRunStatus, +) + +REPORT_ID = "00000000-0000-4000-8000-0000000000e1" +RUN_ID = "00000000-0000-4000-8000-0000000000e2" +SIDE_IDS = [f"00000000-0000-4000-8000-0000000000f{digit}" for digit in "1234"] +POLICY_ID = "00000000-0000-4000-8000-0000000000c5" +# The body the Public API sends (``simulation_reports.report_request``). +REPORT_BODY = { + "country": "us", + "report_type": "economy", + "label": "A larger child tax credit", + "description": None, + "year": 2026, + "region_code": "us", + "dataset_name": None, + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": POLICY_ID, "dynamic_id": None}, + "include_cliffs": False, +} +SIMULATION_BODY = { + "country": "uk", + "simulation_type": "economy", + "year": 2026, + "region_code": "country/wales", + "dataset_name": None, + "policy_id": None, + "dynamic_id": None, +} +CREATED_REPORT = { + "report_id": REPORT_ID, + "report_run_id": RUN_ID, + "status": "pending", + "baseline_simulation_id": SIDE_IDS[0], + "baseline_simulation_run_id": SIDE_IDS[1], + "baseline_reused": True, + "reform_simulation_id": SIDE_IDS[2], + "reform_simulation_run_id": SIDE_IDS[3], + "reform_reused": False, +} + + +def test_a_report_request_keeps_omitted_spm_settings_omitted(): + """Omitting spm uses the bundle's defaults; null is not a selection.""" + omitted = EconomyReportRequest.from_dict(REPORT_BODY).to_dict() + given = EconomyReportRequest.from_dict( + {**REPORT_BODY, "spm": {"geography_kind": "national"}} + ).to_dict() + + assert "spm" not in omitted + assert omitted["reform"] == {"policy_id": POLICY_ID, "dynamic_id": None} + assert given["spm"] == {"geography_kind": "national"} + + +def _client_answering(status_code: int, body: dict) -> AuthenticatedClient: + transport = httpx.MockTransport( + lambda request: httpx.Response(status_code, json=body) + ) + client = AuthenticatedClient(base_url="https://simulation.example", token="test") + client.set_httpx_client( + httpx.Client(base_url="https://simulation.example", transport=transport) + ) + return client + + +def _create_report(status_code: int, body: dict): + return create_report.sync_detailed( + client=_client_answering(status_code, body), + body=EconomyReportRequest.from_dict(REPORT_BODY), + ) + + +def test_a_created_report_is_read_with_its_runs(): + response = _create_report(201, CREATED_REPORT) + + assert isinstance(response.parsed, EconomyReportCreated) + assert response.parsed.status is ReportRunStatus.PENDING + assert response.parsed.to_dict() == CREATED_REPORT + + +def test_a_created_simulation_is_read_with_its_run(): + created = { + "simulation_id": REPORT_ID, + "simulation_run_id": RUN_ID, + "status": "pending", + } + + response = create_simulation.sync_detailed( + client=_client_answering(201, created), + body=EconomySimulationRequest.from_dict(SIMULATION_BODY), + ) + + assert isinstance(response.parsed, EconomySimulationCreated) + assert response.parsed.to_dict() == created + + +@pytest.mark.parametrize( + ("status_code", "body"), + [ + (400, {"code": "unknown_region", "detail": "no region us"}), + ( + 503, + { + "code": "modal_call_refused", + "detail": "Modal did not accept the report's call.", + "report_id": REPORT_ID, + }, + ), + (500, {"code": "internal_error", "detail": "The service could not answer."}), + ], +) +def test_an_error_is_read_with_its_code_and_any_record(status_code, body): + response = _create_report(status_code, body) + + assert isinstance(response.parsed, EconomyError) + assert response.parsed.to_dict() == body + + +def test_stored_values_that_cannot_be_calculated_are_read_with_their_code(): + body = { + "detail": [ + { + "loc": ["body", "policy_id", "gov.a"], + "msg": "no such parameter", + "type": "missing_parameter", + } + ], + "code": "missing_parameters", + } + + response = _create_report(422, body) + + assert isinstance(response.parsed, EconomyRequestUnprocessable) + assert response.parsed.to_dict() == body + + +def test_a_report_and_a_simulation_are_read_with_their_newest_runs(): + report = { + "report_id": REPORT_ID, + "country": "us", + "baseline_simulation_id": SIDE_IDS[0], + "reform_simulation_id": SIDE_IDS[2], + "run": { + "run_id": RUN_ID, + "status": "succeeded", + "country_package_version": "2.2.1", + "policyengine_version": "6.17.1", + "error_message": None, + "created_at": "2026-11-12T15:00:00Z", + "started_at": "2026-11-12T15:00:05Z", + "completed_at": "2026-11-12T15:04:00Z", + "baseline_simulation_run_id": SIDE_IDS[1], + "reform_simulation_run_id": SIDE_IDS[3], + }, + } + simulation = {"simulation_id": REPORT_ID, "run": None} + + report_response = get_report.sync_detailed( + UUID(REPORT_ID), client=_client_answering(200, report) + ) + simulation_response = get_simulation.sync_detailed( + UUID(REPORT_ID), client=_client_answering(200, simulation) + ) + + assert isinstance(report_response.parsed, EconomyReportState) + run = report_response.parsed.run + assert run is not None + assert run.status is ReportRunStatus.SUCCEEDED + assert isinstance(simulation_response.parsed, EconomySimulationState) + assert simulation_response.parsed.to_dict() == simulation diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index e467e87b8..5972daaff 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -116,6 +116,32 @@ be a US household stored in the staging Public API with a default year; each run adds one `household_jobs` row for it. An empty or malformed variable fails the check. +## Economy reports and simulations + +`POST /reports` and `POST /simulations`, and their reads, use the same +database, Modal token and v2 manifest settings as the household routes, and +the deploy's `V2_SIMULATION_OUTPUT_BUCKET`, the `beta` or `prod` variable of +that name: the bucket that keeps v2 simulation output files for good. A +report reuses only a succeeded simulation run whose file is in it; the +service compares output URIs with it and never reads the bucket, so it needs +no storage permission. A missing or invalid bucket name stops the deploy. +The coordinator and workers get the same bucket from the Modal secret the +Stage 12 secret sync writes. + +### The staging economy check + +Every staging deployment runs `staging_economy_check` when the household +check runs: `.github/scripts/simulation-staging-economy-check.sh` creates a +small report, current law against current law for `country/wales` in 2026, +through the candidate's tagged URL with the authenticated tests' client +credentials, follows it through `GET /reports/{report_id}`, and requires its +run to succeed within 30 minutes. `deployment_ready` requires the check to +pass in staging, so a failure stops the production deployment. Each run adds +a report and its runs to the staging database. Once a .py version's first +check has succeeded, later checks on that version reuse its simulation runs, +so they check the coordinator's use of stored outputs rather than the +workers', which the v2 deployment's own validation calls. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 2e3a3f6ba..bc32cc1f6 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -42,6 +42,10 @@ Auth0 permission: | `GET /internal/stage12/reports/{evaluation_id}` | `stage12:operate` | | `POST /household-calculations` | `simulation:submit` | | `GET /household-calculations/{household_job_id}` | `simulation:read` | +| `POST /reports` | `simulation:submit` | +| `GET /reports/{report_id}` | `simulation:read` | +| `POST /simulations` | `simulation:submit` | +| `GET /simulations/{simulation_id}` | `simulation:read` | `GET /health`, `GET /ready`, `GET /versions`, `GET /versions/{kind}` and `POST /ping` stay public. The entrypoint answers `GET /health` and `POST /ping` @@ -102,3 +106,53 @@ failure or a request the service built and the household contract refuses, is the JSON error with code `internal_error`. `DEPLOYMENT.md` describes the deployment settings and the staging check. +## Economy reports and simulations + +`POST /reports` creates an economy report: a baseline and a reform +simulation of one region's data for one year. It names the `country`, +`report_type` (`economy`), `label`, `description`, `year`, `region_code` +and `dataset_name` (a region and dataset of the country's catalog version +for the default worker's .py version; null for the region's default +dataset), each side's stored `policy_id` and `dynamic_id` (null for current +law), `include_cliffs`, and `spm` (SPM settings; omit it, since null is +refused). `POST /simulations` creates a standalone economy simulation from +the same fields for one policy and dynamic. Neither names a .py version or a +record ID: the service runs the v2 version manifest's default worker and +chooses every ID. + +The service reads the rows itself and, in one transaction, creates the +records and their first runs, `pending` with no Modal call ID. A report +side reuses a simulation run that already succeeded instead of creating +one, when it is the newest run of the same policy, dynamic, catalog +version, dataset, region and year on the same country package and .py +versions, its output file is in the output bucket +(`V2_SIMULATION_OUTPUT_BUCKET`), a succeeded report run with the same +options used it, and no report run has refused to reuse it; the two sides +never name the same run. The service then starts the run's Modal call by +name: the report run's coordinator, given the report and its run IDs and no +Stage 12 evaluation record, or a standalone run's simulation function, +given its run's ID. It records the call's ID, and answers 201 with the +record's ID, its run's ID and status, and, for a report, each side's +simulation and run and whether it was reused. + +A request is refused before anything is written: with 400 and a `code` when +a row it names is missing (`unknown_region`, `unknown_dataset` and +`unknown_dynamic` name the caller's own choice; `unknown_catalog_version`, +`unknown_policy`, `no_economy_worker`, `dataset_not_in_bundle` and +`missing_reference` are the service's state, such as a catalog version +missing while a release deploys), and with 422 and a `code` when a policy or +dynamic sets parameters the catalog version lacks (`missing_parameters`) or +both set one (`overlapping_parameters`). The routes answer 503 with +`Retry-After` when the manifest or the database can't be read. When Modal +refuses the call, or doesn't confirm it within ten seconds, the run (and a +report run's new simulation runs) is failed at once and the 503 +(`modal_call_refused`) names the record, as `report_id` or `simulation_id`; +a run claimed meanwhile is answered with 201. Every 500 is the JSON error +with code `internal_error`, and may follow a write. No `detail` holds an +exception's text. + +`GET /reports/{report_id}` and `GET /simulations/{simulation_id}` return the +record and its newest run (created last, ties to the larger ID): its status, +versions, error message and times, and a report run's simulation runs. The +Public API reads results from the tables itself. + diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 2d8bbb804..0c99f05e4 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -41,7 +41,7 @@ from policyengine_simulation_observability.observability import ( init_simulation_observability, ) -from policyengine_simulation_persistence import HouseholdJobStore +from policyengine_simulation_persistence import EconomyRunStore, HouseholdJobStore from starlette.datastructures import MutableHeaders from policyengine_simulation_entry.auth import ( @@ -52,6 +52,18 @@ ) from policyengine_simulation_entry.backend import OldGatewayBackend, SimulationBackend from policyengine_simulation_entry.config import Settings +from policyengine_simulation_entry.economy_routes import ( + EconomySubmissions, + add_economy_routes, +) +from policyengine_simulation_entry.economy_submissions import ( + EconomySubmissionService, + ServiceSettings, +) +from policyengine_simulation_entry.economy_workers import ( + ManifestEconomyWorkers, + ModalEconomyRunCaller, +) from policyengine_simulation_entry.entry_runtime import ( BACKEND_RESPONSE_HEADER, STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS, @@ -110,6 +122,8 @@ SIMULATION_DATABASE_READINESS_TIMEOUT_SECONDS = 3.0 # The database ends each statement of a household route that runs longer. HOUSEHOLD_STATEMENT_TIMEOUT_MILLISECONDS = 5000 +# The database ends each statement of an economy route that runs longer. +ECONOMY_STATEMENT_TIMEOUT_MILLISECONDS = 5000 def _temporary_submission_payload( @@ -523,6 +537,49 @@ def household_calculations_for( ) +def _register_economy_routes( + app: FastAPI, + entry: EntryRuntime, + submissions: EconomySubmissions | 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) + + +def economy_submissions_for( + settings: Settings, database: SimulationDatabaseConnection | None +) -> EconomySubmissions | None: + """Build the economy submission service when the service is configured. + + It needs the simulation database and the v2 Modal environment; without + them the economy routes answer 503. Nothing is read from Modal until the + first submission. Without the output bucket, no report reuses a run. + """ + if not settings.economy_runs_configured or not isinstance( + database, SimulationDatabase + ): + return None + return EconomySubmissionService( + store=EconomyRunStore.from_engine( + database.engine, + statement_timeout_milliseconds=ECONOMY_STATEMENT_TIMEOUT_MILLISECONDS, + ), + workers=ManifestEconomyWorkers.from_modal( + manifest_name=settings.v2_manifest_name, + environment=settings.v2_modal_environment, + ), + caller=ModalEconomyRunCaller(settings.v2_modal_environment), + settings=ServiceSettings( + environment=settings.environment, + modal_environment=settings.v2_modal_environment, + output_bucket=settings.v2_simulation_output_bucket or None, + ), + ) + + def _register_service_routes(app: FastAPI, entry: EntryRuntime) -> None: """Add the public version, health, readiness and ping routes.""" database_readiness = ( @@ -678,6 +735,13 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: submitter=simulation_submitter, reader=simulation_reader, ) + _register_economy_routes( + app, + entry, + economy_submissions_for(runtime_settings, runtime_database), + submitter=simulation_submitter, + reader=simulation_reader, + ) _register_service_routes(app, entry) return app diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index 883b3c077..5b6ae75dd 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -102,6 +102,10 @@ class Settings: # manifest, and the manifest's name. Empty outside deployed environments. v2_modal_environment: str = "" v2_manifest_name: str = V2_VERSION_MANIFEST_NAME + # The bucket that keeps v2 simulation output files for good. An economy + # report reuses only a succeeded run whose file is in it; without it, no + # run is reused. Empty outside deployed environments. + v2_simulation_output_bucket: str = "" @classmethod def from_env(cls) -> Settings: @@ -157,6 +161,7 @@ def from_env(cls) -> Settings: v2_manifest_name=os.getenv( "SIMULATION_V2_MANIFEST_NAME", V2_VERSION_MANIFEST_NAME ), + v2_simulation_output_bucket=os.getenv("V2_SIMULATION_OUTPUT_BUCKET", ""), ) @property @@ -180,6 +185,11 @@ def household_calculations_configured(self) -> bool: """Whether the database and the v2 Modal environment are both set.""" return bool(self.simulation_database_url and self.v2_modal_environment) + @property + def economy_runs_configured(self) -> bool: + """Whether economy reports and simulations can run: as household ones.""" + return self.household_calculations_configured + def validate(self) -> None: """Raise ConfigurationError unless the service can start safely. diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py new file mode 100644 index 000000000..b161c91e7 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_api_models.py @@ -0,0 +1,295 @@ +"""Request and response models of the economy simulation and report routes. + +The Public API names a region and dataset of the country's catalog by code +and name, and each simulation's stored policy and dynamic by ID; the +Simulation Entrypoint reads them from the Public API's tables itself, and +calculates on the v2 version manifest's default worker. A request names no +.py version and no record ID: the service runs its default worker's version +and chooses every ID. + +Answers the Public API reads: + +- 201 with the created record's ID and its first run's ID; +- 400 with ``code`` and ``detail``: one of ``CallerRowRefusalCode`` when + the region, dataset or a dynamic the request names does not exist, or + another code for the service's own state, such as a catalog version + missing while a release deploys. Nothing is written; +- 422 when the request's fields, or the stored values it names, cannot be + calculated. Nothing is written; +- 503 with ``code`` and ``detail`` when the service cannot answer now. Only + ``modal_call_refused`` follows a write, and it names the record created + (``report_id`` or ``simulation_id``), whose run was failed; +- 500 with ``code`` ``internal_error``, possibly after a write: an + unexpected failure, or a commit that failed after the rows were inserted. + +``detail`` never holds an exception's text, and ``code`` is a fixed +identifier. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Annotated, Any, Literal +from uuid import UUID + +from policyengine_simulation_contract.spm import SPMSelection +from policyengine_simulation_contract.stage12_bundle import CountryId +from policyengine_simulation_persistence import ReportRunStatus, SimulationStatus +from pydantic import BaseModel, ConfigDict, Field, StrictBool, field_validator +from pydantic.json_schema import SkipJsonSchema + +# The bounds the tables accept. +SimulationYear = Annotated[int, Field(ge=1900, le=2200, description="The year")] +CatalogKey = Annotated[str, Field(min_length=1, max_length=255, pattern=r"^\S+$")] +ReportLabel = Annotated[str, Field(min_length=1, max_length=255)] +ReportDescription = Annotated[str, Field(max_length=10_000)] + +# Why a request is refused (400), naming a row the caller chose that doesn't +# exist: the Public API answers these as its caller's error. +CallerRowRefusalCode = Literal["unknown_region", "unknown_dataset", "unknown_dynamic"] +# Why a request is refused (400) for the service's own state, or a row the +# Public API checks before asking. +ServiceStateRefusalCode = Literal[ + # The country has no catalog version for the default worker's .py version. + "unknown_catalog_version", + "unknown_policy", + # The default worker cannot calculate the country's economy runs. + "no_economy_worker", + # The default worker's bundle lacks the region's default dataset. + "dataset_not_in_bundle", + # The database refused a row: a row it names was deleted after the reads. + "missing_reference", +] +EconomyRefusalCode = Literal[CallerRowRefusalCode, ServiceStateRefusalCode] +# Why stored values cannot be calculated (422). +EconomyUnprocessableCode = Literal["missing_parameters", "overlapping_parameters"] +# Why a route cannot answer now (503); retry after ``Retry-After`` seconds. +EconomyUnavailableCode = Literal[ + # The service has no database or v2 Modal environment configured. + "not_configured", + "manifest_unavailable", + "database_unavailable", + # Modal refused the run's call, or did not accept it in time; the record + # named was created and its run failed at once. + "modal_call_refused", +] +# Every code an economy route's error answer carries. +EconomyErrorCode = Literal[ + EconomyRefusalCode, + EconomyUnprocessableCode, + EconomyUnavailableCode, + "internal_error", + "unknown_report", + "unknown_simulation", +] + + +def _omit_null_default(schema: dict[str, Any]) -> None: + """Leave out a field's null default: the field is omitted rather than null.""" + schema.pop("default", None) + + +class _Request(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, allow_inf_nan=False) + + +class PolicyChoice(_Request): + """A simulation's stored policy and dynamic; neither means current law.""" + + policy_id: UUID | None = None + dynamic_id: UUID | None = None + + +class EconomySimulationRequest(PolicyChoice): + """An economy simulation: a policy and dynamic applied to one region's data. + + ``region_code`` is a region of the country's catalog version for the + default worker's .py version (``us``, ``state/ca``, ``country/wales``), + and ``dataset_name`` one of that version's datasets; without it, the + region's default dataset is used. + """ + + country: CountryId + simulation_type: Literal["economy"] + year: SimulationYear + region_code: CatalogKey + dataset_name: CatalogKey | None = None + + +class EconomyReportRequest(_Request): + """An economy report: a baseline and a reform simulation of one region's data. + + Both simulations use the report's year, region and dataset, as + ``EconomySimulationRequest`` describes them. ``include_cliffs`` also + calculates the baseline's and reform's cliff impact. ``spm`` selects the + SPM forecast scenario and geography; omit it, rather than sending null, + to use the bundle's certified defaults. + """ + + country: CountryId + report_type: Literal["economy"] + label: ReportLabel + description: ReportDescription | None = None + year: SimulationYear + region_code: CatalogKey + dataset_name: CatalogKey | None = None + baseline: PolicyChoice + reform: PolicyChoice + include_cliffs: StrictBool = False + spm: SPMSelection | SkipJsonSchema[None] = Field( + default=None, json_schema_extra=_omit_null_default + ) + + @field_validator("spm", mode="before") + @classmethod + def refuse_null_spm(cls, value: object) -> object: + """Refuse an explicit null; omit ``spm`` to send no settings.""" + if value is None: + msg = "spm must be an object when given; omit it to send no settings" + raise ValueError(msg) + return value + + def options(self) -> dict[str, Any]: + """Return the report's options as it stores them (``reports.inputs``). + + ``spm`` holds only the settings the request gave, and is left out + when it gave none. + """ + options: dict[str, Any] = {"include_cliffs": self.include_cliffs} + if self.spm is not None: + options["spm"] = self.spm.model_dump(mode="json") + return options + + +class _Answer(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class EconomySimulationCreated(_Answer): + """The simulation a request created, its first run and that run's status. + + The status is the run's when its call's ID was recorded: ``pending``, or + whatever the simulation function had made it by then if it claimed the + run first. A run whose call Modal didn't confirm, or whose call's ID + couldn't be recorded, is answered ``pending``. + """ + + simulation_id: UUID + simulation_run_id: UUID + status: SimulationStatus + + +class EconomyReportCreated(_Answer): + """The report a request created, its first run and that run's status. + + The status is the run's when its call's ID was recorded: ``pending``, or + whatever the coordinator had made it by then if it claimed the run + first. A run whose call Modal didn't confirm, or whose call's ID couldn't + be recorded, is answered ``pending``. Each side names its + simulation and simulation run; ``…_reused`` is true for a side that + names a run that already succeeded, rather than a new one. + """ + + report_id: UUID + report_run_id: UUID + status: ReportRunStatus + baseline_simulation_id: UUID + baseline_simulation_run_id: UUID + baseline_reused: bool + reform_simulation_id: UUID + reform_simulation_run_id: UUID + reform_reused: bool + + +class SimulationRunRead(_Answer): + """A simulation run's state: its status, versions, error message and times.""" + + run_id: UUID + status: SimulationStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + + +class ReportRunRead(_Answer): + """A report run's state, and the simulation runs it compares.""" + + run_id: UUID + status: ReportRunStatus + country_package_version: str + policyengine_version: str + error_message: str | None + created_at: datetime + started_at: datetime | None + completed_at: datetime | None + baseline_simulation_run_id: UUID + reform_simulation_run_id: UUID | None + + +class EconomyReportState(_Answer): + """A report, its two simulations and its newest run, as they stand. + + The newest run is the one created last, ties going to the larger ID; + ``run`` is null for a report with none. Its results are read from the + Public API. + """ + + report_id: UUID + country: str + baseline_simulation_id: UUID | None + reform_simulation_id: UUID | None + run: ReportRunRead | None + + +class EconomySimulationState(_Answer): + """A simulation and its newest run, as they stand; ``run`` null with none.""" + + simulation_id: UUID + run: SimulationRunRead | None + + +class EconomyError(_Answer): + """Why an economy route refused a request or cannot answer now. + + ``report_id`` (on ``/reports``) or ``simulation_id`` (on + ``/simulations``) is set only with ``modal_call_refused``: it names the + record the service created, whose run it failed because Modal refused + the call or didn't accept it in time. + """ + + code: EconomyErrorCode + detail: str + report_id: UUID | None = None + simulation_id: UUID | None = None + + +class EconomyRequestReason(BaseModel): + """One reason a request cannot be calculated, as FastAPI writes a field error. + + ``loc`` starts with the request field the reason traces to. FastAPI's own + field errors may also carry ``input`` and ``ctx``. + """ + + model_config = ConfigDict(extra="allow", frozen=True) + + loc: list[str | int] + msg: str + type: str + + +class EconomyRequestUnprocessable(_Answer): + """A request that cannot be calculated (422), one reason per refused field. + + Request fields FastAPI refuses carry no ``code``. Stored values that + cannot be calculated carry one: ``missing_parameters`` for a policy or + dynamic setting parameters the catalog version lacks, and + ``overlapping_parameters`` for a policy and dynamic setting the same one. + """ + + detail: list[EconomyRequestReason] + code: EconomyUnprocessableCode | SkipJsonSchema[None] = Field( + default=None, json_schema_extra=_omit_null_default + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py new file mode 100644 index 000000000..f04378af8 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_calculation_inputs.py @@ -0,0 +1,261 @@ +"""What the coordinator or a simulation function is given to calculate a run. + +A report run's coordinator takes a ``ReportExecutionInput`` naming the +report run as its ``evaluation_id`` and each side's simulation run as that +side's ``simulation_execution_id``, and a ``Stage12InvocationContext`` whose +``report_run`` names the three runs. It claims the report run and gives +each simulation function the run it claims; no evaluation record is passed. +A standalone simulation run's function takes that run's +``SimulationExecutionInput`` and a context naming the run +(``simulation_run_id``). + +Both sides of a report calculate on the same dataset, region, year, +options and bundle; only their policies differ. The dataset is the bundle's +dataset whose identity is the catalog dataset's name, the region is the +catalog region's code, which .py's region registry reads, and the policies +are the stored policy and dynamic values by parameter name and period. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final + +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + ReportAggregate, + ReportExecutionInput, + ReportRunIds, + SimulationExecutionInput, + SimulationRole, + Stage12InvocationContext, +) + +if TYPE_CHECKING: + from datetime import datetime + from uuid import UUID + + from policyengine_simulation_contract.stage12_bundle import Stage12Dataset + from policyengine_simulation_contract.stage12_execution import CountryId + + from policyengine_simulation_entry.economy_workers import EconomyWorker + +# Every aggregate a report computes, as the coordinator requires. +REPORT_AGGREGATES: Final = tuple(ReportAggregate) +# How long the coordinator keeps a call's Stage 12 artifacts; the longest the +# contract allows. Simulation output files go to the output bucket instead. +ARTIFACT_RETENTION: Final = timedelta(days=30) + + +class DatasetNotInBundleError(LookupError): + """The default worker's bundle has no dataset with the catalog dataset's name.""" + + +@dataclass(frozen=True) +class CalculationScope: + """What every simulation of a request calculates on, as the worker reads it. + + ``options`` are the simulation options both sides share: ``include_cliffs`` + when true, and the SPM selection when one was given. + """ + + country: CountryId + region_code: str + dataset: Stage12Dataset + year: int + options: dict[str, Any] + + +@dataclass(frozen=True) +class CallOrigin: + """Where a call comes from: the deployment, the request and when it was made.""" + + environment: str + modal_environment: str + request_id: str + created_at: datetime + + +@dataclass(frozen=True) +class RunCalculation: + """The two arguments the run's Modal function is called with.""" + + payload: dict[str, Any] + context: dict[str, Any] + + +def bundle_dataset(worker: EconomyWorker, dataset_name: str) -> Stage12Dataset: + """Return the worker's bundle dataset whose identity is ``dataset_name``. + + Raise DatasetNotInBundleError when it has none. + """ + dataset = next( + (item for item in worker.country.datasets if item.identity == dataset_name), + None, + ) + if dataset is None: + msg = ( + f"the default v2 worker, .py {worker.version}, has no dataset " + f"{dataset_name}" + ) + raise DatasetNotInBundleError(msg) + return dataset + + +def simulation_options( + *, include_cliffs: bool, spm: dict[str, Any] | None +) -> dict[str, Any]: + """Return the options both sides of a calculation share.""" + options: dict[str, Any] = {} + if include_cliffs: + options["include_cliffs"] = True + if spm is not None: + options["spm"] = spm + return options + + +def _simulation_input( # noqa: PLR0913 # the run, its role, its policy and what it calculates + worker: EconomyWorker, + scope: CalculationScope, + *, + evaluation_id: UUID, + simulation_run_id: UUID, + role: SimulationRole, + policy: dict[str, Any], +) -> SimulationExecutionInput: + country = worker.country + dataset = scope.dataset + return SimulationExecutionInput( + evaluation_id=evaluation_id, + simulation_execution_id=simulation_run_id, + role=role, + policy=policy, + population=DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256=dataset.sha256, + ) + ), + year=scope.year, + geography=GeographySelection(country=scope.country, region=scope.region_code), + options=scope.options, + bundle=BundleProvenance( + policyengine_version=worker.policyengine_version, + country_package_name=country.country_package_name, + country_package_version=country.country_package_version, + dataset=DatasetProvenance( + identity=dataset.identity, + uri=dataset.uri, + artifact_revision=dataset.artifact_revision, + data_package_name=country.data_package_name, + data_package_version=country.data_package_version, + ), + bundle_manifest_sha256=worker.bundle_manifest_sha256, + ), + ) + + +def _context( + worker: EconomyWorker, origin: CallOrigin, *, runs: ReportRunIds | UUID +) -> Stage12InvocationContext: + """Return the call's context, naming ``runs``. + + ``runs`` is a report run's IDs, for its coordinator, or a standalone + simulation run's ID, for its simulation function. The artifact prefix + names that report run or simulation run, so the two can't differ. + """ + if isinstance(runs, ReportRunIds): + report_run, simulation_run_id, run_id = runs, None, runs.report_run_id + else: + report_run, simulation_run_id, run_id = None, runs, runs + created_at = origin.created_at + return Stage12InvocationContext( + request_id=origin.request_id, + environment=origin.environment, + modal_environment=origin.modal_environment, + worker_version=worker.version, + modal_application=worker.application_name, + simulation_callable=worker.simulation_callable, + version_manifest_sha256=worker.version_manifest_sha256, + bundle_manifest_sha256=worker.bundle_manifest_sha256, + artifact_prefix=( + f"stage-12-runs/{origin.environment}/{created_at:%Y}/{created_at:%m}/" + f"{run_id}" + ), + created_at=created_at, + retention_expires_at=created_at + ARTIFACT_RETENTION, + report_run=report_run, + simulation_run_id=simulation_run_id, + ) + + +def report_run_calculation( # noqa: PLR0913 # the runs, both policies and what they calculate + worker: EconomyWorker, + scope: CalculationScope, + origin: CallOrigin, + *, + run_ids: ReportRunIds, + baseline_policy: dict[str, Any], + reform_policy: dict[str, Any], +) -> RunCalculation: + """Return the coordinator's report and context for the report run.""" + report_run_id = run_ids.report_run_id + report = ReportExecutionInput( + evaluation_id=report_run_id, + baseline=_simulation_input( + worker, + scope, + evaluation_id=report_run_id, + simulation_run_id=run_ids.baseline_simulation_run_id, + role=SimulationRole.BASELINE, + policy=baseline_policy, + ), + reform=_simulation_input( + worker, + scope, + evaluation_id=report_run_id, + simulation_run_id=run_ids.reform_simulation_run_id, + role=SimulationRole.REFORM, + policy=reform_policy, + ), + requested_aggregates=REPORT_AGGREGATES, + ) + context = _context(worker, origin, runs=run_ids) + return RunCalculation( + payload=report.model_dump(mode="json"), + context=context.model_dump(mode="json"), + ) + + +def simulation_run_calculation( + worker: EconomyWorker, + scope: CalculationScope, + origin: CallOrigin, + *, + simulation_run_id: UUID, + policy: dict[str, Any], +) -> RunCalculation: + """Return the simulation function's input and context for a standalone run. + + The context names the run the function claims and records. + """ + simulation = _simulation_input( + worker, + scope, + evaluation_id=simulation_run_id, + simulation_run_id=simulation_run_id, + role=SimulationRole.STANDALONE, + policy=policy, + ) + context = _context(worker, origin, runs=simulation_run_id) + return RunCalculation( + payload=simulation.model_dump(mode="json"), + context=context.model_dump(mode="json"), + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py new file mode 100644 index 000000000..9dd6e5285 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_routes.py @@ -0,0 +1,586 @@ +"""The economy routes: create a report or a simulation, and start its calculation. + +- ``POST /reports`` (permission ``simulation:submit``) creates an economy + report, its simulations and its first run, and starts the report run's + coordinator, answering 201 with the report's ID, its run's ID and status, + and each side's simulation and simulation run. +- ``POST /simulations`` (``simulation:submit``) creates a standalone + economy simulation and its first run, and starts the run's simulation + function, answering 201 with the simulation's ID, its run's ID and status. +- ``GET /reports/{report_id}`` and ``GET /simulations/{simulation_id}`` + (``simulation:read``) return the record with its newest run's status, + versions, error message and times. The Public API reads results from the + tables itself; these reads let a deploy check follow a run. + +A refused request answers 400 (a row it names, or a function it needs, is +missing; ``code`` says which) or 422 (its fields, or the stored values it +names, cannot be calculated), and writes nothing. Either route answers 503 +with ``Retry-After`` when it is not configured or a dependency cannot be +read, before writing anything, or when Modal refused the run's call: then +the record was created and its run failed, and the answer names it. Every +500 the routes answer is the JSON error with code ``internal_error``, and may +follow a write; a failure outside a handler's ``try`` gets the request +middleware's own plain-text 500, which the Public API reads alike. + +A request gets its observability ID from the caller's header, or a new one, +and passes it to the Modal call it starts in the observability context; the +call's ``request_id`` is the request's own ``X-PolicyEngine-Request-Id``, +or a new UUID when that doesn't fit the call's context. +""" + +from __future__ import annotations + +import logging +import re +from typing import TYPE_CHECKING, Any, Final, Protocol +from uuid import UUID, uuid4 + +from fastapi import APIRouter, Depends, Request, status +from fastapi.responses import JSONResponse + +from policyengine_simulation_entry.economy_api_models import ( + EconomyError, + EconomyReportCreated, + EconomyReportRequest, + EconomyReportState, + EconomyRequestUnprocessable, + EconomySimulationCreated, + EconomySimulationRequest, + EconomySimulationState, + ReportRunRead, + SimulationRunRead, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, +) +from policyengine_simulation_entry.route_support import ( + log_unexpected_failure, + observability_context, +) + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from fastapi import FastAPI + from policyengine_simulation_persistence import ( + EconomyReportState as StoredReport, + ) + from policyengine_simulation_persistence import ( + EconomySimulationState as StoredSimulation, + ) + from pydantic import BaseModel + + from policyengine_simulation_entry.auth import AuthorizationDependency + from policyengine_simulation_entry.economy_submission_errors import RecordName + from policyengine_simulation_entry.entry_runtime import EntryRuntime + +logger = logging.getLogger(__name__) + +# The request IDs the observability runtime accepts, which a Modal call's +# context also carries (``Stage12InvocationContext.request_id`` takes at most +# 255 characters). Any other request ID is replaced by a new one for the call. +CALL_REQUEST_ID: Final = re.compile(r"[A-Za-z0-9._:-]{1,128}") +REPORTS_PATH = "/reports" +REPORT_PATH = "/reports/{report_id}" +SIMULATIONS_PATH = "/simulations" +SIMULATION_PATH = "/simulations/{simulation_id}" +RETRY_AFTER_SECONDS = 10 + +_ERROR = {"model": EconomyError} +_RETRY_AFTER = { + "Retry-After": { + "description": "Seconds to wait before asking again", + "schema": {"type": "integer"}, + } +} +_FORBIDDEN = { + "description": "The caller's token is missing, invalid or lacks the permission" +} +_REFUSED = { + "description": ( + "A row the request names, or a function it needs, is missing; nothing " + "was written. unknown_region, unknown_dataset and unknown_dynamic name " + "the caller's own region, dataset or dynamic; any other code is the " + "service's state, such as a catalog version missing while a release " + "deploys" + ), + **_ERROR, +} +_UNPROCESSABLE = { + "description": ( + "The request's fields, or the stored policies and dynamics it names, " + "cannot be calculated; nothing was written. Stored values that cannot " + "be calculated carry a code: missing_parameters or " + "overlapping_parameters" + ), + "model": EconomyRequestUnprocessable, +} +_INTERNAL_ERROR = { + "description": ( + "The service failed unexpectedly, or could not confirm its rows were " + "written (internal_error), possibly after creating the record" + ), + **_ERROR, +} + + +_READ_UNPROCESSABLE = { + "description": "The ID is not a UUID", + "model": EconomyRequestUnprocessable, +} +_READ_UNAVAILABLE = { + "description": ( + "Economy runs are not configured, or the database cannot be read; " + "retry after Retry-After seconds" + ), + "headers": _RETRY_AFTER, + **_ERROR, +} +_READ_INTERNAL_ERROR = { + "description": "The service failed unexpectedly (internal_error)", + **_ERROR, +} + + +def _unavailable_description(record: RecordName) -> dict[str, Any]: + return { + "description": ( + "Economy runs are not configured, or the v2 version manifest or the " + "database cannot be read, and nothing was written; or Modal refused " + "the run's call, or didn't accept it in time " + f"(modal_call_refused): then the {record} named by {record}_id was " + "created and its run failed. Retry after Retry-After seconds" + ), + "headers": _RETRY_AFTER, + **_ERROR, + } + + +class EconomySubmissions(Protocol): + """Create economy reports and simulations (``EconomySubmissionService``).""" + + async def submit_report( + self, + request: EconomyReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + """Create the report and its runs and start its coordinator.""" + ... + + async def submit_simulation( + self, + request: EconomySimulationRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + """Create the simulation and its run and start the run's function.""" + ... + + async def read_report(self, report_id: UUID) -> StoredReport | None: + """Return the report with its newest run, or None.""" + ... + + async def read_simulation(self, simulation_id: UUID) -> StoredSimulation | None: + """Return the simulation with its newest run, or None.""" + ... + + +def _error_response(status_code: int, error: EconomyError) -> JSONResponse: + headers = ( + {"Retry-After": str(RETRY_AFTER_SECONDS)} + if status_code == status.HTTP_503_SERVICE_UNAVAILABLE + else None + ) + return JSONResponse( + error.model_dump(mode="json", exclude_none=True), + status_code=status_code, + headers=headers, + ) + + +def _not_configured() -> JSONResponse: + return _error_response( + status.HTTP_503_SERVICE_UNAVAILABLE, + EconomyError(code="not_configured", detail="Economy runs are not configured."), + ) + + +_SUBMISSION_FAILURES = ( + EconomySubmissionRefusedError, + EconomySubmissionUnprocessableError, + EconomySubmissionUnavailableError, + EconomySubmissionOutcomeUnknownError, +) + + +def _submission_failure_response(error: Exception, record: RecordName) -> JSONResponse: + """Answer a submission the service refused or could not serve.""" + if isinstance(error, EconomySubmissionRefusedError): + return _error_response( + status.HTTP_400_BAD_REQUEST, + EconomyError(code=error.code, detail=error.detail), + ) + if isinstance(error, EconomySubmissionUnprocessableError): + return JSONResponse( + { + "detail": [ + {"loc": list(reason.loc), "msg": reason.msg, "type": reason.type} + for reason in error.reasons + ], + "code": error.code, + }, + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + ) + if isinstance(error, EconomySubmissionOutcomeUnknownError): + return _outcome_unknown_response(error) + assert isinstance(error, EconomySubmissionUnavailableError) # noqa: S101 # the last of _SUBMISSION_FAILURES + record_ids = {} if error.record_id is None else {f"{record}_id": error.record_id} + return _error_response( + status.HTTP_503_SERVICE_UNAVAILABLE, + EconomyError(code=error.code, detail=error.detail, **record_ids), + ) + + +def call_request_id_for(request: Request) -> str: + """Return the request's ID for the Modal call, or a new one that fits it. + + The request middleware keeps the caller's ``X-PolicyEngine-Request-Id`` + unchecked. A replaced ID matches none of the request's own log records; + the observability ID, which the call's observability context carries, + still joins them. + """ + request_id = getattr(request.state, "request_id", "") + if isinstance(request_id, str) and CALL_REQUEST_ID.fullmatch(request_id): + return request_id + return str(uuid4()) + + +def _outcome_unknown_response( + error: EconomySubmissionOutcomeUnknownError, +) -> JSONResponse: + """Answer a create whose commit can't be confirmed with the JSON 500. + + The Public API reads a 500 as a record that may exist; a 503 would tell + it nothing was created. The log is the fixed event + ``economy_record_outcome_unknown``. + """ + fields = {"record": error.record, "error_type": error.error_type} + logger.error( + "economy_record_outcome_unknown", + extra={**fields, "policyengine_attributes": fields}, + ) + return _error_response( + status.HTTP_500_INTERNAL_SERVER_ERROR, + EconomyError( + code="internal_error", + detail=( + f"The service may have created the {error.record}, but could not " + "confirm it was written." + ), + ), + ) + + +def _unexpected_failure_response( + route: str, error: Exception, *, max_characters: int +) -> JSONResponse: + """Answer an unexpected failure with the JSON 500 the routes declare. + + The log is the fixed event ``economy_route_failed``; see + ``log_unexpected_failure``. + """ + log_unexpected_failure( + logger, + "economy_route_failed", + route=route, + error=error, + max_characters=max_characters, + ) + return _error_response( + status.HTTP_500_INTERNAL_SERVER_ERROR, + EconomyError( + code="internal_error", detail="The service could not answer the request." + ), + ) + + +def _report_state(stored: StoredReport) -> EconomyReportState: + run = stored.run + return EconomyReportState( + report_id=stored.id, + country=stored.country, + baseline_simulation_id=stored.baseline_simulation_id, + reform_simulation_id=stored.reform_simulation_id, + run=None + if run is None + else ReportRunRead( + run_id=run.id, + status=run.status, + country_package_version=run.country_package_version, + policyengine_version=run.policyengine_version, + error_message=run.error_message, + created_at=run.created_at, + started_at=run.started_at, + completed_at=run.completed_at, + baseline_simulation_run_id=run.baseline_simulation_run_id, + reform_simulation_run_id=run.reform_simulation_run_id, + ), + ) + + +def _simulation_state(stored: StoredSimulation) -> EconomySimulationState: + run = stored.run + return EconomySimulationState( + simulation_id=stored.id, + run=None + if run is None + else SimulationRunRead( + run_id=run.id, + status=run.status, + country_package_version=run.country_package_version, + policyengine_version=run.policyengine_version, + error_message=run.error_message, + created_at=run.created_at, + started_at=run.started_at, + completed_at=run.completed_at, + ), + ) + + +def _unknown(record: RecordName, record_id: UUID) -> JSONResponse: + return _error_response( + status.HTTP_404_NOT_FOUND, + EconomyError( + code="unknown_report" if record == "report" else "unknown_simulation", + detail=f"There is no {record} {record_id}.", + ), + ) + + +def add_economy_routes( + app: FastAPI, + entry: EntryRuntime, + submissions: EconomySubmissions | None, + *, + submitter: AuthorizationDependency, + reader: AuthorizationDependency, +) -> None: + """Add the economy routes; without a service they answer 503.""" + log_string_limit = entry.runtime.config.limits.max_string_length + app.include_router( + _submission_routes( + entry, submissions, submitter=submitter, log_string_limit=log_string_limit + ) + ) + app.include_router( + _read_routes(submissions, reader=reader, log_string_limit=log_string_limit) + ) + + +def _submission_routes( + entry: EntryRuntime, + submissions: EconomySubmissions | None, + *, + submitter: AuthorizationDependency, + log_string_limit: int, +) -> APIRouter: + """Return the router of ``POST /reports`` and ``POST /simulations``.""" + router = APIRouter() + + @router.post( + REPORTS_PATH, + summary="Create an economy report", + description=( + "Create an economy report comparing a baseline and a reform " + "simulation of one region's data for one year, on the v2 version " + "manifest's default worker, and start its first run. A side whose " + "simulation already succeeded on the same versions reuses that run." + ), + operation_id="create_report", + status_code=status.HTTP_201_CREATED, + response_model=EconomyReportCreated, + responses={ + 201: {"description": "The report and its first run were created"}, + 400: _REFUSED, + 403: _FORBIDDEN, + 422: _UNPROCESSABLE, + 500: _INTERNAL_ERROR, + 503: _unavailable_description("report"), + }, + dependencies=[Depends(submitter)], + ) + async def create_report( + body: EconomyReportRequest, request: Request + ) -> JSONResponse: + if submissions is None: + return _not_configured() + observability_id = entry.start_workflow_observability(request) + try: + created = await submissions.submit_report( + body, + request_id=call_request_id_for(request), + observability_context=observability_context(entry, observability_id), + ) + except _SUBMISSION_FAILURES as error: + return _submission_failure_response(error, "report") + # Every other failure is still answered as the JSON 500 declared. + except Exception as error: # noqa: BLE001 # answered as internal_error + return _unexpected_failure_response( + REPORTS_PATH, error, max_characters=log_string_limit + ) + return JSONResponse( + created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED + ) + + @router.post( + SIMULATIONS_PATH, + summary="Create an economy simulation", + description=( + "Create a standalone economy simulation of one region's data for " + "one year under one stored policy and dynamic, on the v2 version " + "manifest's default worker, and start its first run." + ), + operation_id="create_simulation", + status_code=status.HTTP_201_CREATED, + response_model=EconomySimulationCreated, + responses={ + 201: {"description": "The simulation and its first run were created"}, + 400: _REFUSED, + 403: _FORBIDDEN, + 422: _UNPROCESSABLE, + 500: _INTERNAL_ERROR, + 503: _unavailable_description("simulation"), + }, + dependencies=[Depends(submitter)], + ) + async def create_simulation( + body: EconomySimulationRequest, request: Request + ) -> JSONResponse: + if submissions is None: + return _not_configured() + observability_id = entry.start_workflow_observability(request) + try: + created = await submissions.submit_simulation( + body, + request_id=call_request_id_for(request), + observability_context=observability_context(entry, observability_id), + ) + except _SUBMISSION_FAILURES as error: + return _submission_failure_response(error, "simulation") + # Every other failure is still answered as the JSON 500 declared. + except Exception as error: # noqa: BLE001 # answered as internal_error + return _unexpected_failure_response( + SIMULATIONS_PATH, error, max_characters=log_string_limit + ) + return JSONResponse( + created.model_dump(mode="json"), status_code=status.HTTP_201_CREATED + ) + + return router + + +async def _read_answer[Stored]( # noqa: PLR0913 # the read, its record and how to answer it + submissions: EconomySubmissions | None, + read: Callable[[EconomySubmissions], Awaitable[Stored | None]], + answer: Callable[[Stored], BaseModel], + *, + record: RecordName, + record_id: UUID, + log_string_limit: int, +) -> JSONResponse: + """Answer a read: the record as it stands, 404, 503 or the JSON 500.""" + if submissions is None: + return _not_configured() + route = REPORT_PATH if record == "report" else SIMULATION_PATH + try: + stored = await read(submissions) + state = None if stored is None else answer(stored) + except EconomySubmissionUnavailableError as error: + return _submission_failure_response(error, record) + # Every other failure is still answered as the JSON 500 declared. + except Exception as error: # noqa: BLE001 # answered as internal_error + return _unexpected_failure_response( + route, error, max_characters=log_string_limit + ) + if state is None: + return _unknown(record, record_id) + return JSONResponse(state.model_dump(mode="json")) + + +def _read_routes( + submissions: EconomySubmissions | None, + *, + reader: AuthorizationDependency, + log_string_limit: int, +) -> APIRouter: + """Return the router of the report and simulation reads.""" + router = APIRouter() + + @router.get( + REPORT_PATH, + summary="Get a report's newest run", + description=( + "Return the report, its two simulations and its newest run's " + "status, versions, error message and times." + ), + operation_id="get_report", + response_model=EconomyReportState, + responses={ + 200: {"description": "The report as it stands"}, + 403: _FORBIDDEN, + 404: {"description": "There is no such report (unknown_report)", **_ERROR}, + 422: _READ_UNPROCESSABLE, + 500: _READ_INTERNAL_ERROR, + 503: _READ_UNAVAILABLE, + }, + dependencies=[Depends(reader)], + ) + async def get_report(report_id: UUID) -> JSONResponse: + return await _read_answer( + submissions, + lambda service: service.read_report(report_id), + _report_state, + record="report", + record_id=report_id, + log_string_limit=log_string_limit, + ) + + @router.get( + SIMULATION_PATH, + summary="Get a simulation's newest run", + description=( + "Return the simulation and its newest run's status, versions, error " + "message and times." + ), + operation_id="get_simulation", + response_model=EconomySimulationState, + responses={ + 200: {"description": "The simulation as it stands"}, + 403: _FORBIDDEN, + 404: { + "description": "There is no such simulation (unknown_simulation)", + **_ERROR, + }, + 422: _READ_UNPROCESSABLE, + 500: _READ_INTERNAL_ERROR, + 503: _READ_UNAVAILABLE, + }, + dependencies=[Depends(reader)], + ) + async def get_simulation(simulation_id: UUID) -> JSONResponse: + return await _read_answer( + submissions, + lambda service: service.read_simulation(simulation_id), + _simulation_state, + record="simulation", + record_id=simulation_id, + log_string_limit=log_string_limit, + ) + + return router diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py new file mode 100644 index 000000000..be55aed4b --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submission_errors.py @@ -0,0 +1,243 @@ +"""Why an economy submission is refused, cannot be calculated, or cannot be served. + +The submission service raises these, and the routes answer each with its +status: EconomySubmissionRefusedError 400, EconomySubmissionUnprocessableError +422, EconomySubmissionUnavailableError 503 and +EconomySubmissionOutcomeUnknownError 500. ``read_refusal`` and +``reference_refusal`` turn the store's errors into them, writing each detail +from the request rather than from an error's text, since the Public API +passes a refusal's detail on to its callers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal + +from policyengine_simulation_persistence import ( + MissingPolicyParametersError, + OverlappingReformError, + ReportSideError, + UnknownCatalogVersionError, + UnknownDatasetError, + UnknownDynamicError, + UnknownPolicyError, + UnknownRegionError, +) + +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_simulation_persistence import ( + EconomyRecordReferenceError, + EconomyRequest, + ReportSide, + ) + + from policyengine_simulation_entry.economy_api_models import ( + EconomyRefusalCode, + EconomyUnavailableCode, + EconomyUnprocessableCode, + ) + +# Each column whose foreign key can refuse a submission's rows: the code the +# refusal answers with and the row it names. A region, dataset or dynamic is +# the caller's own choice; any other row is one the Public API checked first, +# or one the service chose, so its refusal is the service's state. +REFERENCE_REFUSALS: dict[str, tuple[EconomyRefusalCode, str]] = { + "region_id": ("unknown_region", "region"), + "dataset_id": ("unknown_dataset", "dataset"), + "dynamic_id": ("unknown_dynamic", "dynamic"), + "policy_id": ("unknown_policy", "policy"), + "tax_benefit_model_version_id": ("unknown_catalog_version", "catalog version"), + "tax_benefit_model_id": ("unknown_catalog_version", "tax-benefit model"), +} +# What a submission creates. +RecordName = Literal["report", "simulation"] + + +class EconomySubmissionRefusedError(Exception): + """A row the submission names, or a function it needs, is missing (400).""" + + def __init__(self, *, code: EconomyRefusalCode, detail: str) -> None: + """Keep the reason's code and its description.""" + super().__init__(detail) + self.code: EconomyRefusalCode = code + self.detail = detail + + +@dataclass(frozen=True) +class UnprocessableReason: + """One reason stored values cannot be calculated, shaped as a FastAPI field error. + + ``loc`` starts with ``body`` and the request field the reason traces to. + """ + + loc: tuple[str | int, ...] + msg: str + type: str + + +class EconomySubmissionUnprocessableError(Exception): + """The stored policies or dynamics the submission names cannot be calculated.""" + + def __init__( + self, + *, + code: EconomyUnprocessableCode, + reasons: tuple[UnprocessableReason, ...], + ) -> None: + """Keep the code and every reason, in order.""" + super().__init__("; ".join(reason.msg for reason in reasons)) + self.code: EconomyUnprocessableCode = code + self.reasons = reasons + + +class EconomySubmissionOutcomeUnknownError(Exception): + """The submission's rows may have been written, but that can't be confirmed. + + The commit failed after the rows were inserted. The route answers 500, + which the Public API reads as a record that may exist, never a 503, + which it reads as nothing created. + """ + + def __init__(self, record: RecordName, *, error_type: str) -> None: + """Keep what the submission creates, and the commit error's type.""" + super().__init__(f"the {record} may have been created") + self.record: RecordName = record + self.error_type = error_type + + +class EconomySubmissionUnavailableError(Exception): + """A dependency could not be read or refused the call; ask again later. + + ``record_id`` names the report or simulation that was created and whose + run was failed, when one was. + """ + + def __init__( + self, + *, + code: EconomyUnavailableCode, + detail: str, + record_id: UUID | None = None, + ) -> None: + """Keep the reason's code, its description and any created record.""" + super().__init__(detail) + self.code: EconomyUnavailableCode = code + self.detail = detail + self.record_id = record_id + + +def read_refusal( + error: Exception, request: EconomyRequest +) -> EconomySubmissionRefusedError | EconomySubmissionUnprocessableError | None: + """Return the refusal a store read's error means, or None for any other error. + + A missing catalog version, region, dataset, policy or dynamic is refused + (400); a policy or dynamic setting parameters the catalog version lacks, + or both setting one, cannot be calculated (422). A report side's refusal + (ReportSideError) is located on that side's fields. Every detail is + written from ``request``, never from the error's text. + """ + side: ReportSide | None = None + if isinstance(error, ReportSideError): + side, error = error.side, error.error + return _missing_row_refusal(error, request, side) or _unprocessable( + error, request, side + ) + + +def _missing_row_refusal( + error: Exception, request: EconomyRequest, side: ReportSide | None +) -> EconomySubmissionRefusedError | None: + """Return the 400 a missing row means, or None for another error.""" + catalog = ( + f"{request.country}'s catalog version for .py {request.policyengine_version}" + ) + named_by = "the request" if side is None else f"the request's {side}" + refusal: tuple[EconomyRefusalCode, str] | None = None + if isinstance(error, UnknownCatalogVersionError): + refusal = ( + "unknown_catalog_version", + f"{request.country} has no catalog version for .py " + f"{request.policyengine_version}, the default v2 worker's version.", + ) + elif isinstance(error, UnknownRegionError): + refusal = ("unknown_region", f"{catalog} has no region {request.region_code}.") + elif isinstance(error, UnknownDatasetError): + dataset = ( + f"default dataset for region {request.region_code}" + if request.dataset_name is None + else f"dataset {request.dataset_name}" + ) + refusal = ("unknown_dataset", f"{catalog} has no {dataset}.") + elif isinstance(error, UnknownPolicyError): + refusal = ( + "unknown_policy", + f"{request.country} has no stored policy {named_by} names.", + ) + elif isinstance(error, UnknownDynamicError): + refusal = ( + "unknown_dynamic", + f"No stored dynamic matches the one {named_by} names.", + ) + if refusal is None: + return None + code, detail = refusal + return EconomySubmissionRefusedError(code=code, detail=detail) + + +def _unprocessable( + error: Exception, request: EconomyRequest, side: ReportSide | None +) -> EconomySubmissionUnprocessableError | None: + """Return the 422 refused values mean, or None for another error. + + Each reason is located at the field that names the values: the request's + own, or the report side's. + """ + location: tuple[str, ...] = ("body",) if side is None else ("body", side) + if isinstance(error, MissingPolicyParametersError): + field_name = "policy_id" if error.source == "policy" else "dynamic_id" + return EconomySubmissionUnprocessableError( + code="missing_parameters", + reasons=tuple( + UnprocessableReason( + loc=(*location, field_name, name), + msg=( + f"the catalog version for .py {request.policyengine_version} " + "has no such parameter" + ), + type="missing_parameter", + ) + for name in error.parameter_names + ), + ) + if isinstance(error, OverlappingReformError): + return EconomySubmissionUnprocessableError( + code="overlapping_parameters", + reasons=tuple( + UnprocessableReason( + loc=(*location, "dynamic_id", name), + msg="the policy and the dynamic both set this parameter", + type="overlapping_parameter", + ) + for name in error.parameter_names + ), + ) + return None + + +def reference_refusal( + error: EconomyRecordReferenceError, +) -> EconomySubmissionRefusedError: + """Return the refusal of rows the database refused for naming a missing row. + + It names the kind of row, never its ID, which the caller sent. + """ + code, row = REFERENCE_REFUSALS.get( + error.column_name or "", ("missing_reference", "row") + ) + return EconomySubmissionRefusedError( + code=code, detail=f"a {row} the request names no longer exists" + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py new file mode 100644 index 000000000..3247e5ecb --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_submissions.py @@ -0,0 +1,691 @@ +"""Create economy reports and simulations, and start the calls that calculate them. + +A report submission names a region, a dataset, a year and each side's +stored policy and dynamic. The service: + +1. finds the v2 version manifest's default worker's economy functions for + the country, and its bundle; +2. reads, in one transaction, the country's catalog version for that + worker's .py version, the region and dataset, each side's policy and + dynamic values, and, when the output bucket is configured, each side's + reusable run: a run of the same simulation that succeeded on the same + versions, used by a succeeded report run with the same options, and + never refused for reuse; +3. creates, in one transaction, each side's new simulation and pending run + (none for a reused side), the report with its options, and its first + report run, pending with no Modal call ID; +4. starts the coordinator's call with ``spawn``, passing the report and its + run IDs, then records the call's ID while the report run is still + pending with none (the coordinator may have claimed it already). When + Modal refuses the call, or does not confirm it in time, it fails the + report run and its new simulation runs at once instead, under the same + condition; a run the coordinator claimed meanwhile, because Modal + accepted the call after all, is left as it is and answered ``pending``. + +A standalone simulation is created and started the same way, with one +simulation run calculated by the country's simulation function. + +Whatever is missing or cannot be calculated is refused in steps 1 and 2, +before anything is written. The coordinator and the simulation functions +then claim their runs and record their results or failures; the stuck-run +sweep fails a pending run no call was recorded for. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID, uuid4 + +from policyengine_simulation_contract.stage12_execution import ReportRunIds +from policyengine_simulation_persistence import ( + EconomyRecordOutcomeUnknownError, + EconomyRecordReferenceError, + EconomyRequest, + NewEconomyReport, + NewEconomySimulation, + ReportRunStatus, + ReportSimulation, + RunReuse, + SimulationChoice, + SimulationStatus, +) +from sqlalchemy.exc import SQLAlchemyError + +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportCreated, + EconomySimulationCreated, +) +from policyengine_simulation_entry.economy_calculation_inputs import ( + CalculationScope, + CallOrigin, + DatasetNotInBundleError, + bundle_dataset, + report_run_calculation, + simulation_options, + simulation_run_calculation, +) +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + RecordName, + read_refusal, + reference_refusal, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, +) +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from policyengine_simulation_persistence import ( + EconomyCatalogRows, + EconomyReportInputs, + EconomyReportState, + EconomySimulationInputs, + EconomySimulationState, + SimulationInputs, + ) + + from policyengine_simulation_entry.economy_api_models import ( + EconomyRefusalCode, + EconomyReportRequest, + EconomySimulationRequest, + PolicyChoice, + ) + from policyengine_simulation_entry.economy_calculation_inputs import ( + RunCalculation, + ) + from policyengine_simulation_entry.economy_workers import ( + EconomyRunCaller, + EconomyWorker, + EconomyWorkerSource, + ) + +logger = logging.getLogger(__name__) + +# How long a submission waits for Modal to accept the run's call before +# failing the run. +CALL_START_TIMEOUT_SECONDS = 10.0 + + +class EconomyRuns(Protocol): + """The economy reads and writes (``EconomyRunStore``).""" + + def read_report_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + """Read a report's catalog rows, both sides' values and reusable runs.""" + ... + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + """Read a standalone simulation's catalog rows and values.""" + ... + + def create_report(self, report: NewEconomyReport) -> None: + """Create the report, its new simulations and its runs.""" + ... + + def create_simulation(self, simulation: NewEconomySimulation) -> None: + """Create the simulation and its first run.""" + ... + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + """Record the report run's call ID while it is pending with none.""" + ... + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + """Fail the report run, and its new simulation runs, while unstarted.""" + ... + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + """Record the simulation run's call ID while it is pending with none.""" + ... + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + """Fail the simulation run while it is pending with no call ID.""" + ... + + def read_report(self, report_id: UUID) -> EconomyReportState | None: + """Return the report with its newest run.""" + ... + + def read_simulation(self, simulation_id: UUID) -> EconomySimulationState | None: + """Return the simulation with its newest run.""" + ... + + +@dataclass(frozen=True) +class ServiceSettings: + """The deployment the service runs in. + + ``environment`` labels the calls' artifacts, ``modal_environment`` holds + the v2 workers, and ``output_bucket`` is the bucket that keeps + simulation output files for good. Without it, no report reuses a run. + """ + + environment: str + modal_environment: str + output_bucket: str | None = None + + @property + def reuse_prefix(self) -> str | None: + """Return ``gs:///``, under which reusable output files lie.""" + return None if not self.output_bucket else f"gs://{self.output_bucket}/" + + +def _database_unavailable() -> EconomySubmissionUnavailableError: + return EconomySubmissionUnavailableError( + code="database_unavailable", detail="The simulation database could not be read." + ) + + +def _choice(policy: PolicyChoice) -> SimulationChoice: + return SimulationChoice(policy_id=policy.policy_id, dynamic_id=policy.dynamic_id) + + +class EconomySubmissionService: + """Create economy reports and simulations, and start their calls. + + The database and the manifest are read on threads, so the event loop + never waits on them. + """ + + def __init__( # noqa: PLR0913 # each seam tests replace + self, + *, + store: EconomyRuns, + workers: EconomyWorkerSource, + caller: EconomyRunCaller, + settings: ServiceSettings, + call_start_timeout_seconds: float = CALL_START_TIMEOUT_SECONDS, + new_id: Callable[[], UUID] = uuid4, + clock: Callable[[], datetime] = lambda: datetime.now(UTC), + ) -> None: + """Use these seams; tests replace each of them.""" + self._store = store + self._workers = workers + self._caller = caller + self._settings = settings + self._call_start_timeout_seconds = call_start_timeout_seconds + self._new_id = new_id + self._clock = clock + + async def submit_report( + self, + request: EconomyReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + """Create the report and its runs and start its coordinator. + + Raise EconomySubmissionRefusedError (400) or + EconomySubmissionUnprocessableError (422) before writing anything; + EconomySubmissionUnavailableError (503) when the manifest or + database cannot be read, or, naming the report, when Modal doesn't + confirm the call and the report run was failed; and + EconomySubmissionOutcomeUnknownError (500) when the create's commit + failed after its rows were inserted. + """ + worker = await self._default_worker(request.country) + options = request.options() + reuse_prefix = self._settings.reuse_prefix + reuse = ( + None + if reuse_prefix is None + else RunReuse( + versions=worker.versions, + inputs=options, + output_uri_prefix=reuse_prefix, + ) + ) + economy_request = _economy_request(request, worker) + inputs = await self._read_inputs( + lambda: self._store.read_report_inputs( + economy_request, + baseline=_choice(request.baseline), + reform=_choice(request.reform), + reuse=reuse, + ), + request=economy_request, + ) + scope = _scope(request, worker, inputs.catalog, spm=options.get("spm")) + report = NewEconomyReport( + id=self._new_id(), + run_id=self._new_id(), + country=request.country, + label=request.label, + description=request.description, + tax_benefit_model_id=inputs.catalog.model_id, + population=inputs.catalog.population(request.year), + versions=worker.versions, + inputs=options, + baseline=self._side(inputs.baseline), + reform=self._side(inputs.reform), + ) + calculation = report_run_calculation( + worker, + scope, + self._origin(request_id), + run_ids=ReportRunIds( + report_run_id=report.run_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_run_id=report.reform.simulation_run_id, + ), + baseline_policy=inputs.baseline.reform, + reform_policy=inputs.reform.reform, + ) + await self._create("report", lambda: self._store.create_report(report)) + status = await self._start_report_run( + worker, report, calculation, observability_context=observability_context + ) + return EconomyReportCreated( + report_id=report.id, + report_run_id=report.run_id, + status=status, + baseline_simulation_id=report.baseline.simulation_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + baseline_reused=report.baseline.reused, + reform_simulation_id=report.reform.simulation_id, + reform_simulation_run_id=report.reform.simulation_run_id, + reform_reused=report.reform.reused, + ) + + async def submit_simulation( + self, + request: EconomySimulationRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + """Create the simulation and its first run, and start the run's function. + + Raise the errors ``submit_report`` names, naming the simulation when + its run was failed. + """ + worker = await self._default_worker(request.country) + economy_request = _economy_request(request, worker) + choice = SimulationChoice( + policy_id=request.policy_id, dynamic_id=request.dynamic_id + ) + inputs = await self._read_inputs( + lambda: self._store.read_simulation_inputs(economy_request, choice), + request=economy_request, + ) + scope = _scope(request, worker, inputs.catalog, spm=None) + simulation = NewEconomySimulation( + id=self._new_id(), + run_id=self._new_id(), + choice=inputs.simulation.choice, + population=inputs.catalog.population(request.year), + versions=worker.versions, + ) + calculation = simulation_run_calculation( + worker, + scope, + self._origin(request_id), + simulation_run_id=simulation.run_id, + policy=inputs.simulation.reform, + ) + await self._create( + "simulation", lambda: self._store.create_simulation(simulation) + ) + status = await self._start_simulation_run( + worker, + simulation, + calculation, + observability_context=observability_context, + ) + return EconomySimulationCreated( + simulation_id=simulation.id, + simulation_run_id=simulation.run_id, + status=status, + ) + + async def read_report(self, report_id: UUID) -> EconomyReportState | None: + """Return the report with its newest run, or None when there is none. + + Raise EconomySubmissionUnavailableError when it cannot be read. + """ + return await self._read_state(lambda: self._store.read_report(report_id)) + + async def read_simulation( + self, simulation_id: UUID + ) -> EconomySimulationState | None: + """Return the simulation with its newest run, or None when there is none. + + Raise EconomySubmissionUnavailableError when it cannot be read. + """ + return await self._read_state( + lambda: self._store.read_simulation(simulation_id) + ) + + async def _read_state[State](self, read: Callable[[], State]) -> State: + try: + return await asyncio.to_thread(read) + except SQLAlchemyError as error: + logger.warning( + "economy_state_unreadable", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + + def _side(self, inputs: SimulationInputs) -> ReportSimulation: + """Return a side naming its reusable run, or a new simulation and run.""" + if inputs.reusable_run is not None: + return ReportSimulation( + simulation_id=inputs.reusable_run.simulation_id, + simulation_run_id=inputs.reusable_run.simulation_run_id, + choice=inputs.choice, + reused=True, + ) + return ReportSimulation( + simulation_id=self._new_id(), + simulation_run_id=self._new_id(), + choice=inputs.choice, + reused=False, + ) + + def _origin(self, request_id: str) -> CallOrigin: + return CallOrigin( + environment=self._settings.environment, + modal_environment=self._settings.modal_environment, + request_id=request_id, + created_at=self._clock(), + ) + + async def _default_worker(self, country: str) -> EconomyWorker: + try: + return await asyncio.to_thread(self._workers.default_worker, country) + except EconomyWorkerUnavailableError: + raise EconomySubmissionRefusedError( + code="no_economy_worker", + detail=( + "The default v2 worker cannot calculate economy runs for " + f"{country}." + ), + ) from None + except VersionManifestUnavailableError as error: + logger.warning( + "economy_manifest_unavailable", + extra={"error_type": type(error.__cause__).__name__}, + ) + raise EconomySubmissionUnavailableError( + code="manifest_unavailable", + detail="The v2 version manifest could not be read.", + ) from None + + async def _read_inputs[Inputs]( + self, read: Callable[[], Inputs], *, request: EconomyRequest + ) -> Inputs: + """Run ``read`` on a thread; answer each refusal as the request's error.""" + try: + return await asyncio.to_thread(read) + except SQLAlchemyError as error: + logger.warning( + "economy_inputs_unreadable", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + except Exception as error: + refusal = read_refusal(error, request) + if refusal is None: + raise + raise refusal from None + + async def _create(self, record: RecordName, create: Callable[[], None]) -> None: + """Create the rows; answer a refusal, an unknown outcome or an outage. + + A commit that failed after the rows were inserted may have written + them, so it is never answered as unavailable, which would tell the + caller nothing was created. + """ + try: + await asyncio.to_thread(create) + except EconomyRecordReferenceError as error: + raise reference_refusal(error) from None + except EconomyRecordOutcomeUnknownError as error: + # The route logs it, with the commit error's type. + raise EconomySubmissionOutcomeUnknownError( + record, error_type=type(error.__cause__).__name__ + ) from None + except SQLAlchemyError as error: + logger.warning( + "economy_records_not_created", + extra={"error_type": type(error).__name__}, + ) + raise _database_unavailable() from None + + async def _start_report_run( + self, + worker: EconomyWorker, + report: NewEconomyReport, + calculation: RunCalculation, + *, + observability_context: dict[str, Any] | None, + ) -> ReportRunStatus: + try: + call_id = await asyncio.wait_for( + self._caller.start_report_run( + worker, + calculation.payload, + calculation.context, + observability_context=observability_context, + ), + timeout=self._call_start_timeout_seconds, + ) + # Modal raises its own errors; any of them, or the timeout, means the + # call was not confirmed, and the run is failed unless it has started. + except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike + failed = await self._fail_unconfirmed( + "report", + report.id, + error, + lambda message: self._store.fail_unstarted_report( + report.run_id, report.new_simulation_run_ids, message + ), + ) + if failed: + refused = _call_refused("report", report.id) + raise refused from None + return ReportRunStatus.PENDING + status = await self._record_call( + "report", + report.id, + lambda: self._store.record_report_call(report.run_id, call_id), + ) + return ReportRunStatus.PENDING if status is None else ReportRunStatus(status) + + async def _start_simulation_run( + self, + worker: EconomyWorker, + simulation: NewEconomySimulation, + calculation: RunCalculation, + *, + observability_context: dict[str, Any] | None, + ) -> SimulationStatus: + try: + call_id = await asyncio.wait_for( + self._caller.start_simulation_run( + worker, + calculation.payload, + calculation.context, + observability_context=observability_context, + ), + timeout=self._call_start_timeout_seconds, + ) + # As for a report run: any error or the timeout leaves it unconfirmed. + except Exception as error: # noqa: BLE001 # every unconfirmed call is answered alike + failed = await self._fail_unconfirmed( + "simulation", + simulation.id, + error, + lambda message: self._store.fail_unstarted_simulation( + simulation.run_id, message + ), + ) + if failed: + refused = _call_refused("simulation", simulation.id) + raise refused from None + return SimulationStatus.PENDING + status = await self._record_call( + "simulation", + simulation.id, + lambda: self._store.record_simulation_call(simulation.run_id, call_id), + ) + return SimulationStatus.PENDING if status is None else SimulationStatus(status) + + async def _fail_unconfirmed( + self, + record: RecordName, + record_id: UUID, + error: BaseException, + fail: Callable[[str], bool], + ) -> bool: + """Fail the run of a call Modal didn't confirm, unless it has started. + + Return whether it was failed. Modal may have accepted the call after + the wait ended, and its function claimed the run first: the failure + write then changes nothing. When the write itself fails, the run is + left pending: the record exists, so it is answered as created, and + the stuck-run sweep fails the run if no call ever claims it. + """ + # The log names the error type only; its text could repeat request data. + logger.error( + "economy_call_unconfirmed", + extra={ + "record": record, + "record_id": str(record_id), + "error_type": type(error).__name__, + }, + ) + message = ( + f"modal_call_refused: Modal did not accept the {record}'s call " + f"({type(error).__name__})." + ) + try: + return await asyncio.to_thread(fail, message) + except SQLAlchemyError as write_error: + logger.warning( + "economy_run_not_failed", + extra={ + "record": record, + "record_id": str(record_id), + "error_type": type(write_error).__name__, + }, + ) + return False + + async def _record_call( + self, + record: RecordName, + record_id: UUID, + recording: Callable[[], ReportRunStatus | SimulationStatus | None], + ) -> str | None: + """Record the call's ID; when that fails, answer the run as created. + + The coordinator or function claims a pending run with no call ID + itself, so the calculation goes ahead either way. + """ + try: + status = await asyncio.to_thread(recording) + except SQLAlchemyError as error: + logger.warning( + "economy_call_not_recorded", + extra={ + "record": record, + "record_id": str(record_id), + "error_type": type(error).__name__, + }, + ) + return None + return None if status is None else str(status) + + +def _call_refused( + record: RecordName, record_id: UUID +) -> EconomySubmissionUnavailableError: + return EconomySubmissionUnavailableError( + code="modal_call_refused", + detail=( + f"Modal did not accept the {record}'s call; the {record} was created " + "and its run failed." + ), + record_id=record_id, + ) + + +def _economy_request( + request: EconomyReportRequest | EconomySimulationRequest, worker: EconomyWorker +) -> EconomyRequest: + return EconomyRequest( + country=request.country, + policyengine_version=worker.policyengine_version, + region_code=request.region_code, + dataset_name=request.dataset_name, + year=request.year, + ) + + +def _scope( + request: EconomyReportRequest | EconomySimulationRequest, + worker: EconomyWorker, + catalog: EconomyCatalogRows, + *, + spm: dict[str, Any] | None, +) -> CalculationScope: + """Return what the calculation runs on, with the bundle's dataset. + + Raise EconomySubmissionRefusedError when the bundle lacks the catalog + dataset: ``unknown_dataset`` when the request named it, and + ``dataset_not_in_bundle`` for the region's default. + """ + try: + dataset = bundle_dataset(worker, catalog.dataset.name) + except DatasetNotInBundleError: + code: EconomyRefusalCode = ( + "dataset_not_in_bundle" + if request.dataset_name is None + else "unknown_dataset" + ) + detail = ( + f"The default v2 worker, .py {worker.version}, has no dataset " + f"{catalog.dataset.name}." + ) + raise EconomySubmissionRefusedError(code=code, detail=detail) from None + return CalculationScope( + country=request.country, + region_code=catalog.region.code, + dataset=dataset, + year=request.year, + options=simulation_options( + include_cliffs=getattr(request, "include_cliffs", False), spm=spm + ), + ) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py new file mode 100644 index 000000000..b07a6e693 --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/economy_workers.py @@ -0,0 +1,245 @@ +"""The default v2 worker's economy functions, and starting their calls. + +The economy routes use only the v2 version manifest's default worker: the +Modal application deployed for one .py version, which names its report +coordinator, each country's simulation function, and the bundle they run. +A report run is calculated by the coordinator, which starts a simulation +function for each simulation run it doesn't reuse; a standalone simulation +run is calculated by the simulation function alone. Each is called by name +with ``spawn``, so a route waits only for Modal to accept the call, and is +given the run IDs it claims and records, never a Stage 12 evaluation +record. + +This is the economy path's own Modal access, apart from the Stage 12 +comparison path's. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +import modal +from policyengine_simulation_contract.stage12_manifest import ( + V2ManifestLoader, + manifest_digest, +) +from policyengine_simulation_persistence import RunVersions + +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) + +if TYPE_CHECKING: + from policyengine_simulation_contract.stage12_bundle import Stage12CountryBundle + from policyengine_simulation_contract.stage12_manifest import ( + V2ManifestStore, + V2VersionManifest, + ) + + +class EconomyWorkerUnavailableError(LookupError): + """The default v2 worker cannot calculate the country's economy runs.""" + + +@dataclass(frozen=True) +class EconomyWorker: + """A country's economy functions in the default worker, and its bundle. + + ``version`` is the worker's .py version, the manifest's default, and + ``version_manifest_sha256`` the digest of the manifest it was read from. + """ + + version: str + version_manifest_sha256: str + application_name: str + coordinator_callable: str + simulation_callable: str + bundle_manifest_sha256: str + policyengine_version: str + country: Stage12CountryBundle + + @property + def versions(self) -> RunVersions: + """Return the versions every run this worker calculates records.""" + return RunVersions( + country_package_version=self.country.country_package_version, + policyengine_version=self.policyengine_version, + ) + + +def default_economy_worker(manifest: V2VersionManifest, country: str) -> EconomyWorker: + """Return the manifest's default worker's economy functions for ``country``. + + Raise EconomyWorkerUnavailableError when it cannot calculate the + country's economy runs. + """ + version = manifest.default_version + worker = manifest.versions[version] + country_worker = worker.economy_worker(country) + country_bundle = next( + (item for item in worker.bundle.countries if item.country == country), None + ) + if ( + country_worker is None + or country_worker.single_simulation_callable is None + or country_bundle is None + ): + msg = ( + f"the default v2 worker, .py {version}, cannot calculate economy " + f"runs for {country}" + ) + raise EconomyWorkerUnavailableError(msg) + return EconomyWorker( + version=version, + version_manifest_sha256=manifest_digest(manifest), + application_name=worker.application_name, + coordinator_callable=worker.report_coordinator_callable, + simulation_callable=country_worker.single_simulation_callable, + bundle_manifest_sha256=worker.bundle_manifest_sha256, + policyengine_version=worker.bundle.policyengine_version, + country=country_bundle, + ) + + +class EconomyWorkerSource(Protocol): + """Find the default worker's economy functions (``ManifestEconomyWorkers``).""" + + def default_worker(self, country: str) -> EconomyWorker: + """Return the country's economy functions, or raise why there are none.""" + ... + + +class ManifestEconomyWorkers: + """Find economy functions in the v2 version manifest's default worker. + + Each lookup reads the manifest, which is small, so a moved default is + used at once. The loader keeps its last valid copy for a stored value + that does not validate; a store that cannot be read is unavailable. + """ + + def __init__(self, store: V2ManifestStore) -> None: + """Read the manifest from ``store``; nothing is read until a lookup.""" + self._loader = V2ManifestLoader(store) + + @classmethod + def from_modal( + cls, *, manifest_name: str, environment: str + ) -> ManifestEconomyWorkers: + """Read the manifest from its Modal Dict, which must already exist.""" + return cls( + modal.Dict.from_name( + manifest_name, environment_name=environment, create_if_missing=False + ) + ) + + def default_worker(self, country: str) -> EconomyWorker: + """Return the default worker's economy functions for ``country``. + + Raise EconomyWorkerUnavailableError when it has none, and + VersionManifestUnavailableError when the manifest cannot be read. + """ + try: + manifest = self._loader.load() + # Modal raises its own errors, and any of them means the same here. + except Exception as error: + msg = "the v2 version manifest could not be read" + raise VersionManifestUnavailableError(msg) from error + return default_economy_worker(manifest, country) + + +class EconomyRunCaller(Protocol): + """Start the Modal call that calculates a run (``ModalEconomyRunCaller``).""" + + async def start_report_run( + self, + worker: EconomyWorker, + report_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the report run's coordinator and return its Modal call ID.""" + ... + + async def start_simulation_run( + self, + worker: EconomyWorker, + simulation_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the standalone simulation run's function; return its call ID.""" + ... + + +class ModalEconomyRunCaller: + """Start economy calls on Modal, by application and function name.""" + + def __init__(self, environment: str) -> None: + """Look up the functions in this Modal environment.""" + self._environment = environment + + async def start_report_run( + self, + worker: EconomyWorker, + report_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the coordinator with the report and its run IDs; return its call ID. + + Raise RuntimeError when Modal returns no call ID. + """ + return await self._spawn( + worker, + worker.coordinator_callable, + report_payload, + context_payload, + observability_context=observability_context, + ) + + async def start_simulation_run( + self, + worker: EconomyWorker, + simulation_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Start the simulation function with its run's ID; return its call ID. + + Raise RuntimeError when Modal returns no call ID. + """ + return await self._spawn( + worker, + worker.simulation_callable, + simulation_payload, + context_payload, + observability_context=observability_context, + ) + + async def _spawn( + self, + worker: EconomyWorker, + function_name: str, + payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + function = modal.Function.from_name( + worker.application_name, + function_name, + environment_name=self._environment, + ) + call = await function.spawn.aio( + payload, context_payload, observability_context=observability_context + ) + call_id = getattr(call, "object_id", None) + if not isinstance(call_id, str) or not call_id: + msg = f"Modal returned no call ID for {function_name}'s call" + raise RuntimeError(msg) + return call_id diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py index 7709bded3..660d4d13b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/household_routes.py @@ -25,9 +25,7 @@ from __future__ import annotations -import contextlib import logging -import traceback from typing import TYPE_CHECKING, Any, Protocol from uuid import UUID @@ -47,6 +45,10 @@ HouseholdCalculationUnavailableError, HouseholdCalculationUnprocessableError, ) +from policyengine_simulation_entry.route_support import ( + log_unexpected_failure, + observability_context, +) if TYPE_CHECKING: from fastapi import FastAPI @@ -207,96 +209,20 @@ def _submission_failure_response(error: Exception) -> JSONResponse: ) -# Separates the frames in the 500 log's ``stack_frames`` attribute. -STACK_FRAME_SEPARATOR = "; " -# The service's own packages, whose frames the 500 log keeps before others. -SERVICE_PACKAGES = frozenset( - {"policyengine_simulation_entry", "policyengine_simulation_persistence"} -) -# The directories a package's files sit under: an installed package's, then -# a checkout's. A frame's path is shortened to its part after one of them. -_PACKAGE_PARENT_DIRECTORIES = ("/site-packages/", "/src/") - - -def _package_path(file_path: str) -> str: - """Return ``file_path`` from its package's directory on. - - That is its part after the last ``site-packages/``, where the image - installs the service and its libraries, or else after the last ``src/``, - where a checkout keeps the service's packages. Any other path is - returned whole. - """ - for parent_directory in _PACKAGE_PARENT_DIRECTORIES: - _, found, package_path = file_path.rpartition(parent_directory) - if found: - return package_path - return file_path - - -def _stack_frames(error: Exception, max_characters: int) -> tuple[str, int]: - """Return the error's frames that fit ``max_characters``, and how many don't. - - Each frame is ``path:line in function``, with its path from its - package's directory on (``_package_path``). The kept frames are written - innermost first, from where the error was raised out to the route's - handler. When they don't all fit, frames are taken in this order until - the next one would pass ``max_characters``: the frame that raised the - error, then the service's own frames (``SERVICE_PACKAGES``) from the - innermost out, then the others from the innermost out. So other - libraries' outer frames are left out first, and left-out frames may sit - between two kept ones. Staying within ``max_characters`` keeps the - string whole: the observability runtime keeps only the start of a longer - string attribute. - """ - frames: list[str] = [] - from_the_service: list[bool] = [] - for frame in reversed(traceback.extract_tb(error.__traceback__)): - package_path = _package_path(frame.filename) - frames.append(f"{package_path}:{frame.lineno} in {frame.name}") - from_the_service.append(package_path.split("/", 1)[0] in SERVICE_PACKAGES) - positions_by_priority = sorted( - range(len(frames)), - key=lambda position: (position > 0, not from_the_service[position], position), - ) - kept_positions: list[int] = [] - length = 0 - for position in positions_by_priority: - added = len(frames[position]) + ( - len(STACK_FRAME_SEPARATOR) if kept_positions else 0 - ) - if length + added > max_characters: - break - kept_positions.append(position) - length += added - kept = [frames[position] for position in sorted(kept_positions)] - return STACK_FRAME_SEPARATOR.join(kept), len(frames) - len(kept) - - def _unexpected_failure_response( route: str, error: Exception, *, max_characters: int ) -> JSONResponse: """Answer an unexpected failure with the JSON 500 the routes declare. - The log is the fixed event ``household_route_failed`` with the route, the - error's type, its stack frames (path, line and function; see - ``_stack_frames``) as one string of at most ``max_characters``, the - observability runtime's string limit, and the number of frames left out. - It never carries the error's message or any value, which could repeat - request or stored data. The fields go in ``policyengine_attributes`` too, - the only extra fields the observability runtime's log handler keeps, and - each is a string or an integer: the runtime keeps only strings, numbers - and booleans. + The log is the fixed event ``household_route_failed``; see + ``log_unexpected_failure``. """ - stack_frames, omitted_frame_count = _stack_frames(error, max_characters) - fields = { - "route": route, - "error_type": type(error).__name__, - "stack_frames": stack_frames, - "stack_frames_omitted": omitted_frame_count, - } - logger.error( + log_unexpected_failure( + logger, "household_route_failed", - extra={**fields, "policyengine_attributes": fields}, + route=route, + error=error, + max_characters=max_characters, ) return _error_response( status.HTTP_500_INTERNAL_SERVER_ERROR, @@ -306,21 +232,6 @@ def _unexpected_failure_response( ) -def _observability_context( - entry: EntryRuntime, observability_id: str -) -> dict[str, Any]: - """Capture the trace context for the household function, with the ID. - - Telemetry is non-fatal: without a captured context, only the ID is sent. - """ - captured: object = None - with contextlib.suppress(Exception): - captured = entry.runtime.capture_context() - context = dict(captured) if isinstance(captured, dict) else {} - context["observability_id"] = observability_id - return context - - def _state(record: HouseholdJobRecord) -> HouseholdCalculationState: return HouseholdCalculationState( job_id=record.id, @@ -395,7 +306,7 @@ async def submit_household_calculation( try: accepted = await calculations.submit( body, - observability_context=_observability_context(entry, observability_id), + observability_context=observability_context(entry, observability_id), ) except _SUBMISSION_FAILURES as error: return _submission_failure_response(error) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.py new file mode 100644 index 000000000..457e467cd --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/route_support.py @@ -0,0 +1,124 @@ +"""What the v2 calculation routes share: their failure logs and trace context. + +The household and economy routes answer an unexpected failure with a JSON +500 and log it as a fixed event with the error's stack frames +(``log_unexpected_failure``), and pass each submission's trace context and +observability ID to the Modal call they start (``observability_context``). +""" + +from __future__ import annotations + +import contextlib +import traceback +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import logging + + from policyengine_simulation_entry.entry_runtime import EntryRuntime + +# Separates the frames in the 500 log's ``stack_frames`` attribute. +STACK_FRAME_SEPARATOR = "; " +# The service's own packages, whose frames the 500 log keeps before others. +SERVICE_PACKAGES = frozenset( + {"policyengine_simulation_entry", "policyengine_simulation_persistence"} +) +# The directories a package's files sit under: an installed package's, then +# a checkout's. A frame's path is shortened to its part after one of them. +_PACKAGE_PARENT_DIRECTORIES = ("/site-packages/", "/src/") + + +def package_path(file_path: str) -> str: + """Return ``file_path`` from its package's directory on. + + That is its part after the last ``site-packages/``, where the image + installs the service and its libraries, or else after the last ``src/``, + where a checkout keeps the service's packages. Any other path is + returned whole. + """ + for parent_directory in _PACKAGE_PARENT_DIRECTORIES: + _, found, path_after_parent = file_path.rpartition(parent_directory) + if found: + return path_after_parent + return file_path + + +def stack_frames(error: Exception, max_characters: int) -> tuple[str, int]: + """Return the error's frames that fit ``max_characters``, and how many don't. + + Each frame is ``path:line in function``, with its path from its + package's directory on (``package_path``). The kept frames are written + innermost first, from where the error was raised out to the route's + handler. When they don't all fit, frames are taken in this order until + the next one would pass ``max_characters``: the frame that raised the + error, then the service's own frames (``SERVICE_PACKAGES``) from the + innermost out, then the others from the innermost out. So other + libraries' outer frames are left out first, and left-out frames may sit + between two kept ones. Staying within ``max_characters`` keeps the + string whole: the observability runtime keeps only the start of a longer + string attribute. + """ + frames: list[str] = [] + from_the_service: list[bool] = [] + for frame in reversed(traceback.extract_tb(error.__traceback__)): + path_in_package = package_path(frame.filename) + frames.append(f"{path_in_package}:{frame.lineno} in {frame.name}") + from_the_service.append(path_in_package.split("/", 1)[0] in SERVICE_PACKAGES) + positions_by_priority = sorted( + range(len(frames)), + key=lambda position: (position > 0, not from_the_service[position], position), + ) + kept_positions: list[int] = [] + length = 0 + for position in positions_by_priority: + added = len(frames[position]) + ( + len(STACK_FRAME_SEPARATOR) if kept_positions else 0 + ) + if length + added > max_characters: + break + kept_positions.append(position) + length += added + kept = [frames[position] for position in sorted(kept_positions)] + return STACK_FRAME_SEPARATOR.join(kept), len(frames) - len(kept) + + +def log_unexpected_failure( + logger: logging.Logger, + event: str, + *, + route: str, + error: Exception, + max_characters: int, +) -> None: + """Log an unexpected failure of ``route`` as the fixed ``event``. + + The record carries the route, the error's type, its stack frames (path, + line and function; see ``stack_frames``) as one string of at most + ``max_characters``, the observability runtime's string limit, and the + number of frames left out. It never carries the error's message or any + value, which could repeat request or stored data. The fields go in + ``policyengine_attributes`` too, the only extra fields the observability + runtime's log handler keeps, and each is a string or an integer: the + runtime keeps only strings, numbers and booleans. + """ + frames, omitted_frame_count = stack_frames(error, max_characters) + fields = { + "route": route, + "error_type": type(error).__name__, + "stack_frames": frames, + "stack_frames_omitted": omitted_frame_count, + } + logger.error(event, extra={**fields, "policyengine_attributes": fields}) + + +def observability_context(entry: EntryRuntime, observability_id: str) -> dict[str, Any]: + """Capture the trace context for the Modal call a submission starts, with the ID. + + Telemetry is non-fatal: without a captured context, only the ID is sent. + """ + captured: object = None + with contextlib.suppress(Exception): + captured = entry.runtime.capture_context() + context = dict(captured) if isinstance(captured, dict) else {} + context["observability_id"] = observability_id + return context diff --git a/projects/policyengine-simulation-entry/tests/economy_fixtures.py b/projects/policyengine-simulation-entry/tests/economy_fixtures.py new file mode 100644 index 000000000..d24a831bd --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/economy_fixtures.py @@ -0,0 +1,49 @@ +"""Test values for the economy report and simulation routes. + +They build on the household path's v2 worker, which can calculate economy +runs for the US and the UK, with a second US dataset that state regions use. +""" + +from __future__ import annotations + +from household_fixtures import POLICYENGINE_VERSION, household_worker_version +from policyengine_simulation_contract.stage12_bundle import Stage12Dataset +from policyengine_simulation_contract.stage12_manifest import ( + V2VersionManifest, + V2WorkerVersion, +) + +STATE_DATASET = Stage12Dataset( + identity="populace_us_2024_acs_local", + uri="hf://policyengine/data/populace_us_2024_acs_local.h5@acs-local-revision", + artifact_revision="acs-local-revision", + sha256="c" * 64, + repo_type="dataset", +) + + +def economy_worker_version( + policyengine_version: str = POLICYENGINE_VERSION, +) -> V2WorkerVersion: + """Return a v2 worker whose US bundle also holds the state dataset.""" + base = household_worker_version(policyengine_version) + countries = tuple( + country.model_copy(update={"datasets": (*country.datasets, STATE_DATASET)}) + if country.country == "us" + else country + for country in base.bundle.countries + ) + return base.model_copy( + update={"bundle": base.bundle.model_copy(update={"countries": countries})} + ) + + +def economy_manifest( + default_worker: V2WorkerVersion | None = None, +) -> V2VersionManifest: + """Return a manifest whose only worker, its default, is ``default_worker``.""" + worker = default_worker or economy_worker_version() + version = worker.bundle.policyengine_version + return V2VersionManifest( + generation="test", default_version=version, versions={version: worker} + ) diff --git a/projects/policyengine-simulation-entry/tests/economy_service_fakes.py b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py new file mode 100644 index 000000000..535873f0f --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/economy_service_fakes.py @@ -0,0 +1,315 @@ +"""Fakes of the economy submission service's store, manifest and Modal calls. + +The store's own reads and writes are tested against PostgreSQL in the +persistence library; here a fake records what the service asks of it. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from datetime import UTC, datetime +from itertools import count +from typing import TYPE_CHECKING, Any +from uuid import UUID + +from economy_fixtures import economy_manifest +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportRequest, + EconomySimulationRequest, +) +from policyengine_simulation_entry.economy_submissions import ( + EconomySubmissionService, + ServiceSettings, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorker, + default_economy_worker, +) +from policyengine_simulation_persistence import ( + CatalogDataset, + CatalogRegion, + EconomyCatalogRows, + EconomyReportInputs, + EconomyReportState, + EconomySimulationInputs, + EconomySimulationState, + ReportRunStatus, + ReusableRun, + SimulationInputs, + SimulationStatus, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from policyengine_simulation_persistence import ( + EconomyRequest, + NewEconomyReport, + NewEconomySimulation, + RunReuse, + SimulationChoice, + ) + +CATALOG_VERSION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +MODEL_ID = UUID("00000000-0000-4000-8000-0000000000c2") +REGION_ID = UUID("00000000-0000-4000-8000-0000000000c3") +DATASET_ID = UUID("00000000-0000-4000-8000-0000000000c4") +POLICY_ID = UUID("00000000-0000-4000-8000-0000000000c5") +DYNAMIC_ID = UUID("00000000-0000-4000-8000-0000000000c6") +REUSED = ReusableRun( + simulation_id=UUID("00000000-0000-4000-8000-0000000000d1"), + simulation_run_id=UUID("00000000-0000-4000-8000-0000000000d2"), +) +CREATED_AT = datetime(2026, 11, 12, 15, tzinfo=UTC) +OUTPUT_BUCKET = "policyengine-v2-simulation-outputs-staging" +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +REFORM = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} +CATALOG = EconomyCatalogRows( + catalog_version_id=CATALOG_VERSION_ID, + model_id=MODEL_ID, + region=CatalogRegion(id=REGION_ID, code="us", default_dataset_id=DATASET_ID), + dataset=CatalogDataset(id=DATASET_ID, name="populace_us_2024"), +) +REPORT_BODY: dict[str, Any] = { + "country": "us", + "report_type": "economy", + "label": "A larger child tax credit", + "description": "Raises the credit to $3,000", + "year": 2026, + "region_code": "us", + "dataset_name": None, + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": str(POLICY_ID), "dynamic_id": str(DYNAMIC_ID)}, + "include_cliffs": True, + "spm": {"geography_kind": "national"}, +} +SIMULATION_BODY: dict[str, Any] = { + "country": "us", + "simulation_type": "economy", + "year": 2026, + "region_code": "us", + "dataset_name": None, + "policy_id": str(POLICY_ID), + "dynamic_id": None, +} + + +def report_request(**changes: Any) -> EconomyReportRequest: + return EconomyReportRequest.model_validate({**REPORT_BODY, **changes}) + + +def simulation_request(**changes: Any) -> EconomySimulationRequest: + return EconomySimulationRequest.model_validate({**SIMULATION_BODY, **changes}) + + +class FakeWorkers: + """The default worker's economy functions, or the lookup's failure.""" + + def __init__(self, failure: Exception | None = None) -> None: + self.failure = failure + + def default_worker(self, country: str) -> EconomyWorker: + if self.failure is not None: + raise self.failure + return default_economy_worker(economy_manifest(), country) + + +@dataclass +class FakeStore: + """Records what the service asks of the database, and answers as told. + + ``claimed`` marks the run as claimed by its coordinator or function: a + failure write then changes nothing. + """ + + catalog: EconomyCatalogRows = CATALOG + reusable: tuple[ReusableRun | None, ReusableRun | None] = (None, None) + read_failure: Exception | None = None + create_failure: Exception | None = None + record_failure: Exception | None = None + fail_failure: Exception | None = None + claimed: bool = False + stored_report: EconomyReportState | None = None + stored_simulation: EconomySimulationState | None = None + report_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) + simulation_reads: list[tuple[EconomyRequest, Any]] = field(default_factory=list) + reports: list[NewEconomyReport] = field(default_factory=list) + simulations: list[NewEconomySimulation] = field(default_factory=list) + recorded_calls: list[tuple[UUID, str]] = field(default_factory=list) + failed_runs: list[tuple[UUID, tuple[UUID, ...], str]] = field(default_factory=list) + + def read_report_inputs( + self, + request: EconomyRequest, + *, + baseline: SimulationChoice, + reform: SimulationChoice, + reuse: RunReuse | None, + ) -> EconomyReportInputs: + self.report_reads.append((request, (baseline, reform, reuse))) + if self.read_failure is not None: + raise self.read_failure + baseline_reusable, reform_reusable = self.reusable + return EconomyReportInputs( + catalog=self.catalog, + baseline=SimulationInputs( + choice=baseline, reform={}, reusable_run=baseline_reusable + ), + reform=SimulationInputs( + choice=reform, reform=REFORM, reusable_run=reform_reusable + ), + ) + + def read_simulation_inputs( + self, request: EconomyRequest, choice: SimulationChoice + ) -> EconomySimulationInputs: + self.simulation_reads.append((request, choice)) + if self.read_failure is not None: + raise self.read_failure + return EconomySimulationInputs( + catalog=self.catalog, + simulation=SimulationInputs(choice=choice, reform=REFORM), + ) + + def read_report(self, report_id: UUID) -> EconomyReportState | None: + if self.read_failure is not None: + raise self.read_failure + stored = self.stored_report + return stored if stored is not None and stored.id == report_id else None + + def read_simulation(self, simulation_id: UUID) -> EconomySimulationState | None: + if self.read_failure is not None: + raise self.read_failure + stored = self.stored_simulation + return stored if stored is not None and stored.id == simulation_id else None + + def create_report(self, report: NewEconomyReport) -> None: + if self.create_failure is not None: + raise self.create_failure + self.reports.append(report) + + def create_simulation(self, simulation: NewEconomySimulation) -> None: + if self.create_failure is not None: + raise self.create_failure + self.simulations.append(simulation) + + def record_report_call( + self, report_run_id: UUID, modal_call_id: str + ) -> ReportRunStatus | None: + if self.record_failure is not None: + raise self.record_failure + self.recorded_calls.append((report_run_id, modal_call_id)) + return ReportRunStatus.RUNNING if self.claimed else ReportRunStatus.PENDING + + def fail_unstarted_report( + self, + report_run_id: UUID, + simulation_run_ids: Sequence[UUID], + error_message: str, + ) -> bool: + if self.fail_failure is not None: + raise self.fail_failure + if self.claimed: + return False + self.failed_runs.append( + (report_run_id, tuple(simulation_run_ids), error_message) + ) + return True + + def record_simulation_call( + self, simulation_run_id: UUID, modal_call_id: str + ) -> SimulationStatus | None: + if self.record_failure is not None: + raise self.record_failure + self.recorded_calls.append((simulation_run_id, modal_call_id)) + return SimulationStatus.RUNNING if self.claimed else SimulationStatus.PENDING + + def fail_unstarted_simulation( + self, simulation_run_id: UUID, error_message: str + ) -> bool: + if self.fail_failure is not None: + raise self.fail_failure + if self.claimed: + return False + self.failed_runs.append((simulation_run_id, (), error_message)) + return True + + +@dataclass +class FakeCaller: + """Starts calls by recording them; ``failure`` is raised instead when set. + + ``hang`` makes a start wait until the service stops waiting for it. + """ + + failure: Exception | None = None + hang: bool = False + report_runs: list[tuple[EconomyWorker, dict, dict, Any]] = field( + default_factory=list + ) + simulation_runs: list[tuple[EconomyWorker, dict, dict, Any]] = field( + default_factory=list + ) + + async def start_report_run( + self, + worker: EconomyWorker, + report_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + await self._maybe_fail() + self.report_runs.append( + (worker, report_payload, context_payload, observability_context) + ) + return "fc-coordinator" + + async def start_simulation_run( + self, + worker: EconomyWorker, + simulation_payload: dict[str, Any], + context_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + await self._maybe_fail() + self.simulation_runs.append( + (worker, simulation_payload, context_payload, observability_context) + ) + return "fc-worker" + + async def _maybe_fail(self) -> None: + if self.hang: + await asyncio.sleep(60) + if self.failure is not None: + raise self.failure + + +def numbered_ids(): + """Return a function giving 00…01, 00…02 and so on, in order.""" + numbers = count(1) + return lambda: UUID(int=next(numbers)) + + +def service( + store: FakeStore, + *, + workers: FakeWorkers | None = None, + caller: FakeCaller | None = None, + output_bucket: str | None = OUTPUT_BUCKET, +) -> EconomySubmissionService: + return EconomySubmissionService( + store=store, + workers=workers or FakeWorkers(), + caller=caller or FakeCaller(), + settings=ServiceSettings( + environment="staging", + modal_environment="staging", + output_bucket=output_bucket, + ), + call_start_timeout_seconds=0.05, + new_id=numbered_ids(), + clock=lambda: CREATED_AT, + ) diff --git a/projects/policyengine-simulation-entry/tests/test_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index 5546a4668..d3b72af5a 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -410,3 +410,26 @@ def test_the_v2_manifest_settings_configure_no_stage12_resource(): settings.validate() assert not settings.stage12_resources_configured + + +def test_the_v2_output_bucket_is_read_from_the_environment(monkeypatch): + from policyengine_simulation_entry.config import Settings + + monkeypatch.setenv("V2_SIMULATION_OUTPUT_BUCKET", "policyengine-v2-outputs-staging") + + assert Settings.from_env().v2_simulation_output_bucket == ( + "policyengine-v2-outputs-staging" + ) + monkeypatch.delenv("V2_SIMULATION_OUTPUT_BUCKET") + assert Settings.from_env().v2_simulation_output_bucket == "" + + +def test_economy_runs_need_the_database_and_the_v2_modal_environment(): + assert not make_settings().economy_runs_configured + assert not make_settings( + simulation_database_url=DATABASE_URL + ).economy_runs_configured + assert not make_settings(v2_modal_environment="staging").economy_runs_configured + assert make_settings( + simulation_database_url=DATABASE_URL, v2_modal_environment="staging" + ).economy_runs_configured diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index b46faf65d..24e2d1a3e 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -215,10 +215,11 @@ def _assert_deployment_environment_settings(reusable_workflow: str) -> None: assert ( # Every job but stage12_v2_ready declares the release environment. # The job that moves latest after the checks reads the prod - # environment's custom-hostname variable, and the staging household + # environment's custom-hostname variable, the staging household # check reads the beta variable SIMULATION_STAGING_HOUSEHOLD_ID and - # the environment's auth secrets. - reusable_workflow.count("environment: ${{ inputs.release_environment }}") == 18 + # the environment's auth secrets, and the staging economy check + # reads those auth secrets too. + reusable_workflow.count("environment: ${{ inputs.release_environment }}") == 19 ) assert "APP_ENVIRONMENT: ${{ inputs.deployment_environment }}" in reusable_workflow assert "STAGE12_ENABLED_VALUE: ${{ vars.STAGE12_ENABLED }}" in reusable_workflow diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py index a0c8ada56..5d8dc1eb0 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets_support.py @@ -644,4 +644,5 @@ def cloud_run_deploy_environment(tmp_path: Path) -> dict[str, str]: "SIMULATION_SERVICE_DATABASE_URL_SECRET_NAME": ( "simulation-staging-service-database-url" ), + "V2_SIMULATION_OUTPUT_BUCKET_VALUE": "policyengine-v2-outputs-staging", } diff --git a/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py b/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py new file mode 100644 index 000000000..722ee115c --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_calculation_inputs.py @@ -0,0 +1,154 @@ +"""What the coordinator or a simulation function is given to calculate a run.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from uuid import UUID + +import pytest +from economy_fixtures import STATE_DATASET, economy_manifest +from household_fixtures import POLICYENGINE_VERSION +from policyengine_simulation_contract.stage12_execution import ( + ReportAggregate, + ReportExecutionInput, + ReportRunIds, + SimulationExecutionInput, + SimulationRole, + Stage12InvocationContext, +) +from policyengine_simulation_entry.economy_calculation_inputs import ( + CalculationScope, + CallOrigin, + DatasetNotInBundleError, + bundle_dataset, + report_run_calculation, + simulation_options, + simulation_run_calculation, +) +from policyengine_simulation_entry.economy_workers import default_economy_worker + +REPORT_RUN_ID = UUID("00000000-0000-4000-8000-000000000001") +BASELINE_RUN_ID = UUID("00000000-0000-4000-8000-000000000002") +REFORM_RUN_ID = UUID("00000000-0000-4000-8000-000000000003") +CREATED_AT = datetime(2026, 11, 12, 15, tzinfo=UTC) +CTC_AMOUNT = "gov.irs.credits.ctc.amount.base[0].amount" +REFORM = {CTC_AMOUNT: {"2026-01-01.2200-12-31": 3_000}} +ORIGIN = CallOrigin( + environment="staging", + modal_environment="staging", + request_id="observed-request", + created_at=CREATED_AT, +) +RUN_IDS = ReportRunIds( + report_run_id=REPORT_RUN_ID, + baseline_simulation_run_id=BASELINE_RUN_ID, + reform_simulation_run_id=REFORM_RUN_ID, +) + + +def _worker(): + return default_economy_worker(economy_manifest(), "us") + + +def _scope(**changes) -> CalculationScope: + values = { + "country": "us", + "region_code": "state/ca", + "dataset": STATE_DATASET, + "year": 2026, + "options": {"include_cliffs": True}, + } + return CalculationScope(**{**values, **changes}) + + +def test_the_bundle_dataset_is_the_one_the_catalog_dataset_names() -> None: + worker = _worker() + + assert bundle_dataset(worker, STATE_DATASET.identity) == STATE_DATASET + with pytest.raises(DatasetNotInBundleError, match="enhanced_cps_2024"): + bundle_dataset(worker, "enhanced_cps_2024") + + +def test_the_options_hold_cliffs_only_when_asked_and_the_spm_selection_given() -> None: + assert simulation_options(include_cliffs=False, spm=None) == {} + assert simulation_options( + include_cliffs=True, spm={"geography_kind": "national"} + ) == {"include_cliffs": True, "spm": {"geography_kind": "national"}} + + +def test_a_report_run_is_named_by_its_run_ids_with_no_evaluation_record() -> None: + worker = _worker() + + calculation = report_run_calculation( + worker, + _scope(), + ORIGIN, + run_ids=RUN_IDS, + baseline_policy={}, + reform_policy=REFORM, + ) + + report = ReportExecutionInput.model_validate(calculation.payload) + assert report.evaluation_id == REPORT_RUN_ID + assert report.requested_aggregates == tuple(ReportAggregate) + for side, role, run_id, policy in ( + (report.baseline, SimulationRole.BASELINE, BASELINE_RUN_ID, {}), + (report.reform, SimulationRole.REFORM, REFORM_RUN_ID, REFORM), + ): + assert (side.role, side.simulation_execution_id, side.policy) == ( + role, + run_id, + policy, + ) + assert side.evaluation_id == REPORT_RUN_ID + assert side.year == 2026 + assert (side.geography.country, side.geography.region) == ("us", "state/ca") + assert side.options == {"include_cliffs": True} + assert side.population.artifact.uri == STATE_DATASET.uri + assert side.population.artifact.content_sha256 == STATE_DATASET.sha256 + assert side.bundle.policyengine_version == POLICYENGINE_VERSION + assert side.bundle.country_package_version == "2.2.1" + assert side.bundle.dataset.identity == STATE_DATASET.identity + assert side.bundle.dataset.artifact_revision == "acs-local-revision" + assert side.bundle.bundle_manifest_sha256 == worker.bundle_manifest_sha256 + + context = Stage12InvocationContext.model_validate(calculation.context) + assert context.report_run == RUN_IDS + assert context.simulation_run_id is None + assert context.production_function_call_id is None + assert (context.environment, context.modal_environment, context.request_id) == ( + "staging", + "staging", + "observed-request", + ) + assert context.worker_version == worker.version + assert context.modal_application == worker.application_name + assert context.simulation_callable == "run_single_simulation_us" + assert context.version_manifest_sha256 == worker.version_manifest_sha256 + assert context.artifact_prefix == f"stage-12-runs/staging/2026/11/{REPORT_RUN_ID}" + assert context.retention_expires_at - context.created_at == timedelta(days=30) + + +def test_a_standalone_run_is_named_in_its_functions_context() -> None: + worker = _worker() + + calculation = simulation_run_calculation( + worker, + _scope(options={}), + ORIGIN, + simulation_run_id=BASELINE_RUN_ID, + policy=REFORM, + ) + + simulation = SimulationExecutionInput.model_validate(calculation.payload) + assert simulation.role is SimulationRole.STANDALONE + assert simulation.simulation_execution_id == BASELINE_RUN_ID + assert simulation.evaluation_id == BASELINE_RUN_ID + assert simulation.policy == REFORM + context = Stage12InvocationContext.model_validate(calculation.context) + assert context.simulation_run_id == BASELINE_RUN_ID + assert context.report_run is None + assert context.records_runs + assert calculation.context["artifact_prefix"] == ( + f"stage-12-runs/staging/2026/11/{BASELINE_RUN_ID}" + ) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_deployment.py b/projects/policyengine-simulation-entry/tests/test_economy_deployment.py new file mode 100644 index 000000000..b500d1be7 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_deployment.py @@ -0,0 +1,84 @@ +"""Deployment wiring for the economy report and simulation routes. + +Every deploy gives the service the bucket that keeps v2 simulation output +files for good (the environment variable ``V2_SIMULATION_OUTPUT_BUCKET``), +so a report reuses only a succeeded run whose file is in it. The service +compares output URIs with it and never reads the bucket. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets_support import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + cloud_run_deploy_environment, + workflow_jobs, +) + + +def _run_deploy(tmp_path, **overrides): + environment = cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + return result, environment + + +@pytest.mark.parametrize("deploy_stage12", ["true", "false"]) +def test_every_deploy_names_the_v2_output_bucket(tmp_path, deploy_stage12): + result, environment = _run_deploy( + tmp_path, + DEPLOY_STAGE12_V2=deploy_stage12, + V2_SIMULATION_OUTPUT_BUCKET_VALUE="policyengine-v2-outputs-staging", + ) + + assert result.returncode == 0, result.stderr + runtime_environment = json.loads( + Path(environment["GCLOUD_ENV_FILE_CAPTURE"]).read_text(encoding="utf-8") + ) + assert runtime_environment["V2_SIMULATION_OUTPUT_BUCKET"] == ( + "policyengine-v2-outputs-staging" + ) + + +@pytest.mark.parametrize( + ("bucket", "message"), + [ + ("", "V2_SIMULATION_OUTPUT_BUCKET_VALUE must be set"), + ( + "gs://policyengine-v2-outputs", + "V2_SIMULATION_OUTPUT_BUCKET_VALUE is invalid", + ), + ("Policyengine", "V2_SIMULATION_OUTPUT_BUCKET_VALUE is invalid"), + ], +) +def test_a_deploy_without_a_valid_output_bucket_stops(tmp_path, bucket, message): + result, environment = _run_deploy( + tmp_path, V2_SIMULATION_OUTPUT_BUCKET_VALUE=bucket + ) + + assert result.returncode != 0 + assert message in result.stderr + assert not Path(environment["GCLOUD_ARGUMENTS_CAPTURE"]).exists() + + +def test_the_deploy_step_passes_the_environments_output_bucket(): + deploy_entrypoint = workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + deploy_step = next( + step + for step in deploy_entrypoint["steps"] + if step.get("name") == "Deploy Cloud Run entrypoint candidate" + ) + + assert deploy_step["env"]["V2_SIMULATION_OUTPUT_BUCKET_VALUE"] == ( + "${{ vars.V2_SIMULATION_OUTPUT_BUCKET }}" + ) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_routes.py b/projects/policyengine-simulation-entry/tests/test_economy_routes.py new file mode 100644 index 000000000..90ca6e108 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_routes.py @@ -0,0 +1,610 @@ +"""The economy routes' answers, with the submission service faked.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from datetime import datetime +from types import SimpleNamespace +from typing import Any, NoReturn +from uuid import UUID + +import pytest +from conftest import FakeBackend, make_settings +from economy_service_fakes import REPORT_BODY, SIMULATION_BODY +from fastapi.testclient import TestClient +from policyengine_observability import REQUEST_ID_HEADER +from policyengine_simulation_entry import app as app_module +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.economy_api_models import ( + EconomyReportCreated, + EconomyReportRequest, + EconomySimulationCreated, + EconomySimulationRequest, +) +from policyengine_simulation_entry.economy_routes import call_request_id_for +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, + UnprocessableReason, +) +from policyengine_simulation_entry.economy_submissions import ( + EconomySubmissionService, +) +from policyengine_simulation_entry.simulation_database import SimulationDatabase +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER +from policyengine_simulation_persistence import ( + EconomyReportState, + EconomySimulationState, + ReportRunState, + ReportRunStatus, + SimulationStatus, +) + +OBSERVABILITY_ID = "6e1f3d4c-0a5b-4c39-9d2e-7f8a9b0c1d2e" +RECORD_ID = UUID("00000000-0000-4000-8000-0000000000e1") +RUN_ID = UUID("00000000-0000-4000-8000-0000000000e2") +SIDE_IDS = [UUID(f"00000000-0000-4000-8000-0000000000f{digit}") for digit in "1234"] +CREATED_REPORT = EconomyReportCreated( + report_id=RECORD_ID, + report_run_id=RUN_ID, + status=ReportRunStatus.PENDING, + baseline_simulation_id=SIDE_IDS[0], + baseline_simulation_run_id=SIDE_IDS[1], + baseline_reused=True, + reform_simulation_id=SIDE_IDS[2], + reform_simulation_run_id=SIDE_IDS[3], + reform_reused=False, +) +CREATED_SIMULATION = EconomySimulationCreated( + simulation_id=RECORD_ID, simulation_run_id=RUN_ID, status=SimulationStatus.PENDING +) +UNEXPECTED_FAILURE_ANSWER = { + "code": "internal_error", + "detail": "The service could not answer the request.", +} + + +def _raise(failure: Exception) -> NoReturn: + raise failure + + +@dataclass +class FakeEconomySubmissions: + """Records each submission, and answers as told.""" + + failure: Exception | None = None + stored_report: EconomyReportState | None = None + stored_simulation: EconomySimulationState | None = None + reports: list[tuple[EconomyReportRequest, str, Any]] = field(default_factory=list) + simulations: list[tuple[EconomySimulationRequest, str, Any]] = field( + default_factory=list + ) + + async def submit_report( + self, + request: EconomyReportRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomyReportCreated: + self.reports.append((request, request_id, observability_context)) + if self.failure is not None: + _raise(self.failure) + return CREATED_REPORT + + async def submit_simulation( + self, + request: EconomySimulationRequest, + *, + request_id: str, + observability_context: dict[str, Any] | None, + ) -> EconomySimulationCreated: + self.simulations.append((request, request_id, observability_context)) + if self.failure is not None: + _raise(self.failure) + return CREATED_SIMULATION + + async def read_report(self, report_id: UUID) -> EconomyReportState | None: + if self.failure is not None: + _raise(self.failure) + stored = self.stored_report + return stored if stored is not None and stored.id == report_id else None + + async def read_simulation( + self, simulation_id: UUID + ) -> EconomySimulationState | None: + if self.failure is not None: + _raise(self.failure) + stored = self.stored_simulation + return stored if stored is not None and stored.id == simulation_id else None + + +@pytest.fixture +def client_for(monkeypatch): + """Return a function building a client whose app uses ``submissions``. + + None leaves the app as an unconfigured service builds it. + """ + + def build(submissions: FakeEconomySubmissions | None) -> TestClient: + if submissions is not None: + monkeypatch.setattr( + app_module, + "economy_submissions_for", + lambda settings, database: submissions, + ) + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=lambda: None, + ) + ) + + return build + + +ROUTES = [("/reports", REPORT_BODY), ("/simulations", SIMULATION_BODY)] + + +def test_a_report_answers_201_with_its_records_and_runs(client_for) -> None: + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post( + "/reports", + json=REPORT_BODY, + headers={OBSERVABILITY_ID_HEADER: OBSERVABILITY_ID}, + ) + + assert result.status_code == 201 + assert result.json() == CREATED_REPORT.model_dump(mode="json") + assert result.headers[OBSERVABILITY_ID_HEADER] == OBSERVABILITY_ID + ((request, request_id, context),) = submissions.reports + assert request.label == REPORT_BODY["label"] + assert request.options() == { + "include_cliffs": True, + "spm": {"geography_kind": "national"}, + } + # The call's request ID is the request's own (see the request ID test). + assert request_id != OBSERVABILITY_ID + assert context["observability_id"] == OBSERVABILITY_ID + + +def test_a_simulation_answers_201_with_its_record_and_run(client_for) -> None: + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post("/simulations", json=SIMULATION_BODY) + + assert result.status_code == 201 + assert result.json() == { + "simulation_id": str(RECORD_ID), + "simulation_run_id": str(RUN_ID), + "status": "pending", + } + ((request, request_id, _),) = submissions.simulations + assert str(request.policy_id) == SIMULATION_BODY["policy_id"] + assert request_id != result.headers[OBSERVABILITY_ID_HEADER] + + +@pytest.mark.parametrize( + ("path", "body"), + [ + ("/reports", {**REPORT_BODY, "report_type": "household"}), + ("/reports", {**REPORT_BODY, "spm": None}), + ("/reports", {**REPORT_BODY, "include_cliffs": "yes"}), + ("/reports", {**REPORT_BODY, "report_id": str(RECORD_ID)}), + ("/reports", {**REPORT_BODY, "label": ""}), + ("/reports", {**REPORT_BODY, "year": 2201}), + ("/reports", {**REPORT_BODY, "policyengine_version": "6.17.1"}), + ("/reports", {**REPORT_BODY, "reform": {"policy_id": "not-a-uuid"}}), + ("/simulations", {**SIMULATION_BODY, "country": "fr"}), + ("/simulations", {**SIMULATION_BODY, "region_code": "state ca"}), + ("/simulations", {**SIMULATION_BODY, "simulation_id": str(RECORD_ID)}), + ("/simulations", {**SIMULATION_BODY, "simulation_type": "household"}), + ], +) +def test_a_malformed_request_is_refused_with_its_reasons(client_for, path, body): + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 422 + assert all({"loc", "msg"} <= set(reason) for reason in result.json()["detail"]) + assert (submissions.reports, submissions.simulations) == ([], []) + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_a_refusal_answers_400_with_its_code(client_for, path, body) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionRefusedError( + code="unknown_region", detail="no region state/zz" + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 400 + assert result.json() == {"code": "unknown_region", "detail": "no region state/zz"} + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_stored_values_that_cannot_be_calculated_answer_422_with_a_code( + client_for, path, body +) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnprocessableError( + code="missing_parameters", + reasons=( + UnprocessableReason( + loc=("body", "policy_id", "gov.a"), + msg="no such parameter", + type="missing_parameter", + ), + ), + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 422 + assert result.json() == { + "detail": [ + { + "loc": ["body", "policy_id", "gov.a"], + "msg": "no such parameter", + "type": "missing_parameter", + } + ], + "code": "missing_parameters", + } + + +@pytest.mark.parametrize( + ("path", "body", "record_field"), + [ + ("/reports", REPORT_BODY, "report_id"), + ("/simulations", SIMULATION_BODY, "simulation_id"), + ], +) +def test_a_refused_call_answers_503_naming_the_record_it_created( + client_for, path, body, record_field +) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnavailableError( + code="modal_call_refused", detail="Modal refused", record_id=RECORD_ID + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 503 + assert result.headers["Retry-After"] == "10" + assert result.json() == { + "code": "modal_call_refused", + "detail": "Modal refused", + record_field: str(RECORD_ID), + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unavailable_dependency_answers_503_naming_no_record( + client_for, path, body +) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnavailableError( + code="database_unavailable", detail="The database could not be read." + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 503 + assert result.json() == { + "code": "database_unavailable", + "detail": "The database could not be read.", + } + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unconfigured_service_answers_503(client_for, path, body) -> None: + with client_for(None) as client: + result = client.post(path, json=body) + + assert result.status_code == 503 + assert result.headers["Retry-After"] == "10" + assert result.json() == { + "code": "not_configured", + "detail": "Economy runs are not configured.", + } + + +def _emitted_records(printed: str, message: str) -> list[dict[str, Any]]: + """Return the records with ``message`` the observability runtime printed.""" + records = [ + json.loads(printed_line) + for printed_line in printed.splitlines() + if printed_line.startswith("{") + ] + return [record for record in records if record.get("message") == message] + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +def test_an_unexpected_failure_answers_a_json_500_without_its_values( + client_for, capsys, path, body +) -> None: + submissions = FakeEconomySubmissions( + failure=RuntimeError("a value that must not be logged") + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 500 + assert result.json() == UNEXPECTED_FAILURE_ANSWER + printed = capsys.readouterr().out + (record,) = _emitted_records(printed, "economy_route_failed") + attributes = record["attributes"] + assert (attributes["route"], attributes["error_type"]) == (path, "RuntimeError") + assert "test_economy_routes.py" in attributes["stack_frames"].split("; ")[0] + assert "a value that must not be logged" not in printed + + +DATABASE_URL = "postgresql://simulation_service:secret@db.example:5432/postgres" + + +def test_a_configured_service_builds_economy_submissions_without_connecting(): + """Nothing reaches the database or Modal until the first submission.""" + database = SimulationDatabase.from_url(DATABASE_URL) + try: + submissions = app_module.economy_submissions_for( + make_settings( + simulation_database_url=DATABASE_URL, + v2_modal_environment="staging", + v2_simulation_output_bucket="policyengine-v2-outputs-staging", + ), + database, + ) + finally: + database.dispose() + + assert isinstance(submissions, EconomySubmissionService) + + +@pytest.mark.parametrize( + "settings_changes", + [ + {}, + {"simulation_database_url": DATABASE_URL}, + {"v2_modal_environment": "staging"}, + ], +) +def test_without_its_database_and_v2_environment_none_is_built(settings_changes): + database = SimulationDatabase.from_url(DATABASE_URL) + try: + submissions = app_module.economy_submissions_for( + make_settings(**settings_changes), database + ) + finally: + database.dispose() + + assert submissions is None + + +READ_AT = "2026-11-12T15:00:00Z" +STORED_REPORT = EconomyReportState( + id=RECORD_ID, + country="uk", + baseline_simulation_id=SIDE_IDS[0], + reform_simulation_id=SIDE_IDS[2], + run=ReportRunState( + id=RUN_ID, + status=ReportRunStatus.FAILED, + country_package_version="2.174.1", + policyengine_version="6.17.1", + error_message="simulation_failed: the reform failed", + created_at=datetime.fromisoformat(READ_AT), + started_at=datetime.fromisoformat(READ_AT), + completed_at=datetime.fromisoformat(READ_AT), + baseline_simulation_run_id=SIDE_IDS[1], + reform_simulation_run_id=SIDE_IDS[3], + ), +) +STORED_SIMULATION = EconomySimulationState(id=RECORD_ID, run=None) + + +def test_a_report_is_read_with_its_newest_runs_status(client_for) -> None: + submissions = FakeEconomySubmissions(stored_report=STORED_REPORT) + + with client_for(submissions) as client: + result = client.get(f"/reports/{RECORD_ID}") + missing = client.get(f"/reports/{RUN_ID}") + + assert result.status_code == 200 + assert result.json() == { + "report_id": str(RECORD_ID), + "country": "uk", + "baseline_simulation_id": str(SIDE_IDS[0]), + "reform_simulation_id": str(SIDE_IDS[2]), + "run": { + "run_id": str(RUN_ID), + "status": "failed", + "country_package_version": "2.174.1", + "policyengine_version": "6.17.1", + "error_message": "simulation_failed: the reform failed", + "created_at": READ_AT, + "started_at": READ_AT, + "completed_at": READ_AT, + "baseline_simulation_run_id": str(SIDE_IDS[1]), + "reform_simulation_run_id": str(SIDE_IDS[3]), + }, + } + assert missing.status_code == 404 + assert missing.json() == { + "code": "unknown_report", + "detail": f"There is no report {RUN_ID}.", + } + + +def test_a_simulation_is_read_with_its_newest_runs_status(client_for) -> None: + submissions = FakeEconomySubmissions(stored_simulation=STORED_SIMULATION) + + with client_for(submissions) as client: + result = client.get(f"/simulations/{RECORD_ID}") + missing = client.get(f"/simulations/{RUN_ID}") + + assert result.status_code == 200 + assert result.json() == {"simulation_id": str(RECORD_ID), "run": None} + assert missing.status_code == 404 + assert missing.json()["code"] == "unknown_simulation" + + +@pytest.mark.parametrize("path", ["/reports/not-a-uuid", "/simulations/not-a-uuid"]) +def test_a_read_of_an_id_that_is_not_a_uuid_is_refused(client_for, path) -> None: + with client_for(FakeEconomySubmissions()) as client: + result = client.get(path) + + assert result.status_code == 422 + + +@pytest.mark.parametrize("record", ["reports", "simulations"]) +def test_a_read_the_database_cannot_answer_is_unavailable(client_for, record) -> None: + submissions = FakeEconomySubmissions( + failure=EconomySubmissionUnavailableError( + code="database_unavailable", detail="The database could not be read." + ) + ) + + with client_for(submissions) as client: + result = client.get(f"/{record}/{RECORD_ID}") + + assert result.status_code == 503 + assert result.headers["Retry-After"] == "10" + assert result.json() == { + "code": "database_unavailable", + "detail": "The database could not be read.", + } + + +@pytest.mark.parametrize("record", ["reports", "simulations"]) +def test_an_unconfigured_service_answers_reads_with_503(client_for, record) -> None: + with client_for(None) as client: + result = client.get(f"/{record}/{RECORD_ID}") + + assert result.status_code == 503 + assert result.json()["code"] == "not_configured" + + +@pytest.mark.parametrize("record", ["reports", "simulations"]) +def test_an_unexpected_read_failure_answers_a_json_500(client_for, record) -> None: + submissions = FakeEconomySubmissions(failure=RuntimeError("must not be logged")) + + with client_for(submissions) as client: + result = client.get(f"/{record}/{RECORD_ID}") + + assert result.status_code == 500 + assert result.json() == UNEXPECTED_FAILURE_ANSWER + + +def test_the_calls_request_id_is_the_requests_own(client_for) -> None: + """The observability ID travels in the observability context instead.""" + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + client.post( + "/reports", + json=REPORT_BODY, + headers={ + REQUEST_ID_HEADER: "request-1", + OBSERVABILITY_ID_HEADER: OBSERVABILITY_ID, + }, + ) + client.post("/simulations", json=SIMULATION_BODY) + + ((_, report_request_id, context),) = submissions.reports + ((_, simulation_request_id, _),) = submissions.simulations + assert report_request_id == "request-1" + assert context["observability_id"] == OBSERVABILITY_ID + assert simulation_request_id not in {"", OBSERVABILITY_ID} + + +@pytest.mark.parametrize( + ("path", "body", "record"), + [ + ("/reports", REPORT_BODY, "report"), + ("/simulations", SIMULATION_BODY, "simulation"), + ], +) +def test_a_create_whose_commit_is_unconfirmed_answers_a_500( + client_for, capsys, path, body, record +) -> None: + """A 500 tells the Public API the record may exist; no 503 may say it doesn't.""" + submissions = FakeEconomySubmissions( + failure=EconomySubmissionOutcomeUnknownError( + record, error_type="OperationalError" + ) + ) + + with client_for(submissions) as client: + result = client.post(path, json=body) + + assert result.status_code == 500 + assert result.json() == { + "code": "internal_error", + "detail": ( + f"The service may have created the {record}, but could not confirm " + "it was written." + ), + } + printed = capsys.readouterr().out + (logged,) = _emitted_records(printed, "economy_record_outcome_unknown") + assert logged["attributes"]["record"] == record + assert logged["attributes"]["error_type"] == "OperationalError" + # One event for one outcome. + assert _emitted_records(printed, "economy_records_outcome_unknown") == [] + + +@pytest.mark.parametrize(("path", "body"), ROUTES) +@pytest.mark.parametrize( + ("request_id", "kept"), + [ + ("r" * 128, True), + ("r" * 129, False), + ("r" * 256, False), + ("has spaces in it", False), + ], + ids=["128-kept", "129-replaced", "256-replaced", "spaces-replaced"], +) +def test_a_request_id_the_call_cannot_carry_is_replaced( + client_for, path, body, request_id, kept +) -> None: + """The call's context takes 1 to 128 of the observability runtime's characters.""" + submissions = FakeEconomySubmissions() + + with client_for(submissions) as client: + result = client.post(path, json=body, headers={REQUEST_ID_HEADER: request_id}) + + assert result.status_code == 201 + ((_, call_request_id, _),) = submissions.reports or submissions.simulations + if kept: + assert call_request_id == request_id + else: + assert str(UUID(call_request_id)) == call_request_id + + +@pytest.mark.parametrize("state", [{"request_id": ""}, {"request_id": None}, {}]) +def test_a_request_without_a_usable_id_gets_a_new_one(state) -> None: + request = SimpleNamespace(state=SimpleNamespace(**state)) + + call_request_id = call_request_id_for(request) + + assert str(UUID(call_request_id)) == call_request_id diff --git a/projects/policyengine-simulation-entry/tests/test_economy_submissions.py b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py new file mode 100644 index 000000000..062fa9b7c --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_submissions.py @@ -0,0 +1,683 @@ +"""Creating economy reports and simulations, with the database and Modal faked.""" + +from __future__ import annotations + +import asyncio +from dataclasses import replace +from uuid import UUID + +import pytest +from economy_fixtures import STATE_DATASET +from economy_service_fakes import ( + CATALOG, + CREATED_AT, + DYNAMIC_ID, + OUTPUT_BUCKET, + POLICY_ID, + REFORM, + REPORT_BODY, + REUSED, + FakeCaller, + FakeStore, + FakeWorkers, + report_request, + service, + simulation_request, +) +from household_fixtures import POLICYENGINE_VERSION +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + ReportRunIds, + Stage12InvocationContext, +) +from policyengine_simulation_entry.economy_api_models import EconomyReportRequest +from policyengine_simulation_entry.economy_submission_errors import ( + EconomySubmissionOutcomeUnknownError, + EconomySubmissionRefusedError, + EconomySubmissionUnavailableError, + EconomySubmissionUnprocessableError, +) +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, +) +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) +from policyengine_simulation_persistence import ( + CatalogDataset, + EconomyRecordOutcomeUnknownError, + EconomyRecordReferenceError, + EconomyReportState, + EconomySimulationState, + MissingPolicyParametersError, + OverlappingReformError, + ReportRunState, + ReportRunStatus, + ReportSideError, + RunReuse, + RunVersions, + SimulationChoice, + SimulationRunState, + SimulationStatus, + UnknownCatalogVersionError, + UnknownDatasetError, + UnknownDynamicError, + UnknownPolicyError, + UnknownRegionError, +) +from sqlalchemy.exc import OperationalError + +DATABASE_ERROR = OperationalError("SELECT 1", {}, Exception("server closed")) +VERSIONS = RunVersions( + country_package_version="2.2.1", policyengine_version=POLICYENGINE_VERSION +) +OBSERVED = {"observability_id": "observed"} + + +def _submit_report(submissions, request=None): + return asyncio.run( + submissions.submit_report( + request or report_request(), + request_id="observed", + observability_context=OBSERVED, + ) + ) + + +def _submit_simulation(submissions, request=None): + return asyncio.run( + submissions.submit_simulation( + request or simulation_request(), + request_id="observed", + observability_context=OBSERVED, + ) + ) + + +def test_a_report_is_created_with_new_runs_and_its_coordinator_started() -> None: + store = FakeStore() + caller = FakeCaller() + + created = _submit_report(service(store, caller=caller)) + + (request, (baseline, reform, reuse)) = store.report_reads[0] + assert (request.country, request.policyengine_version, request.region_code) == ( + "us", + POLICYENGINE_VERSION, + "us", + ) + assert (request.dataset_name, request.year) == (None, 2026) + assert baseline == SimulationChoice(policy_id=None, dynamic_id=None) + assert reform == SimulationChoice(policy_id=POLICY_ID, dynamic_id=DYNAMIC_ID) + options = {"include_cliffs": True, "spm": {"geography_kind": "national"}} + assert reuse == RunReuse( + versions=VERSIONS, inputs=options, output_uri_prefix=f"gs://{OUTPUT_BUCKET}/" + ) + (report,) = store.reports + assert (report.id, report.run_id) == (UUID(int=1), UUID(int=2)) + assert (report.label, report.description, report.country) == ( + "A larger child tax credit", + "Raises the credit to $3,000", + "us", + ) + assert report.tax_benefit_model_id == CATALOG.model_id + assert report.population == CATALOG.population(2026) + assert (report.versions, report.inputs) == (VERSIONS, options) + assert (report.baseline.reused, report.reform.reused) == (False, False) + assert report.new_simulation_run_ids == (UUID(int=4), UUID(int=6)) + assert report.reform.choice == reform + ((_, payload, context, observed),) = caller.report_runs + assert observed == OBSERVED + assert ReportExecutionInput.model_validate(payload).reform.policy == REFORM + assert Stage12InvocationContext.model_validate(context).report_run == ReportRunIds( + report_run_id=report.run_id, + baseline_simulation_run_id=report.baseline.simulation_run_id, + reform_simulation_run_id=report.reform.simulation_run_id, + ) + assert store.recorded_calls == [(report.run_id, "fc-coordinator")] + assert created.model_dump() == { + "report_id": report.id, + "report_run_id": report.run_id, + "status": ReportRunStatus.PENDING, + "baseline_simulation_id": report.baseline.simulation_id, + "baseline_simulation_run_id": report.baseline.simulation_run_id, + "baseline_reused": False, + "reform_simulation_id": report.reform.simulation_id, + "reform_simulation_run_id": report.reform.simulation_run_id, + "reform_reused": False, + } + + +def test_a_side_with_a_reusable_run_names_it_instead_of_a_new_one() -> None: + store = FakeStore(reusable=(REUSED, None)) + caller = FakeCaller() + + created = _submit_report(service(store, caller=caller)) + + (report,) = store.reports + assert report.baseline.reused is True + assert (report.baseline.simulation_id, report.baseline.simulation_run_id) == ( + REUSED.simulation_id, + REUSED.simulation_run_id, + ) + assert report.new_simulation_run_ids == (report.reform.simulation_run_id,) + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.baseline.simulation_execution_id == REUSED.simulation_run_id + assert (created.baseline_reused, created.reform_reused) == (True, False) + assert created.baseline_simulation_run_id == REUSED.simulation_run_id + + +def test_without_the_output_bucket_no_run_is_looked_for() -> None: + store = FakeStore() + + _submit_report(service(store, output_bucket=None)) + + assert store.report_reads[0][1][2] is None + + +def test_the_report_stores_only_the_spm_settings_given() -> None: + store = FakeStore() + caller = FakeCaller() + + _submit_report( + service(store, caller=caller), + report_request(include_cliffs=False, spm={"scenario": "zero"}), + ) + + assert store.reports[0].inputs == { + "include_cliffs": False, + "spm": {"scenario": "zero"}, + } + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.baseline.options == {"spm": {"scenario": "zero"}} + + +def test_a_report_without_spm_settings_stores_none() -> None: + store = FakeStore() + caller = FakeCaller() + without_spm = {key: value for key, value in REPORT_BODY.items() if key != "spm"} + + _submit_report( + service(store, caller=caller), + EconomyReportRequest.model_validate(without_spm), + ) + + assert store.reports[0].inputs == {"include_cliffs": True} + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.reform.options == {"include_cliffs": True} + + +@pytest.mark.parametrize( + ("error", "code"), + [ + ( + UnknownCatalogVersionError("store text: no catalog version"), + "unknown_catalog_version", + ), + (UnknownRegionError("store text: no region"), "unknown_region"), + (UnknownDatasetError("store text: no dataset"), "unknown_dataset"), + (UnknownPolicyError("store text: no policy"), "unknown_policy"), + (UnknownDynamicError("store text: no dynamic"), "unknown_dynamic"), + ], +) +def test_a_missing_row_is_refused_before_anything_is_written(error, code) -> None: + store = FakeStore(read_failure=error) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionRefusedError) as report_refusal: + _submit_report(service(store, caller=caller)) + with pytest.raises(EconomySubmissionRefusedError) as simulation_refusal: + _submit_simulation(service(store, caller=caller)) + + assert report_refusal.value.code == simulation_refusal.value.code == code + assert str(error) not in report_refusal.value.detail + assert (store.reports, store.simulations, caller.report_runs) == ([], [], []) + + +def test_values_the_catalog_version_lacks_cannot_be_calculated() -> None: + error = MissingPolicyParametersError(["gov.a", "gov.b"], source="dynamic") + store = FakeStore(read_failure=error) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_simulation(service(store)) + + assert refusal.value.code == "missing_parameters" + assert [reason.loc for reason in refusal.value.reasons] == [ + ("body", "dynamic_id", "gov.a"), + ("body", "dynamic_id", "gov.b"), + ] + assert POLICYENGINE_VERSION in refusal.value.reasons[0].msg + assert store.simulations == [] + + +def test_a_parameter_the_policy_and_dynamic_both_set_cannot_be_calculated() -> None: + store = FakeStore(read_failure=OverlappingReformError(["gov.a"])) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_simulation(service(store)) + + assert refusal.value.code == "overlapping_parameters" + assert [reason.type for reason in refusal.value.reasons] == [ + "overlapping_parameter" + ] + + +@pytest.mark.parametrize( + ("dataset_name", "code"), + [(None, "dataset_not_in_bundle"), ("enhanced_cps_2024", "unknown_dataset")], +) +def test_a_catalog_dataset_the_bundle_lacks_is_refused(dataset_name, code) -> None: + catalog = replace( + CATALOG, dataset=CatalogDataset(id=CATALOG.dataset.id, name="enhanced_cps_2024") + ) + store = FakeStore(catalog=catalog) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(store), report_request(dataset_name=dataset_name)) + + assert refusal.value.code == code + assert store.reports == [] + + +def test_a_named_dataset_in_the_bundle_is_the_one_calculated() -> None: + catalog = replace( + CATALOG, + dataset=CatalogDataset(id=CATALOG.dataset.id, name=STATE_DATASET.identity), + ) + store = FakeStore(catalog=catalog) + caller = FakeCaller() + + _submit_report( + service(store, caller=caller), + report_request(dataset_name=STATE_DATASET.identity), + ) + + payload = ReportExecutionInput.model_validate(caller.report_runs[0][1]) + assert payload.baseline.bundle.dataset.identity == STATE_DATASET.identity + + +def test_a_country_the_default_worker_cannot_calculate_is_refused() -> None: + store = FakeStore() + workers = FakeWorkers(EconomyWorkerUnavailableError("cannot calculate uk")) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(store, workers=workers)) + + assert refusal.value.code == "no_economy_worker" + assert store.report_reads == [] + + +def test_a_manifest_that_cannot_be_read_is_unavailable() -> None: + failure = VersionManifestUnavailableError("unreadable") + failure.__cause__ = ConnectionError("Modal") + store = FakeStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_report(service(store, workers=FakeWorkers(failure))) + + assert (unavailable.value.code, unavailable.value.record_id) == ( + "manifest_unavailable", + None, + ) + + +@pytest.mark.parametrize("failing", ["read", "create"]) +def test_a_database_that_cannot_be_used_is_unavailable_before_any_call( + failing, +) -> None: + store = FakeStore(**{f"{failing}_failure": DATABASE_ERROR}) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_report(service(store, caller=caller)) + + assert (unavailable.value.code, unavailable.value.record_id) == ( + "database_unavailable", + None, + ) + assert caller.report_runs == [] + + +@pytest.mark.parametrize( + ("column", "code"), + [ + ("region_id", "unknown_region"), + ("dataset_id", "unknown_dataset"), + ("dynamic_id", "unknown_dynamic"), + ("policy_id", "unknown_policy"), + ("tax_benefit_model_version_id", "unknown_catalog_version"), + ("baseline_simulation_run_id", "missing_reference"), + (None, "missing_reference"), + ], +) +def test_a_row_deleted_after_the_reads_is_refused_by_its_column(column, code) -> None: + store = FakeStore( + create_failure=EconomyRecordReferenceError( + None if column is None else "simulations", column + ) + ) + caller = FakeCaller() + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(store, caller=caller)) + + assert refusal.value.code == code + assert "no longer exists" in refusal.value.detail + assert caller.report_runs == [] + + +@pytest.mark.parametrize( + "caller", + [FakeCaller(failure=ConnectionError("refused")), FakeCaller(hang=True)], + ids=["refused", "not accepted in time"], +) +def test_a_call_modal_does_not_confirm_fails_the_report_run_and_names_it( + caller, +) -> None: + store = FakeStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_report(service(store, caller=caller)) + + (report,) = store.reports + assert (unavailable.value.code, unavailable.value.record_id) == ( + "modal_call_refused", + report.id, + ) + ((run_id, simulation_run_ids, message),) = store.failed_runs + assert run_id == report.run_id + assert simulation_run_ids == report.new_simulation_run_ids + assert message.startswith("modal_call_refused: Modal did not accept the report") + assert store.recorded_calls == [] + + +def test_a_report_run_its_coordinator_claimed_meanwhile_is_answered_created() -> None: + store = FakeStore(claimed=True) + + created = _submit_report( + service(store, caller=FakeCaller(failure=TimeoutError("late"))) + ) + + assert store.failed_runs == [] + assert created.report_id == store.reports[0].id + assert created.status is ReportRunStatus.PENDING + + +@pytest.mark.parametrize("failing", ["fail", "record"]) +def test_a_run_whose_state_cannot_be_written_is_answered_created_pending( + failing, +) -> None: + store = FakeStore(**{f"{failing}_failure": DATABASE_ERROR}) + caller = FakeCaller( + failure=ConnectionError("refused") if failing == "fail" else None + ) + + created = _submit_report(service(store, caller=caller)) + + assert created.report_id == store.reports[0].id + assert created.status is ReportRunStatus.PENDING + + +def test_a_simulation_is_created_and_its_run_started_alone() -> None: + store = FakeStore() + caller = FakeCaller() + + created = _submit_simulation(service(store, caller=caller)) + + (request, choice) = store.simulation_reads[0] + assert (request.policyengine_version, request.year) == (POLICYENGINE_VERSION, 2026) + assert choice == SimulationChoice(policy_id=POLICY_ID, dynamic_id=None) + (simulation,) = store.simulations + assert (simulation.id, simulation.run_id) == (UUID(int=1), UUID(int=2)) + assert simulation.population == CATALOG.population(2026) + assert simulation.versions == VERSIONS + ((_, payload, context, observed),) = caller.simulation_runs + assert payload["simulation_execution_id"] == str(simulation.run_id) + assert payload["policy"] == REFORM + assert context["simulation_run_id"] == str(simulation.run_id) + assert observed == OBSERVED + assert store.recorded_calls == [(simulation.run_id, "fc-worker")] + assert created.model_dump() == { + "simulation_id": simulation.id, + "simulation_run_id": simulation.run_id, + "status": SimulationStatus.PENDING, + } + + +def test_a_simulation_run_modal_refuses_is_failed_and_the_simulation_named() -> None: + store = FakeStore() + + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + _submit_simulation( + service(store, caller=FakeCaller(failure=ConnectionError("refused"))) + ) + + (simulation,) = store.simulations + assert (unavailable.value.code, unavailable.value.record_id) == ( + "modal_call_refused", + simulation.id, + ) + assert store.failed_runs[0][0] == simulation.run_id + assert store.failed_runs[0][2].startswith( + "modal_call_refused: Modal did not accept the simulation" + ) + + +def test_the_calls_context_records_when_and_where_it_was_made() -> None: + caller = FakeCaller() + + _submit_report(service(FakeStore(), caller=caller)) + + context = Stage12InvocationContext.model_validate(caller.report_runs[0][2]) + assert (context.created_at, context.request_id, context.environment) == ( + CREATED_AT, + "observed", + "staging", + ) + + +REPORT_RUN = ReportRunState( + id=UUID(int=11), + status=ReportRunStatus.SUCCEEDED, + country_package_version="2.2.1", + policyengine_version=POLICYENGINE_VERSION, + error_message=None, + created_at=CREATED_AT, + started_at=CREATED_AT, + completed_at=CREATED_AT, + baseline_simulation_run_id=UUID(int=12), + reform_simulation_run_id=UUID(int=13), +) +STORED_REPORT = EconomyReportState( + id=UUID(int=10), + country="us", + baseline_simulation_id=UUID(int=14), + reform_simulation_id=UUID(int=15), + run=REPORT_RUN, +) +STORED_SIMULATION = EconomySimulationState( + id=UUID(int=20), + run=SimulationRunState( + id=UUID(int=21), + status=SimulationStatus.FAILED, + country_package_version="2.2.1", + policyengine_version=POLICYENGINE_VERSION, + error_message="simulation_failed: no", + created_at=CREATED_AT, + started_at=None, + completed_at=CREATED_AT, + ), +) + + +def test_a_report_and_a_simulation_are_read_as_they_stand() -> None: + submissions = service( + FakeStore(stored_report=STORED_REPORT, stored_simulation=STORED_SIMULATION) + ) + + assert asyncio.run(submissions.read_report(STORED_REPORT.id)) == STORED_REPORT + assert ( + asyncio.run(submissions.read_simulation(STORED_SIMULATION.id)) + == STORED_SIMULATION + ) + assert asyncio.run(submissions.read_report(UUID(int=99))) is None + assert asyncio.run(submissions.read_simulation(UUID(int=99))) is None + + +def test_a_read_the_database_cannot_answer_is_unavailable() -> None: + submissions = service(FakeStore(read_failure=DATABASE_ERROR)) + + for read in (submissions.read_report, submissions.read_simulation): + with pytest.raises(EconomySubmissionUnavailableError) as unavailable: + asyncio.run(read(UUID(int=1))) + assert unavailable.value.code == "database_unavailable" + + +@pytest.mark.parametrize( + ("side", "source", "field"), + [("reform", "policy", "policy_id"), ("baseline", "dynamic", "dynamic_id")], +) +def test_a_reports_missing_parameters_are_located_on_their_side( + side, source, field +) -> None: + error = ReportSideError( + side, MissingPolicyParametersError(["gov.a"], source=source) + ) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_report(service(FakeStore(read_failure=error))) + + assert refusal.value.code == "missing_parameters" + assert [reason.loc for reason in refusal.value.reasons] == [ + ("body", side, field, "gov.a") + ] + + +def test_a_reports_overlapping_parameters_are_located_on_their_side() -> None: + error = ReportSideError("baseline", OverlappingReformError(["gov.a"])) + + with pytest.raises(EconomySubmissionUnprocessableError) as refusal: + _submit_report(service(FakeStore(read_failure=error))) + + assert [reason.loc for reason in refusal.value.reasons] == [ + ("body", "baseline", "dynamic_id", "gov.a") + ] + + +# Text a store error carries, which no answer may pass on. +ERROR_TEXT = "internal text that must not be passed on" + + +@pytest.mark.parametrize( + ("error", "request_fields", "code", "detail"), + [ + ( + UnknownCatalogVersionError(ERROR_TEXT), + {}, + "unknown_catalog_version", + f"us has no catalog version for .py {POLICYENGINE_VERSION}, the " + "default v2 worker's version.", + ), + ( + UnknownRegionError(ERROR_TEXT), + {"region_code": "state/zz"}, + "unknown_region", + f"us's catalog version for .py {POLICYENGINE_VERSION} has no region " + "state/zz.", + ), + ( + UnknownDatasetError(ERROR_TEXT), + {"dataset_name": "acs_2099"}, + "unknown_dataset", + f"us's catalog version for .py {POLICYENGINE_VERSION} has no dataset " + "acs_2099.", + ), + ( + UnknownDatasetError(ERROR_TEXT), + {}, + "unknown_dataset", + f"us's catalog version for .py {POLICYENGINE_VERSION} has no default " + "dataset for region us.", + ), + ( + ReportSideError("reform", UnknownPolicyError(ERROR_TEXT)), + {}, + "unknown_policy", + "us has no stored policy the request's reform names.", + ), + ( + ReportSideError("baseline", UnknownDynamicError(ERROR_TEXT)), + {}, + "unknown_dynamic", + "No stored dynamic matches the one the request's baseline names.", + ), + ], +) +def test_a_refusals_detail_is_built_from_the_request_not_the_error( + error, request_fields, code, detail +) -> None: + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report( + service(FakeStore(read_failure=error)), report_request(**request_fields) + ) + + assert (refusal.value.code, refusal.value.detail) == (code, detail) + assert ERROR_TEXT not in refusal.value.detail + + +def test_a_standalone_simulations_unknown_rows_name_its_own_fields() -> None: + for error, detail in ( + (UnknownPolicyError(ERROR_TEXT), "us has no stored policy the request names."), + ( + UnknownDynamicError(ERROR_TEXT), + "No stored dynamic matches the one the request names.", + ), + ): + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_simulation(service(FakeStore(read_failure=error))) + assert refusal.value.detail == detail + + +def test_a_country_without_an_economy_worker_is_refused_in_the_services_words() -> None: + workers = FakeWorkers(EconomyWorkerUnavailableError(ERROR_TEXT)) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(FakeStore(), workers=workers)) + + assert refusal.value.detail == ( + "The default v2 worker cannot calculate economy runs for us." + ) + + +def test_a_dataset_the_bundle_lacks_is_named_in_the_services_words() -> None: + catalog = replace( + CATALOG, dataset=CatalogDataset(id=CATALOG.dataset.id, name="enhanced_cps_2024") + ) + + with pytest.raises(EconomySubmissionRefusedError) as refusal: + _submit_report(service(FakeStore(catalog=catalog))) + + assert refusal.value.detail == ( + f"The default v2 worker, .py {POLICYENGINE_VERSION}, has no dataset " + "enhanced_cps_2024." + ) + + +@pytest.mark.parametrize("record", ["report", "simulation"]) +def test_a_create_whose_commit_is_unconfirmed_may_have_written_the_record( + record, +) -> None: + lost_commit = EconomyRecordOutcomeUnknownError("lost") + lost_commit.__cause__ = DATABASE_ERROR + store = FakeStore(create_failure=lost_commit) + caller = FakeCaller() + submit = _submit_report if record == "report" else _submit_simulation + + with pytest.raises(EconomySubmissionOutcomeUnknownError) as unknown: + submit(service(store, caller=caller)) + + assert (unknown.value.record, unknown.value.error_type) == ( + record, + "OperationalError", + ) + assert (caller.report_runs, caller.simulation_runs) == ([], []) diff --git a/projects/policyengine-simulation-entry/tests/test_economy_workers.py b/projects/policyengine-simulation-entry/tests/test_economy_workers.py new file mode 100644 index 000000000..9192e4f27 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_economy_workers.py @@ -0,0 +1,196 @@ +"""The default v2 worker's economy functions, and starting their calls on Modal.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +import pytest +from economy_fixtures import economy_manifest, economy_worker_version +from household_fixtures import ( + BUNDLE_MANIFEST_SHA256, + POLICYENGINE_VERSION, + household_worker_version, + manifest_store, +) +from policyengine_simulation_contract.stage12_manifest import ( + V2CountryWorker, + manifest_digest, + v2_application_name, +) +from policyengine_simulation_entry import economy_workers +from policyengine_simulation_entry.economy_workers import ( + EconomyWorkerUnavailableError, + ManifestEconomyWorkers, + ModalEconomyRunCaller, + default_economy_worker, +) +from policyengine_simulation_entry.household_workers import ( + VersionManifestUnavailableError, +) +from policyengine_simulation_persistence import RunVersions + + +def test_the_default_workers_economy_functions_and_bundle_are_found() -> None: + manifest = economy_manifest() + + worker = default_economy_worker(manifest, "us") + + assert worker.version == POLICYENGINE_VERSION + assert worker.version_manifest_sha256 == manifest_digest(manifest) + assert worker.application_name == v2_application_name(POLICYENGINE_VERSION) + assert worker.coordinator_callable == "coordinate_report" + assert worker.simulation_callable == "run_single_simulation_us" + assert worker.bundle_manifest_sha256 == BUNDLE_MANIFEST_SHA256 + assert worker.country.country == "us" + assert worker.versions == RunVersions( + country_package_version="2.2.1", policyengine_version=POLICYENGINE_VERSION + ) + + +def test_a_country_the_default_worker_cannot_calculate_is_refused() -> None: + base = economy_worker_version() + household_only = base.model_copy( + update={ + "countries": tuple( + V2CountryWorker( + country=item.country, + economy_capable=False, + single_simulation_callable=None, + household_callable=item.household_callable, + ) + if item.country == "uk" + else item + for item in base.countries + ), + "validation": base.validation.model_copy( + update={"country_validation_invocation_ids": {"us": "validation-us"}} + ), + } + ) + manifest = economy_manifest(household_only) + + with pytest.raises(EconomyWorkerUnavailableError, match="uk"): + default_economy_worker(manifest, "uk") + assert default_economy_worker(manifest, "us").country.country == "us" + + +def test_the_manifest_store_is_read_for_the_default_worker() -> None: + newer = household_worker_version("6.18.0") + workers = ManifestEconomyWorkers(manifest_store(household_worker_version(), newer)) + + worker = workers.default_worker("uk") + + assert worker.version == POLICYENGINE_VERSION + assert worker.simulation_callable == "run_single_simulation_uk" + + +class UnreadableStore: + def get(self, key: str, default: object = None) -> object: + msg = "Modal is unreachable" + raise ConnectionError(msg) + + def __setitem__(self, key: str, value: object) -> None: + raise AssertionError + + +def test_a_manifest_that_cannot_be_read_is_unavailable() -> None: + with pytest.raises(VersionManifestUnavailableError) as raised: + ManifestEconomyWorkers(UnreadableStore()).default_worker("us") + + assert isinstance(raised.value.__cause__, ConnectionError) + + +class FakeFunction: + """A Modal function looked up by name, recording each spawn.""" + + def __init__(self, object_id: object) -> None: + self.object_id = object_id + self.spawns: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.spawn = SimpleNamespace(aio=self._spawn) + + async def _spawn(self, *arguments: Any, **keywords: Any) -> SimpleNamespace: + self.spawns.append((arguments, keywords)) + return SimpleNamespace(object_id=self.object_id) + + +@pytest.fixture +def modal_functions(monkeypatch): + """Replace Modal's lookup by name; return the lookups and the function.""" + looked_up: list[tuple[str, str, str]] = [] + function = FakeFunction("fc-started") + + def from_name(application: str, name: str, *, environment_name: str): + looked_up.append((application, name, environment_name)) + return function + + monkeypatch.setattr( + economy_workers.modal.Function, "from_name", staticmethod(from_name) + ) + return looked_up, function + + +def test_a_report_run_starts_the_coordinator_with_its_report_and_context( + modal_functions, +) -> None: + looked_up, function = modal_functions + worker = default_economy_worker(economy_manifest(), "us") + + call_id = asyncio.run( + ModalEconomyRunCaller("staging").start_report_run( + worker, + {"report": "payload"}, + {"report_run": "context"}, + observability_context={"observability_id": "observed"}, + ) + ) + + assert call_id == "fc-started" + assert looked_up == [(worker.application_name, "coordinate_report", "staging")] + # No evaluation record follows the context. + assert function.spawns == [ + ( + ({"report": "payload"}, {"report_run": "context"}), + {"observability_context": {"observability_id": "observed"}}, + ) + ] + + +def test_a_standalone_run_starts_the_countrys_simulation_function( + modal_functions, +) -> None: + looked_up, function = modal_functions + worker = default_economy_worker(economy_manifest(), "uk") + + call_id = asyncio.run( + ModalEconomyRunCaller("main").start_simulation_run( + worker, + {"simulation": 1}, + {"simulation_run_id": "run"}, + observability_context=None, + ) + ) + + assert call_id == "fc-started" + assert looked_up == [(worker.application_name, "run_single_simulation_uk", "main")] + assert function.spawns == [ + ( + ({"simulation": 1}, {"simulation_run_id": "run"}), + {"observability_context": None}, + ) + ] + + +@pytest.mark.parametrize("object_id", [None, ""]) +def test_a_call_modal_gives_no_id_is_an_error(modal_functions, object_id) -> None: + _, function = modal_functions + function.object_id = object_id + worker = default_economy_worker(economy_manifest(), "us") + + with pytest.raises(RuntimeError, match="no call ID"): + asyncio.run( + ModalEconomyRunCaller("staging").start_report_run( + worker, {}, {}, observability_context=None + ) + ) diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 631db41d6..30f8022af 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -33,12 +33,16 @@ class OpenAPIDocument(TypedDict): ("GET", "/health"), ("GET", "/jobs/{job_id}"), ("GET", "/ready"), + ("GET", "/reports/{report_id}"), + ("GET", "/simulations/{simulation_id}"), ("GET", "/versions"), ("GET", "/versions/{kind}"), ("POST", "/household-calculations"), ("POST", "/ping"), + ("POST", "/reports"), ("POST", "/simulate/economy/budget-window"), ("POST", "/simulate/economy/comparison"), + ("POST", "/simulations"), } @@ -77,11 +81,37 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: ) +# The economy report and simulation routes, which only the Cloud Run +# entrypoint has. +ECONOMY_PATHS = ( + "/reports", + "/reports/{report_id}", + "/simulations", + "/simulations/{simulation_id}", +) +ECONOMY_SCHEMAS = ( + "EconomyError", + "EconomyReportCreated", + "EconomyReportRequest", + "EconomyReportState", + "EconomyRequestReason", + "EconomyRequestUnprocessable", + "EconomySimulationCreated", + "EconomySimulationRequest", + "EconomySimulationState", + "PolicyChoice", + "ReportRunRead", + "ReportRunStatus", + "SimulationRunRead", + "SimulationStatus", +) + + def normalized_compatibility_paths(spec: OpenAPIDocument) -> OpenAPIPathMap: """Remove intentional Cloud Run and additive SPM error documentation.""" paths = deepcopy(spec["paths"]) paths.pop("/ready", None) - for path in HOUSEHOLD_PATHS: + for path in (*HOUSEHOLD_PATHS, *ECONOMY_PATHS): paths.pop(path, None) for operations in paths.values(): for method, operation in operations.items(): @@ -212,7 +242,7 @@ def test_normalized_contract_matches_old_gateway(): cloud_run_schemas.pop("SimulationEntrypointReadinessResponse") cloud_run_schemas.pop("SimulationDatabaseReadiness") cloud_run_schemas.pop("SimulationErrorResponse") - for schema_name in HOUSEHOLD_SCHEMAS: + for schema_name in (*HOUSEHOLD_SCHEMAS, *ECONOMY_SCHEMAS): cloud_run_schemas.pop(schema_name) assert cloud_run_schemas == gateway_spec["components"]["schemas"] gateway_operation_ids = operation_ids(gateway_spec) @@ -323,3 +353,106 @@ def test_a_household_calculation_request_names_its_fields_only(): "succeeded", "failed", ] + + +def test_economy_routes_declare_every_answer(): + spec = create_app().openapi() + error = {"$ref": "#/components/schemas/EconomyError"} + unprocessable = {"$ref": "#/components/schemas/EconomyRequestUnprocessable"} + + for path, created in ( + ("/reports", "EconomyReportCreated"), + ("/simulations", "EconomySimulationCreated"), + ): + responses = _responses(spec, "post", path) + assert set(responses) == {"201", "400", "403", "422", "500", "503"} + assert _response_schema(responses["201"]) == { + "$ref": f"#/components/schemas/{created}" + } + for status_code in ("400", "500", "503"): + assert _response_schema(responses[status_code]) == error + assert _response_schema(responses["422"]) == unprocessable + record = path.removeprefix("/").removesuffix("s") + assert f"{record}_id" in responses["503"]["description"] + assert "nothing was written" in responses["400"]["description"] + error_properties = spec["components"]["schemas"]["EconomyError"]["properties"] + assert {"code", "detail", "report_id", "simulation_id"} == set(error_properties) + codes = set(error_properties["code"]["enum"]) + assert { + "unknown_region", + "unknown_dataset", + "unknown_dynamic", + "unknown_catalog_version", + "modal_call_refused", + "internal_error", + } <= codes + + +def test_every_economy_503_declares_its_retry_after_header(): + spec = create_app().openapi() + + for path in ECONOMY_PATHS: + method = "get" if path.endswith("}") else "post" + headers = _responses(spec, method, path)["503"]["headers"] + assert headers["Retry-After"]["schema"] == {"type": "integer"} + + +def test_economy_routes_name_their_operations(): + names = operation_ids(cast("OpenAPIDocument", create_app().openapi())) + + assert names[("POST", "/reports")] == "create_report" + assert names[("POST", "/simulations")] == "create_simulation" + + +def test_economy_requests_name_no_version_and_no_record_id(): + """The service runs its default worker's version and chooses every ID.""" + schemas = create_app().openapi()["components"]["schemas"] + + assert set(schemas["EconomyReportRequest"]["properties"]) == { + "country", + "report_type", + "label", + "description", + "year", + "region_code", + "dataset_name", + "baseline", + "reform", + "include_cliffs", + "spm", + } + assert set(schemas["EconomySimulationRequest"]["properties"]) == { + "country", + "simulation_type", + "year", + "region_code", + "dataset_name", + "policy_id", + "dynamic_id", + } + # Omitted, not null, when the caller gives no SPM settings. + spm = schemas["EconomyReportRequest"]["properties"]["spm"] + assert spm["$ref"] == "#/components/schemas/SPMSelection" + assert "default" not in spm + assert "spm" not in schemas["EconomyReportRequest"].get("required", []) + + +def test_economy_reads_declare_every_answer(): + spec = create_app().openapi() + error = {"$ref": "#/components/schemas/EconomyError"} + + for path, state in ( + ("/reports/{report_id}", "EconomyReportState"), + ("/simulations/{simulation_id}", "EconomySimulationState"), + ): + responses = _responses(spec, "get", path) + assert set(responses) == {"200", "403", "404", "422", "500", "503"} + assert _response_schema(responses["200"]) == { + "$ref": f"#/components/schemas/{state}" + } + for status_code in ("404", "500", "503"): + assert _response_schema(responses[status_code]) == error + assert "Modal" not in responses["503"]["description"] + names = operation_ids(cast("OpenAPIDocument", spec)) + assert names[("GET", "/reports/{report_id}")] == "get_report" + assert names[("GET", "/simulations/{simulation_id}")] == "get_simulation" diff --git a/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py new file mode 100644 index 000000000..097ab1974 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_staging_economy_check.py @@ -0,0 +1,471 @@ +"""The staging economy check, run on every staging deploy. + +``simulation-staging-economy-check.sh`` creates a small economy report, +current law against current law for Wales in 2026, through the deployed +candidate's ``POST /reports``, then follows its run through +``GET /reports/{report_id}`` and requires it to succeed. It runs after the +v2 manifest's default has moved to the release's worker, so the report runs +on the release's own coordinator. Once a .py version's first check has +succeeded, later checks on it reuse those simulation runs, and check the +coordinator's use of their stored output. +""" + +from __future__ import annotations + +import json + +import pytest +from test_deployment_assets_support import ( + REUSABLE_DEPLOY_WORKFLOW, + evaluate_github_expression, + routing_context, + workflow_jobs, +) +from workflow_script_support import FakeCommands, run_script + +CHECK_SCRIPT = "simulation-staging-economy-check.sh" +READY_SCRIPT = "simulation-staging-economy-ready.sh" +CHECK_JOB = "staging_economy_check" +BASE_URL = "https://policyengine-simulation-entry-staging-s-123.a.run.app" +REPORT_ID = "00000000-0000-4000-8000-0000000000e1" +REPORT_RUN_ID = "00000000-0000-4000-8000-0000000000e2" +OLDER_RUN_ID = "00000000-0000-4000-8000-0000000000e3" +ACCESS_TOKEN = "test-access-token" + +# The fake curl answers the token request, the submission and each poll in +# turn, and keeps each request body it was sent on standard input. +FAKE_CURL = r""" +output="" +url="" +method="GET" +data="" +while (($#)); do + case "$1" in + --output) output="$2"; shift 2 ;; + --request) method="$2"; shift 2 ;; + --data) data="$2"; shift 2 ;; + --header | --write-out | --max-time) shift 2 ;; + *) url="$1"; shift ;; + esac +done +calls="$(($(cat "${FAKE_STATE}/calls" 2>/dev/null || echo 0) + 1))" +echo "${calls}" > "${FAKE_STATE}/calls" +if [[ "${data}" == "@-" ]]; then + cat > "${FAKE_STATE}/body-${calls}.json" +fi +case "${method} ${url}" in + "POST https://issuer.example/oauth/token") + printf '{"access_token": "%s"}' "${FAKE_ACCESS_TOKEN}" > "${output}" + printf '200' + ;; + "POST ${FAKE_BASE_URL}/reports") + printf '%s' "${FAKE_SUBMISSION_BODY}" > "${output}" + printf '%s' "${FAKE_SUBMISSION_STATUS}" + ;; + "GET ${FAKE_BASE_URL}/reports/"*) + echo "${url}" > "${FAKE_STATE}/read-url" + polls="$(($(cat "${FAKE_STATE}/polls" 2>/dev/null || echo 0) + 1))" + echo "${polls}" > "${FAKE_STATE}/polls" + answer="$(sed -n "${polls}p" "${FAKE_STATE}/answers")" + if [[ "${answer}" == "TRANSPORT" ]]; then + # curl's exit status when the request times out. + exit 28 + fi + printf '%s' "${answer#* }" > "${output}" + printf '%s' "${answer%% *}" + ;; + *) + printf 'unexpected request: %s %s\n' "${method}" "${url}" >&2 + exit 7 + ;; +esac +""" +CREATED = { + "report_id": REPORT_ID, + "report_run_id": REPORT_RUN_ID, + "status": "pending", + "baseline_simulation_id": "00000000-0000-4000-8000-0000000000f1", + "baseline_simulation_run_id": "00000000-0000-4000-8000-0000000000f2", + "baseline_reused": True, + "reform_simulation_id": "00000000-0000-4000-8000-0000000000f3", + "reform_simulation_run_id": "00000000-0000-4000-8000-0000000000f4", + "reform_reused": False, +} + + +def _report(status: str, *, run_id: str = REPORT_RUN_ID, **run_fields) -> str: + return json.dumps( + { + "report_id": REPORT_ID, + "country": "uk", + "baseline_simulation_id": CREATED["baseline_simulation_id"], + "reform_simulation_id": CREATED["reform_simulation_id"], + "run": { + "run_id": run_id, + "status": status, + "country_package_version": "2.174.1", + "policyengine_version": "6.17.1", + "error_message": None, + "created_at": "2026-11-12T15:00:00Z", + "started_at": None, + "completed_at": None, + "baseline_simulation_run_id": CREATED["baseline_simulation_run_id"], + "reform_simulation_run_id": CREATED["reform_simulation_run_id"], + **run_fields, + }, + } + ) + + +@pytest.fixture +def fakes(tmp_path) -> FakeCommands: + commands = FakeCommands.create(tmp_path) + commands.add("curl", FAKE_CURL) + commands.add("sleep") + return commands + + +def _run(fakes, tmp_path, answers: list[str], **overrides): + """Run the check; ``answers`` are the polls' " " in order.""" + state = tmp_path / "state" + state.mkdir() + (state / "answers").write_text("\n".join(answers) + "\n", encoding="utf-8") + environment = { + "SIMULATION_TEST_AUTH_ISSUER": "https://issuer.example/", + "SIMULATION_TEST_AUTH_AUDIENCE": "simulation-entry", + "SIMULATION_TEST_AUTH_CLIENT_ID": "staging-test-client", + "SIMULATION_TEST_AUTH_CLIENT_SECRET": "staging-test-secret", + "FAKE_STATE": str(state), + "FAKE_BASE_URL": BASE_URL, + "FAKE_ACCESS_TOKEN": ACCESS_TOKEN, + "FAKE_SUBMISSION_STATUS": "201", + "FAKE_SUBMISSION_BODY": json.dumps(CREATED), + **overrides, + } + result = run_script(CHECK_SCRIPT, BASE_URL + "/", env=fakes.env(**environment)) + return result, state + + +def test_the_check_creates_a_small_report_and_waits_for_its_run_to_succeed( + fakes, tmp_path +): + result, state = _run( + fakes, + tmp_path, + [ + f"200 {_report('pending')}", + f"200 {_report('running')}", + f"200 {_report('succeeded', completed_at='2026-11-12T15:05:00Z')}", + ], + ) + + assert result.returncode == 0, result.stderr + assert REPORT_ID in result.stdout + assert "reused" in result.stdout + token_request = json.loads((state / "body-1.json").read_text(encoding="utf-8")) + assert token_request["grant_type"] == "client_credentials" + submission = json.loads((state / "body-2.json").read_text(encoding="utf-8")) + assert submission == { + "country": "uk", + "report_type": "economy", + "label": "Staging economy check", + "year": 2026, + "region_code": "country/wales", + "baseline": {"policy_id": None, "dynamic_id": None}, + "reform": {"policy_id": None, "dynamic_id": None}, + "include_cliffs": False, + } + assert (state / "read-url").read_text(encoding="utf-8").strip() == ( + f"{BASE_URL}/reports/{REPORT_ID}" + ) + assert [call[1:] for call in fakes.calls("sleep")] == [["15"], ["15"]] + + +def test_no_secret_reaches_a_command_line(fakes, tmp_path): + """The token is sent from a file, and the client secret on standard input.""" + result, _ = _run(fakes, tmp_path, [f"200 {_report('succeeded')}"]) + + assert result.returncode == 0, result.stderr + arguments = [argument for call in fakes.calls("curl") for argument in call] + assert not any(ACCESS_TOKEN in argument for argument in arguments) + assert not any("staging-test-secret" in argument for argument in arguments) + + +def test_a_failed_run_fails_the_check_with_its_error_message(fakes, tmp_path): + failed = _report( + "failed", error_message="simulation_failed: the reform ran out of memory" + ) + + result, _ = _run(fakes, tmp_path, [f"200 {failed}"]) + + assert result.returncode == 1 + assert "simulation_failed: the reform ran out of memory" in result.stderr + assert fakes.calls("sleep") == [] + + +def test_a_newest_run_other_than_the_created_one_fails_the_check(fakes, tmp_path): + result, _ = _run( + fakes, tmp_path, [f"200 {_report('succeeded', run_id=OLDER_RUN_ID)}"] + ) + + assert result.returncode == 1 + assert OLDER_RUN_ID in result.stderr + + +@pytest.mark.parametrize( + ("status", "body"), + [ + ("400", {"code": "unknown_region", "detail": "no region country/wales"}), + ("503", {"code": "modal_call_refused", "detail": "x", "report_id": REPORT_ID}), + ], +) +def test_a_refused_submission_fails_the_check_with_its_answer( + fakes, tmp_path, status, body +): + result, state = _run( + fakes, + tmp_path, + [], + FAKE_SUBMISSION_STATUS=status, + FAKE_SUBMISSION_BODY=json.dumps(body), + ) + + assert result.returncode == 1 + assert status in result.stderr + assert body["code"] in result.stderr + assert not (state / "polls").exists() + + +@pytest.mark.parametrize("status", ["429", "500", "502", "504"]) +def test_a_read_cloud_run_answers_in_the_services_place_is_asked_again( + fakes, tmp_path, status +): + result, _ = _run(fakes, tmp_path, [f"{status} {{}}", f"200 {_report('succeeded')}"]) + + assert result.returncode == 0, result.stderr + assert status in result.stderr + + +def test_an_unavailable_read_is_asked_again(fakes, tmp_path): + unavailable = json.dumps({"code": "database_unavailable", "detail": "x"}) + + result, _ = _run( + fakes, tmp_path, [f"503 {unavailable}", f"200 {_report('succeeded')}"] + ) + + assert result.returncode == 0, result.stderr + + +def test_a_read_that_gets_no_answer_is_asked_again(fakes, tmp_path): + result, _ = _run( + fakes, tmp_path, ["TRANSPORT", "TRANSPORT", f"200 {_report('succeeded')}"] + ) + + assert result.returncode == 0, result.stderr + assert "no answer" in result.stderr + + +def test_every_request_has_a_time_limit(fakes, tmp_path): + """The submission, which isn't asked again, may take Cloud Run's full minute.""" + result, _ = _run(fakes, tmp_path, [f"200 {_report('succeeded')}"]) + + assert result.returncode == 0, result.stderr + limits = { + call[-1]: call[call.index("--max-time") + 1] for call in fakes.calls("curl") + } + assert limits == { + "https://issuer.example/oauth/token": "30", + f"{BASE_URL}/reports": "60", + f"{BASE_URL}/reports/{REPORT_ID}": "30", + } + + +def test_any_other_read_answer_fails_the_check(fakes, tmp_path): + result, _ = _run(fakes, tmp_path, ['404 {"code": "unknown_report"}']) + + assert result.returncode == 1 + assert "404" in result.stderr + + +@pytest.mark.parametrize( + ("answer", "last_read"), + [ + (f"200 {_report('running')}", "the run was running"), + ('500 {"code": "internal_error", "detail": "x"}', "it answered 500"), + ("TRANSPORT", "it got no answer"), + ], +) +def test_the_deadline_failure_says_what_the_last_read_found( + fakes, tmp_path, answer, last_read +): + result, _ = _run(fakes, tmp_path, [answer], ECONOMY_CHECK_DEADLINE_SECONDS="0") + + assert result.returncode == 1 + assert f"last read: {last_read}" in result.stderr + + +def test_the_deadline_failure_names_the_latest_read_not_the_first(fakes, tmp_path): + """A running run, then the service answering 503 until the deadline. + + Bash counts the deadline in whole seconds, so a 2-second deadline can't be + reached on the first read, whenever its second starts. + """ + unavailable = '503 {"code": "database_unavailable", "detail": "x"}' + answers = [f"200 {_report('running')}", *([unavailable] * 5000)] + + result, _ = _run(fakes, tmp_path, answers, ECONOMY_CHECK_DEADLINE_SECONDS="2") + + assert result.returncode == 1 + assert "(last read: it answered 503)" in result.stderr + + +def test_a_run_that_does_not_finish_in_time_fails_the_check(fakes, tmp_path): + result, _ = _run( + fakes, + tmp_path, + [f"200 {_report('pending')}"], + ECONOMY_CHECK_DEADLINE_SECONDS="0", + ) + + assert result.returncode == 1 + assert "did not finish within 0 seconds" in result.stderr + + +def test_the_check_needs_the_test_client_credentials(fakes, tmp_path): + result, _ = _run(fakes, tmp_path, [], SIMULATION_TEST_AUTH_CLIENT_ID="") + + assert result.returncode == 1 + assert "SIMULATION_TEST_AUTH_CLIENT_ID" in result.stderr + assert fakes.calls("curl") == [] + + +# --- The workflow job -------------------------------------------------------- + + +def _jobs() -> dict: + return workflow_jobs(REUSABLE_DEPLOY_WORKFLOW) + + +def _context(**overrides): + context = routing_context() + context["inputs"]["deployment_environment"] = "staging" + context["inputs"]["deploy_stage12_v2"] = True + context["needs"]["deploy_entrypoint"] = {"result": "success"} + context["needs"]["publish_stage12_v2_default"] = {"result": "success"} + for path, value in overrides.items(): + target = context + *parents, last = path.split("__") + for part in parents: + target = target[part] + target[last] = value + return context + + +def test_the_check_job_creates_the_report_through_the_candidate(): + job = _jobs()[CHECK_JOB] + (step,) = [step for step in job["steps"] if "run" in step] + + assert set(job["needs"]) == { + "deploy_entrypoint", + "authenticated_test", + "publish_stage12_v2_default", + } + assert job["environment"] == "${{ inputs.release_environment }}" + # Longer than the script's own 30-minute deadline, so the script fails first. + assert job["timeout-minutes"] == 40 + assert step["run"] == ( + f".github/scripts/{CHECK_SCRIPT} " + '"${{ needs.deploy_entrypoint.outputs.candidate_url }}"' + ) + assert step["env"] == { + "SIMULATION_TEST_AUTH_ISSUER": ( + "${{ secrets.SIMULATION_ENTRYPOINT_AUTH_ISSUER }}" + ), + "SIMULATION_TEST_AUTH_AUDIENCE": ( + "${{ secrets.SIMULATION_ENTRYPOINT_AUTH_AUDIENCE }}" + ), + "SIMULATION_TEST_AUTH_CLIENT_ID": ( + "${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_ID }}" + ), + "SIMULATION_TEST_AUTH_CLIENT_SECRET": ( + "${{ secrets.SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET }}" + ), + } + + +@pytest.mark.parametrize( + ("overrides", "runs"), + [ + ({}, True), + ({"inputs__deployment_environment": "production"}, False), + ({"needs__deploy_entrypoint__result": "failure"}, False), + ({"needs__authenticated_test__result": "failure"}, False), + ({"needs__publish_stage12_v2_default__result": "failure"}, False), + ({"needs__publish_stage12_v2_default__result": "skipped"}, False), + ( + { + "inputs__deploy_stage12_v2": False, + "needs__publish_stage12_v2_default__result": "skipped", + }, + True, + ), + ({"cancelled": True}, False), + ], + ids=[ + "staging", + "production", + "entrypoint-failed", + "authenticated-tests-failed", + "default-move-failed", + "default-move-skipped", + "no-stage12", + "cancelled", + ], +) +def test_the_check_runs_on_every_staging_deploy_after_the_default_moves( + overrides, runs +): + job = _jobs()[CHECK_JOB] + + assert ( + bool(evaluate_github_expression(job["if"], _context(**overrides), job_if=True)) + is runs + ) + + +def test_deployment_ready_requires_the_staging_economy_check(): + ready = _jobs()["deployment_ready"] + (step,) = [ + step + for step in ready["steps"] + if step.get("name") == "Check the staging economy check result" + ] + + assert CHECK_JOB in ready["needs"] + assert step["run"] == f".github/scripts/{READY_SCRIPT}" + assert step["env"] == { + "REQUESTED": "${{ inputs.deployment_environment == 'staging' }}", + "CHECK_RESULT": "${{ needs.staging_economy_check.result }}", + } + + +@pytest.mark.parametrize( + ("requested", "check_result", "passes"), + [ + ("true", "success", True), + ("true", "failure", False), + ("true", "skipped", False), + ("true", "cancelled", False), + ("false", "skipped", True), + ("false", "success", False), + ], +) +def test_the_ready_check_requires_the_check_exactly_in_staging( + fakes, requested, check_result, passes +): + result = run_script( + READY_SCRIPT, + env=fakes.env(REQUESTED=requested, CHECK_RESULT=check_result), + ) + + assert result.returncode == (0 if passes else 1)