setup/upload_ckpt_hf.py

#!/usr/bin/env python3
"""Watch a training output directory and push each new checkpoint to the Hub.

Runs alongside training rather than inside it, so a Hub outage or a rate limit
cannot take the training job down with it. Each checkpoint goes to its own
`step-<N>` subfolder, so nothing is overwritten and a partially uploaded step is
obvious from what is missing.

Usage:
    python upload_ckpt_hf.py --watch-dir <checkpoint dir> --repo-id <user/name> [--once]
"""

from __future__ import annotations

import argparse
import json
import os
import pathlib
import sys
import time


def find_steps(watch: pathlib.Path) -> list[tuple[int, pathlib.Path]]:
    """Return (step, dir) for every complete checkpoint, oldest first.

    train_pytorch.py writes to `tmp_<step>` and renames to `<step>` on success, so
    a bare numeric directory containing model.safetensors is finished.
    """
    out = []
    for d in watch.iterdir():
        if not d.is_dir() or not d.name.isdigit():
            continue
        if (d / "model.safetensors").exists():
            out.append((int(d.name), d))
    return sorted(out)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--watch-dir", required=True)
    ap.add_argument("--repo-id", required=True)
    ap.add_argument("--interval", type=int, default=300, help="seconds between scans")
    ap.add_argument("--once", action="store_true", help="upload what exists, then exit")
    ap.add_argument("--private", action="store_true")
    args = ap.parse_args()

    # An invalid HF_TOKEN is exported in this environment and shadows the token
    # file; see the environment gotchas in HANDOFF.md.
    os.environ.pop("HF_TOKEN", None)
    from huggingface_hub import HfApi

    api = HfApi()
    watch = pathlib.Path(args.watch_dir)

    api.create_repo(args.repo_id, repo_type="model", private=args.private, exist_ok=True)
    print(f"repo ready: https://huggingface.co/{args.repo_id} (private={args.private})", flush=True)

    state = watch.parent / ".hf_uploaded.json"
    done: set[int] = set(json.loads(state.read_text())) if state.exists() else set()

    while True:
        if not watch.exists():
            print(f"waiting for {watch} to appear...", flush=True)
        else:
            for step, d in find_steps(watch):
                if step in done:
                    continue
                print(f"uploading step {step} from {d} ...", flush=True)
                try:
                    api.upload_folder(
                        folder_path=str(d),
                        path_in_repo=f"step-{step}",
                        repo_id=args.repo_id,
                        repo_type="model",
                        commit_message=f"checkpoint at step {step}",
                    )
                except Exception as e:  # noqa: BLE001 — keep watching after a Hub failure
                    print(f"  FAILED step {step}: {type(e).__name__}: {e}", flush=True)
                    continue
                done.add(step)
                state.write_text(json.dumps(sorted(done)))
                print(f"  done: https://huggingface.co/{args.repo_id}/tree/main/step-{step}", flush=True)

        if args.once:
            return 0
        time.sleep(args.interval)


if __name__ == "__main__":
    sys.exit(main())