a13labs.infra/ansible/roles/gpu_monitoring/files/llama_exporter.py

#!/usr/bin/env python3
"""LLM inference metrics exporter for Prometheus.

Scrape targets:
    - llama.cpp HTTP API servers (standard /metrics or custom health)
"""

import json
import os
import time
import logging
import urllib.request
import urllib.error
from http.server import HTTPServer, BaseHTTPRequestHandler

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

POLL_INTERVAL = int(os.environ.get("LLAMA_EXPORTER_INTERVAL", "15"))
SCRAPE_TIMEOUT = int(os.environ.get("LLAMA_EXPORTER_TIMEOUT", "5"))

LLAMA_CPP_DEFAULTS = [
    {"name": "llama-cpp-8012", "url": "http://localhost:8012"},
]


def _fetch_json(url, timeout=SCRAPE_TIMEOUT):
    try:
        req = urllib.request.Request(url, headers={"Accept": "application/json"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except Exception as e:
        logger.debug("Failed to fetch %s: %s", url, e)
        return None


def _collect_llama_cpp(target):
    metrics = {}

    def _set_metric(metric_name, labels, value):
        label_key = tuple(sorted(labels.items()))
        metrics[(metric_name, label_key)] = value

    base_url = target["url"].rstrip("/")

    # Check health endpoint
    health_url = f"{base_url}/health"
    health = _fetch_json(health_url)
    if health and isinstance(health, dict):
        status = health.get("status", "unknown")
        _set_metric("llama_server_health", {"server": target["name"], "model": health.get("model", "unknown")}, 1.0 if status == "ok" else 0.0)
    else:
        _set_metric("llama_server_health", {"server": target["name"], "model": "unknown"}, 0.0)

    # Check /metrics endpoint for llama.cpp built-in metrics
    metrics_url = f"{base_url}/metrics"
    try:
        req = urllib.request.Request(metrics_url, headers={"Accept": "text/plain"})
        with urllib.request.urlopen(req, timeout=SCRAPE_TIMEOUT) as resp:
            body = resp.read().decode("utf-8")
            for line in body.splitlines():
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                # Parse simple key value pairs
                if "{" in line:
                    key_part = line.split("{")[0]
                    labels_part = line.split("{")[1].rstrip("}").split("}")[0] if "}" in line else ""
                    value = line.split()[-1] if line.split() else "1"
                    label_str = ""
                    if labels_part:
                        label_str = ",".join(f'{{{k}="{v}"}}' for k, v in (
                            pair.split("=") for pair in labels_part.split(",") if "=" in pair
                        ))
                    _set_metric(key_part.strip(), {"server": target["name"], "raw": labels_part}, float(value) if value.replace(".", "").replace("-", "").isdigit() else value)
                else:
                    parts = line.split()
                    if len(parts) >= 2:
                        try:
                            float(parts[1])
                            _set_metric(parts[0], {"server": target["name"]}, parts[1])
                        except ValueError:
                            pass
    except Exception as e:
        logger.debug("Failed to scrape %s/metrics: %s", metrics_url, e)

    # Check /model endpoint for model info
    model_url = f"{base_url}/models"
    models = _fetch_json(model_url)
    if models and isinstance(models, dict):
        model_list = models.get("data", [models]) if "data" in models else [models]
        for m in model_list:
            if not isinstance(m, dict):
                continue
            model_id = m.get("id", "unknown")
            _set_metric("llama_models_loaded", {"server": target["name"], "model": model_id}, 1.0)
            ctx_max = m.get("context_size", m.get("n_ctx", 0))
            if ctx_max:
                _set_metric("llama_model_context_max", {"server": target["name"], "model": model_id}, int(ctx_max))
            trained = m.get("train_tokens", m.get("n_train", 0))
            if trained:
                _set_metric("llama_model_tokens_trained", {"server": target["name"], "model": model_id}, int(trained))

    return metrics



class MetricsHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path != "/metrics":
            self.send_response(404)
            self.end_headers()
            return

        all_metrics = {}

        # Scrape llama.cpp targets
        targets = []
        env_targets = os.environ.get("LLAMA_TARGETS", "")
        if env_targets:
            try:
                targets = json.loads(env_targets)
            except json.JSONDecodeError:
                logger.warning("Invalid LLAMA_TARGETS JSON, using defaults")
                targets = []

        if not targets:
            targets = LLAMA_CPP_DEFAULTS

        for target in targets:
            if "url" not in target or "name" not in target:
                target["name"] = target.get("name", "default")
                target["url"] = target.get("url", "http://localhost:8012")
            try:
                all_metrics.update(_collect_llama_cpp(target))
            except Exception as e:
                logger.error("Error scraping llama.cpp target %s: %s", target.get("name", "unknown"), e)

        lines = []
        for (metric_type, label_pairs), value in sorted(all_metrics.items()):
            labels = dict(label_pairs)
            if not labels:
                line = f"{metric_type} {value}"
            else:
                label_str = ",".join(f'{k}="{v}"' for k, v in sorted(labels.items()))
                line = f"{metric_type}{{{label_str}}} {value}"
            lines.append(line)

        body = "\n".join(lines) + "\n" if lines else "# no metrics available\n"

        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def log_message(self, format, *args):
        logger.debug("%s - - %s", self.address_string(), format % args)


def main():
    port = int(os.environ.get("LLAMA_EXPORTER_PORT", "9550"))
    host = os.environ.get("LLAMA_EXPORTER_BIND", "0.0.0.0")

    server = HTTPServer((host, port), MetricsHandler)
    logger.info("Starting Llama Exporter on %s:%d", host, port)
    logger.info("Polling every %d seconds, timeout %ds", POLL_INTERVAL, SCRAPE_TIMEOUT)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("Shutting down")
        server.shutdown()


if __name__ == "__main__":
    main()