test/app/core/device.py

52 lines
1.4 KiB
Python

# -*- coding: utf-8 -*-
"""Centralized device detection utility.
This module keeps the historic device helpers while delegating hardware
probing to ``app.core.accelerator``.
"""
from app.core.accelerator import get_accelerator_info
def detect_device(configured: str = "auto") -> str:
"""Resolve a device configuration string to a concrete PyTorch device.
Priority for ``"auto"``: configured accelerator > detected GPU > CPU.
Args:
configured: Value from ``settings.DEVICE`` or caller override.
Accepted: ``"auto"``, ``"cpu"``, ``"cuda:0"``, ``"npu:0"``, etc.
Returns:
A device string ready for ``torch.device()`` / FunASR / ModelScope.
"""
device = configured.strip().lower()
if device == "auto":
return get_accelerator_info().device
# Normalize bare "cuda" to "cuda:0"
if device == "cuda":
return "cuda:0"
if device == "mps":
return "cpu"
return device
def is_cuda() -> bool:
"""True when the active runtime exposes a CUDA-compatible device."""
info = get_accelerator_info()
return info.available and info.device.startswith("cuda")
def has_gpu() -> bool:
"""True when a supported accelerator is available."""
return get_accelerator_info().is_gpu
def get_vram_gb() -> float:
"""Return usable accelerator memory in GB."""
return get_accelerator_info().total_memory_gb