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] 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"