"""GPT-U-20M pretraining on 2x RTX 3060 with DDP.

    scripts/run_train.sh                          # full run in tmux (auto-resumes from the newest checkpoint)
    SESSION=smoke scripts/run_train.sh --smoke    # 200-step smoke test on ~30M pilot tokens

bf16 autocast over fp32 master weights, AdamW + WSD schedule, z-loss, grad clip, torch.compile with an eager fallback.
Batches are drawn from a seeded permutation of token blocks indexed by (step, rank, micro-step), so resuming only needs
the step. Rank 0 appends to the jsonl log every 10 steps, checkpoints atomically every 500 steps, evaluates every
1,000, rewrites STATUS.json at each checkpoint and redraws the loss curve every 2,000.
"""
import argparse
import json
import math
import os
import shutil
import signal
import subprocess
import threading
import time
from collections import deque
from contextlib import nullcontext
from datetime import datetime, timedelta
from pathlib import Path

os.environ.setdefault("NCCL_P2P_DISABLE", "1")  # no NVLink: P2P under WSL2 hangs the first all-reduce silently
os.environ.setdefault("NCCL_IB_DISABLE", "1")
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")  # CUDA ids == nvidia-smi ids (temperature logging)
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

from common import CKPT_DIR, DATA_DIR, EXPECTED_PARAMS, LOG_DIR, OUT_DIR, ROOT, build_model
from evaluate import val_losses
from plot import plot_loss

# ---------------------------------------------------------------- spec (closed)
SEQ = 1024
MICRO_BATCH = 32
TOKENS_PER_STEP = 131_072
SEQS_PER_STEP = TOKENS_PER_STEP // SEQ
TOTAL_STEPS = 19_840
WARMUP_STEPS = 500
STABLE_END = 15_872
PEAK_LR, MIN_LR = 1.8e-3, 1.8e-4
BETAS, EPS = (0.9, 0.95), 1e-8
WEIGHT_DECAY = 0.1
GRAD_CLIP = 1.0
Z_LOSS = 1e-4
DATA_SEED = 20_453_760

# ---------------------------------------------------------------- run safety
LOG_EVERY = 10
KEEP_LAST = 2
INIT_LOSS_RANGE = (9.2, 10.2)  # ln(16384) = 9.70
SPIKE_WINDOW, SPIKE_DELTA, SPIKE_SKIP_STEPS, MAX_SPIKES = 100, 1.0, 200, 3
MIN_TOK_S, SLOW_STEPS = 45_000, 200
HOT_C, HOT_STEPS = 83, 100
COMPILE_TIMEOUT_S = 600

RUNS = {
    "train": dict(data_dir=DATA_DIR, max_blocks=None, steps=TOTAL_STEPS, ckpt_every=500, eval_every=1000,
                  plot_every=2000, milestones={500, 5_000, 10_000, STABLE_END, TOTAL_STEPS}, ckpt_dir=CKPT_DIR,
                  log=LOG_DIR / "train.jsonl", status=ROOT / "STATUS.json", plot=OUT_DIR / "loss.png"),
    "smoke": dict(data_dir=ROOT / "data" / "pilot", max_blocks=29_297, steps=200, ckpt_every=100, eval_every=100,
                  plot_every=200, milestones=set(), ckpt_dir=OUT_DIR / "smoke" / "ckpt",
                  log=LOG_DIR / "smoke.jsonl", status=OUT_DIR / "smoke" / "STATUS.json",
                  plot=OUT_DIR / "smoke" / "loss.png"),
}


def now() -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


def lr_at(step: int) -> float:
    """WSD: linear warmup for 500 steps, constant until 15,872, linear decay to 1.8e-4 at 19,840."""
    if step < WARMUP_STEPS:
        return PEAK_LR * (step + 1) / WARMUP_STEPS
    if step < STABLE_END:
        return PEAK_LR
    return PEAK_LR + (MIN_LR - PEAK_LR) * (step - STABLE_END) / (TOTAL_STEPS - STABLE_END)


