scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 312

Merge ec7d9c61de into a392a39627

Browse source
This commit is contained in:
Oliver Yu 2026-10-31 06:52:17 +00:00 • committed by GitHub
commit b3038fae9a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1312 additions and 241 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

420
litellm/router.py
View file

@ -286,7 +286,15 @@ from litellm.router_strategy.complexity_router.context_compaction import (
reject_recursive_compactor,
surface_for_call,
)
from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
from litellm.router_strategy.least_busy import (
CALL_IDENTITY,
SELECTION_RESERVATION_MAX_KEYS,
LeastBusyLoggingHandler,
call_id_of,
owning_picks,
pick_reserves,
picks_reserve,
)
from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler
@ -1402,13 +1410,14 @@ class Router(DeprecatedRouterMembers):
(``pass_through``).
"""
if kind == "pick":
return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts
model=model,
request_kwargs=request_kwargs if request_kwargs is not None else {},
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
with picks_reserve(False):
return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts
model=model,
request_kwargs=request_kwargs if request_kwargs is not None else {},
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
if kind == "healthy":
return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts
model=model,
@ -1678,6 +1687,9 @@ class Router(DeprecatedRouterMembers):
self._core_cache = DualCache(
redis_cache=redis_cache, in_memory_cache=InMemoryCache()
) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc.
self._core_least_busy_reservations: InMemoryCache = InMemoryCache(
max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS
)
self._claude_code_session_router_cache: DualCache = DualCache(
redis_cache=redis_cache,
in_memory_cache=InMemoryCache(),
@ -2185,7 +2197,9 @@ class Router(DeprecatedRouterMembers):
selector: RouterStrategySelector | None = None
match self._normalize_strategy(strategy):
case RoutingStrategy.LEAST_BUSY.value:
selector = LeastBusyLoggingHandler(router_cache=self._core_cache)
selector = LeastBusyLoggingHandler(
router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations
)
case RoutingStrategy.USAGE_BASED_ROUTING.value:
selector = LowestTPMLoggingHandler(
router_cache=self._core_cache,
@ -2651,6 +2665,7 @@ class Router(DeprecatedRouterMembers):
messages: list[dict[str, str]] | None,
input: str | list | None,
request_kwargs: dict | None,
reserve: bool,
) -> Any | None:
"""
Asks the strategy selector for a deployment. Caller handles
@ -2664,6 +2679,8 @@ class Router(DeprecatedRouterMembers):
return await selector.async_get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
reserve=reserve,
call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})),
)
case "usage-based-routing":
# `LowestTPMLoggingHandler` (v1) only exposes the sync
@ -2704,6 +2721,7 @@ class Router(DeprecatedRouterMembers):
messages: list[dict[str, str]] | None,
input: str | list | None,
request_kwargs: dict | None,
reserve: bool,
) -> Any | None:
"""
Sync sibling of `_select_deployment_async`. Caller handles
@ -2720,6 +2738,8 @@ class Router(DeprecatedRouterMembers):
return selector.get_available_deployments(
model_group=model,
healthy_deployments=healthy_deployments,
reserve=reserve,
call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})),
)
case "usage-based-routing" | "usage-based-routing-v2":
return selector.get_available_deployments(
@ -3416,12 +3436,13 @@ class Router(DeprecatedRouterMembers):
)
# pick the one that is available (lowest TPM/RPM)
deployment = self.get_available_deployment(
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with picks_reserve(True):
deployment = self.get_available_deployment(
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._drop_unsupported_classifier_reasoning_effort(
deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment
model=model,
@ -5953,12 +5974,13 @@ class Router(DeprecatedRouterMembers):
model_name: Final = ""
try:
verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
deployment: Final = self.get_available_deployment(
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with picks_reserve(True):
deployment: Final = self.get_available_deployment(
model=model,
messages=[{"role": "user", "content": "prompt"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
data: Final = deployment["litellm_params"].copy()
@ -6326,37 +6348,38 @@ class Router(DeprecatedRouterMembers):
**kwargs,
):
messages: Final = [{"role": "user", "content": prompt}]
try:
kwargs["model"] = model
kwargs["prompt"] = prompt
kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
kwargs.setdefault("metadata", {}).update({"model_group": model})
with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
try:
kwargs["model"] = model
kwargs["prompt"] = prompt
kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries)
kwargs.setdefault("metadata", {}).update({"model_group": model})
# pick the one that is available (lowest TPM/RPM)
with endpoint_selection("text_completion", model=model):
deployment: Final = self.get_available_deployment(
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
# pick the one that is available (lowest TPM/RPM)
with endpoint_selection("text_completion", model=model), picks_reserve(True):
deployment: Final = self.get_available_deployment(
model=model,
messages=messages,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
data: Final = deployment["litellm_params"].copy()
for k, v in self._core_default_litellm_params.items():
if k not in kwargs: # prioritize model-specific params > default router params
kwargs[k] = v
elif k == "metadata":
kwargs[k].update(v)
pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data))
if pinned_model_id is not None:
kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults
# call via litellm.completion()
return litellm.text_completion(
**{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs}
)
data: Final = deployment["litellm_params"].copy()
for k, v in self._core_default_litellm_params.items():
if k not in kwargs: # prioritize model-specific params > default router params
kwargs[k] = v
elif k == "metadata":
kwargs[k].update(v)
pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data))
if pinned_model_id is not None:
kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults
# call via litellm.completion()
return litellm.text_completion(
**{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs}
)
except Exception as e:
raise e
except Exception as e:
raise e
async def atext_completion(
self,
@ -7471,61 +7494,66 @@ class Router(DeprecatedRouterMembers):
"""
handler_name: Final = original_function.__name__
metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call")
try:
verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model)
self._update_kwargs_before_fallbacks(
model=model,
kwargs=kwargs,
metadata_variable_name=metadata_variable_name,
)
with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name):
deployment: Final = self.get_available_deployment(
with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
try:
verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model)
self._update_kwargs_before_fallbacks(
model=model,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
kwargs=kwargs,
metadata_variable_name=metadata_variable_name,
)
with (
self._video_call_pin(handler_name, kwargs),
self._route_capability_pin(handler_name),
picks_reserve(True),
):
deployment: Final = self.get_available_deployment(
model=model,
messages=kwargs.get("messages", None),
input=kwargs.get("input", None),
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(
deployment=deployment,
kwargs=kwargs,
function_name="generic_api_call",
splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
)
self._update_kwargs_with_deployment(
deployment=deployment,
kwargs=kwargs,
function_name="generic_api_call",
splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
)
data: Final = deployment["litellm_params"].copy()
model_name: Final = data["model"]
data: Final = deployment["litellm_params"].copy()
model_name: Final = data["model"]
self.total_calls[model_name] += 1
self.total_calls[model_name] += 1
# For passthrough routes, use the actual model from deployment
# and swap model name in endpoint if present
if "endpoint" in kwargs and kwargs["endpoint"]:
kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name)
kwargs["model"] = model_name
# For passthrough routes, use the actual model from deployment
# and swap model name in endpoint if present
if "endpoint" in kwargs and kwargs["endpoint"]:
kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name)
kwargs["model"] = model_name
# Perform pre-call checks for routing strategy
self.routing_strategy_pre_call_checks(deployment=deployment)
# Perform pre-call checks for routing strategy
self.routing_strategy_pre_call_checks(deployment=deployment)
custom_llm_provider: Final = provider_for_generic_call(data)
custom_llm_provider: Final = provider_for_generic_call(data)
response: Final = original_function(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self._core_cache_responses,
**kwargs,
}
)
response: Final = original_function(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self._core_cache_responses,
**kwargs,
}
)
self.success_calls[model_name] += 1
verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name)
return response
except Exception as e:
verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e)
if model is not None:
self.fail_calls[model] += 1
raise e
self.success_calls[model_name] += 1
verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name)
return response
except Exception as e:
verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e)
if model is not None:
self.fail_calls[model] += 1
raise e
def embedding(
self,
@ -7548,7 +7576,7 @@ class Router(DeprecatedRouterMembers):
model_name = None
try:
verbose_router_logger.debug("Inside embedding()- model: %s", model)
with embedding_deployment_selection():
with embedding_deployment_selection(), picks_reserve(True):
deployment: Final = self.get_available_deployment(
model=model,
input=input,
@ -7850,68 +7878,71 @@ class Router(DeprecatedRouterMembers):
Returns:
VectorStoreCreateResponse
"""
try:
# If model is None, use the factory function approach (direct SDK call)
if model is None:
from litellm.vector_stores.main import acreate
with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
try:
# If model is None, use the factory function approach (direct SDK call)
if model is None:
from litellm.vector_stores.main import acreate
# Use the factory function to handle the call
factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create")
return await factory_fn(**kwargs)
# Use the factory function to handle the call
factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create")
return await factory_fn(**kwargs)
from litellm.vector_stores import acreate as avector_store_create_sdk
from litellm.vector_stores import acreate as avector_store_create_sdk
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "vector-store-api-fake-text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
data: Final = deployment["litellm_params"].copy()
model_name: Final = data["model"]
self._update_kwargs_with_deployment(
deployment=deployment,
kwargs=kwargs,
function_name="avector_store_create",
)
model_client: Final = self._get_async_openai_model_client(
deployment=deployment,
kwargs=kwargs,
)
self.total_calls[model_name] += 1
# Get custom provider from deployment params
custom_llm_provider = data.get("custom_llm_provider")
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
model=data["model"],
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
response = await avector_store_create_sdk(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self._core_cache_responses,
"client": model_client,
**kwargs,
}
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "vector-store-api-fake-text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
data: Final = deployment["litellm_params"].copy()
model_name: Final = data["model"]
self._update_kwargs_with_deployment(
deployment=deployment,
kwargs=kwargs,
function_name="avector_store_create",
)
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
model_client: Final = self._get_async_openai_model_client(
deployment=deployment,
kwargs=kwargs,
)
self.total_calls[model_name] += 1
return response
except Exception as e:
verbose_router_logger.exception(
"litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e
)
if model is not None:
self.fail_calls[model] += 1
raise e
# Get custom provider from deployment params
custom_llm_provider = data.get("custom_llm_provider")
_, inferred_custom_llm_provider, _, _ = get_llm_provider(
model=data["model"],
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider
async with self._deployment_slot(
deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span
):
response = await avector_store_create_sdk(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self._core_cache_responses,
"client": model_client,
**kwargs,
}
)
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
return response
except Exception as e:
verbose_router_logger.exception(
"litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e
)
if model is not None:
self.fail_calls[model] += 1
raise e
def _override_vector_store_methods_for_router(self):
"""
@ -8159,12 +8190,13 @@ class Router(DeprecatedRouterMembers):
try:
verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self._core_async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "batch-api-fake-text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with picks_reserve(False):
deployment: Final = await self._core_async_get_available_deployment(
model=model,
messages=[{"role": "user", "content": "batch-api-fake-text"}],
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id")
data: Final = deployment["litellm_params"].copy()
@ -9298,38 +9330,39 @@ class Router(DeprecatedRouterMembers):
mock_timeout: Final = kwargs.pop("mock_timeout", None)
try:
self._handle_mock_testing_fallbacks(
kwargs=kwargs,
model_group=model_group,
fallbacks=fallbacks,
context_window_fallbacks=context_window_fallbacks,
content_policy_fallbacks=content_policy_fallbacks,
)
with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))):
try:
self._handle_mock_testing_fallbacks(
kwargs=kwargs,
model_group=model_group,
fallbacks=fallbacks,
context_window_fallbacks=context_window_fallbacks,
content_policy_fallbacks=content_policy_fallbacks,
)
if mock_timeout is not None:
response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout)
else:
response = await self.async_function_with_retries(*args, **kwargs)
if verbose_router_logger.isEnabledFor(logging.DEBUG):
verbose_router_logger.debug("Async Response: %s", response)
response = add_fallback_headers_to_response(
response=response,
attempted_fallbacks=0,
)
return response
except Exception as e:
return await self.async_function_with_fallbacks_common_utils(
e,
disable_fallbacks,
fallbacks,
context_window_fallbacks,
content_policy_fallbacks,
model_group,
args,
kwargs,
include_fallback_errors=include_fallback_errors,
)
if mock_timeout is not None:
response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout)
else:
response = await self.async_function_with_retries(*args, **kwargs)
if verbose_router_logger.isEnabledFor(logging.DEBUG):
verbose_router_logger.debug("Async Response: %s", response)
response = add_fallback_headers_to_response(
response=response,
attempted_fallbacks=0,
)
return response
except Exception as e:
return await self.async_function_with_fallbacks_common_utils(
e,
disable_fallbacks,
fallbacks,
context_window_fallbacks,
content_policy_fallbacks,
model_group,
args,
kwargs,
include_fallback_errors=include_fallback_errors,
)
def _handle_mock_testing_fallbacks(
self,
@ -10618,6 +10651,7 @@ class Router(DeprecatedRouterMembers):
messages=None,
input=None,
request_kwargs=routing_kwargs,
reserve=False,
)
if selected is None:
no_deployment: Final = await async_raise_no_deployment_exception(
@ -15836,6 +15870,7 @@ class Router(DeprecatedRouterMembers):
messages=messages,
input=input,
specific_deployment=specific_deployment,
reserve=pick_reserves(default=True),
)
async def _async_get_available_deployment(
@ -15845,6 +15880,7 @@ class Router(DeprecatedRouterMembers):
messages: list[dict[str, str]] | None,
input: str | list | None,
specific_deployment: bool | None,
reserve: bool,
):
if (
self._core_routing_strategy != "usage-based-routing-v2"
@ -15941,6 +15977,7 @@ class Router(DeprecatedRouterMembers):
messages=messages,
input=input,
request_kwargs=request_kwargs,
reserve=reserve and request_kwargs not in (None, {}),
)
if deployment is None:
exception: Final = await async_raise_no_deployment_exception(
@ -16110,6 +16147,7 @@ class Router(DeprecatedRouterMembers):
messages=messages,
input=input,
request_kwargs=request_kwargs,
reserve=request_kwargs not in (None, {}),
)
if deployment is None:
@ -16983,6 +17021,7 @@ class Router(DeprecatedRouterMembers):
messages=messages,
input=input,
request_kwargs=request_kwargs,
reserve=pick_reserves(default=False) and request_kwargs not in (None, {}),
)
if deployment is None:
@ -17163,6 +17202,7 @@ class Router(DeprecatedRouterMembers):
messages=messages,
input=input,
request_kwargs=request_kwargs,
reserve=request_kwargs not in (None, {}),
)
if deployment is None:

