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