File size: 7,116 Bytes
32f554b | 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | #!/usr/bin/env python3
"""Auto-sync the newest GPU training checkpoint to the Hugging Face repo (add-only, never deletes).
Every --interval seconds (default 3 h):
1. pick the active run dir: the subdir of checkpoints/ (excluding init*, smoke*, quickpod*) whose newest *.pt
is the most recent (so a new OUT dir, e.g. a bilingual resume, is picked up automatically);
2. take its newest checkpoint (highest step), skip it if it is still being written
(size/mtime must be stable for --stable-secs and the zip central directory must be readable);
3. if checkpoints/gpu_<run_dir>/<file> is not on the Hub yet, upload it plus the run's progress.tsv in one commit
whose message carries the step and the matching val_ce / val_field_alone.
Network errors are retried with backoff; a failed cycle never stops the loop.
Auth: huggingface_hub reads ~/.cache/huggingface/token (mode 600); the token is never printed or logged.
Run: nohup nice -n 19 ionice -c3 ~/hfsync-venv/bin/python -u scripts/hf_ckpt_sync.py --loop >> logs/hf_sync.log 2>&1 < /dev/null &
(or: bash scripts/hf_ckpt_sync.sh start|status|stop|once)
"""
import argparse, csv, datetime, fcntl, glob, json, os, re, sys, time, zipfile
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("OMP_NUM_THREADS", "1")
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
EXCLUDE = re.compile(r"^(init|smoke|quickpod)", re.I)
STEP = re.compile(r"atom_native_step_(\d+)_")
def log(msg):
print(f"{datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='seconds')} {msg}", flush=True)
def step_of(path):
m = STEP.search(os.path.basename(path))
return int(m.group(1)) if m else -1
def active_run_dir(ckroot):
best = None
for d in sorted(glob.glob(os.path.join(ckroot, "*/"))):
name = os.path.basename(d.rstrip("/"))
if EXCLUDE.match(name):
continue
pts = glob.glob(os.path.join(d, "*.pt"))
if not pts:
continue
newest = max(os.path.getmtime(p) for p in pts)
if best is None or newest > best[0]:
best = (newest, d.rstrip("/"))
return best[1] if best else None
def newest_ckpt(run_dir):
pts = glob.glob(os.path.join(run_dir, "*.pt"))
return max(pts, key=lambda p: (step_of(p), os.path.getmtime(p))) if pts else None
def is_complete(path, stable_secs):
try:
s1 = os.stat(path)
time.sleep(stable_secs)
s2 = os.stat(path)
if (s1.st_size, s1.st_mtime) != (s2.st_size, s2.st_mtime) or s2.st_size == 0:
return False
with zipfile.ZipFile(path) as z: # torch .pt = zip; readable central directory => fully written
z.namelist()
return True
except Exception as e: # noqa: BLE001
log(f"not ready: {os.path.basename(path)} ({type(e).__name__})")
return False
def val_for_step(tsv, step):
"""(val_ce, val_field_alone, tick) of the row for `step` (else the last row)."""
try:
with open(tsv, newline="", encoding="utf-8") as f:
rows = list(csv.DictReader(f, delimiter="\t"))
except OSError:
return None, None, None
if not rows:
return None, None, None
row = next((r for r in rows if r.get("tick") == str(step)), rows[-1])
return row.get("val_ce"), row.get("val_field_alone"), row.get("tick")
def retry(fn, what, tries=6, base=30):
for i in range(tries):
try:
return fn()
except Exception as e: # noqa: BLE001
msg = str(e).splitlines()[0][:300] if str(e) else ""
log(f"{what} failed (attempt {i + 1}/{tries}): {type(e).__name__}: {msg}")
if i == tries - 1:
raise
time.sleep(min(base * 2 ** i, 1800))
def cycle(args, api):
ckroot = os.path.join(ROOT, "checkpoints")
run_dir = active_run_dir(ckroot)
if not run_dir:
log("no run dir with checkpoints yet")
return
ck = newest_ckpt(run_dir)
run = os.path.basename(run_dir)
dest_dir = f"checkpoints/gpu_{run}"
dest = f"{dest_dir}/{os.path.basename(ck)}"
remote = set(retry(lambda: api.list_repo_files(args.repo), "list_repo_files"))
if dest in remote:
log(f"up to date: {dest} already on the Hub")
return
if not is_complete(ck, args.stable_secs):
log(f"skip this cycle: {os.path.basename(ck)} still being written")
return
from huggingface_hub import CommitOperationAdd
step = step_of(ck)
tsv = os.path.join(run_dir, "progress.tsv")
vce, vfa, row_tick = val_for_step(tsv, step)
ops = [CommitOperationAdd(path_in_repo=dest, path_or_fileobj=ck)]
if os.path.exists(tsv):
ops.append(CommitOperationAdd(path_in_repo=f"{dest_dir}/progress.tsv", path_or_fileobj=tsv))
vtxt = f"val_ce {vce} / field-alone {vfa} (progress row tick {row_tick})" if vce else "val_ce n/a"
msg = f"GPU auto-sync: {run} step {step}, {vtxt}"
size = os.path.getsize(ck)
t0 = time.time()
info = retry(lambda: api.create_commit(args.repo, operations=ops, commit_message=msg), "create_commit")
remote = set(retry(lambda: api.list_repo_files(args.repo), "list_repo_files (verify)"))
ok = dest in remote
log(f"{'UPLOADED' if ok else 'COMMIT MADE BUT FILE MISSING ON HUB (check .gitignore)'}: {dest} "
f"({size} B, {time.time() - t0:.0f}s) step={step} {vtxt} commit={getattr(info, 'oid', '')}")
with open(os.path.join(ROOT, "logs", "hf_sync_state.json"), "w") as f:
json.dump({"last": dest, "step": step, "val_ce": vce, "val_field_alone": vfa, "ok": ok,
"time": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds")}, f)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--repo", default="thefinalboss/atom-ai")
ap.add_argument("--interval", type=int, default=3 * 3600)
ap.add_argument("--stable-secs", type=int, default=30)
ap.add_argument("--loop", action="store_true", help="run forever (default: one cycle)")
args = ap.parse_args()
try:
os.nice(19)
except OSError:
pass
os.makedirs(os.path.join(ROOT, "logs"), exist_ok=True)
lock = open(os.path.join(ROOT, "logs", "hf_sync.lock"), "w")
try:
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
except OSError:
log("another hf_ckpt_sync holds logs/hf_sync.lock; exiting")
sys.exit(1)
if args.loop:
with open(os.path.join(ROOT, "logs", "hf_sync.pid"), "w") as f:
f.write(str(os.getpid()))
from huggingface_hub import HfApi
api = HfApi()
who = retry(lambda: api.whoami(), "whoami")
log(f"start pid={os.getpid()} repo={args.repo} user={who.get('name')} interval={args.interval}s loop={args.loop}")
while True:
try:
cycle(args, api)
except Exception as e: # noqa: BLE001
log(f"cycle error: {type(e).__name__}: {str(e).splitlines()[0][:300] if str(e) else ''}")
if not args.loop:
break
time.sleep(args.interval)
if __name__ == "__main__":
main()
|