Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion rust/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ itertools = "0.14.0"
libc = "0.2.177"
llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "f0985ef65967615db2c79279aa07818499301bfd", default-features = false, features = ["native-tls"] }
mimalloc = "0.1.52"
minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] }
minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order", "fuel"] }
minijinja-contrib = { version = "2.24", features = ["pycompat"] }
native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] }
ndarray = { version = "0.17", features = ["serde"] }
Expand Down
7 changes: 6 additions & 1 deletion rust/src/chat/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,12 @@ impl ChatRequestProcessor {
async fn prepare_text_request(&self, request: ChatRequest) -> Result<TextRequest> {
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
let arrival_time = vllm_llm::current_unix_timestamp_secs();
let mut rendered = self.backend.chat_renderer().render(&request)?;
let renderer = self.backend.chat_renderer();
let (mut rendered, request) = tokio::task::spawn_blocking(move || {
renderer.render(&request).map(|rendered| (rendered, request))
})
.await
.map_err(|e| Error::ChatTemplate(format!("render task join error: {e}")))??;
let chat_template_kwargs = std::mem::take(&mut rendered.effective_template_kwargs);
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
Ok(TextRequest {
Expand Down
7 changes: 7 additions & 0 deletions rust/src/chat/src/renderer/hf/template.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,20 @@ use crate::renderer::hf::{TemplateMessage, TemplateTool};

type Result<T> = std::result::Result<T, TemplateError>;

/// Default fuel budget for template rendering. Prevents unbounded loops and
/// string operations from exhausting CPU. Each template instruction consumes
/// one unit of fuel; when depleted, rendering aborts with an error.
const DEFAULT_TEMPLATE_FUEL: u64 = 500_000;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

PTAL @BugenZhao


/// Build a pre-configured environment with the given template string.
fn build_environment(template: String) -> Result<Environment<'static>> {
let mut env = Environment::new();

env.set_trim_blocks(true);
env.set_lstrip_blocks(true);

env.set_fuel(Some(DEFAULT_TEMPLATE_FUEL));

env.add_template_owned("chat".to_owned(), template)?;

env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
Expand Down
126 changes: 126 additions & 0 deletions tests/renderers/test_executor_replace.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import asyncio
from dataclasses import dataclass
from typing import Any

import pytest
import torch

from vllm.renderers.base import _SwappableExecutor
from vllm.renderers.hf import HfRenderer
from vllm.renderers.params import TokenizeParams
from vllm.utils.async_utils import make_async

MODEL_NAME = "openai-community/gpt2"


@dataclass
class MockHFConfig:
model_type: str = "any"


@dataclass
class MockModelConfig:
runner_type = "generate"
model: str = MODEL_NAME
tokenizer: str = MODEL_NAME
trust_remote_code: bool = False
tokenizer_revision = None
tokenizer_mode = "auto"
hf_config = MockHFConfig()
encoder_config: dict[str, Any] | None = None
enable_prompt_embeds: bool = False
skip_tokenizer_init: bool = False
is_encoder_decoder: bool = False
is_multimodal_model: bool = False
renderer_num_workers: int = 1
hidden_size: int = 768
dtype: torch.dtype = torch.float32

def get_hidden_size(self) -> int:
return self.hidden_size


@dataclass
class MockParallelConfig:
_api_process_rank: int = 0


@dataclass
class MockVllmConfig:
model_config: MockModelConfig
parallel_config: MockParallelConfig


@dataclass
class DummyTokenizer:
truncation_side: str = "left"
max_chars_per_token: int = 1
is_fast: bool = False

def decode(self, tokens: list[int], **kwargs):
return str(tokens)

def encode(self, text: str, **kwargs):
return list(range(len(text)))

def __call__(self, text: str, **kwargs):
return {"input_ids": self.encode(text, **kwargs)}


def _build_renderer() -> HfRenderer:
return HfRenderer(
MockVllmConfig(MockModelConfig(), parallel_config=MockParallelConfig()),
tokenizer=DummyTokenizer(),
)


def test_swappable_executor_keeps_make_async_wrappers_alive():
pool = _SwappableExecutor(max_workers=1)
async_add = make_async(lambda x: x + 1, executor=pool)

async def _run():
assert await async_add(1) == 2
old_inner = pool._inner
pool.replace_inner()

with pytest.raises(RuntimeError, match="cannot schedule new futures"):
old_inner.submit(lambda: None)

assert await async_add(40) == 41

try:
asyncio.run(_run())
finally:
pool.shutdown(wait=False)


def test_replace_executor_does_not_break_tokenize_or_decode():
renderer = _build_renderer()
executor = renderer._executor
old_inner = executor._inner

async def _run():
assert await renderer._tokenize_prompt_async(
{"prompt": "ab"},
TokenizeParams(max_total_tokens=100),
)
renderer._replace_executor()
assert renderer._executor is executor
assert executor._inner is not old_inner

with pytest.raises(RuntimeError, match="cannot schedule new futures"):
old_inner.submit(lambda: None)

tokenized = await renderer._tokenize_prompt_async(
{"prompt": "abc"},
TokenizeParams(max_total_tokens=100),
)
assert tokenized["prompt_token_ids"] == [0, 1, 2]
assert await renderer._async_tokenizer_decode([1, 2]) == "[1, 2]"

try:
asyncio.run(_run())
finally:
renderer.shutdown()
10 changes: 10 additions & 0 deletions vllm/entrypoints/openai/responses/serving.py
Original file line number Diff line number Diff line change
Expand Up @@ -334,6 +334,16 @@ async def _create_responses(
if maybe_validation_error is not None:
return maybe_validation_error

maybe_template_error = self.online_renderer.validate_chat_template(
request_chat_template=None,
chat_template_kwargs=request.chat_template_kwargs,
trust_request_chat_template=(
self.online_renderer.trust_request_chat_template
),
)
if maybe_template_error is not None:
return maybe_template_error

self._preflight()

if request.store and not self.enable_store:
Expand Down
6 changes: 6 additions & 0 deletions vllm/envs.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
CUDA_VISIBLE_DEVICES: str | None = None
VLLM_ENGINE_ITERATION_TIMEOUT_S: int = 60
VLLM_ENGINE_READY_TIMEOUT_S: int = 600
VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT: float = 30.0
VLLM_API_KEY: str | None = None
VLLM_DEBUG_LOG_API_SERVER_RESPONSE: bool = False
S3_ACCESS_KEY_ID: str | None = None
Expand Down Expand Up @@ -812,6 +813,11 @@ def _resolve_rust_cli_path() -> str | None:
"VLLM_ENGINE_READY_TIMEOUT_S": lambda: int(
os.environ.get("VLLM_ENGINE_READY_TIMEOUT_S", "600")
),
# Maximum wall-clock seconds allowed for a single chat template render.
# Set to 0 to disable the timeout.
"VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT": lambda: float(
os.environ.get("VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT", "30")
),
# API key for vLLM API server
"VLLM_API_KEY": lambda: os.environ.get("VLLM_API_KEY", None),
# Whether to log responses from API Server for debugging
Expand Down
29 changes: 27 additions & 2 deletions vllm/renderers/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import weakref
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from concurrent.futures import Future, ThreadPoolExecutor
from concurrent.futures import Executor, Future, ThreadPoolExecutor
from contextlib import ExitStack
from dataclasses import replace
from functools import cached_property
Expand Down Expand Up @@ -77,6 +77,31 @@
_T = TypeVar("_T", bound=TokenizerLike, default=TokenizerLike)


class _SwappableExecutor(Executor):
"""Executor whose inner pool can be replaced without changing identity.

``make_async`` captures the executor object at wrap time. Replacing
``self._executor`` with a new ``ThreadPoolExecutor`` would leave those
wrappers (tokenize, decode, pooling, derender) bound to a shutdown pool.
"""

def __init__(self, max_workers: int) -> None:
super().__init__()
self._max_workers = max_workers
self._inner = ThreadPoolExecutor(max_workers=max_workers)

def submit(self, fn, /, *args, **kwargs):
return self._inner.submit(fn, *args, **kwargs)

def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
self._inner.shutdown(wait=wait, cancel_futures=cancel_futures)

def replace_inner(self) -> None:
old = self._inner
old.shutdown(wait=False)
self._inner = ThreadPoolExecutor(max_workers=self._max_workers)


class BaseRenderer(ABC, Generic[_T]):
def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
super().__init__()
Expand Down Expand Up @@ -115,7 +140,7 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
# multimodal processor receives a deep-copied tokenizer (see #36557)
# so it is safe to run tokenization and MM preprocessing concurrently.
pool_workers = config.model_config.renderer_num_workers
self._executor = ThreadPoolExecutor(max_workers=pool_workers)
self._executor = _SwappableExecutor(max_workers=pool_workers)
self._resources.callback(self._executor.shutdown, wait=False)

# Separate single-worker executor so tokenization never queues behind
Expand Down
52 changes: 47 additions & 5 deletions vllm/renderers/hf.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations

import asyncio
import copy
import inspect
import itertools
Expand All @@ -20,6 +21,7 @@
import torch
from typing_extensions import override

from vllm import envs
from vllm.entrypoints.chat_utils import (
PROMPT_EMBEDS_PLACEHOLDER_TOKEN,
ChatTemplateResolutionError,
Expand Down Expand Up @@ -969,6 +971,44 @@ def __init__(
self.tokenizer, config.model_config.renderer_num_workers + 1
)

def _replace_executor(self) -> None:
"""Replace the inner render pool after a stuck worker.

Called after a render timeout so that the stuck thread (which CPython
cannot interrupt) does not permanently block all subsequent renders.
The executor *object* is kept so ``make_async`` wrappers bound at
init (tokenize, decode, pooling, derender) keep submitting work.
The old inner pool is abandoned with wait=False — its thread will
eventually finish or be reaped on process exit.
"""
self._executor.replace_inner()
logger.warning(
"Chat template render timed out — executor pool replaced "
"(%d worker(s)). The stuck thread may continue consuming "
"CPU until the process exits.",
self.model_config.renderer_num_workers,
)

async def _render_with_timeout(self, coro):
"""Run a template-rendering coroutine with a timeout.

On timeout, replaces the executor so subsequent renders are not
permanently blocked by the stuck thread.
"""
timeout = envs.VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT
if timeout <= 0:
return await coro
try:
return await asyncio.wait_for(coro, timeout=timeout)
except asyncio.TimeoutError:
self._replace_executor()
raise TimeoutError(
f"Chat template rendering timed out after "
f"{timeout}s. This may indicate a malicious or "
f"excessively complex template. Adjust "
f"VLLM_CHAT_TEMPLATE_RENDER_TIMEOUT to override."
) from None

def _can_produce_offsets(self) -> bool:
# Only fast tokenizers expose offset_mapping.
return self.tokenizer is not None and self.tokenizer.is_fast
Expand Down Expand Up @@ -1124,11 +1164,13 @@ async def render_messages_async(
logger.warning_once(_TOKENIZE_OVERRIDE_WARNING)
chat_template_kwargs["tokenize"] = True

prompt_raw = await self._apply_chat_template_async(
model_config,
tokenizer,
conversation,
**chat_template_kwargs,
prompt_raw = await self._render_with_timeout(
self._apply_chat_template_async(
model_config,
tokenizer,
conversation,
**chat_template_kwargs,
)
)

# NOTE: use_unified_vision_chunk is currently specific to Kimi-K2.5
Expand Down
Loading