From 250fc4e1cfda5913dbec8ce98c5c3a3f12180242 Mon Sep 17 00:00:00 2001 From: oliver Date: Thu, 29 Oct 2026 16:58:48 +0800 Subject: [PATCH] fix(least-busy): keep pick reservations in the shared router cache Selectors that share a router cache dedupe to one registered callback, so only that instance saw log_pre_api_call. Reservations held on the other instances were never released and skewed picks for the full TTL. Storing them in the router cache's in-memory layer lets any selector release them, and the per-key TTL replaces the full-map rebuild on every write. Tests now hold the burst with a barrier instead of a sleep and use a manual clock object. --- litellm/router_strategy/least_busy.py | 22 ++++--- tests/unit/router_strategy/test_least_busy.py | 62 ++++++++++--------- 2 files changed, 47 insertions(+), 37 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 3c18ff5836c..1488d1b8bdf 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,7 +1,6 @@ import logging import time from collections.abc import Callable, Mapping, Sequence -from types import MappingProxyType from typing import Final from pydantic import TypeAdapter, ValidationError @@ -46,6 +45,7 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) _MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) +_RESERVATION_TIMES: Final = TypeAdapter(tuple[float, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: @@ -69,6 +69,10 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: return model_group, str(deployment_id) +def _reservation_key(request_count_key: str) -> str: + return f"{request_count_key}:reservations" + + def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: return tuple( _request_count_key(model_group, str(deployment["model_info"]["id"])) @@ -123,15 +127,15 @@ class LeastBusyLoggingHandler(CustomLogger): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) self._clock = clock - self._reservations: Mapping[str, tuple[float, ...]] = MappingProxyType({}) def _live_reservations(self, key: str, now: float) -> tuple[float, ...]: - return _unexpired(self._reservations.get(key, ()), now) + raw: Final = self.router_cache.get_cache(_reservation_key(key), local_only=True) + return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), 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 _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 + ) def _select( self, @@ -146,7 +150,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._store_reservations(keys[index], (*live[index], now), now) + self._store_reservations(keys[index], (*live[index], now)) return healthy_deployments[index] def _release_reservation(self, kwargs: Mapping[str, object]) -> None: @@ -155,7 +159,7 @@ class LeastBusyLoggingHandler(CustomLogger): return key: Final = _request_count_key(*ref) now: Final = self._clock() - self._store_reservations(key, self._live_reservations(key, now)[1:], now) + self._store_reservations(key, self._live_reservations(key, now)[1:]) @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 ba3a20d2fea..18fe9a8cbae 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -402,9 +402,11 @@ async def test_router_spreads_requests_that_start_together( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + burst: Final = 2 * len(ROUTER_DEPLOYMENT_IDS) + all_started: Final = asyncio.Barrier(burst) - async def slow_completion(_: httpx.Request) -> httpx.Response: - await asyncio.sleep(0.05) + async def completion_after_all_started(_: httpx.Request) -> httpx.Response: + await all_started.wait() return httpx.Response( 200, json={ @@ -417,32 +419,37 @@ async def test_router_spreads_requests_that_start_together( }, ) - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=slow_completion) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=completion_after_all_started + ) router: Final = _least_busy_router() await asyncio.gather( - *( - router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) - for _ in range(2 * len(ROUTER_DEPLOYMENT_IDS)) - ) + *(router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(burst)) ) - assert sorted(call.request.headers["authorization"] for call in route.calls) == sorted( - f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS for _ in range(2) - ) + one_per_deployment: Final = tuple(f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS) + assert sorted(call.request.headers["authorization"] for call in route.calls) == sorted(one_per_deployment * 2) + + +class _ManualClock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now def test_a_reservation_that_never_reaches_the_api_call_expires() -> None: - now: list[float] = [0.0] - worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=lambda: now[0]) + clock: Final = _ManualClock() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=clock) - picks: Final = [] - for at in (0.0, 1.0, 2.0): - now[0] = at - picks.append(worker.get_available_deployments(GROUP, HEALTHY, reserve=True)) - now[0] = SELECTION_RESERVATION_TTL_SECONDS + 0.5 - - assert picks == [DEPLOYMENT_A, DEPLOYMENT_B, DEPLOYMENT_A] + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = 1.0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + clock.now = 2.0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.5 assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A @@ -456,14 +463,13 @@ 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_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]) +def test_a_request_started_through_another_selector_on_the_same_cache_releases_the_reservation() -> None: + cache: Final = DualCache() + picker: Final = LeastBusyLoggingHandler(router_cache=cache) + registered: Final = LeastBusyLoggingHandler(router_cache=cache) - 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 picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + registered.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - assert len(worker._reservations) == 1 + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A