fix: prevents nvml runtime error (#292)

This commit prevents an nvml runtime error to be thrown when the nvmllibrary is not found. It also simplifies the hardware info logic.

Author: rickstaa

runner/app/main.py

@@ -2,12 +2,10 @@
 import os
 from contextlib import asynccontextmanager
 
-import app
 from app.routes import health, hardware
 from fastapi import FastAPI
 from fastapi.routing import APIRoute
-from app.utils.hardware import get_gpu_info
-from app.utils.nvml_manager import nvml_manager
+from app.utils.hardware import HardwareInfo
 
 logger = logging.getLogger(__name__)
 
@@ -16,7 +14,8 @@
 async def lifespan(app: FastAPI):
     config_logging()
 
-    nvml_manager.initialize()
+    # Create application wide hardware info service.
+    app.hardware_info_service = HardwareInfo()
 
     app.include_router(health.router)
     app.include_router(hardware.router)
@@ -27,13 +26,11 @@ async def lifespan(app: FastAPI):
     app.pipeline = load_pipeline(pipeline, model_id)
     app.include_router(load_route(pipeline))
 
-    print_cuda_devices()
+    app.hardware_info_service.log_gpu_compute_info()
     logger.info(f"Started up with pipeline {app.pipeline}")
 
     yield
 
-    nvml_manager.shutdown()
-
     logger.info("Shutting down")
 
 
@@ -143,13 +140,6 @@ def config_logging():
     )
 
 
-def print_cuda_devices():
-    devices = get_gpu_info()
-    logger.info("Cuda devices available:")
-    for device in devices:
-        logger.info(devices[device])
-
-
 def use_route_names_as_operation_ids(app: FastAPI) -> None:
     for route in app.routes:
         if isinstance(route, APIRoute):

runner/app/routes/hardware.py

@@ -2,12 +2,10 @@
 from typing import Dict
 
 from app.utils.hardware import (
-    GpuComputeInfo,
-    GpuUtilizationInfo,
-    get_gpu_info,
-    get_gpu_stats,
+    GPUComputeInfo,
+    GPUUtilizationInfo
 )
-from fastapi import APIRouter
+from fastapi import APIRouter, Request
 from pydantic import BaseModel
 
 router = APIRouter()
@@ -18,15 +16,15 @@ class HardwareInformation(BaseModel):
 
     pipeline: str
     model_id: str
-    gpu_info: Dict[int, GpuComputeInfo]
+    gpu_info: Dict[int, GPUComputeInfo]
 
 
 class HardwareStats(BaseModel):
     """Response model for real-time GPU statistics."""
 
     pipeline: str
     model_id: str
-    gpu_stats: Dict[int, GpuUtilizationInfo]
+    gpu_stats: Dict[int, GPUUtilizationInfo]
 
 
 @router.get(
@@ -39,11 +37,11 @@ class HardwareStats(BaseModel):
     response_model=HardwareInformation,
     include_in_schema=False,
 )
-async def hardware_info():
+async def hardware_info(request: Request):
     return HardwareInformation(
         pipeline=os.environ["PIPELINE"],
         model_id=os.environ["MODEL_ID"],
-        gpu_info=get_gpu_info(),
+        gpu_info=request.app.hardware_info_service.get_gpu_compute_info(),
     )
 
 
@@ -57,9 +55,9 @@ async def hardware_info():
     response_model=HardwareStats,
     include_in_schema=False,
 )
-async def hardware_stats():
+async def hardware_stats(request: Request):
     return HardwareStats(
         pipeline=os.environ["PIPELINE"],
         model_id=os.environ["MODEL_ID"],
-        gpu_stats=get_gpu_stats(),
+        gpu_stats=request.app.hardware_info_service.get_gpu_utilization_stats(),
     )

runner/app/utils/hardware.py

@@ -4,11 +4,12 @@
 from pydantic import BaseModel
 import logging
 import pynvml
+import atexit
 
 logger = logging.getLogger(__name__)
 
 
-class GpuBaseInfo(BaseModel):
+class GPUBaseInfo(BaseModel):
     """Model for general GPU information."""
 
     id: str
@@ -17,88 +18,134 @@ class GpuBaseInfo(BaseModel):
     memory_free: int
 
 
-class GpuComputeInfo(GpuBaseInfo):
+class GPUComputeInfo(GPUBaseInfo):
     """Model for detailed GPU compute information."""
 
     major: int
     minor: int
 
 
-class GpuUtilizationInfo(GpuBaseInfo):
-    """Model for real-time GPU utilization statistics."""
+class GPUUtilizationInfo(GPUBaseInfo):
+    """Model for GPU utilization statistics."""
 
     utilization_compute: int
     utilization_memory: int
 
 
-class GpuInfo(GpuComputeInfo, GpuUtilizationInfo):
-    """Model for full CUDA device information."""
+class GPUInfo(GPUComputeInfo, GPUUtilizationInfo):
+    """Model for full GPU device information."""
 
     pass
 
 
