diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1652e81c5853..e7ce18708a8f 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -478,6 +478,7 @@ /jenkins/license_cpp.json @NVIDIA/trt-llm-infra-devs @NVIDIA/trt-llm-oss-compliance /pyproject.toml @NVIDIA/trt-llm-oss-compliance /requirements-dev.txt @NVIDIA/trt-llm-oss-compliance +/requirements-grpc-smg.txt @NVIDIA/trt-llm-oss-compliance /requirements.txt @NVIDIA/trt-llm-oss-compliance /setup.py @NVIDIA/trt-llm-oss-compliance /tests/unittest/api_stability/ @NVIDIA/trt-llm-noncommitted-api-review-committee diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 828859395043..838c731e9a17 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -110,7 +110,7 @@ COPY scripts scripts COPY tensorrt_llm tensorrt_llm COPY triton_kernels triton_kernels COPY 3rdparty 3rdparty -COPY .gitmodules setup.py requirements.txt requirements-dev.txt constraints.txt README.md ./ +COPY .gitmodules setup.py requirements.txt requirements-dev.txt requirements-grpc-smg.txt constraints.txt README.md ./ ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index a3a8a6910166..291bf59a7fb1 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -4227,6 +4227,20 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO sh "cd ${llmSrc} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" } trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmSrc} && pip3 install -r requirements-dev.txt") + // Gateway adapters are opt-in extras excluded from requirements.txt; + // each gateway declares its pins in a dedicated + // requirements-.txt, and a test stage installs exactly + // zero or one gateway file so every adapter is tested under the + // dependency set its real opt-in users receive. A gateway whose + // pins co-resolve with the default environment (SMG today) is + // installed in the shared stages so its unit tests run from the + // regular shard pool instead of being skipped at collection; a + // gateway whose pins conflict with the default environment (for + // example a protobuf major-version floor or a custom package + // index) must instead install its file behind a dedicated stage + // guard and skip this one (see the Ray install below for the + // stage-scoped pattern). + trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmSrc} && pip3 install -r requirements-grpc-smg.txt") trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install opencv-python-headless") if (stageName.contains("-Ray-")) { trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install ray[default]==2.55.1") @@ -4764,6 +4778,7 @@ def runLLMBuild( } trtllm_utils.llmExecStepWithRetry(pipeline, script: "#!/bin/bash \n" + "cd tensorrt_llm/ && pip3 install -r requirements-dev.txt") + trtllm_utils.llmExecStepWithRetry(pipeline, script: "#!/bin/bash \n" + "cd tensorrt_llm/ && pip3 install -r requirements-grpc-smg.txt") if (env.alternativeTRT) { trtllm_utils.replaceWithAlternativeTRT(env.alternativeTRT, cpver) } diff --git a/jenkins/scripts/slurm_install.sh b/jenkins/scripts/slurm_install.sh index 27c591760f5c..91ca8145bec6 100644 --- a/jenkins/scripts/slurm_install.sh +++ b/jenkins/scripts/slurm_install.sh @@ -31,6 +31,7 @@ slurm_install_setup() { fi retry_command --timeout 2700 bash -c "pip3 install --retries 10 opencv-python-headless" retry_command --timeout 2700 bash -c "cd $llmSrcNode && pip3 install --retries 10 -r requirements-dev.txt" + retry_command --timeout 2700 bash -c "cd $llmSrcNode && pip3 install --retries 10 -r requirements-grpc-smg.txt" retry_command --timeout 2700 bash -c "cd $resourcePathNode && pip3 install --retries 10 --force-reinstall --no-deps TensorRT-LLM/tensorrt_llm-*.whl" gpuUuids=$(nvidia-smi -q | grep "GPU UUID" | awk '{print $4}' | tr '\n' ',' || true) hostNodeName="${HOST_NODE_NAME:-$(hostname -f || hostname)}" diff --git a/requirements-grpc-smg.txt b/requirements-grpc-smg.txt new file mode 100644 index 000000000000..94c9efb9fcd0 --- /dev/null +++ b/requirements-grpc-smg.txt @@ -0,0 +1 @@ +smg-grpc-proto>=0.4.2 diff --git a/requirements.txt b/requirements.txt index d63bc1a93be2..24a2ca1f3fc3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -92,8 +92,9 @@ llist cuda-tile>=1.0.1 nvidia-cuda-tileiras>=13.1,<13.2 etcd-sdk-python==0.0.7 +# etcd-sdk-python imports google.protobuf but omits it from its package metadata. +protobuf>=5.27.2 python-multipart -smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 librosa msgpack diff --git a/setup.py b/setup.py index deaeacc793e1..d4cac2103ed7 100644 --- a/setup.py +++ b/setup.py @@ -141,6 +141,13 @@ def has_ext_modules(self): Path("requirements-dev-windows.txt" if on_windows else "requirements-dev.txt")) mx_deps = ["modelexpress==0.4.1"] +# Gateway protocol adapters are opt-in extras: the default installation must +# not carry any gateway protobuf package. Each gateway owns a dedicated +# requirements-.txt as the single source of truth for its pins; CI +# stages that exercise a gateway install that file explicitly, and the file +# may carry gateway-specific options (such as an --extra-index-url) without +# affecting the default dependency graph. +grpc_smg_deps, _ = parse_requirements(Path("requirements-grpc-smg.txt")) constraints_file = Path("constraints.txt") if constraints_file.exists(): constraints, _ = parse_requirements(constraints_file) @@ -484,8 +491,9 @@ def extract_from_precompiled(precompiled_location: str, package_data: list[str], }, scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'], extras_require={ - "devel": devel_deps, + "devel": devel_deps + grpc_smg_deps, "mx": mx_deps, + "grpc-smg": grpc_smg_deps, }, zip_safe=True, install_requires=required_deps, diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index e5fc8cdb9f75..212e85344e6d 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -13,6 +13,7 @@ import sys import time import uuid +from importlib.util import find_spec from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Sequence, Set @@ -60,9 +61,6 @@ # Global variable to store the Popen object of the child process _child_p_global: Optional[subprocess.Popen] = None -# Bound gRPC messages while leaving room for multimodal image payloads. -_GRPC_MAX_MESSAGE_LENGTH_BYTES = 32 * 1024 * 1024 - def _pop_bool_config_option(config: dict[str, Any], key: str) -> bool: return validate_config_bool(config.pop(key, False), key) @@ -637,129 +635,6 @@ def launch_server( _terminate_attached_frontends(frontend_children) -def launch_grpc_server(host: str, - port: int, - llm_args: dict, - served_model_name: Optional[str] = None): - """ - 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. - - Args: - host: Host to bind to - port: Port to bind to - llm_args: Arguments for LLM initialization (from get_llm_args) - served_model_name: Custom model name for API responses (defaults to model path) - """ - import grpc - - try: - from grpc_reflection.v1alpha import reflection - REFLECTION_AVAILABLE = True - except ImportError: - REFLECTION_AVAILABLE = False - - from tensorrt_llm.grpc import trtllm_service_pb2, trtllm_service_pb2_grpc - from tensorrt_llm.grpc.grpc_request_manager import GrpcRequestManager - from tensorrt_llm.grpc.grpc_servicer import TrtllmServiceServicer - - async def serve_grpc_async(): - logger.info("Initializing TensorRT-LLM gRPC server...") - - backend = llm_args.get("backend") - model_path = served_model_name or llm_args.get("model", "") - - if backend == "pytorch": - llm_args.pop("build_config", None) - llm = PyTorchLLM(**llm_args) - elif backend == "_autodeploy": - from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM - llm_args.pop("build_config", None) - llm = AutoDeployLLM(**llm_args) - else: - raise click.BadParameter( - f"{backend} is not a known backend, check help for available options.", - param_hint="backend") - - logger.info("Model loaded successfully") - - # Create request manager - request_manager = GrpcRequestManager(llm) - - # Create servicer - servicer = TrtllmServiceServicer(request_manager, model_path=model_path) - - # Create gRPC server - server = grpc.aio.server( - options=[ - ("grpc.max_send_message_length", - _GRPC_MAX_MESSAGE_LENGTH_BYTES), - ("grpc.max_receive_message_length", - _GRPC_MAX_MESSAGE_LENGTH_BYTES), - ("grpc.keepalive_time_ms", 30000), # 30s keepalive - ("grpc.keepalive_timeout_ms", 10000), # 10s timeout - ("grpc.keepalive_permit_without_calls", True), - ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000), - ], ) - - # Add servicer to server - trtllm_service_pb2_grpc.add_TrtllmServiceServicer_to_server( - servicer, server) - - # Enable reflection for grpcurl and other tools - if REFLECTION_AVAILABLE: - service_names = ( - trtllm_service_pb2.DESCRIPTOR.services_by_name["TrtllmService"]. - full_name, - reflection.SERVICE_NAME, - ) - reflection.enable_server_reflection(service_names, server) - logger.info("gRPC reflection enabled") - - # Bind to address - address = f"{host}:{port}" - server.add_insecure_port(address) - - # Start server - await server.start() - logger.info(f"TensorRT-LLM gRPC server started on {address}") - logger.info("Server is ready to accept requests") - - # Handle shutdown signals - loop = asyncio.get_running_loop() - stop_event = asyncio.Event() - - def signal_handler(): - logger.info("Received shutdown signal") - stop_event.set() - - for sig in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(sig, signal_handler) - - # Serve until shutdown signal - try: - await stop_event.wait() - except KeyboardInterrupt: - logger.info("Interrupted by user") - finally: - logger.info("Shutting down TensorRT-LLM gRPC server...") - - # Stop gRPC server - await server.stop(grace=5.0) - logger.info("gRPC server stopped") - - # Shutdown LLM - if hasattr(llm, "shutdown"): - llm.shutdown() - logger.info("LLM engine stopped") - - logger.info("Shutdown complete") - - uvloop.run(serve_grpc_async()) - - def launch_mm_encoder_server( host: str, port: int, @@ -1252,7 +1127,8 @@ def launch_visual_gen_server( is_flag=True, default=False, help="Run gRPC server instead of OpenAI HTTP server. " - "gRPC server accepts pre-tokenized requests and returns raw token IDs.", + "gRPC server accepts pre-tokenized requests and returns raw token IDs. " + "Requires the tensorrt_llm[grpc-smg] extra.", status="prototype") @stability_option( "--served_model_name", @@ -1494,10 +1370,18 @@ def _serve_llm(): f"Argument '{name}' is not supported when running in gRPC mode. " f"The gRPC server is designed for use with external routers that handle " f"these features (e.g., tool parsing, chat templates).") - launch_grpc_server(host, - port, - llm_args, - served_model_name=served_model_name) + if find_spec("smg_grpc_proto") is None: + raise ValueError( + "gRPC serving with the SMG protocol requires the optional " + "'smg-grpc-proto' package. Install it with: " + 'pip install "tensorrt_llm[grpc-smg]"') + + from tensorrt_llm.grpc.smg.server import launch_smg_server + + launch_smg_server(host, + port, + llm_args, + served_model_name=served_model_name) else: # Default: launch OpenAI HTTP server launch_server( diff --git a/tensorrt_llm/grpc/__init__.py b/tensorrt_llm/grpc/__init__.py index d75315bd0605..4900715bc879 100644 --- a/tensorrt_llm/grpc/__init__.py +++ b/tensorrt_llm/grpc/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,75 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -r"""TensorRT-LLM gRPC module for high-performance communication with external routers. - -This module provides a gRPC server interface that accepts pre-tokenized requests -and returns raw token IDs, enabling efficient binary communication with Rust-based -routers like sgl-router. - -Key Features: -- Pre-tokenized input (no Python tokenization overhead) -- Raw token ID output (no Python detokenization overhead) -- Streaming support with delta tokens -- Full sampling parameter support -- Guided decoding (JSON schema, regex, grammar) -- LoRA and prompt tuning support -- Disaggregated inference support - -Proto definitions are provided by the smg-grpc-proto package (pip install smg-grpc-proto). - -Usage: - python -m tensorrt_llm.commands.serve /path/to/model \ - --grpc \ - --host 0.0.0.0 \ - --port 50051 -""" - -# Try to import generated protobuf modules from smg-grpc-proto package -try: - from smg_grpc_proto.generated import trtllm_service_pb2, trtllm_service_pb2_grpc - - PROTOS_AVAILABLE = True -except ImportError: - PROTOS_AVAILABLE = False - trtllm_service_pb2 = None - trtllm_service_pb2_grpc = None - -# Try to import request manager -try: - from .grpc_request_manager import ( - GrpcRequestManager, - create_disaggregated_params_from_proto, - create_lora_request_from_proto, - create_sampling_params_from_proto, - ) - - REQUEST_MANAGER_AVAILABLE = True -except ImportError: - REQUEST_MANAGER_AVAILABLE = False - GrpcRequestManager = None - create_sampling_params_from_proto = None - create_lora_request_from_proto = None - create_disaggregated_params_from_proto = None - -# Try to import servicer -try: - from .grpc_servicer import TrtllmServiceServicer - - SERVICER_AVAILABLE = True -except ImportError: - SERVICER_AVAILABLE = False - TrtllmServiceServicer = None - -__all__ = [ - "PROTOS_AVAILABLE", - "REQUEST_MANAGER_AVAILABLE", - "SERVICER_AVAILABLE", - "trtllm_service_pb2", - "trtllm_service_pb2_grpc", - "GrpcRequestManager", - "TrtllmServiceServicer", - "create_sampling_params_from_proto", - "create_lora_request_from_proto", - "create_disaggregated_params_from_proto", -] +"""gRPC protocol integrations for TensorRT-LLM.""" diff --git a/tensorrt_llm/grpc/smg/__init__.py b/tensorrt_llm/grpc/smg/__init__.py new file mode 100644 index 000000000000..b78fe0dff49d --- /dev/null +++ b/tensorrt_llm/grpc/smg/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SMG gRPC protocol integration for TensorRT-LLM.""" diff --git a/tensorrt_llm/grpc/smg/bindings.py b/tensorrt_llm/grpc/smg/bindings.py new file mode 100644 index 000000000000..3bc24ad0e87c --- /dev/null +++ b/tensorrt_llm/grpc/smg/bindings.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generated protobuf bindings used by the SMG adapter.""" + +try: + from smg_grpc_proto.generated import trtllm_service_pb2, trtllm_service_pb2_grpc +except ModuleNotFoundError as e: + if e.name != "smg_grpc_proto": + raise + raise ModuleNotFoundError( + "The SMG gRPC adapter requires the optional 'smg-grpc-proto' package, " + "which is not part of the default TensorRT-LLM installation. Install it " + 'with: pip install "tensorrt_llm[grpc-smg]"', + name=e.name, + ) from e + +__all__ = ["trtllm_service_pb2", "trtllm_service_pb2_grpc"] diff --git a/tensorrt_llm/grpc/grpc_request_manager.py b/tensorrt_llm/grpc/smg/request_manager.py similarity index 95% rename from tensorrt_llm/grpc/grpc_request_manager.py rename to tensorrt_llm/grpc/smg/request_manager.py index ca2e59eb1f36..40e7ec09f822 100644 --- a/tensorrt_llm/grpc/grpc_request_manager.py +++ b/tensorrt_llm/grpc/smg/request_manager.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""gRPC Request Manager for TensorRT-LLM. +"""Request manager for the TensorRT-LLM SMG gRPC adapter. Manages request lifecycle for gRPC requests, converting between protobuf and TensorRT-LLM types. Designed for high-performance communication with @@ -35,7 +35,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams -from . import trtllm_service_pb2 as pb2 +from .bindings import trtllm_service_pb2 as pb2 class GrpcRequestManager: @@ -211,8 +211,18 @@ def get_model_config(self) -> Dict[str, Any]: # Try to get tokenizer info if hasattr(self.llm, "tokenizer") and self.llm.tokenizer is not None: - if hasattr(self.llm.tokenizer, "vocab_size"): - config["vocab_size"] = self.llm.tokenizer.vocab_size + # TransformersTokenizer wraps the HF tokenizer without + # overriding vocab_size, so the wrapper attribute resolves to + # PreTrainedTokenizerBase's abstract property (raises + # NotImplementedError). Read from the inner tokenizer. + tokenizer = self.llm.tokenizer + inner_tokenizer = getattr(tokenizer, "tokenizer", tokenizer) + try: + vocab_size = inner_tokenizer.vocab_size + except (AttributeError, NotImplementedError): + vocab_size = None + if vocab_size: + config["vocab_size"] = int(vocab_size) # Try to get max context length from various sources if hasattr(self.llm, "args") and self.llm.args is not None: diff --git a/tensorrt_llm/grpc/smg/server.py b/tensorrt_llm/grpc/smg/server.py new file mode 100644 index 000000000000..608baa47428f --- /dev/null +++ b/tensorrt_llm/grpc/smg/server.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Server lifecycle for the TensorRT-LLM SMG gRPC adapter.""" + +import asyncio +import signal +from typing import Any + +import click +import grpc +import uvloop + +from tensorrt_llm import LLM as PyTorchLLM +from tensorrt_llm.logger import logger + +from .bindings import trtllm_service_pb2, trtllm_service_pb2_grpc +from .request_manager import GrpcRequestManager +from .servicer import TrtllmServiceServicer + +_GRPC_MAX_MESSAGE_LENGTH_BYTES = 32 * 1024 * 1024 + + +def launch_smg_server( + host: str, + port: int, + llm_args: dict[str, Any], + served_model_name: str | None = None, +) -> None: + """Launch the SMG gRPC server. + + Args: + host: Host to bind to. + port: Port to bind to. + llm_args: Arguments for LLM initialization. + served_model_name: Model name returned by discovery RPCs. Defaults to + the model path. + """ + try: + from grpc_reflection.v1alpha import reflection + except ModuleNotFoundError as e: + if e.name != "grpc_reflection": + raise + reflection = None + + async def serve_grpc_async() -> None: + logger.info("Initializing TensorRT-LLM SMG gRPC server...") + + backend = llm_args.get("backend") + model_path = served_model_name or llm_args.get("model", "") + + if backend == "pytorch": + llm_args.pop("build_config", None) + llm = PyTorchLLM(**llm_args) + elif backend == "_autodeploy": + from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM + + llm_args.pop("build_config", None) + llm = AutoDeployLLM(**llm_args) + else: + raise click.BadParameter( + f"{backend} is not a known backend, check help for available options.", + param_hint="backend", + ) + + logger.info("Model loaded successfully") + + server = None + try: + request_manager = GrpcRequestManager(llm) + servicer = TrtllmServiceServicer(request_manager, model_path=model_path) + + server = grpc.aio.server( + options=[ + ("grpc.max_send_message_length", _GRPC_MAX_MESSAGE_LENGTH_BYTES), + ("grpc.max_receive_message_length", _GRPC_MAX_MESSAGE_LENGTH_BYTES), + ("grpc.keepalive_time_ms", 30000), + ("grpc.keepalive_timeout_ms", 10000), + ("grpc.keepalive_permit_without_calls", True), + ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000), + ] + ) + trtllm_service_pb2_grpc.add_TrtllmServiceServicer_to_server(servicer, server) + + if reflection is not None: + service_names = ( + trtllm_service_pb2.DESCRIPTOR.services_by_name["TrtllmService"].full_name, + reflection.SERVICE_NAME, + ) + reflection.enable_server_reflection(service_names, server) + logger.info("gRPC reflection enabled") + + address = f"{host}:{port}" + server.add_insecure_port(address) + await server.start() + logger.info(f"TensorRT-LLM SMG gRPC server started on {address}") + logger.info("Server is ready to accept requests") + + loop = asyncio.get_running_loop() + stop_event = asyncio.Event() + + def signal_handler() -> None: + logger.info("Received shutdown signal") + stop_event.set() + + for sig in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(sig, signal_handler) + + await stop_event.wait() + finally: + logger.info("Shutting down TensorRT-LLM SMG gRPC server...") + try: + if server is not None: + await server.stop(grace=5.0) + logger.info("gRPC server stopped") + finally: + if hasattr(llm, "shutdown"): + llm.shutdown() + logger.info("LLM engine stopped") + logger.info("Shutdown complete") + + uvloop.run(serve_grpc_async()) + + +__all__ = ["launch_smg_server"] diff --git a/tensorrt_llm/grpc/grpc_servicer.py b/tensorrt_llm/grpc/smg/servicer.py similarity index 99% rename from tensorrt_llm/grpc/grpc_servicer.py rename to tensorrt_llm/grpc/smg/servicer.py index 4fd05aab59b9..5b7d0561dbd3 100644 --- a/tensorrt_llm/grpc/grpc_servicer.py +++ b/tensorrt_llm/grpc/smg/servicer.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""gRPC Servicer for TensorRT-LLM. +"""Servicer for the TensorRT-LLM SMG gRPC adapter. Implements the TrtllmService gRPC service for high-performance communication with external routers (e.g., sgl-router) using pre-tokenized input. @@ -31,8 +31,8 @@ from tensorrt_llm.inputs.media_io import _load_and_convert_image from tensorrt_llm.logger import logger -from . import trtllm_service_pb2, trtllm_service_pb2_grpc -from .grpc_request_manager import ( +from .bindings import trtllm_service_pb2, trtllm_service_pb2_grpc +from .request_manager import ( GrpcRequestManager, create_disaggregated_params_from_proto, create_lora_request_from_proto, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 4f968d521eaf..cafc76489901 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -47,6 +47,8 @@ l0_a10: - unittest/others/test_cache_transceiver_precheck_run.py - unittest/others/test_lora_manager.py - unittest/others/test_tracing.py + - unittest/grpc/test_grpc_optional.py + - unittest/grpc/smg/test_smg.py - unittest/disaggregated/test_coordinator_e2e.py - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_agent_multi_backends.py diff --git a/tests/unittest/llmapi/test_grpc.py b/tests/unittest/grpc/smg/test_smg.py similarity index 96% rename from tests/unittest/llmapi/test_grpc.py rename to tests/unittest/grpc/smg/test_smg.py index cf69a99a0514..ca813d40b6d6 100644 --- a/tests/unittest/llmapi/test_grpc.py +++ b/tests/unittest/grpc/smg/test_smg.py @@ -12,31 +12,41 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for gRPC server components.""" +"""Unit tests for the SMG gRPC adapter.""" import asyncio import io import os import sys -import grpc import pytest import torch from PIL import Image from tensorrt_llm import LLM -from tensorrt_llm.grpc import trtllm_service_pb2 as pb2 -from tensorrt_llm.grpc.grpc_request_manager import ( +from tensorrt_llm.llmapi import KvCacheConfig + +# The SMG adapter depends on the optional 'smg-grpc-proto' package +# (pip install "tensorrt_llm[grpc-smg]"). Skip the whole module cleanly when it +# is absent so collection does not fail in an environment without the dependency. +pytest.importorskip( + "smg_grpc_proto", + reason='SMG gRPC adapter dependency not installed (pip install "tensorrt_llm[grpc-smg]")', +) + +import grpc # noqa: E402 + +from tensorrt_llm.grpc.smg.bindings import trtllm_service_pb2 as pb2 # noqa: E402 +from tensorrt_llm.grpc.smg.request_manager import ( # noqa: E402 GrpcRequestManager, create_disaggregated_params_from_proto, create_lora_request_from_proto, create_sampling_params_from_proto, ) -from tensorrt_llm.grpc.grpc_servicer import TrtllmServiceServicer -from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.grpc.smg.servicer import TrtllmServiceServicer # noqa: E402 # isort: off -sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..") +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../..") from utils.llm_data import llm_models_root # isort: on @@ -642,12 +652,15 @@ def get_model_path(model_name): return str(llm_models_root() / model_name) -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def grpc_service(): """Create a real LLM, request manager, and servicer for e2e testing. Uses TinyLlama-1.1B for minimal GPU resource usage. - Shared across all tests in this module. + Shared across all tests in the class; class scope (not module) so the + LLM is shut down and its GPU memory released before the multimodal + class below creates its own LLM — with module scope both models are + alive at once and the second one OOMs on A10. """ model_path = get_model_path(default_model_name) llm = LLM( @@ -820,12 +833,13 @@ def get_test_image_path(): return str(llm_models_root() / "multimodals" / "test_data" / "seashore.png") -@pytest.fixture(scope="module") +@pytest.fixture(scope="class") def grpc_vlm_service(): """Create a real VLM LLM, request manager, and servicer for multimodal e2e testing. Uses Qwen3-VL-8B-Instruct for vision-language model testing. - Shared across all tests in this module. + Shared across all tests in the class; class scope so this LLM does not + coexist with the one from grpc_service (see note there). """ model_path = get_model_path(vlm_model_name) llm = LLM( diff --git a/tests/unittest/grpc/test_grpc_optional.py b/tests/unittest/grpc/test_grpc_optional.py new file mode 100644 index 000000000000..16ae17f941f0 --- /dev/null +++ b/tests/unittest/grpc/test_grpc_optional.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Optional-dependency and lifecycle tests for the SMG gRPC adapter. + +These run correctly with or without the dependency installed: the "missing" case is +simulated so it is meaningful in every environment, and the "present" case is +guarded with ``importorskip``. +""" + +import asyncio +import builtins +import importlib +import sys +import types +from unittest.mock import AsyncMock, MagicMock + +import pytest + + +def test_smg_bindings_missing_gives_actionable_error(monkeypatch): + """A missing 'smg-grpc-proto' yields an actionable install hint. + + Importing the SMG bindings must fail with a ``pip install + "tensorrt_llm[grpc-smg]"`` hint rather than a bare ImportError. + """ + real_import = builtins.__import__ + + def import_without_smg(name, globals=None, locals=None, fromlist=(), level=0): + if name == "smg_grpc_proto.generated": + raise ModuleNotFoundError("No module named 'smg_grpc_proto'", name="smg_grpc_proto") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.delitem(sys.modules, "smg_grpc_proto", raising=False) + monkeypatch.delitem(sys.modules, "smg_grpc_proto.generated", raising=False) + monkeypatch.delitem(sys.modules, "tensorrt_llm.grpc.smg.bindings", raising=False) + monkeypatch.setattr(builtins, "__import__", import_without_smg) + + with pytest.raises(ModuleNotFoundError, match=r"tensorrt_llm\[grpc-smg\]") as exc_info: + import tensorrt_llm.grpc.smg.bindings # noqa: F401 + + assert exc_info.value.name == "smg_grpc_proto" + + +def test_smg_bindings_preserves_unrelated_import_error(monkeypatch): + """A transitive dependency failure is not rewritten as an install hint.""" + smg_package = types.ModuleType("smg_grpc_proto") + smg_package.__path__ = [] + generated_package = types.ModuleType("smg_grpc_proto.generated") + + def missing_protobuf(_name): + raise ModuleNotFoundError("No module named 'google.protobuf'", name="google.protobuf") + + generated_package.__getattr__ = missing_protobuf + monkeypatch.setitem(sys.modules, "smg_grpc_proto", smg_package) + monkeypatch.setitem(sys.modules, "smg_grpc_proto.generated", generated_package) + monkeypatch.delitem(sys.modules, "tensorrt_llm.grpc.smg.bindings", raising=False) + + with pytest.raises(ModuleNotFoundError) as exc_info: + importlib.import_module("tensorrt_llm.grpc.smg.bindings") + + assert exc_info.value.name == "google.protobuf" + assert "tensorrt_llm[grpc-smg]" not in str(exc_info.value) + + +def test_smg_bindings_present_smoke(): + """When the dependency is installed, the bindings import cleanly. + + They must expose the pb2 modules the SMG adapter depends on. + """ + pytest.importorskip( + "smg_grpc_proto", + reason="SMG gRPC adapter dependency not installed", + ) + from tensorrt_llm.grpc.smg import bindings + + assert bindings.trtllm_service_pb2 is not None + assert bindings.trtllm_service_pb2_grpc is not None + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("failure_point", ["bind", "start"]) +def test_smg_server_startup_failure_cleans_up(monkeypatch, failure_point): + """Binding and startup failures stop both the gRPC server and LLM.""" + pytest.importorskip( + "smg_grpc_proto", + reason="SMG gRPC adapter dependency not installed", + ) + from tensorrt_llm.grpc.smg import server as server_module + + llm = MagicMock() + grpc_server = MagicMock() + grpc_server.start = AsyncMock() + grpc_server.stop = AsyncMock() + + if failure_point == "bind": + grpc_server.add_insecure_port.side_effect = RuntimeError("bind failed") + else: + grpc_server.add_insecure_port.return_value = 8000 + grpc_server.start.side_effect = RuntimeError("start failed") + + monkeypatch.setitem(sys.modules, "grpc_reflection", None) + monkeypatch.delitem(sys.modules, "grpc_reflection.v1alpha", raising=False) + monkeypatch.setattr(server_module.uvloop, "run", asyncio.run) + monkeypatch.setattr(server_module, "PyTorchLLM", MagicMock(return_value=llm)) + monkeypatch.setattr(server_module, "GrpcRequestManager", MagicMock()) + monkeypatch.setattr(server_module, "TrtllmServiceServicer", MagicMock()) + monkeypatch.setattr(server_module.grpc.aio, "server", MagicMock(return_value=grpc_server)) + monkeypatch.setattr( + server_module.trtllm_service_pb2_grpc, + "add_TrtllmServiceServicer_to_server", + MagicMock(), + ) + + with pytest.raises(RuntimeError, match=f"{failure_point} failed"): + server_module.launch_smg_server( + host="127.0.0.1", + port=8000, + llm_args={"backend": "pytorch", "model": "test-model"}, + ) + + grpc_server.stop.assert_awaited_once_with(grace=5.0) + llm.shutdown.assert_called_once_with()