Document the 3-hourly GPU checkpoint auto-sync (gpu_* folders); add the uploader script
32f554b verified | #!/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() | |