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