A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Return the tokenizer from maybe_make_thread_pool so it survives pickling (#45460)
Signed-off-by: Wayne Chiu <waynehacking8@gmail.com>
WEI CHENG CHIU committed
Jun 13, 2026 at 14:01 UTC
5b2943f5a6c5fb267d1b2029c666f9d6b0e4ebd6
2 files changed
+28
-1
tests/tokenizers_/test_hf.py
+25
-1
@@ -7,7 +7,11 @@ import pytest
7
from transformers import AutoTokenizer
8
9
from vllm.tokenizers import TokenizerLike
10
-from vllm.tokenizers.hf import get_cached_tokenizer
10
+from vllm.tokenizers.hf import (
11
+ ThreadSafeHFTokenizerMixin,
12
+ get_cached_tokenizer,
13
+ maybe_make_thread_pool,
14
+)
15
16
17
@pytest.mark.parametrize("model_id", ["gpt2", "zai-org/chatglm3-6b"])
@@ -41,3 +45,23 @@ def _check_consistency(target: TokenizerLike, expected: TokenizerLike):
45
)
46
47
assert target.encode("prompt") == expected.encode("prompt")
48
+
49
+
50
+@pytest.mark.parametrize("model_id", ["gpt2"])
51
+def test_thread_pool_tokenizer_pickle(model_id: str):
52
+ """Regression test for issue #45433: the thread-pool tokenizer wrapper
53
+ reconstructs through maybe_make_thread_pool on unpickling, which used to
54
+ fall off the end and return None."""
55
+ reference_tokenizer = AutoTokenizer.from_pretrained(model_id)
56
+
57
+ pooled_tokenizer = maybe_make_thread_pool(deepcopy(reference_tokenizer))
58
+ assert pooled_tokenizer is not None
59
+ assert isinstance(pooled_tokenizer, ThreadSafeHFTokenizerMixin)
60
+
61
+ unpickled_tokenizer = pickle.loads(pickle.dumps(pooled_tokenizer))
62
+ assert unpickled_tokenizer is not None
63
+ assert isinstance(unpickled_tokenizer, ThreadSafeHFTokenizerMixin)
64
+ assert unpickled_tokenizer.encode("prompt") == reference_tokenizer.encode("prompt")
65
+
66
+ # Idempotence: wrapping an already-pooled tokenizer returns it unchanged.
67
+ assert maybe_make_thread_pool(pooled_tokenizer) is pooled_tokenizer
vllm/tokenizers/hf.py
+3
@@ -99,6 +99,9 @@ def maybe_make_thread_pool(tokenizer: _T, copies: int = 1):
99
TokenizerPool.__name__ = f"TokenizerPool{og_tokenizer.__class__.__name__}"
100
101
tokenizer.__class__ = TokenizerPool
102
+ # Return the tokenizer: TokenizerPool.__reduce__ reconstructs through this
103
+ # function, so falling off the end would unpickle to None (issue #45433).
104
+ return tokenizer
105
106
107
def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: