polling stats

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__":