fn_registry/python/functions/ml/cuda_available.py

"""Detecta disponibilidad de CUDA via torch sin lanzar excepcion si torch no esta instalado."""

from __future__ import annotations


def cuda_available() -> dict:
    """Detecta si CUDA esta disponible y devuelve info de los dispositivos GPU.

    No requiere torch instalado: si no esta presente, devuelve
    `torch_version='not_installed'` y `available=False`.

    Returns:
        dict con claves:
            available (bool): True si torch.cuda.is_available().
            device_count (int): numero de GPUs detectadas (0 si no hay CUDA).
            devices (list[str]): nombres de cada GPU (ej. "NVIDIA RTX 4090").
            torch_version (str): version de torch o "not_installed".
            cuda_version (str | None): version de CUDA usada por torch, o None.
    """
    try:
        import torch
    except ImportError:
        return {
            "available": False,
            "device_count": 0,
            "devices": [],
            "torch_version": "not_installed",
            "cuda_version": None,
        }

    available = torch.cuda.is_available()
    device_count = torch.cuda.device_count() if available else 0
    devices = [torch.cuda.get_device_name(i) for i in range(device_count)]
    cuda_version = torch.version.cuda if available else None

    return {
        "available": available,
        "device_count": device_count,
        "devices": devices,
        "torch_version": torch.__version__,
        "cuda_version": cuda_version,
    }