diff --git a/litellm/router.py b/litellm/router.py index aadbdaa8520..ac0933e158b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -242,6 +242,15 @@ from litellm.router_utils.cooldown_handlers import ( is_video_lookup_not_found, set_cooldown_deployments, ) +from litellm.router_utils.embedding_batch_size_check import ( + async_log_rejection_before_provider_call, + clear_rejected_before_provider_call, + deployments_admitting_embedding_batch, + deployments_when_none_admits_embedding_batch, + embedding_batch_too_large_error, + embedding_deployment_selection, + log_rejection_before_provider_call, +) from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, AttemptedFallbackTargets, @@ -6861,19 +6870,33 @@ class Router: except Exception as e: raise e - def _embedding(self, input: str | list, model: str, **kwargs): + def _embedding(self, input: str | list[object], model: str, **kwargs): model_name = None try: verbose_router_logger.debug("Inside embedding()- model: %s; kwargs: %s", model, kwargs) - deployment: Final = self.get_available_deployment( - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with embedding_deployment_selection(): + deployment: Final = self.get_available_deployment( + model=model, + input=input, + 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() model_name = data["model"] + batch_error: Final = embedding_batch_too_large_error( + deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), + input=input, + request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), + model=model, + ) + logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") + if batch_error is not None: + if isinstance(logging_obj, LiteLLMLogging): + log_rejection_before_provider_call(logging_obj, batch_error) + raise batch_error + if isinstance(logging_obj, LiteLLMLogging): + clear_rejected_before_provider_call(logging_obj) potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="sync") # check if provided keys == client keys # @@ -6936,20 +6959,34 @@ class Router: ) raise e - async def _aembedding(self, input: str | list, model: str, **kwargs): + async def _aembedding(self, input: str | list[object], model: str, **kwargs): model_name = None try: verbose_router_logger.debug("Inside _aembedding()- model: %s; kwargs: %s", model, kwargs) parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self.async_get_available_deployment( - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with embedding_deployment_selection(): + deployment: Final = await self.async_get_available_deployment( + model=model, + input=input, + 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() model_name = data["model"] + batch_error: Final = embedding_batch_too_large_error( + deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), + input=input, + request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), + model=model, + ) + logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") + if batch_error is not None: + if isinstance(logging_obj, LiteLLMLogging): + await async_log_rejection_before_provider_call(logging_obj, batch_error) + raise batch_error + if isinstance(logging_obj, LiteLLMLogging): + clear_rejected_before_provider_call(logging_obj) model_client: Final = self._get_async_openai_model_client( deployment=deployment, kwargs=kwargs, @@ -14680,10 +14717,34 @@ class Router: batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts list[dict[str, object]], healthy_deployments ) + embedding_batch_admitting: Final = deployments_admitting_embedding_batch( + batch_candidates, + _OPAQUE_VALUE_ADAPTER.validate_python(input), + _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), + ) + embedding_batch_routed: Final = ( + embedding_batch_admitting + if embedding_batch_admitting is not None + else deployments_when_none_admits_embedding_batch( + batch_candidates, + cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts + await self._async_routable_cooled_deployment_ids( + model=model, + request_deployments=request_deployments, + cooldown_deployments=cooldown_deployments, + request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs + has_request_input=True, + ), + _OPAQUE_VALUE_ADAPTER.validate_python(input), + _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), + ) + ) healthy_deployments = litellm.utils.get_order_filtered_deployments( list( self._narrow_to_video_deployment( - without_deployments_that_cannot_batch(batch_candidates, caller_deployments=lambda: batch_candidates) + without_deployments_that_cannot_batch( + embedding_batch_routed, caller_deployments=lambda: batch_candidates + ) ) ), target_order=_target_order, @@ -15890,11 +15951,34 @@ class Router: batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts list[dict[str, object]], healthy_deployments ) + embedding_batch_admitting: Final = deployments_admitting_embedding_batch( + batch_candidates, + _OPAQUE_VALUE_ADAPTER.validate_python(input), + routing_kwargs, + ) + embedding_batch_routed: Final = ( + embedding_batch_admitting + if embedding_batch_admitting is not None + else deployments_when_none_admits_embedding_batch( + batch_candidates, + cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts + self._deployment_ids( + self._region_routable_deployments( + request_deployments, + cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs + True, + ) + ) + & frozenset(cooldown_deployments), + _OPAQUE_VALUE_ADAPTER.validate_python(input), + routing_kwargs, + ) + ) healthy_deployments = litellm.utils.get_order_filtered_deployments( list( self._narrow_to_video_deployment( without_deployments_that_cannot_batch( - batch_candidates, + embedding_batch_routed, caller_deployments=lambda: deployments_the_caller_may_use(batch_candidates, routing_kwargs), ) ) diff --git a/litellm/router_utils/embedding_batch_size_check.py b/litellm/router_utils/embedding_batch_size_check.py new file mode 100644 index 00000000000..6000bc9766d --- /dev/null +++ b/litellm/router_utils/embedding_batch_size_check.py @@ -0,0 +1,199 @@ +import traceback +from collections.abc import Generator, Iterable, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING, Final + +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 + +import litellm +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.types.utils import CallTypes + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_SELECTING_FOR_EMBEDDING: Final[ContextVar[bool]] = ContextVar("selecting_for_embedding", default=False) + + +def _is_str_keyed_mapping( + value: object, +) -> TypeIs[Mapping[str, object]]: # guard-ok: model_info, litellm_params and extra_body are str-keyed config objects + return isinstance(value, Mapping) + + +def _is_embedding_input( + value: object, +) -> TypeIs[str | Sequence[object]]: # guard-ok: only the length and first item are read + return isinstance(value, Sequence) + + +def count_embedding_inputs(input: str | Sequence[object]) -> int: + if isinstance(input, str): + return 1 + if input and isinstance(input[0], int): + return 1 + return len(input) + + +def effective_embedding_input( + input: object, + extra_body: object, +) -> str | Sequence[object] | None: + """Return the input shape used for batch-size checks, or None when neither is an embedding input. + + Some providers merge ``extra_body`` over the request body (OpenAI-compatible paths); + others ignore ``extra_body.input`` and use the top-level input (e.g. Bedrock Titan). + Use whichever of the two yields the larger input count so limits cannot be bypassed. + """ + override: Final = extra_body.get("input") if _is_str_keyed_mapping(extra_body) else None + candidates: Final = tuple(value for value in (input, override) if _is_embedding_input(value)) + return max(candidates, key=count_embedding_inputs, default=None) + + +def _litellm_param(deployment: Mapping[str, object], key: str) -> object: + litellm_params: Final = deployment.get("litellm_params") + return litellm_params.get(key) if _is_str_keyed_mapping(litellm_params) else None + + +def _embedding_batch_limit(deployment: Mapping[str, object]) -> int | None: + model_info: Final = deployment.get("model_info") + limit: Final = model_info.get("max_embedding_batch_size") if _is_str_keyed_mapping(model_info) else None + return limit if isinstance(limit, int) else None + + +def _requested_embedding_count( + deployment: Mapping[str, object], + input: object, + request_kwargs: Mapping[str, object], +) -> int: + extra_body: Final = request_kwargs.get("extra_body", _litellm_param(deployment, "extra_body")) + effective_input: Final = effective_embedding_input(input, extra_body) + return 0 if effective_input is None else count_embedding_inputs(effective_input) + + +def _deployment_admits_embedding_batch( + deployment: Mapping[str, object], + input: object, + request_kwargs: Mapping[str, object], +) -> bool: + limit: Final = _embedding_batch_limit(deployment) + return limit is None or _requested_embedding_count(deployment, input, request_kwargs) <= limit + + +def _str_keyed_or_empty(value: object) -> Mapping[str, object]: + return value if _is_str_keyed_mapping(value) else {} + + +def _deployment_id(deployment: Mapping[str, object]) -> str | None: + model_info: Final = deployment.get("model_info") + deployment_id: Final = model_info.get("id") if _is_str_keyed_mapping(model_info) else None + return None if deployment_id is None else str(deployment_id) + + +def _is_paused(deployment: Mapping[str, object]) -> bool: + model_info: Final = deployment.get("model_info") + return _is_str_keyed_mapping(model_info) and model_info.get("blocked") is True + + +@contextmanager +def embedding_deployment_selection() -> Generator[None]: + """Apply the input limit to deployment selection inside this block, which only embedding calls enter.""" + token: Final = _SELECTING_FOR_EMBEDDING.set(True) + try: + yield + finally: + _SELECTING_FOR_EMBEDDING.reset(token) + + +def deployments_admitting_embedding_batch( + deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + input: object, + request_kwargs: object, +) -> list[dict[str, object]] | None: # mutable-ok: returns a mutable deployment list + """Keep the deployments whose input limit admits the request, all of them outside embedding selection, else None""" + if not _SELECTING_FOR_EMBEDDING.get(): + return deployments + kwargs: Final = _str_keyed_or_empty(request_kwargs) + admitting: Final = [ + deployment for deployment in deployments if _deployment_admits_embedding_batch(deployment, input, kwargs) + ] + return admitting or None + + +def deployments_when_none_admits_embedding_batch( + deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + before_cooldown: Iterable[Mapping[str, object]], + routable_cooling_ids: frozenset[str], + input: object, + request_kwargs: object, +) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + """None while an unpaused cooling deployment the request routes to admits it, so it waits out the cooldown, + else the deployments with the largest limit, so the 400 names the most the group accepts""" + kwargs: Final = _str_keyed_or_empty(request_kwargs) + routable_cooling: Final = ( + deployment + for deployment in before_cooldown + if _deployment_id(deployment) in routable_cooling_ids and not _is_paused(deployment) + ) + if any(_deployment_admits_embedding_batch(deployment, input, kwargs) for deployment in routable_cooling): + return [] + largest_limit: Final = max((_embedding_batch_limit(deployment) or 0 for deployment in deployments), default=0) + return [deployment for deployment in deployments if _embedding_batch_limit(deployment) == largest_limit] + + +def _mark_litellm_params(logging_obj: "LiteLLMLoggingObj", *keys: str) -> None: + litellm_params: Final = logging_obj.model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + for key in keys: + litellm_params[key] = True + + +def log_rejection_before_provider_call(logging_obj: "LiteLLMLoggingObj", error: Exception) -> None: + """Mark the request as never reaching a provider, so deployment health metrics skip its failure, then log it""" + _mark_litellm_params(logging_obj, PROXY_REJECTED_BEFORE_ROUTING_KEY) + logging_obj.failure_handler(error, "".join(traceback.format_exception(error))) + + +async def async_log_rejection_before_provider_call(logging_obj: "LiteLLMLoggingObj", error: Exception) -> None: + """Log it like an async provider failure: sync handlers inline, before a fallback swaps the deployment they read, + and the call marked async, so a logger with both hooks records it once""" + _mark_litellm_params(logging_obj, PROXY_REJECTED_BEFORE_ROUTING_KEY, CallTypes.aembedding.value) + traceback_text: Final = "".join(traceback.format_exception(error)) + logging_obj.failure_handler(error, traceback_text) + await logging_obj.async_failure_handler(error, traceback_text) + + +def clear_rejected_before_provider_call(logging_obj: "LiteLLMLoggingObj") -> None: + """Drop the mark an earlier rejected attempt left, since a fallback or retry shares the logging object""" + litellm_params: Final = logging_obj.model_call_details.get("litellm_params") + if isinstance(litellm_params, dict) and PROXY_REJECTED_BEFORE_ROUTING_KEY in litellm_params: + del litellm_params[PROXY_REJECTED_BEFORE_ROUTING_KEY] + + +def embedding_batch_too_large_error( + deployment: object, + input: object, + request_kwargs: object, + model: str, +) -> litellm.BadRequestError | None: + if not _is_str_keyed_mapping(deployment): + return None + limit: Final = _embedding_batch_limit(deployment) + if limit is None: + return None + + input_count: Final = _requested_embedding_count(deployment, input, _str_keyed_or_empty(request_kwargs)) + if input_count <= limit: + return None + + return litellm.BadRequestError( + message=( + f"Embedding request has {input_count} inputs, which exceeds " + f"max_embedding_batch_size={limit} for model={model}. " + "Split the input into smaller batches." + ), + model=model, + llm_provider="", + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 6cd758d864e..1f159cb8b16 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -323,6 +323,8 @@ class ModelInfo(MirroredPricingParams): # in the spend log row's metadata. Set it on every deployment of the group. internal_router_model: bool | None = None + max_embedding_batch_size: Annotated[int, Field(gt=0)] | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/tests/unit/router_utils/test_embedding_batch_size_check.py b/tests/unit/router_utils/test_embedding_batch_size_check.py new file mode 100644 index 00000000000..d5986d9e6f7 --- /dev/null +++ b/tests/unit/router_utils/test_embedding_batch_size_check.py @@ -0,0 +1,1024 @@ +import datetime +import json +import threading +from collections.abc import Iterator +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import respx +from prometheus_client import REGISTRY + +import litellm +from litellm import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.types.router import RouterRateLimitError +from litellm.router_utils.embedding_batch_size_check import ( + count_embedding_inputs, + deployments_admitting_embedding_batch, + deployments_when_none_admits_embedding_batch, + effective_embedding_input, + embedding_batch_too_large_error, + embedding_deployment_selection, +) + +_MODEL: Final = "embed-group" +_PROVIDER: Final = "openai" +_OPENAI_EMBED_BASE: Final = "https://batch-limit-embed.local/v1" +_LOGGING_START: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) +_VLLM_EMBED_BASE: Final = "https://batch-limit-vllm.local/v1" +_UNCAPPED_EMBED_BASE: Final = "https://batch-limit-uncapped.local/v1" + + +def _model_info(max_embedding_batch_size: int | None) -> dict[str, object]: + if max_embedding_batch_size is None: + return {} + return {"max_embedding_batch_size": max_embedding_batch_size} + + +def _deployment( + limit: int | None, + extra_body: dict[str, object] | None = None, + deployment_id: str = "embed-deployment", +) -> dict[str, object]: + return { + "model_info": {"id": deployment_id, **_model_info(limit)}, + "litellm_params": {} if extra_body is None else {"extra_body": extra_body}, + } + + +def _admitted( + deployments: list[dict[str, object]], + input_value: object, + request_kwargs: dict[str, object] | None = None, + cooling: tuple[dict[str, object], ...] = (), + routable_cooling_ids: frozenset[str] = frozenset(), +) -> list[dict[str, object]]: + admitting: Final = deployments_admitting_embedding_batch(deployments, input_value, request_kwargs) + if admitting is not None: + return admitting + return deployments_when_none_admits_embedding_batch( + deployments, (*deployments, *cooling), routable_cooling_ids, input_value, request_kwargs + ) + + +def _batch_error( + input_value: object, + limit: int | None, + request_kwargs: dict[str, object] | None = None, +) -> litellm.BadRequestError | None: + return embedding_batch_too_large_error( + deployment={"model_info": _model_info(limit), "litellm_params": {"custom_llm_provider": _PROVIDER}}, + input=input_value, + request_kwargs=request_kwargs or {}, + model=_MODEL, + ) + + +def _embedding_router( + *, + provider_model: str, + api_base: str, + max_embedding_batch_size: int | None, + deployment_id: str, + litellm_params_extra: dict[str, object] | None = None, +) -> Router: + litellm_params: dict[str, object] = { + "model": provider_model, + "api_key": "sk-fake", + "api_base": api_base, + } + if litellm_params_extra: + litellm_params.update(litellm_params_extra) + return Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": litellm_params, + "model_info": {"id": deployment_id, **_model_info(max_embedding_batch_size)}, + } + ], + num_retries=0, + ) + + +async def _call_embedding(router: Router, *, sync: bool, **kwargs: object) -> None: + if sync: + router.embedding(**kwargs) + else: + await router.aembedding(**kwargs) + + +def _vllm_embedding_response(input_count: int) -> httpx.Response: + return httpx.Response( + 200, + json={ + "object": "list", + "data": [{"object": "embedding", "index": i, "embedding": [0.1, 0.2]} for i in range(input_count)], + "model": "bge-m3", + "usage": {"prompt_tokens": input_count, "total_tokens": input_count}, + }, + ) + + +@pytest.mark.parametrize( + ("input_value", "expected_count"), + [ + ("hello", 1), + (["a", "b"], 2), + ([1, 2, 3, 4, 5], 1), + ([[1, 2, 3], [4, 5, 6]], 2), + ], +) +def test_openai_input_shapes(input_value: object, expected_count: int) -> None: + assert count_embedding_inputs(input_value) == expected_count + + +def test_batch_count_is_independent_of_string_length() -> None: + long_text: Final = "word " * 500 + assert count_embedding_inputs([long_text, long_text]) == 2 + + +_MIXED_EXTRA_BODY: Final = { + "truncate": "END", + "encoding_format": "float", + "input": ["a", "b", "c"], + "user": "batch-job", +} + + +@pytest.mark.parametrize( + ("top_level", "extra_body", "expected"), + [ + ("ok", None, "ok"), + ("ok", {"input": ["a", "b", "c"]}, ["a", "b", "c"]), + (["a", "b"], {"truncate": "END"}, ["a", "b"]), + ("ignored-top-level", _MIXED_EXTRA_BODY, ["a", "b", "c"]), + ( + ["would", "be", "three"], + {"dimensions": 256, "input": "single-string", "model": "ignored"}, + ["would", "be", "three"], + ), + ("ok", {"truncate": "NONE", "input": ("a", "b")}, ("a", "b")), + (["x", "y", "z"], {"input": [[1, 2], [3, 4]], "foo": "bar"}, ["x", "y", "z"]), + ( + ["a", "b", "c"], + {"input": "ok"}, + ["a", "b", "c"], + ), + (["a", "b"], {}, ["a", "b"]), + (["a", "b", "c"], {"input": None, "truncate": "END"}, ["a", "b", "c"]), + (["a", "b"], {"input": {"nested": "dict"}, "truncate": "END"}, ["a", "b"]), + (["a", "b"], {"input": 42, "dimensions": 128}, ["a", "b"]), + (None, {"input": ["a", "b"]}, ["a", "b"]), + (7, {"input": ["a", "b"]}, ["a", "b"]), + (None, None, None), + ], +) +def test_resolves_provider_bound_input( + top_level: object, + extra_body: dict[str, object] | None, + expected: object, +) -> None: + assert effective_embedding_input(top_level, extra_body) == expected + + +@pytest.mark.parametrize( + "extra_body", + ["not-a-mapping", ["input", "list"]], +) +def test_non_mapping_extra_body_keeps_top_level(extra_body: object) -> None: + top_level: Final = ["a", "b", "c"] + assert effective_embedding_input(top_level, extra_body) == top_level + + +def test_readonly_mapping_extra_body_with_input() -> None: + extra_body: Final = MappingProxyType({"input": ["a", "b"], "encoding_format": "float"}) + assert effective_embedding_input("ok", extra_body) == ["a", "b"] + + +class TestDeploymentsAdmittingEmbeddingBatch: + @pytest.fixture(autouse=True) + def _embedding_selection(self) -> Iterator[None]: + with embedding_deployment_selection(): + yield + + @pytest.mark.parametrize( + ("input_value", "limit"), + [ + ("hello", 1), + (["a", "b"], 2), + (["a", "b", "c"], 3), + ([1, 2, 3, 4, 5], 1), + ([[1, 2, 3], [4, 5, 6]], 2), + ], + ) + def test_keeps_deployments_whose_limit_admits_the_batch(self, input_value: str | list, limit: int) -> None: + capped: Final = _deployment(limit) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], input_value) == [capped, uncapped] + + def test_drops_deployments_whose_limit_is_below_the_batch(self) -> None: + capped: Final = _deployment(2) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], ["a", "b", "c"]) == [uncapped] + + def test_keeps_the_largest_limit_deployments_when_none_admits_the_batch(self) -> None: + smaller: Final = _deployment(1, deployment_id="smaller") + small: Final = _deployment(2, deployment_id="small") + other_small: Final = _deployment(2, deployment_id="other-small") + assert _admitted([small, smaller, other_small], ["a", "b", "c"]) == [small, other_small] + + def test_request_without_embedding_input_keeps_every_deployment(self) -> None: + capped: Final = _deployment(1) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], None) == [capped, uncapped] + + def test_request_extra_body_replaces_the_deployment_extra_body(self) -> None: + capped: Final = _deployment(2, extra_body={"input": ["a", "b", "c"]}) + uncapped: Final = _deployment(None) + assert _admitted([capped, uncapped], "ok") == [uncapped] + assert _admitted([capped, uncapped], "ok", {"extra_body": {}}) == [capped, uncapped] + + @pytest.mark.parametrize("top_level", [None, 7, {"text": "a"}]) + def test_extra_body_input_routes_when_top_level_input_is_not_a_list(self, top_level: object) -> None: + capped: Final = _deployment(2) + uncapped: Final = _deployment(None) + request_kwargs: Final = {"extra_body": {"input": ["a", "b", "c"]}} + assert _admitted([capped, uncapped], top_level, request_kwargs) == [uncapped] + + def test_keeps_no_deployment_while_only_a_cooling_one_admits_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + cooling: Final = _deployment(None, deployment_id="cooling") + assert _admitted([capped], ["a", "b", "c"], cooling=(cooling,), routable_cooling_ids=frozenset({"cooling"})) == [] + + def test_cooling_deployment_the_request_cannot_route_to_does_not_hold_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + cooling: Final = _deployment(None, deployment_id="cooling") + assert _admitted([capped], ["a", "b", "c"], cooling=(cooling,)) == [capped] + + def test_paused_cooling_deployment_does_not_hold_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + paused: Final = {"model_info": {"id": "cooling", "blocked": True}, "litellm_params": {}} + assert _admitted([capped], ["a", "b", "c"], cooling=(paused,), routable_cooling_ids=frozenset({"cooling"})) == [ + capped + ] + + def test_keeps_every_deployment_when_the_cooling_one_also_rejects_the_batch(self) -> None: + capped: Final = _deployment(2, deployment_id="capped") + cooling: Final = _deployment(1, deployment_id="cooling") + assert _admitted([capped], ["a", "b", "c"], cooling=(cooling,), routable_cooling_ids=frozenset({"cooling"})) == [ + capped + ] + + +@pytest.mark.parametrize( + ("input_value", "limit", "message_match"), + [ + (["a", "b", "c", "d"], 3, "4 inputs"), + ([[1], [2], [3]], 2, None), + ( + effective_embedding_input( + "ok", + {"truncate": "END", "input": ["a", "b", "c", "d"], "user": "x"}, + ), + 3, + "4 inputs", + ), + ( + effective_embedding_input( + ["a", "b", "c"], + {"input": {"bad": "shape"}, "encoding_format": "float"}, + ), + 2, + None, + ), + ], +) +def test_rejects_inputs_above_limit( + input_value: str | list, + limit: int, + message_match: str | None, +) -> None: + error: Final = _batch_error(input_value, limit) + assert isinstance(error, litellm.BadRequestError) + assert error.status_code == 400 + if message_match is not None: + assert message_match in error.message + assert f"max_embedding_batch_size={limit}" in error.message + + +@pytest.mark.parametrize("top_level", [None, 7, {"text": "a"}]) +def test_extra_body_override_counts_when_top_level_input_is_not_a_list(top_level: object) -> None: + error: Final = _batch_error(top_level, 2, {"extra_body": {"input": ["a", "b", "c"]}}) + assert error is not None + assert "3 inputs" in error.message + + +class TestRouterEmbeddingBatchLimit: + @pytest.fixture(autouse=True) + def _disable_aiohttp(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_top_level_input_rejects_before_provider_call(self, sync: bool) -> None: + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + embeddings_url: Final = f"{_OPENAI_EMBED_BASE}/embeddings" + empty_ok: Final = httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=empty_ok) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding( + router, + sync=sync, + model="embed", + input=["one", "two", "three"], + ) + assert route.call_count == 0 + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_rejection_reaches_failure_callbacks(self, sync: bool) -> None: + failures: Final[list[tuple[Exception, object]]] = [] + + class FailureRecorder(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + failures.append((kwargs["exception"], kwargs["litellm_params"].get(PROXY_REJECTED_BEFORE_ROUTING_KEY))) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + failures.append((kwargs["exception"], kwargs["litellm_params"].get(PROXY_REJECTED_BEFORE_ROUTING_KEY))) + + recorder: Final = FailureRecorder() + logging_obj: Final = LitellmLogging( + model="embed", + messages=[], + stream=False, + call_type="embedding" if sync else "aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-rejected", + function_id="embed-batch-rejected", + dynamic_failure_callbacks=[recorder] if sync else None, + dynamic_async_failure_callbacks=None if sync else [recorder], + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}, "proxy_server_request": {}} + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + ) + with pytest.raises(litellm.BadRequestError, match="3 inputs") as excinfo: + await _call_embedding( + router, + sync=sync, + model="embed", + input=["one", "two", "three"], + litellm_logging_obj=logging_obj, + ) + assert route.call_count == 0 + assert failures == [(excinfo.value, True)] + + @pytest.mark.asyncio + async def test_async_rejection_logs_each_failure_hook_once_before_raising(self) -> None: + hook_events: Final[list[str]] = [] + legacy_callback_threads: Final[list[int]] = [] + legacy_callback_ran: Final = threading.Event() + + class BothHooks(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + hook_events.append("sync") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + hook_events.append("async") + + def legacy_failure_callback(kwargs, completion_response, start_time, end_time) -> None: + legacy_callback_threads.append(threading.get_ident()) + legacy_callback_ran.set() + + both_hooks: Final = BothHooks() + logging_obj: Final = LitellmLogging( + model="embed", + messages=[], + stream=False, + call_type="aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-logged-once", + function_id="embed-batch-logged-once", + dynamic_failure_callbacks=[legacy_failure_callback, both_hooks], + dynamic_async_failure_callbacks=[both_hooks], + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}, "proxy_server_request": {}} + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await router.aembedding(model="embed", input=["one", "two", "three"], litellm_logging_obj=logging_obj) + + assert legacy_callback_ran.wait(timeout=5) + assert legacy_callback_threads == [threading.get_ident()] + assert hook_events == ["async"] + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_fallback_after_a_rejection_logs_its_provider_call_as_reached(self, sync: bool) -> None: + litellm_params: Final[dict[str, object]] = {"model": "openai/text-embedding-3-small", "api_key": "sk-fake"} + router: Final = Router( + model_list=[ + { + "model_name": "embed-capped", + "litellm_params": {**litellm_params, "api_base": _OPENAI_EMBED_BASE}, + "model_info": {"id": "capped-box", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed-fallback", + "litellm_params": {**litellm_params, "api_base": _VLLM_EMBED_BASE}, + "model_info": {"id": "fallback-box"}, + }, + ], + fallbacks=[{"embed-capped": ["embed-fallback"]}], + num_retries=0, + ) + logging_obj: Final = LitellmLogging( + model="embed-capped", + messages=[], + stream=False, + call_type="embedding" if sync else "aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-fallback", + function_id="embed-batch-fallback", + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {}}, optional_params={} + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + fallback_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + await _call_embedding( + router, + sync=sync, + model="embed-capped", + input=["one", "two", "three"], + litellm_logging_obj=logging_obj, + ) + assert (capped_route.call_count, fallback_route.call_count) == (0, 1) + assert PROXY_REJECTED_BEFORE_ROUTING_KEY not in logging_obj.model_call_details["litellm_params"] + + @pytest.fixture + def _isolated_prometheus_registry(self) -> Iterator[None]: + def unregister_all() -> None: + for collector in tuple(REGISTRY._collector_to_names): + REGISTRY.unregister(collector) + + unregister_all() + yield + unregister_all() + + @pytest.mark.asyncio + @pytest.mark.usefixtures("_isolated_prometheus_registry") + async def test_oversized_batch_rejection_leaves_the_deployment_healthy_in_prometheus(self) -> None: + prometheus: Final = PrometheusLogger() + metadata: Final[dict[str, object]] = {} + logging_obj: Final = LitellmLogging( + model="embed", + messages=[], + stream=False, + call_type="aembedding", + start_time=_LOGGING_START, + litellm_call_id="embed-batch-prometheus", + function_id="embed-batch-prometheus", + dynamic_async_failure_callbacks=[prometheus], + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": metadata, "proxy_server_request": {}} + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped", + ) + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + ) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await router.aembedding( + model="embed", input=["one", "two", "three"], metadata=metadata, litellm_logging_obj=logging_obj + ) + assert prometheus.litellm_deployment_state._metrics == {} + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_no_limit_configured_forwards_any_batch(self, sync: bool) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=None, + deployment_id="vllm-embed-uncapped", + ) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(100) + ) + await _call_embedding(router, sync=sync, model="embed", input=["a"] * 100) + assert len(json.loads(route.calls.last.request.content)["input"]) == 100 + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_goes_to_the_group_deployment_that_admits_it(self, sync: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "weight": 1000, + }, + "model_info": {"id": "vllm-embed-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "weight": 1, + }, + "model_info": {"id": "vllm-embed-uncapped"}, + }, + ], + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + for _ in range(5): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 5) + + @pytest.mark.asyncio + async def test_non_embedding_call_with_input_ignores_embedding_batch_limit(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-fake", + "api_base": _OPENAI_EMBED_BASE, + "order": 1, + }, + "model_info": {"id": "capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "order": 2, + }, + "model_info": {"id": "uncapped"}, + }, + ], + num_retries=0, + ) + responses_input: Final = [{"role": "user", "content": str(turn)} for turn in range(3)] + + async_pick: Final = await router.async_get_available_deployment( + model="embed", request_kwargs={}, input=responses_input + ) + sync_pick: Final = router.get_available_deployment(model="embed", request_kwargs={}, input=responses_input) + + assert (async_pick["model_info"]["id"], sync_pick["model_info"]["id"]) == ("capped", "capped") + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_waits_for_a_cooling_down_deployment_that_admits_it(self, sync: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + }, + "model_info": {"id": "vllm-embed-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + }, + "model_info": {"id": "vllm-embed-uncapped"}, + }, + ], + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + with pytest.raises(RouterRateLimitError): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_while_a_paused_admitting_deployment_cools_down( + self, sync: bool + ) -> None: + uncapped_params: Final = {"model": "hosted_vllm/bge-m3", "api_key": "sk-fake", "api_base": _UNCAPPED_EMBED_BASE} + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + }, + "model_info": {"id": "vllm-embed-capped", "max_embedding_batch_size": 2}, + }, + {"model_name": "embed", "litellm_params": uncapped_params, "model_info": {"id": "vllm-embed-uncapped"}}, + ], + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + router.upsert_deployment( + Deployment( + model_name="embed", + litellm_params=LiteLLM_Params(**uncapped_params), + model_info=ModelInfo(id="vllm-embed-uncapped", blocked=True), + ) + ) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_extra_body_input_override_rejects_before_provider_call(self, sync: bool) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="vllm-embed-batch-capped", + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + extra_body: Final = {"truncate": "END", "input": ["a", "b", "c"]} + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(3)) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await _call_embedding( + router, + sync=sync, + model="embed", + input="ok", + extra_body=extra_body, + ) + assert route.call_count == 0 + + @pytest.mark.asyncio + async def test_deployment_extra_body_input_counts_toward_batch_limit(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "extra_body": {"input": ["a", "b", "c"]}, + }, + "model_info": { + "id": "vllm-embed-deployment-extra-body", + "max_embedding_batch_size": 2, + }, + } + ], + num_retries=0, + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(3)) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await router.aembedding(model="embed", input="ok") + assert route.call_count == 0 + + @pytest.mark.asyncio + async def test_extra_body_input_at_limit_reaches_provider(self) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="vllm-embed-batch-capped", + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(2)) + await router.aembedding(model="embed", input="ok", extra_body={"input": ["a", "b"]}) + assert route.call_count == 1 + assert json.loads(route.calls.last.request.content)["input"] == ["a", "b"] + + @pytest.mark.asyncio + async def test_unrelated_extra_body_does_not_bypass_top_level_batch_limit(self) -> None: + router: Final = _embedding_router( + provider_model="hosted_vllm/bge-m3", + api_base=_VLLM_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="vllm-embed-batch-capped", + ) + embeddings_url: Final = f"{_VLLM_EMBED_BASE}/embeddings" + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=_vllm_embedding_response(3)) + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await _call_embedding( + router, + sync=False, + model="embed", + input=["a", "b", "c"], + extra_body={"truncate": "END", "encoding_format": "float"}, + ) + assert route.call_count == 0 + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_bedrock_ignores_extra_body_input_override_for_provider_but_guard_uses_top_level( + self, + sync: bool, + ) -> None: + router: Final = _embedding_router( + provider_model="bedrock/amazon.titan-embed-text-v1", + api_base="", + max_embedding_batch_size=2, + deployment_id="bedrock-titan-embed-batch-capped", + litellm_params_extra={ + "aws_access_key_id": "AKIAFAKE", + "aws_secret_access_key": "fake-secret", + "aws_region_name": "us-east-1", + "custom_llm_provider": "bedrock", + }, + ) + + with pytest.raises(litellm.BadRequestError, match="3 inputs"): + await _call_embedding( + router, + sync=sync, + model="embed", + input=["a", "b", "c"], + extra_body={"input": "ok"}, + ) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_explicit_custom_llm_provider_in_deployment_rejects_oversized_batch( + self, + sync: bool, + ) -> None: + router: Final = _embedding_router( + provider_model="openai/text-embedding-3-small", + api_base=_OPENAI_EMBED_BASE, + max_embedding_batch_size=2, + deployment_id="embed-batch-capped-explicit-provider", + litellm_params_extra={"custom_llm_provider": "openai"}, + ) + embeddings_url: Final = f"{_OPENAI_EMBED_BASE}/embeddings" + empty_ok: Final = httpx.Response(200, json={"object": "list", "data": [], "model": "x", "usage": {}}) + + with respx.mock(assert_all_called=False) as respx_mock: + route: Final = respx_mock.post(embeddings_url).mock(return_value=empty_ok) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding( + router, + sync=sync, + model="embed", + input=["one", "two", "three"], + ) + assert route.call_count == 0 + + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_when_tags_exclude_the_admitting_deployment(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "tags": ["team-a"], + }, + "model_info": {"id": "team-a-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "tags": ["team-b"], + }, + "model_info": {"id": "team-b-uncapped"}, + }, + ], + enable_tag_filtering=True, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await router.aembedding(model="embed", input=["a", "b", "c"], metadata={"tags": ["team-a"]}) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 0) + + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_while_another_tags_admitting_deployment_cools_down(self) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "tags": ["team-a"], + }, + "model_info": {"id": "team-a-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "tags": ["team-b"], + }, + "model_info": {"id": "team-b-uncapped"}, + }, + ], + enable_tag_filtering=True, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await router.aembedding(model="embed", input=["a", "b", "c"], metadata={"tags": ["team-b"]}) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await router.aembedding(model="embed", input=["a", "b", "c"], metadata={"tags": ["team-a"]}) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_names_the_cap_while_another_regions_admitting_deployment_cools_down( + self, sync: bool + ) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _VLLM_EMBED_BASE, + "region_name": "eu", + }, + "model_info": {"id": "eu-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + "region_name": "us", + }, + "model_info": {"id": "us-uncapped"}, + }, + ], + enable_pre_call_checks=True, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_VLLM_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + uncapped_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=httpx.Response(429, json={"error": {"message": "rate limited"}}) + ) + with pytest.raises(litellm.RateLimitError): + await _call_embedding( + router, sync=sync, model="embed", input=["a", "b", "c"], allowed_model_region="us" + ) + with pytest.raises(litellm.BadRequestError, match="max_embedding_batch_size=2"): + await _call_embedding( + router, sync=sync, model="embed", input=["a", "b", "c"], allowed_model_region="eu" + ) + assert (capped_route.call_count, uncapped_route.call_count) == (0, 1) + + @pytest.mark.parametrize("sync", [False, True], ids=["async", "sync"]) + @pytest.mark.asyncio + async def test_oversized_batch_rejection_falls_back_with_explicit_provider(self, sync: bool) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-fake", + "api_base": _OPENAI_EMBED_BASE, + "custom_llm_provider": "openai", + }, + "model_info": {"id": "explicit-provider-capped", "max_embedding_batch_size": 2}, + }, + { + "model_name": "embed-fallback", + "litellm_params": { + "model": "hosted_vllm/bge-m3", + "api_key": "sk-fake", + "api_base": _UNCAPPED_EMBED_BASE, + }, + "model_info": {"id": "fallback-uncapped"}, + }, + ], + fallbacks=[{"embed": ["embed-fallback"]}], + fallback_on_unsupported_params=False, + num_retries=0, + ) + + with respx.mock(assert_all_called=False) as respx_mock: + capped_route: Final = respx_mock.post(f"{_OPENAI_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + fallback_route: Final = respx_mock.post(f"{_UNCAPPED_EMBED_BASE}/embeddings").mock( + return_value=_vllm_embedding_response(3) + ) + await _call_embedding(router, sync=sync, model="embed", input=["a", "b", "c"]) + assert (capped_route.call_count, fallback_route.call_count) == (0, 1) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 992a7d91b0d..2b6898e39d9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50107,6 +50107,8 @@ export interface components { input_cost_per_token?: number | null; /** Internal Router Model */ internal_router_model?: boolean | null; + /** Max Embedding Batch Size */ + max_embedding_batch_size?: number | null; /** * Member Auto Router * @default false