fix(least-busy): keep reservations out of the shared router cache and count requests that start during a Redis read

Least-busy now reads a call's deployment from the routing metadata the router sets, so /v1/decisions logging params stay as on main and failed decisions calls no longer start cooldowns.

litellm/decisions/call.py

@@ -557,7 +557,6 @@ def prepare_call(
             litellm_params={
                 "litellm_call_id": kwargs.get("litellm_call_id"),
                 "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
-                **({} if (model_info := kwargs.get("model_info")) is None else {"model_info": model_info}),
             },
             custom_llm_provider=provider,
         )

litellm/router.py

@@ -200,7 +200,7 @@ from litellm.router_strategy.complexity_router.context_compaction import (
     reject_recursive_compactor,
     surface_for_call,
 )
-from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
+from litellm.router_strategy.least_busy import SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler
 from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
 from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
 from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler
@@ -1463,6 +1463,7 @@ class Router:
         self.cache = DualCache(
             redis_cache=redis_cache, in_memory_cache=InMemoryCache()
         )  # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc.
+        self._least_busy_reservations: InMemoryCache = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS)
         self._claude_code_session_router_cache: DualCache = DualCache(
             redis_cache=redis_cache,
             in_memory_cache=InMemoryCache(),
@@ -1967,7 +1968,9 @@ class Router:
         selector: RouterStrategySelector | None = None
         match self._normalize_strategy(strategy):
             case RoutingStrategy.LEAST_BUSY.value:
-                selector = LeastBusyLoggingHandler(router_cache=self.cache)
+                selector = LeastBusyLoggingHandler(
+                    router_cache=self.cache, reservation_store=self._least_busy_reservations
+                )
             case RoutingStrategy.USAGE_BASED_ROUTING.value:
                 selector = LowestTPMLoggingHandler(
                     router_cache=self.cache,

litellm/router_strategy/least_busy.py

@@ -9,12 +9,14 @@ from typing_extensions import ReadOnly, TypedDict
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.caching.redis_cache import log_redis_failure
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
 
 IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
 SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0
+SELECTION_RESERVATION_MAX_KEYS: Final = 10_000
 
 
 class _ModelInfo(TypedDict, total=False):
@@ -23,6 +25,7 @@ class _ModelInfo(TypedDict, total=False):
 
 class _Metadata(TypedDict, total=False):
     model_group: ReadOnly[str | None]
+    model_info: ReadOnly[_ModelInfo | None]
 
 
 class _LitellmParams(TypedDict, total=False):
@@ -57,6 +60,14 @@ def _model_group(metadata: _Metadata | None) -> str | None:
     return metadata.get("model_group") if metadata else None
 
 
+def _deployment_id(model_info: _ModelInfo | None) -> str | int | None:
+    return model_info.get("id") if model_info else None
+
+
+def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None:
+    return _deployment_id(metadata.get("model_info")) if metadata else None
+
+
 def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
     if is_batch_retrieve_call_type(kwargs.get("call_type")):
         return None
@@ -65,11 +76,19 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
     except ValidationError:
         return None
     litellm_params: Final = call.get("litellm_params") or _LitellmParams()
-    model_info: Final = litellm_params.get("model_info")
-    model_group: Final = _model_group(litellm_params.get("litellm_metadata")) or _model_group(
-        litellm_params.get("metadata")
+    routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata"))
+    model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None)
+    deployment_id: Final = next(
+        (
+            deployment_id
+            for deployment_id in (
+                _deployment_id(litellm_params.get("model_info")),
+                *map(_routed_deployment_id, routing_metadata),
+            )
+            if deployment_id is not None
+        ),
+        None,
     )
-    deployment_id: Final = model_info.get("id") if model_info else None
     if model_group is None or deployment_id is None:
         return None
     return model_group, str(deployment_id)
@@ -137,24 +156,30 @@ class LeastBusyLoggingHandler(CustomLogger):
     logged_success: int = 0
     logged_failure: int = 0
 
-    def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.monotonic):
+    def __init__(
+        self,
+        router_cache: DualCache,
+        clock: Callable[[], float] = time.monotonic,
+        reservation_store: _ReservationStore | None = None,
+    ):
         self.router_cache = router_cache
         self.router_cache_id = str(id(router_cache))
         self._clock = clock
-
-    def _reservation_store(self) -> _ReservationStore:
-        return self.router_cache.in_memory_cache
+        self._reservations: Final[_ReservationStore] = (
+            reservation_store
+            if reservation_store is not None
+            else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock)
+        )
 
     def _live_reservations(self, key: str, now: float) -> tuple[float, ...]:
