#!/usr/bin/env python3 """Upload scratch/gandalf staging files to Hugging Face. Runs on the login node (python3 + huggingface_hub). The feanor sidecar only needs to land files in --stage-dir; this process does the HTTP/LFS upload. """ import argparse import json import os import sys import time import traceback def _load_json(path): try: with open(path) as f: return json.load(f) except Exception: return {} def _write_json(path, payload): tmp = path + ".tmp" with open(tmp, "w") as f: json.dump(payload, f, indent=2, sort_keys=True) f.write("\n") os.replace(tmp, path) _KEEP_HB = ( "event", "step", "stage", "loss", "val_acc", "val_mass", "val_excess", "val_spread", "val_mass_multi", "val_excess_multi", "val_out_multi", "val_kl_multi", "val_hbound", "val_hgap", "val_hbound_puz", "val_ceq", "val_hq", "val_ceq_ok", "val_ceq_gap", "loc_wave", "loc_coverage", "loc_lcs", "loc_dup", "puzzle_acc", "last_ckpt_step", ) def _slim_heartbeat(src, dest): raw = _load_json(src) slim = {k: raw[k] for k in _KEEP_HB if k in raw} _write_json(dest, slim) return dest def _slim_log(src, dest): """Keep only eval / promote lines: step, stage, loss, accuracies.""" keep = [] try: with open(src, errors="replace") as f: for line in f: s = line.strip() if not s: continue if s.startswith("[curriculum]"): keep.append(s + "\n") continue # "2000 stage 1 loss 1.47 loc_acc 0.03 val_acc 0.11 inset 0.66" if " loss " in s and "stage" in s and not s.startswith("I"): keep.append(s + "\n") except OSError: return None with open(dest, "w") as f: f.writelines(keep) return dest def _upload(api, token, local, repo, dest): api.upload_file( path_or_fileobj=local, path_in_repo=dest, repo_id=repo, repo_type="model", token=token, commit_message=f"sync {dest}", ) print(f"[hf_push] uploaded {dest} ({os.path.getsize(local)} bytes)", flush=True) def sync_once(args, state, api, token): prefix = args.prefix.strip("/") repo = args.repo for local, dest, key in ( (args.heartbeat, f"{prefix}/heartbeat.json", "heartbeat"), (args.log, f"{prefix}/train.log", "log"), ): if not local or not os.path.isfile(local): continue mtime = os.path.getmtime(local) size = os.path.getsize(local) last = state.get(key, {}) min_age = 0 if key == "heartbeat" else 300 due = (time.time() - float(last.get("uploaded_at", 0))) >= min_age if due and (mtime > float(last.get("mtime", 0)) or size != last.get("size", -1)): upload_path = local if key == "heartbeat": upload_path = _slim_heartbeat( local, os.path.join(os.path.dirname(local) or ".", "_hf_heartbeat_slim.json")) elif key == "log": upload_path = _slim_log( local, os.path.join(os.path.dirname(local) or ".", "_hf_train_slim.log")) if upload_path: _upload(api, token, upload_path, repo, dest) state[key] = {"mtime": mtime, "size": size, "uploaded_at": time.time()} if args.stage_dir and os.path.isdir(args.stage_dir): for name in sorted(os.listdir(args.stage_dir)): if name in ("hf_push_state.json", "heartbeat.json"): continue if not name.endswith(".tar") and name not in ("latest.json",): continue local = os.path.join(args.stage_dir, name) if not os.path.isfile(local): continue # step_/stage_ names are unique per checkpoint, so once uploaded # they never need revisiting. latest.tar/latest.json keep the same # name and are rewritten every save, so skipping them by name would # freeze the rolling checkpoint at whichever step happened to land # first; track those by mtime instead. rolling = name in ("latest.tar", "latest.json") mtime = os.path.getmtime(local) if rolling: if mtime <= float(state.get("rolling", {}).get(name, 0)): continue elif name in state.get("staged", []): continue dest = f"{prefix}/{name}" if name.startswith("stage"): dest = f"{prefix}/stages/{name}" elif name.startswith("step_"): dest = f"{prefix}/steps/{name}" _upload(api, token, local, repo, dest) if rolling: state.setdefault("rolling", {})[name] = mtime continue # keep the file; the sidecar overwrites it in place state.setdefault("staged", []).append(name) # Drop the local copy after a successful upload to save quota. try: os.remove(local) except OSError: pass return state def main(): ap = argparse.ArgumentParser() ap.add_argument("--repo", default="Avra98/Sudoku_superposition") ap.add_argument("--prefix", default="runs/w12_inst_latent") ap.add_argument("--heartbeat", default="") ap.add_argument("--log", default="") ap.add_argument("--stage-dir", default="") ap.add_argument("--token-file", default="/scratch/users/gatmiry/.hf_token") ap.add_argument("--state", default="") ap.add_argument("--interval", type=int, default=60) ap.add_argument("--once", action="store_true") args = ap.parse_args() token = "" if args.token_file and os.path.isfile(args.token_file): with open(args.token_file) as f: token = f.read().strip() token = token or os.environ.get("HF_TOKEN", "") if not token: print("[hf_push] no token", file=sys.stderr) return 1 from huggingface_hub import HfApi api = HfApi(token=token) state_path = args.state or os.path.join( args.stage_dir or "/tmp/sudoku_hf_uploads", "hf_push_state.json") os.makedirs(os.path.dirname(state_path) or ".", exist_ok=True) state = _load_json(state_path) print(f"[hf_push] watching heartbeat={args.heartbeat} log={args.log} " f"stage={args.stage_dir}", flush=True) while True: try: state = sync_once(args, state, api, token) _write_json(state_path, state) except Exception: traceback.print_exc() print("[hf_push] cycle failed; will retry", flush=True) if args.once: break time.sleep(max(5, args.interval)) return 0 if __name__ == "__main__": sys.exit(main())