"""Benchmark #2: end-to-end step time on the REAL 7B, triton_fa2 vs SDPA.

Run (GPU 3 window; pin BEFORE torch import):
    CUDA_VISIBLE_DEVICES=3 TMPDIR=/apps/development/rlhf/nano-mistral-o1/.tmp \
        uv run python scripts/bench_7b_end_to_end.py

WHAT THIS MEASURES, AND WHAT IT DELIBERATELY DOES NOT
-----------------------------------------------------
A full-parameter 7B training step does NOT fit on one A10G: bf16 weights are
13.5 GiB and bf16 parameter gradients are another 13.5 GiB, so 27 GiB against a
23 GiB card. Real SFT gets around this with 4-way FSDP plus an offloaded
optimizer (section-07/11); that is not reproducible inside a single-GPU window,
and pretending otherwise would be measuring a different thing than the label.

So this benchmark freezes every parameter and requires grad on `inputs_embeds`
instead. Backward still flows through every attention and MLP in all 32 layers,
computing the full activation-gradient chain -- dQ/dK/dV included, which is
exactly the part the attention backend controls. What is skipped is the
per-parameter weight-gradient GEMM, which is IDENTICAL for both backends.

The consequence is worth stating precisely, because it decides how to read the
result: the absolute per-step delta (triton minus sdpa, in ms) is exactly what
it would be in a full step, since the omitted work is backend-independent and
additive. The delta as a PERCENTAGE is inflated, because the denominator is
missing that same backend-independent work. **This benchmark therefore reports
an UPPER BOUND on the true percentage cost of choosing triton_fa2.** A number
that bounds the answer from the safe side is worth more here than a number that
requires four GPUs and a training run to obtain.

Both no-checkpointing and activation-checkpointing configurations are measured:
under checkpointing the attention forward runs a second time during backward,
so the backend-attributable share is larger, and SFT is expected to checkpoint.
"""

from __future__ import annotations

import glob
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from align.utils.guard import apply_resource_guards

# Light job by the box rules: no heavy lock (other GPU jobs are live), no
# dataloaders, no multi-GB writes.
apply_resource_guards(cpu_threads=8, heavy=False, label="bench-7b-attn")

import torch
from torch.utils.checkpoint import checkpoint as ckpt_fn

from align.model import checkpoint as ckpt
from align.model.mistral import NmoMistralForCausalLM

REPO = "/apps/development/rlhf/nano-mistral-o1"
OUT = f"{REPO}/docs/bench_7b_end_to_end.json"
SNAPSHOT = (
    "/apps/voxtral/.cache/huggingface/hub/"
    "models--mistralai--Mistral-7B-v0.3/snapshots/"
    "caa1feb0e54d415e2df31207e5f4e273e33509b1"
)

SEQ = 1024
BATCH = 1
#: Hard wall-clock budget. The GPU window is shared and SFT launches the moment
#: the golden gate passes; a benchmark that overstays blocks a training run.
DEADLINE_S = float(os.environ.get("NMO_BENCH_DEADLINE_S", 1500))
_START = time.time()


def _left() -> float:
    return DEADLINE_S - (time.time() - _START)


def _load_7b(attn_impl: str, state_dict, cfg):
    model = NmoMistralForCausalLM.from_state_dict(cfg, state_dict, attn_implementation=attn_impl)
    model = model.to(torch.bfloat16).to("cuda").eval()
    # Freeze: no 13.5 GiB parameter-gradient buffer. See module docstring.
    for p in model.parameters():
        p.requires_grad_(False)
    return model


def _blocks(model):
    inner = model.model if hasattr(model, "model") else model
    for name in ("layers", "blocks", "h"):
        if hasattr(inner, name):
            return getattr(inner, name)
    raise AttributeError("cannot find the decoder block list on the model")


def _timed(fn, iters, warmup):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / iters * 1e3


def _run_config(model, embeds, labels, use_ckpt: bool, iters: int):
    """fwd-only and fwd+bwd timings for one model/checkpointing combination."""
    blocks = _blocks(model)
    originals = [b.forward for b in blocks]
    if use_ckpt:
        for b in blocks:
            fwd = b.forward
            b.forward = lambda *a, _f=fwd, **kw: ckpt_fn(_f, *a, use_reentrant=False, **kw)

    try:

        def fwd_only():
            with torch.no_grad():
                model(inputs_embeds=embeds.detach())

        def fwd_bwd():
            if embeds.grad is not None:
                embeds.grad = None
            logits = model(inputs_embeds=embeds)
            loss = torch.nn.functional.cross_entropy(
                logits.float().view(-1, logits.shape[-1]), labels.view(-1)
            )
            loss.backward()

        torch.cuda.reset_peak_memory_stats()
        f = _timed(fwd_only, iters, 2)
        fb = _timed(fwd_bwd, iters, 2)
        peak = torch.cuda.max_memory_allocated() / 2**30
    finally:
        # `b.forward = original_bound_method` would leave the instance holding a
        # bound method that holds the instance -- a reference cycle the caching
        # allocator cannot reclaim on `empty_cache()`, which OOM'd the second
        # model load. Deleting the instance attribute restores the class method
        # and leaves no cycle.
        for b in blocks:
            if "forward" in b.__dict__:
                del b.__dict__["forward"]
        del originals
    return {"fwd_ms": f, "fwd_bwd_ms": fb, "peak_gib": peak}


