fix(scheduler): remove orphan entries from queue - causing memory leak. (#20866)

* fix(scheduler): remove timed-out requests from queue to prevent memory leak

Fixes #20059

* fix(scheduler): use actual model param instead of hardcoded gpt-3.5-turbo in schedule_acompletion

* trigger CLA recheck

---------

Co-authored-by: Piyush Bhawsar <piyush100x@Piyushs-MacBook-Pro-3.local>

litellm/router.py

@@ -2279,7 +2279,7 @@ class Router:
         item = FlowItem(
             priority=priority,  # 👈 SET PRIORITY FOR REQUEST
             request_id=_request_id,  # 👈 SET REQUEST ID
-            model_name="gpt-3.5-turbo",  # 👈 SAME as 'Router'
+            model_name=model,  # 👈 SAME as 'Router'
         )
         ### [fin] ###
 
@@ -2321,6 +2321,10 @@ class Router:
                 setattr(e, "priority", priority)
                 raise e
         else:
+            # Clean up the request from the scheduler queue also before raising the timeout exception
+            await self.scheduler.remove_request(
+                request_id=item.request_id, model_name=item.model_name
+            )
             raise litellm.Timeout(
                 message="Request timed out while polling queue",
                 model=model,
@@ -2382,6 +2386,10 @@ class Router:
                 setattr(e, "priority", priority)
                 raise e
         else:
+            # Clean up the request from the scheduler queue also before raising the timeout exception
+            await self.scheduler.remove_request(
+                request_id=item.request_id, model_name=item.model_name
+            )
             raise litellm.Timeout(
                 message="Request timed out while polling queue",
                 model=model,

litellm/scheduler.py

@@ -92,6 +92,17 @@ class Scheduler:
 
         return True
 
+    async def remove_request(self, request_id: str, model_name: str) -> None:
+        """
+        Remove a specific request from the priority queue for a model.
+        Used when a request times out while waiting in the queue.
+        """
+        queue = await self.get_queue(model_name=model_name)
+        filtered_queue = [item for item in queue if item[1] != request_id]
+        heapq.heapify(filtered_queue)  # restore heap invariant after filtering
+        await self.save_queue(queue=filtered_queue, model_name=model_name)
+        print_verbose(f"Removed request_id: {request_id} from queue for model: {model_name}")
+
     async def peek(self, id: str, model_name: str, health_deployments: list) -> bool:
         """Return if the id is at the top of the queue. Don't pop the value from heap."""
         queue = await self.get_queue(model_name=model_name)

tests/local_testing/test_scheduler.py

@@ -117,3 +117,41 @@ async def test_scheduler_prioritized_requests(p0, p1, healthy_deployments):
             )
             == False
         )
+
+
+@pytest.mark.asyncio
+async def test_scheduler_queue_cleanup_on_timeout():
+    """
+    Test that a timed-out request is properly removed from the queue.
+    This prevents memory leaks from accumulating timed-out requests.
+    """
+    scheduler = Scheduler()
+
+    # Add multiple requests with different priorities
+    item1 = FlowItem(priority=0, request_id="req-0", model_name="gpt-3.5-turbo")
+    item2 = FlowItem(priority=1, request_id="req-1", model_name="gpt-3.5-turbo")
+    item3 = FlowItem(priority=2, request_id="req-2", model_name="gpt-3.5-turbo")
+
+    await scheduler.add_request(item1)
+    await scheduler.add_request(item2)
+    await scheduler.add_request(item3)
+
+    # Verify initial queue size
+    queue_before = await scheduler.get_queue(model_name="gpt-3.5-turbo")
+    assert len(queue_before) == 3, f"Expected 3 items in queue, got {len(queue_before)}"
+
+    # Simulate timeout cleanup - remove a non-front request (item2)
+    await scheduler.remove_request(request_id="req-1", model_name="gpt-3.5-turbo")
+
+    # Verify queue was cleaned up
+    queue_after = await scheduler.get_queue(model_name="gpt-3.5-turbo")
+    assert len(queue_after) == 2, f"Expected 2 items after cleanup, got {len(queue_after)}"
+
+    # Verify the correct request was removed
+    remaining_ids = [item[1] for item in queue_after]
+    assert "req-1" not in remaining_ids, "Expected req-1 to be removed"
+    assert "req-0" in remaining_ids, "Expected req-0 to remain"
+    assert "req-2" in remaining_ids, "Expected req-2 to remain"
+
+    # Verify remaining items are in correct priority order (0 should be first)
+    assert queue_after[0][1] == "req-0", "Expected req-0 (priority 0) to be at front"