diff --git a/litellm/router.py b/litellm/router.py index d9de7e7fc5..952987e9d5 100644 --- a/litellm/router.py +++ b/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, diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 5f3dd4cbf6..0221e24984 100644 --- a/litellm/scheduler.py +++ b/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) diff --git a/tests/local_testing/test_scheduler.py b/tests/local_testing/test_scheduler.py index f5b4422485..f198e572b2 100644 --- a/tests/local_testing/test_scheduler.py +++ b/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"