Files
compute/tools/scrape_vllm_metrics.py
2026-08-31 16:09:08 -07:00

95 lines
2.9 KiB
Python
Executable File

#!/usr/bin/env python3
"""Scrape vLLM's own Prometheus /metrics on each locally-served model and append
one JSON snapshot line per reachable model to vllm_snapshots.jsonl.
vLLM counters (prompt_tokens_total, generation_tokens_total, request_success_total)
reset to zero on every process restart -- this is what turns them into a durable
history. A model that's down is skipped for this tick, not an error: run this
every few minutes from cron and it just accumulates whatever was actually up.
"""
import json
import os
import re
import time
import urllib.request
USAGE_DIR = os.path.expanduser("~/lumbridge/compute/.compute/usage")
SNAPSHOT_FILE = os.path.join(USAGE_DIR, "vllm_snapshots.jsonl")
# model_name -> port, per registry/models.yaml
MODELS = {
"brain": 8001,
"embed": 8012,
"ocr": 8013,
}
# Prometheus exposition line: metric{labels} value
# vLLM metric names contain a colon (vllm:prompt_tokens_total), which \w does not match.
LINE_RE = re.compile(r'^([\w:]+)(\{[^}]*\})?\s+([0-9eE+\-.]+)\s*$')
COUNTERS = (
"vllm:prompt_tokens_total",
"vllm:generation_tokens_total",
"vllm:num_requests_running",
)
# request_success_total is split by finished_reason -- sum all reasons.
SUCCESS_METRIC = "vllm:request_success_total"
def _fetch(port):
url = f"http://127.0.0.1:{port}/metrics"
with urllib.request.urlopen(url, timeout=3) as resp:
return resp.read().decode()
def _parse(text):
values = {name: 0.0 for name in COUNTERS}
success_total = 0.0
for line in text.splitlines():
if not line or line.startswith("#"):
continue
m = LINE_RE.match(line)
if not m:
continue
name, value = m.group(1), m.group(3)
try:
value = float(value)
except ValueError:
continue
if name in values:
values[name] = value
elif name == SUCCESS_METRIC:
success_total += value
values[SUCCESS_METRIC] = success_total
return values
def main():
os.makedirs(USAGE_DIR, exist_ok=True)
ts = time.time()
lines = []
for model, port in MODELS.items():
try:
text = _fetch(port)
except Exception:
continue # model is down -- skip silently, not an error
parsed = _parse(text)
record = {
"ts": ts,
"model": model,
"prompt_tokens_total": parsed["vllm:prompt_tokens_total"],
"generation_tokens_total": parsed["vllm:generation_tokens_total"],
"request_success_total": parsed[SUCCESS_METRIC],
"num_requests_running": parsed["vllm:num_requests_running"],
}
lines.append(json.dumps(record))
if lines:
with open(SNAPSHOT_FILE, "a") as f:
for line in lines:
f.write(line + "\n")
print(f"scraped {len(lines)}/{len(MODELS)} models reachable")
if __name__ == "__main__":
main()