Merge e2b8bc001a into bc0f598db5

litellm/router.py

@@ -2372,6 +2372,7 @@ class Router:
                 return await selector.async_get_available_deployments(
                     model_group=model,
                     healthy_deployments=healthy_deployments,
+                    reserve=request_kwargs not in (None, {}),
                 )
             case "usage-based-routing":
                 # `LowestTPMLoggingHandler` (v1) only exposes the sync
@@ -2428,6 +2429,7 @@ class Router:
                 return selector.get_available_deployments(
                     model_group=model,
                     healthy_deployments=healthy_deployments,
+                    reserve=request_kwargs not in (None, {}),
                 )
             case "usage-based-routing" | "usage-based-routing-v2":
                 return selector.get_available_deployments(

litellm/router_strategy/least_busy.py

@@ -1,5 +1,7 @@
 import logging
-from collections.abc import Mapping, Sequence
+import time
+from collections.abc import Callable, Mapping, Sequence
+from types import MappingProxyType
 from typing import Final
 
 from pydantic import TypeAdapter, ValidationError
@@ -13,6 +15,7 @@ from litellm.integrations.custom_logger import CustomLogger
 from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 
 IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
+SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0
 
 
 class _ModelInfo(TypedDict, total=False):
@@ -84,12 +87,12 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]:
     return _as_counts(values)
 
 
-def _least_busy(
-    healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...]
-) -> Mapping[str, object] | None:
-    if not healthy_deployments:
-        return None
-    return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])]
+def _unexpired(times: Sequence[float], now: float) -> tuple[float, ...]:
+    return tuple(t for t in times if now - t < SELECTION_RESERVATION_TTL_SECONDS)
+
+
+def _least_busy_index(counts: tuple[int, ...]) -> int:
+    return min(range(len(counts)), key=lambda index: counts[index])
 
 
 def _warn_unreadable(model_group: str, error: Exception) -> None:
@@ -116,12 +119,47 @@ class LeastBusyLoggingHandler(CustomLogger):
     logged_success: int = 0
     logged_failure: int = 0
 
-    def __init__(self, router_cache: DualCache):
+    def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.monotonic):
         self.router_cache = router_cache
         self.router_cache_id = str(id(router_cache))
+        self._clock = clock
+        self._reservations: Mapping[str, tuple[float, ...]] = MappingProxyType({})
+
+    def _live_reservations(self, key: str, now: float) -> tuple[float, ...]:
+        return _unexpired(self._reservations.get(key, ()), now)
+
+    def _store_reservations(self, key: str, times: tuple[float, ...], now: float) -> None:
+        # Drop expired and empty entries on every write so the map only holds the last TTL window.
+        merged: Final = {**self._reservations, key: times}
+        self._reservations = MappingProxyType({k: live for k, v in merged.items() if (live := _unexpired(v, now))})
+
+    def _select(
+        self,
+        healthy_deployments: Sequence[Mapping[str, object]],
+        keys: tuple[str, ...],
+        counts: tuple[int, ...],
+        reserve: bool,
+    ) -> Mapping[str, object] | None:
+        if not healthy_deployments:
+            return None
+        now: Final = self._clock()
+        live: Final = tuple(self._live_reservations(key, now) for key in keys)
+        index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True)))
+        if reserve:
+            self._store_reservations(keys[index], (*live[index], now), now)
+        return healthy_deployments[index]
+
+    def _release_reservation(self, kwargs: Mapping[str, object]) -> None:
+        ref: Final = _deployment_ref(kwargs)
+        if ref is None:
+            return
+        key: Final = _request_count_key(*ref)
+        now: Final = self._clock()
+        self._store_reservations(key, self._live_reservations(key, now)[1:], now)
 
     @with_service_target("router_usage")
     def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
+        self._release_reservation(kwargs)
         self._increment(kwargs, 1)
 
     @with_service_target("router_usage")
@@ -158,7 +196,7 @@ class LeastBusyLoggingHandler(CustomLogger):
 
     @with_service_target("router_usage")
     def get_available_deployments(
-        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
+        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False
     ) -> Mapping[str, object] | None:
         keys: Final = _request_count_keys(model_group, healthy_deployments)
         redis_cache: Final = self.router_cache.redis_cache
@@ -168,13 +206,13 @@ class LeastBusyLoggingHandler(CustomLogger):
             except Exception as e:
                 _warn_unreadable(model_group, e)
             else:
-                return _least_busy(healthy_deployments, shared)
+                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 _least_busy(healthy_deployments, local)
+        return self._select(healthy_deployments, keys, local, reserve)
 
     @with_service_target("router_usage")
     async def async_get_available_deployments(
-        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
+        self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False
     ) -> Mapping[str, object] | None:
         keys: Final = _request_count_keys(model_group, healthy_deployments)
         redis_cache: Final = self.router_cache.redis_cache
