diff --git a/litellm/router.py b/litellm/router.py index c76e4ec843f..bea6b95e843 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2372,6 +2372,7 @@ class Router: return await selector.async_get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, + reserve=request_kwargs not in (None, {}), ) case "usage-based-routing": # `LowestTPMLoggingHandler` (v1) only exposes the sync @@ -2428,6 +2429,7 @@ class Router: return selector.get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, + reserve=request_kwargs not in (None, {}), ) case "usage-based-routing" | "usage-based-routing-v2": return selector.get_available_deployments( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 6f3e0936641..53dff6f2154 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,5 +1,7 @@ import logging -from collections.abc import Mapping, Sequence +import time +from collections.abc import Callable, Mapping, Sequence +from types import MappingProxyType from typing import Final from pydantic import TypeAdapter, ValidationError @@ -13,6 +15,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 +SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 class _ModelInfo(TypedDict, total=False): @@ -84,12 +87,12 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) -def _least_busy( - healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] -) -> Mapping[str, object] | None: - if not healthy_deployments: - return None - return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] +def _unexpired(times: Sequence[float], now: float) -> tuple[float, ...]: + return tuple(t for t in times if now - t < SELECTION_RESERVATION_TTL_SECONDS) + + +def _least_busy_index(counts: tuple[int, ...]) -> int: + return min(range(len(counts)), key=lambda index: counts[index]) def _warn_unreadable(model_group: str, error: Exception) -> None: @@ -116,12 +119,47 @@ class LeastBusyLoggingHandler(CustomLogger): logged_success: int = 0 logged_failure: int = 0 - def __init__(self, router_cache: DualCache): + def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.monotonic): 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) + + 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 _select( + self, + healthy_deployments: Sequence[Mapping[str, object]], + keys: tuple[str, ...], + counts: tuple[int, ...], + reserve: bool, + ) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + now: Final = self._clock() + 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) + return healthy_deployments[index] + + def _release_reservation(self, kwargs: Mapping[str, object]) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + now: Final = self._clock() + self._store_reservations(key, self._live_reservations(key, now)[1:], now) @with_service_target("router_usage") def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + self._release_reservation(kwargs) self._increment(kwargs, 1) @with_service_target("router_usage") @@ -158,7 +196,7 @@ class LeastBusyLoggingHandler(CustomLogger): @with_service_target("router_usage") def get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False ) -> Mapping[str, object] | None: keys: Final = _request_count_keys(model_group, healthy_deployments) redis_cache: Final = self.router_cache.redis_cache @@ -168,13 +206,13 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) + return self._select(healthy_deployments, keys, shared, reserve) local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) - return _least_busy(healthy_deployments, local) + return self._select(healthy_deployments, keys, local, reserve) @with_service_target("router_usage") async def async_get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False ) -> Mapping[str, object] | None: keys: Final = _request_count_keys(model_group, healthy_deployments) redis_cache: Final = self.router_cache.redis_cache @@ -184,9 +222,9 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) + return self._select(healthy_deployments, keys, shared, reserve) local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) - return _least_busy(healthy_deployments, local) + return self._select(healthy_deployments, keys, local, reserve) def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: ref: Final = _deployment_ref(kwargs) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 5b3af388ae5..ba3a20d2fea 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -14,7 +14,11 @@ from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import ( + IN_FLIGHT_COUNT_TTL_SECONDS, + SELECTION_RESERVATION_TTL_SECONDS, + LeastBusyLoggingHandler, +) from litellm.types.utils import ModelResponse GROUP: Final = "least-busy-group" @@ -391,3 +395,75 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts assert response._hidden_params["model_id"] == "1" assert route.calls.last.request.headers["authorization"] == "Bearer key-1" assert _router_counts(router) == seeded + + +@pytest.mark.asyncio +async def test_router_spreads_requests_that_start_together( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + async def slow_completion(_: httpx.Request) -> httpx.Response: + await asyncio.sleep(0.05) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-least-busy", + "object": "chat.completion", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=slow_completion) + 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)) + ) + ) + + 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) + ) + + +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]) + + 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 + + +def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> None: + worker: Final = _worker(None) + + 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")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + 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]) + + 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 len(worker._reservations) == 1 diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 7235333f59f..5c0df8e4ff4 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -1202,12 +1202,12 @@ async def test_least_busy_override_sees_the_overriding_request_in_flight(): stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) busy = stream._hidden_params["model_id"] idle = "d2" if busy == "d1" else "d1" - assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] + assert await _async_override_pick(router, "least-busy") == idle async for _ in stream: pass await _async_wait_until(lambda: _in_flight(router, busy) == 0) - assert await _async_override_pick(router, "least-busy") == "d1" + assert await _async_override_pick(router, "least-busy") == busy assert _selector_is_not_global(router._override_selectors["least-busy"]) @@ -1217,12 +1217,12 @@ def test_sync_least_busy_override_sees_the_overriding_request_in_flight(): stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) busy = stream._hidden_params["model_id"] idle = "d2" if busy == "d1" else "d1" - assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] + assert _sync_override_pick(router, "least-busy") == idle for _ in stream: pass _sync_wait_until(lambda: _in_flight(router, busy) == 0) - assert _sync_override_pick(router, "least-busy") == "d1" + assert _sync_override_pick(router, "least-busy") == busy assert _selector_is_not_global(router._override_selectors["least-busy"])