From 6a7d32d74c2920c63a1dd09e303825188c2f4fe8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:02:10 -0700 Subject: [PATCH] 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 --- litellm/utils.py | 19 +++-- .../litellm_core_utils/test_token_counter.py | 72 ++++++++++++++++--- 2 files changed, 74 insertions(+), 17 deletions(-) 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():