diff --git a/litellm/router.py b/litellm/router.py index 535072d5fcf..bc1a0f5ce57 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -286,7 +286,15 @@ 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 ( + CALL_IDENTITY, + SELECTION_RESERVATION_MAX_KEYS, + LeastBusyLoggingHandler, + call_id_of, + owning_picks, + 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 @@ -1402,13 +1410,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, - ) + 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, @@ -1678,6 +1687,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._claude_code_session_router_cache: DualCache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache(), @@ -2185,7 +2197,9 @@ class Router(DeprecatedRouterMembers): selector: RouterStrategySelector | None = None match self._normalize_strategy(strategy): case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler(router_cache=self._core_cache) + selector = LeastBusyLoggingHandler( + router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations + ) case RoutingStrategy.USAGE_BASED_ROUTING.value: selector = LowestTPMLoggingHandler( router_cache=self._core_cache, @@ -2651,6 +2665,7 @@ class Router(DeprecatedRouterMembers): 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 @@ -2664,6 +2679,8 @@ class Router(DeprecatedRouterMembers): return await selector.async_get_available_deployments( 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 @@ -2704,6 +2721,7 @@ class Router(DeprecatedRouterMembers): 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 @@ -2720,6 +2738,8 @@ class Router(DeprecatedRouterMembers): return selector.get_available_deployments( 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( @@ -3416,12 +3436,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, @@ -5953,12 +5974,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() @@ -6326,37 +6348,38 @@ 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): - 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) + pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) + if pinned_model_id is not None: + kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults + + # 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) - pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults - - # 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, @@ -7471,61 +7494,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): - 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, @@ -7548,7 +7576,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, @@ -7850,68 +7878,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 + # 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): """ @@ -8159,12 +8190,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, - ) + 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() @@ -9298,38 +9330,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, - ) + 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, @@ -10618,6 +10651,7 @@ class Router(DeprecatedRouterMembers): messages=None, input=None, request_kwargs=routing_kwargs, + reserve=False, ) if selected is None: no_deployment: Final = await async_raise_no_deployment_exception( @@ -15836,6 +15870,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, + reserve=pick_reserves(default=True), ) async def _async_get_available_deployment( @@ -15845,6 +15880,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" @@ -15941,6 +15977,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=reserve and request_kwargs not in (None, {}), ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( @@ -16110,6 +16147,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=request_kwargs not in (None, {}), ) if deployment is None: @@ -16983,6 +17021,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17163,6 +17202,7 @@ class Router(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=request_kwargs not in (None, {}), ) if deployment is None: diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 0b328cc3e6b..756d54595da 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 8e0c3d6a7e7..1250dcd19c4 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,6 +1,11 @@ +import contextlib import logging -from collections.abc import Mapping, Sequence -from typing import Final +import threading +import time +import uuid +from collections.abc import Callable, Generator, Mapping, Sequence +from contextvars import ContextVar +from typing import Final, Protocol, TypeAlias from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict @@ -8,11 +13,30 @@ 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 +_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 +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): @@ -21,10 +45,12 @@ 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): metadata: ReadOnly[_Metadata | None] + litellm_metadata: ReadOnly[_Metadata | None] model_info: ReadOnly[_ModelInfo | None] @@ -32,6 +58,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] @@ -41,14 +71,31 @@ 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) +_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | 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: 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_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 @@ -56,23 +103,52 @@ 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 - deployment_id: Final = model_info.get("id") if model_info else None + litellm_params: Final = call.get("litellm_params") or _LitellmParams() + 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, + ) if model_group is None or deployment_id is None: return None return model_group, str(deployment_id) -def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: +def _reservation_key(deployment_id: str) -> str: + return f"least_busy_reservations:{deployment_id}" + + +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( - _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) @@ -84,12 +160,21 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: return _as_counts(values) -def _least_busy( - healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] -) -> Mapping[str, object] | None: - if not healthy_deployments: - return None - return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] +def _unexpired(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 = 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) + return (*picks[:index], *picks[index + 1 :]) + + +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: @@ -111,24 +196,148 @@ 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[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)) + + +@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 logged_failure: int = 0 - def __init__(self, router_cache: DualCache): + 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 + 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, deployment_id: str, now: float) -> tuple[_Pick, ...]: + return _unexpired(_picks(self._reservations, deployment_id), now) + + def _store_reservations(self, deployment_id: str, picks: tuple[_Pick, ...]) -> None: + _put(self._reservations, _reservation_key(deployment_id), picks, SELECTION_RESERVATION_TTL_SECONDS) + + 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 _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, + healthy_deployments: Sequence[Mapping[str, object]], + deployment_ids: tuple[str, ...], + counts: tuple[int, ...], + reserve: bool, + call_id: str | None, + ) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + with _RESERVATIONS_LOCK: + now: Final = self._clock() + 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: + 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: + with _RESERVATIONS_LOCK: + now: Final = self._clock() + self._take_reservation(ref[1], call_id, now) + if call_id is not None: + key: Final = _started_attempts_key(call_id, ref) + 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: str | None) -> bool: + if call_id is None: + return True + key: Final = _started_attempts_key(call_id, ref) + with _RESERVATIONS_LOCK: + started: Final = _attempt_times(self._reservations, key) + if started: + _put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) + else: + self._take_reservation(ref[1], call_id, 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._increment(kwargs, 1) + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + self._start_attempt(ref, call_id_of(kwargs)) + self._increment(_request_count_key(*ref), 1) @with_service_target("router_usage") def log_success_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_success += 1 @@ -136,7 +345,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 @@ -144,7 +353,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: - await self._async_increment(kwargs, -1) + await self._async_settle(kwargs) if self.test_flag: self.logged_success += 1 @@ -152,15 +361,20 @@ 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 @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, + call_id: str | None = None, ) -> 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: @@ -168,15 +382,21 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) - 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, 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, call_id) @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, + call_id: str | None = None, ) -> 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: @@ -184,15 +404,30 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return _least_busy(healthy_deployments, shared) - 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, 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, call_id) - def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + def _settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None: + if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): return - key: Final = _request_count_key(*ref) + 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, call_id_of(kwargs)): + return + 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) + + 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, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache try: local: Final = self.router_cache.increment_cache( @@ -206,11 +441,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/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/decisions/test_main.py b/tests/unit/decisions/test_main.py index 0c357837d31..c8113d638a6 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 @@ -297,6 +299,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(2026, 1, 1), + 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 5b3af388ae5..796bd5b6ad2 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,9 +1,16 @@ import asyncio +import contextlib +import contextvars +import datetime import json import logging +import threading +from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor from typing import Final import httpx +import openai import pytest import respx @@ -11,10 +18,18 @@ 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.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, LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import ( + IN_FLIGHT_COUNT_TTL_SECONDS, + SELECTION_RESERVATION_TTL_SECONDS, + LeastBusyLoggingHandler, + owning_picks, + picks_reserve, + release_unstarted_picks, +) from litellm.types.utils import ModelResponse GROUP: Final = "least-busy-group" @@ -23,8 +38,11 @@ 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]: - return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} +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: @@ -58,6 +76,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) @@ -81,6 +113,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) @@ -98,6 +140,32 @@ 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}, + }, + "litellm_call_id": "call-1", + } + + +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 @@ -259,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}, } @@ -391,3 +459,517 @@ 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) + burst: Final = 2 * len(ROUTER_DEPLOYMENT_IDS) + all_started: Final = asyncio.Barrier(burst) + + async def completion_after_all_started(_: httpx.Request) -> httpx.Response: + await all_started.wait() + 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=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(burst)) + ) + + 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: + clock: Final = _ManualClock() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=clock) + + 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 + + +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) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None: + cache: Final = DualCache() + 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")) + + assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + 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) -> Generator[None, None, 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[object, ...], *, 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, group: str = ROUTER_GROUP) -> str: + deployment: Final = router.get_available_deployment(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() + + 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 + + +@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 +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", "cached-call"), "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, _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 = {**_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 + + +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_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_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 + + 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(2026, 1, 1), + 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") + + 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) + await GLOBAL_LOGGING_WORKER.flush() + + +def _called_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) == _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"}} +_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, + 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", **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) + + +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.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 +) -> 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) == _called_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"] diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 6b5c052618b..9aba894eb36 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -21,6 +21,7 @@ from litellm import Router from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext +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 @@ -482,6 +483,38 @@ 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"}], + ) + + request: Final = _proxy_shaped_request(model="filtered-model") + with picks_reserve(True): + 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": request["litellm_call_id"], + } + 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 @@ -614,6 +647,7 @@ async def test_select_deployment_async_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -628,6 +662,7 @@ def test_select_deployment_sync_returns_none_without_selector(): messages=None, input=None, request_kwargs=None, + reserve=False, ) assert result is None @@ -662,6 +697,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 @@ -708,6 +744,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.py b/tests/unit/test_router/test_router.py index 7ac27e22a93..9039f0fb93b 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -29146,11 +29146,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 8929c5fd345..c3814371c26 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -599,6 +599,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: Final = 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._core_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(