Coverage for orchestrant/monitoring/plotting.py: 97%

1"""Visualization module for plotting system metrics.""" 

2 

3from __future__ import annotations 

4 

5import datetime as dt 

6from pathlib import Path 

7from typing import TYPE_CHECKING 

8 

9from loguru import logger 

10 

11 

12if TYPE_CHECKING: 

13 from collections.abc import Sequence 

14 from types import ModuleType 

15 

16 import matplotlib.pyplot as plt 

17 

18 from orchestrant.monitoring.system import SystemMetrics 

19 

20 

21def _pyplot() -> ModuleType: 

22 """Import pyplot on first use.""" 

23 # pyplot was 90 % of `import orchestrant`: 0.72 of 0.80 s on x64, and that import took 31-37 s under riscv64 QEMU. 

24 import matplotlib.pyplot as plt # noqa: PLC0415 

25 

26 return plt 

27 

28 

29class MetricsPlotter: 

30 """Create visualizations for system monitoring metrics. 

31 

32 Example: 

33 >>> from orchestrant.monitoring import SystemMonitor 

34 >>> from orchestrant.monitoring import MetricsPlotter 

35 >>> monitor = SystemMonitor() 

36 >>> monitor.start() 

37 >>> # ... collect metrics ... 

38 >>> monitor.stop() 

39 >>> plotter = MetricsPlotter(monitor.get_metrics()) 

40 >>> plotter.plot_all() 

41 >>> plotter.save_figure("metrics.png") 

42 """ 

43 

44 def __init__(self, metrics: Sequence[SystemMetrics]) -> None: 

45 """Initialize the plotter with metrics data. 

46 

47 Args: 

48 metrics: List of SystemMetrics to visualize 

49 """ 

50 if not metrics: 

51 message = "No metrics provided for plotting" 

52 raise ValueError(message) 

53 

54 self.metrics = list(metrics) 

55 self.fig = None 

56 self.axes = None 

57 

58 # Check if GPU metrics are available 

59 self.has_gpu = any(m.gpu_utilization is not None for m in self.metrics) 

60 

61 logger.info("Initialized plotter with {} samples", len(self.metrics)) 

62 if self.has_gpu: 

63 logger.info("GPU metrics available") 

64 

65 def _prepare_time_data(self) -> tuple[list[float], list[dt.datetime]]: 

66 """Extract timestamps and convert to relative time.""" 

67 timestamps = [m.timestamp for m in self.metrics] 

68 start_time = timestamps[0] 

69 relative_times = [t - start_time for t in timestamps] 

70 utc_tz = getattr(dt, "UTC", dt.timezone(dt.timedelta(0))) 

71 datetime_objects = [dt.datetime.fromtimestamp(t, tz=utc_tz) for t in timestamps] 

72 return relative_times, datetime_objects 

73 

74 def prepare_time_data(self) -> tuple[list[float], list[dt.datetime]]: 

75 """Public wrapper for preparing time-series data.""" 

76 return self._prepare_time_data() 

77 

78 def plot_cpu_memory(self, ax: plt.Axes | None = None) -> plt.Axes: 

79 """Plot CPU and memory usage over time. 

80 

81 Args: 

82 ax: Matplotlib axes to plot on. If None, creates new figure. 

83 

84 Returns: 

85 The axes object used for plotting 

86 """ 

87 if ax is None: 

88 _fig, ax = _pyplot().subplots(figsize=(12, 4)) 

89 

90 relative_times, _ = self._prepare_time_data() 

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

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

93 

94 ax.plot( 

95 relative_times, cpu_values, label="CPU Usage", linewidth=2, color="#1f77b4" 

96 ) 

97 ax.plot( 

98 relative_times, 

99 mem_values, 

100 label="Memory Usage", 

101 linewidth=2, 

102 color="#ff7f0e", 

103 ) 

104 

105 ax.set_xlabel("Time (seconds)", fontsize=11) 

106 ax.set_ylabel("Usage (%)", fontsize=11) 

107 ax.set_title("CPU and Memory Usage Over Time", fontsize=13, fontweight="bold") 

108 ax.legend(loc="upper right") 

109 ax.grid(visible=True, alpha=0.3) 

110 ax.set_ylim(0, 100) 

111 

112 logger.debug("CPU and memory plot created") 

