[Bugfix][Router] release request-stats slot in finally, not just on success

src/tests/test_request_stats_leak.py

@@ -0,0 +1,105 @@
+# test_request_stats_leak.py
+"""
+Regression tests for the request-tracking leaks fixed alongside
+production-stack#707: RequestStatsMonitor.on_request_complete() previously
+always decremented in_decoding_requests regardless of whether the request
+ever reached decode, and never released request_start_time/first_token_time,
+leaking one entry per request for the life of the process.
+"""
+import unittest
+
+from vllm_router.stats.request_stats import RequestStatsMonitor, SingletonMeta
+
+
+class TestOnRequestCompleteBucketRelease(unittest.TestCase):
+    def setUp(self):
+        if RequestStatsMonitor in SingletonMeta._instances:
+            del SingletonMeta._instances[RequestStatsMonitor]
+        self.monitor = RequestStatsMonitor(sliding_window_size=10.0)
+        self.engine_url = "http://engine-1:8000"
+        self.request_id = "req-1"
+
+    def test_reached_decode_true_releases_decoding_bucket(self):
+        self.monitor.on_new_request(self.engine_url, self.request_id, timestamp=0.0)
+        self.monitor.on_request_response(
+            self.engine_url, self.request_id, timestamp=0.1
+        )
+        self.assertEqual(self.monitor.in_prefill_requests[self.engine_url], 0)
+        self.assertEqual(self.monitor.in_decoding_requests[self.engine_url], 1)
+
+        self.monitor.on_request_complete(
+            self.engine_url, self.request_id, timestamp=0.2, reached_decode=True
+        )
+        self.assertEqual(self.monitor.in_decoding_requests[self.engine_url], 0)
+        self.assertEqual(self.monitor.in_prefill_requests[self.engine_url], 0)
+
+    def test_reached_decode_false_releases_prefill_bucket_not_decoding(self):
+        # A request that fails/disconnects before its first token never calls
+        # on_request_response(), so it's still sitting in in_prefill_requests.
+        self.monitor.on_new_request(self.engine_url, self.request_id, timestamp=0.0)
+        self.assertEqual(self.monitor.in_prefill_requests[self.engine_url], 1)
+
+        self.monitor.on_request_complete(
+            self.engine_url, self.request_id, timestamp=0.2, reached_decode=False
+        )
+        self.assertEqual(self.monitor.in_prefill_requests[self.engine_url], 0)
+        # in_decoding_requests was never touched for this engine -- confirms the
+        # old behavior (always decrementing in_decoding_requests) would have
+        # silently no-opped here (clamped at 0) while leaking in_prefill_requests.
+        self.assertNotIn(self.engine_url, self.monitor.in_decoding_requests)
+
+    def test_default_reached_decode_is_true_for_backward_compatibility(self):
+        # Existing callers (e.g. the multipart/transcription proxy path) call
+        # on_request_complete() without the new kwarg -- must keep behaving
+        # exactly as before.
+        self.monitor.on_new_request(self.engine_url, self.request_id, timestamp=0.0)
+        self.monitor.on_request_response(
+            self.engine_url, self.request_id, timestamp=0.1
+        )
+
+        self.monitor.on_request_complete(
+            self.engine_url, self.request_id, timestamp=0.2
+        )
+        self.assertEqual(self.monitor.in_decoding_requests[self.engine_url], 0)
+
+    def test_on_request_complete_pops_leaking_dicts(self):
+        key = (self.engine_url, self.request_id)
+        self.monitor.on_new_request(self.engine_url, self.request_id, timestamp=0.0)
+        self.monitor.on_request_response(
+            self.engine_url, self.request_id, timestamp=0.1
+        )
+        self.assertIn(key, self.monitor.request_start_time)
+        self.assertIn(key, self.monitor.first_token_time)
+
+        self.monitor.on_request_complete(
+            self.engine_url, self.request_id, timestamp=0.2, reached_decode=True
+        )
+        self.assertNotIn(
+            key,
+            self.monitor.request_start_time,
+            "request_start_time entry must be released on completion, or it "
+            "leaks one entry per request for the life of the process "
+            "(production-stack#707).",
+        )
+        self.assertNotIn(
+            key,
+            self.monitor.first_token_time,
+            "first_token_time entry must be released on completion for the "
+            "same reason.",
+        )
+
+    def test_on_request_complete_pops_request_start_time_even_without_first_token(self):
+        # A request that never reaches decode still populated request_start_time
+        # in on_new_request(); on_request_complete() must still release it.
+        key = (self.engine_url, self.request_id)
+        self.monitor.on_new_request(self.engine_url, self.request_id, timestamp=0.0)
+        self.assertIn(key, self.monitor.request_start_time)
+
+        self.monitor.on_request_complete(
+            self.engine_url, self.request_id, timestamp=0.2, reached_decode=False
+        )
+        self.assertNotIn(key, self.monitor.request_start_time)
+
+
+if __name__ == "__main__":
+    unittest.main()

