115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
"""Hardware / platform detection. Stdlib only, safe to run anywhere."""
|
|
|
|
import os
|
|
import platform
|
|
import re
|
|
import subprocess
|
|
|
|
from .util import which, IS_WINDOWS, IS_MAC, IS_LINUX
|
|
|
|
|
|
def _run_quiet(cmd):
|
|
try:
|
|
out = subprocess.run(
|
|
cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=15
|
|
)
|
|
if out.returncode != 0:
|
|
return None
|
|
return out.stdout.decode("utf-8", errors="replace")
|
|
except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
|
|
return None
|
|
|
|
|
|
def detect_nvidia():
|
|
"""Returns dict with gpus + driver info, or None if no working nvidia-smi."""
|
|
smi = which("nvidia-smi")
|
|
if not smi:
|
|
return None
|
|
fields = "name,memory.total,driver_version,compute_cap"
|
|
csv = _run_quiet([smi, "--query-gpu=" + fields, "--format=csv,noheader"])
|
|
if not csv:
|
|
# older drivers don't know the compute_cap field
|
|
csv = _run_quiet(
|
|
[
|
|
smi,
|
|
"--query-gpu=name,memory.total,driver_version",
|
|
"--format=csv,noheader",
|
|
]
|
|
)
|
|
if not csv:
|
|
return None
|
|
gpus = []
|
|
driver = None
|
|
for line in csv.strip().splitlines():
|
|
parts = [p.strip() for p in line.split(",")]
|
|
if len(parts) >= 3:
|
|
gpu = {"name": parts[0], "memory": parts[1]}
|
|
if len(parts) >= 4:
|
|
gpu["compute_cap"] = parts[3]
|
|
gpus.append(gpu)
|
|
driver = parts[2]
|
|
if not gpus:
|
|
return None
|
|
# Max CUDA version the driver supports only appears in the banner output.
|
|
# Windows/WDDM labels it "CUDA UMD Version:" on recent drivers, Linux just
|
|
# "CUDA Version:" — without the optional word we silently fall through to
|
|
# the "driver present but version unknown, assume current" path and can
|
|
# hand an old driver cu130 wheels it cannot run.
|
|
banner = _run_quiet([smi]) or ""
|
|
m = re.search(r"CUDA(?:\s+[A-Z]+)?\s+Version:\s*([0-9]+\.[0-9]+)", banner)
|
|
cuda_version = m.group(1) if m else None
|
|
return {"gpus": gpus, "driver": driver, "cuda_version": cuda_version}
|
|
|
|
|
|
def detect_rocm():
|
|
"""Returns dict if an AMD ROCm stack is present, else None."""
|
|
if not IS_LINUX:
|
|
return None
|
|
smi = which("rocm-smi")
|
|
if not smi and not os.path.isdir("/opt/rocm"):
|
|
return None
|
|
gpus = []
|
|
out = _run_quiet([smi, "--showproductname"]) if smi else None
|
|
if out:
|
|
for m in re.finditer(r"Card [Ss]eries:\s*(.+)", out):
|
|
gpus.append({"name": m.group(1).strip()})
|
|
return {"gpus": gpus}
|
|
|
|
|
|
def detect():
|
|
"""Full platform detection. Returns a plain dict (json-serializable)."""
|
|
system = platform.system() # Linux / Darwin / Windows
|
|
arch = platform.machine().lower() # x86_64 / amd64 / arm64 / aarch64
|
|
if arch == "amd64":
|
|
arch = "x86_64"
|
|
if arch == "arm64" and not IS_MAC:
|
|
arch = "aarch64"
|
|
|
|
result = {
|
|
"os": {"Linux": "linux", "Darwin": "mac", "Windows": "windows"}.get(
|
|
system, system.lower()
|
|
),
|
|
"arch": arch,
|
|
"python": platform.python_version(),
|
|
"nvidia": None,
|
|
"rocm": None,
|
|
"backend": "cpu",
|
|
}
|
|
|
|
nvidia = detect_nvidia()
|
|
if nvidia:
|
|
result["nvidia"] = nvidia
|
|
result["backend"] = "cuda"
|
|
else:
|
|
rocm = detect_rocm()
|
|
if rocm:
|
|
result["rocm"] = rocm
|
|
result["backend"] = "rocm"
|
|
|
|
if IS_MAC:
|
|
result["backend"] = "mps" if arch == "arm64" else "cpu"
|
|
|
|
# DGX OS / Grace (GB10, DGX Spark): NVIDIA GPU on aarch64 Linux
|
|
result["is_dgx"] = bool(result["os"] == "linux" and arch == "aarch64" and nvidia)
|
|
return result
|