From 170fece7db9b4997fab48cd716811975e43f11f3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:08:25 -0700 Subject: [PATCH 01/10] fix(token_counter): release the GIL for HuggingFace counts and cap exact counting per string Both proxy token counting endpoints already count in a worker thread, but the HuggingFace tokenizer's encode holds the GIL for the whole call, so a 600k-token count on a Claude model still froze the event loop for up to 0.8 s and every other request with it. Count through encode_batch_fast, which releases the GIL, and tokenize at most TOKEN_COUNTER_MAX_EXACT_CHARS characters of any one string (default 4,000,000), scaling the exact count of that prefix by the string's length above it so the largest payloads stay bounded. --- litellm/constants.py | 6 ++ litellm/litellm_core_utils/token_counter.py | 25 ++++++- .../litellm_core_utils/test_token_counter.py | 71 ++++++++++++++++++- tests/test_litellm/proxy/test_proxy_server.py | 34 ++++++++- 4 files changed, 132 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index d53686e5e5b..a9112c90a54 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,6 +391,12 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range( minimum=1, maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS, ) +TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_EXACT_CHARS", + default=4_000_000, + minimum=1, + maximum=1_000_000_000, +) MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..d526f0f2996 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -8,6 +8,7 @@ from typing import Final, Literal, cast import httpx import tiktoken +from tokenizers import Tokenizer import litellm from litellm import verbose_logger @@ -21,6 +22,7 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, + TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.litellm_core_utils.url_utils import safe_get @@ -538,9 +540,28 @@ def _count_extra( return num_tokens +def _get_extrapolating_count_function( + count_exactly: TokenCounterFunction, + max_exact_chars: int = TOKEN_COUNTER_MAX_EXACT_CHARS, +) -> TokenCounterFunction: + def count_tokens(text: str) -> int: + if len(text) <= max_exact_chars: + return count_exactly(text) + return round(count_exactly(text[:max_exact_chars]) * len(text) / max_exact_chars) + + return count_tokens + + def _get_count_function( model: str | None, custom_tokenizer: dict | SelectTokenizerResponse | None = None, +) -> TokenCounterFunction: + return _get_extrapolating_count_function(_get_exact_count_function(model, custom_tokenizer)) + + +def _get_exact_count_function( + model: str | None, + custom_tokenizer: dict | SelectTokenizerResponse | None = None, ) -> TokenCounterFunction: """ Get the function to count tokens based on the model and custom tokenizer.""" @@ -549,10 +570,10 @@ def _get_count_function( if model is not None or custom_tokenizer is not None: tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model) if tokenizer_json["type"] == "huggingface_tokenizer": + tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - enc: Final = tokenizer_json["tokenizer"].encode(text) - return len(enc.ids) + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 4694fa8fbed..fddeacc37f9 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1,8 +1,10 @@ #### What this tests #### # This tests litellm.token_counter.token_counter() function +import asyncio import importlib import time import traceback +from typing import Final from unittest.mock import MagicMock import pytest @@ -14,7 +16,12 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.litellm_core_utils.token_counter import _get_tiktoken_count_function +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import ( + _get_exact_count_function, + _get_extrapolating_count_function, + _get_tiktoken_count_function, +) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text from tests.test_litellm.litellm_core_utils.messages_with_counts import ( @@ -120,6 +127,68 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch): importlib.reload(litellm.constants) +async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): + counted: Final = asyncio.Event() + + async def count_off_the_loop() -> tuple[int, float]: + started: Final = time.perf_counter() + try: + tokens: Final = await asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + return tokens, time.perf_counter() - started + finally: + counted.set() + + (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + + assert tokens > 0 + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" + + +@pytest.mark.parametrize( + ("max_exact_chars", "expected"), + [(1_000, 10_000), (5_000, 6_000), (10_000, 6_000)], +) +def test_count_above_the_cap_scales_the_exact_count_of_the_prefix(max_exact_chars, expected): + def count_exactly(chunk: str) -> int: + return chunk.count("a") + len(chunk) + + count_tokens: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars) + + assert count_tokens("a" * 1_000 + "b" * 4_000) == expected + + +def test_token_counter_applies_the_default_cap(): + max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS + prefix: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prefix + "a" * 200_000 + scaled: Final = round(token_counter_new(model="gpt-5.6", text=prefix) * len(over_the_cap) / max_exact_chars) + + assert token_counter_new(model="gpt-5.6", text=over_the_cap) == scaled + assert _get_exact_count_function("gpt-5.6")(over_the_cap) != scaled + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], +) +def test_max_exact_chars_config_is_honoured(monkeypatch, configured, expected): + monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_EXACT_CHARS") + importlib.reload(litellm.constants) + + def test_token_counter_with_prefix(): messages = [ {"role": "user", "content": "Who won the world cup in 2022?"}, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b7bb58378d4..20b41cf1530 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6,6 +6,7 @@ import os import re import socket import subprocess +import time import types from datetime import datetime, timedelta, timezone from pathlib import Path @@ -28,7 +29,7 @@ from litellm.caching.caching import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, TokenCountRequest, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.proxy_server import app, initialize from litellm.utils import _invalidate_model_cost_lowercase_map @@ -12857,3 +12858,34 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() import litellm.proxy.proxy_server as ps assert ps.general_settings["enable_openai_websocket_passthrough"] is False + + +async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): + from tests.large_text import text + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + counted: Final = asyncio.Event() + + async def count_off_the_loop() -> tuple[int, float]: + started: Final = time.perf_counter() + try: + response: Final = await proxy_server_module.token_counter( + TokenCountRequest(model="claude-fable-5", prompt=text * 100) + ) + return response.total_tokens, time.perf_counter() - started + finally: + counted.set() + + (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + + assert tokens > 0 + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" From 3829ebdce54a0330ca2c9d60a5bbee94c2b09382 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:26:23 -0700 Subject: [PATCH 02/10] fix(token_counter): estimate over-cap strings from evenly spaced samples instead of the prefix A string above TOKEN_COUNTER_MAX_EXACT_CHARS was counted from its first cap characters and scaled, so a string whose start tokenizes unlike its end got a skewed count, and that count reaches fallback billing when the provider sends no usage. The estimate now tokenizes 16 evenly spaced samples that together total the cap and scales their sum by the string's length, keeping the same bound on work while tracking the whole string --- litellm/litellm_core_utils/token_counter.py | 16 ++++++++- .../litellm_core_utils/test_token_counter.py | 36 +++++++++++-------- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index d526f0f2996..e3a0d4c785a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -315,6 +315,8 @@ def calculate_img_tokens( TokenCounterFunction = Callable[[str], int] + +EXTRAPOLATION_SAMPLES: Final = 16 """ Type for a function that counts tokens in a string. """ @@ -547,11 +549,23 @@ def _get_extrapolating_count_function( def count_tokens(text: str) -> int: if len(text) <= max_exact_chars: return count_exactly(text) - return round(count_exactly(text[:max_exact_chars]) * len(text) / max_exact_chars) + samples: Final = _evenly_spaced_samples(text, max_exact_chars) + sampled_chars: Final = sum(len(sample) for sample in samples) + return round(sum(count_exactly(sample) for sample in samples) * len(text) / sampled_chars) return count_tokens +def _evenly_spaced_samples(text: str, total_chars: int) -> tuple[str, ...]: + sample_count: Final = min(EXTRAPOLATION_SAMPLES, total_chars) + sample_chars: Final = total_chars // sample_count + last_start: Final = len(text) - sample_chars + return tuple( + text[start : start + sample_chars] + for start in (last_start * index // max(sample_count - 1, 1) for index in range(sample_count)) + ) + + def _get_count_function( model: str | None, custom_tokenizer: dict | SelectTokenizerResponse | None = None, diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index fddeacc37f9..5da8413ac9b 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -153,27 +153,35 @@ async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free() assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" -@pytest.mark.parametrize( - ("max_exact_chars", "expected"), - [(1_000, 10_000), (5_000, 6_000), (10_000, 6_000)], -) -def test_count_above_the_cap_scales_the_exact_count_of_the_prefix(max_exact_chars, expected): - def count_exactly(chunk: str) -> int: - return chunk.count("a") + len(chunk) +@pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars): + count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) + front_heavy: Final = "a" * 1_000 + "b" * 4_000 + exact: Final = 1_000 + len(front_heavy) - count_tokens: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars) + estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) - assert count_tokens("a" * 1_000 + "b" * 4_000) == expected + assert abs(estimate - exact) <= exact // 100 + assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars + + +def test_count_at_or_below_the_cap_is_exact(): + count_exactly: Final = MagicMock(side_effect=len) + + assert _get_extrapolating_count_function(count_exactly, max_exact_chars=5_000)("a" * 5_000) == 5_000 + assert count_exactly.call_args_list == [(("a" * 5_000,),)] def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS - prefix: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] - over_the_cap: Final = prefix + "a" * 200_000 - scaled: Final = round(token_counter_new(model="gpt-5.6", text=prefix) * len(over_the_cap) / max_exact_chars) + prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] + over_the_cap: Final = prose + "a" * 200_000 + exact: Final = _get_exact_count_function("gpt-5.6")(over_the_cap) - assert token_counter_new(model="gpt-5.6", text=over_the_cap) == scaled - assert _get_exact_count_function("gpt-5.6")(over_the_cap) != scaled + estimate: Final = token_counter_new(model="gpt-5.6", text=over_the_cap) + + assert estimate != exact + assert abs(estimate - exact) <= exact // 100 @pytest.mark.parametrize( From 0d89873daf9e82a692825820adca2a47c0aab8dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:36:09 -0700 Subject: [PATCH 03/10] test(token_counter): type the parametrized cap test arguments --- tests/test_litellm/litellm_core_utils/test_token_counter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 5da8413ac9b..db87bc52600 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -154,7 +154,7 @@ async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free() @pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) -def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars): +def test_count_above_the_cap_samples_the_whole_string_and_scales(max_exact_chars: int): count_exactly: Final = MagicMock(side_effect=lambda chunk: chunk.count("a") + len(chunk)) front_heavy: Final = "a" * 1_000 + "b" * 4_000 exact: Final = 1_000 + len(front_heavy) @@ -188,7 +188,7 @@ def test_token_counter_applies_the_default_cap(): ("configured", "expected"), [("2048", 2048), ("0", 4_000_000), ("not-an-int", 4_000_000)], ) -def test_max_exact_chars_config_is_honoured(monkeypatch, configured, expected): +def test_max_exact_chars_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): monkeypatch.setenv("TOKEN_COUNTER_MAX_EXACT_CHARS", configured) try: assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_EXACT_CHARS == expected From dcd38ab9f0235c18fc429eeb5802519fde26c264 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:15:36 -0700 Subject: [PATCH 04/10] fix(proxy): count failure and rate-limit input tokens off the event loop The failure hook's usage estimate and the project ITPM reservation both called litellm.token_counter inline on the event loop, so a large request that failed or hit the limiter stalled the gateway the same way the count_tokens endpoints did. Both now run through asyncify. The loop-lag probe the existing tests used moves into a shared helper that warms the tokenizer first, and two new tests fail when either count runs inline --- .../hooks/parallel_request_limiter_v3.py | 3 +- litellm/proxy/utils.py | 3 +- .../litellm_core_utils/event_loop_lag.py | 39 +++++++++++++++++++ .../litellm_core_utils/test_token_counter.py | 30 +++++--------- .../proxy/hooks/test_tpm_concurrent.py | 37 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 34 ++++++---------- tests/test_litellm/proxy/test_proxy_utils.py | 38 ++++++++++++++++++ 7 files changed, 139 insertions(+), 45 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/event_loop_lag.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 31437af7770..74266801e50 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -28,6 +28,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -3303,7 +3304,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - raw_estimated_input_tokens: Final = self._estimate_precise_input_tokens( + raw_estimated_input_tokens: Final = await asyncify(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type ) estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddf31cb1d8a..1014c2b9b6a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,6 +91,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, independent_snapshot, @@ -2569,7 +2570,7 @@ class ProxyLogging: original_exception=original_exception, ) - request_data.update(_failure_fields_to_lift(request_data)) + request_data.update(await asyncify(_failure_fields_to_lift)(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/test_litellm/litellm_core_utils/event_loop_lag.py new file mode 100644 index 00000000000..c2bbe2204c7 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/event_loop_lag.py @@ -0,0 +1,39 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable +from typing import Final, TypeVar + +import litellm + +T = TypeVar("T") + + +def warm_tokenizer(model: str) -> None: + litellm.token_counter(model=model, text="load the tokenizer before anything is timed") + + +async def loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: + async def wake_lag() -> float: + started: Final = time.perf_counter() + await asyncio.sleep(0.001) + return time.perf_counter() - started - 0.001 + + return tuple([await wake_lag() for _ in iter(until.is_set, True)]) + + +async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, float, tuple[float, ...]]: + finished: Final = asyncio.Event() + + async def timed() -> tuple[T, float]: + started: Final = time.perf_counter() + try: + return await run(), time.perf_counter() - started + finally: + finished.set() + + (result, took), lags = await asyncio.gather(timed(), loop_wake_lags(finished)) + return result, took, lags + + +def assert_loop_stayed_free(took: float, lags: tuple[float, ...]) -> None: + assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index db87bc52600..e3dd5feb18d 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -24,6 +24,11 @@ from litellm.litellm_core_utils.token_counter import ( ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text +from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, +) from tests.test_litellm.litellm_core_utils.messages_with_counts import ( MESSAGES_TEXT, MESSAGES_WITH_IMAGES, @@ -127,30 +132,15 @@ def test_valid_chunk_size_config_is_honoured(monkeypatch): importlib.reload(litellm.constants) -async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: - async def wake_lag() -> float: - started: Final = time.perf_counter() - await asyncio.sleep(0.001) - return time.perf_counter() - started - 0.001 - - return tuple([await wake_lag() for _ in iter(until.is_set, True)]) - - async def test_huggingface_count_in_a_worker_thread_leaves_the_event_loop_free(): - counted: Final = asyncio.Event() + warm_tokenizer("claude-fable-5") - async def count_off_the_loop() -> tuple[int, float]: - started: Final = time.perf_counter() - try: - tokens: Final = await asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) - return tokens, time.perf_counter() - started - finally: - counted.set() - - (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + tokens, took, lags = await timed_with_loop_lags( + lambda: asyncify(token_counter_new)(model="claude-fable-5", text=text * 100) + ) assert tokens > 0 - assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" + assert_loop_stayed_free(took, lags) @pytest.mark.parametrize("max_exact_chars", [64, 1_000, 2_500]) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 2839acab6b0..398abf60b54 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3670,5 +3670,42 @@ async def test_post_call_success_hook_contains_header_merge_failures( ) +@pytest.mark.asyncio +async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loop(rate_limiter): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + handler, _cache = rate_limiter + stash = get_or_create_request_stash() + warm_tokenizer("claude-fable-5") + data: Dict[str, Any] = { + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + itpm_descriptor = { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "proj-loop:claude-fable-5", + "rate_limit": {"tokens_per_unit": 10_000_000, "window_size": 60}, + } + + _, took, lags = await timed_with_loop_lags( + lambda: handler._reserve_project_io_tokens_or_raise( + descriptors=[itpm_descriptor], + data=data, + requested_model="claude-fable-5", + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-itpm-loop"), project_id="proj-loop"), + tpm_reservation_scopes=[], + tpm_reservation_amount=0, + ) + ) + + assert stash.rate_limit_response is not None + assert_loop_stayed_free(took, lags) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 20b41cf1530..7fa3d5dd246 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12860,32 +12860,20 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() assert ps.general_settings["enable_openai_websocket_passthrough"] is False -async def _loop_wake_lags(until: asyncio.Event) -> tuple[float, ...]: - async def wake_lag() -> float: - started: Final = time.perf_counter() - await asyncio.sleep(0.001) - return time.perf_counter() - started - 0.001 - - return tuple([await wake_lag() for _ in iter(until.is_set, True)]) - - async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) - counted: Final = asyncio.Event() + warm_tokenizer("claude-fable-5") - async def count_off_the_loop() -> tuple[int, float]: - started: Final = time.perf_counter() - try: - response: Final = await proxy_server_module.token_counter( - TokenCountRequest(model="claude-fable-5", prompt=text * 100) - ) - return response.total_tokens, time.perf_counter() - started - finally: - counted.set() + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=text * 100)) + ) - (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) - - assert tokens > 0 - assert max(lags) < took / 4, f"the event loop stalled {max(lags):.3f}s during a {took:.3f}s count" + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9462f2c8eb0..b78ec7dcff6 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1823,6 +1823,44 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): assert lifted["standard_logging_object"] == {"id": "log-1"} +@pytest.mark.asyncio +async def test_a_dispatched_failure_is_counted_off_the_event_loop(): + from unittest.mock import AsyncMock, patch + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("claude-fable-5") + request_data = { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "call_type": "acompletion", + "model": "claude-fable-5", + "messages": [{"role": "user", "content": text * 100}], + } + ), + "metadata": {}, + } + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.alert_types = [] + with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): + _, took, lags = await timed_with_loop_lags( + lambda: proxy_logging_obj.post_call_failure_hook( + request_data=request_data, + original_exception=Exception("boom"), + user_api_key_dict=UserAPIKeyAuth(), + ) + ) + + assert request_data["combined_usage_object"].prompt_tokens > 0 + assert_loop_stayed_free(took, lags) + + @pytest.mark.asyncio async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): """Regression for LIT-6043: an expected 4xx must not format a traceback for From 83d500bac9c375221120b4ae09e0e3cbd55f96b1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:29:26 -0700 Subject: [PATCH 05/10] fix(router): count deployment itpm reservations off the event loop --- litellm/litellm_core_utils/token_counter.py | 4 +-- .../io_token_rate_limit_check.py | 3 +- .../test_router/test_io_token_rate_limits.py | 28 +++++++++++++++++++ 3 files changed, 32 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e3a0d4c785a..e3af64fd3bf 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -315,12 +315,12 @@ def calculate_img_tokens( TokenCounterFunction = Callable[[str], int] - -EXTRAPOLATION_SAMPLES: Final = 16 """ Type for a function that counts tokens in a string. """ +EXTRAPOLATION_SAMPLES: Final = 16 + def _get_tiktoken_count_function( encode_length: Callable[[str], int], diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 01d42627001..dca5a2be845 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -21,6 +21,7 @@ import litellm from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -466,7 +467,7 @@ async def async_io_token_pre_call_check( request_kwargs: Final = get_io_token_rate_limit_request_kwargs() _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" - estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model) + estimated_input: Final = await asyncify(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) dt: Final = get_utc_datetime() diff --git a/tests/test_litellm/test_router/test_io_token_rate_limits.py b/tests/test_litellm/test_router/test_io_token_rate_limits.py index a5a68271111..3cef1c7bb63 100644 --- a/tests/test_litellm/test_router/test_io_token_rate_limits.py +++ b/tests/test_litellm/test_router/test_io_token_rate_limits.py @@ -1039,3 +1039,31 @@ class TestContextSlotRetention: assert deployment is not None router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs) assert get_io_token_rate_limit_request_kwargs() is kwargs + + +@pytest.mark.asyncio +async def test_the_deployment_itpm_reservation_counts_the_request_off_the_event_loop(): + from litellm.utils import get_utc_datetime + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + warm_tokenizer("anthropic/claude-fable-5") + deployment = { + "litellm_params": {"model": "anthropic/claude-fable-5", "itpm": 10_000_000}, + "model_info": {"id": "io-loop-id"}, + "model_name": "claude", + } + set_io_token_rate_limit_request_kwargs({"messages": [{"role": "user", "content": text * 100}], "metadata": {}}) + + _, took, lags = await timed_with_loop_lags(lambda: check.async_pre_call_check(deployment)) + + minute = get_utc_datetime().strftime("%H-%M") + reserved = await dual_cache.async_get_cache(key=f"global_router:io-loop-id:anthropic/claude-fable-5:itpm:{minute}") + assert reserved > 100_000 + assert_loop_stayed_free(took, lags) From d202885f8b4ce1b0d56b7757ebb0c5849aeb3cb8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:21:45 -0700 Subject: [PATCH 06/10] fix(proxy): run prompt caching counts and custom tokenizer loads off the event loop --- litellm/proxy/proxy_server.py | 4 +- .../prompt_caching_deployment_check.py | 5 +- litellm/utils.py | 10 +-- .../test_custom_tokenizer_bug.py | 10 ++- .../litellm_core_utils/event_loop_lag.py | 1 + tests/test_litellm/proxy/test_proxy_server.py | 37 +++++++++++ .../test_prompt_caching_deployment_check.py | 62 +++++++++++++++++++ 7 files changed, 111 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0915b8dd1b9..77354e9885f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12713,7 +12713,9 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = litellm.utils._select_tokenizer(model=model_to_use, custom_tokenizer=custom_tokenizer) + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( + model=model_to_use, custom_tokenizer=custom_tokenizer + ) tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 70362e60495..9ea528b9700 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -14,6 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -61,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger): if request_kwargs is not None and request_kwargs.get("_target_order") is not None: return healthy_deployments - if messages is not None and is_prompt_caching_valid_prompt( + if messages is not None and await asyncify(is_prompt_caching_valid_prompt)( messages=messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), @@ -139,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider - if is_prompt_caching_valid_prompt( + if await asyncify(is_prompt_caching_valid_prompt)( model=model, messages=cast(list[AllMessageValues], messages), ): diff --git a/litellm/utils.py b/litellm/utils.py index d0e11bc9551..ab905381b22 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2284,15 +2284,7 @@ def create_pretrained_tokenizer(identifier: str, revision="main", auth_token: st dict: A dictionary with the tokenizer and its type. """ - try: - tokenizer = Tokenizer.from_pretrained( - identifier, - revision=revision, - auth_token=auth_token, - ) - except Exception as e: - verbose_logger.error("Error creating pretrained tokenizer: %s. Defaulting to version without 'auth_token'.", e) - tokenizer = Tokenizer.from_pretrained(identifier, revision=revision) + tokenizer: Final = Tokenizer.from_pretrained(identifier, revision=revision, token=auth_token) return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py index 89899d3e762..c4b1f4f3afd 100644 --- a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py +++ b/tests/proxy_unit_tests/test_custom_tokenizer_bug.py @@ -23,9 +23,9 @@ from litellm.proxy.proxy_server import token_counter def _fake_hf_tokenizer(num_tokens: int) -> MagicMock: encoding = MagicMock() - encoding.ids = list(range(num_tokens)) + encoding.__len__.return_value = num_tokens tokenizer = MagicMock() - tokenizer.encode.return_value = encoding + tokenizer.encode_batch_fast.return_value = [encoding] return tokenizer @@ -68,13 +68,11 @@ async def test_custom_tokenizer_from_model_info_is_used(monkeypatch): ) ) - mock_tokenizer_cls.from_pretrained.assert_called_once_with( - "my-org/custom-tokenizer", revision="v2", auth_token=None - ) + mock_tokenizer_cls.from_pretrained.assert_called_once_with("my-org/custom-tokenizer", revision="v2", token=None) assert response.tokenizer_type == "huggingface_tokenizer" assert response.request_model == "my-embedding-model" assert response.model_used == "self-hosted-embedder" - assert response.total_tokens > 0 + assert response.total_tokens >= 7 @pytest.mark.asyncio diff --git a/tests/test_litellm/litellm_core_utils/event_loop_lag.py b/tests/test_litellm/litellm_core_utils/event_loop_lag.py index c2bbe2204c7..1cac0365547 100644 --- a/tests/test_litellm/litellm_core_utils/event_loop_lag.py +++ b/tests/test_litellm/litellm_core_utils/event_loop_lag.py @@ -25,6 +25,7 @@ async def timed_with_loop_lags(run: Callable[[], Awaitable[T]]) -> tuple[T, floa finished: Final = asyncio.Event() async def timed() -> tuple[T, float]: + await asyncio.sleep(0) started: Final = time.perf_counter() try: return await run(), time.perf_counter() - started diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7fa3d5dd246..6fa99587dbf 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -12877,3 +12877,40 @@ async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_coun assert response.total_tokens > 0 assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from tests.test_litellm.litellm_core_utils.event_loop_lag import assert_loop_stayed_free, timed_with_loop_lags + + claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + + class SlowHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + time.sleep(0.3) + return claude_tokenizer + + monkeypatch.setattr(litellm.utils, "Tokenizer", SlowHubTokenizer) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + { + "model_name": "self-hosted", + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": {"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": "main", "auth_token": None}}, + } + ] + ), + ) + + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me off the loop")) + ) + + assert response.tokenizer_type == "huggingface_tokenizer" + assert response.total_tokens > 0 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 030bdfe03e9..333e7b2ff31 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -477,3 +477,65 @@ async def test_wildcard_route_resolves_underlying_model_minimum(local_model_cost assert deployments[0]["litellm_params"]["model"] == "anthropic/claude-opus-4-6" assert _get_min_token_count_for_deployments(deployments) == 4096 + + +@pytest.mark.asyncio +async def test_async_filter_deployments_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + check = PromptCachingDeploymentCheck(cache=DualCache()) + deployments = _deployments("anthropic/claude-fable-5") + messages = cast(List[AllMessageValues], [{"role": "user", "content": text * 100}]) + + result, took, lags = await timed_with_loop_lags( + lambda: check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=messages + ) + ) + + assert result == deployments + assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("anthropic/claude-fable-5") + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + messages = cast( + List[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": text * 100, "cache_control": {"type": "ephemeral"}}]}], + ) + standard_logging_object = { + "call_type": "acompletion", + "model": "anthropic/claude-fable-5", + "messages": messages, + "model_id": "dep-1", + } + + _, took, lags = await timed_with_loop_lags( + lambda: check.async_log_success_event( + kwargs={"standard_logging_object": standard_logging_object}, + response_obj=None, + start_time=None, + end_time=None, + ) + ) + + assert await PromptCachingCache(cache=cache).async_get_model_id(messages=messages, tools=None) == { + "model_id": "dep-1" + } + assert_loop_stayed_free(took, lags) From c6a0381595b3085937e514c7af3f54c2ceb77598 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:38:01 -0700 Subject: [PATCH 07/10] fix(token_counter): bound concurrent HuggingFace encodes so a burst of large counts cannot exhaust memory --- litellm/constants.py | 1 + litellm/litellm_core_utils/token_counter.py | 6 +++- .../litellm_core_utils/test_token_counter.py | 31 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index a9112c90a54..05d442dd070 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,6 +397,7 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) +TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES: Final = 4 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e3af64fd3bf..8624e79faea 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,6 +3,7 @@ import base64 import io import struct +import threading from collections.abc import Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast @@ -22,6 +23,7 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, + TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES, TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.default_encoding import encoding as default_encoding @@ -320,6 +322,7 @@ Type for a function that counts tokens in a string. """ EXTRAPOLATION_SAMPLES: Final = 16 +_HF_ENCODE_SLOTS: Final = threading.BoundedSemaphore(TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES) def _get_tiktoken_count_function( @@ -587,7 +590,8 @@ def _get_exact_count_function( tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - return len(tokenizer.encode_batch_fast([text])[0]) + with _HF_ENCODE_SLOTS: + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index e3dd5feb18d..655be5146cd 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -2,8 +2,10 @@ # This tests litellm.token_counter.token_counter() function import asyncio import importlib +import threading import time import traceback +from concurrent.futures import ThreadPoolExecutor from typing import Final from unittest.mock import MagicMock @@ -16,6 +18,7 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( _get_exact_count_function, @@ -162,6 +165,34 @@ def test_count_at_or_below_the_cap_is_exact(): assert count_exactly.call_args_list == [(("a" * 5_000,),)] +class _SlowEncoder: + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self.in_flight = 0 + self.peak_in_flight = 0 + + def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: + with self._lock: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + time.sleep(0.1) + with self._lock: + self.in_flight -= 1 + return [[0] * len(text) for text in texts] + + +def test_huggingface_counts_run_at_most_the_configured_number_at_once(): + encoder: Final = _SlowEncoder() + count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + + with ThreadPoolExecutor(max_workers=burst) as pool: + counts: Final = tuple(pool.map(count, ["abc"] * burst)) + + assert counts == (3,) * burst + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + + def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars] From b5d40bef2b830b6d5ccbb3f3fae74eb10968b9cd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:25:04 -0700 Subject: [PATCH 08/10] fix(token_counter): bound offloaded token counts with a dedicated capacity limiter Replace the thread semaphore around HuggingFace encodes with an anyio CapacityLimiter applied at every offloaded count site through offload_token_count, so waiting counts no longer hold slots in the shared 40-thread pool and inline counts on the event loop never block on the bound. Rename TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES to TOKEN_COUNTER_MAX_CONCURRENT_COUNTS --- litellm/constants.py | 2 +- litellm/litellm_core_utils/token_counter.py | 21 ++++++++++----- .../hooks/parallel_request_limiter_v3.py | 4 +-- litellm/proxy/proxy_server.py | 6 ++--- litellm/proxy/utils.py | 4 +-- litellm/router.py | 5 ++-- .../complexity_router/complexity_router.py | 6 ++--- .../io_token_rate_limit_check.py | 4 +-- .../prompt_caching_deployment_check.py | 6 ++--- .../litellm_core_utils/test_token_counter.py | 26 +++++++++++++------ 10 files changed, 52 insertions(+), 32 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 05d442dd070..0bd7def539d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,7 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) -TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES: Final = 4 +TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = 4 MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 8624e79faea..e71ce5dea9c 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,13 +3,14 @@ import base64 import io import struct -import threading -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast +import anyio import httpx import tiktoken from tokenizers import Tokenizer +from typing_extensions import ParamSpec, TypeVar import litellm from litellm import verbose_logger @@ -23,9 +24,10 @@ from litellm.constants import ( MAX_TILE_HEIGHT, MAX_TILE_WIDTH, TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS, - TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES, + TOKEN_COUNTER_MAX_CONCURRENT_COUNTS, TOKEN_COUNTER_MAX_EXACT_CHARS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.default_encoding import encoding as default_encoding from litellm.litellm_core_utils.url_utils import safe_get from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -322,7 +324,15 @@ Type for a function that counts tokens in a string. """ EXTRAPOLATION_SAMPLES: Final = 16 -_HF_ENCODE_SLOTS: Final = threading.BoundedSemaphore(TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES) +T_ParamSpec: Final = ParamSpec("T_ParamSpec") +T_Retval = TypeVar("T_Retval") +_COUNT_OFFLOAD_LIMITER: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + + +def offload_token_count( + function: Callable[T_ParamSpec, T_Retval], +) -> Callable[T_ParamSpec, Awaitable[T_Retval]]: + return asyncify(function, limiter=_COUNT_OFFLOAD_LIMITER) def _get_tiktoken_count_function( @@ -590,8 +600,7 @@ def _get_exact_count_function( tokenizer: Final[Tokenizer] = tokenizer_json["tokenizer"] def count_tokens(text: str) -> int: - with _HF_ENCODE_SLOTS: - return len(tokenizer.encode_batch_fast([text])[0]) + return len(tokenizer.encode_batch_fast([text])[0]) return count_tokens elif tokenizer_json["type"] == "openai_tokenizer": diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 74266801e50..117be770953 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -28,10 +28,10 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( ESTIMATED_OUTPUT_TOKENS_FIELD, @@ -3304,7 +3304,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, ) - raw_estimated_input_tokens: Final = await asyncify(self._estimate_precise_input_tokens)( + raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type ) estimated_input_tokens: Final = max(raw_estimated_input_tokens, 1) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77354e9885f..0e3e5c251c8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -67,6 +67,7 @@ from litellm.litellm_core_utils.litellm_logging import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import ( UI_TEAM_ID, CallbackDelete, @@ -274,7 +275,6 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -12713,7 +12713,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( + _tokenizer_used: Final = await offload_token_count(litellm.utils._select_tokenizer)( model=model_to_use, custom_tokenizer=custom_tokenizer ) @@ -12728,7 +12728,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None ) - total_tokens: Final = await asyncify(litellm.token_counter)( + total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, text=prompt, messages=counted_messages, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1014c2b9b6a..59c09081890 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,6 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, independent_snapshot, @@ -100,6 +99,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms import load_guardrail_translation_mappings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( @@ -2570,7 +2570,7 @@ class ProxyLogging: original_exception=original_exception, ) - request_data.update(await asyncify(_failure_fields_to_lift)(request_data)) + request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data)) # Remove before callbacks iterate — not serialisable request_data.pop("litellm_logging_obj", None) diff --git a/litellm/router.py b/litellm/router.py index 95cabfad4bd..73abd97d80d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -55,7 +55,7 @@ from litellm.constants import ( SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import asyncify, run_async_function +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, coerce_token_limit, @@ -86,6 +86,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( mask_credentials_in_payload, mask_sensitive_structure, ) +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, @@ -11958,7 +11959,7 @@ class Router: try: if not self._pre_call_checks_need_token_count(model, healthy_deployments): return None - return await asyncify(self._count_pre_call_check_tokens)( + return await offload_token_count(self._count_pre_call_check_tokens)( messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter request_kwargs=request_kwargs, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 7d4497fb6f7..cc1bb54255b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2486,14 +2486,14 @@ class ComplexityRouter(CustomLogger): """Real-tokenizer count of the resolved messages plus the out-of-band carriers, off the event loop; None when counting fails, and the gate then leaves the placement alone.""" import litellm - from litellm.litellm_core_utils.asyncify import asyncify + from litellm.litellm_core_utils.token_counter import offload_token_count out_of_band: Final = self._out_of_band_request_text(request_kwargs) try: - counted: Final = await asyncify(litellm.token_counter)( + counted: Final = await offload_token_count(litellm.token_counter)( messages=cast(list, resolved_messages) # cast-ok: token_counter only iterates the sequence ) - return counted + (await asyncify(litellm.token_counter)(text=out_of_band) if out_of_band else 0) + return counted + (await offload_token_count(litellm.token_counter)(text=out_of_band) if out_of_band else 0) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request verbose_router_logger.debug("ComplexityRouter: context-window token count failed. Got - %s", e) return None diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index dca5a2be845..fbd3e18e357 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -21,7 +21,7 @@ import litellm from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.router import RouterCacheEnum, RouterErrors from litellm.utils import get_utc_datetime @@ -467,7 +467,7 @@ async def async_io_token_pre_call_check( request_kwargs: Final = get_io_token_rate_limit_request_kwargs() _model: Final = (deployment.get("litellm_params") or {}).get("model") or "" - estimated_input: Final = await asyncify(_estimate_input_tokens)(request_kwargs, model=_model) + estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model) max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment) dt: Final = get_utc_datetime() diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 9ea528b9700..0589e290b47 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -14,7 +14,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -62,7 +62,7 @@ class PromptCachingDeploymentCheck(CustomLogger): if request_kwargs is not None and request_kwargs.get("_target_order") is not None: return healthy_deployments - if messages is not None and await asyncify(is_prompt_caching_valid_prompt)( + if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( messages=messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), @@ -140,7 +140,7 @@ class PromptCachingDeploymentCheck(CustomLogger): return ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider - if await asyncify(is_prompt_caching_valid_prompt)( + if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, messages=cast(list[AllMessageValues], messages), ): diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 655be5146cd..7cc7c6c763c 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -5,10 +5,10 @@ import importlib import threading import time import traceback -from concurrent.futures import ThreadPoolExecutor from typing import Final from unittest.mock import MagicMock +import anyio.to_thread import pytest import tiktoken @@ -18,12 +18,13 @@ import litellm from litellm import create_pretrained_tokenizer, decode, encode, get_modified_max_tokens from litellm import token_counter as token_counter_old import litellm.constants -from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.token_counter import ( _get_exact_count_function, _get_extrapolating_count_function, _get_tiktoken_count_function, + offload_token_count, ) from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from tests.large_text import text @@ -181,16 +182,25 @@ class _SlowEncoder: return [[0] * len(text) for text in texts] -def test_huggingface_counts_run_at_most_the_configured_number_at_once(): +@pytest.mark.asyncio +async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): encoder: Final = _SlowEncoder() count: Final = _get_exact_count_function(None, {"type": "huggingface_tokenizer", "tokenizer": encoder}) - burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + shared_pool: Final = anyio.to_thread.current_default_thread_limiter() + burst: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS - with ThreadPoolExecutor(max_workers=burst) as pool: - counts: Final = tuple(pool.map(count, ["abc"] * burst)) + async def shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) -> tuple[int, ...]: + if counting.done(): + return () + await asyncio.sleep(0.01) + return (shared_pool.borrowed_tokens, *await shared_pool_borrowed_until_done(counting)) - assert counts == (3,) * burst - assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_HF_ENCODES + counting: Final = asyncio.ensure_future(asyncio.gather(*(offload_token_count(count)("abc") for _ in range(burst)))) + borrowed: Final = await shared_pool_borrowed_until_done(counting) + + assert await counting == [3] * burst + assert len(borrowed) > 1 and max(borrowed) == 0 + assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS def test_token_counter_applies_the_default_cap(): From babf18f37f383b99e697344019482ff97dfea1c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:20:56 -0700 Subject: [PATCH 09/10] test(token_counter): type the project ITPM lag test's request body --- tests/test_litellm/proxy/hooks/test_tpm_concurrent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 398abf60b54..bdaca9ffc2d 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -3682,7 +3682,7 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo handler, _cache = rate_limiter stash = get_or_create_request_stash() warm_tokenizer("claude-fable-5") - data: Dict[str, Any] = { + data: dict[str, object] = { "model": "claude-fable-5", "messages": [{"role": "user", "content": text * 100}], } From 47eb7f434974604a488e5a11193b3ef43410259c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:17:32 -0700 Subject: [PATCH 10/10] fix(token_counter): give each event loop its own count limiter A single process-wide anyio.CapacityLimiter wakes waiters with the asyncio.Event of whichever loop created it, so a fifth loop in another thread (asyncio.run per request under Celery, gunicorn sync workers, or run_async_function from function_with_fallbacks) waited forever once four counts were in flight, and building it at import time raised AsyncLibraryNotFoundError on anyio below 4.2. The limiter now lives in a RunVar and is created on first use inside the running loop. TOKEN_COUNTER_MAX_CONCURRENT_COUNTS reads from the environment like TOKEN_COUNTER_MAX_EXACT_CHARS, and the proxy's tokenizer lookup runs on the shared thread pool so a Hub download no longer holds a count slot. --- litellm/constants.py | 7 +++- litellm/litellm_core_utils/token_counter.py | 20 +++++++++- litellm/proxy/proxy_server.py | 3 +- .../litellm_core_utils/test_token_counter.py | 38 +++++++++++++++++++ 4 files changed, 64 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0bd7def539d..d8c83876e87 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,12 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) -TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = 4 +TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range( + "TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", + default=4, + minimum=1, + maximum=256, +) MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e71ce5dea9c..70562748c6a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -7,6 +7,7 @@ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from typing import Final, Literal, cast import anyio +import anyio.lowlevel import httpx import tiktoken from tokenizers import Tokenizer @@ -326,13 +327,28 @@ Type for a function that counts tokens in a string. EXTRAPOLATION_SAMPLES: Final = 16 T_ParamSpec: Final = ParamSpec("T_ParamSpec") T_Retval = TypeVar("T_Retval") -_COUNT_OFFLOAD_LIMITER: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) +_COUNT_OFFLOAD_LIMITER: Final = anyio.lowlevel.RunVar[anyio.CapacityLimiter]("litellm_count_offload_limiter") + + +def _count_offload_limiter_for_this_loop() -> anyio.CapacityLimiter: + existing: Final = _COUNT_OFFLOAD_LIMITER.get(None) + if existing is not None: + return existing + created: Final = anyio.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + _COUNT_OFFLOAD_LIMITER.set(created) + return created def offload_token_count( function: Callable[T_ParamSpec, T_Retval], ) -> Callable[T_ParamSpec, Awaitable[T_Retval]]: - return asyncify(function, limiter=_COUNT_OFFLOAD_LIMITER) + async def offloaded( + *args: T_ParamSpec.args, + **kwargs: T_ParamSpec.kwargs, # kwargs-ok: ParamSpec keeps the wrapped function's own keyword contract + ) -> T_Retval: + return await asyncify(function, limiter=_count_offload_limiter_for_this_loop())(*args, **kwargs) + + return offloaded def _get_tiktoken_count_function( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e3e5c251c8..453f57131b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -62,6 +62,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) @@ -12713,7 +12714,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await offload_token_count(litellm.utils._select_tokenizer)( + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( model=model_to_use, custom_tokenizer=custom_tokenizer ) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 7cc7c6c763c..7da04d12569 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -5,6 +5,7 @@ import importlib import threading import time import traceback +from concurrent.futures import Future, wait from typing import Final from unittest.mock import MagicMock @@ -203,6 +204,43 @@ async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: + def slow_count(counted: str) -> int: + time.sleep(0.1) + return len(counted) + + result.set_result(asyncio.run(offload_token_count(slow_count)(text))) + + +def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): + loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + results: Final = tuple(Future[int]() for _ in range(loops)) + threads: Final = tuple( + threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) + for size, result in enumerate(results, start=1) + ) + for thread in threads: + thread.start() + + _, pending = wait(results, timeout=5) + + assert not pending + assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("8", 8), ("0", 4), ("not-an-int", 4)], +) +def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") + importlib.reload(litellm.constants) + + def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]