class LossModel(torch.nn.Module):
    """Model + next-token cross-entropy + z-loss in one module, so torch.compile fuses the fp32 logit math."""

    def __init__(self, model: torch.nn.Module):
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        logits = self.model(input_ids=x, use_cache=False).logits.float()
        lse = torch.logsumexp(logits, dim=-1)
        ce = (lse - logits.gather(-1, y.unsqueeze(-1)).squeeze(-1)).mean()
        return ce + Z_LOSS * lse.square().mean(), ce


def build_optimizer(model: torch.nn.Module) -> torch.optim.Optimizer:
    """Weight decay only on 2D matrices; norms and the (tied) embedding get none."""
    embed = model.model.embed_tokens.weight
    decay = [p for p in model.parameters() if p.ndim == 2 and p is not embed]
    no_decay = [p for p in model.parameters() if p.ndim != 2 or p is embed]
    return torch.optim.AdamW([{"params": decay, "weight_decay": WEIGHT_DECAY},
                              {"params": no_decay, "weight_decay": 0.0}],
                             lr=PEAK_LR, betas=BETAS, eps=EPS, fused=True)


class BlockSampler:
    """Deterministic batches: a seeded permutation of all blocks, addressed by (data step, rank, micro-step).

    Blocks are read with one pread each instead of mmap: on this machine the shards live on the NTFS drive behind
    DrvFs (the WSL disk sits on an exFAT USB drive whose driver bugchecks when the disk grows), and DrvFs page
    faults are far slower than a single read call.
    """

    def __init__(self, data_dir: Path, max_blocks: int | None):
        index = json.loads((data_dir / "index.json").read_text())
        self.fds = [os.open(data_dir / s["file"], os.O_RDONLY) for s in index["shards"]]
        self.cum = np.cumsum([0] + [s["blocks"] for s in index["shards"]])
        self.total = int(self.cum[-1]) if max_blocks is None else min(max_blocks, int(self.cum[-1]))
        self.perm = np.random.default_rng(DATA_SEED).permutation(self.total)

    def get(self, data_step: int, rank: int, micro: int, accum: int) -> tuple[torch.Tensor, bool]:
        start = data_step * SEQS_PER_STEP + (rank * accum + micro) * MICRO_BATCH
        ids = self.perm[np.arange(start, start + MICRO_BATCH) % self.total]
        shard = np.searchsorted(self.cum, ids, side="right") - 1
        offsets = (ids - self.cum[shard]) * SEQ
        rows = np.empty((MICRO_BATCH, SEQ + 1), dtype=np.int64)
        for i, (s, o) in enumerate(zip(shard, offsets)):
            rows[i] = np.frombuffer(os.pread(self.fds[s], 2 * (SEQ + 1), 2 * int(o)), dtype=np.uint16)
        return torch.from_numpy(rows).pin_memory(), start + MICRO_BATCH > self.total


class GpuMonitor(threading.Thread):
    """Polls nvidia-smi in the background so logging never blocks on it."""

    def __init__(self, every_s: float = 5.0):
        super().__init__(daemon=True)
        self.every_s = every_s
        self.exe = shutil.which("nvidia-smi") or "/usr/lib/wsl/lib/nvidia-smi"
        self.temps, self.used_mb = [None, None], [None, None]

    def run(self) -> None:
        while True:
            try:
                out = subprocess.run([self.exe, "--query-gpu=temperature.gpu,memory.used", "--format=csv,noheader,nounits"],
                                     capture_output=True, text=True, timeout=30).stdout
                rows = [[int(v) for v in line.split(",")] for line in out.strip().splitlines()]
                self.temps, self.used_mb = [r[0] for r in rows], [r[1] for r in rows]
            except Exception:
                pass
            time.sleep(self.every_s)


class RunLog:
    def __init__(self, path: Path):
        path.parent.mkdir(parents=True, exist_ok=True)
        self.f = open(path, "a", encoding="utf-8")

    def write(self, rec: dict) -> None:
        self.f.write(json.dumps(rec) + "\n")
        self.f.flush()
        os.fsync(self.f.fileno())  # this machine has crashed under load; keep every logged line
        print(json.dumps(rec), flush=True)

    def event(self, name: str, **fields) -> None:
        self.write({"event": name, "wall_clock": now(), **fields})


def latest_checkpoint(ckpt_dir: Path) -> Path | None:
    paths = sorted(ckpt_dir.glob("step_*.pt"))
    return paths[-1] if paths else None


