fix(token_counter): release the GIL for HuggingFace counts and cap exact counting per string

litellm/constants.py

@@ -398,6 +398,18 @@
     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,
+)
+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))

litellm/litellm_core_utils/token_counter.py

@@ -3,11 +3,15 @@
 import base64
 import io
 import struct
-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 anyio.lowlevel
 import httpx
 import tiktoken
+from tokenizers import Tokenizer
+from typing_extensions import ParamSpec, TypeVar
 
 import litellm
 from litellm import verbose_logger
@@ -21,7 +25,10 @@
     MAX_TILE_HEIGHT,
     MAX_TILE_WIDTH,
     TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS,
+    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
@@ -317,6 +324,32 @@ def calculate_img_tokens(
 Type for a function that counts tokens in a string.
 """
 
+EXTRAPOLATION_SAMPLES: Final = 16
+T_ParamSpec: Final = ParamSpec("T_ParamSpec")
+T_Retval = TypeVar("T_Retval")
+_COUNT_OFFLOAD_LIMITER: Final = anyio.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]]:
+    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(
     encode_length: Callable[[str], int],
@@ -538,9 +571,40 @@ 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)
+        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,
+) -> 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 +613,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":

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -31,6 +31,7 @@
 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,
@@ -3307,7 +3308,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 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

@@ -63,11 +63,13 @@
     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,
 )
 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,
@@ -272,7 +274,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,
@@ -12816,7 +12817,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)
@@ -12829,7 +12832,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

@@ -101,6 +101,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 (
@@ -2900,7 +2901,7 @@ async def post_call_failure_hook(
                 original_exception=original_exception,
             )
 
-        request_data.update(_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

@@ -67,7 +67,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,
@@ -98,6 +98,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,
@@ -12113,7 +12114,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,

litellm/router_strategy/complexity_router/complexity_router.py

@@ -2568,14 +2568,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

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.token_counter import offload_token_count
 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 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,6 +14,7 @@
     AnthropicCacheControlHook,
 )
 from litellm.integrations.custom_logger import CustomLogger, Span
+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
@@ -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 offload_token_count(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 offload_token_count(is_prompt_caching_valid_prompt)(
             model=model,
             messages=cast(list[AllMessageValues], messages),
         ):

litellm/utils.py

@@ -2293,15 +2293,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}
 
 

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

tests/test_litellm/litellm_core_utils/event_loop_lag.py

@@ -0,0 +1,40 @@
+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]:
+        await asyncio.sleep(0)
+        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"