# llmgw's LoRA trainer (BACKLOG F4), run by the node agent with an installed
# engine's Python (its torch and transformers) and a small overlay (peft,
# accelerate). Reads a job, trains, writes the PEFT adapter, and says how it
# goes as one JSON object a line on standard output.
import json, math, os, random, sys, threading, subprocess, time

def say(**kw):
    print(json.dumps(kw), flush=True)

job = json.load(open(sys.argv[1], encoding="utf-8"))
p = {"rank": 16, "alpha": 32, "dropout": 0.05, "epochs": 6, "lr": 2e-4, "batch": 4, "max_len": 512, "seed": 0}
p.update(job.get("params") or {})

say(state="preparing", message="loading libraries")
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

random.seed(p["seed"]); torch.manual_seed(p["seed"])
if not torch.cuda.is_available():
    raise SystemExit("no CUDA device is visible to the trainer")

say(state="preparing", message="fetching " + job["base_repo"])
base_dir = snapshot_download(job["base_repo"], revision=job.get("base_revision") or None,
                             allow_patterns=["*.json", "*.safetensors", "*.txt", "*.model", "*.tiktoken"])
say(state="preparing", message="loading the base model", base_dir=base_dir)

# bf16 where the card has it; fp16 with a loss scaler where it does not
# (Turing and older).
bf16 = torch.cuda.is_bf16_supported()
dtype = torch.bfloat16 if bf16 else torch.float16
tok = AutoTokenizer.from_pretrained(base_dir)
if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(base_dir, dtype=dtype).to("cuda")

rows = [json.loads(l) for l in open(job["data"], encoding="utf-8") if l.strip()]
rows = [r for r in rows if r.get("q") and r.get("a")]
if not rows:
    raise SystemExit("no {q, a} pairs in the data")

def encode(r):
    prompt = tok.apply_chat_template([{"role": "user", "content": r["q"]}], add_generation_prompt=True, tokenize=False)
    full = tok.apply_chat_template([{"role": "user", "content": r["q"]}, {"role": "assistant", "content": r["a"]}], tokenize=False)
    pi = tok(prompt, add_special_tokens=False)["input_ids"]
    fi = tok(full, add_special_tokens=False)["input_ids"][: p["max_len"]]
    return fi, ([-100] * len(pi) + fi[len(pi):])[: len(fi)]

data = [encode(r) for r in rows]

model.gradient_checkpointing_enable()
model.enable_input_require_grads()
model = get_peft_model(model, LoraConfig(r=int(p["rank"]), lora_alpha=int(p["alpha"]), lora_dropout=float(p["dropout"]),
    task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
trainable = [x for x in model.parameters() if x.requires_grad]
for x in trainable:  # fp32 trainable weights, so a loss scaler can unscale them
    x.data = x.data.float()

peak = {"mib": 0}
stop = threading.Event()
def watch():
    while not stop.is_set():
        try:
            out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"],
                                 capture_output=True, text=True, timeout=5).stdout
            for line in out.splitlines():
                pid, mib = [v.strip() for v in line.split(",")]
                if int(pid) == os.getpid():
                    peak["mib"] = max(peak["mib"], int(mib))
        except Exception:
            pass
        time.sleep(0.5)
threading.Thread(target=watch, daemon=True).start()

opt = torch.optim.AdamW(trainable, lr=float(p["lr"]), weight_decay=0.0)
scaler = torch.amp.GradScaler("cuda", enabled=not bf16)
bs, epochs = int(p["batch"]), int(p["epochs"])
steps = epochs * math.ceil(len(data) / bs)
warm = max(1, steps // 10)
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / warm) * max(0.0, 1 - s / steps))
model.train(); step = 0; last = time.time()
for ep in range(epochs):
    random.shuffle(data)
    for i in range(0, len(data), bs):
        batch = data[i:i + bs]; n = max(len(x[0]) for x in batch)
        ids = torch.tensor([x[0] + [tok.pad_token_id] * (n - len(x[0])) for x in batch], device="cuda")
        lab = torch.tensor([x[1] + [-100] * (n - len(x[1])) for x in batch], device="cuda")
        att = torch.tensor([[1] * len(x[0]) + [0] * (n - len(x[0])) for x in batch], device="cuda")
        with torch.autocast("cuda", dtype=dtype):
            loss = model(input_ids=ids, attention_mask=att, labels=lab).loss
        scaler.scale(loss).backward(); scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(trainable, 1.0)
        scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True); sched.step(); step += 1
        if time.time() - last > 2 or step == steps:
            last = time.time()
            say(state="running", message="training", step=step, steps=steps, loss=round(loss.item(), 4), peak_mib=peak["mib"])

model.save_pretrained(job["out"])
stop.set()
say(state="trained", step=step, steps=steps, peak_mib=peak["mib"], base_dir=base_dir)