def fsync_dir(path: Path) -> None:
    try:
        fd = os.open(path, os.O_RDONLY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)
    except OSError:  # DrvFs (NTFS through 9p) may refuse fsync on a directory; the file itself was synced
        pass


def save_checkpoint(ckpt_dir: Path, model, opt, state: dict, milestones: set) -> Path:
    """Atomic and durable: write + fsync a temp file, rename, fsync the directory."""
    path = ckpt_dir / f"step_{state['step']:05d}.pt"
    tmp = path.with_suffix(".pt.tmp")
    with open(tmp, "wb") as f:
        torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "state": state,
                    "rng": {"torch": torch.get_rng_state(), "cuda": torch.cuda.get_rng_state(),
                            "numpy": np.random.get_state()}}, f)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)
    fsync_dir(ckpt_dir)
    steps = sorted(int(p.stem.removeprefix("step_")) for p in ckpt_dir.glob("step_*.pt"))
    keep = set(steps[-KEEP_LAST:]) | (set(steps) & milestones)
    for s in steps:
        if s not in keep:
            (ckpt_dir / f"step_{s:05d}.pt").unlink()
    return path


def write_status(path: Path, **fields) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".json.tmp")
    with open(tmp, "w") as f:
        f.write(json.dumps({"updated": now(), **fields}, indent=2))
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)
    fsync_dir(path.parent)


def measure_gpu_flops(device, n: int = 8192, iters: int = 10) -> float:
    """Peak bf16 matmul throughput of this GPU in FLOPs/s (2·n³ per n×n product): the MFU denominator.
    Measured rather than taken from a spec sheet, since consumer-card FP32 figures understate bf16 kernels."""
    a = torch.randn(n, n, device=device, dtype=torch.bfloat16)
    b = torch.randn(n, n, device=device, dtype=torch.bfloat16)
    for _ in range(3):
        a @ b
    torch.cuda.synchronize(device)
    t0 = time.perf_counter()
    for _ in range(iters):
        a @ b
    torch.cuda.synchronize(device)
    flops = 2 * n**3 * iters / (time.perf_counter() - t0)
    del a, b
    torch.cuda.empty_cache()
    return flops


