osdc/hf-cache: scripts/hf-cache-seed.py to seed models into bucket(s) (#833)

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0)
(oldest at bottom):
* __->__ #833

Standalone seeding tool under osdc/scripts (moved out of the justfile).
Downloads
each HF model once locally, then `aws s3 sync`s it to one or more
clusters'
per-region buckets pytorch-hf-model-cache-<cluster_id>, or --all
clusters that
enable the hf-cache module (synced in parallel). Writes straight to S3
with the
operator's AWS creds — independent of the cluster mount. Run once per
model.

uv run scripts/hf-cache-seed.py -c meta-staging-aws-ue1
Qwen/Qwen2.5-7B-Instruct
  uv run scripts/hf-cache-seed.py --all Qwen/Qwen2.5-7B-Instruct

Author: huydhn

osdc/scripts/hf-cache-seed.py

@@ -0,0 +1,178 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.12"
+# dependencies = ["huggingface_hub>=0.24", "pyyaml>=6.0"]
+# ///
+"""Seed HuggingFace models into OSDC clusters' hf-cache S3 buckets.
+
+Each model is downloaded once locally, then ``aws s3 sync``ed to every target
+cluster's per-region bucket ``pytorch-hf-model-cache-<cluster_id>`` (the layout
+the runners read at ``/mnt/hf_cache/hub``). Target one or more clusters, or
+``--all`` (every cluster that enables the hf-cache module); clusters are synced
+in parallel. Writes straight to S3 — independent of the mount on the cluster.
+
+Requires ``aws`` on PATH (run with mise active, e.g. from the osdc/ dir).
+
+Usage:
+  uv run scripts/hf-cache-seed.py -c meta-staging-aws-ue1 Qwen/Qwen2.5-7B-Instruct
+  uv run scripts/hf-cache-seed.py --all Qwen/Qwen2.5-7B-Instruct meta-llama/Llama-3.1-8B
+"""
+
+from __future__ import annotations
+
+import argparse
+import concurrent.futures
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import yaml
+
+BUCKET_PREFIX = "pytorch-hf-model-cache-"
+HF_MODULE = "hf-cache"
+CLUSTERS_YAML = Path(os.environ.get("CLUSTERS_YAML", Path(__file__).resolve().parent.parent / "clusters.yaml"))
+
+
+def load_clusters() -> dict:
+    with open(CLUSTERS_YAML) as f:
+        return yaml.safe_load(f).get("clusters", {}) or {}
+
+
+def resolve_targets(clusters: dict, requested: list[str] | None, all_clusters: bool) -> list[str]:
+    """Return the cluster ids to seed. --all selects every hf-cache-enabled cluster."""
+    if all_clusters:
+        targets = [cid for cid, c in clusters.items() if HF_MODULE in (c.get("modules") or [])]
+        if not targets:
+            raise SystemExit(f"No clusters enable the '{HF_MODULE}' module in {CLUSTERS_YAML}")
+        return targets
+    targets = []
+    for cid in requested or []:
+        if cid not in clusters:
+            raise SystemExit(f"Unknown cluster '{cid}'. Known: {', '.join(clusters)}")
+        if HF_MODULE not in (clusters[cid].get("modules") or []):
+            print(
+                f"WARNING: cluster '{cid}' does not enable '{HF_MODULE}' — its bucket may not exist.", file=sys.stderr
+            )
+        targets.append(cid)
+    return targets
+
+
+def bucket_for(cid: str) -> str:
+    return f"{BUCKET_PREFIX}{cid}"
+
+
+# config + tokenizer only (no weight files) — enough for benchmarks that init
+# with random weights, e.g. vLLM's ``--load-format dummy`` runs. Excludes
+# *.safetensors / *.bin / *.pth / *.gguf, which for those models are huge/gated.
+CONFIG_ONLY_PATTERNS = ["*.json", "*.txt", "*.model", "*.jinja"]
+
+
+def download_models(models: list[str], staging: Path, config_only: bool = False) -> list[str]:
+    """Download each model once into staging/hub (shared across all target buckets).
+
+    With ``config_only`` only config + tokenizer files are fetched (no weights) —
+    for benchmarks that random-init weights (e.g. vLLM ``--load-format dummy``).
+
+    Returns the models that failed to download (e.g. a gated repo without an
+    HF_TOKEN that has access) — these are skipped so one bad model doesn't abort
+    the whole batch; the rest still download and sync.
+    """
+    from huggingface_hub import snapshot_download
+
+    hub = staging / "hub"
+    hub.mkdir(parents=True, exist_ok=True)
+    allow = CONFIG_ONLY_PATTERNS if config_only else None
+    failed = []
+    for i, model in enumerate(models, 1):
+        suffix = " (config only)" if config_only else ""
+        print(f"-> [{i}/{len(models)}] downloading {model}{suffix} ...", flush=True)
+        try:
+            path = snapshot_download(model, cache_dir=str(hub), allow_patterns=allow)
+            print(f"   {model} -> {path}", flush=True)
+        except Exception as e:
+            print(f"   !! SKIP {model}: {type(e).__name__}: {str(e)[:140]}", flush=True)
+            failed.append(model)
+    return failed
+
+
+def sync_to_cluster(cid: str, region: str, staging: Path) -> tuple[str, bool, str]:
+    cmd = [
+        "aws",
+        "s3",
+        "sync",
+        str(staging / "hub"),
+        f"s3://{bucket_for(cid)}/hub",
+        "--region",
+        region,
+        "--no-progress",
+    ]
+    # Stream aws's per-file "upload: ..." lines live (prefixed with the cluster id
+    # so parallel syncs stay readable) instead of capturing — otherwise the upload
+    # phase looks hung with no output. --no-progress keeps it to one line per file.
+    print(f"[{cid}] syncing -> s3://{bucket_for(cid)}/hub", flush=True)
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
+    for line in proc.stdout or []:
+        print(f"[{cid}] {line}", end="", flush=True)
+    rc = proc.wait()
+    return cid, rc == 0, "" if rc == 0 else f"aws s3 sync exited {rc} (see [{cid}] output above)"
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Seed HF models into OSDC hf-cache S3 buckets.")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("-c", "--cluster", action="append", metavar="CLUSTER_ID", help="target cluster (repeatable)")
+    group.add_argument("--all", action="store_true", help="target every cluster that enables the hf-cache module")
+    parser.add_argument("models", nargs="+", help="HF model id(s), e.g. Qwen/Qwen2.5-7B-Instruct")
+    parser.add_argument("-j", "--jobs", type=int, default=0, help="parallel cluster syncs (default: one per cluster)")
+    parser.add_argument(
+        "--config-only",
+        action="store_true",
+        help="download only config+tokenizer files, no weights (for vLLM --load-format dummy benchmarks)",
+    )
+    args = parser.parse_args(argv)
+
+    clusters = load_clusters()
+    targets = resolve_targets(clusters, args.cluster, args.all)
+    print(f"Targets ({len(targets)}): {', '.join(targets)}")
+    print(f"Models  ({len(args.models)}): {', '.join(args.models)}")
+
+    if shutil.which("aws") is None:
+        raise SystemExit("ERROR: 'aws' not found on PATH — run with mise active (e.g. from the osdc/ dir).")
+
+    staging = Path(tempfile.mkdtemp(prefix="hf-cache-seed-"))
+    try:
+        dl_failed = download_models(args.models, staging, config_only=args.config_only)
+        workers = args.jobs or len(targets)
+        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex:
+            futs = [ex.submit(sync_to_cluster, cid, clusters[cid]["region"], staging) for cid in targets]
+            results = [f.result() for f in concurrent.futures.as_completed(futs)]
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
+
+    print("\n=== results ===")
+    if dl_failed:
+        print(f"  SKIPPED {len(dl_failed)} model(s) — download failed: {', '.join(dl_failed)}")
+    failed = 0
+    for cid, ok, detail in sorted(results):
+        if ok:
+            print(f"  OK   {cid}  -> s3://{bucket_for(cid)}/hub")
+        else:
+            failed += 1
+            tail = "\n        ".join(detail.splitlines()[-5:])
+            print(f"  FAIL {cid}\n        {tail}")
+    seeded = len(args.models) - len(dl_failed)
+    if failed or dl_failed:
+        if failed:
+            print(f"\n{failed}/{len(results)} cluster(s) failed.")
+        if dl_failed:
+            print(f"{len(dl_failed)} model(s) skipped (see above).")
+        return 1
+    print(f"\nSeeded {seeded} model(s) into {len(results)} cluster(s).")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

