From 473ca526e0428aaa5bede3e4d2c1c6251b0f1dec Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Tue, 20 Oct 2026 01:56:07 +0000 Subject: [PATCH 1/2] fix(least-busy): reserve the picked deployment so a burst of requests spreads out The in-flight count only went up in log_pre_api_call, which runs after selection and the pre-call hooks. Requests selected in that window all saw the same counts and landed on the same deployment (16 simultaneous requests over 4 deployments went as lopsided as 12/4/0/0 in the proxy). When the router selects for a request it is about to send, least-busy now records a reservation for the chosen deployment in the same synchronous step that reads the reservations, and log_pre_api_call turns it into the normal in-flight count. Selections without request kwargs (token counting, probes) do not reserve. A reservation whose request never reaches the API call expires after 30 s. The shared and per-worker counts and their success/failure decrements are unchanged. --- litellm/router.py | 2 + litellm/router_strategy/least_busy.py | 57 ++++++++++++---- tests/unit/router_strategy/test_least_busy.py | 65 ++++++++++++++++++- .../test_router_routing_groups.py | 8 +-- 4 files changed, 113 insertions(+), 19 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d8da69f7789..a9165e80f81 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2199,6 +2199,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 @@ -2255,6 +2256,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( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 6f3e0936641..45abc00e1c5 100644 --- a/litellm/router_strategy/least_busy.py +++ b/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,8 @@ 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 _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 +115,42 @@ 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 tuple(t for t in self._reservations.get(key, ()) if now - t < SELECTION_RESERVATION_TTL_SECONDS) + + 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._reservations = MappingProxyType({**self._reservations, keys[index]: (*live[index], 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) + live: Final = self._live_reservations(key, self._clock()) + self._reservations = MappingProxyType({**self._reservations, key: live[1:]}) @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 +187,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 +197,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 +213,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) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 5b3af388ae5..f8a0af37f09 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/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,62 @@ 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 diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 7e4b470f1c1..ed90b846644 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/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"]) From e2b8bc001a38ff2858a1bbca50cb8af6cd20994d Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Tue, 20 Oct 2026 05:33:59 +0000 Subject: [PATCH 2/2] fix(least-busy): prune expired and empty selection reservations The reservation map kept a key per model group forever, so a wildcard route with caller-chosen model names grew it without bound. Drop expired and empty entries on every write. Signed-off-by: Oliver Yu --- litellm/router_strategy/least_busy.py | 17 +++++++++++++---- tests/unit/router_strategy/test_least_busy.py | 13 +++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 45abc00e1c5..53dff6f2154 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -87,6 +87,10 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) +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]) @@ -122,7 +126,12 @@ class LeastBusyLoggingHandler(CustomLogger): self._reservations: Mapping[str, tuple[float, ...]] = MappingProxyType({}) def _live_reservations(self, key: str, now: float) -> tuple[float, ...]: - return tuple(t for t in self._reservations.get(key, ()) if now - t < SELECTION_RESERVATION_TTL_SECONDS) + 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, @@ -137,7 +146,7 @@ class LeastBusyLoggingHandler(CustomLogger): 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._reservations = MappingProxyType({**self._reservations, keys[index]: (*live[index], now)}) + self._store_reservations(keys[index], (*live[index], now), now) return healthy_deployments[index] def _release_reservation(self, kwargs: Mapping[str, object]) -> None: @@ -145,8 +154,8 @@ class LeastBusyLoggingHandler(CustomLogger): if ref is None: return key: Final = _request_count_key(*ref) - live: Final = self._live_reservations(key, self._clock()) - self._reservations = MappingProxyType({**self._reservations, key: live[1:]}) + 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: diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index f8a0af37f09..ba3a20d2fea 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -454,3 +454,16 @@ 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_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