From e2b8bc001a38ff2858a1bbca50cb8af6cd20994d Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Tue, 20 Oct 2026 05:33:59 +0000 Subject: [PATCH] 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