#!/usr/bin/env python3
"""GPU P2P transfer integrity check.

Tests every ordered GPU pair: host -> src GPU -> dst GPU -> host, and
compares against the host copy. ANY mismatch means the P2P DMA path is
corrupting data and tensor-parallel inference MUST run with P2P disabled
(NCCL_P2P_DISABLE=1 and, for vLLM, --disable-custom-all-reduce).

Run somewhere with torch + all GPUs visible, e.g.:
    docker exec vllm python3 /models/../diag/p2p_check.py
or  docker run --rm --gpus all -v $PWD/diag:/diag vllm/vllm-openai:v0.27.1 \
        python3 /diag/p2p_check.py
"""
import sys
import torch


def main():
    n = torch.cuda.device_count()
    print(f"torch={torch.__version__} gpus={n}")
    if n < 2:
        print("need >= 2 GPUs"); return 2
    for i in range(n):
        for j in range(n):
            if i != j:
                print(f"p2p {i}->{j} supported: {torch.cuda.can_device_access_peer(i, j)}")
    x = torch.arange(16 * 1024 * 1024, dtype=torch.float32, device="cpu")
    ref = x.clone()
    bad_pairs = []
    for src in range(n):
        for dst in range(n):
            if src == dst:
                continue
            try:
                a = x.to(f"cuda:{src}")
                b = a.to(f"cuda:{dst}")
                c = b.cpu()
                torch.cuda.synchronize()
                n_bad = int((c != ref).sum().item())
                status = "OK" if n_bad == 0 else f"CORRUPT ({n_bad}/{ref.numel()} elements wrong)"
                print(f"{src}->{dst}: {status}")
                if n_bad:
                    bad_pairs.append((src, dst))
                del a, b, c
                torch.cuda.empty_cache()
            except Exception as e:
                print(f"{src}->{dst}: FAILED {type(e).__name__} {str(e)[:160]}")
                bad_pairs.append((src, dst))
    if bad_pairs:
        print(f"\nRESULT: {len(bad_pairs)} corrupt/failing P2P paths: {bad_pairs}")
        print("Keep NCCL_P2P_DISABLE=1 (+ --disable-custom-all-reduce for vLLM).")
        return 1
    print("\nRESULT: all P2P paths clean.")
    return 0


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