jbilcke-hf's picture
jbilcke-hf HF Staff
add gpu tracking
0d34ea8
"""
System monitoring service for Video Model Studio.
Tracks system resources like CPU, memory, and other metrics.
"""
import os
import time
import logging
import platform
import threading
from datetime import datetime, timedelta
from collections import deque
from typing import Dict, List, Optional, Tuple, Any
import psutil
# Force the use of the Agg backend which is thread-safe
import matplotlib
matplotlib.use('Agg') # Must be before importing pyplot
import matplotlib.pyplot as plt
import numpy as np
from vms.ui.monitoring.services.gpu import GPUMonitoringService
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
class MonitoringService:
"""Service for monitoring system resources and performance"""
def __init__(self, history_minutes: int = 10, sample_interval: int = 5):
"""Initialize the monitoring service
Args:
history_minutes: How many minutes of history to keep
sample_interval: How many seconds between samples
"""
self.history_minutes = history_minutes
self.sample_interval = sample_interval
self.max_samples = (history_minutes * 60) // sample_interval
# Initialize data structures for metrics
self.timestamps = deque(maxlen=self.max_samples)
self.cpu_percent = deque(maxlen=self.max_samples)
self.memory_percent = deque(maxlen=self.max_samples)
self.memory_used = deque(maxlen=self.max_samples)
self.memory_available = deque(maxlen=self.max_samples)
# CPU temperature history (might not be available on all systems)
self.cpu_temp = deque(maxlen=self.max_samples)
# Per-core CPU history
self.cpu_cores_percent = {}
# Initialize GPU monitoring service
self.gpu = GPUMonitoringService(history_minutes=history_minutes, sample_interval=sample_interval)
# Track if the monitoring thread is running
self.is_running = False
self.thread = None
# Initialize with current values
self.collect_metrics()
def collect_metrics(self) -> Dict[str, Any]:
"""Collect current system metrics
Returns:
Dictionary of current metrics
"""
metrics = {
'timestamp': datetime.now(),
'cpu_percent': psutil.cpu_percent(interval=0.1),
'memory_percent': psutil.virtual_memory().percent,
'memory_used': psutil.virtual_memory().used / (1024**3), # GB
'memory_available': psutil.virtual_memory().available / (1024**3), # GB
'cpu_temp': None,
'per_cpu_percent': psutil.cpu_percent(interval=0.1, percpu=True)
}
# Try to get CPU temperature (platform specific)
try:
if platform.system() == 'Linux':
# Try to get temperature from psutil
temps = psutil.sensors_temperatures()
for name, entries in temps.items():
if name.startswith(('coretemp', 'k10temp', 'cpu_thermal')):
metrics['cpu_temp'] = entries[0].current
break
elif platform.system() == 'Darwin': # macOS
# On macOS, we could use SMC reader but it requires additional dependencies
# Leaving as None for now
pass
elif platform.system() == 'Windows':
# Windows might require WMI, leaving as None for simplicity
pass
except (AttributeError, KeyError, IndexError, NotImplementedError):
# Sensors not available
pass
return metrics
def update_history(self, metrics: Dict[str, Any]) -> None:
"""Update metric history with new values
Args:
metrics: New metrics to add to history
"""
self.timestamps.append(metrics['timestamp'])
self.cpu_percent.append(metrics['cpu_percent'])
self.memory_percent.append(metrics['memory_percent'])
self.memory_used.append(metrics['memory_used'])
self.memory_available.append(metrics['memory_available'])
if metrics['cpu_temp'] is not None:
self.cpu_temp.append(metrics['cpu_temp'])
# Update per-core CPU metrics
for i, percent in enumerate(metrics['per_cpu_percent']):
if i not in self.cpu_cores_percent:
self.cpu_cores_percent[i] = deque(maxlen=self.max_samples)
self.cpu_cores_percent[i].append(percent)
def start_monitoring(self) -> None:
"""Start background thread for collecting metrics"""
if self.is_running:
logger.warning("Monitoring thread already running")
return
self.is_running = True
# Start GPU monitoring if available
self.gpu.start_monitoring()
def _monitor_loop():
while self.is_running:
try:
metrics = self.collect_metrics()
self.update_history(metrics)
time.sleep(self.sample_interval)
except Exception as e:
logger.error(f"Error in monitoring thread: {str(e)}", exc_info=True)
time.sleep(self.sample_interval)
self.thread = threading.Thread(target=_monitor_loop, daemon=True)
self.thread.start()
logger.info("System monitoring thread started")
def stop_monitoring(self) -> None:
"""Stop the monitoring thread"""
if not self.is_running:
return
self.is_running = False
# Stop GPU monitoring
self.gpu.stop_monitoring()
if self.thread:
self.thread.join(timeout=1.0)
logger.info("System monitoring thread stopped")
def get_current_metrics(self) -> Dict[str, Any]:
"""Get current system metrics
Returns:
Dictionary with current system metrics
"""
return self.collect_metrics()
def get_system_info(self) -> Dict[str, Any]:
"""Get general system information
Returns:
Dictionary with system details
"""
cpu_info = {
'cores_physical': psutil.cpu_count(logical=False),
'cores_logical': psutil.cpu_count(logical=True),
'current_frequency': None,
'architecture': platform.machine(),
}
# Try to get CPU frequency
try:
cpu_freq = psutil.cpu_freq()
if cpu_freq:
cpu_info['current_frequency'] = cpu_freq.current
except Exception:
pass
memory_info = {
'total': psutil.virtual_memory().total / (1024**3), # GB
'available': psutil.virtual_memory().available / (1024**3), # GB
'used': psutil.virtual_memory().used / (1024**3), # GB
'percent': psutil.virtual_memory().percent
}
disk_info = {}
for part in psutil.disk_partitions(all=False):
if os.name == 'nt' and ('cdrom' in part.opts or part.fstype == ''):
# Skip CD-ROM drives on Windows
continue
try:
usage = psutil.disk_usage(part.mountpoint)
disk_info[part.mountpoint] = {
'total': usage.total / (1024**3), # GB
'used': usage.used / (1024**3), # GB
'free': usage.free / (1024**3), # GB
'percent': usage.percent
}
except PermissionError:
continue
sys_info = {
'system': platform.system(),
'version': platform.version(),
'platform': platform.platform(),
'processor': platform.processor(),
'hostname': platform.node(),
'python_version': platform.python_version(),
'uptime': time.time() - psutil.boot_time()
}
return {
'cpu': cpu_info,
'memory': memory_info,
'disk': disk_info,
'system': sys_info,
}
def generate_cpu_plot(self) -> plt.Figure:
"""Generate a plot of CPU usage over time
Returns:
Matplotlib figure with CPU usage plot
"""
plt.close('all') # Close all existing figures
fig, ax = plt.subplots(figsize=(10, 5))
if not self.timestamps:
ax.set_title("No CPU data available yet")
return fig
x = [t.strftime('%H:%M:%S') for t in self.timestamps]
if len(x) > 10:
# Show fewer x-axis labels for readability
step = len(x) // 10
ax.set_xticks(range(0, len(x), step))
ax.set_xticklabels([x[i] for i in range(0, len(x), step)])
ax.plot(x, list(self.cpu_percent), 'b-', label='CPU Usage %')
if self.cpu_temp and len(self.cpu_temp) > 0:
# Plot temperature on a secondary y-axis if available
ax2 = ax.twinx()
ax2.plot(x[:len(self.cpu_temp)], list(self.cpu_temp), 'r-', label='CPU Temp °C')
ax2.set_ylabel('Temperature (°C)', color='r')
ax2.tick_params(axis='y', colors='r')
ax.set_title('CPU Usage Over Time')
ax.set_xlabel('Time')
ax.set_ylabel('Usage %')
ax.grid(True, alpha=0.3)
ax.set_ylim(0, 100)
# Add legend
lines, labels = ax.get_legend_handles_labels()
if hasattr(locals(), 'ax2'):
lines2, labels2 = ax2.get_legend_handles_labels()
ax.legend(lines + lines2, labels + labels2, loc='upper left')
else:
ax.legend(loc='upper left')
plt.tight_layout()
return fig
def generate_memory_plot(self) -> plt.Figure:
"""Generate a plot of memory usage over time
Returns:
Matplotlib figure with memory usage plot
"""
plt.close('all') # Close all existing figures
fig, ax = plt.subplots(figsize=(10, 5))
if not self.timestamps:
ax.set_title("No memory data available yet")
return fig
x = [t.strftime('%H:%M:%S') for t in self.timestamps]
if len(x) > 10:
# Show fewer x-axis labels for readability
step = len(x) // 10
ax.set_xticks(range(0, len(x), step))
ax.set_xticklabels([x[i] for i in range(0, len(x), step)])
ax.plot(x, list(self.memory_percent), 'g-', label='Memory Usage %')
# Add secondary y-axis for absolute memory values
ax2 = ax.twinx()
ax2.plot(x, list(self.memory_used), 'm--', label='Used (GB)')
ax2.plot(x, list(self.memory_available), 'c--', label='Available (GB)')
ax2.set_ylabel('Memory (GB)')
ax.set_title('Memory Usage Over Time')
ax.set_xlabel('Time')
ax.set_ylabel('Usage %')
ax.grid(True, alpha=0.3)
ax.set_ylim(0, 100)
# Add legend
lines, labels = ax.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax.legend(lines + lines2, labels + labels2, loc='upper left')
plt.tight_layout()
return fig
def generate_per_core_plot(self) -> plt.Figure:
"""Generate a plot of per-core CPU usage
Returns:
Matplotlib figure with per-core CPU usage
"""
num_cores = len(self.cpu_cores_percent)
if num_cores == 0:
# No data yet
plt.close('all') # Close all existing figures
fig, ax = plt.subplots(figsize=(10, 5))
ax.set_title("No per-core CPU data available yet")
return fig
# Determine grid layout based on number of cores
if num_cores <= 4:
rows, cols = 2, 2
elif num_cores <= 6:
rows, cols = 2, 3
elif num_cores <= 9:
rows, cols = 3, 3
elif num_cores <= 12:
rows, cols = 3, 4
else:
rows, cols = 4, 4
fig, axes = plt.subplots(rows, cols, figsize=(12, 8), sharex=True, sharey=True)
axes = axes.flatten()
x = [t.strftime('%H:%M:%S') for t in self.timestamps]
if len(x) > 5:
# Show fewer x-axis labels for readability
step = len(x) // 5
else:
step = 1
for i, (core_id, percentages) in enumerate(self.cpu_cores_percent.items()):
if i >= len(axes):
break
ax = axes[i]
ax.plot(x[:len(percentages)], list(percentages), 'b-')
ax.set_title(f'Core {core_id}')
ax.set_ylim(0, 100)
ax.grid(True, alpha=0.3)
# Add x-axis labels sparingly for readability
if i >= len(axes) - cols: # Only for bottom row
ax.set_xticks(range(0, len(x), step))
ax.set_xticklabels([x[i] for i in range(0, len(x), step)], rotation=45)
# Hide unused subplots
for i in range(num_cores, len(axes)):
axes[i].set_visible(False)
plt.tight_layout()
return fig