From 473ca526e0428aaa5bede3e4d2c1c6251b0f1dec Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Tue, 20 Oct 2026 01:56:07 +0000 Subject: [PATCH 01/18] fix(least-busy): reserve the picked deployment so a burst of requests spreads out The in-flight count only went up in log_pre_api_call, which runs after selection and the pre-call hooks. Requests selected in that window all saw the same counts and landed on the same deployment (16 simultaneous requests over 4 deployments went as lopsided as 12/4/0/0 in the proxy). When the router selects for a request it is about to send, least-busy now records a reservation for the chosen deployment in the same synchronous step that reads the reservations, and log_pre_api_call turns it into the normal in-flight count. Selections without request kwargs (token counting, probes) do not reserve. A reservation whose request never reaches the API call expires after 30 s. The shared and per-worker counts and their success/failure decrements are unchanged. --- litellm/router.py | 2 + litellm/router_strategy/least_busy.py | 57 ++++++++++++---- tests/unit/router_strategy/test_least_busy.py | 65 ++++++++++++++++++- .../test_router_routing_groups.py | 8 +-- 4 files changed, 113 insertions(+), 19 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d8da69f7789..a9165e80f81 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2199,6 +2199,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 @@ -2255,6 +2256,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..45abc00e1c5 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,8 @@ 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 _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 +115,42 @@ 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 tuple(t for t in self._reservations.get(key, ()) if now - t < SELECTION_RESERVATION_TTL_SECONDS) + + 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._reservations = MappingProxyType({**self._reservations, keys[index]: (*live[index], 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) + live: Final = self._live_reservations(key, self._clock()) + self._reservations = MappingProxyType({**self._reservations, key: live[1:]}) @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 +187,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 +197,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 +213,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..f8a0af37f09 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,62 @@ 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 diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 7e4b470f1c1..ed90b846644 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"]) From e2b8bc001a38ff2858a1bbca50cb8af6cd20994d Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Tue, 20 Oct 2026 05:33:59 +0000 Subject: [PATCH 02/18] fix(least-busy): prune expired and empty selection reservations The reservation map kept a key per model group forever, so a wildcard route with caller-chosen model names grew it without bound. Drop expired and empty entries on every write. Signed-off-by: Oliver Yu --- litellm/router_strategy/least_busy.py | 17 +++++++++++++---- tests/unit/router_strategy/test_least_busy.py | 13 +++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 45abc00e1c5..53dff6f2154 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -87,6 +87,10 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) +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]) @@ -122,7 +126,12 @@ class LeastBusyLoggingHandler(CustomLogger): self._reservations: Mapping[str, tuple[float, ...]] = MappingProxyType({}) def _live_reservations(self, key: str, now: float) -> tuple[float, ...]: - return tuple(t for t in self._reservations.get(key, ()) if now - t < SELECTION_RESERVATION_TTL_SECONDS) + 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, @@ -137,7 +146,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._reservations = MappingProxyType({**self._reservations, keys[index]: (*live[index], now)}) + self._store_reservations(keys[index], (*live[index], now), now) return healthy_deployments[index] def _release_reservation(self, kwargs: Mapping[str, object]) -> None: @@ -145,8 +154,8 @@ class LeastBusyLoggingHandler(CustomLogger): if ref is None: return key: Final = _request_count_key(*ref) - live: Final = self._live_reservations(key, self._clock()) - self._reservations = MappingProxyType({**self._reservations, key: live[1:]}) + 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: diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index f8a0af37f09..ba3a20d2fea 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -454,3 +454,16 @@ 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_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 From 250fc4e1cfda5913dbec8ce98c5c3a3f12180242 Mon Sep 17 00:00:00 2001 From: oliver Date: Thu, 29 Oct 2026 16:58:48 +0800 Subject: [PATCH 03/18] 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 From d499a032fcfc3b8f62cd9af4b65967557d66c6a5 Mon Sep 17 00:00:00 2001 From: Oliver Yu Date: Thu, 29 Oct 2026 09:43:58 +0000 Subject: [PATCH 04/18] 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 --- litellm/router_strategy/least_busy.py | 8 +++++--- tests/unit/router_strategy/test_least_busy.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1488d1b8bdf..5d9a451770c 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -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, diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 18fe9a8cbae..6a95e825e6c 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -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) From 6ad23ff70ddc2b757cad9ae6677b08a797efbb0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:09:08 -0700 Subject: [PATCH 05/18] fix(router): count least-busy decisions calls and release reservations on cache hits Least-busy now reads the model group from litellm_metadata before metadata, so /v1/decisions calls that carry caller metadata are counted in flight and release their reservation at pre-call. A response cache hit releases its reservation instead of decrementing another request's in-flight count, and Responses WebSocket turn admission selects without reserving. The reservation store is reached through a typed protocol. --- litellm/router.py | 11 ++- litellm/router_strategy/least_busy.py | 49 +++++++--- .../test_least_busy_concurrent_decisions.py | 90 +++++++++++++++++++ tests/unit/router_strategy/test_least_busy.py | 53 ++++++++++- .../test_router_routing_groups.py | 31 +++++++ .../test_router/test_router_helper_utils.py | 28 ++++++ 6 files changed, 243 insertions(+), 19 deletions(-) create mode 100644 tests/integration/routing/test_least_busy_concurrent_decisions.py diff --git a/litellm/router.py b/litellm/router.py index 43a5a6fddb7..62024bc9e74 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2431,6 +2431,7 @@ class Router: messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, + reserve: bool, ) -> Any | None: """ Asks the strategy selector for a deployment. Caller handles @@ -2444,7 +2445,7 @@ class Router: return await selector.async_get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, - reserve=request_kwargs not in (None, {}), + reserve=reserve, ) case "usage-based-routing": # `LowestTPMLoggingHandler` (v1) only exposes the sync @@ -2485,6 +2486,7 @@ class Router: messages: list[dict[str, str]] | None, input: str | list | None, request_kwargs: dict | None, + reserve: bool, ) -> Any | None: """ Sync sibling of `_select_deployment_async`. Caller handles @@ -2501,7 +2503,7 @@ class Router: return selector.get_available_deployments( model_group=model, healthy_deployments=healthy_deployments, - reserve=request_kwargs not in (None, {}), + reserve=reserve, ) case "usage-based-routing" | "usage-based-routing-v2": return selector.get_available_deployments( @@ -10257,6 +10259,7 @@ class Router: messages=None, input=None, request_kwargs=routing_kwargs, + reserve=False, ) if selected is None: no_deployment: Final = await async_raise_no_deployment_exception( @@ -15550,6 +15553,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + reserve=request_kwargs not in (None, {}), ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( @@ -15719,6 +15723,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + reserve=request_kwargs not in (None, {}), ) if deployment is None: @@ -16581,6 +16586,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + reserve=request_kwargs not in (None, {}), ) if deployment is None: @@ -16763,6 +16769,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + reserve=request_kwargs not in (None, {}), ) if deployment is None: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 5d9a451770c..29fb99ff63d 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,7 +1,7 @@ import logging import time from collections.abc import Callable, Mapping, Sequence -from typing import Final +from typing import Final, Protocol from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -27,6 +27,7 @@ class _Metadata(TypedDict, total=False): class _LitellmParams(TypedDict, total=False): metadata: ReadOnly[_Metadata | None] + litellm_metadata: ReadOnly[_Metadata | None] model_info: ReadOnly[_ModelInfo | None] @@ -44,14 +45,18 @@ 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) +_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None) +_RESERVATION_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: return f"{model_group}_request_count:{deployment_id}" +def _model_group(metadata: _Metadata | None) -> str | None: + return metadata.get("model_group") if metadata else None + + def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: if is_batch_retrieve_call_type(kwargs.get("call_type")): return None @@ -59,10 +64,11 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: call: Final = _CALL_KWARGS.validate_python(kwargs) except ValidationError: return None - litellm_params: Final = call.get("litellm_params") - metadata: Final = litellm_params.get("metadata") if litellm_params else None - model_info: Final = litellm_params.get("model_info") if litellm_params else None - model_group: Final = metadata.get("model_group") if metadata else None + litellm_params: Final = call.get("litellm_params") or _LitellmParams() + model_info: Final = litellm_params.get("model_info") + model_group: Final = _model_group(litellm_params.get("litellm_metadata")) or _model_group( + litellm_params.get("metadata") + ) deployment_id: Final = model_info.get("id") if model_info else None if model_group is None or deployment_id is None: return None @@ -118,6 +124,14 @@ def _warn_unwritable(key: str, error: Exception) -> None: ) +class _ReservationStore(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: tuple[float, ...], *, ttl: float) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -128,16 +142,19 @@ class LeastBusyLoggingHandler(CustomLogger): self.router_cache_id = str(id(router_cache)) self._clock = clock + def _reservation_store(self) -> _ReservationStore: + return self.router_cache.in_memory_cache + def _live_reservations(self, key: str, now: float) -> tuple[float, ...]: - raw: Final = self.router_cache.get_cache(_reservation_key(key), local_only=True) + raw: Final = self._reservation_store().get_cache(_reservation_key(key)) return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now) def _store_reservations(self, key: str, times: tuple[float, ...]) -> None: # 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)) + store: Final = self._reservation_store() + store.delete_cache(_reservation_key(key)) if times: - local.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) + store.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) def _select( self, @@ -172,7 +189,10 @@ class LeastBusyLoggingHandler(CustomLogger): def log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self._increment(kwargs, -1) + if kwargs.get("cache_hit") is True: + self._release_reservation(kwargs) + else: + self._increment(kwargs, -1) if self.test_flag: self.logged_success += 1 @@ -188,7 +208,10 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - await self._async_increment(kwargs, -1) + if kwargs.get("cache_hit") is True: + self._release_reservation(kwargs) + else: + await self._async_increment(kwargs, -1) if self.test_flag: self.logged_success += 1 diff --git a/tests/integration/routing/test_least_busy_concurrent_decisions.py b/tests/integration/routing/test_least_busy_concurrent_decisions.py new file mode 100644 index 00000000000..bb923534db4 --- /dev/null +++ b/tests/integration/routing/test_least_busy_concurrent_decisions.py @@ -0,0 +1,90 @@ +import json +import threading +from collections import Counter +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +pytestmark = pytest.mark.timeout(240) + +_GROUP: Final = "least-busy-decider" +_DEPLOYMENTS: Final = ("least-busy-decider-a", "least-busy-decider-b") +_BURST: Final = 4 +_REPLY: Final = json.dumps( + { + "model": "Qwen/Qwen3-0.6B", + "answers": {"billing": {"type": "noul", "noul": 0.91}}, + "usage": {"input_tokens": 367, "output_tokens": 3}, + } +).encode() +_DECISION: Final = { + "model": _GROUP, + "input": "I was charged twice for order 4417.", + "questions": [{"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}], +} + + +def _held_until_the_burst_arrives(arrived: threading.Barrier) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + arrived.wait() + return Reply(body=_REPLY) + + return respond + + +def _write_config(path: Path, provider_url: str) -> Path: + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": _GROUP, + "litellm_params": {"model": "hosted_vllm/Qwen/Qwen3-0.6B", "api_base": provider_url}, + "model_info": {"id": deployment_id}, + } + for deployment_id in _DEPLOYMENTS + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_model_info_refresh": True, + }, + "router_settings": {"routing_strategy": "least-busy", "num_retries": 0, "disable_cooldowns": True}, + } + ) + ) + return path + + +def _decide(proxy: Gateway) -> httpx.Response: + return proxy.client.post( + "/v1/decisions", json=_DECISION, headers={"Authorization": f"Bearer {proxy.key}"}, timeout=120 + ) + + +def test_least_busy_spreads_a_burst_of_decisions_calls_across_the_group(tmp_path: Path) -> None: + arrived: Final = threading.Barrier(_BURST, timeout=60) + with ( + wire_server(_held_until_the_burst_arrives(arrived)) as upstream, + gateway_from_environment() as rig, + owned_proxy( + rig, + tmp_path, + {"STORE_MODEL_IN_DB": "False"}, + config=_write_config(tmp_path / "least-busy.yaml", upstream.url), + ) as proxy, + ThreadPoolExecutor(max_workers=_BURST) as pool, + ): + responses: Final = tuple(pool.map(lambda _: _decide(proxy), range(_BURST))) + assert [response.status_code for response in responses] == [200] * _BURST, [r.text for r in responses] + assert Counter(response.headers["x-litellm-model-id"] for response in responses) == Counter( + {deployment_id: _BURST // len(_DEPLOYMENTS) for deployment_id in _DEPLOYMENTS} + ) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 6a95e825e6c..46964bdb9f8 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -102,6 +102,31 @@ def test_sync_pick_reads_the_shared_counts() -> None: assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]: + return { + "litellm_params": { + "metadata": {"tags": ["caller-tag"]}, + "litellm_metadata": {"model_group": GROUP}, + "model_info": {"id": deployment_id}, + } + } + + +def test_a_request_routed_through_litellm_metadata_counts_while_in_flight() -> None: + shared: Final = SharedRedisCounters() + calling_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + assert calling_worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + calling_worker.log_pre_api_call(model="m", messages=[], kwargs=_routed_through_litellm_metadata("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + calling_worker.log_success_event(_routed_through_litellm_metadata("dep-a"), None, None, None) + + assert calling_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that stuck count an hour after the key was created, which only works while nothing writes the TTL @@ -211,7 +236,9 @@ class OpenBreakerRedis(SharedRedisCounters): @pytest.mark.asyncio -async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None: +async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( + caplog: pytest.LogCaptureFixture, +) -> None: worker: Final = _worker(OpenBreakerRedis()) with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): @@ -383,9 +410,7 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts if use_async else router.get_available_deployment(model=ROUTER_GROUP, messages=None) ) - response: Final = await router.acompletion( - model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}] - ) + response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(10): await asyncio.sleep(0) await GLOBAL_LOGGING_WORKER.flush() @@ -488,3 +513,23 @@ def test_a_request_started_through_another_selector_on_the_same_cache_releases_t assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +@pytest.mark.asyncio +async def test_a_cache_hit_releases_the_pick_it_reserved() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + await worker.async_log_success_event({**_call_kwargs("dep-b"), "cache_hit": True}, None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + worker.log_success_event({**_call_kwargs("dep-a"), "cache_hit": True}, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 5c0df8e4ff4..c43662a84e2 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -481,6 +481,33 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + +def _filtered_model_pick(router: Router) -> str: + deployment: Final = router.get_available_deployment( + "filtered-model", request_kwargs=_proxy_shaped_request(model="filtered-model") + ) + return deployment["model_info"]["id"] + + +def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = _build_router( + routing_strategy="least-busy", + routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}], + ) + + first: Final = _filtered_model_pick(router) + kwargs: Final = {"litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}} + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert _filtered_model_pick(router) == first + def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): """ Least-busy hangs its counting off litellm's global callback lists, and those lists keep one @@ -613,6 +640,7 @@ async def test_select_deployment_async_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -627,6 +655,7 @@ def test_select_deployment_sync_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -661,6 +690,7 @@ def test_select_deployment_sync_does_not_dispatch_cost_based_routing(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -707,6 +737,7 @@ async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1(): messages=[{"role": "user", "content": "hi"}], input=None, request_kwargs={}, + reserve=False, ) assert v1_spy.called, "async dispatch must route v1 strategy through sync method" diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 513136ae150..f9eff9b5ae9 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -597,6 +597,34 @@ async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_b await admit_turn() +@pytest.mark.asyncio +async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy(): + router = Router( + model_list=[ + { + "model_name": "ws-least-busy", + "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("ws-least-busy-a", "ws-least-busy-b") + ], + routing_strategy="least-busy", + ) + admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"] + assert callable(admit_turn) + for _ in range(3): + await admit_turn() + + picks: Final = [ + (await router.async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ + "model_info" + ]["id"] + for _ in range(2) + ] + + assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"] + + @pytest.mark.asyncio async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): router = Router( From d9357038e8678d69a31f4bd3ddd0258b9427c35b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:53:31 -0700 Subject: [PATCH 06/18] fix(decisions): count proxy /v1/decisions calls in least-busy's in-flight totals --- litellm/decisions/call.py | 1 + tests/unit/decisions/test_main.py | 39 +++++++++++++++++++ tests/unit/router_strategy/test_least_busy.py | 8 ++-- 3 files changed, 44 insertions(+), 4 deletions(-) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 3b73f65c917..415d5d546e9 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -557,6 +557,7 @@ def prepare_call( litellm_params={ "litellm_call_id": kwargs.get("litellm_call_id"), "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model), + **({} if (model_info := kwargs.get("model_info")) is None else {"model_info": model_info}), }, custom_llm_provider=provider, ) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index e5501310556..a8c2d4e7805 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -6,6 +6,8 @@ import copy import json import re import threading +import uuid +from datetime import datetime from collections.abc import Mapping, Sequence from concurrent.futures import wait from types import MappingProxyType @@ -296,6 +298,43 @@ async def test_router_dispatches_typesafe_decisions_without_api_base( assert response.answers["sentiment"].choice == "positive" +@pytest.mark.asyncio +async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_finishes( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "jev", + "litellm_params": {"model": "typesafe/jev-latest", "api_key": "k"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("jev-a", "jev-b") + ], + routing_strategy="least-busy", + ) + + async def call_through_proxy_logging() -> object: + logging_obj, _ = litellm.utils.function_setup( + original_function="adecisions", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="jev", + litellm_call_id=str(uuid.uuid4()), + ) + response: Final = await router.adecisions( + model="jev", + input="Is this a defect?", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + litellm_logging_obj=logging_obj, + ) + await _drain_logging_worker() + return response._hidden_params["model_id"] + + assert (await call_through_proxy_logging(), await call_through_proxy_logging()) == ("jev-a", "jev-a") + + def test_decisions_uses_the_same_wire_contract_for_sync_calls(respx_mock: respx.MockRouter) -> None: route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 46964bdb9f8..3dbd90b814b 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -236,9 +236,7 @@ class OpenBreakerRedis(SharedRedisCounters): @pytest.mark.asyncio -async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( - caplog: pytest.LogCaptureFixture, -) -> None: +async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None: worker: Final = _worker(OpenBreakerRedis()) with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): @@ -410,7 +408,9 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts if use_async else router.get_available_deployment(model=ROUTER_GROUP, messages=None) ) - response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) + response: Final = await router.acompletion( + model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}] + ) for _ in range(10): await asyncio.sleep(0) await GLOBAL_LOGGING_WORKER.flush() From b7330ea77fd0faa22dbf846da3feda6790c8c11a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 11:36:14 -0700 Subject: [PATCH 07/18] fix(least-busy): keep reservations out of the shared router cache and count requests that start during a Redis read Least-busy now reads a call's deployment from the routing metadata the router sets, so /v1/decisions logging params stay as on main and failed decisions calls no longer start cooldowns. --- litellm/decisions/call.py | 1 - litellm/router.py | 7 +- litellm/router_strategy/least_busy.py | 65 ++++++++++++++----- tests/unit/router_strategy/test_least_busy.py | 49 ++++++++++++-- 4 files changed, 97 insertions(+), 25 deletions(-) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 415d5d546e9..3b73f65c917 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -557,7 +557,6 @@ def prepare_call( litellm_params={ "litellm_call_id": kwargs.get("litellm_call_id"), "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model), - **({} if (model_info := kwargs.get("model_info")) is None else {"model_info": model_info}), }, custom_llm_provider=provider, ) diff --git a/litellm/router.py b/litellm/router.py index 62024bc9e74..b34b621c189 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -200,7 +200,7 @@ from litellm.router_strategy.complexity_router.context_compaction import ( reject_recursive_compactor, surface_for_call, ) -from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler @@ -1463,6 +1463,7 @@ class Router: self.cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._least_busy_reservations: InMemoryCache = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS) self._claude_code_session_router_cache: DualCache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache(), @@ -1967,7 +1968,9 @@ class Router: selector: RouterStrategySelector | None = None match self._normalize_strategy(strategy): case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler(router_cache=self.cache) + selector = LeastBusyLoggingHandler( + router_cache=self.cache, reservation_store=self._least_busy_reservations + ) case RoutingStrategy.USAGE_BASED_ROUTING.value: selector = LowestTPMLoggingHandler( router_cache=self.cache, diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 29fb99ff63d..b69620aea23 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -9,12 +9,14 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 +SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 class _ModelInfo(TypedDict, total=False): @@ -23,6 +25,7 @@ class _ModelInfo(TypedDict, total=False): class _Metadata(TypedDict, total=False): model_group: ReadOnly[str | None] + model_info: ReadOnly[_ModelInfo | None] class _LitellmParams(TypedDict, total=False): @@ -57,6 +60,14 @@ def _model_group(metadata: _Metadata | None) -> str | None: return metadata.get("model_group") if metadata else None +def _deployment_id(model_info: _ModelInfo | None) -> str | int | None: + return model_info.get("id") if model_info else None + + +def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None: + return _deployment_id(metadata.get("model_info")) if metadata else None + + def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: if is_batch_retrieve_call_type(kwargs.get("call_type")): return None @@ -65,11 +76,19 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: except ValidationError: return None litellm_params: Final = call.get("litellm_params") or _LitellmParams() - model_info: Final = litellm_params.get("model_info") - model_group: Final = _model_group(litellm_params.get("litellm_metadata")) or _model_group( - litellm_params.get("metadata") + routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata")) + model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None) + deployment_id: Final = next( + ( + deployment_id + for deployment_id in ( + _deployment_id(litellm_params.get("model_info")), + *map(_routed_deployment_id, routing_metadata), + ) + if deployment_id is not None + ), + None, ) - deployment_id: Final = model_info.get("id") if model_info else None if model_group is None or deployment_id is None: return None return model_group, str(deployment_id) @@ -137,24 +156,30 @@ class LeastBusyLoggingHandler(CustomLogger): logged_success: int = 0 logged_failure: int = 0 - def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.monotonic): + def __init__( + self, + router_cache: DualCache, + clock: Callable[[], float] = time.monotonic, + reservation_store: _ReservationStore | None = None, + ): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) self._clock = clock - - def _reservation_store(self) -> _ReservationStore: - return self.router_cache.in_memory_cache + self._reservations: Final[_ReservationStore] = ( + reservation_store + if reservation_store is not None + else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock) + ) def _live_reservations(self, key: str, now: float) -> tuple[float, ...]: - raw: Final = self._reservation_store().get_cache(_reservation_key(key)) + raw: Final = self._reservations.get_cache(_reservation_key(key)) return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now) def _store_reservations(self, key: str, times: tuple[float, ...]) -> None: # InMemoryCache keeps a live key's first TTL, so drop the key to restart it - store: Final = self._reservation_store() - store.delete_cache(_reservation_key(key)) + self._reservations.delete_cache(_reservation_key(key)) if times: - store.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) + self._reservations.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) def _select( self, @@ -235,9 +260,8 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - 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 self._select(healthy_deployments, keys, local, reserve) + return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve) + return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve) @with_service_target("router_usage") async def async_get_available_deployments( @@ -251,9 +275,14 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - 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 self._select(healthy_deployments, keys, local, reserve) + return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve) + return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve) + + def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]: + return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + + def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]: + return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True)) 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 3dbd90b814b..177a3968f76 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,6 +1,7 @@ import asyncio import json import logging +from collections.abc import Callable from typing import Final import httpx @@ -11,7 +12,7 @@ import litellm from litellm import Router from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, 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 ( @@ -62,6 +63,20 @@ class SharedRedisCounters: return self.batch_get_counts(key_list) +class SlowRedisRead(SharedRedisCounters): + """Answers a count read with the counts from when the read started, after letting other work run.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + counts_when_asked: Final = self.batch_get_counts(key_list) + await asyncio.sleep(0) + self.while_reading() + return counts_when_asked + + def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double return LeastBusyLoggingHandler(router_cache=cache) @@ -85,6 +100,16 @@ async def test_worker_routes_around_a_request_another_worker_started() -> None: assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +@pytest.mark.asyncio +async def test_a_request_that_starts_while_redis_is_read_still_counts_on_its_deployment() -> None: + redis: Final = SlowRedisRead() + worker: Final = _worker(redis) + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + redis.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + def test_sync_pick_reads_the_shared_counts() -> None: shared: Final = SharedRedisCounters() streaming_worker: Final = _worker(shared) @@ -503,10 +528,11 @@ 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_a_request_started_through_another_selector_on_the_same_cache_releases_the_reservation() -> None: +def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None: cache: Final = DualCache() - picker: Final = LeastBusyLoggingHandler(router_cache=cache) - registered: Final = LeastBusyLoggingHandler(router_cache=cache) + reservations: Final = InMemoryCache() + picker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + registered: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) 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")) @@ -515,6 +541,21 @@ def test_a_request_started_through_another_selector_on_the_same_cache_releases_t assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A +def _router_pick(router: Router) -> str: + deployment: Final = router.get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}) + return deployment["model_info"]["id"] + + +def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None: + router: Final = _least_busy_router() + + first: Final = _router_pick(router) + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): + router.cache.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60) + + assert _router_pick(router) != first + + @pytest.mark.asyncio async def test_a_cache_hit_releases_the_pick_it_reserved() -> None: worker: Final = _worker(None) From b6e27dd5dc5da63b98842b2145a1570a91008c2d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:23:20 -0700 Subject: [PATCH 08/18] style(router): wrap the least-busy reservation store line --- litellm/router.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index c4934d0fa9e..b3a9524fd0e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1656,7 +1656,9 @@ class Router(DeprecatedRouterMembers): self._core_cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. - self._core_least_busy_reservations: InMemoryCache = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS) + self._core_least_busy_reservations: InMemoryCache = InMemoryCache( + max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS + ) self._claude_code_session_router_cache: DualCache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache(), From 854231f823408c17a66c92e363b5288997882d67 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 13:59:46 -0700 Subject: [PATCH 09/18] fix(least-busy): lock the reservation read-and-update so threaded sync calls keep each other's picks Router.completion from a thread pool and sync success callbacks can read and rewrite a deployment's reservations at the same time, so one thread's pick or release could be lost. --- litellm/router_strategy/least_busy.py | 18 +++--- tests/unit/router_strategy/test_least_busy.py | 60 ++++++++++++++++++- 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index b69620aea23..1efd4a6e7a5 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,4 +1,5 @@ import logging +import threading import time from collections.abc import Callable, Mapping, Sequence from typing import Final, Protocol @@ -17,6 +18,7 @@ from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 +_RESERVATIONS_LOCK: Final = threading.Lock() class _ModelInfo(TypedDict, total=False): @@ -190,11 +192,12 @@ class LeastBusyLoggingHandler(CustomLogger): ) -> 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)) + with _RESERVATIONS_LOCK: + 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)) return healthy_deployments[index] def _release_reservation(self, kwargs: Mapping[str, object]) -> None: @@ -202,8 +205,9 @@ class LeastBusyLoggingHandler(CustomLogger): 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:]) + with _RESERVATIONS_LOCK: + now: Final = self._clock() + 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 63614e384ff..a0be7741aa9 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,7 +1,10 @@ import asyncio +import contextlib import json import logging -from collections.abc import Callable +import threading +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx @@ -541,6 +544,61 @@ def test_a_request_started_through_another_selector_of_the_router_releases_the_r assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A +class _StoreWhereThreadsReadBeforeWriting: + def __init__(self) -> None: + self._store: Final = InMemoryCache() + self._first_reads: Final = threading.Barrier(2, timeout=0.5) + self._thread_state: Final = threading.local() + self._meeting: Final = threading.Event() + + @contextlib.contextmanager + def first_reads_meet(self) -> Iterator[None]: + self._meeting.set() + yield + self._meeting.clear() + + def get_cache(self, key: str) -> object: + value: Final = self._store.get_cache(key) + if self._meeting.is_set() and not getattr(self._thread_state, "has_read", False): + self._thread_state.has_read = True + with contextlib.suppress(threading.BrokenBarrierError): + self._first_reads.wait() + return value + + def set_cache(self, key: str, value: tuple[float, ...], *, ttl: float) -> None: + self._store.set_cache(key, value, ttl=ttl) + + def delete_cache(self, key: str) -> None: + self._store.delete_cache(key) + + +def test_two_threads_picking_at_once_reserve_different_deployments() -> None: + store: Final = _StoreWhereThreadsReadBeforeWriting() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) + + with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: + picks: Final = tuple(pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2))) + + assert sorted(pick["model_info"]["id"] for pick in picks if pick is not None) == ["dep-a", "dep-b"] + + +def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> None: + store: Final = _StoreWhereThreadsReadBeforeWriting() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b")) + + with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: + reserved: Final = pool.submit(worker.get_available_deployments, GROUP, HEALTHY, True) + started: Final = pool.submit(worker.log_pre_api_call, "m", [], _call_kwargs("dep-a")) + assert reserved.result() is DEPLOYMENT_A + started.result() + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + def _router_pick(router: Router) -> str: deployment: Final = router.get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}) return deployment["model_info"]["id"] From 67bc8093939b9dfe3cb2a50648769ff9d8861b3c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:02:26 -0700 Subject: [PATCH 10/18] fix(least-busy): release a pick whose call ends before reaching the provider A call that fails or returns before the provider handoff (a 400 for an unsupported param, token counting, a proxy-gate rejection without a start log, batch create) now drops its reservation instead of decrementing a count it never added. Router.select picks and batch cancel no longer reserve, since neither makes a call that logs. --- litellm/router.py | 7 +- litellm/router_strategy/least_busy.py | 33 +++++-- tests/unit/router_strategy/test_least_busy.py | 98 ++++++++++++++++++- 3 files changed, 125 insertions(+), 13 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b3a9524fd0e..8c095360301 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1387,6 +1387,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, + reserve=False, ) if kind == "healthy": return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts @@ -8063,6 +8064,7 @@ class Router(DeprecatedRouterMembers): messages=[{"role": "user", "content": "batch-api-fake-text"}], specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, + reserve=False, ) selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id") @@ -15695,6 +15697,7 @@ class Router(DeprecatedRouterMembers): messages: list[dict[str, str]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, + reserve: bool = True, ): """ Async implementation of 'get_available_deployments'. @@ -15708,6 +15711,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, + reserve=reserve, ) async def _async_get_available_deployment( @@ -15717,6 +15721,7 @@ class Router(DeprecatedRouterMembers): messages: list[dict[str, str]] | None, input: str | list | None, specific_deployment: bool | None, + reserve: bool, ): if ( self._core_routing_strategy != "usage-based-routing-v2" @@ -15813,7 +15818,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, - reserve=request_kwargs not in (None, {}), + reserve=reserve and request_kwargs not in (None, {}), ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1efd4a6e7a5..f920e427d41 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -12,6 +12,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import log_redis_failure +from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type @@ -96,6 +97,12 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: return model_group, str(deployment_id) +def _counted_in_flight(kwargs: Mapping[str, object]) -> bool: + return kwargs.get("cache_hit") is not True and ( + kwargs.get("api_call_start_time") is not None or LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL in kwargs + ) + + def _reservation_key(request_count_key: str) -> str: return f"{request_count_key}:reservations" @@ -218,10 +225,7 @@ class LeastBusyLoggingHandler(CustomLogger): def log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - if kwargs.get("cache_hit") is True: - self._release_reservation(kwargs) - else: - self._increment(kwargs, -1) + self._settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -229,7 +233,7 @@ class LeastBusyLoggingHandler(CustomLogger): def log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - self._increment(kwargs, -1) + self._settle(kwargs) if self.test_flag: self.logged_failure += 1 @@ -237,10 +241,7 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - if kwargs.get("cache_hit") is True: - self._release_reservation(kwargs) - else: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -248,7 +249,7 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_failure += 1 @@ -282,6 +283,18 @@ class LeastBusyLoggingHandler(CustomLogger): return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve) return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve) + def _settle(self, kwargs: Mapping[str, object]) -> None: + if not _counted_in_flight(kwargs): + self._release_reservation(kwargs) + return + self._increment(kwargs, -1) + + async def _async_settle(self, kwargs: Mapping[str, object]) -> None: + if not _counted_in_flight(kwargs): + self._release_reservation(kwargs) + return + await self._async_increment(kwargs, -1) + def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]: return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index a0be7741aa9..4b7cbe253bd 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import datetime import json import logging import threading @@ -8,6 +9,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx +import openai import pytest import respx @@ -17,6 +19,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router_strategy.least_busy import ( IN_FLIGHT_COUNT_TTL_SECONDS, @@ -31,10 +34,14 @@ DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] -def _call_kwargs(deployment_id: str) -> dict[str, object]: +def _routed_call_kwargs(deployment_id: str) -> dict[str, object]: return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} +def _call_kwargs(deployment_id: str) -> dict[str, object]: + return {**_routed_call_kwargs(deployment_id), "api_call_start_time": datetime.datetime.now()} + + class SharedRedisCounters: """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when the key is created, and ordered reads that raise rather than invent a value.""" @@ -136,7 +143,8 @@ def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]: "metadata": {"tags": ["caller-tag"]}, "litellm_metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}, - } + }, + "api_call_start_time": datetime.datetime.now(), } @@ -632,3 +640,89 @@ def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> No worker.log_success_event({**_call_kwargs("dep-a"), "cache_hit": True}, None, None, None) assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.parametrize("outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event]) +def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick( + outcome: Callable[[LeastBusyLoggingHandler, dict[str, object], None, None, None], None], +) -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + outcome(worker, _routed_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behind() -> None: + worker: Final = _worker(None) + rejected: Final = {**_routed_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True} + + worker.log_pre_api_call(model="m", messages=[], kwargs=rejected) + worker.log_failure_event(rejected, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +async def _drain_logging() -> None: + for _ in range(10): + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() + + +def _batch_deployment(route: respx.Route) -> str: + return route.calls.last.request.headers["authorization"].removeprefix("Bearer key-") + + +_BATCH: Final = { + "id": "batch_least_busy", + "object": "batch", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-least-busy", + "completion_window": "24h", + "status": "validating", + "created_at": 1, +} +_REJECTED_BATCH: Final = {"error": {"message": "Invalid input file", "type": "invalid_request_error"}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)]) +async def test_a_batch_create_does_not_keep_its_deployment_busy( + provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/batches").mock(return_value=provider_response) + router: Final = _least_busy_router() + + with contextlib.suppress(openai.BadRequestError): + await router.acreate_batch( + model=ROUTER_GROUP, completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-least-busy" + ) + await _drain_logging() + + assert _router_pick(router) == _batch_deployment(route) + + +@pytest.mark.asyncio +async def test_a_batch_cancel_does_not_keep_its_deployment_busy( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/batches/batch_least_busy/cancel").mock( + return_value=httpx.Response(200, json={**_BATCH, "status": "cancelling"}) + ) + router: Final = _least_busy_router() + + await router.acancel_batch(model=ROUTER_GROUP, batch_id="batch_least_busy") + + assert _router_pick(router) == _batch_deployment(route) + + +@pytest.mark.asyncio +async def test_picking_a_deployment_without_calling_it_does_not_keep_it_busy() -> None: + router: Final = _least_busy_router() + + picked: Final = await router.select(ROUTER_GROUP, {"metadata": {}}) + + assert _router_pick(router) == picked["model_info"]["id"] From 15f8086684a62d6d6f92e65c0eee2080fe6219c5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:40:39 -0700 Subject: [PATCH 11/18] test(least-busy): give callback fixtures a fixed call-start time The routing-group tests fed least-busy a finished call without the start time litellm stamps before counting it, so the count stayed up after the fix. Fixtures now use a fixed timestamp instead of the clock. --- tests/unit/router_strategy/test_least_busy.py | 5 +++-- .../router_strategy/test_router_routing_groups.py | 14 +++++++++++--- tests/unit/test_router/test_router_helper_utils.py | 2 +- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 4b7cbe253bd..2eee767aa06 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -29,6 +29,7 @@ from litellm.router_strategy.least_busy import ( from litellm.types.utils import ModelResponse GROUP: Final = "least-busy-group" +CALL_STARTED_AT: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] @@ -39,7 +40,7 @@ def _routed_call_kwargs(deployment_id: str) -> dict[str, object]: def _call_kwargs(deployment_id: str) -> dict[str, object]: - return {**_routed_call_kwargs(deployment_id), "api_call_start_time": datetime.datetime.now()} + return {**_routed_call_kwargs(deployment_id), "api_call_start_time": CALL_STARTED_AT} class SharedRedisCounters: @@ -144,7 +145,7 @@ def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]: "litellm_metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}, }, - "api_call_start_time": datetime.datetime.now(), + "api_call_start_time": CALL_STARTED_AT, } diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index c43662a84e2..70be27c2645 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -56,6 +56,9 @@ def _model_list(): ] +CALL_STARTED_AT: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + + def _build_router(routing_strategy="simple-shuffle", routing_groups=None): return Router( model_list=_model_list(), @@ -468,7 +471,8 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): "litellm_params": { "metadata": {"model_group": "filtered-model"}, "model_info": {"id": "deploy-1"}, - } + }, + "api_call_start_time": CALL_STARTED_AT, } for callback in litellm.input_callback: @@ -498,7 +502,10 @@ def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkey ) first: Final = _filtered_model_pick(router) - kwargs: Final = {"litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}} + kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}, + "api_call_start_time": CALL_STARTED_AT, + } for callback in litellm.input_callback: if isinstance(callback, CustomLogger): callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) @@ -525,7 +532,8 @@ def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): "litellm_params": { "metadata": {"model_group": "filtered-model"}, "model_info": {"id": "deploy-1"}, - } + }, + "api_call_start_time": CALL_STARTED_AT, } for callback in litellm.input_callback: diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 84e495d1be6..5cc62ce6450 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -599,7 +599,7 @@ async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_b @pytest.mark.asyncio async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy(): - router = Router( + router: Final = Router( model_list=[ { "model_name": "ws-least-busy", From 2c4bae982bf666ec185f962f9537336441570484 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:30:42 -0700 Subject: [PATCH 12/18] fix(least-busy): free tier-group picks, stop lookups reserving, and count only started attempts Reservations are keyed by deployment id, so a complexity, auto, adaptive or quality router frees its tier pick when the call starts. Each started attempt is recorded under its call id, and only a recorded start lowers the count when the attempt ends, so a fallback hop rejected before the provider frees its pick. Whether a pick reserves now comes from a context variable set by the routed call sites, so get_available_deployment and the deprecated async_get_available_deployment no longer hold a deployment, and custom routing strategies and Router subclasses keep their method signatures. --- litellm/router.py | 78 +++++----- litellm/router_deprecated.py | 20 ++- litellm/router_strategy/least_busy.py | 136 +++++++++++------- tests/unit/router_strategy/test_least_busy.py | 85 ++++++++--- .../test_router_routing_groups.py | 25 ++-- tests/unit/test_router/test_router.py | 12 +- .../test_router/test_router_helper_utils.py | 2 +- 7 files changed, 238 insertions(+), 120 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 8c095360301..8fccce53cc7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -219,7 +219,12 @@ from litellm.router_strategy.complexity_router.context_compaction import ( reject_recursive_compactor, surface_for_call, ) -from litellm.router_strategy.least_busy import SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import ( + SELECTION_RESERVATION_MAX_KEYS, + LeastBusyLoggingHandler, + pick_reserves, + picks_reserve, +) from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler @@ -1381,14 +1386,14 @@ class Router(DeprecatedRouterMembers): (``pass_through``). """ if kind == "pick": - return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - reserve=False, - ) + with picks_reserve(False): + return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts + model=model, + request_kwargs=request_kwargs if request_kwargs is not None else {}, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) if kind == "healthy": return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts model=model, @@ -3404,12 +3409,13 @@ class Router(DeprecatedRouterMembers): ) # pick the one that is available (lowest TPM/RPM) - deployment = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._drop_unsupported_classifier_reasoning_effort( deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment model=model, @@ -5884,12 +5890,13 @@ class Router(DeprecatedRouterMembers): model_name: Final = "" try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - deployment: Final = self.get_available_deployment( - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment: Final = self.get_available_deployment( + model=model, + messages=[{"role": "user", "content": "prompt"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) data: Final = deployment["litellm_params"].copy() @@ -6237,7 +6244,7 @@ class Router(DeprecatedRouterMembers): kwargs.setdefault("metadata", {}).update({"model_group": model}) # pick the one that is available (lowest TPM/RPM) - with endpoint_selection("text_completion", model=model): + with endpoint_selection("text_completion", model=model), picks_reserve(True): deployment: Final = self.get_available_deployment( model=model, messages=messages, @@ -7378,7 +7385,11 @@ class Router(DeprecatedRouterMembers): kwargs=kwargs, metadata_variable_name=metadata_variable_name, ) - with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): + with ( + self._video_call_pin(handler_name, kwargs), + self._route_capability_pin(handler_name), + picks_reserve(True), + ): deployment: Final = self.get_available_deployment( model=model, messages=kwargs.get("messages", None), @@ -7448,7 +7459,7 @@ class Router(DeprecatedRouterMembers): model_name = None try: verbose_router_logger.debug("Inside embedding()- model: %s", model) - with embedding_deployment_selection(): + with embedding_deployment_selection(), picks_reserve(True): deployment: Final = self.get_available_deployment( model=model, input=input, @@ -8059,13 +8070,13 @@ class Router(DeprecatedRouterMembers): try: verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "batch-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - reserve=False, - ) + with picks_reserve(False): + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "batch-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id") data: Final = deployment["litellm_params"].copy() @@ -15697,7 +15708,6 @@ class Router(DeprecatedRouterMembers): messages: list[dict[str, str]] | None = None, input: str | list | None = None, specific_deployment: bool | None = False, - reserve: bool = True, ): """ Async implementation of 'get_available_deployments'. @@ -15711,7 +15721,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, - reserve=reserve, + reserve=pick_reserves(default=True), ) async def _async_get_available_deployment( @@ -16864,7 +16874,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, - reserve=request_kwargs not in (None, {}), + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 046fb8c452a..b7df0677a10 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -27,11 +27,12 @@ import inspect import sys import warnings from collections import Counter -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload +from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload from litellm.router_api import SETTING_FIELDS, ChangeSet +from litellm.router_strategy.least_busy import picks_reserve if TYPE_CHECKING: from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged @@ -58,6 +59,7 @@ if TYPE_CHECKING: from litellm.types.utils import ModelInfo as ModelMapInfo _T: Final = TypeVar("_T") +_P: Final = ParamSpec("_P") DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: the live tally _note increments @@ -163,6 +165,18 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib return _DeprecatedAttribute(replacement, read) +def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coroutine[object, object, _T]]: + @functools.wraps(lookup) + async def looked_up( + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: forwards the wrapped lookup's ParamSpec unchanged + ) -> _T: + with picks_reserve(False): + return await lookup(*args, **kwargs) + + return looked_up + + # (id(router), member name) pairs whose subclass override is running in this context. _OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar( "router_overrides_running", default=frozenset() @@ -325,7 +339,7 @@ class DeprecatedRouterMembers: deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map) async_get_available_deployment = _attr( "select(model, request_kwargs)", - lambda r: r._core_async_get_available_deployment, + lambda r: _without_reserving(r._core_async_get_available_deployment), ) # -- observers ------------------------------------------------------------------------------ diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index f920e427d41..dff4815f71c 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,7 +1,9 @@ +import contextlib import logging import threading import time -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence +from contextvars import ContextVar from typing import Final, Protocol from pydantic import TypeAdapter, ValidationError @@ -12,7 +14,6 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import log_redis_failure -from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type @@ -20,6 +21,20 @@ IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 _RESERVATIONS_LOCK: Final = threading.Lock() +_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves") + + +@contextlib.contextmanager +def picks_reserve(reserve: bool) -> Generator[None, None, None]: + token: Final = _PICK_RESERVES.set(reserve) + try: + yield + finally: + _PICK_RESERVES.reset(token) + + +def pick_reserves(default: bool) -> bool: + return _PICK_RESERVES.get(default) class _ModelInfo(TypedDict, total=False): @@ -97,23 +112,24 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: return model_group, str(deployment_id) -def _counted_in_flight(kwargs: Mapping[str, object]) -> bool: - return kwargs.get("cache_hit") is not True and ( - kwargs.get("api_call_start_time") is not None or LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL in kwargs - ) +def _reservation_key(deployment_id: str) -> str: + return f"least_busy_reservations:{deployment_id}" -def _reservation_key(request_count_key: str) -> str: - return f"{request_count_key}:reservations" +def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str: + return f"least_busy_started:{call_id}:{_request_count_key(*ref)}" -def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: +def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: return tuple( - _request_count_key(model_group, str(deployment["model_info"]["id"])) - for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) ) +def _request_count_keys(model_group: str, deployment_ids: tuple[str, ...]) -> tuple[str, ...]: + return tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids) + + def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: return tuple(0 if value is None else int(value) for value in values) @@ -180,20 +196,28 @@ class LeastBusyLoggingHandler(CustomLogger): else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock) ) - def _live_reservations(self, key: str, now: float) -> tuple[float, ...]: - raw: Final = self._reservations.get_cache(_reservation_key(key)) - return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now) + def _times(self, key: str) -> tuple[float, ...]: + return _RESERVATION_TIMES.validate_python(self._reservations.get_cache(key)) or () - def _store_reservations(self, key: str, times: tuple[float, ...]) -> None: + def _store_times(self, key: str, times: tuple[float, ...], ttl: float) -> None: # InMemoryCache keeps a live key's first TTL, so drop the key to restart it - self._reservations.delete_cache(_reservation_key(key)) + self._reservations.delete_cache(key) if times: - self._reservations.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) + self._reservations.set_cache(key, times, ttl=ttl) + + def _live_reservations(self, deployment_id: str, now: float) -> tuple[float, ...]: + return _unexpired(self._times(_reservation_key(deployment_id)), now) + + def _store_reservations(self, deployment_id: str, times: tuple[float, ...]) -> None: + self._store_times(_reservation_key(deployment_id), times, SELECTION_RESERVATION_TTL_SECONDS) + + def _drop_oldest_reservation(self, deployment_id: str, now: float) -> None: + self._store_reservations(deployment_id, self._live_reservations(deployment_id, now)[1:]) def _select( self, healthy_deployments: Sequence[Mapping[str, object]], - keys: tuple[str, ...], + deployment_ids: tuple[str, ...], counts: tuple[int, ...], reserve: bool, ) -> Mapping[str, object] | None: @@ -201,25 +225,39 @@ class LeastBusyLoggingHandler(CustomLogger): return None with _RESERVATIONS_LOCK: now: Final = self._clock() - live: Final = tuple(self._live_reservations(key, now) for key in keys) + live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids) 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)) + self._store_reservations(deployment_ids[index], (*live[index], 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) + def _start_attempt(self, ref: tuple[str, str], call_id: object) -> None: with _RESERVATIONS_LOCK: now: Final = self._clock() - self._store_reservations(key, self._live_reservations(key, now)[1:]) + self._drop_oldest_reservation(ref[1], now) + if isinstance(call_id, str): + key: Final = _started_attempts_key(call_id, ref) + self._store_times(key, (*self._times(key), now), IN_FLIGHT_COUNT_TTL_SECONDS) + + def _end_attempt(self, ref: tuple[str, str], call_id: object) -> bool: + if not isinstance(call_id, str): + return True + key: Final = _started_attempts_key(call_id, ref) + with _RESERVATIONS_LOCK: + started: Final = self._times(key) + if started: + self._store_times(key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) + else: + self._drop_oldest_reservation(ref[1], self._clock()) + return bool(started) @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) + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + self._start_attempt(ref, kwargs.get("litellm_call_id")) + self._increment(_request_count_key(*ref), 1) @with_service_target("router_usage") def log_success_event( @@ -257,7 +295,8 @@ class LeastBusyLoggingHandler(CustomLogger): def get_available_deployments( 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) + deployment_ids: Final = _deployment_ids(healthy_deployments) + keys: Final = _request_count_keys(model_group, deployment_ids) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: @@ -265,14 +304,17 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve) - return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve) + return self._select( + healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve + ) + return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve) @with_service_target("router_usage") async def async_get_available_deployments( 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) + deployment_ids: Final = _deployment_ids(healthy_deployments) + keys: Final = _request_count_keys(model_group, deployment_ids) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: @@ -280,20 +322,22 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve) - return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve) + return self._select( + healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve + ) + return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve) def _settle(self, kwargs: Mapping[str, object]) -> None: - if not _counted_in_flight(kwargs): - self._release_reservation(kwargs) + ref: Final = _deployment_ref(kwargs) + if ref is None or not self._end_attempt(ref, kwargs.get("litellm_call_id")): return - self._increment(kwargs, -1) + self._increment(_request_count_key(*ref), -1) async def _async_settle(self, kwargs: Mapping[str, object]) -> None: - if not _counted_in_flight(kwargs): - self._release_reservation(kwargs) + ref: Final = _deployment_ref(kwargs) + if ref is None or not self._end_attempt(ref, kwargs.get("litellm_call_id")): return - await self._async_increment(kwargs, -1) + await self._async_increment(_request_count_key(*ref), -1) def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]: return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) @@ -301,11 +345,7 @@ class LeastBusyLoggingHandler(CustomLogger): def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]: return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True)) - def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: - ref: Final = _deployment_ref(kwargs) - if ref is None: - return - key: Final = _request_count_key(*ref) + def _increment(self, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache try: local: Final = self.router_cache.increment_cache( @@ -319,11 +359,7 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unwritable(key, e) - async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: - ref: Final = _deployment_ref(kwargs) - if ref is None: - return - key: Final = _request_count_key(*ref) + async def _async_increment(self, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache try: local: Final = await self.router_cache.async_increment_cache( diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 2eee767aa06..daecaee6a39 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -4,7 +4,7 @@ import datetime import json import logging import threading -from collections.abc import Callable, Iterator +from collections.abc import Callable, Generator from concurrent.futures import ThreadPoolExecutor from typing import Final @@ -25,22 +25,21 @@ from litellm.router_strategy.least_busy import ( IN_FLIGHT_COUNT_TTL_SECONDS, SELECTION_RESERVATION_TTL_SECONDS, LeastBusyLoggingHandler, + picks_reserve, ) from litellm.types.utils import ModelResponse GROUP: Final = "least-busy-group" -CALL_STARTED_AT: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] -def _routed_call_kwargs(deployment_id: str) -> dict[str, object]: - return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} - - -def _call_kwargs(deployment_id: str) -> dict[str, object]: - return {**_routed_call_kwargs(deployment_id), "api_call_start_time": CALL_STARTED_AT} +def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]: + return { + "litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}, + "litellm_call_id": call_id, + } class SharedRedisCounters: @@ -145,7 +144,7 @@ def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]: "litellm_metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}, }, - "api_call_start_time": CALL_STARTED_AT, + "litellm_call_id": "call-1", } @@ -561,7 +560,7 @@ class _StoreWhereThreadsReadBeforeWriting: self._meeting: Final = threading.Event() @contextlib.contextmanager - def first_reads_meet(self) -> Iterator[None]: + def first_reads_meet(self) -> Generator[None, None, None]: self._meeting.set() yield self._meeting.clear() @@ -616,11 +615,25 @@ def _router_pick(router: Router) -> str: def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None: router: Final = _least_busy_router() - first: Final = _router_pick(router) - for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): - router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60) + with picks_reserve(True): + first: Final = _router_pick(router) + for index in range(DEFAULT_MAX_SIZE_IN_MEMORY): + router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60) - assert _router_pick(router) != first + assert _router_pick(router) != first + + +@pytest.mark.asyncio +async def test_looking_up_a_deployment_does_not_hold_it_for_later_picks() -> None: + router: Final = _least_busy_router() + + sync_picks: Final = [_router_pick(router) for _ in range(3)] + async_picks: Final = [ + (await router.async_get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}))["model_info"]["id"] + for _ in range(3) + ] + + assert sync_picks + async_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6 @pytest.mark.asyncio @@ -638,7 +651,7 @@ def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> No worker: Final = _worker(None) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - worker.log_success_event({**_call_kwargs("dep-a"), "cache_hit": True}, None, None, None) + worker.log_success_event({**_call_kwargs("dep-a", "cached-call"), "cache_hit": True}, None, None, None) assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B @@ -650,14 +663,14 @@ def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick( worker: Final = _worker(None) assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A - outcome(worker, _routed_call_kwargs("dep-a"), None, None, None) + outcome(worker, _call_kwargs("dep-a"), None, None, None) assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behind() -> None: worker: Final = _worker(None) - rejected: Final = {**_routed_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True} + rejected: Final = {**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True} worker.log_pre_api_call(model="m", messages=[], kwargs=rejected) worker.log_failure_event(rejected, None, None, None) @@ -665,6 +678,44 @@ def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behin assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments("tier-group", HEALTHY) is DEPLOYMENT_A + + +def test_a_fallback_attempt_that_fails_before_reaching_its_provider_frees_its_pick() -> None: + worker: Final = _worker(None) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "busy-call")) + first_attempt_started_at: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + first_attempt: Final = {**_call_kwargs("dep-a", "fallback-call"), "api_call_start_time": first_attempt_started_at} + worker.log_pre_api_call(model="m", messages=[], kwargs=first_attempt) + worker.log_failure_event(first_attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + second_attempt: Final = {**_call_kwargs("dep-b", "fallback-call"), "api_call_start_time": first_attempt_started_at} + worker.log_failure_event(second_attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_a_retry_that_starts_before_the_failed_attempt_is_logged_counts_each_attempt_once() -> None: + worker: Final = _worker(None) + attempt: Final = _call_kwargs("dep-a", "retried-call") + + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_failure_event(attempt, None, None, None) + worker.log_success_event(attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + async def _drain_logging() -> None: for _ in range(10): await asyncio.sleep(0) diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 70be27c2645..c13019b4859 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -20,6 +20,7 @@ import litellm from litellm import Router from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger +from litellm.router_strategy.least_busy import picks_reserve from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -56,9 +57,6 @@ def _model_list(): ] -CALL_STARTED_AT: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) - - def _build_router(routing_strategy="simple-shuffle", routing_groups=None): return Router( model_list=_model_list(), @@ -471,8 +469,7 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): "litellm_params": { "metadata": {"model_group": "filtered-model"}, "model_info": {"id": "deploy-1"}, - }, - "api_call_start_time": CALL_STARTED_AT, + } } for callback in litellm.input_callback: @@ -485,7 +482,6 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 - def _filtered_model_pick(router: Router) -> str: deployment: Final = router.get_available_deployment( "filtered-model", request_kwargs=_proxy_shaped_request(model="filtered-model") @@ -501,10 +497,11 @@ def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkey routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}], ) - first: Final = _filtered_model_pick(router) + with picks_reserve(True): + first: Final = _filtered_model_pick(router) kwargs: Final = { "litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}, - "api_call_start_time": CALL_STARTED_AT, + "litellm_call_id": "call-1", } for callback in litellm.input_callback: if isinstance(callback, CustomLogger): @@ -515,6 +512,7 @@ def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkey assert _filtered_model_pick(router) == first + def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): """ Least-busy hangs its counting off litellm's global callback lists, and those lists keep one @@ -532,8 +530,7 @@ def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): "litellm_params": { "metadata": {"model_group": "filtered-model"}, "model_info": {"id": "deploy-1"}, - }, - "api_call_start_time": CALL_STARTED_AT, + } } for callback in litellm.input_callback: @@ -1241,12 +1238,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") == idle + assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] async for _ in stream: pass await _async_wait_until(lambda: _in_flight(router, busy) == 0) - assert await _async_override_pick(router, "least-busy") == busy + assert await _async_override_pick(router, "least-busy") == "d1" assert _selector_is_not_global(router._override_selectors["least-busy"]) @@ -1256,12 +1253,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") == idle + assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] for _ in stream: pass _sync_wait_until(lambda: _in_flight(router, busy) == 0) - assert _sync_override_pick(router, "least-busy") == busy + assert _sync_override_pick(router, "least-busy") == "d1" assert _selector_is_not_global(router._override_selectors["least-busy"]) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a0a7dd78193..cf4f258afeb 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -29096,11 +29096,21 @@ def test_reset_custom_routing_strategy(): assert "get_available_deployment" not in router.__dict__ assert "async_get_available_deployment" not in router.__dict__ - assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment + assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment router._reset_custom_routing_strategy() +@pytest.mark.asyncio +async def test_select_picks_through_a_custom_routing_strategy() -> None: + router: Final = _create_custom_routing_router() + router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router)) + + picked: Final = await router.select("azure-model", {"metadata": {}}) + + assert picked["model_info"]["id"] == "very-special-endpoint" + + @pytest.mark.asyncio async def test_azure_router_acompletion_passes_deployment_credentials_for_plain_and_stream( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5cc62ce6450..031f0ac533a 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -616,7 +616,7 @@ async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_ await admit_turn() picks: Final = [ - (await router.async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ + (await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ "model_info" ]["id"] for _ in range(2) From 5cdcb8954d603ff7d33397755598f07302f1d6d7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:39:26 -0700 Subject: [PATCH 13/18] fix(least-busy): free a request's unstarted picks when it ends Each pick now carries its call id, so a start frees only its own call's pick, and the router frees every pick a request never started once that request ends. A fallback hop rejected before the provider no longer leaves its pick reserved for 30 seconds. --- litellm/router.py | 9 ++ litellm/router_strategy/least_busy.py | 144 +++++++++++++----- tests/unit/router_strategy/test_least_busy.py | 78 +++++++++- .../test_router_routing_groups.py | 5 +- 4 files changed, 195 insertions(+), 41 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 8fccce53cc7..4ca90044464 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -220,10 +220,13 @@ from litellm.router_strategy.complexity_router.context_compaction import ( surface_for_call, ) from litellm.router_strategy.least_busy import ( + CALL_IDENTITY, SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler, + call_id_of, pick_reserves, picks_reserve, + release_unstarted_picks, ) from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler @@ -2655,6 +2658,7 @@ class Router(DeprecatedRouterMembers): model_group=model, healthy_deployments=healthy_deployments, reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), ) case "usage-based-routing": # `LowestTPMLoggingHandler` (v1) only exposes the sync @@ -2713,6 +2717,7 @@ class Router(DeprecatedRouterMembers): model_group=model, healthy_deployments=healthy_deployments, reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), ) case "usage-based-routing" | "usage-based-routing-v2": return selector.get_available_deployments( @@ -9242,6 +9247,10 @@ class Router(DeprecatedRouterMembers): kwargs, include_fallback_errors=include_fallback_errors, ) + finally: + release_unstarted_picks( + self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs)) + ) def _handle_mock_testing_fallbacks( self, diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index dff4815f71c..c01693586a3 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -4,7 +4,7 @@ import threading import time from collections.abc import Callable, Generator, Mapping, Sequence from contextvars import ContextVar -from typing import Final, Protocol +from typing import Final, Protocol, TypeAlias from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -56,6 +56,10 @@ class _CallKwargs(TypedDict, total=False): litellm_params: ReadOnly[_LitellmParams | None] +class CallIdentity(TypedDict, total=False): + litellm_call_id: ReadOnly[object] + + class _DeploymentModelInfo(TypedDict): id: ReadOnly[str | int] @@ -65,9 +69,13 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +CALL_IDENTITY: Final = TypeAdapter(CallIdentity) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) _MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None) -_RESERVATION_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None) +_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None) +_Pick: TypeAlias = tuple[float, str | None] +_PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None) +_PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: @@ -120,6 +128,15 @@ def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str: return f"least_busy_started:{call_id}:{_request_count_key(*ref)}" +def _picked_deployments_key(call_id: str) -> str: + return f"least_busy_picked:{call_id}" + + +def call_id_of(kwargs: Mapping[str, object]) -> str | None: + call_id: Final = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: return tuple( str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) @@ -141,8 +158,17 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) -def _unexpired(times: Sequence[float], now: float) -> tuple[float, ...]: - return tuple(t for t in times if now - t < SELECTION_RESERVATION_TTL_SECONDS) +def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]: + return tuple(pick for pick in picks if now - pick[0] < SELECTION_RESERVATION_TTL_SECONDS) + + +def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]: + tags: Final = tuple(tag for _, tag in picks) + owner: Final = call_id if call_id in tags else None + if owner not in tags: + return picks + index: Final = tags.index(owner) + return (*picks[:index], *picks[index + 1 :]) def _least_busy_index(counts: tuple[int, ...]) -> int: @@ -171,11 +197,46 @@ def _warn_unwritable(key: str, error: Exception) -> None: class _ReservationStore(Protocol): def get_cache(self, key: str) -> object: ... - def set_cache(self, key: str, value: tuple[float, ...], *, ttl: float) -> None: ... + def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: ... def delete_cache(self, key: str) -> None: ... +def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]: + return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or () + + +def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]: + return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or () + + +def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]: + return _PICKED_DEPLOYMENTS.validate_python(store.get_cache(_picked_deployments_key(call_id))) or () + + +def _put(store: _ReservationStore, key: str, values: tuple[object, ...], ttl: float) -> None: + # InMemoryCache keeps a live key's first TTL, so drop the key to restart it + store.delete_cache(key) + if values: + store.set_cache(key, values, ttl=ttl) + + +def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str) -> None: + kept: Final = tuple(pick for pick in _picks(store, deployment_id) if pick[1] != call_id) + _put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS) + + +def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None: + if call_id is None: + return + if not _picked_deployments(store, call_id): + return + with _RESERVATIONS_LOCK: + for deployment_id in frozenset(_picked_deployments(store, call_id)): + _release_call_on(store, deployment_id, call_id) + store.delete_cache(_picked_deployments_key(call_id)) + + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -196,23 +257,18 @@ class LeastBusyLoggingHandler(CustomLogger): else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock) ) - def _times(self, key: str) -> tuple[float, ...]: - return _RESERVATION_TIMES.validate_python(self._reservations.get_cache(key)) or () + def _live_reservations(self, deployment_id: str, now: float) -> tuple[_Pick, ...]: + return _unexpired(_picks(self._reservations, deployment_id), now) - def _store_times(self, key: str, times: tuple[float, ...], ttl: float) -> None: - # InMemoryCache keeps a live key's first TTL, so drop the key to restart it - self._reservations.delete_cache(key) - if times: - self._reservations.set_cache(key, times, ttl=ttl) + def _store_reservations(self, deployment_id: str, picks: tuple[_Pick, ...]) -> None: + _put(self._reservations, _reservation_key(deployment_id), picks, SELECTION_RESERVATION_TTL_SECONDS) - def _live_reservations(self, deployment_id: str, now: float) -> tuple[float, ...]: - return _unexpired(self._times(_reservation_key(deployment_id)), now) + def _take_reservation(self, deployment_id: str, call_id: str | None, now: float) -> None: + self._store_reservations(deployment_id, _without_one_pick(self._live_reservations(deployment_id, now), call_id)) - def _store_reservations(self, deployment_id: str, times: tuple[float, ...]) -> None: - self._store_times(_reservation_key(deployment_id), times, SELECTION_RESERVATION_TTL_SECONDS) - - def _drop_oldest_reservation(self, deployment_id: str, now: float) -> None: - self._store_reservations(deployment_id, self._live_reservations(deployment_id, now)[1:]) + def _record_pick(self, call_id: str, deployment_id: str) -> None: + picked: Final = (*_picked_deployments(self._reservations, call_id), deployment_id) + _put(self._reservations, _picked_deployments_key(call_id), picked, SELECTION_RESERVATION_TTL_SECONDS) def _select( self, @@ -220,6 +276,7 @@ class LeastBusyLoggingHandler(CustomLogger): deployment_ids: tuple[str, ...], counts: tuple[int, ...], reserve: bool, + call_id: str | None, ) -> Mapping[str, object] | None: if not healthy_deployments: return None @@ -228,27 +285,30 @@ class LeastBusyLoggingHandler(CustomLogger): live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids) index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True))) if reserve: - self._store_reservations(deployment_ids[index], (*live[index], now)) + self._store_reservations(deployment_ids[index], (*live[index], (now, call_id))) + if call_id is not None: + self._record_pick(call_id, deployment_ids[index]) return healthy_deployments[index] - def _start_attempt(self, ref: tuple[str, str], call_id: object) -> None: + def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None: with _RESERVATIONS_LOCK: now: Final = self._clock() - self._drop_oldest_reservation(ref[1], now) - if isinstance(call_id, str): + self._take_reservation(ref[1], call_id, now) + if call_id is not None: key: Final = _started_attempts_key(call_id, ref) - self._store_times(key, (*self._times(key), now), IN_FLIGHT_COUNT_TTL_SECONDS) + started: Final = (*_attempt_times(self._reservations, key), now) + _put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS) - def _end_attempt(self, ref: tuple[str, str], call_id: object) -> bool: - if not isinstance(call_id, str): + def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool: + if call_id is None: return True key: Final = _started_attempts_key(call_id, ref) with _RESERVATIONS_LOCK: - started: Final = self._times(key) + started: Final = _attempt_times(self._reservations, key) if started: - self._store_times(key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) + _put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) else: - self._drop_oldest_reservation(ref[1], self._clock()) + self._take_reservation(ref[1], call_id, self._clock()) return bool(started) @with_service_target("router_usage") @@ -256,7 +316,7 @@ class LeastBusyLoggingHandler(CustomLogger): ref: Final = _deployment_ref(kwargs) if ref is None: return - self._start_attempt(ref, kwargs.get("litellm_call_id")) + self._start_attempt(ref, call_id_of(kwargs)) self._increment(_request_count_key(*ref), 1) @with_service_target("router_usage") @@ -293,7 +353,11 @@ class LeastBusyLoggingHandler(CustomLogger): @with_service_target("router_usage") def get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False + self, + model_group: str, + healthy_deployments: Sequence[Mapping[str, object]], + reserve: bool = False, + call_id: str | None = None, ) -> Mapping[str, object] | None: deployment_ids: Final = _deployment_ids(healthy_deployments) keys: Final = _request_count_keys(model_group, deployment_ids) @@ -305,13 +369,17 @@ class LeastBusyLoggingHandler(CustomLogger): _warn_unreadable(model_group, e) else: return self._select( - healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve + healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve, call_id ) - return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve) + return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve, call_id) @with_service_target("router_usage") async def async_get_available_deployments( - self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False + self, + model_group: str, + healthy_deployments: Sequence[Mapping[str, object]], + reserve: bool = False, + call_id: str | None = None, ) -> Mapping[str, object] | None: deployment_ids: Final = _deployment_ids(healthy_deployments) keys: Final = _request_count_keys(model_group, deployment_ids) @@ -323,19 +391,19 @@ class LeastBusyLoggingHandler(CustomLogger): _warn_unreadable(model_group, e) else: return self._select( - healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve + healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve, call_id ) - return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve) + return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve, call_id) def _settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None or not self._end_attempt(ref, kwargs.get("litellm_call_id")): + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): return self._increment(_request_count_key(*ref), -1) async def _async_settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None or not self._end_attempt(ref, kwargs.get("litellm_call_id")): + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): return await self._async_increment(_request_count_key(*ref), -1) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index daecaee6a39..b5d2087879e 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -26,6 +26,7 @@ from litellm.router_strategy.least_busy import ( SELECTION_RESERVATION_TTL_SECONDS, LeastBusyLoggingHandler, picks_reserve, + release_unstarted_picks, ) from litellm.types.utils import ModelResponse @@ -573,7 +574,7 @@ class _StoreWhereThreadsReadBeforeWriting: self._first_reads.wait() return value - def set_cache(self, key: str, value: tuple[float, ...], *, ttl: float) -> None: + def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: self._store.set_cache(key, value, ttl=ttl) def delete_cache(self, key: str) -> None: @@ -703,6 +704,81 @@ def test_a_fallback_attempt_that_fails_before_reaching_its_provider_frees_its_pi assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B +def test_ending_one_request_frees_only_the_picks_it_never_started() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="ended-call") is DEPLOYMENT_A + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="running-call") is DEPLOYMENT_A + + release_unstarted_picks(reservations, "ended-call") + + assert [worker.get_available_deployments(GROUP, HEALTHY, reserve=True) for _ in range(2)] == [ + DEPLOYMENT_A, + DEPLOYMENT_B, + ] + + +def test_a_call_started_without_a_pick_of_its_own_leaves_another_calls_pick_reserved() -> None: + worker: Final = _worker(None) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="picked-call") is DEPLOYMENT_A + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "looked-up-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +def _fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": "unreachable-group", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "key-unreachable"}, + "model_info": {"id": "unreachable"}, + }, + *( + { + "model_name": ROUTER_GROUP, + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": f"key-{deployment_id}"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ROUTER_DEPLOYMENT_IDS + ), + ], + routing_strategy="least-busy", + num_retries=0, + fallbacks=[{"unreachable-group": [ROUTER_GROUP]}], + ) + + +@pytest.mark.asyncio +async def test_a_fallback_hop_rejected_before_its_provider_frees_its_pick_when_every_hop_shares_one_logging_object( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(500, json={"error": {"message": "upstream down"}}) + ) + router: Final = _fallback_router() + logging_obj, request = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.datetime.now(), + model="unreachable-group", + messages=[{"role": "user", "content": "hi"}], + logit_bias={"1": 1}, + litellm_call_id="fallback-call", + ) + + with pytest.raises(litellm.InternalServerError): + await router.acompletion(**request, litellm_logging_obj=logging_obj) + await _drain_logging() + + assert {call.request.headers["authorization"] for call in route.calls} == {"Bearer key-unreachable"} + assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + + def test_a_retry_that_starts_before_the_failed_attempt_is_logged_counts_each_attempt_once() -> None: worker: Final = _worker(None) attempt: Final = _call_kwargs("dep-a", "retried-call") diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index c13019b4859..ba35b64b1a4 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -497,11 +497,12 @@ def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkey routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}], ) + request: Final = _proxy_shaped_request(model="filtered-model") with picks_reserve(True): - first: Final = _filtered_model_pick(router) + first: Final = router.get_available_deployment("filtered-model", request_kwargs=request)["model_info"]["id"] kwargs: Final = { "litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}, - "litellm_call_id": "call-1", + "litellm_call_id": request["litellm_call_id"], } for callback in litellm.input_callback: if isinstance(callback, CustomLogger): From 252012a4694c65e6eb17daf1acf67116b0ff5145 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 04:26:41 -0700 Subject: [PATCH 14/18] fix(router): free a least-busy pick when a vector store create ends --- litellm/router.py | 4 +++ tests/unit/router_strategy/test_least_busy.py | 36 +++++++++++++++++-- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 4ca90044464..b21dc5cc212 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7828,6 +7828,10 @@ class Router(DeprecatedRouterMembers): if model is not None: self.fail_calls[model] += 1 raise e + finally: + release_unstarted_picks( + self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs)) + ) def _override_vector_store_methods_for_router(self): """ diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index b5d2087879e..12e220e6294 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -798,7 +798,7 @@ async def _drain_logging() -> None: await GLOBAL_LOGGING_WORKER.flush() -def _batch_deployment(route: respx.Route) -> str: +def _called_deployment(route: respx.Route) -> str: return route.calls.last.request.headers["authorization"].removeprefix("Bearer key-") @@ -829,7 +829,37 @@ async def test_a_batch_create_does_not_keep_its_deployment_busy( ) await _drain_logging() - assert _router_pick(router) == _batch_deployment(route) + assert _router_pick(router) == _called_deployment(route) + + +_VECTOR_STORE: Final = { + "id": "vs_least_busy", + "object": "vector_store", + "created_at": 1, + "name": "least-busy", + "usage_bytes": 0, + "file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0}, + "status": "completed", +} +_REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)] +) +async def test_a_vector_store_create_does_not_keep_its_deployment_busy( + provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response) + router: Final = _least_busy_router() + + with contextlib.suppress(openai.BadRequestError): + await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", litellm_call_id="vector-store-call") + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) @pytest.mark.asyncio @@ -844,7 +874,7 @@ async def test_a_batch_cancel_does_not_keep_its_deployment_busy( await router.acancel_batch(model=ROUTER_GROUP, batch_id="batch_least_busy") - assert _router_pick(router) == _batch_deployment(route) + assert _router_pick(router) == _called_deployment(route) @pytest.mark.asyncio From 59924d3b89de8c802406131417ae2fd19260caa0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:46:25 -0700 Subject: [PATCH 15/18] fix(router): free least-busy picks of calls made without a call id --- litellm/router.py | 188 +++++++++--------- litellm/router_strategy/least_busy.py | 22 +- tests/unit/router_strategy/test_least_busy.py | 53 ++++- 3 files changed, 161 insertions(+), 102 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b21dc5cc212..845c8eba37b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -224,9 +224,9 @@ from litellm.router_strategy.least_busy import ( SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler, call_id_of, + owning_picks, pick_reserves, picks_reserve, - release_unstarted_picks, ) from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler @@ -7766,72 +7766,71 @@ class Router(DeprecatedRouterMembers): Returns: VectorStoreCreateResponse """ - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + # If model is None, use the factory function approach (direct SDK call) + if model is None: + from litellm.vector_stores.main import acreate - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) + # Use the factory function to handle the call + factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") + return await factory_fn(**kwargs) - from litellm.vector_stores import acreate as avector_store_create_sdk + from litellm.vector_stores import acreate as avector_store_create_sdk - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } + parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) + deployment: Final = await self._core_async_get_available_deployment( + model=model, + messages=[{"role": "user", "content": "vector-store-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="avector_store_create", ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + model_client: Final = self._get_async_openai_model_client( + deployment=deployment, + kwargs=kwargs, + ) + self.total_calls[model_name] += 1 - return response - except Exception as e: - verbose_router_logger.exception( - "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - finally: - release_unstarted_picks( - self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs)) - ) + # Get custom provider from deployment params + custom_llm_provider = data.get("custom_llm_provider") + _, inferred_custom_llm_provider, _, _ = get_llm_provider( + model=data["model"], + custom_llm_provider=custom_llm_provider, + ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider + + async with self._deployment_slot( + deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span + ): + response = await avector_store_create_sdk( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + "client": model_client, + **kwargs, + } + ) + + self.success_calls[model_name] += 1 + verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + + return response + except Exception as e: + verbose_router_logger.exception( + "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e + ) + if model is not None: + self.fail_calls[model] += 1 + raise e def _override_vector_store_methods_for_router(self): """ @@ -9219,42 +9218,39 @@ class Router(DeprecatedRouterMembers): mock_timeout: Final = kwargs.pop("mock_timeout", None) - try: - self._handle_mock_testing_fallbacks( - kwargs=kwargs, - model_group=model_group, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + self._handle_mock_testing_fallbacks( + kwargs=kwargs, + model_group=model_group, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + ) - if mock_timeout is not None: - response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) - else: - response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Async Response: %s", response) - response = add_fallback_headers_to_response( - response=response, - attempted_fallbacks=0, - ) - return response - except Exception as e: - return await self.async_function_with_fallbacks_common_utils( - e, - disable_fallbacks, - fallbacks, - context_window_fallbacks, - content_policy_fallbacks, - model_group, - args, - kwargs, - include_fallback_errors=include_fallback_errors, - ) - finally: - release_unstarted_picks( - self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs)) - ) + if mock_timeout is not None: + response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) + else: + response = await self.async_function_with_retries(*args, **kwargs) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug("Async Response: %s", response) + response = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=0, + ) + return response + except Exception as e: + return await self.async_function_with_fallbacks_common_utils( + e, + disable_fallbacks, + fallbacks, + context_window_fallbacks, + content_policy_fallbacks, + model_group, + args, + kwargs, + include_fallback_errors=include_fallback_errors, + ) def _handle_mock_testing_fallbacks( self, diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index c01693586a3..1250dcd19c4 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -2,6 +2,7 @@ import contextlib import logging import threading import time +import uuid from collections.abc import Callable, Generator, Mapping, Sequence from contextvars import ContextVar from typing import Final, Protocol, TypeAlias @@ -22,6 +23,7 @@ SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 _RESERVATIONS_LOCK: Final = threading.Lock() _PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves") +_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None) @contextlib.contextmanager @@ -164,7 +166,7 @@ def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]: def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]: tags: Final = tuple(tag for _, tag in picks) - owner: Final = call_id if call_id in tags else None + owner: Final = next((tag for tag in (call_id, _PICK_OWNER.get()) if tag is not None and tag in tags), None) if owner not in tags: return picks index: Final = tags.index(owner) @@ -237,6 +239,17 @@ def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> No store.delete_cache(_picked_deployments_key(call_id)) +@contextlib.contextmanager +def owning_picks(store: _ReservationStore, call_id: str | None) -> Generator[None, None, None]: + owner: Final = call_id or _PICK_OWNER.get() or uuid.uuid4().hex + token: Final = _PICK_OWNER.set(owner) + try: + yield + finally: + _PICK_OWNER.reset(token) + release_unstarted_picks(store, owner) + + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -285,9 +298,10 @@ class LeastBusyLoggingHandler(CustomLogger): live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids) index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True))) if reserve: - self._store_reservations(deployment_ids[index], (*live[index], (now, call_id))) - if call_id is not None: - self._record_pick(call_id, deployment_ids[index]) + tag: Final = call_id if call_id is not None else _PICK_OWNER.get() + self._store_reservations(deployment_ids[index], (*live[index], (now, tag))) + if tag is not None: + self._record_pick(tag, deployment_ids[index]) return healthy_deployments[index] def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None: diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 12e220e6294..7d84d6ca997 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import contextvars import datetime import json import logging @@ -25,6 +26,7 @@ from litellm.router_strategy.least_busy import ( IN_FLIGHT_COUNT_TTL_SECONDS, SELECTION_RESERVATION_TTL_SECONDS, LeastBusyLoggingHandler, + owning_picks, picks_reserve, release_unstarted_picks, ) @@ -719,6 +721,31 @@ def test_ending_one_request_frees_only_the_picks_it_never_started() -> None: ] +def test_a_call_without_a_call_id_turns_its_own_pick_into_its_in_flight_count() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + + with owning_picks(reservations, 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", "attempt-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_another_call_starting_leaves_the_pick_of_a_call_without_a_call_id_reserved() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + other_call: Final = contextvars.copy_context() + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + other_call.run(worker.log_pre_api_call, model="m", messages=[], kwargs=_call_kwargs("dep-a", "other-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + def test_a_call_started_without_a_pick_of_its_own_leaves_another_calls_pick_reserved() -> None: worker: Final = _worker(None) assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="picked-call") is DEPLOYMENT_A @@ -842,21 +869,43 @@ _VECTOR_STORE: Final = { "status": "completed", } _REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}} +_CALL_IDENTITIES: Final = ({"litellm_call_id": "least-busy-call"}, {}) @pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) @pytest.mark.parametrize( "provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)] ) async def test_a_vector_store_create_does_not_keep_its_deployment_busy( - provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + provider_response: httpx.Response, + call_identity: dict[str, str], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response) router: Final = _least_busy_router() with contextlib.suppress(openai.BadRequestError): - await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", litellm_call_id="vector-store-call") + await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", **call_identity) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +async def test_listing_fine_tuning_jobs_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.get(url__startswith="https://api.openai.com/v1/fine_tuning/jobs").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "has_more": False}) + ) + router: Final = _least_busy_router() + + await router.alist_fine_tuning_jobs(model=ROUTER_GROUP, **call_identity) await _drain_logging() assert _router_pick(router) == _called_deployment(route) From 2336c0a6edec76fab03a011c83713e12a3267764 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 06:17:48 -0700 Subject: [PATCH 16/18] fix(router): free the least-busy pick of a sync text completion --- litellm/router.py | 53 ++++++++++--------- tests/unit/router_strategy/test_least_busy.py | 36 +++++++++++-- 2 files changed, 59 insertions(+), 30 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 845c8eba37b..d15c0611498 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6242,34 +6242,35 @@ class Router(DeprecatedRouterMembers): **kwargs, ): messages: Final = [{"role": "user", "content": prompt}] - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - kwargs.setdefault("metadata", {}).update({"model_group": model}) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + kwargs["model"] = model + kwargs["prompt"] = prompt + kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) + kwargs.setdefault("metadata", {}).update({"model_group": model}) - # pick the one that is available (lowest TPM/RPM) - with endpoint_selection("text_completion", model=model), picks_reserve(True): - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, + # pick the one that is available (lowest TPM/RPM) + with endpoint_selection("text_completion", model=model), picks_reserve(True): + deployment: Final = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + + data: Final = deployment["litellm_params"].copy() + for k, v in self._core_default_litellm_params.items(): + if k not in kwargs: # prioritize model-specific params > default router params + kwargs[k] = v + elif k == "metadata": + kwargs[k].update(v) + + # call via litellm.completion() + return litellm.text_completion( + **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} ) - - data: Final = deployment["litellm_params"].copy() - for k, v in self._core_default_litellm_params.items(): - if k not in kwargs: # prioritize model-specific params > default router params - kwargs[k] = v - elif k == "metadata": - kwargs[k].update(v) - - # call via litellm.completion() - return litellm.text_completion( - **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} - ) - except Exception as e: - raise e + except Exception as e: + raise e async def atext_completion( self, diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 7d84d6ca997..212dc571cf5 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -327,11 +327,11 @@ ROUTER_GROUP: Final = "least-busy-router-group" ROUTER_DEPLOYMENT_IDS: Final = ("1", "2", "3") -def _least_busy_router() -> Router: +def _least_busy_router(group: str = ROUTER_GROUP) -> Router: return Router( model_list=[ { - "model_name": ROUTER_GROUP, + "model_name": group, "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": f"key-{deployment_id}"}, "model_info": {"id": deployment_id}, } @@ -610,8 +610,8 @@ def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> N assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B -def _router_pick(router: Router) -> str: - deployment: Final = router.get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}) +def _router_pick(router: Router, group: str = ROUTER_GROUP) -> str: + deployment: Final = router.get_available_deployment(group, request_kwargs={"metadata": {}}) return deployment["model_info"]["id"] @@ -911,6 +911,34 @@ async def test_listing_fine_tuning_jobs_does_not_keep_its_deployment_busy( assert _router_pick(router) == _called_deployment(route) +TEXT_GROUP: Final = "openai/least-busy-text-group" + + +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +def test_a_sync_text_completion_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://api.openai.com/v1/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "cmpl-least-busy", + "object": "text_completion", + "created": 1, + "model": "least-busy-text-group", + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + ) + ) + router: Final = _least_busy_router(TEXT_GROUP) + + router.text_completion(model=TEXT_GROUP, prompt="Say ok", **call_identity) + + assert _router_pick(router, TEXT_GROUP) == _called_deployment(route) + + @pytest.mark.asyncio async def test_a_batch_cancel_does_not_keep_its_deployment_busy( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch From c60a56a1c78d3dd67482956bb8c7a5c6147e6366 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:18:38 -0700 Subject: [PATCH 17/18] test(router): use fixed start times in least-busy fixtures --- tests/unit/decisions/test_main.py | 2 +- tests/unit/router_strategy/test_least_busy.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index 73fda703f88..c8113d638a6 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -320,7 +320,7 @@ async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_ logging_obj, _ = litellm.utils.function_setup( original_function="adecisions", rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), model="jev", litellm_call_id=str(uuid.uuid4()), ) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 212dc571cf5..6cccec1d8ad 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -791,7 +791,7 @@ async def test_a_fallback_hop_rejected_before_its_provider_frees_its_pick_when_e logging_obj, request = litellm.utils.function_setup( original_function="acompletion", rules_obj=litellm.utils.Rules(), - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), model="unreachable-group", messages=[{"role": "user", "content": "hi"}], logit_bias={"1": 1}, From c4e873191584dff4ca5b7778b2f3ed652bda346c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:21:41 -0700 Subject: [PATCH 18/18] fix(router): free the least-busy pick of a sync generic call that fails before its provider --- litellm/router.py | 101 +++++++++--------- tests/unit/router_strategy/test_least_busy.py | 12 +++ 2 files changed, 63 insertions(+), 50 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 244c7aac747..846cec8f05b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7392,65 +7392,66 @@ class Router(DeprecatedRouterMembers): """ handler_name: Final = original_function.__name__ metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") - try: - verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - with ( - self._video_call_pin(handler_name, kwargs), - self._route_capability_pin(handler_name), - picks_reserve(True), - ): - deployment: Final = self.get_available_deployment( + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) + self._update_kwargs_before_fallbacks( model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, + kwargs=kwargs, + metadata_variable_name=metadata_variable_name, + ) + with ( + self._video_call_pin(handler_name, kwargs), + self._route_capability_pin(handler_name), + picks_reserve(True), + ): + deployment: Final = self.get_available_deployment( + model=model, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="generic_api_call", + splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, ) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="generic_api_call", - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] - self.total_calls[model_name] += 1 + self.total_calls[model_name] += 1 - # For passthrough routes, use the actual model from deployment - # and swap model name in endpoint if present - if "endpoint" in kwargs and kwargs["endpoint"]: - kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) - kwargs["model"] = model_name + # For passthrough routes, use the actual model from deployment + # and swap model name in endpoint if present + if "endpoint" in kwargs and kwargs["endpoint"]: + kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) + kwargs["model"] = model_name - # Perform pre-call checks for routing strategy - self.routing_strategy_pre_call_checks(deployment=deployment) + # Perform pre-call checks for routing strategy + self.routing_strategy_pre_call_checks(deployment=deployment) - custom_llm_provider: Final = provider_for_generic_call(data) + custom_llm_provider: Final = provider_for_generic_call(data) - response: Final = original_function( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - **kwargs, - } - ) + response: Final = original_function( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + **kwargs, + } + ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) - return response - except Exception as e: - verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) - if model is not None: - self.fail_calls[model] += 1 - raise e + self.success_calls[model_name] += 1 + verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) + return response + except Exception as e: + verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) + if model is not None: + self.fail_calls[model] += 1 + raise e def embedding( self, diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 6cccec1d8ad..796bd5b6ad2 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -939,6 +939,18 @@ def test_a_sync_text_completion_does_not_keep_its_deployment_busy( assert _router_pick(router, TEXT_GROUP) == _called_deployment(route) + +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its_pick( + call_identity: dict[str, str], +) -> None: + router: Final = _least_busy_router() + + with pytest.raises(litellm.InternalServerError): + router.responses(model=ROUTER_GROUP, input="hi", api_base="not-a-url", **call_identity) + + assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + @pytest.mark.asyncio async def test_a_batch_cancel_does_not_keep_its_deployment_busy( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch