From 9c3eba55db4f3d0f2b1ba14ccae4476d923daecc Mon Sep 17 00:00:00 2001
From: fegger <flo.egger@icloud.com>
Date: Fri, 4 Sep 2026 14:06:49 +0200
Subject: [PATCH] Add GPU stats server for AMD ROCm

Small dependency-free HTTP server that exposes rocm-smi metrics
(utilization, VRAM, temperature, power) as JSON on port 9101.
Runs as the gpu-stats compose service with /opt/rocm mounted and
/dev/kfd + /dev/dri passed through.
---
 gpu_stats.py | 101 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 101 insertions(+)
 create mode 100644 gpu_stats.py

diff --git a/gpu_stats.py b/gpu_stats.py
new file mode 100644
index 0000000..fbec426
--- /dev/null
+++ b/gpu_stats.py
@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""
+GPU stats HTTP server for AMD ROCm (rocm-smi).
+No external Python dependencies — uses rocm-smi subprocess for data.
+
+Usage (host):
+  python3 /opt/rocm/libexec/rocm_smi/../../.. && python3 gpu_stats.py
+  # or just: python3 gpu_stats.py  (if /opt/rocm/bin is in PATH or ROCM_SMI set)
+
+Usage (Docker — see docker-compose.yml gpu-stats service):
+  Requires /opt/rocm mounted and /dev/kfd + /dev/dri devices passed through.
+
+Dashboard config (config.json):
+  "gpuStats": { "api": "/api/gpu-stats" }
+
+Then uncomment the /api/gpu-stats block in nginx.conf.
+"""
+import json
+import os
+import subprocess
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+PORT     = 9101
+ROCM_SMI = os.environ.get("ROCM_SMI_PATH", "/opt/rocm/bin/rocm-smi")
+
+
+def _run(extra_flags):
+    result = subprocess.run(
+        [ROCM_SMI] + extra_flags + ["--json"],
+        capture_output=True, text=True, timeout=5,
+    )
+    if result.returncode != 0:
+        return {}
+    try:
+        return json.loads(result.stdout).get("card0", {})
+    except Exception:
+        return {}
+
+
+def query_gpu():
+    try:
+        stats = _run(["--showuse", "--showtemp", "--showpower", "--showid"])
+        vram_raw = subprocess.run(
+            [ROCM_SMI, "--showmeminfo", "vram", "--json"],
+            capture_output=True, text=True, timeout=5,
+        )
+        vram = {}
+        if vram_raw.returncode == 0:
+            vram = json.loads(vram_raw.stdout).get("card0", {})
+
+        if not stats:
+            return None
+
+        vram_total_b = int(vram.get("VRAM Total Memory (B)", 0))
+        vram_used_b  = int(vram.get("VRAM Total Used Memory (B)", 0))
+
+        raw_name = stats.get("Card Series", stats.get("Device Name", ""))
+        name     = raw_name if raw_name and raw_name != "N/A" else os.environ.get("GPU_NAME", "AMD GPU")
+        return {
+            "name":        name,
+            "utilization": int(float(stats.get("GPU use (%)", 0))),
+            "memory_used":  vram_used_b  // (1024 * 1024),   # MB
+            "memory_total": vram_total_b // (1024 * 1024),   # MB
+            "temperature": int(float(
+                stats.get("Temperature (Sensor junction) (C)",
+                stats.get("Temperature (Sensor edge) (C)", 0))
+            )),
+            "power_w":     float(stats.get("Average Graphics Package Power (W)", 0)),
+            "power_max_w": float(stats.get("Max Graphics Package Power (W)", 0)),
+        }
+    except Exception:
+        return None
+
+
+class Handler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path not in ("/", "/stats"):
+            self.send_response(404)
+            self.end_headers()
+            return
+        data = query_gpu()
+        if data is None:
+            body = json.dumps({"error": "rocm-smi unavailable"}).encode()
+            self.send_response(503)
+        else:
+            body = json.dumps(data).encode()
+            self.send_response(200)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Access-Control-Allow-Origin", "*")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def log_message(self, format, *args):  # noqa: A002
+        pass  # suppress per-request noise
+
+
+if __name__ == "__main__":
+    server = HTTPServer(("0.0.0.0", PORT), Handler)
+    print(f"GPU stats server (AMD ROCm) listening on :{PORT}")
+    server.serve_forever()