GPUMonitor

def __init__(self, gpu_id: int = 0, period: int = 1, enabled: bool = True, dri_node: str = ""):
    self.period = max(1, int(period))
    self.enabled = enabled
    self.gpu_id = gpu_id
    self.dri_node = dri_node
    self.stop_event = asyncio.Event()
    self.task: Optional[asyncio.Task] = None
    self.on_stats: Optional[Callable[..., Awaitable[None]]] = None

def start(self) -> None:
    """Starts the sampling task; no-op when disabled."""
    if not self.enabled:
        return
    self.stop_event.clear()
    self.task = asyncio.create_task(self._monitor_loop())
    logger_gpu.info("GPU monitor started")

def _get_gpu_stats(self) -> Optional[Tuple]:
    """Returns `(load, mem_total, mem_used)` for the target GPU; blocking.

    A `dri_node` match already narrows `get_gpus` to the pipeline's card;
    `gpu_id` indexes only the unfiltered list.

    Returns:
        The stats tuple, or None when the GPU cannot be found or queried.
    """
    try:
        gpus = gpu_stats.get_gpus(self.dri_node)
        idx = 0 if (self.dri_node and len(gpus) == 1) else self.gpu_id
        if not gpus or idx >= len(gpus):
            return None
        gpu = gpus[idx]
        return (gpu.load, gpu.memoryTotal, gpu.memoryUsed)
    except Exception as e:
        logger_gpu.warning(f"Error while fetching GPU stats: {e}")
        return None

async def _monitor_loop(self) -> None:
    """Samples until stopped; exits at once when the first probe finds no GPU.

    Nothing is substituted for a missing GPU: CPU load and system memory
    are `SystemMonitor`'s, and the GPU gauge contract (fractional load, MiB
    memory) cannot carry them without unit errors.
    """
    try:
        if await asyncio.to_thread(self._get_gpu_stats) is None:
            logger_gpu.info(
                f"No GPU with ID {self.gpu_id} found; GPU stats disabled "
                "(CPU and system memory are reported by the system monitor)."
            )
            return
        while not self.stop_event.is_set():
            stats = await asyncio.to_thread(self._get_gpu_stats)
            if stats is not None and self.on_stats:
                load, mem_total, mem_used = stats
                await self.on_stats(load, mem_total, mem_used)
            try:
                await asyncio.wait_for(self.stop_event.wait(), timeout=self.period)
            except asyncio.TimeoutError:
                pass
    except asyncio.CancelledError:
        pass
    except Exception as e:
        logger_gpu.error(f"GPU monitor error: {e}", exc_info=True)
    finally:
        logger_gpu.debug("GPU monitor loop exited")

async def stop(self) -> None:
    """Signals the loop to exit and waits for the task to finish."""
    self.stop_event.set()
    if self.task:
        await self.task
    logger_gpu.info("GPU monitor stopped")