#!/usr/bin/env python3

import json
import os
import statistics

import torch
import torch.distributed as dist


def main() -> None:
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")

    # 16 tokens x 4096 hidden values in BF16: 128 KiB per rank.
    tensor = torch.ones((16, 4096), dtype=torch.bfloat16, device="cuda")
    for _ in range(10):
        dist.all_reduce(tensor)
    torch.cuda.synchronize()

    samples_ms = []
    for _ in range(50):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        dist.all_reduce(tensor)
        end.record()
        end.synchronize()
        samples_ms.append(float(start.elapsed_time(end)))

    if dist.get_rank() == 0:
        ordered = sorted(samples_ms)
        result = {
            "backend": "nccl",
            "collective": "all_reduce",
            "dtype": "bfloat16",
            "elements_per_rank": tensor.numel(),
            "bytes_per_rank": tensor.numel() * tensor.element_size(),
            "world_size": dist.get_world_size(),
            "warmup_iterations": 10,
            "measured_iterations": len(samples_ms),
            "mean_ms": statistics.fmean(samples_ms),
            "p50_ms": statistics.median(samples_ms),
            "p95_ms": ordered[int(0.95 * (len(ordered) - 1))],
            "min_ms": min(samples_ms),
            "max_ms": max(samples_ms),
        }
        print(json.dumps(result, sort_keys=True))

    dist.destroy_process_group()


if __name__ == "__main__":
    main()