scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(least-busy): keep pick reservations in the shared router cache

Browse source 
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.
This commit is contained in:
oliver 2026-10-29 16:58:48 +08:00
parent 9d7be79659
commit 250fc4e1cf
2 changed files with 47 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/router_strategy/least_busy.py
View file

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

62
tests/unit/router_strategy/test_least_busy.py
View file

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