src/vllm_router/services/request_service/request.py

@@ -301,6 +301,24 @@ async def process_request(
 
     request_status = "success"
     http_status_code = None
+    stats_released = False
+
+    def _release_in_flight_slot():
+        # Always release this request's in-flight slot exactly once, regardless of
+        # how the request's lifecycle ends: normal completion, an HTTP error status,
+        # a timeout, or a raw connection drop (e.g. aiohttp.ServerDisconnectedError)
+        # when a backend pod is torn down mid-stream. Previously this only ran on
+        # the success path below, inside the try block -- any exception talking to
+        # the backend left the request's slot permanently counted in the router's
+        # own in_prefill_requests/in_decoding_requests bookkeeping, which is what
+        # vllm:num_requests_running reports (production-stack#707-adjacent leak).
+        nonlocal stats_released
+        if stats_released:
+            return
+        stats_released = True
+        request.app.state.request_stats_monitor.on_request_complete(
+            backend_url, request_id, time.time(), reached_decode=first_token
+        )
 
     try:
         async with request.app.state.aiohttp_client_wrapper().request(
@@ -330,11 +348,6 @@ async def process_request(
                     full_response.extend(chunk)
                 yield chunk
 
-        end_time = time.time()
-        request.app.state.request_stats_monitor.on_request_complete(
-            backend_url, request_id, end_time
-        )
-
         if http_status_code is not None and http_status_code >= 400:
             request_status = "error"
 
@@ -374,6 +387,7 @@ async def process_request(
         end_span(span, error=e) if tracing_active else None
         raise
     finally:
+        _release_in_flight_slot()
         request_latency_seconds.labels(
             server=backend_url, model=model_name, status=request_status
         ).observe(time.time() - start_time)

src/vllm_router/stats/request_stats.py

@@ -200,27 +200,49 @@ def on_request_response(self, engine_url: str, request_id: str, timestamp: float
         ttft = timestamp - self.request_start_time[(engine_url, request_id)]
         self.ttft_monitors[engine_url].update(timestamp, ttft)
 
-    def on_request_complete(self, engine_url: str, request_id: str, timestamp: float):
+    def on_request_complete(
+        self,
+        engine_url: str,
+        request_id: str,
+        timestamp: float,
+        reached_decode: bool = True,
+    ):
         """
         Tell the monitor that a request has been completed.
 
         Args:
             engine_url: The URL of the serving engine
             request_id: The global request ID
             timestamp: The timestamp when the request was completed
+            reached_decode: Whether the request ever received a first token
+                (i.e. on_request_response() was called for it). A request
+                that fails or disconnects before its first token is still
+                counted in in_prefill_requests, never in_decoding_requests --
+                decrementing the wrong bucket silently no-ops (clamped at 0)
+                and leaves the real bucket leaked forever.
         """
         if engine_url not in self.finished_requests:
             self.finished_requests[engine_url] = 0
-        self.in_decoding_requests[engine_url] = max(
-            0, self.in_decoding_requests.get(engine_url, 1) - 1
-        )
+        if reached_decode:
+            self.in_decoding_requests[engine_url] = max(
+                0, self.in_decoding_requests.get(engine_url, 1) - 1
+            )
+        else:
+            self.in_prefill_requests[engine_url] = max(
+                0, self.in_prefill_requests.get(engine_url, 1) - 1
+            )
         self.finished_requests[engine_url] += 1
 
         if request_start_time := self.request_start_time.get((engine_url, request_id)):
             self.latency_monitors[engine_url].update(
                 timestamp, time.time() - request_start_time
             )
 
+        # Bound these dicts -- they otherwise grow by one entry per request
+        # for the life of the process (production-stack#707).
+        self.request_start_time.pop((engine_url, request_id), None)
+        self.first_token_time.pop((engine_url, request_id), None)
+
     def on_request_swapped(self, engine_url: str, request_id: str, timestamp: float):
         # This function should be called if a request is determined to be swapped from GPU to CPU.
         """