[anth-volk]

Fixes #878

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

US state and congressional-district requests no longer download the 9.8 GB ACS-local dataset, or prepare it for every state, on each new worker. The precompute prepares one year file per state for 2025–2027, and the image includes them. Each commit is a separate change:

1. RegionResolution.state_code. Resolving a state or district region now reports the two-letter code of its state. National, UK and region-group resolutions report none. This is a shared helper, and Stage 12 is unaffected.
2. _load_dataset(data_folder=...). Callers can choose the folder a year file is read from. If they don't, the folder is chosen as before. This is a shared helper.
3. Content key sds1 and folder layout.
   • The key covers the regional dataset's name, its revision-pinned URI and certified SHA-256, the year, the state code and FIPS code, the filter variable, the model and .py versions, and the SPM selection.
   • Store path: state-datasets/us/<digest>/<stem>_year_<year>.h5.
   • Image folder: state-datasets/<dataset>/<STATE>.
4. Plan. state_dataset_targets() resolves every state region the way a request for it would be resolved. It keeps only the states whose dataset is not the default. The plan lists one entry per year and state, in a fixed order. A bundle that maps states back to the default dataset plans none.
5. Build. build_state_datasets_impl:
   • re-derives each key inside the container;
   • has .py prepare the full-country year file once, then releases that memory and reloads the one year;
   • refuses data in which a household's district code names a registered district of another state (an unassigned code such as 0 matches no district and is allowed);
   • for each state, filters with the state's own scoping strategy and writes the file with .py's save();
   • copies the stamp, which save() drops, and confirms .py's reuse checks accept the file;
   • uploads the file write-once. Weights are unchanged.
6. Image fetch. A manifest entry with a folder is placed in that folder. Absolute or unsafe folders and filenames are refused before anything downloads, as are repeated files. The fetch timeout rises from 900 s to 3,600 s.
7. Precompute run. build_state_datasets runs with 8 CPUs, 96 GiB and a two-hour limit, one container per year and dataset. It starts alongside the dataset builds. The run waits for it only before publishing the manifest, so baselines never wait. For each build it logs the files built and uploaded, the bytes uploaded and the time taken.
8. Manifest. The manifest gains state_dataset entries, each with folder and state_code. They use their own model, chosen by type. A plan with no per-state files publishes byte for byte the same manifest as before.
9. Runtime. A state or district request reads its state's baked file when the image has it for the requested year and .py's reuse checks accept it (or can be run at all). Otherwise it prepares the dataset during the request, as before, so .py never rebuilds the regional dataset inside a baked folder. The run records dataset_source as baked_state or on_demand.

Each state file holds exactly the rows the run-time state filter would keep from the full-country file, and carries the same stamp and records. Results are therefore unchanged.

First deploy. The precompute builds 153 per-state files (51 states × 3 years), each from a 9.8 GB source, and the image copies in several GB more.

Testing

• Test-first. Each commit was written test-first; every new test was run and seen to fail for the expected reason before the code was written.

• Real .py year files. The per-state build tests write a three-household year file with .py's own save(), version stamp and SPM record. They show that .py's real ensure_datasets reuses a state file without rebuilding it. Re-applying the state filter keeps the same rows, and weights are unchanged.

• Executor suite. uv run pytest tests/ in projects/policyengine-simulation-executor: 1,640 passed, 22 skipped.

• Code review. The review-code review found no bugs. Its points are addressed in the last four commits, each test-first:

  • the district-code check narrowed to codes that name another state's district;
  • builds grouped by year and dataset;
  • the log reports bytes actually uploaded and build times;
  • a district request is tested from region resolution to the folder it loads;
  • each state file is compared with .py's filter applied to the full source data;
  • the target's row filter is typed, and a target without one is refused.

  A second round found no bugs. Its points are addressed in the last two commits:

  • a state without a row filter is refused before the year file is prepared;
  • if a .py upgrade moves the private helper the reuse checks import, requests load on demand instead of failing;
  • docstrings corrected.

• Quality checks. python scripts/run_quality_guards.py: passed.

• Pyright. python -m scripts.rules pyright (pinned 1.1.414) passes for the executor. Main started running Pyright (Run Pyright in CI for every package, in its own locked environment #864) after this branch was written, and the last commit fixes the seven errors the branch added:

  • the run's spans now open with scopes.span(runtime, ...), as docs/engineering/skills/type-checking.md says, and _build_simulation is annotated with the StageReportingSimulation it returns;
  • the precompute app casts its Modal functions to RemoteFunction;
  • the manifest list and the district-code check are typed correctly.

  This also clears 54 existing errors, so the executor baseline drops from 178 to 124 entries. That changes rules/, which needs code-owner approval.

• Ruff 0.15.22 (the version CI pins): ruff check . and ruff format --check projects/*/src libs/*/src are clean.

• Not run locally: the Modal image smoke, the Docker build and the local integration tests. CI runs them.

• Planned after the beta deploy:

  • check the precompute log for 153 uploads and their total size;
  • check the image fetch finishes within its time limit;
  • run Utah and California requests and confirm dataset_source=baked_state with dataset_load in seconds;
  • compare those outputs with the on-demand path.

🤖 Generated with Claude Code