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)