diff --git a/rust/Cargo.toml b/rust/Cargo.toml index f3d3053ed5..0002796d40 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -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"] } diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 95c4332f74..86b94b6b9c 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -180,7 +180,12 @@ impl ChatRequestProcessor { async fn prepare_text_request(&self, request: ChatRequest) -> Result { // 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 { diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index efcfd420ed..3a496d957d 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -29,6 +29,11 @@ use crate::renderer::hf::{TemplateMessage, TemplateTool}; type Result = std::result::Result; +/// 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; + /// Build a pre-configured environment with the given template string. fn build_environment(template: String) -> Result> { let mut env = Environment::new(); @@ -36,6 +41,8 @@ fn build_environment(template: String) -> Result> { 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); diff --git a/tests/renderers/test_executor_replace.py b/tests/renderers/test_executor_replace.py new file mode 100644 index 0000000000..fc0175e1e2 --- /dev/null +++ b/tests/renderers/test_executor_replace.py @@ -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() diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 34f6fb44a2..66e0f2640d 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -334,6 +334,16 @@ class OpenAIServingResponses(GenerateBaseServing): 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: diff --git a/vllm/envs.py b/vllm/envs.py index 67e0207492..5daa4fd3e6 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: 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 @@ -812,6 +813,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "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 diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 87d11b7b92..0936dd021e 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -5,7 +5,7 @@ import time 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 @@ -77,6 +77,31 @@ logger = init_logger(__name__) _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__() @@ -115,7 +140,7 @@ class BaseRenderer(ABC, Generic[_T]): # 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 diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 11f562f74f..4301c69176 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations +import asyncio import copy import inspect import itertools @@ -20,6 +21,7 @@ import jinja2.sandbox import torch from typing_extensions import override +from vllm import envs from vllm.entrypoints.chat_utils import ( PROMPT_EMBEDS_PLACEHOLDER_TOKEN, ChatTemplateResolutionError, @@ -969,6 +971,44 @@ class HfRenderer(BaseRenderer[HfTokenizer]): 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 @@ -1124,11 +1164,13 @@ class HfRenderer(BaseRenderer[HfTokenizer]): 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