#!/usr/bin/env python3
"""What the host bridge costs, as a number — the PHB receipt README §A2.8 requires.

    torchrun --nproc_per_node 2 phb_allreduce.py --out receipts/phb-allreduce.json

benchbox's two RTX 3090s have **no NVLink** (`nvidia-smi nvlink -s`: all links
inActive) and sit at **PHB** (`topo -m`): every tensor-parallel all-reduce crosses
the host bridge, and card 1 is on a **four-lane** link. A bench that reports a
decision latency over that link without measuring the link is reporting a number
it cannot explain — so this runs first, on the same two boards, with nothing else
on them.

WHY THESE MESSAGE SIZES. A tensor-parallel rank all-reduces a `[tokens, hidden]`
activation twice per layer. OpenJev is hidden 5,120 in bfloat16, so one token is
**10 KiB** on the wire and a 16,384-token prefill is **160 MiB**. The sweep
therefore runs from one token to a full prefill, and the two ends answer two
different questions:

  * the SMALL end is **decode**: one token, 128 all-reduces per output token
    (2 per layer x 64 layers), and it is pure latency — bandwidth cannot help it;
  * the LARGE end is **prefill**: bandwidth-bound, and it is where an x4 link
    shows its width.

Reported per size: median and p95 per-call latency, the algorithm bandwidth
(bytes moved / time) and the **bus** bandwidth (`algbw * 2(n-1)/n`, the ring
all-reduce's own figure, which is what a link is compared on).
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import statistics
import sys

import torch
import torch.distributed as dist

HIDDEN = 5120          # openjev/openjev-FP8 text_config.hidden_size
ELEM = 2               # bfloat16 activations on the wire


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="phb-allreduce.json")
    ap.add_argument("--iters", type=int, default=50)
    ap.add_argument("--warmup", type=int, default=10)
    a = ap.parse_args()

    dist.init_process_group("nccl")
    rank, world = dist.get_rank(), dist.get_world_size()
    torch.cuda.set_device(rank)
    dev = torch.device("cuda", rank)

    #: tokens per all-reduce, from one (decode) to a full 16,384-token prefill
    token_counts = [1, 4, 16, 64, 256, 1024, 4096, 8192, 16384]
    results = []
    for toks in token_counts:
        x = torch.ones(toks, HIDDEN, dtype=torch.bfloat16, device=dev)
        nbytes = x.numel() * ELEM
        for _ in range(a.warmup):
            dist.all_reduce(x)
        torch.cuda.synchronize()
        times = []
        for _ in range(a.iters):
            s, e = torch.cuda.Event(True), torch.cuda.Event(True)
            dist.barrier()
            s.record()
            dist.all_reduce(x)
            e.record()
            torch.cuda.synchronize()
            times.append(s.elapsed_time(e) / 1e3)      # seconds
        med = statistics.median(times)
        p95 = sorted(times)[min(len(times) - 1, int(round(0.95 * (len(times) - 1))))]
        algbw = nbytes / med
        busbw = algbw * 2 * (world - 1) / world
        results.append({
            "tokens": toks, "bytes": nbytes, "iters": a.iters,
            "median_us": med * 1e6, "p95_us": p95 * 1e6, "min_us": min(times) * 1e6,
            "algbw_GBs": algbw / 1e9, "busbw_GBs": busbw / 1e9,
            "what": ("one decode step" if toks == 1
                     else "a %d-token prefill chunk" % toks),
        })
        if rank == 0:
            r = results[-1]
            print("%8d tok  %10.2f KiB  median %9.1f us  p95 %9.1f us  "
                  "algbw %6.2f GB/s  busbw %6.2f GB/s"
                  % (toks, nbytes / 1024, r["median_us"], r["p95_us"],
                     r["algbw_GBs"], r["busbw_GBs"]), flush=True)
        del x
        torch.cuda.empty_cache()

    if rank == 0:
        one = next(r for r in results if r["tokens"] == 1)
        out = {
            "stamp": dt.datetime.now(dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            "box": "benchbox", "world_size": world,
            "torch": torch.__version__, "nccl": ".".join(map(str, torch.cuda.nccl.version())),
            "devices": [torch.cuda.get_device_name(i) for i in range(world)],
            "hidden_size": HIDDEN, "dtype": "bfloat16",
            "topology": "PHB, no NVLink; card 1 on a four-lane link",
            "env": {k: v for k, v in os.environ.items() if k.startswith("NCCL_")},
            "results": results,
            "decode_cost": {
                "one_allreduce_us": one["median_us"],
                "allreduces_per_output_token": 128,
                "note": ("2 all-reduces per layer x 64 layers; this is the floor "
                         "the link puts under a written token, before the model "
                         "computes anything"),
                "per_output_token_ms": one["median_us"] * 128 / 1e3,
            },
        }
        with open(a.out, "w") as fh:
            json.dump(out, fh, indent=1)
        print("\nwrote %s" % a.out)
        print("PHB floor under one written token: %.1f ms (128 x %.1f us)"
              % (out["decode_cost"]["per_output_token_ms"], one["median_us"]))
    dist.barrier()
    dist.destroy_process_group()
    return 0


if __name__ == "__main__":
    sys.exit(main())