#!/usr/bin/env python3
"""Per-GPU basic integrity: H2D/D2H copy fidelity, matmul correctness,
repeated-copy stress. Does NOT test P2P — use p2p_check.py for that.

Run:  docker exec vllm python3 /path/to/gpu_integrity.py
"""
import torch

print("torch", torch.__version__, "cuda", torch.version.cuda, "devs", torch.cuda.device_count())
for dev in range(torch.cuda.device_count()):
    try:
        torch.cuda.set_device(dev)
        x = torch.randint(0, 2**31 - 1, (128 * 1024 * 1024,), dtype=torch.int32)
        y = x.to(f"cuda:{dev}")
        z = y.cpu()
        copy_ok = torch.equal(x, z)
        a = torch.randn(2048, 2048, dtype=torch.float32)
        b = torch.randn(2048, 2048, dtype=torch.float32)
        ref = (a.double() @ b.double()).float()
        c = (a.to(f"cuda:{dev}") @ b.to(f"cuda:{dev}")).cpu()
        err = (c - ref).abs().max().item()
        stress_ok = True
        for _ in range(20):
            s = torch.randint(0, 2**31 - 1, (16 * 1024 * 1024,), dtype=torch.int32)
            if not torch.equal(s, s.to(f"cuda:{dev}").cpu()):
                stress_ok = False
                break
        print(f"GPU{dev}: copy_ok={copy_ok} matmul_max_err={err:.3e} "
              f"stress_ok={stress_ok} name={torch.cuda.get_device_name(dev)}")
        del x, y, z, a, b, c
        torch.cuda.empty_cache()
    except Exception as e:
        print(f"GPU{dev}: ERROR {type(e).__name__}: {e}")