-def retrieve_cuda_info() -> Dict[int, GpuInfo]:
-    """Retrieve CUDA device information.
-
-    Returns:
-        CUDA device information.
-    """
-    devices = {}
-    for i in range(pynvml.nvmlDeviceGetCount()):
-        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
-        uuid = pynvml.nvmlDeviceGetUUID(handle)
-        name = pynvml.nvmlDeviceGetName(handle)
-        memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
-        major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
-        utilization_rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
-        devices[i] = GpuInfo(
-            id=uuid,
-            name=name,
-            memory_total=memory_info.total,
-            memory_free=memory_info.free,
-            major=major,
-            minor=minor,
-            utilization_compute=utilization_rates.gpu,
-            utilization_memory=utilization_rates.memory,
-        )
-    return devices
-
-
-def get_gpu_info() -> Dict[int, GpuComputeInfo]:
-    """Get detailed GPU compute information.
-
-    Returns:
-        The detailed GPU compute information.
-    """
-    basic_info = retrieve_cuda_info()
-    return {
-        i: GpuComputeInfo(
-            id=info.id,
-            name=info.name,
-            memory_total=info.memory_total,
-            memory_free=info.memory_free,
-            major=info.major,
-            minor=info.minor,
-        )
-        for i, info in basic_info.items()
-    }
-
-
-def get_gpu_stats() -> Dict[int, GpuUtilizationInfo]:
-    """Get real-time GPU utilization statistics.
-
-    Returns:
-        The real-time GPU utilization statistics.
-    """
-    basic_info = retrieve_cuda_info()
-    return {
-        i: GpuUtilizationInfo(
-            id=info.id,
-            name=info.name,
-            memory_total=info.memory_total,
-            memory_free=info.memory_free,
-            utilization_compute=info.utilization_compute,
-            utilization_memory=info.utilization_memory,
-        )
-        for i, info in basic_info.items()
-    }
+class HardwareInfo:
+    """Class used to retrieve hardware information about the host machine."""
+
+    def __init__(self):
+        """Initialize the HardwareInfo class and hardware info retrieval services."""
+        self._initialized = False
+        self._initialize_nvml()
+        atexit.register(self._shutdown_nvml)
+
+    def _initialize_nvml(self) -> None:
+        """Initialize NVML."""
+        if not self._initialized:
+            try:
+                pynvml.nvmlInit()
+                self._initialized = True
+                logger.info("NVML initialized successfully.")
+            except pynvml.NVMLError as e:
+                logger.error(f"Failed to initialize NVML: {e}")
+
+    def _shutdown_nvml(self) -> None:
+        """Shutdown NVML."""
+        if self._initialized:
+            try:
+                pynvml.nvmlShutdown()
+                self._initialized = False
+                logger.info("NVML shutdown successfully.")
+            except pynvml.NVMLError as e:
+                logger.error(f"Failed to shutdown NVML: {e}")
+
+    def get_cuda_info(self) -> Dict[int, GPUInfo]:
+        """Retrieve CUDA device information.
+
+        Returns:
+            A dictionary mapping GPU device IDs to their information.
+        """
+        devices = {}
+        if not self._initialized:
+            logger.warning(
+                "NVML is not initialized. Cannot retrieve CUDA device information."
+            )
+            return devices
+
+        try:
+            for i in range(pynvml.nvmlDeviceGetCount()):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                uuid = pynvml.nvmlDeviceGetUUID(handle)
+                name = pynvml.nvmlDeviceGetName(handle)
+                memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
+                utilization_rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
+                devices[i] = GPUInfo(
+                    id=uuid,
+                    name=name,
+                    memory_total=memory_info.total,
+                    memory_free=memory_info.free,
+                    major=major,
+                    minor=minor,
+                    utilization_compute=utilization_rates.gpu,
+                    utilization_memory=utilization_rates.memory,
+                )
+        except pynvml.NVMLError as e:
+            logger.warning(f"Failed to retrieve CUDA device information: {e}")
+        return devices
+
+    def get_gpu_compute_info(self) -> Dict[int, GPUComputeInfo]:
+        """Get detailed GPU compute information.
+
+        Returns:
+            A dictionary mapping GPU device IDs to their compute information.
+        """
+        basic_info = self.get_cuda_info()
+        return {
+            i: GPUComputeInfo(
+                id=info.id,
+                name=info.name,
+                memory_total=info.memory_total,
+                memory_free=info.memory_free,
+                major=info.major,
+                minor=info.minor,
+            )
+            for i, info in basic_info.items()
+        }
+
+    def log_gpu_compute_info(self):
+        """Log detailed GPU compute information."""
+        devices = self.get_gpu_compute_info()
+        if devices:
+            logger.info("CUDA devices available:")
+            for device_id, device_info in devices.items():
+                logger.info(f"Device {device_id}: {device_info}")
+        else:
+            logger.info("No CUDA devices available.")
+
+    def get_gpu_utilization_stats(self) -> Dict[int, GPUUtilizationInfo]:
+        """Get GPU utilization statistics.
+
+        Returns:
+            A dictionary mapping GPU device IDs to their utilization statistics.
+        """
+        basic_info = self.get_cuda_info()
+        return {
+            i: GPUUtilizationInfo(
+                id=info.id,
+                name=info.name,
+                memory_total=info.memory_total,
+                memory_free=info.memory_free,
+                utilization_compute=info.utilization_compute,
+                utilization_memory=info.utilization_memory,
+            )
+            for i, info in basic_info.items()
+        }