device.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Resolve compute device: CUDA (incl. ROCm), Metal (MPS), or CPU."""
 
from __future__ import annotations
 
import json
import os
import platform
import shutil
import subprocess
from typing import Any
 
_VALID_OVERRIDES = frozenset({"cuda", "mps", "cpu", "auto"})
 
 
def _import_torch():
    import torch
 
    return torch
 
 
def cuda_available() -> bool:
    try:
        return bool(_import_torch().cuda.is_available())
    except Exception:
        return False
 
 
def mps_available() -> bool:
    try:
        torch = _import_torch()
        mps = getattr(torch.backends, "mps", None)
        return bool(mps and mps.is_available())
    except Exception:
        return False
 
 
def is_rocm_build() -> bool:
    try:
        torch = _import_torch()
        return getattr(torch.version, "hip", None) is not None
    except Exception:
        return False
 
 
def cuda_backend_label() -> str:
    """PyTorch still uses device type 'cuda' for ROCm builds."""
    return "rocm" if is_rocm_build() else "cuda"
 
 
def _env_override() -> str | None:
    raw = (os.environ.get("AQUIN_DEVICE") or "").strip().lower()
    if not raw or raw == "auto":
        return None
    if raw not in _VALID_OVERRIDES:
        return None
    return raw
 
 
def resolve_compute_device(*, prefer: str | None = None) -> str:
    """
    Pick torch device string: cuda → mps → cpu.
    Override with AQUIN_DEVICE=cuda|mps|cpu|auto or `prefer`.
    """
    override = (prefer or _env_override() or "").strip().lower() or None
 
    if override == "cpu":
        return "cpu"
    if override == "cuda":
        if cuda_available():
            return "cuda"
        if mps_available():
            return "mps"
        return "cpu"
    if override == "mps":
        if mps_available():
            return "mps"
        if cuda_available():
            return "cuda"
        return "cpu"
 
    if cuda_available():
        return "cuda"
    if mps_available():
        return "mps"
    return "cpu"
 
 
def default_dtype_for_device(device: str | None = None) -> Any:
    torch = _import_torch()
    dev = device or resolve_compute_device()
    if dev == "cuda":
        return torch.bfloat16
    if dev == "mps":
        return torch.float16
    return torch.float32
 
 
def device_and_dtype(*, prefer: str | None = None) -> tuple[str, Any]:
    """Convenience for train/eval modules that need both together."""
    dev = resolve_compute_device(prefer=prefer)
    return dev, default_dtype_for_device(dev)
 
 
def accelerator_available() -> bool:
    return resolve_compute_device() != "cpu"
 
 
def allow_cpu_load() -> bool:
    return os.environ.get("AQUIN_ALLOW_CPU", "").strip().lower() in (
        "1",
        "true",
        "yes",
        "on",
    )
 
 
def empty_device_cache(device: str | None = None) -> None:
    dev = device or resolve_compute_device()
    try:
        torch = _import_torch()
        if dev == "cuda" and torch.cuda.is_available():
            torch.cuda.empty_cache()
        elif dev == "mps" and mps_available() and hasattr(torch, "mps"):
            empty = getattr(torch.mps, "empty_cache", None)
            if callable(empty):
                empty()
    except Exception:
        pass
 
 
def synchronize_device(device: str | None = None) -> None:
    dev = device or resolve_compute_device()
    try:
        torch = _import_torch()
        if dev == "cuda" and torch.cuda.is_available():
            torch.cuda.synchronize()
        elif dev == "mps" and mps_available() and hasattr(torch, "mps"):
            sync = getattr(torch.mps, "synchronize", None)
            if callable(sync):
                sync()
    except Exception:
        pass
 
 
def is_oom_error(exc: BaseException) -> bool:
    if exc.__class__.__name__ in ("OutOfMemoryError", "CUDAOutOfMemoryError"):
        return True
    msg = str(exc).lower()
    if "out of memory" in msg:
        return True
    if "cuda out of memory" in msg:
        return True
    if "mps" in msg and "memory" in msg:
        return True
    return False
 
 
def _cuda_gpu_entries() -> list[dict[str, Any]]:
    entries: list[dict[str, Any]] = []
    try:
        torch = _import_torch()
        if not torch.cuda.is_available():
            return entries
        for idx in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(idx)
            entries.append({
                "index": idx,
                "name": props.name,
                "vram_gb": round(props.total_memory / (1024 ** 3), 1),
            })
    except Exception:
        pass
    return entries
 
 
def _is_windows() -> bool:
    return platform.system().lower() in ("windows", "win32") or os.name == "nt"
 
 
def _rocm_install_hint() -> str | None:
    if platform.system().lower() != "linux":
        return None
    if cuda_available():
        return None
    return (
        "Linux AMD GPU: install PyTorch with ROCm "
        "(https://pytorch.org/get-started/locally/) — pip's default torch is often CPU-only."
    )
 
 
def torch_version_label() -> str | None:
    try:
        return str(_import_torch().__version__)
    except Exception:
        return None
 
 
def torch_is_cpu_wheel() -> bool:
    label = (torch_version_label() or "").lower()
    return "+cpu" in label
 
 
def nvidia_smi_gpu_names() -> list[str]:
    if not shutil.which("nvidia-smi"):
        return []
    try:
        proc = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True,
            text=True,
            timeout=8,
            check=False,
        )
        if proc.returncode != 0:
            return []
        return [line.strip() for line in proc.stdout.splitlines() if line.strip()]
    except Exception:
        return []
 
 
def _nvidia_cpu_torch_hint() -> str | None:
    names = nvidia_smi_gpu_names()
    if not names or cuda_available():
        return None
    ver = torch_version_label() or "unknown"
    if torch_is_cpu_wheel() or not cuda_available():
        return (
            f"NVIDIA GPU detected ({names[0]}) but PyTorch is CPU-only ({ver}). "
            "Reinstall the CUDA build of torch/torchvision from pytorch.org."
        )
    return (
        f"NVIDIA GPU detected ({names[0]}) but torch.cuda.is_available() is false — "
        "check NVIDIA driver / CUDA runtime."
    )
 
 
def _best_fix_hint() -> str | None:
    return (
        _nvidia_cpu_torch_hint()
        or _darwin_cpu_hint()
        or _rocm_install_hint()
        or (_windows_amd_hint() if not nvidia_smi_gpu_names() else None)
        or "Set AQUIN_ALLOW_CPU=1 for CPU-only smoke tests."
    )
 
 
def _windows_amd_hint() -> str | None:
    if not _is_windows():
        return None
    if cuda_available() or mps_available():
        return None
    return (
        "Windows AMD GPU: use WSL2 + Linux ROCm PyTorch, or CPU with AQUIN_ALLOW_CPU=1."
    )
 
 
def _darwin_cpu_hint() -> str | None:
    if platform.system().lower() != "darwin":
        return None
    if mps_available():
        return None
    return "Apple Silicon needs PyTorch with MPS (Metal); reinstall compute extras in ~/.aquin/venv."
 
 
def load_blocked_message(model_id: str) -> str:
    dev = resolve_compute_device()
    if dev != "cpu":
        return f"Cannot load '{model_id}' on device {dev}."
    hints = []
    primary = _best_fix_hint()
    if primary:
        hints.append(primary)
    if not primary or "AQUIN_ALLOW_CPU" not in primary:
        hints.append("For slow CPU smoke tests: AQUIN_ALLOW_CPU=1")
    hint_lines = [h for h in hints if h]
    body = (
        f"Cannot load '{model_id}': no GPU backend (CUDA/ROCm or Metal MPS).\n"
        + "\n".join(f"  · {h}" for h in hint_lines)
    )
    return body
 
 
def require_load_device(model_id: str, cfg: dict[str, Any] | None = None) -> str:
    """
    Device for aquin load model. Raises RuntimeError when only CPU and not allowed.
    """
    device = resolve_compute_device()
    if device != "cpu":
        return device
    if allow_cpu_load():
        return "cpu"
    raise RuntimeError(load_blocked_message(model_id))
 
 
def probe_backend() -> dict[str, Any]:
    """JSON-friendly backend snapshot for status / diagnostics."""
    selected = resolve_compute_device()
    cuda = cuda_available()
    mps = mps_available()
    gpus = _cuda_gpu_entries()
    backend = "cpu"
    if selected == "cuda":
        backend = cuda_backend_label()
    elif selected == "mps":
        backend = "mps"
 
    nvidia_hw = nvidia_smi_gpu_names()
    torch_ver = torch_version_label()
 
    summary_parts = [f"device={selected}", f"backend={backend}"]
    if gpus:
        g = gpus[0]
        summary_parts.append(f"{g.get('name', 'GPU')} · {g.get('vram_gb', '?')} GB")
        if len(gpus) > 1:
            summary_parts.append(f"+{len(gpus) - 1} more")
    elif nvidia_hw and not cuda:
        summary_parts.append(
            f"{nvidia_hw[0]} visible to nvidia-smi · torch={torch_ver or '?'}"
        )
    elif selected == "mps":
        summary_parts.append("Apple Metal (MPS)")
    else:
        summary_parts.append("no GPU accelerator")
 
    fix = None
    if not accelerator_available():
        fix = _best_fix_hint()
 
    return {
        "selected": selected,
        "backend": backend,
        "cuda_available": cuda,
        "mps_available": mps,
        "rocm": is_rocm_build(),
        "accelerator_available": accelerator_available(),
        "allow_cpu": allow_cpu_load(),
        "gpus": gpus,
        "nvidia_smi_gpus": nvidia_hw,
        "torch_version": torch_ver,
        "torch_cpu_wheel": torch_is_cpu_wheel(),
        "platform": platform.system(),
        "summary": " · ".join(summary_parts),
        "fix": fix,
    }
 
 
def format_backend_summary() -> str:
    return probe_backend()["summary"]
 
 
def probe_backend_json() -> str:
    return json.dumps(probe_backend(), ensure_ascii=False)