258 lines
7.7 KiB
Python
258 lines
7.7 KiB
Python
"""
|
||
Hardware detection — run once at startup, read everywhere.
|
||
|
||
Usage:
|
||
# At FastAPI lifespan startup:
|
||
from utils.hardware import detect_hardware
|
||
detect_hardware()
|
||
|
||
# Anywhere else:
|
||
from utils.hardware import DEVICE, DeviceType, is_apple_silicon
|
||
if DEVICE == DeviceType.CUDA:
|
||
import torch
|
||
...
|
||
"""
|
||
import platform
|
||
import logging
|
||
from enum import Enum
|
||
from typing import Optional, Dict, Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ========== Device Enum ==========
|
||
|
||
class DeviceType(str, Enum):
|
||
"""Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
|
||
CUDA = "cuda"
|
||
MLX = "mlx"
|
||
CPU = "cpu"
|
||
|
||
|
||
# ========== Global State (set once by detect_hardware) ==========
|
||
|
||
DEVICE: Optional[DeviceType] = None
|
||
|
||
|
||
# ========== Detection ==========
|
||
|
||
def is_apple_silicon() -> bool:
|
||
"""Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
|
||
return platform.system() == "Darwin" and platform.machine() == "arm64"
|
||
|
||
|
||
def _has_torch() -> bool:
|
||
"""Check if PyTorch is importable."""
|
||
try:
|
||
import torch
|
||
return True
|
||
except ImportError:
|
||
return False
|
||
|
||
|
||
def _has_mlx() -> bool:
|
||
"""Check if MLX is importable."""
|
||
try:
|
||
import mlx.core
|
||
return True
|
||
except ImportError:
|
||
return False
|
||
|
||
|
||
def detect_hardware() -> DeviceType:
|
||
"""
|
||
Detect the best available compute device and set the module-level DEVICE global.
|
||
|
||
Should be called exactly once during FastAPI lifespan startup.
|
||
Safe to call multiple times (idempotent).
|
||
|
||
Detection order:
|
||
1. CUDA (NVIDIA GPU, requires torch)
|
||
2. MLX (Apple Silicon via MLX framework)
|
||
3. CPU (fallback)
|
||
"""
|
||
global DEVICE
|
||
|
||
# --- CUDA: try PyTorch ---
|
||
if _has_torch():
|
||
import torch
|
||
if torch.cuda.is_available():
|
||
DEVICE = DeviceType.CUDA
|
||
device_name = torch.cuda.get_device_properties(0).name
|
||
logger.info(f"Hardware detected: CUDA — {device_name}")
|
||
return DEVICE
|
||
|
||
# --- MLX: Apple Silicon ---
|
||
if is_apple_silicon() and _has_mlx():
|
||
DEVICE = DeviceType.MLX
|
||
chip = platform.processor() or platform.machine()
|
||
logger.info(f"Hardware detected: MLX — Apple Silicon ({chip})")
|
||
return DEVICE
|
||
|
||
# --- Fallback ---
|
||
DEVICE = DeviceType.CPU
|
||
logger.info("Hardware detected: CPU (no GPU backend available)")
|
||
return DEVICE
|
||
|
||
|
||
# ========== Convenience helpers ==========
|
||
|
||
def get_device() -> DeviceType:
|
||
"""
|
||
Return the detected device. Auto-detects if detect_hardware() hasn't been called yet.
|
||
Prefer calling detect_hardware() explicitly at startup instead.
|
||
"""
|
||
global DEVICE
|
||
if DEVICE is None:
|
||
detect_hardware()
|
||
return DEVICE
|
||
|
||
|
||
def clear_gpu_cache():
|
||
"""
|
||
Clear GPU memory cache for the current device.
|
||
Safe to call on any platform — no-ops gracefully.
|
||
"""
|
||
import gc
|
||
gc.collect()
|
||
|
||
device = get_device()
|
||
|
||
if device == DeviceType.CUDA:
|
||
import torch
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.ipc_collect()
|
||
elif device == DeviceType.MLX:
|
||
# MLX manages memory automatically; no explicit cache clear needed.
|
||
# mlx.core has no empty_cache equivalent — gc.collect() above is enough.
|
||
pass
|
||
|
||
|
||
def get_gpu_memory_info() -> Dict[str, Any]:
|
||
"""
|
||
Get GPU memory information.
|
||
Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only environments.
|
||
"""
|
||
device = get_device()
|
||
|
||
# ---- CUDA path ----
|
||
if device == DeviceType.CUDA:
|
||
try:
|
||
import torch
|
||
idx = torch.cuda.current_device()
|
||
props = torch.cuda.get_device_properties(idx)
|
||
|
||
total = props.total_memory
|
||
allocated = torch.cuda.memory_allocated(idx)
|
||
reserved = torch.cuda.memory_reserved(idx)
|
||
|
||
return {
|
||
"available": True,
|
||
"backend": device.value,
|
||
"device": idx,
|
||
"device_name": props.name,
|
||
"total_gb": total / (1024**3),
|
||
"allocated_gb": allocated / (1024**3),
|
||
"reserved_gb": reserved / (1024**3),
|
||
"free_gb": (total - allocated) / (1024**3),
|
||
"utilization_pct": (allocated / total) * 100,
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Error getting CUDA GPU info: {e}")
|
||
return {"available": False, "backend": device.value, "error": str(e)}
|
||
|
||
# ---- MLX path (Apple Silicon) ----
|
||
if device == DeviceType.MLX:
|
||
try:
|
||
import mlx.core as mx
|
||
import psutil
|
||
|
||
# MLX uses unified memory — report system memory as the pool
|
||
total = psutil.virtual_memory().total
|
||
# MLX doesn't expose per-process GPU allocation; report 0 as allocated
|
||
allocated = 0
|
||
|
||
return {
|
||
"available": True,
|
||
"backend": device.value,
|
||
"device": 0,
|
||
"device_name": f"Apple Silicon ({platform.processor() or platform.machine()})",
|
||
"total_gb": total / (1024**3),
|
||
"allocated_gb": allocated / (1024**3),
|
||
"reserved_gb": 0,
|
||
"free_gb": (total - allocated) / (1024**3),
|
||
"utilization_pct": (allocated / total) * 100 if total else 0,
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Error getting MLX GPU info: {e}")
|
||
return {"available": False, "backend": device.value, "error": str(e)}
|
||
|
||
# ---- CPU-only ----
|
||
return {"available": False, "backend": "cpu"}
|
||
|
||
|
||
def log_gpu_memory(context: str):
|
||
"""Log GPU memory usage with context."""
|
||
memory_info = get_gpu_memory_info()
|
||
if memory_info.get("available"):
|
||
backend = memory_info.get("backend", "unknown").upper()
|
||
device_name = memory_info.get("device_name", "")
|
||
label = f"{backend}" + (f" ({device_name})" if device_name else "")
|
||
logger.info(
|
||
f"GPU Memory [{context}] {label}: "
|
||
f"{memory_info['allocated_gb']:.2f}GB/{memory_info['total_gb']:.2f}GB "
|
||
f"({memory_info['utilization_pct']:.1f}% used, "
|
||
f"{memory_info['free_gb']:.2f}GB free)"
|
||
)
|
||
else:
|
||
logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)")
|
||
|
||
|
||
# ========== GPU Summary & Package Versions ==========
|
||
|
||
def get_gpu_summary() -> Dict[str, Any]:
|
||
"""
|
||
Return a compact summary of the primary GPU.
|
||
|
||
Returns dict with keys:
|
||
gpu_name – e.g. "NVIDIA L4" (or None)
|
||
vram_total_gb – e.g. 22.17 (or None)
|
||
"""
|
||
mem = get_gpu_memory_info()
|
||
if mem.get("available"):
|
||
return {
|
||
"gpu_name": mem.get("device_name"),
|
||
"vram_total_gb": round(mem.get("total_gb", 0), 2),
|
||
}
|
||
return {"gpu_name": None, "vram_total_gb": None}
|
||
|
||
|
||
def get_package_versions() -> Dict[str, Optional[str]]:
|
||
"""
|
||
Return the installed versions of key ML packages.
|
||
|
||
Uses importlib.metadata (stdlib) so no subprocess is needed.
|
||
CUDA version comes from torch.version.cuda.
|
||
|
||
Returns dict with keys: unsloth, torch, transformers, cuda.
|
||
Missing packages yield None.
|
||
"""
|
||
from importlib.metadata import version as pkg_version, PackageNotFoundError
|
||
|
||
packages = ("unsloth", "torch", "transformers")
|
||
versions: Dict[str, Optional[str]] = {}
|
||
|
||
for name in packages:
|
||
try:
|
||
versions[name] = pkg_version(name)
|
||
except PackageNotFoundError:
|
||
versions[name] = None
|
||
|
||
# CUDA toolkit version bundled with torch
|
||
try:
|
||
import torch
|
||
versions["cuda"] = getattr(torch.version, "cuda", None)
|
||
except Exception:
|
||
versions["cuda"] = None
|
||
|
||
return versions
|