dgenerate-ultralytics-headless 8.3.137__py3-none-any.whl → 8.3.224__py3-none-any.whl

ultralytics/utils/autobatch.py

@@ -1,6 +1,8 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """Functions for estimating the best YOLO batch size to use a fraction of the available CUDA memory in PyTorch."""
 
+from __future__ import annotations
+
 import os
 from copy import deepcopy
 
@@ -11,15 +13,20 @@ from ultralytics.utils import DEFAULT_CFG, LOGGER, colorstr
 from ultralytics.utils.torch_utils import autocast, profile_ops
 
 
-def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):
-    """
-    Compute optimal YOLO training batch size using the autobatch() function.
+def check_train_batch_size(
+    model: torch.nn.Module,
+    imgsz: int = 640,
+    amp: bool = True,
+    batch: int | float = -1,
+    max_num_obj: int = 1,
+) -> int:
+    """Compute optimal YOLO training batch size using the autobatch() function.
 
     Args:
         model (torch.nn.Module): YOLO model to check batch size for.
         imgsz (int, optional): Image size used for training.
         amp (bool, optional): Use automatic mixed precision if True.
-        batch (float, optional): Fraction of GPU memory to use. If -1, use default.
+        batch (int | float, optional): Fraction of GPU memory to use. If -1, use default.
         max_num_obj (int, optional): The maximum number of objects from dataset.
 
     Returns:
@@ -35,9 +42,14 @@ def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):
         )
 
 
-def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch, max_num_obj=1):
-    """
-    Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
+def autobatch(
+    model: torch.nn.Module,
+    imgsz: int = 640,
+    fraction: float = 0.60,
+    batch_size: int = DEFAULT_CFG.batch,
+    max_num_obj: int = 1,
+) -> int:
+    """Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
 
     Args:
         model (torch.nn.Module): YOLO model to compute batch size for.

ultralytics/utils/autodevice.py

@@ -1,38 +1,56 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
+from __future__ import annotations
+
+from typing import Any
+
 from ultralytics.utils import LOGGER
 from ultralytics.utils.checks import check_requirements
 
 
 class GPUInfo:
-    """
-    Manages NVIDIA GPU information via pynvml with robust error handling.
+    """Manages NVIDIA GPU information via pynvml with robust error handling.
 
-    Provides methods to query detailed GPU statistics (utilization, memory, temp, power) and select the most idle
-    GPUs based on configurable criteria. It safely handles the absence or initialization failure of the pynvml
-    library by logging warnings and disabling related features, preventing application crashes.
+    Provides methods to query detailed GPU statistics (utilization, memory, temp, power) and select the most idle GPUs
+    based on configurable criteria. It safely handles the absence or initialization failure of the pynvml library by
+    logging warnings and disabling related features, preventing application crashes.
 
     Includes fallback logic using `torch.cuda` for basic device counting if NVML is unavailable during GPU
     selection. Manages NVML initialization and shutdown internally.
 
     Attributes:
         pynvml (module | None): The `pynvml` module if successfully imported and initialized, otherwise `None`.
-        nvml_available (bool): Indicates if `pynvml` is ready for use. True if import and `nvmlInit()` succeeded,
-            False otherwise.
-        gpu_stats (list[dict]): A list of dictionaries, each holding stats for one GPU. Populated on initialization
-            and by `refresh_stats()`. Keys include: 'index', 'name', 'utilization' (%), 'memory_used' (MiB),
-            'memory_total' (MiB), 'memory_free' (MiB), 'temperature' (C), 'power_draw' (W),
-            'power_limit' (W or 'N/A'). Empty if NVML is unavailable or queries fail.
+        nvml_available (bool): Indicates if `pynvml` is ready for use. True if import and `nvmlInit()` succeeded, False
+            otherwise.
+        gpu_stats (list[dict[str, Any]]): A list of dictionaries, each holding stats for one GPU, populated on
+        initialization and by `refresh_stats()`. Keys include: 'index', 'name', 'utilization' (%), 'memory_used' (MiB),
+            'memory_total' (MiB), 'memory_free' (MiB), 'temperature' (C), 'power_draw' (W), 'power_limit' (W or 'N/A').
+            Empty if NVML is unavailable or queries fail.
+
+    Methods:
+        refresh_stats: Refresh the internal gpu_stats list by querying NVML.
+        print_status: Print GPU status in a compact table format using current stats.
+        select_idle_gpu: Select the most idle GPUs based on utilization and free memory.
+        shutdown: Shut down NVML if it was initialized.
+
+    Examples:
+        Initialize GPUInfo and print status
+        >>> gpu_info = GPUInfo()
+        >>> gpu_info.print_status()
+
+        Select idle GPUs with minimum memory requirements
+        >>> selected = gpu_info.select_idle_gpu(count=2, min_memory_fraction=0.2)
+        >>> print(f"Selected GPU indices: {selected}")
     """
 
     def __init__(self):
-        """Initializes GPUInfo, attempting to import and initialize pynvml."""
-        self.pynvml = None
-        self.nvml_available = False
-        self.gpu_stats = []
+        """Initialize GPUInfo, attempting to import and initialize pynvml."""
+        self.pynvml: Any | None = None
+        self.nvml_available: bool = False
+        self.gpu_stats: list[dict[str, Any]] = []
 
         try:
-            check_requirements("pynvml>=12.0.0")
+            check_requirements("nvidia-ml-py>=12.0.0")
             self.pynvml = __import__("pynvml")
             self.pynvml.nvmlInit()
             self.nvml_available = True
@@ -41,11 +59,11 @@ class GPUInfo:
             LOGGER.warning(f"Failed to initialize pynvml, GPU stats disabled: {e}")
 
     def __del__(self):
-        """Ensures NVML is shut down when the object is garbage collected."""
+        """Ensure NVML is shut down when the object is garbage collected."""
         self.shutdown()
 
     def shutdown(self):
-        """Shuts down NVML if it was initialized."""
+        """Shut down NVML if it was initialized."""
         if self.nvml_available and self.pynvml:
             try:
                 self.pynvml.nvmlShutdown()
@@ -54,21 +72,20 @@ class GPUInfo:
             self.nvml_available = False
 
     def refresh_stats(self):
-        """Refreshes the internal gpu_stats list by querying NVML."""
+        """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))
+            self.gpu_stats.extend(self._get_device_stats(i) for i in range(device_count))
         except Exception as e:
             LOGGER.warning(f"Error during device query: {e}")
             self.gpu_stats = []
 
