"""Run on the allocated worker, never substitute a CPU smoke for this check."""
import argparse
import json
import platform
import socket
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    import torch
    if not torch.cuda.is_available():
        raise RuntimeError("A GPU allocation is required; CUDA is unavailable")
    if not torch.cuda.is_bf16_supported():
        raise RuntimeError("This pilot requires a GPU with BF16 support")
    x = torch.randn((256, 256), device="cuda", dtype=torch.bfloat16)
    y = x @ x.T
    torch.cuda.synchronize()
    assert torch.isfinite(y).all().item()
    result = {"hostname": socket.gethostname(), "python": platform.python_version(),
              "torch": torch.__version__, "cuda": torch.version.cuda,
              "gpu": torch.cuda.get_device_name(), "devices": torch.cuda.device_count(),
              "memory_bytes": torch.cuda.get_device_properties(0).total_memory,
              "bf16_matmul": "passed"}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()