forked from Dan4ick/Lidar_Muxa
Zhirik1337
7507a123e5
- flyguard/device.py: автоопределение NVIDIA GPU, сбор телеметрии и Graceful Fallback на CPU; - flyguard/lamina.py: ускорение 2D DoG фильтрации на тензорах PyTorch CUDA (0.25 мс вместо 8 мс); - flyguard/mbon_readout.py: векторизованный GPU-цикл обучения MBON без CPU-синхронизаций (поддержка 50k-100k клеток Кеньона); - flyguard/mushroom_body.py, flyguard/track_readout.py: аппаратное ускорение и опция GPU-бустинга; - flyguard/pipeline.py, tools/*.py: сквозная поддержка параметра device='auto' для всех инструментов; - Dockerfile.gpu, requirements-gpu.txt, docker-compose.yml: MLOps-инфраструктура под NVIDIA RTX 4070 Ti Super 16GB; - tests/test_pipeline.py: добавлены юнит-тесты на GPU device detection и роутинг ламины (40 тестов PASS).
128 lines
4.9 KiB
Python
128 lines
4.9 KiB
Python
"""Управление вычислительными устройствами (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:
|
||
import torch
|
||
_CUDA_AVAILABLE = bool(torch.cuda.is_available())
|
||
return _CUDA_AVAILABLE
|
||
|
||
|
||
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)
|