-    def _get_device_stats(self, index):
-        """Gets stats for a single GPU device."""
+    def _get_device_stats(self, index: int) -> dict[str, Any]:
+        """Get stats for a single GPU device."""
         handle = self.pynvml.nvmlDeviceGetHandleByIndex(index)
         memory = self.pynvml.nvmlDeviceGetMemoryInfo(handle)
         util = self.pynvml.nvmlDeviceGetUtilizationRates(handle)
@@ -86,16 +103,16 @@ class GPUInfo:
             "index": index,
             "name": self.pynvml.nvmlDeviceGetName(handle),
             "utilization": util.gpu if util else -1,
-            "memory_used": memory.used >> 20 if memory else -1,
+            "memory_used": memory.used >> 20 if memory else -1,  # Convert bytes to MiB
             "memory_total": memory.total >> 20 if memory else -1,
             "memory_free": memory.free >> 20 if memory else -1,
             "temperature": safe_get(self.pynvml.nvmlDeviceGetTemperature, handle, temp_type),
-            "power_draw": safe_get(self.pynvml.nvmlDeviceGetPowerUsage, handle, divisor=1000),
+            "power_draw": safe_get(self.pynvml.nvmlDeviceGetPowerUsage, handle, divisor=1000),  # Convert mW to W
             "power_limit": safe_get(self.pynvml.nvmlDeviceGetEnforcedPowerLimit, handle, divisor=1000),
         }
 
     def print_status(self):
-        """Prints GPU status in a compact table format using current stats."""
+        """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.")
@@ -116,22 +133,28 @@ class GPUInfo:
 
         LOGGER.info(f"{'-' * len(hdr)}\n")
 
-    def select_idle_gpu(self, count=1, min_memory_mb=0):
-        """
-        Selects the 'count' most idle GPUs based on utilization and free memory.
+    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. Defaults to 1.
-            min_memory_mb (int): Minimum free memory required (MiB). Defaults to 0.
+            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.
+            (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.
         """
-        LOGGER.info(f"Searching for {count} idle GPUs with >= {min_memory_mb} MiB free memory...")
+        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 []
@@ -145,7 +168,8 @@ class GPUInfo:
         eligible_gpus = [
             gpu
             for gpu in self.gpu_stats
-            if gpu.get("memory_free", -1) >= min_memory_mb and gpu.get("utilization", -1) != -1
+            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)))
 
@@ -155,20 +179,26 @@ class GPUInfo:
         if selected:
             LOGGER.info(f"Selected idle CUDA devices {selected}")
         else:
-            LOGGER.warning(f"No GPUs met criteria (Util != -1, Free Mem >= {min_memory_mb} MiB).")
+            LOGGER.warning(
+                f"No GPUs met criteria (Free Mem >= {min_memory_fraction * 100:.1f}% and Free Util >= {min_util_fraction * 100:.1f}%)."
+            )
 
         return selected
 
 
 if __name__ == "__main__":
-    required_free_mem = 2048  # Require 2GB free VRAM
+    required_free_mem_fraction = 0.2  # Require 20% free VRAM
+    required_free_util_fraction = 0.2  # Require 20% free utilization
     num_gpus_to_select = 1
 
     gpu_info = GPUInfo()
     gpu_info.print_status()
 
-    selected = gpu_info.select_idle_gpu(count=num_gpus_to_select, min_memory_mb=required_free_mem)
-    if selected:
+    if selected := gpu_info.select_idle_gpu(
+        count=num_gpus_to_select,
+        min_memory_fraction=required_free_mem_fraction,
+        min_util_fraction=required_free_util_fraction,
+    ):
         print(f"\n==> Using selected GPU indices: {selected}")
         devices = [f"cuda:{idx}" for idx in selected]
         print(f"    Target devices: {devices}")