mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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 <tc3oliver@gmail.com>
This commit is contained in:
parent
250fc4e1cf
commit
d499a032fc
2 changed files with 20 additions and 3 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue