fix(scheduler): preserve request outcomes and back off queue conflicts

litellm/scheduler.py

@@ -2,6 +2,7 @@ import asyncio
 import enum
 import heapq
 import logging as _logging
+import random as _random
 import time
 from collections.abc import Awaitable, Callable, Sequence
 from functools import partial as _partial
@@ -106,7 +107,7 @@ async def _mutate_redis_queue(
     cache: RedisCache, key: str, model_name: str, ttl: int | None, change: Callable[[_Queue], _QueueResult]
 ) -> _QueueResult:
     script: Final = cache.async_register_script(_QUEUE_COMPARE_SET)
-    for _ in range(100):
+    for attempt in range(100):
         raw, queue = await _read_redis_queue(cache, key)
         admitted, updated = change(queue)
         if updated == queue:
@@ -123,6 +124,8 @@ async def _mutate_redis_queue(
             )
         ):
             return admitted, updated
+        if attempt < 99:
+            await asyncio.sleep(_random.uniform(0, min(0.05, 0.001 * 2 ** min(attempt, 6))))
     raise Timeout(
         message="Scheduler queue remained busy after concurrent updates", model=model_name, llm_provider="openai"
     )
@@ -251,9 +254,15 @@ class Scheduler:
                     return
                 await asyncio.sleep(self.polling_interval)
         finally:
-            await asyncio.shield(self.remove_request(request_id=request.request_id, model_name=request.model_name))
+            await asyncio.shield(self._cleanup_request(request))
         raise Timeout(message="Request timed out while polling queue", model=request.model_name, llm_provider="openai")
 
+    async def _cleanup_request(self, request: FlowItem) -> None:
+        try:
+            await self.remove_request(request_id=request.request_id, model_name=request.model_name)
+        except Exception as error:  # noqa: BLE001  # Cleanup must preserve every request outcome, including cancellation
+            _log_redis_failure(_verbose_router_logger, _logging.WARNING, "Scheduler queue cleanup failed", error)
+
     async def remove_request(self, request_id: str, model_name: str) -> None:
         """
         Remove a specific request from the priority queue for a model.

tests/integration/sdk/test_dual_cache_redis.py

@@ -3,7 +3,7 @@ import os
 import threading
 import uuid
 from collections.abc import AsyncIterator, Iterator
-from typing import Final
+from typing import Final, Literal
 from unittest.mock import patch
 
 import pytest
@@ -244,3 +244,64 @@ async def test_shared_scheduler_retains_priority_and_request_id(
         )
         <= 17
     )
+
+
+@pytest.mark.parametrize("outcome", ["admitted", "failed", "cancelled"])
+async def test_shared_scheduler_cleanup_failure_preserves_request_outcome(
+    shared_scheduler_cache: RedisCache,
+    outcome: Literal["admitted", "failed", "cancelled"],
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache)
+    original: Final = RuntimeError("health lookup failed")
+
+    async def health_after_corrupting_queue() -> tuple[object, ...]:
+        await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", {"invalid": "queue"})
+        match outcome:
+            case "failed":
+                raise original
+            case "cancelled":
+                raise asyncio.CancelledError
+            case "admitted":
+                return ("healthy",)
+
+    async def wait() -> None:
+        await scheduler.wait_for_turn(
+            FlowItem(priority=1, request_id="cleanup", model_name="shared"),
+            timeout=5,
+            get_healthy_deployments=health_after_corrupting_queue,
+        )
+
+    match outcome:
+        case "admitted":
+            await wait()
+        case "failed":
+            with pytest.raises(RuntimeError) as caught:
+                await wait()
+            assert caught.value is original
+        case "cancelled":
+            with pytest.raises(asyncio.CancelledError):
+                await wait()
+    assert "Scheduler queue cleanup failed" in caplog.text
+    assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == {"invalid": "queue"}
+
+
+async def test_shared_scheduler_handles_three_hundred_concurrent_enqueues_and_removals(
+    shared_scheduler_cache: RedisCache,
+) -> None:
+    schedulers: Final = tuple(Scheduler(redis_cache=shared_scheduler_cache) for _ in range(300))
+    for scheduler in schedulers:
+        scheduler.cache.update_cache_ttl(default_in_memory_ttl=60, default_redis_ttl=None)
+    requests: Final = tuple(
+        FlowItem(priority=i % 8, request_id=f"burst-{i:03}", model_name="shared") for i in range(300)
+    )
+
+    await asyncio.gather(*(scheduler.add_request(request) for scheduler, request in zip(schedulers, requests)))
+
+    assert await schedulers[0].get_queue("shared") == sorted(
+        (request.priority, request.request_id) for request in requests
+    )
+    await asyncio.gather(
+        *(scheduler.remove_request(request.request_id, "shared") for scheduler, request in zip(schedulers, requests))
+    )
+    assert await schedulers[0].get_queue("shared") == []