113 return ax 

114 

115 def plot_gpu_utilization(self, ax: plt.Axes | None = None) -> plt.Axes | None: 

116 """Plot GPU utilization over time. 

117 

118 Args: 

119 ax: Matplotlib axes to plot on. If None, creates new figure. 

120 

121 Returns: 

122 The axes object used for plotting, or None if no GPU data 

123 """ 

124 if not self.has_gpu: 

125 logger.warning("No GPU metrics available for plotting") 

126 return None 

127 

128 if ax is None: 

129 _fig, ax = _pyplot().subplots(figsize=(12, 4)) 

130 

131 relative_times, _ = self._prepare_time_data() 

132 gpu_values = [ 

133 m.gpu_utilization if m.gpu_utilization is not None else 0 

134 for m in self.metrics 

135 ] 

136 

137 ax.plot( 

138 relative_times, gpu_values, label="GPU Usage", linewidth=2, color="#2ca02c" 

139 ) 

140 ax.fill_between(relative_times, gpu_values, alpha=0.3, color="#2ca02c") 

141 

142 ax.set_xlabel("Time (seconds)", fontsize=11) 

143 ax.set_ylabel("Utilization (%)", fontsize=11) 

144 ax.set_title("GPU Utilization Over Time", fontsize=13, fontweight="bold") 

145 ax.legend(loc="upper right") 

146 ax.grid(visible=True, alpha=0.3) 

147 ax.set_ylim(0, 100) 

148 

149 logger.debug("GPU utilization plot created") 

150 return ax 

151 

152 def plot_gpu_memory(self, ax: plt.Axes | None = None) -> plt.Axes | None: 

153 """Plot GPU memory usage over time. 

154 

155 Args: 

156 ax: Matplotlib axes to plot on. If None, creates new figure. 

157 

158 Returns: 

159 The axes object used for plotting, or None if no GPU data 

160 """ 

161 if not self.has_gpu: 

162 logger.warning("No GPU metrics available for plotting") 

163 return None 

164 

165 if ax is None: 

166 _fig, ax = _pyplot().subplots(figsize=(12, 4)) 

167 

168 relative_times, _ = self._prepare_time_data() 

169 gpu_mem_used = [ 

170 m.gpu_memory_used_mb if m.gpu_memory_used_mb is not None else 0 

171 for m in self.metrics 

172 ] 

173 gpu_mem_total = self.metrics[0].gpu_memory_total_mb or 1 

174 

175 ax.plot( 

176 relative_times, 

177 gpu_mem_used, 

178 label="GPU Memory Used", 

179 linewidth=2, 

180 color="#d62728", 

181 ) 

182 ax.axhline( 

183 y=gpu_mem_total, 

184 color="gray", 

185 linestyle="--", 

186 label=f"Total: {gpu_mem_total:.0f}MB", 

187 alpha=0.7, 

188 ) 

189 

190 ax.set_xlabel("Time (seconds)", fontsize=11) 

191 ax.set_ylabel("Memory (MB)", fontsize=11) 

192 ax.set_title("GPU Memory Usage Over Time", fontsize=13, fontweight="bold") 

193 ax.legend(loc="upper right") 

194 ax.grid(visible=True, alpha=0.3) 

195 

196 logger.debug("GPU memory plot created") 

197 return ax 

198 

199 def plot_gpu_temperature(self, ax: plt.Axes | None = None) -> plt.Axes | None: 

200 """Plot GPU temperature over time. 

201 

202 Args: 

203 ax: Matplotlib axes to plot on. If None, creates new figure. 

204 

205 Returns: 

206 The axes object used for plotting, or None if no GPU data 

207 """ 

208 if not self.has_gpu: 

209 logger.warning("No GPU metrics available for plotting") 

210 return None 

211 

212 if ax is None: 

213 _fig, ax = _pyplot().subplots(figsize=(12, 4)) 

214 

215 relative_times, _ = self._prepare_time_data() 

216 gpu_temps = [ 

217 m.gpu_temperature if m.gpu_temperature is not None else 0 

218 for m in self.metrics 

219 ] 

220 

221 ax.plot( 

222 relative_times, 

223 gpu_temps, 

224 label="GPU Temperature", 

225 linewidth=2, 

226 color="#ff4500", 

227 ) 

