diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 48be61283df..9faa3422ef8 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -126,6 +126,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/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 11798542088..9f7218a97d6 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -686,7 +686,7 @@ def test_served_request_leaves_no_queue_entry_for_the_other_instance(pair: Pair) group: Final = PAIR_GROUPS[0] markers: Final = (new_marker(), new_marker(), new_marker()) assert_served(post(pair.first, "/v1/chat/completions", chat_body(group, markers[0], priority=1)), markers[0]) - assert queue_entries(pair.cache, group) == [] + assert_drained(pair.cache, group) assert_served(post(pair.second, "/v1/chat/completions", chat_body(group, markers[1], priority=1)), markers[1]) assert_served(post(pair.first, "/v1/chat/completions", chat_body(group, markers[2], priority=1)), markers[2]) assert_once(received_markers(pair.wire), markers) @@ -1441,7 +1441,6 @@ def test_concurrent_burst_without_priority_across_instances_and_endpoints_is_ser @pytest.mark.timeout(PAIR_CELL_TIMEOUT) def test_concurrent_burst_without_priority_across_instances_leaves_no_queue_entry(defaulted: Defaulted) -> None: - pytest.skip("BUG: concurrent queue writes across workers can leave a served request's entry in the queue") group: Final = DEFAULT_GROUPS[0] serve_burst_without_priority(defaulted.pair, group) assert_drained(defaulted.pair.cache, group) diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index db99818ebc4..8bb8b852cf9 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 @pytest.mark.parametrize( @@ -212,13 +212,13 @@ def test_queue_chat_completions_queues_at_the_priority_header( ) -> None: router: Final = _queue_router(monkeypatch, default_priority=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, headers={"x-litellm-priority": priority_header} ) 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) @@ -227,7 +227,7 @@ def test_queue_chat_completions_rejects_a_malformed_priority_header_like_a_malfo ) -> None: router: Final = _queue_router(monkeypatch, default_priority=5) payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "priority": 1} - 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, headers={"x-litellm-priority": "high"} @@ -236,7 +236,7 @@ def test_queue_chat_completions_rejects_a_malformed_priority_header_like_a_malfo error: Final = response.json()["error"] assert (error["type"], error["param"]) == ("invalid_request_error", "priority") assert "priority must be an integer" 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 c16c1c62202..2235615f1b2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -28230,20 +28230,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", {}) @@ -28307,12 +28307,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 @@ -28339,11 +28339,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 @@ -28387,7 +28387,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, @@ -28395,7 +28395,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 @@ -28423,12 +28423,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 @@ -28455,9 +28455,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", {}) @@ -28474,12 +28474,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 @@ -28503,10 +28503,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) @@ -28516,9 +28516,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 fbc3c0e47c3..09b11713cbe 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -282,6 +282,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: