"""Sample a vLLM server's Prometheus /metrics and report peak in-flight sequences.

vLLM exposes `vllm:num_requests_running` (currently executing on GPU) and
`vllm:num_requests_waiting` (queued). Peak running vs VLLM_MAX_NUM_SEQS tells us
whether we are saturating or overflowing the GPU's sequence slots.

Run this in the background during a benchmark, then read its printed summary.

Usage:
    python scripts/poll_vllm_metrics.py \
        --url http://127.0.0.1:8000/metrics --interval 0.25 --duration 120
"""
from __future__ import annotations

import argparse
import re
import time

import requests

_RUNNING_RE = re.compile(r"^vllm:num_requests_running\S*\s+([0-9.]+)", re.MULTILINE)
_WAITING_RE = re.compile(r"^vllm:num_requests_waiting\S*\s+([0-9.]+)", re.MULTILINE)


def _scrape(url: str) -> tuple[float, float]:
    text = requests.get(url, timeout=5).text
    running = max((float(m) for m in _RUNNING_RE.findall(text)), default=0.0)
    waiting = max((float(m) for m in _WAITING_RE.findall(text)), default=0.0)
    return running, waiting


def main() -> None:
    parser = argparse.ArgumentParser(description="Poll vLLM /metrics for peak in-flight sequences.")
    parser.add_argument("--url", default="http://127.0.0.1:8000/metrics")
    parser.add_argument("--interval", type=float, default=0.25)
    parser.add_argument("--duration", type=float, default=120)
    args = parser.parse_args()

    peak_running = 0.0
    peak_waiting = 0.0
    deadline = time.perf_counter() + args.duration
    while time.perf_counter() < deadline:
        try:
            running, waiting = _scrape(args.url)
            peak_running = max(peak_running, running)
            peak_waiting = max(peak_waiting, waiting)
        except Exception:
            pass
        time.sleep(args.interval)
    print(f"peak_running={peak_running} peak_waiting={peak_waiting}")


if __name__ == "__main__":
    main()