Merge 3ec953d9de into 8fcb016dc6

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

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)

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(

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
 

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: