Karti Tripathi
103 lines
4.1 KiB
Python
103 lines
4.1 KiB
Python
"""Speculative-decoding counter parsing.
|
|
|
|
The parser reads a live Prometheus payload, so its failure modes are all "the server said
|
|
something slightly different than expected". Each of these is a shape a real /metrics endpoint
|
|
actually produces.
|
|
|
|
python3 -m unittest discover -s tests -v
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
|
|
from kbench.specdec import SpecCounters, metrics_url, parse_metrics # noqa: E402
|
|
|
|
PAYLOAD = """\
|
|
# HELP vllm:spec_decode_num_draft_tokens_total Number of draft tokens.
|
|
# TYPE vllm:spec_decode_num_draft_tokens_total counter
|
|
vllm:spec_decode_num_draft_tokens_total{model_name="qwen"} 1000.0
|
|
# HELP vllm:spec_decode_num_accepted_tokens_total Number accepted.
|
|
# TYPE vllm:spec_decode_num_accepted_tokens_total counter
|
|
vllm:spec_decode_num_accepted_tokens_total{model_name="qwen"} 720.0
|
|
vllm:num_requests_running{model_name="qwen"} 4.0
|
|
"""
|
|
|
|
|
|
class Parsing(unittest.TestCase):
|
|
def test_reads_both_counters(self):
|
|
c = parse_metrics(PAYLOAD)
|
|
self.assertEqual((c.draft, c.accepted), (1000.0, 720.0))
|
|
self.assertEqual(c.acceptance_rate, 0.72)
|
|
|
|
def test_ignores_comments_and_unrelated_series(self):
|
|
c = parse_metrics(PAYLOAD)
|
|
self.assertEqual(c.draft, 1000.0, "an unrelated metric leaked into the total")
|
|
|
|
def test_sums_across_label_sets(self):
|
|
# A server hosting two models exposes one series each. Summing is correct, and taking
|
|
# the first would silently measure only one of them.
|
|
two = PAYLOAD + (
|
|
'vllm:spec_decode_num_draft_tokens_total{model_name="b"} 500.0\n'
|
|
'vllm:spec_decode_num_accepted_tokens_total{model_name="b"} 250.0\n'
|
|
)
|
|
c = parse_metrics(two)
|
|
self.assertEqual((c.draft, c.accepted), (1500.0, 970.0))
|
|
|
|
def test_no_speculative_counters_means_none_not_zero(self):
|
|
# An engine without speculative decoding must report "no data". Zero would render as a
|
|
# 0% acceptance rate, which reads as a broken draft model rather than an absent one.
|
|
self.assertIsNone(parse_metrics("vllm:num_requests_running 1.0\n"))
|
|
|
|
def test_a_missing_half_is_not_usable(self):
|
|
self.assertIsNone(
|
|
parse_metrics('vllm:spec_decode_num_draft_tokens_total{m="x"} 10.0\n')
|
|
)
|
|
|
|
def test_malformed_values_do_not_raise(self):
|
|
payload = (
|
|
'vllm:spec_decode_num_draft_tokens_total{m="x"} not-a-number\n'
|
|
'vllm:spec_decode_num_accepted_tokens_total{m="x"} 5.0\n'
|
|
)
|
|
self.assertIsNone(parse_metrics(payload))
|
|
|
|
def test_unlabelled_series_are_read(self):
|
|
payload = (
|
|
"vllm:spec_decode_num_draft_tokens_total 100.0\n"
|
|
"vllm:spec_decode_num_accepted_tokens_total 50.0\n"
|
|
)
|
|
self.assertEqual(parse_metrics(payload).acceptance_rate, 0.5)
|
|
|
|
|
|
class WindowedRate(unittest.TestCase):
|
|
def test_the_delta_is_what_gets_reported(self):
|
|
# Lifetime counters. Reporting the total would fold warm-up and every earlier
|
|
# concurrency level into this point's number.
|
|
before = SpecCounters(accepted=700.0, draft=1000.0)
|
|
after = SpecCounters(accepted=1600.0, draft=2000.0)
|
|
window = after - before
|
|
self.assertEqual(window.acceptance_rate, 0.9)
|
|
self.assertNotEqual(window.acceptance_rate, after.acceptance_rate)
|
|
|
|
def test_no_drafting_in_the_window_is_none(self):
|
|
same = SpecCounters(accepted=10.0, draft=10.0)
|
|
self.assertIsNone((same - same).acceptance_rate)
|
|
|
|
|
|
class MetricsUrl(unittest.TestCase):
|
|
def test_strips_the_openai_prefix(self):
|
|
# base_url points at the OpenAI-compatible surface; /metrics is at the server root.
|
|
self.assertEqual(metrics_url("http://host:8001/v1"), "http://host:8001/metrics")
|
|
self.assertEqual(metrics_url("http://host:8001/v1/"), "http://host:8001/metrics")
|
|
|
|
def test_leaves_a_bare_root_alone(self):
|
|
self.assertEqual(metrics_url("http://host:8001"), "http://host:8001/metrics")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|