20
litellm/router_deprecated.py
View file

@ -27,11 +27,12 @@ import inspect
import sys
import warnings
from collections import Counter
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload
from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload
from litellm.router_api import SETTING_FIELDS, ChangeSet
from litellm.router_strategy.least_busy import picks_reserve
if TYPE_CHECKING:
from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged
@ -58,6 +59,7 @@ if TYPE_CHECKING:
from litellm.types.utils import ModelInfo as ModelMapInfo
_T: Final = TypeVar("_T")
_P: Final = ParamSpec("_P")
DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: the live tally _note increments
@ -163,6 +165,18 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib
return _DeprecatedAttribute(replacement, read)
def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coroutine[object, object, _T]]:
@functools.wraps(lookup)
async def looked_up(
*args: _P.args,
**kwargs: _P.kwargs, # kwargs-ok: forwards the wrapped lookup's ParamSpec unchanged
) -> _T:
with picks_reserve(False):
return await lookup(*args, **kwargs)
return looked_up
# (id(router), member name) pairs whose subclass override is running in this context.
_OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar(
"router_overrides_running", default=frozenset()
@ -325,7 +339,7 @@ class DeprecatedRouterMembers:
deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map)
async_get_available_deployment = _attr(
"select(model, request_kwargs)",
lambda r: r._core_async_get_available_deployment,
lambda r: _without_reserving(r._core_async_get_available_deployment),
)
# -- observers ------------------------------------------------------------------------------

