| |
| """Watch a training run's hf/ dir and upload each new step-N export to a private HF |
| repo, so ephemeral-disk checkpoints survive a machine kill and are eval-able anywhere. |
| |
| train_sft_qwen3.py writes a bf16 hf/step-N/ every steps_per_hf_export. This uploads each |
| one (once it has finished writing) to <repo> under path_in_repo=step-N/. One repo per |
| experiment; steps are subfolders. Idempotent via a local .pushed_steps.json marker, so a |
| restart of this watcher (or of the whole run) re-uploads nothing. |
| |
| Run in the background alongside training (run_sft.sh does this for you): |
| python push_exports.py --out <rundir> --repo fzzhang/<exp>_sft --until-step 2000 |
| |
| HF auth: uses your `huggingface-cli login` token (~/.cache/huggingface/token). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
|
|
| def step_dirs(hf_dir: Path) -> dict[int, Path]: |
| out: dict[int, Path] = {} |
| if not hf_dir.is_dir(): |
| return out |
| for d in hf_dir.glob("step-*"): |
| if d.is_dir(): |
| try: |
| out[int(d.name.split("-", 1)[1])] = d |
| except (IndexError, ValueError): |
| pass |
| return out |
|
|
|
|
| def is_stable(d: Path, stable_sec: float) -> bool: |
| """True if no file in the dir was modified within the last stable_sec (write finished).""" |
| files = [p for p in d.rglob("*") if p.is_file()] |
| if not files: |
| return False |
| return (time.time() - max(p.stat().st_mtime for p in files)) > stable_sec |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description="Upload hf/step-N exports to HF as they land") |
| ap.add_argument("--out", required=True, help="training run dir (contains hf/)") |
| ap.add_argument("--repo", required=True, help="target HF repo id, e.g. fzzhang/<exp>_sft") |
| ap.add_argument("--poll", type=int, default=120, help="seconds between scans") |
| ap.add_argument("--stable-sec", type=float, default=60.0, |
| help="a step dir must be idle this long before upload (guards partial writes)") |
| ap.add_argument("--every", type=int, default=1, |
| help="upload every Nth export (exports are every 100 steps; 1=all). " |
| "The --until-step export is always uploaded regardless.") |
| ap.add_argument("--until-step", type=int, default=0, |
| help="exit after this step is pushed (0 = run forever / until killed)") |
| ap.add_argument("--once", action="store_true", |
| help="scan once, push all unpushed steps (ignore --stable-sec), then exit") |
| args = ap.parse_args() |
|
|
| from huggingface_hub import HfApi |
|
|
| api = HfApi() |
| api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True) |
|
|
| hf_dir = Path(args.out) / "hf" |
| marker = Path(args.out) / ".pushed_steps.json" |
| pushed = set(json.loads(marker.read_text())) if marker.exists() else set() |
| print(f"[push] {hf_dir} -> {args.repo} (private); already pushed: {sorted(pushed)}", flush=True) |
|
|
| while True: |
| for step in sorted(step_dirs(hf_dir)): |
| if step in pushed: |
| continue |
| if args.every > 1 and (step // 100) % args.every != 0 and step != args.until_step: |
| continue |
| d = step_dirs(hf_dir)[step] |
| if not args.once and not is_stable(d, args.stable_sec): |
| continue |
| print(f"[push] uploading step-{step} ...", flush=True) |
| try: |
| api.upload_folder(folder_path=str(d), repo_id=args.repo, repo_type="model", |
| path_in_repo=f"step-{step}", commit_message=f"add step-{step}") |
| pushed.add(step) |
| marker.write_text(json.dumps(sorted(pushed))) |
| print(f"[push] DONE step-{step} -> {args.repo}/step-{step}", flush=True) |
| except Exception as e: |
| print(f"[push] FAILED step-{step}: {type(e).__name__}: {e} (retry next poll)", flush=True) |
| if args.once: |
| print("[push] --once sweep complete; exiting.", flush=True) |
| return |
| if args.until_step and args.until_step in pushed: |
| print(f"[push] final step-{args.until_step} pushed; exiting.", flush=True) |
| return |
| time.sleep(args.poll) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
