| |
| """Disk-guarded snapshot download of an open Lean-prover teacher into HF_HOME (/apps only). |
| |
| source /apps/development/rlhf/env/apps_only.sh |
| PY=/apps/development/rlhf/mistral-prm-ppo/.venv/bin/python # any venv with huggingface_hub |
| align/.venv/bin/python scripts/hostguard.py run --cores 2 --mem-gb 8 --label dl --wait -- \ |
| $PY scripts/teachers/download_teacher.py REPO_ID [--revision SHA] [--min-free-gb 60] |
| |
| Refuses (exit 2) when free space on /apps minus the repo's weight bytes would fall below --min-free-gb. Sizes come from the |
| hub API (`blobs=true`), so the check happens BEFORE any byte is fetched. 2 workers, no xet. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| import shutil |
| import sys |
|
|
| os.environ.setdefault("HF_HUB_DISABLE_XET", "1") |
| if not os.environ.get("HF_HOME", "").startswith("/apps/"): |
| sys.exit("HF_HOME must be under /apps (source env/apps_only.sh)") |
|
|
| from huggingface_hub import HfApi, snapshot_download |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("repo") |
| ap.add_argument("--revision", default=None) |
| ap.add_argument("--min-free-gb", type=float, default=60.0) |
| a = ap.parse_args() |
| info = HfApi().model_info(a.repo, revision=a.revision, files_metadata=True) |
| total = sum(s.size or 0 for s in info.siblings) |
| |
| blobs = os.path.join(os.environ["HF_HOME"], "hub", "models--" + a.repo.replace("/", "--"), "blobs") |
| have = sum(e.stat().st_size for e in os.scandir(blobs)) if os.path.isdir(blobs) else 0 |
| need = max(total - have, 0) |
| free = shutil.disk_usage("/apps").free |
| after = (free - need) / 1e9 |
| print(f"{a.repo}@{info.sha}: total {total / 1e9:.1f} GB, have {have / 1e9:.1f}, need {need / 1e9:.1f} GB, free {free / 1e9:.1f} GB, free after {after:.1f} GB (floor {a.min_free_gb})", flush=True) |
| if after < a.min_free_gb: |
| print("REFUSED: would cross the free-space floor", file=sys.stderr) |
| return 2 |
| p = snapshot_download(a.repo, revision=info.sha, max_workers=2) |
| print("snapshot:", p, flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
