From 876fa5510d28b5e67c861371eeed6c34e7fb70a3 Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Thu, 24 Sep 2026 18:23:24 -0400 Subject: [PATCH 1/6] fix(scheduler): remove a request's queue entry once it stops waiting Requests admitted while a healthy deployment existed never left the priority queue, and neither did requests cancelled while waiting. The stale entries blocked later requests during cooldown and, with Redis, made add_request raise a TypeError on queues read back as JSON lists Both scheduling paths now share one polling helper that removes the entry in a finally block, whether the request was admitted, timed out or cancelled Related to #43059 --- litellm/router.py | 138 ++++++++---------------------- tests/test_litellm/test_router.py | 59 +++++++++++++ 2 files changed, 96 insertions(+), 101 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 8960cd92cd8..7a4da72a2b8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4313,57 +4313,17 @@ 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' + await self._wait_for_scheduler_turn( + model=model, priority=priority, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs) ) - ### [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", - ) + 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 async def _schedule_factory( self, @@ -4373,60 +4333,36 @@ class Router: args: tuple[object, ...], kwargs: dict[str, object], ): - 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' + await self._wait_for_scheduler_turn( + model=model, priority=priority, parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs) ) - ### [fin] ### + try: + _response: Final = await original_function(*args, **kwargs) + if isinstance(_response._hidden_params, dict): + _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 - ## ADDS REQUEST TO QUEUE ## + async def _wait_for_scheduler_turn(self, model: str, priority: int, parent_otel_span: Span | None) -> None: + item: Final = FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model) 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 original_function(*args, **kwargs) - if isinstance(_response._hidden_params, dict): - _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", - ) + try: + end_time: Final = time.monotonic() + self.timeout + while time.monotonic() < end_time: + healthy_deployments, _ = await self._async_get_healthy_deployments( + model=model, parent_otel_span=parent_otel_span + ) + if await self.scheduler.poll( + id=item.request_id, model_name=model, health_deployments=healthy_deployments + ): + return + await asyncio.sleep(self.scheduler.polling_interval) + finally: + await self.scheduler.remove_request(request_id=item.request_id, model_name=model) + raise litellm.Timeout(message="Request timed out while polling queue", model=model, llm_provider="openai") def _is_prompt_management_model(self, model: str) -> bool: model_list: Final = self.get_model_list(model_name=model) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index f985b212b01..69f499a020c 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -50,6 +50,7 @@ from litellm.router_strategy import simple_shuffle from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute +from litellm.scheduler import FlowItem from litellm.types.llms.openai import ChatCompletionRequest from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy @@ -17862,3 +17863,61 @@ def test_access_windows_filter_reserved_deployments_method(): request_team_id="team-a", ) ] == ["reserved-deployment", "open-deployment"] + + +def _scheduled_router(timeout: float) -> Router: + return Router( + model_list=[ + { + "model_name": "sched-model", + "litellm_params": {"model": "openai/sched-model", "api_key": "sk-fake", "mock_response": "hi"}, + "model_info": {"id": "sched-deployment"}, + } + ], + timeout=timeout, + ) + + +async def _send_scheduled_chat(router: Router, priority: int) -> object: + return await router.acompletion( + model="sched-model", messages=[{"role": "user", "content": "hi"}], priority=priority + ) + + +async def _send_scheduled_text(router: Router, priority: int) -> object: + return await router.atext_completion(model="sched-model", prompt="hi", priority=priority) + + +@pytest.mark.parametrize( + "send", [_send_scheduled_chat, _send_scheduled_text], ids=["schedule_acompletion", "schedule_factory"] +) +@pytest.mark.asyncio +async def test_admitted_prioritized_request_does_not_block_later_request_during_cooldown( + send: Callable[[Router, int], Awaitable[object]], +): + from litellm.types.router import RouterRateLimitError + + router: Final = _scheduled_router(timeout=1) + await send(router, 1) + _cool_down(router, "sched-deployment") + + with pytest.raises(RouterRateLimitError, match="cooldown"): + await send(router, 2) + + +@pytest.mark.parametrize("stop_waiting", ["cancel", "timeout"]) +@pytest.mark.asyncio +async def test_prioritized_request_leaves_queue_when_it_stops_waiting(stop_waiting: Literal["cancel", "timeout"]): + router: Final = _scheduled_router(timeout=0.5) + _cool_down(router, "sched-deployment") + await router.scheduler.add_request(FlowItem(priority=0, request_id="head-of-queue", model_name="sched-model")) + waiting: Final = asyncio.create_task(_send_scheduled_chat(router, 5)) + await asyncio.sleep(0.05) + assert len(await router.scheduler.get_queue("sched-model")) == 2 + + if stop_waiting == "cancel": + waiting.cancel() + with pytest.raises(asyncio.CancelledError if stop_waiting == "cancel" else litellm.Timeout): + await waiting + + assert await router.scheduler.get_queue("sched-model") == [(0, "head-of-queue")] From da4dc271cf19f95882c24ea60468388f1c4bd395 Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Thu, 24 Sep 2026 18:37:54 -0400 Subject: [PATCH 2/6] test(router): allowlist _wait_for_scheduler_turn in the router coverage check The coverage script only counts direct calls in test files. The helper is exercised through prioritized acompletion and atext_completion in test_router.py, like the other allowlisted entries --- tests/code_coverage_tests/router_code_coverage.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a11f015743b..5dc581e7d13 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -88,6 +88,7 @@ ignored_function_names = [ "_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py "_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py "_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name) + "_wait_for_scheduler_turn", # Tested through prioritized acompletion and atext_completion in test_router.py ] From 629b9ef1de343dc5d3e8e7fe72c64c2eafd6ad12 Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Thu, 24 Sep 2026 18:50:07 -0400 Subject: [PATCH 3/6] fix(scheduler): admit healthy requests before reading the queue, and clean up cancelled enqueues poll() raised on an empty queue before checking for healthy deployments. With the cleanup now rewriting the queue after every admission, a concurrent write from another replica can erase a waiting request's entry, and that request then failed while a deployment was healthy. poll() now admits as soon as a deployment is healthy and only reads the queue during cooldown add_request also moved inside the try block, so a request cancelled while its queue write is in flight still has its entry removed --- litellm/router.py | 2 +- litellm/scheduler.py | 26 ++++++++++--------------- tests/test_litellm/test_router.py | 29 +++++++++++++++++++++++++++- tests/test_litellm/test_scheduler.py | 25 ++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 18 deletions(-) create mode 100644 tests/test_litellm/test_scheduler.py diff --git a/litellm/router.py b/litellm/router.py index 7a4da72a2b8..dcc9870565a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4348,8 +4348,8 @@ class Router: async def _wait_for_scheduler_turn(self, model: str, priority: int, parent_otel_span: Span | None) -> None: item: Final = FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model) - await self.scheduler.add_request(request=item) try: + await self.scheduler.add_request(request=item) end_time: Final = time.monotonic() + self.timeout while time.monotonic() < end_time: healthy_deployments, _ = await self._async_get_healthy_deployments( diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 028b5d085e2..942bd5ea395 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -61,27 +61,21 @@ class Scheduler: * If no healthy deployments available * AND request not at the top of queue """ + print_verbose(f"len(health_deployments): {len(health_deployments)}") + if len(health_deployments) > 0: + return True + queue: Final = await self.get_queue(model_name=model_name) if not queue: raise Exception(f"Incorrectly setup. Queue is invalid. Queue={queue}") - # ------------ - # Setup values - # ------------ - - print_verbose(f"len(health_deployments): {len(health_deployments)}") - if len(health_deployments) == 0: - print_verbose(f"queue: {queue}, seeking id={id}") - # Check if the id is at the top of the heap - if queue[0][1] == id: - # Remove the item from the queue - heapq.heappop(queue) - await self.save_queue(queue=queue, model_name=model_name) - print_verbose(f"Popped id: {id}") - return True - else: - return False + print_verbose(f"queue: {queue}, seeking id={id}") + if queue[0][1] != id: + return False + heapq.heappop(queue) + await self.save_queue(queue=queue, model_name=model_name) + print_verbose(f"Popped id: {id}") return True async def remove_request(self, request_id: str, model_name: str) -> None: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 69f499a020c..df466752951 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -50,7 +50,7 @@ from litellm.router_strategy import simple_shuffle from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute -from litellm.scheduler import FlowItem +from litellm.scheduler import FlowItem, Scheduler from litellm.types.llms.openai import ChatCompletionRequest from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy @@ -17921,3 +17921,30 @@ async def test_prioritized_request_leaves_queue_when_it_stops_waiting(stop_waiti await waiting assert await router.scheduler.get_queue("sched-model") == [(0, "head-of-queue")] + + +class _PausesAfterEnqueueScheduler(Scheduler): + def __init__(self) -> None: + super().__init__() + self.enqueued: Final = asyncio.Event() + + async def add_request(self, request: FlowItem) -> None: + await super().add_request(request) + self.enqueued.set() + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_prioritized_request_cancelled_while_enqueueing_leaves_queue(): + router: Final = _scheduled_router(timeout=5) + scheduler: Final = _PausesAfterEnqueueScheduler() + router.scheduler = scheduler + waiting: Final = asyncio.create_task(_send_scheduled_chat(router, 1)) + await scheduler.enqueued.wait() + assert len(await scheduler.get_queue("sched-model")) == 1 + + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + + assert await scheduler.get_queue("sched-model") == [] diff --git a/tests/test_litellm/test_scheduler.py b/tests/test_litellm/test_scheduler.py new file mode 100644 index 00000000000..bf088642604 --- /dev/null +++ b/tests/test_litellm/test_scheduler.py @@ -0,0 +1,25 @@ +from typing import Final + +import pytest + +from litellm.scheduler import FlowItem, Scheduler + + +@pytest.mark.asyncio +async def test_poll_admits_request_missing_from_queue_while_a_deployment_is_healthy(): + scheduler: Final = Scheduler() + + assert await scheduler.poll( + id="erased-by-concurrent-write", model_name="sched-model", health_deployments=[{"model_info": {"id": "a"}}] + ) + + +@pytest.mark.asyncio +async def test_poll_during_cooldown_admits_only_the_head_of_the_queue(): + scheduler: Final = Scheduler() + await scheduler.add_request(FlowItem(priority=2, request_id="later", model_name="sched-model")) + await scheduler.add_request(FlowItem(priority=1, request_id="head", model_name="sched-model")) + + assert not await scheduler.poll(id="later", model_name="sched-model", health_deployments=[]) + assert await scheduler.poll(id="head", model_name="sched-model", health_deployments=[]) + assert await scheduler.get_queue("sched-model") == [(2, "later")] From f5383a87ddf69f48461045451db07ee2cd0f1866 Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Thu, 24 Sep 2026 19:07:35 -0400 Subject: [PATCH 4/6] refactor(scheduler): move the wait loop into Scheduler.wait_for_turn The router passes a healthy-deployments callable into the scheduler, so tests inject a Scheduler directly instead of replacing the router's scheduler attribute. The cancelled-mid-enqueue test moves to test_scheduler.py with the other scheduler tests, and poll() takes the deployments as a Sequence since it only checks whether any are healthy --- litellm/router.py | 25 +++++++------------ litellm/scheduler.py | 27 +++++++++++++++++++- tests/test_litellm/test_router.py | 28 +-------------------- tests/test_litellm/test_scheduler.py | 37 ++++++++++++++++++++++++++++ 4 files changed, 73 insertions(+), 44 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index dcc9870565a..d5fb6457199 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4347,22 +4347,15 @@ class Router: raise e async def _wait_for_scheduler_turn(self, model: str, priority: int, parent_otel_span: Span | None) -> None: - item: Final = FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model) - try: - await self.scheduler.add_request(request=item) - end_time: Final = time.monotonic() + self.timeout - while time.monotonic() < end_time: - healthy_deployments, _ = await self._async_get_healthy_deployments( - model=model, parent_otel_span=parent_otel_span - ) - if await self.scheduler.poll( - id=item.request_id, model_name=model, health_deployments=healthy_deployments - ): - return - await asyncio.sleep(self.scheduler.polling_interval) - finally: - await self.scheduler.remove_request(request_id=item.request_id, model_name=model) - raise litellm.Timeout(message="Request timed out while polling queue", model=model, llm_provider="openai") + async def healthy_deployments() -> Sequence[object]: + deployments, _ = await self._async_get_healthy_deployments(model=model, parent_otel_span=parent_otel_span) + return deployments + + await self.scheduler.wait_for_turn( + request=FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model), + timeout=self.timeout, + get_healthy_deployments=healthy_deployments, + ) def _is_prompt_management_model(self, model: str) -> bool: model_list: Final = self.get_model_list(model_name=model) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 942bd5ea395..66976310485 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -1,5 +1,8 @@ +import asyncio import enum import heapq +import time +from collections.abc import Awaitable, Callable, Sequence from typing import Final from pydantic import BaseModel @@ -7,6 +10,7 @@ from pydantic import BaseModel from litellm import print_verbose from litellm.caching.caching import DualCache, RedisCache from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL +from litellm.exceptions import Timeout class SchedulerCacheKeys(enum.Enum): @@ -49,7 +53,7 @@ class Scheduler: # save the queue await self.save_queue(queue=queue, model_name=request.model_name) - async def poll(self, id: str, model_name: str, health_deployments: list) -> bool: + async def poll(self, id: str, model_name: str, health_deployments: Sequence[object]) -> bool: """ Return if request can be processed. @@ -78,6 +82,27 @@ class Scheduler: print_verbose(f"Popped id: {id}") return True + async def wait_for_turn( + self, + request: FlowItem, + timeout: float, + get_healthy_deployments: Callable[[], Awaitable[Sequence[object]]], + ) -> None: + try: + await self.add_request(request=request) + end_time: Final = time.monotonic() + timeout + while time.monotonic() < end_time: + if await self.poll( + id=request.request_id, + model_name=request.model_name, + health_deployments=await get_healthy_deployments(), + ): + return + await asyncio.sleep(self.polling_interval) + finally: + await self.remove_request(request_id=request.request_id, model_name=request.model_name) + raise Timeout(message="Request timed out while polling queue", model=request.model_name, llm_provider="openai") + async def remove_request(self, request_id: str, model_name: str) -> None: """ Remove a specific request from the priority queue for a model. diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index df466752951..b56b32b078d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -50,7 +50,7 @@ from litellm.router_strategy import simple_shuffle from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute -from litellm.scheduler import FlowItem, Scheduler +from litellm.scheduler import FlowItem from litellm.types.llms.openai import ChatCompletionRequest from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy @@ -17922,29 +17922,3 @@ async def test_prioritized_request_leaves_queue_when_it_stops_waiting(stop_waiti assert await router.scheduler.get_queue("sched-model") == [(0, "head-of-queue")] - -class _PausesAfterEnqueueScheduler(Scheduler): - def __init__(self) -> None: - super().__init__() - self.enqueued: Final = asyncio.Event() - - async def add_request(self, request: FlowItem) -> None: - await super().add_request(request) - self.enqueued.set() - await asyncio.Event().wait() - - -@pytest.mark.asyncio -async def test_prioritized_request_cancelled_while_enqueueing_leaves_queue(): - router: Final = _scheduled_router(timeout=5) - scheduler: Final = _PausesAfterEnqueueScheduler() - router.scheduler = scheduler - waiting: Final = asyncio.create_task(_send_scheduled_chat(router, 1)) - await scheduler.enqueued.wait() - assert len(await scheduler.get_queue("sched-model")) == 1 - - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - - assert await scheduler.get_queue("sched-model") == [] diff --git a/tests/test_litellm/test_scheduler.py b/tests/test_litellm/test_scheduler.py index bf088642604..6b4c5f2fe80 100644 --- a/tests/test_litellm/test_scheduler.py +++ b/tests/test_litellm/test_scheduler.py @@ -1,3 +1,5 @@ +import asyncio +from collections.abc import Sequence from typing import Final import pytest @@ -23,3 +25,38 @@ async def test_poll_during_cooldown_admits_only_the_head_of_the_queue(): assert not await scheduler.poll(id="later", model_name="sched-model", health_deployments=[]) assert await scheduler.poll(id="head", model_name="sched-model", health_deployments=[]) assert await scheduler.get_queue("sched-model") == [(2, "later")] + + +class _PausesAfterEnqueueScheduler(Scheduler): + def __init__(self) -> None: + super().__init__() + self.enqueued: Final = asyncio.Event() + + async def add_request(self, request: FlowItem) -> None: + await super().add_request(request) + self.enqueued.set() + await asyncio.Event().wait() + + +async def _no_healthy_deployments() -> Sequence[object]: + return () + + +@pytest.mark.asyncio +async def test_wait_for_turn_removes_entry_when_cancelled_mid_enqueue(): + scheduler: Final = _PausesAfterEnqueueScheduler() + waiting: Final = asyncio.create_task( + scheduler.wait_for_turn( + request=FlowItem(priority=1, request_id="cancelled", model_name="sched-model"), + timeout=5, + get_healthy_deployments=_no_healthy_deployments, + ) + ) + await scheduler.enqueued.wait() + assert await scheduler.get_queue("sched-model") == [(1, "cancelled")] + + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + + assert await scheduler.get_queue("sched-model") == [] From e7db39186eef0ef532427888ee004c894fcd607d Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Tue, 29 Sep 2026 13:54:18 -0400 Subject: [PATCH 5/6] test(scheduler): move scheduler tests into tests/unit --- tests/{test_litellm => unit}/test_scheduler.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{test_litellm => unit}/test_scheduler.py (100%) diff --git a/tests/test_litellm/test_scheduler.py b/tests/unit/test_scheduler.py similarity index 100% rename from tests/test_litellm/test_scheduler.py rename to tests/unit/test_scheduler.py From 0c9ab36172fe7e9d3e4f70c2dd025afcb161ebb8 Mon Sep 17 00:00:00 2001 From: RachelHuangZW Date: Tue, 29 Sep 2026 13:59:23 -0400 Subject: [PATCH 6/6] fix(scheduler): finish the queue removal when cancellation is delivered again during cleanup A second cancel, or an anyio cancel scope that re-cancels on every await, interrupted remove_request mid-write and left the entry in Redis. --- litellm/scheduler.py | 2 +- tests/unit/test_scheduler.py | 40 ++++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 66976310485..cffed3acce5 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -100,7 +100,7 @@ class Scheduler: return await asyncio.sleep(self.polling_interval) finally: - await self.remove_request(request_id=request.request_id, model_name=request.model_name) + await asyncio.shield(self.remove_request(request_id=request.request_id, model_name=request.model_name)) raise Timeout(message="Request timed out while polling queue", model=request.model_name, llm_provider="openai") async def remove_request(self, request_id: str, model_name: str) -> None: diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index 6b4c5f2fe80..2b643558089 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -60,3 +60,43 @@ async def test_wait_for_turn_removes_entry_when_cancelled_mid_enqueue(): await waiting assert await scheduler.get_queue("sched-model") == [] + + +class _HeldRemovalScheduler(Scheduler): + def __init__(self) -> None: + super().__init__() + self.removing: Final = asyncio.Event() + self.finish_removal: Final = asyncio.Event() + + async def remove_request(self, request_id: str, model_name: str) -> None: + self.removing.set() + await self.finish_removal.wait() + await super().remove_request(request_id=request_id, model_name=model_name) + + +@pytest.mark.asyncio +async def test_wait_for_turn_finishes_removal_when_cancelled_again_during_cleanup(): + scheduler: Final = _HeldRemovalScheduler() + await scheduler.add_request(FlowItem(priority=0, request_id="head", model_name="sched-model")) + polling: Final = asyncio.Event() + + async def no_healthy_deployments() -> Sequence[object]: + polling.set() + return () + + waiting: Final = asyncio.create_task( + scheduler.wait_for_turn( + request=FlowItem(priority=1, request_id="cancelled", model_name="sched-model"), + timeout=5, + get_healthy_deployments=no_healthy_deployments, + ) + ) + await polling.wait() + waiting.cancel() + await scheduler.removing.wait() + waiting.cancel() + scheduler.finish_removal.set() + with pytest.raises(asyncio.CancelledError): + await waiting + + assert await scheduler.get_queue("sched-model") == [(0, "head")]