From 9f3658934c7ecf744b68811945c8486f5a6147e0 Mon Sep 17 00:00:00 2001 From: Kris Hung <43719498+krishung5@users.noreply.github.com> Date: Tue, 3 Nov 2026 18:38:54 -0800 Subject: [PATCH] feat(trtllm): support audio through the OpenEngine sidecar (#16871) Signed-off-by: krishung5 Signed-off-by: krishung5 <43719498+krishung5@users.noreply.github.com> (cherry picked from commit bb160083f62edb30c3bcfbb8fd7c4af1ecb09d68) Signed-off-by: krishung5 --- .../shared-test-trtllm-nightly-image.yml | 6 +- lib/sidecar/testkit/tests/support/trtllm.rs | 35 +--- .../testkit/tests/trtllm_multimodal.rs | 150 ++++++++++++------ lib/sidecar/trtllm/README.md | 41 ++++- lib/sidecar/trtllm/ci/nightly.json | 6 +- lib/sidecar/trtllm/ci/run_image_test.py | 96 +++++++---- .../trtllm/ci/tests/test_agg_launch.py | 102 ++++++++++++ .../trtllm/ci/tests/test_run_media_test.py | 40 +++++ lib/sidecar/trtllm/ci/whisper.yaml | 20 +++ lib/sidecar/trtllm/launch/agg.sh | 30 ++-- lib/sidecar/trtllm/launch/common.sh | 52 ++++++ lib/sidecar/trtllm/src/client.rs | 23 +++ lib/sidecar/trtllm/src/convert.rs | 64 ++++++-- .../trtllm/src/convert/request_tests.rs | 114 +++++++++++-- lib/sidecar/trtllm/src/test_fixtures.rs | 18 +++ lib/sidecar/trtllm/src/tests/e2e.rs | 84 ++++++++++ 16 files changed, 741 insertions(+), 140 deletions(-) create mode 100644 lib/sidecar/trtllm/ci/tests/test_agg_launch.py create mode 100644 lib/sidecar/trtllm/ci/whisper.yaml diff --git a/.github/workflows/shared-test-trtllm-nightly-image.yml b/.github/workflows/shared-test-trtllm-nightly-image.yml index 7409afc5f7bb..0382a42d7ee6 100644 --- a/.github/workflows/shared-test-trtllm-nightly-image.yml +++ b/.github/workflows/shared-test-trtllm-nightly-image.yml @@ -27,7 +27,9 @@ jobs: steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Test nightly selection and media harness - run: python3 -m unittest discover -s lib/sidecar/trtllm/ci/tests -v + run: | + python3 -m pip install 'PyYAML==6.0.3' + python3 -m unittest discover -s lib/sidecar/trtllm/ci/tests -v - name: Resolve latest upstream nightly id: config run: >- @@ -61,7 +63,7 @@ jobs: strategy: fail-fast: false matrix: - media: [image, video] + media: [image, video, audio] needs: build runs-on: [gpu-1, amd64, v2, us-west-2] timeout-minutes: 30 diff --git a/lib/sidecar/testkit/tests/support/trtllm.rs b/lib/sidecar/testkit/tests/support/trtllm.rs index 00f030d9d297..4fcb88c8189d 100644 --- a/lib/sidecar/testkit/tests/support/trtllm.rs +++ b/lib/sidecar/testkit/tests/support/trtllm.rs @@ -42,9 +42,7 @@ pub struct Discovery { pub calls: AtomicUsize, pub unavailable: AtomicBool, pub hold: AtomicBool, - pub portable_multimodal: AtomicBool, - pub advertises_image: AtomicBool, - pub advertises_video: AtomicBool, + pub model_info: Mutex>, pub omit_context_length: AtomicBool, pub error: Mutex>, release: tokio::sync::Notify, @@ -213,32 +211,15 @@ impl Control for ControlledControl { if let Some(code) = *self.discovery.error.lock().unwrap() { return Err(Status::new(code, "model discovery rejected")); } - let mut response = self.inner.get_model_info(request).await?; - let info = response.get_mut(); + let model_info = self.discovery.model_info.lock().unwrap().clone(); + let mut response = if let Some(info) = model_info { + Response::new(info) + } else { + self.inner.get_model_info(request).await? + }; if self.discovery.omit_context_length.load(Ordering::SeqCst) { - info.max_context_length = None; + response.get_mut().max_context_length = None; } - info.supports_multimodal = Some(self.discovery.portable_multimodal.load(Ordering::SeqCst)); - let modalities: Vec<_> = [ - ( - self.discovery.advertises_image.load(Ordering::SeqCst), - "image", - ), - ( - self.discovery.advertises_video.load(Ordering::SeqCst), - "video", - ), - ] - .into_iter() - .filter_map(|(enabled, modality)| enabled.then_some(modality)) - .collect(); - info.extra = Some( - dynamo_sidecar_common::json_to_struct( - serde_json::json!({"trtllm_supported_modalities": modalities}), - "test capabilities", - ) - .unwrap(), - ); Ok(response) } diff --git a/lib/sidecar/testkit/tests/trtllm_multimodal.rs b/lib/sidecar/testkit/tests/trtllm_multimodal.rs index 9f77dae3c1b5..0206c877299c 100644 --- a/lib/sidecar/testkit/tests/trtllm_multimodal.rs +++ b/lib/sidecar/testkit/tests/trtllm_multimodal.rs @@ -72,15 +72,23 @@ fn finished() -> pb::GenerateResponse { } } -fn advertise_images(fixture: &Fixture) { - fixture - .discovery - .portable_multimodal - .store(true, Ordering::SeqCst); - fixture - .discovery - .advertises_image - .store(true, Ordering::SeqCst); +fn advertise_media(fixture: &Fixture, modalities: &[&str]) { + *fixture.discovery.model_info.lock().unwrap() = Some(pb::ModelInfo { + model_id: "mocker-model".into(), + served_model_name: "mocker-model".into(), + max_context_length: Some(32768), + supports_multimodal: Some(true), + supports_text_input: Some(true), + supports_token_ids_input: Some(true), + extra: Some( + dynamo_sidecar_common::json_to_struct( + serde_json::json!({"trtllm_supported_modalities": modalities}), + "test capabilities", + ) + .unwrap(), + ), + ..Default::default() + }); } async fn wait_for_calls(fixture: &Fixture, count: usize) { @@ -96,7 +104,7 @@ async fn wait_for_calls(fixture: &Fixture, count: usize) { async fn images_preserve_order_uuids_and_engine_expanded_usage() { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; - advertise_images(&fixture); + advertise_media(&fixture, &["image"]); let engine = fixture.engine().await; engine.start(0).await.unwrap(); let ctx = mock_context(); @@ -193,7 +201,7 @@ async fn formatter_regrouped_tool_media_reaches_openengine_in_rendered_order() { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; - advertise_images(&fixture); + advertise_media(&fixture, &["image"]); let engine = fixture.engine().await; engine.start(0).await.unwrap(); let ctx = mock_context(); @@ -232,23 +240,19 @@ async fn both_media_capability_signals_are_required() { for (portable, advertised) in [(true, false), (false, true)] { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; + let mut modalities = vec![if video { "image" } else { "video" }]; + if advertised { + modalities.push(if video { "video" } else { "image" }); + } + advertise_media(&fixture, &modalities); fixture .discovery - .portable_multimodal - .store(portable, Ordering::SeqCst); - let (target, other) = if video { - ( - &fixture.discovery.advertises_video, - &fixture.discovery.advertises_image, - ) - } else { - ( - &fixture.discovery.advertises_image, - &fixture.discovery.advertises_video, - ) - }; - target.store(advertised, Ordering::SeqCst); - other.store(true, Ordering::SeqCst); + .model_info + .lock() + .unwrap() + .as_mut() + .unwrap() + .supports_multimodal = Some(portable); let engine = fixture.engine().await; engine.start(0).await.unwrap(); let ctx = mock_context(); @@ -278,14 +282,7 @@ async fn videos_recover_discovery_and_preserve_sources_uuids_and_usage() { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; fixture.discovery.unavailable.store(true, Ordering::SeqCst); - fixture - .discovery - .portable_multimodal - .store(true, Ordering::SeqCst); - fixture - .discovery - .advertises_video - .store(true, Ordering::SeqCst); + advertise_media(&fixture, &["video"]); let mut startup = tokio::spawn(bootstrap(&fixture, Some(32768), 5)); wait_for_calls(&fixture, 2).await; assert!(!startup.is_finished()); @@ -359,11 +356,7 @@ async fn startup_recovers_media_capabilities_before_registration() { for omit_window in [false, true] { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; - advertise_images(&fixture); - fixture - .discovery - .advertises_video - .store(true, Ordering::SeqCst); + advertise_media(&fixture, &["image", "video"]); fixture .discovery .omit_context_length @@ -508,11 +501,7 @@ async fn shutdown_cancels_retrying_and_stalled_bootstrap_discovery() { async fn media_window_override_is_rejected_before_dispatch() { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; - advertise_images(&fixture); - fixture - .discovery - .advertises_video - .store(true, Ordering::SeqCst); + advertise_media(&fixture, &["image", "video"]); let engine = fixture.engine_with_context_length(4096); engine.start(0).await.unwrap(); for (video, max_tokens) in [ @@ -538,6 +527,57 @@ async fn media_window_override_is_rejected_before_dispatch() { fixture.shutdown().await; } +#[tokio::test] +async fn text_only_audio_preserves_whisper_task_prompt_with_empty_token_ids() { + let control = Controller::::default(); + let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; + advertise_media(&fixture, &["audio"]); + fixture + .discovery + .model_info + .lock() + .unwrap() + .as_mut() + .unwrap() + .supports_token_ids_input = Some(false); + let engine = fixture.engine().await; + engine.start(0).await.unwrap(); + let uri = "data:audio/wav;base64,YXVkaW8="; + for prompt in [ + "", + "<|startoftranscript|><|fr|><|translate|><|notimestamps|>", + ] { + let mut req = request("mocker-model", vec![], 16); + req.extra_args = Some(serde_json::json!({"formatted_prompt": prompt})); + req.multi_modal_data = Some(HashMap::from([( + "audio_url".into(), + vec![MultimodalData::RawUrl(uri.into())], + )])); + let ctx = mock_context(); + let handle = control.request(ctx.id(), RequestPlan::default()); + fixture.respond(ctx.id(), vec![finished()]); + let outputs = collect(&engine, req, GenerateContext::new(ctx, None)).await; + assert_eq!( + outputs.last().unwrap().as_ref().unwrap().finish_reason, + Some(FinishReason::Stop) + ); + let native = handle.native_request().unwrap(); + assert_eq!( + native.input, + Some(pb::generate_request::Input::Prompt(prompt.into())) + ); + assert_eq!(native.media.len(), 1); + assert_eq!(native.media[0].modality, pb::Modality::Audio as i32); + assert_eq!( + native.media[0].source, + Some(pb::media_item::Source::DataUri(uri.into())) + ); + } + assert_eq!(fixture.discovery.calls.load(Ordering::SeqCst), 1); + engine.cleanup().await.unwrap(); + fixture.shutdown().await; +} + #[test] fn remote_url_policy_is_enforced_before_engine_dispatch() { // Isolate startup environment changes from the other parallel tests. @@ -590,11 +630,7 @@ fn remote_url_policy_is_enforced_before_engine_dispatch() { async fn check_remote_url_policy(allow_remote: bool) { let control = Controller::::default(); let mut fixture = Fixture::start(control.clone(), FixtureConfig::default()).await; - advertise_images(&fixture); - fixture - .discovery - .advertises_video - .store(true, Ordering::SeqCst); + advertise_media(&fixture, &["image", "video", "audio"]); let engine = fixture.engine().await; engine.start(0).await.unwrap(); for (key, media) in [ @@ -630,6 +666,22 @@ async fn check_remote_url_policy(allow_remote: bool) { "video_url", MultimodalData::RawUrl("data:video/mp4;base64,dmlkZW8=".into()), ), + ( + "audio_url", + MultimodalData::RawUrl("http://example.com/audio.wav".into()), + ), + ( + "audio_url", + MultimodalData::RawUrl("HTTPS://example.com/audio.wav".into()), + ), + ( + "audio_url", + MultimodalData::Url("https://example.com/audio.wav".parse().unwrap()), + ), + ( + "audio_url", + MultimodalData::RawUrl("data:audio/wav;base64,YXVkaW8=".into()), + ), ] { let inline = matches!(&media, MultimodalData::RawUrl(url) if url.starts_with("data:")); let mut req = image_request(); diff --git a/lib/sidecar/trtllm/README.md b/lib/sidecar/trtllm/README.md index 5e94f4b78d4d..93e0156b5d09 100644 --- a/lib/sidecar/trtllm/README.md +++ b/lib/sidecar/trtllm/README.md @@ -31,7 +31,7 @@ supported, even when `--context-length` is supplied. > The standalone launch and deployment examples use > `nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc3.dev202610290007`, which includes > image/video URL passthrough, native KVCM2 event streaming, and OpenEngine routing load support. -> The nightly CI job tests image and video passthrough against the latest regular upstream nightly; +> The nightly CI job tests image, video, and audio passthrough against the latest regular upstream nightly; > it also supports manual dispatch. The examples below stay pinned to the validated version. ## Support matrix @@ -67,6 +67,45 @@ curl -s localhost:8000/v1/chat/completions \ -d '{"model":"Qwen/Qwen3-0.6B","messages":[{"role":"user","content":"Hello"}],"max_tokens":32}' ``` +### Audio (Whisper) + +Audio requires a runtime containing upstream audio support, such as the validated +`nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc3.dev202610300406`; the older image/video +example pin above predates audio. The nightly audio lane runs real Whisper-tiny +inference and checks the transcript and authoritative token usage. Full frontend +HTTP CI is tracked separately from that adapter-to-engine test. + +Whisper-tiny has no chat template and accepts text task prefixes rather than +token-ID input. Supply `--custom-jinja-template` to the sidecar. An empty template +selects the checkpoint's default transcription task; a template may instead emit +a complete language/task prefix, for example +`<|startoftranscript|><|fr|><|translate|><|notimestamps|>`. + +Run these processes in separate terminals with the appropriate Dynamo and +TensorRT-LLM environments, using `DYN_DISCOVERY_BACKEND=file` on a single host: + +```bash +python3 -m dynamo.frontend --migration-limit 0 + +# In the audio-capable TensorRT-LLM environment, from the repository root: +python3 -m tensorrt_llm.commands.serve openai/whisper-tiny \ + --grpc --grpc-protocol openengine --host 127.0.0.1 --port 30001 \ + --max_seq_len 1500 --extra_llm_api_options lib/sidecar/trtllm/ci/whisper.yaml + +# In the Dynamo environment: +WHISPER_TEMPLATE=$(mktemp --suffix=.jinja) +dynamo-trtllm-sidecar --grpc-endpoint 127.0.0.1:30001 \ + --model-path openai/whisper-tiny --custom-jinja-template "$WHISPER_TEMPLATE" +``` + +Send audio using a chat content item with `"type":"audio_url"` and +`"audio_url":{"url":"https://.../speech.wav"}` (or a data URI). +For text-input models, **keep frontend `--migration-limit 0`**: retry migration +currently extends token IDs, while these engines receive the original formatted +text, which can restart a partially delivered transcript. Worker-advertised +migration restrictions need separate frontend support. The aggregated launcher +preserves explicit engine settings, including `guided_decoding_backend: null`. + ## Deploy on Kubernetes The manifests pin compatible Dynamo frontend and TensorRT-LLM runtime nightlies diff --git a/lib/sidecar/trtllm/ci/nightly.json b/lib/sidecar/trtllm/ci/nightly.json index 67b4986ed17c..40f4fd76f6e6 100644 --- a/lib/sidecar/trtllm/ci/nightly.json +++ b/lib/sidecar/trtllm/ci/nightly.json @@ -1,5 +1,9 @@ { "repository": "nvcr.io/nvidia/tensorrt-llm/release", "model": "Qwen/Qwen3-VL-2B-Instruct", - "model_revision": "89644892e4d85e24eaac8bacfd4f463576704203" + "model_revision": "89644892e4d85e24eaac8bacfd4f463576704203", + "audio": { + "model": "openai/whisper-tiny", + "model_revision": "169d4a4341b33bc18d8881c4b69c2e104e1cc0af" + } } diff --git a/lib/sidecar/trtllm/ci/run_image_test.py b/lib/sidecar/trtllm/ci/run_image_test.py index 09de9a343801..bd230552fe61 100755 --- a/lib/sidecar/trtllm/ci/run_image_test.py +++ b/lib/sidecar/trtllm/ci/run_image_test.py @@ -17,6 +17,7 @@ import sys import time import traceback +import urllib.request import xml.etree.ElementTree as ET from pathlib import Path @@ -73,12 +74,40 @@ ], "context_length": 2048, }, + "audio": { + "prompt": "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>", + "token_ids": [], + "context_length": 1500, + }, } VIDEO_PATH = ( Path(__file__).resolve().parents[3] / "llm/tests/data/media/triangle_240p_10.mp4" ) VIDEO_SHA256 = "04171a85ee3ba5ac9118b0a5858adf0892b49474cdf15320e32599a3d1ab9919" +AUDIO_URL = ( + "https://raw.githubusercontent.com/yuekaizhang/Triton-ASR-Client/" + "8e35ed77aeece087fa53abf2a68a282dc13efc8c/" + "datasets/mini_en/wav/1221-135766-0002.wav" +) +AUDIO_SHA256 = "65b222837919ccbb924a4e1077413ea7cc6af3e68b663b012a9539d5c05850f0" + + +def prepare_audio(results): + with urllib.request.urlopen(AUDIO_URL, timeout=30) as response: + data = response.read(1024 * 1024 + 1) + if hashlib.sha256(data).hexdigest() != AUDIO_SHA256: + raise ValueError("Audio fixture does not match the validated speech sample") + path = results / "audio-data-uri.txt" + path.write_text("data:audio/wav;base64," + base64.b64encode(data).decode("ascii")) + return path + + +def model_config(config, media): + if media == "audio": + return {**config, **config["audio"]} + return config + def video_data_uri(path=VIDEO_PATH): data = path.read_bytes() @@ -95,6 +124,9 @@ def validate_response(media, caption, response, input_token_count): colors = [word for word in words if word in ("red", "blue")] if colors[:2] != ["red", "blue"]: raise AssertionError(f"Expected red, then blue, got: {caption}") + elif media == "audio": + if "thoughts affected hester" not in " ".join(words): + raise AssertionError(f"Unexpected speech transcription: {caption}") elif not {"triangle", "triangular"}.intersection(words): raise AssertionError( f"Expected a triangle in the video description, got: {caption}" @@ -166,7 +198,7 @@ def run(args, results): fixture = FIXTURES[args.media] test = f"tests::e2e::e2e_real_openengine_{args.media}" - config = json.loads(args.config.read_text()) + config = model_config(json.loads(args.config.read_text()), args.media) identity = { **config, "media": args.media, @@ -180,6 +212,7 @@ def run(args, results): ) video = video_data_uri() if args.media == "video" else None + audio_path = prepare_audio(results) if args.media == "audio" else None binary = args.binary.resolve() binary.chmod(binary.stat().st_mode | 0o111) @@ -195,40 +228,44 @@ def run(args, results): config["model"], revision=config["model_revision"] ) tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True) - # The fixed Rust fixture uses Qwen3-VL's rendering of this exact message. - # Check it against the pinned tokenizer so a template change fails loudly. - token_ids = tokenizer.apply_chat_template( - [ - { - "role": "user", - "content": [ - {"type": args.media}, - {"type": "text", "text": fixture["prompt"]}, - ], - } - ], - tokenize=True, - add_generation_prompt=True, - ) expected_ids = fixture["token_ids"] - token_ids = ( - token_ids["input_ids"] - if isinstance(token_ids, dict) or hasattr(token_ids, "input_ids") - else token_ids - ) - if token_ids != expected_ids: - raise RuntimeError(f"{args.media} fixture tokenizer mismatch: {token_ids}") + if args.media != "audio": + # The fixed Rust fixture uses Qwen3-VL's rendering of this exact message. + # Check it against the pinned tokenizer so a template change fails loudly. + token_ids = tokenizer.apply_chat_template( + [ + { + "role": "user", + "content": [ + {"type": args.media}, + {"type": "text", "text": fixture["prompt"]}, + ], + } + ], + tokenize=True, + add_generation_prompt=True, + ) + token_ids = ( + token_ids["input_ids"] + if isinstance(token_ids, dict) or hasattr(token_ids, "input_ids") + else token_ids + ) + if token_ids != expected_ids: + raise RuntimeError(f"{args.media} fixture tokenizer mismatch: {token_ids}") with socket.socket() as listener: listener.bind(("127.0.0.1", 0)) port = listener.getsockname()[1] options = results / "engine.yaml" context_length = fixture["context_length"] - options.write_text( - f"max_batch_size: 1\nmax_num_tokens: {context_length}\n" - f"kv_cache_config:\n max_tokens: {context_length}\n" - " free_gpu_memory_fraction: 0.15\n" - ) + if args.media == "audio": + options.write_text(Path(__file__).with_name("whisper.yaml").read_text()) + else: + options.write_text( + f"max_batch_size: 1\nmax_num_tokens: {context_length}\n" + f"kv_cache_config:\n max_tokens: {context_length}\n" + " free_gpu_memory_fraction: 0.15\n" + ) result_path = results / "response.json" engine = None try: @@ -264,6 +301,9 @@ def run(args, results): ) if video is not None: env["TRTLLM_E2E_VIDEO_DATA_URI"] = video + if audio_path is not None: + env["TRTLLM_E2E_AUDIO_DATA_URI_FILE"] = str(audio_path) + env["TRTLLM_E2E_AUDIO_PROMPT"] = fixture["prompt"] run_rust_test(binary, env, test, results) finally: stop_engine(engine) diff --git a/lib/sidecar/trtllm/ci/tests/test_agg_launch.py b/lib/sidecar/trtllm/ci/tests/test_agg_launch.py new file mode 100644 index 000000000000..7991d9516623 --- /dev/null +++ b/lib/sidecar/trtllm/ci/tests/test_agg_launch.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import signal +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +import yaml + +ROOT = Path(__file__).resolve().parents[5] + + +class AggregatedLaunchTests(unittest.TestCase): + def test_engine_receives_base_config_defaults_and_profiler_overrides(self): + cases = [ + (None, False, "xgrammar"), + ("max_batch_size: 1\n", False, "xgrammar"), + ("guided_decoding_backend: guidance\n", False, "guidance"), + ((ROOT / "lib/sidecar/trtllm/ci/whisper.yaml").read_text(), False, None), + ("guided_decoding_backend: guidance\n", True, None), + ] + for config, disabled, expected in cases: + with self.subTest( + config=config, disabled=disabled + ), tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + wrapper = root / "python3" + wrapper.write_text( + f"#!{sys.executable}\n" + "import json, os, signal, sys\n" + "from pathlib import Path\n" + "import yaml\n" + "args = sys.argv[1:]\n" + "if args[:1] == ['-c']:\n" + " sys.exit(0)\n" + "if args[:2] == ['-m', 'tensorrt_llm.commands.serve']:\n" + " config = Path(args[args.index('--extra_llm_api_options') + 1])\n" + " Path(os.environ['CAPTURE']).write_text(json.dumps(yaml.safe_load(config.read_text())))\n" + " sys.exit(0)\n" + "if args[:1] == ['-m'] or '--grpc-endpoint' in args:\n" + " signal.pause()\n" + f"os.execv({sys.executable!r}, [{sys.executable!r}, *args])\n" + ) + wrapper.chmod(0o755) + capture = root / "captured.json" + env = { + key: value + for key, value in os.environ.items() + if not key.startswith(("TRTLLM_", "DYN_", "_PROFILE_OVERRIDE_")) + and key not in ("AGG_ENGINE_ARGS", "MODEL") + } + env.update( + { + "PATH": f"{root}:{os.environ['PATH']}", + "TRTLLM_PYTHON": str(wrapper), + "DYNAMO_SIDECAR_BIN": str(wrapper), + "CAPTURE": str(capture), + "TMPDIR": str(root), + "_PROFILE_OVERRIDE_TRTLLM_MAX_TOTAL_TOKENS": "512", + } + ) + if config is not None: + base = root / "base.yaml" + base.write_text(config) + env["AGG_ENGINE_ARGS"] = str(base) + if disabled: + env["TRTLLM_DISABLE_GUIDED_DECODING"] = "1" + process = subprocess.Popen( + ["bash", str(ROOT / "lib/sidecar/trtllm/launch/agg.sh")], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = process.communicate(timeout=10) + finally: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5) + self.assertEqual(process.returncode, 0, stdout + stderr) + actual = json.loads(capture.read_text()) + self.assertEqual(actual["guided_decoding_backend"], expected) + self.assertEqual(actual["kv_cache_config"]["max_tokens"], 512) + if config is not None: + for key, value in yaml.safe_load(config).items(): + if key == "kv_cache_config": + self.assertEqual(actual[key], {**value, "max_tokens": 512}) + elif key != "guided_decoding_backend": + self.assertEqual(actual[key], value) + + +if __name__ == "__main__": + unittest.main() diff --git a/lib/sidecar/trtllm/ci/tests/test_run_media_test.py b/lib/sidecar/trtllm/ci/tests/test_run_media_test.py index 4237da921f35..642004292218 100644 --- a/lib/sidecar/trtllm/ci/tests/test_run_media_test.py +++ b/lib/sidecar/trtllm/ci/tests/test_run_media_test.py @@ -45,6 +45,46 @@ def test_video_requires_expanded_prompt_and_completion_usage(self): ): runner.validate_response("video", "A rotating triangle.", response, 21) + def test_audio_requires_expected_transcription_and_authoritative_usage(self): + response = {"prompt_tokens": 4, "completion_tokens": 17} + runner.validate_response("audio", "Thoughts affected Hester.", response, 0) + with self.assertRaisesRegex(AssertionError, "transcription"): + runner.validate_response("audio", "An unrelated sentence.", response, 0) + for usage in ( + {"prompt_tokens": 0, "completion_tokens": 17}, + {"prompt_tokens": 4, "completion_tokens": 0}, + ): + with self.subTest(usage=usage), self.assertRaisesRegex( + AssertionError, "token usage" + ): + runner.validate_response("audio", "Thoughts affected Hester.", usage, 0) + + def test_audio_selects_whisper_without_changing_image_identity(self): + config = { + "model": "vision", + "model_revision": "vision-revision", + "image": "official@sha256:digest", + "upstream_commit": "commit", + "audio": {"model": "whisper", "model_revision": "audio-revision"}, + } + selected = runner.model_config(config, "audio") + self.assertEqual(selected["model"], "whisper") + self.assertEqual(selected["model_revision"], "audio-revision") + self.assertEqual(selected["image"], config["image"]) + self.assertEqual(config["model"], "vision") + + def test_changed_audio_fixture_is_rejected(self): + with tempfile.TemporaryDirectory() as tmp: + with ( + patch.object(runner.urllib.request, "urlopen") as fetch, + self.assertRaisesRegex(ValueError, "Audio fixture"), + ): + fetch.return_value.__enter__.return_value.read.return_value = ( + b"wrong audio" + ) + runner.prepare_audio(Path(tmp)) + self.assertFalse((Path(tmp) / "audio-data-uri.txt").exists()) + def test_video_data_uri_contains_the_checked_in_clip(self): uri = runner.video_data_uri() header, encoded = uri.split(",", 1) diff --git a/lib/sidecar/trtllm/ci/whisper.yaml b/lib/sidecar/trtllm/ci/whisper.yaml new file mode 100644 index 000000000000..a66d9fb01331 --- /dev/null +++ b/lib/sidecar/trtllm/ci/whisper.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +tensor_parallel_size: 1 +max_batch_size: 1 +max_input_len: 1500 +max_num_tokens: 1500 +backend: pytorch +disable_overlap_scheduler: true +enable_chunked_prefill: false + +scheduler_config: + use_python_scheduler: true + +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.2 + cross_kv_cache_fraction: 0.5 + use_kv_cache_manager_v2: true +guided_decoding_backend: null diff --git a/lib/sidecar/trtllm/launch/agg.sh b/lib/sidecar/trtllm/launch/agg.sh index c5a9f817342c..5b01ffee0551 100755 --- a/lib/sidecar/trtllm/launch/agg.sh +++ b/lib/sidecar/trtllm/launch/agg.sh @@ -19,7 +19,7 @@ source "$SCRIPT_DIR/../../../../examples/common/gpu_utils.sh" # build_trtllm_o # shellcheck disable=SC1091 # Resolved relative to this script at runtime. source "$SCRIPT_DIR/../../../../examples/common/launch_utils.sh" # print_launch_banner, wait_any_exit # shellcheck disable=SC1091 # Resolved relative to this script at runtime. -source "$SCRIPT_DIR/common.sh" # trtllm_ensure_openengine_bindings, trtllm_resolve_context_length +source "$SCRIPT_DIR/common.sh" # OpenEngine bindings, context length, config merge MODEL="${MODEL:-Qwen/Qwen3-0.6B}" @@ -48,8 +48,11 @@ while [[ $# -gt 0 ]]; do echo " DYN_HTTP_PORT Dynamo frontend port (default: 8000)" echo " DYN_SYSTEM_PORT Dynamo sidecar system port (default: 8081)" echo " TRTLLM_GRPC_PORT TensorRT-LLM gRPC port (default: 30001)" + echo " AGG_ENGINE_ARGS Base TensorRT-LLM YAML/JSON options file" echo " TRTLLM_CONTEXT_LENGTH Model context length, applied to both the engine and the" echo " sidecar (default: 4096; unset when --max_seq_len is given)" + echo " TRTLLM_DISABLE_GUIDED_DECODING=1" + echo " Disable xgrammar for models that do not support it" exit 0 ;; *) @@ -76,15 +79,13 @@ TRTLLM_PYTHON="${TRTLLM_PYTHON:-python3}" TRTLLM_GRPC_PORT=$(dyn_port DYN_ENGINE_GRPC_PORT 1 "${TRTLLM_GRPC_PORT:-30001}") CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" -# `--extra_llm_api_options` (alias `--config`) is last-wins, not additive, so a -# forwarded copy would drop the guided-decoding backend and a required or named -# tool_choice would fail again. Refuse it rather than silently losing the setting. +# `--extra_llm_api_options` (alias `--config`) is last-wins, not additive. The +# launcher reserves it for required and profiler overrides; callers can supply +# base options through AGG_ENGINE_ARGS, which are merged below. for arg in "${EXTRA_ARGS[@]}"; do case "$arg" in --extra_llm_api_options|--extra_llm_api_options=*|--config|--config=*) - echo "Cannot forward ${arg%%=*}: this launcher needs it for" >&2 - echo "guided_decoding_backend. Add 'guided_decoding_backend: xgrammar'" >&2 - echo "to your file and run the engine yourself." >&2 + echo "Cannot forward ${arg%%=*}: use AGG_ENGINE_ARGS instead." >&2 exit 1 ;; esac @@ -98,9 +99,18 @@ HTTP_PORT="${DYN_HTTP_PORT:-8000}" # Dynamo sends a required or named tool_choice as a JSON schema, which # TensorRT-LLM enforces only with a guided-decoding backend. TRTLLM_EXTRA_CONFIG=$(mktemp "${TMPDIR:-/tmp}/dynamo-trtllm-sidecar.XXXXXX.yaml") -build_trtllm_override_args_with_mem \ - --merge-with-json '{"guided_decoding_backend": "xgrammar"}' \ - > "$TRTLLM_EXTRA_CONFIG" +TRTLLM_REQUIRED_OVERRIDES='{}' +if [[ "${TRTLLM_DISABLE_GUIDED_DECODING:-0}" == "1" ]]; then + TRTLLM_REQUIRED_OVERRIDES='{"guided_decoding_backend": null}' +fi +TRTLLM_OVERRIDES=$(build_trtllm_override_args_with_mem \ + --merge-with-json "$TRTLLM_REQUIRED_OVERRIDES") +trtllm_write_merged_config \ + "${AGG_ENGINE_ARGS:-}" \ + "$TRTLLM_OVERRIDES" \ + "$TRTLLM_EXTRA_CONFIG" \ + "$TRTLLM_PYTHON" \ + '{"guided_decoding_backend":"xgrammar"}' print_launch_banner "Launching TensorRT-LLM OpenEngine-gRPC Sidecar (1 GPU)" "$MODEL" "$HTTP_PORT" \ "TensorRT-LLM gRPC: 127.0.0.1:${TRTLLM_GRPC_PORT}" \ diff --git a/lib/sidecar/trtllm/launch/common.sh b/lib/sidecar/trtllm/launch/common.sh index 21e4a901dacf..c622336527af 100755 --- a/lib/sidecar/trtllm/launch/common.sh +++ b/lib/sidecar/trtllm/launch/common.sh @@ -9,6 +9,7 @@ # Source it, then call: # trtllm_ensure_openengine_bindings "$TRTLLM_PYTHON" # trtllm_resolve_context_length "${EXTRA_ARGS[@]}" +# trtllm_write_merged_config "$BASE" "$OVERRIDES" "$OUTPUT" "$TRTLLM_PYTHON" # # `trtllm_resolve_context_length` sets TRTLLM_MAX_SEQ_LEN_ARGS, # TRTLLM_CONTEXT_LENGTH_ARGS, and TRTLLM_CONTEXT_LENGTH. @@ -78,3 +79,54 @@ trtllm_resolve_context_length() { TRTLLM_CONTEXT_LENGTH_ARGS=(--context-length "$TRTLLM_CONTEXT_LENGTH") fi } + +# Write a TensorRT-LLM options file by recursively applying JSON overrides to +# an optional YAML/JSON base. Recursive merging preserves sibling settings in +# nested blocks such as kv_cache_config when the test profiler supplies a +# max_tokens override. Optional fifth argument supplies defaults for absent keys. +trtllm_write_merged_config() { + local base_config="$1" + local overrides_json="$2" + local output="$3" + local python="$4" + local defaults_json='{}' + if [[ $# -ge 5 ]]; then + defaults_json="$5" + fi + + if [[ -n "$base_config" && ! -r "$base_config" ]]; then + echo "TensorRT-LLM engine config is not readable: $base_config" >&2 + return 1 + fi + + "$python" - "$base_config" "$overrides_json" "$output" "$defaults_json" <<'PY' +import json +import sys + +import yaml + + +def merge(base, overrides): + for key, value in overrides.items(): + if isinstance(value, dict) and isinstance(base.get(key), dict): + merge(base[key], value) + else: + base[key] = value + + +base_path, overrides_json, output_path, defaults_json = sys.argv[1:] +config = None +if base_path: + with open(base_path, encoding="utf-8") as source: + config = yaml.safe_load(source) +if config is None: + config = {} +if not isinstance(config, dict): + raise TypeError(f"TensorRT-LLM engine config must be a mapping: {base_path}") +for key, value in json.loads(defaults_json).items(): + config.setdefault(key, value) +merge(config, json.loads(overrides_json)) +with open(output_path, "w", encoding="utf-8") as destination: + yaml.safe_dump(config, destination, sort_keys=False) +PY +} diff --git a/lib/sidecar/trtllm/src/client.rs b/lib/sidecar/trtllm/src/client.rs index 6dfb47f690e6..a035117ff052 100644 --- a/lib/sidecar/trtllm/src/client.rs +++ b/lib/sidecar/trtllm/src/client.rs @@ -40,6 +40,14 @@ const LOAD_RPC_TIMEOUT: Duration = Duration::from_secs(2); pub(crate) struct SupportedMedia { pub(crate) image: bool, pub(crate) video: bool, + pub(crate) audio: bool, +} + +/// Input forms explicitly advertised by `Control.GetModelInfo`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct SupportedInputs { + pub(crate) text: bool, + pub(crate) token_ids: bool, } impl SupportedMedia { @@ -47,6 +55,7 @@ impl SupportedMedia { match modality { "image" => self.image, "video" => self.video, + "audio" => self.audio, _ => false, } } @@ -64,6 +73,9 @@ pub(crate) struct ModelLimits { /// Cap on generated tokens, when the engine reports one. A context window /// alone can imply a larger budget than the engine will accept. pub(crate) max_output_tokens: Option, + /// `None` preserves token-ID behavior for older servers that omit both + /// optional input-capability fields. + pub(crate) supported_inputs: Option, /// Modalities advertised by both the portable and engine-specific capabilities. /// `None` is unresolved metadata; startup always resolves this. pub(crate) supported_media: Option, @@ -213,6 +225,15 @@ impl TrtllmClient { .await? .into_inner(); let mut supported_media = SupportedMedia::default(); + let supported_inputs = + if info.supports_text_input.is_none() && info.supports_token_ids_input.is_none() { + None + } else { + Some(SupportedInputs { + text: info.supports_text_input == Some(true), + token_ids: info.supports_token_ids_input == Some(true), + }) + }; if info.supports_multimodal == Some(true) && let Some(prost_types::value::Kind::ListValue(modalities)) = info .extra @@ -228,6 +249,7 @@ impl TrtllmClient { match modality.as_str() { "image" => supported_media.image = true, "video" => supported_media.video = true, + "audio" => supported_media.audio = true, _ => {} } } @@ -237,6 +259,7 @@ impl TrtllmClient { context_length: info.max_context_length.filter(|len| *len > 0), engine_context_length: info.max_context_length.filter(|len| *len > 0), max_output_tokens: info.max_output_tokens.filter(|cap| *cap > 0), + supported_inputs, supported_media: Some(supported_media), }, served_model_name: info.served_model_name, diff --git a/lib/sidecar/trtllm/src/convert.rs b/lib/sidecar/trtllm/src/convert.rs index 9f6fab3448d0..4c5aeca4b625 100644 --- a/lib/sidecar/trtllm/src/convert.rs +++ b/lib/sidecar/trtllm/src/convert.rs @@ -4,7 +4,7 @@ //! Conversion between Dynamo's `PreprocessedRequest` / `LLMEngineOutput` and the //! TensorRT-LLM OpenEngine (`openengine.v1`) protobuf messages. //! -//! Supports aggregated and disaggregated image and video generation. Audio, +//! Supports aggregated and disaggregated image, video, and audio generation. //! LoRA, beam search, and `n > 1` are rejected before dispatch. //! //! Disaggregation is driven by [`DisaggregationMode`]: a prefill worker marks @@ -51,10 +51,6 @@ pub(crate) fn build_generate_request_with_routing( let request_has_media = has_media(&request); let media = media(&mut request, limits)?; - let sampling = &request.sampling_options; - let stop = &request.stop_conditions; - let output = &request.output_options; - // A prefill worker only needs the context phase; TensorRT-LLM still requires // a positive budget, and one token is what the context phase produces. // No minimum check here: the minimum applies to the decode worker's budget, @@ -105,16 +101,17 @@ pub(crate) fn build_generate_request_with_routing( }) .map(str::to_string); let include_stop_in_output = has_visible_stop_tokens(&request).then_some(true); - let token_ids = Arc::unwrap_or_clone(request.token_ids); + let input = request_input(&mut request, limits)?; + let sampling = &request.sampling_options; + let stop = &request.stop_conditions; + let output = &request.output_options; Ok(pb::GenerateRequest { request_id: request_id.to_string(), // The OpenEngine server rejects an empty model; any non-empty name is // served by the loaded model (single-model server). model: model.to_string(), - input: Some(pb::generate_request::Input::TokenIds(pb::TokenIds { - ids: token_ids, - })), + input: Some(input), sampling: Some(pb::SamplingParams { temperature: sampling.temperature.map(f64::from), top_p: sampling.top_p.map(f64::from), @@ -168,6 +165,44 @@ pub(crate) fn build_generate_request_with_routing( }) } +fn request_input( + request: &mut PreprocessedRequest, + limits: Option, +) -> Result { + match limits.and_then(|limits| limits.supported_inputs) { + Some(inputs) if inputs.token_ids => token_ids_input(request), + Some(inputs) if inputs.text => { + let prompt = request + .extra_args + .as_ref() + .and_then(|extra| extra.get("formatted_prompt")) + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + client::invalid_argument( + "the loaded TensorRT-LLM model requires text input, but the request has no string extra_args.formatted_prompt", + ) + })?; + Ok(pb::generate_request::Input::Prompt(prompt.to_string())) + } + Some(_) => Err(client::invalid_argument( + "the loaded TensorRT-LLM model advertises neither text nor token-ID input support", + )), + // Older OpenEngine servers leave both optional capability fields unset. + None => token_ids_input(request), + } +} + +fn token_ids_input( + request: &mut PreprocessedRequest, +) -> Result { + if request.token_ids.is_empty() { + return Err(client::invalid_argument("token_ids must not be empty")); + } + Ok(pb::generate_request::Input::TokenIds(pb::TokenIds { + ids: Arc::unwrap_or_clone(std::mem::take(&mut request.token_ids)), + })) +} + // Derive omitted text max_tokens from the context window. For media, only // TensorRT-LLM knows the prompt length after expansion, so leave it unset. // Remove this fallback when NVIDIA/TensorRT-LLM#16549 lands; context_length @@ -208,6 +243,7 @@ fn max_tokens( const IMAGE_KEY: &str = "image_url"; const VIDEO_KEY: &str = "video_url"; +const AUDIO_KEY: &str = "audio_url"; /// The engine fetches passthrough media, so enforce Dynamo's policy before dispatch. pub(crate) fn validate_remote_media_urls( @@ -247,7 +283,7 @@ struct MediaKind { modality: pb::Modality, } -const MEDIA_KINDS: [MediaKind; 2] = [ +const MEDIA_KINDS: [MediaKind; 3] = [ MediaKind { key: IMAGE_KEY, name: "image", @@ -258,6 +294,11 @@ const MEDIA_KINDS: [MediaKind; 2] = [ name: "video", modality: pb::Modality::Video, }, + MediaKind { + key: AUDIO_KEY, + name: "audio", + modality: pb::Modality::Audio, + }, ]; fn media_kind(key: &str) -> Option { @@ -666,9 +707,6 @@ fn validate_request( mode: DisaggregationMode, allow_dp_rank: bool, ) -> Result<(), DynamoError> { - if request.token_ids.is_empty() { - return Err(client::invalid_argument("token_ids must not be empty")); - } if request.prompt_embeds.is_some() { return Err(client::invalid_argument( "prompt embeddings are not supported by the TensorRT-LLM sidecar", diff --git a/lib/sidecar/trtllm/src/convert/request_tests.rs b/lib/sidecar/trtllm/src/convert/request_tests.rs index aa0ef13f47de..07c8f49bfc0b 100644 --- a/lib/sidecar/trtllm/src/convert/request_tests.rs +++ b/lib/sidecar/trtllm/src/convert/request_tests.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use super::*; -use crate::client::SupportedMedia; +use crate::client::{SupportedInputs, SupportedMedia}; use crate::test_fixtures::*; use dynamo_backend_common::MultimodalItemRef; use serde_json::json; @@ -78,6 +78,65 @@ fn request_maps_sampling_stop_and_output_fields() { } } +#[test] +fn text_input_model_receives_the_formatted_audio_prompt() { + let mut req = request(); + req.token_ids = Arc::new(Vec::new()); + req.multi_modal_data = Some(HashMap::from([( + "audio_url".to_string(), + vec![MultimodalData::RawUrl( + "https://example.com/audio.wav".to_string(), + )], + )])); + let task_prompt = "<|startoftranscript|><|fr|><|translate|><|notimestamps|>"; + req.extra_args = Some(json!({"formatted_prompt": task_prompt})); + let mut model_limits = all_media_limits(4096).expect("limits"); + model_limits.supported_inputs = Some(SupportedInputs { + text: true, + token_ids: false, + }); + + let proto = build_generate_request(req, "req", "model", Some(model_limits), AGG) + .expect("text-input audio request"); + assert_eq!( + proto.input, + Some(pb::generate_request::Input::Prompt(task_prompt.to_string())) + ); +} + +#[test] +fn text_input_model_accepts_an_empty_formatted_prompt() { + let mut req = request(); + req.token_ids = Arc::new(Vec::new()); + req.extra_args = Some(json!({"formatted_prompt": ""})); + let mut model_limits = limits(4096).expect("limits"); + model_limits.supported_inputs = Some(SupportedInputs { + text: true, + token_ids: false, + }); + + let proto = build_generate_request(req, "req", "model", Some(model_limits), AGG) + .expect("an empty prompt lets the engine select its model default"); + assert_eq!( + proto.input, + Some(pb::generate_request::Input::Prompt(String::new())) + ); +} + +#[test] +fn text_input_model_requires_a_formatted_prompt() { + let req = request(); + let mut model_limits = limits(4096).expect("limits"); + model_limits.supported_inputs = Some(SupportedInputs { + text: true, + token_ids: false, + }); + + let error = build_generate_request(req, "req", "model", Some(model_limits), AGG) + .expect_err("a token-only preprocessed request cannot serve a text-input model"); + assert!(error.to_string().contains("extra_args.formatted_prompt")); +} + #[test] fn image_urls_data_uris_and_uuids_are_mapped_in_order() { let mut req = request(); @@ -145,7 +204,7 @@ fn inline_media_conversion_preserves_the_payload_allocation() { } #[test] -fn mixed_image_and_video_media_follow_prompt_order() { +fn mixed_media_follow_prompt_order() { let mut req = request(); req.multi_modal_data = Some(HashMap::from([ ( @@ -162,6 +221,12 @@ fn mixed_image_and_video_media_follow_prompt_order() { MultimodalData::RawUrl("data:video/mp4;base64,dmlkZW8=".to_string()), ], ), + ( + "audio_url".to_string(), + vec![MultimodalData::RawUrl( + "data:audio/wav;base64,YXVkaW8=".to_string(), + )], + ), ])); req.multi_modal_uuids = Some(HashMap::from([ ( @@ -172,6 +237,7 @@ fn mixed_image_and_video_media_follow_prompt_order() { "video_url".to_string(), vec![Some("video-0".to_string()), Some("video-1".to_string())], ), + ("audio_url".to_string(), vec![Some("audio-0".to_string())]), ])); req.multi_modal_order = Some(vec![ MultimodalItemRef { @@ -182,6 +248,10 @@ fn mixed_image_and_video_media_follow_prompt_order() { modality: "image_url".to_string(), index: 0, }, + MultimodalItemRef { + modality: "audio_url".to_string(), + index: 0, + }, MultimodalItemRef { modality: "video_url".to_string(), index: 1, @@ -192,7 +262,7 @@ fn mixed_image_and_video_media_follow_prompt_order() { }, ]); - let proto = build_generate_request(req, "req", "model", image_video_limits(4096), AGG) + let proto = build_generate_request(req, "req", "model", all_media_limits(4096), AGG) .expect("mixed media request"); assert_eq!( @@ -204,6 +274,7 @@ fn mixed_image_and_video_media_follow_prompt_order() { [ pb::Modality::Video as i32, pb::Modality::Image as i32, + pb::Modality::Audio as i32, pb::Modality::Video as i32, pb::Modality::Image as i32, ] @@ -214,7 +285,7 @@ fn mixed_image_and_video_media_follow_prompt_order() { .iter() .map(|item| item.uuid.as_str()) .collect::>(), - ["video-0", "image-0", "video-1", ""] + ["video-0", "image-0", "audio-0", "video-1", ""] ); } @@ -279,7 +350,11 @@ fn mm_routing_hashes_are_not_used_as_engine_cache_uuids() { #[test] fn media_requires_the_loaded_models_advertised_capability() { - for (key, modality) in [("image_url", "image"), ("video_url", "video")] { + for (key, modality) in [ + ("image_url", "image"), + ("video_url", "video"), + ("audio_url", "audio"), + ] { let mut req = request(); req.multi_modal_data = Some(HashMap::from([( key.to_string(), @@ -310,6 +385,7 @@ fn images_reject_unknown_or_overridden_engine_windows() { supported_media: Some(SupportedMedia { image: true, video: true, + ..Default::default() }), ..Default::default() }); @@ -321,18 +397,37 @@ fn images_reject_unknown_or_overridden_engine_windows() { } #[test] -fn unsupported_media_modalities_and_sources_are_rejected() { +fn image_video_models_reject_audio() { let mut req = request(); req.multi_modal_data = Some(HashMap::from([( "audio_url".to_string(), + vec![MultimodalData::RawUrl( + "https://example.com/audio.wav".to_string(), + )], + )])); + let error = build_generate_request(req, "req", "model", image_video_limits(4096), AGG) + .expect_err("a model without audio support must reject audio"); + assert!( + error + .to_string() + .contains("does not advertise audio support"), + "unexpected error: {error}" + ); +} + +#[test] +fn unsupported_media_modalities_and_sources_are_rejected() { + let mut req = request(); + req.multi_modal_data = Some(HashMap::from([( + "future_media".to_string(), vec![MultimodalData::RawUrl( "https://example.com/media".to_string(), )], )])); - let error = build_generate_request(req, "req", "model", image_limits(4096), AGG) - .expect_err("audio is not supported yet"); + let error = build_generate_request(req, "req", "model", all_media_limits(4096), AGG) + .expect_err("unsupported modalities must be rejected"); assert!( - error.to_string().contains("audio_url"), + error.to_string().contains("future_media"), "unexpected error: {error}" ); @@ -532,6 +627,7 @@ fn the_derived_max_tokens_respects_the_engines_output_cap() { Some(ModelLimits { context_length: Some(4096), max_output_tokens: Some(256), + supported_inputs: None, supported_media: Some(SupportedMedia::default()), ..Default::default() }), diff --git a/lib/sidecar/trtllm/src/test_fixtures.rs b/lib/sidecar/trtllm/src/test_fixtures.rs index 554c1cfc7cc7..dad103c5375f 100644 --- a/lib/sidecar/trtllm/src/test_fixtures.rs +++ b/lib/sidecar/trtllm/src/test_fixtures.rs @@ -56,6 +56,7 @@ pub(crate) fn limits(context_length: u32) -> Option { context_length: Some(context_length), engine_context_length: Some(context_length), max_output_tokens: None, + supported_inputs: None, supported_media: Some(SupportedMedia::default()), }) } @@ -65,6 +66,7 @@ pub(crate) fn image_limits(context_length: u32) -> Option { context_length: Some(context_length), engine_context_length: Some(context_length), max_output_tokens: None, + supported_inputs: None, supported_media: Some(SupportedMedia { image: true, ..Default::default() @@ -77,9 +79,25 @@ pub(crate) fn image_video_limits(context_length: u32) -> Option { context_length: Some(context_length), engine_context_length: Some(context_length), max_output_tokens: None, + supported_inputs: None, supported_media: Some(SupportedMedia { image: true, video: true, + ..Default::default() + }), + }) +} + +pub(crate) fn all_media_limits(context_length: u32) -> Option { + Some(ModelLimits { + context_length: Some(context_length), + engine_context_length: Some(context_length), + max_output_tokens: None, + supported_inputs: None, + supported_media: Some(SupportedMedia { + image: true, + video: true, + audio: true, }), }) } diff --git a/lib/sidecar/trtllm/src/tests/e2e.rs b/lib/sidecar/trtllm/src/tests/e2e.rs index dd03683375c6..c7c8487aab46 100644 --- a/lib/sidecar/trtllm/src/tests/e2e.rs +++ b/lib/sidecar/trtllm/src/tests/e2e.rs @@ -513,3 +513,87 @@ async fn e2e_real_openengine_video() { "the engine must report the expanded video prompt" ); } + +/// Set `TRTLLM_E2E_AUDIO_DATA_URI` or `TRTLLM_E2E_AUDIO_DATA_URI_FILE` to a data URI. +#[tokio::test] +#[ignore = "requires a live audio-capable TensorRT-LLM OpenEngine server"] +async fn e2e_real_openengine_audio() { + let endpoint = std::env::var("TRTLLM_E2E_ENDPOINT") + .expect("set TRTLLM_E2E_ENDPOINT, e.g. http://127.0.0.1:50051"); + let model = + std::env::var("TRTLLM_E2E_MODEL").unwrap_or_else(|_| "openai/whisper-tiny".to_string()); + let audio = match std::env::var("TRTLLM_E2E_AUDIO_DATA_URI_FILE") { + Ok(path) => std::fs::read_to_string(path).expect("read audio data URI"), + Err(_) => std::env::var("TRTLLM_E2E_AUDIO_DATA_URI") + .expect("set TRTLLM_E2E_AUDIO_DATA_URI or TRTLLM_E2E_AUDIO_DATA_URI_FILE"), + }; + let prompt = std::env::var("TRTLLM_E2E_AUDIO_PROMPT") + .unwrap_or_else(|_| "<|startoftranscript|><|en|><|transcribe|><|notimestamps|>".into()); + let engine = TrtllmSidecarEngine::new( + GrpcEndpoint::parse(&endpoint, "--grpc-endpoint").expect("valid endpoint"), + transport(1), + ConfiguredModel { + source: model.clone(), + context_length: None, + ..Default::default() + }, + AGG, + ); + engine + .start(0) + .await + .expect("start against the live audio-capable server"); + + let mut request = PreprocessedRequest::builder() + .model(model) + .token_ids(vec![]) + .stop_conditions(StopConditions { + max_tokens: Some(96), + ..Default::default() + }) + .sampling_options(SamplingOptions { + temperature: Some(0.0), + ..Default::default() + }) + .output_options(OutputOptions::default()) + .build() + .expect("request"); + request.extra_args = Some(json!({"formatted_prompt": prompt})); + request.multi_modal_data = Some(HashMap::from([( + "audio_url".to_string(), + vec![MultimodalData::RawUrl(audio)], + )])); + + let outputs = collect(&engine, request).await; + let generated = outputs + .iter() + .flat_map(|output| output.token_ids.iter().copied()) + .collect::>(); + let terminal = outputs + .iter() + .find(|output| output.finish_reason.is_some()) + .expect("terminal output"); + let usage = terminal.completion_usage.as_ref().expect("terminal usage"); + eprintln!( + "[e2e-audio] generated={generated:?}; usage: prompt={}, completion={}", + usage.prompt_tokens, usage.completion_tokens + ); + + if let Ok(path) = std::env::var("TRTLLM_E2E_RESULT") { + std::fs::write( + path, + serde_json::to_vec(&json!({ + "token_ids": generated, + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + })) + .expect("serialize audio result"), + ) + .expect("write audio result"); + } + assert!( + usage.completion_tokens > 0, + "expected completion token usage" + ); + assert!(!generated.is_empty(), "expected audio-conditioned output"); +}