313
litellm/router_strategy/least_busy.py
View file

@ -1,6 +1,11 @@
import contextlib
import logging
from collections.abc import Mapping, Sequence
from typing import Final
import threading
import time
import uuid
from collections.abc import Callable, Generator, Mapping, Sequence
from contextvars import ContextVar
from typing import Final, Protocol, TypeAlias
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
@ -8,11 +13,30 @@ from typing_extensions import ReadOnly, TypedDict
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import log_redis_failure
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0
SELECTION_RESERVATION_MAX_KEYS: Final = 10_000
_RESERVATIONS_LOCK: Final = threading.Lock()
_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves")
_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None)
@contextlib.contextmanager
def picks_reserve(reserve: bool) -> Generator[None, None, None]:
token: Final = _PICK_RESERVES.set(reserve)
try:
yield
finally:
_PICK_RESERVES.reset(token)
def pick_reserves(default: bool) -> bool:
return _PICK_RESERVES.get(default)
class _ModelInfo(TypedDict, total=False):
@ -21,10 +45,12 @@ class _ModelInfo(TypedDict, total=False):
class _Metadata(TypedDict, total=False):
model_group: ReadOnly[str | None]
model_info: ReadOnly[_ModelInfo | None]
class _LitellmParams(TypedDict, total=False):
metadata: ReadOnly[_Metadata | None]
litellm_metadata: ReadOnly[_Metadata | None]
model_info: ReadOnly[_ModelInfo | None]
@ -32,6 +58,10 @@ class _CallKwargs(TypedDict, total=False):
litellm_params: ReadOnly[_LitellmParams | None]
class CallIdentity(TypedDict, total=False):
litellm_call_id: ReadOnly[object]
class _DeploymentModelInfo(TypedDict):
id: ReadOnly[str | int]
@ -41,14 +71,31 @@ class _Deployment(TypedDict):
_CALL_KWARGS: Final = TypeAdapter(_CallKwargs)
CALL_IDENTITY: Final = TypeAdapter(CallIdentity)
_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None)
_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None)
_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None)
_Pick: TypeAlias = tuple[float, str | None]
_PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None)
_PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None)
def _request_count_key(model_group: str, deployment_id: str) -> str:
return f"{model_group}_request_count:{deployment_id}"
def _model_group(metadata: _Metadata | None) -> str | None:
return metadata.get("model_group") if metadata else None
def _deployment_id(model_info: _ModelInfo | None) -> str | int | None:
return model_info.get("id") if model_info else None
def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None:
return _deployment_id(metadata.get("model_info")) if metadata else None
def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return None
@ -56,23 +103,52 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
call: Final = _CALL_KWARGS.validate_python(kwargs)
except ValidationError:
return None
litellm_params: Final = call.get("litellm_params")
metadata: Final = litellm_params.get("metadata") if litellm_params else None
model_info: Final = litellm_params.get("model_info") if litellm_params else None
model_group: Final = metadata.get("model_group") if metadata else None
deployment_id: Final = model_info.get("id") if model_info else None
litellm_params: Final = call.get("litellm_params") or _LitellmParams()
routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata"))
model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None)
deployment_id: Final = next(
(
deployment_id
for deployment_id in (
_deployment_id(litellm_params.get("model_info")),
*map(_routed_deployment_id, routing_metadata),
)
if deployment_id is not None
),
None,
)
if model_group is None or deployment_id is None:
return None
return model_group, str(deployment_id)
def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
def _reservation_key(deployment_id: str) -> str:
return f"least_busy_reservations:{deployment_id}"
def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str:
return f"least_busy_started:{call_id}:{_request_count_key(*ref)}"
def _picked_deployments_key(call_id: str) -> str:
return f"least_busy_picked:{call_id}"
def call_id_of(kwargs: Mapping[str, object]) -> str | None:
call_id: Final = kwargs.get("litellm_call_id")
return call_id if isinstance(call_id, str) else None
def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
return tuple(
_request_count_key(model_group, str(deployment["model_info"]["id"]))
for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
)
def _request_count_keys(model_group: str, deployment_ids: tuple[str, ...]) -> tuple[str, ...]:
return tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids)
def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]:
return tuple(0 if value is None else int(value) for value in values)
@ -84,12 +160,21 @@ def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]:
return _as_counts(values)
def _least_busy(
healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...]
) -> Mapping[str, object] | None:
if not healthy_deployments:
return None
return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])]
def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]:
return tuple(pick for pick in picks if now - pick[0] < SELECTION_RESERVATION_TTL_SECONDS)
def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]:
tags: Final = tuple(tag for _, tag in picks)
owner: Final = next((tag for tag in (call_id, _PICK_OWNER.get()) if tag is not None and tag in tags), None)
if owner not in tags:
return picks
index: Final = tags.index(owner)
return (*picks[:index], *picks[index + 1 :])
def _least_busy_index(counts: tuple[int, ...]) -> int:
return min(range(len(counts)), key=lambda index: counts[index])
def _warn_unreadable(model_group: str, error: Exception) -> None:
@ -111,24 +196,148 @@ def _warn_unwritable(key: str, error: Exception) -> None:
)
class _ReservationStore(Protocol):
def get_cache(self, key: str) -> object: ...
def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None: ...
def delete_cache(self, key: str) -> None: ...
def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]:
return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or ()
def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]:
return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or ()
def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]:
return _PICKED_DEPLOYMENTS.validate_python(store.get_cache(_picked_deployments_key(call_id))) or ()
def _put(store: _ReservationStore, key: str, values: tuple[object, ...], ttl: float) -> None:
# InMemoryCache keeps a live key's first TTL, so drop the key to restart it
store.delete_cache(key)
if values:
store.set_cache(key, values, ttl=ttl)
def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str) -> None:
kept: Final = tuple(pick for pick in _picks(store, deployment_id) if pick[1] != call_id)
_put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS)
def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None:
if call_id is None:
return
if not _picked_deployments(store, call_id):
return
with _RESERVATIONS_LOCK:
for deployment_id in frozenset(_picked_deployments(store, call_id)):
_release_call_on(store, deployment_id, call_id)
store.delete_cache(_picked_deployments_key(call_id))
@contextlib.contextmanager
def owning_picks(store: _ReservationStore, call_id: str | None) -> Generator[None, None, None]:
owner: Final = call_id or _PICK_OWNER.get() or uuid.uuid4().hex
token: Final = _PICK_OWNER.set(owner)
try:
yield
finally:
_PICK_OWNER.reset(token)
release_unstarted_picks(store, owner)
class LeastBusyLoggingHandler(CustomLogger):
test_flag: bool = False
logged_success: int = 0
logged_failure: int = 0
def __init__(self, router_cache: DualCache):
def __init__(
self,
router_cache: DualCache,
clock: Callable[[], float] = time.monotonic,
reservation_store: _ReservationStore | None = None,
):
self.router_cache = router_cache
self.router_cache_id = str(id(router_cache))
self._clock = clock
self._reservations: Final[_ReservationStore] = (
reservation_store
if reservation_store is not None
else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock)
)
def _live_reservations(self, deployment_id: str, now: float) -> tuple[_Pick, ...]:
return _unexpired(_picks(self._reservations, deployment_id), now)
def _store_reservations(self, deployment_id: str, picks: tuple[_Pick, ...]) -> None:
_put(self._reservations, _reservation_key(deployment_id), picks, SELECTION_RESERVATION_TTL_SECONDS)
def _take_reservation(self, deployment_id: str, call_id: str | None, now: float) -> None:
self._store_reservations(deployment_id, _without_one_pick(self._live_reservations(deployment_id, now), call_id))
def _record_pick(self, call_id: str, deployment_id: str) -> None:
picked: Final = (*_picked_deployments(self._reservations, call_id), deployment_id)
_put(self._reservations, _picked_deployments_key(call_id), picked, SELECTION_RESERVATION_TTL_SECONDS)
def _select(
self,
healthy_deployments: Sequence[Mapping[str, object]],
deployment_ids: tuple[str, ...],
counts: tuple[int, ...],
reserve: bool,
call_id: str | None,
) -> Mapping[str, object] | None:
if not healthy_deployments:
return None
with _RESERVATIONS_LOCK:
now: Final = self._clock()
live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids)
index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True)))
if reserve:
tag: Final = call_id if call_id is not None else _PICK_OWNER.get()
self._store_reservations(deployment_ids[index], (*live[index], (now, tag)))
if tag is not None:
self._record_pick(tag, deployment_ids[index])
return healthy_deployments[index]
def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None:
with _RESERVATIONS_LOCK:
now: Final = self._clock()
self._take_reservation(ref[1], call_id, now)
if call_id is not None:
key: Final = _started_attempts_key(call_id, ref)
started: Final = (*_attempt_times(self._reservations, key), now)
_put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS)
def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool:
if call_id is None:
return True
key: Final = _started_attempts_key(call_id, ref)
with _RESERVATIONS_LOCK:
started: Final = _attempt_times(self._reservations, key)
if started:
_put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS)
else:
self._take_reservation(ref[1], call_id, self._clock())
return bool(started)
@with_service_target("router_usage")
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
self._increment(kwargs, 1)
ref: Final = _deployment_ref(kwargs)
if ref is None:
return
self._start_attempt(ref, call_id_of(kwargs))
self._increment(_request_count_key(*ref), 1)
@with_service_target("router_usage")
def log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._increment(kwargs, -1)
self._settle(kwargs)
if self.test_flag:
self.logged_success += 1
@ -136,7 +345,7 @@ class LeastBusyLoggingHandler(CustomLogger):
def log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self._increment(kwargs, -1)
self._settle(kwargs)
if self.test_flag:
self.logged_failure += 1
@ -144,7 +353,7 @@ class LeastBusyLoggingHandler(CustomLogger):
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
await self._async_increment(kwargs, -1)
await self._async_settle(kwargs)
if self.test_flag:
self.logged_success += 1
@ -152,15 +361,20 @@ class LeastBusyLoggingHandler(CustomLogger):
async def async_log_failure_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
await self._async_increment(kwargs, -1)
await self._async_settle(kwargs)
if self.test_flag:
self.logged_failure += 1
@with_service_target("router_usage")
def get_available_deployments(
self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
self,
model_group: str,
healthy_deployments: Sequence[Mapping[str, object]],
reserve: bool = False,
call_id: str | None = None,
) -> Mapping[str, object] | None:
keys: Final = _request_count_keys(model_group, healthy_deployments)
deployment_ids: Final = _deployment_ids(healthy_deployments)
keys: Final = _request_count_keys(model_group, deployment_ids)
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is not None:
try:
@ -168,15 +382,21 @@ class LeastBusyLoggingHandler(CustomLogger):
except Exception as e:
_warn_unreadable(model_group, e)
else:
return _least_busy(healthy_deployments, shared)
local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
return _least_busy(healthy_deployments, local)
return self._select(
healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve, call_id
)
return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve, call_id)
@with_service_target("router_usage")
async def async_get_available_deployments(
self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]]
self,
model_group: str,
healthy_deployments: Sequence[Mapping[str, object]],
reserve: bool = False,
call_id: str | None = None,
) -> Mapping[str, object] | None:
keys: Final = _request_count_keys(model_group, healthy_deployments)
deployment_ids: Final = _deployment_ids(healthy_deployments)
keys: Final = _request_count_keys(model_group, deployment_ids)
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is not None:
try:
@ -184,15 +404,30 @@ class LeastBusyLoggingHandler(CustomLogger):
except Exception as e:
_warn_unreadable(model_group, e)
else:
return _least_busy(healthy_deployments, shared)
local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys)
return _least_busy(healthy_deployments, local)
return self._select(
healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve, call_id
)
return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve, call_id)
def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
def _settle(self, kwargs: Mapping[str, object]) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None:
if ref is None or not self._end_attempt(ref, call_id_of(kwargs)):
return
key: Final = _request_count_key(*ref)
self._increment(_request_count_key(*ref), -1)
async def _async_settle(self, kwargs: Mapping[str, object]) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None or not self._end_attempt(ref, call_id_of(kwargs)):
return
await self._async_increment(_request_count_key(*ref), -1)
def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]:
return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]:
return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True))
def _increment(self, key: str, delta: int) -> None:
redis_cache: Final = self.router_cache.redis_cache
try:
local: Final = self.router_cache.increment_cache(
@ -206,11 +441,7 @@ class LeastBusyLoggingHandler(CustomLogger):
except Exception as e:
_warn_unwritable(key, e)
async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None:
return
key: Final = _request_count_key(*ref)
async def _async_increment(self, key: str, delta: int) -> None:
redis_cache: Final = self.router_cache.redis_cache
try:
local: Final = await self.router_cache.async_increment_cache(

90
tests/integration/routing/test_least_busy_concurrent_decisions.py Normal file
View file

@ -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}
)

