"""Train Dory1.0-4B on Beam traces on Modal (torchrun + DDP on B200s).

  modal run modal_app/train_dory.py::download_dory              # one-time: base weights -> Volume
  modal run modal_app/train_dory.py --gpus 1 --run smoke --extra "--max-steps 20"
  modal run modal_app/train_dory.py --gpus 8 --run r1 --data sft/sft_verified.jsonl

Data files are paths inside the `beam-dory-data` Volume (upload with `modal volume put`).
Checkpoints land in the `beam-dory-ckpt` Volume under runs/<run>/.
"""

import shlex
import subprocess
import time

import modal

APP = modal.App("beam-dory-train")
weights = modal.Volume.from_name("beam-weights", create_if_missing=True, version=2)
data = modal.Volume.from_name("beam-dory-data", create_if_missing=True, version=2)
ckpt = modal.Volume.from_name("beam-dory-ckpt", create_if_missing=True, version=2)
W, D, C = "/weights", "/data", "/ckpt"
DORY_REPO = "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16"

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch==2.14.1", "safetensors", "transformers>=5.20", "jinja2", "numpy",
                 "huggingface_hub[hf_xet]>=0.34")
    .env({"PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false"})
    .add_local_dir("dory", "/proj/dory")
    .add_local_dir("scripts", "/proj/scripts")
)


@APP.function(image=image, volumes={W: weights}, cpu=4, memory=16384, timeout=3600)
def download_dory():
    from huggingface_hub import snapshot_download

    snapshot_download(DORY_REPO, local_dir=f"{W}/dory-4b", allow_patterns=["*.json", "*.jinja", "*.safetensors"])
    weights.commit()
    print(subprocess.run(["ls", "-la", f"{W}/dory-4b"], capture_output=True, text=True).stdout)


@APP.function(image=image, gpu=["B200:8", "H200:8", "B200:4"], volumes={W: weights, D: data, C: ckpt},
              timeout=12 * 3600, cpu=32, memory=262144)
def train(run: str, data_files: list, nproc: int = 0, extra: str = ""):
    """nproc=0 uses every GPU in the container (whichever fallback config was allocated)."""
    if not nproc:
        out_ = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, text=True)
        names = [l for l in out_.stdout.splitlines() if l.strip()]
        nproc = len(names)
        print(f"got {nproc}x {names[0]}", flush=True)
    if "--grad-accum" not in extra:  # keep the global batch of 8 GPUs x 1 sequence on smaller fallbacks
        extra += f" --grad-accum {max(1, 8 // nproc)}"
    out = f"{C}/runs/{run}"
    cmd = ["torchrun", "--standalone", f"--nproc_per_node={nproc}", "/proj/scripts/train_dory.py",
           "--model", f"{W}/dory-4b", "--out", out, "--data", *[f"{D}/{f}" for f in data_files],
           *shlex.split(extra)]
    print("cmd:", " ".join(cmd), flush=True)
    t0 = time.time()
    proc = subprocess.run(cmd, cwd="/proj")
    ckpt.commit()
    print(f"exit {proc.returncode} after {(time.time() - t0) / 60:.1f} min", flush=True)
    if proc.returncode != 0:
        raise RuntimeError(f"training failed with exit code {proc.returncode}")


@APP.local_entrypoint()
def main(run: str, data: str = "sft/sft_verified.jsonl", gpus: int = 8, extra: str = ""):
    """Interactive runs (smoke tests). Long runs: deploy, then spawn via scripts/spawn_train.py."""
    files = [f.strip() for f in data.split(",") if f.strip()]
    train.with_options(gpu=f"B200:{gpus}").remote(run, files, gpus, extra)