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.
This commit is contained in:
+101
@@ -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()
|
||||||
Reference in New Issue
Block a user