A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 10:00 UTC
[Bugfix][Metrics] Fix RayPrometheusMetric.labels() returning shared labeled child (#40840)
When vLLM runs with Ray Prometheus `vllm:request_success{finished_reason=...}` only ever increments the repetition bucket regardless of the request's actual finish reason; stop, length, abort, and error stay at zero. Root cause was `labels()` mutated the wrapped Ray metric's default tags in place and returned self, so every `.labels(...)` call on a given wrapper returned the same object. Co-authored-by: Marwan Sarieddine <sarieddine.marwan@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Marwan Sarieddine <sarieddine.marwan@gmail.com> Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Seiji Eicher committed
May 1, 2026 at 00:43 UTC
a3ec4a35f5943c250974d504706d22297d423468
2 files changed
+175
-12
tests/v1/metrics/test_ray_metrics.py
+154
-1
@@ -1,13 +1,21 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from unittest.mock import MagicMock
5
+
6
import pytest
7
import ray
8
9
from vllm.config.model import ModelDType
10
from vllm.sampling_params import SamplingParams
11
from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM
10
-from vllm.v1.metrics.ray_wrappers import RayPrometheusMetric, RayPrometheusStatLogger
12
+from vllm.v1.metrics.ray_wrappers import (
13
+ RayCounterWrapper,
14
+ RayGaugeWrapper,
15
+ RayHistogramWrapper,
16
+ RayPrometheusMetric,
17
+ RayPrometheusStatLogger,
18
+)
19
20
MODELS = [
21
"distilbert/distilgpt2",
@@ -94,3 +102,148 @@ def test_sanitized_opentelemetry_name():
102
103
# Test empty string
104
assert RayPrometheusMetric._get_sanitized_opentelemetry_name("") == ""
105
+
106
+
107
+def _install_mock_metric(wrapper: RayPrometheusMetric) -> MagicMock:
108
+ """Swap the wrapper's underlying Ray metric for a MagicMock while
109
+ preserving the real metric's ``_tag_keys`` (labels() reads them to
110
+ validate arity)."""
111
+ real_metric = wrapper.metric
112
+ mock = MagicMock()
113
+ mock._tag_keys = real_metric._tag_keys
114
+ wrapper.metric = mock
115
+ return mock
116
+
117
+
118
+def test_ray_counter_labels_returns_independent_children():
119
+ """RayCounterWrapper.labels() must return distinct labeled children that
120
+ each carry their own tag set."""
121
+ base = RayCounterWrapper(
122
+ name="vllm_test_finish_reason",
123
+ documentation="",
124
+ labelnames=["reason"],
125
+ )
126
+
127
+ stop_child = base.labels("stop")
128
+ rep_child = base.labels("repetition")
129
+
130
+ assert stop_child is not rep_child
131
+ assert stop_child._tags["reason"] == "stop"
132
+ assert rep_child._tags["reason"] == "repetition"
133
+ # Mutating one child's tags must not leak into another.
134
+ stop_child._tags["reason"] = "mutated"
135
+ assert rep_child._tags["reason"] == "repetition"
136
+
137
+
138
+def test_ray_counter_inc_forwards_per_child_tags():
139
+ """.inc() on a labeled counter must forward that child's tags to the
140
+ underlying Ray metric (not rely on a shared set_default_tags)."""
141
+ wrapper = RayCounterWrapper(
142
+ name="vllm_test_counter_tag_forward",
143
+ documentation="",
144
+ labelnames=["reason"],
145
+ )
146
+ mock = _install_mock_metric(wrapper)
147
+
148
+ wrapper.labels("stop").inc()
149
+ wrapper.labels("repetition").inc(3)
150
+ wrapper.labels("stop").inc(0) # zero increment must be a no-op.
151
+
152
+ # The zero-increment call should not reach the underlying metric.
153
+ assert mock.inc.call_count == 2
154
+ first, second = mock.inc.call_args_list
155
+ assert first.args == (1.0,)
156
+ assert first.kwargs["tags"]["reason"] == "stop"
157
+ assert second.args == (3,)
158
+ assert second.kwargs["tags"]["reason"] == "repetition"
159
+
160
+
161
+def test_ray_gauge_labels_returns_independent_children_and_forwards_tags():
162
+ wrapper = RayGaugeWrapper(
163
+ name="vllm_test_gauge_tag_forward",
164
+ documentation="",
165
+ labelnames=["kind"],
166
+ )
167
+ mock = _install_mock_metric(wrapper)
168
+
169
+ a = wrapper.labels("a")
170
+ b = wrapper.labels("b")
171
+ assert a is not b
172
+
173
+ a.set(1)
174
+ b.set(2)
175
+ assert mock.set.call_args_list[0].args == (1,)
176
+ assert mock.set.call_args_list[0].kwargs["tags"]["kind"] == "a"
177
+ assert mock.set.call_args_list[1].args == (2,)
178
+ assert mock.set.call_args_list[1].kwargs["tags"]["kind"] == "b"
179
+
180
+
181
+def test_ray_histogram_labels_returns_independent_children_and_forwards_tags():
182
+ wrapper = RayHistogramWrapper(
183
+ name="vllm_test_histogram_tag_forward",
184
+ documentation="",
185
+ labelnames=["bucket"],
186
+ buckets=[1.0, 2.0, 5.0],
187
+ )
188
+ mock = _install_mock_metric(wrapper)
189
+
190
+ x = wrapper.labels("x")
191
+ y = wrapper.labels("y")
192
+ assert x is not y
193
+
194
+ x.observe(0.5)
195
+ y.observe(4.0)
196
+ assert mock.observe.call_args_list[0].args == (0.5,)
197
+ assert mock.observe.call_args_list[0].kwargs["tags"]["bucket"] == "x"
198
+ assert mock.observe.call_args_list[1].args == (4.0,)
199
+ assert mock.observe.call_args_list[1].kwargs["tags"]["bucket"] == "y"
200
+
201
+
202
+def test_ray_counter_labels_accepts_non_string_label_values():
203
+ """RayPrometheusStatLogger passes ``str(idx)`` for engine indexes; this
204
+ covers the coercion path for any caller that passes a non-string label
205
+ value positionally."""
206
+ wrapper = RayCounterWrapper(
207
+ name="vllm_test_nonstr_label",
208
+ documentation="",
209
+ labelnames=["engine", "reason"],
210
+ )
211
+ child = wrapper.labels(0, "stop")
212
+ assert child._tags["engine"] == "0"
213
+ assert child._tags["reason"] == "stop"
214
+
215
+
216
+def test_ray_counter_labels_arity_validation():
217
+ wrapper = RayCounterWrapper(
218
+ name="vllm_test_arity",
219
+ documentation="",
220
+ labelnames=["a", "b"],
221
+ )
222
+ with pytest.raises(ValueError, match="Number of labels must match"):
223
+ wrapper.labels("only-one")
224
+
225
+
226
+def test_unlabeled_inc_carries_replica_id():
227
+ """Recording on an unlabeled metric must still pass ReplicaId — it's a
228
+ declared tag_key and Ray rejects updates that omit any declared key."""
229
+ wrapper = RayCounterWrapper(
230
+ name="vllm_test_unlabeled_replica_id",
231
+ documentation="",
232
+ labelnames=None,
233
+ )
234
+ mock = _install_mock_metric(wrapper)
235
+ wrapper.inc()
236
+ assert mock.inc.call_args.kwargs["tags"] == {"ReplicaId": ""}
237
+
238
+
239
+def test_double_labels_raises():
240
+ """labels() on an already-labeled child should raise, mirroring the
241
+ prometheus_client contract."""
242
+ wrapper = RayCounterWrapper(
243
+ name="vllm_test_double_labels",
244
+ documentation="",
245
+ labelnames=["reason"],
246
+ )
247
+ child = wrapper.labels("stop")
248
+ with pytest.raises(ValueError, match="already-labeled"):
249
+ child.labels("repetition")
vllm/v1/metrics/ray_wrappers.py
+21
-11
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
import time
5
6
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm
@@ -28,10 +29,13 @@ def _get_replica_id() -> str | None:
29
30
31
class RayPrometheusMetric:
32
+ _is_labeled: bool = False
33
+
34
def __init__(self):
35
if ray_metrics is None:
36
raise ImportError("RayPrometheusMetric requires Ray to be installed.")
37
self.metric: Metric = None
38
+ self._tags: dict[str, str] = {"ReplicaId": _get_replica_id() or ""}
39
40
@staticmethod
41
def _get_tag_keys(labelnames: list[str] | None) -> tuple[str, ...]:
@@ -39,7 +43,7 @@ class RayPrometheusMetric:
43
labels.append("ReplicaId")
44
return tuple(labels)
45
42
- def labels(self, *labels, **labelskwargs):
46
+ def _build_tags(self, *labels, **labelskwargs) -> dict[str, str]:
47
if labels:
48
# -1 because ReplicaId was added automatically
49
expected = len(self.metric._tag_keys) - 1
@@ -52,12 +56,15 @@ class RayPrometheusMetric:
56
57
labelskwargs["ReplicaId"] = _get_replica_id() or ""
58
55
- if labelskwargs:
56
- for k, v in labelskwargs.items():
57
- if not isinstance(v, str):
58
- labelskwargs[k] = str(v)
59
- self.metric.set_default_tags(labelskwargs)
60
- return self
59
+ return {k: v if isinstance(v, str) else str(v) for k, v in labelskwargs.items()}
60
+
61
+ def labels(self, *labels, **labelskwargs) -> "RayPrometheusMetric":
62
+ if self._is_labeled:
63
+ raise ValueError("labels() cannot be called on an already-labeled metric.")
64
+ clone = copy.copy(self)
65
+ clone._tags = self._build_tags(*labels, **labelskwargs)
66
+ clone._is_labeled = True
67
+ return clone
68
69
@staticmethod
70
def _get_sanitized_opentelemetry_name(name: str) -> str:
@@ -91,6 +98,7 @@ class RayGaugeWrapper(RayPrometheusMetric):
98
# implemented at the observability layer (Prometheus/Grafana).
99
del multiprocess_mode
100
101
+ super().__init__()
102
tag_keys = self._get_tag_keys(labelnames)
103
name = self._get_sanitized_opentelemetry_name(name)
104
@@ -101,11 +109,11 @@ class RayGaugeWrapper(RayPrometheusMetric):
109
)
110
111
def set(self, value: int | float):
104
- return self.metric.set(value)
112
+ return self.metric.set(value, tags=self._tags)
113
114
def set_to_current_time(self):
115
# ray metrics doesn't have set_to_current time, https://docs.ray.io/en/latest/_modules/ray/util/metrics.html
108
- return self.metric.set(time.time())
116
+ return self.set(time.time())
117
118
119
class RayCounterWrapper(RayPrometheusMetric):
@@ -118,6 +126,7 @@ class RayCounterWrapper(RayPrometheusMetric):
126
documentation: str | None = "",
127
labelnames: list[str] | None = None,
128
):
129
+ super().__init__()
130
tag_keys = self._get_tag_keys(labelnames)
131
name = self._get_sanitized_opentelemetry_name(name)
132
self.metric = ray_metrics.Counter(
@@ -129,7 +138,7 @@ class RayCounterWrapper(RayPrometheusMetric):
138
def inc(self, value: int | float = 1.0):
139
if value == 0:
140
return
132
- return self.metric.inc(value)
141
+ return self.metric.inc(value, tags=self._tags)
142
143
144
class RayHistogramWrapper(RayPrometheusMetric):
@@ -143,6 +152,7 @@ class RayHistogramWrapper(RayPrometheusMetric):
152
labelnames: list[str] | None = None,
153
buckets: list[float] | None = None,
154
):
155
+ super().__init__()
156
tag_keys = self._get_tag_keys(labelnames)
157
name = self._get_sanitized_opentelemetry_name(name)
158
@@ -155,7 +165,7 @@ class RayHistogramWrapper(RayPrometheusMetric):
165
)
166
167
def observe(self, value: int | float):
158
- return self.metric.observe(value)
168
+ return self.metric.observe(value, tags=self._tags)
169
170
171
class RaySpecDecodingProm(SpecDecodingProm):