feat: monitor.py, update main.py

monitor.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""
+Real-time GPU and RAM monitoring utility
+"""
+
+import psutil
+import time
+import threading
+from rich.console import Console
+from rich.live import Live
+from rich.table import Table
+
+try:
+    import pynvml
+    NVIDIA_AVAILABLE = True
+except ImportError:
+    NVIDIA_AVAILABLE = False
+
+console = Console()
+
+class SystemMonitor:
+    def __init__(self):
+        self.running = False
+        self.gpu_available = False
+        
+        if NVIDIA_AVAILABLE:
+            try:
+                pynvml.nvmlInit()
+                self.gpu_count = pynvml.nvmlDeviceGetCount()
+                self.gpu_available = self.gpu_count > 0
+            except:
+                pass
+    
+    def get_ram_usage(self):
+        memory = psutil.virtual_memory()
+        return {
+            "used_gb": memory.used / 1024**3,
+            "total_gb": memory.total / 1024**3,
+            "percent": memory.percent
+        }
+    
+    def get_gpu_usage(self):
+        if not self.gpu_available:
+            return None
+        
+        gpu_info = []
+        try:
+            for i in range(self.gpu_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                
+                gpu_info.append({
+                    "id": i,
+                    "used_gb": info.used / 1024**3,
+                    "total_gb": info.total / 1024**3,
+                    "percent": (info.used / info.total) * 100
+                })
+        except:
+            pass
+        
+        return gpu_info
+    
+    def create_table(self):
+        table = Table(title="System Resources")
+        table.add_column("Resource", style="cyan")
+        table.add_column("Used", style="yellow")
+        table.add_column("Total", style="green")
+        table.add_column("Usage %", style="red")
+        
+        # RAM
+        ram = self.get_ram_usage()
+        table.add_row(
+            "RAM",
+            f"{ram['used_gb']:.2f} GB",
+            f"{ram['total_gb']:.2f} GB",
+            f"{ram['percent']:.1f}%"
+        )
+        
+        # GPU
+        if self.gpu_available:
+            gpus = self.get_gpu_usage()
+            if gpus:
+                for gpu in gpus:
+                    table.add_row(
+                        f"GPU {gpu['id']}",
+                        f"{gpu['used_gb']:.2f} GB",
+                        f"{gpu['total_gb']:.2f} GB",
+                        f"{gpu['percent']:.1f}%"
+                    )
+        else:
+            table.add_row("GPU", "N/A", "N/A", "N/A")
+        
+        return table
+    
+    def start_monitoring(self):
+        self.running = True
+        
+        with Live(self.create_table(), refresh_per_second=1, console=console) as live:
+            while self.running:
+                live.update(self.create_table())
+                time.sleep(1)
+    
+    def stop_monitoring(self):
+        self.running = False
+        if self.gpu_available:
+            try:
+                pynvml.nvmlShutdown()
+            except:
+                pass
+
+def monitor_resources():
+    """Start monitoring in a separate thread"""
+    monitor = SystemMonitor()
+    thread = threading.Thread(target=monitor.start_monitoring, daemon=True)
+    thread.start()
+    return monitor
+
+if __name__ == "__main__":
+    console.print("[bold cyan]System Resource Monitor[/bold cyan]")
+    console.print("Press Ctrl+C to stop\n")
+    
+    monitor = SystemMonitor()
+    try:
+        monitor.start_monitoring()
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Monitoring stopped.[/yellow]")
+        monitor.stop_monitoring()