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 d53686e5e5bb..a9112c90a548 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,6 +391,12 @@ 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 3732ffd734ca..d526f0f2996a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -8,6 +8,7 @@ import httpx import tiktoken +from tokenizers import Tokenizer import litellm from litellm import verbose_logger @@ -21,6 +22,7 @@ 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 4694fa8fbed0..fddeacc37f94 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 @@ 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 b7bb58378d40..20b41cf15309 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -6,6 +6,7 @@ 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.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 d526f0f2996a..e3a0d4c785a4 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 fddeacc37f94..5da8413ac9b6 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 count_off_the_loop() -> tuple[int, float]: 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) + + estimate: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars)(front_heavy) - count_tokens: Final = _get_extrapolating_count_function(count_exactly, max_exact_chars=max_exact_chars) + assert abs(estimate - exact) <= exact // 100 + assert sum(len(call.args[0]) for call in count_exactly.call_args_list) <= max_exact_chars - assert count_tokens("a" * 1_000 + "b" * 4_000) == expected + +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) + + estimate: Final = token_counter_new(model="gpt-5.6", text=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 + 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 5da8413ac9b6..db87bc526009 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 count_off_the_loop() -> tuple[int, float]: @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 31437af7770a..74266801e505 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._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 @@ async def _reserve_project_io_tokens_or_raise( 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 ddf31cb1d8ad..1014c2b9b6a3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,6 +91,7 @@ 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 @@ async def post_call_failure_hook( 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 000000000000..c2bbe2204c73 --- /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 db87bc526009..e3dd5feb18d1 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 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 2839acab6b03..398abf60b546 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 @@ def raise_on_merge(**_kwargs): ) +@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 20b41cf15309..7fa3d5dd2465 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() - - (tokens, took), lags = await asyncio.gather(count_off_the_loop(), _loop_wake_lags(counted)) + response, took, lags = await timed_with_loop_lags( + lambda: proxy_server_module.token_counter(TokenCountRequest(model="claude-fable-5", prompt=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 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 9462f2c8eb0d..b78ec7dcff6e 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 e3a0d4c785a4..e3af64fd3bf3 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 01d42627001c..dca5a2be845b 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 @@ 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 a5a682711114..3cef1c7bb632 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 @@ async def test_router_pins_kwargs_for_io_limited_deployment(self): 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 0915b8dd1b94..77354e9885f2 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 70362e604951..9ea528b97007 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 @@ 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 @@ async def async_filter_deployments( 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 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti 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 d0e11bc95518..ab905381b223 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 89899d3e762b..c4b1f4f3afdb 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 @@ 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 c2bbe2204c73..1cac03655471 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 7fa3d5dd2465..6fa99587dbf5 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 030bdfe03e90..333e7b2ff318 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 a9112c90a548..05d442dd070b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,6 +397,7 @@ 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 e3af64fd3bf3..8624e79faeac 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 @@ 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 @@ def calculate_img_tokens( """ 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 e3dd5feb18d1..655be5146cde 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 @@ 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 05d442dd070b..0bd7def539de 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,7 @@ 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 8624e79faeac..e71ce5dea9cc 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 @@ 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 @@ def calculate_img_tokens( """ 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 74266801e505..117be7709535 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._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 @@ async def _reserve_project_io_tokens_or_raise( 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 77354e9885f2..0e3e5c251c81 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -67,6 +67,7 @@ ) 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 @@ def generate_feedback_box(): 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 1014c2b9b6a3..59c090818905 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,6 @@ 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.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 @@ async def post_call_failure_hook( 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 95cabfad4bd9..73abd97d80d9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -55,7 +55,7 @@ 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 @@ 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 @@ async def _acount_pre_call_check_tokens( 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 7d4497fb6f74..cc1bb54255ba 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -2486,14 +2486,14 @@ async def _counted_request_tokens( """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 dca5a2be845b..fbd3e18e3574 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 @@ 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 9ea528b97007..0589e290b47e 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 @@ 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 @@ async def async_filter_deployments( 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 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti 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 655be5146cde..7cc7c6c763cc 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 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 @@ 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 @@ def encode_batch_fast(self, texts: list[str]) -> list[list[int]]: 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 398abf60b546..bdaca9ffc2d4 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 0bd7def539de..d8c83876e878 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,12 @@ 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 e71ce5dea9cc..70562748c6aa 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -7,6 +7,7 @@ from typing import Final, Literal, cast import anyio +import anyio.lowlevel import httpx import tiktoken from tokenizers import Tokenizer @@ -326,13 +327,28 @@ def calculate_img_tokens( 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 0e3e5c251c81..453f57131b59 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -62,6 +62,7 @@ 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 7cc7c6c763cc..7da04d12569d 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 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 shared_pool_borrowed_until_done(counting: asyncio.Future[list[int]]) - 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]