39
tests/unit/decisions/test_main.py
View file

@ -6,6 +6,8 @@ import copy
import json
import re
import threading
import uuid
from datetime import datetime
from collections.abc import Mapping, Sequence
from concurrent.futures import wait
from types import MappingProxyType
@ -297,6 +299,43 @@ async def test_router_dispatches_typesafe_decisions_without_api_base(
assert response.answers["sentiment"].choice == "positive"
@pytest.mark.asyncio
async def test_least_busy_frees_a_decisions_deployment_once_a_proxy_logged_call_finishes(
respx_mock: respx.MockRouter,
) -> None:
respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE)
router: Final = litellm.Router(
model_list=[
{
"model_name": "jev",
"litellm_params": {"model": "typesafe/jev-latest", "api_key": "k"},
"model_info": {"id": deployment_id},
}
for deployment_id in ("jev-a", "jev-b")
],
routing_strategy="least-busy",
)
async def call_through_proxy_logging() -> object:
logging_obj, _ = litellm.utils.function_setup(
original_function="adecisions",
rules_obj=litellm.utils.Rules(),
start_time=datetime(2026, 1, 1),
model="jev",
litellm_call_id=str(uuid.uuid4()),
)
response: Final = await router.adecisions(
model="jev",
input="Is this a defect?",
questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
litellm_logging_obj=logging_obj,
)
await _drain_logging_worker()
return response._hidden_params["model_id"]
assert (await call_through_proxy_logging(), await call_through_proxy_logging()) == ("jev-a", "jev-a")
def test_decisions_uses_the_same_wire_contract_for_sync_calls(respx_mock: respx.MockRouter) -> None:
route: Final = respx_mock.post("https://api.perplexity.ai/v1/decisions").respond(json=_RESPONSE)

