"""Train Puck's character LoRA on Holotron-12B, on Modal.

TRL SFT + PEFT LoRA over the 162-example chat curriculum (build_dataset.py).
Text-only LoRA on the language side of the VLM — character/voice, not vision.
Adapter is saved to a Modal volume (no HF token in env); publish later.

  cd molt && uv run build_dataset.py
  modal run --detach train_modal.py              # .spawn() inside → truly detached
  modal volume get puck-lora /puck-holotron-12b-lora ./out

⚠️ BLOCKED (2026-06-07): Hcompany/Holotron-12B's published trust_remote_code
modeling.py imports a `_fully_shard.py` that isn't in the repo — an H Company
packaging bug in their *training* path (vLLM inference is unaffected, which is
why the vision endpoint works). transformers can't load it for TRL.
Options when revisiting: (a) wait for H Company to publish the missing file;
(b) train the character LoRA on the BASE nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL
(still Nemotron, still ≤32B, dataset is base-agnostic — voice doesn't need
Holotron's CUA tuning); (c) stub _fully_shard.py if it's a no-op FSDP helper.
Character already lands well via the enriched prompt, so the LoRA (Well-Tuned
badge) is lower priority than vision.

Hybrid-Mamba caveat once unblocked: target_modules='all-linear', gradient
checkpointing off (Mamba layers dislike it)."""

import json
from pathlib import Path

import modal

MODEL = "Hcompany/Holotron-12B"
HERE = Path(__file__).resolve().parent

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch",
        "transformers>=4.48",
        "trl>=0.12",
        "peft>=0.14",
        "datasets",
        "accelerate",
        "huggingface_hub[hf_transfer]",
        "sentencepiece",
        "einops",  # hybrid-Mamba modeling code often needs it
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

vol = modal.Volume.from_name("puck-lora", create_if_missing=True)
hf_cache = modal.Volume.from_name("puck-hf-cache", create_if_missing=True)

app = modal.App("puck-train")


@app.function(
    image=image,
    gpu="A100-80GB",
    timeout=60 * 60,
    volumes={"/adapter": vol, "/root/.cache/huggingface": hf_cache},
)
def train(records: list[dict]):
    import torch
    from datasets import Dataset
    from peft import LoraConfig
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from trl import SFTConfig, SFTTrainer

    # conversational dataset → TRL applies the chat template itself
    ds = Dataset.from_list([{"messages": r["messages"]} for r in records])

    tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto"
    )

    peft_cfg = LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules="all-linear",  # robust across the hybrid's linear layers
    )

    cfg = SFTConfig(
        output_dir="/adapter/run",
        num_train_epochs=4,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=2e-4,
        warmup_ratio=0.05,
        logging_steps=5,
        save_strategy="epoch",
        bf16=True,
        gradient_checkpointing=False,  # Mamba layers + checkpointing don't mix
        max_length=1024,
        report_to="none",
    )

    trainer = SFTTrainer(model=model, args=cfg, train_dataset=ds, peft_config=peft_cfg)
    trainer.train()

    out = "/adapter/puck-holotron-12b-lora"
    trainer.save_model(out)
    tok.save_pretrained(out)
    vol.commit()
    print(f"saved adapter → {out}")
    return out


@app.local_entrypoint()
def main():
    records = [
        json.loads(line) for line in (HERE / "data" / "sft.jsonl").read_text().splitlines()
    ]
    print(f"training on {len(records)} examples")
    # .spawn() so a detached run survives the local caller disconnecting
    # (.remote() is synchronous and Modal cancels it when the CLI exits).
    call = train.spawn(records)
    print(f"spawned: {call.object_id}  —  modal volume get puck-lora /puck-holotron-12b-lora ./out")