File size: 882 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | """Validate the requested H200 allocation before smoke or full evaluation."""
import argparse
import json
import os
from pathlib import Path
import torch
parser = argparse.ArgumentParser()
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--expected-gpus", type=int, default=8)
args = parser.parse_args()
devices = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
if len(devices) != args.expected_gpus or any("H200" not in name for name in devices):
raise RuntimeError(f"Evaluation requires {args.expected_gpus} H200 GPUs; allocated devices: {devices}")
args.output.mkdir(parents=True, exist_ok=True)
record = {"devices": devices, "torch": torch.__version__, "slurm_job_id": os.environ.get("SLURM_JOB_ID")}
(args.output / "node.json").write_text(json.dumps(record, indent=2) + "\n")
print(json.dumps(record), flush=True)
|