scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

feat(embeddings): enforce per-request input batch limit (#44574)

Browse source 
* feat(embeddings): add max_batch_inputs pre-call validation

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>

* add unit tests for max embedding batch size

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>

* feat(embeddings): enforce batch limit on extra_body input override

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>

* chore(dashboard): update schema.d.ts for embedding batch size model field

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>

* fix(embeddings): close Bedrock batch bypass and satisfy ruff PT012

Count max(top-level input, extra_body.input) for batch limits so providers
that ignore extra_body cannot bypass max_embedding_batch_size. Add router
tests and refactor pytest.raises blocks for ruff-tests.

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>

* fix(embeddings): drop unchecked casts in router batch-size check

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>

* fix(embeddings): route oversized batches to a group deployment that admits them

* fix(embeddings): count extra_body input for any top-level input and prefer admitting deployments before cooldown

* fix(embeddings): log cap rejections to failure callbacks and scope batch routing to embedding calls

* fix(embeddings): apply the batch cap after tag and pre-call filters and fall back on cap rejections

* fix(router): keep embedding cap rejections out of deployment health and hold only for routable cooling deployments

* fix(router): clear the embedding cap rejection mark before a fallback's provider call

A fallback or retry shares the request's logging object, so the mark a rejected attempt set made telemetry skip the fallback deployment's real call. Unit tests now pass a fixed logging start time instead of reading the clock.

* fix(router): log embedding batch rejections once, before a fallback runs

* refactor(router): drop docstrings that restate the code

---------

Signed-off-by: Vincent Gimenes <vincent.gimenes@gmail.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Vincent Gimenes 2026-10-20 13:35:18 +02:00 • committed by GitHub
parent 963333d654
commit 808333682e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 1327 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

116
litellm/router.py
View file

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

199
litellm/router_utils/embedding_batch_size_check.py Normal file
View file

@ -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="",
)

2
litellm/types/router.py
View file

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

1024
tests/unit/router_utils/test_embedding_batch_size_check.py Normal file
View file

File diff suppressed because it is too large Load diff

2
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

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