From 0c611e63c86bad89568e60b6a82c9bc866c82471 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 09:12:49 +0000 Subject: [PATCH 1/2] fix(utils): cache custom HuggingFace tokenizers across /utils/token_counter requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 9 ++- tests/test_litellm/proxy/test_proxy_server.py | 57 +++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 18df5e2abf7..af22b11224b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2202,15 +2202,20 @@ def _is_streaming_request( def _select_tokenizer(model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None = None): if custom_tokenizer is not None: - _tokenizer: Final = create_pretrained_tokenizer( + return _select_custom_tokenizer_helper( identifier=custom_tokenizer["identifier"], revision=custom_tokenizer["revision"], auth_token=custom_tokenizer["auth_token"], ) - return _tokenizer return _select_tokenizer_helper(model=model) +@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) +def _select_custom_tokenizer_helper(identifier: str, revision: str, auth_token: str | None) -> SelectTokenizerResponse: + verbose_logger.debug("Loading custom HuggingFace tokenizer %s (revision %s)", identifier, revision) + return create_pretrained_tokenizer(identifier=identifier, revision=revision, auth_token=auth_token) + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse: if litellm.disable_hf_tokenizer_download is True: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 04173ced776..552044e2598 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13521,3 +13521,60 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp assert response.tokenizer_type == "huggingface_tokenizer" assert response.total_tokens > 0 assert_loop_stayed_free(took, lags) + + +async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch): + from tokenizers import Tokenizer + + from litellm import Router + from litellm.types.router import DeploymentTypedDict + + claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + loads: Final[list[tuple[str, str, str | None]]] = [] + + class CountingHubTokenizer: + @staticmethod + def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: + loads.append((identifier, revision, token)) + return claude_tokenizer + + def deployment(model_name: str, revision: str, auth_token: str | None) -> DeploymentTypedDict: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"}, + "model_info": { + "custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": revision, "auth_token": auth_token} + }, + } + + monkeypatch.setattr(litellm.utils, "Tokenizer", CountingHubTokenizer) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + deployment("self-hosted", "main", None), + deployment("self-hosted-pinned", "v2", None), + deployment("self-hosted-private", "main", "hf_test_token"), + ] + ), + ) + litellm.utils._select_custom_tokenizer_helper.cache_clear() + try: + responses: Final = [ + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once")) + for _ in range(3) + ] + assert loads == [("my-org/tokenizer", "main", None)] + assert all(response.tokenizer_type == "huggingface_tokenizer" for response in responses) + assert len({response.total_tokens for response in responses}) == 1 + assert responses[0].total_tokens > 0 + + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once")) + await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once")) + assert loads == [ + ("my-org/tokenizer", "main", None), + ("my-org/tokenizer", "v2", None), + ("my-org/tokenizer", "main", "hf_test_token"), + ] + finally: + litellm.utils._select_custom_tokenizer_helper.cache_clear() From b64e430e93c7fb5a197845b4f5f8f10654d4c42d Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 15 Sep 2026 09:37:53 +0000 Subject: [PATCH 2/2] test(proxy): record custom tokenizer loads with a mock instead of a mutable list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_server.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 552044e2598..42af8e0af21 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13529,14 +13529,8 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi from litellm import Router from litellm.types.router import DeploymentTypedDict - claude_tokenizer: Final = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] - loads: Final[list[tuple[str, str, str | None]]] = [] - - class CountingHubTokenizer: - @staticmethod - def from_pretrained(identifier: str, revision: str = "main", token: str | None = None) -> Tokenizer: - loads.append((identifier, revision, token)) - return claude_tokenizer + claude_tokenizer: Final[Tokenizer] = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"] + from_pretrained: Final = MagicMock(return_value=claude_tokenizer) def deployment(model_name: str, revision: str, auth_token: str | None) -> DeploymentTypedDict: return { @@ -13547,7 +13541,7 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi }, } - monkeypatch.setattr(litellm.utils, "Tokenizer", CountingHubTokenizer) + monkeypatch.setattr(litellm.utils, "Tokenizer", MagicMock(from_pretrained=from_pretrained)) monkeypatch.setattr( "litellm.proxy.proxy_server.llm_router", Router( @@ -13564,17 +13558,17 @@ async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revisi await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once")) for _ in range(3) ] - assert loads == [("my-org/tokenizer", "main", None)] + assert from_pretrained.call_args_list == [mock.call("my-org/tokenizer", revision="main", token=None)] assert all(response.tokenizer_type == "huggingface_tokenizer" for response in responses) assert len({response.total_tokens for response in responses}) == 1 assert responses[0].total_tokens > 0 await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once")) await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once")) - assert loads == [ - ("my-org/tokenizer", "main", None), - ("my-org/tokenizer", "v2", None), - ("my-org/tokenizer", "main", "hf_test_token"), + assert from_pretrained.call_args_list == [ + mock.call("my-org/tokenizer", revision="main", token=None), + mock.call("my-org/tokenizer", revision="v2", token=None), + mock.call("my-org/tokenizer", revision="main", token="hf_test_token"), ] finally: litellm.utils._select_custom_tokenizer_helper.cache_clear()