test/app/core/accelerator.py

# -*- 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, ""