| 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) |