594
tests/unit/router_strategy/test_least_busy.py
View file

@ -1,9 +1,16 @@
import asyncio
import contextlib
import contextvars
import datetime
import json
import logging
import threading
from collections.abc import Callable, Generator
from concurrent.futures import ThreadPoolExecutor
from typing import Final
import httpx
import openai
import pytest
import respx
@ -11,10 +18,18 @@ import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler
from litellm.router_strategy.least_busy import (
IN_FLIGHT_COUNT_TTL_SECONDS,
SELECTION_RESERVATION_TTL_SECONDS,
LeastBusyLoggingHandler,
owning_picks,
picks_reserve,
release_unstarted_picks,
)
from litellm.types.utils import ModelResponse
GROUP: Final = "least-busy-group"
@ -23,8 +38,11 @@ DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}}
HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B]
def _call_kwargs(deployment_id: str) -> dict[str, object]:
return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}}
def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]:
return {
"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}},
"litellm_call_id": call_id,
}
class SharedRedisCounters:
@ -58,6 +76,20 @@ class SharedRedisCounters:
return self.batch_get_counts(key_list)
class SlowRedisRead(SharedRedisCounters):
"""Answers a count read with the counts from when the read started, after letting other work run."""
def __init__(self) -> None:
super().__init__()
self.while_reading: Callable[[], None] = lambda: None
async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
counts_when_asked: Final = self.batch_get_counts(key_list)
await asyncio.sleep(0)
self.while_reading()
return counts_when_asked
def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler:
cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double
return LeastBusyLoggingHandler(router_cache=cache)
@ -81,6 +113,16 @@ async def test_worker_routes_around_a_request_another_worker_started() -> None:
assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
@pytest.mark.asyncio
async def test_a_request_that_starts_while_redis_is_read_still_counts_on_its_deployment() -> None:
redis: Final = SlowRedisRead()
worker: Final = _worker(redis)
assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
redis.while_reading = lambda: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
def test_sync_pick_reads_the_shared_counts() -> None:
shared: Final = SharedRedisCounters()
streaming_worker: Final = _worker(shared)
@ -98,6 +140,32 @@ def test_sync_pick_reads_the_shared_counts() -> None:
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]:
return {
"litellm_params": {
"metadata": {"tags": ["caller-tag"]},
"litellm_metadata": {"model_group": GROUP},
"model_info": {"id": deployment_id},
},
"litellm_call_id": "call-1",
}
def test_a_request_routed_through_litellm_metadata_counts_while_in_flight() -> None:
shared: Final = SharedRedisCounters()
calling_worker: Final = _worker(shared)
picking_worker: Final = _worker(shared)
assert calling_worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
calling_worker.log_pre_api_call(model="m", messages=[], kwargs=_routed_through_litellm_metadata("dep-a"))
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
calling_worker.log_success_event(_routed_through_litellm_metadata("dep-a"), None, None, None)
assert calling_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_the_handler_never_pushes_a_counters_ttl_forward() -> None:
"""A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that
stuck count an hour after the key was created, which only works while nothing writes the TTL
@ -259,11 +327,11 @@ ROUTER_GROUP: Final = "least-busy-router-group"
ROUTER_DEPLOYMENT_IDS: Final = ("1", "2", "3")
def _least_busy_router() -> Router:
def _least_busy_router(group: str = ROUTER_GROUP) -> Router:
return Router(
model_list=[
{
"model_name": ROUTER_GROUP,
"model_name": group,
"litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": f"key-{deployment_id}"},
"model_info": {"id": deployment_id},
}
@ -391,3 +459,517 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts
assert response._hidden_params["model_id"] == "1"
assert route.calls.last.request.headers["authorization"] == "Bearer key-1"
assert _router_counts(router) == seeded
@pytest.mark.asyncio
async def test_router_spreads_requests_that_start_together(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
burst: Final = 2 * len(ROUTER_DEPLOYMENT_IDS)
all_started: Final = asyncio.Barrier(burst)
async def completion_after_all_started(_: httpx.Request) -> httpx.Response:
await all_started.wait()
return httpx.Response(
200,
json={
"id": "chatcmpl-least-busy",
"object": "chat.completion",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
},
)
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
side_effect=completion_after_all_started
)
router: Final = _least_busy_router()
await asyncio.gather(
*(router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(burst))
)
one_per_deployment: Final = tuple(f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS)
assert sorted(call.request.headers["authorization"] for call in route.calls) == sorted(one_per_deployment * 2)
class _ManualClock:
def __init__(self) -> None:
self.now = 0.0
def __call__(self) -> float:
return self.now
def test_a_reservation_that_never_reaches_the_api_call_expires() -> None:
clock: Final = _ManualClock()
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), clock=clock)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
clock.now = 1.0
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
clock.now = 2.0
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.5
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
def test_a_fresh_reservation_outlives_the_ttl_of_the_first_one_on_its_deployment() -> None:
clock: Final = _ManualClock()
cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
worker: Final = LeastBusyLoggingHandler(router_cache=cache, clock=clock)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
worker.log_success_event(kwargs=_call_kwargs("dep-a"), response_obj=None, start_time=None, end_time=None)
clock.now = SELECTION_RESERVATION_TTL_SECONDS - 0.1
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
clock.now = SELECTION_RESERVATION_TTL_SECONDS + 0.1
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> None:
worker: Final = _worker(None)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None:
cache: Final = DualCache()
reservations: Final = InMemoryCache()
picker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations)
registered: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations)
assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
registered.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
class _StoreWhereThreadsReadBeforeWriting:
def __init__(self) -> None:
self._store: Final = InMemoryCache()
self._first_reads: Final = threading.Barrier(2, timeout=0.5)
self._thread_state: Final = threading.local()
self._meeting: Final = threading.Event()
@contextlib.contextmanager
def first_reads_meet(self) -> Generator[None, None, None]:
self._meeting.set()
yield
self._meeting.clear()
def get_cache(self, key: str) -> object:
value: Final = self._store.get_cache(key)
if self._meeting.is_set() and not getattr(self._thread_state, "has_read", False):
self._thread_state.has_read = True
with contextlib.suppress(threading.BrokenBarrierError):
self._first_reads.wait()
return value
def set_cache(self, key: str, value: tuple[object, ...], *, ttl: float) -> None:
self._store.set_cache(key, value, ttl=ttl)
def delete_cache(self, key: str) -> None:
self._store.delete_cache(key)
def test_two_threads_picking_at_once_reserve_different_deployments() -> None:
store: Final = _StoreWhereThreadsReadBeforeWriting()
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store)
with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool:
picks: Final = tuple(pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2)))
assert sorted(pick["model_info"]["id"] for pick in picks if pick is not None) == ["dep-a", "dep-b"]
def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> None:
store: Final = _StoreWhereThreadsReadBeforeWriting()
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b"))
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b"))
with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool:
reserved: Final = pool.submit(worker.get_available_deployments, GROUP, HEALTHY, True)
started: Final = pool.submit(worker.log_pre_api_call, "m", [], _call_kwargs("dep-a"))
assert reserved.result() is DEPLOYMENT_A
started.result()
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
def _router_pick(router: Router, group: str = ROUTER_GROUP) -> str:
deployment: Final = router.get_available_deployment(group, request_kwargs={"metadata": {}})
return deployment["model_info"]["id"]
def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None:
router: Final = _least_busy_router()
with picks_reserve(True):
first: Final = _router_pick(router)
for index in range(DEFAULT_MAX_SIZE_IN_MEMORY):
router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60)
assert _router_pick(router) != first
@pytest.mark.asyncio
async def test_looking_up_a_deployment_does_not_hold_it_for_later_picks() -> None:
router: Final = _least_busy_router()
sync_picks: Final = [_router_pick(router) for _ in range(3)]
async_picks: Final = [
(await router.async_get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}))["model_info"]["id"]
for _ in range(3)
]
assert sync_picks + async_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6
@pytest.mark.asyncio
async def test_a_cache_hit_releases_the_pick_it_reserved() -> None:
worker: Final = _worker(None)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
await worker.async_log_success_event({**_call_kwargs("dep-b"), "cache_hit": True}, None, None, None)
assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> None:
worker: Final = _worker(None)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
worker.log_success_event({**_call_kwargs("dep-a", "cached-call"), "cache_hit": True}, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
@pytest.mark.parametrize("outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event])
def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick(
outcome: Callable[[LeastBusyLoggingHandler, dict[str, object], None, None, None], None],
) -> None:
worker: Final = _worker(None)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
outcome(worker, _call_kwargs("dep-a"), None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behind() -> None:
worker: Final = _worker(None)
rejected: Final = {**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True}
worker.log_pre_api_call(model="m", messages=[], kwargs=rejected)
worker.log_failure_event(rejected, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None:
worker: Final = _worker(None)
assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
worker.log_success_event(_call_kwargs("dep-a"), None, None, None)
assert worker.get_available_deployments("tier-group", HEALTHY) is DEPLOYMENT_A
def test_a_fallback_attempt_that_fails_before_reaching_its_provider_frees_its_pick() -> None:
worker: Final = _worker(None)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "busy-call"))
first_attempt_started_at: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
first_attempt: Final = {**_call_kwargs("dep-a", "fallback-call"), "api_call_start_time": first_attempt_started_at}
worker.log_pre_api_call(model="m", messages=[], kwargs=first_attempt)
worker.log_failure_event(first_attempt, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
second_attempt: Final = {**_call_kwargs("dep-b", "fallback-call"), "api_call_start_time": first_attempt_started_at}
worker.log_failure_event(second_attempt, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
def test_ending_one_request_frees_only_the_picks_it_never_started() -> None:
reservations: Final = InMemoryCache()
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call"))
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="ended-call") is DEPLOYMENT_A
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="running-call") is DEPLOYMENT_A
release_unstarted_picks(reservations, "ended-call")
assert [worker.get_available_deployments(GROUP, HEALTHY, reserve=True) for _ in range(2)] == [
DEPLOYMENT_A,
DEPLOYMENT_B,
]
def test_a_call_without_a_call_id_turns_its_own_pick_into_its_in_flight_count() -> None:
reservations: Final = InMemoryCache()
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call"))
with owning_picks(reservations, None):
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "attempt-call"))
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_another_call_starting_leaves_the_pick_of_a_call_without_a_call_id_reserved() -> None:
reservations: Final = InMemoryCache()
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations)
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call"))
other_call: Final = contextvars.copy_context()
with owning_picks(reservations, None):
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
other_call.run(worker.log_pre_api_call, model="m", messages=[], kwargs=_call_kwargs("dep-a", "other-call"))
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
def test_a_call_started_without_a_pick_of_its_own_leaves_another_calls_pick_reserved() -> None:
worker: Final = _worker(None)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="picked-call") is DEPLOYMENT_A
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "looked-up-call"))
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
def _fallback_router() -> Router:
return Router(
model_list=[
{
"model_name": "unreachable-group",
"litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "key-unreachable"},
"model_info": {"id": "unreachable"},
},
*(
{
"model_name": ROUTER_GROUP,
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": f"key-{deployment_id}"},
"model_info": {"id": deployment_id},
}
for deployment_id in ROUTER_DEPLOYMENT_IDS
),
],
routing_strategy="least-busy",
num_retries=0,
fallbacks=[{"unreachable-group": [ROUTER_GROUP]}],
)
@pytest.mark.asyncio
async def test_a_fallback_hop_rejected_before_its_provider_frees_its_pick_when_every_hop_shares_one_logging_object(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(500, json={"error": {"message": "upstream down"}})
)
router: Final = _fallback_router()
logging_obj, request = litellm.utils.function_setup(
original_function="acompletion",
rules_obj=litellm.utils.Rules(),
start_time=datetime.datetime(2026, 1, 1),
model="unreachable-group",
messages=[{"role": "user", "content": "hi"}],
logit_bias={"1": 1},
litellm_call_id="fallback-call",
)
with pytest.raises(litellm.InternalServerError):
await router.acompletion(**request, litellm_logging_obj=logging_obj)
await _drain_logging()
assert {call.request.headers["authorization"] for call in route.calls} == {"Bearer key-unreachable"}
assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0]
def test_a_retry_that_starts_before_the_failed_attempt_is_logged_counts_each_attempt_once() -> None:
worker: Final = _worker(None)
attempt: Final = _call_kwargs("dep-a", "retried-call")
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
worker.log_failure_event(attempt, None, None, None)
worker.log_success_event(attempt, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
async def _drain_logging() -> None:
for _ in range(10):
await asyncio.sleep(0)
await GLOBAL_LOGGING_WORKER.flush()
def _called_deployment(route: respx.Route) -> str:
return route.calls.last.request.headers["authorization"].removeprefix("Bearer key-")
_BATCH: Final = {
"id": "batch_least_busy",
"object": "batch",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-least-busy",
"completion_window": "24h",
"status": "validating",
"created_at": 1,
}
_REJECTED_BATCH: Final = {"error": {"message": "Invalid input file", "type": "invalid_request_error"}}
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)])
async def test_a_batch_create_does_not_keep_its_deployment_busy(
provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.openai.com/v1/batches").mock(return_value=provider_response)
router: Final = _least_busy_router()
with contextlib.suppress(openai.BadRequestError):
await router.acreate_batch(
model=ROUTER_GROUP, completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-least-busy"
)
await _drain_logging()
assert _router_pick(router) == _called_deployment(route)
_VECTOR_STORE: Final = {
"id": "vs_least_busy",
"object": "vector_store",
"created_at": 1,
"name": "least-busy",
"usage_bytes": 0,
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
"status": "completed",
}
_REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}}
_CALL_IDENTITIES: Final = ({"litellm_call_id": "least-busy-call"}, {})
@pytest.mark.asyncio
@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES)
@pytest.mark.parametrize(
"provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)]
)
async def test_a_vector_store_create_does_not_keep_its_deployment_busy(
provider_response: httpx.Response,
call_identity: dict[str, str],
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response)
router: Final = _least_busy_router()
with contextlib.suppress(openai.BadRequestError):
await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", **call_identity)
await _drain_logging()
assert _router_pick(router) == _called_deployment(route)
@pytest.mark.asyncio
@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES)
async def test_listing_fine_tuning_jobs_does_not_keep_its_deployment_busy(
call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.get(url__startswith="https://api.openai.com/v1/fine_tuning/jobs").mock(
return_value=httpx.Response(200, json={"object": "list", "data": [], "has_more": False})
)
router: Final = _least_busy_router()
await router.alist_fine_tuning_jobs(model=ROUTER_GROUP, **call_identity)
await _drain_logging()
assert _router_pick(router) == _called_deployment(route)
TEXT_GROUP: Final = "openai/least-busy-text-group"
@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES)
def test_a_sync_text_completion_does_not_keep_its_deployment_busy(
call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.openai.com/v1/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "cmpl-least-busy",
"object": "text_completion",
"created": 1,
"model": "least-busy-text-group",
"choices": [{"text": "ok", "index": 0, "finish_reason": "stop", "logprobs": None}],
"usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3},
},
)
)
router: Final = _least_busy_router(TEXT_GROUP)
router.text_completion(model=TEXT_GROUP, prompt="Say ok", **call_identity)
assert _router_pick(router, TEXT_GROUP) == _called_deployment(route)
@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES)
def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its_pick(
call_identity: dict[str, str],
) -> None:
router: Final = _least_busy_router()
with pytest.raises(litellm.InternalServerError):
router.responses(model=ROUTER_GROUP, input="hi", api_base="not-a-url", **call_identity)
assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0]
@pytest.mark.asyncio
async def test_a_batch_cancel_does_not_keep_its_deployment_busy(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post("https://api.openai.com/v1/batches/batch_least_busy/cancel").mock(
return_value=httpx.Response(200, json={**_BATCH, "status": "cancelling"})
)
router: Final = _least_busy_router()
await router.acancel_batch(model=ROUTER_GROUP, batch_id="batch_least_busy")
assert _router_pick(router) == _called_deployment(route)
@pytest.mark.asyncio
async def test_picking_a_deployment_without_calling_it_does_not_keep_it_busy() -> None:
router: Final = _least_busy_router()
picked: Final = await router.select(ROUTER_GROUP, {"metadata": {}})
assert _router_pick(router) == picked["model_info"]["id"]

37
tests/unit/router_strategy/test_router_routing_groups.py
View file

@ -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"

12
tests/unit/test_router/test_router.py
View file

@ -29146,11 +29146,21 @@ def test_reset_custom_routing_strategy():
assert "get_available_deployment" not in router.__dict__
assert "async_get_available_deployment" not in router.__dict__
assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment
assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment
router._reset_custom_routing_strategy()
@pytest.mark.asyncio
async def test_select_picks_through_a_custom_routing_strategy() -> None:
router: Final = _create_custom_routing_router()
router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router))
picked: Final = await router.select("azure-model", {"metadata": {}})
assert picked["model_info"]["id"] == "very-special-endpoint"
@pytest.mark.asyncio
async def test_azure_router_acompletion_passes_deployment_credentials_for_plain_and_stream(
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter

28
tests/unit/test_router/test_router_helper_utils.py
View file

@ -599,6 +599,34 @@ async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_b
await admit_turn()
@pytest.mark.asyncio
async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy():
router: Final = Router(
model_list=[
{
"model_name": "ws-least-busy",
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"},
"model_info": {"id": deployment_id},
}
for deployment_id in ("ws-least-busy-a", "ws-least-busy-b")
],
routing_strategy="least-busy",
)
admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"]
assert callable(admit_turn)
for _ in range(3):
await admit_turn()
picks: Final = [
(await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[
"model_info"
]["id"]
for _ in range(2)
]
assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"]
@pytest.mark.asyncio
async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm():
router = Router(