fegger
9c3eba55db
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.
102 lines
3.4 KiB
Python
102 lines
3.4 KiB
Python
#!/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()
|