Coverage for orchestrant/monitoring/gpu.py: 67%

1"""Vendor-neutral GPU probing (NVML for NVIDIA, gpu_amd for AMD) for the system monitors.""" 

2 

3from __future__ import annotations 

4 

5import importlib 

6from contextlib import suppress 

7from typing import TYPE_CHECKING, Self 

8 

9from loguru import logger 

10 

11from orchestrant.monitoring.gpu_amd import AMD_AVAILABLE, AmdGpuProbe 

12from orchestrant.monitoring.gpu_types import GPUSnapshot 

13 

14 

15if TYPE_CHECKING: 

16 from types import TracebackType 

17 

18 

19__all__ = ["AMD_AVAILABLE", "PYNVML_AVAILABLE", "GPUProbe", "GPUSnapshot"] 

20 

21 

22try: 

23 # Optional extra; debug, not warning, below: on AMD hosts the AMD backend takes over. 

24 pynvml = importlib.import_module("pynvml") 

25except ImportError: 

26 # importlib above, not `import`: ty would type pynvml as the module and reject None. 

27 pynvml = None 

28 logger.debug("nvidia-ml-py not installed; NVIDIA GPU monitoring disabled") 

29 

30 

31# Patched by tests; call sites still test `pynvml is not None`, since a bool narrows nothing. 

32PYNVML_AVAILABLE = pynvml is not None 

33 

34 

35class GPUProbe: 

36 """Shared GPU probe resolving one NVML or AMD device; also a context manager.""" 

37 

38 def __init__(self, gpu_index: int = 0, *, vendor: str | None = None) -> None: 

39 """Probe the requested vendor (or whichever answers) at ``gpu_index``. 

40 

41 Args: 

42 gpu_index: AMD devices are ordered largest dedicated VRAM first. 

43 vendor: ``"nvidia"`` or ``"amd"``; ``None`` tries NVIDIA, then AMD. 

44 """ 

45 self.gpu_index = gpu_index 

46 self.vendor = "none" 

47 self.backend = "none" 

48 self._handle: object | None = None 

49 self.gpu_name = "N/A" 

50 self.memory_type = "" 

51 self.memory_bandwidth_mbps = 0 

52 self.available = False 

53 self._amd: AmdGpuProbe | None = None 

54 

55 if vendor in (None, "nvidia") and PYNVML_AVAILABLE and pynvml is not None: 

56 self._init_nvidia() 

57 if not self.available and vendor in (None, "amd") and AMD_AVAILABLE: 

58 self._init_amd() 

59 if not self.available: 

60 logger.debug("No GPU backend answered for device index {}", gpu_index) 

61 

62 def _init_nvidia(self) -> None: 

63 """Initialize the NVML device handle for ``gpu_index``.""" 

64 if pynvml is None: # ty needs the identity test in this frame too 

65 return 

66 try: 

67 pynvml.nvmlInit() 

68 self._handle = pynvml.nvmlDeviceGetHandleByIndex(self.gpu_index) 

69 name = pynvml.nvmlDeviceGetName(self._handle) 

70 if isinstance(name, bytes): 

71 name = name.decode("utf-8") 

72 self.gpu_name = name 

73 self.vendor = "nvidia" 

74 self.backend = "nvml" 

75 self.available = True 

76 logger.success("GPU monitoring initialized: {}", self.gpu_name) 

77 except (pynvml.NVMLError, RuntimeError) as exc: 

78 logger.debug("NVIDIA GPU monitoring unavailable: {}", exc) 

79 self._handle = None 

80 

81 def _init_amd(self) -> None: 

82 """Initialize the platform AMD backend for ``gpu_index``.""" 

83 try: 

84 probe = AmdGpuProbe(self.gpu_index) 

85 except Exception as exc: # pragma: no cover - defensive: driver/DLL quirks 

86 logger.debug("AMD GPU monitoring unavailable: {}", exc) 

87 return 

88 if not probe.available: 

89 return 

90 self._amd = probe 

91 self.vendor = "amd" 

92 self.backend = probe.backend 

93 self.gpu_name = probe.name 

94 self.memory_type = probe.memory_type 

95 self.memory_bandwidth_mbps = probe.memory_bandwidth_mbps 

96 self._handle = probe.handle 

97 self.available = True 

98 

99 def read(self) -> GPUSnapshot | None: 

100 """Read current GPU metrics, or None if unavailable.""" 

101 if self._amd is not None: 

102 return self._amd.read() 

103 if not self.available or self._handle is None or pynvml is None: 

104 return None 

105 

106 try: 

107 util = pynvml.nvmlDeviceGetUtilizationRates(self._handle) 

108 mem = pynvml.nvmlDeviceGetMemoryInfo(self._handle) 

109 temp = float( 

110 pynvml.nvmlDeviceGetTemperature( 

111 self._handle, pynvml.NVML_TEMPERATURE_GPU 

112 ) 

113 ) 

114 

115 power = 0.0 

116 with suppress(Exception): 

117 power = pynvml.nvmlDeviceGetPowerUsage(self._handle) / 1000.0 

118 

119 return GPUSnapshot( 

120 utilization=float(util.gpu), 

121 memory_used_bytes=int(mem.used), 

122 memory_total_bytes=int(mem.total), 

123 temperature_celsius=temp, 

124 power_watts=power, 

125 vendor="nvidia", 

126 ) 

127 except (pynvml.NVMLError, RuntimeError) as exc: 

128 logger.debug("Error reading GPU metrics: {}", exc) 

129 return None 

130 

131 def shutdown(self) -> None: 

132 """Release GPU resources; idempotent, and read() returns None afterwards.""" 

133 if self._amd is not None: 

134 self._amd.shutdown() 

135 self._amd = None 

136 self._handle = None 

137 self.available = False 

138 return 

139 if ( 

140 self._handle is not None 

141 and PYNVML_AVAILABLE 

142 and self.available 

143 and pynvml is not None 

144 ): 

145 with suppress(Exception): 

146 pynvml.nvmlShutdown() 

147 logger.debug("GPU monitoring shutdown complete") 

148 self._handle = None 

149 self.available = False 

150 

151 def __enter__(self) -> Self: 

152 """Enter context manager.""" 

153 return self 

154 

155 def __exit__( 

156 self, 

157 _exc_type: type[BaseException] | None, 

158 _exc_val: BaseException | None, 

159 _exc_tb: TracebackType | None, 

160 ) -> None: 

161 """Exit context manager and ensure cleanup.""" 

162 self.shutdown()