diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 1652e81c5853..93ccb357d07a 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-openengine.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..43201731d820 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-openengine.txt constraints.txt README.md ./ ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel diff --git a/requirements-openengine.txt b/requirements-openengine.txt new file mode 100644 index 000000000000..3ac09eceb170 --- /dev/null +++ b/requirements-openengine.txt @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Source: https://github.com/ai-dynamo/openengine/releases/tag/v0.1.0 +# BSR module: https://buf.build/openengine/openengine +# Immutable BSR commit: 768a93c7b44e40f28c692ad0b471a8f2 +--extra-index-url https://buf.build/gen/python +openengine-openengine-grpc-python==1.67.1.2.20260730172104+768a93c7b44e +openengine-openengine-protocolbuffers-python==31.1.0.2.20260730172104+768a93c7b44e +openengine-openengine-protocolbuffers-pyi==31.1.0.2.20260730172104+768a93c7b44e +grpcio>=1.67.1,<2 +protobuf>=6.31.1,<7 diff --git a/setup.py b/setup.py index deaeacc793e1..0ba2ff15220e 100644 --- a/setup.py +++ b/setup.py @@ -140,6 +140,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(): @@ -485,6 +486,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 65691acbf91c..954ee57a485e 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -57,9 +57,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) @@ -540,7 +537,7 @@ def launch_server( num_input_processor_workers: int = 8, num_media_load_workers: int = 8, multi_frontend_enabled: bool = True, - internal_disagg_auth_key: Optional[str] = None): + internal_disagg_auth_key: Optional[str] = None) -> None: backend = llm_args["backend"] model = served_model_name or llm_args["model"] @@ -631,129 +628,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, @@ -1237,9 +1111,13 @@ def launch_visual_gen_server( "--grpc", 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.", + help="Run the selected gRPC protocol instead of the OpenAI HTTP server.", status="prototype") +@stability_option("--grpc-protocol", + type=click.Choice(["smg", "openengine"]), + default="smg", + help="Protocol used when --grpc is enabled.", + status="prototype") @stability_option( "--served_model_name", type=str, @@ -1297,14 +1175,18 @@ def serve(model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str], agent_types: Optional[str], video_pruning_rate: Optional[float], 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]): + middleware: tuple[str, ...], grpc: bool, grpc_protocol: str, + enable_visual_gen: bool, served_model_name: Optional[str], + visual_gen_args: Optional[str]) -> None: """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path """ logger.set_level(log_level) + if not grpc and grpc_protocol != "smg": + raise click.UsageError("--grpc-protocol requires --grpc") + if moe_cluster_parallel_size is not None: logger.warning( "--moe_cluster_parallel_size / --cluster_size is deprecated and " @@ -1451,6 +1333,10 @@ def _serve_llm(): media_io_kwargs=parsed_media_io_kwargs) if grpc: + if num_serve_frontends != 1: + raise click.UsageError( + "--num_serve_frontends must be 1 when --grpc is enabled.") + # gRPC mode: launch gRPC server instead of OpenAI HTTP server # Check for unsupported arguments that are silently ignored in gRPC mode unsupported_args = { @@ -1478,10 +1364,27 @@ 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 grpc_protocol == "smg": + from tensorrt_llm.grpc.smg.server import \ + launch_server as launch_grpc_server + + launch_grpc_server(host, + port, + llm_args, + served_model_name=served_model_name) + else: + try: + from tensorrt_llm.grpc.openengine.server import \ + launch_server as launch_grpc_server + except ImportError as error: + raise click.ClickException( + f"Failed to import OpenEngine support: {error}. " + "Install the optional Python bindings with `python -m " + "pip install --extra-index-url " + "https://buf.build/gen/python " + "\"tensorrt_llm[openengine]\"`.") from error + + launch_grpc_server(host, port) else: # Default: launch OpenAI HTTP server launch_server( @@ -1514,6 +1417,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 grpc: + raise click.UsageError( + "--grpc is not supported by the VisualGen server") _serve_visual_gen() else: _serve_llm() diff --git a/tensorrt_llm/grpc/__init__.py b/tensorrt_llm/grpc/__init__.py index d75315bd0605..3f062fec0e7e 100644 --- a/tensorrt_llm/grpc/__init__.py +++ b/tensorrt_llm/grpc/__init__.py @@ -1,87 +1,6 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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. -r"""TensorRT-LLM gRPC module for high-performance communication with external routers. +"""Protocol-specific gRPC integrations for TensorRT-LLM.""" -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", -] +__all__ = [] diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md new file mode 100644 index 000000000000..da8d1f915bd5 --- /dev/null +++ b/tensorrt_llm/grpc/openengine/README.md @@ -0,0 +1,46 @@ + + + +# TensorRT-LLM OpenEngine stub server + +`trtllm-serve` can expose an experimental OpenEngine gRPC server instead of its normal OpenAI HTTP server. SMG remains the default gRPC protocol. + +Install the optional Python bindings from the Buf Schema Registry: + +```bash +python -m pip install \ + --extra-index-url https://buf.build/gen/python \ + "tensorrt_llm[openengine]" +``` + +Then select OpenEngine when starting the gRPC server: + +```bash +trtllm-serve \ + --grpc \ + --grpc-protocol openengine \ + --host 0.0.0.0 \ + --port 50051 +``` + +Existing `--grpc` invocations continue to select SMG. OpenEngine and VisualGen cannot be enabled together. + +This initial integration is a protocol stub. Every OpenEngine RPC returns gRPC status `UNIMPLEMENTED`; no request reaches the TensorRT-LLM engine. OpenEngine and SMG are independent protocol integrations. This integration does not make a replacement or convergence decision between them. + +## Dependency provenance + +The schema source is the Apache-2.0-licensed [`ai-dynamo/openengine`](https://github.com/ai-dynamo/openengine) repository at signed Git tag [`v0.1.0`](https://github.com/ai-dynamo/openengine/releases/tag/v0.1.0). That release maps to the public [`buf.build/openengine/openengine`](https://buf.build/openengine/openengine) module at immutable BSR commit `768a93c7b44e40f28c692ad0b471a8f2`. + +The BSR generated the pinned wheels from that module commit: + +| Package | Generator | Version | SHA-256 | +| --- | --- | --- | --- | +| `openengine-openengine-grpc-python` | [`grpc/python`](https://buf.build/grpc/python) | `1.67.1.2.20260730172104+768a93c7b44e` | `1485aed9799c4eb9367d1a261ca5cc5319f1e9b8d950ac98a26f3cb3641b8cf6` | +| `openengine-openengine-protocolbuffers-python` | [`protocolbuffers/python`](https://buf.build/protocolbuffers/python) | `31.1.0.2.20260730172104+768a93c7b44e` | `6eae12c3d8d06147fccf608da9772d6391139031fabdafdb7cf4c71a19c1f25e` | +| `openengine-openengine-protocolbuffers-pyi` | [`protocolbuffers/pyi`](https://buf.build/protocolbuffers/pyi) | `31.1.0.2.20260730172104+768a93c7b44e` | `8b0a054dbdaaa67459b3fa4786f13d8f6f4d30cf30be325f5416dbd97aba46a6` | + +Buf documents the package naming and version format in its [Python-generated SDK guide](https://buf.build/docs/bsr/generated-sdks/python/). The final version segment is the BSR commit prefix. The exact requirements are pinned in `requirements-openengine.txt`. + +## Maintenance boundary + +The OpenEngine contributor community owns this adapter, its tests, protocol version updates, and integration bugs. TensorRT-LLM internal APIs do not provide compatibility guarantees to protocol adapters. Adapter updates must follow core runtime changes and must not block normal TensorRT-LLM development or releases. diff --git a/tensorrt_llm/grpc/openengine/__init__.py b/tensorrt_llm/grpc/openengine/__init__.py new file mode 100644 index 000000000000..d1049d02ea13 --- /dev/null +++ b/tensorrt_llm/grpc/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 gRPC integration for TensorRT-LLM.""" + +from .server import OpenEngineServer, launch_server + +__all__ = ["OpenEngineServer", "launch_server"] diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py new file mode 100644 index 000000000000..641ff7a90eb0 --- /dev/null +++ b/tensorrt_llm/grpc/openengine/server.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine gRPC server lifecycle for TensorRT-LLM.""" + +import asyncio +import signal + +import grpc +import uvloop +from openengine.v1 import openengine_pb2_grpc + +from tensorrt_llm.logger import logger + +__all__ = ["OpenEngineServer", "launch_server"] + + +def _format_bind_address(host: str, port: int) -> str: + """Format a host and port as a gRPC bind address.""" + if ":" in host and not (host.startswith("[") and host.endswith("]")): + host = f"[{host}]" + return f"{host}:{port}" + + +class OpenEngineServer: + """OpenEngine gRPC server with intentionally unimplemented RPCs. + + Args: + host: Interface on which the server listens. + port: Port on which the server listens. Use zero to select a free port. + """ + + def __init__(self, host: str, port: int) -> None: + self.host = host + self.port = port + self._server = grpc.aio.server() + openengine_pb2_grpc.add_InferenceServicer_to_server( + openengine_pb2_grpc.InferenceServicer(), self._server + ) + openengine_pb2_grpc.add_ControlServicer_to_server( + openengine_pb2_grpc.ControlServicer(), self._server + ) + self._bind_address = _format_bind_address(host, port) + bound_port = self._server.add_insecure_port(self._bind_address) + if bound_port == 0: + raise RuntimeError(f"Failed to bind OpenEngine server to {self._bind_address}") + if port == 0: + self.port = bound_port + + async def start(self) -> None: + """Start accepting OpenEngine requests.""" + await self._server.start() + address = _format_bind_address(self.host, self.port) + logger.info(f"OpenEngine stub server started on {address}") + + async def stop(self, grace: float = 5.0) -> None: + """Stop accepting OpenEngine requests. + + Args: + grace: Maximum time in seconds to allow active RPCs to finish. + """ + await self._server.stop(grace=grace) + logger.info("OpenEngine stub server stopped") + + async def wait_for_termination(self) -> None: + """Wait until the OpenEngine server terminates.""" + await self._server.wait_for_termination() + + +def launch_server(host: str, port: int) -> None: + """Launch the dedicated OpenEngine gRPC server. + + Args: + host: Interface on which the server listens. + port: Port on which the server listens. + """ + + async def serve() -> None: + server = OpenEngineServer(host=host, port=port) + 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) + + try: + logger.warning( + "OpenEngine protocol support is a stub: no model is loaded and " + "all RPCs return UNIMPLEMENTED." + ) + await server.start() + await stop_event.wait() + finally: + await server.stop() + + uvloop.run(serve()) diff --git a/tensorrt_llm/grpc/smg/__init__.py b/tensorrt_llm/grpc/smg/__init__.py new file mode 100644 index 000000000000..197f1dec1141 --- /dev/null +++ b/tensorrt_llm/grpc/smg/__init__.py @@ -0,0 +1,95 @@ +# 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"); +# 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. + +r"""SMG integration 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 + +# Import protobuf-dependent modules only when the generated package is present. +if PROTOS_AVAILABLE: + try: + from .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: + from .servicer import TrtllmServiceServicer + + SERVICER_AVAILABLE = True + except ImportError: + SERVICER_AVAILABLE = False + TrtllmServiceServicer = None +else: + REQUEST_MANAGER_AVAILABLE = False + GrpcRequestManager = None + create_sampling_params_from_proto = None + create_lora_request_from_proto = None + create_disaggregated_params_from_proto = None + SERVICER_AVAILABLE = False + TrtllmServiceServicer = None + +__all__ = [ + "PROTOS_AVAILABLE", + "REQUEST_MANAGER_AVAILABLE", + "SERVICER_AVAILABLE", + "GrpcRequestManager", + "TrtllmServiceServicer", + "create_disaggregated_params_from_proto", + "create_lora_request_from_proto", + "create_sampling_params_from_proto", + "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 99% rename from tensorrt_llm/grpc/grpc_request_manager.py rename to tensorrt_llm/grpc/smg/request_manager.py index ca2e59eb1f36..48c4d8967faf 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. +"""SMG request manager for TensorRT-LLM. Manages request lifecycle for gRPC requests, converting between protobuf and TensorRT-LLM types. Designed for high-performance communication with diff --git a/tensorrt_llm/grpc/smg/server.py b/tensorrt_llm/grpc/smg/server.py new file mode 100644 index 000000000000..6f28fd429333 --- /dev/null +++ b/tensorrt_llm/grpc/smg/server.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SMG gRPC server lifecycle for TensorRT-LLM.""" + +import asyncio +import signal + +import click +import grpc +import uvloop + +from tensorrt_llm import LLM as PyTorchLLM +from tensorrt_llm.logger import logger + +from . import PROTOS_AVAILABLE, trtllm_service_pb2, trtllm_service_pb2_grpc + +__all__ = ["launch_server"] + +_MAX_MESSAGE_LENGTH_BYTES = 32 * 1024 * 1024 + + +def launch_server( + host: str, + port: int, + llm_args: dict, + 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 used to initialize the LLM. + served_model_name: Model name exposed by the server. Defaults to the model path. + """ + if not PROTOS_AVAILABLE: + raise click.ClickException( + "SMG gRPC support requires smg-grpc-proto. Install it with " + "`python -m pip install smg-grpc-proto`." + ) + + from .request_manager import GrpcRequestManager + from .servicer import TrtllmServiceServicer + + try: + from grpc_reflection.v1alpha import reflection + + reflection_available = True + except ImportError: + reflection_available = False + + 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") + + request_manager = GrpcRequestManager(llm) + servicer = TrtllmServiceServicer(request_manager, model_path=model_path) + server = grpc.aio.server( + options=[ + ("grpc.max_send_message_length", _MAX_MESSAGE_LENGTH_BYTES), + ("grpc.max_receive_message_length", _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_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") + + address = f"{host}:{port}" + bound_port = server.add_insecure_port(address) + if bound_port == 0: + try: + await server.stop(grace=0) + finally: + if hasattr(llm, "shutdown"): + llm.shutdown() + raise RuntimeError(f"Failed to bind SMG gRPC server to {address}") + await server.start() + logger.info(f"TensorRT-LLM SMG gRPC server started on {host}:{bound_port}") + 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) + + try: + await stop_event.wait() + except KeyboardInterrupt: + logger.info("Interrupted by user") + finally: + logger.info("Shutting down TensorRT-LLM SMG gRPC server...") + await server.stop(grace=5.0) + logger.info("gRPC server stopped") + + if hasattr(llm, "shutdown"): + llm.shutdown() + logger.info("LLM engine stopped") + logger.info("Shutdown complete") + + uvloop.run(serve_grpc_async()) 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..c18f8dd9f150 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. +"""SMG servicer for TensorRT-LLM. Implements the TrtllmService gRPC service for high-performance communication with external routers (e.g., sgl-router) using pre-tokenized input. @@ -32,7 +32,7 @@ from tensorrt_llm.logger import logger from . import trtllm_service_pb2, trtllm_service_pb2_grpc -from .grpc_request_manager import ( +from .request_manager import ( GrpcRequestManager, create_disaggregated_params_from_proto, create_lora_request_from_proto, diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index c3430712d780..6fcac26248dc 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -184,6 +184,15 @@ commands: is_flag: true flags: - "--grpc" + grpc_protocol: + type: Choice(['smg', 'openengine']) + default: smg + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--grpc-protocol" host: type: str default: localhost diff --git a/tests/unittest/llmapi/test_grpc.py b/tests/unittest/grpc/smg/test_server.py similarity index 99% rename from tests/unittest/llmapi/test_grpc.py rename to tests/unittest/grpc/smg/test_server.py index cf69a99a0514..9665917e2258 100644 --- a/tests/unittest/llmapi/test_grpc.py +++ b/tests/unittest/grpc/smg/test_server.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 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"); @@ -12,7 +12,7 @@ # 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 SMG gRPC server components.""" import asyncio import io @@ -25,18 +25,18 @@ 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.grpc.smg import trtllm_service_pb2 as pb2 +from tensorrt_llm.grpc.smg.request_manager import ( 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.grpc.smg.servicer import TrtllmServiceServicer from tensorrt_llm.llmapi import KvCacheConfig # 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