From fcfe57e7a7889271a4b2468e8f4f0cf5b473ebfc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:01:22 -0700 Subject: [PATCH] fix(token_counter): let a process that exits during a background tokenizer retry exit cleanly instead of aborting --- litellm/utils.py | 6 ++--- .../litellm_core_utils/test_token_counter.py | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 16b51b9b711..803f426a8b2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2684,14 +2684,14 @@ _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() +def _start_in_thread_joined_at_exit(work: Callable[[], None]) -> None: + threading.Thread(target=work, name="litellm-huggingface-tokenizer-retry", daemon=False).start() @dataclass(slots=True) class _HuggingFaceTokenizerLoadOutcomes: clock: Callable[[], float] = time.monotonic - start_retry: Callable[[Callable[[], None]], None] = _start_in_daemon_thread + start_retry: Callable[[Callable[[], None]], None] = _start_in_thread_joined_at_exit pid: Callable[[], int] = os.getpid lock: threading.Lock = field(default_factory=threading.Lock) by_tokenizer: Mapping[_HuggingFaceTokenizerKey, _HuggingFaceTokenizerLoadOutcome] = field( diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 59f9188264a..78cbff4f951 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -2611,6 +2611,27 @@ tokens = litellm.encode(model="meta-llama/Llama-2-7b-chat", text=text) print(json.dumps({"tokens": tokens, "decoded": litellm.decode(model="meta-llama/Llama-2-7b-chat", tokens=tokens), "requested": sorted(set(requested))})) """ +HUB_RETRY_DURING_EXIT_SCRIPT: Final = HUB_SETUP_SCRIPT + """ +import threading +import time +from litellm.utils import _HuggingFaceTokenizerLoadOutcomes, _return_huggingface_tokenizer +retry_reached_hub = threading.Event() +def answer_once_the_process_is_exiting(request): + retry_reached_hub.set() + while threading.main_thread().is_alive(): + time.sleep(0.01) + time.sleep(0.5) + return handle(request) +now = [0.0] +outcomes = _HuggingFaceTokenizerLoadOutcomes(clock=lambda: now[0]) +model = "meta-llama/llama-3-70b-instruct" +_return_huggingface_tokenizer(model, outcomes) +huggingface_hub.set_client_factory(lambda: httpx.Client(transport=httpx.MockTransport(answer_once_the_process_is_exiting))) +now[0] = 3600.0 +_return_huggingface_tokenizer(model, outcomes) +print(json.dumps({"retry_reached_hub": retry_reached_hub.wait(timeout=60)})) +""" + def _word_level_tokenizer_json( pre_tokenizer: pre_tokenizers.PreTokenizer, vocab: Mapping[str, int] | None = None @@ -2673,6 +2694,10 @@ def test_token_counter_uses_the_tokenizer_of_each_model_family_and_of_a_custom_t } +def test_process_exits_cleanly_while_a_background_tokenizer_retry_is_downloading(tmp_path: Path) -> None: + assert _run_in_memory_hub(HUB_RETRY_DURING_EXIT_SCRIPT, {}, "", tmp_path) == {"retry_reached_hub": True} + + def test_empty_custom_tokenizer_uses_model_tokenizer() -> None: text_value: Final = "A tokenizer fallback should preserve the model encoding." custom_count: Final = _get_exact_count_function("gpt-3.5-turbo", {})(text_value)