cortex/hfsync.py

"""Hugging Face Hub checkpoint sync for the Phase-1 Kaggle run (9h-session friendly).

A session trains until its wall-clock budget, then uploads EVERYTHING needed to resume
in the next session to one model repo:

  <repo>/<step>/...          Orbax checkpoint -- fp32 MASTER weights + optimizer state
  <repo>/model.safetensors   bf16 INFERENCE weights (flat, ~half the size)
  <repo>/resume.json         step, per-source HF stream positions, RNG, elapsed, val history
  <repo>/val/*.npy           cached held-out validation batches

Resume downloads the whole repo back into out_dir: Orbax restores master+optimizer, the
loader replays resume.json's per-source state_dicts (HF-native fast resume), and the val
cache is reused. Tokens are read from the environment (HF_TOKEN) -- never hardcoded.
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import numpy as np


def hf_login() -> str | None:
    """Authenticate from $HF_TOKEN (set via Kaggle Secrets). Returns the token or None.
    Public dataset streaming works without it; uploads need it."""
    tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not tok:
        print("[hfsync] HF_TOKEN not set -> HF upload/download disabled (public streaming still works)")
        return None
    from huggingface_hub import login
    login(token=tok, add_to_git_credential=False)
    return tok


# ---- bf16 inference weights ----------------------------------------------------------
def _flat(params) -> dict:
    """Flatten an nnx param pytree to {stable_path: leaf} (sorted-stable key order)."""
    from jax.tree_util import tree_flatten_with_path, keystr
    leaves, _ = tree_flatten_with_path(params)
    return {f"{i:04d}|{keystr(path)}": leaf for i, (path, leaf) in enumerate(leaves)}


def save_bf16_safetensors(params, path) -> None:
    import jax.numpy as jnp
    from safetensors.flax import save_file
    # jnp.asarray: leaves may be host numpy (host-master offload), not jax arrays
    save_file({k: jnp.asarray(v).astype(jnp.bfloat16) for k, v in _flat(params).items()}, str(path))


def load_bf16_into(params, path):
    """Load bf16 safetensors back into a params pytree of matching structure (inference)."""
    import jax
    from safetensors.flax import load_file
    flat = load_file(str(path))
    leaves, treedef = jax.tree_util.tree_flatten(params)
    ordered = [flat[k] for k in sorted(flat)]                  # keys are zero-padded by leaf index
    return jax.tree_util.tree_unflatten(treedef, ordered)


# ---- resume metadata -----------------------------------------------------------------
def _jsonable(o):
    if isinstance(o, dict):
        return {k: _jsonable(v) for k, v in o.items()}
    if isinstance(o, (list, tuple)):
        return [_jsonable(v) for v in o]
    if isinstance(o, (np.integer,)):
        return int(o)
    if isinstance(o, (np.floating,)):
        return float(o)
    if isinstance(o, np.ndarray):
        return o.tolist()
    return o


def write_resume(out_dir, info: dict) -> None:
    Path(out_dir, "resume.json").write_text(json.dumps(_jsonable(info), indent=2))


def read_resume(out_dir) -> dict | None:
    p = Path(out_dir, "resume.json")
    return json.loads(p.read_text()) if p.exists() else None


# ---- hub up/download -----------------------------------------------------------------
def upload(repo: str, out_dir, *, msg: str = "phase1 checkpoint") -> bool:
    if not os.environ.get("HF_TOKEN") and not os.environ.get("HUGGING_FACE_HUB_TOKEN"):
        print("[hfsync] no token -> skipping upload"); return False
    from huggingface_hub import HfApi, create_repo
    create_repo(repo, repo_type="model", exist_ok=True, private=True)
    HfApi().upload_folder(folder_path=str(out_dir), repo_id=repo, repo_type="model",
                          commit_message=msg, ignore_patterns=["*.tmp", "**/*.orbax-checkpoint-tmp-*"])
    print(f"[hfsync] uploaded {out_dir} -> {repo}")
    return True


def download(repo: str, out_dir) -> bool:
    """Pull ONLY the latest checkpoint into out_dir: the highest-numbered Orbax step folder
    + resume.json + the val cache (NOT model.safetensors, NOT older step folders).  Pulling the
    whole repo scaled with history and OOM'd Kaggle's /kaggle/working once several sessions'
    checkpoints had piled up (v32 downloaded ~16GB of old folders -> no room for the next save).
    Returns False if nothing is there yet."""
    from huggingface_hub import snapshot_download, list_repo_files
    from huggingface_hub.errors import RepositoryNotFoundError
    try:
        files = list_repo_files(repo_id=repo, repo_type="model")
        steps = sorted({int(f.split("/", 1)[0]) for f in files if f.split("/", 1)[0].isdigit()})
        allow = ["resume.json", "val/*", "val/**"]
        if steps:
            allow += [f"{steps[-1]}/*", f"{steps[-1]}/**"]      # latest Orbax step folder only
        snapshot_download(repo_id=repo, repo_type="model", local_dir=str(out_dir),
                          allow_patterns=allow)
        print(f"[hfsync] restored latest checkpoint (step {steps[-1] if steps else None}) "
              f"from {repo} -> {out_dir}")
        return True
    except RepositoryNotFoundError:
        print(f"[hfsync] repo {repo} not found yet -> fresh start")
        return False
    except Exception as e:
        print(f"[hfsync] download failed ({type(e).__name__}: {str(e)[:100]}) -> fresh start")
        return False


def purge_old_steps(repo: str, keep: int = 2) -> None:
    """Keep only the `keep` highest-numbered Orbax step folders in the repo; delete older ones so
    the repo (and the next resume's download) stays bounded.  Best-effort -- never fatal."""
    if not os.environ.get("HF_TOKEN") and not os.environ.get("HUGGING_FACE_HUB_TOKEN"):
        return
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        files = api.list_repo_files(repo, repo_type="model")
        steps = sorted({int(f.split("/", 1)[0]) for f in files if f.split("/", 1)[0].isdigit()})
        for s in steps[:-keep] if keep > 0 else steps:
            api.delete_folder(path_in_repo=str(s), repo_id=repo, repo_type="model",
                              commit_message=f"purge old checkpoint {s}")
            print(f"[hfsync] purged old checkpoint {s} from {repo}")
    except Exception as e:
        print(f"[hfsync] purge skipped ({type(e).__name__}: {str(e)[:100]})")