forked from Karylab-cklius/vllm
[Frontend] Remove AsyncMicrobatchTokenizer. (#45759)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
This commit is contained in:
wang.yuqi
+12
-16
@@ -38,10 +38,7 @@ from vllm.multimodal.processing import BaseMultiModalProcessor
|
||||
from vllm.multimodal.processing import ProcessorInputs as MMProcessorInputs
|
||||
from vllm.multimodal.registry import MultiModalTimingRegistry
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.utils.async_utils import (
|
||||
AsyncMicrobatchTokenizer,
|
||||
make_async,
|
||||
)
|
||||
from vllm.utils.async_utils import make_async
|
||||
from vllm.utils.counter import AtomicCounter
|
||||
from vllm.utils.torch_utils import set_default_torch_num_threads
|
||||
from vllm.v1.metrics.stats import MultiModalCacheStats
|
||||
@@ -92,8 +89,9 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
# to keep the asyncio event loop responsive under concurrent load.
|
||||
self._mm_executor: Executor = self._executor
|
||||
|
||||
# Lazy initialization since offline LLM doesn't use async
|
||||
self._async_tokenizer: AsyncMicrobatchTokenizer | None = None
|
||||
# Offloading tokenizer encode & decode to thread pool.
|
||||
self._async_tokenizer_encode = make_async(self._encode, executor=self._executor)
|
||||
self._async_tokenizer_decode = make_async(self._decode, executor=self._executor)
|
||||
|
||||
self.mm_processor: BaseMultiModalProcessor | None = None
|
||||
self._readonly_mm_processor: BaseMultiModalProcessor | None = None
|
||||
@@ -146,13 +144,11 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
|
||||
return tokenizer
|
||||
|
||||
def get_async_tokenizer(self) -> AsyncMicrobatchTokenizer:
|
||||
if self._async_tokenizer is None:
|
||||
self._async_tokenizer = AsyncMicrobatchTokenizer(
|
||||
self.get_tokenizer(), executor=self._executor
|
||||
)
|
||||
def _decode(self, *args, **kwargs):
|
||||
return self.get_tokenizer().decode(*args, **kwargs)
|
||||
|
||||
return self._async_tokenizer
|
||||
def _encode(self, *args, **kwargs):
|
||||
return self.get_tokenizer().encode(*args, **kwargs)
|
||||
|
||||
def get_mm_processor(self) -> "BaseMultiModalProcessor":
|
||||
if self.mm_processor is None:
|
||||
@@ -436,8 +432,7 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
prompt: TextPrompt,
|
||||
params: TokenizeParams,
|
||||
) -> TokensPrompt:
|
||||
tokenizer = self.get_async_tokenizer()
|
||||
prompt_token_ids = await tokenizer.encode(
|
||||
prompt_token_ids = await self._async_tokenizer_encode(
|
||||
prompt["prompt"],
|
||||
**params.get_encode_kwargs(),
|
||||
)
|
||||
@@ -451,8 +446,9 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
return prompt
|
||||
|
||||
async def _detokenize_prompt_async(self, prompt: TokensPrompt) -> TokensPrompt:
|
||||
tokenizer = self.get_async_tokenizer()
|
||||
prompt["prompt"] = await tokenizer.decode(prompt["prompt_token_ids"])
|
||||
prompt["prompt"] = await self._async_tokenizer_decode(
|
||||
prompt["prompt_token_ids"]
|
||||
)
|
||||
|
||||
return prompt
|
||||
|
||||
|
||||
@@ -14,215 +14,12 @@ from concurrent.futures import Executor, ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from transformers.tokenization_utils_base import BatchEncoding
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
P = ParamSpec("P")
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class AsyncMicrobatchTokenizer:
|
||||
"""Asynchronous tokenizer with micro-batching.
|
||||
|
||||
Pulls pending encode/decode requests from a queue and batches them
|
||||
up to reduce overhead. A single-thread ThreadPoolExecutor is used
|
||||
so the event loop stays responsive.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer,
|
||||
max_batch_size: int = 32,
|
||||
batch_wait_timeout_s: float = 0.002,
|
||||
executor: ThreadPoolExecutor | None = None,
|
||||
) -> None:
|
||||
self.tokenizer = tokenizer
|
||||
self.max_batch_size = max_batch_size
|
||||
self.batch_wait_timeout_s = batch_wait_timeout_s
|
||||
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._queues: dict[
|
||||
tuple,
|
||||
asyncio.Queue[tuple[str, dict, Future] | tuple[list[int], Future]],
|
||||
] = {}
|
||||
self._batcher_tasks: list[Task] = []
|
||||
|
||||
# Single-thread executor for blocking tokenizer calls.
|
||||
# Accept an external executor to serialize with other tokenizer users.
|
||||
self._executor = executor or ThreadPoolExecutor(max_workers=1)
|
||||
|
||||
# === Public async API ===
|
||||
async def __call__(self, prompt, **kwargs) -> BatchEncoding:
|
||||
result_future: Future = self._loop.create_future()
|
||||
key = self._queue_key("encode", kwargs)
|
||||
queue = self._get_queue(self._loop, key)
|
||||
await queue.put((prompt, kwargs, result_future))
|
||||
return await result_future
|
||||
|
||||
async def encode(self, prompt, **kwargs) -> list[int]:
|
||||
return (await self(prompt, **kwargs)).input_ids
|
||||
|
||||
async def decode(self, token_ids, **kwargs) -> str:
|
||||
result_future: Future = self._loop.create_future()
|
||||
key = self._queue_key("decode", kwargs)
|
||||
queue = self._get_queue(self._loop, key)
|
||||
await queue.put((token_ids, result_future))
|
||||
return await result_future
|
||||
|
||||
# === Internal helpers ===
|
||||
def _get_queue(
|
||||
self, loop: asyncio.AbstractEventLoop, key: tuple
|
||||
) -> asyncio.Queue[tuple[str, dict, Future] | tuple[list[int], Future]]:
|
||||
"""Get the request queue for the given operation key, creating a new
|
||||
queue and batcher task if needed."""
|
||||
queue = self._queues.get(key)
|
||||
if queue is None:
|
||||
self._queues[key] = queue = asyncio.Queue()
|
||||
if key[0] == "encode":
|
||||
can_batch = key[1] != "other"
|
||||
coro = self._batch_encode_loop(queue, can_batch)
|
||||
else:
|
||||
assert key[0] == "decode", f"Unknown operation type: {key[0]}."
|
||||
coro = self._batch_decode_loop(queue)
|
||||
self._batcher_tasks.append(loop.create_task(coro))
|
||||
return queue
|
||||
|
||||
async def _batch_encode_loop(self, queue: asyncio.Queue, can_batch: bool):
|
||||
"""Batch incoming encode requests for efficiency."""
|
||||
while True:
|
||||
prompt, kwargs, result_future = await queue.get()
|
||||
prompts = [prompt]
|
||||
kwargs_list = [kwargs]
|
||||
result_futures = [result_future]
|
||||
deadline = self._loop.time() + self.batch_wait_timeout_s
|
||||
|
||||
while len(prompts) < self.max_batch_size:
|
||||
timeout = deadline - self._loop.time()
|
||||
if timeout <= 0:
|
||||
break
|
||||
try:
|
||||
prompt, kwargs, result_future = await asyncio.wait_for(
|
||||
queue.get(), timeout
|
||||
)
|
||||
prompts.append(prompt)
|
||||
result_futures.append(result_future)
|
||||
if not can_batch:
|
||||
kwargs_list.append(kwargs)
|
||||
except asyncio.TimeoutError:
|
||||
break
|
||||
|
||||
try:
|
||||
# If every request uses identical kwargs we can run a single
|
||||
# batched tokenizer call for a big speed-up.
|
||||
if can_batch and len(prompts) > 1:
|
||||
batch_encode_fn = partial(self.tokenizer, prompts, **kwargs)
|
||||
results = await self._loop.run_in_executor(
|
||||
self._executor, batch_encode_fn
|
||||
)
|
||||
|
||||
for i, fut in enumerate(result_futures):
|
||||
if not fut.done():
|
||||
data = {k: v[i] for k, v in results.items()}
|
||||
fut.set_result(BatchEncoding(data))
|
||||
else:
|
||||
encode_fn = lambda prompts=prompts, kwargs=kwargs_list: [
|
||||
self.tokenizer(p, **kw) for p, kw in zip(prompts, kwargs)
|
||||
]
|
||||
results = await self._loop.run_in_executor(
|
||||
self._executor, encode_fn
|
||||
)
|
||||
|
||||
for fut, res in zip(result_futures, results):
|
||||
if not fut.done():
|
||||
fut.set_result(res)
|
||||
except Exception as e:
|
||||
for fut in result_futures:
|
||||
if not fut.done():
|
||||
fut.set_exception(e)
|
||||
|
||||
async def _batch_decode_loop(self, queue: asyncio.Queue):
|
||||
"""Batch incoming decode requests for efficiency."""
|
||||
while True:
|
||||
token_ids, result_future = await queue.get()
|
||||
token_ids_list = [token_ids]
|
||||
result_futures = [result_future]
|
||||
deadline = self._loop.time() + self.batch_wait_timeout_s
|
||||
|
||||
while len(token_ids_list) < self.max_batch_size:
|
||||
timeout = deadline - self._loop.time()
|
||||
if timeout <= 0:
|
||||
break
|
||||
try:
|
||||
token_ids, result_future = await asyncio.wait_for(
|
||||
queue.get(), timeout
|
||||
)
|
||||
token_ids_list.append(token_ids)
|
||||
result_futures.append(result_future)
|
||||
except asyncio.TimeoutError:
|
||||
break
|
||||
|
||||
try:
|
||||
# Perform a single batched decode call for all requests
|
||||
results = await self._loop.run_in_executor(
|
||||
self._executor, self.tokenizer.batch_decode, token_ids_list
|
||||
)
|
||||
for fut, res in zip(result_futures, results):
|
||||
if not fut.done():
|
||||
fut.set_result(res)
|
||||
except Exception as e:
|
||||
for fut in result_futures:
|
||||
if not fut.done():
|
||||
fut.set_exception(e)
|
||||
|
||||
def _queue_key(self, op: str, kwargs: dict) -> tuple:
|
||||
"""
|
||||
Return a normalized key describing operation + kwargs.
|
||||
|
||||
- `add_special_tokens`: {True/False}
|
||||
- `truncation`: {True/False}
|
||||
- If `truncation` is False (`max_length` is None),
|
||||
returns a key for a can_batch queue.
|
||||
- If `truncation` is True and `max_length` is None or equals
|
||||
`tokenizer.model_max_length`, returns a key for a can_batch queue.
|
||||
- Otherwise, returns a key for a cannot_batch queue.
|
||||
|
||||
Examples:
|
||||
- Decode: ("decode",)
|
||||
- Encode typical:
|
||||
("encode", add_special_tokens, bool_truncation, max_length_label)
|
||||
- Fallback: ("encode", "other")
|
||||
"""
|
||||
|
||||
if op == "decode":
|
||||
return ("decode",)
|
||||
|
||||
add_special_tokens = kwargs.get("add_special_tokens", True)
|
||||
truncation = kwargs.get("truncation", False)
|
||||
max_length = kwargs.get("max_length")
|
||||
|
||||
if not truncation:
|
||||
return "encode", add_special_tokens, False, None
|
||||
|
||||
model_max = getattr(self.tokenizer, "model_max_length", None)
|
||||
if max_length is None or (model_max is not None and max_length == model_max):
|
||||
return "encode", add_special_tokens, True, "model_max"
|
||||
|
||||
return "encode", "other"
|
||||
|
||||
def __del__(self):
|
||||
if (
|
||||
(tasks := getattr(self, "_batcher_tasks", None))
|
||||
and (loop := getattr(self, "_loop", None))
|
||||
and not loop.is_closed()
|
||||
):
|
||||
|
||||
def cancel_tasks():
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
|
||||
loop.call_soon_threadsafe(cancel_tasks)
|
||||
|
||||
|
||||
def cancel_task_threadsafe(task: Task):
|
||||
if task and not task.done():
|
||||
run_in_loop(task.get_loop(), task.cancel)
|
||||
|
||||
Reference in New Issue
Block a user