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 = {