@@ -184,9 +222,9 @@ class LeastBusyLoggingHandler(CustomLogger):
             except Exception as e:
                 _warn_unreadable(model_group, e)
             else:
-                return _least_busy(healthy_deployments, shared)
+                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 _least_busy(healthy_deployments, local)
+        return self._select(healthy_deployments, keys, local, reserve)
 
     def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
         ref: Final = _deployment_ref(kwargs)

tests/unit/router_strategy/test_least_busy.py

@@ -14,7 +14,11 @@ from litellm.caching.caching import DualCache
 from litellm.caching.in_memory_cache import 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 IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler
+from litellm.router_strategy.least_busy import (
+    IN_FLIGHT_COUNT_TTL_SECONDS,
+    SELECTION_RESERVATION_TTL_SECONDS,
+    LeastBusyLoggingHandler,
+)
 from litellm.types.utils import ModelResponse
 
 GROUP: Final = "least-busy-group"
@@ -391,3 +395,75 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts
     assert response._hidden_params["model_id"] == "1"
     assert route.calls.last.request.headers["authorization"] == "Bearer key-1"
     assert _router_counts(router) == seeded
+
+
+@pytest.mark.asyncio
+async def test_router_spreads_requests_that_start_together(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+
+    async def slow_completion(_: httpx.Request) -> httpx.Response:
+        await asyncio.sleep(0.05)
+        return httpx.Response(
+            200,
+            json={
+                "id": "chatcmpl-least-busy",
+                "object": "chat.completion",
+                "created": 1,
+                "model": "gpt-4.1-mini",
+                "choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}],
+                "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
+            },
+        )
+
+    route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=slow_completion)
+    router: Final = _least_busy_router()
+
+    await asyncio.gather(
+        *(
+            router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}])
+            for _ in range(2 * len(ROUTER_DEPLOYMENT_IDS))
+        )
+    )
+
+    assert sorted(call.request.headers["authorization"] for call in route.calls) == sorted(
+        f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS for _ in range(2)
+    )
+
+
+def test_a_reservation_that_never_reaches_the_api_call_expires() -> None:
+    now: list[float] = [0.0]
+    worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=lambda: now[0])
+
+    picks: Final = []
+    for at in (0.0, 1.0, 2.0):
+        now[0] = at
+        picks.append(worker.get_available_deployments(GROUP, HEALTHY, reserve=True))
+    now[0] = SELECTION_RESERVATION_TTL_SECONDS + 0.5
+
+    assert picks == [DEPLOYMENT_A, DEPLOYMENT_B, DEPLOYMENT_A]
+    assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
+
+
+def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> None:
+    worker: Final = _worker(None)
+
+    assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
+    worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
+
+    assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
+    assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
+
+
+def test_expired_and_released_reservations_do_not_accumulate_across_model_groups() -> None:
+    now: list[float] = [0.0]
+    worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=lambda: now[0])
+
+    for i in range(50):
+        worker.get_available_deployments(f"group-{i}", HEALTHY, reserve=True)
+    worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
+    now[0] = SELECTION_RESERVATION_TTL_SECONDS + 0.5
+    worker.get_available_deployments("group-new", HEALTHY, reserve=True)
+
+    assert len(worker._reservations) == 1

tests/unit/router_strategy/test_router_routing_groups.py

@@ -1202,12 +1202,12 @@ async def test_least_busy_override_sees_the_overriding_request_in_flight():
     stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True))
     busy = stream._hidden_params["model_id"]
     idle = "d2" if busy == "d1" else "d1"
-    assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle]
+    assert await _async_override_pick(router, "least-busy") == idle
 
     async for _ in stream:
         pass
     await _async_wait_until(lambda: _in_flight(router, busy) == 0)
-    assert await _async_override_pick(router, "least-busy") == "d1"
+    assert await _async_override_pick(router, "least-busy") == busy
     assert _selector_is_not_global(router._override_selectors["least-busy"])
 
 
@@ -1217,12 +1217,12 @@ def test_sync_least_busy_override_sees_the_overriding_request_in_flight():
     stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True))
     busy = stream._hidden_params["model_id"]
     idle = "d2" if busy == "d1" else "d1"
-    assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle]
+    assert _sync_override_pick(router, "least-busy") == idle
 
     for _ in stream:
         pass
     _sync_wait_until(lambda: _in_flight(router, busy) == 0)
-    assert _sync_override_pick(router, "least-busy") == "d1"
+    assert _sync_override_pick(router, "least-busy") == busy
     assert _selector_is_not_global(router._override_selectors["least-busy"])