[#21102][feat] Add an opt-in pygrpc-native server for OpenEngine

setup.py

@@ -727,6 +727,7 @@ def get_build_state_options():
     extras_require={
         "devel": devel_deps + grpc_smg_deps,
         "openengine": [],
+        "openengine-native": ["pygrpc-native>=0.5.0,<1"],
         "mx": mx_deps,
         "grpc-smg": grpc_smg_deps,
     },

tensorrt_llm/commands/serve.py

@@ -1209,6 +1209,14 @@ def launch_visual_gen_server(
     help=
     "Expose per-rank KV occupancy for OpenEngine routing, independently of KV events.",
     status="prototype")
+@stability_option(
+    "--openengine-grpc-server",
+    type=click.Choice(["grpcio", "pygrpc-native"]),
+    default="grpcio",
+    help="gRPC server for OpenEngine. pygrpc-native runs the HTTP/2 and "
+    "gRPC work in Rust outside the GIL and needs "
+    "`pip install \"tensorrt_llm[openengine-native]\"`.",
+    status="prototype")
 @stability_option(
     "--served_model_name",
     type=str,
@@ -1310,6 +1318,7 @@ def serve(
     grpc: bool,
     grpc_protocol: str,
     openengine_enable_load_metrics: bool,
+    openengine_grpc_server: str,
     enable_visual_gen: bool,
     served_model_name: Optional[str],
     visual_gen_args: Optional[str],
@@ -1342,6 +1351,11 @@ def serve(
         raise click.UsageError(
             "--openengine-enable-load-metrics requires --grpc --grpc-protocol openengine"
         )
+    if openengine_grpc_server != "grpcio" and (not grpc or grpc_protocol
+                                               != "openengine"):
+        raise click.UsageError(
+            "--openengine-grpc-server requires --grpc --grpc-protocol openengine"
+        )
     if moe_cluster_parallel_size is not None:
         logger.warning(
             "--moe_cluster_parallel_size / --cluster_size is deprecated and "
@@ -1601,7 +1615,8 @@ def _serve_llm():
                         served_model_name=served_model_name,
                         report_failure=_report_observed_child_failure,
                         enable_load_metrics=openengine_enable_load_metrics,
-                        subagent_affinity_auth_key=internal_disagg_auth_key)
+                        subagent_affinity_auth_key=internal_disagg_auth_key,
+                        grpc_server=openengine_grpc_server)
         else:
             # Default: launch OpenAI HTTP server
             launch_server(

tensorrt_llm/grpc/openengine/README.md

@@ -168,6 +168,10 @@ Frontends send heartbeats once per second. A missing heartbeat for 10 seconds wi
 
 Multiple frontends require the default classic IPC executor and a nonzero port. SMG and AutoDeploy do not support this mode. Frontends share host CPU resources, so provision sufficient CPU and benchmark the intended model, router, and connection pool before choosing a frontend count. More processes do not necessarily increase throughput once GPU execution becomes the bottleneck.
 
+## gRPC server implementation
+
+The server runs on `grpc.aio` by default. `--openengine-grpc-server pygrpc-native` runs it on [pygrpc-native](https://github.com/pygrpc/pygrpc-native) instead, a drop-in replacement for `grpc.aio.server()` that does the HTTP/2 and gRPC work in Rust outside the GIL. The servicers, health service and server options stay the same. Install it with `pip install "tensorrt_llm[openengine-native]"`; the server checks for it before loading the model. It works with multiple frontend processes. The internal coordinator between frontends always uses `grpc.aio`.
+
 ## Schema and binding 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), Git commit `b5f2bd93721f7b888d3e2440679e0ae7012939d1`. That release maps to the public [`buf.build/openengine/openengine`](https://buf.build/openengine/openengine) module at immutable BSR release `768a93c7b44e40f28c692ad0b471a8f2`.

tensorrt_llm/grpc/openengine/server.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import gc
+import importlib.metadata
 import ipaddress
 import os
 import signal
@@ -20,6 +21,7 @@
 import grpc
 import uvloop
 from grpc_health.v1 import health_pb2, health_pb2_grpc
+from packaging.version import Version
 
 from tensorrt_llm import LLM as PyTorchLLM
 from tensorrt_llm.grpc._protobuf import _warn_if_python_protobuf
@@ -58,11 +60,36 @@
     ("grpc.http2.max_pings_without_data", 0),
     ("grpc.http2.min_ping_interval_without_data_ms", 10000),
 ]
+# The first release that honors grpc.so_reuseport, which multiple frontends need.
+_PYGRPC_NATIVE_MIN_VERSION = "0.5.0"
 
 
 _format_bind_address = _format_host_port
 
 
+def _check_pygrpc_native() -> None:
+    """Fail before the engine loads if pygrpc-native is missing or too old."""
+    try:
+        installed = importlib.metadata.version("pygrpc-native")
+    except importlib.metadata.PackageNotFoundError:
+        installed = None
+    if installed is None or Version(installed) < Version(_PYGRPC_NATIVE_MIN_VERSION):
+        raise click.ClickException(
+            f"--openengine-grpc-server pygrpc-native requires "
+            f"pygrpc-native>={_PYGRPC_NATIVE_MIN_VERSION}, found {installed or 'none'}. "
+            'Install it with: pip install "tensorrt_llm[openengine-native]"'
+        )
+
+
+def _create_grpc_server(grpc_server: str, options: list[tuple[str, Any]]) -> Any:
+    """Build the grpc.aio server, or pygrpc-native's drop-in replacement for it."""
+    if grpc_server == "pygrpc-native":
+        import pygrpc_native.aio
+
+        return pygrpc_native.aio.server(options=options)
+    return grpc.aio.server(options=options)
+
+
 def _is_loopback(host: str) -> bool:
     """Whether `host` resolves to a loopback address."""
     cleaned = host.strip("[]")
@@ -99,6 +126,7 @@ class OpenEngineServer:
         llm: Initialized TensorRT-LLM LLM instance.
         model: Model name accepted by Generate requests.
         subagent_affinity_auth_key: Shared key for authenticated parent affinity hints.
+        grpc_server: gRPC server implementation, "grpcio" or "pygrpc-native".
     """
 
     def __init__(
@@ -111,6 +139,7 @@ def __init__(
         frontend: FrontendClient | None = None,
         instance_id: str | None = None,
         subagent_affinity_auth_key: str | None = None,
+        grpc_server: str = "grpcio",
     ) -> None:
         self.host = host
         self.port = port
@@ -125,7 +154,7 @@ def __init__(
         self._fatal_engine_failure = False
         options = list(_SERVER_OPTIONS)
         options.append(("grpc.so_reuseport", 1 if frontend is not None else 0))
-        self._server = grpc.aio.server(options=options)
+        self._server = _create_grpc_server(grpc_server, options)
         self._input_executor = ThreadPoolExecutor(
             max_workers=4, thread_name_prefix="openengine_input"
         )
@@ -167,7 +196,7 @@ def __init__(
                 "The listener is unauthenticated and unencrypted: restrict it to a trusted "
                 "network or front it with an authenticating TLS proxy."
             )
-        # grpc.aio raises RuntimeError itself when the bind fails, naming the
+        # Both servers raise RuntimeError themselves when the bind fails, naming the
         # address, so there is no zero return to check for.
         try:
             bound_port = self._server.add_insecure_port(bind_address)
@@ -307,6 +336,7 @@ def launch_server(
     report_failure: Callable[[int, str, str], None] | None = None,
     enable_load_metrics: bool = False,
     subagent_affinity_auth_key: str | None = None,
+    grpc_server: str = "grpcio",
 ) -> None:
     """Launch the dedicated OpenEngine gRPC server.
 
@@ -318,8 +348,11 @@ def launch_server(
         report_failure: Records a child exit before startup readiness.
         enable_load_metrics: Enable scheduler-owned routing-load snapshots, independently of KV events.
         subagent_affinity_auth_key: Shared key for authenticated parent affinity hints.
+        grpc_server: gRPC server implementation, "grpcio" or "pygrpc-native".
     """
     _warn_if_python_protobuf()
+    if grpc_server == "pygrpc-native":
+        _check_pygrpc_native()
 
     async def serve() -> None:
         logger.info("Initializing TensorRT-LLM OpenEngine server...")
@@ -491,6 +524,7 @@ def monitor_done(task: asyncio.Task) -> None:
                 frontend=frontend,
                 instance_id=instance_id,
                 subagent_affinity_auth_key=subagent_affinity_auth_key,
+                grpc_server=grpc_server,
             )
             _disable_gc_if_requested()
             await server.start()

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -211,6 +211,7 @@ l0_cpu:
   - unittest/grpc/test_grpc_optional.py::test_protobuf_warning_skips_unavailable_api
   - unittest/grpc/test_grpc_optional.py::test_protobuf_warning_does_not_block_model_startup
   - unittest/grpc/test_grpc_optional.py::test_smg_server_startup_failure_cleans_up
+  - unittest/grpc/test_grpc_optional.py::test_openengine_pygrpc_native_is_checked_before_engine_start
   - unittest/inputs
   - unittest/kv_cache_manager_v2_tests
   - unittest/llmapi/test_deepseek_v32_encoding.py

tests/unittest/api_stability/references/trtllm_serve_cli.yaml

@@ -378,6 +378,15 @@ commands:
         is_flag: true
         flags:
         - "--openengine-enable-load-metrics"
+      openengine_grpc_server:
+        type: Choice(['grpcio', 'pygrpc-native'])
+        default: grpcio
+        status: prototype
+        required: false
+        multiple: false
+        is_flag: false
+        flags:
+        - "--openengine-grpc-server"
       otlp_traces_endpoint:
         type: str
         default: null