#!/bin/bash
# Full SFT + GRPO pipeline inside HF Job. Repo cloned at /w.
# Requires secrets: HF_TOKEN, WANDB_API_KEY
set -e
cd /w
export PYTHONPATH="/w:${PYTHONPATH}"

# Force single-GPU training. On a100-large (1x A100-80GB) this is a no-op,
# but kept as a defensive guardrail in case we ever fall back to a 2x flavor:
# TRL's SFTTrainer entropy logging path on multi-GPU crashes with
# `RuntimeError: tensor a (2) must match tensor b (4) at non-singleton
# dimension 0` because per-rank logits get mixed with the full-batch
# attention mask. Hiding extra GPUs completely sidesteps this.
export CUDA_VISIBLE_DEVICES=0

# Install huggingface_hub first so we can download data files (parity with job_smoke.sh)
pip install -q 'huggingface_hub>=0.28'

# Fetch large data files from HF (uploaded separately via hf upload)
mkdir -p data
python -c "from huggingface_hub import hf_hub_download; [hf_hub_download(repo_id='Pratham-math/fathom-code', filename=f'data/{f}', local_dir='/w', token='${HF_TOKEN}') for f in ['train.jsonl','eval.jsonl','sft_traces.jsonl']]"
ls -la data/

# Base runtime deps (these don't conflict)
pip install -q openenv-core fastapi 'uvicorn[standard]' pydantic RestrictedPython tiktoken httpx
pip install -q hydra-core omegaconf wandb 'huggingface_hub>=0.28' tyro

# Core ML deps.
# NOTE: trl 1.2 imports `is_trackio_available` from transformers, which is not
# present in 4.49.0. Use a newer transformers line in jobs.
# bitsandbytes pin notes:
#  - 0.45.1: ships cu124 only — breaks under vllm's cu128 torch.
#  - 0.47.x: cu128 + cu124 dual-wheel, but vllm's BitsAndBytesLinearMethod
#    gates on `bnb >= 0.48.1` (vllm/model_executor/layers/quantization/
#    bitsandbytes.py:_check_bitsandbytes_version) and raises ImportError
#    during `LLM(...)` init when used with a 4-bit model.
#  - >=0.48.1: satisfies both vllm's gate and our cu128 wheel needs.
pip install -q transformers==4.56.2 accelerate==1.5.2 peft==0.14.0 'bitsandbytes>=0.48.1'
pip install -q datasets==4.7.0
# Keep TRL pinned for OpenEnv path but bypass resolver deadlock with datasets pin.
pip install -q --no-deps trl==1.2.0
pip install -q vllm==0.18.0

# Optional transitive deps often required by quantized loaders / datasets stack
pip install -q safetensors sentencepiece einops scipy xxhash protobuf pyyaml fsspec aiohttp dill multiprocess pyarrow requests filelock packaging tokenizers regex tqdm

# Plotting deps for scripts/make_plots.py — NOT in the base pytorch image.
# Without these the plot step crashes with ModuleNotFoundError and judges lose
# the reward/loss curves on the model repo (they're regenerable locally from
# the job log via scripts/parse_log_to_plots.py, but remote-side is the
# happy path we want to keep working).
pip install -q matplotlib

# flash-attn removed: source build is the largest single memory spike during install.
# vLLM + transformers fall back to PyTorch SDPA without it (small throughput cost).
echo "flash-attn intentionally skipped to avoid build-time OOM"

nohup uvicorn env.server.app:app --host 0.0.0.0 --port 8001 > env.log 2>&1 &
sleep 10
python -c "import urllib.request, sys; sys.exit(0 if urllib.request.urlopen('http://localhost:8001/healthz', timeout=5).status==200 else 1)" || (echo "env down"; cat env.log; exit 1)

# Smoke first — fail fast if pipeline broken
python -m train.smoke_test --env-url http://localhost:8001
grep -q "VERDICT: GO" outputs/smoke/SMOKE_RESULT.md || (cat outputs/smoke/SMOKE_RESULT.md; exit 1)

# SFT (Path 1 — production 1.5B). qwen_1_5b config points at
# unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit per STACK §3.1.
# NOTE: use initialize_config_dir (absolute path) — heredocs have no caller file,
# so the relative `../configs` resolves against CWD and breaks. Absolute is robust.
python <<'PY'
from hydra import initialize_config_dir, compose
from train.model_load import load_model_and_tokenizer
from train.sft import run_sft
with initialize_config_dir(config_dir="/w/configs", version_base="1.3"):
    cfg = compose(config_name="config", overrides=["model=qwen_1_5b","train=sft"])
m, t = load_model_and_tokenizer(cfg)
print("SFT adapter:", run_sft(cfg, m, t))
PY

# GRPO 200 steps (Path 1 — production 1.5B on A100-80GB).
# Memory budget on a100-large (~80GB):
#  - 1.5B 4-bit + LoRA-r16 weights + optim state ≈ 6-8 GB
#  - vLLM colocate KV cache @ num_generations=8, max_prompt=4096, max_completion=1024 ≈ 25-35 GB
#  - vllm_gpu_memory_utilization=0.50 → leaves ~30 GB for trainer + activations
# Ramping max_steps from 50 → 200 to give the reward curve room to climb
# (50 steps was a smoke; 200 is enough to show the gate-crossing transition
# from format_gate=0 → format_gate=1 → composite reward rising).
python <<'PY'
from hydra import initialize_config_dir, compose
from train.model_load import load_model_and_tokenizer
from train.grpo import run_grpo
from rewards.compose import make_reward_fn
with initialize_config_dir(config_dir="/w/configs", version_base="1.3"):
    cfg = compose(config_name="config", overrides=[
        "model=qwen_1_5b",
        "train=grpo",
        "train.max_steps=200",
        "train.num_generations=8",
        "train.max_prompt_length=4096",
        "train.max_completion_length=1024",
        "train.vllm_gpu_memory_utilization=0.50",
    ])
m, t = load_model_and_tokenizer(cfg)
print("GRPO merged:", run_grpo(cfg, m, t, make_reward_fn(cfg.reward), "http://localhost:8001"))
PY

# Generate reward + loss curve PNGs from the training run (judges need these)
python /w/scripts/make_plots.py || echo "plot generation failed (non-fatal)"

# Push artifacts + plots to a PUBLIC model repo (judges must access)
HF_USER=$(python -c "from huggingface_hub import HfApi; print(HfApi(token='${HF_TOKEN}').whoami()['name'])")
python <<PY
from huggingface_hub import HfApi
api = HfApi(token="${HF_TOKEN}")
repo = "${HF_USER}/fathom-1.5b-grpo"
api.create_repo(repo, repo_type="model", exist_ok=True, private=False)
import os
for sub in ["sft_adapter", "grpo_merged_16bit", "plots"]:
    if os.path.isdir(f"outputs/{sub}"):
        api.upload_folder(folder_path=f"outputs/{sub}", path_in_repo=sub, repo_id=repo, repo_type="model")
        print(f"Uploaded outputs/{sub} -> {repo}/{sub}")
print(f"Trained model + plots: https://huggingface.co/{repo}")
PY
echo "DONE"