Coverage for orchestr_ant_ion / 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 

9import matplotlib.pyplot as plt 

10from loguru import logger 

11 

12 

13if TYPE_CHECKING: 

14 from collections.abc import Sequence 

15 

16 from orchestr_ant_ion.monitoring.system import SystemMetrics 

17 

18 

19class MetricsPlotter: 

20 """Create visualizations for system monitoring metrics. 

21 

22 Example: 

23 >>> from orchestr_ant_ion.monitoring import SystemMonitor 

24 >>> from orchestr_ant_ion.monitoring import MetricsPlotter 

25 >>> monitor = SystemMonitor() 

26 >>> monitor.start() 

27 >>> # ... collect metrics ... 

28 >>> monitor.stop() 

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

30 >>> plotter.plot_all() 

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

32 """ 

33 

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

35 """Initialize the plotter with metrics data. 

36 

37 Args: 

38 metrics: List of SystemMetrics to visualize 

39 """ 

40 if not metrics: 

41 message = "No metrics provided for plotting" 

42 raise ValueError(message) 

43 

44 self.metrics = list(metrics) 

45 self.fig = None 

46 self.axes = None 

47 

48 # Check if GPU metrics are available 

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

50 

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

52 if self.has_gpu: 

53 logger.info("GPU metrics available") 

54 

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

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

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

58 start_time = timestamps[0] 

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

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

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

62 return relative_times, datetime_objects 

63 

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

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

66 return self._prepare_time_data() 

67 

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

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

70 

71 Args: 

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

73 

74 Returns: 

75 The axes object used for plotting 

76 """ 

77 if ax is None: 

78 _fig, ax = plt.subplots(figsize=(12, 4)) 

79 

80 relative_times, _ = self._prepare_time_data() 

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

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

83 

84 ax.plot( 

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

86 ) 

87 ax.plot( 

88 relative_times, 

89 mem_values, 

90 label="Memory Usage", 

91 linewidth=2, 

92 color="#ff7f0e", 

93 ) 

94 

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

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

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

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

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

100 ax.set_ylim(0, 100) 

101 

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

103 return ax 

104 

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

106 """Plot GPU utilization over time. 

107 

108 Args: 

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

110 

111 Returns: 

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

113 """ 

114 if not self.has_gpu: 

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

116 return None 

117 

118 if ax is None: 

119 _fig, ax = plt.subplots(figsize=(12, 4)) 

120 

121 relative_times, _ = self._prepare_time_data() 

122 gpu_values = [ 

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

124 for m in self.metrics 

125 ] 

126 

127 ax.plot( 

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

129 ) 

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

131 

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

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

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

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

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

137 ax.set_ylim(0, 100) 

138 

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

140 return ax 

141 

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

143 """Plot GPU memory usage over time. 

144 

145 Args: 

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

147 

148 Returns: 

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

150 """ 

151 if not self.has_gpu: 

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

153 return None 

154 

155 if ax is None: 

156 _fig, ax = plt.subplots(figsize=(12, 4)) 

157 

158 relative_times, _ = self._prepare_time_data() 

159 gpu_mem_used = [ 

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

161 for m in self.metrics 

162 ] 

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

164 

165 ax.plot( 

166 relative_times, 

167 gpu_mem_used, 

168 label="GPU Memory Used", 

169 linewidth=2, 

170 color="#d62728", 

171 ) 

172 ax.axhline( 

173 y=gpu_mem_total, 

174 color="gray", 

175 linestyle="--", 

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

177 alpha=0.7, 

178 ) 

179 

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

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

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

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

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

185 

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

187 return ax 

188 

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

190 """Plot GPU temperature over time. 

191 

192 Args: 

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

194 

195 Returns: 

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

197 """ 

198 if not self.has_gpu: 

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

200 return None 

201 

202 if ax is None: 

203 _fig, ax = plt.subplots(figsize=(12, 4)) 

204 

205 relative_times, _ = self._prepare_time_data() 

206 gpu_temps = [ 

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

208 for m in self.metrics 

209 ] 

210 

211 ax.plot( 

212 relative_times, 

213 gpu_temps, 

214 label="GPU Temperature", 

215 linewidth=2, 

216 color="#ff4500", 

217 ) 

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

219 

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

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

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

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

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

225 

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

227 return ax 

228 

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

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

231 

232 Args: 

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

234 

235 Returns: 

236 The created figure object 

237 """ 

238 n_plots = 2 if not self.has_gpu else 5 

239 

240 self.fig, self.axes = plt.subplots(n_plots, 1, figsize=figsize) 

241 self.fig.suptitle( 

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

243 ) 

244 

245 axes_list = self.axes 

246 

247 # Plot CPU and Memory 

248 self.plot_cpu_memory(axes_list[0]) 

249 

250 # Plot memory in MB 

251 relative_times, _ = self._prepare_time_data() 

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

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

254 

255 axes_list[1].plot( 

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

257 ) 

258 axes_list[1].plot( 

259 relative_times, 

260 mem_available_mb, 

261 label="Available", 

262 linewidth=2, 

263 color="#2ca02c", 

264 ) 

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

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

267 axes_list[1].set_title( 

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

269 ) 

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

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

272 

273 # GPU plots if available 

274 if self.has_gpu: 

275 self.plot_gpu_utilization(axes_list[2]) 

276 self.plot_gpu_memory(axes_list[3]) 

277 self.plot_gpu_temperature(axes_list[4]) 

278 

279 plt.tight_layout() 

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

281 return self.fig 

282 

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

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

285 

286 Args: 

287 filepath: Path where to save the figure 

288 dpi: Resolution in dots per inch 

289 """ 

290 if self.fig is None: 

291 logger.error( 

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

293 ) 

294 return 

295 

296 output_path = Path(filepath) 

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

298 

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

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

301 

302 def show(self) -> None: 

303 """Display the plot interactively.""" 

304 if self.fig is None: 

305 logger.error( 

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

307 ) 

308 return 

309 

310 plt.show() 

311 logger.debug("Displaying plot interactively") 

312 

313 

314def quick_plot( 

315 metrics: list[SystemMetrics], 

316 output_path: str | None = None, 

317 *, 

318 show: bool = True, 

319) -> None: 

320 """Convenience function to quickly plot metrics. 

321 

322 Args: 

323 metrics: List of SystemMetrics to plot 

324 output_path: Optional path to save the figure 

325 show: Whether to display the plot interactively 

326 """ 

327 plotter = MetricsPlotter(metrics) 

328 plotter.plot_all() 

329 

330 if output_path: 

331 plotter.save_figure(output_path) 

332 

333 if show: 

334 plotter.show()