def compile_with_fallback(ddp: DDP, sampler: BlockSampler, state: dict, rank: int, accum: int, device, opt):
    """Compile and warm up on one micro-batch (no grad sync); on error or > 10 min, every rank falls back to eager."""

    def on_timeout(signum, frame):
        raise TimeoutError(f"torch.compile exceeded {COMPILE_TIMEOUT_S}s")

    compiled = torch.compile(ddp)
    t0, ok, err = time.time(), 1, None
    signal.signal(signal.SIGALRM, on_timeout)
    signal.alarm(COMPILE_TIMEOUT_S)
    try:
        x, _ = sampler.get(state["step"] + state["data_skip"], rank, 0, accum)
        x = x.to(device)
        with ddp.no_sync():
            with torch.autocast("cuda", dtype=torch.bfloat16):
                loss, _ = compiled(x[:, :-1], x[:, 1:])
            loss.backward()
        torch.cuda.synchronize(device)
    except Exception as e:
        ok, err = 0, f"{type(e).__name__}: {e}"[:1000]
    finally:
        signal.alarm(0)
    opt.zero_grad(set_to_none=True)
    flag = torch.tensor([ok], device=device)
    dist.all_reduce(flag, op=dist.ReduceOp.MIN)
    if flag.item():
        return compiled, "compile", time.time() - t0, None
    torch._dynamo.reset()
    return ddp, "eager", time.time() - t0, err or "compile failed on another rank"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true", help="200-step smoke test on the pilot data")
    ap.add_argument("--stop-at", type=int, help="exit cleanly once this step is reached (resume testing)")
    args = ap.parse_args()
    run_name = "smoke" if args.smoke else "train"
    cfg = RUNS[run_name]

    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device("cuda", local_rank)
    torch.cuda.set_device(device)
    dist.init_process_group("nccl", timeout=timedelta(minutes=30), device_id=device)
    rank, world = dist.get_rank(), dist.get_world_size()
    master = rank == 0
    accum = SEQS_PER_STEP // (world * MICRO_BATCH)
    assert accum * world * MICRO_BATCH == SEQS_PER_STEP, "tokens/step must stay 131,072"
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    log = RunLog(cfg["log"]) if master else None

    def event(name: str, **fields) -> None:
        if master:
            log.event(name, **fields)

    torch.manual_seed(0)
    raw = build_model().to(device)
    n_params = sum(p.numel() for p in raw.parameters())
    assert n_params == EXPECTED_PARAMS, f"param count {n_params:,} != {EXPECTED_PARAMS:,}"
    opt = build_optimizer(raw)
    state = {"step": 0, "data_skip": 0, "lr_scale": 1.0, "spikes": 0, "loss_hist": [], "alerts": [],
             "last_val": None, "compile": None}
    cfg["ckpt_dir"].mkdir(parents=True, exist_ok=True)
    resume = latest_checkpoint(cfg["ckpt_dir"])
    if resume:
        ck = torch.load(resume, map_location=device, weights_only=False)
        raw.load_state_dict(ck["model"])
        opt.load_state_dict(ck["optimizer"])
        state.update(ck["state"])
        torch.set_rng_state(ck["rng"]["torch"].cpu())  # map_location moved it to the GPU; the CPU RNG needs it back
        del ck
        event("resume", checkpoint=resume.name, step=state["step"], data_skip=state["data_skip"],
              lr_scale=state["lr_scale"], loss_guard_events=state["spikes"])
    else:
        event("start", run=run_name, params=n_params, world=world, micro_batch=MICRO_BATCH, grad_accum=accum,
              tokens_per_step=TOKENS_PER_STEP, steps=cfg["steps"], torch=torch.__version__)

    sampler = BlockSampler(cfg["data_dir"], cfg["max_blocks"])
    peak = torch.tensor([measure_gpu_flops(device)], device=device, dtype=torch.float64)
    peaks = [torch.zeros_like(peak) for _ in range(world)]
    dist.all_gather(peaks, peak)
    peak_flops = sum(p.item() for p in peaks)
    event("gpu_flops", tflops_per_gpu=[round(p.item() / 1e12, 1) for p in peaks], kernel="bf16 matmul 8192x8192")

    ddp = DDP(LossModel(raw), device_ids=[local_rank], gradient_as_bucket_view=True, broadcast_buffers=False)
    runner, mode, compile_s, compile_err = compile_with_fallback(ddp, sampler, state, rank, accum, device, opt)
    state["compile"] = mode
    event("compile", mode=mode, seconds=round(compile_s, 1), **({"error": compile_err} if compile_err else {}))

    if resume is None and master:
        save_checkpoint(cfg["ckpt_dir"], raw, opt, state, cfg["milestones"])  # fallback for the loss guard
    dist.barrier()

    gpu = GpuMonitor() if master else None
    if gpu:
        gpu.start()
    flops_per_token = 6 * n_params + 12 * raw.config.num_hidden_layers * raw.config.num_attention_heads \
        * raw.config.head_dim * SEQ
    step = state["step"]
    step_times, window, all_dts = deque(maxlen=100), [], []
    slow = hot = 0
    wrapped_logged = False
    mark_t, mark_step = time.time(), step
    maxima = {"temp": [0, 0], "smi_mb": [0, 0], "peak_alloc": [0.0] * world, "peak_reserved": [0.0] * world}
    init_loss = last_loss = None
    lr = lr_at(step) * state["lr_scale"]

    def status(state_label: str) -> None:
        if not master:
            return
        hist = state["loss_hist"]
        sec_per_step = (time.time() - mark_t) / max(step - mark_step, 1)
        tok_s = TOKENS_PER_STEP * len(step_times) / sum(step_times) if step_times else None
        write_status(cfg["status"], run=run_name, state=state_label, step=step, total_steps=cfg["steps"],
                     percent=round(100 * step / cfg["steps"], 2),
                     loss_smoothed=round(sum(hist) / len(hist), 4) if hist else None,
                     tok_s=round(tok_s) if tok_s else None, lr=lr, lr_scale=state["lr_scale"],
                     eta_hours=round((cfg["steps"] - step) * sec_per_step / 3600, 2),
                     eta_at=datetime.fromtimestamp(time.time() + (cfg["steps"] - step) * sec_per_step)
                     .astimezone().isoformat(timespec="minutes"),
                     tokens_seen=step * TOKENS_PER_STEP, data_skip_steps=state["data_skip"],
                     loss_guard_events=state["spikes"], compile=mode, last_val=state["last_val"],
                     alerts=state["alerts"][-20:])

    def alert(kind: str, message: str, **fields) -> None:
        if master:
            rec = {"kind": kind, "step": step, "message": message, "wall_clock": now(), **fields}
            state["alerts"].append(rec)
            log.write({"event": "alert", **rec})
            status("training")

    def stop(reason: str, code: int) -> None:
        event("stopped", step=step, reason=reason)
        status(f"stopped: {reason}")
        dist.destroy_process_group()
        raise SystemExit(code)

    while step < cfg["steps"] and (args.stop_at is None or step < args.stop_at):
        lr = lr_at(step) * state["lr_scale"]
        for group in opt.param_groups:
            group["lr"] = lr
        t0 = time.perf_counter()
        acc = torch.zeros(2, device=device)
        for micro in range(accum):
            x, wrapped = sampler.get(step + state["data_skip"], rank, micro, accum)
            if wrapped and not wrapped_logged:
                event("data_wrap", step=step, blocks=sampler.total)
                wrapped_logged = True
            x = x.to(device, non_blocking=True)
            with ddp.no_sync() if micro < accum - 1 else nullcontext():
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    loss, ce = runner(x[:, :-1], x[:, 1:])
                (loss / accum).backward()
            acc += torch.stack([ce.detach(), loss.detach()]) / accum
        dist.all_reduce(acc, op=dist.ReduceOp.AVG)
        ce_val, loss_val = acc.tolist()

        hist = state["loss_hist"]
        moving_avg = sum(hist) / len(hist) if hist else None
        if not math.isfinite(ce_val) or (len(hist) >= SPIKE_WINDOW and ce_val > moving_avg + SPIKE_DELTA):
            state["spikes"] += 1
            event("loss_guard", step=step, loss=ce_val, moving_avg=moving_avg, count=state["spikes"])
            if state["spikes"] >= MAX_SPIKES:
                stop(f"loss guard triggered {MAX_SPIKES} times", 3)
            path = latest_checkpoint(cfg["ckpt_dir"])
            ck = torch.load(path, map_location=device, weights_only=False)
            raw.load_state_dict(ck["model"])
            opt.load_state_dict(ck["optimizer"])
            keep = {k: state[k] for k in ("spikes", "alerts", "compile")}
            skip, scale = state["data_skip"], state["lr_scale"]
            state.update(ck["state"])
            state.update(keep, data_skip=skip + SPIKE_SKIP_STEPS, lr_scale=scale * 0.5)
            del ck
            opt.zero_grad(set_to_none=True)
            step = state["step"]
            event("reload", checkpoint=path.name, step=step, data_skip=state["data_skip"], lr_scale=state["lr_scale"])
            continue
        if step == 0:
            init_loss = ce_val
            event("init_loss", loss=round(ce_val, 4), expected=round(math.log(raw.config.vocab_size), 4))
            if not INIT_LOSS_RANGE[0] <= ce_val <= INIT_LOSS_RANGE[1]:
                stop(f"step-0 loss {ce_val:.3f} outside {INIT_LOSS_RANGE}: init or target shift is wrong", 2)

        grad_norm = torch.nn.utils.clip_grad_norm_(raw.parameters(), GRAD_CLIP)
        opt.step()
        opt.zero_grad(set_to_none=True)
        step += 1
        state["step"] = step
        hist.append(ce_val)
        del hist[:-SPIKE_WINDOW]
        last_loss = ce_val
        dt = time.perf_counter() - t0
        step_times.append(dt)
        window.append(dt)
        all_dts.append(dt)

        if master:
            slow = slow + 1 if TOKENS_PER_STEP / dt < MIN_TOK_S else 0
            if slow == SLOW_STEPS:
                alert("throughput", f"tok/s below {MIN_TOK_S:,} for {SLOW_STEPS} consecutive steps (I/O or thermal?)",
                      tok_s=round(TOKENS_PER_STEP * len(step_times) / sum(step_times)), temps=gpu.temps)

        if step % LOG_EVERY == 0:
            mem = torch.tensor([torch.cuda.memory_allocated(device), torch.cuda.max_memory_allocated(device),
                                torch.cuda.max_memory_reserved(device)], device=device, dtype=torch.float64) / 2**30
            mems = [torch.zeros_like(mem) for _ in range(world)]
            dist.all_gather(mems, mem)
            if master:
                tok_s = TOKENS_PER_STEP * len(window) / sum(window)
                temps = gpu.temps
                log.write({"step": step, "loss": round(ce_val, 4), "loss_total": round(loss_val, 4), "lr": lr,
                           "grad_norm": round(grad_norm.item(), 4), "tokens_seen": step * TOKENS_PER_STEP,
                           "tok_s": round(tok_s), "mfu": round(flops_per_token * tok_s / peak_flops, 4),
                           "gpu0_temp": temps[0], "gpu1_temp": temps[1] if len(temps) > 1 else None,
                           "mem_alloc": [round(m[0].item(), 2) for m in mems],
                           "mem_peak": [round(m[1].item(), 2) for m in mems], "wall_clock": now()})
                for i, t in enumerate(temps):
                    if t is not None:
                        maxima["temp"][i] = max(maxima["temp"][i], t)
                for i, used in enumerate(gpu.used_mb):
                    if used is not None:
                        maxima["smi_mb"][i] = max(maxima["smi_mb"][i], used)
                maxima["peak_alloc"] = [round(m[1].item(), 2) for m in mems]
                maxima["peak_reserved"] = [round(m[2].item(), 2) for m in mems]
                hot = hot + LOG_EVERY if any(t is not None and t > HOT_C for t in temps) else 0
                if hot == HOT_STEPS:
                    alert("thermal", f"GPU above {HOT_C}°C for {HOT_STEPS} steps (training continues)", temps=temps)
            window = []

        if step % cfg["eval_every"] == 0 or step == cfg["steps"]:
            results = val_losses(raw, device, rank, world)
            if master:
                state["last_val"] = {"step": step, **{d: {k: round(v, 4) for k, v in r.items()}
                                                      for d, r in results.items()}}
                log.write({"event": "eval", "wall_clock": now(), **state["last_val"]})

        if step % cfg["ckpt_every"] == 0 or step in cfg["milestones"] or step == cfg["steps"]:
            if master:
                path = save_checkpoint(cfg["ckpt_dir"], raw, opt, state, cfg["milestones"])
                event("checkpoint", step=step, file=path.name)
                status("training")
            dist.barrier()

        if master and (step % cfg["plot_every"] == 0 or step == cfg["steps"]):
            try:
                plot_loss(cfg["log"], cfg["plot"], cfg["steps"])
            except Exception as e:  # a plotting bug must never kill the run
                event("plot_error", error=f"{type(e).__name__}: {e}")

    if step >= cfg["steps"]:
        event("finished", step=step, hours=round((time.time() - mark_t) / 3600, 2))
        status("finished")
    else:
        event("stopped_at", step=step)
        status(f"stopped at step {step} (--stop-at)")

    if master and args.smoke:
        steady = all_dts[20:] or all_dts
        mean_dt = sum(steady) / len(steady)
        summary = {"params": n_params, "compile": mode, "compile_seconds": round(compile_s, 1),
                   "compile_error": compile_err, "init_loss": init_loss, "last_step": step, "loss_at_last_step": last_loss,
                   "tok_s_mean": round(TOKENS_PER_STEP / mean_dt), "sec_per_step": round(mean_dt, 3),
                   "eta_full_run_hours": round(TOTAL_STEPS * mean_dt / 3600, 2),
                   "peak_mem_alloc_gb": maxima["peak_alloc"], "peak_mem_reserved_gb": maxima["peak_reserved"],
                   "nvidia_smi_used_mb_max": maxima["smi_mb"], "temp_max_c": maxima["temp"],
                   "last_val": state["last_val"]}
        (OUT_DIR / "smoke" / "summary.json").write_text(json.dumps(summary, indent=2))
        event("smoke_summary", **summary)
    dist.destroy_process_group()


if __name__ == "__main__":
    main()