122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
import threading
|
|
import time
|
|
from typing import Optional
|
|
import xml.etree.ElementTree as ET
|
|
import psutil
|
|
|
|
class SystemMetrics:
|
|
"""
|
|
Tracks system metrics including CPU, GPU, memory and disk usage.
|
|
Maintains rolling averages for CPU and GPU between context generations.
|
|
|
|
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 context
|
|
_gpu_samples: List of GPU usage samples since last context
|
|
_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 generate_context(self, context_usage: float) -> ET.Element:
|
|
"""
|
|
Generate XML context element with current system metrics.
|
|
Clears usage samples after generating averages.
|
|
|
|
Args:
|
|
context_usage: Current context size usage (0-1)
|
|
|
|
Returns:
|
|
ET.Element: Context element with system metrics
|
|
"""
|
|
# Create context element
|
|
context = ET.Element("context")
|
|
|
|
# Add timestamp
|
|
context.set("time", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
|
|
|
|
# Add usage metrics
|
|
with self._metrics_lock:
|
|
# CPU average (or 0 if no samples)
|
|
cpu_avg = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
|
|
context.set("cpu", str(cpu_avg))
|
|
self._cpu_samples.clear()
|
|
|
|
# GPU average (or 0 if no samples)
|
|
gpu_avg = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
|
|
context.set("gpu", str(gpu_avg))
|
|
self._gpu_samples.clear()
|
|
|
|
# Memory usage in bytes
|
|
memory = psutil.virtual_memory()
|
|
context.set("memory_used", str(memory.used))
|
|
context.set("memory_total", str(memory.total))
|
|
|
|
# Disk usage in bytes (root partition)
|
|
disk = psutil.disk_usage('/')
|
|
context.set("disk_used", str(disk.used))
|
|
context.set("disk_total", str(disk.total))
|
|
|
|
# Context usage (0-100)
|
|
context.set("context", str(round(context_usage * 100)))
|
|
|
|
# Standard input buffer size
|
|
context.set("stdin", "0") # Set by agent
|
|
|
|
return context
|
|
|
|
def stop(self):
|
|
"""Stop the monitoring thread and clean up."""
|
|
self._stop_event.set()
|
|
self._monitor_thread.join() |