From 5de851a68032e2dee97a9e48d1d5225049989138 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:21:31 -0700 Subject: [PATCH] fix(utils): keep a loaded Hugging Face tokenizer over a racing failure and retry once per window A download that fails after another count loaded the same tokenizer no longer hides it, the first count after the retry window claims the retry while the others keep using tiktoken, and HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS defaults to 60s to match the proxy's tokenizer reuse window --- litellm/constants.py | 2 +- litellm/utils.py | 72 +++++++++++---- .../litellm_core_utils/test_token_counter.py | 92 ++++++++++++++++--- .../unit/litellm_core_utils/test_tokenizer.py | 12 +-- .../test_huggingface_chat_transformation.py | 6 +- 5 files changed, 146 insertions(+), 38 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6a503463a8c..6b94c24cc7e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -157,7 +157,7 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_ DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250) HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range( - "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 300, 1, 86400 + "HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS", 60, 1, 86400 ) # Maximum wall-clock seconds a streaming response is allowed to run. diff --git a/litellm/utils.py b/litellm/utils.py index 379eddb3155..77e48f77f7d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -29,7 +29,7 @@ import threading import time import traceback from contextlib import nullcontext -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import lru_cache, wraps from importlib import resources from inspect import iscoroutine @@ -2645,9 +2645,15 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None: return None +@dataclass(frozen=True, slots=True) +class _HuggingFaceTokenizerLoaded: + pass + + @dataclass(frozen=True, slots=True) class _HuggingFaceTokenizerLoadFailure: reason: str + retry_at: float def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: @@ -2664,22 +2670,44 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _Hugg if kind is None: return None backend: Final = _huggingface_tokenizer_backend() - key: Final = _huggingface_tokenizer_failure_key(kind, backend) + key: Final = _huggingface_tokenizer_outcome_key(kind, backend) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): - cached_failure: Final = _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.get_cache(key) - if isinstance(cached_failure, _HuggingFaceTokenizerLoadFailure): - return cached_failure - tokenizer: Final = _load_huggingface_tokenizer_recording_failure(kind, backend, key) + known: Final = _claim_huggingface_tokenizer_load(key) + if isinstance(known, _HuggingFaceTokenizerLoadFailure): + return known + tokenizer: Final = ( + _load_huggingface_tokenizer(kind, backend) + if isinstance(known, _HuggingFaceTokenizerLoaded) + else _load_huggingface_tokenizer_remembering_outcome(kind, backend, key) + ) if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) return tokenizer return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} -def _huggingface_tokenizer_failure_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: +def _huggingface_tokenizer_outcome_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: return f"{kind}:{backend.name}" +def _claim_huggingface_tokenizer_load( + key: str, +) -> _HuggingFaceTokenizerLoaded | _HuggingFaceTokenizerLoadFailure | None: + with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: + known: Final = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.get_cache(key) + now: Final = time.monotonic() + match known: + case _HuggingFaceTokenizerLoadFailure(retry_at=retry_at) if now >= retry_at: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache( + key, replace(known, retry_at=now + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) + ) + return None + case _HuggingFaceTokenizerLoaded() | _HuggingFaceTokenizerLoadFailure(): + return known + case _: + return None + + @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decision) -> HuggingFace: """One tokenizer per kind and backend; `backend` is the cache key, the dispatch re-derives it.""" @@ -2702,19 +2730,30 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio assert_never(kind) -def _load_huggingface_tokenizer_recording_failure( +def _load_huggingface_tokenizer_remembering_outcome( kind: HuggingFaceTokenizerKind, backend: Decision, key: str ) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: try: - return _load_huggingface_tokenizer(kind, backend) + tokenizer: Final = _load_huggingface_tokenizer(kind, backend) except Exception as error: from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - failure: Final = _HuggingFaceTokenizerLoadFailure(reason=str(error)) - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.set_cache(key, failure, ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) - return failure + return _remember_huggingface_tokenizer_failure(key, str(error)) + with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache(key, _HuggingFaceTokenizerLoaded()) + return tokenizer + + +def _remember_huggingface_tokenizer_failure(key: str, reason: str) -> _HuggingFaceTokenizerLoadFailure: + failure: Final = _HuggingFaceTokenizerLoadFailure( + reason=reason, retry_at=time.monotonic() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + ) + with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: + if not isinstance(_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.get_cache(key), _HuggingFaceTokenizerLoaded): + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache(key, failure) + return failure def encode(model="", text="", custom_tokenizer: dict | None = None): @@ -8332,7 +8371,7 @@ def trim_messages( from litellm.caching.in_memory_cache import InMemoryCache -class _HuggingFaceTokenizerLoadFailureCache(Protocol): +class _HuggingFaceTokenizerLoadOutcomeCache(Protocol): def get_cache(self, key: str) -> object | None: ... def set_cache(self, key: str, value: object, ttl: float | None = None) -> None: ... @@ -8340,16 +8379,17 @@ class _HuggingFaceTokenizerLoadFailureCache(Protocol): def flush_cache(self) -> None: ... -_HUGGINGFACE_TOKENIZER_LOAD_FAILURES: Final[_HuggingFaceTokenizerLoadFailureCache] = ( +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES: Final[_HuggingFaceTokenizerLoadOutcomeCache] = ( cast( # cast-ok: cache methods are untyped - _HuggingFaceTokenizerLoadFailureCache, + _HuggingFaceTokenizerLoadOutcomeCache, InMemoryCache( max_size_in_memory=16, - default_ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, + default_ttl=2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, clock=lambda: time.monotonic(), ), ) ) +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: Final = threading.Lock() class AvailableModelsCache(InMemoryCache): diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 194d7995735..93deda720ff 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -941,7 +941,7 @@ def test_token_counter(): import unittest from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, @@ -950,7 +950,7 @@ from litellm.utils import ( # Clear the cache at module load to ensure clean state _load_huggingface_tokenizer.cache_clear() -_HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() class TestTokenizerSelection(unittest.TestCase): @@ -962,7 +962,7 @@ class TestTokenizerSelection(unittest.TestCase): --dist=loadscope (tests from same file run on same worker). """ _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @patch("litellm.utils.tokenizer_dispatch.from_pretrained") def test_llama3_tokenizer_api_failure(self, mock_from_pretrained): @@ -1051,7 +1051,7 @@ class TestTokenizerSelection(unittest.TestCase): self, mock_from_pretrained: MagicMock ): _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: with patch("litellm.utils.time.monotonic", return_value=0.0): mock_from_pretrained.side_effect = RuntimeError("HF unavailable") @@ -1069,7 +1069,7 @@ class TestTokenizerSelection(unittest.TestCase): assert second_count == expected finally: _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @patch("litellm.utils.time.monotonic", return_value=299.0) @patch("litellm.utils.tokenizer_dispatch.from_pretrained") @@ -1077,7 +1077,7 @@ class TestTokenizerSelection(unittest.TestCase): self, mock_from_pretrained: MagicMock, mock_monotonic: MagicMock ): _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: def from_pretrained_side_effect(_: str) -> object: if mock_from_pretrained.call_count == 1: @@ -1107,7 +1107,7 @@ class TestTokenizerSelection(unittest.TestCase): assert mock_from_pretrained.call_count == 2 finally: _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() _CODE_SAMPLE: Final = " for item in items:\n total += item.price * item.quantity\n" * 20 @@ -1175,7 +1175,7 @@ def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch: def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() start: Final = threading.Barrier(8) def first_tokenizer(_: int) -> object: @@ -1190,7 +1190,7 @@ def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplog: pytest.LogCaptureFixture): _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" try: with ( @@ -1203,14 +1203,82 @@ def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplo assert len(fallback_warnings) == 1 finally: _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def test_download_that_fails_after_another_count_loaded_the_tokenizer_does_not_hide_it(): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + slow_download_started: Final = threading.Event() + loading_count_finished: Final = threading.Event() + + def download(_: str) -> object: + if threading.current_thread().name.startswith("slow-download"): + slow_download_started.set() + loading_count_finished.wait(timeout=5) + raise RuntimeError("HF rate limited") + return sentinel.tokenizer + + try: + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download): + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="slow-download") as pool: + slow_count: Final = pool.submit(_select_tokenizer_helper, model) + assert slow_download_started.wait(timeout=5) + loading_count: Final = _select_tokenizer_helper(model) + loading_count_finished.set() + slow_count_type: Final = slow_count.result(timeout=5)["type"] + later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3)) + assert slow_count_type == "openai_tokenizer" + assert loading_count["tokenizer"] is sentinel.tokenizer + assert tuple(count["tokenizer"] for count in later_counts) == (sentinel.tokenizer,) * 3 + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def test_counts_during_a_retry_after_the_window_use_tiktoken_instead_of_downloading_again(): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + retry_started: Final = threading.Event() + other_counts_finished: Final = threading.Event() + + def download(_: str) -> object: + if threading.current_thread().name.startswith("retry"): + retry_started.set() + other_counts_finished.wait(timeout=5) + return sentinel.tokenizer + raise RuntimeError("HF unavailable") + + try: + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, + patch("litellm.utils.time.monotonic", return_value=0.0) as clock, + ): + first_type: Final = _select_tokenizer_helper(model)["type"] + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1.0 + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="retry") as pool: + retry: Final = pool.submit(_select_tokenizer_helper, model) + assert retry_started.wait(timeout=5) + types_during_retry: Final = tuple(_select_tokenizer_helper(model)["type"] for _ in range(3)) + other_counts_finished.set() + retried: Final = retry.result(timeout=5) + downloads: Final = from_pretrained.call_count + assert first_type == "openai_tokenizer" + assert types_during_retry == ("openai_tokenizer",) * 3 + assert downloads == 2 + assert retried["tokenizer"] is sentinel.tokenizer + finally: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered(): from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" forked: Final = ForkedAfterNativeRuntimeStarted("native runtime started before fork") try: @@ -1223,7 +1291,7 @@ def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered assert loaded["tokenizer"] is sentinel.tokenizer finally: _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() def test_token_counter_with_anthropic_tool_use(): diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 61e797d7166..cdb2764ecd5 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -408,14 +408,14 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp from unittest.mock import patch from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _load_huggingface_tokenizer, _select_tokenizer_helper, ) monkeypatch.setenv("LITELLM_RUST", "false") _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): result = _select_tokenizer_helper("llama-2") @@ -425,7 +425,7 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") finally: _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) @@ -483,14 +483,14 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): from unittest.mock import patch from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _load_huggingface_tokenizer, _select_tokenizer_helper, ) monkeypatch.setenv("LITELLM_RUST", "false") _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model = "llama-2\r\nforged-model\x1b[31m\u2028\u2029" secret = "sk-" + "x" * 48 failure = OSError("download failed\r\nforged-error\x1b[31m api_key=" + secret) @@ -518,7 +518,7 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): assert "install tokenizers and huggingface-hub" in message finally: _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) diff --git a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py index 04c2ce153c6..57032e5bacc 100644 --- a/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py +++ b/tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py @@ -10,7 +10,7 @@ import respx import litellm from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.huggingface.common_utils import _fetch_inference_provider_mapping -from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, _load_huggingface_tokenizer +from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _load_huggingface_tokenizer HUGGINGFACE_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct" HUGGINGFACE_ROUTER_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct-Turbo" @@ -85,12 +85,12 @@ def isolate_huggingface_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator monkeypatch.delenv("HUGGINGFACE_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_KEY", raising=False) _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() _fetch_inference_provider_mapping.cache_clear() yield _fetch_inference_provider_mapping.cache_clear() _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() PROVIDER_MAPPING_RESPONSE = {