-        raw: Final = self._reservation_store().get_cache(_reservation_key(key))
+        raw: Final = self._reservations.get_cache(_reservation_key(key))
         return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now)
 
     def _store_reservations(self, key: str, times: tuple[float, ...]) -> None:
         # InMemoryCache keeps a live key's first TTL, so drop the key to restart it
-        store: Final = self._reservation_store()
-        store.delete_cache(_reservation_key(key))
+        self._reservations.delete_cache(_reservation_key(key))
         if times:
-            store.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS)
+            self._reservations.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS)
 
     def _select(
         self,
@@ -235,9 +260,8 @@ class LeastBusyLoggingHandler(CustomLogger):
             except Exception as e:
                 _warn_unreadable(model_group, e)
             else:
-                return self._select(healthy_deployments, keys, shared, reserve)
-        local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
-        return self._select(healthy_deployments, keys, local, reserve)
+                return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve)
+        return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve)
 
     @with_service_target("router_usage")
     async def async_get_available_deployments(
@@ -251,9 +275,14 @@ class LeastBusyLoggingHandler(CustomLogger):
             except Exception as e:
                 _warn_unreadable(model_group, e)
             else:
-                return self._select(healthy_deployments, keys, shared, reserve)
-        local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys)
-        return self._select(healthy_deployments, keys, local, reserve)
+                return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve)
+        return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve)
+
+    def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]:
+        return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
+
+    def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]:
+        return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True))
 
     def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
         ref: Final = _deployment_ref(kwargs)

tests/unit/router_strategy/test_least_busy.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+from collections.abc import Callable
 from typing import Final
 
 import httpx
@@ -11,7 +12,7 @@ import litellm
 from litellm import Router
 
 from litellm.caching.caching import DualCache
-from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache
 from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.router_strategy.least_busy import (
@@ -62,6 +63,20 @@ class SharedRedisCounters:
         return self.batch_get_counts(key_list)
 
 
+class SlowRedisRead(SharedRedisCounters):
+    """Answers a count read with the counts from when the read started, after letting other work run."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.while_reading: Callable[[], None] = lambda: None
+
+    async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
+        counts_when_asked: Final = self.batch_get_counts(key_list)
+        await asyncio.sleep(0)
+        self.while_reading()
+        return counts_when_asked
+
+
 def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler:
     cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared)  # pyright: ignore[reportArgumentType]  # duck-typed Redis double
     return LeastBusyLoggingHandler(router_cache=cache)
@@ -85,6 +100,16 @@ async def test_worker_routes_around_a_request_another_worker_started() -> None:
     assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
 
 
+@pytest.mark.asyncio
+async def test_a_request_that_starts_while_redis_is_read_still_counts_on_its_deployment() -> None:
+    redis: Final = SlowRedisRead()
+    worker: Final = _worker(redis)
+    assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
+    redis.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
+
+    assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
+
+
 def test_sync_pick_reads_the_shared_counts() -> None:
     shared: Final = SharedRedisCounters()
     streaming_worker: Final = _worker(shared)
@@ -503,10 +528,11 @@ def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> N
     assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
 
 
-def test_a_request_started_through_another_selector_on_the_same_cache_releases_the_reservation() -> None:
+def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None:
     cache: Final = DualCache()
-    picker: Final = LeastBusyLoggingHandler(router_cache=cache)
-    registered: Final = LeastBusyLoggingHandler(router_cache=cache)
+    reservations: Final = InMemoryCache()
+    picker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations)
+    registered: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations)
 
     assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
     registered.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
@@ -515,6 +541,21 @@ def test_a_request_started_through_another_selector_on_the_same_cache_releases_t
     assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
 
 
+def _router_pick(router: Router) -> str:
+    deployment: Final = router.get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}})
+    return deployment["model_info"]["id"]
+
+
+def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None:
+    router: Final = _least_busy_router()
+
+    first: Final = _router_pick(router)
+    for index in range(DEFAULT_MAX_SIZE_IN_MEMORY):
+        router.cache.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60)
+
+    assert _router_pick(router) != first
+
+
 @pytest.mark.asyncio
 async def test_a_cache_hit_releases_the_pick_it_reserved() -> None:
     worker: Final = _worker(None)