fn_registry/python/functions/ml/gpu_info.py

"""Consulta informacion de GPUs NVIDIA via nvidia-smi."""

from __future__ import annotations

import csv
import subprocess


def gpu_info() -> list[dict]:
    """Devuelve informacion de todas las GPUs NVIDIA detectadas por nvidia-smi.

    Consulta nvidia-smi via subprocess. Si nvidia-smi no esta disponible o
    falla, devuelve lista vacia sin lanzar excepcion.

    Returns:
        Lista de dicts, uno por GPU, con claves:
            index (int): indice de la GPU (0, 1, ...).
            name (str): nombre del modelo (ej. "NVIDIA GeForce RTX 4090").
            vram_total_mb (int): memoria total en MB.
            vram_free_mb (int): memoria libre en MB.
            driver_version (str): version del driver NVIDIA.
            cuda_version (str): version maxima de CUDA soportada por el driver.
        Lista vacia si nvidia-smi no esta disponible.
    """
    try:
        result = subprocess.run(
            [
                "nvidia-smi",
                "--query-gpu=index,name,memory.total,memory.free,driver_version,compute_cap",
                "--format=csv,noheader,nounits",
            ],
            capture_output=True,
            text=True,
            timeout=10,
        )
    except FileNotFoundError:
        return []
    except subprocess.TimeoutExpired:
        return []
    except OSError:
        return []

    if result.returncode != 0:
        return []

    gpus = []
    reader = csv.reader(result.stdout.strip().splitlines())
    for row in reader:
        if len(row) < 5:
            continue
        try:
            index = int(row[0].strip())
            name = row[1].strip()
            vram_total_mb = int(row[2].strip())
            vram_free_mb = int(row[3].strip())
            driver_version = row[4].strip()
            # compute_cap (ej. "8.9") como aproximacion de cuda_version soportada
            cuda_version = row[5].strip() if len(row) > 5 else ""
        except (ValueError, IndexError):
            continue

        gpus.append(
            {
                "index": index,
                "name": name,
                "vram_total_mb": vram_total_mb,
                "vram_free_mb": vram_free_mb,
                "driver_version": driver_version,
                "cuda_version": cuda_version,
            }
        )

    return gpus