Make target_compilation survivable across worker restarts (#299 prereq 1; implements #217)

[MaxGhenis]

Summary

Makes tools/build_us_fiscal_refresh_release.py target_compilation survivable across worker restarts, so a preempted or restarted Build E run reuses already-completed JCT reform materializations instead of recomputing all of them. Two independent, code-confirmed fragilities let Build B silently lose its per-reform cache and forced full recomputation; this PR closes both, adds the finer reform-vector cache key from #217, and documents a rigorous peak-memory bound for Modal box sizing.

Prerequisite 1 of #299 ("fix target_compilation OOM first"). Also implements #217 (reusable materialization checkpoints across calibration-only runs).

Root cause (code-confirmed)

Per-reform materialization checkpointing already existed and is robust in isolation: _read_reform_income_tax_cache / _write_reform_income_tax_cache write an atomic (os.replace), sha256-verified, shape-checked JSON+.npy entry per reform; the reform loop reads at loop top, computes on miss, writes immediately after each reform, and dels + GCs the dense vector between reforms. A kill after 2 of 5 reforms leaves 2 valid entries and a restart should resume at reform 3.

Build B still lost the whole cache because of two compounding, code-level fragilities:

1. Ephemeral out-dir binding (primary). target_materialization_cache_dir and target_frame_checkpoint_path defaulted to <out>/artifacts/…. When Modal preempted B and auto-restarted into a fresh --out, the default cache path moved with it, orphaning the two completed entries on the dead ephemeral dir — so all reforms recomputed. The --target-materialization-cache-dir / --target-frame-checkpoint overrides existed but nothing forced a stable location; one missing flag (or out-dir templating) reproduces B exactly.

2. Over-coarse reform-vector key (Make target materialization checkpoints reusable across calibration-only runs #217, compounding). The per-reform cache identity included build_commit + target_registry_version. A restart that re-resolves HEAD to a newer commit, or any registry bump, invalidated all reform entries even though the reforms and PE-US were unchanged.

What this PR changes

Fix A — durable checkpoint root, decoupled from the ephemeral out-dir.
New --checkpoint-root PATH. When set, the per-reform materialization cache and the full target-frame checkpoint default to <checkpoint-root>/target_materialization_cache and <checkpoint-root>/target_frame_checkpoint.h5 instead of under <out>/artifacts. Explicit --target-materialization-cache-dir / --target-frame-checkpoint overrides and the --no-* disable flags keep their existing precedence. Point --checkpoint-root at a persistent Modal volume that is stable across worker restarts and a fresh --out can no longer orphan completed reforms. The path resolution is factored into a pure, unit-testable helper _resolve_checkpoint_paths(args, artifact_root=…).

Fix B — finer reform-vector cache key (#217).
The per-reform income-tax cache identity now projects the build context onto only the inputs that determine a per-household reform estimate:
(base_dataset_sha256, policyengine_us_version, target_period, congressional_district_vintage_crosswalk_sha256) plus the reform's (measure, neutralized_variable) and n_households. build_commit, seed, and target_registry_version are deliberately excluded (see REFORM_VECTOR_CACHE_CONTEXT_KEYS). A restart at a newer commit, or a calibration-only / registry-only rerun, now reuses completed reforms; a change to base-H5, reform vector, PE-US version, period, or geography still invalidates (no stale poisoning). The cache schema version is bumped 1 → 2 so old coarse-key entries live under different filenames and never collide. The whole-frame checkpoint keeps its own separate, coarser identity (_target_frame_checkpoint_identity, unchanged) — safe because it already gates on target-surface / registry identity independently. Diagnostics still record the per-reform checkpoint key and hit/miss/write status for every reform via compilation["target_materialization_cache"], satisfying #217 AC4.

Satisfies #217 acceptance criteria 1–4:

• AC1 (calibration-only rerun does not rerun reforms): test__given_only_build_commit_changed__then_reform_cache_is_reused.
• AC2 (target-surface semantic change still invalidates the frame/matrix): unchanged whole-frame key already gates on this.
• AC3 (two different H5s never share incompatible per-household vectors): test__given_changed_frame_identity__then_stale_checkpoint_is_not_reused.
• AC4 (diagnostics record which checkpoint keys were used): cache_stats["entries"] flows into release calibration diagnostics.

Fix C — memory bound (analysis; no base-sim rewrite). See below.

Fix D — TDD on a tiny synthetic frame. Four new tests over a 2-household / 3-tax-unit synthetic (fully faked PE-US, no data download):

• test__given_kill_after_two_reforms__then_restart_only_recomputes_the_third — raise mid reform 3, assert exactly 2 durable entries on disk, restart, assert reforms 1–2 load from cache and only reform 3 recomputes (hits==2, misses==1, writes==1), and all three reform vectors are numerically correct.
• test__given_changed_reform_vector__then_stale_checkpoint_is_not_reused — same measure name, different neutralized variable ⇒ recompute (no stale reuse).
• test__given_changed_frame_identity__then_stale_checkpoint_is_not_reused — different base-H5 sha ⇒ recompute.
• test__given_only_build_commit_changed__then_reform_cache_is_reused — only build_commit/seed/registry change ⇒ pure cache hit, zero reform sims.

Memory-bound analysis (for Modal box sizing)

The peak-memory question in prong C resolves per-component. target_compilation = base-sim materialization + the household-batched reform loop, and it does not hold the calibration matrix (a later, separately-timed stage). Sizing the pool from the Build B/D failure diagnosis (224,026 households, ~2.5 persons/hh ⇒ ~0.5–0.6M persons; PE-US ships 5,648 variables, measured):

• Base-sim materialization (one unbatched Microsimulation over the full pool + full entity-table copy + top-level variable caches held to teardown): the PE-US holder cache dominates. Its hard ceiling — every one of the 5,648 PE-US variables cached simultaneously at person scale — is 5,648 × 605k × 8 B ≈ 25 GiB, and that ceiling is physically unreachable (most variables are unrelated benefit programs that income_tax's closure never touches). Realistic footprint for income_tax + the ~dozen top-level vars is ~1,000–2,000 distinct arrays of mixed scale ≈ 5–9 GiB.
• Reform loop: the base sim is torn down (_invalidate_all_caches() + del + GC) before the loop begins, so peaks do not stack. Each iteration builds one Microsimulation over batch_size households; at the default batch=5,000 (~13.5k persons) the per-batch ceiling is < 0.6 GiB, plus the resident output tables (~0.25 GiB, constant). Only ≤1 reform's dense vector is held at a time.
• Conclusion: the base-sim + reform-loop peak is bounded at ~25 GiB (hard ceiling), ~5–10 GiB realistic — it is not a 220 GiB risk. The endorsed decision to not batch/rewrite the base-sim in this survivability PR is therefore correct on the merits (batching the hottest, correctness-sensitive path would buy no headroom), but the rationale differs from the working hypothesis: the base-sim is not the plausible OOM locus.
• Where the ~100 GiB actually lives: the only structure of that size in the whole build is the calibration dense loss/estimate matrix, 2 × n_targets × n_households × 8 B = 2 × 32,637 × 224,026 × 8 ≈ 109 GiB, which runs after target_compilation in the calibrate() stage. If D genuinely died inside target_compilation (reform 3 of 5), the cause was most likely a timeout on the single unresumed pass rather than a memory ceiling; if it died at calibration, box sizing must be driven by that 109 GiB matrix plus optimizer working set, not by the base sim.

Box-sizing recommendation for Build E: --checkpoint-root <persistent volume>, commit the volume after each reform so a preemption resumes at the next reform; --maximum-microsim-batch-size 1000 (cheap insurance, negligible cost); target_compilation fits comfortably in ≤32 GiB, but size the box for the calibration stage (≥300 GiB nonpreemptible remains the safe envelope for the 109 GiB dense matrix + solver).

Testing

• uv run pytest (workspace, after uv sync --all-packages): 1406 passed, 51 skipped, exit 0.
• uv run ruff check .: clean.

Notes for the reviewer

• Change vs. the locked design: the design's memory rationale named the base-sim as "the most plausible D OOM locus." The per-component arithmetic here does not support that (25 GiB hard ceiling); the base-sim is not batched (matching the design's decision) but the PR states the corrected bound and points box-sizing at the calibration matrix instead.
• No @-mentions. Does not touch HF upload paths, the calibration solver, or the whole-frame checkpoint identity.

Prerequisite 1 for #299. Implements #217.