Monitor CPU memory usage for the current process.
| 4346 | |
| 4347 | |
| 4348 | class CPUMemoryMonitor: |
| 4349 | """Monitor CPU memory usage for the current process.""" |
| 4350 | |
| 4351 | def __init__(self): |
| 4352 | self.device_name = "CPU" |
| 4353 | self._peak_rss = 0 |
| 4354 | self._process = None |
| 4355 | self.total_memory = 0 |
| 4356 | self.total_memory_gib = 0 |
| 4357 | |
| 4358 | if is_psutil_available(): |
| 4359 | import psutil |
| 4360 | |
| 4361 | self._process = psutil.Process(os.getpid()) |
| 4362 | mem_info = psutil.virtual_memory() |
| 4363 | self.total_memory = mem_info.total |
| 4364 | self.total_memory_gib = self._to_gib(self.total_memory) |
| 4365 | |
| 4366 | def _to_gib(self, memory_in_bytes: int) -> float: |
| 4367 | """Convert bytes to GiB.""" |
| 4368 | return memory_in_bytes / (1024 * 1024 * 1024) |
| 4369 | |
| 4370 | def _to_pct(self, memory_in_bytes: int) -> float: |
| 4371 | """Convert bytes to percentage of total memory.""" |
| 4372 | if self.total_memory == 0: |
| 4373 | return 0.0 |
| 4374 | return 100.0 * memory_in_bytes / self.total_memory |
| 4375 | |
| 4376 | def _update_peak(self) -> None: |
| 4377 | """Update peak memory tracking.""" |
| 4378 | if self._process is not None: |
| 4379 | current_rss = self._process.memory_info().rss |
| 4380 | self._peak_rss = max(self._peak_rss, current_rss) |
| 4381 | |
| 4382 | def get_stats(self) -> MemoryStats: |
| 4383 | """Get current memory statistics.""" |
| 4384 | if not is_psutil_available(): |
| 4385 | return MemoryStats(0, 0, 0, 0, 0, 0, 0) |
| 4386 | |
| 4387 | import psutil |
| 4388 | |
| 4389 | self._update_peak() |
| 4390 | |
| 4391 | mem_info = self._process.memory_info() |
| 4392 | sys_mem = psutil.virtual_memory() |
| 4393 | |
| 4394 | return MemoryStats( |
| 4395 | rss_gib=self._to_gib(mem_info.rss), |
| 4396 | rss_pct=self._to_pct(mem_info.rss), |
| 4397 | vms_gib=self._to_gib(mem_info.vms), |
| 4398 | peak_rss_gib=self._to_gib(self._peak_rss), |
| 4399 | peak_rss_pct=self._to_pct(self._peak_rss), |
| 4400 | available_gib=self._to_gib(sys_mem.available), |
| 4401 | total_gib=self._to_gib(sys_mem.total), |
| 4402 | ) |
| 4403 | |
| 4404 | def reset_peak_stats(self) -> None: |
| 4405 | """Reset peak memory tracking.""" |
no outgoing calls