|
|
@@ -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'"));
|