[anth-volk]

Fixes #991

Same-Repository Draft PR Check

• This PR is a draft.
• This PR is from a branch in PolicyEngine/policyengine-sim-api, not a personal fork.

Summary

The v2 Modal report coordinator and its US and UK simulation workers can now record a report's run in the Public API's run tables, with the Modal worker role's own database credentials. When a call's invocation context names a report run and its baseline and reform simulation runs, the coordinator and workers write those runs' status, each simulation's output file and the report's economy results. The report run is marked succeeded in the same transaction as its result rows, so a call that Modal restarts after pre-emption cannot record its result twice. A call that names no report run behaves exactly as today. Nothing sends such a call yet: the simulation service's /simulations and /reports routes, a later change, will.

How a call names its runs (policyengine_simulation_contract)

Stage12InvocationContext.report_run (new, optional): a ReportRunIds with report_run_id, baseline_simulation_run_id and reform_simulation_run_id, three different IDs. A shadow run compared with a production call cannot carry it. Comparison runs leave it unset, and their child calls receive exactly the context fields they received before.

Coexisting with the comparison path

The smallest change: a call that names a report run still runs through the Stage 12 comparison records (stage12_evaluation_reports, stage12_evaluation_simulations) and the Stage 12 result bucket, which coordinate its workers as before, and also records the Public API's runs. Those writes go through a second store (ReportRunStore, in policyengine_stage12_persistence), connected as the Modal worker role.

• Coordinator: checks with one SELECT that the report run names the call's two simulation runs as its baseline and reform, then claims the report run under its own Modal call ID, then its Stage 12 report. The claim takes a pending or running report run only while it records no Modal call or this call's own: the simulation service records the call it starts while the run is still pending, and Modal restarts a pre-empted call under the same ID. A report run that records another call, or that is already finished, makes the call return a deduplicated summary and start nothing. It marks its Stage 12 report succeeded or failed first and its report run second; the report run's success is written in one transaction with its result rows. Any failure, including one claiming the Stage 12 report, fails the report run with a redacted code and summary.
• Workers: a wrapper around the comparison store (SimulationRunRecordingStore) writes each Stage 12 child record first and then mirrors its claim, success and failure onto that child's simulation run: the worker claims its simulation run under its own Modal call ID, records the output file as the run succeeds, and a recorded failure, including the coordinator's cancellation of a child, fails the simulation run. A simulation run that is already failed refuses the worker's claim, which it treats as a cancellation.
• Lost second writes: before it completes or fails a report run, the coordinator gives each of its two simulation runs its child record's outcome, with the same conditional writes, so a write a worker lost (a dropped connection, or a worker stopped between its two writes) is made then, and a report run never succeeds beside a simulation run still running. A worker whose second write fails after its child record succeeded logs it and returns its result.
• Every write is conditional: each status write is one UPDATE … WHERE id = … AND status = … (and, for the report run, AND modal_call_id = …), returning whether it changed the run and the run's status, so a repeated write changes nothing that is final.

Restarts and cancellation

Modal pre-empts a call with SIGINT, received as KeyboardInterrupt, and runs it again under the same call ID; it cancels a call with InputCancellation, and never runs a cancelled call again. Neither is an Exception.

• Pre-emption: a coordinator whose call names a report run cancels its started, unfinished child calls and re-raises, recording no failure. When Modal restarts it, it resumes the Stage 12 claim it holds under that call ID (instead of returning a duplicate), cancels the child calls its earlier attempt started, and starts the children again.
• A restart between the two success writes: the Stage 12 record is written first, so a restarted call finds it finished. A coordinator whose Stage 12 report succeeded under the same call ID reads the stored aggregate, checks its digest, and completes the report run from it without calculating again. A worker whose child record succeeded completes its simulation run from that record, so both name the same output file. A restart after both writes starts nothing.
• Cancellation: a coordinator whose call is cancelled records its unfinished children, its Stage 12 report and its report run as failed (report_cancelled), because nothing would finish them otherwise; so does one cancelled while it claims its Stage 12 report.
• Segments: a segmented US worker that Modal stops now cancels its region-group segment calls, on every path; Modal does not cancel the calls a stopped call spawned. This also applies to comparison runs and changes none of their records.
• Artifacts: such calls write their artifacts under keys that name the content's SHA-256 (inputs/<role>-<sha256>.json, simulations/<role>-<sha256>.parquet, reports/aggregate-<sha256>.json), so a restarted call that calculates different last bits stores them beside its earlier attempt's instead of failing the write-once check. Comparison runs keep their fixed keys.
• Time limits are unchanged.

Results as rows

The report's result is the single-year economy output Stage 12 already builds, in the format the Public API and v3 read today. Each part becomes rows, with the meaning .py's row models give the result tables' columns, in the same transaction as the report run's success:

