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