File size: 5,093 Bytes
8819e91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | #!/usr/bin/env python3
"""One-shot push of the code and the newest checkpoint per run to the Hub.
hf_sync.py (the sidecar on the compute node) cannot upload: huggingface_hub is
not importable there, so it only scp's tarballs to the login node's staging
dirs. This script runs on the LOGIN node, where huggingface_hub does import, and
does the actual upload:
code/... every .py/.sh of the run directory, flat mirror
runs/<tag>/steps/ the newest N checkpoint tarballs found in
/tmp/sudoku_hf_<tag>/
Tarball names (step_*/stage*) are unique per checkpoint, so anything already in
the repo is skipped; the rolling latest.tar is always re-sent.
"""
import argparse
import os
import sys
import time
def code_files(root):
out = []
for name in sorted(os.listdir(root)):
p = os.path.join(root, name)
if os.path.isfile(p) and name.endswith((".py", ".sh")):
out.append((p, f"code/{name}"))
sub = os.path.join(root, "train")
if os.path.isdir(sub):
for name in sorted(os.listdir(sub)):
p = os.path.join(sub, name)
if os.path.isfile(p) and name.endswith(".py"):
out.append((p, f"code/train/{name}"))
return out
def run_dirs(stage_root, prefix, only):
out = []
for name in sorted(os.listdir(stage_root)):
if not name.startswith(prefix):
continue
tag = name[len(prefix):]
if only and tag not in only:
continue
d = os.path.join(stage_root, name)
if os.path.isdir(d):
out.append((tag, d))
return out
def newest_tars(d, n):
tars = []
for name in os.listdir(d):
if not name.endswith(".tar"):
continue
p = os.path.join(d, name)
if os.path.isfile(p):
tars.append((os.path.getmtime(p), name, p))
tars.sort()
return [(name, p) for _, name, p in tars[-n:]]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="Avra98/Sudoku_superposition")
ap.add_argument("--code-dir", default=os.path.dirname(
os.path.abspath(__file__)))
ap.add_argument("--stage-root", default="/tmp")
ap.add_argument("--stage-prefix", default="sudoku_hf_")
ap.add_argument("--tags", default="",
help="comma-separated run tags; empty = every staging dir")
ap.add_argument("--per-run", type=int, default=1,
help="how many of the newest tarballs to push per run")
ap.add_argument("--token-file", default="/scratch/users/gatmiry/.hf_token")
ap.add_argument("--skip-code", action="store_true")
ap.add_argument("--dry-run", action="store_true")
args = ap.parse_args()
token = os.environ.get("HF_TOKEN", "")
if not token and os.path.isfile(args.token_file):
with open(args.token_file) as f:
token = f.read().strip()
if not token:
print("no HF token", file=sys.stderr)
return 1
from huggingface_hub import HfApi
api = HfApi(token=token)
try:
existing = set(api.list_repo_files(args.repo, repo_type="model"))
except Exception as e:
print(f"could not list repo ({e}); nothing will be skipped")
existing = set()
plan = []
if not args.skip_code:
plan += code_files(args.code_dir)
only = {t for t in args.tags.split(",") if t}
for tag, d in run_dirs(args.stage_root, args.stage_prefix, only):
for name, p in newest_tars(d, args.per_run):
sub = "stages" if name.startswith("stage") else "steps"
dest = f"runs/{tag}/{sub}/{name}"
if name == "latest.tar":
dest = f"runs/{tag}/latest.tar"
elif dest in existing:
print(f"skip {dest} (already in repo)")
continue
plan.append((p, dest))
total = sum(os.path.getsize(p) for p, _ in plan)
print(f"{len(plan)} files, {total / 1e9:.2f} GB", flush=True)
if args.dry_run:
for p, dest in plan:
print(f" {dest} <- {p} ({os.path.getsize(p) / 1e6:.1f} MB)")
return 0
done_bytes = 0
for i, (p, dest) in enumerate(plan, 1):
sz = os.path.getsize(p)
t0 = time.time()
for attempt in (1, 2, 3):
try:
api.upload_file(path_or_fileobj=p, path_in_repo=dest,
repo_id=args.repo, repo_type="model",
commit_message=f"push {dest}")
break
except Exception as e:
print(f" attempt {attempt} failed for {dest}: {e}", flush=True)
if attempt == 3:
print(f"GIVING UP on {dest}", flush=True)
else:
time.sleep(10 * attempt)
done_bytes += sz
print(f"[{i}/{len(plan)}] {dest} {sz / 1e6:.1f} MB in "
f"{time.time() - t0:.1f}s ({done_bytes / 1e9:.2f} GB done)",
flush=True)
print("push complete", flush=True)
return 0
if __name__ == "__main__":
sys.exit(main())
|