scripts/train_on_hf.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "transformers==4.56.2",
#   "trl==0.22.2",
#   "datasets",
#   "peft",
#   "accelerate",
#   "bitsandbytes",
#   "unsloth",
#   "openenv-core",
#   "fastapi",
#   "uvicorn",
#   "pydantic",
#   "matplotlib",
#   "huggingface_hub",
# ]
# ///
"""
End-to-end training job for HF Jobs.

Submit from local machine with:
    hf jobs uv run --flavor a10g-large --secrets HF_TOKEN scripts/train_on_hf.py

What it does (no babysitting required):
  1. Clone rhythm_env from HF Space (gets latest meta-RL code from main)
  2. Generate dataset (continuous profiles, hint_fraction=0.15)
  3. Train Qwen 2.5-3B + LoRA rank 8 via GRPO (1500 steps)
  4. Run eval on all 3 conditions (discrete, in-dist, OOD)
  5. Generate all 5 plots from log_history
  6. Upload trained model + plots + eval JSON to a new HF Hub model repo

Override defaults via env vars:
    MAX_STEPS, NUM_EPISODES, LORA_RANK, BETA, MODEL_REPO

Estimated cost on a10g-large at $1.50/hr: ~$3 for 1500 steps (~2h).
"""

import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

# ---------------------------------------------------------------------------
# Config (overridable via env vars)
# ---------------------------------------------------------------------------
REPO_URL = os.environ.get("REPO_URL", "https://huggingface.co/spaces/InosLihka/rhythm_env")
WORK_DIR = "/tmp/rhythm_env"
OUTPUT_DIR = "/tmp/rhythm_env/outputs/rhythmenv_meta_trained"
PLOTS_DIR = "/tmp/rhythm_env/plots"

# FAST_MODE preset: ~10-15 min iteration on A100 large.
# Use for hyperparameter sweeps and pipeline debugging.
FAST_MODE = os.environ.get("FAST_MODE", "0") == "1"

if FAST_MODE:
    # Iter 3 preset: 800 steps + 8 generations + LoRA 16 to escape mode collapse for real
    DEFAULTS = dict(MAX_STEPS=800, NUM_EPISODES=200, MAX_SAMPLES=2000,
                    NUM_GENERATIONS=8, LORA_RANK=16, BETA=0.04,
                    LEARNING_RATE=5e-5, EVAL_EPISODES=2)
else:
    DEFAULTS = dict(MAX_STEPS=2000, NUM_EPISODES=400, MAX_SAMPLES=4000,
                    NUM_GENERATIONS=8, LORA_RANK=16, BETA=0.04,
                    LEARNING_RATE=5e-5, EVAL_EPISODES=5)

MAX_STEPS = int(os.environ.get("MAX_STEPS", str(DEFAULTS["MAX_STEPS"])))
NUM_EPISODES = int(os.environ.get("NUM_EPISODES", str(DEFAULTS["NUM_EPISODES"])))
MAX_SAMPLES = int(os.environ.get("MAX_SAMPLES", str(DEFAULTS["MAX_SAMPLES"])))
NUM_GENERATIONS = int(os.environ.get("NUM_GENERATIONS", str(DEFAULTS["NUM_GENERATIONS"])))
LORA_RANK = int(os.environ.get("LORA_RANK", str(DEFAULTS["LORA_RANK"])))
BETA = float(os.environ.get("BETA", str(DEFAULTS["BETA"])))
LEARNING_RATE = float(os.environ.get("LEARNING_RATE", str(DEFAULTS["LEARNING_RATE"])))
EVAL_EPISODES = int(os.environ.get("EVAL_EPISODES", str(DEFAULTS["EVAL_EPISODES"])))

# Each iteration uploads to a unique repo if MODEL_REPO_SUFFIX is set
SUFFIX = os.environ.get("MODEL_REPO_SUFFIX", "")
DEFAULT_REPO = "InosLihka/rhythm-env-meta-trained" + (f"-{SUFFIX}" if SUFFIX else "")
MODEL_REPO = os.environ.get("MODEL_REPO", DEFAULT_REPO)

print(f"=== Run config ===")
print(f"  FAST_MODE: {FAST_MODE}")
print(f"  MAX_STEPS={MAX_STEPS}, NUM_EPISODES={NUM_EPISODES}, MAX_SAMPLES={MAX_SAMPLES}")
print(f"  NUM_GENERATIONS={NUM_GENERATIONS}, LORA_RANK={LORA_RANK}, BETA={BETA}")
print(f"  LEARNING_RATE={LEARNING_RATE}, EVAL_EPISODES={EVAL_EPISODES}")
print(f"  MODEL_REPO={MODEL_REPO}")
print()


