#!/usr/bin/env python3 """Prometheus exporter for an AMD accelerator, built on rocm-smi. Purpose: publish accelerator power and memory for Track X in the shape Prometheus scrapes, because AMD ships no equivalent of NVIDIA's DCGM exporter for a Ryzen AI Max+ desktop. It runs rocm-smi, reads two numbers out of its JSON, and serves them at /metrics. Platform: strix (AMD Ryzen AI Max+ 395 and other ROCm machines). Not for the other tracks: Tracks S and N use the DCGM exporter, Track M uses mac-exporter.py in this directory. Minimum memory: 8 GB, which is the service being watched; this process needs almost none. Assumes: rocm-smi on PATH and a user who may run it. The JSON key names rocm-smi uses have changed between ROCm releases and AMD now describes amd-smi as the successor to rocm-smi, so this script does not hard-code a single spelling: it matches keys by pattern and tells you plainly which ones it found. Run it once with --once first and read the output before you point Prometheus at it. Usage: python3 rocm-exporter.py --once python3 rocm-exporter.py --port 9401 ROCM_SMI=/opt/rocm/bin/rocm-smi python3 rocm-exporter.py --port 9401 """ import argparse import json import os import re import shutil import subprocess import sys from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer ACCELERATOR = "amd" # Key patterns, in the order they are preferred. rocm-smi's JSON keys are human sentences # rather than identifiers, and they differ between releases, so each metric is described by # a list of case-insensitive patterns and the first key that matches wins. PATTERNS = { "power_watts": [r"average.*power.*\(w\)", r"\bpower\b.*\(w\)", r"socket.*power"], "memory_total_bytes": [r"vram total memory", r"total (vram|memory).*\(b\)"], "memory_used_bytes": [r"vram total used memory", r"used (vram|memory).*\(b\)"], } state = {"scrapes": 0, "errors": 0, "last_error": ""} def rocm_smi_binary(): return os.environ.get("ROCM_SMI", "rocm-smi") def read_rocm_smi(timeout=15): """Run rocm-smi once and return its parsed JSON, or None with the reason recorded.""" binary = rocm_smi_binary() if shutil.which(binary) is None: state["last_error"] = f"{binary} is not on PATH" return None cmd = [binary, "--showpower", "--showmemuse", "--json"] try: out = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, check=False) except (OSError, subprocess.TimeoutExpired) as exc: state["last_error"] = f"running {binary} failed: {exc}" return None if out.returncode != 0: state["last_error"] = f"{binary} exited {out.returncode}: {out.stderr.strip()[:200]}" return None try: return json.loads(out.stdout) except json.JSONDecodeError as exc: state["last_error"] = f"{binary} did not print JSON: {exc}" return None def to_number(value): """rocm-smi prints numbers as strings, sometimes with a unit stuck on the end.""" if isinstance(value, (int, float)): return float(value) match = re.search(r"-?\d+(?:\.\d+)?", str(value)) return float(match.group(0)) if match else None def extract(card_fields): """Pull the three numbers out of one card's fields, by pattern rather than by name.""" found = {} for metric, patterns in PATTERNS.items(): for pattern in patterns: for key, value in card_fields.items(): if re.search(pattern, key, re.IGNORECASE): number = to_number(value) if number is not None: found[metric] = (key, number) break if metric in found: break return found def collect(): """Return a list of (card, metric, source key, value) rows, one set per card.""" data = read_rocm_smi() if data is None: state["errors"] += 1 return [] rows = [] for card, fields in data.items(): if not isinstance(fields, dict): continue for metric, (key, value) in extract(fields).items(): rows.append((card, metric, key, value)) if not rows: state["errors"] += 1 state["last_error"] = ("rocm-smi answered but none of its keys matched the patterns " "this script looks for; run with --once to see the keys") return rows def render(rows): """Prometheus text exposition format. No client library, so it is written out here.""" state["scrapes"] += 1 lines = [ "# HELP local_llm_accelerator_power_watts Accelerator power draw reported by rocm-smi.", "# TYPE local_llm_accelerator_power_watts gauge", "# HELP local_llm_accelerator_memory_used_bytes Accelerator memory in use.", "# TYPE local_llm_accelerator_memory_used_bytes gauge", "# HELP local_llm_accelerator_memory_total_bytes Accelerator memory fitted.", "# TYPE local_llm_accelerator_memory_total_bytes gauge", "# HELP local_llm_exporter_scrape_errors_total Failed collections since start.", "# TYPE local_llm_exporter_scrape_errors_total counter", "# HELP local_llm_exporter_scrapes_total Collections served since start.", "# TYPE local_llm_exporter_scrapes_total counter", ] names = { "power_watts": "local_llm_accelerator_power_watts", "memory_used_bytes": "local_llm_accelerator_memory_used_bytes", "memory_total_bytes": "local_llm_accelerator_memory_total_bytes", } for card, metric, _key, value in rows: lines.append(f'{names[metric]}{{accelerator="{ACCELERATOR}",device="{card}"}} {value}') lines.append(f"local_llm_exporter_scrape_errors_total {state['errors']}") lines.append(f"local_llm_exporter_scrapes_total {state['scrapes']}") return "\n".join(lines) + "\n" class Handler(BaseHTTPRequestHandler): """Two paths: /metrics for Prometheus, anything else for a human who guessed.""" def do_GET(self): # noqa: N802 - the name is fixed by BaseHTTPRequestHandler if self.path.split("?")[0] != "/metrics": self.send_response(404) self.send_header("Content-Type", "text/plain; charset=utf-8") self.end_headers() self.wfile.write(b"Nothing here. The metrics are at /metrics.\n") return body = render(collect()).encode("utf-8") self.send_response(200) self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def log_message(self, fmt, *args): """Quiet by default: a scrape every fifteen seconds is not news.""" if os.environ.get("EXPORTER_ACCESS_LOG"): sys.stderr.write("%s - %s\n" % (self.address_string(), fmt % args)) def print_once(): """What the reader should run first: show the raw keys and what was matched.""" data = read_rocm_smi() if data is None: print(f"rocm-smi could not be read: {state['last_error']}", file=sys.stderr) return 1 print("Keys rocm-smi returned, per card:") for card, fields in data.items(): if not isinstance(fields, dict): continue print(f" {card}") for key, value in fields.items(): print(f" {key} = {value}") print() rows = collect() if not rows: print("No key matched the patterns this script looks for. Edit PATTERNS at the top", file=sys.stderr) print("of this file to match the key names your ROCm version prints.", file=sys.stderr) return 1 print("Matched:") for card, metric, key, value in rows: print(f" {card} {metric:<20} from {key!r} = {value}") print() print("What Prometheus would receive:") print(render(rows)) return 0 def main(): parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--port", type=int, default=9401, help="port to serve /metrics on") parser.add_argument("--host", default="127.0.0.1", help="address to bind; keep the default unless you know why not") parser.add_argument("--once", action="store_true", help="print what rocm-smi says and what was matched, then exit") args = parser.parse_args() if args.once: sys.exit(print_once()) server = ThreadingHTTPServer((args.host, args.port), Handler) print(f"rocm-exporter listening on http://{args.host}:{args.port}/metrics") print("Stop it with Ctrl-C. Run with --once first if a panel stays empty.") try: server.serve_forever() except KeyboardInterrupt: print("\nstopped") finally: server.server_close() if __name__ == "__main__": main()