code/wavecurriculum_run/hf_push_login.py

#!/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())