first commit

callbacks/system.py

@@ -0,0 +1,86 @@
+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=0.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)}")
+