From 8a99b36e1ea8ef122de6faba97ba218514630edc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 6 Nov 2026 10:07:24 -0800 Subject: [PATCH] fix(router): count long embedding batches in slices so chat counts do not wait behind them A text batch is now tokenized in slices of 65,536 characters, each its own offloaded call, so the FIFO count limiter can serve a chat count between slices. The no-OpenAI-tokenizer rule for text lists moved into the pre-call check, so RoutingContext.count_pre_call_check_tokens counts a text list by its longest text. --- litellm/router_rust/steps/precall.py | 123 +++++++++++++++++++------- tests/unit/test_router/test_router.py | 45 ++++++++++ 2 files changed, 134 insertions(+), 34 deletions(-) diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py index b8361653de2..c40a1e597b0 100644 --- a/litellm/router_rust/steps/precall.py +++ b/litellm/router_rust/steps/precall.py @@ -13,7 +13,9 @@ deployment records rather than a step, and those need not be in the Rust registr """ from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from functools import partial +from itertools import accumulate, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol @@ -38,6 +40,17 @@ if TYPE_CHECKING: # The chat params a deployment must take when the request sends them (`_PRE_CALL_CHECK_FILTERED_PARAMS`). _FILTERED_PARAMS: Final = frozenset({"response_format"}) +_CHARS_PER_OFFLOADED_COUNT: Final = 1 << 16 + + +@dataclass(frozen=True, slots=True) +class _TextBatch: + texts: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class _TokenIdBatch: + longest: int class PreCallRegistry(Protocol): @@ -81,10 +94,14 @@ def count_tokens( LiteLLMCompletionResponsesConfig, ) - text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) - embedding_tokens: Final = _longest_embedding_input_tokens(host, text_input, tokenizer_model, fits_within) - if embedding_tokens is not None: - return embedding_tokens + text_input: Final = _text_input(input, deployments) + match _embedding_batch(text_input): + case _TextBatch(texts): + return _longest_text_tokens(counter, tokenizer_model, texts, fits_within) + case _TokenIdBatch(longest): + return _token_id_count(host, longest) + case None: + pass input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it responses_api_request={"instructions": instructions} if instructions is not None else {}, @@ -94,30 +111,39 @@ def count_tokens( raise ValueError("Either messages or input must be provided to count tokens") -def _longest_embedding_input_tokens( - host: "PythonHost", value: object, tokenizer_model: str, fits_within: int -) -> int | None: - items: Final = _items(value) if _is_list(value) else () +def _text_input(input: object, deployments: Sequence[Mapping[str, object]]) -> object: # the request field's name + text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) + return text_input + + +def _embedding_batch(text_input: object) -> _TextBatch | _TokenIdBatch | None: + items: Final = _items(text_input) if _is_list(text_input) else () if not items: return None texts: Final = tuple(item for item in items if isinstance(item, str)) if len(texts) == len(items): - return _longest_text_tokens(host, texts, tokenizer_model, fits_within) + return _TextBatch(texts) if all(_is_token_id(item) for item in items): - return _token_id_count(host, len(items)) + return _TokenIdBatch(len(items)) if all(_is_token_id_list(item) for item in items): - return _token_id_count(host, max(len(_items(item)) for item in items)) + return _TokenIdBatch(max(len(_items(item)) for item in items)) return None -def _longest_text_tokens(host: "PythonHost", texts: Sequence[str], tokenizer_model: str, fits_within: int) -> int: - if not tokenizer_model: - # the default tokenizer overcounts non-English text, so only an OpenAI tokenizer may refuse a batch - return 0 - counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens +def _longest_text_tokens( + counter: Callable[..., int], tokenizer_model: str, texts: Sequence[str], fits_within: int +) -> int: return max(_text_tokens(counter, tokenizer_model, text, fits_within) for text in texts) +def _offload_slices(texts: Sequence[str]) -> tuple[tuple[str, ...], ...]: + ends: Final = accumulate(len(text) for text in texts) + slices: Final = groupby( + zip(ends, texts, strict=True), key=lambda end_and_text: (end_and_text[0] - 1) // _CHARS_PER_OFFLOADED_COUNT + ) + return tuple(tuple(text for _, text in texts_in_slice) for _, texts_in_slice in slices) + + def _text_tokens(counter: Callable[..., int], tokenizer_model: str, text: str, fits_within: int) -> int: size: Final = len(text.encode(errors="surrogatepass")) return size if size <= fits_within else counter(model=tokenizer_model, text=text) @@ -294,17 +320,31 @@ class PreCallFacts: def _fits_within(self) -> int: return min((limit for limit in self._windows.values() if limit is not None), default=0) + def _text_batch(self) -> _TextBatch | None: + if self._messages is not None or self._input is None: + return None + batch: Final = _embedding_batch(_text_input(self._input, self._candidate_records())) + return batch if isinstance(batch, _TextBatch) else None + def _count(self) -> int | None: try: - return count_tokens( - self._host, - self._messages, - self._input, - self._request, - self._tokenizer_model(), - self._candidate_records(), - self._fits_within(), - ) + tokenizer_model: Final = self._tokenizer_model() + match self._text_batch(): + case _TextBatch() if not tokenizer_model: + return 0 + case _TextBatch(texts): + counter: Final = self._host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens + return _longest_text_tokens(counter, tokenizer_model, texts, self._fits_within()) + case None: + return count_tokens( + self._host, + self._messages, + self._input, + self._request, + tokenizer_model, + self._candidate_records(), + self._fits_within(), + ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e @@ -313,15 +353,30 @@ class PreCallFacts: async def _count_async(self) -> int | None: try: - return await offload_token_count(count_tokens)( - self._host, - self._messages, - self._input, - self._request, - self._tokenizer_model(), - self._candidate_records(), - self._fits_within(), - ) + tokenizer_model: Final = self._tokenizer_model() + match self._text_batch(): + case _TextBatch() if not tokenizer_model: + return 0 + case _TextBatch(texts): + counter: Final = self._host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # see count_tokens + fits_within: Final = self._fits_within() + count_slice: Final = offload_token_count(_longest_text_tokens) + return max( + [ + await count_slice(counter, tokenizer_model, texts_in_slice, fits_within) + for texts_in_slice in _offload_slices(texts) + ] + ) + case None: + return await offload_token_count(count_tokens)( + self._host, + self._messages, + self._input, + self._request, + tokenizer_model, + self._candidate_records(), + self._fits_within(), + ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.error( "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e45368107d2..377186b8aeb 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -4376,6 +4376,51 @@ async def test_aembedding_skips_the_context_window_check_when_token_counting_is_ assert response.data +@pytest.mark.asyncio +async def test_a_chat_pre_call_count_does_not_wait_behind_long_embedding_batches(): + from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + + window: Final = {"max_input_tokens": 10000} + router: Final = litellm.Router( + model_list=[ + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_key": "sk-test", + "mock_response": [0.1], + }, + "model_info": window, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"}, + "model_info": window, + }, + ], + enable_pre_call_checks=True, + ) + batch: Final = ["word " * 2100] * 76 + batches: Final = tuple( + asyncio.create_task(router.aembedding(model="embed", input=batch)) + for _ in range(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + ) + await asyncio.sleep(0.05) + chat: Final = asyncio.create_task(router.acompletion(model="chat", messages=[{"role": "user", "content": "hi"}])) + + first_done, _ = await asyncio.wait((chat, *batches), return_when=asyncio.FIRST_COMPLETED) + await asyncio.gather(chat, *batches) + + assert first_done == {chat} + + +def test_a_routing_plugin_counts_an_embedding_text_batch_by_its_longest_text(): + long_text: Final = "word " * 4000 + context: Final = RouterContext(_embedding_router(max_input_tokens=10)) + + assert context.count_pre_call_check_tokens(None, ["one", long_text]) == litellm.token_counter(text=long_text) + + @pytest.mark.asyncio @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") @pytest.mark.parametrize(