[#21032][feat] Expose OpenEngine Prometheus metrics

tensorrt_llm/commands/serve.py

@@ -1202,6 +1202,19 @@ def launch_visual_gen_server(
                   default="smg",
                   help="Protocol used when --grpc is enabled.",
                   status="prototype")
+@stability_option(
+    "--openengine-metrics-port",
+    type=click.IntRange(1, 65535),
+    default=None,
+    help="OpenEngine Prometheus HTTP port (default: 8001). "
+    "Enabled by return_perf_metrics or perf_metrics_output_dir in the engine configuration.",
+    status="prototype")
+@stability_option(
+    "--openengine-metrics-host",
+    type=str,
+    default=None,
+    help="Bind address for OpenEngine Prometheus (default: --host).",
+    status="prototype")
 @stability_option(
     "--openengine-enable-load-metrics",
     is_flag=True,
@@ -1310,6 +1323,8 @@ def serve(
     grpc: bool,
     grpc_protocol: str,
     openengine_enable_load_metrics: bool,
+    openengine_metrics_port: Optional[int],
+    openengine_metrics_host: Optional[str],
     enable_visual_gen: bool,
     served_model_name: Optional[str],
     visual_gen_args: Optional[str],
@@ -1342,6 +1357,19 @@ def serve(
         raise click.UsageError(
             "--openengine-enable-load-metrics requires --grpc --grpc-protocol openengine"
         )
+    if openengine_metrics_port is not None and (not grpc or grpc_protocol
+                                                != "openengine"):
+        raise click.UsageError(
+            "--openengine-metrics-port requires --grpc --grpc-protocol openengine"
+        )
+    if openengine_metrics_host is not None and (not grpc or grpc_protocol
+                                                != "openengine"):
+        raise click.UsageError(
+            "--openengine-metrics-host requires --grpc --grpc-protocol openengine"
+        )
+    if openengine_metrics_port is not None and openengine_metrics_port == port:
+        raise click.UsageError(
+            "--openengine-metrics-port must differ from --port")
 
     if moe_cluster_parallel_size is not None:
         logger.warning(
@@ -1601,6 +1629,9 @@ def _serve_llm():
                         served_model_name=served_model_name,
                         report_failure=_report_observed_child_failure,
                         enable_load_metrics=openengine_enable_load_metrics,
+                        metrics_host=openengine_metrics_host,
+                        metrics_port=8001 if openengine_metrics_port is None
+                        else openengine_metrics_port,
                         subagent_affinity_auth_key=internal_disagg_auth_key)
         else:
             # Default: launch OpenAI HTTP server

tensorrt_llm/executor/proxy.py

@@ -1282,7 +1282,10 @@ def get_stats(self, timeout: float) -> List[dict]:
             logger.warning("RPC client not initialized, cannot get stats")
             return []
 
-        stats = self.rpc_client.fetch_stats_wait_async(timeout=timeout).remote()
+        # Allow transport overhead beyond the worker's queue wait, while
+        # bounding the client even if the worker never replies.
+        stats = self.rpc_client.fetch_stats_wait_async(timeout=timeout).remote(
+            timeout=None if timeout is None else timeout + 1.0)
         return [json.loads(s) if isinstance(s, str) else s for s in stats]
 
     def get_kv_cache_capacity(self) -> dict:
@@ -1390,7 +1393,8 @@ def aget_stats(self, timeout: float) -> IterationResult:
         # Fetch stats via RPC and populate the result
         try:
             stats = self.rpc_client.fetch_stats_wait_async(
-                timeout=timeout).remote()
+                timeout=timeout).remote(
+                    timeout=None if timeout is None else timeout + 1.0)
         except Exception as e:
             logger.debug(f"Error fetching stats via RPC: {e}")
             stats = []

tensorrt_llm/executor/rpc/rpc_client.py

@@ -125,6 +125,7 @@ def __init__(self,
 
         self._server_stopped = False
         self._closed = False
+        self._call_lock = threading.Lock()
         self._loop = None
         self._loop_thread = None
         self._reader_asyncio_task = None  # Track the asyncio task for proper cancellation
@@ -157,9 +158,10 @@ def shutdown_server(self):
     def close(self):
         """Gracefully close the client, cleaning up background tasks."""
 
-        if self._closed:
-            return
-        self._closed = True
+        with self._call_lock:
+            if self._closed:
+                return
+            self._closed = True
 
         logger_debug("[client] RPC Client closing")
 
@@ -168,29 +170,35 @@ def close(self):
         # loop/executor because they may depend on the loop to complete.
         self._broadcast_streaming_error(RPCCancelled("RPC client closed"))
 
-        # 1. Cancel the reader task
-        if self._reader_task and not self._reader_task.done():
-            if self._loop and self._loop.is_running(
-            ) and self._reader_asyncio_task:
-                try:
-
-                    async def cancel_reader_task():
-                        if self._reader_asyncio_task and not self._reader_asyncio_task.done(
-                        ):
-                            self._reader_asyncio_task.cancel()
-                            try:
-                                await self._reader_asyncio_task
-                            except asyncio.CancelledError:
-                                pass
-
-                    cancel_future = asyncio.run_coroutine_threadsafe(
-                        cancel_reader_task(), self._loop)
-                    cancel_future.result(timeout=2.0)
-                    logger_debug("[client] Reader task cancelled successfully")
-                except concurrent.futures.TimeoutError:
-                    logger.warning("Reader task did not exit gracefully")
-                except Exception as e:
-                    logger_debug(f"[client] Reader task cleanup: {e}")
+        # 1. Settle regular calls before stopping the loop they depend on.
+        if self._loop and self._loop.is_running():
+            try:
+
+                async def cancel_reader_task():
+                    pending = [
+                        future for future, loop in tuple(
+                            self._pending_futures.values())
+                        if loop is asyncio.get_running_loop()
+                    ]
+                    self._fail_pending_futures(
+                        RPCCancelled("RPC client closed"))
+                    await asyncio.gather(*pending, return_exceptions=True)
+                    if self._reader_asyncio_task and not self._reader_asyncio_task.done(
+                    ):
+                        self._reader_asyncio_task.cancel()
+                        try:
+                            await self._reader_asyncio_task
+                        except asyncio.CancelledError:
+                            pass
+
+                cancel_future = asyncio.run_coroutine_threadsafe(
+                    cancel_reader_task(), self._loop)
+                cancel_future.result(timeout=2.0)
+                logger_debug("[client] Reader task cancelled successfully")
+            except concurrent.futures.TimeoutError:
+                logger.warning("Reader task did not exit gracefully")
+            except Exception as e:
+                logger_debug(f"[client] Reader task cleanup: {e}")
 
         # 2. Stop the event loop
         if self._loop and self._loop.is_running():
@@ -296,8 +304,12 @@ async def _handle_reader_exception(self, exception: Exception):
         """
         logger.error(f"Exception in RPC response reader: {exception}")
 
-        # Propagate exception to all pending futures
-        for (future, target_loop) in self._pending_futures.values():
+        self._fail_pending_futures(exception)
+        self._broadcast_streaming_error(exception)
+
+    def _fail_pending_futures(self, exception: Exception) -> None:
+        """Wake regular callers on their owning loops before RPC teardown."""
+        for future, target_loop in tuple(self._pending_futures.values()):
             if not future.done():
 
                 def safe_set_exception(f=future, exc=exception):
@@ -311,9 +323,6 @@ def safe_set_exception(f=future, exc=exception):
 
                 target_loop.call_soon_threadsafe(safe_set_exception)
 
-        # Propagate to streaming queues via common helper
-        self._broadcast_streaming_error(exception)
-
     async def _wait_for_response(self) -> RPCResponse:
         """Wait for a response from the socket.
 
@@ -457,7 +466,7 @@ async def _call_async(self, method_name, *args, **kwargs):
         nvtx_mark_debug(f"RPC.async.{method_name}",
                         color="yellow",
                         category="RPC")
-        if self._server_stopped:
+        if self._closed or self._server_stopped:
             raise RPCCancelled("Server is shutting down, request cancelled")
 
         rpc_params = kwargs.pop("__rpc_params", RPCParams())
@@ -474,12 +483,14 @@ async def _call_async(self, method_name, *args, **kwargs):
         await self._client_socket.put_async(request)
         logger_debug(f"[client] RPC Client sent request: {request}")
 
-        if not need_response:
-            return None
-
-        loop = asyncio.get_running_loop()
-        future = loop.create_future()
-        self._pending_futures[request_id] = (future, loop)
+        with self._call_lock:
+            if self._closed:
+                raise RPCCancelled("RPC client closed")
+            if not need_response:
+                return None
+            loop = asyncio.get_running_loop()
+            future = loop.create_future()
+            self._pending_futures[request_id] = (future, loop)
 
         try:
             # If timeout, the remote call should return a timeout error timely,
@@ -532,8 +543,11 @@ def _call_sync(self, method_name, *args, **kwargs):
         nvtx_mark_debug(f"RPC.sync.{method_name}",
                         color="green",
                         category="RPC")
-        future = asyncio.run_coroutine_threadsafe(
-            self._call_async(method_name, *args, **kwargs), self._loop)
+        with self._call_lock:
+            if self._closed:
+                raise RPCCancelled("RPC client closed")
+            future = asyncio.run_coroutine_threadsafe(
+                self._call_async(method_name, *args, **kwargs), self._loop)
         result = future.result()
         return result
 
@@ -552,12 +566,7 @@ def _call_future(self, name: str, *args,
         """
         nvtx_mark_debug(f"RPC.future.{name}", color="blue", category="RPC")
 
-        def _async_to_sync():
-            future = asyncio.run_coroutine_threadsafe(
-                self._call_async(name, *args, **kwargs), self._loop)
-            return future.result()
-
-        return self._executor.submit(_async_to_sync)
+        return self._executor.submit(self._call_sync, name, *args, **kwargs)
 
     async def _call_streaming(self, name: str, *args,
                               **kwargs) -> AsyncIterator[Any]:
@@ -570,7 +579,7 @@ async def _call_streaming(self, name: str, *args,
         """
         nvtx_mark_debug(f"RPC.streaming.{name}", color="red", category="RPC")
 
-        if self._server_stopped:
+        if self._closed or self._server_stopped:
             raise RPCCancelled("Server is shutting down, request cancelled")
         rpc_params = kwargs.pop("__rpc_params", RPCParams())
         timeout = rpc_params.timeout if rpc_params.timeout is not None else self._timeout

tensorrt_llm/executor/rpc_proxy.py

@@ -109,7 +109,8 @@ def get_stats(self, timeout: float) -> List[dict]:
         """
         try:
             stats = self.rpc_client.fetch_stats_wait_async(
-                timeout=timeout).remote()
+                timeout=timeout).remote(
+                    timeout=None if timeout is None else timeout + 1.0)
             return [json.loads(s) if isinstance(s, str) else s for s in stats]
         except Exception as e:
             logger.debug(f"Error fetching stats via RPC: {e}")
@@ -181,7 +182,8 @@ def aget_stats(self, timeout: float) -> IterationResult:
         # Fetch stats via RPC and populate the result
         try:
             stats = self.rpc_client.fetch_stats_wait_async(
-                timeout=timeout).remote()
+                timeout=timeout).remote(
+                    timeout=None if timeout is None else timeout + 1.0)
         except Exception:
             stats = []
 

tensorrt_llm/grpc/openengine/README.md

@@ -91,6 +91,66 @@ Consumers must track rank-local sequence numbers and resynchronize on gaps. Publ
 
 Direct ZMQ sockets are unauthenticated. Keep them on a trusted network and restrict access to intended subscribers.
 
+## Prometheus metrics
+
+OpenEngine serves metrics on a separate HTTP port, using the same collection settings and `trtllm_*` metric families as the HTTP server.
+
+Save this as `metrics.yaml` to collect request timing, token counts, and engine statistics such as KV-cache usage:
+
+```yaml
+return_perf_metrics: true
+enable_iter_perf_stats: true
+```
+
+Start the server and check the metrics endpoint:
+
+```bash
+trtllm-serve Qwen/Qwen3-0.6B \
+  --config metrics.yaml \
+  --grpc --grpc-protocol openengine --host 127.0.0.1 --port 50051
+
+curl --fail http://127.0.0.1:8001/prometheus/metrics
+```
+
+Metrics use port `8001` and inherit `--host`. Override them with `--openengine-metrics-port` and `--openengine-metrics-host`; the metrics port must differ from the gRPC port. These options change the address, not whether collection is enabled.
+
+For Prometheus, add a separate scrape target:
+
+```yaml
+scrape_configs:
+  - job_name: trtllm-openengine
+    scrape_interval: 15s
+    metrics_path: /prometheus/metrics
+    static_configs:
+      - targets: ["engine-host:8001"]
+```
+
+Replace `engine-host` with an address Prometheus can reach. For remote scraping, bind metrics to a reachable interface. The endpoint has no authentication; restrict access to a trusted network. Use `Control.Health` to check inference health; a successful scrape only checks the exporter.
+
+### Collection settings
+
+- `return_perf_metrics: true` enables the endpoint and request metrics. Both collection settings default to `false`.
+- `enable_iter_perf_stats: true` adds engine statistics. Set it to `false` for request-only collection; it does not enable the endpoint by itself.
+- `perf_metrics_output_dir` also enables collection and the endpoint, matching HTTP enablement, but OpenEngine does not write HTTP per-request JSONL files.
+- KV events and routing load remain separate: metrics do not enable event publishing or `--openengine-enable-load-metrics`.
+
+### Reading the metrics
+
+- `trtllm_request_success_total` counts completed output candidates by finish reason. One successful request with `num_sequences=2` contributes two candidates.
+- `trtllm_prompt_tokens_total` counts the prompt once per request; `trtllm_generation_tokens_total` sums tokens across outputs.
+- `trtllm_openengine_rpc_outcomes_total{phase,outcome}` counts one outcome per `Generate` call, including rejection, cancellation, and engine errors. Use this to distinguish RPC outcomes from candidate completion. Completion means the server consumed the response stream, not that the client acknowledged receipt.
+- `trtllm_openengine_metrics_errors_total{source}` reports collection failures without failing inference.
+
+Request latency histograms measure this worker's execution phase. For separate prefill and decode workers, they do not measure the client's full request latency. Rejected or interrupted requests may have no candidate timing observations. Health probes are excluded from request metrics, but their engine work can appear in iteration statistics.
+
+Engine statistics are collected approximately once per second. Gauges retain the last engine sample when no new sample is available. Served-model labels and `prometheus_metrics_config` histogram buckets match HTTP; OpenEngine's engine gauges omit `pid` labels.
+
+### Performance and deployment notes
+
+- Measure overhead with your workload and scrape interval. To sample KV statistics less often, increase `kv_cache_config.iteration_stats_interval`; this can delay updates. Keep `iter_perf_stats_interval: 1` for complete iteration counters, including speculative-decoding counts.
+- With multiple frontends, one process serves metrics and collects engine statistics; request metrics are combined across all frontends.
+- Metrics storage is managed automatically. If setting `PROMETHEUS_MULTIPROC_DIR`, use an existing writable parent directory; the server creates and cleans up a private child. Forced termination can leave that child behind. When embedding the server, initialize the environment before importing `prometheus_client`.
+
 ## Transport security
 
 The listener is plaintext h2c with no authentication. Any client that can reach