scripts/deploy_training.py

"""Deploy FATHOM training job to a GPU-powered HF Space.

Creates a Docker Space with L4 GPU ($0.80/hr) that:
1. Installs deps
2. Generates dataset (1000 train / 200 eval / 500 SFT)
3. Runs SFT warm-start on 0.5B model
4. Runs GRPO training (400 steps)
5. Pushes fine-tuned model to HF Hub

Usage: python scripts/deploy_training.py
Cost: ~$2-4 for 0.5B smoke, ~$8-15 for 1.5B full run
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

if sys.platform == "win32":
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")

HF_TOKEN = os.getenv("HF_TOKEN")
SPACE_NAME = os.environ.get("FATHOM_TRAIN_SPACE", "Pratham-math/fathom-train")
REPO_ROOT = Path(__file__).parent.parent


def deploy():
    if not HF_TOKEN:
        print("ERROR: Set HF_TOKEN environment variable first")
        sys.exit(1)

    from huggingface_hub import HfApi

    api = HfApi(token=HF_TOKEN)
    me = api.whoami()
    print(f"Logged in as: {me['name']}")

    # Create GPU Space with L4 ($0.80/hr, 24GB VRAM)
    print(f"Creating GPU Space: {SPACE_NAME} ...")
    api.create_repo(
        repo_id=SPACE_NAME,
        repo_type="space",
        space_sdk="docker",
        private=False,
        exist_ok=True,
    )

    # Set Space hardware to L4 GPU
    try:
        api.request_space_hardware(
            repo_id=SPACE_NAME,
            hardware="l4x1",  # L4 24GB - $0.80/hr
        )
        print("Hardware set to L4 (24GB VRAM, $0.80/hr)")
    except Exception as e:
        print(f"Hardware request note: {e}")
        print("You may need to set hardware manually at https://huggingface.co/spaces/{SPACE_NAME}/settings")

    # Set HF_TOKEN as a Space secret (needed to push the fine-tuned model)
    try:
        api.add_space_secret(repo_id=SPACE_NAME, key="HF_TOKEN", value=HF_TOKEN)
        print("HF_TOKEN secret set")
    except Exception as e:
        print(f"Secret set note: {e}")

    # Collect ALL files needed for training
    uploads = []

    # Training Dockerfile
    uploads.append((REPO_ROOT / "space" / "Dockerfile.train", "Dockerfile"))

    # Space README with metadata
    uploads.append((REPO_ROOT / "space" / "README_train.md", "README.md"))

    # Run script
    uploads.append((REPO_ROOT / "scripts" / "run_training.py", "run_training.py"))

    # pyproject.toml
    uploads.append((REPO_ROOT / "pyproject.toml", "pyproject.toml"))

    # All Python packages
    for pkg in ["train", "rewards", "data", "env", "configs"]:
        pkg_dir = REPO_ROOT / pkg
        if not pkg_dir.exists():
            continue
        for p in pkg_dir.rglob("*"):
            if p.is_file() and not p.name.startswith(".") and "__pycache__" not in str(p):
                rel = str(p.relative_to(REPO_ROOT)).replace("\\", "/")
                uploads.append((p, rel))

    # seeds.json
    seeds = REPO_ROOT / "data" / "seeds.json"
    if seeds.exists():
        uploads.append((seeds, "data/seeds.json"))

    print(f"Uploading {len(uploads)} files...")
    for local, remote in uploads:
        if local.exists():
            api.upload_file(
                path_or_fileobj=str(local),
                path_in_repo=remote,
                repo_id=SPACE_NAME,
                repo_type="space",
                commit_message=f"Train deploy: {remote}",
            )
            print(f"  OK  {remote}")
        else:
            print(f"  SKIP  {local}")

    hf_url = f"https://huggingface.co/spaces/{SPACE_NAME}"
    print("=" * 60)
    print("Training Space deployed!")
    print(f"  Monitor: {hf_url}")
    print(f"  Logs:    {hf_url}?logs=container")
    print()
    print("The Space will:")
    print("  1. Build Docker image (~3 min)")
    print("  2. Generate dataset")
    print("  3. Run SFT warm-start on 0.5B model (~10 min)")
    print("  4. Run GRPO 400 steps (~2-3 hrs on L4)")
    print("  5. Push fine-tuned model to Pratham-math/fathom-0.5b-grpo")
    print()
    print("Estimated cost: $2-4 for 0.5B run")
    print("=" * 60)


if __name__ == "__main__":
    deploy()