main
py 127 lines 4.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import prometheus_client
5 import pytest
6
7 from vllm.v1.metrics.reader import (
8 Counter,
9 Gauge,
10 Histogram,
11 Vector,
12 get_metrics_snapshot,
13 )
14
15 pytestmark = pytest.mark.cpu_test
16
17
18 @pytest.fixture(autouse=True)
19 def test_registry(monkeypatch):
20 # Use a custom registry for tests
21 test_registry = prometheus_client.CollectorRegistry(auto_describe=True)
22 monkeypatch.setattr("vllm.v1.metrics.reader.REGISTRY", test_registry)
23 return test_registry
24
25
26 @pytest.mark.parametrize("num_engines", [1, 4])
27 def test_gauge_metric(test_registry, num_engines):
28 g = prometheus_client.Gauge(
29 "vllm:test_gauge",
30 "Test gauge metric",
31 labelnames=["model", "engine_index"],
32 registry=test_registry,
33 )
34 for i in range(num_engines):
35 g.labels(model="foo", engine_index=str(i)).set(98.5)
36
37 metrics = get_metrics_snapshot()
38 assert len(metrics) == num_engines
39 engine_labels = [str(i) for i in range(num_engines)]
40 for m in metrics:
41 assert isinstance(m, Gauge)
42 assert m.name == "vllm:test_gauge"
43 assert m.value == 98.5
44 assert m.labels["model"] == "foo"
45 assert m.labels["engine_index"] in engine_labels
46 engine_labels.remove(m.labels["engine_index"])
47
48
49 @pytest.mark.parametrize("num_engines", [1, 4])
50 def test_counter_metric(test_registry, num_engines):
51 c = prometheus_client.Counter(
52 "vllm:test_counter",
53 "Test counter metric",
54 labelnames=["model", "engine_index"],
55 registry=test_registry,
56 )
57 for i in range(num_engines):
58 c.labels(model="bar", engine_index=str(i)).inc(19)
59
60 metrics = get_metrics_snapshot()
61 assert len(metrics) == num_engines
62 engine_labels = [str(i) for i in range(num_engines)]
63 for m in metrics:
64 assert isinstance(m, Counter)
65 assert m.name == "vllm:test_counter"
66 assert m.value == 19
67 assert m.labels["model"] == "bar"
68 assert m.labels["engine_index"] in engine_labels
69 engine_labels.remove(m.labels["engine_index"])
70
71
72 @pytest.mark.parametrize("num_engines", [1, 4])
73 def test_histogram_metric(test_registry, num_engines):
74 h = prometheus_client.Histogram(
75 "vllm:test_histogram",
76 "Test histogram metric",
77 labelnames=["model", "engine_index"],
78 buckets=[10, 20, 30, 40, 50],
79 registry=test_registry,
80 )
81 for i in range(num_engines):
82 hist = h.labels(model="blaa", engine_index=str(i))
83 hist.observe(42)
84 hist.observe(21)
85 hist.observe(7)
86
87 metrics = get_metrics_snapshot()
88 assert len(metrics) == num_engines
89 engine_labels = [str(i) for i in range(num_engines)]
90 for m in metrics:
91 assert isinstance(m, Histogram)
92 assert m.name == "vllm:test_histogram"
93 assert m.count == 3
94 assert m.sum == 70
95 assert m.buckets["10.0"] == 1
96 assert m.buckets["20.0"] == 1
97 assert m.buckets["30.0"] == 2
98 assert m.buckets["40.0"] == 2
99 assert m.buckets["50.0"] == 3
100 assert m.labels["model"] == "blaa"
101 assert m.labels["engine_index"] in engine_labels
102 engine_labels.remove(m.labels["engine_index"])
103
104
105 @pytest.mark.parametrize("num_engines", [1, 4])
106 def test_vector_metric(test_registry, num_engines):
107 c = prometheus_client.Counter(
108 "vllm:spec_decode_num_accepted_tokens_per_pos",
109 "Vector-like counter metric",
110 labelnames=["position", "model", "engine_index"],
111 registry=test_registry,
112 )
113 for i in range(num_engines):
114 c.labels(position="0", model="llama", engine_index=str(i)).inc(10)
115 c.labels(position="1", model="llama", engine_index=str(i)).inc(5)
116 c.labels(position="2", model="llama", engine_index=str(i)).inc(1)
117
118 metrics = get_metrics_snapshot()
119 assert len(metrics) == num_engines
120 engine_labels = [str(i) for i in range(num_engines)]
121 for m in metrics:
122 assert isinstance(m, Vector)
123 assert m.name == "vllm:spec_decode_num_accepted_tokens_per_pos"
124 assert m.values == [10, 5, 1]
125 assert m.labels["model"] == "llama"
126 assert m.labels["engine_index"] in engine_labels
127 engine_labels.remove(m.labels["engine_index"])