Spaces:
Running
Running
| """Move data and artefacts between this repository and a private Hugging Face dataset repo (plan §0.6: the Mac and | |
| Colab share one store). | |
| python -m experiments.hub_sync push --bundles data,embeddings,cache,packs,runs | |
| python -m experiments.hub_sync pull --bundles data,embeddings,cache,packs,runs | |
| python -m experiments.hub_sync status --bundles data,embeddings,cache,packs,runs # read-only comparison | |
| Bundles are defined in configs/base.yaml (`hub.bundles`): include/exclude patterns relative to the repository root | |
| (fnmatch, `*` crosses folders), or `archive: <folder>` for a folder of many small files, which travels as one | |
| tar.gz under archives/. Pull never replaces a local file that differs from the hub's copy and was modified after | |
| the hub's last commit of it (it is kept and reported; push it). Push uploads only files whose content differs from the hub's copy (git blob sha1 for small | |
| files, sha256 for LFS files), in commits of bounded size, and never deletes anything on the hub. runs/log.csv is | |
| merged rather than overwritten (push_log / pull_log): the hub keeps the union of every machine's rows, and each | |
| machine appends the rows it lacks. Credentials come from the logged-in Hugging Face client (HF_TOKEN or `hf auth | |
| login`); none is read or printed here. Symlinked folders (a git worktree) are followed. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import fnmatch | |
| import hashlib | |
| import io | |
| import os | |
| import tarfile | |
| import time | |
| from pathlib import Path | |
| from typing import Any | |
| REPO = Path(__file__).resolve().parent.parent | |
| LOG_IN_HUB = "logs/runs_log.csv" | |
| MAX_COMMIT_BYTES = 512 * 1024 ** 2 # a failed upload (network timeout) loses at most one such chunk | |
| MAX_COMMIT_FILES = 500 | |
| COMMIT_RETRIES = 4 | |
| RETRY_WAIT_S = 30.0 | |
| JUNK_FILES = {".DS_Store", "Thumbs.db"} | |
| def git_blob_sha1_bytes(data: bytes) -> str: | |
| return hashlib.sha1(b"blob %d\x00" % len(data) + data).hexdigest() | |
| def _file_hashes(path: Path) -> tuple[str, str]: | |
| """(git blob sha1, sha256) of a file, streamed.""" | |
| size = path.stat().st_size | |
| g, s = hashlib.sha1(b"blob %d\x00" % size), hashlib.sha256() | |
| with open(path, "rb") as f: | |
| for chunk in iter(lambda: f.read(1 << 20), b""): | |
| g.update(chunk) | |
| s.update(chunk) | |
| return g.hexdigest(), s.hexdigest() | |
| def _matches(rel: str, patterns: list[str]) -> bool: | |
| return any(fnmatch.fnmatch(rel, p) or fnmatch.fnmatch(rel, p.rstrip("/*") + "/*") for p in patterns) | |
| def select(root: Path, bundle: dict[str, Any]) -> list[str]: | |
| """Repository-relative paths (posix) of the bundle's files, following symlinked folders.""" | |
| include, exclude = bundle.get("include", []), bundle.get("exclude", []) | |
| out = [] | |
| for dirpath, dirnames, filenames in os.walk(root, followlinks=True): | |
| rel_dir = Path(dirpath).relative_to(root).as_posix() | |
| dirnames[:] = [d for d in dirnames if d not in (".git", ".venv", "__pycache__") | |
| and not _matches(f"{rel_dir}/{d}".lstrip("./"), exclude)] | |
| for name in filenames: | |
| if name in JUNK_FILES: | |
| continue | |
| rel = f"{rel_dir}/{name}".lstrip("./") if rel_dir != "." else name | |
| if _matches(rel, include) and not _matches(rel, exclude): | |
| out.append(rel) | |
| return sorted(out) | |
| def _remote(api: Any, repo_id: str) -> dict[str, Any]: | |
| try: | |
| return {e.path: e for e in api.list_repo_tree(repo_id, repo_type="dataset", recursive=True, expand=True) | |
| if hasattr(e, "size")} | |
| except Exception: # an empty repo has no tree yet | |
| return {} | |
| def _same(path: Path, r: Any) -> bool: | |
| """A local file has the content of the hub's entry (git blob sha1 for small files, sha256 for LFS files).""" | |
| if r.size != path.stat().st_size: | |
| return False | |
| g, s = _file_hashes(path) | |
| return getattr(r, "blob_id", None) == g or bool(getattr(r, "lfs", None) and r.lfs.sha256 == s) | |
| def _literal(rel: str) -> str: | |
| """An fnmatch pattern matching exactly this path (snapshot_download filters with fnmatch).""" | |
| return "".join(f"[{c}]" if c in "*?[" else c for c in rel) | |
| def newer_local(root: Path, remote: dict[str, Any], bundle: dict[str, Any]) -> list[str]: | |
| """Bundle files on the hub whose local copy differs and was modified after the hub's last commit of the file: a | |
| pull must not replace them (2026-10-20: a push during a running job put a partial file on the hub, and a later | |
| pull replaced the finished local file with it). Push them instead. An older local copy is replaced as before.""" | |
| include, exclude = bundle.get("include", []), bundle.get("exclude", []) | |
| keep = [] | |
| for rel, r in sorted(remote.items()): | |
| p = root / rel | |
| if not (_matches(rel, include) and not _matches(rel, exclude)) or not p.is_file(): | |
| continue | |
| date = getattr(getattr(r, "last_commit", None), "date", None) | |
| if date is not None and p.stat().st_mtime > date.timestamp() and not _same(p, r): | |
| keep.append(rel) | |
| return keep | |
| def _archive(folder: Path) -> bytes: | |
| import gzip | |
| buf = io.BytesIO() | |
| # gzip writes the current time into its header; mtime=0 keeps the archive byte-identical for identical files | |
| with gzip.GzipFile(fileobj=buf, mode="wb", compresslevel=6, mtime=0) as gz, tarfile.open(fileobj=gz, mode="w") as tar: | |
| for p in sorted(folder.rglob("*")): | |
| if p.is_file(): | |
| info = tar.gettarinfo(str(p), arcname=p.relative_to(folder).as_posix()) | |
| info.mtime, info.uid, info.gid, info.uname, info.gname = 0, 0, 0, "", "" # same bytes for the same files | |
| with open(p, "rb") as f: | |
| tar.addfile(info, f) | |
| return buf.getvalue() | |
| def _plan(api: Any, repo_id: str, root: Path, bundles: dict[str, Any], names: list[str] | |
| ) -> tuple[list[tuple[str, Any, int, bool]], int]: | |
| """Files whose content differs from the hub's copy: [(path in repo, local source, size, missing on hub)], and the | |
| number of files that already match. Archives are staged under .cache/hub_sync so the hub's chunked upload works.""" | |
| remote = _remote(api, repo_id) | |
| todo: list[tuple[str, Any, int, bool]] = [] | |
| matched = 0 | |
| for name in names: | |
| b = bundles[name] | |
| if "archive" in b: | |
| folder = root / b["archive"] | |
| if not folder.is_dir(): | |
| continue | |
| data = _archive(folder) | |
| path = f"archives/{Path(b['archive']).name}.tar.gz" | |
| r = remote.get(path) | |
| if r is None or getattr(r, "blob_id", None) != git_blob_sha1_bytes(data) and not ( | |
| getattr(r, "lfs", None) and r.lfs.sha256 == hashlib.sha256(data).hexdigest()): | |
| staged = root / ".cache" / "hub_sync" / path # a file, so the hub's chunked upload can be used | |
| staged.parent.mkdir(parents=True, exist_ok=True) | |
| staged.write_bytes(data) | |
| todo.append((path, str(staged), len(data), r is None)) | |
| else: | |
| matched += 1 | |
| continue | |
| for rel in select(root, b): | |
| p = (root / rel).resolve() | |
| r = remote.get(rel) | |
| if r is not None and _same(p, r): | |
| matched += 1 | |
| continue | |
| todo.append((rel, str(p), p.stat().st_size, r is None)) | |
| return todo, matched | |
| def status(api: Any, repo_id: str, root: Path, bundles: dict[str, Any], names: list[str]) -> dict[str, Any]: | |
| """Read-only: which local bundle files differ from (or are missing on) the hub. Nothing is uploaded.""" | |
| todo, matched = _plan(api, repo_id, root, bundles, names) | |
| return {"matched": matched, "differs": len(todo), "missing_on_hub": sum(m for *_, m in todo), | |
| "bytes": sum(t[2] for t in todo), "files": [t[0] for t in todo]} | |
| def push(api: Any, repo_id: str, root: Path, bundles: dict[str, Any], names: list[str]) -> dict[str, Any]: | |
| from huggingface_hub import CommitOperationAdd | |
| api.create_repo(repo_id, repo_type="dataset", private=True, exist_ok=True) | |
| todo, _ = _plan(api, repo_id, root, bundles, names) | |
| ops = [CommitOperationAdd(path_in_repo=rel, path_or_fileobj=src) for rel, src, _, _ in todo] | |
| sizes = [size for _, _, size, _ in todo] | |
| print(f"[hub_sync] plan: {len(ops)} files, {sum(sizes) / 1e9:.2f} GB to upload", flush=True) | |
| commits, start = 0, 0 | |
| while start < len(ops): # bounded commits: a failed upload loses one chunk, not everything | |
| end, total = start, 0 | |
| while end < len(ops) and end - start < MAX_COMMIT_FILES and (total + sizes[end] <= MAX_COMMIT_BYTES or end == start): | |
| total += sizes[end] | |
| end += 1 | |
| for attempt in range(COMMIT_RETRIES + 1): # network timeouts happen on long uploads: retry the chunk | |
| try: | |
| api.create_commit(repo_id=repo_id, repo_type="dataset", operations=ops[start:end], | |
| commit_message=f"sync {','.join(names)}: {end - start} files") | |
| break | |
| except Exception as e: | |
| if attempt == COMMIT_RETRIES: | |
| raise | |
| wait = RETRY_WAIT_S * 2 ** attempt | |
| print(f"[hub_sync] commit of {end - start} files failed ({type(e).__name__}); retry {attempt + 1}/" | |
| f"{COMMIT_RETRIES} in {wait:.0f}s", flush=True) | |
| time.sleep(wait) | |
| commits += 1 | |
| print(f"[hub_sync] committed {end}/{len(ops)} files", flush=True) | |
| start = end | |
| return {"uploaded": len(ops), "bytes": sum(sizes), "commits": commits} | |
| def _read_log(text: str) -> tuple[list[str], list[dict[str, str]]]: | |
| r = csv.DictReader(io.StringIO(text)) | |
| return list(r.fieldnames or []), list(r) | |
| def _key(row: dict[str, str]) -> tuple[str, str]: | |
| return row.get("timestamp", ""), row.get("run_id", "") | |
| def push_log(api: Any, repo_id: str, root: Path) -> int: | |
| """Union of the hub's run log and this machine's runs/log.csv, written back to the hub; returns rows added.""" | |
| from huggingface_hub import CommitOperationAdd | |
| local_path = root / "runs" / "log.csv" | |
| if not local_path.is_file(): | |
| return 0 | |
| fields, local = _read_log(local_path.read_text(encoding="utf-8")) | |
| hub_fields, hub = fields, [] | |
| try: | |
| hub_fields, hub = _read_log(Path(api.hf_hub_download(repo_id, LOG_IN_HUB, repo_type="dataset", | |
| local_dir=root / ".cache" / "hub_sync")).read_text()) | |
| except Exception: | |
| pass | |
| have = {_key(r) for r in hub} | |
| new = [r for r in local if _key(r) not in have] | |
| if not new and hub: | |
| return 0 | |
| rows = sorted(hub + new, key=_key) | |
| buf = io.StringIO() | |
| w = csv.DictWriter(buf, fieldnames=hub_fields or fields) | |
| w.writeheader() | |
| w.writerows(rows) | |
| api.create_commit(repo_id=repo_id, repo_type="dataset", commit_message=f"run log: +{len(new)} rows", | |
| operations=[CommitOperationAdd(path_in_repo=LOG_IN_HUB, | |
| path_or_fileobj=io.BytesIO(buf.getvalue().encode()))]) | |
| return len(new) | |
| def pull_log(api: Any, repo_id: str, root: Path) -> int: | |
| """Append to runs/log.csv the hub's rows this machine lacks (never rewrites existing rows); returns rows added.""" | |
| try: | |
| path = Path(api.hf_hub_download(repo_id, LOG_IN_HUB, repo_type="dataset", local_dir=root / ".cache" / "hub_sync")) | |
| except Exception: | |
| return 0 | |
| _, hub = _read_log(path.read_text(encoding="utf-8")) | |
| local_path = root / "runs" / "log.csv" | |
| fields, local = _read_log(local_path.read_text(encoding="utf-8")) | |
| have = {_key(r) for r in local} | |
| new = [r for r in hub if _key(r) not in have] | |
| if new: | |
| with open(local_path, "a", newline="", encoding="utf-8") as f: | |
| csv.DictWriter(f, fieldnames=fields, extrasaction="ignore").writerows(new) | |
| return len(new) | |
| def _snapshot(repo_id: str, root: Path, include: list[str], exclude: list[str]) -> None: | |
| from huggingface_hub import snapshot_download | |
| snapshot_download(repo_id, repo_type="dataset", local_dir=root, allow_patterns=include, | |
| ignore_patterns=exclude or None) | |
| def pull(api: Any, repo_id: str, root: Path, bundles: dict[str, Any], names: list[str]) -> dict[str, Any]: | |
| """Download the bundles. A local file newer than the hub's differing copy is kept and reported (newer_local).""" | |
| pulled, kept = [], [] | |
| remote = None | |
| for name in names: | |
| b = bundles[name] | |
| if "archive" in b: | |
| path = f"archives/{Path(b['archive']).name}.tar.gz" | |
| local = Path(api.hf_hub_download(repo_id, path, repo_type="dataset", local_dir=root / ".cache" / "hub_sync")) | |
| dest = root / b["archive"] | |
| dest.mkdir(parents=True, exist_ok=True) | |
| with tarfile.open(local, mode="r:gz") as tar: | |
| tar.extractall(dest, filter="data") # the "data" filter refuses absolute paths and links out of dest | |
| pulled.append(name) | |
| continue | |
| remote = _remote(api, repo_id) if remote is None else remote | |
| keep = newer_local(root, remote, b) | |
| kept += keep | |
| _snapshot(repo_id, root, list(b.get("include", [])), list(b.get("exclude", [])) + [_literal(k) for k in keep]) | |
| pulled.append(name) | |
| if kept: | |
| print(f"[hub_sync] kept {len(kept)} local files newer than the hub's copy (push them): {kept[:20]}", flush=True) | |
| return {"pulled": pulled, "kept_local": kept} | |
| def main() -> None: | |
| import yaml | |
| from huggingface_hub import HfApi | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("action", choices=["push", "pull", "status"]) | |
| ap.add_argument("--bundles", default="data,embeddings,cache,packs,runs") | |
| ap.add_argument("--config", default="configs/base.yaml") | |
| args = ap.parse_args() | |
| hub = yaml.safe_load((REPO / args.config).read_text(encoding="utf-8"))["hub"] | |
| names = [n.strip() for n in args.bundles.split(",") if n.strip()] | |
| api = HfApi() | |
| if args.action == "status": | |
| s = status(api, hub["repo"], REPO, hub["bundles"], names) | |
| print({k: v for k, v in s.items() if k != "files"}, flush=True) | |
| for f in s["files"][:50]: | |
| print(" differs:", f) | |
| elif args.action == "push": | |
| print(push(api, hub["repo"], REPO, hub["bundles"], names), flush=True) | |
| if "runs" in names: | |
| print({"log_rows_added_to_hub": push_log(api, hub["repo"], REPO)}) | |
| else: | |
| print(pull(api, hub["repo"], REPO, hub["bundles"], names), flush=True) | |
| if "runs" in names: | |
| print({"log_rows_added_here": pull_log(api, hub["repo"], REPO)}) | |
| if __name__ == "__main__": | |
| main() | |