Bläddra i källkod

staging local-view: fix the shared-cache symlink root, nest one view per repo, adopt a pre-staged tree

Three bugs found on the first end-to-end run:
- the symlink root was the view's own mount (/root/.cache/hf-view), so every
  non-local file dangled — it must point at the SHARED cache mount
  (/root/.cache/huggingface), which is where those files actually are;
- the view is laid out as dst/<repo>/ so  can be <mount>/<repo>
  (a bare dst mount needs the repo dir inside it);
- with no ledger yet the apply pass re-copied 190 GB instead of hashing what an
  operator had already rsync'd in — it now hashes the existing file and adopts
  it; the  mode likewise subtracts bytes already in place, so a
  correctly staged node is not reported short.

Measured result (1017 TP4, DSpark k5, stock vLLM 0.32): prefill 750 -> 3,957
t/s on a unique 12.6K prompt, shared-mount reads 0 MB during the run (was
+258 MB per short completion), decode c1 52.9 -> 59.7 t/s.
Alexander Korolev 4 timmar sedan
förälder
incheckning
01011f077d
3 ändrade filer med 56 tillägg och 17 borttagningar
  1. 7 0
      specs/tech-architecture/staging-feature.md
  2. 43 15
      src/staging/scripts/local_view.py
  3. 6 2
      src/staging/view.rs

+ 7 - 0
specs/tech-architecture/staging-feature.md

@@ -194,6 +194,13 @@ 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.
 
