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 7, 06:30 UTC
[Feature] Add basic metrics for /realtime endpoint (#35500)
Signed-off-by: Thomas Pouget-Abadie <thomaspou@microsoft.com> Signed-off-by: pougetat <thomas.pougetabadie@gmail.com> Co-authored-by: Thomas Pouget-Abadie <thomaspou@microsoft.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
pougetat committed
Mar 4, 2026 at 03:56 UTC
1659b2e05804ff5c8a5b8dc3feaa4e66945592f4
2 files changed
+86
vllm/entrypoints/openai/api_server.py
+8
@@ -264,6 +264,14 @@ def build_app(
264
# Add scaling middleware to check for scaling state
265
app.add_middleware(ScalingMiddleware)
266
267
+ if "realtime" in supported_tasks:
268
+ # Add WebSocket metrics middleware
269
+ from vllm.entrypoints.openai.realtime.metrics import (
270
+ WebSocketMetricsMiddleware,
271
+ )
272
+
273
+ app.add_middleware(WebSocketMetricsMiddleware)
274
+
275
if envs.VLLM_DEBUG_LOG_API_SERVER_RESPONSE:
276
logger.warning(
277
"CAUTION: Enabling log response in the API Server. "
vllm/entrypoints/openai/realtime/metrics.py
new
+78
@@ -0,0 +1,78 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ASGI middleware for WebSocket Prometheus metrics.
4
+
5
+Modeled after prometheus-fastapi-instrumentator, this middleware
6
+transparently instruments WebSocket endpoints with standard metrics
7
+without requiring changes to handler code.
8
+
9
+NOTE: This module intentionally has zero vllm imports so that it can
10
+be extracted into a standalone package (similar to
11
+prometheus-fastapi-instrumentator) in the future. Please keep it that way.
12
+"""
13
+
14
+import time
15
+from collections.abc import Awaitable
16
+
17
+from prometheus_client import Counter, Gauge, Histogram
18
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
19
+
20
+# Standard WebSocket metric names (not vllm-specific, following
21
+# the same convention as prometheus-fastapi-instrumentator).
22
+_active_sessions = Gauge(
23
+ name="vllm:websocket_connections_active",
24
+ documentation="Number of currently active WebSocket connections.",
25
+ multiprocess_mode="livesum",
26
+)
27
+
28
+_total_sessions = Counter(
29
+ name="vllm:websocket_connections_total",
30
+ documentation="Total number of WebSocket connections.",
31
+)
32
+
33
+_session_duration = Histogram(
34
+ name="vllm:websocket_connection_duration_seconds",
35
+ documentation="Duration of WebSocket connections in seconds.",
36
+ buckets=[0.5, 1, 2.5, 5, 10, 30, 60, 120, 300, 600, 1800],
37
+)
38
+
39
+
40
+class WebSocketMetricsMiddleware:
41
+ """Pure ASGI middleware that instruments WebSocket connections.
42
+
43
+ Tracks active connections (gauge), total connections (counter),
44
+ and connection duration (histogram) for all WebSocket endpoints.
45
+
46
+ Usage::
47
+
48
+ app.add_middleware(WebSocketMetricsMiddleware)
49
+ """
50
+
51
+ def __init__(self, app: ASGIApp) -> None:
52
+ self.app = app
53
+
54
+ def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]:
55
+ if scope["type"] != "websocket":
56
+ return self.app(scope, receive, send)
57
+
58
+ return self._handle_websocket(scope, receive, send)
59
+
60
+ async def _handle_websocket(
61
+ self, scope: Scope, receive: Receive, send: Send
62
+ ) -> None:
63
+ start_time: float | None = None
64
+
65
+ async def send_wrapper(message: Message) -> None:
66
+ nonlocal start_time
67
+ if message["type"] == "websocket.accept":
68
+ start_time = time.monotonic()
69
+ _active_sessions.inc()
70
+ _total_sessions.inc()
71
+ await send(message)
72
+
73
+ try:
74
+ await self.app(scope, receive, send_wrapper)
75
+ finally:
76
+ if start_time is not None:
77
+ _active_sessions.dec()
78
+ _session_duration.observe(time.monotonic() - start_time)