Reference for ultralytics/utils/autodevice.py

def __del__(self):
    """Ensure NVML is shut down when the object is garbage collected."""
    self.shutdown()

def print_status(self):
    """Print GPU status in a compact table format using current stats."""
    self.refresh_stats()
    if not self.gpu_stats:
        LOGGER.warning("No GPU stats available.")
        return

    stats = self.gpu_stats
    name_len = max(len(gpu.get("name", "N/A")) for gpu in stats)
    hdr = f"{'Idx':<3} {'Name':<{name_len}} {'Util':>6} {'Mem (MiB)':>15} {'Temp':>5} {'Pwr (W)':>10}"
    LOGGER.info(f"\n--- GPU Status ---\n{hdr}\n{'-' * len(hdr)}")

    for gpu in stats:
        u = f"{gpu['utilization']:>5}%" if gpu["utilization"] >= 0 else " N/A "
        m = f"{gpu['memory_used']:>6}/{gpu['memory_total']:<6}" if gpu["memory_used"] >= 0 else " N/A / N/A "
        t = f"{gpu['temperature']}C" if gpu["temperature"] >= 0 else " N/A "
        p = f"{gpu['power_draw']:>3}/{gpu['power_limit']:<3}" if gpu["power_draw"] >= 0 else " N/A "

        LOGGER.info(f"{gpu.get('index'):<3d} {gpu.get('name', 'N/A'):<{name_len}} {u:>6} {m:>15} {t:>5} {p:>10}")

    LOGGER.info(f"{'-' * len(hdr)}\n")

def refresh_stats(self):
    """Refresh the internal gpu_stats list by querying NVML."""
    self.gpu_stats = []
    if not self.nvml_available or not self.pynvml:
        return

    try:
        device_count = self.pynvml.nvmlDeviceGetCount()
        for i in range(device_count):
            self.gpu_stats.append(self._get_device_stats(i))
    except Exception as e:
        LOGGER.warning(f"Error during device query: {e}")
        self.gpu_stats = []

def select_idle_gpu(
    self, count: int = 1, min_memory_fraction: float = 0, min_util_fraction: float = 0
) -> List[int]:
    """
    Select the most idle GPUs based on utilization and free memory.

    Args:
        count (int): The number of idle GPUs to select.
        min_memory_fraction (float): Minimum free memory required as a fraction of total memory.
        min_util_fraction (float): Minimum free utilization rate required from 0.0 - 1.0.

    Returns:
        (List[int]): Indices of the selected GPUs, sorted by idleness (lowest utilization first).

    Notes:
         Returns fewer than 'count' if not enough qualify or exist.
         Returns basic CUDA indices if NVML fails. Empty list if no GPUs found.
    """
    assert min_memory_fraction <= 1.0, f"min_memory_fraction must be <= 1.0, got {min_memory_fraction}"
    assert min_util_fraction <= 1.0, f"min_util_fraction must be <= 1.0, got {min_util_fraction}"
    LOGGER.info(
        f"Searching for {count} idle GPUs with free memory >= {min_memory_fraction * 100:.1f}% and free utilization >= {min_util_fraction * 100:.1f}%..."
    )

    if count <= 0:
        return []

    self.refresh_stats()
    if not self.gpu_stats:
        LOGGER.warning("NVML stats unavailable.")
        return []

    # Filter and sort eligible GPUs
    eligible_gpus = [
        gpu
        for gpu in self.gpu_stats
        if gpu.get("memory_free", 0) / gpu.get("memory_total", 1) >= min_memory_fraction
        and (100 - gpu.get("utilization", 100)) >= min_util_fraction * 100
    ]
    eligible_gpus.sort(key=lambda x: (x.get("utilization", 101), -x.get("memory_free", 0)))

    # Select top 'count' indices
    selected = [gpu["index"] for gpu in eligible_gpus[:count]]

    if selected:
        LOGGER.info(f"Selected idle CUDA devices {selected}")
    else:
        LOGGER.warning(
            f"No GPUs met criteria (Free Mem >= {min_memory_fraction * 100:.1f}% and Free Util >= {min_util_fraction * 100:.1f}%)."
        )

    return selected

def shutdown(self):
    """Shut down NVML if it was initialized."""
    if self.nvml_available and self.pynvml:
        try:
            self.pynvml.nvmlShutdown()
        except Exception:
            pass
        self.nvml_available = False