diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 871cff4d241..c8cca47f6ae 100644 --- a/litellm/scheduler.py +++ b/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. diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index e189a90c687..ae9bda04e94 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/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") == []