Spaces:
Sleeping
Sleeping
File size: 3,879 Bytes
8ad9950 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """Verify the inference device available to FactoryFlow.
Run as a module: ``python -m src.inference.rocm_check``
Selection order (override with ``AMD_DEVICE`` env var):
1. AMD ROCm GPU (torch built with ROCm exposes ``torch.version.hip``)
2. NVIDIA CUDA GPU (useful for local dev on non-AMD boxes)
3. Apple MPS (MacBook fallback)
4. CPU
"""
from __future__ import annotations
import os
import platform
from dataclasses import dataclass
import structlog
import torch
log = structlog.get_logger()
@dataclass
class DeviceInfo:
backend: str # "rocm" | "cuda" | "mps" | "cpu"
torch_device: str # value to pass to ``.to(...)``
name: str
vram_gb: float | None
runtime_version: str | None
def as_dict(self) -> dict[str, object]:
return {
"backend": self.backend,
"torch_device": self.torch_device,
"name": self.name,
"vram_gb": self.vram_gb,
"runtime_version": self.runtime_version,
}
def _detect_rocm() -> DeviceInfo | None:
if not torch.cuda.is_available():
return None
hip_version = getattr(torch.version, "hip", None)
if not hip_version:
return None
props = torch.cuda.get_device_properties(0)
return DeviceInfo(
backend="rocm",
torch_device="cuda", # ROCm exposes the CUDA-compatible API
name=props.name,
vram_gb=round(props.total_memory / 1024**3, 1),
runtime_version=hip_version,
)
def _detect_cuda() -> DeviceInfo | None:
if not torch.cuda.is_available():
return None
if getattr(torch.version, "hip", None):
return None # already handled by ROCm path
props = torch.cuda.get_device_properties(0)
return DeviceInfo(
backend="cuda",
torch_device="cuda",
name=props.name,
vram_gb=round(props.total_memory / 1024**3, 1),
runtime_version=torch.version.cuda,
)
def _detect_mps() -> DeviceInfo | None:
mps = getattr(torch.backends, "mps", None)
if mps is None or not mps.is_available():
return None
return DeviceInfo(
backend="mps",
torch_device="mps",
name=f"Apple MPS ({platform.processor() or platform.machine()})",
vram_gb=None,
runtime_version=None,
)
def _detect_cpu() -> DeviceInfo:
return DeviceInfo(
backend="cpu",
torch_device="cpu",
name=platform.processor() or platform.machine() or "CPU",
vram_gb=None,
runtime_version=None,
)
def detect_device() -> DeviceInfo:
"""Return the best available device, honoring ``AMD_DEVICE`` override."""
override = os.getenv("AMD_DEVICE", "").strip().lower()
if override == "cpu":
return _detect_cpu()
info = _detect_rocm() or _detect_cuda() or _detect_mps() or _detect_cpu()
log.info(
"device_detected",
component="inference.rocm_check",
backend=info.backend,
torch_device=info.torch_device,
name=info.name,
vram_gb=info.vram_gb,
runtime_version=info.runtime_version,
torch_version=torch.__version__,
)
return info
def main() -> None:
info = detect_device()
print(f"torch: {torch.__version__}")
print(f"backend: {info.backend}")
print(f"torch_device: {info.torch_device}")
print(f"name: {info.name}")
print(f"vram_gb: {info.vram_gb}")
print(f"runtime_version: {info.runtime_version}")
if info.backend == "rocm":
print("✓ AMD ROCm GPU detected — ready for MI300X demo run.")
elif info.backend in ("cuda", "mps"):
print(f"⚠ Running on {info.backend} — fine for local dev, swap to ROCm for the demo.")
else:
print("⚠ CPU only — inference will be slow; use for unit tests only.")
if __name__ == "__main__":
main()
|