228 ax.fill_between(relative_times, gpu_temps, alpha=0.3, color="#ff4500") 

229 

230 ax.set_xlabel("Time (seconds)", fontsize=11) 

231 ax.set_ylabel("Temperature (°C)", fontsize=11) 

232 ax.set_title("GPU Temperature Over Time", fontsize=13, fontweight="bold") 

233 ax.legend(loc="upper right") 

234 ax.grid(visible=True, alpha=0.3) 

235 

236 logger.debug("GPU temperature plot created") 

237 return ax 

238 

239 def plot_all(self, figsize: tuple[int, int] = (14, 10)) -> plt.Figure: 

240 """Create a comprehensive dashboard with all available metrics. 

241 

242 Args: 

243 figsize: Figure size as (width, height) tuple 

244 

245 Returns: 

246 The created figure object 

247 """ 

248 n_plots = 2 if not self.has_gpu else 5 

249 

250 self.fig, self.axes = _pyplot().subplots(n_plots, 1, figsize=figsize) 

251 self.fig.suptitle( 

252 "System Metrics Dashboard", fontsize=16, fontweight="bold", y=0.995 

253 ) 

254 

255 axes_list = self.axes 

256 

257 # Plot CPU and Memory 

258 self.plot_cpu_memory(axes_list[0]) 

259 

260 # Plot memory in MB 

261 relative_times, _ = self._prepare_time_data() 

262 mem_used_mb = [m.memory_used_mb for m in self.metrics] 

263 mem_available_mb = [m.memory_available_mb for m in self.metrics] 

264 

265 axes_list[1].plot( 

266 relative_times, mem_used_mb, label="Used", linewidth=2, color="#ff7f0e" 

267 ) 

268 axes_list[1].plot( 

269 relative_times, 

270 mem_available_mb, 

271 label="Available", 

272 linewidth=2, 

273 color="#2ca02c", 

274 ) 

275 axes_list[1].set_xlabel("Time (seconds)", fontsize=11) 

276 axes_list[1].set_ylabel("Memory (MB)", fontsize=11) 

277 axes_list[1].set_title( 

278 "Memory Usage (Absolute)", fontsize=13, fontweight="bold" 

279 ) 

280 axes_list[1].legend(loc="upper right") 

281 axes_list[1].grid(visible=True, alpha=0.3) 

282 

283 # GPU plots if available 

284 if self.has_gpu: 

285 self.plot_gpu_utilization(axes_list[2]) 

286 self.plot_gpu_memory(axes_list[3]) 

287 self.plot_gpu_temperature(axes_list[4]) 

288 

289 _pyplot().tight_layout() 

290 logger.info("Complete metrics dashboard created") 

291 return self.fig 

292 

293 def save_figure(self, filepath: str, dpi: int = 150) -> None: 

294 """Save the current figure to a file. 

295 

296 Args: 

297 filepath: Path where to save the figure 

298 dpi: Resolution in dots per inch 

299 """ 

300 if self.fig is None: 

301 logger.error( 

302 "No figure to save. Call plot_all() or other plot methods first." 

303 ) 

304 return 

305 

306 output_path = Path(filepath) 

307 output_path.parent.mkdir(parents=True, exist_ok=True) 

308 

309 self.fig.savefig(filepath, dpi=dpi, bbox_inches="tight") 

310 logger.info("Figure saved to: {}", filepath) 

311 

312 def show(self) -> None: 

313 """Display the plot interactively.""" 

314 if self.fig is None: 

315 logger.error( 

316 "No figure to show. Call plot_all() or other plot methods first." 

317 ) 

318 return 

319 

320 _pyplot().show() 

321 logger.debug("Displaying plot interactively") 

322 

323 

324def quick_plot( 

325 metrics: list[SystemMetrics], 

326 output_path: str | None = None, 

327 *, 

328 show: bool = True, 

329) -> None: 

330 """Convenience function to quickly plot metrics. 

331 

332 Args: 

333 metrics: List of SystemMetrics to plot 

334 output_path: Optional path to save the figure 

335 show: Whether to display the plot interactively 

336 """ 

337 plotter = MetricsPlotter(metrics) 

338 plotter.plot_all() 

339 

340 if output_path: 

341 plotter.save_figure(output_path) 

342 

343 if show: 

344 plotter.show()