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:
Oliver Yu 2026-10-29 09:43:58 +00:00
parent 250fc4e1cf
commit d499a032fc
2 changed files with 20 additions and 3 deletions

View file

@ -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,

View file

@ -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)