fix(frontend): evict async tokenizer executors when the tokenizer is retired

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(

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