34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
"""GPU hardware detection with graceful CPU fallback."""
|
|
|
|
import subprocess
|
|
|
|
|
|
def detect_hardware() -> dict:
|
|
"""Detect GPU model and available VRAM. Returns hardware profile dict."""
|
|
try:
|
|
import torch # type: ignore[import]
|
|
if torch.cuda.is_available():
|
|
idx = torch.cuda.current_device()
|
|
name = torch.cuda.get_device_name(idx)
|
|
props = torch.cuda.get_device_properties(idx)
|
|
vram_mb = props.total_memory // (1024 * 1024)
|
|
return {"device": "cuda", "gpu_name": name, "vram_mb": vram_mb}
|
|
except ImportError:
|
|
pass
|
|
|
|
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 {"device": "cuda", "gpu_name": gpu_name, "vram_mb": vram_mb}
|
|
except (FileNotFoundError, subprocess.TimeoutExpired, ValueError, IndexError):
|
|
pass
|
|
|
|
return {"device": "cpu", "gpu_name": None, "vram_mb": 0}
|