Node-local env staging: pack envs to single images, extract to local disk at spawn

[OwenPriceSkelly]

Closes #180. Subsumes #179 (the rootstock stage CLI is the one-shot job-prologue warmer).

Why

Cold worker starts on network filesystems (worst on Delta's HDD-backed Lustre, #160/#167) are bounded by two costs the page-cache prewarm (#168/#171/#178) structurally can't remove:

• metadata RPCs — every cold stat of a 35k-file env tree is a round trip, and
• warmth eviction — cached pages are rented, not owned (cgroup limits, memory pressure), and a resubmitted job lands on a cold node anyway.

The 2026-07-29/30 field data (3–4 concurrent prewarms at 946–1322 s/env; 8 of 45 sync verifies blowing the 600 s connect timeout on freshly rebuilt envs) motivated the prior-art survey, which landed squarely on this shape: single-image env + extract to node-local disk (CSC Tykky, NERSC Shifter, OLCF sbcast+NVMe, Meta XAR).

What

Pack — install now archives each env it builds into {root}/images/<env>-<sha12>.tar.zst (env tree + its .python/ interpreter, root-relative paths, streamed tar | zstd with in-flight hashing), inside the same transaction, best-effort (--no-pack to skip; a failed pack degrades spawns to the prewarm path). rootstock pack backfills existing envs. Manifest schema bumps to v7 with an optional per-env image record; currency is packed_at >= built_at, stamped in the same refresh as built_at, so a rebuild without a repack can never serve a stale image.

Stage — at spawn, worker spawns resolve a node-local base (ROOTSTOCK_STAGE_DIR > layout.json stage_dir > Cluster.stage_dir registry fallback; ROOTSTOCK_NO_STAGE=1 disables; the dir must exist, be writable, and be a different filesystem than the root), then:

• extract the image content-addressed by archive sha — O_EXCL lock so one extractor runs per node while concurrent same-env spawns wait for the atomic rename (the committee-demo contention case), disk preflight with LRU eviction (age-gated), dead-partial sweeping;
• fix up the venv (bin/python* symlink targets, pyvenv.cfg home =) so the interpreter and stdlib are local too — targets are written pre-rename in final-path terms, so the rename stays the atomic publication point;
• overlay the checkpoint's recorded weight files (Track where each checkpoint's weights live: record weight files in the manifest at add/verify time #177) into a local cache mirror: recorded files copied, everything else symlinked back to the shared cache so unrecorded side files (HF config.json — capture only sees mmap'd files) still resolve; HF-hub snapshots//refs/ are recreated so their relative links land on the local blobs;
• repoint HOME/XDG_CACHE_HOME/HF_HOME at the mirror only when the overlay succeeds, and set ROOTSTOCK_NO_PREWARM=1 only when nothing shared remains to warm.

Safety boundaries: download spawns never stage (they must write the shared cache); weights_capture spawns (add/verify) stage the env but keep caches shared, so captured records can't be poisoned with mirror paths — while first-verify-after-rebuild still gets the fast path that the 600 s timeout data asked for. Graceful degradation is load-bearing: any missing piece (no config, no image, stale image, no zstd, no space, lock timeout, any error) falls back to today's prewarm, and an unconfigured cluster behaves exactly as before.

Config/CLI: rootstock init --stage-dir '$SLURM_TMPDIR' writes the declaration to layout.json (env vars expand on the node at spawn time); rootstock stage <checkpoint-id> ... warms envs+weights up front in job prologues and degrades to a sequential prewarm pass where staging isn't configured. Docs in cluster-setup.md.

Testing

35 new stage/pack tests (real tar|zstd round trips, venv fixups, warm reuse, content-addressed rebuild-beside, every fallback edge, the HF-hub overlay indirection, spawn integration gates) plus manifest v7 migration/round-trip/currency coverage. Full suite: 997 passed; ruff/format/ty clean.

Follow-ups (not in this PR)

• Orphaned-image GC in prune (an env pruned after packing leaves its image behind; repacks already clean their own superseded images).
• Per-cluster stage_dir recon + rollout, Delta first (node-local path + wipe policy still unconfirmed there), then A/B against the prewarm path with ROOTSTOCK_NO_STAGE — the metric is the multi-worker congested-afternoon case.

🤖 Generated with Claude Code

[OwenPriceSkelly]

Code-review follow-up status — all 16 verified findings (top-10 report + 6 overflow) are applied: overflow six in 128610d, top-10 in 8b5c7bc. Per-finding outcomes for the top-10, most-severe first:

1. Fixed — _mirror_current treats a fallthrough symlink as never-current; _copy_file_atomic's rename-over replaces the link. The completion marker can't mask it: post-fix, markers are only written after an overlay whose currency checks were symlink-aware.
2. Fixed — staged envs carry .users/<pid> markers written by the client process at every stage_env success; _evict_lru skips dirs with a live registered pid (dead pidfiles pruned in the scan), min-age kept as fallback shield.
3. Fixed — {base}/rootstock is created explicitly and chmod'd sticky-1777 (usage-spool recipe) before the 0700 per-user leaf.
4. Fixed — pack's partial sweep is pid+age aware; the finished-image rename failure surfaces as PackError. _pid_alive moved to pack.py (stage imports it).
5. Fixed — the bare-sweep filter is _image_usable = current-by-timestamp AND archive on disk; a purged images/ dir reads as needs-repack.
6. Fixed — both pack_environments modes update the manifest for successful packs before raising an aggregated failure error.
7. Fixed — _remap_into_stage resolves absolute targets against the resolved root (install-time alias spellings); deterministic fixup failures raise _FixupError and are cached per archive+client-version in {sha}.failed; transient extraction failures deliberately not cached.
8. Fixed — stage_env/stage_weights are never-raise wrappers (preamble failures degrade to prewarm with a log line); _sweep_partials stat guarded against rename races.
9. Fixed — superseded-image cleanup matches ^<env>-[0-9a-f]{12}\.tar\.zst$ exactly; demo-tuned-* survives packing demo.
10. Fixed — plan_prune gains an image GC tier: archives no manifest env records (and crashed packs' partials) are collected, age-guarded like the other tiers.

Verification: 1010 tests passing (11 new), ruff/format/ty clean.

🤖 Generated with Claude Code