From 733b7764039711cc9f7ee103c8e17b7e9f4a47d3 Mon Sep 17 00:00:00 2001 From: Oliver Yu <57586652+tc3oliver@users.noreply.github.com> Date: Thu, 5 Nov 2026 08:30:12 +0800 Subject: [PATCH] fix(least-busy): reserve the picked deployment so a burst of requests spreads out (#47484) * fix(least-busy): reserve the picked deployment so a burst of requests spreads out The in-flight count only went up in log_pre_api_call, which runs after selection and the pre-call hooks. Requests selected in that window all saw the same counts and landed on the same deployment (16 simultaneous requests over 4 deployments went as lopsided as 12/4/0/0 in the proxy). When the router selects for a request it is about to send, least-busy now records a reservation for the chosen deployment in the same synchronous step that reads the reservations, and log_pre_api_call turns it into the normal in-flight count. Selections without request kwargs (token counting, probes) do not reserve. A reservation whose request never reaches the API call expires after 30 s. The shared and per-worker counts and their success/failure decrements are unchanged. * fix(least-busy): prune expired and empty selection reservations The reservation map kept a key per model group forever, so a wildcard route with caller-chosen model names grew it without bound. Drop expired and empty entries on every write. Signed-off-by: Oliver Yu * fix(least-busy): keep pick reservations in the shared router cache Selectors that share a router cache dedupe to one registered callback, so only that instance saw log_pre_api_call. Reservations held on the other instances were never released and skewed picks for the full TTL. Storing them in the router cache's in-memory layer lets any selector release them, and the per-key TTL replaces the full-map rebuild on every write. Tests now hold the burst with a barrier instead of a sleep and use a manual clock object. * fix(least-busy): restart the reservation TTL on every write InMemoryCache keeps the TTL a live key was first given, so a reservation key that kept being rewritten still expired 30 s after its first pick and took the newer reservations with it. Drop the key before each write so the TTL counts from the latest one, and leave it deleted once it is empty. Signed-off-by: Oliver Yu * fix(router): count least-busy decisions calls and release reservations on cache hits Least-busy now reads the model group from litellm_metadata before metadata, so /v1/decisions calls that carry caller metadata are counted in flight and release their reservation at pre-call. A response cache hit releases its reservation instead of decrementing another request's in-flight count, and Responses WebSocket turn admission selects without reserving. The reservation store is reached through a typed protocol. * fix(decisions): count proxy /v1/decisions calls in least-busy's in-flight totals * fix(least-busy): keep reservations out of the shared router cache and count requests that start during a Redis read Least-busy now reads a call's deployment from the routing metadata the router sets, so /v1/decisions logging params stay as on main and failed decisions calls no longer start cooldowns. * style(router): wrap the least-busy reservation store line * fix(least-busy): lock the reservation read-and-update so threaded sync calls keep each other's picks Router.completion from a thread pool and sync success callbacks can read and rewrite a deployment's reservations at the same time, so one thread's pick or release could be lost. * fix(least-busy): release a pick whose call ends before reaching the provider A call that fails or returns before the provider handoff (a 400 for an unsupported param, token counting, a proxy-gate rejection without a start log, batch create) now drops its reservation instead of decrementing a count it never added. Router.select picks and batch cancel no longer reserve, since neither makes a call that logs. * test(least-busy): give callback fixtures a fixed call-start time The routing-group tests fed least-busy a finished call without the start time litellm stamps before counting it, so the count stayed up after the fix. Fixtures now use a fixed timestamp instead of the clock. * 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. * fix(least-busy): free a request's unstarted picks when it ends Each pick now carries its call id, so a start frees only its own call's pick, and the router frees every pick a request never started once that request ends. A fallback hop rejected before the provider no longer leaves its pick reserved for 30 seconds. * fix(router): free a least-busy pick when a vector store create ends * fix(router): free least-busy picks of calls made without a call id * fix(router): free the least-busy pick of a sync text completion * test(router): use fixed start times in least-busy fixtures * fix(router): free the least-busy pick of a sync generic call that fails before its provider * fix(router): read in-flight counts and start attempts under one lock so a busy deployment never looks free * fix(router): keep pass-through lookups from reserving least-busy picks * test(least-busy): wait for the helper thread to start before the timed join * fix(least-busy): keep realtime sessions out of in-flight counts --------- Signed-off-by: Oliver Yu Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router_deprecated.py | 20 +- litellm/router_python_core.py | 441 ++++++------ litellm/router_strategy/least_busy.py | 341 +++++++-- .../test_least_busy_concurrent_decisions.py | 90 +++ tests/unit/decisions/test_main.py | 39 + tests/unit/router_strategy/test_least_busy.py | 672 +++++++++++++++++- .../test_router_routing_groups.py | 37 + tests/unit/test_router/test_router.py | 12 +- .../test_router/test_router_helper_utils.py | 28 + 9 files changed, 1416 insertions(+), 264 deletions(-) create mode 100644 tests/integration/routing/test_least_busy_concurrent_decisions.py diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 83a1251d3c0..0a02408a717 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 ASSIGNABLE_ALIASES, 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 @@ -164,6 +166,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() @@ -333,7 +347,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_python_core.py b/litellm/router_python_core.py index c0c5d74a88a..4b5ecc82581 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -282,7 +282,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 @@ -1462,13 +1470,14 @@ class PythonCore(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, @@ -1480,13 +1489,14 @@ class PythonCore(DeprecatedRouterMembers): health_check_probe=health_check_probe, ) if kind == "pass_through": - return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) + with picks_reserve(False): + return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts + model=model, + messages=messages, + input=input, + specific_deployment=specific_deployment, + request_kwargs=request_kwargs, + ) raise ValueError(f"unknown select kind {kind!r}") def apply(self, changes: ChangeSet) -> ApplyResult: @@ -1768,6 +1778,9 @@ class PythonCore(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(), @@ -2271,7 +2284,9 @@ class PythonCore(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, @@ -2736,6 +2751,7 @@ class PythonCore(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 @@ -2749,6 +2765,8 @@ class PythonCore(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 @@ -2789,6 +2807,7 @@ class PythonCore(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 @@ -2805,6 +2824,8 @@ class PythonCore(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( @@ -3524,12 +3545,13 @@ class PythonCore(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, @@ -6104,12 +6126,13 @@ class PythonCore(DeprecatedRouterMembers): deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with picks_reserve(True): + deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init + 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() @@ -6496,37 +6519,38 @@ class PythonCore(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 self._host.invoke("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 self._host.invoke("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, @@ -7671,64 +7695,69 @@ class PythonCore(DeprecatedRouterMembers): handler_name: Final = original_function.__name__ metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - 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 = self.get_available_deployment( # rebind-ok: set on success, see pre-init + 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 = self.get_available_deployment( # rebind-ok: set on success, see pre-init + 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 - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - 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 + self._stamp_failed_attempt( + e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None + ) + raise e def embedding( self, @@ -7752,7 +7781,7 @@ class PythonCore(DeprecatedRouterMembers): deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: verbose_router_logger.debug("Inside embedding()- model: %s", model) - with embedding_deployment_selection(): + with embedding_deployment_selection(), picks_reserve(True): deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init model=model, input=input, @@ -8061,68 +8090,71 @@ class PythonCore(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): """ @@ -8370,12 +8402,13 @@ class PythonCore(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() @@ -9520,38 +9553,39 @@ class PythonCore(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, @@ -10856,6 +10890,7 @@ class PythonCore(DeprecatedRouterMembers): messages=None, input=None, request_kwargs=routing_kwargs, + reserve=False, ) if selected is None: no_deployment: Final = await async_raise_no_deployment_exception( @@ -16214,6 +16249,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, specific_deployment=specific_deployment, + reserve=pick_reserves(default=True), ) async def _async_get_available_deployment( @@ -16223,6 +16259,7 @@ class PythonCore(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" @@ -16319,6 +16356,7 @@ class PythonCore(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( @@ -16488,6 +16526,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17365,6 +17404,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: @@ -17545,6 +17585,7 @@ class PythonCore(DeprecatedRouterMembers): messages=messages, input=input, request_kwargs=request_kwargs, + reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), ) if deployment is None: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 8e0c3d6a7e7..22159079f98 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,31 @@ 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 +from litellm.types.utils import CallTypes 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 +46,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 +59,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,38 +72,85 @@ 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")): + call_type: Final = kwargs.get("call_type") + if is_batch_retrieve_call_type(call_type) or call_type == CallTypes.arealtime.value: return None try: 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 +162,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 +198,156 @@ 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, ...], + keys: tuple[str, ...], + shared: tuple[int, ...] | None, + reserve: bool, + call_id: str | None, + ) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + with _RESERVATIONS_LOCK: + now: Final = self._clock() + counts: Final = ( + self._this_workers_counts(keys) if shared is None else self._at_least_this_workers(keys, shared) + ) + 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) + self._count_locally(_request_count_key(*ref), 1) + 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: + with _RESERVATIONS_LOCK: + self._count_locally(_request_count_key(*ref), -1) + 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) + self._count_locally(_request_count_key(*ref), -1) + 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._count_shared(_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 +355,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 +363,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 +371,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 +392,19 @@ 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, keys, shared, reserve, call_id) + return self._select(healthy_deployments, deployment_ids, keys, None, 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,42 +412,51 @@ 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, keys, shared, reserve, call_id) + return self._select(healthy_deployments, deployment_ids, keys, None, 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) - redis_cache: Final = self.router_cache.redis_cache + self._count_shared(_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_count_shared(_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 _count_locally(self, key: str, delta: int) -> None: try: local: Final = self.router_cache.increment_cache( key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS ) if local < 0: self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: - return + except TypeError as e: + _warn_unwritable(key, e) + + def _count_shared(self, key: str, delta: int) -> None: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return + try: redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) 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_count_shared(self, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return try: - local: Final = await self.router_cache.async_increment_cache( - key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS - ) - if local is not None and local < 0: - await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: - return await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) 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 1e42e74806d..0e2813e0a0d 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -6,8 +6,10 @@ import copy import json import re import threading +import uuid from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import wait +from datetime import datetime from functools import partial from types import MappingProxyType from typing import Final, TypeAlias @@ -300,6 +302,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..f156a5ee405 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,11 +18,19 @@ 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.types.utils import ModelResponse +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 CallTypes, ModelResponse GROUP: Final = "least-busy-group" DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} @@ -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,44 @@ 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_a_closed_realtime_session_leaves_its_deployment_idle() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + session_event: Final = {**_routed_through_litellm_metadata("dep-a"), "call_type": CallTypes.arealtime.value} + + for _ in range(5): + worker.log_pre_api_call(model="m", messages=[], kwargs=session_event) + worker.log_success_event(session_event, None, None, None) + + assert 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,12 +339,16 @@ 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, *, use_in_pass_through: bool = False) -> Router: return Router( model_list=[ { - "model_name": ROUTER_GROUP, - "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": f"key-{deployment_id}"}, + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": f"key-{deployment_id}", + "use_in_pass_through": use_in_pass_through, + }, "model_info": {"id": deployment_id}, } for deployment_id in ROUTER_DEPLOYMENT_IDS @@ -391,3 +475,575 @@ 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 + + + +class _SlowLocalRead(InMemoryCache): + """Answers a count read with the counts from when the read started, once another thread has begun its request.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + self.other_thread: threading.Thread | None = None + + def batch_get_cache(self, keys: list[str], **kwargs: object) -> list[object]: + counts_when_asked: Final = super().batch_get_cache(keys, **kwargs) + began: Final = threading.Event() + + def begin_then_run() -> None: + began.set() + self.while_reading() + + self.other_thread = threading.Thread(target=begin_then_run) + self.other_thread.start() + began.wait() + self.other_thread.join(timeout=0.2) + return counts_when_asked + + +def test_a_request_that_starts_on_another_thread_while_counts_are_read_still_counts_on_its_deployment() -> None: + local: Final = _SlowLocalRead() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(in_memory_cache=local)) + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + local.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + picked_while_the_other_thread_started: Final = worker.get_available_deployments(GROUP, HEALTHY, reserve=True) + if local.other_thread is not None: + local.other_thread.join() + + assert picked_while_the_other_thread_started 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_looking_up_a_pass_through_deployment_does_not_hold_it_for_later_picks() -> None: + router: Final = _least_busy_router(use_in_pass_through=True) + + selected: Final = [ + (await router.select(ROUTER_GROUP, {"metadata": {}}, kind="pass_through"))["model_info"]["id"] for _ in range(2) + ] + async_picks: Final = [ + (await router.async_get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}}))[ + "model_info" + ]["id"] + for _ in range(2) + ] + sync_picks: Final = [ + router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"]["id"] + for _ in range(2) + ] + + assert selected + async_picks + sync_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 afcd1e2a7ff..951182051c4 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 c16c1c62202..f5962cb4933 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -30102,11 +30102,21 @@ def test_reset_custom_routing_strategy(): assert "get_available_deployment" not in router._core.__dict__ assert "_core_async_get_available_deployment" not in router._core.__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 5ad6eb819ef..3bbb69413d9 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -844,6 +844,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(