+**Measured 2026-10-23 (vLLM 0.32 stock, 1017 TP4, DSpark k5):** Engram on the
+shared NAS → prefill **750 t/s**; the same lane with the two Engram shards staged
+locally → **3,957 t/s** on a unique 12.6K-token prompt (5.3×), and the shared
+mount's read counter stays at **0 MB** across the whole measurement (it was
++258 MB per short completion before). Decode c1 went 52.9 → 59.7 t/s. This was
+the whole remaining gap against the production lane (2,800–3,100 t/s prefill).
+
 `local-view` makes that staging declarative — no manual copy dance:
 
 ```toml

+ 43 - 15
src/staging/scripts/local_view.py

@@ -32,7 +32,13 @@ SAMPLE_BYTES = 1 << 20
 
 
 def repo_dir(src_root, repo):
-    return os.path.join(src_root, "models--" + repo.replace("/", "--"))
+    """The repo's cache dir: HF hub layout is `<root>/hub/models--org--name`,
+    but a caller may hand us the hub dir directly."""
+    name = "models--" + repo.replace("/", "--")
+    for cand in (os.path.join(src_root, "hub", name), os.path.join(src_root, name)):
+        if os.path.isdir(cand):
+            return cand
+    return os.path.join(src_root, "hub", name)
 
 
 def snapshot_dir(src_root, repo, revision=None):
@@ -79,13 +85,26 @@ def sample_hash(path):
     return h.hexdigest()[:32]
 
 
-def local_bytes(snapshot, globs):
+def local_bytes(snapshot, globs, dst=None):
+    """Bytes the `local` copies cost — minus whatever is already in place.
+
+    Already-staged nodes must not be reported as short: a file that exists in
+    `dst` as a regular file of the same size counts as done here (the `apply`
+    pass still verifies it by sampled content hash and re-copies on mismatch,
+    so this is only the preflight's arithmetic, not the verification)."""
     need = 0
     files = []
     for rel in list_files(snapshot):
-        if matches(rel, globs):
-            need += os.path.getsize(os.path.join(snapshot, rel))
-            files.append(rel)
+        if not matches(rel, globs):
+            continue
+        files.append(rel)
+        size = os.path.getsize(os.path.join(snapshot, rel))
+        if dst:
+            existing = os.path.join(dst, rel)
+            if os.path.isfile(existing) and not os.path.islink(existing):
+                if os.path.getsize(existing) == size:
+                    continue
+        need += size
     return need, files
 
 
@@ -145,7 +164,8 @@ def main():
     a = ap.parse_args()
 
     snapshot, revision = snapshot_dir(a.src, a.repo)
-    need, files = local_bytes(snapshot, a.local)
+    root = os.path.join(a.dst, a.repo)
+    need, files = local_bytes(snapshot, a.local, root)
     free = free_bytes(a.dst)
     # A preflight may ask for extra headroom on top of the copies themselves.
     floor = parse_size(a.min_free) if a.min_free else 0
@@ -160,14 +180,16 @@ def main():
         print(f"NEED {required} FREE {free} DST {a.dst}", file=sys.stderr)
         return 3
 
-    os.makedirs(a.dst, exist_ok=True)
-    ledger = load_ledger(a.dst)
+    # One view per repo under dst — `model` becomes <container-root>/<repo>.
+    root = os.path.join(a.dst, a.repo)
+    os.makedirs(root, exist_ok=True)
+    ledger = load_ledger(root)
     seen = ledger.setdefault("files", {})
     copied = kept = 0
     for rel in list_files(snapshot):
         src = os.path.join(snapshot, rel)
-        dst = os.path.join(a.dst, rel)
-        os.makedirs(os.path.dirname(dst) or a.dst, exist_ok=True)
+        dst = os.path.join(root, rel)
+        os.makedirs(os.path.dirname(dst) or root, exist_ok=True)
         if not matches(rel, a.local):
             target = container_path(a.container_root, a.repo, revision, rel)
             if os.path.islink(dst) and os.readlink(dst) == target:
@@ -178,10 +200,16 @@ def main():
             seen.pop(rel, None)
             continue
         want = None if a.verify == "presence" else sample_hash(src)
-        have = seen.get(rel)
-        if not a.force and os.path.isfile(dst) and not os.path.islink(dst) and have == want:
-            kept += 1
-            continue
+        if not a.force and os.path.isfile(dst) and not os.path.islink(dst):
+            # No ledger entry (a hand-staged tree, e.g. rsync'd by the operator):
+            # hash what is already there and adopt it instead of re-copying.
+            have = seen.get(rel)
+            if have is None and want is not None:
+                have = sample_hash(dst)
+            if have == want:
+                seen[rel] = have
+                kept += 1
+                continue
         tmp = dst + ".sm121-tmp"
         shutil.copyfile(src, tmp)
         os.replace(tmp, dst)
@@ -190,7 +218,7 @@ def main():
     ledger["repo"] = a.repo
     ledger["revision"] = revision
     ledger["local"] = a.local
-    save_ledger(a.dst, ledger)
+    save_ledger(root, ledger)
     print(f"local-view {a.repo}@{revision[:12]}: {copied} copied, {kept} verified, "
           f"{need} bytes local, {free} free")
     return 0

+ 6 - 2
src/staging/view.rs

@@ -23,6 +23,10 @@ pub const SCRIPT_BODY: &str = include_str!("scripts/local_view.py");
 /// Container-side path the view is mounted at (the model references it).
 pub const CONTAINER_ROOT: &str = "/root/.cache/hf-view";
 
+/// Container path of the *shared* HF cache — where every non-local file's
+/// symlink points, so it resolves inside the container.
+pub const SHARED_CONTAINER_ROOT: &str = "/root/.cache/huggingface";
+
 fn script_path(id: &str) -> String {
     format!("/tmp/sm121-local-view-{id}.py")
 }
@@ -59,7 +63,7 @@ pub fn command_for(id: &str, mode: &str, src: &str, dst: &str, spec: &ViewSpec<'
         crate::docker::shq(src),
         crate::docker::shq(spec.repo),
         crate::docker::shq(dst),
-        crate::docker::shq(CONTAINER_ROOT),
+        crate::docker::shq(SHARED_CONTAINER_ROOT),
         spec.verify,
     );
     for g in spec.local {
@@ -278,7 +282,7 @@ mod tests {
         // shq() shell-quotes arguments, so the glob may come back quoted
         assert!(cmd.contains("model-00047-of-00048.safetensors"), "{cmd}");
         assert!(cmd.contains("--local"), "{cmd}");
-        assert!(cmd.contains(CONTAINER_ROOT), "{cmd}");
+        assert!(cmd.contains(SHARED_CONTAINER_ROOT), "{cmd}");
         assert!(cmd.contains("--min-free '10G'"), "{cmd}");
         assert!(cmd.ends_with("--force"), "{cmd}");
     }