diff --git a/litellm/utils.py b/litellm/utils.py index 38f56765389..69d70d434ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2670,13 +2670,14 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None: @dataclass(frozen=True, slots=True) class _HuggingFaceTokenizerLoaded: - pass + tokenizer: HuggingFace @dataclass(frozen=True, slots=True) class _HuggingFaceTokenizerLoadFailure: reason: str retry_at: float + retrying: bool = False _HuggingFaceTokenizerKey: TypeAlias = tuple[HuggingFaceTokenizerKind, Decision] @@ -2695,25 +2696,33 @@ class _HuggingFaceTokenizerLoadOutcomes: known: Final = self.by_tokenizer.get(key) now: Final = time.monotonic() match known: - case _HuggingFaceTokenizerLoadFailure(retry_at=retry_at) if now >= retry_at: - self._remember(key, replace(known, retry_at=now + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS)) + case _HuggingFaceTokenizerLoadFailure(retry_at=retry_at, retrying=False) if now >= retry_at: + self._remember(key, replace(known, retrying=True)) return None case _: return known - def remember_loaded(self, key: _HuggingFaceTokenizerKey) -> None: + def remember_loaded(self, key: _HuggingFaceTokenizerKey, tokenizer: HuggingFace) -> None: with self.lock: - self._remember(key, _HuggingFaceTokenizerLoaded()) + self._remember(key, _HuggingFaceTokenizerLoaded(tokenizer)) - def remember_failure(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadFailure: + def remember_failure(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadOutcome: failure: Final = _HuggingFaceTokenizerLoadFailure( reason=reason, retry_at=time.monotonic() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS ) with self.lock: - if not isinstance(self.by_tokenizer.get(key), _HuggingFaceTokenizerLoaded): - self._remember(key, failure) + known: Final = self.by_tokenizer.get(key) + if isinstance(known, _HuggingFaceTokenizerLoaded): + return known + self._remember(key, failure) return failure + def end_retry(self, key: _HuggingFaceTokenizerKey) -> None: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if isinstance(known, _HuggingFaceTokenizerLoadFailure) and known.retrying: + self._remember(key, replace(known, retrying=False)) + def flush_cache(self) -> None: with self.lock: self.by_tokenizer = MappingProxyType({}) @@ -2744,7 +2753,7 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _Hugg if isinstance(known, _HuggingFaceTokenizerLoadFailure): return known tokenizer: Final = ( - _load_huggingface_tokenizer(kind, backend) + known.tokenizer if isinstance(known, _HuggingFaceTokenizerLoaded) else _load_huggingface_tokenizer_remembering_outcome(kind, backend) ) @@ -2754,9 +2763,7 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _Hugg return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} -@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.""" +def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> HuggingFace: match kind: case "cohere": return tokenizer_dispatch.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer") @@ -2779,16 +2786,21 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio def _load_huggingface_tokenizer_remembering_outcome( kind: HuggingFaceTokenizerKind, backend: Decision ) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: + key: Final = (kind, backend) try: - tokenizer: Final = _load_huggingface_tokenizer(kind, backend) + tokenizer: Final = _load_huggingface_tokenizer(kind) except Exception as error: from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - return _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.remember_failure((kind, backend), str(error)) - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.remember_loaded((kind, backend)) - return tokenizer + outcome: Final = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.remember_failure(key, str(error)) + return outcome.tokenizer if isinstance(outcome, _HuggingFaceTokenizerLoaded) else outcome + else: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.remember_loaded(key, tokenizer) + return tokenizer + finally: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.end_retry(key) def encode(model="", text="", custom_tokenizer: dict | None = None): diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 89150fe0c04..a9fe2b9e9b9 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -978,14 +978,12 @@ import unittest from litellm.utils import ( _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, - _load_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding, ) # Clear the cache at module load to ensure clean state -_load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @@ -997,7 +995,6 @@ class TestTokenizerSelection(unittest.TestCase): @lru_cache, which can cause cache hits from previous tests when running with --dist=loadscope (tests from same file run on same worker). """ - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @patch("litellm.utils.tokenizer_dispatch.from_pretrained") @@ -1086,7 +1083,6 @@ class TestTokenizerSelection(unittest.TestCase): def test_failed_huggingface_tokenizer_download_is_not_retried_on_every_count( self, mock_from_pretrained: MagicMock ): - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: with patch("litellm.utils.time.monotonic", return_value=0.0): @@ -1104,7 +1100,6 @@ class TestTokenizerSelection(unittest.TestCase): assert first_count == expected assert second_count == expected finally: - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @patch("litellm.utils.time.monotonic", return_value=299.0) @@ -1112,7 +1107,6 @@ class TestTokenizerSelection(unittest.TestCase): def test_failed_huggingface_tokenizer_download_is_retried_after_retry_window( self, mock_from_pretrained: MagicMock, mock_monotonic: MagicMock ): - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: def from_pretrained_side_effect(_: str) -> object: @@ -1142,7 +1136,6 @@ class TestTokenizerSelection(unittest.TestCase): assert retried_result["tokenizer"] is sentinel.tokenizer assert mock_from_pretrained.call_count == 2 finally: - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @@ -1211,7 +1204,6 @@ def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch: @pytest.mark.parametrize("model", ["llama-2-7b", "meta-llama/llama-3-8b-instruct"]) def test_offline_mode_counts_downloaded_tokenizer_models_with_tiktoken(monkeypatch: pytest.MonkeyPatch, model: str): - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") @@ -1231,7 +1223,6 @@ def test_offline_mode_keeps_deepseek_v4_on_its_bundled_tokenizer(monkeypatch: py def test_concurrent_first_deepseek_counts_share_one_loaded_tokenizer(): - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() start: Final = threading.Barrier(8) @@ -1246,7 +1237,6 @@ 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_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" try: @@ -1259,12 +1249,10 @@ def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplo assert types == ("openai_tokenizer",) * 3 assert len(fallback_warnings) == 1 finally: - _load_huggingface_tokenizer.cache_clear() _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() +def test_download_that_fails_after_another_count_loaded_the_tokenizer_counts_with_it(): _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" slow_download_started: Final = threading.Event() @@ -1284,18 +1272,33 @@ def test_download_that_fails_after_another_count_loaded_the_tokenizer_does_not_h 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"] + slow_count_tokenizer: Final = slow_count.result(timeout=5)["tokenizer"] later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3)) - assert slow_count_type == "openai_tokenizer" + assert slow_count_tokenizer is sentinel.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_loaded_tokenizer_serves_later_counts_without_downloading_it_again(): + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + try: + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", return_value=sentinel.tokenizer): + loaded: Final = _select_tokenizer_helper(model) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF rate limited") + ) as from_pretrained: + later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3)) + assert loaded["tokenizer"] is sentinel.tokenizer + assert tuple(count["tokenizer"] for count in later_counts) == (sentinel.tokenizer,) * 3 + assert from_pretrained.call_count == 0 + finally: _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() @@ -1327,12 +1330,70 @@ def test_counts_during_a_retry_after_the_window_use_tiktoken_instead_of_download assert downloads == 2 assert retried["tokenizer"] is sentinel.tokenizer finally: - _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def test_retry_slower_than_the_window_stays_the_only_download_until_it_ends(): + _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, + ): + _select_tokenizer_helper(model) + 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) + clock.return_value = 5 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + types_while_retry_runs: Final = tuple(_select_tokenizer_helper(model)["type"] for _ in range(3)) + downloads_while_retry_runs: Final = from_pretrained.call_count + other_counts_finished.set() + retried: Final = retry.result(timeout=5) + assert types_while_retry_runs == ("openai_tokenizer",) * 3 + assert downloads_while_retry_runs == 2 + assert retried["tokenizer"] is sentinel.tokenizer + finally: + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +def test_fork_guard_error_during_a_retry_leaves_the_next_count_free_to_retry(): + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + try: + with ( + patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained, + patch("litellm.utils.time.monotonic", return_value=0.0) as clock, + ): + _select_tokenizer_helper(model) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1.0 + from_pretrained.side_effect = ForkedAfterNativeRuntimeStarted("native runtime started before fork") + with pytest.raises(ForkedAfterNativeRuntimeStarted): + _select_tokenizer_helper(model) + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + retried: Final = _select_tokenizer_helper(model) + assert retried["tokenizer"] is sentinel.tokenizer + finally: _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() def test_counts_after_an_idle_stretch_of_several_windows_still_share_one_retry_download(): - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" count_during_retry: Final = Future[SelectTokenizerResponse]() @@ -1358,14 +1419,12 @@ def test_counts_after_an_idle_stretch_of_several_windows_still_share_one_retry_d 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_OUTCOMES.flush_cache() model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" forked: Final = ForkedAfterNativeRuntimeStarted("native runtime started before fork") @@ -1378,7 +1437,6 @@ def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered assert loaded["type"] == "huggingface_tokenizer" assert loaded["tokenizer"] is sentinel.tokenizer finally: - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 0ab04500493..343b67699d3 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -403,12 +403,10 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp from litellm.utils import ( _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, - _load_huggingface_tokenizer, _select_tokenizer_helper, ) monkeypatch.setenv("LITELLM_RUST", "false") - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() try: with caplog.at_level(log_level, logger="LiteLLM"), patch.dict(sys.modules, {"tokenizers": None}): @@ -418,7 +416,6 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp assert ("token counts may be approximate" in caplog.text) is (log_level == "WARNING") assert ("install tokenizers" in caplog.text) is (log_level == "WARNING") finally: - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @@ -478,12 +475,10 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): from litellm.utils import ( _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, - _load_huggingface_tokenizer, _select_tokenizer_helper, ) monkeypatch.setenv("LITELLM_RUST", "false") - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() model = "llama-2\r\nforged-model\x1b[31m\u2028\u2029" secret = "sk-" + "x" * 48 @@ -511,7 +506,6 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): assert "REDACTED" in message assert "install tokenizers and huggingface-hub" in message finally: - _load_huggingface_tokenizer.cache_clear() _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() 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 57032e5bacc..35e66c9742b 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_OUTCOMES, _load_huggingface_tokenizer +from litellm.utils import _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES HUGGINGFACE_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct" HUGGINGFACE_ROUTER_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct-Turbo" @@ -84,12 +84,10 @@ def isolate_huggingface_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator monkeypatch.delenv("HF_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_BASE", raising=False) monkeypatch.delenv("HUGGINGFACE_API_KEY", raising=False) - _load_huggingface_tokenizer.cache_clear() _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_OUTCOMES.flush_cache()