From d499a032fcfc3b8f62cd9af4b65967557d66c6a5 Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Thu, 29 Oct 2026 09:43:58 +0000 Subject: [PATCH] fix(least-busy): restart the reservation TTL on every write InMemoryCache keeps the TTL a live key was first given, so a reservation key that kept being rewritten still expired 30 s after its first pick and took the newer reservations with it. Drop the key before each write so the TTL counts from the latest one, and leave it deleted once it is empty. Signed-off-by: Oliver Yu --- litellm/router_strategy/least_busy.py | 8 +++++--- tests/unit/router_strategy/test_least_busy.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1488d1b8bdf..5d9a451770c 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -133,9 +133,11 @@ class LeastBusyLoggingHandler(CustomLogger): return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now) def _store_reservations(self, key: str, times: tuple[float, ...]) -> None: - self.router_cache.set_cache( - _reservation_key(key), times, local_only=True, ttl=SELECTION_RESERVATION_TTL_SECONDS - ) + # InMemoryCache keeps a live key's first TTL, so drop the key to restart it + local: Final = self.router_cache.in_memory_cache + local.delete_cache(_reservation_key(key)) + if times: + local.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) def _select( self, diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 18fe9a8cbae..6a95e825e6c 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -453,6 +453,21 @@ def test_a_reservation_that_never_reaches_the_api_call_expires() -> None: assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A +def test_a_fresh_reservation_outlives_the_ttl_of_the_first_one_on_its_deployment() -> None: + clock: Final = _ManualClock() + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + worker: Final = LeastBusyLoggingHandler(router_cache=cache, clock=clock) + + 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")) + worker.log_success_event(kwargs=_call_kwargs("dep-a"), response_obj=None, start_time=None, end_time=None) + clock.now = SELECTION_RESERVATION_TTL_SECONDS - 0.1 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.1 + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> None: worker: Final = _worker(None)