117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
import threading
|
|
import time
|
|
from typing import Optional, Dict
|
|
import psutil
|
|
|
|
class SystemMetrics:
|
|
"""
|
|
Tracks system metrics including CPU, GPU, memory and disk usage.
|
|
Maintains rolling averages for CPU and GPU.
|
|
|
|
Attributes:
|
|
_stop_event: Threading event to signal monitor thread to stop
|
|
_monitor_thread: Background thread that collects usage data
|
|
_cpu_samples: List of CPU usage samples since last reading
|
|
_gpu_samples: List of GPU usage samples since last reading
|
|
_metrics_lock: Lock for thread-safe access to samples
|
|
"""
|
|
|
|
def __init__(self, sample_interval: float = 0.1):
|
|
"""
|
|
Initialize system metrics monitoring.
|
|
|
|
Args:
|
|
sample_interval: Seconds between usage samples
|
|
"""
|
|
self._stop_event = threading.Event()
|
|
self._cpu_samples = []
|
|
self._gpu_samples = []
|
|
self._metrics_lock = threading.Lock()
|
|
self._sample_interval = sample_interval
|
|
|
|
# Start monitoring thread
|
|
self._monitor_thread = threading.Thread(
|
|
target=self._monitor_loop,
|
|
daemon=True
|
|
)
|
|
self._monitor_thread.start()
|
|
|
|
def _get_gpu_usage(self) -> Optional[float]:
|
|
"""
|
|
Get current GPU usage percentage.
|
|
Returns None if GPU monitoring not available.
|
|
"""
|
|
try:
|
|
import pynvml
|
|
pynvml.nvmlInit()
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
|
info = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
|
pynvml.nvmlShutdown()
|
|
return info.gpu
|
|
except:
|
|
return None
|
|
|
|
def _monitor_loop(self):
|
|
"""
|
|
Background loop that collects system usage samples.
|
|
Runs until stop event is set.
|
|
"""
|
|
while not self._stop_event.is_set():
|
|
with self._metrics_lock:
|
|
# Get CPU usage across all cores
|
|
self._cpu_samples.append(psutil.cpu_percent())
|
|
|
|
# Get GPU usage if available
|
|
gpu_usage = self._get_gpu_usage()
|
|
if gpu_usage is not None:
|
|
self._gpu_samples.append(gpu_usage)
|
|
|
|
time.sleep(self._sample_interval)
|
|
|
|
def get_metrics(self) -> Dict:
|
|
"""
|
|
Get current system metrics.
|
|
Clears usage samples after calculating averages.
|
|
|
|
Returns:
|
|
Dict containing:
|
|
- cpu: Average CPU usage percentage
|
|
- gpu: Average GPU usage percentage
|
|
- memory_used: Used memory in bytes
|
|
- memory_total: Total memory in bytes
|
|
- disk_used: Used disk space in bytes
|
|
- disk_total: Total disk space in bytes
|
|
- timestamp: Current UTC timestamp in ISO format
|
|
"""
|
|
metrics = {}
|
|
|
|
# Add timestamp
|
|
metrics["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
|
|
# Add usage metrics
|
|
with self._metrics_lock:
|
|
# CPU average (or 0 if no samples)
|
|
metrics["cpu"] = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
|
|
self._cpu_samples.clear()
|
|
|
|
# GPU average (or 0 if no samples)
|
|
metrics["gpu"] = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
|
|
self._gpu_samples.clear()
|
|
|
|
# Memory usage in bytes
|
|
memory = psutil.virtual_memory()
|
|
metrics["memory_used"] = memory.used
|
|
metrics["memory_total"] = memory.total
|
|
|
|
# Disk usage in bytes (root partition)
|
|
disk = psutil.disk_usage('/')
|
|
metrics["disk_used"] = disk.used
|
|
metrics["disk_total"] = disk.total
|
|
|
|
return metrics
|
|
|
|
def stop(self):
|
|
"""Stop the monitoring thread and clean up."""
|
|
self._stop_event.set()
|
|
self._monitor_thread.join()
|