fn_registry/python/functions/ml/torch_device_select.py

"""Selecciona el mejor torch device disponible segun preferencia."""

from __future__ import annotations

import warnings


def _cuda_available() -> bool:
    """Retorna True si torch esta instalado y CUDA disponible."""
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False


def _mps_available() -> bool:
    """Retorna True si torch esta instalado y MPS (Apple Silicon) disponible."""
    try:
        import torch
        return (
            hasattr(torch.backends, "mps")
            and torch.backends.mps.is_available()
        )
    except ImportError:
        return False


def _cuda_device_count() -> int:
    """Retorna el numero de dispositivos CUDA disponibles."""
    try:
        import torch
        return torch.cuda.device_count() if torch.cuda.is_available() else 0
    except ImportError:
        return 0


def torch_device_select(preference: str = "auto") -> str:
    """Selecciona el torch device optimo segun preferencia y disponibilidad.

    Con preference='auto': elige CUDA si disponible, luego MPS (Apple M1/M2),
    luego CPU. Para preferencias explicitas, valida disponibilidad y hace
    fallback a CPU con advertencia si el device solicitado no esta disponible.

    Args:
        preference: 'auto' | 'cuda' | 'cuda:N' | 'mps' | 'cpu'.
            'auto': detecta automaticamente el mejor device.
            'cuda': usa cuda:0 si disponible, fallback a cpu.
            'cuda:N': usa el dispositivo N si existe, fallback a cpu.
            'mps': usa MPS si disponible (Mac Apple Silicon), fallback a cpu.
            'cpu': siempre retorna 'cpu'.

    Returns:
        String de device para torch: 'cuda:0', 'cuda:N', 'mps' o 'cpu'.
    """
    if preference == "cpu":
        return "cpu"

    if preference == "auto":
        if _cuda_available():
            return "cuda:0"
        if _mps_available():
            return "mps"
        return "cpu"

    if preference == "mps":
        if _mps_available():
            return "mps"
        warnings.warn(
            "MPS no esta disponible en este sistema. Usando 'cpu'.",
            stacklevel=2,
        )
        return "cpu"

    if preference == "cuda":
        if _cuda_available():
            return "cuda:0"
        warnings.warn(
            "CUDA no esta disponible en este sistema. Usando 'cpu'.",
            stacklevel=2,
        )
        return "cpu"

    if preference.startswith("cuda:"):
        try:
            device_idx = int(preference.split(":")[1])
        except (IndexError, ValueError):
            warnings.warn(
                f"Formato de device no valido: '{preference}'. Usando 'cpu'.",
                stacklevel=2,
            )
            return "cpu"

        count = _cuda_device_count()
        if _cuda_available() and device_idx < count:
            return preference
        warnings.warn(
            f"Device '{preference}' no disponible "
            f"(cuda_count={count}). Usando 'cpu'.",
            stacklevel=2,
        )
        return "cpu"

    warnings.warn(
        f"Preferencia desconocida: '{preference}'. Usando 'cpu'.",
        stacklevel=2,
    )
    return "cpu"