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 15, 09:45 UTC
[Misc] Add Ray Prometheus logger to V1 (#17925)
Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Seiji Eicher committed
May 16, 2025 at 03:02 UTC
541817670cfa6101b135cb12428bd8f875364432
4 files changed
+223
-35
tests/v1/metrics/test_ray_metrics.py
new
+57
@@ -0,0 +1,57 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import pytest
3
+import ray
4
+
5
+from vllm.sampling_params import SamplingParams
6
+from vllm.v1.engine.async_llm import AsyncEngineArgs, AsyncLLM
7
+from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
8
+
9
+
10
+@pytest.fixture(scope="function", autouse=True)
11
+def use_v1_only(monkeypatch):
12
+ """
13
+ The change relies on V1 APIs, so set VLLM_USE_V1=1.
14
+ """
15
+ monkeypatch.setenv('VLLM_USE_V1', '1')
16
+
17
+
18
+MODELS = [
19
+ "distilbert/distilgpt2",
20
+]
21
+
22
+
23
+@pytest.mark.parametrize("model", MODELS)
24
+@pytest.mark.parametrize("dtype", ["half"])
25
+@pytest.mark.parametrize("max_tokens", [16])
26
+def test_engine_log_metrics_ray(
27
+ example_prompts,
28
+ model: str,
29
+ dtype: str,
30
+ max_tokens: int,
31
+) -> None:
32
+ """ Simple smoke test, verifying this can be used without exceptions.
33
+ Need to start a Ray cluster in order to verify outputs."""
34
+
35
+ @ray.remote(num_gpus=1)
36
+ class EngineTestActor:
37
+
38
+ async def run(self):
39
+ engine_args = AsyncEngineArgs(
40
+ model=model,
41
+ dtype=dtype,
42
+ disable_log_stats=False,
43
+ )
44
+
45
+ engine = AsyncLLM.from_engine_args(
46
+ engine_args, stat_loggers=[RayPrometheusStatLogger])
47
+
48
+ for i, prompt in enumerate(example_prompts):
49
+ engine.generate(
50
+ request_id=f"request-id-{i}",
51
+ prompt=prompt,
52
+ sampling_params=SamplingParams(max_tokens=max_tokens),
53
+ )
54
+
55
+ # Create the actor and call the async method
56
+ actor = EngineTestActor.remote() # type: ignore[attr-defined]
57
+ ray.get(actor.run.remote())
vllm/v1/metrics/loggers.py
+29
-25
@@ -138,6 +138,10 @@ class LoggingStatLogger(StatLoggerBase):
138
139
140
class PrometheusStatLogger(StatLoggerBase):
141
+ _gauge_cls = prometheus_client.Gauge
142
+ _counter_cls = prometheus_client.Counter
143
+ _histogram_cls = prometheus_client.Histogram
144
+ _spec_decoding_cls = SpecDecodingProm
145
146
def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
147
self._unregister_vllm_metrics()
@@ -156,18 +160,18 @@ class PrometheusStatLogger(StatLoggerBase):
160
161
max_model_len = vllm_config.model_config.max_model_len
162
159
- self.spec_decoding_prom = SpecDecodingProm(
163
+ self.spec_decoding_prom = self._spec_decoding_cls(
164
vllm_config.speculative_config, labelnames, labelvalues)
165
166
#
167
# Scheduler state
168
#
165
- self.gauge_scheduler_running = prometheus_client.Gauge(
169
+ self.gauge_scheduler_running = self._gauge_cls(
170
name="vllm:num_requests_running",
171
documentation="Number of requests in model execution batches.",
172
labelnames=labelnames).labels(*labelvalues)
173
170
- self.gauge_scheduler_waiting = prometheus_client.Gauge(
174
+ self.gauge_scheduler_waiting = self._gauge_cls(
175
name="vllm:num_requests_waiting",
176
documentation="Number of requests waiting to be processed.",
177
labelnames=labelnames).labels(*labelvalues)
@@ -175,18 +179,18 @@ class PrometheusStatLogger(StatLoggerBase):
179
#
180
# GPU cache
181
#
178
- self.gauge_gpu_cache_usage = prometheus_client.Gauge(
182
+ self.gauge_gpu_cache_usage = self._gauge_cls(
183
name="vllm:gpu_cache_usage_perc",
184
documentation="GPU KV-cache usage. 1 means 100 percent usage.",
185
labelnames=labelnames).labels(*labelvalues)
186
183
- self.counter_gpu_prefix_cache_queries = prometheus_client.Counter(
187
+ self.counter_gpu_prefix_cache_queries = self._counter_cls(
188
name="vllm:gpu_prefix_cache_queries",
189
documentation=
190
"GPU prefix cache queries, in terms of number of queried tokens.",
191
labelnames=labelnames).labels(*labelvalues)
192
189
- self.counter_gpu_prefix_cache_hits = prometheus_client.Counter(
193
+ self.counter_gpu_prefix_cache_hits = self._counter_cls(
194
name="vllm:gpu_prefix_cache_hits",
195
documentation=
196
"GPU prefix cache hits, in terms of number of cached tokens.",
@@ -195,24 +199,24 @@ class PrometheusStatLogger(StatLoggerBase):
199
#
200
# Counters
201
#
198
- self.counter_num_preempted_reqs = prometheus_client.Counter(
202
+ self.counter_num_preempted_reqs = self._counter_cls(
203
name="vllm:num_preemptions_total",
204
documentation="Cumulative number of preemption from the engine.",
205
labelnames=labelnames).labels(*labelvalues)
206
203
- self.counter_prompt_tokens = prometheus_client.Counter(
207
+ self.counter_prompt_tokens = self._counter_cls(
208
name="vllm:prompt_tokens_total",
209
documentation="Number of prefill tokens processed.",
210
labelnames=labelnames).labels(*labelvalues)
211
208
- self.counter_generation_tokens = prometheus_client.Counter(
212
+ self.counter_generation_tokens = self._counter_cls(
213
name="vllm:generation_tokens_total",
214
documentation="Number of generation tokens processed.",
215
labelnames=labelnames).labels(*labelvalues)
216
217
self.counter_request_success: dict[FinishReason,
218
prometheus_client.Counter] = {}
215
- counter_request_success_base = prometheus_client.Counter(
219
+ counter_request_success_base = self._counter_cls(
220
name="vllm:request_success_total",
221
documentation="Count of successfully processed requests.",
222
labelnames=labelnames + ["finished_reason"])
@@ -225,21 +229,21 @@ class PrometheusStatLogger(StatLoggerBase):
229
# Histograms of counts
230
#
231
self.histogram_num_prompt_tokens_request = \
228
- prometheus_client.Histogram(
232
+ self._histogram_cls(
233
name="vllm:request_prompt_tokens",
234
documentation="Number of prefill tokens processed.",
235
buckets=build_1_2_5_buckets(max_model_len),
236
labelnames=labelnames).labels(*labelvalues)
237
238
self.histogram_num_generation_tokens_request = \
235
- prometheus_client.Histogram(
239
+ self._histogram_cls(
240
name="vllm:request_generation_tokens",
241
documentation="Number of generation tokens processed.",
242
buckets=build_1_2_5_buckets(max_model_len),
243
labelnames=labelnames).labels(*labelvalues)
244
245
self.histogram_iteration_tokens = \
242
- prometheus_client.Histogram(
246
+ self._histogram_cls(
247
name="vllm:iteration_tokens_total",
248
documentation="Histogram of number of tokens per engine_step.",
249
buckets=[
@@ -249,7 +253,7 @@ class PrometheusStatLogger(StatLoggerBase):
253
labelnames=labelnames).labels(*labelvalues)
254
255
self.histogram_max_num_generation_tokens_request = \
252
- prometheus_client.Histogram(
256
+ self._histogram_cls(
257
name="vllm:request_max_num_generation_tokens",
258
documentation=
259
"Histogram of maximum number of requested generation tokens.",
@@ -257,14 +261,14 @@ class PrometheusStatLogger(StatLoggerBase):
261
labelnames=labelnames).labels(*labelvalues)
262
263
self.histogram_n_request = \
260
- prometheus_client.Histogram(
264
+ self._histogram_cls(
265
name="vllm:request_params_n",
266
documentation="Histogram of the n request parameter.",
267
buckets=[1, 2, 5, 10, 20],
268
labelnames=labelnames).labels(*labelvalues)
269
270
self.histogram_max_tokens_request = \
267
- prometheus_client.Histogram(
271
+ self._histogram_cls(
272
name="vllm:request_params_max_tokens",
273
documentation="Histogram of the max_tokens request parameter.",
274
buckets=build_1_2_5_buckets(max_model_len),
@@ -274,7 +278,7 @@ class PrometheusStatLogger(StatLoggerBase):
278
# Histogram of timing intervals
279
#
280
self.histogram_time_to_first_token = \
277
- prometheus_client.Histogram(
281
+ self._histogram_cls(
282
name="vllm:time_to_first_token_seconds",
283
documentation="Histogram of time to first token in seconds.",
284
buckets=[
@@ -285,7 +289,7 @@ class PrometheusStatLogger(StatLoggerBase):
289
labelnames=labelnames).labels(*labelvalues)
290
291
self.histogram_time_per_output_token = \
288
- prometheus_client.Histogram(
292
+ self._histogram_cls(
293
name="vllm:time_per_output_token_seconds",
294
documentation="Histogram of time per output token in seconds.",
295
buckets=[
@@ -299,34 +303,34 @@ class PrometheusStatLogger(StatLoggerBase):
303
40.0, 50.0, 60.0, 120.0, 240.0, 480.0, 960.0, 1920.0, 7680.0
304
]
305
self.histogram_e2e_time_request = \
302
- prometheus_client.Histogram(
306
+ self._histogram_cls(
307
name="vllm:e2e_request_latency_seconds",
308
documentation="Histogram of e2e request latency in seconds.",
309
buckets=request_latency_buckets,
310
labelnames=labelnames).labels(*labelvalues)
311
self.histogram_queue_time_request = \
308
- prometheus_client.Histogram(
312
+ self._histogram_cls(
313
name="vllm:request_queue_time_seconds",
314
documentation=
315
"Histogram of time spent in WAITING phase for request.",
316
buckets=request_latency_buckets,
317
labelnames=labelnames).labels(*labelvalues)
318
self.histogram_inference_time_request = \
315
- prometheus_client.Histogram(
319
+ self._histogram_cls(
320
name="vllm:request_inference_time_seconds",
321
documentation=
322
"Histogram of time spent in RUNNING phase for request.",
323
buckets=request_latency_buckets,
324
labelnames=labelnames).labels(*labelvalues)
325
self.histogram_prefill_time_request = \
322
- prometheus_client.Histogram(
326
+ self._histogram_cls(
327
name="vllm:request_prefill_time_seconds",
328
documentation=
329
"Histogram of time spent in PREFILL phase for request.",
330
buckets=request_latency_buckets,
331
labelnames=labelnames).labels(*labelvalues)
332
self.histogram_decode_time_request = \
329
- prometheus_client.Histogram(
333
+ self._histogram_cls(
334
name="vllm:request_decode_time_seconds",
335
documentation=
336
"Histogram of time spent in DECODE phase for request.",
@@ -343,7 +347,7 @@ class PrometheusStatLogger(StatLoggerBase):
347
self.labelname_running_lora_adapters = "running_lora_adapters"
348
self.max_lora = vllm_config.lora_config.max_loras
349
self.gauge_lora_info = \
346
- prometheus_client.Gauge(
350
+ self._gauge_cls(
351
name="vllm:lora_requests_info",
352
documentation="Running stats on lora requests.",
353
labelnames=[
@@ -365,7 +369,7 @@ class PrometheusStatLogger(StatLoggerBase):
369
# Info type metrics are syntactic sugar for a gauge permanently set to 1
370
# Since prometheus multiprocessing mode does not support Info, emulate
371
# info here with a gauge.
368
- info_gauge = prometheus_client.Gauge(
372
+ info_gauge = self._gauge_cls(
373
name=name,
374
documentation=documentation,
375
labelnames=metrics_info.keys()).labels(**metrics_info)
vllm/v1/metrics/ray_wrappers.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import time
3
+from typing import Optional, Union
4
+
5
+from vllm.config import VllmConfig
6
+from vllm.v1.metrics.loggers import PrometheusStatLogger
7
+from vllm.v1.spec_decode.metrics import SpecDecodingProm
8
+
9
+try:
10
+ from ray.util import metrics as ray_metrics
11
+ from ray.util.metrics import Metric
12
+except ImportError:
13
+ ray_metrics = None
14
+
15
+
16
+class RayPrometheusMetric:
17
+
18
+ def __init__(self):
19
+ if ray_metrics is None:
20
+ raise ImportError(
21
+ "RayPrometheusMetric requires Ray to be installed.")
22
+
23
+ self.metric: Metric = None
24
+
25
+ def labels(self, *labels, **labelskwargs):
26
+ if labelskwargs:
27
+ for k, v in labelskwargs.items():
28
+ if not isinstance(v, str):
29
+ labelskwargs[k] = str(v)
30
+
31
+ self.metric.set_default_tags(labelskwargs)
32
+
33
+ return self
34
+
35
+
36
+class RayGaugeWrapper(RayPrometheusMetric):
37
+ """Wraps around ray.util.metrics.Gauge to provide same API as
38
+ prometheus_client.Gauge"""
39
+
40
+ def __init__(self,
41
+ name: str,
42
+ documentation: Optional[str] = "",
43
+ labelnames: Optional[list[str]] = None):
44
+ labelnames_tuple = tuple(labelnames) if labelnames else None
45
+ self.metric = ray_metrics.Gauge(name=name,
46
+ description=documentation,
47
+ tag_keys=labelnames_tuple)
48
+
49
+ def set(self, value: Union[int, float]):
50
+ return self.metric.set(value)
51
+
52
+ def set_to_current_time(self):
53
+ # ray metrics doesn't have set_to_current time, https://docs.ray.io/en/latest/_modules/ray/util/metrics.html
54
+ return self.metric.set(time.time())
55
+
56
+
57
+class RayCounterWrapper(RayPrometheusMetric):
58
+ """Wraps around ray.util.metrics.Counter to provide same API as
59
+ prometheus_client.Counter"""
60
+
61
+ def __init__(self,
62
+ name: str,
63
+ documentation: Optional[str] = "",
64
+ labelnames: Optional[list[str]] = None):
65
+ labelnames_tuple = tuple(labelnames) if labelnames else None
66
+ self.metric = ray_metrics.Counter(name=name,
67
+ description=documentation,
68
+ tag_keys=labelnames_tuple)
69
+
70
+ def inc(self, value: Union[int, float] = 1.0):
71
+ if value == 0:
72
+ return
73
+ return self.metric.inc(value)
74
+
75
+
76
+class RayHistogramWrapper(RayPrometheusMetric):
77
+ """Wraps around ray.util.metrics.Histogram to provide same API as
78
+ prometheus_client.Histogram"""
79
+
80
+ def __init__(self,
81
+ name: str,
82
+ documentation: Optional[str] = "",
83
+ labelnames: Optional[list[str]] = None,
84
+ buckets: Optional[list[float]] = None):
85
+ labelnames_tuple = tuple(labelnames) if labelnames else None
86
+ boundaries = buckets if buckets else []
87
+ self.metric = ray_metrics.Histogram(name=name,
88
+ description=documentation,
89
+ tag_keys=labelnames_tuple,
90
+ boundaries=boundaries)
91
+
92
+ def observe(self, value: Union[int, float]):
93
+ return self.metric.observe(value)
94
+
95
+
96
+class RaySpecDecodingProm(SpecDecodingProm):
97
+ """
98
+ RaySpecDecodingProm is used by RayMetrics to log to Ray metrics.
99
+ Provides the same metrics as SpecDecodingProm but uses Ray's
100
+ util.metrics library.
101
+ """
102
+
103
+ _counter_cls = RayCounterWrapper
104
+
105
+
106
+class RayPrometheusStatLogger(PrometheusStatLogger):
107
+ """RayPrometheusStatLogger uses Ray metrics instead."""
108
+
109
+ _gauge_cls = RayGaugeWrapper
110
+ _counter_cls = RayCounterWrapper
111
+ _histogram_cls = RayHistogramWrapper
112
+ _spec_decoding_cls = RaySpecDecodingProm
113
+
114
+ def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
115
+ super().__init__(vllm_config, engine_index)
116
+
117
+ @staticmethod
118
+ def _unregister_vllm_metrics():
119
+ # No-op on purpose
120
+ pass
vllm/v1/spec_decode/metrics.py
+17
-10
@@ -120,24 +120,30 @@ class SpecDecodingProm:
120
vllm:spec_decode_num_drafts[$interval]
121
"""
122
123
- def __init__(self, speculative_config: Optional[SpeculativeConfig],
124
- labelnames: list[str], labelvalues: list[str]):
123
+ _counter_cls = prometheus_client.Counter
124
+
125
+ def __init__(
126
+ self,
127
+ speculative_config: Optional[SpeculativeConfig],
128
+ labelnames: list[str],
129
+ labelvalues: list[str],
130
+ ):
131
self.spec_decoding_enabled = speculative_config is not None
132
if not self.spec_decoding_enabled:
133
return
134
135
self.counter_spec_decode_num_drafts = \
130
- prometheus_client.Counter(
136
+ self._counter_cls(
137
name="vllm:spec_decode_num_drafts_total",
138
documentation="Number of spec decoding drafts.",
139
labelnames=labelnames).labels(*labelvalues)
140
self.counter_spec_decode_num_draft_tokens = \
135
- prometheus_client.Counter(
141
+ self._counter_cls(
142
name="vllm:spec_decode_num_draft_tokens_total",
143
documentation="Number of draft tokens.",
138
- labelnames=labelnames).labels(*labelvalues)
144
+ labelnames=labelnames,).labels(*labelvalues)
145
self.counter_spec_decode_num_accepted_tokens = \
140
- prometheus_client.Counter(
146
+ self._counter_cls(
147
name="vllm:spec_decode_num_accepted_tokens_total",
148
documentation="Number of accepted tokens.",
149
labelnames=labelnames).labels(*labelvalues)
@@ -146,12 +152,13 @@ class SpecDecodingProm:
152
num_spec_tokens = (speculative_config.num_speculative_tokens
153
if self.spec_decoding_enabled else 0)
154
pos_labelnames = labelnames + ["position"]
149
- base_counter = prometheus_client.Counter(
155
+ base_counter = self._counter_cls(
156
name="vllm:spec_decode_num_accepted_tokens_per_pos",
157
documentation="Accepted tokens per draft position.",
152
- labelnames=pos_labelnames)
153
- self.counter_spec_decode_num_accepted_tokens_per_pos: \
154
- list[prometheus_client.Counter] = []
158
+ labelnames=pos_labelnames,
159
+ )
160
+ self.counter_spec_decode_num_accepted_tokens_per_pos: list[
161
+ prometheus_client.Counter] = []
162
for pos in range(num_spec_tokens):
163
pos_labelvalues = labelvalues + [str(pos)]
164
self.counter_spec_decode_num_accepted_tokens_per_pos.append(