# -*- coding: utf-8 -*- """Unified accelerator detection and runtime validation.""" from __future__ import annotations import os from functools import lru_cache from typing import Optional from .accelerators import ( AcceleratorInfo, IluvatarAcceleratorAdapter, MetaxAcceleratorAdapter, MThreadsAcceleratorAdapter, NvidiaAcceleratorAdapter, ) SUPPORTED_ACCELERATORS = {"auto", "cpu", "nvidia", "metax", "iluvatar", "mthreads"} def _cpu_info(reason: str = "") -> AcceleratorInfo: return AcceleratorInfo( vendor="cpu", runtime="cpu", device="cpu", available=True, reason=reason, ) def _normalize_configured(configured: Optional[str]) -> str: value = (configured or os.getenv("ACCELERATOR") or "auto").strip().lower() if value == "cuda": return "nvidia" if value in {"maca", "muxi", "mx"}: return "metax" if value in {"ix", "tianshu", "天数"}: return "iluvatar" if value in {"mthreads", "musa", "moorethreads", "摩尔线程"}: return "mthreads" if value not in SUPPORTED_ACCELERATORS: raise ValueError( f"Unsupported ACCELERATOR={value!r}; expected one of: " f"{', '.join(sorted(SUPPORTED_ACCELERATORS))}" ) return value def _detect_uncached(configured: Optional[str]) -> AcceleratorInfo: accelerator = _normalize_configured(configured) if accelerator == "cpu": return _cpu_info("forced by ACCELERATOR=cpu") if accelerator == "nvidia": return NvidiaAcceleratorAdapter().detect() if accelerator == "metax": return MetaxAcceleratorAdapter().detect() if accelerator == "iluvatar": return IluvatarAcceleratorAdapter().detect() if accelerator == "mthreads": return MThreadsAcceleratorAdapter().detect() # Auto mode prefers vendor SMI commands before NVIDIA. This avoids a vendor # PyTorch build exposing torch.cuda and being mistaken for NVIDIA. metax = MetaxAcceleratorAdapter().detect() if metax.available: return metax iluvatar = IluvatarAcceleratorAdapter().detect() if iluvatar.available: return iluvatar mthreads = MThreadsAcceleratorAdapter().detect() if mthreads.available: return mthreads nvidia = NvidiaAcceleratorAdapter().detect() if nvidia.available: return nvidia return _cpu_info("no supported accelerator detected") @lru_cache(maxsize=8) def _detect_cached(configured: str) -> AcceleratorInfo: return _detect_uncached(configured) def detect_accelerator(configured: Optional[str] = None, *, refresh: bool = False) -> AcceleratorInfo: """Detect the active accelerator. Args: configured: Optional override matching ACCELERATOR values. refresh: Clear cached detection before probing. """ normalized = _normalize_configured(configured) if refresh: _detect_cached.cache_clear() return _detect_cached(normalized) def get_accelerator_info(*, refresh: bool = False) -> AcceleratorInfo: try: from app.core.config import settings configured = settings.ACCELERATOR except Exception: configured = os.getenv("ACCELERATOR", "auto") return detect_accelerator(configured, refresh=refresh) def validate_accelerator_runtime() -> tuple[bool, str]: """Validate explicit accelerator selections before model loading.""" try: from app.core.config import settings configured = _normalize_configured(settings.ACCELERATOR) except Exception: configured = _normalize_configured(os.getenv("ACCELERATOR", "auto")) info = detect_accelerator(configured, refresh=True) if configured in {"auto", "cpu", ""}: return True, "" if not info.available: return False, f"ACCELERATOR={configured} requested but {info.reason}" if configured == "metax" and not info.metadata.get("torch_cuda_available"): return ( False, "ACCELERATOR=metax detected mx-smi devices, but the active Python " "environment does not expose torch.cuda. Run ./scripts/sync_metax_env.sh " "or install the MetaX MACA PyTorch stack.", ) if configured == "iluvatar" and not info.metadata.get("torch_cuda_available"): return ( False, "ACCELERATOR=iluvatar detected ixsmi devices, but the active Python " "environment does not expose torch.cuda. Use the official Iluvatar " "vLLM image as the base image or install the Iluvatar PyTorch stack.", ) if configured == "mthreads" and not info.metadata.get("torch_cuda_available"): return ( False, "ACCELERATOR=mthreads detected mthreads-gmi devices, but the active Python " "environment does not expose torch.cuda. Use the official Moore Threads " "MUSA vLLM image as the base image or install the matching MUSA PyTorch stack.", ) if configured == "nvidia" and not info.metadata.get("torch_cuda_available"): return ( False, "ACCELERATOR=nvidia detected NVIDIA devices, but torch.cuda is not " "available in the active Python environment. Run ./scripts/sync_gpu_env.sh.", ) return True, ""