diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 73196d24e4d..e4ff6e3a287 100644 --- a/litellm/scheduler.py +++ b/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 diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index a86314e3063..ec969842c21 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/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( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d9b655a0aff..c6886735557 100644 --- a/tests/unit/test_router/test_router.py +++ b/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 diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index e0968e9b2cf..8688fc1ba5d 100644 --- a/tests/unit/test_scheduler.py +++ b/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: