diff --git a/components/src/dynamo/frontend/prepost.py b/components/src/dynamo/frontend/prepost.py index e6a6f0b5f1b2..0a560c545ed3 100644 --- a/components/src/dynamo/frontend/prepost.py +++ b/components/src/dynamo/frontend/prepost.py @@ -6,11 +6,12 @@ import json import os import re +import weakref from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from dataclasses import dataclass -from functools import lru_cache +from functools import lru_cache, partial from typing import TYPE_CHECKING, Any, Protocol, TypeGuard, cast from vllm.entrypoints.chat_utils import make_tool_call_id @@ -69,10 +70,35 @@ class PreprocessResult: uses_dynamo_json_tool_call_fallback: bool = False -_ASYNC_TOKENIZER_POOL: dict[int, Callable[..., Awaitable[Any]]] = {} +# One executor per live tokenizer. The id-keyed registry stores only a weak +# reference to the tokenizer, and every lookup verifies object identity. This +# avoids both id reuse and WeakKeyDictionary's referent-based equality. +_ASYNC_TOKENIZER_EXECUTORS: dict[ + int, tuple[weakref.ReferenceType[TokenizerLike], ThreadPoolExecutor] +] = {} +# Fallback for tokenizers that do not support weak references; retains +# entries for the process lifetime (the previous behavior for all tokenizers). +# The tokenizer is stored alongside its executor so its id() cannot be +# recycled by a different tokenizer while the entry lives. +_STRONG_ASYNC_TOKENIZER_EXECUTORS: dict[ + int, tuple[TokenizerLike, ThreadPoolExecutor] +] = {} SKIP_REQUEST_VALIDATION = os.getenv("DYN_VLLM_SKIP_REQUEST_VALIDATION", "1") == "1" +def _evict_async_tokenizer_executor( + key: int, + executor: ThreadPoolExecutor, + tokenizer_ref: weakref.ReferenceType[TokenizerLike], +) -> None: + entry = _ASYNC_TOKENIZER_EXECUTORS.get(key) + if entry is not None and entry[0] is tokenizer_ref: + del _ASYNC_TOKENIZER_EXECUTORS[key] + # The tokenizer's last reference may be released by this executor's own + # worker, so shutdown must not try to join the current thread. + executor.shutdown(wait=False) + + def _reject_non_finite_json(value: str) -> Any: raise ValueError(f"non-finite JSON number: {value}") @@ -332,14 +358,30 @@ def _build_assistant_guided_decoding( def _get_async_tokenizer(tokenizer: TokenizerLike) -> Callable[..., Awaitable[Any]]: - key = id(tokenizer) - async_tokenizer = _ASYNC_TOKENIZER_POOL.get(key) - if async_tokenizer is None: - async_tokenizer = make_async( - tokenizer, executor=ThreadPoolExecutor(max_workers=1) - ) - _ASYNC_TOKENIZER_POOL[key] = async_tokenizer - return async_tokenizer + try: + tokenizer_ref = weakref.ref(tokenizer) + except TypeError: + # Tokenizer does not support weak references. + key = id(tokenizer) + entry = _STRONG_ASYNC_TOKENIZER_EXECUTORS.get(key) + if entry is None: + executor = ThreadPoolExecutor(max_workers=1) + _STRONG_ASYNC_TOKENIZER_EXECUTORS[key] = (tokenizer, executor) + else: + executor = entry[1] + else: + key = id(tokenizer) + entry = _ASYNC_TOKENIZER_EXECUTORS.get(key) + if entry is None or entry[0]() is not tokenizer: + executor = ThreadPoolExecutor(max_workers=1) + tokenizer_ref = weakref.ref( + tokenizer, + partial(_evict_async_tokenizer_executor, key, executor), + ) + _ASYNC_TOKENIZER_EXECUTORS[key] = (tokenizer_ref, executor) + else: + executor = entry[1] + return make_async(tokenizer, executor=executor) def _materialize_assistant_tool_calls( diff --git a/tests/frontend/test_prepost_tokenizer_pool.py b/tests/frontend/test_prepost_tokenizer_pool.py new file mode 100644 index 000000000000..d9eb6729cd86 --- /dev/null +++ b/tests/frontend/test_prepost_tokenizer_pool.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the async tokenizer executor pool in dynamo.frontend.prepost.""" + +# mypy seems to be running both sides of the HAS_VLLM if statement +# mypy: ignore-errors + +import asyncio +import gc +import threading +import weakref +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +import pytest + +from .common import check_module_available + +HAS_VLLM = check_module_available("vllm.entrypoints.openai.chat_completion.protocol") +if HAS_VLLM: + from dynamo.frontend import prepost + + +pytestmark = [ + pytest.mark.vllm, + pytest.mark.core, + pytest.mark.gpu_0, # "Hardware" + pytest.mark.pre_merge, # "Lifecyle" + pytest.mark.unit, # "Test Type" + pytest.mark.skipif(not HAS_VLLM, reason="requires vllm"), +] + + +class _DummyTokenizer: + def __init__(self): + self.prompts = [] + + def __call__(self, prompt, **kwargs): + self.prompts.append(prompt) + return SimpleNamespace(input_ids=[1, 2, 3]) + + +class _BlockingTokenizer(_DummyTokenizer): + def __init__(self): + super().__init__() + self.started = threading.Event() + self.release = threading.Event() + + def __call__(self, prompt, **kwargs): + self.started.set() + if not self.release.wait(timeout=5): + raise TimeoutError("test tokenizer was not released") + return super().__call__(prompt, **kwargs) + + +class _SlotsTokenizer: + # No __weakref__: weakref.ref() raises TypeError. + __slots__ = () + + def __call__(self, prompt, **kwargs): + return SimpleNamespace(input_ids=[4, 5, 6]) + + +class _EqualTokenizer(_DummyTokenizer): + def __eq__(self, other): + return isinstance(other, _EqualTokenizer) + + def __hash__(self): + return 1 + + +class _RecordingExecutor(ThreadPoolExecutor): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.shutdown_calls = 0 + self.shutdown_waits = [] + + def shutdown(self, wait=True, *, cancel_futures=False): + self.shutdown_calls += 1 + self.shutdown_waits.append(wait) + super().shutdown(wait, cancel_futures=cancel_futures) + + +@pytest.fixture(autouse=True) +def _clean_pools(monkeypatch): + monkeypatch.setattr(prepost, "ThreadPoolExecutor", _RecordingExecutor) + prepost._ASYNC_TOKENIZER_EXECUTORS.clear() + prepost._STRONG_ASYNC_TOKENIZER_EXECUTORS.clear() + yield + for _, executor in list(prepost._ASYNC_TOKENIZER_EXECUTORS.values()): + executor.shutdown() + for _, executor in list(prepost._STRONG_ASYNC_TOKENIZER_EXECUTORS.values()): + executor.shutdown() + prepost._ASYNC_TOKENIZER_EXECUTORS.clear() + prepost._STRONG_ASYNC_TOKENIZER_EXECUTORS.clear() + + +async def _call_tokenizer(tokenizer, prompt): + # make_async returns a Future, so the call must be awaited inside a + # running loop rather than passed to asyncio.run directly. + return await prepost._get_async_tokenizer(tokenizer)(prompt) + + +def test_same_tokenizer_reuses_one_executor(): + tokenizer = _DummyTokenizer() + prepost._get_async_tokenizer(tokenizer) + prepost._get_async_tokenizer(tokenizer) + + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 1 + tokenizer_ref, executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(tokenizer)] + assert tokenizer_ref() is tokenizer + assert isinstance(executor, _RecordingExecutor) + + result = asyncio.run(_call_tokenizer(tokenizer, "hello")) + assert result.input_ids == [1, 2, 3] + assert tokenizer.prompts == ["hello"] + + +def test_executor_evicted_and_shut_down_when_tokenizer_collected(): + tokenizer = _DummyTokenizer() + prepost._get_async_tokenizer(tokenizer) + _, executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(tokenizer)] + tokenizer_ref = weakref.ref(tokenizer) + + # No cached wrapper may survive this del: it would pin the tokenizer. + del tokenizer + gc.collect() + + assert tokenizer_ref() is None + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 0 + assert executor.shutdown_calls == 1 + + +def test_new_tokenizer_after_eviction_gets_fresh_executor(): + first = _DummyTokenizer() + prepost._get_async_tokenizer(first) + _, first_executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(first)] + del first + gc.collect() + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 0 + + second = _DummyTokenizer() + prepost._get_async_tokenizer(second) + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 1 + _, second_executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(second)] + assert second_executor is not first_executor + + +def test_equal_tokenizers_have_independent_executor_lifetimes(): + first = _EqualTokenizer() + second = _EqualTokenizer() + assert first == second + assert first is not second + + prepost._get_async_tokenizer(first) + prepost._get_async_tokenizer(second) + _, first_executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(first)] + _, second_executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(second)] + + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 2 + assert first_executor is not second_executor + + first_ref = weakref.ref(first) + del first + gc.collect() + + assert first_ref() is None + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 1 + assert first_executor.shutdown_calls == 1 + assert second_executor.shutdown_calls == 0 + + result = asyncio.run(_call_tokenizer(second, "still live")) + assert result.input_ids == [1, 2, 3] + assert second.prompts == ["still live"] + + +async def _run_inflight_tokenizer(tokenizer): + tokenizer_ref = weakref.ref(tokenizer) + async_tokenizer = prepost._get_async_tokenizer(tokenizer) + _, executor = prepost._ASYNC_TOKENIZER_EXECUTORS[id(tokenizer)] + future = async_tokenizer("in flight") + release = tokenizer.release + assert await asyncio.to_thread(tokenizer.started.wait, 5) + + try: + # The running work item is the only remaining strong reference to the + # tokenizer after the wrapper and caller references are released. + del async_tokenizer + del tokenizer + gc.collect() + assert tokenizer_ref() is not None + + release.set() + result = await future + for _ in range(100): + gc.collect() + if tokenizer_ref() is None: + break + await asyncio.sleep(0.01) + + return result, executor, tokenizer_ref() + finally: + release.set() + + +def test_inflight_tokenizer_is_evicted_from_worker_thread(): + result, executor, tokenizer = asyncio.run( + _run_inflight_tokenizer(_BlockingTokenizer()) + ) + + assert result.input_ids == [1, 2, 3] + assert tokenizer is None + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 0 + assert executor.shutdown_calls == 1 + assert executor.shutdown_waits == [False] + + +def test_non_weakrefable_tokenizer_falls_back_to_strong_pool(): + tokenizer = _SlotsTokenizer() + result = asyncio.run(_call_tokenizer(tokenizer, "hi")) + assert result.input_ids == [4, 5, 6] + + assert len(prepost._ASYNC_TOKENIZER_EXECUTORS) == 0 + assert len(prepost._STRONG_ASYNC_TOKENIZER_EXECUTORS) == 1 + + # The fallback entry must retain the tokenizer itself so its id() cannot + # be recycled by a different tokenizer while the entry lives. + retained_tokenizer, _ = prepost._STRONG_ASYNC_TOKENIZER_EXECUTORS[id(tokenizer)] + assert retained_tokenizer is tokenizer + + prepost._get_async_tokenizer(tokenizer) + assert len(prepost._STRONG_ASYNC_TOKENIZER_EXECUTORS) == 1