From 72f828a4d08243e63638fb86aee586de0aa0479c Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 15:34:39 -0700 Subject: [PATCH 01/20] feat(openengine): add sibling protocol server Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 1 + scripts/install_openengine.py | 81 ++ .../_torch/models/modeling_gemma4mm.py | 6 + .../_torch/models/modeling_nemotron_nano.py | 6 + tensorrt_llm/_torch/models/modeling_phi4mm.py | 33 +- .../_torch/models/modeling_qwen3vl.py | 6 + tensorrt_llm/_utils.py | 5 +- tensorrt_llm/commands/serve.py | 150 +- tensorrt_llm/grpc/grpc_request_manager.py | 54 +- tensorrt_llm/inputs/registry.py | 118 +- tensorrt_llm/openengine/README.md | 21 + tensorrt_llm/openengine/__init__.py | 8 + tensorrt_llm/openengine/_schema_pin.py | 8 + tensorrt_llm/openengine/converters.py | 392 ++++++ tensorrt_llm/openengine/lora_registry.py | 120 ++ tensorrt_llm/openengine/server.py | 129 ++ tensorrt_llm/openengine/servicer.py | 1226 +++++++++++++++++ .../kv_cache_manager_v2/_event_manager.py | 9 + tensorrt_llm/serve/kv_event_fanout.py | 134 ++ tensorrt_llm/serve/openai_server.py | 41 +- tensorrt_llm/serve/request_tracker.py | 156 +++ tensorrt_llm/serve/stats_fanout.py | 81 ++ .../_torch/modeling/test_modeling_phi4mm.py | 17 + .../references/trtllm_serve_cli.yaml | 18 + tests/unittest/openengine/test_converters.py | 225 +++ tests/unittest/openengine/test_fanout.py | 105 ++ .../unittest/openengine/test_lora_registry.py | 89 ++ .../openengine/test_native_grpc_compat.py | 49 + .../openengine/test_request_tracker.py | 73 + tests/unittest/openengine/test_server.py | 86 ++ tests/unittest/openengine/test_servicer.py | 775 +++++++++++ 31 files changed, 4117 insertions(+), 105 deletions(-) create mode 100644 OPENENGINE_COMMIT create mode 100644 scripts/install_openengine.py create mode 100644 tensorrt_llm/openengine/README.md create mode 100644 tensorrt_llm/openengine/__init__.py create mode 100644 tensorrt_llm/openengine/_schema_pin.py create mode 100644 tensorrt_llm/openengine/converters.py create mode 100644 tensorrt_llm/openengine/lora_registry.py create mode 100644 tensorrt_llm/openengine/server.py create mode 100644 tensorrt_llm/openengine/servicer.py create mode 100644 tensorrt_llm/serve/kv_event_fanout.py create mode 100644 tensorrt_llm/serve/request_tracker.py create mode 100644 tensorrt_llm/serve/stats_fanout.py create mode 100644 tests/unittest/openengine/test_converters.py create mode 100644 tests/unittest/openengine/test_fanout.py create mode 100644 tests/unittest/openengine/test_lora_registry.py create mode 100644 tests/unittest/openengine/test_native_grpc_compat.py create mode 100644 tests/unittest/openengine/test_request_tracker.py create mode 100644 tests/unittest/openengine/test_server.py create mode 100644 tests/unittest/openengine/test_servicer.py diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT new file mode 100644 index 000000000000..2fb5de7752ae --- /dev/null +++ b/OPENENGINE_COMMIT @@ -0,0 +1 @@ +cea19cb06acf03c911b84d5c147e519b60dd92a6 diff --git a/scripts/install_openengine.py b/scripts/install_openengine.py new file mode 100644 index 000000000000..89f8f3e3ba93 --- /dev/null +++ b/scripts/install_openengine.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Verify and install the exact local OpenEngine Python sibling package.""" + +import argparse +import re +import runpy +import subprocess +import sys +from pathlib import Path + + +def _run(*args: str, cwd: Path) -> str: + return subprocess.run( + args, + cwd=cwd, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--sibling", + type=Path, + help="OpenEngine checkout (default: ../openengine-trtllm)", + ) + parser.add_argument( + "--verify-only", + action="store_true", + help="Verify the sibling without invoking pip", + ) + args = parser.parse_args() + + root = Path(__file__).resolve().parents[1] + sibling = (args.sibling or root.parent / "openengine-trtllm").resolve() + expected = (root / "OPENENGINE_COMMIT").read_text(encoding="utf-8").strip() + if re.fullmatch(r"[0-9a-f]{40}", expected) is None: + raise RuntimeError("OPENENGINE_COMMIT must contain one full lowercase Git SHA") + packaged = runpy.run_path(str(root / "tensorrt_llm" / "openengine" / "_schema_pin.py"))[ + "OPENENGINE_COMMIT" + ] + if packaged != expected: + raise RuntimeError( + f"Packaged OpenEngine pin is {packaged}, but OPENENGINE_COMMIT contains {expected}" + ) + + actual = _run("git", "rev-parse", "HEAD", cwd=sibling) + if actual != expected: + raise RuntimeError(f"OpenEngine sibling is at {actual}, but TensorRT-LLM pins {expected}") + dirty = _run( + "git", + "status", + "--porcelain", + "--", + "packages/python", + "proto", + cwd=sibling, + ) + if dirty: + raise RuntimeError("OpenEngine Python/proto sources have uncommitted changes") + + package = sibling / "packages" / "python" + if not (package / "pyproject.toml").is_file(): + raise RuntimeError(f"OpenEngine Python package is missing: {package}") + if not args.verify_only: + subprocess.run( + [sys.executable, "-m", "pip", "install", "-e", str(package)], + cwd=root, + check=True, + ) + + print(f"Verified OpenEngine {expected}") + print(f"export OPENENGINE_SCHEMA_RELEASE={expected}") + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 0f5a340c0a9a..16b2129442eb 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -209,6 +209,12 @@ class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInpu # image, when present) can be split across chunks safely. mm_bidirectional_blocks = False + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image", "audio") + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + def __init__( self, model_path: str, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py index 386b30716d08..e39376945942 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py @@ -954,6 +954,12 @@ def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tenso class NanoV2VLInputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): supports_token_id_mm_expansion: ClassVar[bool] = True + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image", "video", "audio") + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + def __init__( self, model_path: str, diff --git a/tensorrt_llm/_torch/models/modeling_phi4mm.py b/tensorrt_llm/_torch/models/modeling_phi4mm.py index 2fbe96ac03c1..5e45078e47f6 100644 --- a/tensorrt_llm/_torch/models/modeling_phi4mm.py +++ b/tensorrt_llm/_torch/models/modeling_phi4mm.py @@ -36,6 +36,7 @@ ExtraProcessedInputs, MultimodalPlaceholderMetadata, MultimodalPlaceholderPlacement, TextPrompt, register_input_processor) +from ...inputs.registry import MultimodalLoraSpec from ...logger import logger from ...lora_helper import LoraConfig from ...sampling_params import SamplingParams @@ -760,6 +761,36 @@ def forward(self, multimodal_params: List[MultimodalParams], class Phi4MMInputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder): + def get_model_owned_lora_identities(self) -> dict[str, int]: + return {"vision-lora": 0, "speech-lora": 1} + + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image", "audio") + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + + def get_required_lora_spec( + self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + requested = set(modalities) + if {"image", "audio"}.issubset(requested): + raise ValueError( + "Phi-4 multimodal requests cannot combine image and audio because they " + "require different built-in LoRA adapters") + if "image" in requested: + return MultimodalLoraSpec( + name="vision-lora", + adapter_id=0, + path=os.path.join(self._model_path, "vision-lora"), + ) + if "audio" in requested: + return MultimodalLoraSpec( + name="speech-lora", + adapter_id=1, + path=os.path.join(self._model_path, "speech-lora"), + ) + return None + def __init__(self, model_path: str, config: transformers.PretrainedConfig, @@ -1092,7 +1123,7 @@ def forward( else: raise NotImplementedError( "Phi-4-multimodal does not support disaggregated inference yet. Please unset " - f"the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." + "the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'." ) mm_embedding = find_input_mm_embeds( mm_embedding, multimodal_params[:num_context_requests]) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 95da74e5d3bd..3943560deef0 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -402,6 +402,12 @@ def get_preferred_media_io_kwargs(self) -> Dict[str, Dict[str, Any]]: # the per-frame CHW-float conversion in the IO loader. return {"video": {"format": "np"}} + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image", "video") + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return ("image", "video") + def build_disagg_prefill_multimodal_inputs( self, inputs: TextPrompt, mm_handles: List[Dict[str, Any]] ) -> DisaggPrefillMultimodalInputs: diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index cf242cd205c9..c7605056a923 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1223,7 +1223,10 @@ def _stored_block_to_json(data): KVCacheEventSerializer._unique_tokens_to_json(token) for token in data.tokens ], - # "lora_id": data.lora_id, # TODO (shreyasm): enable serialization of lora_id + "lora_id": + getattr(data, "lora_id", None), + "lora_name": + getattr(data, "lora_name", None), "cache_salt": data.cache_salt, "cache_level": diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index bc949b3db948..8817a918adc6 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -122,8 +122,7 @@ def _signal_handler_cleanup_child(signum, frame): def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): - """ - Check if a parameter should be explicitly included in llm_args. + """Check if a parameter should be explicitly included in llm_args. Returns True if parameter is either: 1. Always required (core params that must be present), OR @@ -387,7 +386,9 @@ def launch_server( served_model_name: Optional[str] = None, allow_request_chat_template: bool = False, num_input_processor_workers: int = 8, - num_media_load_workers: int = 8): + num_media_load_workers: int = 8, + openengine_host: str = "127.0.0.1", + openengine_port: Optional[int] = None): backend = llm_args["backend"] model = served_model_name or llm_args["model"] @@ -427,6 +428,51 @@ def launch_server( f"{backend} is not a known backend, check help for available options.", param_hint="backend") + request_tracker = None + kv_event_fanout = None + stats_fanout = None + openengine_server_cls = None + resolved_openengine_role = None + if openengine_port is not None: + from fastapi.responses import JSONResponse + + try: + from tensorrt_llm.openengine.server import (OpenEngineServer, + openengine_role) + except ModuleNotFoundError as error: + if error.name and error.name.startswith("openengine"): + llm.shutdown() + raise click.ClickException( + "OpenEngine support requires openengine-proto 0.2.0. " + "Run scripts/install_openengine.py against the pinned " + "local OpenEngine sibling checkout.") from error + raise + from tensorrt_llm.serve.kv_event_fanout import KvEventFanout + from tensorrt_llm.serve.request_tracker import (RequestTracker, + track_http_response) + from tensorrt_llm.serve.stats_fanout import StatsFanout + + request_tracker = RequestTracker(llm) + kv_cache_config = getattr(llm.args, "kv_cache_config", None) + event_buffer_size = getattr(kv_cache_config, + "event_buffer_max_size", 0) + if event_buffer_size > 0: + kv_event_fanout = KvEventFanout(llm, + buffer_size=event_buffer_size) + stats_buffer_size = getattr(llm.args, "iter_stats_max_iterations", + 1000) + if stats_buffer_size is None or stats_buffer_size == 0: + stats_buffer_size = 1000 + elif stats_buffer_size < 0: + stats_buffer_size = None + stats_fanout = StatsFanout(llm, buffer_size=stats_buffer_size) + openengine_server_cls = OpenEngineServer + try: + resolved_openengine_role = openengine_role(server_role) + except ValueError: + llm.shutdown() + raise + server = OpenAIServer( generator=llm, model=model, @@ -438,22 +484,76 @@ def launch_server( chat_template=chat_template, allow_request_chat_template=allow_request_chat_template, input_processor_workers=num_input_processor_workers, - media_load_workers=num_media_load_workers) + media_load_workers=num_media_load_workers, + kv_event_fanout=kv_event_fanout, + stats_fanout=stats_fanout, + shutdown_generator=openengine_server_cls is None) _apply_fastapi_middlewares(server.app, middleware) + if request_tracker is not None: + + @server.app.middleware("http") + async def openengine_process_admission(request, call_next): + if request.method != "POST": + return await call_next(request) + try: + request_tracker.begin_external() + if stats_fanout is not None: + stats_fanout.wake() + except RuntimeError: + return JSONResponse( + status_code=503, + content={"error": "TensorRT-LLM is draining"}) + response = None + try: + response = await call_next(request) + finally: + if response is None: + await request_tracker.finish_external() + return await track_http_response(response, request_tracker) + # Optionally disable GC (default: not disabled) if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": gc.disable() - uvloop.run(server(host, port, sockets=[s])) + if openengine_server_cls is None: + uvloop.run(server(host, port, sockets=[s])) + else: + + async def serve_http_and_openengine(): + openengine_server = None + try: + openengine_server = openengine_server_cls( + llm=llm, + model=model, + role=resolved_openengine_role, + host=openengine_host, + port=openengine_port, + tracker=request_tracker, + media_config=multimodal_server_config, + reasoning_parser=getattr(llm.args, "reasoning_parser", + None), + tool_parser=tool_parser, + kv_event_fanout=kv_event_fanout, + stats_fanout=stats_fanout, + ) + await openengine_server.start() + await server(host, port, sockets=[s]) + finally: + try: + if openengine_server is not None: + await openengine_server.stop() + finally: + llm.shutdown() + + uvloop.run(serve_http_and_openengine()) def launch_grpc_server(host: str, port: int, llm_args: dict, served_model_name: Optional[str] = None): - """ - Launch a gRPC server for TensorRT-LLM. + """Launch a gRPC server for TensorRT-LLM. This provides a high-performance gRPC interface designed for external routers (e.g., sgl-router) using pre-tokenized input and raw token ID output. @@ -1075,6 +1175,18 @@ def convert(self, value: Any, param: Optional["click.Parameter"], help="Run gRPC server instead of OpenAI HTTP server. " "gRPC server accepts pre-tokenized requests and returns raw token IDs.", status="prototype") +@stability_option("--openengine-host", + type=str, + default="127.0.0.1", + help="Host for the optional OpenEngine sibling gRPC server.", + status="prototype") +@stability_option( + "--openengine-port", + type=int, + default=None, + help= + "Port for the optional OpenEngine sibling gRPC server. Disabled when absent.", + status="prototype") @stability_option( "--served_model_name", type=str, @@ -1132,13 +1244,18 @@ def serve( telemetry: bool, custom_module_dirs: list[Path], chat_template: Optional[str], allow_request_chat_template: bool, middleware: tuple[str, ...], grpc: bool, enable_visual_gen: bool, - served_model_name: Optional[str], visual_gen_args: Optional[str]): + served_model_name: Optional[str], visual_gen_args: Optional[str], + openengine_host: str, openengine_port: Optional[int]): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path """ logger.set_level(log_level) + if grpc and openengine_port is not None: + raise click.UsageError( + "--openengine-port cannot be combined with dedicated --grpc mode") + if moe_cluster_parallel_size is not None: logger.warning( "--moe_cluster_parallel_size / --cluster_size is deprecated and " @@ -1277,6 +1394,12 @@ def _serve_llm(): multimodal_server_config = MultimodalServerConfig( media_io_kwargs=parsed_media_io_kwargs) + if openengine_port is not None and server_role is not None: + role_name = getattr(server_role, "name", str(server_role)).upper() + if role_name in ("MM_ENCODER", "VISUAL_GEN"): + raise click.UsageError( + f"OpenEngine does not support server role {role_name!r}") + if grpc: # gRPC mode: launch gRPC server instead of OpenAI HTTP server # Check for unsupported arguments that are silently ignored in gRPC mode @@ -1323,7 +1446,9 @@ def _serve_llm(): served_model_name=served_model_name, allow_request_chat_template=allow_request_chat_template, num_input_processor_workers=num_input_processor_workers, - num_media_load_workers=num_media_load_workers) + num_media_load_workers=num_media_load_workers, + openengine_host=openengine_host, + openengine_port=openengine_port) def _serve_visual_gen(): parsed_visual_gen_args = (VisualGenArgs.from_yaml(visual_gen_args) @@ -1338,6 +1463,9 @@ def _serve_visual_gen(): is_visual_gen = (enable_visual_gen or visual_gen_args is not None or get_is_diffusion_only_model(model)) if is_visual_gen: + if openengine_port is not None: + raise click.UsageError( + "--openengine-port is not supported by the VisualGen server") _serve_visual_gen() else: _serve_llm() @@ -1665,7 +1793,6 @@ def disaggregated( schedule_style: str, ): """Running server in disaggregated mode""" - logger.set_level(log_level) if metrics_log_interval != 0: @@ -1756,7 +1883,6 @@ def set_cuda_device(): status="beta") def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): """Launching disaggregated MPI worker""" - from tensorrt_llm._utils import mpi_rank if os.environ.get(DisaggLauncherEnvs. TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT) != "1": @@ -1944,7 +2070,7 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str, f"Child process {_child_p_global.pid} failed to be killed even after 30s." ) assert _child_p_global.poll( - ) is not None, f"the subprocess should be terminated" + ) is not None, "the subprocess should be terminated" # Check if the process was launched and assert it's terminated if _child_p_global and hasattr(_child_p_global, diff --git a/tensorrt_llm/grpc/grpc_request_manager.py b/tensorrt_llm/grpc/grpc_request_manager.py index ca2e59eb1f36..9cb4868682c6 100644 --- a/tensorrt_llm/grpc/grpc_request_manager.py +++ b/tensorrt_llm/grpc/grpc_request_manager.py @@ -34,6 +34,7 @@ from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams +from tensorrt_llm.serve.request_tracker import RequestTracker from . import trtllm_service_pb2 as pb2 @@ -59,8 +60,9 @@ def __init__(self, llm: Any): llm: The TensorRT-LLM LLM instance (tensorrt_llm.LLM or tensorrt_llm._tensorrt_engine.LLM) """ self.llm = llm - # Track active requests: request_id -> GenerationResult - self._rid_to_result: Dict[str, GenerationResult] = {} + self._request_tracker = RequestTracker(llm) + # Kept as an alias for callers which inspect native gRPC state. + self._rid_to_result = self._request_tracker.active_requests logger.info("GrpcRequestManager initialized") @@ -93,6 +95,7 @@ async def generate( GenerationResult objects containing token IDs (text will be empty because detokenize=False) """ + admitted = False try: # Submit to LLM.generate_async which returns a GenerationResult # that is an async iterator @@ -111,7 +114,15 @@ async def generate( ) # Track the result for potential abort - self._rid_to_result[request_id] = gen_result + try: + self._request_tracker.admit(request_id, gen_result) + except (RuntimeError, ValueError): + try: + gen_result.abort() + except (RuntimeError, AssertionError) as error: + logger.warning("Failed to abort rejected request %s: %s", request_id, error) + raise + admitted = True # Iterate over the async generator # GenerationResult implements __aiter__ and __anext__ @@ -130,7 +141,8 @@ async def generate( raise finally: # Cleanup tracking - self._rid_to_result.pop(request_id, None) + if admitted: + await self._request_tracker.finish(request_id) async def abort(self, request_id: str) -> bool: """Abort a running request. @@ -141,22 +153,13 @@ async def abort(self, request_id: str) -> bool: Returns: True if request was found and aborted, False otherwise """ - gen_result = self._rid_to_result.get(request_id) - - if gen_result is None: + if request_id not in self._rid_to_result: logger.debug(f"Abort: request {request_id} not found (may have already completed)") return False - - try: - # GenerationResult has an abort() method - gen_result.abort() - self._rid_to_result.pop(request_id, None) + aborted = await self._request_tracker.abort(request_id) + if aborted: logger.info(f"Request {request_id} aborted") - return True - except Exception as e: - logger.error(f"Error aborting request {request_id}: {e}") - self._rid_to_result.pop(request_id, None) - return False + return aborted async def health_check(self) -> Tuple[bool, str]: """Check if the engine is healthy. @@ -168,22 +171,7 @@ async def health_check(self) -> Tuple[bool, str]: if self.llm is None: return False, "LLM not initialized" - # Check executor health (includes error queue, MPI worker - # liveness, and fatal error state — not just doing_shutdown). - if hasattr(self.llm, "_executor"): - if self.llm._executor is None: - return False, "Executor is not available" - if not self.llm._executor.check_health(): - error_msg = "Executor is unhealthy" - if self.llm._executor._fatal_error is not None: - exc = self.llm._executor._fatal_error - lines = str(exc).splitlines() - short = (lines[0] if lines else type(exc).__name__)[:200] - error_msg = f"{type(exc).__name__}: {short}" - logger.error(f"Health check fatal error: {repr(exc)}") - return False, error_msg - - return True, "OK" + return await self._request_tracker.health() except Exception as e: logger.error(f"Health check error: {e}") return False, f"Error: {e}" diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 58ca102d35be..ec4b6df03478 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -28,9 +28,17 @@ ExtraProcessedInputs = Dict[str, Any] +@dataclass(frozen=True) +class MultimodalLoraSpec: + """Model-required LoRA adapter for a set of input modalities.""" + + name: str + adapter_id: int + path: str + + class InputProcessor(Protocol): - """ - Protocol for InputProcessor classes. + """Protocol for InputProcessor classes. InputProcessor's functions are more relevant to multimodal use cases: - Preprocess: extra steps to manipulate the prompts. - Forward: the main logic to process the inputs. In multimodal cases, this may run a multimodal encoder model. @@ -121,8 +129,7 @@ def __call__( class BaseMultimodalInputProcessor(ABC): - """ - Base class for multimodal input processors with default implementations + """Base class for multimodal input processors with default implementations of get_num_tokens_per_image and get_num_tokens_per_video methods. This class provides default implementations that work with most AutoProcessor-based @@ -150,6 +157,24 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False + def get_openengine_modalities(self) -> tuple[str, ...]: + """Modalities this processor accepts through aggregated OpenEngine.""" + return ("image", ) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + """Modalities supported in context-first OpenEngine P/D requests.""" + return () + + def get_required_lora_spec( + self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + """Return a model-owned adapter required to process these modalities.""" + del modalities + return None + + def get_model_owned_lora_identities(self) -> dict[str, int]: + """Names and IDs reserved for adapters selected by the model processor.""" + return {} + def __init__(self, model_path, config, @@ -373,16 +398,14 @@ class flag and implement both of the following hooks; otherwise raises @property def use_fast(self) -> bool: - """ - Whether to use fast tokenizer for AutoProcessor. + """Whether to use fast tokenizer for AutoProcessor. Default is True for most multimodal models. """ return self._use_fast @property def multimodal_hashing_supported(self) -> Optional[bool]: - """ - Whether multimodal hashing is supported for this processor. + """Whether multimodal hashing is supported for this processor. Returns None if unknown (will be detected at runtime), True if supported, False if not supported. @@ -474,8 +497,7 @@ def get_mm_special_token_ids(self) -> Optional[Tensor]: @property def get_num_multimodal_tokens(self): - """ - Get the Hugging Face processor's '_get_num_multimodal_tokens' method. + """Get the Hugging Face processor's '_get_num_multimodal_tokens' method. """ if hasattr(self.processor, '_get_num_multimodal_tokens'): return self.processor._get_num_multimodal_tokens @@ -491,8 +513,7 @@ def get_num_tokens_per_image( image: Union[Image.Image, torch.Tensor], **kwargs, ): - """ - Calculate the number of tokens generated for an image. + """Calculate the number of tokens generated for an image. Delegates to the Hugging Face processor's ``_get_num_multimodal_tokens``. @@ -521,8 +542,7 @@ def get_num_tokens_per_video( video_audio: Optional[Any] = None, **kwargs, ): - """ - Calculate the number of tokens generated for a video. + """Calculate the number of tokens generated for a video. Delegates to the Hugging Face processor's ``_get_num_multimodal_tokens``; a fallback treats the video as a stack of frames if the HF processor @@ -643,10 +663,9 @@ def get_dummy_mm_data_for_tokens( class MultimodalPlaceholderPlacement(enum.Enum): - """ - The placement of the multimodal placeholder in the prompt. Valid values are: - - BEFORE_TEXT: the placeholders are placed before the text prompt. - - AFTER_TEXT: the placeholders are placed after the text prompt. + """The placement of the multimodal placeholder in the prompt. Valid values are: + - BEFORE_TEXT: the placeholders are placed before the text prompt. + - AFTER_TEXT: the placeholders are placed after the text prompt. """ INVALID = -1 BEFORE_TEXT = 0 @@ -655,31 +674,30 @@ class MultimodalPlaceholderPlacement(enum.Enum): @dataclass(frozen=True) class MultimodalPlaceholderMetadata: - """ - Metadata for the multimodal placeholder. It has 5 components: - - placeholder_map: - A mapping from modality to placeholder string. - Modality can be "image", "video", "audio", etc. - - placeholder_placement: - The placement of the placeholders, e.g. before or after the text prompt. - Only used when interleave_placeholders is False (the default). - Ignored when interleave_placeholders is True. - - placeholders_separator: - The separator between the placeholders, e.g. some models use "\n" to separate the placeholders. - - content_format: - Optional override for the content format expected by the chat template. - ContentFormat.OPENAI means the template handles multimodal content dicts natively. - ContentFormat.STRING means the template expects plain string content. - ContentFormat.PASSTHROUGH skips chat template rendering entirely. - None means auto-detect at runtime via Jinja AST analysis. - - interleave_placeholders: - When True and content_parts is available, placeholders are inserted - at the exact media positions within the text (interleaved). - In this mode, placeholder_placement is ignored - the position of - each placeholder is determined by where the media appears in the - user's message. - When False (default), placeholders are bulk-prepended or appended - according to placeholder_placement. + """Metadata for the multimodal placeholder. It has 5 components: + - placeholder_map: + A mapping from modality to placeholder string. + Modality can be "image", "video", "audio", etc. + - placeholder_placement: + The placement of the placeholders, e.g. before or after the text prompt. + Only used when interleave_placeholders is False (the default). + Ignored when interleave_placeholders is True. + - placeholders_separator: + The separator between the placeholders, e.g. some models use "\n" to separate the placeholders. + - content_format: + Optional override for the content format expected by the chat template. + ContentFormat.OPENAI means the template handles multimodal content dicts natively. + ContentFormat.STRING means the template expects plain string content. + ContentFormat.PASSTHROUGH skips chat template rendering entirely. + None means auto-detect at runtime via Jinja AST analysis. + - interleave_placeholders: + When True and content_parts is available, placeholders are inserted + at the exact media positions within the text (interleaved). + In this mode, placeholder_placement is ignored - the position of + each placeholder is determined by where the media appears in the + user's message. + When False (default), placeholders are bulk-prepended or appended + according to placeholder_placement. """ placeholder_map: Dict[str, str] = field(default_factory=dict) placeholder_placement: MultimodalPlaceholderPlacement = MultimodalPlaceholderPlacement.AFTER_TEXT @@ -689,8 +707,7 @@ class MultimodalPlaceholderMetadata: class MultimodalPlaceholderRegistry: - """ - Registry for the multimodal models to keep track of the placeholder information. + """Registry for the multimodal models to keep track of the placeholder information. """ def __init__(self) -> None: @@ -806,8 +823,7 @@ def __init__(self) -> None: def support_multimodal_disaggregated(model_cls: Type[nn.Module]): - """ - Model-class decorator to declare support for multimodal disaggregated inputs. + """Model-class decorator to declare support for multimodal disaggregated inputs. Apply this to a model class AFTER its input processor has been registered via @register_input_processor. The decorator will locate the processor class, @@ -840,10 +856,9 @@ def register_input_processor( processor_cls: Type[InputProcessor], model_type: str, placeholder_metadata: MultimodalPlaceholderMetadata = None): - """ - Register an input processor to a model class. + """Register an input processor to a model class. - NOTE: + Note: 1. Since this API is only used for multimodal models, we are checking the model type only for that. 2. If this is used for other models in the future, this logic needs to be @@ -1066,8 +1081,7 @@ def create_input_processor_with_hash( def multimodal_hashing_process( inputs: TextPrompt, sampling_params: SamplingParams ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: - """ - Process multimodal hashing for media tokens if possible. + """Process multimodal hashing for media tokens if possible. Delegates the raw `(token_ids, extra_processed_inputs)` production to the input processor's `__call__` (which itself dispatches between the diff --git a/tensorrt_llm/openengine/README.md b/tensorrt_llm/openengine/README.md new file mode 100644 index 000000000000..4a7063ea4de6 --- /dev/null +++ b/tensorrt_llm/openengine/README.md @@ -0,0 +1,21 @@ + + + +# TensorRT-LLM OpenEngine sibling server + +`trtllm-serve` can expose OpenEngine alongside its normal HTTP server with +`--openengine-port`. The feature remains disabled when the flag is absent and +keeps the OpenEngine bindings out of TensorRT-LLM's required dependencies. + +The exact sibling source revision is recorded in `OPENENGINE_COMMIT`. Verify +and install its generated Python package with: + +```bash +python scripts/install_openengine.py +python -m pip install -e . +``` + +The installer rejects a different or dirty sibling package/proto checkout. It +prints the required `OPENENGINE_SCHEMA_RELEASE` export; sibling startup fails +closed unless that value exactly matches `OPENENGINE_COMMIT`. No registry +publication is assumed. diff --git a/tensorrt_llm/openengine/__init__.py b/tensorrt_llm/openengine/__init__.py new file mode 100644 index 000000000000..ca230e3fecdd --- /dev/null +++ b/tensorrt_llm/openengine/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine server integration for TensorRT-LLM.""" + +from .server import OpenEngineServer + +__all__ = ["OpenEngineServer"] diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py new file mode 100644 index 000000000000..8a5cfa30ece3 --- /dev/null +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine source identity embedded in TensorRT-LLM distributions.""" + +OPENENGINE_COMMIT = "cea19cb06acf03c911b84d5c147e519b60dd92a6" + +__all__ = ["OPENENGINE_COMMIT"] diff --git a/tensorrt_llm/openengine/converters.py b/tensorrt_llm/openengine/converters.py new file mode 100644 index 000000000000..e934a5059138 --- /dev/null +++ b/tensorrt_llm/openengine/converters.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Conversions between OpenEngine messages and TensorRT-LLM API objects.""" + +import asyncio +import base64 +import hashlib +import json +import re +from dataclasses import asdict, is_dataclass +from typing import Any + +from google.protobuf.json_format import MessageToDict +from openengine.v1 import generation_pb2, input_pb2, kv_pb2 + +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.executor.result import Logprob +from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams + +HANDOFF_ATTRIBUTE = "tensorrt_llm.disaggregated_params.v1" +_MODALITY_NAMES = { + input_pb2.MODALITY_UNSPECIFIED: "image", + input_pb2.MODALITY_IMAGE: "image", + input_pb2.MODALITY_VIDEO: "video", + input_pb2.MODALITY_AUDIO: "audio", +} + + +def _optional(message: object, field: str) -> Any | None: + return getattr(message, field) if message.HasField(field) else None + + +def _candidate_count(selection: object, enabled: bool) -> int | None: + if not enabled: + return None + kind = selection.WhichOneof("selection") + if kind in (None, "top_n"): + return selection.top_n if kind == "top_n" else 0 + raise ValueError("TensorRT-LLM OpenEngine supports top_n logprob selection only") + + +def to_sampling_params(request: generation_pb2.GenerateRequest) -> SamplingParams: + """Build TRT-LLM sampling params without inventing wire defaults.""" + sampling = request.sampling + stopping = request.stopping + response = request.response + kwargs: dict[str, Any] = { + "max_tokens": (32 if _optional(stopping, "max_tokens") is None else stopping.max_tokens), + "detokenize": True, + } + field_map = { + "temperature": "temperature", + "top_p": "top_p", + "top_k": "top_k", + "min_p": "min_p", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "repetition_penalty": "repetition_penalty", + "seed": "seed", + } + for proto_name, trt_name in field_map.items(): + value = _optional(sampling, proto_name) + if value is not None: + kwargs[trt_name] = value + num_sequences = _optional(sampling, "num_sequences") + if num_sequences is not None: + kwargs["n"] = num_sequences + kwargs["best_of"] = num_sequences + min_tokens = _optional(stopping, "min_tokens") + if min_tokens is not None: + kwargs["min_tokens"] = min_tokens + ignore_eos = _optional(stopping, "ignore_eos") + if ignore_eos is not None: + kwargs["ignore_eos"] = ignore_eos + include_stop = _optional(stopping, "include_stop_in_output") + if include_stop is not None: + kwargs["include_stop_str_in_output"] = include_stop + stop_text: list[str] = [] + stop_ids: list[int] = [] + for condition in stopping.conditions: + kind = condition.WhichOneof("condition") + if kind == "stop_text": + stop_text.append(condition.stop_text) + elif kind == "stop_token_id": + stop_ids.append(condition.stop_token_id) + if stop_text: + kwargs["stop"] = stop_text + if stop_ids: + kwargs["stop_token_ids"] = stop_ids + + prompt_requested = bool(_optional(response, "return_prompt_logprobs")) + output_requested = bool(_optional(response, "return_output_logprobs")) + prompt_count = _candidate_count(response.prompt_candidates, prompt_requested) + output_count = _candidate_count(response.output_candidates, output_requested) + if prompt_count is not None: + kwargs["prompt_logprobs"] = prompt_count + if output_count is not None: + kwargs["logprobs"] = output_count + prompt_start = _optional(response, "prompt_logprob_start") + if prompt_start not in (None, 0): + raise ValueError("TensorRT-LLM does not support prompt_logprob_start through OpenEngine") + + guide_kind = request.guided.WhichOneof("guide") + if guide_kind: + guide_kwargs: dict[str, Any] + if guide_kind == "json_schema": + guide_kwargs = {"json": request.guided.json_schema} + elif guide_kind == "regex": + guide_kwargs = {"regex": request.guided.regex} + elif guide_kind == "ebnf_grammar": + guide_kwargs = {"grammar": request.guided.ebnf_grammar} + elif guide_kind == "structural_tag": + guide_kwargs = {"structural_tag": request.guided.structural_tag} + elif guide_kind == "choice": + choices = [f"(?:{re.escape(choice)})" for choice in request.guided.choice.choices] + guide_kwargs = {"regex": "^(?:" + "|".join(choices) + ")$"} + elif guide_kind == "json_object": + guide_kwargs = {"json_object": True} + else: + raise ValueError(f"Unsupported guided decoding mode {guide_kind!r}") + kwargs["guided_decoding"] = GuidedDecodingParams(**guide_kwargs) + return SamplingParams(**kwargs) + + +def to_priority(request: generation_pb2.GenerateRequest) -> float: + """Map signed OpenEngine ordering into TRT-LLM's bounded priority domain.""" + if not request.HasField("priority"): + return 0.5 + priority = request.priority + return 0.5 + 0.5 * priority / (1 + abs(priority)) + + +def _message_struct(struct_message: object) -> dict[str, Any]: + if not struct_message.fields: + return {} + return MessageToDict(struct_message, preserving_proto_field_name=True) + + +def validate_media_options(options: dict[str, Any]) -> None: + unknown = set(options).difference(MEDIA_IO_REGISTRY) + if unknown: + raise ValueError(f"Unknown media option modalities: {sorted(unknown)}") + invalid = [name for name, value in options.items() if not isinstance(value, dict)] + if invalid: + raise ValueError(f"Media options must be objects for modalities: {sorted(invalid)}") + + +async def load_media( + media: list[object], + media_options: object, + server_config: MultimodalServerConfig | None, +) -> dict[str, list[Any]] | None: + """Decode ordered OpenEngine media with TRT-LLM's media-I/O merge rules.""" + if not media: + return None + request_options = _message_struct(media_options) + validate_media_options(request_options) + server_options = (server_config.media_io_kwargs if server_config is not None else None) or {} + output: dict[str, list[Any]] = {} + pending: list[tuple[str, asyncio.Future[Any] | asyncio.Task[Any] | Any]] = [] + for item in media: + modality = _MODALITY_NAMES.get(item.modality) + if modality is None: + raise ValueError(f"Unsupported media modality {item.modality}") + media_io = MEDIA_IO_REGISTRY[modality].create( + server_options.get(modality), request_options.get(modality) + ) + source = item.WhichOneof("source") + if source in ("url", "data_uri"): + pending.append((modality, media_io.async_load(getattr(item, source)))) + elif source == "raw_bytes": + pending.append( + (modality, media_io._run_in_executor(media_io.load_bytes, item.raw_bytes)) + ) + else: + raise ValueError("Each media item must carry exactly one source") + decoded = await asyncio.gather(*(awaitable for _, awaitable in pending)) + for (modality, _), value in zip(pending, decoded): + output.setdefault(modality, []).append(value) + return output + + +def stable_request_id(request_id: str) -> int: + """Map arbitrary wire request IDs into TRT-LLM's positive int64 domain.""" + digest = hashlib.sha256(request_id.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "big") & ((1 << 63) - 1) + + +def _json_value(value: Any) -> Any: + if is_dataclass(value): + return _json_value(asdict(value)) + if isinstance(value, dict): + return {str(key): _json_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_value(item) for item in value] + if isinstance(value, bytes): + return base64.b64encode(value).decode("ascii") + if hasattr(value, "logprob"): + output = {"logprob": float(value.logprob)} + if getattr(value, "rank", None) is not None: + output["rank"] = int(value.rank) + return output + if hasattr(value, "item"): + return value.item() + return value + + +def encode_handoff( + params: DisaggregatedParams, *, requires_decode_media: bool = False +) -> kv_pb2.KvSessionRef: + """Encode the revisioned context-first TRT handoff profile.""" + unsupported = { + "first-generation logits": params.first_gen_logits, + "multimodal embedding handles": params.multimodal_embedding_handles, + "multimodal hashes": params.multimodal_hashes, + } + if not requires_decode_media: + unsupported.update( + { + "mrope position IDs": params.mrope_position_ids_handle, + "mrope position deltas": params.mrope_position_deltas_handle, + } + ) + present = [name for name, value in unsupported.items() if value is not None] + if present: + raise ValueError("OpenEngine context-first handoff does not support " + ", ".join(present)) + if params.schedule_style not in (None, DisaggScheduleStyle.CONTEXT_FIRST): + raise ValueError("Generation-first handoff is not supported") + payload = { + "first_gen_tokens": _json_value(params.first_gen_tokens), + "first_gen_log_probs": _json_value(params.first_gen_log_probs), + "ctx_request_id": None if params.ctx_request_id is None else str(params.ctx_request_id), + "disagg_request_id": None + if params.disagg_request_id is None + else str(params.disagg_request_id), + "ctx_dp_rank": params.ctx_dp_rank, + "ctx_info_endpoint": params.ctx_info_endpoint, + "draft_tokens": _json_value(params.draft_tokens), + "ctx_usage": _json_value(params.ctx_usage), + "conversation_id": params.conversation_id, + "schedule_style": "context_first", + "requires_decode_media": requires_decode_media, + "opaque_state": ( + None + if params.opaque_state is None + else base64.b64encode(params.opaque_state).decode("ascii") + ), + } + canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + session = kv_pb2.KvSessionRef( + session_id=str(params.disagg_request_id or params.ctx_request_id or ""), + transfer_backend="tensorrt_llm", + dp_rank=params.ctx_dp_rank or 0, + ) + session.attributes_struct[HANDOFF_ATTRIBUTE] = canonical + return session + + +def _decode_handoff_payload(session: kv_pb2.KvSessionRef) -> dict[str, Any]: + if HANDOFF_ATTRIBUTE not in session.attributes_struct: + raise ValueError(f"KV session is missing {HANDOFF_ATTRIBUTE!r}") + encoded = session.attributes_struct[HANDOFF_ATTRIBUTE] + if not isinstance(encoded, str): + raise ValueError("TensorRT-LLM handoff attribute must be a JSON string") + try: + payload = json.loads(encoded) + except json.JSONDecodeError as error: + raise ValueError("TensorRT-LLM handoff is not valid JSON") from error + if not isinstance(payload, dict): + raise ValueError("TensorRT-LLM handoff must contain a JSON object") + return payload + + +def handoff_requires_decode_media(session: kv_pb2.KvSessionRef) -> bool: + """Return whether decode must recompute transient MM state from raw media.""" + payload = _decode_handoff_payload(session) + value = payload.get("requires_decode_media", False) + if not isinstance(value, bool): + raise ValueError("requires_decode_media must be a boolean") + return value + + +def decode_handoff(session: kv_pb2.KvSessionRef) -> DisaggregatedParams: + """Decode and validate the revisioned context-first TRT handoff profile.""" + payload = _decode_handoff_payload(session) + if payload.get("schedule_style", "context_first") != "context_first": + raise ValueError("Generation-first handoff is not supported") + if any( + payload.get(key) is not None + for key in ( + "first_gen_logits", + "multimodal_embedding_handles", + "multimodal_hashes", + "mrope_position_ids_handle", + "mrope_position_deltas_handle", + ) + ): + raise ValueError("Encoder-stage and first-generation logits handles are not supported") + + def decimal_id(name: str) -> int | None: + value = payload.get(name) + if value is None: + return None + if not isinstance(value, str) or not value.isdecimal(): + raise ValueError(f"{name} must be a decimal string") + return int(value) + + opaque = payload.get("opaque_state") + try: + opaque_state = None if opaque is None else base64.b64decode(opaque, validate=True) + except (ValueError, TypeError) as error: + raise ValueError("opaque_state must be canonical base64") from error + + ctx_usage = payload.get("ctx_usage") + if ctx_usage is not None: + if not isinstance(ctx_usage, dict): + raise ValueError("ctx_usage must be an object") + prompt_tokens = ctx_usage.get("prompt_tokens") + if ( + not isinstance(prompt_tokens, int) + or isinstance(prompt_tokens, bool) + or prompt_tokens < 0 + ): + raise ValueError("ctx_usage.prompt_tokens must be a non-negative integer") + prompt_details = ctx_usage.get("prompt_tokens_details") + if prompt_details is not None: + if not isinstance(prompt_details, dict): + raise ValueError("ctx_usage.prompt_tokens_details must be an object") + cached_tokens = prompt_details.get("cached_tokens", 0) + if ( + not isinstance(cached_tokens, int) + or isinstance(cached_tokens, bool) + or cached_tokens < 0 + ): + raise ValueError( + "ctx_usage.prompt_tokens_details.cached_tokens must be a non-negative integer" + ) + + def logprobs() -> list[Any] | None: + encoded_logprobs = payload.get("first_gen_log_probs") + if encoded_logprobs is None: + return None + if not isinstance(encoded_logprobs, list): + raise ValueError("first_gen_log_probs must be a list") + decoded: list[Any] = [] + for position in encoded_logprobs: + if isinstance(position, (int, float)): + decoded.append(float(position)) + continue + if not isinstance(position, dict): + raise ValueError("first_gen_log_probs entries must be candidate objects or numbers") + candidates: dict[int, Logprob] = {} + for token_id, value in position.items(): + if not isinstance(value, dict) or "logprob" not in value: + raise ValueError("first_gen_log_probs contains an invalid candidate") + try: + decoded_token_id = int(token_id) + except (TypeError, ValueError) as error: + raise ValueError( + "first_gen_log_probs contains a non-integer token ID" + ) from error + candidates[decoded_token_id] = Logprob( + logprob=float(value["logprob"]), + rank=(None if value.get("rank") is None else int(value["rank"])), + ) + decoded.append(candidates) + return decoded + + return DisaggregatedParams( + request_type="generation_only", + first_gen_tokens=payload.get("first_gen_tokens"), + first_gen_log_probs=logprobs(), + ctx_request_id=decimal_id("ctx_request_id"), + disagg_request_id=decimal_id("disagg_request_id"), + ctx_dp_rank=payload.get("ctx_dp_rank", session.dp_rank), + ctx_info_endpoint=payload.get("ctx_info_endpoint"), + draft_tokens=payload.get("draft_tokens"), + ctx_usage=ctx_usage, + conversation_id=payload.get("conversation_id"), + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + opaque_state=opaque_state, + ) + + +def to_lora_request(adapter: object | None) -> LoRARequest | None: + if adapter is None: + return None + return LoRARequest(adapter.lora_name, adapter.lora_id, adapter.source_path) diff --git a/tensorrt_llm/openengine/lora_registry.py b/tensorrt_llm/openengine/lora_registry.py new file mode 100644 index 000000000000..22ec8af8efdd --- /dev/null +++ b/tensorrt_llm/openengine/lora_registry.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lazy, logical LoRA registry for the OpenEngine control plane.""" + +import asyncio +import json +from pathlib import Path + +from openengine.v1 import lora_pb2 + +from tensorrt_llm.executor.request import LoRARequest + + +class LoraRegistry: + """Validate adapters now and let TRT-LLM load weights on first use.""" + + def __init__(self, reserved_identities: dict[str, int] | None = None) -> None: + reserved_identities = dict(reserved_identities or {}) + if any(not name for name in reserved_identities): + raise ValueError("Reserved LoRA names must not be empty") + if any(adapter_id < 0 for adapter_id in reserved_identities.values()): + raise ValueError("Reserved LoRA IDs must be non-negative") + if len(set(reserved_identities.values())) != len(reserved_identities): + raise ValueError("Reserved LoRA IDs must be unique") + self._by_name: dict[str, lora_pb2.LoraAdapter] = {} + self._identities_by_name: dict[str, lora_pb2.LoraAdapter] = {} + self._names_by_id: dict[int, str] = {} + self._names_by_path: dict[str, str] = {} + self._reserved_by_name = reserved_identities + self._reserved_by_id = { + adapter_id: name for name, adapter_id in reserved_identities.items() + } + self._lock = asyncio.Lock() + + @staticmethod + def _validated(adapter: lora_pb2.LoraAdapter) -> lora_pb2.LoraAdapter: + if not adapter.lora_name: + raise ValueError("lora_name must not be empty") + if adapter.lora_id < 0: + raise ValueError("lora_id must be non-negative") + path = Path(adapter.source_path).expanduser().resolve() + if not path.is_dir(): + raise ValueError(f"LoRA source path is not a directory: {path}") + if not (path / "adapter_config.json").is_file(): + raise ValueError(f"LoRA directory is missing adapter_config.json: {path}") + try: + config = json.loads((path / "adapter_config.json").read_text()) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"LoRA adapter_config.json is invalid: {path}") from error + if not isinstance(config, dict): + raise ValueError(f"LoRA adapter_config.json must contain an object: {path}") + if not any( + (path / filename).is_file() + for filename in ("adapter_model.safetensors", "adapter_model.bin") + ): + raise ValueError(f"LoRA directory is missing adapter weights: {path}") + return lora_pb2.LoraAdapter( + lora_id=adapter.lora_id, lora_name=adapter.lora_name, source_path=str(path) + ) + + async def load(self, adapter: lora_pb2.LoraAdapter) -> tuple[lora_pb2.LoraAdapter, bool]: + validated = self._validated(adapter) + async with self._lock: + reserved_id = self._reserved_by_name.get(validated.lora_name) + if reserved_id is not None: + raise ValueError( + f"LoRA name {validated.lora_name!r} is reserved for a model-owned adapter" + ) + reserved_name = self._reserved_by_id.get(validated.lora_id) + if reserved_name is not None: + raise ValueError( + f"LoRA ID {validated.lora_id} is reserved for model-owned adapter " + f"{reserved_name!r}" + ) + identity = self._identities_by_name.get(validated.lora_name) + if identity is not None: + if identity != validated: + raise ValueError( + f"LoRA name {validated.lora_name!r} is permanently bound to different " + "attributes" + ) + if validated.lora_name in self._by_name: + return identity, True + self._by_name[validated.lora_name] = identity + return identity, False + registered_name = self._names_by_id.get(validated.lora_id) + if registered_name is not None: + raise ValueError( + f"LoRA ID {validated.lora_id} is permanently bound to {registered_name!r}" + ) + registered_name = self._names_by_path.get(validated.source_path) + if registered_name is not None: + raise ValueError( + f"LoRA path {validated.source_path!r} is permanently bound to " + f"{registered_name!r}" + ) + self._by_name[validated.lora_name] = validated + self._identities_by_name[validated.lora_name] = validated + self._names_by_id[validated.lora_id] = validated.lora_name + self._names_by_path[validated.source_path] = validated.lora_name + return validated, False + + async def unload(self, name: str) -> lora_pb2.LoraAdapter: + async with self._lock: + adapter = self._by_name.pop(name, None) + if adapter is None: + raise KeyError(name) + return adapter + + async def list(self) -> list[lora_pb2.LoraAdapter]: + async with self._lock: + return [self._by_name[name] for name in sorted(self._by_name)] + + async def request(self, name: str) -> LoRARequest: + async with self._lock: + adapter = self._by_name.get(name) + if adapter is None: + raise KeyError(name) + return LoRARequest(adapter.lora_name, adapter.lora_id, adapter.source_path) diff --git a/tensorrt_llm/openengine/server.py b/tensorrt_llm/openengine/server.py new file mode 100644 index 000000000000..4604047ad1e2 --- /dev/null +++ b/tensorrt_llm/openengine/server.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lifecycle wrapper for the optional OpenEngine sibling gRPC server.""" + +import grpc +from openengine.v1 import engine_pb2, openengine_pb2_grpc + +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.request_tracker import RequestTracker +from tensorrt_llm.serve.stats_fanout import StatsFanout + +from ._schema_pin import OPENENGINE_COMMIT +from .servicer import OpenEngineServicer, schema_release + +_MAX_MESSAGE_LENGTH = 256 * 1024 * 1024 + + +def validate_schema_release(schema_release: str) -> str: + """Require an immutable OpenEngine source identity before binding.""" + if schema_release == OPENENGINE_COMMIT: + return schema_release + raise RuntimeError( + "OPENENGINE_SCHEMA_RELEASE must exactly match the pinned OPENENGINE_COMMIT " + f"({OPENENGINE_COMMIT})" + ) + + +class OpenEngineServer: + """An OpenEngine gRPC server which never owns or shuts down the LLM.""" + + def __init__( + self, + llm: object, + model: str, + role: int, + host: str, + port: int, + tracker: RequestTracker, + media_config: MultimodalServerConfig | None = None, + reasoning_parser: str | None = None, + tool_parser: str | None = None, + kv_event_fanout: KvEventFanout | None = None, + stats_fanout: StatsFanout | None = None, + ) -> None: + validate_schema_release(schema_release()) + self.host = host + self.port = port + self._kv_event_fanout = kv_event_fanout + self._stats_fanout = stats_fanout + self._server = grpc.aio.server( + options=[ + ("grpc.max_send_message_length", _MAX_MESSAGE_LENGTH), + ("grpc.max_receive_message_length", _MAX_MESSAGE_LENGTH), + ("grpc.keepalive_time_ms", 30_000), + ("grpc.keepalive_timeout_ms", 10_000), + ] + ) + self.servicer = OpenEngineServicer( + llm=llm, + model=model, + role=role, + tracker=tracker, + media_config=media_config, + reasoning_parser=reasoning_parser, + tool_parser=tool_parser, + kv_event_fanout=kv_event_fanout, + stats_fanout=stats_fanout, + event_host=host, + event_port=port, + ) + openengine_pb2_grpc.add_OpenEngineServicer_to_server(self.servicer, self._server) + bound = self._server.add_insecure_port(f"{host}:{port}") + if bound == 0: + raise RuntimeError(f"Failed to bind OpenEngine server to {host}:{port}") + self.port = bound if port == 0 else port + advertised_host = { + "0.0.0.0": "127.0.0.1", + "::": "::1", + "[::]": "::1", + }.get(host, host) + self.servicer.event_host = advertised_host + self.servicer.event_port = self.port + + async def start(self) -> None: + if self._kv_event_fanout is not None: + self._kv_event_fanout.start() + if self._stats_fanout is not None: + self._stats_fanout.start() + try: + await self._server.start() + except BaseException: + if self._kv_event_fanout is not None: + await self._kv_event_fanout.stop() + if self._stats_fanout is not None: + await self._stats_fanout.stop() + self.servicer.close() + raise + logger.info("OpenEngine sibling server started on %s:%d", self.host, self.port) + + async def stop(self, grace: float = 5.0) -> None: + try: + await self._server.stop(grace=grace) + finally: + self.servicer.close() + if self._kv_event_fanout is not None: + await self._kv_event_fanout.stop() + if self._stats_fanout is not None: + await self._stats_fanout.stop() + logger.info("OpenEngine sibling server stopped") + + +def openengine_role(server_role: object | None) -> int: + """Map TRT-LLM serve roles onto the supported OpenEngine milestone.""" + if server_role is None: + return engine_pb2.ENGINE_ROLE_AGGREGATED + name = getattr(server_role, "name", str(server_role)).upper() + if name == "CONTEXT": + return engine_pb2.ENGINE_ROLE_PREFILL + if name == "GENERATION": + return engine_pb2.ENGINE_ROLE_DECODE + if name in ("MM_ENCODER", "VISUAL_GEN"): + raise ValueError(f"OpenEngine does not support TRT-LLM server role {name!r}") + raise ValueError(f"Unknown TRT-LLM server role {name!r}") + + +__all__ = ["OpenEngineServer", "openengine_role", "validate_schema_release"] diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py new file mode 100644 index 000000000000..0835a0d0ac97 --- /dev/null +++ b/tensorrt_llm/openengine/servicer.py @@ -0,0 +1,1226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine gRPC servicer backed by one TensorRT-LLM LLM instance.""" + +import asyncio +import os +import time +import uuid +from collections.abc import AsyncGenerator +from dataclasses import replace +from pathlib import Path +from typing import Any + +import grpc +from openengine import MINIMUM_CLIENT_REVISION, SCHEMA_RELEASE, SCHEMA_REVISION +from openengine.v1 import ( + engine_pb2, + error_pb2, + generation_pb2, + input_pb2, + kv_pb2, + lifecycle_pb2, + lora_pb2, + model_pb2, + observability_pb2, + openengine_pb2_grpc, +) + +import tensorrt_llm +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_hash import ( + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + get_effective_kv_cache_event_hash_algo, +) +from tensorrt_llm.scheduling_params import SchedulingParams +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.request_tracker import RequestTracker +from tensorrt_llm.serve.stats_fanout import StatsFanout + +from .converters import ( + decode_handoff, + encode_handoff, + handoff_requires_decode_media, + load_media, + stable_request_id, + to_priority, + to_sampling_params, +) +from .lora_registry import LoraRegistry + + +def schema_release() -> str: + """Return the immutable OpenEngine source identity configured at startup.""" + return os.getenv("OPENENGINE_SCHEMA_RELEASE", SCHEMA_RELEASE) + + +def _arg(llm: object, name: str, default: Any = None) -> Any: + args = getattr(llm, "args", None) + value = getattr(args, name, None) + if value is not None: + return value + parallel = getattr(args, "parallel_config", None) + return getattr(parallel, name, default) + + +def _token_text(llm: object, token_id: int) -> str: + tokenizer = getattr(llm, "tokenizer", None) + if tokenizer is None: + return "" + try: + return tokenizer.decode([token_id], skip_special_tokens=False) + except (ValueError, TypeError, AttributeError): + return "" + + +def _signed_i64(value: int) -> int: + if value >= 1 << 63: + return value - (1 << 64) + if value < -(1 << 63): + return ((value + (1 << 63)) % (1 << 64)) - (1 << 63) + return value + + +def _block_hash(value: object, hash_algo: str | None) -> kv_pb2.KvBlockHash: + if isinstance(value, bytes): + return kv_pb2.KvBlockHash(value=value, encoding=hash_algo or "bytes") + if isinstance(value, int): + return kv_pb2.KvBlockHash( + value=str(_signed_i64(value)).encode("ascii"), encoding="decimal_int64" + ) + text = str(value) + try: + encoded = bytes.fromhex(text) + except ValueError: + return kv_pb2.KvBlockHash(value=text.encode("utf-8"), encoding="utf8") + return kv_pb2.KvBlockHash(value=encoded, encoding=hash_algo or "hex") + + +class OpenEngineServicer(openengine_pb2_grpc.OpenEngineServicer): + """Engine-neutral service surface over TensorRT-LLM's Python LLM API.""" + + def __init__( + self, + llm: object, + model: str, + role: int, + tracker: RequestTracker, + media_config: MultimodalServerConfig | None = None, + reasoning_parser: str | None = None, + tool_parser: str | None = None, + kv_event_fanout: KvEventFanout | None = None, + stats_fanout: StatsFanout | None = None, + instance_id: str | None = None, + event_host: str = "127.0.0.1", + event_port: int = 0, + kv_session_ttl_seconds: float = 300.0, + post_abort_cleanup_timeout_seconds: float = 1.0, + ) -> None: + self.llm = llm + self.model = model + self.role = role + self.tracker = tracker + self.media_config = media_config + self.reasoning_parser = reasoning_parser or "" + self.tool_parser = tool_parser or "" + processor = getattr(llm, "input_processor", None) + identity_getter = getattr(processor, "get_model_owned_lora_identities", None) + model_owned_loras = identity_getter() if callable(identity_getter) else {} + self.loras = LoraRegistry(model_owned_loras) + self._lora_names_by_id: dict[int, str] = { + adapter_id: name for name, adapter_id in model_owned_loras.items() + } + self.kv_events = kv_event_fanout + self.stats = stats_fanout + self._kv_session_requests: dict[str, str] = {} + self._kv_session_timers: dict[str, asyncio.TimerHandle] = {} + if kv_session_ttl_seconds <= 0: + raise ValueError("kv_session_ttl_seconds must be positive") + self._kv_session_ttl_seconds = kv_session_ttl_seconds + if post_abort_cleanup_timeout_seconds <= 0: + raise ValueError("post_abort_cleanup_timeout_seconds must be positive") + self._post_abort_cleanup_timeout_seconds = post_abort_cleanup_timeout_seconds + self._partial_block_hashes: dict[int, set[object]] = {} + self.instance_id = instance_id or str(uuid.uuid4()) + self.event_host = event_host + self.event_port = event_port + + async def Generate( + self, request: generation_pb2.GenerateRequest, context: grpc.aio.ServicerContext + ) -> AsyncGenerator[generation_pb2.GenerateResponse, None]: + result = None + consumed_session_id = None + try: + self._validate_generate(request) + params = to_sampling_params(request) + media = await load_media(list(request.media), request.media_options, self.media_config) + input_kind = request.WhichOneof("input") + inputs: dict[str, Any] + if input_kind == "prompt": + inputs = {"prompt": request.prompt} + else: + inputs = {"prompt_token_ids": list(request.token_ids.ids)} + if media: + inputs["multi_modal_data"] = media + lora_request = self._required_multimodal_lora(request) + if lora_request is not None and request.lora_name: + raise ValueError( + f"Modalities in this request require the model-owned " + f"{lora_request.lora_name!r} adapter and cannot be combined with " + "a user-selected LoRA" + ) + if lora_request is None and request.lora_name: + lora_request = await self.loras.request(request.lora_name) + disaggregated = self._disaggregated_params(request) + context_usage = ( + disaggregated.ctx_usage + if disaggregated is not None and disaggregated.request_type == "generation_only" + else None + ) + scheduling = self._scheduling_params(request) + priority = to_priority(request) + trace_headers = { + key.lower(): value + for key, value in request.metadata.items() + if key.lower() in ("traceparent", "tracestate", "baggage") + } + result = self.llm.generate_async( + inputs=inputs, + sampling_params=params, + lora_request=lora_request, + streaming=True, + disaggregated_params=disaggregated, + scheduling_params=scheduling, + cache_salt=(request.kv.cache_salt if request.kv.HasField("cache_salt") else None), + trace_headers=trace_headers or None, + priority=priority, + ) + self.tracker.admit(request.request_id, result) + if self.stats is not None: + self.stats.wake() + if request.kv.HasField("session") and request.kv.session.session_id: + consumed_session_id = request.kv.session.session_id + self._track_kv_session(consumed_session_id, request.request_id) + except KeyError as error: + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, f"LoRA is not loaded: {error.args[0]}" + ) + return + except (ValueError, TypeError, RuntimeError) as error: + if ( + result is not None + and self.tracker.active_requests.get(request.request_id) is not result + ): + try: + result.abort() + except (RuntimeError, AssertionError): + logger.warning( + "Failed to abort rejected duplicate request %s", + request.request_id, + ) + code = ( + grpc.StatusCode.UNAVAILABLE + if self.tracker.draining + else grpc.StatusCode.INVALID_ARGUMENT + ) + await context.abort(code, str(error)) + return + + sent_tokens: dict[int, int] = {} + sent_text: dict[int, int] = {} + prompt_sent = False + completed = False + try: + async for current in result: + if context.cancelled(): + await self.tracker.abort(request.request_id) + return + if not prompt_sent: + prompt = self._prompt_output(current) + if prompt is not None: + prompt_sent = True + yield generation_pb2.GenerateResponse( + request_id=request.request_id, prompt=prompt + ) + if self.role == engine_pb2.ENGINE_ROLE_PREFILL: + if current.finished: + handoff = current.disaggregated_params + if handoff is None and current.outputs: + handoff = current.outputs[0].disaggregated_params + if handoff is None: + raise RuntimeError( + "Context-only result did not return disaggregated parameters" + ) + handoff = replace(handoff, ctx_usage=self._usage_payload(current)) + session = encode_handoff(handoff, requires_decode_media=bool(request.media)) + self._track_kv_session(session.session_id, request.request_id) + yield generation_pb2.GenerateResponse( + request_id=request.request_id, + prefill_ready=generation_pb2.PrefillReady(kv_session=session), + usage=self._usage(current), + ) + completed = True + continue + for output in current.outputs: + token_start = sent_tokens.get(output.index, 0) + text_start = sent_text.get(output.index, 0) + token_ids = list(output.token_ids or []) + delta_ids = token_ids[token_start:] + delta_text = (output.text or "")[text_start:] + if delta_ids or delta_text: + logprobs = list(output.logprobs or [])[token_start:] + yield generation_pb2.GenerateResponse( + request_id=request.request_id, + token=generation_pb2.TokenOutput( + output_index=output.index, + tokens=self._token_infos(delta_ids, logprobs), + text=delta_text, + ), + ) + sent_tokens[output.index] = len(token_ids) + sent_text[output.index] = len(output.text or "") + if current.finished: + outputs = current.outputs + for index, output in enumerate(outputs): + response = generation_pb2.GenerateResponse( + request_id=request.request_id, + finished=self._finished(output), + ) + if index == len(outputs) - 1: + response.usage.CopyFrom(self._usage(current, context_usage)) + yield response + completed = True + except asyncio.CancelledError: + await self.tracker.abort(request.request_id) + raise + except (RuntimeError, ValueError, TypeError) as error: + logger.error("OpenEngine request %s failed: %s", request.request_id, error) + yield generation_pb2.GenerateResponse( + request_id=request.request_id, + error=error_pb2.EngineError( + code=error_pb2.ERROR_CODE_INTERNAL, + message=str(error), + retryable=False, + ), + ) + finally: + if consumed_session_id is not None: + self._release_kv_session(consumed_session_id) + if not completed and request.request_id in self.tracker.active_requests: + await self.tracker.abort(request.request_id) + else: + await self.tracker.finish(request.request_id) + if self.stats is not None: + self.stats.wake() + + def _track_kv_session(self, session_id: str, request_id: str) -> None: + self._release_kv_session(session_id) + self._kv_session_requests[session_id] = request_id + loop = asyncio.get_running_loop() + self._kv_session_timers[session_id] = loop.call_later( + self._kv_session_ttl_seconds, self._release_kv_session, session_id + ) + + def _release_kv_session(self, session_id: str) -> bool: + released = self._kv_session_requests.pop(session_id, None) is not None + timer = self._kv_session_timers.pop(session_id, None) + if timer is not None: + timer.cancel() + return released + + def _release_all_kv_sessions(self) -> None: + for timer in self._kv_session_timers.values(): + timer.cancel() + self._kv_session_timers.clear() + self._kv_session_requests.clear() + + def close(self) -> None: + """Release protocol-owned timers without shutting down the shared LLM.""" + self._release_all_kv_sessions() + + def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: + if not request.request_id: + raise ValueError("request_id must not be empty") + if request.model and request.model != self.model: + raise ValueError(f"Unknown model {request.model!r}") + if request.WhichOneof("input") is None: + raise ValueError("Generate requires prompt or token_ids input") + if request.request_id in self.tracker.active_requests: + raise ValueError(f"Request {request.request_id!r} is already active") + if self.tracker.draining: + raise RuntimeError("TensorRT-LLM is draining") + configured_backend = getattr( + getattr(self.llm, "args", None), "guided_decoding_backend", None + ) + if request.guided.backend and request.guided.backend != configured_backend: + raise ValueError("Per-request guided decoding backend selection is not supported") + if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: + raise ValueError("Prefix-cache bypass is not supported by TensorRT-LLM") + if request.kv.HasField("data_parallel_rank"): + self._scheduling_params(request) + if request.media: + processor = getattr(self.llm, "input_processor", None) + aggregate = self._available_modalities(processor) + prefill_decode = self._available_modalities(processor, prefill_decode=True) + names = { + input_pb2.MODALITY_UNSPECIFIED: "image", + input_pb2.MODALITY_IMAGE: "image", + input_pb2.MODALITY_VIDEO: "video", + input_pb2.MODALITY_AUDIO: "audio", + } + allowed = ( + aggregate if self.role == engine_pb2.ENGINE_ROLE_AGGREGATED else prefill_decode + ) + unsupported = { + names.get(item.modality, "unspecified") + for item in request.media + if names.get(item.modality) not in allowed + } + if unsupported: + raise ValueError( + f"Media modalities are not supported for this role: {sorted(unsupported)}" + ) + if isinstance(processor, BaseMultimodalInputProcessor): + requested_modalities = tuple( + dict.fromkeys(names[item.modality] for item in request.media) + ) + processor.get_required_lora_spec(requested_modalities) + if self.role == engine_pb2.ENGINE_ROLE_DECODE: + if not request.kv.HasField("session"): + raise ValueError("Decode requests require a prefill KV session") + requires_media = handoff_requires_decode_media(request.kv.session) + if requires_media and not request.media: + raise ValueError("Decode request must resend the ordered context-phase media") + if request.media and not requires_media: + raise ValueError("Decode KV session does not require raw media") + elif self.role == engine_pb2.ENGINE_ROLE_PREFILL: + if request.kv.HasField("session"): + raise ValueError("Prefill requests cannot consume a KV session") + elif request.kv.HasField("session"): + raise ValueError("Aggregated requests cannot consume a KV session") + + def _disaggregated_params( + self, request: generation_pb2.GenerateRequest + ) -> DisaggregatedParams | None: + if self.role == engine_pb2.ENGINE_ROLE_AGGREGATED: + return None + if self.role == engine_pb2.ENGINE_ROLE_PREFILL: + return DisaggregatedParams( + request_type="context_only", + disagg_request_id=stable_request_id(request.request_id), + ctx_dp_rank=( + request.kv.data_parallel_rank + if request.kv.HasField("data_parallel_rank") + else None + ), + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + conversation_id=request.metadata.get("conversation_id"), + ) + return decode_handoff(request.kv.session) + + def _scheduling_params( + self, request: generation_pb2.GenerateRequest + ) -> SchedulingParams | None: + if not request.kv.HasField("data_parallel_rank"): + return None + rank = request.kv.data_parallel_rank + data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) + if rank >= data_parallel_size: + raise ValueError( + f"data_parallel_rank {rank} is outside the configured DP size {data_parallel_size}" + ) + attention_dp = bool(_arg(self.llm, "enable_attention_dp", False)) + if not attention_dp or getattr(self.llm, "_on_trt_backend", False): + if rank == 0: + return None + raise ValueError( + "A nonzero data_parallel_rank requires the PyTorch backend with attention DP " + "enabled" + ) + return SchedulingParams(attention_dp_rank=rank, attention_dp_relax=False) + + def _required_multimodal_lora( + self, request: generation_pb2.GenerateRequest + ) -> LoRARequest | None: + if not request.media: + return None + processor = getattr(self.llm, "input_processor", None) + if not isinstance(processor, BaseMultimodalInputProcessor): + return None + names = { + input_pb2.MODALITY_UNSPECIFIED: "image", + input_pb2.MODALITY_IMAGE: "image", + input_pb2.MODALITY_VIDEO: "video", + input_pb2.MODALITY_AUDIO: "audio", + } + modalities = tuple( + dict.fromkeys(names[item.modality] for item in request.media if item.modality in names) + ) + spec = processor.get_required_lora_spec(modalities) + if spec is None: + return None + if not self._supports_lora(): + raise ValueError( + f"Modalities in this request require the model-owned {spec.name!r} adapter, " + "but TensorRT-LLM was not configured with LoRA support" + ) + self._lora_names_by_id[spec.adapter_id] = spec.name + return LoRARequest(spec.name, spec.adapter_id, spec.path) + + def _available_modalities( + self, processor: object, *, prefill_decode: bool = False + ) -> tuple[str, ...]: + if not isinstance(processor, BaseMultimodalInputProcessor): + return () + declared = ( + processor.get_openengine_prefill_decode_modalities() + if prefill_decode + else processor.get_openengine_modalities() + ) + available = [] + for modality in declared: + spec = processor.get_required_lora_spec((modality,)) + if spec is not None and (not self._supports_lora() or not Path(spec.path).is_dir()): + continue + available.append(modality) + return tuple(available) + + def _token_infos( + self, token_ids: list[int], logprobs: list[Any] + ) -> list[generation_pb2.TokenInfo]: + infos: list[generation_pb2.TokenInfo] = [] + for index, token_id in enumerate(token_ids): + info = generation_pb2.TokenInfo( + token_id=token_id, + token=_token_text(self.llm, token_id), + ) + if index < len(logprobs): + value = logprobs[index] + if isinstance(value, dict): + sampled = value.get(token_id) + if sampled is not None: + info.logprob = sampled.logprob + if sampled.rank is not None: + info.rank = sampled.rank + for candidate_id, candidate in value.items(): + candidate_proto = info.candidates.add( + token_id=candidate_id, + logprob=candidate.logprob, + token=_token_text(self.llm, candidate_id), + ) + if candidate.rank is not None: + candidate_proto.rank = candidate.rank + elif isinstance(value, (int, float)): + info.logprob = float(value) + infos.append(info) + return infos + + def _prompt_output(self, result: object) -> generation_pb2.PromptOutput | None: + if not result.outputs: + return None + prompt_logprobs = result.outputs[0].prompt_logprobs + if not prompt_logprobs: + return None + token_ids = list(result.prompt_token_ids) + return generation_pb2.PromptOutput( + tokens=self._token_infos(token_ids, list(prompt_logprobs)) + ) + + @staticmethod + def _usage(result: object, context_usage: dict[str, Any] | None = None) -> generation_pb2.Usage: + prompt = len(result.prompt_token_ids) + completion = sum(len(output.token_ids or []) for output in result.outputs) + cached = getattr(result, "cached_tokens", None) + if context_usage is not None: + prompt = int(context_usage["prompt_tokens"]) + details = context_usage.get("prompt_tokens_details") or {} + cached = int(details.get("cached_tokens", 0)) + usage = generation_pb2.Usage( + prompt_tokens=prompt, completion_tokens=completion, total_tokens=prompt + completion + ) + if cached is not None: + usage.cached_prompt_tokens = cached + return usage + + @classmethod + def _usage_payload(cls, result: object) -> dict[str, Any]: + usage = cls._usage(result) + return { + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + "total_tokens": usage.total_tokens, + "prompt_tokens_details": {"cached_tokens": usage.cached_prompt_tokens}, + } + + @staticmethod + def _finished(output: object) -> generation_pb2.GenerationFinished: + reason_map = { + "stop": generation_pb2.FINISH_REASON_STOP, + "length": generation_pb2.FINISH_REASON_LENGTH, + "cancelled": generation_pb2.FINISH_REASON_CANCELLED, + "timeout": generation_pb2.FINISH_REASON_CANCELLED, + } + finished = generation_pb2.GenerationFinished( + output_index=output.index, + reason=reason_map.get(output.finish_reason, generation_pb2.FINISH_REASON_STOP), + ) + if isinstance(output.stop_reason, int): + finished.stop_match.stop_token_id = output.stop_reason + elif isinstance(output.stop_reason, str): + finished.stop_match.stop_text = output.stop_reason + return finished + + async def GetEngineInfo( + self, request: engine_pb2.GetEngineInfoRequest, context: grpc.aio.ServicerContext + ) -> engine_pb2.EngineInfo: + del request, context + tp = _arg(self.llm, "tensor_parallel_size", 1) + pp = _arg(self.llm, "pipeline_parallel_size", 1) + dp = _arg(self.llm, "data_parallel_size", 1) + return engine_pb2.EngineInfo( + engine_name="tensorrt_llm", + engine_version=getattr(tensorrt_llm, "__version__", "unknown"), + role=self.role, + instance_id=self.instance_id, + supported_models=[self.model], + parallelism=engine_pb2.ParallelismInfo( + tensor_parallel_size=tp, + pipeline_parallel_size=pp, + data_parallel_size=dp, + data_parallel_rank=_arg(self.llm, "data_parallel_rank", 0), + ), + kv_connector=self._kv_connector_info(), + schema_revision=SCHEMA_REVISION, + minimum_client_revision=MINIMUM_CLIENT_REVISION, + schema_release=schema_release(), + ) + + async def GetModelInfo( + self, request: model_pb2.GetModelInfoRequest, context: grpc.aio.ServicerContext + ) -> model_pb2.ModelInfo: + if request.model and request.model != self.model: + await context.abort(grpc.StatusCode.NOT_FOUND, f"Unknown model {request.model!r}") + raise ValueError(f"Unknown model {request.model!r}") + input_processor = getattr(self.llm, "input_processor", None) + aggregate = self._available_modalities(input_processor) + prefill_decode = self._available_modalities(input_processor, prefill_decode=True) + modality_enum = { + "image": input_pb2.MODALITY_IMAGE, + "video": input_pb2.MODALITY_VIDEO, + "audio": input_pb2.MODALITY_AUDIO, + } + args = getattr(self.llm, "args", None) + guided_backend = getattr(args, "guided_decoding_backend", None) + guided_modes = [] + if guided_backend is not None: + guided_modes = [ + model_pb2.GUIDED_DECODING_MODE_JSON_SCHEMA, + model_pb2.GUIDED_DECODING_MODE_REGEX, + model_pb2.GUIDED_DECODING_MODE_EBNF_GRAMMAR, + model_pb2.GUIDED_DECODING_MODE_CHOICE, + model_pb2.GUIDED_DECODING_MODE_JSON_OBJECT, + ] + if guided_backend == "xgrammar": + guided_modes.append(model_pb2.GUIDED_DECODING_MODE_STRUCTURAL_TAG) + max_context = getattr(args, "max_seq_len", None) or getattr(args, "max_input_len", None) + max_requests = getattr(args, "max_batch_size", None) + max_tokens = getattr(args, "max_num_tokens", None) + kv_capacity = self._kv_capacity() + info = model_pb2.ModelInfo( + model_id=self.model, + served_model_name=self.model, + served_model_aliases=[self.model], + tokenizer_modes=["auto"], + supports_text_input=True, + supports_token_ids_input=True, + supports_lora=self._supports_lora(), + supports_multimodal=bool(aggregate), + reasoning_parser=self.reasoning_parser, + tool_call_parser=self.tool_parser, + generation=model_pb2.GenerationCapabilities( + prompt_logprobs=model_pb2.LogprobCapabilities( + supported=True, + candidate_selection_modes=[model_pb2.CANDIDATE_TOKEN_SELECTION_MODE_TOP_N], + max_top_n=20, + ), + output_logprobs=model_pb2.LogprobCapabilities( + supported=True, + candidate_selection_modes=[model_pb2.CANDIDATE_TOKEN_SELECTION_MODE_TOP_N], + max_top_n=20, + ), + guided_decoding=model_pb2.GuidedDecodingCapabilities( + supported=guided_backend is not None, + modes=guided_modes, + ), + max_num_sequences=getattr(args, "max_beam_width", 1), + supports_priority=True, + supports_stop_in_output=True, + supports_cache_salt=True, + supports_prefix_cache_bypass=False, + ), + multimodal_capabilities=model_pb2.MultimodalCapabilities( + aggregate_modalities=[modality_enum[name] for name in aggregate], + prefill_decode_modalities=[modality_enum[name] for name in prefill_decode], + source_types=[ + input_pb2.MEDIA_SOURCE_TYPE_URL, + input_pb2.MEDIA_SOURCE_TYPE_DATA_URI, + input_pb2.MEDIA_SOURCE_TYPE_RAW_BYTES, + ], + supports_per_request_media_options=True, + ), + ) + if max_context is not None: + info.max_context_length = max_context + if max_requests is not None: + info.max_running_requests = max_requests + if max_tokens is not None: + info.max_batched_tokens = max_tokens + if kv_capacity.get("tokensPerBlock") is not None: + info.kv_block_size = kv_capacity["tokensPerBlock"] + if kv_capacity.get("maxNumBlocks") is not None: + info.total_kv_blocks = kv_capacity["maxNumBlocks"] + return info + + def _kv_capacity(self) -> dict[str, int]: + capacity: dict[str, int] = {} + config = getattr(getattr(self.llm, "args", None), "kv_cache_config", None) + tokens_per_block = getattr(config, "tokens_per_block", None) + if tokens_per_block is not None: + capacity["tokensPerBlock"] = int(tokens_per_block) + getter = getattr(self.llm, "get_kv_cache_capacity", None) + if callable(getter): + try: + discovered = getter() + except (RuntimeError, AttributeError, TypeError) as error: + logger.debug("KV capacity discovery unavailable: %s", error) + else: + if isinstance(discovered, dict): + for key in ("maxNumBlocks", "tokensPerBlock", "maxNumTokens"): + if discovered.get(key) is not None: + capacity[key] = int(discovered[key]) + return capacity + + async def GetLoad( + self, request: observability_pb2.GetLoadRequest, context: grpc.aio.ServicerContext + ) -> observability_pb2.LoadInfo: + del context + latest = self.stats.latest_by_rank() if self.stats is not None else {} + rank_infos = [] + running_requests = 0 + queued_requests = 0 + used_kv_blocks = 0 + total_kv_blocks = 0 + prefill_batch_size = 0 + decode_batch_size = 0 + attributes = {"source": "shared_stats_fanout" if latest else "shared_request_tracker"} + tokens_per_block: set[int] = set() + for rank, stat in sorted(latest.items()): + kv_stats = stat.get("kvCacheStats") or {} + ifb_stats = stat.get("inflightBatchingStats") or {} + rank_running = int(stat.get("numActiveRequests", 0)) + rank_queued = int(stat.get("numQueuedRequests", 0)) + rank_total = int(kv_stats.get("maxNumBlocks", 0)) + rank_used = int( + kv_stats.get( + "usedNumBlocks", + max(0, rank_total - int(kv_stats.get("freeNumBlocks", rank_total))), + ) + ) + rank_prefill = int(ifb_stats.get("numContextRequests", 0)) + rank_decode = int(ifb_stats.get("numGenRequests", 0)) + rank_tokens_per_block = kv_stats.get("tokensPerBlock") + if rank_tokens_per_block is not None: + rank_tokens_per_block = int(rank_tokens_per_block) + tokens_per_block.add(rank_tokens_per_block) + attributes[f"rank.{rank}.kv_tokens_per_block"] = str(rank_tokens_per_block) + running_requests += rank_running + queued_requests += rank_queued + used_kv_blocks += rank_used + total_kv_blocks += rank_total + prefill_batch_size += rank_prefill + decode_batch_size += rank_decode + if request.include_per_rank: + rank_infos.append( + observability_pb2.RankLoadInfo( + data_parallel_rank=rank, + running_requests=rank_running, + queued_requests=rank_queued, + used_kv_blocks=rank_used, + total_kv_blocks=rank_total, + prefill_batch_size=rank_prefill, + decode_batch_size=rank_decode, + ) + ) + + capacity = self._kv_capacity() + running_requests = max(running_requests, self.tracker.active_count) + if not latest: + total_kv_blocks = capacity.get("maxNumBlocks", 0) + if capacity.get("tokensPerBlock") is not None: + tokens_per_block.add(capacity["tokensPerBlock"]) + if len(tokens_per_block) == 1: + attributes["kv_tokens_per_block"] = str(next(iter(tokens_per_block))) + response = observability_pb2.LoadInfo( + instance_id=self.instance_id, + timestamp_unix_nanos=time.time_ns(), + running_requests=running_requests, + queued_requests=queued_requests, + active_kv_sessions=len(self._kv_session_requests), + ranks=rank_infos, + attributes=attributes, + ) + if used_kv_blocks or latest: + response.used_kv_blocks = used_kv_blocks + if total_kv_blocks: + response.total_kv_blocks = total_kv_blocks + if prefill_batch_size or latest: + response.prefill_batch_size = prefill_batch_size + if decode_batch_size or latest: + response.decode_batch_size = decode_batch_size + return response + + async def Health( + self, request: lifecycle_pb2.HealthRequest, context: grpc.aio.ServicerContext + ) -> lifecycle_pb2.HealthResponse: + if request.model and request.model != self.model: + await context.abort(grpc.StatusCode.NOT_FOUND, f"Unknown model {request.model!r}") + raise ValueError(f"Unknown model {request.model!r}") + if request.role not in (engine_pb2.ENGINE_ROLE_UNSPECIFIED, self.role): + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "Requested health role does not match this engine", + ) + raise ValueError("Requested health role does not match this engine") + if request.include_inference_probe: + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "Role-safe inference probes are not implemented by this server", + ) + raise NotImplementedError("Role-safe inference probes are not implemented") + healthy, message = await self.tracker.health() + state = ( + lifecycle_pb2.HEALTH_STATE_DRAINING + if self.tracker.draining + else lifecycle_pb2.HEALTH_STATE_READY + if healthy + else lifecycle_pb2.HEALTH_STATE_NOT_READY + ) + checks = [lifecycle_pb2.HealthCheck(name="scheduler", state=state, message=message)] + return lifecycle_pb2.HealthResponse(state=state, checks=checks) + + async def Abort( + self, request: lifecycle_pb2.AbortRequest, context: grpc.aio.ServicerContext + ) -> lifecycle_pb2.AbortResponse: + target = request.WhichOneof("target") + if target is None: + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "Abort target must be specified") + raise ValueError("Abort target must be specified") + if target == "all_requests": + count = await self.tracker.abort_all() + self._release_all_kv_sessions() + return lifecycle_pb2.AbortResponse( + status=lifecycle_pb2.ABORT_STATUS_ABORTED, message=f"Aborted {count} requests" + ) + request_id = request.request_id + released_session = False + if target == "kv_session": + session_id = request.kv_session.session_id + request_id = self._kv_session_requests.get(session_id, session_id) + released_session = self._release_kv_session(session_id) + aborted = await self.tracker.abort(request_id) + return lifecycle_pb2.AbortResponse( + status=( + lifecycle_pb2.ABORT_STATUS_ABORTED + if aborted or released_session + else lifecycle_pb2.ABORT_STATUS_ALREADY_FINISHED + ), + message=( + f"Aborted {request_id}" + if aborted + else f"Released KV session {request.kv_session.session_id}" + if released_session + else f"Request {request_id} is not active" + ), + ) + + async def Drain( + self, request: lifecycle_pb2.DrainRequest, context: grpc.aio.ServicerContext + ) -> AsyncGenerator[lifecycle_pb2.DrainResponse, None]: + del context + if request.stop_accepting_new_requests: + await self.tracker.start_drain() + yield lifecycle_pb2.DrainResponse( + state=lifecycle_pb2.DRAIN_STATE_STARTED, + in_flight_requests=self.tracker.active_count, + open_kv_sessions=len(self._kv_session_requests), + message="Process-wide drain started", + ) + timeout = request.deadline_ms / 1000.0 if request.HasField("deadline_ms") else None + empty = await self.tracker.wait_empty(timeout) + if not empty and request.abort_after_deadline: + await self.tracker.abort_all() + empty = await self.tracker.wait_empty(self._post_abort_cleanup_timeout_seconds) + if not empty: + yield lifecycle_pb2.DrainResponse( + error=error_pb2.EngineError( + code=error_pb2.ERROR_CODE_INTERNAL, + message="Drain deadline expired with active requests", + retryable=False, + ), + in_flight_requests=self.tracker.active_count, + open_kv_sessions=len(self._kv_session_requests), + ) + return + self._release_all_kv_sessions() + yield lifecycle_pb2.DrainResponse( + state=lifecycle_pb2.DRAIN_STATE_COMPLETE, + in_flight_requests=0, + open_kv_sessions=0, + message="Drain complete", + ) + + async def LoadLora( + self, request: lora_pb2.LoadLoraRequest, context: grpc.aio.ServicerContext + ) -> lora_pb2.LoadLoraResponse: + if not self._supports_lora(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "TensorRT-LLM was not configured with LoRA support", + ) + raise RuntimeError("TensorRT-LLM was not configured with LoRA support") + try: + adapter, already_loaded = await self.loras.load(request.adapter) + except ValueError as error: + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) + raise + self._lora_names_by_id[adapter.lora_id] = adapter.lora_name + return lora_pb2.LoadLoraResponse(adapter=adapter, already_loaded=already_loaded) + + async def UnloadLora( + self, request: lora_pb2.UnloadLoraRequest, context: grpc.aio.ServicerContext + ) -> lora_pb2.UnloadLoraResponse: + if not self._supports_lora(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "TensorRT-LLM was not configured with LoRA support", + ) + raise RuntimeError("TensorRT-LLM was not configured with LoRA support") + try: + adapter = await self.loras.unload(request.lora_name) + except KeyError: + await context.abort( + grpc.StatusCode.NOT_FOUND, f"LoRA {request.lora_name!r} is not loaded" + ) + raise + return lora_pb2.UnloadLoraResponse(adapter=adapter) + + async def ListLoras( + self, request: lora_pb2.ListLorasRequest, context: grpc.aio.ServicerContext + ) -> lora_pb2.ListLorasResponse: + del request + if not self._supports_lora(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "TensorRT-LLM was not configured with LoRA support", + ) + raise RuntimeError("TensorRT-LLM was not configured with LoRA support") + return lora_pb2.ListLorasResponse(adapters=await self.loras.list()) + + def _supports_lora(self) -> bool: + args = getattr(self.llm, "args", None) + if getattr(args, "lora_config", None) is not None or bool( + getattr(args, "enable_lora", False) + ): + return True + build_config = getattr(args, "build_config", None) + plugin_config = getattr(build_config, "plugin_config", None) + return bool(getattr(plugin_config, "lora_plugin", False)) + + def _kv_connector_info(self) -> kv_pb2.KvConnectorInfo: + enabled = self.role in (engine_pb2.ENGINE_ROLE_PREFILL, engine_pb2.ENGINE_ROLE_DECODE) + return kv_pb2.KvConnectorInfo( + enabled=enabled, + transfer_backend="tensorrt_llm" if enabled else "", + supported_protocols=["nixl"] if enabled else [], + supports_remote_prefill=enabled, + supports_decode_pull=enabled, + supports_abort_cleanup=enabled, + supports_drain=enabled, + schema_version=1, + ) + + def _kv_events_enabled(self) -> bool: + args = getattr(self.llm, "args", None) + config = getattr(args, "kv_cache_config", None) + if self.kv_events is None or getattr(config, "event_buffer_max_size", 0) <= 0: + return False + effective_hash_algo = get_effective_kv_cache_event_hash_algo( + getattr(config, "kv_cache_event_hash_algo", "auto"), + bool(getattr(config, "use_kv_cache_manager_v2", False)), + ) + return effective_hash_algo in (KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2_SHA256_64) + + async def GetKvConnectorInfo( + self, request: kv_pb2.GetKvConnectorInfoRequest, context: grpc.aio.ServicerContext + ) -> kv_pb2.KvConnectorInfo: + del request, context + return self._kv_connector_info() + + async def GetKvEventSources( + self, request: kv_pb2.GetKvEventSourcesRequest, context: grpc.aio.ServicerContext + ) -> kv_pb2.GetKvEventSourcesResponse: + if not self._kv_events_enabled(): + return kv_pb2.GetKvEventSourcesResponse() + data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) + available_ranks = set(range(data_parallel_size)) + requested_ranks = set(request.data_parallel_ranks) + invalid_ranks = requested_ranks - available_ranks + if invalid_ranks: + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + f"Unknown data-parallel ranks: {sorted(invalid_ranks)}", + ) + return kv_pb2.GetKvEventSourcesResponse() + selected_ranks = sorted(requested_ranks or available_ranks) + return kv_pb2.GetKvEventSourcesResponse( + sources=[ + kv_pb2.KvEventSource( + transport="grpc", + endpoint_addr=kv_pb2.KvEndpoint( + host=self.event_host, port=self.event_port, protocol="grpc" + ), + data_parallel_rank=rank, + encoding="protobuf", + schema_version=1, + max_queue_size=self.kv_events.buffer_size, + ) + for rank in selected_ranks + ] + ) + + async def SubscribeKvEvents( + self, request: kv_pb2.SubscribeKvEventsRequest, context: grpc.aio.ServicerContext + ) -> AsyncGenerator[kv_pb2.SubscribeKvEventsResponse, None]: + if not self._kv_events_enabled(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, "KV event collection is disabled" + ) + return + if request.include_snapshot or request.start_sequence_number: + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "KV event snapshot and replay subscriptions are not implemented", + ) + return + selected_ranks = set(request.data_parallel_ranks) + available_ranks = set(range(int(_arg(self.llm, "data_parallel_size", 1)))) + invalid_ranks = selected_ranks - available_ranks + if invalid_ranks: + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + f"Unknown data-parallel ranks: {sorted(invalid_ranks)}", + ) + return + async for sequence, raw in self.kv_events.subscribe(selected_ranks): + if context.cancelled(): + return + rank = raw.get("attention_dp_rank", 0) + if selected_ranks and rank not in selected_ranks: + continue + yield kv_pb2.SubscribeKvEventsResponse(batch=self._kv_batch(raw, sequence)) + + def _kv_batch(self, raw: dict[str, Any], sequence: int) -> kv_pb2.KvEventBatch: + rank = raw.get("attention_dp_rank", 0) + group_idx = raw.get("layer_group_id", 0) + block_size = self._kv_capacity().get("tokensPerBlock") + hash_algo = raw.get("hash_algo") + data = raw.get("data", raw) + event_type = data.get("type") + events: list[kv_pb2.KvEvent] = [] + if event_type == "stored": + if block_size is None or block_size <= 0: + raise ValueError("Cannot bridge KV events without a configured tokens_per_block") + parent = data.get("parent_hash") + blocks = data.get("blocks", []) + complete_blocks = [] + for block in blocks: + block_hash = block.get("block_hash") + tokens = block.get("tokens", []) + if len(tokens) > block_size: + raise ValueError("TRT-LLM KV event block exceeds tokens_per_block") + if len(tokens) < block_size: + self._partial_block_hashes.setdefault(rank, set()).add(block_hash) + break + if block.get("cache_salt") not in (None, ""): + return self._kv_fail_closed_batch( + rank, + sequence, + "Salted TRT-LLM KV event cannot be represented by OpenEngine", + ) + if any(token.get("token_extra_id", 0) for token in tokens): + return self._kv_fail_closed_batch( + rank, + sequence, + "Eagle TRT-LLM KV event cannot be represented by OpenEngine", + ) + token_ids = [token.get("token_id") for token in tokens] + if any( + not isinstance(token_id, int) or not 0 <= token_id < (1 << 32) + for token_id in token_ids + ): + return self._kv_fail_closed_batch( + rank, + sequence, + "Non-text TRT-LLM KV event contains non-integer tokens", + ) + complete_blocks.append(block) + if complete_blocks: + lora_ids = { + int(block["lora_id"]) + for block in complete_blocks + if block.get("lora_id") is not None + } + lora_names = { + str(block.get("lora_name") or data.get("lora_name")) + for block in complete_blocks + if block.get("lora_name") or data.get("lora_name") + } + if len(lora_ids) > 1 or len(lora_names) > 1: + return self._kv_fail_closed_batch( + rank, + sequence, + "TRT-LLM KV event mixes LoRA identities", + ) + lora_id = next(iter(lora_ids), None) + lora_name = next( + iter(lora_names), + self._lora_names_by_id.get(lora_id, "") if lora_id is not None else "", + ) + mm_keys = [] + for block_index, block in enumerate(complete_blocks): + for key in block.get("mm_keys", []): + mm_hash = key.get("hash") + if ( + key.get("type") != "mm_key" + or not isinstance(mm_hash, str) + or not mm_hash + ): + return self._kv_fail_closed_batch( + rank, + sequence, + "TRT-LLM KV event contains malformed multimodal metadata", + ) + try: + mm_hash_prefix = str(int(mm_hash[:16], 16)) + except ValueError: + return self._kv_fail_closed_batch( + rank, + sequence, + "TRT-LLM KV event contains a non-hex multimodal hash", + ) + mm_keys.append( + kv_pb2.OpaqueKeyTuple( + values=[ + "trt_mm_v1", + str(block_index), + mm_hash_prefix, + str(key.get("start_offset", 0)), + ] + ) + ) + stored = kv_pb2.BlockStored( + block_hashes=[ + _block_hash(block.get("block_hash"), hash_algo) for block in complete_blocks + ], + token_ids=[ + token["token_id"] + for block in complete_blocks + for token in block.get("tokens", []) + ], + block_size=block_size, + lora_id=0 if lora_id is None else int(lora_id), + lora_name=lora_name, + medium=kv_pb2.STORAGE_MEDIUM_GPU, + extra_keys=mm_keys, + group_idx=group_idx, + ) + if parent is not None: + stored.parent_block_hash.CopyFrom(_block_hash(parent, hash_algo)) + events.append(kv_pb2.KvEvent(block_stored=stored)) + else: + stored = kv_pb2.BlockStored( + block_size=block_size, + medium=kv_pb2.STORAGE_MEDIUM_GPU, + group_idx=group_idx, + ) + if parent is not None: + stored.parent_block_hash.CopyFrom(_block_hash(parent, hash_algo)) + events.append(kv_pb2.KvEvent(block_stored=stored)) + elif event_type == "removed": + partial_hashes = self._partial_block_hashes.setdefault(rank, set()) + removed_hashes = [] + for value in data.get("block_hashes", []): + if value in partial_hashes: + partial_hashes.discard(value) + else: + removed_hashes.append(value) + events.append( + kv_pb2.KvEvent( + block_removed=kv_pb2.BlockRemoved( + block_hashes=[_block_hash(value, hash_algo) for value in removed_hashes], + medium=kv_pb2.STORAGE_MEDIUM_GPU, + group_idx=group_idx, + ) + ) + ) + elif event_type == "all_cleared": + self._partial_block_hashes.pop(rank, None) + events.append(kv_pb2.KvEvent(all_blocks_cleared=kv_pb2.AllBlocksCleared())) + return kv_pb2.KvEventBatch( + sequence_number=sequence, + timestamp_unix_nanos=time.time_ns(), + data_parallel_rank=rank, + events=events, + ) + + @staticmethod + def _kv_fail_closed_batch(rank: int, sequence: int, message: str) -> kv_pb2.KvEventBatch: + logger.error("%s; clearing the advertised KV index", message) + return kv_pb2.KvEventBatch( + sequence_number=sequence, + timestamp_unix_nanos=time.time_ns(), + data_parallel_rank=rank, + events=[kv_pb2.KvEvent(all_blocks_cleared=kv_pb2.AllBlocksCleared())], + ) + + async def Embed(self, request: object, context: grpc.aio.ServicerContext) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Embedding is not implemented by this server" + ) + + async def Classify(self, request: object, context: grpc.aio.ServicerContext) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Classification is not implemented by this server" + ) + + async def Score(self, request: object, context: grpc.aio.ServicerContext) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Scoring is not implemented by this server" + ) + + async def SubscribeRuntimeEvents( + self, request: object, context: grpc.aio.ServicerContext + ) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Runtime event streaming is not implemented" + ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py index c308790f9c29..114e414ed808 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -91,6 +91,8 @@ class KVCacheStoredBlockData: priority: int mm_keys: list[MmKey] = field(default_factory=list) cache_salt: str | None = None + lora_id: int | None = None + lora_name: str | None = None @dataclass(slots=True, frozen=True) @@ -552,12 +554,19 @@ def _stored_block_from_radix_block( if life_cycle_ids is not None and not found_page: return None + root = block + while self._is_radix_block(root): + root = root.prev + lora_id, cache_salt = self._root_attrs_from_root_block(root) + return KVCacheStoredBlockData( block_hash=self._hash_from_radix_block(block), tokens=[self._normalize_token(token) for token in block.tokens], cache_level=int(cache_level), priority=int(priority), mm_keys=[], + cache_salt=None if cache_salt is None else str(cache_salt), + lora_id=lora_id, ) @staticmethod diff --git a/tensorrt_llm/serve/kv_event_fanout.py b/tensorrt_llm/serve/kv_event_fanout.py new file mode 100644 index 000000000000..60fce03c8a31 --- /dev/null +++ b/tensorrt_llm/serve/kv_event_fanout.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-consumer TRT-LLM KV-event pump with in-process fanout.""" + +import asyncio +from collections import deque +from collections.abc import AsyncGenerator +from typing import Any + +from tensorrt_llm.logger import logger + + +class KvEventFanout: + """Own TRT-LLM's event queue and copy events to every protocol.""" + + def __init__(self, llm: object, buffer_size: int = 1024) -> None: + if buffer_size <= 0: + raise ValueError("buffer_size must be positive") + self._llm = llm + self._buffer: deque[dict[str, Any]] = deque(maxlen=buffer_size) + self._subscribers: dict[asyncio.Queue[tuple[int, dict[str, Any]]], frozenset[int]] = {} + self._task: asyncio.Task[None] | None = None + self._sequence_numbers: dict[int, int] = {} + self._last_engine_event_ids: dict[int, int] = {} + self._max_window_size: int | None = None + self._processing_initial_created_events = True + + @property + def buffer_size(self) -> int: + return self._buffer.maxlen or 0 + + def start(self) -> None: + if self._task is None: + self._task = asyncio.create_task(self._pump()) + + async def stop(self) -> None: + if self._task is None: + return + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + def drain_http_buffer(self) -> list[dict[str, Any]]: + events = list(self._buffer) + self._buffer.clear() + return events + + async def subscribe( + self, data_parallel_ranks: set[int] | None = None + ) -> AsyncGenerator[tuple[int, dict[str, Any]], None]: + self.start() + queue: asyncio.Queue[tuple[int, dict[str, Any]]] = asyncio.Queue(maxsize=self.buffer_size) + self._subscribers[queue] = frozenset(data_parallel_ranks or ()) + try: + while True: + yield await queue.get() + finally: + self._subscribers.pop(queue, None) + + def _is_routable_event(self, event: dict[str, Any]) -> bool: + event_type = event.get("data", event).get("type") + if event_type == "created" and self._processing_initial_created_events: + window_size = event.get("window_size") + if window_size is not None: + self._max_window_size = max(self._max_window_size or 0, int(window_size)) + return False + self._processing_initial_created_events = False + if event_type not in ("stored", "removed"): + return False + window_size = event.get("window_size") + return ( + window_size is None + or self._max_window_size is None + or int(window_size) == self._max_window_size + ) + + def _publish(self, event: dict[str, Any]) -> None: + rank = int(event.get("attention_dp_rank", 0)) + event_id = event.get("event_id") + if isinstance(event_id, int) and not isinstance(event_id, bool): + previous_event_id = self._last_engine_event_ids.get(rank) + if (previous_event_id is None and event_id != 0) or ( + previous_event_id is not None and event_id != previous_event_id + 1 + ): + logger.warning( + "TRT-LLM KV event gap on DP rank %d: expected %d, received %d; " + "clearing the advertised KV index", + rank, + 0 if previous_event_id is None else previous_event_id + 1, + event_id, + ) + self._publish_routable( + { + "attention_dp_rank": rank, + "data": {"type": "all_cleared"}, + }, + include_http_raw_event=False, + ) + self._last_engine_event_ids[rank] = event_id + if not self._is_routable_event(event): + return + self._publish_routable(event) + + def _publish_routable( + self, event: dict[str, Any], *, include_http_raw_event: bool = True + ) -> None: + rank = int(event.get("attention_dp_rank", 0)) + sequence = self._sequence_numbers.get(rank, 0) + 1 + self._sequence_numbers[rank] = sequence + sequenced_event = (sequence, event) + if include_http_raw_event: + self._buffer.append(event) + for queue, selected_ranks in tuple(self._subscribers.items()): + if selected_ranks and rank not in selected_ranks: + continue + try: + queue.put_nowait(sequenced_event) + except asyncio.QueueFull: + logger.warning("Dropping KV event for a slow OpenEngine subscriber") + + async def _pump(self) -> None: + while True: + try: + async for event in self._llm.get_kv_cache_events_async(1): + self._publish(event) + except (IndexError, asyncio.QueueEmpty): + await asyncio.sleep(0) + except (RuntimeError, AttributeError) as error: + logger.debug("KV event pump unavailable: %s", error) + await asyncio.sleep(1) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index fd5fb55c25c0..030f80dfb552 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -65,6 +65,7 @@ from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterWorker from tensorrt_llm.serve.encode_batcher import (EncodeBatcher, InputTooLongError, QueueFullError) +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, @@ -92,6 +93,7 @@ from tensorrt_llm.serve.responses_utils import get_steady_clock_now_in_seconds from tensorrt_llm.serve.responses_utils import \ request_preprocess as responses_api_request_preprocess +from tensorrt_llm.serve.stats_fanout import StatsFanout from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import \ build_visual_gen_timing_headers @@ -265,7 +267,10 @@ def __init__( embedding_max_queue_delay: float = 0.005, embedding_max_queue_size: int = 2048, input_processor_workers: int = 8, - media_load_workers: int = 8): + media_load_workers: int = 8, + kv_event_fanout: Optional[KvEventFanout] = None, + stats_fanout: Optional[StatsFanout] = None, + shutdown_generator: bool = True): self.generator = generator self._is_visual_gen = isinstance(generator, VisualGen) self._embedding_max_queue_delay = embedding_max_queue_delay @@ -290,6 +295,9 @@ def __init__( self.multimodal_server_config = cfg self.allow_request_chat_template = allow_request_chat_template self.server_role = server_role + self.kv_event_fanout = kv_event_fanout + self.stats_fanout = stats_fanout + self.shutdown_generator = shutdown_generator # Will be set in __call__ self.binding_addr = None self.host = None @@ -346,6 +354,8 @@ def __init__( @asynccontextmanager async def lifespan(app: FastAPI): + if self.kv_event_fanout is not None: + self.kv_event_fanout.start() if self.metadata_server is not None: metadata = { "model": self.model, @@ -384,7 +394,12 @@ async def lifespan(app: FastAPI): # Start background iteration stats collector if metrics are enabled # The args for pytorch and autodeploy backend has attribute `enable_iter_perf_stats` while # tensorrt backend does not have this attribute but it always has iter stats enabled. - if self.metrics_collector and getattr( + if self.stats_fanout is not None and getattr( + self.generator.args, "enable_iter_perf_stats", True): + consumer = (self.metrics_collector.log_iteration_stats + if self.metrics_collector is not None else None) + self.stats_fanout.start(consumer) + elif self.metrics_collector and getattr( self.generator.args, "enable_iter_perf_stats", True): # The background loop becomes the sole consumer of the # engine stats queue; /metrics reads from a tee buffer @@ -432,7 +447,12 @@ async def lifespan(app: FastAPI): await self.disagg_cluster_worker.deregister_worker() if self.resource_governor is not None: self.resource_governor.close() - self.generator.shutdown() + if self.kv_event_fanout is not None: + await self.kv_event_fanout.stop() + if self.stats_fanout is not None: + await self.stats_fanout.stop() + if self.shutdown_generator: + self.generator.shutdown() self.app = FastAPI(lifespan=lifespan) if _MSGSPEC_ENABLED: @@ -1193,6 +1213,9 @@ async def get_iteration_stats(self) -> JSONResponse: stats.append(stat) return JSONResponse(content=stats) + if self.stats_fanout is not None: + return JSONResponse(content=self.stats_fanout.drain_http_buffer()) + # When the background collector loop is active it is the sole # consumer of the engine stats queue; serve /metrics from the tee # buffer it populates so we do not race it for queue items. Racing @@ -1316,6 +1339,9 @@ async def get_perf_metrics(self) -> JSONResponse: return JSONResponse(content=list(perf_metrics)) async def get_kv_cache_events(self) -> JSONResponse: + if self.kv_event_fanout is not None: + return JSONResponse( + content=self.kv_event_fanout.drain_http_buffer()) events = [] try: async for event in self.generator.get_kv_cache_events_async(0): @@ -1337,10 +1363,13 @@ async def _extract_metrics(self, res: RequestOutput, raw_request: Request): # (e.g. PostprocWorker). self.metrics_collector.log_request_metrics_dict( res.metrics_dict) - # Note: Iteration stats are collected by the background _iteration_stats_collector_loop task - # Wake up the stats collector to drain iteration stats - if getattr(self.generator.args, "enable_iter_perf_stats", True): + # Note: Iteration stats are collected by the background collector. + if (self.stats_fanout is None and getattr( + self.generator.args, "enable_iter_perf_stats", True)): self._iteration_stats_wakeup_event.set() + if self.stats_fanout is not None and getattr( + self.generator.args, "enable_iter_perf_stats", True): + self.stats_fanout.wake() if self.generator.args.return_perf_metrics: output = res.outputs[0] item = { diff --git a/tensorrt_llm/serve/request_tracker.py b/tensorrt_llm/serve/request_tracker.py new file mode 100644 index 000000000000..8c6ebe3c8674 --- /dev/null +++ b/tensorrt_llm/serve/request_tracker.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Protocol-neutral tracking for requests submitted to an LLM instance.""" + +import asyncio +from collections.abc import Iterable +from typing import Protocol + +from tensorrt_llm.logger import logger + + +class AbortableResult(Protocol): + def abort(self) -> None: + """Abort the associated engine request.""" + + +class RequestTracker: + """Track active engine results for admission, cancellation, and drain.""" + + def __init__(self, llm: object) -> None: + self.llm = llm + self._requests: dict[str, AbortableResult] = {} + self._external_requests = 0 + self._draining = False + self._changed = asyncio.Condition() + + @property + def draining(self) -> bool: + return self._draining + + @property + def active_count(self) -> int: + return len(self._requests) + self._external_requests + + @property + def active_requests(self) -> dict[str, AbortableResult]: + """Compatibility view for protocol request managers.""" + return self._requests + + def admit(self, request_id: str, result: AbortableResult) -> None: + if self._draining: + raise RuntimeError("TensorRT-LLM is draining") + if request_id in self._requests: + raise ValueError(f"Request {request_id!r} is already active") + self._requests[request_id] = result + + async def finish(self, request_id: str) -> None: + async with self._changed: + self._requests.pop(request_id, None) + self._changed.notify_all() + + def begin_external(self) -> None: + if self._draining: + raise RuntimeError("TensorRT-LLM is draining") + self._external_requests += 1 + + async def finish_external(self) -> None: + async with self._changed: + self._external_requests = max(0, self._external_requests - 1) + self._changed.notify_all() + + async def abort(self, request_id: str) -> bool: + result = self._requests.get(request_id) + if result is None: + return False + try: + result.abort() + except (RuntimeError, AssertionError) as error: + logger.warning("Failed to abort request %s: %s", request_id, error) + return False + finally: + await self.finish(request_id) + return True + + async def abort_all(self) -> int: + tracked_by_identity = {id(result): result for result in self._requests.values()} + executor = getattr(self.llm, "_executor", None) + executor_results = getattr(executor, "_results", None) + if isinstance(executor_results, dict): + tracked_by_identity.update( + (id(result), result) for result in tuple(executor_results.values()) + ) + + aborted = 0 + for result in tracked_by_identity.values(): + try: + result.abort() + except (RuntimeError, AssertionError) as error: + logger.warning("Failed to abort engine request: %s", error) + else: + aborted += 1 + + async with self._changed: + self._requests.clear() + self._changed.notify_all() + return aborted + + async def start_drain(self) -> int: + async with self._changed: + self._draining = True + self._changed.notify_all() + return self.active_count + + async def wait_empty(self, timeout: float | None = None) -> bool: + async def _wait() -> None: + async with self._changed: + await self._changed.wait_for(lambda: self.active_count == 0) + + if self.active_count == 0: + return True + try: + await asyncio.wait_for(_wait(), timeout=timeout) + except asyncio.TimeoutError: + return False + return True + + async def health(self) -> tuple[bool, str]: + if not hasattr(self.llm, "_executor"): + return True, "OK" + executor = getattr(self.llm, "_executor", None) + if executor is None: + return False, "Executor is not available" + try: + healthy = executor.check_health() + except (RuntimeError, AttributeError) as error: + return False, f"Executor health check failed: {error}" + if healthy: + return True, "OK" + fatal_error = getattr(executor, "_fatal_error", None) + if fatal_error is None: + return False, "Executor is unhealthy" + lines = str(fatal_error).splitlines() + short = (lines[0] if lines else type(fatal_error).__name__)[:200] + return False, f"{type(fatal_error).__name__}: {short}" + + def iter_results(self) -> Iterable[AbortableResult]: + return tuple(self._requests.values()) + + +async def track_http_response(response: object, tracker: RequestTracker) -> object: + """Keep HTTP admission active until a streaming response body closes.""" + body_iterator = getattr(response, "body_iterator", None) + if body_iterator is None: + await tracker.finish_external() + return response + + async def tracked_body(): + try: + async for chunk in body_iterator: + yield chunk + finally: + await tracker.finish_external() + + response.body_iterator = tracked_body() + return response diff --git a/tensorrt_llm/serve/stats_fanout.py b/tensorrt_llm/serve/stats_fanout.py new file mode 100644 index 000000000000..7b2093efe3eb --- /dev/null +++ b/tensorrt_llm/serve/stats_fanout.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-consumer TRT-LLM iteration-stat pump with shared snapshots.""" + +import asyncio +from collections import deque +from collections.abc import Callable +from typing import Any + +from tensorrt_llm.logger import logger + + +class StatsFanout: + """Own the engine stats queue and tee it to HTTP, metrics, and OpenEngine.""" + + def __init__(self, llm: object, buffer_size: int | None = 1000) -> None: + self._llm = llm + self._buffer: deque[dict[str, Any]] = deque(maxlen=buffer_size) + self._latest_by_rank: dict[int, dict[str, Any]] = {} + self._consumer: Callable[[dict[str, Any]], None] | None = None + self._wake_event: asyncio.Event | None = None + self._task: asyncio.Task[None] | None = None + + def start(self, consumer: Callable[[dict[str, Any]], None] | None = None) -> None: + self._consumer = consumer + if self._task is None: + self._wake_event = asyncio.Event() + self._task = asyncio.create_task(self._pump()) + + def wake(self) -> None: + if self._wake_event is not None: + self._wake_event.set() + + async def stop(self) -> None: + if self._task is None: + return + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + self._wake_event = None + + def drain_http_buffer(self) -> list[dict[str, Any]]: + stats = list(self._buffer) + self._buffer.clear() + return stats + + def latest_by_rank(self) -> dict[int, dict[str, Any]]: + return {rank: dict(stat) for rank, stat in self._latest_by_rank.items()} + + def _publish(self, stat: dict[str, Any]) -> None: + rank = int(stat.get("attentionDpRank", stat.get("attention_dp_rank", 0))) + self._latest_by_rank[rank] = stat + self._buffer.append(stat) + if self._consumer is not None: + try: + self._consumer(stat) + except (RuntimeError, ValueError, TypeError) as error: + logger.error("Iteration-stat consumer failed: %s", error) + + async def _pump(self) -> None: + assert self._wake_event is not None + while True: + await self._wake_event.wait() + self._wake_event.clear() + try: + async for stat in self._llm.get_stats_async(timeout=0.5): + self._publish(stat) + except ( + RuntimeError, + AttributeError, + IndexError, + TypeError, + ValueError, + asyncio.QueueEmpty, + ) as error: + logger.error("Error collecting iteration stats: %s", error) + await asyncio.sleep(0.1) diff --git a/tests/unittest/_torch/modeling/test_modeling_phi4mm.py b/tests/unittest/_torch/modeling/test_modeling_phi4mm.py index 58c7de77f938..42013d661070 100644 --- a/tests/unittest/_torch/modeling/test_modeling_phi4mm.py +++ b/tests/unittest/_torch/modeling/test_modeling_phi4mm.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import List +import pytest import torch import transformers from test_modeling_multimodal import MultimodalScenario, TestModelingMultimodal, llm_models_root @@ -10,9 +11,25 @@ _AUDIO_SPECIAL_TOKEN_ID, _IMAGE_SPECIAL_TOKEN_ID, Phi4MMForCausalLM, + Phi4MMInputProcessor, ) from tensorrt_llm.inputs import default_multimodal_input_loader, prompt_inputs + +def test_phi4mm_required_multimodal_lora_selection(tmp_path) -> None: + processor = object.__new__(Phi4MMInputProcessor) + processor._model_path = str(tmp_path) + + assert processor.get_model_owned_lora_identities() == { + "vision-lora": 0, + "speech-lora": 1, + } + assert processor.get_required_lora_spec(("image",)).name == "vision-lora" + assert processor.get_required_lora_spec(("audio",)).name == "speech-lora" + with pytest.raises(ValueError, match="cannot combine"): + processor.get_required_lora_spec(("image", "audio")) + + PHI4MM_CONFIG = { "_name_or_path": str(os.path.join(llm_models_root(), "multimodals/Phi-4-multimodal-instruct")), "architectures": ["Phi4MMForCausalLM"], diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 4a80a5095240..67c595b6f8fe 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -321,6 +321,24 @@ commands: is_flag: false flags: - "--num_media_load_workers" + openengine_host: + type: str + default: 127.0.0.1 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--openengine-host" + openengine_port: + type: int + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--openengine-port" otlp_traces_endpoint: type: str default: null diff --git a/tests/unittest/openengine/test_converters.py b/tests/unittest/openengine/test_converters.py new file mode 100644 index 000000000000..b1dc4236eb4f --- /dev/null +++ b/tests/unittest/openengine/test_converters.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import pytest +from openengine.v1 import generation_pb2, input_pb2, kv_pb2 + +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.executor.result import Logprob +from tensorrt_llm.openengine.converters import ( + HANDOFF_ATTRIBUTE, + decode_handoff, + encode_handoff, + handoff_requires_decode_media, + load_media, + to_priority, + to_sampling_params, +) + + +def test_context_first_handoff_round_trip_large_ids_and_binary() -> None: + params = DisaggregatedParams( + request_type="context_only", + first_gen_tokens=[11, 12], + first_gen_log_probs=[ + { + 11: Logprob(logprob=-0.25, rank=1), + 91: Logprob(logprob=-2.0, rank=2), + }, + {12: Logprob(logprob=-0.5, rank=1)}, + ], + ctx_request_id=2**63 - 2, + disagg_request_id=2**63 - 1, + ctx_dp_rank=7, + ctx_info_endpoint="tcp://context:1234", + draft_tokens=[21, 22], + ctx_usage={"prompt_tokens": 9}, + conversation_id="conversation", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + opaque_state=b"\x00\xffbinary", + ) + + session = encode_handoff(params) + payload = json.loads(session.attributes_struct[HANDOFF_ATTRIBUTE]) + assert payload["ctx_request_id"] == str(2**63 - 2) + assert payload["disagg_request_id"] == str(2**63 - 1) + + restored = decode_handoff(session) + assert restored.request_type == "generation_only" + assert restored.ctx_request_id == 2**63 - 2 + assert restored.disagg_request_id == 2**63 - 1 + assert restored.opaque_state == b"\x00\xffbinary" + assert restored.first_gen_log_probs[0][11] == Logprob(-0.25, 1) + assert restored.first_gen_log_probs[0][91] == Logprob(-2.0, 2) + assert restored.first_gen_log_probs[1][12] == Logprob(-0.5, 1) + + +def test_raw_media_handoff_omits_transient_mrope_and_requires_decode_media() -> None: + params = DisaggregatedParams( + request_type="context_only", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + mrope_position_ids_handle={"tensor": "ids"}, + mrope_position_deltas_handle={"tensor": "deltas"}, + ) + + with pytest.raises(ValueError, match="mrope"): + encode_handoff(params) + + session = encode_handoff(params, requires_decode_media=True) + assert handoff_requires_decode_media(session) + restored = decode_handoff(session) + assert restored.mrope_position_ids_handle is None + assert restored.mrope_position_deltas_handle is None + + +@pytest.mark.parametrize( + "params, message", + [ + ( + DisaggregatedParams( + request_type="context_only", + schedule_style=DisaggScheduleStyle.GENERATION_FIRST, + ), + "Generation-first", + ), + (DisaggregatedParams(request_type="context_only", first_gen_logits=[object()]), "logits"), + ( + DisaggregatedParams( + request_type="context_only", + multimodal_embedding_handles=[{"handle": "encoder"}], + ), + "embedding handles", + ), + ], +) +def test_handoff_rejects_out_of_scope_topologies(params: DisaggregatedParams, message: str) -> None: + with pytest.raises(ValueError, match=message): + encode_handoff(params) + + +def test_decode_rejects_non_decimal_identifier() -> None: + session = kv_pb2.KvSessionRef() + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( + { + "schedule_style": "context_first", + "ctx_request_id": 9007199254740993, + } + ) + with pytest.raises(ValueError, match="decimal string"): + decode_handoff(session) + + +def test_decode_rejects_multimodal_hashes_even_without_embedding_handles() -> None: + session = kv_pb2.KvSessionRef() + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( + { + "schedule_style": "context_first", + "multimodal_hashes": [[1, 2, 3, 4, 5, 6, 7, 8]], + } + ) + + with pytest.raises(ValueError, match="handles"): + decode_handoff(session) + + +def test_decode_rejects_malformed_context_usage() -> None: + session = kv_pb2.KvSessionRef() + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( + { + "schedule_style": "context_first", + "ctx_usage": {"prompt_tokens": "2"}, + } + ) + + with pytest.raises(ValueError, match="prompt_tokens"): + decode_handoff(session) + + +def test_sampling_conversion_preserves_requested_controls() -> None: + request = generation_pb2.GenerateRequest( + sampling=generation_pb2.SamplingParams( + temperature=0.25, + top_p=0.8, + top_k=16, + seed=2**63, + num_sequences=2, + ), + stopping=generation_pb2.StoppingOptions( + max_tokens=8, + min_tokens=2, + conditions=[ + generation_pb2.StopCondition(stop_text="done"), + generation_pb2.StopCondition(stop_token_id=42), + ], + ignore_eos=True, + include_stop_in_output=True, + ), + response=generation_pb2.ResponseOptions( + return_output_logprobs=True, + output_candidates=generation_pb2.CandidateTokenSelection(top_n=4), + ), + guided=generation_pb2.GuidedDecoding(regex="[a-z]+"), + ) + params = to_sampling_params(request) + assert params.max_tokens == 8 + assert params.n == 2 + assert params.best_of == 2 + assert params.stop == ["done"] + assert params.stop_token_ids == [42] + assert params.logprobs == 4 + assert params.guided_decoding.regex == "[a-z]+" + + +def test_choice_guidance_escapes_literal_regex_characters() -> None: + request = generation_pb2.GenerateRequest() + request.guided.choice.choices.extend(["a.b", "c+(d)"]) + params = to_sampling_params(request) + assert params.guided_decoding.regex == r"^(?:(?:a\.b)|(?:c\+\(d\)))$" + + +def test_priority_mapping_is_centered_bounded_and_strictly_monotonic() -> None: + absent = generation_pb2.GenerateRequest() + negative = generation_pb2.GenerateRequest(priority=-1000) + zero = generation_pb2.GenerateRequest(priority=0) + positive = generation_pb2.GenerateRequest(priority=1000) + + assert 0 < to_priority(negative) < to_priority(zero) + assert to_priority(zero) == to_priority(absent) == 0.5 + assert to_priority(zero) < to_priority(positive) < 1 + + +@pytest.mark.asyncio +async def test_media_preserves_per_modality_order_and_merge_inputs(monkeypatch) -> None: + calls = [] + + class _MediaIO: + @classmethod + def create(cls, defaults, request): + calls.append((defaults, request)) + return cls() + + async def _run_in_executor(self, function, data): + return function(data) + + def load_bytes(self, data): + return data.decode() + + monkeypatch.setattr( + "tensorrt_llm.openengine.converters.MEDIA_IO_REGISTRY", + {"image": _MediaIO, "video": _MediaIO, "audio": _MediaIO}, + ) + options = generation_pb2.GenerateRequest().media_options + options.update({"image": {"format": "pil"}}) + config = type("Config", (), {"media_io_kwargs": {"image": {"device": "cpu"}}})() + media = [ + input_pb2.MediaItem(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"one"), + input_pb2.MediaItem(modality=input_pb2.MODALITY_VIDEO, raw_bytes=b"video"), + input_pb2.MediaItem(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"two"), + ] + + decoded = await load_media(media, options, config) + + assert decoded == {"image": ["one", "two"], "video": ["video"]} + assert calls[0] == ({"device": "cpu"}, {"format": "pil"}) diff --git a/tests/unittest/openengine/test_fanout.py b/tests/unittest/openengine/test_fanout.py new file mode 100644 index 000000000000..623910b7c0fa --- /dev/null +++ b/tests/unittest/openengine/test_fanout.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.stats_fanout import StatsFanout + + +def test_kv_fanout_sequence_exposes_subscriber_queue_drops() -> None: + fanout = KvEventFanout(object(), buffer_size=1) + queue = asyncio.Queue(maxsize=1) + fanout._subscribers[queue] = frozenset({0}) + + first = {"data": {"type": "stored"}, "event": 1, "attention_dp_rank": 0} + second = {"data": {"type": "stored"}, "event": 2, "attention_dp_rank": 0} + third = {"data": {"type": "stored"}, "event": 3, "attention_dp_rank": 0} + other_rank = {"data": {"type": "stored"}, "event": 20, "attention_dp_rank": 1} + fanout._publish(first) + fanout._publish(second) + assert queue.get_nowait() == (1, first) + + fanout._publish(other_rank) + fanout._publish(third) + assert queue.get_nowait() == (3, third) + assert fanout.drain_http_buffer() == [third] + assert fanout._sequence_numbers == {0: 3, 1: 1} + + +def test_kv_fanout_filters_non_global_attention_without_sequence_gaps() -> None: + fanout = KvEventFanout(object()) + fanout._publish({"window_size": 128, "data": {"type": "created"}}) + fanout._publish({"window_size": 4096, "data": {"type": "created"}}) + fanout._publish({"window_size": 128, "data": {"type": "stored"}}) + global_event = {"window_size": 4096, "data": {"type": "stored"}} + fanout._publish(global_event) + + assert fanout._sequence_numbers == {0: 1} + assert fanout.drain_http_buffer() == [global_event] + + +def test_kv_fanout_detects_raw_event_id_gap_before_attention_filtering() -> None: + fanout = KvEventFanout(object()) + queue = asyncio.Queue() + fanout._subscribers[queue] = frozenset({0}) + fanout._publish({"event_id": 0, "window_size": 128, "data": {"type": "created"}}) + fanout._publish({"event_id": 1, "window_size": 4096, "data": {"type": "created"}}) + fanout._publish({"event_id": 2, "window_size": 128, "data": {"type": "stored"}}) + global_event = {"event_id": 3, "window_size": 4096, "data": {"type": "stored"}} + fanout._publish(global_event) + fanout._publish({"event_id": 5, "window_size": 128, "data": {"type": "stored"}}) + next_global = {"event_id": 6, "window_size": 4096, "data": {"type": "removed"}} + fanout._publish(next_global) + + reset = {"attention_dp_rank": 0, "data": {"type": "all_cleared"}} + assert fanout._last_engine_event_ids == {0: 6} + assert fanout._sequence_numbers == {0: 3} + assert [queue.get_nowait() for _ in range(3)] == [ + (1, global_event), + (2, reset), + (3, next_global), + ] + assert fanout.drain_http_buffer() == [global_event, next_global] + + +def test_synthetic_gap_reset_does_not_change_http_raw_event_contract() -> None: + fanout = KvEventFanout(object()) + queue = asyncio.Queue() + fanout._subscribers[queue] = frozenset() + fanout._publish({"event_id": 0, "data": {"type": "created"}}) + stored = {"event_id": 2, "data": {"type": "stored"}} + fanout._publish(stored) + + assert queue.get_nowait()[1]["data"]["type"] == "all_cleared" + assert queue.get_nowait() == (2, stored) + assert fanout.drain_http_buffer() == [stored] + + +@pytest.mark.asyncio +async def test_stats_fanout_is_single_consumer_for_http_metrics_and_load() -> None: + class _Llm: + calls = 0 + + async def get_stats_async(self, timeout: float): + assert timeout == 0.5 + self.calls += 1 + yield {"attentionDpRank": 0, "numActiveRequests": 1} + yield {"attentionDpRank": 1, "numActiveRequests": 2} + + llm = _Llm() + consumed = [] + fanout = StatsFanout(llm, buffer_size=4) + fanout.start(consumed.append) + fanout.wake() + for _ in range(20): + if len(consumed) == 2: + break + await asyncio.sleep(0) + await fanout.stop() + + assert llm.calls == 1 + assert consumed == fanout.drain_http_buffer() + assert set(fanout.latest_by_rank()) == {0, 1} diff --git a/tests/unittest/openengine/test_lora_registry.py b/tests/unittest/openengine/test_lora_registry.py new file mode 100644 index 000000000000..81477be16f47 --- /dev/null +++ b/tests/unittest/openengine/test_lora_registry.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from pathlib import Path + +import pytest +from openengine.v1 import lora_pb2 + +from tensorrt_llm.openengine.lora_registry import LoraRegistry + + +def _adapter(path: Path, name: str = "adapter", adapter_id: int = 7) -> lora_pb2.LoraAdapter: + path.mkdir() + (path / "adapter_config.json").write_text("{}") + (path / "adapter_model.safetensors").write_bytes(b"weights") + return lora_pb2.LoraAdapter(lora_name=name, lora_id=adapter_id, source_path=str(path)) + + +@pytest.mark.asyncio +async def test_lora_registry_is_lazy_idempotent_and_logically_unloads( + tmp_path: Path, +) -> None: + registry = LoraRegistry() + adapter = _adapter(tmp_path / "adapter") + + registered, already_loaded = await registry.load(adapter) + assert not already_loaded + repeated, already_loaded = await registry.load(adapter) + assert already_loaded + assert repeated == registered + + request = await registry.request("adapter") + assert request.lora_int_id == 7 + assert request.lora_path == str((tmp_path / "adapter").resolve()) + + await registry.unload("adapter") + with pytest.raises(KeyError): + await registry.request("adapter") + + reloaded, already_loaded = await registry.load(adapter) + assert not already_loaded + assert reloaded == registered + + +@pytest.mark.asyncio +async def test_lora_registry_rejects_conflicting_identity(tmp_path: Path) -> None: + registry = LoraRegistry() + await registry.load(_adapter(tmp_path / "one", "one", 1)) + with pytest.raises(ValueError, match="ID 1"): + await registry.load(_adapter(tmp_path / "two", "two", 1)) + + +@pytest.mark.asyncio +async def test_lora_registry_tombstones_name_id_and_path_after_unload(tmp_path: Path) -> None: + registry = LoraRegistry() + original = _adapter(tmp_path / "one", "one", 1) + await registry.load(original) + await registry.unload("one") + + with pytest.raises(ValueError, match="name 'one'.*permanently"): + await registry.load(_adapter(tmp_path / "renamed-path", "one", 2)) + with pytest.raises(ValueError, match="ID 1.*permanently"): + await registry.load(_adapter(tmp_path / "reused-id", "two", 1)) + with pytest.raises(ValueError, match="path.*permanently"): + await registry.load( + lora_pb2.LoraAdapter(lora_name="two", lora_id=2, source_path=original.source_path) + ) + + +@pytest.mark.asyncio +async def test_lora_registry_reserves_model_owned_names_and_ids(tmp_path: Path) -> None: + registry = LoraRegistry({"vision-lora": 0, "speech-lora": 1}) + + with pytest.raises(ValueError, match="name 'vision-lora'.*model-owned"): + await registry.load(_adapter(tmp_path / "reserved-name", "vision-lora", 7)) + with pytest.raises(ValueError, match="ID 0.*model-owned"): + await registry.load(_adapter(tmp_path / "reserved-id", "user-adapter", 0)) + + +@pytest.mark.asyncio +async def test_lora_registry_serializes_concurrent_idempotent_loads(tmp_path: Path) -> None: + registry = LoraRegistry() + adapter = _adapter(tmp_path / "adapter") + + results = await asyncio.gather(registry.load(adapter), registry.load(adapter)) + + assert sorted(already_loaded for _, already_loaded in results) == [False, True] + assert len(await registry.list()) == 1 diff --git a/tests/unittest/openengine/test_native_grpc_compat.py b/tests/unittest/openengine/test_native_grpc_compat.py new file mode 100644 index 000000000000..5c9a0ebc1151 --- /dev/null +++ b/tests/unittest/openengine/test_native_grpc_compat.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from tensorrt_llm.grpc.grpc_request_manager import GrpcRequestManager + + +class _Result: + def __init__(self) -> None: + self.aborted = False + + def abort(self) -> None: + self.aborted = True + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + +class _Llm: + def __init__(self) -> None: + self.last_result = None + + def generate_async(self, *args, **kwargs): + del args, kwargs + self.last_result = _Result() + return self.last_result + + +@pytest.mark.asyncio +async def test_native_grpc_duplicate_aborts_rejected_engine_result() -> None: + llm = _Llm() + manager = GrpcRequestManager(llm) + manager._request_tracker.admit("duplicate", _Result()) + + with pytest.raises(ValueError, match="already active"): + _ = [ + result + async for result in manager.generate( + request_id="duplicate", + prompt_token_ids=[1], + sampling_params=object(), + ) + ] + + assert llm.last_result.aborted diff --git a/tests/unittest/openengine/test_request_tracker.py b/tests/unittest/openengine/test_request_tracker.py new file mode 100644 index 000000000000..5ab44e5a3e5d --- /dev/null +++ b/tests/unittest/openengine/test_request_tracker.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from tensorrt_llm.serve.request_tracker import RequestTracker, track_http_response + + +class _Result: + def __init__(self) -> None: + self.aborted = False + + def abort(self) -> None: + self.aborted = True + + +@pytest.mark.asyncio +async def test_tracker_abort_drain_and_external_admission() -> None: + tracker = RequestTracker(object()) + result = _Result() + tracker.admit("request", result) + tracker.begin_external() + assert tracker.active_count == 2 + + await tracker.start_drain() + with pytest.raises(RuntimeError, match="draining"): + tracker.begin_external() + assert await tracker.abort("request") + assert result.aborted + assert not await tracker.wait_empty(0) + await tracker.finish_external() + assert await tracker.wait_empty(0) + + +@pytest.mark.asyncio +async def test_streaming_http_admission_lives_until_body_closes() -> None: + tracker = RequestTracker(object()) + tracker.begin_external() + + async def body(): + yield b"first" + yield b"second" + + response = type("Response", (), {})() + response.body_iterator = body() + await track_http_response(response, tracker) + + assert tracker.active_count == 1 + chunks = [chunk async for chunk in response.body_iterator] + assert chunks == [b"first", b"second"] + assert tracker.active_count == 0 + + +@pytest.mark.asyncio +async def test_abort_all_reaches_untracked_http_engine_results() -> None: + result = _Result() + llm = type("Llm", (), {})() + llm._executor = type("Executor", (), {"_results": {1: result}})() + tracker = RequestTracker(llm) + tracker.begin_external() + + async def finish_http() -> None: + while not result.aborted: + await asyncio.sleep(0) + await tracker.finish_external() + + finish_task = asyncio.create_task(finish_http()) + assert await tracker.abort_all() == 1 + assert await tracker.wait_empty() + await finish_task + assert result.aborted diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py new file mode 100644 index 000000000000..213046595d45 --- /dev/null +++ b/tests/unittest/openengine/test_server.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest + +from tensorrt_llm.openengine._schema_pin import OPENENGINE_COMMIT +from tensorrt_llm.openengine.server import OpenEngineServer, validate_schema_release +from tensorrt_llm.openengine.servicer import schema_release + + +def test_schema_release_accepts_exact_pinned_identity() -> None: + assert validate_schema_release(OPENENGINE_COMMIT) == OPENENGINE_COMMIT + + +def test_packaged_schema_pin_matches_repository_pin() -> None: + repository_pin = ( + (Path(__file__).resolve().parents[3] / "OPENENGINE_COMMIT") + .read_text(encoding="utf-8") + .strip() + ) + assert OPENENGINE_COMMIT == repository_pin + + +@pytest.mark.parametrize( + "release", + [ + "", + "unreleased", + "main", + "cea19cb", + "latest", + "v0.2.0", + "signed-tag:1.2.3", + "a" * 40, + "a" * 64, + ], +) +def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> None: + with pytest.raises(RuntimeError, match="exactly match"): + validate_schema_release(release) + + +def test_schema_release_reads_launch_environment(monkeypatch) -> None: + commit = "cea19cb06acf03c911b84d5c147e519b60dd92a6" + monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) + assert schema_release() == commit + + +@pytest.mark.asyncio +async def test_server_stop_cleans_protocol_state_and_stats_on_grpc_failure() -> None: + calls = [] + + class _GrpcServer: + async def stop(self, grace: float) -> None: + calls.append(("grpc", grace)) + raise RuntimeError("stop failed") + + class _Servicer: + def close(self) -> None: + calls.append(("servicer", None)) + + class _Stats: + async def stop(self) -> None: + calls.append(("stats", None)) + + class _KvEvents: + async def stop(self) -> None: + calls.append(("kv_events", None)) + + server = object.__new__(OpenEngineServer) + server._server = _GrpcServer() + server.servicer = _Servicer() + server._kv_event_fanout = _KvEvents() + server._stats_fanout = _Stats() + + with pytest.raises(RuntimeError, match="stop failed"): + await server.stop(grace=1.5) + + assert calls == [ + ("grpc", 1.5), + ("servicer", None), + ("kv_events", None), + ("stats", None), + ] diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py new file mode 100644 index 000000000000..2c7dac52bbe0 --- /dev/null +++ b/tests/unittest/openengine/test_servicer.py @@ -0,0 +1,775 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from types import SimpleNamespace + +import pytest +from openengine.v1 import ( + engine_pb2, + generation_pb2, + input_pb2, + kv_pb2, + lifecycle_pb2, + lora_pb2, + model_pb2, + observability_pb2, +) + +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.inputs.registry import MultimodalLoraSpec +from tensorrt_llm.openengine.converters import encode_handoff +from tensorrt_llm.openengine.servicer import OpenEngineServicer +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.request_tracker import RequestTracker +from tensorrt_llm.serve.stats_fanout import StatsFanout + + +class _Tokenizer: + def decode(self, token_ids: list[int], **_: object) -> str: + return "".join(chr(96 + token_id) for token_id in token_ids) + + +class _Result: + def __init__(self, handoff: DisaggregatedParams | None = None) -> None: + self.prompt_token_ids = [1, 2] + self.cached_tokens = 1 + self.disaggregated_params = handoff + self.outputs = [ + SimpleNamespace( + index=0, + token_ids=[3, 4], + text="cd", + logprobs=[], + prompt_logprobs=[], + finish_reason="length", + stop_reason=None, + disaggregated_params=handoff, + ) + ] + self.finished = True + self._yielded = False + self.aborted = False + + def __aiter__(self) -> "_Result": + return self + + async def __anext__(self) -> "_Result": + if self._yielded: + raise StopAsyncIteration + self._yielded = True + return self + + def abort(self) -> None: + self.aborted = True + + +class _Llm: + def __init__(self, handoff: DisaggregatedParams | None = None) -> None: + self.tokenizer = _Tokenizer() + self.args = SimpleNamespace() + self.handoff = handoff + self.kwargs = None + self.result = None + + def generate_async(self, **kwargs: object) -> _Result: + self.kwargs = kwargs + self.result = _Result(self.handoff) + return self.result + + +class _Context: + def __init__(self, cancelled: bool = False) -> None: + self._cancelled = cancelled + + def cancelled(self) -> bool: + return self._cancelled + + async def abort(self, code: object, message: str) -> None: + raise AssertionError(f"unexpected abort {code}: {message}") + + +async def _collect_async(iterator): + return [item async for item in iterator] + + +@pytest.mark.asyncio +async def test_aggregate_streams_deltas_finish_and_terminal_usage() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + servicer = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, tracker) + request = generation_pb2.GenerateRequest( + request_id="request", + model="model", + prompt="hello", + stopping=generation_pb2.StoppingOptions(max_tokens=2), + metadata={ + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", + "ignored": "not-a-trace-header", + }, + ) + + responses = [response async for response in servicer.Generate(request, _Context())] + + assert [response.WhichOneof("event") for response in responses] == [ + "token", + "finished", + ] + assert list(responses[0].token.tokens)[0].token_id == 3 + assert responses[0].token.text == "cd" + assert responses[-1].usage.prompt_tokens == 2 + assert responses[-1].usage.completion_tokens == 2 + assert llm.kwargs["trace_headers"] == { + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", + } + assert tracker.active_count == 0 + + +@pytest.mark.asyncio +async def test_generate_cancellation_aborts_and_cleans_tracking() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, tracker) + request = generation_pb2.GenerateRequest(request_id="cancel", model="model", prompt="hello") + + responses = [response async for response in service.Generate(request, _Context(True))] + + assert responses == [] + assert llm.result.aborted + assert tracker.active_count == 0 + + +def test_generate_rejects_unsupported_cache_bypass_and_nonzero_dp_placement() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") + request.kv.bypass_prefix_cache = True + with pytest.raises(ValueError, match="Prefix-cache bypass"): + service._validate_generate(request) + + request.kv.bypass_prefix_cache = False + request.kv.data_parallel_rank = 0 + service._validate_generate(request) + assert service._scheduling_params(request) is None + + llm.args = SimpleNamespace(data_parallel_size=2) + request.kv.data_parallel_rank = 1 + with pytest.raises(ValueError, match="attention DP"): + service._validate_generate(request) + + +@pytest.mark.asyncio +async def test_generate_applies_strict_attention_dp_placement() -> None: + llm = _Llm() + llm.args = SimpleNamespace(enable_attention_dp=True, data_parallel_size=4) + llm._on_trt_backend = False + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") + request.kv.data_parallel_rank = 2 + + _ = [response async for response in service.Generate(request, _Context())] + scheduling = llm.kwargs["scheduling_params"] + + assert scheduling.attention_dp_rank == 2 + assert not scheduling.attention_dp_relax + + +def test_prefill_preserves_routing_rank_in_context_handoff() -> None: + llm = _Llm() + llm.args = SimpleNamespace(enable_attention_dp=True, data_parallel_size=4) + llm._on_trt_backend = False + service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_PREFILL, RequestTracker(llm)) + request = generation_pb2.GenerateRequest(request_id="prefill", model="model", prompt="hello") + request.kv.data_parallel_rank = 3 + + params = service._disaggregated_params(request) + + assert params.ctx_dp_rank == 3 + + +@pytest.mark.asyncio +async def test_generate_selects_model_owned_multimodal_lora(monkeypatch, tmp_path) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("audio",) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + assert modalities == ("audio",) + return MultimodalLoraSpec("speech-lora", 1, str(tmp_path)) + + async def _load_media(*_args, **_kwargs): + return {"audio": [object()]} + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) + monkeypatch.setattr("tensorrt_llm.openengine.servicer.load_media", _load_media) + llm = _Llm() + llm.args = SimpleNamespace(lora_config=object()) + llm.input_processor = _Processor() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="audio", model="model", prompt="transcribe") + request.media.add(modality=input_pb2.MODALITY_AUDIO, raw_bytes=b"audio") + + _ = [response async for response in service.Generate(request, _Context())] + + assert llm.kwargs["lora_request"].lora_name == "speech-lora" + assert llm.kwargs["lora_request"].lora_int_id == 1 + assert await service.loras.list() == [] + + +@pytest.mark.asyncio +async def test_model_info_suppresses_modalities_when_required_lora_is_unavailable( + monkeypatch, tmp_path +) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("audio",) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + assert modalities == ("audio",) + return MultimodalLoraSpec("speech-lora", 1, str(tmp_path / "missing")) + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) + llm = _Llm() + llm.input_processor = _Processor() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + + info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + + assert not info.supports_multimodal + assert list(info.multimodal_capabilities.aggregate_modalities) == [] + + +@pytest.mark.asyncio +async def test_context_then_generation_round_trips_handoff() -> None: + context_handoff = DisaggregatedParams( + request_type="context_only", + first_gen_tokens=[7], + ctx_request_id=100, + disagg_request_id=101, + ctx_info_endpoint="tcp://context:1", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + opaque_state=b"state", + ) + context_llm = _Llm(context_handoff) + context_service = OpenEngineServicer( + context_llm, + "model", + engine_pb2.ENGINE_ROLE_PREFILL, + RequestTracker(context_llm), + ) + context_request = generation_pb2.GenerateRequest( + request_id="request", + model="model", + token_ids=input_pb2.TokenIds(ids=[1, 2]), + ) + context_responses = [ + response async for response in context_service.Generate(context_request, _Context()) + ] + assert len(context_responses) == 1 + assert context_responses[0].WhichOneof("event") == "prefill_ready" + assert len(context_service._kv_session_requests) == 1 + + decode_llm = _Llm() + decode_service = OpenEngineServicer( + decode_llm, + "model", + engine_pb2.ENGINE_ROLE_DECODE, + RequestTracker(decode_llm), + ) + decode_request = generation_pb2.GenerateRequest( + request_id="decode", + model="model", + token_ids=input_pb2.TokenIds(ids=[1, 2]), + ) + decode_request.kv.session.CopyFrom(context_responses[0].prefill_ready.kv_session) + decode_responses = [ + response async for response in decode_service.Generate(decode_request, _Context()) + ] + decoded = decode_llm.kwargs["disaggregated_params"] + assert decoded.request_type == "generation_only" + assert decoded.disagg_request_id == 101 + assert decoded.opaque_state == b"state" + assert decoded.ctx_usage == { + "prompt_tokens": 2, + "completion_tokens": 2, + "total_tokens": 4, + "prompt_tokens_details": {"cached_tokens": 1}, + } + terminal_usage = decode_responses[-1].usage + assert terminal_usage.prompt_tokens == 2 + assert terminal_usage.cached_prompt_tokens == 1 + assert terminal_usage.completion_tokens == 2 + assert terminal_usage.total_tokens == 4 + + abort_response = await context_service.Abort( + lifecycle_pb2.AbortRequest(kv_session=context_responses[0].prefill_ready.kv_session), + _Context(), + ) + assert abort_response.status == lifecycle_pb2.ABORT_STATUS_ABORTED + assert not context_service._kv_session_requests + + +@pytest.mark.asyncio +async def test_prefill_session_expires_and_close_cancels_timers() -> None: + handoff = DisaggregatedParams( + request_type="context_only", + disagg_request_id=101, + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + ) + llm = _Llm(handoff) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_PREFILL, + RequestTracker(llm), + kv_session_ttl_seconds=0.01, + ) + request = generation_pb2.GenerateRequest(request_id="prefill", model="model", prompt="hello") + + _ = [response async for response in service.Generate(request, _Context())] + assert service._kv_session_requests + assert service._kv_session_timers + await asyncio.sleep(0.02) + assert not service._kv_session_requests + assert not service._kv_session_timers + + service._track_kv_session("second", "request") + timer = service._kv_session_timers["second"] + service.close() + assert timer.cancelled() + assert not service._kv_session_requests + + +def test_decode_media_is_required_only_for_marked_handoff(monkeypatch) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image", "video") + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return ("image", "video") + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + del modalities + return None + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) + llm = _Llm() + llm.input_processor = _Processor() + service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) + session = encode_handoff( + DisaggregatedParams( + request_type="context_only", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + mrope_position_ids_handle={"tensor": "ids"}, + mrope_position_deltas_handle={"tensor": "deltas"}, + ), + requires_decode_media=True, + ) + request = generation_pb2.GenerateRequest( + request_id="decode", model="model", token_ids=input_pb2.TokenIds(ids=[1, 2]) + ) + request.kv.session.CopyFrom(session) + + with pytest.raises(ValueError, match="must resend"): + service._validate_generate(request) + + request.media.add(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"image") + service._validate_generate(request) + + +def test_kv_batch_preserves_chain_geometry_and_source_sequence() -> None: + llm = _Llm() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + service._lora_names_by_id[7] = "adapter" + batch = service._kv_batch( + { + "attention_dp_rank": 2, + "layer_group_id": 3, + "window_size": 4096, + "hash_algo": "v2_sha256", + "data": { + "type": "stored", + "parent_hash": "00" * 32, + "blocks": [ + { + "block_hash": "11" * 32, + "tokens": [{"token_id": 1}, {"token_id": 2}], + "mm_keys": [ + { + "type": "mm_key", + "hash": "0123456789abcdef0011", + "start_offset": 4, + } + ], + "lora_id": 7, + }, + { + "block_hash": "22" * 32, + "tokens": [{"token_id": 3}, {"token_id": 4}], + "lora_id": 7, + }, + ], + }, + }, + sequence=9, + ) + + assert batch.sequence_number == 9 + assert batch.data_parallel_rank == 2 + assert len(batch.events) == 1 + stored = batch.events[0].block_stored + assert len(stored.block_hashes) == 2 + assert stored.parent_block_hash.value == bytes.fromhex("00" * 32) + assert list(stored.token_ids) == [1, 2, 3, 4] + assert stored.block_size == 2 + assert stored.lora_id == 7 + assert stored.lora_name == "adapter" + assert stored.group_idx == 3 + assert list(stored.extra_keys[0].values) == [ + "trt_mm_v1", + "0", + "81985529216486895", + "4", + ] + + +def test_kv_batch_stops_at_partial_block_and_uses_decimal_int_hashes() -> None: + llm = _Llm() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + batch = service._kv_batch( + { + "attention_dp_rank": 0, + "data": { + "type": "stored", + "parent_hash": 2**64 - 1, + "blocks": [ + { + "block_hash": 2**63, + "tokens": [{"token_id": 1}, {"token_id": 2}], + }, + {"block_hash": 9, "tokens": [{"token_id": 3}]}, + { + "block_hash": 10, + "tokens": [{"token_id": 4}, {"token_id": 5}], + }, + ], + }, + }, + sequence=3, + ) + + assert len(batch.events) == 1 + stored = batch.events[0].block_stored + assert stored.parent_block_hash.encoding == "decimal_int64" + assert stored.parent_block_hash.value == b"-1" + assert stored.block_hashes[0].value == str(-(2**63)).encode() + assert 9 in service._partial_block_hashes[0] + + +def test_kv_batch_preserves_model_owned_lora_id_zero_and_gap_reset() -> None: + class _Processor: + @staticmethod + def get_model_owned_lora_identities() -> dict[str, int]: + return {"vision-lora": 0} + + llm = _Llm() + llm.input_processor = _Processor() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + batch = service._kv_batch( + { + "data": { + "type": "stored", + "blocks": [ + { + "block_hash": 2, + "tokens": [{"token_id": 1}, {"token_id": 2}], + "lora_id": 0, + } + ], + } + }, + sequence=1, + ) + service._partial_block_hashes[0] = {3} + reset = service._kv_batch({"attention_dp_rank": 0, "data": {"type": "all_cleared"}}, sequence=2) + + assert batch.events[0].block_stored.lora_id == 0 + assert batch.events[0].block_stored.lora_name == "vision-lora" + assert reset.events[0].WhichOneof("event") == "all_blocks_cleared" + assert 0 not in service._partial_block_hashes + + +@pytest.mark.parametrize( + "unsupported", + [ + {"cache_salt": "private"}, + {"tokens": [{"token_id": 1, "token_extra_id": 1}, {"token_id": 2}]}, + {"mm_keys": [{"type": "mm_key", "hash": "not-hex"}]}, + ], +) +def test_kv_batch_fails_closed_for_unrepresentable_cache_namespace(unsupported) -> None: + llm = _Llm() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + block = { + "block_hash": 2, + "tokens": [{"token_id": 1}, {"token_id": 2}], + **unsupported, + } + + batch = service._kv_batch( + {"data": {"type": "stored", "parent_hash": 1, "blocks": [block]}}, + sequence=1, + ) + + assert len(batch.events) == 1 + assert batch.events[0].WhichOneof("event") == "all_blocks_cleared" + + +@pytest.mark.asyncio +async def test_discovery_and_load_use_config_and_shared_stats() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + guided_decoding_backend=None, + lora_config=None, + enable_lora=False, + kv_cache_config=SimpleNamespace(tokens_per_block=64), + ) + llm.get_kv_cache_capacity = lambda: { + "maxNumBlocks": 100, + "tokensPerBlock": 64, + "maxNumTokens": 6400, + } + stats = StatsFanout(llm) + stats._publish( + { + "attentionDpRank": 1, + "numActiveRequests": 3, + "numQueuedRequests": 2, + "kvCacheStats": { + "maxNumBlocks": 100, + "freeNumBlocks": 25, + "tokensPerBlock": 64, + }, + "inflightBatchingStats": {"numContextRequests": 1, "numGenRequests": 2}, + } + ) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + stats_fanout=stats, + ) + + model_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert model_info.kv_block_size == 64 + assert model_info.total_kv_blocks == 100 + assert not model_info.generation.guided_decoding.supported + assert not model_info.supports_lora + + load = await service.GetLoad( + observability_pb2.GetLoadRequest(include_per_rank=True), _Context() + ) + assert load.running_requests == 3 + assert load.queued_requests == 2 + assert load.used_kv_blocks == 75 + assert load.total_kv_blocks == 100 + assert load.prefill_batch_size == 1 + assert load.decode_batch_size == 2 + assert load.ranks[0].data_parallel_rank == 1 + assert load.attributes["source"] == "shared_stats_fanout" + assert load.attributes["kv_tokens_per_block"] == "64" + assert load.attributes["rank.1.kv_tokens_per_block"] == "64" + + llm.args.guided_decoding_backend = "xgrammar" + guided_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert guided_info.generation.guided_decoding.supported + assert ( + model_pb2.GUIDED_DECODING_MODE_STRUCTURAL_TAG + in guided_info.generation.guided_decoding.modes + ) + + +@pytest.mark.asyncio +async def test_load_never_undercounts_live_tracker_from_stale_stats() -> None: + llm = _Llm() + stats = StatsFanout(llm) + stats._publish({"attentionDpRank": 0, "numActiveRequests": 0}) + tracker = RequestTracker(llm) + tracker.admit("active", _Result()) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + tracker, + stats_fanout=stats, + ) + + load = await service.GetLoad(observability_pb2.GetLoadRequest(), _Context()) + + assert load.running_requests == 1 + + +@pytest.mark.asyncio +async def test_kv_event_sources_are_rank_scoped() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + data_parallel_size=2, + kv_cache_config=SimpleNamespace(event_buffer_max_size=8), + ) + fanout = KvEventFanout(llm, buffer_size=8) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=fanout, + ) + + response = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + + assert [source.data_parallel_rank for source in response.sources] == [0, 1] + + +@pytest.mark.asyncio +async def test_kv_event_discovery_requires_decimal_compatible_hashes() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + kv_cache_config=SimpleNamespace( + event_buffer_max_size=8, + kv_cache_event_hash_algo="v2_sha256", + use_kv_cache_manager_v2=True, + ) + ) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=KvEventFanout(llm), + ) + + response = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + assert list(response.sources) == [] + + llm.args.kv_cache_config.kv_cache_event_hash_algo = "v2_sha256_64" + response = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + assert len(response.sources) == 1 + + +@pytest.mark.asyncio +async def test_direct_kv_subscription_rejects_unknown_rank() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + data_parallel_size=2, + kv_cache_config=SimpleNamespace(event_buffer_max_size=8), + ) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=KvEventFanout(llm), + ) + request = kv_pb2.SubscribeKvEventsRequest(data_parallel_ranks=[2]) + + with pytest.raises(AssertionError, match="INVALID_ARGUMENT"): + await _collect_async(service.SubscribeKvEvents(request, _Context())) + + +@pytest.mark.asyncio +async def test_drain_deadline_does_not_wait_forever_for_external_http() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + tracker.begin_external() + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + tracker, + post_abort_cleanup_timeout_seconds=0.01, + ) + request = lifecycle_pb2.DrainRequest( + stop_accepting_new_requests=True, + deadline_ms=1, + abort_after_deadline=True, + ) + + responses = await asyncio.wait_for( + _collect_async(service.Drain(request, _Context())), timeout=0.1 + ) + + assert responses[-1].WhichOneof("event") == "error" + await tracker.finish_external() + + +@pytest.mark.asyncio +async def test_model_info_omits_unknown_kv_geometry() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert not info.HasField("kv_block_size") + assert not info.HasField("total_kv_blocks") + + +@pytest.mark.asyncio +async def test_model_info_unknown_model_maps_to_not_found() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + + with pytest.raises(AssertionError, match="NOT_FOUND"): + await service.GetModelInfo(model_pb2.GetModelInfoRequest(model="other"), _Context()) + + +@pytest.mark.asyncio +async def test_health_probe_and_unconfigured_lora_fail_explicitly(tmp_path) -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + with pytest.raises(AssertionError, match="UNIMPLEMENTED"): + await service.Health(lifecycle_pb2.HealthRequest(include_inference_probe=True), _Context()) + + adapter_dir = tmp_path / "adapter" + adapter_dir.mkdir() + with pytest.raises(AssertionError, match="FAILED_PRECONDITION"): + await service.LoadLora( + lora_pb2.LoadLoraRequest( + adapter=lora_pb2.LoraAdapter( + lora_id=1, lora_name="adapter", source_path=str(adapter_dir) + ) + ), + _Context(), + ) From 461b120a4b2144ad684113b9409c2f3d60e0f3a6 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 22:58:43 -0700 Subject: [PATCH 02/20] build(ucx): avoid conflicting target link features Signed-off-by: Connor Carpenter --- .../executor/cache_transmission/ucx_utils/CMakeLists.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt index a59d85d77923..dae43e989694 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt @@ -31,7 +31,9 @@ if(ENABLE_UCX) target_link_libraries(${UCX_WRAPPER_TARGET} PRIVATE $) - target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ucxx::ucxx ucx::ucs) + # ucxx is already linked with WHOLE_ARCHIVE above. Repeating the same target + # without a feature is rejected by newer CMake versions. + target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ucx::ucs) target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ${CUDA_RT_LIB}) target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ${TORCH_LIBRARIES}) target_link_libraries(${UCX_WRAPPER_TARGET} PRIVATE ${ZMQ_LIBRARIES}) From 90386d69b6d6c65e9ee3dd5a0d6963fed962e652 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 22:58:54 -0700 Subject: [PATCH 03/20] fix(disaggregation): prepare lazy LoRA on decode Signed-off-by: Connor Carpenter --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 3 ++- .../_torch/executor/test_benchmark_disagg.py | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f69041fa7ab8..33a8675b9efb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5405,7 +5405,8 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests): for resource_mgr_type in ( ResourceManagerType.KV_CACHE_MANAGER, ResourceManagerType.SPEC_RESOURCE_MANAGER, - ResourceManagerType.DRAFT_KV_CACHE_MANAGER): + ResourceManagerType.DRAFT_KV_CACHE_MANAGER, + ResourceManagerType.PEFT_CACHE_MANAGER): if (resource_mgr_type in self.resource_manager.resource_managers and self.resource_manager. resource_managers[resource_mgr_type] is not None): diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index dabb2b910d55..b1d2c829f86d 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -30,6 +30,7 @@ import pytest from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests # --------------------------------------------------------------------------- @@ -649,6 +650,32 @@ def test_fetch_called_once_even_in_benchmark_disagg(self): mock_fetch.assert_called_once() +class TestPrepareDisaggGenInitResources: + def test_prepares_peft_cache_for_lazy_lora(self): + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + ex = object.__new__(PyExecutor) + managers = { + kind: Mock() + for kind in ( + ResourceManagerType.KV_CACHE_MANAGER, + ResourceManagerType.SPEC_RESOURCE_MANAGER, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER, + ResourceManagerType.PEFT_CACHE_MANAGER, + ) + } + ex.resource_manager = Mock(resource_managers=managers) + ex._recv_disagg_gen_cache = Mock() + request = Mock() + + ex._prepare_disagg_gen_init([request]) + + for manager in managers.values(): + prepared = manager.prepare_resources.call_args.args[0] + assert prepared.context_requests == [request] + ex._recv_disagg_gen_cache.assert_called_once_with([request]) + + # --------------------------------------------------------------------------- # Benchmark fill admission flow control # --------------------------------------------------------------------------- From cee72f7da8e6f112237e2d8a59ccc37f42e0b63b Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 22:59:55 -0700 Subject: [PATCH 04/20] fix(openengine): complete context handoff accounting Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/servicer.py | 37 ++++++++++++++++++++-- tests/unittest/openengine/test_servicer.py | 34 +++++++++++++++++++- 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index 0835a0d0ac97..effd562b4495 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -257,7 +257,13 @@ async def Generate( raise RuntimeError( "Context-only result did not return disaggregated parameters" ) - handoff = replace(handoff, ctx_usage=self._usage_payload(current)) + handoff = replace( + handoff, + ctx_usage=self._usage_payload(current), + ctx_info_endpoint=( + handoff.ctx_info_endpoint or self._context_info_endpoint() + ), + ) session = encode_handoff(handoff, requires_decode_media=bool(request.media)) self._track_kv_session(session.session_id, request.request_id) yield generation_pb2.GenerateResponse( @@ -334,12 +340,39 @@ def _release_kv_session(self, session_id: str) -> bool: timer.cancel() return released + def _context_info_endpoint(self) -> str | None: + """Return the context worker endpoint needed for direct NIXL discovery.""" + params = getattr(self.llm, "disaggregated_params", None) + if not isinstance(params, dict): + return None + endpoint = params.get("ctx_info_endpoint") + if isinstance(endpoint, str): + return endpoint or None + if isinstance(endpoint, (list, tuple)): + return next( + (item for item in endpoint if isinstance(item, str) and item), + None, + ) + return None + def _release_all_kv_sessions(self) -> None: for timer in self._kv_session_timers.values(): timer.cancel() self._kv_session_timers.clear() self._kv_session_requests.clear() + def _active_kv_session_count(self) -> int: + """Count sessions whose owning request is still engine-active. + + Prefill handoff handles remain abortable until their TTL expires, but + their context request has already finished and must not inflate load. + Drain continues to report and release all open handles separately. + """ + return sum( + request_id in self.tracker.active_requests + for request_id in self._kv_session_requests.values() + ) + def close(self) -> None: """Release protocol-owned timers without shutting down the shared LLM.""" self._release_all_kv_sessions() @@ -771,7 +804,7 @@ async def GetLoad( timestamp_unix_nanos=time.time_ns(), running_requests=running_requests, queued_requests=queued_requests, - active_kv_sessions=len(self._kv_session_requests), + active_kv_sessions=self._active_kv_session_count(), ranks=rank_infos, attributes=attributes, ) diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index 2c7dac52bbe0..6c8e34bd3502 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -18,7 +18,7 @@ from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from tensorrt_llm.inputs.registry import MultimodalLoraSpec -from tensorrt_llm.openengine.converters import encode_handoff +from tensorrt_llm.openengine.converters import decode_handoff, encode_handoff from tensorrt_llm.openengine.servicer import OpenEngineServicer from tensorrt_llm.serve.kv_event_fanout import KvEventFanout from tensorrt_llm.serve.request_tracker import RequestTracker @@ -284,6 +284,10 @@ async def test_context_then_generation_round_trips_handoff() -> None: assert len(context_responses) == 1 assert context_responses[0].WhichOneof("event") == "prefill_ready" assert len(context_service._kv_session_requests) == 1 + context_load = await context_service.GetLoad( + observability_pb2.GetLoadRequest(), _Context() + ) + assert context_load.active_kv_sessions == 0 decode_llm = _Llm() decode_service = OpenEngineServicer( @@ -325,6 +329,34 @@ async def test_context_then_generation_round_trips_handoff() -> None: assert not context_service._kv_session_requests +@pytest.mark.asyncio +async def test_prefill_adds_context_endpoint_from_llm_discovery() -> None: + handoff = DisaggregatedParams( + request_type="context_only", + first_gen_tokens=[7], + disagg_request_id=101, + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + ) + llm = _Llm(handoff) + llm.disaggregated_params = {"ctx_info_endpoint": ["tcp://context:1234"]} + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_PREFILL, + RequestTracker(llm), + ) + request = generation_pb2.GenerateRequest( + request_id="prefill", + model="model", + token_ids=input_pb2.TokenIds(ids=[1, 2]), + ) + + responses = [response async for response in service.Generate(request, _Context())] + decoded = decode_handoff(responses[0].prefill_ready.kv_session) + + assert decoded.ctx_info_endpoint == "tcp://context:1234" + + @pytest.mark.asyncio async def test_prefill_session_expires_and_close_cancels_timers() -> None: handoff = DisaggregatedParams( From 81a106e5647e5e7c57fd140158df07918bb1f392 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 21 Jul 2026 15:46:57 -0700 Subject: [PATCH 05/20] fix(openengine): update server protocol contract Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 2 +- tensorrt_llm/openengine/_schema_pin.py | 2 +- tensorrt_llm/openengine/converters.py | 5 +- tensorrt_llm/openengine/server.py | 11 +- tensorrt_llm/openengine/servicer.py | 155 +++++++++++++------ tests/unittest/openengine/test_converters.py | 11 +- tests/unittest/openengine/test_servicer.py | 124 +++++++++------ 7 files changed, 194 insertions(+), 116 deletions(-) diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT index 2fb5de7752ae..3e11d4b670be 100644 --- a/OPENENGINE_COMMIT +++ b/OPENENGINE_COMMIT @@ -1 +1 @@ -cea19cb06acf03c911b84d5c147e519b60dd92a6 +7d6ac3848bba9c9a2b1aabcd4fab5f1a4ea04558 diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py index 8a5cfa30ece3..39a1375a0c77 100644 --- a/tensorrt_llm/openengine/_schema_pin.py +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -3,6 +3,6 @@ """OpenEngine source identity embedded in TensorRT-LLM distributions.""" -OPENENGINE_COMMIT = "cea19cb06acf03c911b84d5c147e519b60dd92a6" +OPENENGINE_COMMIT = "7d6ac3848bba9c9a2b1aabcd4fab5f1a4ea04558" __all__ = ["OPENENGINE_COMMIT"] diff --git a/tensorrt_llm/openengine/converters.py b/tensorrt_llm/openengine/converters.py index e934a5059138..98a36d8fbd94 100644 --- a/tensorrt_llm/openengine/converters.py +++ b/tensorrt_llm/openengine/converters.py @@ -126,11 +126,10 @@ def to_sampling_params(request: generation_pb2.GenerateRequest) -> SamplingParam return SamplingParams(**kwargs) -def to_priority(request: generation_pb2.GenerateRequest) -> float: +def to_priority(priority: int | None) -> float: """Map signed OpenEngine ordering into TRT-LLM's bounded priority domain.""" - if not request.HasField("priority"): + if priority is None: return 0.5 - priority = request.priority return 0.5 + 0.5 * priority / (1 + abs(priority)) diff --git a/tensorrt_llm/openengine/server.py b/tensorrt_llm/openengine/server.py index 4604047ad1e2..54987ef791f0 100644 --- a/tensorrt_llm/openengine/server.py +++ b/tensorrt_llm/openengine/server.py @@ -4,7 +4,7 @@ """Lifecycle wrapper for the optional OpenEngine sibling gRPC server.""" import grpc -from openengine.v1 import engine_pb2, openengine_pb2_grpc +from openengine.v1 import openengine_pb2_grpc, server_pb2 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.logger import logger @@ -71,7 +71,8 @@ def __init__( event_host=host, event_port=port, ) - openengine_pb2_grpc.add_OpenEngineServicer_to_server(self.servicer, self._server) + openengine_pb2_grpc.add_InferenceServicer_to_server(self.servicer, self._server) + openengine_pb2_grpc.add_ControlServicer_to_server(self.servicer, self._server) bound = self._server.add_insecure_port(f"{host}:{port}") if bound == 0: raise RuntimeError(f"Failed to bind OpenEngine server to {host}:{port}") @@ -115,12 +116,12 @@ async def stop(self, grace: float = 5.0) -> None: def openengine_role(server_role: object | None) -> int: """Map TRT-LLM serve roles onto the supported OpenEngine milestone.""" if server_role is None: - return engine_pb2.ENGINE_ROLE_AGGREGATED + return server_pb2.ENGINE_ROLE_AGGREGATED name = getattr(server_role, "name", str(server_role)).upper() if name == "CONTEXT": - return engine_pb2.ENGINE_ROLE_PREFILL + return server_pb2.ENGINE_ROLE_PREFILL if name == "GENERATION": - return engine_pb2.ENGINE_ROLE_DECODE + return server_pb2.ENGINE_ROLE_DECODE if name in ("MM_ENCODER", "VISUAL_GEN"): raise ValueError(f"OpenEngine does not support TRT-LLM server role {name!r}") raise ValueError(f"Unknown TRT-LLM server role {name!r}") diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index effd562b4495..4a99df07ed25 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -15,7 +15,6 @@ import grpc from openengine import MINIMUM_CLIENT_REVISION, SCHEMA_RELEASE, SCHEMA_REVISION from openengine.v1 import ( - engine_pb2, error_pb2, generation_pb2, input_pb2, @@ -25,6 +24,7 @@ model_pb2, observability_pb2, openengine_pb2_grpc, + server_pb2, ) import tensorrt_llm @@ -102,7 +102,10 @@ def _block_hash(value: object, hash_algo: str | None) -> kv_pb2.KvBlockHash: return kv_pb2.KvBlockHash(value=encoded, encoding=hash_algo or "hex") -class OpenEngineServicer(openengine_pb2_grpc.OpenEngineServicer): +class OpenEngineServicer( + openengine_pb2_grpc.InferenceServicer, + openengine_pb2_grpc.ControlServicer, +): """Engine-neutral service surface over TensorRT-LLM's Python LLM API.""" def __init__( @@ -151,13 +154,66 @@ def __init__( self.event_host = event_host self.event_port = event_port + @staticmethod + def _request_metadata(context: grpc.aio.ServicerContext) -> dict[str, str]: + metadata: dict[str, str] = {} + for item in context.invocation_metadata(): + try: + key, value = item + except (TypeError, ValueError): + key, value = item.key, item.value + normalized_key = str(key).lower() + if normalized_key.startswith("openengine-") and normalized_key in metadata: + raise ValueError(f"Duplicate reserved gRPC metadata key {normalized_key!r}") + if isinstance(value, bytes): + try: + normalized_value = value.decode("ascii") + except UnicodeDecodeError as error: + raise ValueError( + f"gRPC metadata value for {normalized_key!r} must be ASCII" + ) from error + else: + normalized_value = str(value) + metadata[normalized_key] = normalized_value + return metadata + + @staticmethod + def _metadata_int(metadata: dict[str, str], key: str, minimum: int, maximum: int) -> int | None: + if key not in metadata: + return None + value = metadata[key] + digits = value[1:] if value.startswith("-") and minimum < 0 else value + if not digits or not digits.isdecimal(): + raise ValueError(f"gRPC metadata {key!r} must be a base-10 integer") + try: + parsed = int(value, 10) + except ValueError as error: + raise ValueError(f"gRPC metadata {key!r} must be a base-10 integer") from error + if not minimum <= parsed <= maximum: + raise ValueError(f"gRPC metadata {key!r} must be in the range [{minimum}, {maximum}]") + return parsed + async def Generate( self, request: generation_pb2.GenerateRequest, context: grpc.aio.ServicerContext ) -> AsyncGenerator[generation_pb2.GenerateResponse, None]: result = None consumed_session_id = None try: - self._validate_generate(request) + metadata = self._request_metadata(context) + priority_value = self._metadata_int( + metadata, "openengine-priority", -(1 << 31), (1 << 31) - 1 + ) + target_dp_rank = self._metadata_int( + metadata, "openengine-target-dp-rank", 0, (1 << 32) - 1 + ) + if self.role == server_pb2.ENGINE_ROLE_DECODE and request.kv.HasField("session"): + session_dp_rank = request.kv.session.dp_rank + if target_dp_rank is not None and target_dp_rank != session_dp_rank: + raise ValueError( + "openengine-target-dp-rank does not match the KV session dp_rank" + ) + target_dp_rank = session_dp_rank + self._validate_generate(request, target_dp_rank) params = to_sampling_params(request) media = await load_media(list(request.media), request.media_options, self.media_config) input_kind = request.WhichOneof("input") @@ -177,17 +233,17 @@ async def Generate( ) if lora_request is None and request.lora_name: lora_request = await self.loras.request(request.lora_name) - disaggregated = self._disaggregated_params(request) + disaggregated = self._disaggregated_params(request, target_dp_rank, metadata) context_usage = ( disaggregated.ctx_usage if disaggregated is not None and disaggregated.request_type == "generation_only" else None ) - scheduling = self._scheduling_params(request) - priority = to_priority(request) + scheduling = self._scheduling_params(target_dp_rank) + priority = to_priority(priority_value) trace_headers = { key.lower(): value - for key, value in request.metadata.items() + for key, value in metadata.items() if key.lower() in ("traceparent", "tracestate", "baggage") } result = self.llm.generate_async( @@ -248,7 +304,7 @@ async def Generate( yield generation_pb2.GenerateResponse( request_id=request.request_id, prompt=prompt ) - if self.role == engine_pb2.ENGINE_ROLE_PREFILL: + if self.role == server_pb2.ENGINE_ROLE_PREFILL: if current.finished: handoff = current.disaggregated_params if handoff is None and current.outputs: @@ -377,7 +433,9 @@ def close(self) -> None: """Release protocol-owned timers without shutting down the shared LLM.""" self._release_all_kv_sessions() - def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: + def _validate_generate( + self, request: generation_pb2.GenerateRequest, target_dp_rank: int | None = None + ) -> None: if not request.request_id: raise ValueError("request_id must not be empty") if request.model and request.model != self.model: @@ -395,8 +453,8 @@ def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: raise ValueError("Per-request guided decoding backend selection is not supported") if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: raise ValueError("Prefix-cache bypass is not supported by TensorRT-LLM") - if request.kv.HasField("data_parallel_rank"): - self._scheduling_params(request) + if target_dp_rank is not None: + self._scheduling_params(target_dp_rank) if request.media: processor = getattr(self.llm, "input_processor", None) aggregate = self._available_modalities(processor) @@ -408,7 +466,7 @@ def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: input_pb2.MODALITY_AUDIO: "audio", } allowed = ( - aggregate if self.role == engine_pb2.ENGINE_ROLE_AGGREGATED else prefill_decode + aggregate if self.role == server_pb2.ENGINE_ROLE_AGGREGATED else prefill_decode ) unsupported = { names.get(item.modality, "unspecified") @@ -424,7 +482,7 @@ def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: dict.fromkeys(names[item.modality] for item in request.media) ) processor.get_required_lora_spec(requested_modalities) - if self.role == engine_pb2.ENGINE_ROLE_DECODE: + if self.role == server_pb2.ENGINE_ROLE_DECODE: if not request.kv.HasField("session"): raise ValueError("Decode requests require a prefill KV session") requires_media = handoff_requires_decode_media(request.kv.session) @@ -432,37 +490,34 @@ def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: raise ValueError("Decode request must resend the ordered context-phase media") if request.media and not requires_media: raise ValueError("Decode KV session does not require raw media") - elif self.role == engine_pb2.ENGINE_ROLE_PREFILL: + elif self.role == server_pb2.ENGINE_ROLE_PREFILL: if request.kv.HasField("session"): raise ValueError("Prefill requests cannot consume a KV session") elif request.kv.HasField("session"): raise ValueError("Aggregated requests cannot consume a KV session") def _disaggregated_params( - self, request: generation_pb2.GenerateRequest + self, + request: generation_pb2.GenerateRequest, + target_dp_rank: int | None = None, + metadata: dict[str, str] | None = None, ) -> DisaggregatedParams | None: - if self.role == engine_pb2.ENGINE_ROLE_AGGREGATED: + if self.role == server_pb2.ENGINE_ROLE_AGGREGATED: return None - if self.role == engine_pb2.ENGINE_ROLE_PREFILL: + if self.role == server_pb2.ENGINE_ROLE_PREFILL: return DisaggregatedParams( request_type="context_only", disagg_request_id=stable_request_id(request.request_id), - ctx_dp_rank=( - request.kv.data_parallel_rank - if request.kv.HasField("data_parallel_rank") - else None - ), + ctx_dp_rank=target_dp_rank, schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, - conversation_id=request.metadata.get("conversation_id"), + conversation_id=(metadata or {}).get("conversation_id"), ) return decode_handoff(request.kv.session) - def _scheduling_params( - self, request: generation_pb2.GenerateRequest - ) -> SchedulingParams | None: - if not request.kv.HasField("data_parallel_rank"): + def _scheduling_params(self, target_dp_rank: int | None) -> SchedulingParams | None: + if target_dp_rank is None: return None - rank = request.kv.data_parallel_rank + rank = target_dp_rank data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) if rank >= data_parallel_size: raise ValueError( @@ -609,26 +664,43 @@ def _finished(output: object) -> generation_pb2.GenerationFinished: finished.stop_match.stop_text = output.stop_reason return finished - async def GetEngineInfo( - self, request: engine_pb2.GetEngineInfoRequest, context: grpc.aio.ServicerContext - ) -> engine_pb2.EngineInfo: + async def GetServerInfo( + self, request: server_pb2.GetServerInfoRequest, context: grpc.aio.ServicerContext + ) -> server_pb2.ServerInfo: del request, context tp = _arg(self.llm, "tensor_parallel_size", 1) pp = _arg(self.llm, "pipeline_parallel_size", 1) dp = _arg(self.llm, "data_parallel_size", 1) - return engine_pb2.EngineInfo( + args = getattr(self.llm, "args", None) + kv_capacity = self._kv_capacity() + capacity = server_pb2.DeploymentCapacity() + if kv_capacity.get("tokensPerBlock") is not None: + capacity.kv_block_size = kv_capacity["tokensPerBlock"] + if kv_capacity.get("maxNumBlocks") is not None: + capacity.total_kv_blocks = kv_capacity["maxNumBlocks"] + max_requests = getattr(args, "max_batch_size", None) + if max_requests is not None: + capacity.max_running_requests = max_requests + max_tokens = getattr(args, "max_num_tokens", None) + if max_tokens is not None: + capacity.max_batched_tokens = max_tokens + max_loras = getattr(getattr(args, "lora_config", None), "max_loras", None) + if max_loras is not None: + capacity.max_loras = max_loras + return server_pb2.ServerInfo( engine_name="tensorrt_llm", engine_version=getattr(tensorrt_llm, "__version__", "unknown"), - role=self.role, + engine_role=self.role, instance_id=self.instance_id, supported_models=[self.model], - parallelism=engine_pb2.ParallelismInfo( + parallelism=server_pb2.ParallelismInfo( tensor_parallel_size=tp, pipeline_parallel_size=pp, data_parallel_size=dp, data_parallel_rank=_arg(self.llm, "data_parallel_rank", 0), ), kv_connector=self._kv_connector_info(), + capacity=capacity, schema_revision=SCHEMA_REVISION, minimum_client_revision=MINIMUM_CLIENT_REVISION, schema_release=schema_release(), @@ -662,9 +734,6 @@ async def GetModelInfo( if guided_backend == "xgrammar": guided_modes.append(model_pb2.GUIDED_DECODING_MODE_STRUCTURAL_TAG) max_context = getattr(args, "max_seq_len", None) or getattr(args, "max_input_len", None) - max_requests = getattr(args, "max_batch_size", None) - max_tokens = getattr(args, "max_num_tokens", None) - kv_capacity = self._kv_capacity() info = model_pb2.ModelInfo( model_id=self.model, served_model_name=self.model, @@ -710,14 +779,6 @@ async def GetModelInfo( ) if max_context is not None: info.max_context_length = max_context - if max_requests is not None: - info.max_running_requests = max_requests - if max_tokens is not None: - info.max_batched_tokens = max_tokens - if kv_capacity.get("tokensPerBlock") is not None: - info.kv_block_size = kv_capacity["tokensPerBlock"] - if kv_capacity.get("maxNumBlocks") is not None: - info.total_kv_blocks = kv_capacity["maxNumBlocks"] return info def _kv_capacity(self) -> dict[str, int]: @@ -824,7 +885,7 @@ async def Health( if request.model and request.model != self.model: await context.abort(grpc.StatusCode.NOT_FOUND, f"Unknown model {request.model!r}") raise ValueError(f"Unknown model {request.model!r}") - if request.role not in (engine_pb2.ENGINE_ROLE_UNSPECIFIED, self.role): + if request.role not in (server_pb2.ENGINE_ROLE_UNSPECIFIED, self.role): await context.abort( grpc.StatusCode.INVALID_ARGUMENT, "Requested health role does not match this engine", @@ -976,7 +1037,7 @@ def _supports_lora(self) -> bool: return bool(getattr(plugin_config, "lora_plugin", False)) def _kv_connector_info(self) -> kv_pb2.KvConnectorInfo: - enabled = self.role in (engine_pb2.ENGINE_ROLE_PREFILL, engine_pb2.ENGINE_ROLE_DECODE) + enabled = self.role in (server_pb2.ENGINE_ROLE_PREFILL, server_pb2.ENGINE_ROLE_DECODE) return kv_pb2.KvConnectorInfo( enabled=enabled, transfer_backend="tensorrt_llm" if enabled else "", diff --git a/tests/unittest/openengine/test_converters.py b/tests/unittest/openengine/test_converters.py index b1dc4236eb4f..5a5c9853ab64 100644 --- a/tests/unittest/openengine/test_converters.py +++ b/tests/unittest/openengine/test_converters.py @@ -180,14 +180,9 @@ def test_choice_guidance_escapes_literal_regex_characters() -> None: def test_priority_mapping_is_centered_bounded_and_strictly_monotonic() -> None: - absent = generation_pb2.GenerateRequest() - negative = generation_pb2.GenerateRequest(priority=-1000) - zero = generation_pb2.GenerateRequest(priority=0) - positive = generation_pb2.GenerateRequest(priority=1000) - - assert 0 < to_priority(negative) < to_priority(zero) - assert to_priority(zero) == to_priority(absent) == 0.5 - assert to_priority(zero) < to_priority(positive) < 1 + assert 0 < to_priority(-1000) < to_priority(0) + assert to_priority(0) == to_priority(None) == 0.5 + assert to_priority(0) < to_priority(1000) < 1 @pytest.mark.asyncio diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index 6c8e34bd3502..ff60a71aeb89 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -6,7 +6,6 @@ import pytest from openengine.v1 import ( - engine_pb2, generation_pb2, input_pb2, kv_pb2, @@ -14,6 +13,7 @@ lora_pb2, model_pb2, observability_pb2, + server_pb2, ) from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle @@ -79,12 +79,16 @@ def generate_async(self, **kwargs: object) -> _Result: class _Context: - def __init__(self, cancelled: bool = False) -> None: + def __init__(self, cancelled: bool = False, metadata: tuple[tuple[str, str], ...] = ()) -> None: self._cancelled = cancelled + self._metadata = metadata def cancelled(self) -> bool: return self._cancelled + def invocation_metadata(self) -> tuple[tuple[str, str], ...]: + return self._metadata + async def abort(self, code: object, message: str) -> None: raise AssertionError(f"unexpected abort {code}: {message}") @@ -97,20 +101,26 @@ async def _collect_async(iterator): async def test_aggregate_streams_deltas_finish_and_terminal_usage() -> None: llm = _Llm() tracker = RequestTracker(llm) - servicer = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, tracker) + servicer = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, tracker) request = generation_pb2.GenerateRequest( request_id="request", model="model", prompt="hello", stopping=generation_pb2.StoppingOptions(max_tokens=2), - metadata={ - "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", - "tracestate": "vendor=value", - "ignored": "not-a-trace-header", - }, ) - responses = [response async for response in servicer.Generate(request, _Context())] + context = _Context( + metadata=( + ( + "traceparent", + "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + ), + ("tracestate", "vendor=value"), + ("openengine-priority", "7"), + ("ignored", "not-a-trace-header"), + ) + ) + responses = [response async for response in servicer.Generate(request, context)] assert [response.WhichOneof("event") for response in responses] == [ "token", @@ -124,6 +134,7 @@ async def test_aggregate_streams_deltas_finish_and_terminal_usage() -> None: "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", "tracestate": "vendor=value", } + assert 0.5 < llm.kwargs["priority"] < 1 assert tracker.active_count == 0 @@ -131,7 +142,7 @@ async def test_aggregate_streams_deltas_finish_and_terminal_usage() -> None: async def test_generate_cancellation_aborts_and_cleans_tracking() -> None: llm = _Llm() tracker = RequestTracker(llm) - service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, tracker) + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, tracker) request = generation_pb2.GenerateRequest(request_id="cancel", model="model", prompt="hello") responses = [response async for response in service.Generate(request, _Context(True))] @@ -144,7 +155,7 @@ async def test_generate_cancellation_aborts_and_cleans_tracking() -> None: def test_generate_rejects_unsupported_cache_bypass_and_nonzero_dp_placement() -> None: llm = _Llm() service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") request.kv.bypass_prefix_cache = True @@ -152,14 +163,12 @@ def test_generate_rejects_unsupported_cache_bypass_and_nonzero_dp_placement() -> service._validate_generate(request) request.kv.bypass_prefix_cache = False - request.kv.data_parallel_rank = 0 - service._validate_generate(request) - assert service._scheduling_params(request) is None + service._validate_generate(request, 0) + assert service._scheduling_params(0) is None llm.args = SimpleNamespace(data_parallel_size=2) - request.kv.data_parallel_rank = 1 with pytest.raises(ValueError, match="attention DP"): - service._validate_generate(request) + service._validate_generate(request, 1) @pytest.mark.asyncio @@ -168,27 +177,41 @@ async def test_generate_applies_strict_attention_dp_placement() -> None: llm.args = SimpleNamespace(enable_attention_dp=True, data_parallel_size=4) llm._on_trt_backend = False service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") - request.kv.data_parallel_rank = 2 - _ = [response async for response in service.Generate(request, _Context())] + context = _Context(metadata=(("openengine-target-dp-rank", "2"),)) + _ = [response async for response in service.Generate(request, context)] scheduling = llm.kwargs["scheduling_params"] assert scheduling.attention_dp_rank == 2 assert not scheduling.attention_dp_relax +@pytest.mark.asyncio +async def test_generate_rejects_non_decimal_priority_metadata() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") + context = _Context(metadata=(("openengine-priority", "1_000"),)) + + with pytest.raises(AssertionError, match="INVALID_ARGUMENT.*base-10 integer"): + _ = [response async for response in service.Generate(request, context)] + + assert llm.kwargs is None + + def test_prefill_preserves_routing_rank_in_context_handoff() -> None: llm = _Llm() llm.args = SimpleNamespace(enable_attention_dp=True, data_parallel_size=4) llm._on_trt_backend = False - service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_PREFILL, RequestTracker(llm)) + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_PREFILL, RequestTracker(llm)) request = generation_pb2.GenerateRequest(request_id="prefill", model="model", prompt="hello") - request.kv.data_parallel_rank = 3 - params = service._disaggregated_params(request) + params = service._disaggregated_params(request, 3) assert params.ctx_dp_rank == 3 @@ -215,7 +238,7 @@ async def _load_media(*_args, **_kwargs): llm.args = SimpleNamespace(lora_config=object()) llm.input_processor = _Processor() service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) request = generation_pb2.GenerateRequest(request_id="audio", model="model", prompt="transcribe") request.media.add(modality=input_pb2.MODALITY_AUDIO, raw_bytes=b"audio") @@ -246,7 +269,7 @@ def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraS llm = _Llm() llm.input_processor = _Processor() service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) @@ -270,7 +293,7 @@ async def test_context_then_generation_round_trips_handoff() -> None: context_service = OpenEngineServicer( context_llm, "model", - engine_pb2.ENGINE_ROLE_PREFILL, + server_pb2.ENGINE_ROLE_PREFILL, RequestTracker(context_llm), ) context_request = generation_pb2.GenerateRequest( @@ -284,16 +307,14 @@ async def test_context_then_generation_round_trips_handoff() -> None: assert len(context_responses) == 1 assert context_responses[0].WhichOneof("event") == "prefill_ready" assert len(context_service._kv_session_requests) == 1 - context_load = await context_service.GetLoad( - observability_pb2.GetLoadRequest(), _Context() - ) + context_load = await context_service.GetLoad(observability_pb2.GetLoadRequest(), _Context()) assert context_load.active_kv_sessions == 0 decode_llm = _Llm() decode_service = OpenEngineServicer( decode_llm, "model", - engine_pb2.ENGINE_ROLE_DECODE, + server_pb2.ENGINE_ROLE_DECODE, RequestTracker(decode_llm), ) decode_request = generation_pb2.GenerateRequest( @@ -342,7 +363,7 @@ async def test_prefill_adds_context_endpoint_from_llm_discovery() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_PREFILL, + server_pb2.ENGINE_ROLE_PREFILL, RequestTracker(llm), ) request = generation_pb2.GenerateRequest( @@ -368,7 +389,7 @@ async def test_prefill_session_expires_and_close_cancels_timers() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_PREFILL, + server_pb2.ENGINE_ROLE_PREFILL, RequestTracker(llm), kv_session_ttl_seconds=0.01, ) @@ -403,7 +424,7 @@ def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraS monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) llm = _Llm() llm.input_processor = _Processor() - service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) session = encode_handoff( DisaggregatedParams( request_type="context_only", @@ -429,7 +450,7 @@ def test_kv_batch_preserves_chain_geometry_and_source_sequence() -> None: llm = _Llm() llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) service._lora_names_by_id[7] = "adapter" batch = service._kv_batch( @@ -488,7 +509,7 @@ def test_kv_batch_stops_at_partial_block_and_uses_decimal_int_hashes() -> None: llm = _Llm() llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) batch = service._kv_batch( { @@ -530,7 +551,7 @@ def get_model_owned_lora_identities() -> dict[str, int]: llm.input_processor = _Processor() llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) batch = service._kv_batch( { @@ -568,7 +589,7 @@ def test_kv_batch_fails_closed_for_unrepresentable_cache_namespace(unsupported) llm = _Llm() llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) block = { "block_hash": 2, @@ -616,14 +637,15 @@ async def test_discovery_and_load_use_config_and_shared_stats() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_AGGREGATED, + server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm), stats_fanout=stats, ) + server_info = await service.GetServerInfo(server_pb2.GetServerInfoRequest(), _Context()) + assert server_info.capacity.kv_block_size == 64 + assert server_info.capacity.total_kv_blocks == 100 model_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) - assert model_info.kv_block_size == 64 - assert model_info.total_kv_blocks == 100 assert not model_info.generation.guided_decoding.supported assert not model_info.supports_lora @@ -660,7 +682,7 @@ async def test_load_never_undercounts_live_tracker_from_stale_stats() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_AGGREGATED, + server_pb2.ENGINE_ROLE_AGGREGATED, tracker, stats_fanout=stats, ) @@ -681,7 +703,7 @@ async def test_kv_event_sources_are_rank_scoped() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_AGGREGATED, + server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm), kv_event_fanout=fanout, ) @@ -704,7 +726,7 @@ async def test_kv_event_discovery_requires_decimal_compatible_hashes() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_AGGREGATED, + server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm), kv_event_fanout=KvEventFanout(llm), ) @@ -727,7 +749,7 @@ async def test_direct_kv_subscription_rejects_unknown_rank() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_AGGREGATED, + server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm), kv_event_fanout=KvEventFanout(llm), ) @@ -745,7 +767,7 @@ async def test_drain_deadline_does_not_wait_forever_for_external_http() -> None: service = OpenEngineServicer( llm, "model", - engine_pb2.ENGINE_ROLE_AGGREGATED, + server_pb2.ENGINE_ROLE_AGGREGATED, tracker, post_abort_cleanup_timeout_seconds=0.01, ) @@ -764,21 +786,21 @@ async def test_drain_deadline_does_not_wait_forever_for_external_http() -> None: @pytest.mark.asyncio -async def test_model_info_omits_unknown_kv_geometry() -> None: +async def test_server_info_omits_unknown_kv_geometry() -> None: llm = _Llm() service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) - info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) - assert not info.HasField("kv_block_size") - assert not info.HasField("total_kv_blocks") + info = await service.GetServerInfo(server_pb2.GetServerInfoRequest(), _Context()) + assert not info.capacity.HasField("kv_block_size") + assert not info.capacity.HasField("total_kv_blocks") @pytest.mark.asyncio async def test_model_info_unknown_model_maps_to_not_found() -> None: llm = _Llm() service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) with pytest.raises(AssertionError, match="NOT_FOUND"): @@ -789,7 +811,7 @@ async def test_model_info_unknown_model_maps_to_not_found() -> None: async def test_health_probe_and_unconfigured_lora_fail_explicitly(tmp_path) -> None: llm = _Llm() service = OpenEngineServicer( - llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) with pytest.raises(AssertionError, match="UNIMPLEMENTED"): await service.Health(lifecycle_pb2.HealthRequest(include_inference_probe=True), _Context()) From 74ce7a33c24038a501d85af30332edc413813655 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 21 Jul 2026 15:56:41 -0700 Subject: [PATCH 06/20] test(openengine): refresh schema release fixture Signed-off-by: Connor Carpenter --- tests/unittest/openengine/test_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index 213046595d45..bf98fc3f4a78 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -43,7 +43,7 @@ def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> Non def test_schema_release_reads_launch_environment(monkeypatch) -> None: - commit = "cea19cb06acf03c911b84d5c147e519b60dd92a6" + commit = "7d6ac3848bba9c9a2b1aabcd4fab5f1a4ea04558" monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) assert schema_release() == commit From 20270d26e5bb7ada2be722bbfa415e423f358f14 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 21 Jul 2026 16:07:37 -0700 Subject: [PATCH 07/20] build(openengine): advance pinned protocol release Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 2 +- tensorrt_llm/openengine/_schema_pin.py | 2 +- tests/unittest/openengine/test_server.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT index 3e11d4b670be..0b9175ed0b2c 100644 --- a/OPENENGINE_COMMIT +++ b/OPENENGINE_COMMIT @@ -1 +1 @@ -7d6ac3848bba9c9a2b1aabcd4fab5f1a4ea04558 +f1a7189311770f8aa1f0dd787561df809847595d diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py index 39a1375a0c77..9d40d1049069 100644 --- a/tensorrt_llm/openengine/_schema_pin.py +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -3,6 +3,6 @@ """OpenEngine source identity embedded in TensorRT-LLM distributions.""" -OPENENGINE_COMMIT = "7d6ac3848bba9c9a2b1aabcd4fab5f1a4ea04558" +OPENENGINE_COMMIT = "f1a7189311770f8aa1f0dd787561df809847595d" __all__ = ["OPENENGINE_COMMIT"] diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index bf98fc3f4a78..467dfc0fd218 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -43,7 +43,7 @@ def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> Non def test_schema_release_reads_launch_environment(monkeypatch) -> None: - commit = "7d6ac3848bba9c9a2b1aabcd4fab5f1a4ea04558" + commit = "f1a7189311770f8aa1f0dd787561df809847595d" monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) assert schema_release() == commit From c2890940df52ec46bfd3f60f8d42644e8f6194f5 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 21 Jul 2026 17:55:38 -0700 Subject: [PATCH 08/20] refactor(openengine): remove unused RPC handlers Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 2 +- tensorrt_llm/openengine/_schema_pin.py | 2 +- tensorrt_llm/openengine/servicer.py | 26 ------------------------ tests/unittest/openengine/test_server.py | 2 +- 4 files changed, 3 insertions(+), 29 deletions(-) diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT index 0b9175ed0b2c..c19080c682e4 100644 --- a/OPENENGINE_COMMIT +++ b/OPENENGINE_COMMIT @@ -1 +1 @@ -f1a7189311770f8aa1f0dd787561df809847595d +df3a9be24a2a36a4ff7a6d4fef9f1d7480ae210d diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py index 9d40d1049069..9be88c7e55dc 100644 --- a/tensorrt_llm/openengine/_schema_pin.py +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -3,6 +3,6 @@ """OpenEngine source identity embedded in TensorRT-LLM distributions.""" -OPENENGINE_COMMIT = "f1a7189311770f8aa1f0dd787561df809847595d" +OPENENGINE_COMMIT = "df3a9be24a2a36a4ff7a6d4fef9f1d7480ae210d" __all__ = ["OPENENGINE_COMMIT"] diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index 4a99df07ed25..fa298ce3a187 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -1292,29 +1292,3 @@ def _kv_fail_closed_batch(rank: int, sequence: int, message: str) -> kv_pb2.KvEv data_parallel_rank=rank, events=[kv_pb2.KvEvent(all_blocks_cleared=kv_pb2.AllBlocksCleared())], ) - - async def Embed(self, request: object, context: grpc.aio.ServicerContext) -> None: - del request - await context.abort( - grpc.StatusCode.UNIMPLEMENTED, "Embedding is not implemented by this server" - ) - - async def Classify(self, request: object, context: grpc.aio.ServicerContext) -> None: - del request - await context.abort( - grpc.StatusCode.UNIMPLEMENTED, "Classification is not implemented by this server" - ) - - async def Score(self, request: object, context: grpc.aio.ServicerContext) -> None: - del request - await context.abort( - grpc.StatusCode.UNIMPLEMENTED, "Scoring is not implemented by this server" - ) - - async def SubscribeRuntimeEvents( - self, request: object, context: grpc.aio.ServicerContext - ) -> None: - del request - await context.abort( - grpc.StatusCode.UNIMPLEMENTED, "Runtime event streaming is not implemented" - ) diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index 467dfc0fd218..8514e4de2801 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -43,7 +43,7 @@ def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> Non def test_schema_release_reads_launch_environment(monkeypatch) -> None: - commit = "f1a7189311770f8aa1f0dd787561df809847595d" + commit = "df3a9be24a2a36a4ff7a6d4fef9f1d7480ae210d" monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) assert schema_release() == commit From e8720db10ce7a3807ece78de77a86227758df7c2 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 21 Jul 2026 23:09:01 -0700 Subject: [PATCH 09/20] feat(openengine): support schema revision 3 Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 2 +- tensorrt_llm/commands/serve.py | 2 +- tensorrt_llm/inputs/registry.py | 16 +++++++++ tensorrt_llm/openengine/README.md | 5 +-- tensorrt_llm/openengine/_schema_pin.py | 2 +- tensorrt_llm/openengine/converters.py | 3 ++ tensorrt_llm/openengine/servicer.py | 32 ++++++++++++++--- .../inputs/test_multimodal_input_processor.py | 16 +++++++++ tests/unittest/openengine/test_converters.py | 9 +++++ tests/unittest/openengine/test_server.py | 2 +- tests/unittest/openengine/test_servicer.py | 34 +++++++++++++++++-- 11 files changed, 111 insertions(+), 12 deletions(-) diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT index c19080c682e4..c15702493507 100644 --- a/OPENENGINE_COMMIT +++ b/OPENENGINE_COMMIT @@ -1 +1 @@ -df3a9be24a2a36a4ff7a6d4fef9f1d7480ae210d +b0cf2a4826d246192dc65b055dab6d2b38d2d67e diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 8dcf393075f4..c94dc8fc2126 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -427,7 +427,7 @@ def launch_server( if error.name and error.name.startswith("openengine"): llm.shutdown() raise click.ClickException( - "OpenEngine support requires openengine-proto 0.2.0. " + "OpenEngine support requires openengine-proto 0.3.0. " "Run scripts/install_openengine.py against the pinned " "local OpenEngine sibling checkout.") from error raise diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 78d8e4a35c79..2598261c0823 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -192,6 +192,22 @@ def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: """Modalities supported in context-first OpenEngine P/D requests.""" return () + def get_openengine_routing_image_token_id(self) -> int | None: + """Return a stable image placeholder token for OpenEngine routing.""" + if "image" not in self.get_openengine_modalities(): + return None + candidates: set[int] = set() + for config in (self._config, getattr(self._config, "text_config", None), + getattr(self._config, "vision_config", None)): + for field_name in ("image_token_id", "image_token_index"): + value = getattr(config, field_name, None) + if isinstance(value, int) and not isinstance(value, bool): + candidates.add(value) + if len(candidates) != 1: + return None + token_id = candidates.pop() + return token_id if 0 <= token_id <= (1 << 32) - 1 else None + def get_required_lora_spec( self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: """Return a model-owned adapter required to process these modalities.""" diff --git a/tensorrt_llm/openengine/README.md b/tensorrt_llm/openengine/README.md index 4a7063ea4de6..af91deacca00 100644 --- a/tensorrt_llm/openengine/README.md +++ b/tensorrt_llm/openengine/README.md @@ -7,8 +7,9 @@ `--openengine-port`. The feature remains disabled when the flag is absent and keeps the OpenEngine bindings out of TensorRT-LLM's required dependencies. -The exact sibling source revision is recorded in `OPENENGINE_COMMIT`. Verify -and install its generated Python package with: +The exact schema-revision-3 sibling source revision and `openengine-proto` +0.3.0 package are recorded in `OPENENGINE_COMMIT`. Verify and install its +generated Python package with: ```bash python scripts/install_openengine.py diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py index 9be88c7e55dc..3c8565fafa54 100644 --- a/tensorrt_llm/openengine/_schema_pin.py +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -3,6 +3,6 @@ """OpenEngine source identity embedded in TensorRT-LLM distributions.""" -OPENENGINE_COMMIT = "df3a9be24a2a36a4ff7a6d4fef9f1d7480ae210d" +OPENENGINE_COMMIT = "b0cf2a4826d246192dc65b055dab6d2b38d2d67e" __all__ = ["OPENENGINE_COMMIT"] diff --git a/tensorrt_llm/openengine/converters.py b/tensorrt_llm/openengine/converters.py index 98a36d8fbd94..900eabbab126 100644 --- a/tensorrt_llm/openengine/converters.py +++ b/tensorrt_llm/openengine/converters.py @@ -254,12 +254,15 @@ def encode_handoff( session_id=str(params.disagg_request_id or params.ctx_request_id or ""), transfer_backend="tensorrt_llm", dp_rank=params.ctx_dp_rank or 0, + handoff_profile=HANDOFF_ATTRIBUTE, ) session.attributes_struct[HANDOFF_ATTRIBUTE] = canonical return session def _decode_handoff_payload(session: kv_pb2.KvSessionRef) -> dict[str, Any]: + if session.handoff_profile and session.handoff_profile != HANDOFF_ATTRIBUTE: + raise ValueError(f"Unsupported TensorRT-LLM handoff profile {session.handoff_profile!r}") if HANDOFF_ATTRIBUTE not in session.attributes_struct: raise ValueError(f"KV session is missing {HANDOFF_ATTRIBUTE!r}") encoded = session.attributes_struct[HANDOFF_ATTRIBUTE] diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index fa298ce3a187..900c99cfd7ba 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -44,6 +44,7 @@ from tensorrt_llm.serve.stats_fanout import StatsFanout from .converters import ( + HANDOFF_ATTRIBUTE, decode_handoff, encode_handoff, handoff_requires_decode_media, @@ -127,6 +128,14 @@ def __init__( ) -> None: self.llm = llm self.model = model + model_id = _arg(llm, "model", model) + self.model_id = str(model_id) if model_id else model + tokenizer_source = _arg(llm, "tokenizer", None) + self.tokenizer_source = str(tokenizer_source) if tokenizer_source else self.model_id + tokenizer_mode = _arg(llm, "tokenizer_mode", "auto") + self.tokenizer_mode = str(tokenizer_mode) if tokenizer_mode else "auto" + self.model_aliases = [self.model_id] if self.model_id != self.model else [] + self._accepted_model_names = {self.model, *self.model_aliases} self.role = role self.tracker = tracker self.media_config = media_config @@ -438,7 +447,7 @@ def _validate_generate( ) -> None: if not request.request_id: raise ValueError("request_id must not be empty") - if request.model and request.model != self.model: + if request.model and request.model not in self._accepted_model_names: raise ValueError(f"Unknown model {request.model!r}") if request.WhichOneof("input") is None: raise ValueError("Generate requires prompt or token_ids input") @@ -485,6 +494,8 @@ def _validate_generate( if self.role == server_pb2.ENGINE_ROLE_DECODE: if not request.kv.HasField("session"): raise ValueError("Decode requests require a prefill KV session") + if request.kv.session.HasField("bootstrap"): + raise ValueError("TensorRT-LLM does not support client-created KV bootstrap") requires_media = handoff_requires_decode_media(request.kv.session) if requires_media and not request.media: raise ValueError("Decode request must resend the ordered context-phase media") @@ -735,10 +746,14 @@ async def GetModelInfo( guided_modes.append(model_pb2.GUIDED_DECODING_MODE_STRUCTURAL_TAG) max_context = getattr(args, "max_seq_len", None) or getattr(args, "max_input_len", None) info = model_pb2.ModelInfo( - model_id=self.model, + model_id=self.model_id, served_model_name=self.model, - served_model_aliases=[self.model], - tokenizer_modes=["auto"], + served_model_aliases=self.model_aliases, + tokenizer_modes=[self.tokenizer_mode], + tokenizer=model_pb2.TokenizerInfo( + source=self.tokenizer_source, + mode=self.tokenizer_mode, + ), supports_text_input=True, supports_token_ids_input=True, supports_lora=self._supports_lora(), @@ -777,6 +792,13 @@ async def GetModelInfo( supports_per_request_media_options=True, ), ) + routing_token_getter = getattr( + input_processor, "get_openengine_routing_image_token_id", None + ) + if "image" in aggregate and callable(routing_token_getter): + routing_image_token_id = routing_token_getter() + if routing_image_token_id is not None: + info.multimodal_capabilities.routing_image_token_id = routing_image_token_id if max_context is not None: info.max_context_length = max_context return info @@ -1047,6 +1069,8 @@ def _kv_connector_info(self) -> kv_pb2.KvConnectorInfo: supports_abort_cleanup=enabled, supports_drain=enabled, schema_version=1, + handoff_profile=HANDOFF_ATTRIBUTE, + supports_client_bootstrap=False, ) def _kv_events_enabled(self) -> bool: diff --git a/tests/unittest/inputs/test_multimodal_input_processor.py b/tests/unittest/inputs/test_multimodal_input_processor.py index d94354cedf0c..496b538ce2dc 100644 --- a/tests/unittest/inputs/test_multimodal_input_processor.py +++ b/tests/unittest/inputs/test_multimodal_input_processor.py @@ -96,6 +96,22 @@ def test_get_num_tokens_per_video_delegates_to_hf_processor(): sub._processor._get_num_multimodal_tokens.assert_called_once() +@pytest.mark.parametrize( + "config, expected", + [ + (SimpleNamespace(image_token_id=151655), 151655), + (SimpleNamespace(image_token_id=151655, image_token_index=151646), None), + ], +) +def test_openengine_routing_image_token_requires_one_stable_config_value( + config: SimpleNamespace, expected: int | None +) -> None: + sub = _make_subclass() + sub._config = config + + assert sub.get_openengine_routing_image_token_id() == expected + + # --------------------------------------------------------------------------- # Dummy contract defaults (BaseMultimodalDummyInputsBuilder). # --------------------------------------------------------------------------- diff --git a/tests/unittest/openengine/test_converters.py b/tests/unittest/openengine/test_converters.py index 5a5c9853ab64..115dbbff92c9 100644 --- a/tests/unittest/openengine/test_converters.py +++ b/tests/unittest/openengine/test_converters.py @@ -42,6 +42,7 @@ def test_context_first_handoff_round_trip_large_ids_and_binary() -> None: ) session = encode_handoff(params) + assert session.handoff_profile == HANDOFF_ATTRIBUTE payload = json.loads(session.attributes_struct[HANDOFF_ATTRIBUTE]) assert payload["ctx_request_id"] == str(2**63 - 2) assert payload["disagg_request_id"] == str(2**63 - 1) @@ -111,6 +112,14 @@ def test_decode_rejects_non_decimal_identifier() -> None: decode_handoff(session) +def test_decode_rejects_foreign_handoff_profile() -> None: + session = kv_pb2.KvSessionRef(handoff_profile="vllm.kv_transfer_params.v1") + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps({"schedule_style": "context_first"}) + + with pytest.raises(ValueError, match="Unsupported TensorRT-LLM handoff profile"): + decode_handoff(session) + + def test_decode_rejects_multimodal_hashes_even_without_embedding_handles() -> None: session = kv_pb2.KvSessionRef() session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index 8514e4de2801..d53efe83f5f6 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -43,7 +43,7 @@ def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> Non def test_schema_release_reads_launch_environment(monkeypatch) -> None: - commit = "df3a9be24a2a36a4ff7a6d4fef9f1d7480ae210d" + commit = "b0cf2a4826d246192dc65b055dab6d2b38d2d67e" monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) assert schema_release() == commit diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index ff60a71aeb89..7ff2615bc065 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -607,14 +607,33 @@ def test_kv_batch_fails_closed_for_unrepresentable_cache_namespace(unsupported) @pytest.mark.asyncio -async def test_discovery_and_load_use_config_and_shared_stats() -> None: +async def test_discovery_and_load_use_config_and_shared_stats(monkeypatch) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image",) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return ("image",) + + def get_openengine_routing_image_token_id(self) -> int: + return 151655 + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> None: + del modalities + return None + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) llm = _Llm() llm.args = SimpleNamespace( + model="Qwen/Qwen3-VL-2B-Instruct", + tokenizer="Qwen/Qwen3-VL-2B-Instruct", + tokenizer_mode="slow", guided_decoding_backend=None, lora_config=None, enable_lora=False, kv_cache_config=SimpleNamespace(tokens_per_block=64), ) + llm.input_processor = _Processor() llm.get_kv_cache_capacity = lambda: { "maxNumBlocks": 100, "tokensPerBlock": 64, @@ -636,16 +655,27 @@ async def test_discovery_and_load_use_config_and_shared_stats() -> None: ) service = OpenEngineServicer( llm, - "model", + "qwen3-vl", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm), stats_fanout=stats, ) server_info = await service.GetServerInfo(server_pb2.GetServerInfoRequest(), _Context()) + assert server_info.schema_revision == 3 + assert server_info.minimum_client_revision == 1 assert server_info.capacity.kv_block_size == 64 assert server_info.capacity.total_kv_blocks == 100 + assert server_info.kv_connector.handoff_profile == "tensorrt_llm.disaggregated_params.v1" + assert server_info.kv_connector.HasField("supports_client_bootstrap") + assert not server_info.kv_connector.supports_client_bootstrap model_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert model_info.model_id == "Qwen/Qwen3-VL-2B-Instruct" + assert model_info.served_model_name == "qwen3-vl" + assert list(model_info.served_model_aliases) == ["Qwen/Qwen3-VL-2B-Instruct"] + assert model_info.tokenizer.source == "Qwen/Qwen3-VL-2B-Instruct" + assert model_info.tokenizer.mode == "slow" + assert model_info.multimodal_capabilities.routing_image_token_id == 151655 assert not model_info.generation.guided_decoding.supported assert not model_info.supports_lora From b51a218501bc35f6632d6590904bfd6f076b5a7b Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 21 Jul 2026 23:59:46 -0700 Subject: [PATCH 10/20] fix(openengine): preserve routing and model identities Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/converters.py | 14 ++++++++++++++ tensorrt_llm/openengine/servicer.py | 9 +++++++-- tests/unittest/openengine/test_servicer.py | 12 ++++++++++-- 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/openengine/converters.py b/tensorrt_llm/openengine/converters.py index 900eabbab126..51d0e4639b24 100644 --- a/tensorrt_llm/openengine/converters.py +++ b/tensorrt_llm/openengine/converters.py @@ -183,6 +183,20 @@ async def load_media( return output +def media_uuids(media: list[object]) -> dict[str, list[str | None]] | None: + """Preserve caller media IDs for engine KV-event key correlation.""" + output: dict[str, list[str | None]] = {} + any_uuid = False + for item in media: + modality = _MODALITY_NAMES.get(item.modality) + if modality is None: + raise ValueError(f"Unsupported media modality {item.modality}") + uuid = item.uuid or None + any_uuid = any_uuid or uuid is not None + output.setdefault(modality, []).append(uuid) + return output if any_uuid else None + + def stable_request_id(request_id: str) -> int: """Map arbitrary wire request IDs into TRT-LLM's positive int64 domain.""" digest = hashlib.sha256(request_id.encode("utf-8")).digest() diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index 900c99cfd7ba..e88d1fa084f5 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -49,6 +49,7 @@ encode_handoff, handoff_requires_decode_media, load_media, + media_uuids, stable_request_id, to_priority, to_sampling_params, @@ -224,7 +225,8 @@ async def Generate( target_dp_rank = session_dp_rank self._validate_generate(request, target_dp_rank) params = to_sampling_params(request) - media = await load_media(list(request.media), request.media_options, self.media_config) + media_items = list(request.media) + media = await load_media(media_items, request.media_options, self.media_config) input_kind = request.WhichOneof("input") inputs: dict[str, Any] if input_kind == "prompt": @@ -233,6 +235,9 @@ async def Generate( inputs = {"prompt_token_ids": list(request.token_ids.ids)} if media: inputs["multi_modal_data"] = media + uuids = media_uuids(media_items) + if uuids is not None: + inputs["multi_modal_uuids"] = uuids lora_request = self._required_multimodal_lora(request) if lora_request is not None and request.lora_name: raise ValueError( @@ -720,7 +725,7 @@ async def GetServerInfo( async def GetModelInfo( self, request: model_pb2.GetModelInfoRequest, context: grpc.aio.ServicerContext ) -> model_pb2.ModelInfo: - if request.model and request.model != self.model: + if request.model and request.model not in self._accepted_model_names: await context.abort(grpc.StatusCode.NOT_FOUND, f"Unknown model {request.model!r}") raise ValueError(f"Unknown model {request.model!r}") input_processor = getattr(self.llm, "input_processor", None) diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index 7ff2615bc065..d6ae3df83c99 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -241,12 +241,17 @@ async def _load_media(*_args, **_kwargs): llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) ) request = generation_pb2.GenerateRequest(request_id="audio", model="model", prompt="transcribe") - request.media.add(modality=input_pb2.MODALITY_AUDIO, raw_bytes=b"audio") + request.media.add( + modality=input_pb2.MODALITY_AUDIO, + raw_bytes=b"audio", + uuid="0123456789abcdef", + ) _ = [response async for response in service.Generate(request, _Context())] assert llm.kwargs["lora_request"].lora_name == "speech-lora" assert llm.kwargs["lora_request"].lora_int_id == 1 + assert llm.kwargs["inputs"]["multi_modal_uuids"] == {"audio": ["0123456789abcdef"]} assert await service.loras.list() == [] @@ -669,7 +674,10 @@ def get_required_lora_spec(self, modalities: tuple[str, ...]) -> None: assert server_info.kv_connector.handoff_profile == "tensorrt_llm.disaggregated_params.v1" assert server_info.kv_connector.HasField("supports_client_bootstrap") assert not server_info.kv_connector.supports_client_bootstrap - model_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + model_info = await service.GetModelInfo( + model_pb2.GetModelInfoRequest(model="Qwen/Qwen3-VL-2B-Instruct"), + _Context(), + ) assert model_info.model_id == "Qwen/Qwen3-VL-2B-Instruct" assert model_info.served_model_name == "qwen3-vl" assert list(model_info.served_model_aliases) == ["Qwen/Qwen3-VL-2B-Instruct"] From ffc9e8f178a395914384687bcae1a6ec6a5f031d Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 22 Jul 2026 00:14:16 -0700 Subject: [PATCH 11/20] feat(openengine): emit structured execution evidence Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/servicer.py | 127 +++++++++++++++++++++ tests/unittest/openengine/test_servicer.py | 27 ++++- 2 files changed, 153 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index e88d1fa084f5..b2ddf203ecb7 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -4,6 +4,7 @@ """OpenEngine gRPC servicer backed by one TensorRT-LLM LLM instance.""" import asyncio +import json import os import time import uuid @@ -89,6 +90,16 @@ def _signed_i64(value: int) -> int: return value +def _seconds(value: object) -> float: + converter = getattr(value, "total_seconds", None) + if callable(converter): + return float(converter()) + try: + return float(value) + except (TypeError, ValueError): + return 0.0 + + def _block_hash(value: object, hash_algo: str | None) -> kv_pb2.KvBlockHash: if isinstance(value, bytes): return kv_pb2.KvBlockHash(value=value, encoding=hash_algo or "bytes") @@ -164,6 +175,100 @@ def __init__( self.event_host = event_host self.event_port = event_port + def _log_handoff( + self, + request: generation_pb2.GenerateRequest, + phase: str, + session: kv_pb2.KvSessionRef | None = None, + result: object | None = None, + ) -> None: + if self.role == server_pb2.ENGINE_ROLE_AGGREGATED: + return + if session is None: + if not request.kv.HasField("session"): + return + session = request.kv.session + if not session.session_id: + return + + record: dict[str, Any] = { + "phase": phase, + "role": server_pb2.EngineRole.Name(self.role), + "request_id": request.request_id, + "session_id": session.session_id, + "handoff_profile": session.handoff_profile, + "dp_rank": session.dp_rank, + } + try: + handoff = decode_handoff(session) + except (TypeError, ValueError): + handoff = None + if handoff is not None: + record["tensorrt_llm"] = { + "ctx_request_id": ( + str(handoff.ctx_request_id) if handoff.ctx_request_id is not None else None + ), + "disagg_request_id": ( + str(handoff.disagg_request_id) + if handoff.disagg_request_id is not None + else None + ), + "ctx_dp_rank": handoff.ctx_dp_rank, + "ctx_info_endpoint": handoff.ctx_info_endpoint, + "conversation_id": handoff.conversation_id, + "schedule_style": ( + handoff.schedule_style.name + if isinstance(handoff.schedule_style, DisaggScheduleStyle) + else str(handoff.schedule_style) + ), + } + + if result is not None: + for output in getattr(result, "outputs", ()): + metrics = getattr(output, "request_perf_metrics", None) + timing = getattr(metrics, "timing_metrics", None) + if timing is None: + continue + kv_bytes = int(getattr(timing, "kv_cache_size", 0) or 0) + start = getattr(timing, "kv_cache_transfer_start", None) + end = getattr(timing, "kv_cache_transfer_end", None) + start_seconds = _seconds(start) + end_seconds = _seconds(end) + record["kv_transfer"] = { + "bytes": kv_bytes, + "duration_seconds": max(0.0, end_seconds - start_seconds), + } + break + + logger.info("OpenEngine handoff %s", json.dumps(record, sort_keys=True)) + + def _log_lora_selection( + self, + request: generation_pb2.GenerateRequest, + phase: str, + session: kv_pb2.KvSessionRef | None = None, + ) -> None: + if not request.lora_name: + return + session_id = "" + if session is not None: + session_id = session.session_id + elif request.kv.HasField("session"): + session_id = request.kv.session.session_id + logger.info( + "OpenEngine LoRA selection %s", + json.dumps( + { + "phase": phase, + "role": server_pb2.EngineRole.Name(self.role), + "request_id": request.request_id, + "session_id": session_id, + "lora_name": request.lora_name, + }, + sort_keys=True, + ), + ) + @staticmethod def _request_metadata(context: grpc.aio.ServicerContext) -> dict[str, str]: metadata: dict[str, str] = {} @@ -208,6 +313,10 @@ async def Generate( ) -> AsyncGenerator[generation_pb2.GenerateResponse, None]: result = None consumed_session_id = None + produced_session = None + admitted = False + selected_lora = False + selection_logged = False try: metadata = self._request_metadata(context) priority_value = self._metadata_int( @@ -272,6 +381,10 @@ async def Generate( priority=priority, ) self.tracker.admit(request.request_id, result) + admitted = True + self._log_handoff(request, "admitted") + if request.lora_name: + selected_lora = True if self.stats is not None: self.stats.wake() if request.kv.HasField("session") and request.kv.session.session_id: @@ -311,6 +424,9 @@ async def Generate( if context.cancelled(): await self.tracker.abort(request.request_id) return + if selected_lora and not selection_logged: + self._log_lora_selection(request, "selected") + selection_logged = True if not prompt_sent: prompt = self._prompt_output(current) if prompt is not None: @@ -335,7 +451,11 @@ async def Generate( ), ) session = encode_handoff(handoff, requires_decode_media=bool(request.media)) + produced_session = session self._track_kv_session(session.session_id, request.request_id) + self._log_handoff(request, "complete", session, current) + if selection_logged: + self._log_lora_selection(request, "complete", session) yield generation_pb2.GenerateResponse( request_id=request.request_id, prefill_ready=generation_pb2.PrefillReady(kv_session=session), @@ -362,6 +482,9 @@ async def Generate( sent_tokens[output.index] = len(token_ids) sent_text[output.index] = len(output.text or "") if current.finished: + self._log_handoff(request, "complete", result=current) + if selection_logged: + self._log_lora_selection(request, "complete") outputs = current.outputs for index, output in enumerate(outputs): response = generation_pb2.GenerateResponse( @@ -386,6 +509,10 @@ async def Generate( ), ) finally: + if admitted and not completed: + self._log_handoff(request, "aborted", produced_session) + if selection_logged: + self._log_lora_selection(request, "aborted", produced_session) if consumed_session_id is not None: self._release_kv_session(consumed_session_id) if not completed and request.request_id in self.tracker.active_requests: diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index d6ae3df83c99..fd47c69d59cf 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import asyncio +import json from types import SimpleNamespace import pytest @@ -284,7 +285,13 @@ def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraS @pytest.mark.asyncio -async def test_context_then_generation_round_trips_handoff() -> None: +async def test_context_then_generation_round_trips_handoff(monkeypatch) -> None: + log_messages: list[str] = [] + + def capture_info(message: str, *args: object) -> None: + log_messages.append(message % args) + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.logger.info", capture_info) context_handoff = DisaggregatedParams( request_type="context_only", first_gen_tokens=[7], @@ -347,6 +354,24 @@ async def test_context_then_generation_round_trips_handoff() -> None: assert terminal_usage.completion_tokens == 2 assert terminal_usage.total_tokens == 4 + handoff_records = [ + json.loads(message.split("OpenEngine handoff ", 1)[1]) + for message in log_messages + if "OpenEngine handoff " in message + ] + completed = [record for record in handoff_records if record["phase"] == "complete"] + assert {record["role"] for record in completed} == { + "ENGINE_ROLE_PREFILL", + "ENGINE_ROLE_DECODE", + } + assert {record["session_id"] for record in completed} == { + context_responses[0].prefill_ready.kv_session.session_id + } + assert {record["handoff_profile"] for record in completed} == { + "tensorrt_llm.disaggregated_params.v1" + } + assert {record["tensorrt_llm"]["disagg_request_id"] for record in completed} == {"101"} + abort_response = await context_service.Abort( lifecycle_pb2.AbortRequest(kv_session=context_responses[0].prefill_ready.kv_session), _Context(), From 9b8deed6ec66fe40511f5a921342cea984eba01a Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 22 Jul 2026 00:28:16 -0700 Subject: [PATCH 12/20] fix(openengine): avoid wildcard security false positive Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/server.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/openengine/server.py b/tensorrt_llm/openengine/server.py index 54987ef791f0..f13e845bb694 100644 --- a/tensorrt_llm/openengine/server.py +++ b/tensorrt_llm/openengine/server.py @@ -3,6 +3,8 @@ """Lifecycle wrapper for the optional OpenEngine sibling gRPC server.""" +import ipaddress + import grpc from openengine.v1 import openengine_pb2_grpc, server_pb2 @@ -16,6 +18,7 @@ from .servicer import OpenEngineServicer, schema_release _MAX_MESSAGE_LENGTH = 256 * 1024 * 1024 +_IPV4_UNSPECIFIED = str(ipaddress.IPv4Address(0)) def validate_schema_release(schema_release: str) -> str: @@ -78,7 +81,7 @@ def __init__( raise RuntimeError(f"Failed to bind OpenEngine server to {host}:{port}") self.port = bound if port == 0 else port advertised_host = { - "0.0.0.0": "127.0.0.1", + _IPV4_UNSPECIFIED: "127.0.0.1", "::": "::1", "[::]": "::1", }.get(host, host) From 6953e49df5ebe16e710d8c2edcbb43696a64000b Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 22 Jul 2026 01:24:27 -0700 Subject: [PATCH 13/20] fix(openengine): preserve request ownership through cleanup Signed-off-by: Connor Carpenter --- tensorrt_llm/grpc/grpc_request_manager.py | 7 +- tensorrt_llm/openengine/server.py | 4 + tensorrt_llm/openengine/servicer.py | 99 ++++++++++---- tensorrt_llm/serve/kv_event_fanout.py | 35 ++++- tensorrt_llm/serve/request_tracker.py | 70 +++++++++- tests/unittest/openengine/test_fanout.py | 40 +++--- .../openengine/test_request_tracker.py | 100 ++++++++++++++ tests/unittest/openengine/test_server.py | 8 ++ tests/unittest/openengine/test_servicer.py | 125 +++++++++++++++++- 9 files changed, 429 insertions(+), 59 deletions(-) diff --git a/tensorrt_llm/grpc/grpc_request_manager.py b/tensorrt_llm/grpc/grpc_request_manager.py index 9cb4868682c6..61b17d734e00 100644 --- a/tensorrt_llm/grpc/grpc_request_manager.py +++ b/tensorrt_llm/grpc/grpc_request_manager.py @@ -134,7 +134,6 @@ async def generate( except asyncio.CancelledError: logger.info(f"Request {request_id} cancelled by client") - await self.abort(request_id) raise except Exception as e: logger.error(f"Error in generate for {request_id}: {e}") @@ -142,7 +141,11 @@ async def generate( finally: # Cleanup tracking if admitted: - await self._request_tracker.finish(request_id) + if gen_result.finished: + await self._request_tracker.finish(request_id, gen_result) + else: + await self._request_tracker.abort(request_id) + self._request_tracker.reap(request_id, gen_result) async def abort(self, request_id: str) -> bool: """Abort a running request. diff --git a/tensorrt_llm/openengine/server.py b/tensorrt_llm/openengine/server.py index f13e845bb694..6876dcbcfccb 100644 --- a/tensorrt_llm/openengine/server.py +++ b/tensorrt_llm/openengine/server.py @@ -108,6 +108,10 @@ async def stop(self, grace: float = 5.0) -> None: try: await self._server.stop(grace=grace) finally: + if not await self.servicer.tracker.close_reapers(timeout=grace): + logger.warning( + "OpenEngine request cleanup did not finish within %.1f seconds", grace + ) self.servicer.close() if self._kv_event_fanout is not None: await self._kv_event_fanout.stop() diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index b2ddf203ecb7..ac43b471cfe6 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -40,7 +40,7 @@ get_effective_kv_cache_event_hash_algo, ) from tensorrt_llm.scheduling_params import SchedulingParams -from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout, KvEventSubscriberOverflow from tensorrt_llm.serve.request_tracker import RequestTracker from tensorrt_llm.serve.stats_fanout import StatsFanout @@ -163,6 +163,7 @@ def __init__( self.kv_events = kv_event_fanout self.stats = stats_fanout self._kv_session_requests: dict[str, str] = {} + self._kv_session_owners: dict[str, object] = {} self._kv_session_timers: dict[str, asyncio.TimerHandle] = {} if kv_session_ttl_seconds <= 0: raise ValueError("kv_session_ttl_seconds must be positive") @@ -313,6 +314,7 @@ async def Generate( ) -> AsyncGenerator[generation_pb2.GenerateResponse, None]: result = None consumed_session_id = None + consumed_session_owner = None produced_session = None admitted = False selected_lora = False @@ -369,6 +371,11 @@ async def Generate( for key, value in metadata.items() if key.lower() in ("traceparent", "tracestate", "baggage") } + if request.kv.HasField("session") and request.kv.session.session_id: + consumed_session_id = request.kv.session.session_id + consumed_session_owner = self._track_kv_session( + consumed_session_id, request.request_id, expires=False + ) result = self.llm.generate_async( inputs=inputs, sampling_params=params, @@ -387,15 +394,20 @@ async def Generate( selected_lora = True if self.stats is not None: self.stats.wake() - if request.kv.HasField("session") and request.kv.session.session_id: - consumed_session_id = request.kv.session.session_id - self._track_kv_session(consumed_session_id, request.request_id) + except asyncio.CancelledError: + if consumed_session_owner is not None: + self._release_kv_session(consumed_session_id, consumed_session_owner) + raise except KeyError as error: + if consumed_session_owner is not None: + self._release_kv_session(consumed_session_id, consumed_session_owner) await context.abort( grpc.StatusCode.FAILED_PRECONDITION, f"LoRA is not loaded: {error.args[0]}" ) return except (ValueError, TypeError, RuntimeError) as error: + if consumed_session_owner is not None: + self._release_kv_session(consumed_session_id, consumed_session_owner) if ( result is not None and self.tracker.active_requests.get(request.request_id) is not result @@ -420,9 +432,10 @@ async def Generate( prompt_sent = False completed = False try: + if context.cancelled(): + return async for current in result: if context.cancelled(): - await self.tracker.abort(request.request_id) return if selected_lora and not selection_logged: self._log_lora_selection(request, "selected") @@ -496,7 +509,6 @@ async def Generate( yield response completed = True except asyncio.CancelledError: - await self.tracker.abort(request.request_id) raise except (RuntimeError, ValueError, TypeError) as error: logger.error("OpenEngine request %s failed: %s", request.request_id, error) @@ -513,25 +525,49 @@ async def Generate( self._log_handoff(request, "aborted", produced_session) if selection_logged: self._log_lora_selection(request, "aborted", produced_session) - if consumed_session_id is not None: - self._release_kv_session(consumed_session_id) - if not completed and request.request_id in self.tracker.active_requests: - await self.tracker.abort(request.request_id) - else: - await self.tracker.finish(request.request_id) + if admitted: + terminal = bool(getattr(result, "finished", False)) + if not terminal: + await self.tracker.abort(request.request_id) + reaper = self.tracker.reap(request.request_id, result) + if consumed_session_id is not None: + reaper.add_done_callback( + lambda _task, + session_id=consumed_session_id, + owner=consumed_session_owner: ( + self._release_kv_session(session_id, owner) + ) + ) + else: + if consumed_session_id is not None: + self._release_kv_session(consumed_session_id, consumed_session_owner) + await self.tracker.finish(request.request_id, result) if self.stats is not None: self.stats.wake() - def _track_kv_session(self, session_id: str, request_id: str) -> None: - self._release_kv_session(session_id) + def _track_kv_session( + self, session_id: str, request_id: str, *, expires: bool = True + ) -> object: + if session_id in self._kv_session_requests: + raise ValueError(f"KV session {session_id!r} is already in use") + owner = object() self._kv_session_requests[session_id] = request_id - loop = asyncio.get_running_loop() - self._kv_session_timers[session_id] = loop.call_later( - self._kv_session_ttl_seconds, self._release_kv_session, session_id - ) + self._kv_session_owners[session_id] = owner + if expires: + loop = asyncio.get_running_loop() + self._kv_session_timers[session_id] = loop.call_later( + self._kv_session_ttl_seconds, + self._release_kv_session, + session_id, + owner, + ) + return owner - def _release_kv_session(self, session_id: str) -> bool: + def _release_kv_session(self, session_id: str, owner: object | None = None) -> bool: + if owner is not None and self._kv_session_owners.get(session_id) is not owner: + return False released = self._kv_session_requests.pop(session_id, None) is not None + self._kv_session_owners.pop(session_id, None) timer = self._kv_session_timers.pop(session_id, None) if timer is not None: timer.cancel() @@ -557,6 +593,7 @@ def _release_all_kv_sessions(self) -> None: timer.cancel() self._kv_session_timers.clear() self._kv_session_requests.clear() + self._kv_session_owners.clear() def _active_kv_session_count(self) -> int: """Count sessions whose owning request is still engine-active. @@ -967,6 +1004,10 @@ async def GetLoad( prefill_batch_size = 0 decode_batch_size = 0 attributes = {"source": "shared_stats_fanout" if latest else "shared_request_tracker"} + if self.kv_events is not None: + attributes["kv_event_subscriber_overflows"] = str( + self.kv_events.subscriber_overflow_count + ) tokens_per_block: set[int] = set() for rank, stat in sorted(latest.items()): kv_stats = stat.get("kvCacheStats") or {} @@ -1080,7 +1121,8 @@ async def Abort( if target == "kv_session": session_id = request.kv_session.session_id request_id = self._kv_session_requests.get(session_id, session_id) - released_session = self._release_kv_session(session_id) + if request_id not in self.tracker.active_requests: + released_session = self._release_kv_session(session_id) aborted = await self.tracker.abort(request_id) return lifecycle_pb2.AbortResponse( status=( @@ -1277,13 +1319,16 @@ async def SubscribeKvEvents( f"Unknown data-parallel ranks: {sorted(invalid_ranks)}", ) return - async for sequence, raw in self.kv_events.subscribe(selected_ranks): - if context.cancelled(): - return - rank = raw.get("attention_dp_rank", 0) - if selected_ranks and rank not in selected_ranks: - continue - yield kv_pb2.SubscribeKvEventsResponse(batch=self._kv_batch(raw, sequence)) + try: + async for sequence, raw in self.kv_events.subscribe(selected_ranks): + if context.cancelled(): + return + rank = raw.get("attention_dp_rank", 0) + if selected_ranks and rank not in selected_ranks: + continue + yield kv_pb2.SubscribeKvEventsResponse(batch=self._kv_batch(raw, sequence)) + except KvEventSubscriberOverflow as error: + await context.abort(grpc.StatusCode.RESOURCE_EXHAUSTED, str(error)) def _kv_batch(self, raw: dict[str, Any], sequence: int) -> kv_pb2.KvEventBatch: rank = raw.get("attention_dp_rank", 0) diff --git a/tensorrt_llm/serve/kv_event_fanout.py b/tensorrt_llm/serve/kv_event_fanout.py index 60fce03c8a31..12826e8640de 100644 --- a/tensorrt_llm/serve/kv_event_fanout.py +++ b/tensorrt_llm/serve/kv_event_fanout.py @@ -6,11 +6,18 @@ import asyncio from collections import deque from collections.abc import AsyncGenerator -from typing import Any +from typing import Any, cast from tensorrt_llm.logger import logger +class KvEventSubscriberOverflow(RuntimeError): + """Raised when a subscriber can no longer receive a lossless event stream.""" + + +_SUBSCRIBER_OVERFLOW = object() + + class KvEventFanout: """Own TRT-LLM's event queue and copy events to every protocol.""" @@ -19,7 +26,8 @@ def __init__(self, llm: object, buffer_size: int = 1024) -> None: raise ValueError("buffer_size must be positive") self._llm = llm self._buffer: deque[dict[str, Any]] = deque(maxlen=buffer_size) - self._subscribers: dict[asyncio.Queue[tuple[int, dict[str, Any]]], frozenset[int]] = {} + self._subscribers: dict[asyncio.Queue[object], frozenset[int]] = {} + self._subscriber_overflow_count = 0 self._task: asyncio.Task[None] | None = None self._sequence_numbers: dict[int, int] = {} self._last_engine_event_ids: dict[int, int] = {} @@ -30,6 +38,10 @@ def __init__(self, llm: object, buffer_size: int = 1024) -> None: def buffer_size(self) -> int: return self._buffer.maxlen or 0 + @property + def subscriber_overflow_count(self) -> int: + return self._subscriber_overflow_count + def start(self) -> None: if self._task is None: self._task = asyncio.create_task(self._pump()) @@ -53,11 +65,16 @@ async def subscribe( self, data_parallel_ranks: set[int] | None = None ) -> AsyncGenerator[tuple[int, dict[str, Any]], None]: self.start() - queue: asyncio.Queue[tuple[int, dict[str, Any]]] = asyncio.Queue(maxsize=self.buffer_size) + queue: asyncio.Queue[object] = asyncio.Queue(maxsize=self.buffer_size) self._subscribers[queue] = frozenset(data_parallel_ranks or ()) try: while True: - yield await queue.get() + item = await queue.get() + if item is _SUBSCRIBER_OVERFLOW: + raise KvEventSubscriberOverflow( + "OpenEngine KV event subscriber queue overflowed" + ) + yield cast(tuple[int, dict[str, Any]], item) finally: self._subscribers.pop(queue, None) @@ -120,7 +137,15 @@ def _publish_routable( try: queue.put_nowait(sequenced_event) except asyncio.QueueFull: - logger.warning("Dropping KV event for a slow OpenEngine subscriber") + self._subscriber_overflow_count += 1 + self._subscribers.pop(queue, None) + while not queue.empty(): + queue.get_nowait() + queue.put_nowait(_SUBSCRIBER_OVERFLOW) + logger.error( + "Closing slow OpenEngine KV subscriber after queue overflow; " + "lossless routing can no longer be guaranteed" + ) async def _pump(self) -> None: while True: diff --git a/tensorrt_llm/serve/request_tracker.py b/tensorrt_llm/serve/request_tracker.py index 8c6ebe3c8674..487db2613b30 100644 --- a/tensorrt_llm/serve/request_tracker.py +++ b/tensorrt_llm/serve/request_tracker.py @@ -24,6 +24,8 @@ def __init__(self, llm: object) -> None: self._external_requests = 0 self._draining = False self._changed = asyncio.Condition() + self._reapers: set[asyncio.Task[None]] = set() + self._reaper_owners: dict[asyncio.Task[None], tuple[str, AbortableResult]] = {} @property def draining(self) -> bool: @@ -45,11 +47,70 @@ def admit(self, request_id: str, result: AbortableResult) -> None: raise ValueError(f"Request {request_id!r} is already active") self._requests[request_id] = result - async def finish(self, request_id: str) -> None: + async def finish(self, request_id: str, result: AbortableResult | None = None) -> None: async with self._changed: - self._requests.pop(request_id, None) + current = self._requests.get(request_id) + if result is None or current is result: + self._requests.pop(request_id, None) self._changed.notify_all() + def reap(self, request_id: str, result: AbortableResult) -> asyncio.Task[None]: + """Transfer an abandoned stream to a single terminal-response consumer.""" + + async def _reap() -> None: + try: + aresult = getattr(result, "aresult", None) + if callable(aresult): + await aresult() + else: + async for _ in result: # type: ignore[attr-defined] + pass + except asyncio.CancelledError: + raise + except Exception as error: + logger.warning("Request %s cleanup ended with: %s", request_id, error) + finally: + await self.finish(request_id, result) + + task = asyncio.create_task(_reap(), name=f"trtllm-request-reaper-{request_id}") + self._reapers.add(task) + self._reaper_owners[task] = (request_id, result) + + def _forget(completed: asyncio.Task[None]) -> None: + self._reapers.discard(completed) + self._reaper_owners.pop(completed, None) + + task.add_done_callback(_forget) + return task + + async def wait_reapers(self, timeout: float | None = None) -> bool: + """Wait for all currently owned engine-cleanup tasks without cancelling them.""" + if not self._reapers: + return True + _, pending = await asyncio.wait(tuple(self._reapers), timeout=timeout) + return not pending + + async def close_reapers(self, timeout: float | None = None) -> bool: + """Bound cleanup at process shutdown and release every owned task.""" + completed = await self.wait_reapers(timeout) + if completed: + return True + pending = tuple(self._reapers) + owners = tuple( + owner for task in pending if (owner := self._reaper_owners.get(task)) is not None + ) + for task in pending: + task.cancel() + # Give ordinary cancellation-aware result iterators one loop turn to + # run their finally blocks without ever waiting on a resistant task. + await asyncio.sleep(0) + for task in pending: + self._reapers.discard(task) + self._reaper_owners.pop(task, None) + for request_id, result in owners: + await self.finish(request_id, result) + return False + def begin_external(self) -> None: if self._draining: raise RuntimeError("TensorRT-LLM is draining") @@ -69,8 +130,6 @@ async def abort(self, request_id: str) -> bool: except (RuntimeError, AssertionError) as error: logger.warning("Failed to abort request %s: %s", request_id, error) return False - finally: - await self.finish(request_id) return True async def abort_all(self) -> int: @@ -91,9 +150,6 @@ async def abort_all(self) -> int: else: aborted += 1 - async with self._changed: - self._requests.clear() - self._changed.notify_all() return aborted async def start_drain(self) -> int: diff --git a/tests/unittest/openengine/test_fanout.py b/tests/unittest/openengine/test_fanout.py index 623910b7c0fa..1ed8b440cf2d 100644 --- a/tests/unittest/openengine/test_fanout.py +++ b/tests/unittest/openengine/test_fanout.py @@ -5,28 +5,34 @@ import pytest -from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout, KvEventSubscriberOverflow from tensorrt_llm.serve.stats_fanout import StatsFanout -def test_kv_fanout_sequence_exposes_subscriber_queue_drops() -> None: +@pytest.mark.asyncio +async def test_kv_fanout_closes_subscriber_instead_of_hiding_queue_overflow() -> None: fanout = KvEventFanout(object(), buffer_size=1) - queue = asyncio.Queue(maxsize=1) - fanout._subscribers[queue] = frozenset({0}) + subscription = fanout.subscribe({0}) + try: + first_item = asyncio.create_task(subscription.__anext__()) + await asyncio.sleep(0) - first = {"data": {"type": "stored"}, "event": 1, "attention_dp_rank": 0} - second = {"data": {"type": "stored"}, "event": 2, "attention_dp_rank": 0} - third = {"data": {"type": "stored"}, "event": 3, "attention_dp_rank": 0} - other_rank = {"data": {"type": "stored"}, "event": 20, "attention_dp_rank": 1} - fanout._publish(first) - fanout._publish(second) - assert queue.get_nowait() == (1, first) - - fanout._publish(other_rank) - fanout._publish(third) - assert queue.get_nowait() == (3, third) - assert fanout.drain_http_buffer() == [third] - assert fanout._sequence_numbers == {0: 3, 1: 1} + first = {"data": {"type": "stored"}, "event": 1, "attention_dp_rank": 0} + second = {"data": {"type": "stored"}, "event": 2, "attention_dp_rank": 0} + third = {"data": {"type": "stored"}, "event": 3, "attention_dp_rank": 0} + fanout._publish(first) + assert await first_item == (1, first) + fanout._publish(second) + fanout._publish(third) + with pytest.raises(KvEventSubscriberOverflow, match="overflowed"): + await subscription.__anext__() + + assert fanout.subscriber_overflow_count == 1 + assert not fanout._subscribers + assert fanout.drain_http_buffer() == [third] + assert fanout._sequence_numbers == {0: 3} + finally: + await fanout.stop() def test_kv_fanout_filters_non_global_attention_without_sequence_gaps() -> None: diff --git a/tests/unittest/openengine/test_request_tracker.py b/tests/unittest/openengine/test_request_tracker.py index 5ab44e5a3e5d..602e83a84af6 100644 --- a/tests/unittest/openengine/test_request_tracker.py +++ b/tests/unittest/openengine/test_request_tracker.py @@ -30,6 +30,8 @@ async def test_tracker_abort_drain_and_external_admission() -> None: assert await tracker.abort("request") assert result.aborted assert not await tracker.wait_empty(0) + await tracker.finish("request") + assert not await tracker.wait_empty(0) await tracker.finish_external() assert await tracker.wait_empty(0) @@ -71,3 +73,101 @@ async def finish_http() -> None: assert await tracker.wait_empty() await finish_task assert result.aborted + + +@pytest.mark.asyncio +async def test_abort_failure_remains_tracked_until_stream_finishes() -> None: + class _FailingResult: + def abort(self) -> None: + raise RuntimeError("engine still owns request") + + tracker = RequestTracker(object()) + tracker.admit("request", _FailingResult()) + + assert not await tracker.abort("request") + assert tracker.active_count == 1 + assert not await tracker.wait_empty(0) + + await tracker.finish("request") + assert await tracker.wait_empty(0) + + +@pytest.mark.asyncio +async def test_abort_all_waits_for_tracked_stream_termination() -> None: + result = _Result() + tracker = RequestTracker(object()) + tracker.admit("request", result) + + assert await tracker.abort_all() == 1 + assert result.aborted + assert not await tracker.wait_empty(0) + + await tracker.finish("request") + assert await tracker.wait_empty(0) + + +@pytest.mark.asyncio +async def test_reaper_keeps_abort_active_until_engine_terminal() -> None: + class _DelayedResult(_Result): + def __init__(self) -> None: + super().__init__() + self.terminal = asyncio.Event() + + async def aresult(self) -> "_DelayedResult": + await self.terminal.wait() + return self + + tracker = RequestTracker(object()) + result = _DelayedResult() + tracker.admit("request", result) + + assert await tracker.abort("request") + reaper = tracker.reap("request", result) + assert not await tracker.wait_empty(0) + + result.terminal.set() + await reaper + assert await tracker.wait_empty(0) + + +@pytest.mark.asyncio +async def test_late_reaper_cannot_remove_reused_request_id() -> None: + class _DelayedResult(_Result): + def __init__(self) -> None: + super().__init__() + self.terminal = asyncio.Event() + + async def aresult(self) -> "_DelayedResult": + await self.terminal.wait() + return self + + tracker = RequestTracker(object()) + old = _DelayedResult() + tracker.admit("request", old) + reaper = tracker.reap("request", old) + await tracker.finish("request", old) + + replacement = _Result() + tracker.admit("request", replacement) + old.terminal.set() + await reaper + + assert tracker.active_requests["request"] is replacement + await tracker.finish("request", replacement) + + +@pytest.mark.asyncio +async def test_reaper_shutdown_is_bounded_and_releases_tracking() -> None: + class _NeverTerminal(_Result): + async def aresult(self) -> "_NeverTerminal": + await asyncio.Event().wait() + return self + + tracker = RequestTracker(object()) + result = _NeverTerminal() + tracker.admit("request", result) + tracker.reap("request", result) + + assert not await tracker.close_reapers(timeout=0) + assert await tracker.wait_empty(0) + assert not tracker._reapers diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index d53efe83f5f6..fa8d8f3a900a 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -58,6 +58,13 @@ async def stop(self, grace: float) -> None: raise RuntimeError("stop failed") class _Servicer: + class _Tracker: + async def close_reapers(self, timeout: float) -> bool: + calls.append(("reapers", timeout)) + return True + + tracker = _Tracker() + def close(self) -> None: calls.append(("servicer", None)) @@ -80,6 +87,7 @@ async def stop(self) -> None: assert calls == [ ("grpc", 1.5), + ("reapers", 1.5), ("servicer", None), ("kv_events", None), ("stats", None), diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index fd47c69d59cf..ed47709e91b5 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -48,7 +48,7 @@ def __init__(self, handoff: DisaggregatedParams | None = None) -> None: disaggregated_params=handoff, ) ] - self.finished = True + self.finished = False self._yielded = False self.aborted = False @@ -59,6 +59,7 @@ async def __anext__(self) -> "_Result": if self._yielded: raise StopAsyncIteration self._yielded = True + self.finished = True return self def abort(self) -> None: @@ -150,9 +151,65 @@ async def test_generate_cancellation_aborts_and_cleans_tracking() -> None: assert responses == [] assert llm.result.aborted + assert await tracker.wait_empty(0.1) assert tracker.active_count == 0 +@pytest.mark.asyncio +async def test_decode_cancellation_during_media_load_does_not_claim_session( + monkeypatch, +) -> None: + entered = asyncio.Event() + + async def blocked_load_media(*_args, **_kwargs): + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.load_media", blocked_load_media) + llm = _Llm() + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) + request = generation_pb2.GenerateRequest( + request_id="decode-cancel", model="model", prompt="hello" + ) + request.kv.session.CopyFrom( + encode_handoff( + DisaggregatedParams( + request_type="context_only", + disagg_request_id=101, + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + ) + ) + ) + + generation = asyncio.create_task(_collect_async(service.Generate(request, _Context()))) + await entered.wait() + assert not service._kv_session_requests + generation.cancel() + with pytest.raises(asyncio.CancelledError): + await generation + + assert not service._kv_session_requests + assert not service._kv_session_owners + + +@pytest.mark.asyncio +async def test_rejected_duplicate_does_not_finish_original_request() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + original = _Result() + tracker.admit("duplicate", original) + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_AGGREGATED, tracker) + request = generation_pb2.GenerateRequest(request_id="duplicate", model="model", prompt="hello") + + with pytest.raises(AssertionError, match="already active"): + await _collect_async(service.Generate(request, _Context())) + + assert tracker.active_requests["duplicate"] is original + assert not await tracker.wait_empty(0) + assert await tracker.abort("duplicate") + await tracker.finish("duplicate", original) + + def test_generate_rejects_unsupported_cache_bypass_and_nonzero_dp_placement() -> None: llm = _Llm() service = OpenEngineServicer( @@ -439,6 +496,72 @@ async def test_prefill_session_expires_and_close_cancels_timers() -> None: assert not service._kv_session_requests +@pytest.mark.asyncio +async def test_active_decode_session_does_not_expire() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, + "model", + server_pb2.ENGINE_ROLE_DECODE, + RequestTracker(llm), + kv_session_ttl_seconds=0.01, + ) + owner = service._track_kv_session("active-decode", "decode", expires=False) + await asyncio.sleep(0.02) + assert service._kv_session_requests["active-decode"] == "decode" + assert "active-decode" not in service._kv_session_timers + assert service._release_kv_session("active-decode", owner) + + +@pytest.mark.asyncio +async def test_late_kv_session_release_cannot_remove_reused_session_id() -> None: + llm = _Llm() + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) + old_owner = service._track_kv_session("session", "old-request") + release_old = asyncio.Event() + + async def late_release() -> bool: + await release_old.wait() + return service._release_kv_session("session", old_owner) + + late_reaper = asyncio.create_task(late_release()) + assert service._release_kv_session("session", old_owner) + new_owner = service._track_kv_session("session", "new-request") + release_old.set() + + assert not await late_reaper + assert service._kv_session_requests["session"] == "new-request" + assert service._kv_session_owners["session"] is new_owner + assert service._release_kv_session("session", new_owner) + assert not service._kv_session_timers + + +@pytest.mark.asyncio +async def test_kv_session_rejects_concurrent_duplicate_consumers() -> None: + llm = _Llm() + service = OpenEngineServicer(llm, "model", server_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) + session = encode_handoff( + DisaggregatedParams( + request_type="context_only", + disagg_request_id=101, + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + ) + ) + owner = service._track_kv_session(session.session_id, "first-request", expires=False) + request = generation_pb2.GenerateRequest( + request_id="second-request", model="model", prompt="hello" + ) + request.kv.session.CopyFrom(session) + + with pytest.raises(AssertionError, match="already in use"): + await _collect_async(service.Generate(request, _Context())) + + assert llm.result is None + assert service._kv_session_requests[session.session_id] == "first-request" + assert service._kv_session_owners[session.session_id] is owner + assert service._release_kv_session(session.session_id, owner) + + def test_decode_media_is_required_only_for_marked_handoff(monkeypatch) -> None: class _Processor: def get_openengine_modalities(self) -> tuple[str, ...]: From 01138754b3f4b047dfacd714ff121cc7f4df6d67 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 22 Jul 2026 16:51:35 -0700 Subject: [PATCH 14/20] fix(openengine): support advertised video decoding Signed-off-by: Connor Carpenter --- requirements-openengine.txt | 3 ++ scripts/install_openengine.py | 11 +++++- setup.py | 2 + tensorrt_llm/commands/serve.py | 2 +- tensorrt_llm/inputs/multimodal.py | 13 ++++++- tensorrt_llm/openengine/README.md | 8 +++- tensorrt_llm/openengine/server.py | 32 ++++++++++++++- tests/unittest/inputs/test_multimodal.py | 41 ++++++++++++++++++++ tests/unittest/openengine/test_packaging.py | 21 ++++++++++ tests/unittest/openengine/test_server.py | 43 ++++++++++++++++++++- 10 files changed, 170 insertions(+), 6 deletions(-) create mode 100644 requirements-openengine.txt create mode 100644 tests/unittest/openengine/test_packaging.py diff --git a/requirements-openengine.txt b/requirements-openengine.txt new file mode 100644 index 000000000000..7690a28f9485 --- /dev/null +++ b/requirements-openengine.txt @@ -0,0 +1,3 @@ +# OpenCV 4.12 requires NumPy <2.3, while TensorRT-LLM supports NumPy <2.4 and +# current serving images use 2.3. Pin the last compatible video decoder line. +opencv-python-headless==4.11.0.86 diff --git a/scripts/install_openengine.py b/scripts/install_openengine.py index 89f8f3e3ba93..e159285cd0e8 100644 --- a/scripts/install_openengine.py +++ b/scripts/install_openengine.py @@ -68,7 +68,16 @@ def main() -> None: raise RuntimeError(f"OpenEngine Python package is missing: {package}") if not args.verify_only: subprocess.run( - [sys.executable, "-m", "pip", "install", "-e", str(package)], + [ + sys.executable, + "-m", + "pip", + "install", + "-e", + str(package), + "-r", + str(root / "requirements-openengine.txt"), + ], cwd=root, check=True, ) diff --git a/setup.py b/setup.py index e55f60291b89..d5f5ac8f7b75 100644 --- a/setup.py +++ b/setup.py @@ -129,6 +129,7 @@ def has_ext_modules(self): devel_deps, _ = parse_requirements( Path("requirements-dev-windows.txt" if on_windows else "requirements-dev.txt")) +openengine_deps, _ = parse_requirements(Path("requirements-openengine.txt")) mx_deps = ["modelexpress==0.4.1"] constraints_file = Path("constraints.txt") if constraints_file.exists(): @@ -464,6 +465,7 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'], extras_require={ "devel": devel_deps, + "openengine": openengine_deps, "mx": mx_deps, }, zip_safe=True, diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index c94dc8fc2126..cfd0dccf7d12 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -514,7 +514,7 @@ async def serve_http_and_openengine(): host=openengine_host, port=openengine_port, tracker=request_tracker, - media_config=multimodal_server_config, + media_config=server.multimodal_server_config, reasoning_parser=getattr(llm.args, "reasoning_parser", None), tool_parser=tool_parser, diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index fb4b33779d32..2c656cf0c165 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -1026,7 +1026,18 @@ def find_mm_token_lengths( video_metadata = item.metadata video_audio = item.audio item = item.frames - assert isinstance(item, list), "Video must be a list of frames" + if isinstance(item, np.ndarray): + if item.ndim != 4 or item.shape[-1] != 3: + raise ValueError( + "Video NumPy frames must have shape (N, H, W, 3)" + ) + item = list(item) + if not isinstance(item, list): + raise TypeError( + "Video must be a list of frames or a 4D NumPy frame array" + ) + if not item: + raise ValueError("Video frames cannot be empty") call_kwargs = {"video": item} if video_metadata is not None: # Used by per-model overrides that need to account for diff --git a/tensorrt_llm/openengine/README.md b/tensorrt_llm/openengine/README.md index af91deacca00..5e0a2b37a8a3 100644 --- a/tensorrt_llm/openengine/README.md +++ b/tensorrt_llm/openengine/README.md @@ -9,13 +9,19 @@ keeps the OpenEngine bindings out of TensorRT-LLM's required dependencies. The exact schema-revision-3 sibling source revision and `openengine-proto` 0.3.0 package are recorded in `OPENENGINE_COMMIT`. Verify and install its -generated Python package with: +generated Python package and OpenEngine runtime dependencies with: ```bash python scripts/install_openengine.py python -m pip install -e . ``` +The OpenEngine runtime dependencies include the headless OpenCV decoder needed +to serve the video capability advertised by video-capable input processors. The +`openengine` extra installs those runtime dependencies, but the sibling +installer remains required to verify and install the pinned `openengine-proto` +package. + The installer rejects a different or dirty sibling package/proto checkout. It prints the required `OPENENGINE_SCHEMA_RELEASE` export; sibling startup fails closed unless that value exactly matches `OPENENGINE_COMMIT`. No registry diff --git a/tensorrt_llm/openengine/server.py b/tensorrt_llm/openengine/server.py index 6876dcbcfccb..2396b95d463e 100644 --- a/tensorrt_llm/openengine/server.py +++ b/tensorrt_llm/openengine/server.py @@ -3,6 +3,7 @@ """Lifecycle wrapper for the optional OpenEngine sibling gRPC server.""" +import importlib import ipaddress import grpc @@ -31,6 +32,29 @@ def validate_schema_release(schema_release: str) -> str: ) +def validate_runtime_dependencies(llm: object) -> None: + """Fail startup when an advertised modality lacks its runtime decoder.""" + processor = getattr(llm, "input_processor", None) + modalities: set[str] = set() + for getter_name in ( + "get_openengine_modalities", + "get_openengine_prefill_decode_modalities", + ): + getter = getattr(processor, getter_name, None) + if callable(getter): + modalities.update(getter()) + if "video" in modalities: + try: + importlib.import_module("cv2") + except ImportError as exc: + raise RuntimeError( + "TensorRT-LLM OpenEngine advertises video input for this model, but " + "OpenCV is not installed. Install the OpenEngine runtime dependencies " + "with `python scripts/install_openengine.py` or " + "`pip install 'tensorrt-llm[openengine]'`." + ) from exc + + class OpenEngineServer: """An OpenEngine gRPC server which never owns or shuts down the LLM.""" @@ -49,6 +73,7 @@ def __init__( stats_fanout: StatsFanout | None = None, ) -> None: validate_schema_release(schema_release()) + validate_runtime_dependencies(llm) self.host = host self.port = port self._kv_event_fanout = kv_event_fanout @@ -134,4 +159,9 @@ def openengine_role(server_role: object | None) -> int: raise ValueError(f"Unknown TRT-LLM server role {name!r}") -__all__ = ["OpenEngineServer", "openengine_role", "validate_schema_release"] +__all__ = [ + "OpenEngineServer", + "openengine_role", + "validate_runtime_dependencies", + "validate_schema_release", +] diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index d7d36e19fff4..abebaf51ed4d 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -4,6 +4,7 @@ from unittest.mock import patch +import numpy as np import pytest import torch @@ -14,6 +15,7 @@ _find_mm_embedding_lengths_from_masks, find_mm_token_lengths, ) +from tensorrt_llm.inputs.multimodal_data import VideoData from tensorrt_llm.inputs.registry import ( create_input_processor_with_hash, maybe_compute_mm_embed_cumsum, @@ -261,6 +263,45 @@ def get_num_tokens_per_video(self, *, video, video_grid_thw=None): ) == {"video": [1, 2]} +def test_find_mm_token_lengths_accepts_video_data_numpy_frames(): + """Qwen3-VL's preferred uint8 HWC frame array remains hashable.""" + frames = np.zeros((2, 8, 10, 3), dtype=np.uint8) + + class FakeVideoProcessor: + def get_num_tokens_per_video( + self, *, video, video_grid_thw=None, video_metadata=None + ): + assert isinstance(video, list) + assert len(video) == 2 + assert all(isinstance(frame, np.ndarray) for frame in video) + assert all(frame.shape == (8, 10, 3) for frame in video) + assert video_metadata == {} + return 16 + + assert find_mm_token_lengths( + {"video": [VideoData(frames=frames, metadata={})]}, + FakeVideoProcessor(), + ) == {"video": [16]} + + +def test_find_mm_token_lengths_preserves_video_data_tensor_frame_lists(): + """The existing list-of-tensors decoder contract remains unchanged.""" + frames = [torch.zeros((3, 8, 10)), torch.ones((3, 8, 10))] + + class FakeVideoProcessor: + def get_num_tokens_per_video( + self, *, video, video_grid_thw=None, video_metadata=None + ): + assert video is frames + assert video_metadata == {} + return 16 + + assert find_mm_token_lengths( + {"video": [VideoData(frames=frames, metadata={})]}, + FakeVideoProcessor(), + ) == {"video": [16]} + + def test_disagg_prefill_multimodal_inputs_builds_typed_handoff(): """Typed EPD handoff converts back to legacy MultimodalInput fields.""" handoff = DisaggPrefillMultimodalInputs( diff --git a/tests/unittest/openengine/test_packaging.py b/tests/unittest/openengine/test_packaging.py new file mode 100644 index 000000000000..dcd1497a64d6 --- /dev/null +++ b/tests/unittest/openengine/test_packaging.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + + +def _requirement(path: Path, name: str) -> str: + for line in path.read_text(encoding="utf-8").splitlines(): + value = line.split("#", 1)[0].strip() + if value.startswith(name): + return value + raise AssertionError(f"{name} is not declared in {path}") + + +def test_openengine_video_dependency_excludes_numpy_incompatible_opencv() -> None: + root = Path(__file__).resolve().parents[3] + numpy = _requirement(root / "requirements.txt", "numpy") + opencv = _requirement(root / "requirements-openengine.txt", "opencv-python-headless") + + assert numpy == "numpy>=2.0.0,<2.4" + assert opencv == "opencv-python-headless==4.11.0.86" diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index fa8d8f3a900a..bd7604e4f6a3 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -6,7 +6,11 @@ import pytest from tensorrt_llm.openengine._schema_pin import OPENENGINE_COMMIT -from tensorrt_llm.openengine.server import OpenEngineServer, validate_schema_release +from tensorrt_llm.openengine.server import ( + OpenEngineServer, + validate_runtime_dependencies, + validate_schema_release, +) from tensorrt_llm.openengine.servicer import schema_release @@ -48,6 +52,43 @@ def test_schema_release_reads_launch_environment(monkeypatch) -> None: assert schema_release() == commit +def test_video_capability_requires_openengine_decoder(monkeypatch) -> None: + class _Processor: + @staticmethod + def get_openengine_modalities() -> tuple[str, ...]: + return ("image", "video") + + @staticmethod + def get_openengine_prefill_decode_modalities() -> tuple[str, ...]: + return ("video",) + + class _Llm: + input_processor = _Processor() + + def missing_cv2(_: str) -> None: + raise ImportError("cv2 is unavailable") + + monkeypatch.setattr("tensorrt_llm.openengine.server.importlib.import_module", missing_cv2) + with pytest.raises(RuntimeError, match="OpenCV is not installed"): + validate_runtime_dependencies(_Llm()) + + +def test_non_video_capabilities_do_not_require_openengine_decoder(monkeypatch) -> None: + class _Processor: + @staticmethod + def get_openengine_modalities() -> tuple[str, ...]: + return ("image", "audio") + + class _Llm: + input_processor = _Processor() + + monkeypatch.setattr( + "tensorrt_llm.openengine.server.importlib.import_module", + lambda _: pytest.fail("cv2 lookup must be capability-gated"), + ) + validate_runtime_dependencies(_Llm()) + + @pytest.mark.asyncio async def test_server_stop_cleans_protocol_state_and_stats_on_grpc_failure() -> None: calls = [] From 3205e6818ab03f2405f792ab67337744ebf13ae6 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 22 Jul 2026 17:18:20 -0700 Subject: [PATCH 15/20] style(openengine): format video validation Signed-off-by: Connor Carpenter --- tensorrt_llm/inputs/multimodal.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py index 2c656cf0c165..a9653a1a7d76 100644 --- a/tensorrt_llm/inputs/multimodal.py +++ b/tensorrt_llm/inputs/multimodal.py @@ -1029,8 +1029,7 @@ def find_mm_token_lengths( if isinstance(item, np.ndarray): if item.ndim != 4 or item.shape[-1] != 3: raise ValueError( - "Video NumPy frames must have shape (N, H, W, 3)" - ) + "Video NumPy frames must have shape (N, H, W, 3)") item = list(item) if not isinstance(item, list): raise TypeError( From 2234f9704a1731fc954e5105f5381e21f6aa4a58 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 22 Jul 2026 17:29:50 -0700 Subject: [PATCH 16/20] style(openengine): format video tests Signed-off-by: Connor Carpenter --- tests/unittest/inputs/test_multimodal.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index abebaf51ed4d..ba52dac265ed 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -268,9 +268,7 @@ def test_find_mm_token_lengths_accepts_video_data_numpy_frames(): frames = np.zeros((2, 8, 10, 3), dtype=np.uint8) class FakeVideoProcessor: - def get_num_tokens_per_video( - self, *, video, video_grid_thw=None, video_metadata=None - ): + def get_num_tokens_per_video(self, *, video, video_grid_thw=None, video_metadata=None): assert isinstance(video, list) assert len(video) == 2 assert all(isinstance(frame, np.ndarray) for frame in video) @@ -289,9 +287,7 @@ def test_find_mm_token_lengths_preserves_video_data_tensor_frame_lists(): frames = [torch.zeros((3, 8, 10)), torch.ones((3, 8, 10))] class FakeVideoProcessor: - def get_num_tokens_per_video( - self, *, video, video_grid_thw=None, video_metadata=None - ): + def get_num_tokens_per_video(self, *, video, video_grid_thw=None, video_metadata=None): assert video is frames assert video_metadata == {} return 16 From 276675522687441041f6b9b7581813eaecb47ea5 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 23 Jul 2026 17:30:40 -0700 Subject: [PATCH 17/20] fix(torch): preserve attention-dp mapping for tied embeddings Signed-off-by: Connor Carpenter --- tensorrt_llm/_torch/models/modeling_utils.py | 1 + .../modeling/test_modeling_qwen3_embedding.py | 36 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 237ed00af403..6028e14a1d52 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -392,6 +392,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], vocab_size, hidden_size, dtype=config.pretrained_config.torch_dtype, + mapping=config.mapping, ) else: if (hasattr(config, 'lora_config') diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3_embedding.py b/tests/unittest/_torch/modeling/test_modeling_qwen3_embedding.py index f65296f9fe0d..10bec2cf24e5 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3_embedding.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3_embedding.py @@ -18,6 +18,16 @@ model so it routes through the llm.encode() path. """ +from types import SimpleNamespace + +import torch +from torch import nn + +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm.mapping import Mapping + def test_qwen3_text_embedding_is_registered(): # Importing the module triggers @register_auto_model. @@ -33,3 +43,29 @@ def test_qwen3_text_embedding_is_not_generation(): assert ModelConfig.is_generation_model(["Qwen3ForTextEmbedding"]) is False # Sanity: the plain causal LM is still a generation model. assert ModelConfig.is_generation_model(["Qwen3ForCausalLM"]) is True + + +def test_tied_embeddings_use_attention_dp_mapping(): + mapping = Mapping(world_size=2, rank=0, tp_size=2, enable_attention_dp=True) + model = nn.Module() + model.embed_tokens = Embedding( + num_embeddings=16, + embedding_dim=8, + dtype=torch.float32, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) + config = SimpleNamespace( + mapping=mapping, + pretrained_config=SimpleNamespace(torch_dtype=torch.float32, tie_word_embeddings=True), + quant_config_dict=None, + quant_config=None, + lm_head_gather_output=True, + ) + + causal_lm = DecoderModelForCausalLM(model, config=config, hidden_size=8, vocab_size=16) + + assert causal_lm.lm_head.tp_size == model.embed_tokens.tp_size == 2 + assert causal_lm.lm_head.tp_mode == model.embed_tokens.tp_mode + assert causal_lm.lm_head.weight is model.embed_tokens.weight From d0eecfbfb59ad9b422c3ce2c5074458f87a489c8 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 23 Jul 2026 17:48:55 -0700 Subject: [PATCH 18/20] fix(openengine): expose attention-dp ranks Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/servicer.py | 14 ++++++++---- tests/unittest/openengine/test_servicer.py | 26 ++++++++++++++++++++++ 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index ac43b471cfe6..bc905694af08 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -72,6 +72,12 @@ def _arg(llm: object, name: str, default: Any = None) -> Any: return getattr(parallel, name, default) +def _data_parallel_size(llm: object) -> int: + if bool(_arg(llm, "enable_attention_dp", False)): + return max(1, int(_arg(llm, "tensor_parallel_size", _arg(llm, "tp_size", 1)))) + return max(1, int(_arg(llm, "data_parallel_size", 1))) + + def _token_text(llm: object, token_id: int) -> str: tokenizer = getattr(llm, "tokenizer", None) if tokenizer is None: @@ -698,7 +704,7 @@ def _scheduling_params(self, target_dp_rank: int | None) -> SchedulingParams | N if target_dp_rank is None: return None rank = target_dp_rank - data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) + data_parallel_size = _data_parallel_size(self.llm) if rank >= data_parallel_size: raise ValueError( f"data_parallel_rank {rank} is outside the configured DP size {data_parallel_size}" @@ -850,7 +856,7 @@ async def GetServerInfo( del request, context tp = _arg(self.llm, "tensor_parallel_size", 1) pp = _arg(self.llm, "pipeline_parallel_size", 1) - dp = _arg(self.llm, "data_parallel_size", 1) + dp = _data_parallel_size(self.llm) args = getattr(self.llm, "args", None) kv_capacity = self._kv_capacity() capacity = server_pb2.DeploymentCapacity() @@ -1269,7 +1275,7 @@ async def GetKvEventSources( ) -> kv_pb2.GetKvEventSourcesResponse: if not self._kv_events_enabled(): return kv_pb2.GetKvEventSourcesResponse() - data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) + data_parallel_size = _data_parallel_size(self.llm) available_ranks = set(range(data_parallel_size)) requested_ranks = set(request.data_parallel_ranks) invalid_ranks = requested_ranks - available_ranks @@ -1311,7 +1317,7 @@ async def SubscribeKvEvents( ) return selected_ranks = set(request.data_parallel_ranks) - available_ranks = set(range(int(_arg(self.llm, "data_parallel_size", 1)))) + available_ranks = set(range(_data_parallel_size(self.llm))) invalid_ranks = selected_ranks - available_ranks if invalid_ranks: await context.abort( diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index ed47709e91b5..88a7a2da8bde 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -899,6 +899,32 @@ async def test_kv_event_sources_are_rank_scoped() -> None: assert [source.data_parallel_rank for source in response.sources] == [0, 1] +@pytest.mark.asyncio +async def test_attention_dp_uses_tensor_parallel_size_for_discovery_and_placement() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + enable_attention_dp=True, + tensor_parallel_size=2, + data_parallel_size=1, + kv_cache_config=SimpleNamespace(event_buffer_max_size=8), + ) + llm._on_trt_backend = False + service = OpenEngineServicer( + llm, + "model", + server_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=KvEventFanout(llm, buffer_size=8), + ) + + info = await service.GetServerInfo(server_pb2.GetServerInfoRequest(), _Context()) + sources = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + + assert info.parallelism.data_parallel_size == 2 + assert [source.data_parallel_rank for source in sources.sources] == [0, 1] + assert service._scheduling_params(1).attention_dp_rank == 1 + + @pytest.mark.asyncio async def test_kv_event_discovery_requires_decimal_compatible_hashes() -> None: llm = _Llm() From 918e39f81370314ead81128d53d3aa3919edd0a1 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 23 Jul 2026 17:57:38 -0700 Subject: [PATCH 19/20] fix(openengine): retain explicit attention-dp size Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/servicer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index bc905694af08..f79e942a1892 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -74,7 +74,12 @@ def _arg(llm: object, name: str, default: Any = None) -> Any: def _data_parallel_size(llm: object) -> int: if bool(_arg(llm, "enable_attention_dp", False)): - return max(1, int(_arg(llm, "tensor_parallel_size", _arg(llm, "tp_size", 1)))) + size = _arg(llm, "tensor_parallel_size", None) + if size is None: + size = _arg(llm, "tp_size", None) + if size is None: + size = _arg(llm, "data_parallel_size", 1) + return max(1, int(size)) return max(1, int(_arg(llm, "data_parallel_size", 1))) From 58132ce1f654bfede8e530f78e31979596c5c34e Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 23 Jul 2026 23:45:37 -0700 Subject: [PATCH 20/20] build(openengine): update protocol pin Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 2 +- tensorrt_llm/openengine/_schema_pin.py | 2 +- tensorrt_llm/openengine/converters.py | 10 ++--- tensorrt_llm/openengine/servicer.py | 44 ++++++++------------ tests/unittest/openengine/test_converters.py | 8 ++-- tests/unittest/openengine/test_server.py | 2 +- tests/unittest/openengine/test_servicer.py | 31 +++++--------- 7 files changed, 40 insertions(+), 59 deletions(-) diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT index c15702493507..104fa5afd2a1 100644 --- a/OPENENGINE_COMMIT +++ b/OPENENGINE_COMMIT @@ -1 +1 @@ -b0cf2a4826d246192dc65b055dab6d2b38d2d67e +a66ff6f73a65e262a7c3edd5ea6fd0d8701d402f diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py index 3c8565fafa54..f99d123225b6 100644 --- a/tensorrt_llm/openengine/_schema_pin.py +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -3,6 +3,6 @@ """OpenEngine source identity embedded in TensorRT-LLM distributions.""" -OPENENGINE_COMMIT = "b0cf2a4826d246192dc65b055dab6d2b38d2d67e" +OPENENGINE_COMMIT = "a66ff6f73a65e262a7c3edd5ea6fd0d8701d402f" __all__ = ["OPENENGINE_COMMIT"] diff --git a/tensorrt_llm/openengine/converters.py b/tensorrt_llm/openengine/converters.py index 51d0e4639b24..2f1f24c9875e 100644 --- a/tensorrt_llm/openengine/converters.py +++ b/tensorrt_llm/openengine/converters.py @@ -12,7 +12,7 @@ from typing import Any from google.protobuf.json_format import MessageToDict -from openengine.v1 import generation_pb2, input_pb2, kv_pb2 +from openengine.v1 import generation_pb2, kv_pb2 from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from tensorrt_llm.executor.request import LoRARequest @@ -23,10 +23,10 @@ HANDOFF_ATTRIBUTE = "tensorrt_llm.disaggregated_params.v1" _MODALITY_NAMES = { - input_pb2.MODALITY_UNSPECIFIED: "image", - input_pb2.MODALITY_IMAGE: "image", - input_pb2.MODALITY_VIDEO: "video", - input_pb2.MODALITY_AUDIO: "audio", + generation_pb2.MODALITY_UNSPECIFIED: "image", + generation_pb2.MODALITY_IMAGE: "image", + generation_pb2.MODALITY_VIDEO: "video", + generation_pb2.MODALITY_AUDIO: "audio", } diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index f79e942a1892..54d0cdc86903 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -18,12 +18,10 @@ from openengine.v1 import ( error_pb2, generation_pb2, - input_pb2, kv_pb2, lifecycle_pb2, lora_pb2, model_pb2, - observability_pb2, openengine_pb2_grpc, server_pb2, ) @@ -649,10 +647,10 @@ def _validate_generate( aggregate = self._available_modalities(processor) prefill_decode = self._available_modalities(processor, prefill_decode=True) names = { - input_pb2.MODALITY_UNSPECIFIED: "image", - input_pb2.MODALITY_IMAGE: "image", - input_pb2.MODALITY_VIDEO: "video", - input_pb2.MODALITY_AUDIO: "audio", + generation_pb2.MODALITY_UNSPECIFIED: "image", + generation_pb2.MODALITY_IMAGE: "image", + generation_pb2.MODALITY_VIDEO: "video", + generation_pb2.MODALITY_AUDIO: "audio", } allowed = ( aggregate if self.role == server_pb2.ENGINE_ROLE_AGGREGATED else prefill_decode @@ -733,10 +731,10 @@ def _required_multimodal_lora( if not isinstance(processor, BaseMultimodalInputProcessor): return None names = { - input_pb2.MODALITY_UNSPECIFIED: "image", - input_pb2.MODALITY_IMAGE: "image", - input_pb2.MODALITY_VIDEO: "video", - input_pb2.MODALITY_AUDIO: "audio", + generation_pb2.MODALITY_UNSPECIFIED: "image", + generation_pb2.MODALITY_IMAGE: "image", + generation_pb2.MODALITY_VIDEO: "video", + generation_pb2.MODALITY_AUDIO: "audio", } modalities = tuple( dict.fromkeys(names[item.modality] for item in request.media if item.modality in names) @@ -907,9 +905,9 @@ async def GetModelInfo( aggregate = self._available_modalities(input_processor) prefill_decode = self._available_modalities(input_processor, prefill_decode=True) modality_enum = { - "image": input_pb2.MODALITY_IMAGE, - "video": input_pb2.MODALITY_VIDEO, - "audio": input_pb2.MODALITY_AUDIO, + "image": generation_pb2.MODALITY_IMAGE, + "video": generation_pb2.MODALITY_VIDEO, + "audio": generation_pb2.MODALITY_AUDIO, } args = getattr(self.llm, "args", None) guided_backend = getattr(args, "guided_decoding_backend", None) @@ -965,9 +963,9 @@ async def GetModelInfo( aggregate_modalities=[modality_enum[name] for name in aggregate], prefill_decode_modalities=[modality_enum[name] for name in prefill_decode], source_types=[ - input_pb2.MEDIA_SOURCE_TYPE_URL, - input_pb2.MEDIA_SOURCE_TYPE_DATA_URI, - input_pb2.MEDIA_SOURCE_TYPE_RAW_BYTES, + generation_pb2.MEDIA_SOURCE_TYPE_URL, + generation_pb2.MEDIA_SOURCE_TYPE_DATA_URI, + generation_pb2.MEDIA_SOURCE_TYPE_RAW_BYTES, ], supports_per_request_media_options=True, ), @@ -1003,8 +1001,8 @@ def _kv_capacity(self) -> dict[str, int]: return capacity async def GetLoad( - self, request: observability_pb2.GetLoadRequest, context: grpc.aio.ServicerContext - ) -> observability_pb2.LoadInfo: + self, request: server_pb2.GetLoadRequest, context: grpc.aio.ServicerContext + ) -> server_pb2.LoadInfo: del context latest = self.stats.latest_by_rank() if self.stats is not None else {} rank_infos = [] @@ -1047,7 +1045,7 @@ async def GetLoad( decode_batch_size += rank_decode if request.include_per_rank: rank_infos.append( - observability_pb2.RankLoadInfo( + server_pb2.RankLoadInfo( data_parallel_rank=rank, running_requests=rank_running, queued_requests=rank_queued, @@ -1066,7 +1064,7 @@ async def GetLoad( tokens_per_block.add(capacity["tokensPerBlock"]) if len(tokens_per_block) == 1: attributes["kv_tokens_per_block"] = str(next(iter(tokens_per_block))) - response = observability_pb2.LoadInfo( + response = server_pb2.LoadInfo( instance_id=self.instance_id, timestamp_unix_nanos=time.time_ns(), running_requests=running_requests, @@ -1269,12 +1267,6 @@ def _kv_events_enabled(self) -> bool: ) return effective_hash_algo in (KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2_SHA256_64) - async def GetKvConnectorInfo( - self, request: kv_pb2.GetKvConnectorInfoRequest, context: grpc.aio.ServicerContext - ) -> kv_pb2.KvConnectorInfo: - del request, context - return self._kv_connector_info() - async def GetKvEventSources( self, request: kv_pb2.GetKvEventSourcesRequest, context: grpc.aio.ServicerContext ) -> kv_pb2.GetKvEventSourcesResponse: diff --git a/tests/unittest/openengine/test_converters.py b/tests/unittest/openengine/test_converters.py index 115dbbff92c9..b9d09a3a1b4e 100644 --- a/tests/unittest/openengine/test_converters.py +++ b/tests/unittest/openengine/test_converters.py @@ -4,7 +4,7 @@ import json import pytest -from openengine.v1 import generation_pb2, input_pb2, kv_pb2 +from openengine.v1 import generation_pb2, kv_pb2 from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from tensorrt_llm.executor.result import Logprob @@ -218,9 +218,9 @@ def load_bytes(self, data): options.update({"image": {"format": "pil"}}) config = type("Config", (), {"media_io_kwargs": {"image": {"device": "cpu"}}})() media = [ - input_pb2.MediaItem(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"one"), - input_pb2.MediaItem(modality=input_pb2.MODALITY_VIDEO, raw_bytes=b"video"), - input_pb2.MediaItem(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"two"), + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"one"), + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_VIDEO, raw_bytes=b"video"), + generation_pb2.MediaItem(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"two"), ] decoded = await load_media(media, options, config) diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py index bd7604e4f6a3..f392c6ff1c2b 100644 --- a/tests/unittest/openengine/test_server.py +++ b/tests/unittest/openengine/test_server.py @@ -47,7 +47,7 @@ def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> Non def test_schema_release_reads_launch_environment(monkeypatch) -> None: - commit = "b0cf2a4826d246192dc65b055dab6d2b38d2d67e" + commit = "a66ff6f73a65e262a7c3edd5ea6fd0d8701d402f" monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) assert schema_release() == commit diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index 88a7a2da8bde..0bb3627294bc 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -6,16 +6,7 @@ from types import SimpleNamespace import pytest -from openengine.v1 import ( - generation_pb2, - input_pb2, - kv_pb2, - lifecycle_pb2, - lora_pb2, - model_pb2, - observability_pb2, - server_pb2, -) +from openengine.v1 import generation_pb2, kv_pb2, lifecycle_pb2, lora_pb2, model_pb2, server_pb2 from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from tensorrt_llm.inputs.registry import MultimodalLoraSpec @@ -300,7 +291,7 @@ async def _load_media(*_args, **_kwargs): ) request = generation_pb2.GenerateRequest(request_id="audio", model="model", prompt="transcribe") request.media.add( - modality=input_pb2.MODALITY_AUDIO, + modality=generation_pb2.MODALITY_AUDIO, raw_bytes=b"audio", uuid="0123456789abcdef", ) @@ -368,7 +359,7 @@ def capture_info(message: str, *args: object) -> None: context_request = generation_pb2.GenerateRequest( request_id="request", model="model", - token_ids=input_pb2.TokenIds(ids=[1, 2]), + token_ids=generation_pb2.TokenIds(ids=[1, 2]), ) context_responses = [ response async for response in context_service.Generate(context_request, _Context()) @@ -376,7 +367,7 @@ def capture_info(message: str, *args: object) -> None: assert len(context_responses) == 1 assert context_responses[0].WhichOneof("event") == "prefill_ready" assert len(context_service._kv_session_requests) == 1 - context_load = await context_service.GetLoad(observability_pb2.GetLoadRequest(), _Context()) + context_load = await context_service.GetLoad(server_pb2.GetLoadRequest(), _Context()) assert context_load.active_kv_sessions == 0 decode_llm = _Llm() @@ -389,7 +380,7 @@ def capture_info(message: str, *args: object) -> None: decode_request = generation_pb2.GenerateRequest( request_id="decode", model="model", - token_ids=input_pb2.TokenIds(ids=[1, 2]), + token_ids=generation_pb2.TokenIds(ids=[1, 2]), ) decode_request.kv.session.CopyFrom(context_responses[0].prefill_ready.kv_session) decode_responses = [ @@ -456,7 +447,7 @@ async def test_prefill_adds_context_endpoint_from_llm_discovery() -> None: request = generation_pb2.GenerateRequest( request_id="prefill", model="model", - token_ids=input_pb2.TokenIds(ids=[1, 2]), + token_ids=generation_pb2.TokenIds(ids=[1, 2]), ) responses = [response async for response in service.Generate(request, _Context())] @@ -588,14 +579,14 @@ def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraS requires_decode_media=True, ) request = generation_pb2.GenerateRequest( - request_id="decode", model="model", token_ids=input_pb2.TokenIds(ids=[1, 2]) + request_id="decode", model="model", token_ids=generation_pb2.TokenIds(ids=[1, 2]) ) request.kv.session.CopyFrom(session) with pytest.raises(ValueError, match="must resend"): service._validate_generate(request) - request.media.add(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"image") + request.media.add(modality=generation_pb2.MODALITY_IMAGE, raw_bytes=b"image") service._validate_generate(request) @@ -835,9 +826,7 @@ def get_required_lora_spec(self, modalities: tuple[str, ...]) -> None: assert not model_info.generation.guided_decoding.supported assert not model_info.supports_lora - load = await service.GetLoad( - observability_pb2.GetLoadRequest(include_per_rank=True), _Context() - ) + load = await service.GetLoad(server_pb2.GetLoadRequest(include_per_rank=True), _Context()) assert load.running_requests == 3 assert load.queued_requests == 2 assert load.used_kv_blocks == 75 @@ -873,7 +862,7 @@ async def test_load_never_undercounts_live_tracker_from_stale_stats() -> None: stats_fanout=stats, ) - load = await service.GetLoad(observability_pb2.GetLoadRequest(), _Context()) + load = await service.GetLoad(server_pb2.GetLoadRequest(), _Context()) assert load.running_requests == 1