diff --git a/litellm/utils.py b/litellm/utils.py index 69d70d434ee..595bb615c6a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -30,7 +30,7 @@ import time import traceback from contextlib import nullcontext from dataclasses import dataclass, field, replace -from functools import lru_cache, wraps +from functools import lru_cache, partial, wraps from importlib import resources from inspect import iscoroutine from io import StringIO @@ -2684,49 +2684,68 @@ _HuggingFaceTokenizerKey: TypeAlias = tuple[HuggingFaceTokenizerKind, Decision] _HuggingFaceTokenizerLoadOutcome: TypeAlias = _HuggingFaceTokenizerLoaded | _HuggingFaceTokenizerLoadFailure +def _start_in_daemon_thread(work: Callable[[], None]) -> None: + threading.Thread(target=work, name="litellm-huggingface-tokenizer-retry", daemon=True).start() + + @dataclass(slots=True) class _HuggingFaceTokenizerLoadOutcomes: + clock: Callable[[], float] = time.monotonic + start_retry: Callable[[Callable[[], None]], None] = _start_in_daemon_thread lock: threading.Lock = field(default_factory=threading.Lock) by_tokenizer: Mapping[_HuggingFaceTokenizerKey, _HuggingFaceTokenizerLoadOutcome] = field( default_factory=lambda: MappingProxyType({}) ) - def claim(self, key: _HuggingFaceTokenizerKey) -> _HuggingFaceTokenizerLoadOutcome | None: + def known( + self, key: _HuggingFaceTokenizerKey, retry_when_due: Callable[[], None] + ) -> _HuggingFaceTokenizerLoadOutcome | None: with self.lock: known: Final = self.by_tokenizer.get(key) - now: Final = time.monotonic() - match known: - case _HuggingFaceTokenizerLoadFailure(retry_at=retry_at, retrying=False) if now >= retry_at: - self._remember(key, replace(known, retrying=True)) - return None - case _: - return known + if ( + not isinstance(known, _HuggingFaceTokenizerLoadFailure) + or known.retrying + or self.clock() < known.retry_at + ): + return known + retrying: Final = replace(known, retrying=True) + self._remember(key, retrying) + self.start_retry(retry_when_due) + return retrying - def remember_loaded(self, key: _HuggingFaceTokenizerKey, tokenizer: HuggingFace) -> None: + def remember_loaded(self, key: _HuggingFaceTokenizerKey, tokenizer: HuggingFace) -> _HuggingFaceTokenizerLoaded: + loaded: Final = _HuggingFaceTokenizerLoaded(tokenizer) with self.lock: - self._remember(key, _HuggingFaceTokenizerLoaded(tokenizer)) + self._remember(key, loaded) + return loaded - def remember_failure(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadOutcome: - failure: Final = _HuggingFaceTokenizerLoadFailure( - reason=reason, retry_at=time.monotonic() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - ) + 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.retrying: - self._remember(key, replace(known, retrying=False)) + return failure 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}) @@ -2743,24 +2762,25 @@ 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 = (kind, _huggingface_tokenizer_backend()) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): - known: Final = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.claim((kind, backend)) - if isinstance(known, _HuggingFaceTokenizerLoadFailure): - return known - tokenizer: Final = ( - known.tokenizer - if isinstance(known, _HuggingFaceTokenizerLoaded) - else _load_huggingface_tokenizer_remembering_outcome(kind, backend) + 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) ) - if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): - _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) - return tokenizer - return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + match outcome: + case _HuggingFaceTokenizerLoaded(tokenizer=tokenizer): + return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + case _HuggingFaceTokenizerLoadFailure(): + return outcome def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> HuggingFace: @@ -2784,9 +2804,12 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> HuggingFace: def _load_huggingface_tokenizer_remembering_outcome( - kind: HuggingFaceTokenizerKind, backend: Decision -) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: - key: Final = (kind, backend) + model: str, + key: _HuggingFaceTokenizerKey, + outcomes: _HuggingFaceTokenizerLoadOutcomes, + remember_failure: Callable[[_HuggingFaceTokenizerKey, str], _HuggingFaceTokenizerLoadOutcome], +) -> _HuggingFaceTokenizerLoadOutcome: + kind, _ = key try: tokenizer: Final = _load_huggingface_tokenizer(kind) except Exception as error: @@ -2794,13 +2817,22 @@ def _load_huggingface_tokenizer_remembering_outcome( if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - 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) + 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: + 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)) 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 a9fe2b9e9b9..3193cc39993 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -978,6 +978,9 @@ import unittest from litellm.utils import ( _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES, + _HuggingFaceTokenizerLoadFailure, + _HuggingFaceTokenizerLoadOutcomes, + _return_huggingface_tokenizer, _select_tokenizer_helper, claude_json_str, encoding, @@ -1085,56 +1088,19 @@ class TestTokenizerSelection(unittest.TestCase): ): _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: - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.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 - ): - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.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: _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() @@ -1298,128 +1264,212 @@ def test_loaded_tokenizer_serves_later_counts_without_downloading_it_again(): _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() -def test_counts_during_a_retry_after_the_window_use_tiktoken_instead_of_downloading_again(): - _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" - retry_started: Final = threading.Event() - other_counts_finished: Final = threading.Event() + clock: Final = MagicMock(return_value=299.0) + outcomes: Final = _HuggingFaceTokenizerLoadOutcomes(clock=clock, start_retry=_run_now) 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") + 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 - try: - with ( - patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, - patch("litellm.utils.time.monotonic", return_value=0.0) as clock, - ): - first_type: Final = _select_tokenizer_helper(model)["type"] - clock.return_value = HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1.0 - with ThreadPoolExecutor(max_workers=1, thread_name_prefix="retry") as pool: - retry: Final = pool.submit(_select_tokenizer_helper, model) - assert retry_started.wait(timeout=5) - types_during_retry: Final = tuple(_select_tokenizer_helper(model)["type"] for _ in range(3)) - other_counts_finished.set() - retried: Final = retry.result(timeout=5) - downloads: Final = from_pretrained.call_count - assert first_type == "openai_tokenizer" - assert types_during_retry == ("openai_tokenizer",) * 3 - assert downloads == 2 - assert retried["tokenizer"] is sentinel.tokenizer - finally: - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + 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_slower_than_the_window_stays_the_only_download_until_it_ends(): - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() +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() - 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(): - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() - model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" - count_during_retry: Final = Future[SelectTokenizerResponse]() + retry_released: Final = threading.Event() def download(_: str) -> object: if from_pretrained.call_count == 1: raise RuntimeError("HF unavailable") - if from_pretrained.call_count == 2: - count_during_retry.set_result(_select_tokenizer_helper(model)) + retry_started.set() + retry_released.wait(timeout=5) return sentinel.tokenizer - try: - with ( - patch("litellm.utils.tokenizer_dispatch.from_pretrained", side_effect=download) as from_pretrained, - patch("litellm.utils.time.monotonic", return_value=0.0) as clock, - ): - first_type: Final = _select_tokenizer_helper(model)["type"] - clock.return_value = 2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + 1.0 - retried: Final = _select_tokenizer_helper(model) - downloads: Final = from_pretrained.call_count - assert first_type == "openai_tokenizer" - assert count_during_retry.result(timeout=0)["type"] == "openai_tokenizer" - assert downloads == 2 - assert retried["tokenizer"] is sentinel.tokenizer - finally: - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + 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 + + +def test_fork_guard_error_in_a_retry_waits_out_the_next_window_like_any_failed_retry(): + from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted + + 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) + 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) + 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} def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered():