File size: 2,943 Bytes
e0265b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | from __future__ import annotations
import shutil
from pathlib import Path
from adam.models import SystemSnapshot
try:
import psutil
except ImportError: # pragma: no cover - optional runtime dependency
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
|