def main() -> int:
    if torch.cuda.device_count() != 1:
        print(
            f"{torch.cuda.device_count()} devices visible; pin exactly one "
            "(shared box, three other GPU jobs live)",
            file=sys.stderr,
        )
        return 1
    free_gib = torch.cuda.mem_get_info()[0] / 2**30
    print(f"device: {torch.cuda.get_device_name(0)}  free {free_gib:.1f} GiB")
    if free_gib < 20.0:
        print(
            "refusing: need ~20 GiB free for 13.5 GiB of bf16 weights plus "
            "activations; another job is resident on this device.",
            file=sys.stderr,
        )
        return 1

    from align.eval.golden import resolve_model_config

    cfg = resolve_model_config(SNAPSHOT)
    print(
        f"config: {cfg.n_layers}L hidden={cfg.hidden_size} "
        f"{cfg.n_heads}q:{cfg.n_kv_heads}kv head_dim={cfg.head_dim} "
        f"vocab={cfg.vocab_size} sliding_window={cfg.sliding_window}"
    )

    import safetensors.torch as st

    t0 = time.time()
    hf_sd = {}
    for shard in sorted(glob.glob(os.path.join(SNAPSHOT, "*.safetensors"))):
        hf_sd.update(st.load_file(shard))
    nmo_sd = ckpt.hf_to_nmo(hf_sd, expected_n_layers=cfg.n_layers)
    del hf_sd
    print(f"loaded + remapped {len(nmo_sd)} tensors in {time.time() - t0:.1f}s")

    torch.manual_seed(0)
    labels = torch.randint(0, cfg.vocab_size, (BATCH, SEQ), device="cuda")
    results = {}

    # One backend per PROCESS when an arg is given. 13.5 GiB of weights leaves
    # no room for a second model, and process exit is the only VRAM release
    # that cannot be defeated by a stray reference. Partials are merged below.
    only = sys.argv[1] if len(sys.argv) > 1 else None
    impls = (only,) if only else ("sdpa", "triton_fa2")

    for impl in impls:
        if _left() < 240:
            print(f"\ndeadline guard: {_left():.0f}s left, skipping {impl}", file=sys.stderr)
            break
        print(f"\n=== {impl} ===", flush=True)
        model = _load_7b(impl, nmo_sd, cfg)
        embeds = torch.randn(
            BATCH, SEQ, cfg.hidden_size, device="cuda", dtype=torch.bfloat16, requires_grad=True
        )
        results[impl] = {}
        for use_ckpt in (False, True):
            if _left() < 120:
                print(
                    f"deadline guard: {_left():.0f}s left, skipping {impl}/ckpt={use_ckpt}",
                    file=sys.stderr,
                )
                break
            r = _run_config(model, embeds, labels, use_ckpt, iters=5)
            key = "activation_ckpt" if use_ckpt else "no_ckpt"
            results[impl][key] = r
            print(
                f"  {key:16s} fwd {r['fwd_ms']:8.1f} ms   "
                f"fwd+bwd {r['fwd_bwd_ms']:8.1f} ms   peak {r['peak_gib']:5.2f} GiB",
                flush=True,
            )
        del model, embeds
        import gc

        gc.collect()
        torch.cuda.empty_cache()

    # Merge any partial written by a sibling process, so running one backend
    # per process still produces a single complete receipt.
    partial = f"{REPO}/.tmp/bench7b_partial.json"
    if only:
        prev = {}
        if os.path.isfile(partial):
            with open(partial) as f:
                prev = json.load(f)
        prev.update(results)
        with open(partial, "w") as f:
            json.dump(prev, f, indent=2)
        results = prev
        print(f"\npartial results for {sorted(results)} -> {partial}")

    payload = {
        "note": "parameters FROZEN; grad flows via inputs_embeds. Absolute ms "
        "deltas are exact; percentages are an UPPER BOUND on the true "
        "cost, because the omitted weight-grad GEMMs are "
        "backend-independent and would enlarge the denominator only. "
        "See the module docstring.",
        "device": torch.cuda.get_device_name(0),
        "torch": torch.__version__,
        "model": "mistralai/Mistral-7B-v0.3 (real weights)",
        "batch": BATCH,
        "seq": SEQ,
        "dtype": "bfloat16",
        "results": results,
    }

    if "sdpa" in results and "triton_fa2" in results:
        payload["comparison"] = {}
        print("\n=== triton_fa2 vs sdpa on the real 7B ===")
        for key in ("no_ckpt", "activation_ckpt"):
            if key in results["sdpa"] and key in results["triton_fa2"]:
                s = results["sdpa"][key]["fwd_bwd_ms"]
                t = results["triton_fa2"][key]["fwd_bwd_ms"]
                c = {
                    "sdpa_ms": s,
                    "triton_ms": t,
                    "delta_ms": t - s,
                    "ratio": t / s,
                    "upper_bound_pct": 100.0 * (t - s) / s,
                }
                payload["comparison"][key] = c
                print(
                    f"  {key:16s} sdpa {s:8.1f} ms -> triton {t:8.1f} ms "
                    f"({c['delta_ms']:+.1f} ms, {c['upper_bound_pct']:+.1f}% "
                    "UPPER BOUND)"
                )

    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "w") as f:
        json.dump(payload, f, indent=2)
        f.write("\n")
    print(f"\nwrote {OUT}  ({_left():.0f}s of window left)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())