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] 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():