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

fix(utils): keep Hugging Face tokenizer outcomes until their retry time instead of a TTL

Browse source 
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
This commit is contained in:
mateo-berri 2026-10-28 05:29:12 -07:00
parent 610897aca3
commit d1cb37bc4a
2 changed files with 84 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
litellm/utils.py
View file

@ -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)

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

@ -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