Source code for gyoza.worker.worker_context.utils

"""Utilities for auto-detecting worker resources."""

import psutil

from gyoza.worker.worker_context.types import GPU, Resources




def detect_resources() -> Resources:
    """Auto-detect hardware resources available on this machine.

    Returns
    -------
    Resources
        Detected CPU cores, RAM, and GPUs.
    """
    cpu_cores = psutil.cpu_count(logical=True) or 1
    ram_mb = psutil.virtual_memory().total // (1024 * 1024)
    gpus = _detect_gpus()

    return Resources(cpu_cores=cpu_cores, ram_mb=ram_mb, gpus=gpus)



def _detect_gpus() -> list[GPU]:
    """Detect available NVIDIA GPUs using nvidia-ml-py.

    Returns
    -------
    list[GPU]
        List of detected GPUs with VRAM info.
    """
    gpus: list[GPU] = []

    try:
        from pynvml import (
            nvmlDeviceGetCount,
            nvmlDeviceGetHandleByIndex,
            nvmlDeviceGetMemoryInfo,
            nvmlInit,
            nvmlShutdown,
        )

        nvmlInit()
        device_count = nvmlDeviceGetCount()

        for i in range(device_count):
            handle = nvmlDeviceGetHandleByIndex(i)
            memory_info = nvmlDeviceGetMemoryInfo(handle)
            vram_mb = memory_info.total // (1024 * 1024)
            gpus.append(GPU(id=i, vram_mb=vram_mb, tags=["cuda"]))

        nvmlShutdown()

    except Exception:
        # nvidia-ml-py not available or no NVIDIA GPUs
        pass

    return gpus