Merge 399a2d918d into e32d8a4d03

litellm/scheduler.py

@@ -101,6 +101,8 @@ class Scheduler:
         timeout: float,
         get_healthy_deployments: Callable[[], Awaitable[Sequence[object]]],
     ) -> None:
+        if len(await get_healthy_deployments()) > 0:
+            return
         try:
             await self.add_request(request=request)
             end_time: Final = time.monotonic() + timeout

tests/unit/proxy/proxy_server/test_routes_queue.py

@@ -154,11 +154,11 @@ def test_queue_chat_completions_reports_the_priority_the_router_queued_at(
 ) -> None:
     router: Final = _queue_router(monkeypatch, default_priority)
     payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         with auth_as():
             response: Final = client.post("/queue/chat/completions", json=payload)
     assert response.status_code == 200, response.text
-    assert [call.kwargs["request"].priority for call in add_request.await_args_list] == [expected_priority]
+    assert [call.kwargs["request"].priority for call in wait_for_turn.await_args_list] == [expected_priority]
     assert response.headers.get("x-litellm-priority") == str(expected_priority)
 
 
@@ -181,14 +181,14 @@ def test_queue_chat_completions_rejects_a_request_it_cannot_queue_as_an_invalid_
 ) -> None:
     router: Final = _queue_router(monkeypatch, default_priority)
     payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields}
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         with auth_as():
             response: Final = client.post("/queue/chat/completions", json=payload)
     assert response.status_code == 400, response.text
     error: Final = response.json()["error"]
     assert (error["type"], error["param"]) == ("invalid_request_error", "priority")
     assert message in error["message"]
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
 
 
 def test_streamed_queue_responses_report_the_priority_the_router_queued_at(

tests/unit/test_router/test_router.py

@@ -26382,20 +26382,20 @@ async def test_acompletion_queues_the_resolved_priority_exactly_once(
     default_priority: int | None, request_kwargs: dict[str, object], queued_priority: int
 ) -> None:
     router: Final = _priority_router(default_priority)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong", **request_kwargs
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == queued_priority
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == queued_priority
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
 async def test_acompletion_without_any_priority_never_enters_the_scheduler() -> None:
     router: Final = _priority_router(None)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong")
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
     assert "x-litellm-request-prioritization-used" not in response._hidden_params.get("additional_headers", {})
 
 
@@ -26459,12 +26459,12 @@ async def test_acompletion_with_an_integer_priority_sends_no_priority_to_the_pro
 
 async def test_schedule_acompletion_under_a_default_priority_queues_the_explicit_priority_once() -> None:
     router: Final = _priority_router(3)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.schedule_acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, priority=1, mock_response="pong"
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == 1
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == 1
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26491,11 +26491,11 @@ async def test_acompletion_drops_a_non_integer_priority_under_the_router_or_depl
         default_priority=3,
         default_litellm_params=router_defaults,
     )
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong", priority="high"
         )
-    assert add_request.await_args.kwargs["request"].priority == 3
+    assert wait_for_turn.await_args.kwargs["request"].priority == 3
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26539,7 +26539,7 @@ async def test_deployment_drop_params_drop_a_non_integer_priority_on_every_deplo
 ) -> None:
     monkeypatch.setattr(litellm, "drop_params", False)
     router: Final = Router(model_list=[deployment], default_priority=3)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(
             model=requested_model,
             messages=_PRIORITY_MESSAGES,
@@ -26547,7 +26547,7 @@ async def test_deployment_drop_params_drop_a_non_integer_priority_on_every_deplo
             priority="high",
             **routing_fields,
         )
-    assert add_request.await_args.kwargs["request"].priority == 3
+    assert wait_for_turn.await_args.kwargs["request"].priority == 3
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26575,12 +26575,12 @@ async def test_schedule_acompletion_queues_a_prompt_management_model(monkeypatch
         litellm_call_id="prompt-management-priority",
         function_id="f",
     )
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.schedule_acompletion(
             model="cached-chat", messages=_PRIORITY_MESSAGES, priority=1, litellm_logging_obj=logging_obj
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == 1
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == 1
     assert provider.await_count == 1
     assert "priority" not in provider.await_args.kwargs
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
@@ -26607,9 +26607,9 @@ def test_an_invalid_default_priority_is_ignored_with_a_warning(
 
 async def test_an_invalid_default_priority_leaves_every_request_unscheduled() -> None:
     router: Final = _priority_router("high")
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.acompletion(model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong")
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
     assert "x-litellm-request-prioritization-used" not in response._hidden_params.get("additional_headers", {})
 
 
@@ -26626,12 +26626,12 @@ async def test_schedule_acompletion_without_a_priority_or_a_default_is_rejected_
 
 async def test_schedule_acompletion_without_a_priority_queues_the_default_once() -> None:
     router: Final = _priority_router(3)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.schedule_acompletion(
             model="qa-chat", messages=_PRIORITY_MESSAGES, mock_response="pong"
         )
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == 3
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == 3
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
 
 
@@ -26655,10 +26655,10 @@ async def test_atext_completion_queues_the_resolved_priority_exactly_once(
 ) -> None:
     provider: Final = _priority_provider(respx_mock, monkeypatch)
     router: Final = _priority_router(default_priority)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.atext_completion(model="qa-chat", prompt="pong?", **request_kwargs)
-    assert add_request.await_count == 1
-    assert add_request.await_args.kwargs["request"].priority == queued_priority
+    assert wait_for_turn.await_count == 1
+    assert wait_for_turn.await_args.kwargs["request"].priority == queued_priority
     assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
     assert "priority" not in json.loads(provider.calls.last.request.content)
 
@@ -26668,9 +26668,9 @@ async def test_atext_completion_without_any_priority_never_enters_the_scheduler(
 ) -> None:
     provider: Final = _priority_provider(respx_mock, monkeypatch)
     router: Final = _priority_router(None)
-    with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request:
+    with patch.object(router.scheduler, "wait_for_turn", wraps=router.scheduler.wait_for_turn) as wait_for_turn:
         response: Final = await router.atext_completion(model="qa-chat", prompt="pong?")
-    assert add_request.await_count == 0
+    assert wait_for_turn.await_count == 0
     assert "x-litellm-request-prioritization-used" not in response._hidden_params.get("additional_headers", {})
     assert provider.called
 

tests/unit/test_scheduler.py

@@ -227,6 +227,23 @@ async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring():
     assert await scheduler.get_queue("sched-model") == []
 
 
+@pytest.mark.asyncio
+async def test_wait_for_turn_never_queues_a_request_a_healthy_deployment_can_take():
+    scheduler: Final = Scheduler()
+
+    async def one_healthy_deployment_seen_with_an_empty_queue() -> Sequence[object]:
+        queue: Final = await scheduler.get_queue("sched-model")
+        assert queue == [], f"the admitted request was queued while a deployment was healthy: {queue}"
+        return ({"model_info": {"id": "a"}},)
+
+    await scheduler.wait_for_turn(
+        request=FlowItem(priority=1, request_id="admitted", model_name="sched-model"),
+        timeout=5,
+        get_healthy_deployments=one_healthy_deployment_seen_with_an_empty_queue,
+    )
+
+    assert await scheduler.get_queue("sched-model") == []
+
 
 class _JsonRoundTripRedisCache:
     def __init__(self) -> None: