"""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: //... Orbax checkpoint -- fp32 MASTER weights + optimizer state /model.safetensors bf16 INFERENCE weights (flat, ~half the size) /resume.json step, per-source HF stream positions, RNG, elapsed, val history /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]})")