87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
import asyncio
|
|
import psutil
|
|
import pynvml
|
|
import time
|
|
from typing import Dict, Any
|
|
|
|
async def get_system_metrics(args: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Retrieves system metrics including CPU usage, RAM usage, GPU usage (for NVIDIA GPUs),
|
|
system uptime, and network traffic of the primary network interface.
|
|
|
|
Args:
|
|
args: Dictionary of arguments (not used in this implementation).
|
|
|
|
Returns:
|
|
Dictionary with CPU usage (%), RAM used (GB), RAM total (GB), GPU usage (%),
|
|
system uptime (seconds), network traffic (Mbit/s), and input arguments.
|
|
|
|
Raises:
|
|
OSError: If retrieving metrics fails.
|
|
"""
|
|
try:
|
|
# Get CPU usage (averaged over 0.1 seconds for responsiveness)
|
|
cpu_usage = await asyncio.to_thread(psutil.cpu_percent, interval=1)
|
|
|
|
# Get RAM usage
|
|
memory = psutil.virtual_memory()
|
|
ram_used = memory.used / (1024 ** 3) # Convert bytes to GB
|
|
ram_total = memory.total / (1024 ** 3) # Convert bytes to GB
|
|
|
|
# Get system uptime
|
|
uptime = time.time() - psutil.boot_time() # Time since system boot in seconds
|
|
|
|
# Get network traffic for primary interface
|
|
network_traffic = 0.0
|
|
net_io = psutil.net_io_counters(pernic=True)
|
|
primary_interface = None
|
|
max_bytes = 0
|
|
|
|
# Find primary network interface (with most bytes sent/received)
|
|
for interface, stats in net_io.items():
|
|
total_bytes = stats.bytes_sent + stats.bytes_recv
|
|
if total_bytes > max_bytes:
|
|
max_bytes = total_bytes
|
|
primary_interface = interface
|
|
|
|
if primary_interface:
|
|
# Measure traffic over 1 second
|
|
initial_stats = net_io[primary_interface]
|
|
await asyncio.sleep(1)
|
|
final_stats = psutil.net_io_counters(pernic=True)[primary_interface]
|
|
|
|
# Calculate traffic in Mbit/s
|
|
bytes_diff = (final_stats.bytes_sent - initial_stats.bytes_sent +
|
|
final_stats.bytes_recv - initial_stats.bytes_recv)
|
|
network_traffic = (bytes_diff * 8) / (1024 ** 2) # Convert bytes to Mbit/s
|
|
|
|
# Try to get GPU usage (NVIDIA only)
|
|
gpu_usage = 0.0
|
|
try:
|
|
pynvml.nvmlInit()
|
|
device_count = pynvml.nvmlDeviceGetCount()
|
|
if device_count > 0:
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # Use first GPU
|
|
utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
|
gpu_usage = utilization.gpu
|
|
pynvml.nvmlShutdown()
|
|
except pynvml.NVMLError:
|
|
gpu_usage = -1 # Indicate GPU metrics unavailable
|
|
|
|
return {
|
|
"message": "System metrics retrieved successfully",
|
|
"metrics": {
|
|
"cpu_usage_percent": round(cpu_usage, 2),
|
|
"ram_used_gb": round(ram_used, 2),
|
|
"ram_total_gb": round(ram_total, 2),
|
|
"gpu_usage_percent": gpu_usage,
|
|
"uptime_seconds": round(uptime, 2),
|
|
"network_traffic_mbps": round(network_traffic, 2)
|
|
},
|
|
"args": args
|
|
}
|
|
|
|
except Exception as e:
|
|
raise OSError(f"Error while retrieving system metrics: {str(e)}")
|
|
|