Lumbridge Bench

tests/test_specdec.py

@@ -0,0 +1,102 @@
+"""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()