hf_backup.py

"""Background watcher: push the latest training checkpoint + log to the HF Hub
whenever the checkpoint changes (throttled). Run in tmux alongside training so
checkpoints survive pre-emption without SCP/laptop. Resume anywhere via hf_hub_download.

  HF_TOKEN=...  HF_CKPT_REPO=capotej/WikipediaGutenberg-0.5B \
  HF_CKPT_PATH=/root/nanoGPT/out-wiki-gutenberg-530m-1ep/ckpt.pt \
  HF_CKPT_LOG=/root/nanoGPT/train.log  python hf_backup.py
"""
import os, time
from huggingface_hub import HfApi

REPO    = os.environ["HF_CKPT_REPO"]
CKPT    = os.environ.get("HF_CKPT_PATH", "out-wiki-gutenberg-530m-1ep/ckpt.pt")
LOG     = os.environ.get("HF_CKPT_LOG",  "train.log")
MIN_INT = int(os.environ.get("HF_MIN_INTERVAL", "1800"))   # throttle: >=30 min between pushes

api = HfApi(token=os.environ["HF_TOKEN"])
api.create_repo(REPO, private=False, exist_ok=True)        # public repo
last_m, last_up = 0, 0
while True:
    try:
        if not os.path.exists(CKPT):
            time.sleep(60); continue
        m = os.path.getmtime(CKPT); now = time.time()
        if m != last_m and now - last_up >= MIN_INT:
            api.upload_file(path_or_fileobj=CKPT, path_in_repo="ckpt.pt", repo_id=REPO)
            try: api.upload_file(path_or_fileobj=LOG, path_in_repo="train.log", repo_id=REPO)
            except Exception: pass
            last_m, last_up = m, now
            print(f"[{time.ctime()}] pushed ckpt.pt + train.log -> {REPO}", flush=True)
    except Exception as e:
        print(f"[{time.ctime()}] upload error: {e}", flush=True)
    time.sleep(60)