From 47eb7f434974604a488e5a11193b3ef43410259c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:17:32 -0700 Subject: [PATCH] fix(token_counter): give each event loop its own count limiter A single process-wide anyio.CapacityLimiter wakes waiters with the asyncio.Event of whichever loop created it, so a fifth loop in another thread (asyncio.run per request under Celery, gunicorn sync workers, or run_async_function from function_with_fallbacks) waited forever once four counts were in flight, and building it at import time raised AsyncLibraryNotFoundError on anyio below 4.2. The limiter now lives in a RunVar and is created on first use inside the running loop. TOKEN_COUNTER_MAX_CONCURRENT_COUNTS reads from the environment like TOKEN_COUNTER_MAX_EXACT_CHARS, and the proxy's tokenizer lookup runs on the shared thread pool so a Hub download no longer holds a count slot. --- litellm/constants.py | 7 +++- litellm/litellm_core_utils/token_counter.py | 20 +++++++++- litellm/proxy/proxy_server.py | 3 +- .../litellm_core_utils/test_token_counter.py | 38 +++++++++++++++++++ 4 files changed, 64 insertions(+), 4 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0bd7def539d..d8c83876e87 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -397,7 +397,12 @@ TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range( minimum=1, maximum=1_000_000_000, ) -TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = 4 +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)) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index e71ce5dea9c..70562748c6a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -7,6 +7,7 @@ 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 @@ -326,13 +327,28 @@ 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.CapacityLimiter(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) +_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]]: - return asyncify(function, limiter=_COUNT_OFFLOAD_LIMITER) + 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( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e3e5c251c8..453f57131b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -62,6 +62,7 @@ from litellm.constants import ( 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, ) @@ -12713,7 +12714,7 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) CustomHuggingfaceTokenizer | None, model_info.get("custom_tokenizer", None), ) - _tokenizer_used: Final = await offload_token_count(litellm.utils._select_tokenizer)( + _tokenizer_used: Final = await asyncify(litellm.utils._select_tokenizer)( model=model_to_use, custom_tokenizer=custom_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 7cc7c6c763c..7da04d12569 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -5,6 +5,7 @@ import importlib import threading import time import traceback +from concurrent.futures import Future, wait from typing import Final from unittest.mock import MagicMock @@ -203,6 +204,43 @@ async def test_offloaded_counts_do_not_borrow_from_the_shared_thread_pool(): assert 1 < encoder.peak_in_flight <= TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +def _count_in_a_fresh_event_loop(text: str, result: Future[int]) -> None: + def slow_count(counted: str) -> int: + time.sleep(0.1) + return len(counted) + + result.set_result(asyncio.run(offload_token_count(slow_count)(text))) + + +def test_offloaded_counts_finish_in_every_event_loop_that_shares_the_process(): + loops: Final = 2 * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + results: Final = tuple(Future[int]() for _ in range(loops)) + threads: Final = tuple( + threading.Thread(target=_count_in_a_fresh_event_loop, args=("a" * size, result), daemon=True) + for size, result in enumerate(results, start=1) + ) + for thread in threads: + thread.start() + + _, pending = wait(results, timeout=5) + + assert not pending + assert tuple(result.result() for result in results) == tuple(range(1, loops + 1)) + + +@pytest.mark.parametrize( + ("configured", "expected"), + [("8", 8), ("0", 4), ("not-an-int", 4)], +) +def test_max_concurrent_counts_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS", configured) + try: + assert importlib.reload(litellm.constants).TOKEN_COUNTER_MAX_CONCURRENT_COUNTS == expected + finally: + monkeypatch.delenv("TOKEN_COUNTER_MAX_CONCURRENT_COUNTS") + importlib.reload(litellm.constants) + + def test_token_counter_applies_the_default_cap(): max_exact_chars: Final = litellm.constants.TOKEN_COUNTER_MAX_EXACT_CHARS prose: Final = ("The quick brown fox jumps over the lazy dog. " * (max_exact_chars // 45 + 1))[:max_exact_chars]