|
|
@@ -0,0 +1,271 @@
|
|
|
+#!/usr/bin/env python3
|
|
|
+"""sm121 hub-mirror staging engine (spec: staging-feature.md, P1).
|
|
|
+
|
|
|
+Mirror the parts of a shared HF hub cache a recipe needs (repo trees + hub
|
|
|
+blobs) onto a node-local hub root, keeping the canonical layout so refs/main,
|
|
|
+snapshot trees, and composed-model symlinks resolve unchanged:
|
|
|
+
|
|
|
+ <src>/hub/models--org--model/... -> <dst>/hub/models--org--model/...
|
|
|
+
|
|
|
+Idempotency is by CONTENT — never by mtime or by size alone. Stale NFS
|
|
|
+`st_mtime` and size-only "already copied" checks silently skip real copies
|
|
|
+(the failure class of the 2026-09-27 staging saga), so a per-root ledger
|
|
|
+records a sampled content hash per staged file and re-checks the sample
|
|
|
+cheaply. Files whose relative path contains a `--skip` component at any depth
|
|
|
+are left alone (the caller's own binds cover them).
|
|
|
+
|
|
|
+Snapshot symlinks are recreated inside the mirror: link targets under `<src>`
|
|
|
+are staged into their mapped `<dst>` location and the link is rewritten with
|
|
|
+the equivalent relative target; targets outside `<src>` are materialized as
|
|
|
+regular files so the mirror is self-contained.
|
|
|
+
|
|
|
+Markers on stdout (consumed by the launcher):
|
|
|
+ STAGED <repo> links=<n> files=<n> verified=<n> copied=<n> secs=<n>
|
|
|
+ STAGING_DONE | STAGING_FAIL <reason> (exit != 0)
|
|
|
+"""
|
|
|
+import argparse
|
|
|
+import hashlib
|
|
|
+import json
|
|
|
+import os
|
|
|
+import stat
|
|
|
+import sys
|
|
|
+import time
|
|
|
+
|
|
|
+SAMPLE_BLOCK = 1 << 20 # 1 MiB per sampled region
|
|
|
+SAMPLE_REGIONS = 16 # spread over the file -> 8 MiB hashed per big file
|
|
|
+
|
|
|
+
|
|
|
+def fail(msg):
|
|
|
+ print(f"STAGING_FAIL {msg}", flush=True)
|
|
|
+ sys.exit(1)
|
|
|
+
|
|
|
+
|
|
|
+def sample_hash(path):
|
|
|
+ """Digest of 16 spread-out 1 MiB reads (size-salted). Small files hash
|
|
|
+ their tail block; this is exactly the check a ledger skip relies on."""
|
|
|
+ size = os.path.getsize(path)
|
|
|
+ h = hashlib.sha256(str(size).encode())
|
|
|
+ if size <= SAMPLE_BLOCK * SAMPLE_REGIONS:
|
|
|
+ off = 0 if size <= SAMPLE_BLOCK else size - SAMPLE_BLOCK
|
|
|
+ if size:
|
|
|
+ with open(path, "rb") as f:
|
|
|
+ f.seek(off)
|
|
|
+ h.update(f.read(SAMPLE_BLOCK))
|
|
|
+ else:
|
|
|
+ step = (size - SAMPLE_BLOCK) / (SAMPLE_REGIONS - 1)
|
|
|
+ off = 0.0
|
|
|
+ with open(path, "rb") as f:
|
|
|
+ for _ in range(SAMPLE_REGIONS):
|
|
|
+ f.seek(int(off))
|
|
|
+ h.update(f.read(SAMPLE_BLOCK))
|
|
|
+ off += step
|
|
|
+ return h.hexdigest()
|
|
|
+
|
|
|
+
|
|
|
+class Ledger:
|
|
|
+ """Content-addressed check ledger: {dst_path: [size, sample_sha256]}."""
|
|
|
+
|
|
|
+ def __init__(self, path):
|
|
|
+ self.path = path
|
|
|
+ self.data = {}
|
|
|
+ try:
|
|
|
+ with open(path) as f:
|
|
|
+ self.data = json.load(f).get("files", {})
|
|
|
+ except Exception:
|
|
|
+ pass
|
|
|
+
|
|
|
+ def get(self, key):
|
|
|
+ v = self.data.get(key)
|
|
|
+ return tuple(v) if v else None
|
|
|
+
|
|
|
+ def put(self, key, size, sample):
|
|
|
+ self.data[key] = [size, sample]
|
|
|
+
|
|
|
+ def save(self):
|
|
|
+ tmp = self.path + ".tmp"
|
|
|
+ os.makedirs(os.path.dirname(self.path), exist_ok=True)
|
|
|
+ with open(tmp, "w") as f:
|
|
|
+ json.dump({"files": self.data, "ts": int(time.time())}, f)
|
|
|
+ os.replace(tmp, self.path)
|
|
|
+
|
|
|
+
|
|
|
+class Stager:
|
|
|
+ def __init__(self, src, dst, skip, verify, force, ledger):
|
|
|
+ self.src = src
|
|
|
+ self.dst = dst
|
|
|
+ self.skip = set(skip)
|
|
|
+ self.verify = verify
|
|
|
+ self.force = force
|
|
|
+ self.ledger = ledger
|
|
|
+ self.links = self.files = self.verified = self.copied = 0
|
|
|
+
|
|
|
+ def skipped(self, rel_parts):
|
|
|
+ return any(p in self.skip for p in rel_parts)
|
|
|
+
|
|
|
+ def copy_into(self, src_path, dst_path):
|
|
|
+ """Copy src_path -> dst_path via tmp, verify the copy, publish."""
|
|
|
+ size = os.path.getsize(src_path)
|
|
|
+ tmp = dst_path + ".sm121tmp"
|
|
|
+ try:
|
|
|
+ if os.path.exists(tmp):
|
|
|
+ os.remove(tmp)
|
|
|
+ buf = memoryview(bytearray(1 << 22))
|
|
|
+ with open(src_path, "rb") as fin, open(tmp, "wb") as fout:
|
|
|
+ while True:
|
|
|
+ n = fin.readinto(buf)
|
|
|
+ if not n:
|
|
|
+ break
|
|
|
+ fout.write(buf[:n])
|
|
|
+ if os.path.getsize(tmp) != size:
|
|
|
+ raise RuntimeError("short copy")
|
|
|
+ if self.verify == "sample-hash" and sample_hash(tmp) != sample_hash(src_path):
|
|
|
+ raise RuntimeError("verify failed after copy")
|
|
|
+ os.replace(tmp, dst_path)
|
|
|
+ finally:
|
|
|
+ if os.path.exists(tmp):
|
|
|
+ os.remove(tmp)
|
|
|
+ self.copied += 1
|
|
|
+
|
|
|
+ def ensure_file(self, src_path, dst_path):
|
|
|
+ """Stage src_path into dst_path unless its content is already proven.
|
|
|
+ Returns True if bytes were copied."""
|
|
|
+ size = os.path.getsize(src_path)
|
|
|
+ if not self.force and os.path.exists(dst_path):
|
|
|
+ entry = self.ledger.get(dst_path)
|
|
|
+ if entry and entry[0] == size and entry[1] == sample_hash(dst_path):
|
|
|
+ self.verified += 1
|
|
|
+ return False
|
|
|
+ if entry is None and os.path.getsize(dst_path) == size \
|
|
|
+ and sample_hash(dst_path) == sample_hash(src_path):
|
|
|
+ # untracked but currently matching source: adopt into ledger
|
|
|
+ self.ledger.put(dst_path, size, sample_hash(dst_path))
|
|
|
+ self.verified += 1
|
|
|
+ return False
|
|
|
+ os.makedirs(os.path.dirname(dst_path), exist_ok=True)
|
|
|
+ self.copy_into(src_path, dst_path)
|
|
|
+ self.ledger.put(dst_path, size, sample_hash(dst_path))
|
|
|
+ self.files += 1
|
|
|
+ return True
|
|
|
+
|
|
|
+ def stage_tree(self, repo_src, repo_dst):
|
|
|
+ """Walk the repo dir (no symlink following); copy files, make dirs.
|
|
|
+ Returns the list of (link_path, rel_from_src_root) entries."""
|
|
|
+ links = []
|
|
|
+ stack = [repo_src]
|
|
|
+ while stack:
|
|
|
+ root = stack.pop()
|
|
|
+ try:
|
|
|
+ names = sorted(os.listdir(root))
|
|
|
+ except OSError as e:
|
|
|
+ fail(f"cannot list {root}: {e}")
|
|
|
+ for name in names:
|
|
|
+ p = os.path.join(root, name)
|
|
|
+ rel = os.path.relpath(p, self.src)
|
|
|
+ parts = rel.split(os.sep)
|
|
|
+ if self.skipped(parts[:-1]):
|
|
|
+ continue
|
|
|
+ st = os.lstat(p)
|
|
|
+ if stat.S_ISLNK(st.st_mode):
|
|
|
+ links.append((p, rel))
|
|
|
+ elif stat.S_ISDIR(st.st_mode):
|
|
|
+ os.makedirs(os.path.join(self.dst, rel), exist_ok=True)
|
|
|
+ stack.append(p)
|
|
|
+ elif stat.S_ISREG(st.st_mode):
|
|
|
+ self.ensure_file(p, os.path.join(self.dst, rel))
|
|
|
+ return links
|
|
|
+
|
|
|
+ def stage_links(self, links):
|
|
|
+ for p, rel in links:
|
|
|
+ target = os.path.realpath(p)
|
|
|
+ if not os.path.exists(target):
|
|
|
+ fail(f"dangling link in source: {p}")
|
|
|
+ mapped_rel = os.path.relpath(target, self.src)
|
|
|
+ dst_entry = os.path.join(self.dst, rel)
|
|
|
+ if mapped_rel.startswith(".."):
|
|
|
+ # target outside the source tree: materialize as a file
|
|
|
+ os.makedirs(os.path.dirname(dst_entry), exist_ok=True)
|
|
|
+ if os.path.islink(dst_entry):
|
|
|
+ os.remove(dst_entry)
|
|
|
+ self.ensure_file(target, dst_entry)
|
|
|
+ continue
|
|
|
+ mapped_dst = os.path.normpath(os.path.join(self.dst, mapped_rel))
|
|
|
+ if os.path.isdir(target):
|
|
|
+ os.makedirs(mapped_dst, exist_ok=True)
|
|
|
+ else:
|
|
|
+ if os.path.abspath(mapped_dst) != os.path.abspath(dst_entry):
|
|
|
+ self.ensure_file(target, mapped_dst)
|
|
|
+ else:
|
|
|
+ self.ensure_file(target, mapped_dst)
|
|
|
+ rel_target = os.path.relpath(mapped_dst, os.path.dirname(dst_entry))
|
|
|
+ os.makedirs(os.path.dirname(dst_entry), exist_ok=True)
|
|
|
+ if os.path.islink(dst_entry):
|
|
|
+ if os.readlink(dst_entry) == rel_target:
|
|
|
+ self.links += 1
|
|
|
+ continue
|
|
|
+ os.remove(dst_entry)
|
|
|
+ elif os.path.exists(dst_entry):
|
|
|
+ os.remove(dst_entry)
|
|
|
+ try:
|
|
|
+ os.symlink(rel_target, dst_entry)
|
|
|
+ except OSError as e:
|
|
|
+ fail(f"link {dst_entry} -> {rel_target}: {e}")
|
|
|
+ self.links += 1
|
|
|
+
|
|
|
+ def stage_repo(self, repo_dir):
|
|
|
+ repo_src = os.path.join(self.src, "hub", repo_dir)
|
|
|
+ repo_dst = os.path.join(self.dst, "hub", repo_dir)
|
|
|
+ if not os.path.isdir(repo_src):
|
|
|
+ fail(f"source repo missing: {repo_src}")
|
|
|
+ t0 = time.monotonic()
|
|
|
+ self.links = self.files = self.verified = self.copied = 0
|
|
|
+ links = self.stage_tree(repo_src, repo_dst)
|
|
|
+ self.stage_links(links)
|
|
|
+ elapsed = time.monotonic() - t0
|
|
|
+ print(f"STAGED {repo_dir} links={self.links} files={self.files} "
|
|
|
+ f"verified={self.verified} copied={self.copied} secs={elapsed:.0f}", flush=True)
|
|
|
+
|
|
|
+
|
|
|
+def parse_min_free(s):
|
|
|
+ s = s.strip().lower().rstrip("b")
|
|
|
+ mult = 1
|
|
|
+ if s.endswith("g"):
|
|
|
+ mult = 1 << 30
|
|
|
+ elif s.endswith("m"):
|
|
|
+ mult = 1 << 20
|
|
|
+ elif s.endswith("k"):
|
|
|
+ mult = 1 << 10
|
|
|
+ return int(float(s.rstrip("gkm")) * mult)
|
|
|
+
|
|
|
+
|
|
|
+def main():
|
|
|
+ ap = argparse.ArgumentParser()
|
|
|
+ ap.add_argument("--src", required=True, help="shared hub root (…/huggingface)")
|
|
|
+ ap.add_argument("--dst", required=True, help="node-local mirror root")
|
|
|
+ ap.add_argument("--repo", action="append", required=True,
|
|
|
+ help="models--org--name directory (repeatable)")
|
|
|
+ ap.add_argument("--skip", action="append", default=[],
|
|
|
+ help="path component ignored at any depth (repeatable)")
|
|
|
+ ap.add_argument("--verify", default="sample-hash", choices=["sample-hash", "presence"])
|
|
|
+ ap.add_argument("--min-free", default="0", help="required free space at dst")
|
|
|
+ ap.add_argument("--force", action="store_true", help="re-verify everything")
|
|
|
+ args = ap.parse_args()
|
|
|
+
|
|
|
+ if not os.path.isdir(args.src):
|
|
|
+ fail(f"source root not readable: {args.src}")
|
|
|
+ os.makedirs(args.dst, exist_ok=True)
|
|
|
+ need = parse_min_free(args.min_free)
|
|
|
+ free = os.statvfs(args.dst).f_bavail * os.statvfs(args.dst).f_frsize
|
|
|
+ if need and free < need:
|
|
|
+ fail(f"insufficient local space at {args.dst}: {free / 2**30:.0f} GiB free, "
|
|
|
+ f"need {need / 2**30:.0f} GiB")
|
|
|
+
|
|
|
+ ledger = Ledger(os.path.join(args.dst, "hub", ".sm121-staging-ledger.json"))
|
|
|
+ stager = Stager(args.src, args.dst, args.skip, args.verify, args.force, ledger)
|
|
|
+ for repo in args.repo:
|
|
|
+ stager.stage_repo(repo)
|
|
|
+ ledger.save()
|
|
|
+ print("STAGING_DONE", flush=True)
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ main()
|