فهرست منبع

mirror: verify copied repo completeness by content (entry count + du -sb vs source), not just refs/main presence — a failed cp -a (GLM: 41/42 shards after the copy died) left a partial mirror that the old refs/main gate would trust and serve; now the gate re-verifies and re-copies the partial dir, and a post-copy mismatch fails loudly (MIRROR_VERIFY_FAILED) instead of silently serving a broken tree

Alexander Korolev 6 روز پیش
والد
کامیت
59291f31da
1فایلهای تغییر یافته به همراه51 افزوده شده و 10 حذف شده
  1. 51 10
      src/docker/mirror.rs

+ 51 - 10
src/docker/mirror.rs

@@ -66,21 +66,48 @@ pub fn repos_for_recipe(recipe: &Recipe) -> Vec<String> {
 
 /// Remote command that ensures one repo's HF hub dir exists in the mirror,
 /// copying it from the shared cache when missing. Prints exactly one marker
-/// line: `MIRROR_PRESENT`, `MIRROR_COPIED`, or `MIRROR_MISSING_SOURCE`.
-/// A broken/partial mirror dir (`snapshots`/`refs/main` missing) is replaced.
+/// line: `MIRROR_PRESENT`, `MIRROR_COPIED`, `MIRROR_VERIFY_FAILED`, or
+/// `MIRROR_MISSING_SOURCE`.
+///
+/// Completeness is verified by CONTENT, not by a marker file: the mirror dir
+/// must match the source in both entry count (`find | wc -l` — files,
+/// symlinks and dirs) and apparent size (`du -sb`). A partial copy (a killed
+/// or failed `cp -a` leaves the hub layout with `refs/main` already in place
+/// long before the last shard lands) therefore fails the gate, gets
+/// `rm -rf`ed and re-copied on the next serve instead of being served with a
+/// missing weight shard.
 pub fn mirror_copy_command(source: &Path, mirror: &Path, repo: &str) -> String {
     let dir = model_cache_dir_name(repo);
     let src_dir = format!("{}/hub/{}", source.display(), dir);
     let dst_dir = format!("{}/hub/{}", mirror.display(), dir);
     let hub = format!("{}/hub", mirror.display());
+    let (d, r, s, h) = (
+        shq(&dst_dir),
+        shq(&format!("{dst_dir}/refs/main")),
+        shq(&src_dir),
+        shq(&hub),
+    );
+    // Same-tree comparison: entry count + apparent size. Quoted via shq; the
+    // arithmetic is inside command substitutions the remote shell expands.
+    let same = |d: &str, s: &str| {
+        format!(
+            "test \"$(find {d} | wc -l)\" = \"$(find {s} | wc -l)\" && \
+             test \"$(du -sb {d} | cut -f1)\" = \"$(du -sb {s} | cut -f1)\""
+        )
+    };
+    let present = same(&d, &s);
+    let copied = same(&d, &s);
     format!(
-        "if test -d {d} && test -f {r}; then echo MIRROR_PRESENT; \
-         elif test -d {s}; then rm -rf {d}; mkdir -p {hub}; cp -a {s} {hub}/. && echo MIRROR_COPIED; \
+        "if test -d {d} && test -f {r} && {present}; then echo MIRROR_PRESENT; \
+         elif test -d {s}; then rm -rf {d}; mkdir -p {h}; \
+         if cp -a {s} {h}/. && {copied}; then echo MIRROR_COPIED; else echo MIRROR_VERIFY_FAILED; fi; \
          else echo MIRROR_MISSING_SOURCE; fi",
-        d = shq(&dst_dir),
-        r = shq(&format!("{dst_dir}/refs/main")),
-        s = shq(&src_dir),
-        hub = shq(&hub),
+        d = d,
+        r = r,
+        s = s,
+        h = h,
+        present = present,
+        copied = copied,
     )
 }
 
@@ -127,6 +154,11 @@ pub async fn ensure_mirrors(
                     "MIRROR_COPIED" => {
                         info!("{}: copied {} into mirror {}", node.name, repo, mirror.display())
                     }
+                    "MIRROR_VERIFY_FAILED" => bail!(
+                        "{}: mirror copy of {} is incomplete (size/count mismatch vs the shared cache) — it will be re-copied on the next serve",
+                        node.name,
+                        repo
+                    ),
                     "MIRROR_MISSING_SOURCE" => missing.push(repo.clone()),
                     other => warn!("{}: mirror status for {}: {:?}", node.name, repo, other),
                 }
@@ -316,15 +348,24 @@ mod tests {
             Path::new("/home/u/.cache/hf-models"),
             "org/main",
         );
-        // All three markers present, in order.
+        // All four markers present.
         assert!(c.contains("MIRROR_PRESENT"));
         assert!(c.contains("MIRROR_COPIED"));
+        assert!(c.contains("MIRROR_VERIFY_FAILED"));
         assert!(c.contains("MIRROR_MISSING_SOURCE"));
         // Hub layout + model dir name.
         assert!(c.contains("models--org--main"));
         assert!(c.contains("/home/u/.cache/huggingface/hub/"));
         assert!(c.contains("/home/u/.cache/hf-models/hub/"));
-        // refs/main presence gate.
+        // Content verification (entry count + apparent size) runs both in the
+        // presence gate and after the copy (2 checks × 2 sites) — a partial
+        // copy (refs/main already on disk, last shard missing) must not be
+        // trusted.
+        assert_eq!(c.matches("find ").count(), 4);
+        assert_eq!(c.matches("du -sb").count(), 4);
+        // The partial dir is replaced before re-copying.
+        assert!(c.contains("rm -rf"));
+        // refs/main is still part of the presence gate.
         assert!(c.contains("refs/main"));
         // Paths are shell-quoted.
         assert!(c.contains("'/home/u/.cache/hf-models/hub/models--org--main'"));