• budget_summary (entity household): household_tax, household_benefits and, for the US only, household_state_income_tax (change), household_net_income and household_weight (baseline total, the latter the household count); the budgetary impact is tax minus benefits, as the output computes it, and a missing state income tax row is 0.0, as the output reports for the UK.
• program_statistics: the detailed budget, one row per program, with a null entity (the output names none) and the output's is_tax flag.
• decile_impacts: the change in household_net_income by income decile, with a null decile_variable, and (UK) by wealth decile, with decile_variable household_wealth_decile, as .py names them; the relative change is .py's percentage, so the output's fraction is multiplied by 100 (a reader divides again, equal to within floating-point rounding).
• intra_decile_impacts: income and (UK) wealth deciles 1 to 10. Row 0 is .py's population-wide share, which the output doesn't carry: its overall shares are the mean of the deciles, which a reader computes again exactly.
• poverty, per simulation run: the poverty and deep poverty rates under .py's poverty types (spm, spm_deep; relative_bhc, absolute_bhc), for everyone and by .py's groups: filter_variable names the variable that selects a group (age, is_male, race) and filter_group the group (child, adult, senior; male, female; white, black, hispanic, other), both null for everyone. Race covers poverty only, and only in the US, as the output does. Today's output reports 0.0 for a rate .py didn't report and for every rate of a breakdown it failed to calculate, so a stored 0.0 may be such a rate. inequality, per simulation run, under .py's income variables. A simulation run that another report run already stored a row for keeps its row (ON CONFLICT DO NOTHING).
• congressional_district_impacts, with the winner, loser and no-change shares, keyed by the Census state FIPS code and district number (0 for an at-large state), from .py's US state data.
• constituency_impacts and local_authority_impacts, with their disclosure tier and interval bounds.
• labor_supply_responses: one row with the labour supply response's earnings changes from the income and substitution effects and their total, the revenue change, the two relative responses, and the baseline, reform and changed weekly hours with the change split into the two effects. labor_supply_decile_responses: one row per decile, with the income and substitution responses per household and relative to earnings; .py reports relative responses only for deciles above zero, so a decile without one has null relative columns. An output with no labour supply response stores neither.
• On the report run itself, set by the same UPDATE that marks it succeeded: the baseline's and reform's cliff gap and cliff share (four columns, all null when the output has no cliff impact), the SPM receipts (jsonb) and the data version.

Nothing else in the output needs a column: a reader computes the budgetary impact and the overall intra-decile shares again from these rows, and model_version is the report run's country_package_version, which the simulation service sets when it creates the run.

Tests

The store's tests run against a disposable Postgres migrated with the Public API's v2 chain from PolicyEngine/policyengine-api#4162, with that change's declared privileges applied by its own reconcile script, connected as the Modal worker role, whose UPDATE covers only its runs' status, call, error, timing and result columns. That change keeps the foreign keys between the two services' tables, so each test's runs name real catalog rows, simulations and a report. They skip when REPORT_RUN_TEST_DATABASE_URL and REPORT_RUN_TEST_SEED_DATABASE_URL are unset, as they are in CI. They cover the conditional claims, including a pending report run that records another call, completions and failures; a result whose insert fails leaving the run running, with none of its rows and none of its own result columns; a repeated completion storing nothing; a shared simulation run keeping its rows; a round trip of each part of the output, including the UK wealth deciles, the poverty rates by group and the labour supply response by decile; an output without labour supply or cliff impact; a report run that names other simulation runs than its call; and the role being refused an insert into report_runs. The runtime tests use in-memory stand-ins for Modal and both stores, including coordinators pre-empted at each point between their writes and restarted, cancelled coordinators, restarted workers, lost second writes, and a report run through the real worker code.

Environment variables

• MODAL_WORKER_DATABASE_URL (new): read by the coordinator and the US and UK simulation workers, only for a call that names a report run. Provider: the Modal secret policyengine-modal-worker-database, which an operator creates in the staging and main Modal environments (below). Wiring: modal.Secret.from_name(..., required_keys=["MODAL_WORKER_DATABASE_URL"]) attached through secrets= to coordinate_report, run_single_simulation_us and run_single_simulation_uk in src/modal/v2_app.py. Not verified: the secret does not exist yet.
• STAGE12_DATABASE_URL, STAGE12_ARTIFACT_BUCKET (unchanged): the Modal secret stage12-evaluation-runtime, synced from Secret Manager by each deploy's stage12-sync-modal-secrets.sh.
• REPORT_RUN_TEST_DATABASE_URL, REPORT_RUN_TEST_SEED_DATABASE_URL: test-only, set by hand for a local database; nothing sets them in CI.

What this waits for before merging

1. Record each run of a v2 simulation and report, and key results by run policyengine-api#4162 merged and deployed: the run tables and result tables this writes, and the Modal worker role's declared privileges, including its column-level UPDATE.
2. The Modal worker role policyengine_v2_modal_worker created, with a password, in the staging and production v2 databases (an operator step), so the Public API's deploy grants it its privileges.
3. The Modal secret policyengine-modal-worker-database with key MODAL_WORKER_DATABASE_URL in the staging and main Modal environments (an operator step): the URL of that role on that environment's v2 database. Merging deploys the v2 Modal apps with the secret attached, and modal deploy fails without it, so this comes before the merge.
4. An approving review.

The Modal secret (operator steps, for each of staging and main)

• Name policyengine-modal-worker-database, key MODAL_WORKER_DATABASE_URL and no other.

• Value: a postgresql:// URL that logs in as policyengine_v2_modal_worker to database postgres of that environment's v2 database, in the same format as the existing STAGE12_DATABASE_URL: postgresql://policyengine_v2_modal_worker.<project-ref>:<password>@<pooler-host>:5432/postgres?sslmode=require. The scheme must be postgresql://; the code selects the driver.

• Create it from projects/policyengine-simulation-executor, without echoing the value:

  umask 077
  secret_file="$(mktemp)"
  # Write {"MODAL_WORKER_DATABASE_URL": "<the URL>"} into "$secret_file" with an editor.
  uv run --frozen modal secret create policyengine-modal-worker-database --from-json "$secret_file" --env staging
  rm -f "$secret_file"

• Check it without printing it: modal secret list --env staging lists the name; psql "<url>" -Atc 'select current_user' (from a file holding only the URL) prints the role; select count(*) from public.report_runs succeeds and insert into public.report_runs … is refused. The first deploy after merging must pass modal deploy, which checks the required key.

🤖 Generated with Claude Code