diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py index 0330b9d509c6..4d49dd8038d6 100644 --- a/tensorrt_llm/inputs/__init__.py +++ b/tensorrt_llm/inputs/__init__.py @@ -22,6 +22,7 @@ ALL_SUPPORTED_MULTIMODAL_MODELS, ALL_SUPPORTED_VIDEO_MODELS, ConversationMessage, MultimodalData, MultimodalDataTracker, add_multimodal_placeholders, apply_chat_template, + apply_mm_placeholders, async_build_multimodal_prompt, async_load_audio, async_load_image, async_load_video, convert_image_mode, default_multimodal_input_loader, encode_base64_content_from_url, encode_base64_image, @@ -63,6 +64,8 @@ "async_load_image", "async_load_video", "add_multimodal_placeholders", + "apply_mm_placeholders", + "async_build_multimodal_prompt", "apply_chat_template", "convert_image_mode", "default_multimodal_input_loader", diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 8a6d406604be..d7395d0ca86f 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -20,7 +20,9 @@ from tensorrt_llm.inputs.content_format import (ContentFormat, detect_content_format) -from tensorrt_llm.inputs.media_io import (_get_aiohttp_session, +from tensorrt_llm.inputs.data import prompt_inputs +from tensorrt_llm.inputs.media_io import (MEDIA_IO_REGISTRY, + _get_aiohttp_session, _load_and_convert_image, _load_video_by_cv2, _normalize_file_uri, @@ -820,6 +822,132 @@ async def async_apply_chat_template( ) +def apply_mm_placeholders( + model_type: str, + message: ConversationMessage, + placeholder_counts: Dict[str, int], + mm_data_tracker: "MultimodalDataTracker", + *, + item_order_start: int = 0, + content_format: Optional[ContentFormat] = None, +) -> None: + """Insert this message's multimodal placeholders into its text, in place. + + Only `ContentFormat.STRING` templates need them. For `OPENAI`, + `apply_chat_template` rebuilds `content` from `content_parts` via + `_build_openai_content`, so pre-inserting here would render every media + item twice: once as text and once as the template's own content part. + + When the model opts into interleaving and `content_parts` is available the + placeholders keep their original positions; otherwise they are placed in + bulk according to the registered placement. + """ + if not placeholder_counts: + return + + if content_format is None: + registry_format = MULTIMODAL_PLACEHOLDER_REGISTRY.get_content_format( + model_type) + content_format = (registry_format if registry_format is not None else + ContentFormat.STRING) + if content_format != ContentFormat.STRING: + return + + content_parts = message.get("content_parts") + interleave = MULTIMODAL_PLACEHOLDER_REGISTRY.get_interleave_placeholders( + model_type) + if content_parts and interleave: + message["content"] = interleave_mm_placeholders( + model_type, content_parts, placeholder_counts, + mm_data_tracker.placeholder_modalities()) + else: + message["content"] = add_multimodal_placeholders( + model_type, + message["content"], + placeholder_counts, + item_order=mm_data_tracker.item_order()[item_order_start:], + ) + + +async def async_build_multimodal_prompt( + *, + model_type: str, + tokenizer: Union[TransformersTokenizer, TokenizerBase], + processor: ProcessorMixin, + prompt: str, + media: List[str], + modality: str = "image", + add_generation_prompt: bool = True, + chat_template: Optional[str] = None, + chat_template_kwargs: Optional[Dict[str, Any]] = None, + media_io_kwargs: Optional[Dict[str, Dict[str, Any]]] = None, +) -> Dict[str, Any]: + """Package one text prompt plus media URLs into engine inputs. + + For callers whose request format is a plain prompt and a list of media + references rather than OpenAI chat messages -- the Triton `llmapi` backend, + for example. Runs the same steps `trtllm-serve` runs for + `v1/chat/completions`: resolve placeholders for the model's content format, + render the chat template, and attach the fetched media. + + The media fetches and the template rendering are awaited together, so + several items on one request are loaded concurrently. + + Args: + model_type: Top-level HF `model_type`, the multimodal registry key. + prompt: The user's text, without any placeholder tokens. + media: URLs, `data:` URIs or local paths, one per item. + modality: Registered media modality, e.g. `"image"`. + + Returns: + A `TextPrompt` carrying `multi_modal_data`, ready for + `LLM.generate_async`. + """ + media_io_cls = MEDIA_IO_REGISTRY.get(modality) + if media_io_cls is None: + raise ValueError(f"Unsupported modality {modality!r}. " + f"Registered modalities: {list(MEDIA_IO_REGISTRY)}") + media_io = media_io_cls.create((media_io_kwargs or {}).get(modality), None) + + mm_data_tracker = MultimodalDataTracker(model_type) + content_parts: List[Union[str, Dict[str, Any]]] = [] + for index, item in enumerate(media): + # `add_data` takes the un-awaited fetch; they run concurrently below. + mm_data_tracker.add_data(modality, media_io.async_load(item)) + content_parts.append({"type": modality, "media_index": index}) + content_parts.append(prompt) + + message = ConversationMessage(role="user", + content=prompt, + media=[], + content_parts=content_parts) + placeholder_counts = mm_data_tracker.placeholder_counts() + apply_mm_placeholders(model_type, message, placeholder_counts, + mm_data_tracker) + + prompt_task = async_apply_chat_template( + model_type=model_type, + tokenizer=tokenizer, + processor=processor, + conversation=[message], + add_generation_prompt=add_generation_prompt, + mm_placeholder_counts=[placeholder_counts], + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + ) + rendered, (mm_data, + _) = await asyncio.gather(prompt_task, + mm_data_tracker.retrieve_all_async()) + + inputs = prompt_inputs(rendered) + if mm_data: + inputs["multi_modal_data"] = mm_data + item_order = mm_data_tracker.item_order() + if item_order: + inputs["mm_item_order"] = item_order + return inputs + + def default_multimodal_input_loader( *, tokenizer: Optional[Union[TransformersTokenizer, TokenizerBase]], diff --git a/tensorrt_llm/serve/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index d60d93f2ce45..c1f56b98af01 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -17,12 +17,11 @@ from tensorrt_llm.inputs import (ContentFormat, ConversationMessage, MultimodalData, MultimodalDataTracker, - add_multimodal_placeholders, load_base64_image_embeds) from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY, BaseMediaIO from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY -from tensorrt_llm.inputs.utils import interleave_mm_placeholders +from tensorrt_llm.inputs.utils import apply_mm_placeholders from tensorrt_llm.logger import logger @@ -497,25 +496,12 @@ def parse_chat_messages_coroutines( placeholder] = msg_placeholder_counts.get( placeholder, 0) + 1 - if msg_placeholder_counts and content_format == ContentFormat.STRING: - # For STRING format, use interleaving when the model opts in - # and content_parts is available, otherwise fall back to bulk - # prepend/append according to placeholder_placement. - content_parts = parsed_msg.get("content_parts") - interleave = MULTIMODAL_PLACEHOLDER_REGISTRY.get_interleave_placeholders( - model_type) - if content_parts and interleave: - parsed_msg["content"] = interleave_mm_placeholders( - model_type, content_parts, msg_placeholder_counts, - mm_data_tracker.placeholder_modalities()) - else: - msg_item_order = mm_data_tracker.item_order()[item_order_start:] - parsed_msg["content"] = add_multimodal_placeholders( - type(model_config).model_type, - parsed_msg["content"], - msg_placeholder_counts, - item_order=msg_item_order, - ) + apply_mm_placeholders(model_type, + parsed_msg, + msg_placeholder_counts, + mm_data_tracker, + item_order_start=item_order_start, + content_format=content_format) mm_placeholder_counts.append(msg_placeholder_counts) # ``item_order`` is populated synchronously by ``add_data``, so it can diff --git a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py index edd77e1a26a7..c2dd22455053 100755 --- a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py +++ b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.py @@ -35,6 +35,7 @@ from dataclasses import dataclass from random import randint from typing import Any +from urllib.parse import urlparse import numpy as np import pandas as pd @@ -85,6 +86,19 @@ def get_model_config(filename, include_keys=None, exclude_keys=None): return engine_config +# `image_url` is client-controlled, so restrict it to web URLs for security. +ALLOWED_MEDIA_SCHEMES = ("http", "https") + + +def validate_media_urls(urls): + """Reject anything that is not a web URL.""" + for url in urls: + if urlparse(url).scheme not in ALLOWED_MEDIA_SCHEMES: + raise pb_utils.TritonModelException( + f"Unsupported image_url {url!r}: only " + f"{', '.join(ALLOWED_MEDIA_SCHEMES)} URLs are accepted.") + + def get_input_scalar_by_name(request, name, expected_batch_size=1, @@ -251,6 +265,11 @@ def initialize(self, args): self._response_thread = threading.Thread(target=self._response_loop) self._response_thread.start() + self.multimodal_enabled = bool( + triton_config.get("multimodal", False)) + if self.multimodal_enabled: + self._init_multimodal() + self.req_id_to_request_data = {} self.triton_user_id_to_req_ids = {} self.lock = threading.Lock() @@ -479,13 +498,14 @@ async def _execute_single_request(self, request): # Unique request id used to identify each triton request triton_req_id = str(randint(0, sys.maxsize)) + request_registered = False try: from tensorrt_llm import SamplingParams # TODO: [JIRA-4496] Implement when request contains batched prompts (prompt, sampling_params, streaming, output_config, - lora_request) = self._convert_request(request) + lora_request) = await self._convert_request(request) if streaming and not self.decoupled: raise pb_utils.TritonModelException( "Streaming is only supported in decoupled mode.") @@ -508,6 +528,7 @@ async def _execute_single_request(self, request): # TODO: [JIRA-4496] Add all batched request ids to the set self.triton_user_id_to_req_ids[triton_user_id].add( triton_req_id) + request_registered = True async for request_output in response_iterator: # Send each response if streaming. @@ -560,7 +581,7 @@ async def _execute_single_request(self, request): # already sent a COMPLETE_FINAL on this response_sender; they # remove the entry from req_id_to_request_data as the signal. with self.lock: - was_cancelled = (triton_req_id + was_cancelled = (request_registered and triton_req_id not in self.req_id_to_request_data) if not was_cancelled: error = pb_utils.TritonError(f"Error generating request: {e}") @@ -582,7 +603,45 @@ async def _execute_single_request(self, request): if triton_user_id is not None and triton_user_id != "" and triton_user_id in self.triton_user_id_to_req_ids: del self.triton_user_id_to_req_ids[triton_user_id] - def _convert_request(self, request): + def _init_multimodal(self): + """Resolve the tokenizer, HF processor and model type once.""" + from transformers import AutoProcessor + + from tensorrt_llm._torch.pyexecutor.config_utils import \ + load_pretrained_config + + self._mm_tokenizer = self._llm_engine.tokenizer + hf_model_dir = self._llm_engine._hf_model_dir or getattr( + getattr(self._mm_tokenizer, "tokenizer", None), "name_or_path", + None) + if hf_model_dir is None: + raise pb_utils.TritonModelException( + "triton_config.multimodal is enabled but the checkpoint directory " + "could not be resolved from the engine.") + hf_model_dir = str(hf_model_dir) + trust_remote_code = self._llm_engine.args.trust_remote_code + try: + self._mm_processor = AutoProcessor.from_pretrained( + hf_model_dir, trust_remote_code=trust_remote_code) + model_config = load_pretrained_config( + hf_model_dir, + trust_remote_code=trust_remote_code, + checkpoint_format=getattr(self._llm_engine.args, + "checkpoint_format", None)) + except Exception as e: + raise pb_utils.TritonModelException( + f"triton_config.multimodal is enabled but the HF processor/config " + f"for '{hf_model_dir}' could not be loaded: {e}") + + # Composite configs (e.g. Qwen2_5_VLConfig) delegate the instance + # attribute to `text_config`, so prefer the class attribute. + self._mm_model_type = getattr(type(model_config), + "model_type", None) or getattr( + model_config, "model_type", "") + self.logger.log_info("[trtllm] multimodal input enabled for model_type " + f"'{self._mm_model_type}'") + + async def _convert_request(self, request): """Helper function to convert the request into a prompt for LLM.generate_async Args: @@ -593,6 +652,8 @@ def _convert_request(self, request): Notes: - The current implementation only supports text_input being a 1D tensor(a single prompt). + - With `image_url`, prompt becomes a multimodal PromptInputs dict. + `image_url` is only read when `triton_config.multimodal` is set. """ text_input = get_input_tensor_by_name(request, 'text_input') if text_input is None: @@ -608,6 +669,27 @@ def _convert_request(self, request): if isinstance(prompt, bytes): prompt = prompt.decode("utf-8") + # Only read `image_url` when the operator opts in, so a deployment + # already declaring that input keeps its behavior after an upgrade. + if self.multimodal_enabled: + image_url = get_input_tensor_by_name(request, 'image_url') + if image_url is not None and image_url.size > 0: + from tensorrt_llm.inputs import async_build_multimodal_prompt + + media = [ + url.decode("utf-8") if isinstance(url, bytes) else str(url) + for url in image_url.reshape(-1) + ] + validate_media_urls(media) + prompt = await async_build_multimodal_prompt( + model_type=self._mm_model_type, + tokenizer=self._mm_tokenizer, + processor=self._mm_processor, + prompt=prompt, + media=media, + modality="image", + ) + sampling_params = get_sampling_params_from_request(request) output_config = get_output_config_from_request(request) streaming = get_streaming_from_request(request) diff --git a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.yaml b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.yaml index c9ad98f9dfce..6fa9040b026e 100644 --- a/triton_backend/all_models/llmapi/tensorrt_llm/1/model.yaml +++ b/triton_backend/all_models/llmapi/tensorrt_llm/1/model.yaml @@ -17,3 +17,7 @@ pipeline_parallel_size: 1 triton_config: max_batch_size: 0 # The current implementation does not support batching, batch support is tracked in JIRA-4496 decoupled: False + # Accept the optional `image_url` request input and serve multimodal + # models. Off by default so that a deployment already declaring an + # `image_url` input for another purpose is unaffected by an upgrade. + multimodal: False diff --git a/triton_backend/all_models/llmapi/tensorrt_llm/config.pbtxt b/triton_backend/all_models/llmapi/tensorrt_llm/config.pbtxt index a3f995aae882..ea69fc227321 100644 --- a/triton_backend/all_models/llmapi/tensorrt_llm/config.pbtxt +++ b/triton_backend/all_models/llmapi/tensorrt_llm/config.pbtxt @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -51,6 +51,14 @@ input [ data_type: TYPE_STRING dims: [ -1 ] }, + ## Optional http(s) image URLs, one per image. Only read when + ## `triton_config.multimodal` is set in model.yaml. + { + name: "image_url" + data_type: TYPE_STRING + dims: [ -1 ] + optional: true + }, { name: "streaming" data_type: TYPE_BOOL diff --git a/triton_backend/all_models/tests/test_llmapi_python_backend.py b/triton_backend/all_models/tests/test_llmapi_python_backend.py index b6b79e04aeed..2427104fc1f5 100644 --- a/triton_backend/all_models/tests/test_llmapi_python_backend.py +++ b/triton_backend/all_models/tests/test_llmapi_python_backend.py @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -24,7 +24,10 @@ # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +import asyncio +import json import sys +import threading from dataclasses import dataclass from typing import Dict, List, Union from unittest.mock import MagicMock, patch @@ -42,7 +45,7 @@ get_sampling_params_from_request, get_streaming_from_request) # Use PYTHONPATH=../llmapi/tensorrt_llm/1/ -from model import * +from model import TritonPythonModel, validate_media_urls @dataclass @@ -361,3 +364,204 @@ def test_get_parameter(): # Special cases assert get_parameter(model_config, "empty_param") is None assert get_parameter(model_config, "env_var_param") is None + + +def _make_multimodal_model(enabled: bool): + """A bare model with just the multimodal state `_convert_request` reads.""" + model = TritonPythonModel.__new__(TritonPythonModel) + model.multimodal_enabled = enabled + model._mm_tokenizer = "tokenizer" + model._mm_processor = "processor" + model._mm_model_type = "qwen2_5_vl" + return model + + +def test_convert_request_ignores_image_url_when_multimodal_disabled(): + # A deployment that already declares an `image_url` input for another + # purpose must be unaffected until it opts in via triton_config.multimodal. + model = _make_multimodal_model(enabled=False) + request = make_mock_triton_request({ + **inputs(), + "image_url": [b"https://example.com/a.jpg"], + }) + + prompt, _, _, _, _ = asyncio.run(model._convert_request(request)) + + assert prompt == "Tell me a story." + + +def test_convert_request_delegates_to_shared_inputs_helper(): + # The backend must not build the multimodal prompt itself: placeholder + # handling depends on the model's ContentFormat, which the shared helper in + # tensorrt_llm.inputs owns. Assert we hand it the right arguments. + model = _make_multimodal_model(enabled=True) + captured = {} + + async def fake_helper(**kwargs): + captured.update(kwargs) + return { + "prompt": "rendered", + "multi_modal_data": { + "image": ["decoded"] + } + } + + inputs_mod = MagicMock() + inputs_mod.async_build_multimodal_prompt = fake_helper + request = make_mock_triton_request({ + **inputs(), + "image_url": [b"https://example.com/a.jpg", b"/tmp/b.png"], + }) + + with patch.dict(sys.modules, { + "tensorrt_llm": MagicMock(), + "tensorrt_llm.inputs": inputs_mod, + }): + prompt, _, _, _, _ = asyncio.run(model._convert_request(request)) + + assert prompt["multi_modal_data"] == {"image": ["decoded"]} + assert captured["model_type"] == "qwen2_5_vl" + assert captured["tokenizer"] == "tokenizer" + assert captured["processor"] == "processor" + assert captured["modality"] == "image" + assert captured["prompt"] == "Tell me a story." + # Bytes tensors are decoded, order preserved. + assert captured["media"] == ["https://example.com/a.jpg", "/tmp/b.png"] + + +def _bare_model_for_execute(): + """A model with only the state `_execute_single_request` touches.""" + model = TritonPythonModel.__new__(TritonPythonModel) + model.logger = MagicMock() + model.lock = threading.Lock() + model.req_id_to_request_data = {} + model.triton_user_id_to_req_ids = {} + model._ongoing_request_count = 0 + model.decoupled = False + model.output_dtype = np.object_ + return model + + +class _RecordingSender: + + def __init__(self): + self.sent = [] + + def send(self, response, flags=None): + self.sent.append((response, flags)) + + +def test_execute_single_request_reports_preprocessing_failure(): + # A bad image URL fails inside _convert_request, before the request is + # registered in req_id_to_request_data. The error must still reach the + # client with COMPLETE_FINAL, otherwise it waits forever. + model = _bare_model_for_execute() + sender = _RecordingSender() + request = make_mock_triton_request({"text_input": ["describe this"]}) + request.get_response_sender = lambda: sender + request.request_id = lambda: "triton-user-1" + + async def failing_convert(_request): + raise RuntimeError( + "Cannot connect to host example.invalid:443 [Name or service not known]" + ) + + model._convert_request = failing_convert + + with patch.dict(sys.modules, {"tensorrt_llm": MagicMock()}): + with pytest.raises(RuntimeError): + asyncio.run(model._execute_single_request(request)) + + pb_utils = sys.modules["triton_python_backend_utils"] + assert len(sender.sent) == 1, "client must receive exactly one response" + response, flags = sender.sent[0] + assert flags == pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL + assert response.has_error() + assert "example.invalid" in response.error.message + + +def test_execute_single_request_skips_response_when_cancelled(): + # Once the request IS registered, an empty map means the cancellation loop + # already sent COMPLETE_FINAL and removed the entry, so the error handler + # must stay silent rather than send a second final response. + model = _bare_model_for_execute() + sender = _RecordingSender() + request = make_mock_triton_request({"text_input": ["describe this"]}) + request.get_response_sender = lambda: sender + request.request_id = lambda: "triton-user-2" + + async def convert(_request): + return ("a prompt", {}, False, {}, None) + + class CancellingIterator: + + def __aiter__(self): + return self + + async def __anext__(self): + # Stand in for cancellation_loop: it sends COMPLETE_FINAL and drops + # the entry while generation is in flight. + with model.lock: + model.req_id_to_request_data.clear() + raise RuntimeError("request aborted") + + engine = MagicMock() + engine.generate_async.return_value = CancellingIterator() + model._convert_request = convert + model._llm_engine = engine + + with patch.dict(sys.modules, {"tensorrt_llm": MagicMock()}): + with pytest.raises(RuntimeError): + asyncio.run(model._execute_single_request(request)) + + assert sender.sent == [], "must not double-send after cancellation" + + +@pytest.mark.parametrize("url", [ + "/etc/passwd", + "/lustre/private/secret.png", + "file:///etc/passwd", + "relative/path.jpg", + "ftp://example.com/a.jpg", + "data:image/png;base64,iVBORw0KGgo=", +]) +def test_validate_media_urls_rejects_non_web_urls(url): + with pytest.raises(MockTritonModelException) as excinfo: + validate_media_urls([url]) + assert "image_url" in str(excinfo.value) + + +@pytest.mark.parametrize("url", [ + "http://images.example.com/a.jpg", + "https://images.example.com/a.jpg", + "HTTPS://images.example.com/a.jpg", +]) +def test_validate_media_urls_accepts_web_urls(url): + validate_media_urls([url]) + + +def test_convert_request_rejects_local_path_before_loading(): + # The rejection must happen before the media loader is reached. + model = _make_multimodal_model(enabled=True) + called = False + + async def fake_helper(**kwargs): + nonlocal called + called = True + return {"prompt": "rendered"} + + inputs_mod = MagicMock() + inputs_mod.async_build_multimodal_prompt = fake_helper + request = make_mock_triton_request({ + **inputs(), + "image_url": [b"/etc/passwd"], + }) + + with patch.dict(sys.modules, { + "tensorrt_llm": MagicMock(), + "tensorrt_llm.inputs": inputs_mod, + }): + with pytest.raises(MockTritonModelException): + asyncio.run(model._convert_request(request)) + + assert not called, "loader must not be reached for a rejected URL"