staging-feature.md 7.4 KB

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

Status: planned (spec cut 2026-09-27; nothing implemented yet). 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/)

  • 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)

  • 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)

  • Staging roots 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 (and their ledger entries closed) when staging a new lane would cross it. sm121 staging list/evict for visibility.

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 watermark + sm121 staging evict

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.