[anth-volk]

Fixes #884

Rebased on main after #879 merged.

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

Stage 12 per-state datasets

Stage 12 gets its own copy of #879's mechanism. Nothing in it imports the main path's precompute modules.

1. Content key stage12-sds1 and folder layout. The key is built from the regional dataset's entry in the published Stage 12 bundle, plus the state and filter. Files go in state-datasets/<dataset>/<STATE> under the Stage 12 data folder.

2. Plan. state_dataset_targets() resolves each state region the way a Stage 12 worker would resolve it. There is one entry per year and state.

3. Build. build_state_datasets_impl:

   • re-derives each key from the installed bundle, and refuses a state without a row filter;
   • prepares the full-country year file once;
   • refuses data in which a household's district code names a registered district of another state;
   • filters each state with the region's own scoping strategy;
   • copies the version stamp, and checks that .py would reuse the file;
   • uploads the file to the write-once Stage 12 cache.

4. Precompute run. build_state_datasets runs with 8 CPUs, 96 GiB and a two-hour limit, one call per year and dataset. The run waits for it only before publishing the manifest. For each build it logs the files built and uploaded, the bytes uploaded and the time taken.

5. Image fetch and worker check. Every entry is checked before anything downloads:

   • unknown entry types are refused;
   • every year must have files for the same set of states;
   • folders and file names must be safe;
   • each file may be listed once, identified by folder plus file name.

   The fetch timeout rises from 1,800 s to 3,600 s, matching the main path. The worker check expects 63 files plus three per state that uses a non-default dataset. The cache release it reports is unchanged.

6. Manifest. The manifest lists state_dataset entries. A plan with no state files publishes exactly the same manifest as before.

7. Runtime. Stage 12 state and district calculations read their state's baked file when it exists and .py would reuse it. Otherwise they prepare the dataset as before, including when an upgrade has moved the private .py helper the reuse checks import.

Dataset provenance and data_version

8. Main path data_version. US state and district runs now report the loaded dataset's revision (the ACS-local revision) as data_version, not the national release. National runs are unchanged.
9. Published bundle.
   • Stage12CountryBundle.region_datasets names the dataset each region type loads. A templated path such as states/{state_code}.h5 maps to none. A path that names no bundle dataset is refused.
   • region_type_of and Stage12CountryBundle.dataset_for_region read region codes the way the runtime normalises them. The worker and the entry adapter both use them.
   • The bundle digest leaves region_datasets out, so the installed 6.7.5 bundle keeps the digest it has on main.
   • The v2 version manifest is written under active-with-region-datasets. The keys active-with-cache and active keep their current shapes for entry services already running, whose models forbid unknown fields. Loaders read the newest key.
10. Workers.
    • They accept the dataset their region type loads, and, during the rollout, the default dataset that entry services not yet updated still send.
    • Any other dataset is refused.
    • After resolving the dataset, a worker requires the provenance to name it, by name and URI. The only exception is that rollout case.
11. Entry adapter. State and district requests send the regional dataset's provenance and population. National and UK requests, and workers whose manifest predates region datasets, are unchanged.

Deploy effects.

• Stage 12 cache. The bundle digest is unchanged, so existing Stage 12 cache files stay valid. The first deploy builds only the 153 state files (51 states × 3 years) in each environment.
• Deploy order. It is safe in both directions:
  • workers accept both kinds of provenance;
  • entry services send regional provenance only to workers whose manifest lists region datasets;
  • the digest that entry services send still matches a redeployed worker.
• Qualification during the rollout. A qualification of a state report whose provenance still names the default dataset (from an entry service not yet updated) reports a data_version difference, because the main path now reports the ACS-local revision. This stops once the entry service is updated.
• Rollback. Code older than this PR republishes only active-with-cache and active, leaving the active-with-region-datasets copy out of date. Readers and the publisher use that copy only while it equals the active-with-cache manifest once its region datasets are removed. After the older code's publish, entry services stop sending regional provenance, and the next publish keeps the versions the older code added.
  • Rollback window: the older worker is deployed before its publish job runs. Until then, the copy still matches, so updated entry services keep sending regional provenance for state and district runs, and the older worker refuses those Stage 12 comparison runs. Production responses are unaffected. If the older deploy's validation fails, the window lasts until a later publish.
  • To avoid it: delete the active-with-region-datasets key from the simulation-api-v2-version-manifest Modal Dict before deploying the older code. The migration document describes this.

Follow-up: once every entry service sends regional provenance, remove the workers' acceptance of the default dataset for state runs.

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.
• Suites:
  • executor (uv run pytest tests/): 1,753 passed, 22 skipped;
  • entry service: 662 passed;
  • contract library: 274 passed;
  • gateway: 260 passed.
• Real installed bundle (6.7.5). Tests against it show:
  • 153 planned state files, all from the ACS-local dataset;
  • state and district regions mapped to populace_us_2024_acs_local;
  • matching dataset URIs in the release bundle and the Stage 12 bundle.
• Code review. Two rounds of review-code, each finding fixed test-first.
  • Round 1:
    • the bundle digest kept stable;
    • region codes read once, matching the runtime;
    • the fetch timeout raised;
    • the Stage 12 copies of Bake per-state ACS-local year files into the main path's worker image #879's fixes;
    • tests for the cache release report and for UK provenance;
    • the migration document updated.
  • Round 2:
    • the region dataset manifest copy is read only while it is current, making rollback safe;
    • a test that the contract and the runtime read region codes alike;
    • the default dataset is kept when another dataset shares its file;
    • the digest is pinned;
    • the Stage 12 copies of Bake per-state ACS-local year files into the main path's worker image #879's second-round fixes.
  • Round 3: found no bugs. The rollback window and its avoiding step are documented, and an outdated comment and part of the migration document are corrected.
• Quality checks. python scripts/run_quality_guards.py passed. ruff check . and ruff format --check with ruff 0.15.22 are clean.
• Pyright. python -m scripts.rules pyright passes for the executor, the contract library and the entry service, with the executor's errors matching main's baseline (77). The last commit passes isin a list in Stage 12's district-code check. No rules/ file changes.
• Not run locally: the Modal image smoke, the Docker build and the local integration tests. CI runs them.

🤖 Generated with Claude Code