Coverage for orchestrant/monitoring/system.py: 75%

1"""System monitoring module for tracking CPU, GPU, and memory metrics.""" 

2 

3from __future__ import annotations 

4 

5import time 

6from dataclasses import dataclass 

7 

8import psutil 

9from loguru import logger 

10 

11from orchestrant.monitoring.gpu import GPUProbe 

12from orchestrant.monitoring.snapshot import read_raw_snapshot 

13 

14 

15@dataclass 

16class SystemMetrics: 

17 """Container for system metrics at a specific timestamp.""" 

18 

19 timestamp: float 

20 cpu_percent: float 

21 memory_percent: float 

22 memory_used_mb: float 

23 memory_available_mb: float 

24 gpu_utilization: float | None = None 

25 gpu_memory_used_mb: float | None = None 

26 gpu_memory_total_mb: float | None = None 

27 gpu_temperature: float | None = None 

28 

29 

30class SystemMonitor: 

31 """Time-series monitor: collect CPU/RAM/GPU metrics and summarize them. 

32 

33 Keeps a history of :class:`SystemMetrics` across ``record()`` calls. For 

34 a stateless per-frame snapshot in pipeline UIs use 

35 :class:`orchestrant.pipeline.monitoring.system.SystemMonitor`; both 

36 read through :mod:`orchestrant.monitoring.snapshot`. 

37 

38 Example: 

39 >>> monitor = SystemMonitor(interval=1.0) 

40 >>> monitor.start() 

41 >>> # ... do some work ... 

42 >>> monitor.stop() 

43 >>> metrics = monitor.get_metrics() 

44 >>> monitor.print_summary() 

45 """ 

46 

47 def __init__(self, interval: float = 1.0, gpu_index: int = 0) -> None: 

48 """Initialize the system monitor. 

49 

50 Args: 

51 interval: Time interval between measurements in seconds 

52 gpu_index: GPU device index to monitor (default: 0) 

53 """ 

54 self.interval = interval 

55 self.gpu_index = gpu_index 

56 self.metrics: list[SystemMetrics] = [] 

57 self._monitoring = False 

58 self._gpu = GPUProbe(gpu_index) 

59 

60 # Prime psutil's cpu_percent for non-blocking reads 

61 psutil.cpu_percent(interval=None) 

62 

63 def _collect_metrics(self) -> SystemMetrics: 

64 """Collect current system metrics.""" 

65 raw = read_raw_snapshot(self._gpu) 

66 memory = raw.ram 

67 

68 gpu_util = None 

69 gpu_mem_used = None 

70 gpu_mem_total = None 

71 gpu_temp = None 

72 

73 snapshot = raw.gpu 

74 if snapshot is not None: 

75 gpu_util = snapshot.utilization 

76 gpu_mem_used = snapshot.memory_used_bytes / (1024**2) 

77 gpu_mem_total = snapshot.memory_total_bytes / (1024**2) 

78 gpu_temp = snapshot.temperature_celsius 

79 

80 return SystemMetrics( 

81 timestamp=time.time(), 

82 cpu_percent=raw.cpu_percent, 

83 memory_percent=memory.percent, 

84 memory_used_mb=memory.used / (1024**2), 

85 memory_available_mb=memory.available / (1024**2), 

86 gpu_utilization=gpu_util, 

87 gpu_memory_used_mb=gpu_mem_used, 

88 gpu_memory_total_mb=gpu_mem_total, 

89 gpu_temperature=gpu_temp, 

90 ) 

91 

92 def start(self) -> None: 

93 """Start monitoring and collecting metrics.""" 

94 self._monitoring = True 

95 logger.info("System monitoring started (interval: {}s)", self.interval) 

96 self.metrics = [] 

97 

98 def record(self) -> None: 

99 """Record a single metric snapshot.""" 

100 if not self._monitoring: 

101 logger.warning("Monitor not started. Call start() first.") 

102 return 

103 

104 metrics = self._collect_metrics() 

105 self.metrics.append(metrics) 

106 

107 log_msg = f"CPU: {metrics.cpu_percent:.1f}% | RAM: {metrics.memory_percent:.1f}% ({metrics.memory_used_mb:.0f}MB)" 

108 

109 if metrics.gpu_utilization is not None: 

110 log_msg += f" | GPU: {metrics.gpu_utilization:.1f}% | VRAM: {metrics.gpu_memory_used_mb:.0f}/{metrics.gpu_memory_total_mb:.0f}MB | Temp: {metrics.gpu_temperature:.0f}C" 

111 

112 logger.debug(log_msg) 

113 

114 def stop(self) -> None: 

115 """Stop monitoring.""" 