def run(cmd: list[str], **kw):
    """Run subprocess with logging."""
    print(f"\n>>> {' '.join(cmd) if isinstance(cmd, list) else cmd}", flush=True)
    subprocess.run(cmd, check=True, **kw)


def main():
    # ---------------------------------------------------------------
    # 1. Clone the rhythm_env repo
    # ---------------------------------------------------------------
    if Path(WORK_DIR).exists():
        shutil.rmtree(WORK_DIR)
    run(["git", "clone", REPO_URL, WORK_DIR])
    os.chdir(WORK_DIR)
    sys.path.insert(0, WORK_DIR)
    sys.path.insert(0, os.path.join(WORK_DIR, "training"))

    # Verify meta-RL code is present
    dataset_py = Path("training/dataset.py").read_text()
    assert "profile_mode" in dataset_py, "Cloned repo doesn't have meta-RL code"
    print("OK: meta-RL code present in cloned repo")

    # ---------------------------------------------------------------
    # 2. Train
    # ---------------------------------------------------------------
    # MODEL_NAME env var lets us refine an existing trained model (e.g. SFT'd
    # checkpoint on HF Hub) instead of starting from the base Qwen. Default
    # is the original base model.
    base_model = os.environ.get("MODEL_NAME", "unsloth/Qwen2.5-3B-Instruct")

    train_args = [
        "python", "training/train.py",
        "--model_name", base_model,
        "--max_steps", str(MAX_STEPS),
        "--num_episodes", str(NUM_EPISODES),
        "--max_samples", str(MAX_SAMPLES),
        "--num_generations", str(NUM_GENERATIONS),
        "--lora_rank", str(LORA_RANK),
        "--beta", str(BETA),
        "--learning_rate", str(LEARNING_RATE),
        "--output_dir", OUTPUT_DIR,
    ]
    print(f"Starting from model: {base_model}")
    run(train_args)

    # ---------------------------------------------------------------
    # 3. Eval (3 conditions: discrete-3 / in-dist / OOD)
    # ---------------------------------------------------------------
    eval_args = [
        "python", "training/inference_eval.py",
        "--model_path", OUTPUT_DIR,
        "--num_episodes", str(EVAL_EPISODES),
        "--output_file", "eval_results.json",
    ]
    run(eval_args)

    # ---------------------------------------------------------------
    # 4. Generate plots from saved log_history
    # ---------------------------------------------------------------
    Path(PLOTS_DIR).mkdir(exist_ok=True)
    log_path = os.path.join(OUTPUT_DIR, "log_history.json")
    if Path(log_path).exists():
        run(["python", "scripts/plot_from_log.py", "--log", log_path, "--out", PLOTS_DIR])
    else:
        print(f"WARNING: log_history.json not found at {log_path}")

    # ---------------------------------------------------------------
    # 5. Upload everything to HF Hub
    # ---------------------------------------------------------------
    token = os.environ.get("HF_TOKEN")
    if not token:
        print("WARNING: HF_TOKEN not set, skipping upload")
        print(f"Outputs in: {OUTPUT_DIR}")
        return

    from huggingface_hub import HfApi, login
    login(token=token)
    api = HfApi()
    api.create_repo(MODEL_REPO, exist_ok=True, repo_type="model")

    # Upload trained model + config + log_history
    api.upload_folder(
        folder_path=OUTPUT_DIR,
        repo_id=MODEL_REPO,
        repo_type="model",
        commit_message=f"Trained {MAX_STEPS}-step GRPO meta-RL agent",
    )

    # Upload eval JSON
    api.upload_file(
        path_or_fileobj="eval_results.json",
        path_in_repo="eval_results.json",
        repo_id=MODEL_REPO,
        repo_type="model",
    )

    # Upload plots if generated
    if Path(PLOTS_DIR).exists() and any(Path(PLOTS_DIR).iterdir()):
        api.upload_folder(
            folder_path=PLOTS_DIR,
            path_in_repo="plots",
            repo_id=MODEL_REPO,
            repo_type="model",
        )

    print()
    print("=" * 60)
    print("DONE")
    print(f"  Trained model: https://huggingface.co/{MODEL_REPO}")
    print(f"  Eval JSON:     https://huggingface.co/{MODEL_REPO}/blob/main/eval_results.json")
    print(f"  Plots:         https://huggingface.co/{MODEL_REPO}/tree/main/plots")
    print("=" * 60)


if __name__ == "__main__":
    main()