From 8a99efefb24dee974b906eac3b5117a92b75c238 Mon Sep 17 00:00:00 2001 From: fenilfaldu Date: Sat, 24 Jan 2026 19:47:49 +0530 Subject: [PATCH] feat: add async tokenization support with threadpool for large inputs --- litellm/__init__.py | 3 + litellm/litellm_core_utils/token_counter.py | 189 ++++++++++++++++++++ litellm/utils.py | 38 ++++ 3 files changed, 230 insertions(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index e5c09702b9b..6b19f7f55a7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -252,6 +252,9 @@ ssl_ecdh_curve: Optional[str] = ( ) disable_streaming_logging: bool = False disable_token_counter: bool = False +async_tokenizer_threshold_bytes: Optional[int] = None # Run tokenization in threadpool for inputs larger than this (e.g., 500000 for 500KB) +tokenizer_threadpool_max_workers: int = 4 # Maximum threads for async tokenization +tokenizer_timeout_seconds: Optional[float] = None # Timeout for tokenization in threadpool (None = no timeout) disable_add_transform_inline_image_block: bool = False disable_add_user_agent_to_request_tags: bool = False extra_spend_tag_headers: Optional[List[str]] = None diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index a99bd1cd0f3..d93a4224697 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -1,8 +1,10 @@ # What is this? ## Helper utilities for token counting +import asyncio import base64 import io import struct +from concurrent.futures import ThreadPoolExecutor from typing import ( Any, Callable, @@ -19,6 +21,8 @@ import tiktoken import litellm from litellm import verbose_logger + +_tokenizer_threadpool: Optional[ThreadPoolExecutor] = None from litellm.constants import ( DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_TOKEN_COUNT, @@ -313,6 +317,70 @@ Type for a function that counts tokens in a string. """ +def _get_tokenizer_threadpool() -> Optional[ThreadPoolExecutor]: + """ + Get or create the global threadpool for async tokenization. + Returns None if async tokenization is disabled. + """ + global _tokenizer_threadpool + + if litellm.async_tokenizer_threshold_bytes is None: + return None + + if _tokenizer_threadpool is None: + max_workers = litellm.tokenizer_threadpool_max_workers or 4 + _tokenizer_threadpool = ThreadPoolExecutor( + max_workers=max_workers, + thread_name_prefix="litellm_tokenizer" + ) + verbose_logger.debug( + f"Created tokenizer threadpool with {max_workers} workers" + ) + + return _tokenizer_threadpool + + +def _calculate_input_size_bytes( + text: Optional[Union[str, List[str]]], + messages: Optional[List[Union[AllMessageValues, Message]]] +) -> int: + """Calculate approximate size of input for threshold checking""" + size = 0 + + if text is not None: + if isinstance(text, str): + size = len(text.encode('utf-8')) + elif isinstance(text, list): + size = sum(len(str(t).encode('utf-8')) for t in text) + + if messages is not None: + for msg in messages: + if isinstance(msg, dict): + content = msg.get("content", "") + if isinstance(content, str): + size += len(content.encode('utf-8')) + elif isinstance(content, list): + for item in content: + if isinstance(item, dict) and "text" in item: + size += len(str(item["text"]).encode('utf-8')) + + return size + + +def _should_use_async_tokenization(input_size_bytes: int) -> bool: + """ + Determine if we should use async tokenization based on input size. + + Returns True if: + - async_tokenizer_threshold_bytes is set AND + - input size exceeds the threshold + """ + if litellm.async_tokenizer_threshold_bytes is None: + return False + + return input_size_bytes > litellm.async_tokenizer_threshold_bytes + + class _MessageCountParams: """ A class to hold the parameters for counting tokens in messages. @@ -425,6 +493,127 @@ def token_counter( return num_tokens +async def async_token_counter( + model="", + custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None, + text: Optional[Union[str, List[str]]] = None, + messages: Optional[List[Union[AllMessageValues, Message]]] = None, + count_response_tokens: Optional[bool] = False, + tools: Optional[List[ChatCompletionToolParam]] = None, + tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None, + use_default_image_token_count: Optional[bool] = False, + default_token_count: Optional[int] = None, +) -> int: + """ + Async version of token_counter that runs tokenization in a threadpool for large inputs. + + This prevents CPU-intensive tokenization from blocking the async event loop. + For small inputs, falls back to sync tokenization for efficiency. + + Args: + Same as token_counter() + + Returns: + int: The number of tokens in the text. + """ + from litellm.utils import convert_list_message_to_dict + + if litellm.disable_token_counter is True: + return 0 + + input_size = _calculate_input_size_bytes(text, messages) + + if not _should_use_async_tokenization(input_size): + # Input is small, run synchronously (faster for small inputs) + return token_counter( + model=model, + custom_tokenizer=custom_tokenizer, + text=text, + messages=messages, + count_response_tokens=count_response_tokens, + tools=tools, + tool_choice=tool_choice, + use_default_image_token_count=use_default_image_token_count, + default_token_count=default_token_count, + ) + + # Large input - use threadpool to avoid blocking event loop + verbose_logger.debug( + f"Using async tokenization for large input ({input_size:,} bytes > {litellm.async_tokenizer_threshold_bytes:,} threshold)" + ) + + threadpool = _get_tokenizer_threadpool() + + if threadpool is None: + verbose_logger.warning( + "Async tokenization requested but threadpool not available, falling back to sync" + ) + return token_counter( + model=model, + custom_tokenizer=custom_tokenizer, + text=text, + messages=messages, + count_response_tokens=count_response_tokens, + tools=tools, + tool_choice=tool_choice, + use_default_image_token_count=use_default_image_token_count, + default_token_count=default_token_count, + ) + + # Run tokenization in threadpool + loop = asyncio.get_event_loop() + + try: + # Run in threadpool with optional timeout + if litellm.tokenizer_timeout_seconds is not None: + num_tokens = await asyncio.wait_for( + loop.run_in_executor( + threadpool, + token_counter, + model, + custom_tokenizer, + text, + messages, + count_response_tokens, + tools, + tool_choice, + use_default_image_token_count, + default_token_count, + ), + timeout=litellm.tokenizer_timeout_seconds + ) + else: + num_tokens = await loop.run_in_executor( + threadpool, + token_counter, + model, + custom_tokenizer, + text, + messages, + count_response_tokens, + tools, + tool_choice, + use_default_image_token_count, + default_token_count, + ) + + verbose_logger.debug( + f"Async tokenization completed: {num_tokens} tokens" + ) + return num_tokens + + except asyncio.TimeoutError: + verbose_logger.error( + f"Tokenization timed out after {litellm.tokenizer_timeout_seconds}s, returning 0 tokens" + ) + return 0 + except Exception as e: + verbose_logger.error( + f"Error in async tokenization: {e}, returning 0 tokens" + ) + return 0 + + def _count_messages( params: _MessageCountParams, messages: List[AllMessageValues], diff --git a/litellm/utils.py b/litellm/utils.py index ad9e36795af..8e79d821309 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2267,6 +2267,44 @@ def token_counter( ) +async def async_token_counter( + model="", + custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None, + text: Optional[Union[str, List[str]]] = None, + messages: Optional[List] = None, + count_response_tokens: Optional[bool] = False, + tools: Optional[List[ChatCompletionToolParam]] = None, + tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None, + use_default_image_token_count: Optional[bool] = False, + default_token_count: Optional[int] = None, +) -> int: + """ + Async version of token_counter that runs in threadpool for large inputs. + + This prevents CPU-intensive tokenization from blocking the async event loop. + Uses threadpool when input size exceeds async_tokenizer_threshold_bytes. + + Args: + Same as token_counter() + + Returns: + int: The number of tokens in the text. + """ + from litellm.litellm_core_utils.token_counter import async_token_counter as _async_token_counter_impl + + return await _async_token_counter_impl( + model=model, + custom_tokenizer=custom_tokenizer, + text=text, + messages=messages, + count_response_tokens=count_response_tokens, + tools=tools, + tool_choice=tool_choice, + use_default_image_token_count=use_default_image_token_count, + default_token_count=default_token_count, + ) + + def supports_httpx_timeout(custom_llm_provider: str) -> bool: """ Helper function to know if a provider implementation supports httpx timeout