277 lines
9.4 KiB
Python
277 lines
9.4 KiB
Python
"""GPU hardware detection with graceful CPU fallback."""
|
||
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import time
|
||
|
||
|
||
def _detect_ram_mb() -> int:
|
||
"""Return host physical RAM in MB, or 0 if unavailable."""
|
||
try:
|
||
pages = os.sysconf("SC_PHYS_PAGES")
|
||
page_size = os.sysconf("SC_PAGE_SIZE")
|
||
return int((pages * page_size) // (1024 * 1024))
|
||
except (AttributeError, OSError, ValueError):
|
||
pass
|
||
return _detect_windows_ram_mb()
|
||
|
||
|
||
def _detect_windows_ram_mb() -> int:
|
||
"""Return Windows physical RAM in MB, or 0."""
|
||
try:
|
||
import ctypes
|
||
|
||
class _MemoryStatusEx(ctypes.Structure):
|
||
_fields_ = [
|
||
("dwLength", ctypes.c_ulong),
|
||
("dwMemoryLoad", ctypes.c_ulong),
|
||
("ullTotalPhys", ctypes.c_ulonglong),
|
||
("ullAvailPhys", ctypes.c_ulonglong),
|
||
("ullTotalPageFile", ctypes.c_ulonglong),
|
||
("ullAvailPageFile", ctypes.c_ulonglong),
|
||
("ullTotalVirtual", ctypes.c_ulonglong),
|
||
("ullAvailVirtual", ctypes.c_ulonglong),
|
||
("ullAvailExtendedVirtual", ctypes.c_ulonglong),
|
||
]
|
||
|
||
status = _MemoryStatusEx()
|
||
status.dwLength = ctypes.sizeof(_MemoryStatusEx)
|
||
if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
|
||
return int(status.ullTotalPhys // (1024 * 1024))
|
||
except (AttributeError, OSError, ValueError):
|
||
pass
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"powershell",
|
||
"-NoProfile",
|
||
"-Command",
|
||
"(Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory",
|
||
],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=5,
|
||
)
|
||
if result.returncode == 0 and result.stdout.strip():
|
||
return int(result.stdout.strip()) // (1024 * 1024)
|
||
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError):
|
||
pass
|
||
return 0
|
||
|
||
|
||
def _detect_windows_gpu_memory() -> dict | None:
|
||
"""Return Windows GPU memory metadata from Win32_VideoController, if available."""
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"powershell",
|
||
"-NoProfile",
|
||
"-Command",
|
||
(
|
||
"Get-CimInstance Win32_VideoController | "
|
||
"Select-Object Name,AdapterRAM | ConvertTo-Json -Compress"
|
||
),
|
||
],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=5,
|
||
)
|
||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||
return None
|
||
|
||
if result.returncode != 0 or not result.stdout.strip():
|
||
return None
|
||
try:
|
||
raw = json.loads(result.stdout)
|
||
except json.JSONDecodeError:
|
||
return None
|
||
entries = raw if isinstance(raw, list) else [raw]
|
||
best: dict | None = None
|
||
for entry in entries:
|
||
if not isinstance(entry, dict):
|
||
continue
|
||
name = str(entry.get("Name") or "").strip()
|
||
if not name:
|
||
continue
|
||
try:
|
||
adapter_ram = int(entry.get("AdapterRAM") or 0)
|
||
except (TypeError, ValueError):
|
||
adapter_ram = 0
|
||
vram_mb = max(0, adapter_ram // (1024 * 1024))
|
||
if best is None or vram_mb > best["vram_mb"]:
|
||
best = {"gpu_name": name, "vram_mb": vram_mb}
|
||
return best
|
||
|
||
|
||
def _detect_nvidia_smi_gpu_memory() -> dict | None:
|
||
"""Return NVIDIA GPU memory metadata from nvidia-smi, if available."""
|
||
try:
|
||
result = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
|
||
capture_output=True, text=True, timeout=5,
|
||
)
|
||
if result.returncode == 0 and result.stdout.strip():
|
||
line = result.stdout.strip().splitlines()[0]
|
||
parts = line.split(",", 1)
|
||
gpu_name = parts[0].strip()
|
||
vram_mb = int(parts[1].strip()) if len(parts) > 1 else 0
|
||
return {"gpu_name": gpu_name, "vram_mb": max(0, vram_mb)}
|
||
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError):
|
||
pass
|
||
return None
|
||
|
||
|
||
def _detect_torch_cuda_inventory(torch_module) -> dict | None:
|
||
"""Return torch-visible CUDA/HIP GPU metadata without running kernels."""
|
||
try:
|
||
if not torch_module.cuda.is_available() or torch_module.cuda.device_count() < 1:
|
||
return None
|
||
idx = torch_module.cuda.current_device()
|
||
name = torch_module.cuda.get_device_name(idx)
|
||
props = torch_module.cuda.get_device_properties(idx)
|
||
vram_mb = int(props.total_memory // (1024 * 1024))
|
||
gpu = {"gpu_name": name, "vram_mb": max(0, vram_mb)}
|
||
gcn_arch = getattr(props, "gcnArchName", None)
|
||
if gcn_arch:
|
||
gpu["gcn_arch"] = str(gcn_arch)
|
||
return gpu
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _torch_cuda_is_executable(torch_module) -> bool:
|
||
"""Return True only if this Python process can execute a CUDA tensor op."""
|
||
try:
|
||
if not torch_module.cuda.is_available():
|
||
return False
|
||
probe = torch_module.empty((1,), device="cuda")
|
||
probe += 1
|
||
torch_module.cuda.synchronize()
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
|
||
if gpu is None:
|
||
return None
|
||
profile = {
|
||
"device": "cpu",
|
||
"gpu_name": gpu["gpu_name"],
|
||
"vram_mb": gpu["vram_mb"],
|
||
"dedicated_vram_mb": gpu["vram_mb"],
|
||
"shared_vram_mb": max(0, ram_mb // 2),
|
||
"ram_mb": ram_mb,
|
||
"cuda_available": False,
|
||
}
|
||
if gpu.get("gcn_arch"):
|
||
profile["gcn_arch"] = gpu["gcn_arch"]
|
||
return profile
|
||
|
||
|
||
def detect_hardware() -> dict:
|
||
"""Detect GPU model and available VRAM. Returns hardware profile dict."""
|
||
ram_mb = _detect_ram_mb()
|
||
torch_gpu: dict | None = None
|
||
try:
|
||
import torch # type: ignore[import]
|
||
|
||
torch_gpu = _detect_torch_cuda_inventory(torch)
|
||
if _torch_cuda_is_executable(torch):
|
||
if torch_gpu is None:
|
||
torch_gpu = _detect_torch_cuda_inventory(torch)
|
||
name = torch_gpu["gpu_name"] if torch_gpu is not None else "CUDA GPU"
|
||
vram_mb = torch_gpu["vram_mb"] if torch_gpu is not None else 0
|
||
shared_vram_mb = max(0, ram_mb // 2)
|
||
profile = {
|
||
"device": "cuda",
|
||
"gpu_name": name,
|
||
"vram_mb": vram_mb,
|
||
"dedicated_vram_mb": vram_mb,
|
||
"shared_vram_mb": shared_vram_mb,
|
||
"ram_mb": ram_mb,
|
||
"cuda_available": True,
|
||
}
|
||
if torch_gpu is not None and torch_gpu.get("gcn_arch"):
|
||
profile["gcn_arch"] = torch_gpu["gcn_arch"]
|
||
return profile
|
||
except ImportError:
|
||
pass
|
||
|
||
torch_inventory = _gpu_inventory_profile(torch_gpu, ram_mb)
|
||
if torch_inventory is not None:
|
||
return torch_inventory
|
||
|
||
nvidia_gpu = _gpu_inventory_profile(_detect_nvidia_smi_gpu_memory(), ram_mb)
|
||
if nvidia_gpu is not None:
|
||
return nvidia_gpu
|
||
|
||
windows_gpu = _gpu_inventory_profile(_detect_windows_gpu_memory(), ram_mb)
|
||
if windows_gpu is not None:
|
||
return windows_gpu
|
||
|
||
return {
|
||
"device": "cpu",
|
||
"gpu_name": None,
|
||
"vram_mb": 0,
|
||
"dedicated_vram_mb": 0,
|
||
"shared_vram_mb": 0,
|
||
"ram_mb": ram_mb,
|
||
"cuda_available": False,
|
||
}
|
||
|
||
|
||
def benchmark_throughput_checked(device_str: str = "cpu") -> tuple[float, bool, str | None]:
|
||
"""
|
||
Estimate compute throughput via a synthetic transformer GEMM benchmark.
|
||
|
||
Runs hidden_size × (hidden_size*4) matmul — the dominant op in FFN layers —
|
||
and returns iterations/second as a relative speed index. Higher = faster.
|
||
The value is used as benchmark_tokens_per_sec in tracker registration for
|
||
routing tiebreaks; it is not an absolute token rate.
|
||
|
||
Returns (score, ok, error). Score falls back to 1.0 when the requested
|
||
device cannot run the benchmark.
|
||
"""
|
||
try:
|
||
import torch # type: ignore[import]
|
||
|
||
device = torch.device(device_str)
|
||
# bfloat16 on CUDA matches real inference dtype; float32 on CPU avoids
|
||
# precision-downcast surprises on older hardware without bfloat16 support.
|
||
dtype = torch.bfloat16 if device_str == "cuda" else torch.float32
|
||
|
||
# hidden_size=2048 is representative of a mid-sized model; large enough
|
||
# that BLAS finds an efficient kernel on both GPU and CPU.
|
||
hidden_size = 2048
|
||
a = torch.randn(1, hidden_size, dtype=dtype, device=device)
|
||
b = torch.randn(hidden_size, hidden_size * 4, dtype=dtype, device=device)
|
||
|
||
def _sync() -> None:
|
||
if device_str == "cuda":
|
||
torch.cuda.synchronize()
|
||
|
||
# Warmup: prime caches and JIT compilation.
|
||
for _ in range(10):
|
||
torch.matmul(a, b)
|
||
_sync()
|
||
|
||
n_iters = 50
|
||
t0 = time.perf_counter()
|
||
for _ in range(n_iters):
|
||
torch.matmul(a, b)
|
||
_sync()
|
||
elapsed = time.perf_counter() - t0
|
||
|
||
return round(n_iters / max(elapsed, 1e-9), 2), True, None
|
||
except Exception as exc:
|
||
return 1.0, False, f"{type(exc).__name__}: {exc}"
|
||
|
||
|
||
def benchmark_throughput(device_str: str = "cpu") -> float:
|
||
"""Return only the numeric throughput index, preserving the legacy API."""
|
||
score, _ok, _error = benchmark_throughput_checked(device_str)
|
||
return score
|