116 self._monitoring = False 

117 logger.info( 

118 "System monitoring stopped. Collected {} samples.", len(self.metrics) 

119 ) 

120 

121 def get_metrics(self) -> list[SystemMetrics]: 

122 """Get all collected metrics.""" 

123 return self.metrics 

124 

125 @property 

126 def is_monitoring(self) -> bool: 

127 """Return True if monitoring is active.""" 

128 return self._monitoring 

129 

130 @property 

131 def gpu_handle(self) -> object | None: 

132 """Return the GPU handle if initialized.""" 

133 return self._gpu._handle # noqa: SLF001 

134 

135 def print_summary(self) -> None: 

136 """Print a summary of collected metrics.""" 

137 if not self.metrics: 

138 logger.warning("No metrics collected.") 

139 return 

140 

141 cpu_values = [m.cpu_percent for m in self.metrics] 

142 mem_values = [m.memory_percent for m in self.metrics] 

143 

144 cpu_stats = self._compute_stats(cpu_values) 

145 mem_stats = self._compute_stats(mem_values) 

146 

147 logger.info("=" * 60) 

148 logger.info("SYSTEM MONITORING SUMMARY") 

149 logger.info("=" * 60) 

150 logger.info("Samples collected: {}", len(self.metrics)) 

151 logger.info( 

152 "Duration: {:.2f}s", 

153 self.metrics[-1].timestamp - self.metrics[0].timestamp, 

154 ) 

155 logger.info("") 

156 logger.info("CPU Usage:") 

157 logger.info(" Average: {:.1f}%", cpu_stats["avg"]) 

158 logger.info(" Min: {:.1f}%", cpu_stats["min"]) 

159 logger.info(" Max: {:.1f}%", cpu_stats["max"]) 

160 logger.info("") 

161 logger.info("Memory Usage:") 

162 logger.info(" Average: {:.1f}%", mem_stats["avg"]) 

163 logger.info(" Min: {:.1f}%", mem_stats["min"]) 

164 logger.info(" Max: {:.1f}%", mem_stats["max"]) 

165 

166 if self.metrics[0].gpu_utilization is not None: 

167 gpu_util_values = [ 

168 m.gpu_utilization for m in self.metrics if m.gpu_utilization is not None 

169 ] 

170 gpu_mem_values = [ 

171 m.gpu_memory_used_mb 

172 for m in self.metrics 

173 if m.gpu_memory_used_mb is not None 

174 ] 

175 gpu_temp_values = [ 

176 m.gpu_temperature for m in self.metrics if m.gpu_temperature is not None 

177 ] 

178 

179 if gpu_util_values: 

180 gpu_util_stats = self._compute_stats(gpu_util_values) 

181 logger.info("") 

182 logger.info("GPU Usage:") 

183 logger.info(" Average: {:.1f}%", gpu_util_stats["avg"]) 

184 logger.info(" Min: {:.1f}%", gpu_util_stats["min"]) 

185 logger.info(" Max: {:.1f}%", gpu_util_stats["max"]) 

186 

187 if gpu_mem_values: 

188 gpu_mem_stats = self._compute_stats(gpu_mem_values) 

189 logger.info("") 

190 logger.info("GPU Memory:") 

191 logger.info(" Average: {:.0f}MB", gpu_mem_stats["avg"]) 

192 logger.info(" Min: {:.0f}MB", gpu_mem_stats["min"]) 

193 logger.info(" Max: {:.0f}MB", gpu_mem_stats["max"]) 

194 

195 if gpu_temp_values: 

196 gpu_temp_stats = self._compute_stats(gpu_temp_values) 

197 logger.info("") 

198 logger.info("GPU Temperature:") 

199 logger.info(" Average: {:.1f}C", gpu_temp_stats["avg"]) 

200 logger.info(" Min: {:.1f}C", gpu_temp_stats["min"]) 

201 logger.info(" Max: {:.1f}C", gpu_temp_stats["max"]) 

202 

203 logger.info("=" * 60) 

204 

205 @staticmethod 

206 def _compute_stats(values: list[float]) -> dict[str, float]: 

207 """Compute min, max, and average in a single pass.""" 

208 if not values: 

209 return {"avg": 0.0, "min": 0.0, "max": 0.0} 

210 total = 0.0 

211 min_val = float("inf") 

212 max_val = float("-inf") 

213 for v in values: 

214 total += v 

215 min_val = min(min_val, v) 

216 max_val = max(max_val, v) 

217 return {"avg": total / len(values), "min": min_val, "max": max_val} 

218 

219 def __del__(self) -> None: 

220 """Cleanup GPU monitoring on deletion.""" 

221 self._gpu.shutdown()