osdc/scripts/python/test_hf_cache_seed.py

@@ -0,0 +1,157 @@
+"""Unit tests for scripts/hf-cache-seed.py."""
+
+import importlib.util
+import sys
+import types
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+# Filename has a hyphen, so load it via importlib (can't `import`).
+_spec = importlib.util.spec_from_file_location(
+    "hf_cache_seed", str(Path(__file__).resolve().parent.parent / "hf-cache-seed.py")
+)
+_mod = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_mod)
+
+CLUSTERS = {
+    "meta-staging-aws-ue1": {"region": "us-east-1", "modules": ["arc", "hf-cache"]},
+    "meta-staging-aws-uw1": {"region": "us-west-1", "modules": ["arc", "hf-cache"]},
+    "meta-prod-aws-ue1": {"region": "us-east-1", "modules": ["arc", "nodepools"]},
+}
+
+
+def test_bucket_for():
+    assert _mod.bucket_for("meta-staging-aws-ue1") == "pytorch-hf-model-cache-meta-staging-aws-ue1"
+
+
+def test_resolve_targets_all_selects_hf_cache_clusters():
+    assert set(_mod.resolve_targets(CLUSTERS, None, True)) == {"meta-staging-aws-ue1", "meta-staging-aws-uw1"}
+
+
+def test_resolve_targets_all_none_enabled_exits():
+    with pytest.raises(SystemExit):
+        _mod.resolve_targets({"c": {"region": "r", "modules": []}}, None, True)
+
+
+def test_resolve_targets_explicit():
+    assert _mod.resolve_targets(CLUSTERS, ["meta-staging-aws-ue1"], False) == ["meta-staging-aws-ue1"]
+
+
+def test_resolve_targets_unknown_exits():
+    with pytest.raises(SystemExit):
+        _mod.resolve_targets(CLUSTERS, ["nope"], False)
+
+
+def test_resolve_targets_warns_when_module_disabled(capsys):
+    assert _mod.resolve_targets(CLUSTERS, ["meta-prod-aws-ue1"], False) == ["meta-prod-aws-ue1"]
+    assert "does not enable" in capsys.readouterr().err
+
+
+def test_load_clusters(tmp_path, monkeypatch):
+    f = tmp_path / "clusters.yaml"
+    f.write_text("clusters:\n  c1:\n    region: us-east-1\n")
+    monkeypatch.setattr(_mod, "CLUSTERS_YAML", f)
+    assert _mod.load_clusters() == {"c1": {"region": "us-east-1"}}
+
+
+def test_download_models(monkeypatch, tmp_path):
+    fake = types.ModuleType("huggingface_hub")
+    fake.snapshot_download = lambda model, cache_dir, allow_patterns=None: f"{cache_dir}/{model}"
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake)
+    assert _mod.download_models(["org/m"], tmp_path) == []
+    assert (tmp_path / "hub").is_dir()
+
+
+def test_download_models_config_only_passes_allow_patterns(monkeypatch, tmp_path):
+    seen = {}
+
+    def cap(model, cache_dir, allow_patterns=None):
+        seen["allow"] = allow_patterns
+        return f"{cache_dir}/{model}"
+
+    fake = types.ModuleType("huggingface_hub")
+    fake.snapshot_download = cap
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake)
+    _mod.download_models(["org/m"], tmp_path, config_only=True)
+    assert seen["allow"] == _mod.CONFIG_ONLY_PATTERNS
+
+
+def test_download_models_skips_failure(monkeypatch, tmp_path):
+    def boom(model, cache_dir, allow_patterns=None):
+        if model == "org/bad":
+            raise OSError("gated repo")
+        return f"{cache_dir}/{model}"
+
+    fake = types.ModuleType("huggingface_hub")
+    fake.snapshot_download = boom
+    monkeypatch.setitem(sys.modules, "huggingface_hub", fake)
+    assert _mod.download_models(["org/ok", "org/bad"], tmp_path) == ["org/bad"]
+
+
+def _fake_popen(lines, rc):
+    fake = MagicMock()
+    fake.stdout = list(lines)  # iterable + truthy
+    fake.wait.return_value = rc
+    return fake
+
+
+def test_sync_to_cluster_ok():
+    with patch.object(_mod.subprocess, "Popen", return_value=_fake_popen(["upload: a\n"], 0)):
+        cid, ok, _ = _mod.sync_to_cluster("c1", "us-east-1", Path("/tmp/x"))
+    assert ok
+    assert cid == "c1"
+
+
+def test_sync_to_cluster_fail():
+    with patch.object(_mod.subprocess, "Popen", return_value=_fake_popen(["AccessDenied\n"], 1)):
+        _, ok, detail = _mod.sync_to_cluster("c1", "us-east-1", Path("/tmp/x"))
+    assert not ok
+    assert "exited 1" in detail
+
+
+def _run_main(argv, sync=lambda cid, region, staging: (cid, True, "ok")):
+    with (
+        patch.object(_mod, "load_clusters", return_value=CLUSTERS),
+        patch.object(_mod, "download_models", return_value=[]),
+        patch.object(_mod.shutil, "which", return_value="/usr/bin/aws"),
+        patch.object(_mod, "sync_to_cluster", side_effect=sync),
+    ):
+        return _mod.main(argv)
+
+
+def test_main_success(capsys):
+    assert _run_main(["-c", "meta-staging-aws-ue1", "org/m"]) == 0
+    assert "Seeded" in capsys.readouterr().out
+
+
+def test_main_all_success():
+    assert _run_main(["--all", "org/m"]) == 0
+
+
+def test_main_reports_failure(capsys):
+    rc = _run_main(["-c", "meta-staging-aws-ue1", "org/m"], sync=lambda cid, region, staging: (cid, False, "e1\ne2"))
+    assert rc == 1
+    assert "FAIL" in capsys.readouterr().out
+
+
+def test_main_reports_download_skip(capsys):
+    with (
+        patch.object(_mod, "load_clusters", return_value=CLUSTERS),
+        patch.object(_mod, "download_models", return_value=["org/bad"]),
+        patch.object(_mod.shutil, "which", return_value="/usr/bin/aws"),
+        patch.object(_mod, "sync_to_cluster", side_effect=lambda cid, region, staging: (cid, True, "")),
+    ):
+        rc = _mod.main(["-c", "meta-staging-aws-ue1", "org/ok", "org/bad"])
+    assert rc == 1
+    assert "SKIPPED" in capsys.readouterr().out
+
+
+def test_main_requires_aws():
+    with (
+        patch.object(_mod, "load_clusters", return_value=CLUSTERS),
+        patch.object(_mod.shutil, "which", return_value=None),
+        pytest.raises(SystemExit),
+    ):
+        _mod.main(["-c", "meta-staging-aws-ue1", "org/m"])