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

@@ -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))

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

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)

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,

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)

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,

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

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

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),
         ):

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