Brainrot_Muxa/flyguard/device.py

"""Управление вычислительными устройствами (NVIDIA GPU / CUDA и CPU fallback).

Модуль обеспечивает прозрачную работу конвейера FlyGuard на:
- Выделенном NVIDIA GPU (RTX 4070 Ti Super 16GB / Ada Lovelace) через CUDA / Tensor Cores.
- Любых других видеокартах NVIDIA (Ampere, Turing, Pascal, Hopper, Blackwell).
- CPU Linux при отсутствии GPU или в среде жюри (Graceful Fallback).
"""
from __future__ import annotations

import logging
from typing import Any

import numpy as np

logger = logging.getLogger("flyguard.device")

_TORCH_AVAILABLE: bool | None = None
_CUDA_AVAILABLE: bool | None = None


def is_torch_available() -> bool:
    """Проверка наличия установленного пакета PyTorch."""
    global _TORCH_AVAILABLE
    if _TORCH_AVAILABLE is None:
        try:
            import torch  # noqa: F401
            _TORCH_AVAILABLE = True
        except ImportError:
            _TORCH_AVAILABLE = False
    return _TORCH_AVAILABLE


def is_cuda_available() -> bool:
    """Проверка доступности ускорения NVIDIA CUDA."""
    global _CUDA_AVAILABLE
    if _CUDA_AVAILABLE is None:
        if not is_torch_available():
            _CUDA_AVAILABLE = False
        else:
            try:
                import torch
                _CUDA_AVAILABLE = bool(torch.cuda.is_available() and torch.cuda.device_count() > 0)
            except Exception as e:
                logger.warning("Проверка CUDA завершилась ошибкой: %s. Fallback на CPU.", e)
                _CUDA_AVAILABLE = False
    return _CUDA_AVAILABLE


def notify_cuda_error(exc: Exception | None = None) -> None:
    """Зафиксировать сбой CUDA в рантайме и принудительно перевести систему в режим CPU fallback.

    Вызывается, если во время работы на GPU произошёл OOM, таймаут или сбой драйвера.
    Последующие вызовы конвейера будут прозрачно исполняться на CPU.
    """
    global _CUDA_AVAILABLE
    _CUDA_AVAILABLE = False
    logger.warning("Зафиксирован сбой GPU в рантайме (%s). Выполнен необратимый Fallback на CPU.", exc)


def reset_device_cache() -> None:
    """Сбросить кэш состояния устройств (для юнит-тестов)."""
    global _TORCH_AVAILABLE, _CUDA_AVAILABLE
    _TORCH_AVAILABLE = None
    _CUDA_AVAILABLE = None


def get_device(preferred: str = "auto") -> str:
    """Выбрать вычислительное устройство с автоматическим fallback на CPU.

    Args:
        preferred: 'auto' (выбрать cuda при наличии, иначе cpu), 'cuda', 'cuda:0', 'cpu'.

    Returns:
        Строка устройства: 'cuda', 'cuda:X' или 'cpu'.
    """
    pref = (preferred or "auto").strip().lower()
    if pref == "auto":
        return "cuda" if is_cuda_available() else "cpu"
    if pref.startswith("cuda"):
        if is_cuda_available():
            return pref
        logger.warning("Запрошено устройство '%s', но CUDA недоступна. Выполнен fallback на CPU.", pref)
        return "cpu"
    return "cpu"


def get_device_info(device_str: str | None = None) -> dict[str, Any]:
    """Сводная информация об аппаратных ресурсах для инспекции и логирования."""
    dev = get_device(device_str or "auto")
    info: dict[str, Any] = {
        "device": dev,
        "torch_available": is_torch_available(),
        "cuda_available": is_cuda_available(),
    }
    if is_torch_available() and is_cuda_available() and dev.startswith("cuda"):
        import torch
        dev_idx = 0
        if ":" in dev:
            try:
                dev_idx = int(dev.split(":")[1])
            except ValueError:
                dev_idx = 0
        props = torch.cuda.get_device_properties(dev_idx)
        info.update({
            "name": props.name,
            "total_memory_mb": round(props.total_memory / (1024 * 1024), 2),
            "major": props.major,
            "minor": props.minor,
            "multi_processor_count": props.multi_processor_count,
            "allocated_mb": round(torch.cuda.memory_allocated(dev_idx) / (1024 * 1024), 2),
            "reserved_mb": round(torch.cuda.memory_reserved(dev_idx) / (1024 * 1024), 2),
            "cuda_version": torch.version.cuda,
            "cudnn_version": torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else None,
        })
    else:
        import platform
        info.update({
            "name": platform.processor() or "CPU",
            "python_version": platform.python_version(),
        })
    return info


def log_device_info(device_str: str | None = None) -> None:
    """Вывести в лог/консоль активное вычислительное окружение."""
    info = get_device_info(device_str)
    if info.get("cuda_available") and str(info["device"]).startswith("cuda"):
        logger.info("FlyGuard GPU Active: %s (VRAM: %.1f MB, CUDA: %s, SM: %d.%d)",
                    info.get("name"), info.get("total_memory_mb", 0.0),
                    info.get("cuda_version"), info.get("major", 0), info.get("minor", 0))
    else:
        logger.info("FlyGuard CPU Mode (Fallback): %s", info.get("name"))


def empty_cache() -> None:
    """Очистить кэш памяти CUDA для предотвращения фрагментации VRAM."""
    if is_cuda_available():
        import torch
        torch.cuda.empty_cache()


def to_numpy(arr: Any) -> np.ndarray:
    """Преобразовать входной тензор (PyTorch или NumPy) в numpy.ndarray."""
    if isinstance(arr, np.ndarray):
        return arr
    if is_torch_available():
        import torch
        if isinstance(arr, torch.Tensor):
            return arr.detach().cpu().numpy()
    return np.asarray(arr)