From 642f941296390fe655fddd2bca1e1e3d1bab8fdd Mon Sep 17 00:00:00 2001 From: mateo Date: Thu, 13 Aug 2026 19:46:39 +0000 Subject: [PATCH] fix(router): use default_priority when the request omits priority acompletion computed the effective priority but called schedule_acompletion without it, so any proxy with router_settings.default_priority set returned a 500 TypeError for every /v1/chat/completions request. Route through _schedule_factory so the priority is queued and the scheduler no longer re-enters acompletion. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router.py | 64 ++++--------------------- tests/test_litellm/test_router.py | 79 +++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 54 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index fb2af41dcf2..e2a1b7a9f7a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2140,7 +2140,7 @@ class Router: kwargs["original_function"] = self._acompletion self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - request_priority: Final = kwargs.get("priority") or self.default_priority + request_priority: Final = kwargs.pop("priority", None) or self.default_priority start_time: Final = time.time() _is_prompt_management_model: Final = self._is_prompt_management_model(model) @@ -2150,8 +2150,14 @@ class Router: messages=messages, kwargs=kwargs, ) - if request_priority is not None and isinstance(request_priority, int): - response = await self.schedule_acompletion(**kwargs) + if isinstance(request_priority, int): + response = await self._schedule_factory( + model=model, + priority=request_priority, + original_function=self.async_function_with_fallbacks, + args=(), + kwargs=kwargs, + ) else: response = await self.async_function_with_fallbacks(**kwargs) end_time: Final = time.time() @@ -3614,57 +3620,7 @@ class Router: stream=False, **kwargs, ): - parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) - ### FLOW ITEM ### - _request_id: Final = str(uuid.uuid4()) - item: Final = FlowItem( - priority=priority, # 👈 SET PRIORITY FOR REQUEST - request_id=_request_id, # 👈 SET REQUEST ID - model_name=model, # 👈 SAME as 'Router' - ) - ### [fin] ### - - ## ADDS REQUEST TO QUEUE ## - await self.scheduler.add_request(request=item) - - ## POLL QUEUE - end_time: Final = time.monotonic() + self.timeout - curr_time = time.monotonic() - poll_interval: Final = self.scheduler.polling_interval # poll every 3ms - make_request = False - - while curr_time < end_time: - _healthy_deployments, _ = await self._async_get_healthy_deployments( - model=model, parent_otel_span=parent_otel_span - ) - make_request = await self.scheduler.poll( ## POLL QUEUE ## - returns 'True' if there's healthy deployments OR if request is at top of queue - id=item.request_id, - model_name=item.model_name, - health_deployments=_healthy_deployments, - ) - if make_request: ## IF TRUE -> MAKE REQUEST - break - else: ## ELSE -> loop till default_timeout - await asyncio.sleep(poll_interval) - curr_time = time.monotonic() - - if make_request: - try: - _response: Final = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs) - _response._hidden_params.setdefault("additional_headers", {}) - _response._hidden_params["additional_headers"].update({"x-litellm-request-prioritization-used": True}) - return _response - except Exception as e: - 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, - llm_provider="openai", - ) + return await self.acompletion(model=model, messages=messages, stream=stream, priority=priority, **kwargs) async def _schedule_factory( self, diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index bdbf33fb0e1..253c265c2f8 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7954,3 +7954,82 @@ def test_ensure_deployment_affinity_callback_is_idempotent(): finally: for cb in router.optional_callbacks or []: litellm.logging_callback_manager.remove_callback_from_all_lists(cb) + + +def _priority_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "code", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key"}, + } + ], + default_priority=10, + ) + + +@pytest.mark.asyncio +async def test_acompletion_uses_default_priority_when_request_has_none(): + """default_priority must be forwarded to the scheduler instead of blowing up the request""" + router = _priority_router() + queued: list[int] = [] + original_add_request = router.scheduler.add_request + + async def spy(request): + queued.append(request.priority) + return await original_add_request(request=request) + + with patch.object(router.scheduler, "add_request", side_effect=spy): + response = await router.acompletion( + model="code", + messages=[{"role": "user", "content": "Hi"}], + mock_response="Hello", + ) + + assert queued == [10] + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +@pytest.mark.asyncio +async def test_acompletion_request_priority_wins_over_default_priority(): + router = _priority_router() + queued: list[int] = [] + original_add_request = router.scheduler.add_request + + async def spy(request): + queued.append(request.priority) + return await original_add_request(request=request) + + with patch.object(router.scheduler, "add_request", side_effect=spy): + response = await router.acompletion( + model="code", + messages=[{"role": "user", "content": "Hi"}], + priority=5, + mock_response="Hello", + ) + + assert queued == [5] + assert response.choices[0].message.content == "Hello" + + +@pytest.mark.asyncio +async def test_schedule_acompletion_queues_once_with_default_priority_configured(): + """schedule_acompletion must not re-enter the scheduler via acompletion's default_priority""" + router = _priority_router() + queued: list[int] = [] + original_add_request = router.scheduler.add_request + + async def spy(request): + queued.append(request.priority) + return await original_add_request(request=request) + + with patch.object(router.scheduler, "add_request", side_effect=spy): + response = await router.schedule_acompletion( + model="code", + messages=[{"role": "user", "content": "Hi"}], + priority=3, + mock_response="Hello", + ) + + assert queued == [3] + assert response.choices[0].message.content == "Hello"