src/fpgm/utils/gpu.py

"""GPU selection and diagnostics for a shared, multi-tenant box.

``torch`` is deliberately not a package-level import here (it isn't even a
``pyproject.toml`` dependency -- see that file's comment on the cu128
constraint): every function imports it lazily so this module stays importable
before the environment is fully set up.
"""

from __future__ import annotations

from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

# The actual root cause on this host, ~100% of the time a bare "CUDA not
# available" shows up, is the kernel driver (570.172.08) and userspace libs
# (580.173.02) being out of sync -- not a missing GPU or a bad torch install.
# Spelling out the fix inline saves a multi-minute GPU-forensics detour.
_SHIM_REMEDY = (
    "CUDA is not available to torch. If `nvidia-smi` fails with "
    '"Failed to initialize NVML: Driver/library version mismatch", '
    "this host's kernel driver (570.172.08) is older than its userspace "
    "libraries (580.173.02). Fix with no sudo/reboot required:\n"
    "    source scripts/nvidia_lib_shim.sh\n"
    "or, to run a single command under the shim:\n"
    "    ./scripts/nvidia_lib_shim.sh nvidia-smi"
)


def ensure_cuda() -> None:
    """Raise ``RuntimeError`` with the shim remedy if CUDA is unavailable."""
    import torch

    if not torch.cuda.is_available():
        raise RuntimeError(_SHIM_REMEDY)


def select_device(gpu_index: int | None = None):
    """Pick a CUDA device: ``gpu_index`` if given, else the one with most free memory.

    Picking by free memory rather than always defaulting to index 0 matters
    because this box is shared with other jobs -- each with only 11-27 GB free
    per GPU -- so blindly taking GPU 0 would pile every run onto whichever
    device happens to be first regardless of how loaded it already is.

    Returns:
        ``torch.device``.
    """
    import torch

    ensure_cuda()
    if gpu_index is not None:
        device = torch.device(f"cuda:{gpu_index}")
        torch.cuda.set_device(device)
        return device

    best_index, best_free = 0, -1
    for i in range(torch.cuda.device_count()):
        free, _total = torch.cuda.mem_get_info(i)
        if free > best_free:
            best_index, best_free = i, free
    device = torch.device(f"cuda:{best_index}")
    torch.cuda.set_device(device)
    logger.info("selected cuda:%d (%.1f GB free)", best_index, best_free / 1e9)
    return device


def log_gpu_memory() -> None:
    """Log free/total memory for every visible CUDA device."""
    import torch

    if not torch.cuda.is_available():
        logger.warning("no CUDA devices visible")
        return
    for i in range(torch.cuda.device_count()):
        free, total = torch.cuda.mem_get_info(i)
        name = torch.cuda.get_device_name(i)
        logger.info("cuda:%d %s: %.1f / %.1f GB free", i, name, free / 1e9, total / 1e9)