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

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

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(

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
     )
 

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]