From d1cb37bc4a413c64abbdd47e1d6fb2398945fa1f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 05:29:12 -0700 Subject: [PATCH] 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 --- litellm/utils.py | 111 ++++++++---------- .../litellm_core_utils/test_token_counter.py | 34 +++++- 2 files changed, 84 insertions(+), 61 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index cf0a2bb4bf6..2675a56c6b3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2679,6 +2679,52 @@ class _HuggingFaceTokenizerLoadFailure: retry_at: float +_HuggingFaceTokenizerKey: TypeAlias = tuple[HuggingFaceTokenizerKind, Decision] +_HuggingFaceTokenizerLoadOutcome: TypeAlias = _HuggingFaceTokenizerLoaded | _HuggingFaceTokenizerLoadFailure + + +@dataclass(slots=True) +class _HuggingFaceTokenizerLoadOutcomes: + 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: + with self.lock: + known: Final = self.by_tokenizer.get(key) + now: Final = time.monotonic() + match known: + case _HuggingFaceTokenizerLoadFailure(retry_at=retry_at) if now >= retry_at: + self._remember(key, replace(known, retry_at=now + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS)) + return None + case _: + return known + + def remember_loaded(self, key: _HuggingFaceTokenizerKey) -> None: + with self.lock: + self._remember(key, _HuggingFaceTokenizerLoaded()) + + def remember_failure(self, key: _HuggingFaceTokenizerKey, reason: str) -> _HuggingFaceTokenizerLoadFailure: + failure: Final = _HuggingFaceTokenizerLoadFailure( + reason=reason, retry_at=time.monotonic() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS + ) + with self.lock: + if not isinstance(self.by_tokenizer.get(key), _HuggingFaceTokenizerLoaded): + self._remember(key, failure) + return failure + + def flush_cache(self) -> None: + with self.lock: + self.by_tokenizer = MappingProxyType({}) + + 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: verbose_logger.warning( "Falling back to tiktoken for %s; token counts may be approximate. " @@ -2693,15 +2739,14 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _Hugg 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_outcome_key(kind, backend) with _BUNDLED_TOKENIZER_LOAD_LOCKS.get(kind, nullcontext()): - known: Final = _claim_huggingface_tokenizer_load(key) + known: Final = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.claim((kind, backend)) if isinstance(known, _HuggingFaceTokenizerLoadFailure): return known tokenizer: Final = ( _load_huggingface_tokenizer(kind, backend) if isinstance(known, _HuggingFaceTokenizerLoaded) - else _load_huggingface_tokenizer_remembering_outcome(kind, backend, key) + else _load_huggingface_tokenizer_remembering_outcome(kind, backend) ) if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure): _warn_huggingface_tokenizer_fallback(model, tokenizer.reason) @@ -2709,28 +2754,6 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _Hugg return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} -def _huggingface_tokenizer_outcome_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str: - return f"{kind}:{backend.name}" - - -def _claim_huggingface_tokenizer_load( - key: str, -) -> _HuggingFaceTokenizerLoaded | _HuggingFaceTokenizerLoadFailure | None: - with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: - known: Final = _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.get_cache(key) - now: Final = time.monotonic() - match known: - case _HuggingFaceTokenizerLoadFailure(retry_at=retry_at) if now >= retry_at: - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache( - key, replace(known, retry_at=now + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS) - ) - return None - case _HuggingFaceTokenizerLoaded() | _HuggingFaceTokenizerLoadFailure(): - return known - case _: - return None - - @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.""" @@ -2754,7 +2777,7 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio def _load_huggingface_tokenizer_remembering_outcome( - kind: HuggingFaceTokenizerKind, backend: Decision, key: str + kind: HuggingFaceTokenizerKind, backend: Decision ) -> HuggingFace | _HuggingFaceTokenizerLoadFailure: try: tokenizer: Final = _load_huggingface_tokenizer(kind, backend) @@ -2763,22 +2786,11 @@ def _load_huggingface_tokenizer_remembering_outcome( if isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking)): raise - return _remember_huggingface_tokenizer_failure(key, str(error)) - with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache(key, _HuggingFaceTokenizerLoaded()) + return _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.remember_failure((kind, backend), str(error)) + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.remember_loaded((kind, backend)) return tokenizer -def _remember_huggingface_tokenizer_failure(key: str, reason: str) -> _HuggingFaceTokenizerLoadFailure: - failure: Final = _HuggingFaceTokenizerLoadFailure( - reason=reason, retry_at=time.monotonic() + HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS - ) - with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: - if not isinstance(_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.get_cache(key), _HuggingFaceTokenizerLoaded): - _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache(key, failure) - return failure - - def encode(model="", text="", custom_tokenizer: dict | None = None): """ Encodes the given text using the specified model. @@ -8350,27 +8362,6 @@ def trim_messages( from litellm.caching.in_memory_cache import InMemoryCache -class _HuggingFaceTokenizerLoadOutcomeCache(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_OUTCOMES: Final[_HuggingFaceTokenizerLoadOutcomeCache] = ( - cast( # cast-ok: cache methods are untyped - _HuggingFaceTokenizerLoadOutcomeCache, - InMemoryCache( - max_size_in_memory=16, - default_ttl=2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS, - clock=lambda: time.monotonic(), - ), - ) -) -_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: Final = threading.Lock() - - 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 96c4bb19c4d..28b595787b0 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -43,6 +43,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, @@ -1210,7 +1211,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: @@ -1329,6 +1330,37 @@ def test_counts_during_a_retry_after_the_window_use_tiktoken_instead_of_download _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() +def test_counts_after_an_idle_stretch_of_several_windows_still_share_one_retry_download(): + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + count_during_retry: Final = Future[SelectTokenizerResponse]() + + 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)) + 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: + _load_huggingface_tokenizer.cache_clear() + _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache() + + def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered(): from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted