From d037d86a943828a82b975348bc7fa5b870d70f26 Mon Sep 17 00:00:00 2001 From: quietjules <349048734+quietjules@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:50:58 -0700 Subject: [PATCH] fix(token_counter)!: keep a loaded Hugging Face tokenizer when a slower download fails (#48916) * 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 * fix(utils): keep Hugging Face tokenizer outcomes until their retry time instead of a TTL The outcome store no longer expires entries after twice the retry window, so after a long idle stretch the first count still claims the one retry and the other counts keep using tiktoken instead of each downloading again. The store is a small lock-guarded holder of an immutable mapping keyed by tokenizer kind and backend * fix(token_counter): keep the loaded tokenizer in its outcome and run one retry at a time The loaded outcome now holds the tokenizer, so an LRU eviction can no longer send every count back to an unguarded download. A retry blocks further retries until it ends, even when it runs past the window, and a download that fails after another count loaded the tokenizer counts with that tokenizer * fix(token_counter): retry a failed Hugging Face tokenizer download in the background After the retry window, the count that finds the stored failure starts one background download and answers with tiktoken, so no live count waits on a retry. A failure from a download that started before the first failure no longer resets a running retry, and a failed retry never replaces a tokenizer another download loaded. The outcome store takes its clock and retry starter as dependencies, so the tests inject them instead of patching time.monotonic. * fix(token_counter): keep retrying a failed Hugging Face tokenizer download when its retry thread cannot start If starting the background retry thread raises, the count answers with tiktoken and the failure is stored with a fresh retry time, so the next window retries again instead of the tokenizer staying on tiktoken until restart. * fix(token_counter): let a forked process or a crashed retry thread retry a failed Hugging Face tokenizer download again A failed download's in-flight retry mark is now the pid of the process running it, so a process forked during a retry starts its own instead of waiting on a thread it does not have. The retry thread clears its mark in a finally, so a native panic or fork guard error escaping it leaves the next window free to retry. The fork guard error is no longer swallowed in the retry thread and reaches the thread's exception hook, as on the first load * fix(token_counter): let a process that exits during a background tokenizer retry exit cleanly instead of aborting --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 2 +- litellm/utils.py | 177 +++++-- .../litellm_core_utils/test_token_counter.py | 455 +++++++++++++++--- .../unit/litellm_core_utils/test_tokenizer.py | 18 +- .../test_huggingface_chat_transformation.py | 8 +- 5 files changed, 524 insertions(+), 136 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 42ad624026c..7f2d98b7a28 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -158,7 +158,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 26ce38d63b3..aeb25f81841 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -29,8 +29,8 @@ import threading import time import traceback from contextlib import nullcontext -from dataclasses import dataclass, field -from functools import lru_cache, wraps +from dataclasses import dataclass, field, replace +from functools import lru_cache, partial, wraps from importlib import resources from inspect import iscoroutine from io import StringIO @@ -2683,9 +2683,99 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None: return None +@dataclass(frozen=True, slots=True) +class _HuggingFaceTokenizerLoaded: + tokenizer: HuggingFace + + @dataclass(frozen=True, slots=True) class _HuggingFaceTokenizerLoadFailure: reason: str + retry_at: float + retry_pid: int | None = None + + +_HuggingFaceTokenizerKey: TypeAlias = tuple[HuggingFaceTokenizerKind, Decision] +_HuggingFaceTokenizerLoadOutcome: TypeAlias = _HuggingFaceTokenizerLoaded | _HuggingFaceTokenizerLoadFailure + + +def _start_in_thread_joined_at_exit(work: Callable[[], None]) -> None: + threading.Thread(target=work, name="litellm-huggingface-tokenizer-retry", daemon=False).start() + + +@dataclass(slots=True) +class _HuggingFaceTokenizerLoadOutcomes: + clock: Callable[[], float] = time.monotonic + start_retry: Callable[[Callable[[], None]], None] = _start_in_thread_joined_at_exit + pid: Callable[[], int] = os.getpid + lock: threading.Lock = field(default_factory=threading.Lock) + by_tokenizer: Mapping[_HuggingFaceTokenizerKey, _HuggingFaceTokenizerLoadOutcome] = field( + default_factory=lambda: MappingProxyType({}) + ) + + def known( + self, key: _HuggingFaceTokenizerKey, retry_when_due: Callable[[], None] + ) -> _HuggingFaceTokenizerLoadOutcome | None: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if ( + not isinstance(known, _HuggingFaceTokenizerLoadFailure) + or known.retry_pid == self.pid() + or self.clock() < known.retry_at + ): + return known + retrying: Final = replace(known, retry_pid=self.pid()) + self._remember(key, retrying) + try: + self.start_retry(retry_when_due) + except RuntimeError as error: + return self.remember_failed_retry(key, str(error)) + return retrying + + def remember_loaded(self, key: _HuggingFaceTokenizerKey, tokenizer: HuggingFace) -> _HuggingFaceTokenizerLoaded: + loaded: Final = _HuggingFaceTokenizerLoaded(tokenizer) + with self.lock: + self._remember(key, loaded) + return loaded + + def remember_first_failure(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadOutcome: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if known is not None: + return known + failure: Final = self._failure(reason) + self._remember(key, failure) + return failure + + def remember_failed_retry(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadOutcome: + with self.lock: + known: Final = self.by_tokenizer.get(key) + if isinstance(known, _HuggingFaceTokenizerLoaded): + return known + failure: Final = self._failure(reason) + 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.retry_pid is not None: + self._remember(key, self._failure(known.reason)) + + def flush_cache(self) -> None: + with self.lock: + self.by_tokenizer = MappingProxyType({}) + + def _failure(self, reason: str) -> _HuggingFaceTokenizerLoadFailure: + return _HuggingFaceTokenizerLoadFailure( + reason=reason, retry_at=self.clock() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + ) + + def _remember(self, key: _HuggingFaceTokenizerKey, outcome: _HuggingFaceTokenizerLoadOutcome) -> None: + self.by_tokenizer = MappingProxyType({**self.by_tokenizer, key: outcome}) + + +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES: Final = _HuggingFaceTokenizerLoadOutcomes() def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: @@ -2697,30 +2787,28 @@ def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None: ) -def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: +def _return_huggingface_tokenizer( + model: str, outcomes: _HuggingFaceTokenizerLoadOutcomes = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES +) -> SelectTokenizerResponse | _HuggingFaceTokenizerLoadFailure | None: kind: Final = huggingface_tokenizer_kind(model) if kind is None or (kind not in _BUNDLED_TOKENIZER_LOAD_LOCKS and is_offline_mode()): return None - backend: Final = _huggingface_tokenizer_backend() - key: Final = _huggingface_tokenizer_failure_key(kind, backend) + key: Final = (kind, _huggingface_tokenizer_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) - if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): - _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) - return tokenizer - return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + known: Final = outcomes.known(key, partial(_retry_huggingface_tokenizer_load, model, key, outcomes)) + outcome: Final = ( + known + if known is not None + else _load_huggingface_tokenizer_remembering_outcome(model, key, outcomes, outcomes.remember_first_failure) + ) + match outcome: + case _HuggingFaceTokenizerLoaded(tokenizer=tokenizer): + return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + case _HuggingFaceTokenizerLoadFailure(): + return outcome -def _huggingface_tokenizer_failure_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: - return f"{kind}:{backend.name}" - - -@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") @@ -2740,19 +2828,34 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio assert_never(kind) -def _load_huggingface_tokenizer_recording_failure( - kind: HuggingFaceTokenizerKind, backend: Decision, key: str -) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: +def _load_huggingface_tokenizer_remembering_outcome( + model: str, + key: _HuggingFaceTokenizerKey, + outcomes: _HuggingFaceTokenizerLoadOutcomes, + remember_failure: Callable[[_HuggingFaceTokenizerKey, str], _HuggingFaceTokenizerLoadOutcome], +) -> _HuggingFaceTokenizerLoadOutcome: + kind, _ = key try: - return _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 - failure: Final = _HuggingFaceTokenizerLoadFailure(reason=str(error)) - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.set_cache(key, failure, ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) - return failure + outcome: Final = remember_failure(key, str(error)) + if isinstance(outcome, _HuggingFaceTokenizerLoadFailure): + _warn_huggingface_tokenizer_fallback(model, outcome.reason) + return outcome + return outcomes.remember_loaded(key, tokenizer) + + +def _retry_huggingface_tokenizer_load( + model: str, key: _HuggingFaceTokenizerKey, outcomes: _HuggingFaceTokenizerLoadOutcomes +) -> None: + try: + _load_huggingface_tokenizer_remembering_outcome(model, key, outcomes, outcomes.remember_failed_retry) + finally: + outcomes.end_retry(key) def encode(model="", text="", custom_tokenizer: dict | None = None): @@ -8346,26 +8449,6 @@ def trim_messages( from litellm.caching.in_memory_cache import InMemoryCache -class _HuggingFaceTokenizerLoadFailureCache(Protocol): - def get_cache(self, key: str) -> object | None: ... - - def set_cache(self, key: str, value: object, ttl: float | None = None) -> None: ... - - def flush_cache(self) -> None: ... - - -_HUGGINGFACE_TOKENIZER_LOAD_FAILURES: Final[_HuggingFaceTokenizerLoadFailureCache] = ( - cast( # cast-ok: cache methods are untyped - _HuggingFaceTokenizerLoadFailureCache, - InMemoryCache( - max_size_in_memory=16, - default_ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, - clock=lambda: time.monotonic(), - ), - ) -) - - class AvailableModelsCache(InMemoryCache): def __init__(self, ttl_seconds: int = 300, max_size: int = 1000): super().__init__(ttl_seconds, max_size) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 000e3993bc2..d97970753f8 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -14,6 +14,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import Future, ThreadPoolExecutor, wait from itertools import accumulate, chain from pathlib import Path +from queue import SimpleQueue from typing import Final, cast from unittest.mock import MagicMock, sentinel @@ -45,6 +46,7 @@ from litellm.litellm_core_utils.token_counter import ( from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.utils import SelectTokenizerResponse from tests.large_text import text from tests.unit.litellm_core_utils.event_loop_lag import ( assert_loop_stayed_free, @@ -978,16 +980,17 @@ def test_token_counter(): import unittest from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, - _load_huggingface_tokenizer, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, + _HuggingFaceTokenizerLoadFailure, + _HuggingFaceTokenizerLoadOutcomes, + _return_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_FAILURES.flush_cache() +_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() class TestTokenizerSelection(unittest.TestCase): @@ -998,8 +1001,7 @@ 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_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): @@ -1087,64 +1089,23 @@ 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_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") - model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" - text: Final = "Count the tokens for this message." - expected: Final = len(encoding.encode(text, disallowed_special=())) - - first_count: Final = litellm.token_counter(model=model, text=text) - first_attempts: Final = mock_from_pretrained.call_count - second_count: Final = litellm.token_counter(model=model, text=text) - - assert first_attempts > 0 - assert mock_from_pretrained.call_count == first_attempts - assert first_count == expected - assert second_count == expected - finally: - _load_huggingface_tokenizer.cache_clear() - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache() - - @patch("litellm.utils.time.monotonic", return_value=299.0) - @patch("litellm.utils.tokenizer_dispatch.from_pretrained") - 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_FAILURES.flush_cache() - try: - def from_pretrained_side_effect(_: str) -> object: - if mock_from_pretrained.call_count == 1: - mock_monotonic.return_value = 322.0 - raise RuntimeError("HF unavailable") - return sentinel.tokenizer - - mock_from_pretrained.side_effect = from_pretrained_side_effect + mock_from_pretrained.side_effect = RuntimeError("HF unavailable") model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + text: Final = "Count the tokens for this message." + expected: Final = len(encoding.encode(text, disallowed_special=())) - first_result: Final = _select_tokenizer_helper(model) - assert first_result["type"] == "openai_tokenizer" - mock_monotonic.return_value = 323.0 - second_result: Final = _select_tokenizer_helper(model) - assert second_result["type"] == "openai_tokenizer", "tokenizer was retried at t=323" - assert mock_from_pretrained.call_count == 1 + first_count: Final = litellm.token_counter(model=model, text=text) + first_attempts: Final = mock_from_pretrained.call_count + second_count: Final = litellm.token_counter(model=model, text=text) - mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - 1 - third_result: Final = _select_tokenizer_helper(model) - assert third_result["type"] == "openai_tokenizer" - assert mock_from_pretrained.call_count == 1 - - mock_monotonic.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1 - retried_result: Final = _select_tokenizer_helper(model) - assert retried_result["type"] == "huggingface_tokenizer" - assert retried_result["tokenizer"] is sentinel.tokenizer - assert mock_from_pretrained.call_count == 2 + assert first_attempts > 0 + assert mock_from_pretrained.call_count == first_attempts + assert first_count == expected + assert second_count == expected 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 @@ -1212,8 +1173,7 @@ 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_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() monkeypatch.setenv("LITELLM_OFFLINE_MODE", "true") with patch("litellm.utils.tokenizer_dispatch.from_pretrained") as from_pretrained: @@ -1232,8 +1192,7 @@ 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_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() start: Final = threading.Barrier(8) def first_tokenizer(_: int) -> object: @@ -1247,8 +1206,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 ( @@ -1260,15 +1218,346 @@ 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_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + +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() + 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_tokenizer: Final = slow_count.result(timeout=5)["tokenizer"] + later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3)) + 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: + _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 _run_now(retry: Callable[[], None]) -> None: + retry() + + +def _start_retry_resolving(ended: Future[None]) -> Callable[[Callable[[], None]], None]: + def start(retry: Callable[[], None]) -> None: + def run_then_resolve() -> None: + retry() + ended.set_result(None) + + threading.Thread(target=run_then_resolve, daemon=True).start() + + return start + + +def _eventually(probe: Callable[[], bool]) -> bool: + deadline: Final = time.monotonic() + 5 + while time.monotonic() < deadline: + if probe(): + return True + time.sleep(0.01) + return False + + +def test_failed_download_is_retried_once_per_window_until_it_loads(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=299.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=_run_now) + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + clock.return_value = 322.0 + raise RuntimeError("HF unavailable") + if from_pretrained.call_count == 2: + raise RuntimeError("HF rate limited") + return sentinel.tokenizer + + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained: + first: Final = _return_huggingface_tokenizer(model, outcomes) + clock.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - 1 + inside_the_window: Final = _return_huggingface_tokenizer(model, outcomes) + downloads_inside_the_window: Final = from_pretrained.call_count + clock.return_value = 322.0 + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + claimant: Final = _return_huggingface_tokenizer(model, outcomes) + inside_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + downloads_inside_the_next_window: Final = from_pretrained.call_count + clock.return_value = 322.0 + 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + after_the_second_retry: Final = _return_huggingface_tokenizer(model, outcomes) + assert isinstance(first, _HuggingFaceTokenizerLoadFailure) + assert isinstance(inside_the_window, _HuggingFaceTokenizerLoadFailure) + assert downloads_inside_the_window == 1 + assert isinstance(claimant, _HuggingFaceTokenizerLoadFailure) + assert isinstance(inside_the_next_window, _HuggingFaceTokenizerLoadFailure) + assert downloads_inside_the_next_window == 2 + assert from_pretrained.call_count == 3 + assert after_the_second_retry == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + + +def test_retry_after_the_window_runs_in_the_background_while_every_count_answers_at_once(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock) + retry_started: Final = threading.Event() + retry_released: Final = threading.Event() + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + raise RuntimeError("HF unavailable") + retry_started.set() + retry_released.wait(timeout=5) + return sentinel.tokenizer + + with patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = 3 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + counts_while_retry_runs: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(4)) + assert retry_started.wait(timeout=5) + clock.return_value = 9 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + later_counts_while_retry_runs: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(2)) + downloads_while_retry_runs: Final = from_pretrained.call_count + retry_released.set() + retried: Final = _eventually( + lambda: _return_huggingface_tokenizer(model, outcomes) + == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + ) + assert all( + isinstance(count, _HuggingFaceTokenizerLoadFailure) + for count in counts_while_retry_runs + later_counts_while_retry_runs + ) + assert downloads_while_retry_runs == 2 + assert retried + assert from_pretrained.call_count == 2 + + +def test_download_from_before_the_first_failure_ending_during_a_retry_leaves_it_the_only_retry(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock) + early_download_started: Final = threading.Event() + early_download_released: Final = threading.Event() + retry_started: Final = threading.Event() + retry_released: Final = threading.Event() + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + early_download_started.set() + early_download_released.wait(timeout=5) + raise RuntimeError("HF unavailable") + if from_pretrained.call_count == 2: + raise RuntimeError("HF unavailable") + retry_started.set() + retry_released.wait(timeout=5) + return sentinel.tokenizer + + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, + ThreadPoolExecutor(max_workers=1) as early_pool, + ): + early_count: Final = early_pool.submit(_return_huggingface_tokenizer, model, outcomes) + assert early_download_started.wait(timeout=5) + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1.0 + _return_huggingface_tokenizer(model, outcomes) + assert retry_started.wait(timeout=5) + early_download_released.set() + early: Final = early_count.result(timeout=5) + clock.return_value = 5 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + counts_while_retry_runs: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(3)) + downloads_while_retry_runs: Final = from_pretrained.call_count + retry_released.set() + retried: Final = _eventually( + lambda: _return_huggingface_tokenizer(model, outcomes) + == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + ) + assert isinstance(early, _HuggingFaceTokenizerLoadFailure) + assert all(isinstance(count, _HuggingFaceTokenizerLoadFailure) for count in counts_while_retry_runs) + assert downloads_while_retry_runs == 3 + assert retried + + +def test_retry_failing_after_an_earlier_download_loaded_the_tokenizer_keeps_it(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + retry_ended: Final = Future[None]() + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=_start_retry_resolving(retry_ended)) + early_download_started: Final = threading.Event() + early_download_released: Final = threading.Event() + retry_started: Final = threading.Event() + retry_released: Final = threading.Event() + + def download(_: str) -> object: + if from_pretrained.call_count == 1: + early_download_started.set() + early_download_released.wait(timeout=5) + return sentinel.tokenizer + if from_pretrained.call_count == 2: + raise RuntimeError("HF unavailable") + retry_started.set() + retry_released.wait(timeout=5) + raise RuntimeError("HF rate limited") + + with ( + patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, + ThreadPoolExecutor(max_workers=1) as early_pool, + ): + early_count: Final = early_pool.submit(_return_huggingface_tokenizer, model, outcomes) + assert early_download_started.wait(timeout=5) + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + assert retry_started.wait(timeout=5) + early_download_released.set() + early: Final = early_count.result(timeout=5) + retry_released.set() + retry_ended.result(timeout=5) + later_counts: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(3)) + loaded: Final = {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + assert early == loaded + assert from_pretrained.call_count == 3 + assert later_counts == (loaded,) * 3 + + +class _NativePanic(BaseException): + pass + + +def _fork_guard_error() -> BaseException: + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + return ForkedAfterNativeRuntimeStarted("native runtime started before fork") + + +def _native_panic() -> BaseException: + return _NativePanic("tokenizer parser panicked") + + +@pytest.mark.parametrize("escaping_error", (_fork_guard_error, _native_panic), ids=("fork_guard", "native_panic")) +def test_retry_that_raises_past_failure_handling_still_retries_after_the_next_window( + escaping_error: Callable[[], BaseException], +): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + error: Final = escaping_error() + retries: Final = SimpleQueue[Future[None]]() + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="tokenizer-retry") as pool: + + def start_retry(retry: Callable[[], None]) -> None: + retries.put(pool.submit(retry)) + + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=start_retry) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + from_pretrained.side_effect = error + _return_huggingface_tokenizer(model, outcomes) + raised: Final = retries.get(timeout=5).exception(timeout=5) + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + inside_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + no_retry_inside_the_next_window: Final = retries.empty() + clock.return_value = 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + retries.get(timeout=5).result(timeout=5) + after_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + assert raised is error + assert isinstance(inside_the_next_window, _HuggingFaceTokenizerLoadFailure) + assert no_retry_inside_the_next_window + assert from_pretrained.call_count == 3 + assert after_the_next_window == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + + +def test_process_forked_during_a_retry_starts_its_own_retry_instead_of_waiting_on_the_parents(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + pid: Final = MagicMock(return_value=100) + start_retry: Final = MagicMock() + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=start_retry, pid=pid) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + _return_huggingface_tokenizer(model, outcomes) + retries_started_by_the_parent: Final = start_retry.call_count + pid.return_value = 101 + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + _return_huggingface_tokenizer(model, outcomes) + start_retry.call_args.args[0]() + in_the_forked_process: Final = _return_huggingface_tokenizer(model, outcomes) + assert retries_started_by_the_parent == 1 + assert start_retry.call_count == 2 + assert in_the_forked_process == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + + +def test_retry_thread_that_cannot_start_answers_the_count_and_retries_after_the_next_window(): + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + clock: Final = MagicMock(return_value=0.0) + thread_starts: Final = MagicMock(side_effect=(RuntimeError("can't start new thread"), None)) + + def start_thread_then_retry(retry: Callable[[], None]) -> None: + thread_starts() + retry() + + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=start_thread_then_retry) + with patch( + "litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=RuntimeError("HF unavailable") + ) as from_pretrained: + _return_huggingface_tokenizer(model, outcomes) + clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + when_the_thread_cannot_start: Final = _return_huggingface_tokenizer(model, outcomes) + from_pretrained.side_effect = None + from_pretrained.return_value = sentinel.tokenizer + inside_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + clock.return_value = 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + _return_huggingface_tokenizer(model, outcomes) + after_the_next_window: Final = _return_huggingface_tokenizer(model, outcomes) + assert isinstance(when_the_thread_cannot_start, _HuggingFaceTokenizerLoadFailure) + assert isinstance(inside_the_next_window, _HuggingFaceTokenizerLoadFailure) + assert from_pretrained.call_count == 2 + assert after_the_next_window == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} 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: @@ -1280,8 +1569,7 @@ 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_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() def test_token_counter_with_anthropic_tool_use(): @@ -2325,6 +2613,27 @@ tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text) print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))})) """ +HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +import threading +import time +from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer +retry_reached_hub = threading.Event() +def answer_once_the_process_is_exiting(request): + retry_reached_hub.set() + while threading.main_thread().is_alive(): + time.sleep(0.01) + time.sleep(0.5) + return handle(request) +now = [0.0] +outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0]) +model = "meta-llama/llama-3-70b-instruct" +_return_huggingface_tokenizer(model, outcomes) +huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting))) +now[0] = 3600.0 +_return_huggingface_tokenizer(model, outcomes) +print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)})) +""" + def _word_level_tokenizer_json( pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None @@ -2387,6 +2696,10 @@ def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_t } +def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None: + assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True} + + def test_empty_custom_tokenizer_uses_model_tokenizer() -> None: text_value: Final = "A tokenizer fallback should preserve the model encoding." custom_count: Final = _get_exact_count_function("gpt-3.5-turbo", {})(text_value) diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 54185c08c10..343b67699d3 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -402,14 +402,12 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp from unittest.mock import patch from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, - _load_huggingface_tokenizer, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _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") @@ -418,8 +416,7 @@ 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_FAILURES.flush_cache() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @pytest.mark.parametrize("python_installed", [False, True]) @@ -477,14 +474,12 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch): from unittest.mock import patch from litellm.utils import ( - _HUGGINGFACE_TOKENIZER_LOAD_FAILURES, - _load_huggingface_tokenizer, + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, _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) @@ -511,8 +506,7 @@ 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_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..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_FAILURES, _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,13 +84,11 @@ 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_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 = {