staging-feature.md 13 KB

Staging as a first-class sm121 feature — design & plan

Status: P1 + P2 implemented (packed-bin and sparse-copy) (2026-09-27, src/staging/ — hub-mirror + sparse-rows (+format = "sparse-copy" for the aiden16 DSV41_ENGRAM_DIR safetensors-pread layout) + ledger + --stage; P3/P4 planned). Spec cut 2026-09-27. Motivation: weight/engram localization of model data onto node-local NVMe is now what makes the measured difference (EXL3 2x: 13.3 → 35 t/s decode, pp 112 → 795 t/s; Aiden SGLang docs: NFS engram reads cost 6–8 ms/step against 2.8–3 ms local). So far we did it by hand — per-node staging scripts, hooks re-deriving paths, out-of-repo Python — and it is fragile, model-specific, and incompatible with how we actually work: we switch recipes (models) constantly, so "run the staging step once by hand" is not an option.

Requirements (from the operator):

  1. A recipe declares what it needs staged; the operator runs plain sm121 serve --recipe <x> and the launcher stages, verifies, and mounts. No manual prep steps, no per-recipe bespoke scripts.
  2. It must be idempotent and cheap to re-run (skip-complete-by-content, not by mtime — NAS st_mtime is stale over NFS; hardlink-dedup between stores, never silently relink by size alone).
  3. Model switching is the normal case: mirrors coexist per repo (models--… dirs are naturally keyed), and disk pressure needs an eviction policy, not a manual cleanup.
  4. Failures must be loud: verify-or-abort with a clear reason, per node.

Recipe API sketch

# recipes/<lane>/recipe.toml
[[staging]]
id = "exl3-weights"            # ledger key; multiple blocks allowed
kind = "hub-mirror"           # | "sparse-rows" (below)
src = "/mnt/nas/ai/huggingface/hub"       # canonical hub root (NAS)
dst = "~/.cache/hf-local/hub"             # local hub root ("~" per node user)
repos = ["models--coolbho3k--DeepSeek-V4.1-Flash-EXL3-3bpw",
         "models--coolbho3k--DeepSeek-V4.1-Flash-DSpark-EXL3-3bpw"]
skip = ["engram-page15-v1"]               # subpath components to exclude
verify = "sample-hash"                    # size + sampled content per >1G file
min_free = "420G"                         # per-node preflight, else abort

[[staging]]
id = "aiden-engram"
kind = "sparse-rows"                      # per-rank contiguous row ranges
src = "{{model_dir}}"                     # checkpoint with index.json+shards
dst = "/var/tmp/engram-local"
tables = [ { name = "layers.1.engram.embed",  rows = "auto" },   # from safetensors header
           { name = "layers.14.engram.embed", rows = "auto" } ]
partition = "by-rank"                     # rank = launch node order; N contiguous quarters
emit_env = { SPARK_ENGRAM_DIR = "{{dst}}", SPARK_ENGRAM_ROWS = "auto" }
mount = { host = "{{dst}}", container = "/engram-local", read_only = true }

The launcher's existing repo path-ref resolution (repo = in extra_mounts) prefers dst when the mirror for that repo is staged+verified for this lane; otherwise it resolves through src exactly as today. That is the "locality" rule the EXL3 hook implements as fin() — moved into the launcher (the same shape as cluster_state::verdict being shared between health/web/serve).

Implementation plan

P1 — hub-mirror in the launcher (src/staging/)

Implemented 2026-09-27: src/config/staging.rs (block), src/staging.rs (+ mirror.rs, scripts/hub_mirror.py), serve integration in src/docker/mod.rs (staging pass + cache_dir swap), --stage in cli. The script resolves link targets through the hub store and dedups by content sample; ledger lives in the mirror itself (dst/hub/.sm121- staging-ledger.json) so it travels with the mirror and survives node reinstall of /tmp. EXL3 hook migration is the first real adoption (next session, with a verify re-run).

  • Tree walk (links stay links, real files copied, skip pruned), then materialize every snapshot link target resolved through the hub-level store (the lesson that cost the whole six-generation staging saga: snapshot entry → link string → maybe symlink → hub/blobs/xx/<sha>), then a hardlink pass that dedups stores by content sample, never by size.
  • Idempotency ledger per (id, node) in the run ledger dir; resume-safe.
  • Preflight: min_free, source reachability; parallel across launch nodes; streaming progress via the same poll loops as image sync.
  • Flags: --stage=auto|skip|force (auto = verify ledger, copy what is new; skip = assume staged — CI; force = re-verify everything).
  • Migrate EXL3: recipe drops its staging hook + manifest generator (the 171-line bind manifest becomes a derivation of staged files); recipes/dsv41-flash-exl3-2x keeps verify = "sample-hash".

