128 lines
3.5 KiB
Python
128 lines
3.5 KiB
Python
#!/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()
|