| from __future__ import annotations |
|
|
| import shutil |
| from pathlib import Path |
|
|
| from adam.models import SystemSnapshot |
|
|
| try: |
| import psutil |
| except ImportError: |
| psutil = None |
|
|
|
|
| class SystemMonitor: |
| def __init__(self, root: Path) -> None: |
| self.root = root |
| self._nvml = None |
| self._gpu_handle = None |
| try: |
| import pynvml |
|
|
| pynvml.nvmlInit() |
| self._nvml = pynvml |
| if pynvml.nvmlDeviceGetCount() > 0: |
| self._gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(0) |
| except Exception: |
| self._nvml = None |
| self._gpu_handle = None |
|
|
| def snapshot(self) -> SystemSnapshot: |
| snapshot = SystemSnapshot() |
| if psutil is not None: |
| memory = psutil.virtual_memory() |
| disk = psutil.disk_usage(str(self.root.anchor or self.root)) |
| snapshot.cpu_percent = float(psutil.cpu_percent(interval=None)) |
| snapshot.memory_percent = float(memory.percent) |
| snapshot.memory_used_gb = memory.used / (1024**3) |
| snapshot.memory_total_gb = memory.total / (1024**3) |
| snapshot.disk_percent = float(disk.percent) |
| snapshot.disk_used_gb = disk.used / (1024**3) |
| snapshot.disk_total_gb = disk.total / (1024**3) |
| else: |
| total, used, _ = shutil.disk_usage(self.root) |
| snapshot.disk_percent = used / total * 100 if total else 0.0 |
| snapshot.disk_used_gb = used / (1024**3) |
| snapshot.disk_total_gb = total / (1024**3) |
|
|
| if self._nvml is not None and self._gpu_handle is not None: |
| try: |
| name = self._nvml.nvmlDeviceGetName(self._gpu_handle) |
| if isinstance(name, bytes): |
| name = name.decode("utf-8", errors="replace") |
| utilization = self._nvml.nvmlDeviceGetUtilizationRates(self._gpu_handle) |
| memory = self._nvml.nvmlDeviceGetMemoryInfo(self._gpu_handle) |
| snapshot.gpu_name = str(name) |
| snapshot.gpu_percent = float(utilization.gpu) |
| snapshot.vram_used_gb = memory.used / (1024**3) |
| snapshot.vram_total_gb = memory.total / (1024**3) |
| snapshot.vram_percent = ( |
| memory.used / memory.total * 100 if memory.total else 0.0 |
| ) |
| snapshot.gpu_temperature = float( |
| self._nvml.nvmlDeviceGetTemperature( |
| self._gpu_handle, |
| self._nvml.NVML_TEMPERATURE_GPU, |
| ) |
| ) |
| except Exception: |
| pass |
| return snapshot |
|
|
| def close(self) -> None: |
| if self._nvml is not None: |
| try: |
| self._nvml.nvmlShutdown() |
| except Exception: |
| pass |
| self._nvml = None |
|
|
