@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 828523ad8e

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Chore] Separate out `vllm.utils.async_utils` (#26913)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Oct 15, 2025 at 23:33 UTC 828523ad8e26d15e5cec27ba623310a950f53121

17 files changed +364 -354

tests/lora/test_add_lora.py

+1 -1

@@ -12,7 +12,7 @@ from vllm.entrypoints.openai.api_server import (

12 from vllm.inputs import TextPrompt

13 from vllm.lora.request import LoRARequest

14 from vllm.sampling_params import SamplingParams

15 -from vllm.utils import merge_async_iterators

15 +from vllm.utils.async_utils import merge_async_iterators

17 MODEL_PATH = "zai-org/chatglm3-6b"

18 LORA_RANK = 64

tests/utils_/test_async_utils.py new

+42

@@ -0,0 +1,42 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import asyncio

4 +from collections.abc import AsyncIterator

5 +

6 +import pytest

7 +

8 +from vllm.utils.async_utils import merge_async_iterators

9 +

10 +

11 +async def _mock_async_iterator(idx: int):

12 + try:

13 + while True:

14 + yield f"item from iterator {idx}"

15 + await asyncio.sleep(0.1)

16 + except asyncio.CancelledError:

17 + print(f"iterator {idx} cancelled")

18 +

19 +

20 +@pytest.mark.asyncio

21 +async def test_merge_async_iterators():

22 + iterators = [_mock_async_iterator(i) for i in range(3)]

23 + merged_iterator = merge_async_iterators(*iterators)

24 +

25 + async def stream_output(generator: AsyncIterator[tuple[int, str]]):

26 + async for idx, output in generator:

27 + print(f"idx: {idx}, output: {output}")

28 +

29 + task = asyncio.create_task(stream_output(merged_iterator))

30 + await asyncio.sleep(0.5)

31 + task.cancel()

32 + with pytest.raises(asyncio.CancelledError):

33 + await task

34 +

35 + for iterator in iterators:

36 + try:

37 + await asyncio.wait_for(anext(iterator), 1)

38 + except StopAsyncIteration:

39 + # All iterators should be cancelled and print this message.

40 + print("Iterator was cancelled normally")

41 + except (Exception, asyncio.CancelledError) as e:

42 + raise AssertionError() from e

tests/utils_/test_utils.py

-36

@@ -2,14 +2,12 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 # ruff: noqa

5 -import asyncio

5 import hashlib

6 import json

7 import os

8 import pickle

9 import socket

10 import tempfile

12 -from collections.abc import AsyncIterator

11 from pathlib import Path

12 from unittest.mock import patch

@@ -37,7 +35,6 @@ from vllm.utils import (

35 make_zmq_path,

36 make_zmq_socket,

37 memory_profiling,

40 - merge_async_iterators,

38 sha256,

39 split_host_port,

40 split_zmq_path,

@@ -48,39 +45,6 @@ from vllm.utils import (

45 from ..utils import create_new_process_for_each_test

51 -@pytest.mark.asyncio

52 -async def test_merge_async_iterators():

53 - async def mock_async_iterator(idx: int):

54 - try:

55 - while True:

56 - yield f"item from iterator {idx}"

57 - await asyncio.sleep(0.1)

58 - except asyncio.CancelledError:

59 - print(f"iterator {idx} cancelled")

60 -

61 - iterators = [mock_async_iterator(i) for i in range(3)]

62 - merged_iterator = merge_async_iterators(*iterators)

63 -

64 - async def stream_output(generator: AsyncIterator[tuple[int, str]]):

65 - async for idx, output in generator:

66 - print(f"idx: {idx}, output: {output}")

67 -

68 - task = asyncio.create_task(stream_output(merged_iterator))

69 - await asyncio.sleep(0.5)

70 - task.cancel()

71 - with pytest.raises(asyncio.CancelledError):

72 - await task

73 -

74 - for iterator in iterators:

75 - try:

76 - await asyncio.wait_for(anext(iterator), 1)

77 - except StopAsyncIteration:

78 - # All iterators should be cancelled and print this message.

79 - print("Iterator was cancelled normally")

80 - except (Exception, asyncio.CancelledError) as e:

81 - raise AssertionError() from e

82 -

83 -

48 def test_get_open_port(monkeypatch: pytest.MonkeyPatch):

49 with monkeypatch.context() as m:

50 m.setenv("VLLM_PORT", "5678")

vllm/benchmarks/throughput.py

+1 -1

@@ -34,7 +34,7 @@ from vllm.inputs import TextPrompt, TokensPrompt

34 from vllm.lora.request import LoRARequest

35 from vllm.outputs import RequestOutput

36 from vllm.sampling_params import BeamSearchParams

37 -from vllm.utils import merge_async_iterators

37 +from vllm.utils.async_utils import merge_async_iterators

40 def run_vllm(

vllm/entrypoints/openai/serving_completion.py

+2 -1

@@ -34,7 +34,8 @@ from vllm.logprobs import Logprob

34 from vllm.outputs import RequestOutput

35 from vllm.sampling_params import BeamSearchParams, SamplingParams

36 from vllm.transformers_utils.tokenizer import AnyTokenizer

37 -from vllm.utils import as_list, merge_async_iterators

37 +from vllm.utils import as_list

38 +from vllm.utils.async_utils import merge_async_iterators

40 logger = init_logger(__name__)

vllm/entrypoints/openai/serving_embedding.py

+1 -2

@@ -40,6 +40,7 @@ from vllm.outputs import (

40 )

41 from vllm.pooling_params import PoolingParams

42 from vllm.utils import chunk_list

43 +from vllm.utils.async_utils import merge_async_iterators

45 logger = init_logger(__name__)

@@ -387,8 +388,6 @@ class EmbeddingMixin(OpenAIServing):

388 )

389 generators.append(generator)

390

390 - from vllm.utils import merge_async_iterators

391 -

391 ctx.result_generator = merge_async_iterators(*generators)

392

393 return None

vllm/entrypoints/openai/serving_engine.py

+3 -4

@@ -90,14 +90,13 @@ from vllm.tracing import (

90 log_tracing_disabled_warning,

91 )

92 from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer

93 -from vllm.utils import (

93 +from vllm.utils import is_list_of, random_uuid

94 +from vllm.utils.async_utils import (

95 AsyncMicrobatchTokenizer,

96 collect_from_async_generator,

96 - is_list_of,

97 + make_async,

98 merge_async_iterators,

98 - random_uuid,

99 )

100 -from vllm.utils.func import make_async

100 from vllm.v1.engine import EngineCoreRequest

101

102 logger = init_logger(__name__)

vllm/entrypoints/openai/serving_pooling.py

+1 -1

@@ -36,7 +36,7 @@ from vllm.entrypoints.utils import _validate_truncation_size

36 from vllm.logger import init_logger

37 from vllm.outputs import PoolingOutput, PoolingRequestOutput

38 from vllm.tasks import SupportedTask

39 -from vllm.utils import merge_async_iterators

39 +from vllm.utils.async_utils import merge_async_iterators

41 logger = init_logger(__name__)

vllm/entrypoints/openai/serving_score.py

+1 -2

@@ -37,8 +37,7 @@ from vllm.logger import init_logger

37 from vllm.lora.request import LoRARequest

38 from vllm.outputs import PoolingRequestOutput, ScoringRequestOutput

39 from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer

40 -from vllm.utils import merge_async_iterators

41 -from vllm.utils.func import make_async

40 +from vllm.utils.async_utils import make_async, merge_async_iterators

42 logger = init_logger(__name__)

vllm/entrypoints/renderer.py

+1 -1

@@ -17,7 +17,7 @@ from vllm.inputs.data import TextPrompt as EngineTextPrompt

17 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt

18 from vllm.inputs.parse import get_prompt_components, parse_raw_prompts

19 from vllm.transformers_utils.tokenizer import AnyTokenizer

20 -from vllm.utils import AsyncMicrobatchTokenizer

20 +from vllm.utils.async_utils import AsyncMicrobatchTokenizer

23 @dataclass(frozen=True)

vllm/executor/executor_base.py

+1 -1

@@ -17,7 +17,7 @@ from vllm.logger import init_logger

17 from vllm.lora.request import LoRARequest

18 from vllm.sequence import ExecuteModelRequest

19 from vllm.tasks import SupportedTask

20 -from vllm.utils.func import make_async

20 +from vllm.utils.async_utils import make_async

21 from vllm.v1.outputs import SamplerOutput

22 from vllm.v1.worker.worker_base import WorkerBase

vllm/executor/ray_distributed_executor.py

+7 -2

@@ -20,12 +20,11 @@ from vllm.platforms import current_platform

20 from vllm.ray.ray_env import get_env_vars_to_copy

21 from vllm.sequence import ExecuteModelRequest

22 from vllm.utils import (

23 - _run_task_with_lock,

23 get_distributed_init_method,

24 get_ip,

25 get_open_port,

26 )

28 -from vllm.utils.func import make_async

27 +from vllm.utils.async_utils import make_async

28 from vllm.v1.outputs import SamplerOutput

30 if ray is not None:

@@ -748,3 +747,9 @@ class RayDistributedExecutor(DistributedExecutorBase):

747 # Assume that the Ray workers are healthy.

748 # TODO: check the health of the Ray workers

749 return

750 +

751 +

752 +async def _run_task_with_lock(task: Callable, lock: asyncio.Lock, *args, **kwargs):

753 + """Utility function to run async task in a lock"""

754 + async with lock:

755 + return await task(*args, **kwargs)

vllm/utils/__init__.py

-277

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 -import asyncio

4 import contextlib

5 import datetime

6 import enum

@@ -38,10 +37,8 @@ from argparse import (

37 RawDescriptionHelpFormatter,

38 _ArgumentGroup,

39 )

41 -from asyncio import FIRST_COMPLETED, AbstractEventLoop, Task

40 from collections import UserDict, defaultdict

41 from collections.abc import (

44 - AsyncGenerator,

42 Callable,

43 Collection,

44 Generator,

@@ -51,7 +48,6 @@ from collections.abc import (

48 Mapping,

49 Sequence,

50 )

54 -from concurrent.futures import ThreadPoolExecutor

51 from concurrent.futures.process import ProcessPoolExecutor

52 from dataclasses import dataclass, field

53 from functools import cache, lru_cache, partial, wraps

@@ -82,7 +78,6 @@ import zmq.asyncio

78 from packaging import version

79 from packaging.version import Version

80 from torch.library import Library

85 -from transformers.tokenization_utils_base import BatchEncoding

81 from typing_extensions import Never, TypeIs, assert_never

83 import vllm.envs as envs

@@ -223,278 +218,12 @@ def random_uuid() -> str:

218 return str(uuid.uuid4().hex)

219

220

226 -class AsyncMicrobatchTokenizer:

227 - """Asynchronous tokenizer with micro-batching.

228 -

229 - Pulls pending encode/decode requests from a queue and batches them

230 - up to reduce overhead. A single-thread ThreadPoolExecutor is used

231 - so the event loop stays responsive.

232 - """

233 -

234 - def __init__(

235 - self,

236 - tokenizer,

237 - max_batch_size: int = 32,

238 - batch_wait_timeout_s: float = 0.002,

239 - ) -> None:

240 - self.tokenizer = tokenizer

241 - self.max_batch_size = max_batch_size

242 - self.batch_wait_timeout_s = batch_wait_timeout_s

243 -

244 - self._loop = asyncio.get_running_loop()

245 - self._queues: dict[

246 - tuple,

247 - asyncio.Queue[

248 - tuple[str, dict, asyncio.Future] | tuple[list[int], asyncio.Future]

249 - ],

250 - ] = {}

251 - self._batcher_tasks: list[asyncio.Task] = []

252 -

253 - # Single-thread executor for blocking tokenizer calls.

254 - self._executor = ThreadPoolExecutor(max_workers=1)

255 -

256 - # === Public async API ===

257 - async def __call__(self, prompt, **kwargs):

258 - result_future: asyncio.Future = self._loop.create_future()

259 - key = self._queue_key("encode", kwargs)

260 - queue = self._get_queue(self._loop, key)

261 - await queue.put((prompt, kwargs, result_future))

262 - return await result_future

263 -

264 - async def decode(self, token_ids, **kwargs):

265 - result_future: asyncio.Future = self._loop.create_future()

266 - key = self._queue_key("decode", kwargs)

267 - queue = self._get_queue(self._loop, key)

268 - await queue.put((token_ids, result_future))

269 - return await result_future

270 -

271 - # === Internal helpers ===

272 - def _get_queue(

273 - self, loop: asyncio.AbstractEventLoop, key: tuple

274 - ) -> asyncio.Queue[

275 - tuple[str, dict, asyncio.Future] | tuple[list[int], asyncio.Future]

276 - ]:

277 - """Get the request queue for the given operation key, creating a new

278 - queue and batcher task if needed."""

279 - queue = self._queues.get(key)

280 - if queue is None:

281 - self._queues[key] = queue = asyncio.Queue()

282 - if key[0] == "encode":

283 - can_batch = key[1] != "other"

284 - coro = self._batch_encode_loop(queue, can_batch)

285 - else:

286 - assert key[0] == "decode", f"Unknown operation type: {key[0]}."

287 - coro = self._batch_decode_loop(queue)

288 - self._batcher_tasks.append(loop.create_task(coro))

289 - return queue

290 -

291 - async def _batch_encode_loop(self, queue: asyncio.Queue, can_batch: bool):

292 - """Batch incoming encode requests for efficiency."""

293 - while True:

294 - prompt, kwargs, result_future = await queue.get()

295 - prompts = [prompt]

296 - kwargs_list = [kwargs]

297 - result_futures = [result_future]

298 - deadline = self._loop.time() + self.batch_wait_timeout_s

299 -

300 - while len(prompts) < self.max_batch_size:

301 - timeout = deadline - self._loop.time()

302 - if timeout <= 0:

303 - break

304 - try:

305 - prompt, kwargs, result_future = await asyncio.wait_for(

306 - queue.get(), timeout

307 - )

308 - prompts.append(prompt)

309 - result_futures.append(result_future)

310 - if not can_batch:

311 - kwargs_list.append(kwargs)

312 - except asyncio.TimeoutError:

313 - break

314 -

315 - try:

316 - # If every request uses identical kwargs we can run a single

317 - # batched tokenizer call for a big speed-up.

318 - if can_batch and len(prompts) > 1:

319 - batch_encode_fn = partial(self.tokenizer, prompts, **kwargs)

320 - results = await self._loop.run_in_executor(

321 - self._executor, batch_encode_fn

322 - )

323 -

324 - for i, fut in enumerate(result_futures):

325 - if not fut.done():

326 - data = {k: v[i] for k, v in results.items()}

327 - fut.set_result(BatchEncoding(data))

328 - else:

329 - encode_fn = lambda prompts=prompts, kwargs=kwargs_list: [

330 - self.tokenizer(p, **kw) for p, kw in zip(prompts, kwargs)

331 - ]

332 - results = await self._loop.run_in_executor(

333 - self._executor, encode_fn

334 - )

335 -

336 - for fut, res in zip(result_futures, results):

337 - if not fut.done():

338 - fut.set_result(res)

339 - except Exception as e:

340 - for fut in result_futures:

341 - if not fut.done():

342 - fut.set_exception(e)

343 -

344 - async def _batch_decode_loop(self, queue: asyncio.Queue):

345 - """Batch incoming decode requests for efficiency."""

346 - while True:

347 - token_ids, result_future = await queue.get()

348 - token_ids_list = [token_ids]

349 - result_futures = [result_future]

350 - deadline = self._loop.time() + self.batch_wait_timeout_s

351 -

352 - while len(token_ids_list) < self.max_batch_size:

353 - timeout = deadline - self._loop.time()

354 - if timeout <= 0:

355 - break

356 - try:

357 - token_ids, result_future = await asyncio.wait_for(

358 - queue.get(), timeout

359 - )

360 - token_ids_list.append(token_ids)

361 - result_futures.append(result_future)

362 - except asyncio.TimeoutError:

363 - break

364 -

365 - try:

366 - # Perform a single batched decode call for all requests

367 - results = await self._loop.run_in_executor(

368 - self._executor, self.tokenizer.batch_decode, token_ids_list

369 - )

370 - for fut, res in zip(result_futures, results):

371 - if not fut.done():

372 - fut.set_result(res)

373 - except Exception as e:

374 - for fut in result_futures:

375 - if not fut.done():

376 - fut.set_exception(e)

377 -

378 - def _queue_key(self, op: str, kwargs: dict) -> tuple:

379 - """

380 - Return a normalized key describing operation + kwargs.

381 -

382 - - `add_special_tokens`: {True/False}

383 - - `truncation`: {True/False}

384 - - If `truncation` is False (`max_length` is None),

385 - returns a key for a can_batch queue.

386 - - If `truncation` is True and `max_length` is None or equals

387 - `tokenizer.model_max_length`, returns a key for a can_batch queue.

388 - - Otherwise, returns a key for a cannot_batch queue.

389 -

390 - Examples:

391 - - Decode: ("decode",)

392 - - Encode typical:

393 - ("encode", add_special_tokens, bool_truncation, max_length_label)

394 - - Fallback: ("encode", "other")

395 - """

396 -

397 - if op == "decode":

398 - return ("decode",)

399 -

400 - add_special_tokens = kwargs.get("add_special_tokens", True)

401 - truncation = kwargs.get("truncation", False)

402 - max_length = kwargs.get("max_length")

403 -

404 - if not truncation:

405 - return "encode", add_special_tokens, False, None

406 -

407 - model_max = getattr(self.tokenizer, "model_max_length", None)

408 - if max_length is None or (model_max is not None and max_length == model_max):

409 - return "encode", add_special_tokens, True, "model_max"

410 -

411 - return "encode", "other"

412 -

413 - def __del__(self):

414 - if (

415 - (tasks := getattr(self, "_batcher_tasks", None))

416 - and (loop := getattr(self, "_loop", None))

417 - and not loop.is_closed()

418 - ):

419 -

420 - def cancel_tasks():

421 - for task in tasks:

422 - task.cancel()

423 -

424 - loop.call_soon_threadsafe(cancel_tasks)

425 -

426 -

427 -def cancel_task_threadsafe(task: Task):

428 - if task and not task.done():

429 - run_in_loop(task.get_loop(), task.cancel)

430 -

431 -

221 def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):

222 for sock in sockets:

223 if sock is not None:

224 sock.close(linger=0)

225

226

438 -def run_in_loop(loop: AbstractEventLoop, function: Callable, *args):

439 - if in_loop(loop):

440 - function(*args)

441 - elif not loop.is_closed():

442 - loop.call_soon_threadsafe(function, *args)

443 -

444 -

445 -def in_loop(event_loop: AbstractEventLoop) -> bool:

446 - try:

447 - return asyncio.get_running_loop() == event_loop

448 - except RuntimeError:

449 - return False

450 -

451 -

452 -async def merge_async_iterators(

453 - *iterators: AsyncGenerator[T, None],

454 -) -> AsyncGenerator[tuple[int, T], None]:

455 - """Merge multiple asynchronous iterators into a single iterator.

456 -

457 - This method handle the case where some iterators finish before others.

458 - When it yields, it yields a tuple (i, item) where i is the index of the

459 - iterator that yields the item.

460 - """

461 - if len(iterators) == 1:

462 - # Fast-path single iterator case.

463 - async for item in iterators[0]:

464 - yield 0, item

465 - return

466 -

467 - loop = asyncio.get_running_loop()

468 -

469 - awaits = {loop.create_task(anext(it)): (i, it) for i, it in enumerate(iterators)}

470 - try:

471 - while awaits:

472 - done, _ = await asyncio.wait(awaits.keys(), return_when=FIRST_COMPLETED)

473 - for d in done:

474 - pair = awaits.pop(d)

475 - try:

476 - item = await d

477 - i, it = pair

478 - awaits[loop.create_task(anext(it))] = pair

479 - yield i, item

480 - except StopAsyncIteration:

481 - pass

482 - finally:

483 - # Cancel any remaining iterators

484 - for f, (_, it) in awaits.items():

485 - with contextlib.suppress(BaseException):

486 - f.cancel()

487 - await it.aclose()

488 -

489 -

490 -async def collect_from_async_generator(iterator: AsyncGenerator[T, None]) -> list[T]:

491 - """Collect all items from an async generator into a list."""

492 - items = []

493 - async for item in iterator:

494 - items.append(item)

495 - return items

496 -

497 -

227 def get_ip() -> str:

228 host_ip = envs.VLLM_HOST_IP

229 if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ:

@@ -1803,12 +1532,6 @@ class FlexibleArgumentParser(ArgumentParser):

1532 return processed_args

1533

1534

1806 -async def _run_task_with_lock(task: Callable, lock: asyncio.Lock, *args, **kwargs):

1807 - """Utility function to run async task in a lock"""

1808 - async with lock:

1809 - return await task(*args, **kwargs)

1810 -

1811 -

1535 # Using dynamo with vLLM doesn't really work well with PyTorch versions < 2.4.0.

1536 # In particular, the FakeScalarType is not supported for earlier versions of

1537 # PyTorch which breaks dynamo for any ops registered using ScalarType.

vllm/utils/async_utils.py new

+299

@@ -0,0 +1,299 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Contains helpers related to asynchronous code."""

4 +

5 +import asyncio

6 +import contextlib

7 +from asyncio import FIRST_COMPLETED, AbstractEventLoop, Future, Task

8 +from collections.abc import AsyncGenerator, Awaitable, Callable

9 +from concurrent.futures import Executor, ThreadPoolExecutor

10 +from functools import partial

11 +from typing import TypeVar

12 +

13 +from transformers.tokenization_utils_base import BatchEncoding

14 +from typing_extensions import ParamSpec

15 +

16 +P = ParamSpec("P")

17 +T = TypeVar("T")

18 +

19 +

20 +class AsyncMicrobatchTokenizer:

21 + """Asynchronous tokenizer with micro-batching.

22 +

23 + Pulls pending encode/decode requests from a queue and batches them

24 + up to reduce overhead. A single-thread ThreadPoolExecutor is used

25 + so the event loop stays responsive.

26 + """

27 +

28 + def __init__(

29 + self,

30 + tokenizer,

31 + max_batch_size: int = 32,

32 + batch_wait_timeout_s: float = 0.002,

33 + ) -> None:

34 + self.tokenizer = tokenizer

35 + self.max_batch_size = max_batch_size

36 + self.batch_wait_timeout_s = batch_wait_timeout_s

37 +

38 + self._loop = asyncio.get_running_loop()

39 + self._queues: dict[

40 + tuple,

41 + asyncio.Queue[tuple[str, dict, Future] | tuple[list[int], Future]],

42 + ] = {}

43 + self._batcher_tasks: list[Task] = []

44 +

45 + # Single-thread executor for blocking tokenizer calls.

46 + self._executor = ThreadPoolExecutor(max_workers=1)

47 +

48 + # === Public async API ===

49 + async def __call__(self, prompt, **kwargs):

50 + result_future: Future = self._loop.create_future()

51 + key = self._queue_key("encode", kwargs)

52 + queue = self._get_queue(self._loop, key)

53 + await queue.put((prompt, kwargs, result_future))

54 + return await result_future

55 +

56 + async def decode(self, token_ids, **kwargs):

57 + result_future: Future = self._loop.create_future()

58 + key = self._queue_key("decode", kwargs)

59 + queue = self._get_queue(self._loop, key)

60 + await queue.put((token_ids, result_future))

61 + return await result_future

62 +

63 + # === Internal helpers ===

64 + def _get_queue(

65 + self, loop: asyncio.AbstractEventLoop, key: tuple

66 + ) -> asyncio.Queue[tuple[str, dict, Future] | tuple[list[int], Future]]:

67 + """Get the request queue for the given operation key, creating a new

68 + queue and batcher task if needed."""

69 + queue = self._queues.get(key)

70 + if queue is None:

71 + self._queues[key] = queue = asyncio.Queue()

72 + if key[0] == "encode":

73 + can_batch = key[1] != "other"

74 + coro = self._batch_encode_loop(queue, can_batch)

75 + else:

76 + assert key[0] == "decode", f"Unknown operation type: {key[0]}."

77 + coro = self._batch_decode_loop(queue)

78 + self._batcher_tasks.append(loop.create_task(coro))

79 + return queue

80 +

81 + async def _batch_encode_loop(self, queue: asyncio.Queue, can_batch: bool):

82 + """Batch incoming encode requests for efficiency."""

83 + while True:

84 + prompt, kwargs, result_future = await queue.get()

85 + prompts = [prompt]

86 + kwargs_list = [kwargs]

87 + result_futures = [result_future]

88 + deadline = self._loop.time() + self.batch_wait_timeout_s

89 +

90 + while len(prompts) < self.max_batch_size:

91 + timeout = deadline - self._loop.time()

92 + if timeout <= 0:

93 + break

94 + try:

95 + prompt, kwargs, result_future = await asyncio.wait_for(

96 + queue.get(), timeout

97 + )

98 + prompts.append(prompt)

99 + result_futures.append(result_future)

100 + if not can_batch:

101 + kwargs_list.append(kwargs)

102 + except asyncio.TimeoutError:

103 + break

104 +

105 + try:

106 + # If every request uses identical kwargs we can run a single

107 + # batched tokenizer call for a big speed-up.

108 + if can_batch and len(prompts) > 1:

109 + batch_encode_fn = partial(self.tokenizer, prompts, **kwargs)

110 + results = await self._loop.run_in_executor(

111 + self._executor, batch_encode_fn

112 + )

113 +

114 + for i, fut in enumerate(result_futures):

115 + if not fut.done():

116 + data = {k: v[i] for k, v in results.items()}

117 + fut.set_result(BatchEncoding(data))

118 + else:

119 + encode_fn = lambda prompts=prompts, kwargs=kwargs_list: [

120 + self.tokenizer(p, **kw) for p, kw in zip(prompts, kwargs)

121 + ]

122 + results = await self._loop.run_in_executor(

123 + self._executor, encode_fn

124 + )

125 +

126 + for fut, res in zip(result_futures, results):

127 + if not fut.done():

128 + fut.set_result(res)

129 + except Exception as e:

130 + for fut in result_futures:

131 + if not fut.done():

132 + fut.set_exception(e)

133 +

134 + async def _batch_decode_loop(self, queue: asyncio.Queue):

135 + """Batch incoming decode requests for efficiency."""

136 + while True:

137 + token_ids, result_future = await queue.get()

138 + token_ids_list = [token_ids]

139 + result_futures = [result_future]

140 + deadline = self._loop.time() + self.batch_wait_timeout_s

141 +

142 + while len(token_ids_list) < self.max_batch_size:

143 + timeout = deadline - self._loop.time()

144 + if timeout <= 0:

145 + break

146 + try:

147 + token_ids, result_future = await asyncio.wait_for(

148 + queue.get(), timeout

149 + )

150 + token_ids_list.append(token_ids)

151 + result_futures.append(result_future)

152 + except asyncio.TimeoutError:

153 + break

154 +

155 + try:

156 + # Perform a single batched decode call for all requests

157 + results = await self._loop.run_in_executor(

158 + self._executor, self.tokenizer.batch_decode, token_ids_list

159 + )

160 + for fut, res in zip(result_futures, results):

161 + if not fut.done():

162 + fut.set_result(res)

163 + except Exception as e:

164 + for fut in result_futures:

165 + if not fut.done():

166 + fut.set_exception(e)

167 +

168 + def _queue_key(self, op: str, kwargs: dict) -> tuple:

169 + """

170 + Return a normalized key describing operation + kwargs.

171 +

172 + - `add_special_tokens`: {True/False}

173 + - `truncation`: {True/False}

174 + - If `truncation` is False (`max_length` is None),

175 + returns a key for a can_batch queue.

176 + - If `truncation` is True and `max_length` is None or equals

177 + `tokenizer.model_max_length`, returns a key for a can_batch queue.

178 + - Otherwise, returns a key for a cannot_batch queue.

179 +

180 + Examples:

181 + - Decode: ("decode",)

182 + - Encode typical:

183 + ("encode", add_special_tokens, bool_truncation, max_length_label)

184 + - Fallback: ("encode", "other")

185 + """

186 +

187 + if op == "decode":

188 + return ("decode",)

189 +

190 + add_special_tokens = kwargs.get("add_special_tokens", True)

191 + truncation = kwargs.get("truncation", False)

192 + max_length = kwargs.get("max_length")

193 +

194 + if not truncation:

195 + return "encode", add_special_tokens, False, None

196 +

197 + model_max = getattr(self.tokenizer, "model_max_length", None)

198 + if max_length is None or (model_max is not None and max_length == model_max):

199 + return "encode", add_special_tokens, True, "model_max"

200 +

201 + return "encode", "other"

202 +

203 + def __del__(self):

204 + if (

205 + (tasks := getattr(self, "_batcher_tasks", None))

206 + and (loop := getattr(self, "_loop", None))

207 + and not loop.is_closed()

208 + ):

209 +

210 + def cancel_tasks():

211 + for task in tasks:

212 + task.cancel()

213 +

214 + loop.call_soon_threadsafe(cancel_tasks)

215 +

216 +

217 +def cancel_task_threadsafe(task: Task):

218 + if task and not task.done():

219 + run_in_loop(task.get_loop(), task.cancel)

220 +

221 +

222 +def make_async(

223 + func: Callable[P, T],

224 + executor: Executor | None = None,

225 +) -> Callable[P, Awaitable[T]]:

226 + """

227 + Take a blocking function, and run it on in an executor thread.

228 +

229 + This function prevents the blocking function from blocking the

230 + asyncio event loop.

231 + The code in this function needs to be thread safe.

232 + """

233 +

234 + def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Future[T]:

235 + loop = asyncio.get_event_loop()

236 + p_func = partial(func, *args, **kwargs)

237 + return loop.run_in_executor(executor=executor, func=p_func)

238 +

239 + return _async_wrapper

240 +

241 +

242 +def run_in_loop(loop: AbstractEventLoop, function: Callable, *args):

243 + if in_loop(loop):

244 + function(*args)

245 + elif not loop.is_closed():

246 + loop.call_soon_threadsafe(function, *args)

247 +

248 +

249 +def in_loop(event_loop: AbstractEventLoop) -> bool:

250 + try:

251 + return asyncio.get_running_loop() == event_loop

252 + except RuntimeError:

253 + return False

254 +

255 +

256 +async def merge_async_iterators(

257 + *iterators: AsyncGenerator[T, None],

258 +) -> AsyncGenerator[tuple[int, T], None]:

259 + """Merge multiple asynchronous iterators into a single iterator.

260 +

261 + This method handle the case where some iterators finish before others.

262 + When it yields, it yields a tuple (i, item) where i is the index of the

263 + iterator that yields the item.

264 + """

265 + if len(iterators) == 1:

266 + # Fast-path single iterator case.

267 + async for item in iterators[0]:

268 + yield 0, item

269 + return

270 +

271 + loop = asyncio.get_running_loop()

272 +

273 + awaits = {loop.create_task(anext(it)): (i, it) for i, it in enumerate(iterators)}

274 + try:

275 + while awaits:

276 + done, _ = await asyncio.wait(awaits.keys(), return_when=FIRST_COMPLETED)

277 + for d in done:

278 + pair = awaits.pop(d)

279 + try:

280 + item = await d

281 + i, it = pair

282 + awaits[loop.create_task(anext(it))] = pair

283 + yield i, item

284 + except StopAsyncIteration:

285 + pass

286 + finally:

287 + # Cancel any remaining iterators

288 + for f, (_, it) in awaits.items():

289 + with contextlib.suppress(BaseException):

290 + f.cancel()

291 + await it.aclose()

292 +

293 +

294 +async def collect_from_async_generator(iterator: AsyncGenerator[T, None]) -> list[T]:

295 + """Collect all items from an async generator into a list."""

296 + items = []

297 + async for item in iterator:

298 + items.append(item)

299 + return items

vllm/utils/func.py

+1 -23

@@ -6,12 +6,10 @@ Contains helpers that are applied to functions.

6 This is similar in concept to the `functools` module.

7 """

9 -import asyncio

10 -import concurrent.futures

9 import inspect

10 import threading

11 import warnings

14 -from collections.abc import Awaitable, Callable, Mapping

12 +from collections.abc import Callable, Mapping

13 from functools import lru_cache, partial, wraps

14 from typing import Any, TypeVar

@@ -32,26 +30,6 @@ def identity(value: T, **kwargs) -> T:

30 return value

35 -def make_async(

36 - func: Callable[P, T],

37 - executor: concurrent.futures.Executor | None = None,

38 -) -> Callable[P, Awaitable[T]]:

39 - """

40 - Take a blocking function, and run it on in an executor thread.

41 -

42 - This function prevents the blocking function from blocking the

43 - asyncio event loop.

44 - The code in this function needs to be thread safe.

45 - """

46 -

47 - def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> asyncio.Future[T]:

48 - loop = asyncio.get_event_loop()

49 - p_func = partial(func, *args, **kwargs)

50 - return loop.run_in_executor(executor=executor, func=p_func)

51 -

52 - return _async_wrapper

53 -

54 -

33 def run_once(f: Callable[P, None]) -> Callable[P, None]:

34 def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:

35 if wrapper.has_run: # type: ignore[attr-defined]

vllm/v1/engine/async_llm.py

+2 -1

@@ -29,7 +29,8 @@ from vllm.tracing import init_tracer

29 from vllm.transformers_utils.config import maybe_register_config_serialize_by_value

30 from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs

31 from vllm.usage.usage_lib import UsageContext

32 -from vllm.utils import Device, as_list, cancel_task_threadsafe, cdiv

32 +from vllm.utils import Device, as_list, cdiv

33 +from vllm.utils.async_utils import cancel_task_threadsafe

34 from vllm.utils.func import deprecate_kwargs

35 from vllm.v1.engine import EngineCoreRequest

36 from vllm.v1.engine.core_client import EngineCoreClient

vllm/v1/engine/core_client.py

+1 -1

@@ -27,9 +27,9 @@ from vllm.utils import (

27 close_sockets,

28 get_open_port,

29 get_open_zmq_inproc_path,

30 - in_loop,

30 make_zmq_socket,

31 )

32 +from vllm.utils.async_utils import in_loop

33 from vllm.v1.engine import (

34 EngineCoreOutputs,

35 EngineCoreRequest,