diff --git a/litellm/utils.py b/litellm/utils.py index f67f16edaab..16b51b9b711 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2677,7 +2677,7 @@ class _HuggingFaceTokenizerLoaded: class _HuggingFaceTokenizerLoadFailure: reason: str retry_at: float - retrying: bool = False + retry_pid: int | None = None _HuggingFaceTokenizerKey: TypeAlias = tuple[HuggingFaceTokenizerKind, Decision] @@ -2692,6 +2692,7 @@ def _start_in_daemon_thread(work: Callable[[], None]) -> None: class _HuggingFaceTokenizerLoadOutcomes: clock: Callable[[], float] = time.monotonic start_retry: Callable[[Callable[[], None]], None] = _start_in_daemon_thread + pid: Callable[[], int] = os.getpid lock: threading.Lock = field(default_factory=threading.Lock) by_tokenizer: Mapping[_HuggingFaceTokenizerKey, _HuggingFaceTokenizerLoadOutcome] = field( default_factory=lambda: MappingProxyType({}) @@ -2704,11 +2705,11 @@ class _HuggingFaceTokenizerLoadOutcomes: known: Final = self.by_tokenizer.get(key) if ( not isinstance(known, _HuggingFaceTokenizerLoadFailure) - or known.retrying + or known.retry_pid == self.pid() or self.clock() < known.retry_at ): return known - retrying: Final = replace(known, retrying=True) + retrying: Final = replace(known, retry_pid=self.pid()) self._remember(key, retrying) try: self.start_retry(retry_when_due) @@ -2740,6 +2741,12 @@ class _HuggingFaceTokenizerLoadOutcomes: 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({}) @@ -2830,12 +2837,10 @@ def _load_huggingface_tokenizer_remembering_outcome( def _retry_huggingface_tokenizer_load( model: str, key: _HuggingFaceTokenizerKey, outcomes: _HuggingFaceTokenizerLoadOutcomes ) -> None: - from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking - try: _load_huggingface_tokenizer_remembering_outcome(model, key, outcomes, outcomes.remember_failed_retry) - except (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking) as error: - outcomes.remember_failed_retry(key, str(error)) + finally: + 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 fa844a02f6d..59f9188264a 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 @@ -1447,29 +1448,80 @@ def test_retry_failing_after_an_earlier_download_loaded_the_tokenizer_keeps_it() assert later_counts == (loaded,) * 3 -def test_fork_guard_error_in_a_retry_waits_out_the_next_window_like_any_failed_retry(): +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) - outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=_run_now) + 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 - from_pretrained.side_effect = ForkedAfterNativeRuntimeStarted("native runtime started before fork") _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 - inside_the_next_window: Final = tuple(_return_huggingface_tokenizer(model, outcomes) for _ in range(3)) - downloads_inside_the_next_window: Final = from_pretrained.call_count - 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 all(isinstance(count, _HuggingFaceTokenizerLoadFailure) for count in inside_the_next_window) - assert downloads_inside_the_next_window == 2 - assert after_the_next_window == {"type": "huggingface_tokenizer", "tokenizer": sentinel.tokenizer} + 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():