P2 — sparse-rows (engram staging for Aiden)

Implemented 2026-09-27 in packed-bin form: src/staging/rows.rs + scripts/sparse_rows.py (clean-room; format = runtime read contract: 4096 B header <Q magic, layer, lo, hi, rows, 264> + 264 B interleaved records). Per-node rank comes from the cluster launch rank, tp from the recipe (else node count); row splits derive from the checkpoint header. Ledger per block in the dst dir (size + sampled row hashes vs source). First adoption: dsv41-flash-4x-mia (DSV41_PACKED_DIR=/state/engram, dst under the mia-state bind). NOTE: first staging serve adds ~15 min of packing to launch — beyond the 30-min readiness budget when combined with a cold NAS weight load; use --no-wait on the very first run. Implemented 2026-09-27: format = "sparse-copy" produces the aiden16 DSV41_ENGRAM_DIR layout — sparse node-local shard copies at the original byte offsets holding exactly the rank's row ranges, plus engram-local.json ({"layers": {"<L>": [lo, hi]}}). Sample-hash verification adopts a manually staged tree (tools/engram_local.py output) into the ledger instead of re-copying it; corrupted rows trigger a range rebuild; untracked manual trees that don't verify are rebuilt into the declarative format. E2E (fixture, rank1/tp4): pack, verify-only, adoption, corrupt-rebuild all green.

  • Port of tools/engram_local.py semantics to Rust (pread/pwrite ranges, fadvise-DONTNEED, sparse files, row-verify vs source).
  • Per-rank partition is derived by the launcher (it knows each node's launch rank; table row counts read from the checkpoint header — never hardcoded).
  • Per-node env injection for the emit_env keys (head node gets SPARK_ENGRAM_ROWS; others get their copies) — new capability: staging-scoped per-node env.
  • Migrate Aiden 4x once done; until then Aiden runs with NAS engrams (works, measured cost per step ~+3–5 ms/rank) — acceptable, no manual staging.

P3 — cache policy (model switching): data AND images

  • Staging mirrors stay keyed by repo; multiple lanes may share blobs (hardlinks). LRU watermark per node (staging.max_used, default ~75% of the local disk): LRU repo mirrors evicted (ledger closed, local dirs removed, source untouched) when staging a new lane would cross it.
  • Docker images get the same treatment — the disk incident class is the same (499 G of docker in one incident). Every serve/build touches an image LRU ledger (per node, next to the staging ledger): image id + last-use time + lanes that reference it. Auto-GC runs as part of the serve preflight (same poll-loop style, --images=auto|keep opt-out):
    • never evict: images referenced by any non-archived recipe, any container currently running, anything built today;
    • then LRU-down to the watermark (images.max_used, default ~70% of the docker filesystem; target: only the needed + the recent stay on the node);
    • dangling/untagged layers still cleaned by the existing docker GC path;
    • evicted images are always re-pullable (registry push is what keeps them restorable — image GC leans on the registry go-live).
  • One command for visibility over both stores: sm121 staging list/evict / sm121 images list/evict (or one subcommand with a --store= selector).
  • Ordering note: image LRU + mirror LRU share the eviction budget per node — a serve that needs both (new model: weights staged + image pulled) is the exact case P3 exists for; the preflight prints both plans (what is kept, what is dropped, how much freed) before acting, and serve -d shows them.

P4 — probe-and-go, one command

The "run, gather info, re-run with the right parameters" pattern should never require a second invocation by a human. Precedent that already works: NCCL GID auto-detect, image digest checks, the verified stop poll. Fold the rest into the launch pipeline:

  • reachability/power (already), occupancy (already),
  • per-node GPU/host profile (already probed by doctor — surface the deltas),
  • first-boot KV/context-size probes stay operator-visible, but serve captures the "info" into the run's manifest (bench_artifacts-style) instead of requiring a re-run. Any new auto-discovered value becomes: probe → inject → log in the launch line (like NCCL_IB_GID_INDEX=auto), with the config override as opt-out.

Migration map

today becomes
EXL3 recipe hook (fin() manifest) + /tmp/stage-exl3.py staging.kind="hub-mirror", lane ships declarative block only
Aiden engram rows staged by hand staging.kind="sparse-rows", partition="by-rank"
qwen hf-local staging habits hub-mirror with one repo
manual docker system df / prune between model switches P3 image-LRU watermark + sm121 images evict (auto in serve preflight)

Risks / open questions

  • Sparse-rows ownership model: Aiden's engine treats contiguous partitions per rank; if a future engine maps rows by head buckets instead, the derived ranges fall back to source reads (safe but silent). Mitigation: after sparse staging, serve greps the engine log for the reader's "local vs fallback" line and surfaces it in the launch summary.
  • 205G-class mirrors + 119G hosts: min_free must also guard page-cache headroom (the EXL3 incident), not just filesystem free.
  • Staging on serve adds latency to first boot (~1 h worst case per 400G over 10G NAS; ~6 min/node for the sparse-rows kind at 150 MB/s). serve -d must print the full staging plan (what/where/how much, per node) without touching data, exactly like --clear does.

Work packages

  1. src/staging/{mod,mirror,verify}.rs + ledger glue + CLI flags; unit + dry-run integration tests (≤500-line files!).
  2. Resolution rule in docker/path_ref.rs (mirror-first) + per-file bind generation from staged stores; EXL3 recipe migration + re-bench equal.
  3. sparse_rows.rs (port of engram_local.py) + per-node env emission; Aiden recipe migration + A/B bench (NAS vs local engram) to keep the number.
  4. Watermark/LRU + staging subcommand.
  5. Probe-and-go consolidation.

kind = "local-view" — keep the hot side table on local NVMe

Measured 2026-10-23 on the 1017 lane: a single 250-token completion made the shared NFS mount serve +258 MB of reads — the Engram n-gram tables are streamed at ~1 MB per token, straight off the NAS. That is why a vLLM lane with the same checkpoint and the same four Sparks showed 605 t/s prefill where the production lane (which stages Engram locally) shows 2,800–3,100.

local-view makes that staging declarative — no manual copy dance:

[[staging]]
kind = "local-view"
repo = "deepseek-ai/DeepSeek-V4.1-Flash-1017"        # default: the recipe's model
dst = "/home/kilork/models/dsv41-1017-local"          # host dir, per node
local = ["model-00047-of-00048.safetensors",          # REAL local copies
         "model-00048-of-00048.safetensors"]
min_free = "230G"                                     # headroom on top of the copies
verify = "sample-hash"                                # content, never mtime/size

The block materializes dst as a view of the snapshot: files matching local become real node-local copies, everything else becomes a symlink to the shared cache's container path (/root/.cache/huggingface/…). serve then mounts the view at /root/.cache/hf-view and rewrites model to /root/.cache/hf-view/<repo>, so the runtime reads the table from NVMe while the bulk of the weights still resolves through the shared mount.

Verification is a per-view ledger of sampled content hashes (dst/.sm121-view-ledger.json), the same discipline as hub-mirror: a second serve of an unchanged repo costs zero copies, a same-size mutation is detected and re-copied, and a hand-staged tree is adopted instead of re-copied.

Free-space preflight (all staging kinds)

serve now runs staging::preflight before any copy and before the first container starts. For every local-view block it asks each launch node how many bytes the local copies need and how much the target filesystem has (the node script's measure mode, which also honours min_free headroom), and aborts with a message that names the node, the repo, the size and the target:

not enough free space on the stage target(s) — nothing was started:
  dgx7.lan: deepseek-ai/DeepSeek-V4.1-Flash-1017 needs 203.0 GiB in
  /home/kilork/models/dsv41-1017-local, has 167.0 GiB

This is deliberately a pre-start gate: dying halfway through a 200 GB copy (or starting four containers that then find a half-populated directory) is much worse than a refusal that says exactly what to free.