scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(utils): keep a loaded Hugging Face tokenizer over a racing failure and retry once per window

Browse source 
A download that fails after another count loaded the same tokenizer no longer hides it, the first count after the retry window claims the retry while the others keep using tiktoken, and HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS defaults to 60s to match the proxy's tokenizer reuse window
This commit is contained in:
mateo-berri 2026-10-20 19:21:31 -07:00
parent c515f80f82
commit 5de851a680
5 changed files with 146 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -157,7 +157,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.

72
litellm/utils.py
View file

@ -29,7 +29,7 @@ import threading
import time
import traceback
from contextlib import nullcontext
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from functools import lru_cache, wraps
from importlib import resources
from inspect import iscoroutine
@ -2645,9 +2645,15 @@ def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
return None
@dataclass(frozen=True, slots=True)
class _HuggingFaceTokenizerLoaded:
pass
@dataclass(frozen=True, slots=True)
class _HuggingFaceTokenizerLoadFailure:
reason: str
retry_at: float
def _warn_huggingface_tokenizer_fallback(model: str, reason: str) -> None:
@ -2664,22 +2670,44 @@ def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | _Hugg
if kind is None:
return None
backend: Final = _huggingface_tokenizer_backend()
key: Final = _huggingface_tokenizer_failure_key(kind, backend)
key: Final = _huggingface_tokenizer_outcome_key(kind, 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)
known: Final = _claim_huggingface_tokenizer_load(key)
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)
)
if isinstance(tokenizer, _HuggingFaceTokenizerLoadFailure):
_warn_huggingface_tokenizer_fallback(model, tokenizer.reason)
return tokenizer
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
def _huggingface_tokenizer_failure_key(kind: HuggingFaceTokenizerKind, backend: Decision) -> str:
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."""
@ -2702,19 +2730,30 @@ def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind, backend: Decisio
assert_never(kind)
def _load_huggingface_tokenizer_recording_failure(
def _load_huggingface_tokenizer_remembering_outcome(
kind: HuggingFaceTokenizerKind, backend: Decision, key: str
) -> HuggingFace | _HuggingFaceTokenizerLoadFailure:
try:
return _load_huggingface_tokenizer(kind, backend)
tokenizer: Final = _load_huggingface_tokenizer(kind, backend)
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
return _remember_huggingface_tokenizer_failure(key, str(error))
with _HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK:
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.set_cache(key, _HuggingFaceTokenizerLoaded())
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):
@ -8332,7 +8371,7 @@ def trim_messages(
from litellm.caching.in_memory_cache import InMemoryCache
class _HuggingFaceTokenizerLoadFailureCache(Protocol):
class _HuggingFaceTokenizerLoadOutcomeCache(Protocol):
def get_cache(self, key: str) -> object | None: ...
def set_cache(self, key: str, value: object, ttl: float | None = None) -> None: ...
@ -8340,16 +8379,17 @@ class _HuggingFaceTokenizerLoadFailureCache(Protocol):
def flush_cache(self) -> None: ...
_HUGGINGFACE_TOKENIZER_LOAD_FAILURES: Final[_HuggingFaceTokenizerLoadFailureCache] = (
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES: Final[_HuggingFaceTokenizerLoadOutcomeCache] = (
cast( # cast-ok: cache methods are untyped
_HuggingFaceTokenizerLoadFailureCache,
_HuggingFaceTokenizerLoadOutcomeCache,
InMemoryCache(
max_size_in_memory=16,
default_ttl=HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS,
default_ttl=2 * HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS,
clock=lambda: time.monotonic(),
),
)
)
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES_LOCK: Final = threading.Lock()
class AvailableModelsCache(InMemoryCache):

92
tests/unit/litellm_core_utils/test_token_counter.py
View file

@ -941,7 +941,7 @@ def test_token_counter():
import unittest
from litellm.utils import (
_HUGGINGFACE_TOKENIZER_LOAD_FAILURES,
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES,
_load_huggingface_tokenizer,
_select_tokenizer_helper,
claude_json_str,
@ -950,7 +950,7 @@ from litellm.utils import (
# 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):
@ -962,7 +962,7 @@ class TestTokenizerSelection(unittest.TestCase):
--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):
@ -1051,7 +1051,7 @@ class TestTokenizerSelection(unittest.TestCase):
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")
@ -1069,7 +1069,7 @@ class TestTokenizerSelection(unittest.TestCase):
assert second_count == expected
finally:
_load_huggingface_tokenizer.cache_clear()
_HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache()
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache()
@patch("litellm.utils.time.monotonic", return_value=299.0)
@patch("litellm.utils.tokenizer_dispatch.from_pretrained")
@ -1077,7 +1077,7 @@ class TestTokenizerSelection(unittest.TestCase):
self, mock_from_pretrained: MagicMock, mock_monotonic: MagicMock
):
_load_huggingface_tokenizer.cache_clear()
_HUGGINGFACE_TOKENIZER_LOAD_FAILURES.flush_cache()
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache()
try:
def from_pretrained_side_effect(_: str) -> object:
if mock_from_pretrained.call_count == 1:
@ -1107,7 +1107,7 @@ class TestTokenizerSelection(unittest.TestCase):
assert mock_from_pretrained.call_count == 2
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
@ -1175,7 +1175,7 @@ def test_disable_hf_tokenizer_download_keeps_deepseek_v4_on_cl100k(monkeypatch:
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:
@ -1190,7 +1190,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 (
@ -1203,14 +1203,82 @@ def test_remembered_tokenizer_failure_warns_once_instead_of_on_every_count(caplo
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_does_not_hide_it():
_load_huggingface_tokenizer.cache_clear()
_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_type: Final = slow_count.result(timeout=5)["type"]
later_counts: Final = tuple(_select_tokenizer_helper(model) for _ in range(3))
assert slow_count_type == "openai_tokenizer"
assert loading_count["tokenizer"] is sentinel.tokenizer
assert tuple(count["tokenizer"] for count in later_counts) == (sentinel.tokenizer,) * 3
finally:
_load_huggingface_tokenizer.cache_clear()
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache()
def test_counts_during_a_retry_after_the_window_use_tiktoken_instead_of_downloading_again():
_load_huggingface_tokenizer.cache_clear()
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES.flush_cache()
model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct"
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,
):
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:
_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
_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:
@ -1223,7 +1291,7 @@ def test_fork_guard_error_while_loading_a_tokenizer_is_raised_and_not_remembered
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():

12
tests/unit/litellm_core_utils/test_tokenizer.py
View file

@ -408,14 +408,14 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp
from unittest.mock import patch
from litellm.utils import (
_HUGGINGFACE_TOKENIZER_LOAD_FAILURES,
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES,
_load_huggingface_tokenizer,
_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")
@ -425,7 +425,7 @@ def test_missing_python_tokenizer_warns_before_approximate_count(caplog, monkeyp
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])
@ -483,14 +483,14 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch):
from unittest.mock import patch
from litellm.utils import (
_HUGGINGFACE_TOKENIZER_LOAD_FAILURES,
_HUGGINGFACE_TOKENIZER_LOAD_OUTCOMES,
_load_huggingface_tokenizer,
_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)
@ -518,7 +518,7 @@ def test_tokenizer_fallback_logs_safe_diagnostic_context(caplog, monkeypatch):
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])

6
tests/unit/llms/huggingface/chat/test_huggingface_chat_transformation.py
View file

@ -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, _load_huggingface_tokenizer
HUGGINGFACE_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct"
HUGGINGFACE_ROUTER_MODEL: Final = "meta-llama/Meta-Llama-3-8B-Instruct-Turbo"
@ -85,12 +85,12 @@ def isolate_huggingface_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator
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 = {