diff --git a/nvidia_smi.py b/nvidia_smi.py index 765ad22..8dd7de2 100644 --- a/nvidia_smi.py +++ b/nvidia_smi.py @@ -1,4 +1,5 @@ import subprocess +import psutil import time def get_gpu_stats(): @@ -12,16 +13,27 @@ def get_gpu_stats(): return { "index": int(stats[0]), "name": str(stats[1]), - "memory_used": float(stats[2]), # in MB - "memory_free": float(stats[3]), # in MB + "memory_used": int(stats[2]), # in MB + "memory_free": int(stats[3]), # in MB "power_draw": float(stats[4]), # in watts "temperature": float(stats[5]) # in Celsius } +def get_cpu_memory_stats(): + cpu_usage = psutil.cpu_percent(interval=1) + memory_info = psutil.virtual_memory() + return { + "cpu_usage": cpu_usage, + "memory_used": memory_info.used // (1024 ** 2), # in MB + "memory_total": memory_info.total // (1024 ** 2) # in MB + } + def main(): + import json while True: - stats = get_gpu_stats() - print(f"Power: {stats['power_draw']} W, Memory: {stats['memory_used']} MB, Temp: {stats['temperature']}°C") + gpu_stats = get_gpu_stats() + cpu_stats = get_cpu_memory_stats() + print(json.dumps([gpu_stats, cpu_stats])) time.sleep(1) # Update every second if __name__ == "__main__":