From d23a5fd214d9c0480c88335de0f956bf769a687e Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 10:19:15 -0700 Subject: [PATCH 1/8] feat(openengine): expose Prometheus metrics Signed-off-by: Tanmay Verma --- tensorrt_llm/commands/serve.py | 31 ++ tensorrt_llm/grpc/openengine/README.md | 76 +++ tensorrt_llm/grpc/openengine/metrics.py | 328 ++++++++++++ tensorrt_llm/grpc/openengine/server.py | 47 +- tensorrt_llm/grpc/openengine/servicer.py | 59 ++- tensorrt_llm/metrics/collector.py | 141 +++++- tensorrt_llm/serve/_perf_metrics_config.py | 35 ++ tensorrt_llm/serve/openai_server.py | 126 +---- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../references/trtllm_serve_cli.yaml | 18 + .../unittest/grpc/openengine/test_metrics.py | 468 ++++++++++++++++++ .../grpc/openengine/test_openengine.py | 111 +++++ .../unittest/grpc/openengine/test_servicer.py | 56 ++- 13 files changed, 1364 insertions(+), 133 deletions(-) create mode 100644 tensorrt_llm/grpc/openengine/metrics.py create mode 100644 tensorrt_llm/serve/_perf_metrics_config.py create mode 100644 tests/unittest/grpc/openengine/test_metrics.py diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 0873ddcf5251..0f24abb6233f 100644 --- a/tensorrt_llm/commands/serve.py +++ b/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 diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index e1347e9711d7..101a81024b4e 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -91,6 +91,82 @@ 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 + +Select collection with the same engine configuration fields used by the HTTP server. For request timing and token metrics without iteration collection, save this as `metrics.yaml`: + +```yaml +return_perf_metrics: true +enable_iter_perf_stats: false +``` + +Then expose an HTTP scrape endpoint alongside the OpenEngine gRPC listener: + +```bash +trtllm-serve Qwen/Qwen3-0.6B \ + --config metrics.yaml \ + --grpc --grpc-protocol openengine --port 50051 + +curl --fail http://127.0.0.1:8001/prometheus/metrics +``` + +Like the HTTP server, OpenEngine exposes Prometheus when `return_perf_metrics: true`. Configuring `perf_metrics_output_dir` also opts into request-performance collection and Prometheus, matching HTTP's existing enablement behavior. Otherwise, the endpoint stays disabled. Boolean values use the LLM schema's validation, so quoted values such as `"false"`, `"0"`, and `"off"` do not enable it. No metrics-specific enable flag is needed. The listener uses port `8001` and the existing `--host` bind address by default. Use `--openengine-metrics-port` or `--openengine-metrics-host` only to override the address. Specifying either override does not enable metrics. This HTTP endpoint is unauthenticated; restrict its network exposure. The metrics port must differ from the gRPC port. Only `/prometheus/metrics` is exposed on the HTTP listener. + +Iteration collection remains independently configured. Without `perf_metrics_output_dir`, the settings behave as follows: + +| `return_perf_metrics` | `enable_iter_perf_stats` | Collected measurements | +| --- | --- | --- | +| `false` | `false` | Prometheus endpoint disabled | +| `true` | `false` | Also request timing and token metrics; no iteration-statistics reader runs | +| `false` | `true` | Prometheus endpoint disabled; engine iteration collection remains enabled | +| `true` | `true` | Both request and iteration metrics | + +Both engine fields default to `false`. Set both to `true` for full collection. The exporter always collects inexpensive Generate RPC outcomes and static configuration while the listener is enabled. Disabled metric families receive no observations; some shared collector counters may still expose their initial zero value. Collection costs depend on the engine timing paths and workload: iteration-only collection is not necessarily faster than full collection. Benchmark the selected configuration on the target build. + +The separate Prometheus listener does not write HTTP per-request JSONL artifacts; matching endpoint enablement does not add that HTTP frontend feature to OpenEngine. + +Prometheus collection does not enable KV-event publishing or change `GetLoad`: `--openengine-enable-load-metrics` still independently controls routing-load snapshots. Health checks use `Control.Health`; a successful scrape only proves that the exporter is responding. The OpenEngine health RPC is not the standard Kubernetes gRPC health-check service. + +For full collection, the existing KV-cache sampling setting reduces how often KV statistics are read from the engine: + +```yaml +kv_cache_config: + iteration_stats_interval: 4 +``` + +The engine accumulates KV deltas between samples. Reports can therefore be delayed until another collection iteration occurs. This setting is independent of the one-second exporter polling cadence and the Prometheus scrape interval. Leave `iter_perf_stats_interval: 1` when exact counters from every executor iteration are required: increasing that broader sampling interval can undercount counters such as speculative-decoding draft/accepted totals, even though selected completion and KV counters are carried forward. Request latency/token accounting is collected separately at completion. + +The exporter reuses TensorRT-LLM's `trtllm_*` metrics, served-model labels, and `prometheus_metrics_config` histogram buckets: + +| Metric | Counting and timing contract | +| --- | --- | +| `trtllm_request_success_total` | Existing engine finish-reason accounting, per completed output candidate. A successful `num_sequences=2` RPC contributes two candidates. Use the RPC outcome metric below to distinguish cancellation and errors. | +| `trtllm_prompt_tokens_total` | Shared prompt counted once, rather than once per candidate. | +| `trtllm_generation_tokens_total` | Sum of generated tokens across candidates. | +| Request latency histograms | Engine-measured latency for this worker's phase; prefill/decode workers do not measure cross-worker client end-to-end latency. | +| Iteration counters and gauges | Engine statistics collected by one owner approximately once per second, including during active streams. All available iteration deltas are processed; gauges retain the last engine sample when no new sample is produced. Coverage follows the engine's statistics API; these are not a new per-rank reporting API. | +| `trtllm_openengine_rpc_outcomes_total{phase,outcome}` | One outcome per `Generate` RPC. Phases are `aggregated`, `prefill`, and `decode`; outcomes are `completed`, `rejected`, `engine_error`, `cancelled`, `timeout`, `stalled`, and `unavailable`. An in-band engine error is counted even when gRPC transport status is OK. | +| `trtllm_openengine_metrics_errors_total{source}` | Request or iteration metrics recording failures. These failures preserve inference responses. | + +RPC completion is recorded after the response stream has been consumed by the gRPC serving layer; for prefill this includes the handoff response. It does not acknowledge application-level receipt. Engine performance metrics are recorded once when the engine result is final; rejected or interrupted requests may have no candidate timing metrics. Internal `Control.Health` inference probes are excluded from Generate outcomes and candidate metrics, but their work is included in engine-level iteration statistics. + +With `--num_serve_frontends`, frontend zero owns one HTTP listener and the sole iteration-statistics consumer. Frontends contribute their own request counters and histograms to a shared registry. Engine gauges have one writer and export without PID labels. Scrape serialization and engine-statistics reads run outside the gRPC event loop; HTTP concurrency is bounded. A listener startup failure prevents serving readiness, and unexpected listener termination fails serving-group supervision. + +Each serving group uses a fresh temporary metrics directory, inherited by its attached frontends. If `PROMETHEUS_MULTIPROC_DIR` is supplied, it must be an existing writable parent directory; the server creates its own child and never deletes the parent's existing contents. Storage is removed after the group exits normally. A forced process kill may leave a child directory behind, but a new group uses a fresh child and does not aggregate the old samples. Initialize this environment before importing `prometheus_client` when embedding the server; an already initialized single-process client is rejected. + +Example Prometheus configuration: + +```yaml +scrape_configs: + - job_name: trtllm-openengine + scrape_interval: 15s + metrics_path: /prometheus/metrics + static_configs: + - targets: ["engine-host:8001"] +``` + +For deployment acceptance, verify `up{job="trtllm-openengine"} == 1`, run a known generation workload, and query the RPC outcome, token-counter, and latency-histogram changes. Compare metrics disabled, collection enabled without scrapes, and collection enabled at the intended scrape cadence using matched workloads and frontend counts. Use repeated runs to distinguish measurement noise; target at most 2% throughput loss and 5% p99 inter-token latency growth, and record CPU and resident memory as well. GPU inference and performance qualification require a compatible GPU, model weights, and the matching TensorRT-LLM build. + ## Transport security The listener is plaintext h2c with no authentication. Any client that can reach diff --git a/tensorrt_llm/grpc/openengine/metrics.py b/tensorrt_llm/grpc/openengine/metrics.py new file mode 100644 index 000000000000..467649401ab6 --- /dev/null +++ b/tensorrt_llm/grpc/openengine/metrics.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prometheus collection and the optional, leader-owned HTTP listener.""" + +import asyncio +import os +import socket +import tempfile +import threading +import time +from enum import Enum +from typing import TYPE_CHECKING, Any, Literal + +from tensorrt_llm.logger import logger +from tensorrt_llm.metrics.collector import MetricsCollector + +from .coordinator import CoordinationError + +if TYPE_CHECKING: + from tensorrt_llm.disaggregated_params import DisaggregatedParams + from tensorrt_llm.llmapi.llm import RequestOutput + + from .bindings.generation_pb2 import GenerateResponse + +_DIRECTORY_ENV = "TLLM_OPENENGINE_METRICS_DIR" +_COLLECTION_INTERVAL_SECONDS = 1.0 + + +class MetricsDirectory: + """Own a fresh serving-group directory; attached frontends only inherit it. + + The operator's PROMETHEUS_MULTIPROC_DIR is a parent directory, never a + directory whose existing contents this server may delete or aggregate. + """ + + def __init__(self, attached: bool) -> None: + self._directory = None + self._previous = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if attached: + path = os.environ.get(_DIRECTORY_ENV) + if not path or path != self._previous or not os.path.isdir(path): + raise RuntimeError("OpenEngine frontend has no shared Prometheus directory") + else: + self._directory = tempfile.TemporaryDirectory( + prefix="tllm-openengine-metrics-", dir=self._previous + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._directory.name + os.environ[_DIRECTORY_ENV] = self._directory.name + + def close(self) -> None: + """Remove owned storage after all serving-group processes have stopped.""" + if self._directory is not None: + self._directory.cleanup() + os.environ.pop(_DIRECTORY_ENV, None) + if self._previous is None: + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + else: + os.environ["PROMETHEUS_MULTIPROC_DIR"] = self._previous + + +class RpcOutcome(str, Enum): + REJECTED = "rejected" + COMPLETED = "completed" + CANCELLED = "cancelled" + UNAVAILABLE = "unavailable" + ENGINE_ERROR = "engine_error" + TIMEOUT = "timeout" + STALLED = "stalled" + + +class RequestMetrics: + """Own one RPC's classification and exactly-once outcome recording.""" + + def __init__(self, metrics: "OpenEngineMetrics") -> None: + self._metrics = metrics + self._phase: Literal["aggregated", "prefill", "decode"] = "aggregated" + self._finished = False + self._outcome = RpcOutcome.REJECTED + + def begin(self, params: "DisaggregatedParams | None") -> None: + if params is None: + return + if params.request_type == "context_only": + self._phase = "prefill" + elif params.request_type == "generation_only": + self._phase = "decode" + else: + self._phase = "aggregated" + + def accepted(self) -> None: + # An accepted request that closes before a terminal result is cancelled. + self._outcome = RpcOutcome.CANCELLED + + def observe_exception(self, error: BaseException) -> None: + if self._finished: + return + # RPC cancellation and coordination failure describe the eventual + # transport result, even if the engine already completed or failed. + if isinstance(error, asyncio.CancelledError): + self._outcome = RpcOutcome.CANCELLED + elif isinstance(error, CoordinationError): + self._outcome = RpcOutcome.UNAVAILABLE + else: + self._outcome = RpcOutcome.ENGINE_ERROR + + def engine_failed(self) -> None: + """An engine operation failed; its exception type does not change its origin.""" + if not self._finished: + self._outcome = RpcOutcome.ENGINE_ERROR + + def observe_response(self, response: "GenerateResponse") -> None: + if response.HasField("error"): + self.engine_failed() + + def stalled(self) -> None: + if not self._finished: + self._outcome = RpcOutcome.STALLED + self.finish() + + def complete(self, result: "RequestOutput", *, handoff_ready: bool) -> None: + if self._finished: + return + # An error already sent to the client survives normal engine completion. + if self._outcome != RpcOutcome.ENGINE_ERROR: + reasons = {output.finish_reason for output in result.outputs} + if "cancelled" in reasons: + self._outcome = RpcOutcome.CANCELLED + elif "timeout" in reasons: + self._outcome = RpcOutcome.TIMEOUT + elif self._phase == "prefill" and not handoff_ready: + self._outcome = RpcOutcome.ENGINE_ERROR + else: + self._outcome = RpcOutcome.COMPLETED + # Engine candidate accounting is independent of the RPC's outcome. + self._metrics.record_result(result) + + def finish(self) -> None: + """Record at most once, including when a suspended stream resumes later.""" + if self._finished: + return + self._finished = True + self._metrics.record_outcome(self._phase, self._outcome.value) + + +class OpenEngineMetrics: + """Collect frontend metrics; only the engine owner serves and drains stats. + + Args: + llm: Initialized LLM with engine performance statistics enabled. + model: Stable served model label; request-supplied model aliases are ignored. + host: HTTP bind address. + port: HTTP port, or zero for an ephemeral port in embedded use. + owner: Whether this process owns the engine and HTTP endpoint. + """ + + def __init__(self, llm: Any, model: str, host: str, port: int, owner: bool) -> None: + # All imports follow MetricsDirectory creation and precede construction + # of any metric, including in re-executed attached frontends. + from prometheus_client import CollectorRegistry, Counter, make_asgi_app, multiprocess + from prometheus_client.values import ValueClass + + if not ValueClass._multiprocess: + raise RuntimeError( + "OpenEngine metrics require PROMETHEUS_MULTIPROC_DIR before importing " + "prometheus_client; start a fresh trtllm-serve process" + ) + self._llm = llm + self.host = host + self.port = port + self.owner = owner + self._collect_requests = llm.args.return_perf_metrics + self._collect_iterations = llm.args.enable_iter_perf_stats + registry = CollectorRegistry() + self.collector = MetricsCollector.from_args( + llm.args, model, registry=registry, gauge_multiprocess_mode="livemostrecent" + ) + self._outcomes = Counter( + "trtllm_openengine_rpc_outcomes_total", + "Generate RPC outcomes, independent of output-candidate completion metrics.", + ["model_name", "phase", "outcome"], + registry=registry, + ) + self._model = model + self._errors = Counter( + "trtllm_openengine_metrics_errors_total", + "Metrics collection failures; inference results are preserved.", + ["model_name", "source"], + registry=registry, + ) + self._last_error_log = 0.0 + self._http = None + self._thread = None + self._socket = None + self._task = None + self._read_task = None + self._stopping = threading.Event() + if owner: + self.collector.log_config_from_args(llm.args, model) + scrape_registry = CollectorRegistry() + multiprocess.MultiProcessCollector(scrape_registry) + prometheus_app = make_asgi_app(registry=scrape_registry) + + async def app(scope, receive, send) -> None: + if scope["path"] != "/prometheus/metrics": + await send({"type": "http.response.start", "status": 404, "headers": []}) + await send({"type": "http.response.body", "body": b"Not found"}) + return + await prometheus_app(scope, receive, send) + + import uvicorn + + self._http = uvicorn.Server( + uvicorn.Config( + app, + host=host, + port=port, + lifespan="off", + access_log=False, + limit_concurrency=4, + backlog=8, + timeout_keep_alive=5, + timeout_graceful_shutdown=2, + ) + ) + + def _record_error(self, source: str, error: Exception) -> None: + # Telemetry errors must not turn successful generation into engine errors. + if self._stopping.is_set(): + return + try: + self._errors.labels(self._model, source).inc() + except Exception: + pass + now = time.monotonic() + if now - self._last_error_log >= 30: + self._last_error_log = now + logger.warning(f"OpenEngine metrics {source} failed: {error}") + + def record_outcome(self, phase: str, outcome: str) -> None: + """Record an RPC result without affecting inference.""" + try: + self._outcomes.labels(self._model, phase, outcome).inc() + except Exception as error: + self._record_error("request", error) + + def record_result(self, result: Any) -> None: + """Record newly completed candidates, following the HTTP collector contract.""" + if not self._collect_requests: + return + try: + candidates = getattr(result, "candidate_metrics", None) + if candidates: + for metrics in candidates: + self.collector.log_request_metrics_dict(metrics) + elif metrics := getattr(result, "metrics_dict", None): + self.collector.log_request_metrics_dict(metrics) + except Exception as error: + self._record_error("request", error) + + async def start(self) -> None: + """Bind the HTTP listener and start periodic collection in the owner only.""" + if not self.owner: + return + family = socket.getaddrinfo(self.host, self.port, type=socket.SOCK_STREAM)[0][0] + self._socket = socket.create_server((self.host, self.port), family=family, backlog=8) + self.port = self._socket.getsockname()[1] + self._thread = threading.Thread( + target=self._http.run, + kwargs={"sockets": [self._socket]}, + name="openengine-prometheus", + daemon=True, + ) + self._thread.start() + deadline = time.monotonic() + 5 + while not self._http.started: + if not self._thread.is_alive() or time.monotonic() >= deadline: + raise RuntimeError("OpenEngine Prometheus listener failed to start") + await asyncio.sleep(0.01) + if self._collect_iterations: + self._task = asyncio.create_task(self._collect(), name="openengine-metrics") + logger.info(f"OpenEngine Prometheus endpoint: {self.host}:{self.port}/prometheus/metrics") + + def check_health(self) -> None: + """Fail supervision if the requested endpoint or collection task exited.""" + if self.owner and ( + self._thread is None + or not self._thread.is_alive() + or (self._collect_iterations and (self._task is None or self._task.done())) + ): + raise RuntimeError("OpenEngine Prometheus service stopped unexpectedly") + + async def _collect(self) -> None: + while True: + # The proxy's ostensibly async stats API performs a synchronous RPC. + # Both that read and metric-file updates stay off the serving loop. + self._read_task = asyncio.create_task(asyncio.to_thread(self._collect_once)) + await asyncio.shield(self._read_task) + await asyncio.sleep(_COLLECTION_INTERVAL_SECONDS) + + def _collect_once(self) -> None: + try: + for stats in self._llm.get_stats(timeout=0.1): + if self._stopping.is_set(): + return + self.collector.log_iteration_stats(stats) + except Exception as error: + self._record_error("iteration", error) + + async def stop(self) -> None: + """Stop tasks and listener before the launcher removes shared storage.""" + self._stopping.set() + if self._task is not None: + self._task.cancel() + await asyncio.gather(self._task, return_exceptions=True) + if self._http is not None: + self._http.should_exit = True + if self._thread is not None: + await asyncio.to_thread(self._thread.join, 5) + if self._socket is not None: + self._socket.close() + if self._thread is not None and self._thread.is_alive(): + raise RuntimeError("OpenEngine Prometheus listener did not stop") + + async def wait_for_collection(self) -> None: + """Join the last read after engine shutdown releases pending executor RPCs.""" + if self._read_task is not None: + await asyncio.gather(self._read_task, return_exceptions=True) diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index ea888ffd2f6d..ff5b4da74843 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -29,11 +29,13 @@ _spawn_attached_frontends, _terminate_attached_frontends, ) +from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from .bindings import openengine_pb2_grpc from .control import OpenEngineControlServicer from .coordinator import CoordinationError, Coordinator, FrontendClient from .kv_events import KvEventsUnavailableError, _format_host_port, events_config +from .metrics import MetricsDirectory, OpenEngineMetrics from .servicer import OpenEngineInferenceServicer __all__ = ["OpenEngineServer", "launch_server"] @@ -106,10 +108,12 @@ def __init__( frontend: FrontendClient | None = None, instance_id: str | None = None, subagent_affinity_auth_key: str | None = None, + metrics: OpenEngineMetrics | None = None, ) -> None: self.host = host self.port = port self._llm = llm + self._metrics = metrics options = list(_SERVER_OPTIONS) options.append(("grpc.so_reuseport", 1 if frontend is not None else 0)) self._server = grpc.aio.server(options=options) @@ -120,6 +124,7 @@ def __init__( kv_transfer_backend=kv_transfer_backend, frontend=frontend, subagent_affinity_auth_key=subagent_affinity_auth_key, + metrics=metrics, ) openengine_pb2_grpc.add_InferenceServicer_to_server(inference, self._server) # Control shares the inference servicer's in-flight request table so @@ -161,6 +166,8 @@ async def start(self) -> None: "KV events are configured but the loaded engine has no streaming publisher; " "set kv_cache_config.use_kv_cache_manager_v2: true and use a supported model" ) + if self._metrics is not None: + await self._metrics.start() await self._server.start() address = _format_bind_address(self.host, self.port) logger.info(f"OpenEngine server started on {address}") @@ -171,7 +178,11 @@ async def stop(self, grace: float = 5.0) -> None: Args: grace: Maximum time in seconds to allow active RPCs to finish. """ - await self._server.stop(grace=grace) + try: + await self._server.stop(grace=grace) + finally: + if self._metrics is not None: + await self._metrics.stop() logger.info("OpenEngine server stopped") async def wait_for_termination(self) -> None: @@ -199,6 +210,8 @@ def launch_server( report_failure: Callable[[int, str, str], None] | None = None, enable_load_metrics: bool = False, subagent_affinity_auth_key: str | None = None, + metrics_host: str | None = None, + metrics_port: int = 8001, ) -> None: """Launch the dedicated OpenEngine gRPC server. @@ -210,6 +223,8 @@ 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. + metrics_host: HTTP bind address; defaults to the inference host. + metrics_port: HTTP port when performance-metrics configuration enables Prometheus. """ _warn_if_python_protobuf() @@ -224,6 +239,16 @@ async def serve() -> None: param_hint="backend", ) + perf_metrics = resolve_perf_metrics_config( + llm_args.get("return_perf_metrics", False), + llm_args.get("perf_metrics_output_dir"), + ) + metrics_enabled = perf_metrics.collect_request_metrics + if metrics_enabled: + llm_args["return_perf_metrics"] = True + if metrics_enabled and metrics_port == port and port != 0: + raise click.UsageError("--openengine-metrics-port must differ from --port") + mode = _init_multi_frontend_mode(llm_args, enabled=True) if mode.num_frontends > 1 and port == 0: raise click.UsageError("Multiple OpenEngine frontends require a fixed --port") @@ -238,8 +263,10 @@ async def serve() -> None: startup_cancelled = threading.Event() server = llm = coordinator = frontend = monitor = None directory = None + metrics_directory = metrics = None children = [] failure = None + serving = False parent_pid = os.getppid() def fail(reason: str) -> None: @@ -278,6 +305,8 @@ async def supervise() -> None: if not healthy(): fail("OpenEngine engine health check failed") return + if serving and metrics is not None: + metrics.check_health() if frontend is not None: try: status = await frontend.request("heartbeat", frontend=frontend.frontend_id) @@ -300,6 +329,8 @@ def monitor_done(task: asyncio.Task) -> None: fail(f"OpenEngine supervision failed: {error}") try: + if metrics_enabled: + metrics_directory = MetricsDirectory(mode.is_attached_frontend) llm_args["_enable_routing_load"] = enable_load_metrics # A new incarnation fences every node's metadata after a restart. # The executor starts node discovery only when streaming is enabled. @@ -314,6 +345,14 @@ def monitor_done(task: asyncio.Task) -> None: raise RuntimeError(failure) return logger.info("Model loaded successfully") + if metrics_enabled: + metrics = OpenEngineMetrics( + llm, + model, + host if metrics_host is None else metrics_host, + metrics_port, + owner=not mode.is_attached_frontend, + ) instance_id = str(llm.llm_id) if mode.is_launcher: directory = tempfile.TemporaryDirectory(prefix="tllm-openengine-") @@ -371,9 +410,11 @@ def monitor_done(task: asyncio.Task) -> None: frontend=frontend, instance_id=instance_id, subagent_affinity_auth_key=subagent_affinity_auth_key, + metrics=metrics, ) _disable_gc_if_requested() await server.start() + serving = True if any(child.poll() is not None for child in children): raise RuntimeError("An OpenEngine frontend exited during startup") _signal_frontend_ready(mode) @@ -409,8 +450,12 @@ def monitor_done(task: asyncio.Task) -> None: if llm is not None: await asyncio.to_thread(llm.shutdown) finally: + if metrics is not None: + await metrics.wait_for_collection() if directory is not None: directory.cleanup() + if metrics_directory is not None: + metrics_directory.close() for sig, previous in previous_handlers.items(): signal.signal(sig, previous) logger.info("LLM engine stopped") diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index 6a9d6e9018b9..f6826401dfbc 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -27,6 +27,7 @@ from .disagg import disaggregated_params_from_request from .errors import AbortFailedError, UnsupportedFeatureError from .formatting import _engine_error_response, _stop_texts +from .metrics import OpenEngineMetrics, RequestMetrics from .request_mapping import ( _input_from_request, _metadata_from_context, @@ -53,6 +54,7 @@ def __init__( kv_transfer_backend: str = "", frontend: FrontendClient | None = None, subagent_affinity_auth_key: str | None = None, + metrics: OpenEngineMetrics | None = None, ) -> None: self._frontend = frontend self._llm = llm @@ -64,6 +66,7 @@ def __init__( self._guided_backend = llm.args.guided_decoding_backend self.subagent_affinity_auth_key = subagent_affinity_auth_key self._active_requests: dict[str, Any] = {} + self._metrics = metrics def active_request_count(self) -> int: """Requests accepted and not yet completed.""" @@ -121,6 +124,7 @@ async def Generate( context: grpc.aio.ServicerContext, ) -> AsyncIterator[generation_pb2.GenerateResponse]: reservation = None + request_metrics = RequestMetrics(self._metrics) if self._metrics else None try: if self._frontend is not None: if not request.request_id: @@ -128,30 +132,44 @@ async def Generate( grpc.StatusCode.INVALID_ARGUMENT, "request_id must be non-empty" ) reservation = await self._frontend.reserve(request.request_id) - async with aclosing(self._generate(request, context, reservation)) as responses: + async with aclosing( + self._generate(request, context, reservation, request_metrics) + ) as responses: async for response in responses: yield response except DuplicateRequestError as error: await context.abort(grpc.StatusCode.ALREADY_EXISTS, str(error)) except InvalidRequestIdError as error: await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) + except asyncio.CancelledError as error: + if request_metrics is not None: + request_metrics.observe_exception(error) + raise except CoordinationError as error: + if request_metrics is not None: + request_metrics.observe_exception(error) await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) finally: - if reservation is not None: - try: + try: + if reservation is not None: await asyncio.shield(reservation.release()) - except CoordinationError as error: - # Cleanup still fails the serving group through FrontendClient. - # Preserve cancellation or an already selected RPC error. - if not context.cancelled() and context.code() in (None, grpc.StatusCode.OK): - await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) + except CoordinationError as error: + if request_metrics is not None: + request_metrics.observe_exception(error) + # Cleanup still fails the serving group through FrontendClient. + # Preserve cancellation or an already selected RPC error. + if not context.cancelled() and context.code() in (None, grpc.StatusCode.OK): + await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) + finally: + if request_metrics is not None: + request_metrics.finish() async def _generate( self, request: generation_pb2.GenerateRequest, context: grpc.aio.ServicerContext, reservation: Reservation | None = None, + request_metrics: RequestMetrics | None = None, ) -> AsyncIterator[generation_pb2.GenerateResponse]: """Run a server-streaming OpenEngine generation request.""" request_id = request.request_id @@ -182,6 +200,8 @@ async def _generate( trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context) conversation_params = conversation_params_from_request(request) disaggregated_params = disaggregated_params_from_request(request) + if request_metrics is not None: + request_metrics.begin(disaggregated_params) parent_affinity_id = validate_parent_affinity( self.subagent_affinity_auth_key, request, @@ -272,6 +292,8 @@ async def _generate( # router treat a totally broken worker as a fleet of bad clients: it # would not retry elsewhere, eject the engine, or alarm. except TypeError as error: + if request_metrics is not None: + request_metrics.engine_failed() logger.error(f"OpenEngine request {request_id} could not be submitted: {error}") await context.abort(grpc.StatusCode.INTERNAL, str(error)) return @@ -279,12 +301,16 @@ async def _generate( # cannot enumerate here; any such failure must become a clean INTERNAL # status for the client rather than an unhandled servicer crash. except Exception as error: + if request_metrics is not None: + request_metrics.engine_failed() logger.error(f"Failed to submit OpenEngine request {request_id}: {error}") await context.abort(grpc.StatusCode.INTERNAL, str(error)) return if reservation is not None: reservation.bind(result_handle) + if request_metrics is not None: + request_metrics.accepted() engine_terminal = False abort_requested = False @@ -305,6 +331,8 @@ def abort_request(reason: str) -> None: ) def on_stall() -> None: + if request_metrics is not None: + request_metrics.stalled() logger.warning( f"OpenEngine request {request_id} stopped consuming responses " f"for {RESPONSE_STALL_TIMEOUT_SECONDS:g} seconds" @@ -347,6 +375,8 @@ def on_rpc_done(_: asyncio.Task) -> None: return responses = _format_result(result, format_state) for resp in responses or (): + if request_metrics is not None: + request_metrics.observe_response(resp) watchdog.pending() yield resp watchdog.delivered() @@ -365,6 +395,10 @@ def on_rpc_done(_: asyncio.Task) -> None: engine_terminal = True return if result.finished: + if request_metrics is not None: + request_metrics.complete( + result, handoff_ready=bool(format_state.prefill_ready_sent) + ) engine_terminal = True return @@ -373,18 +407,23 @@ def on_rpc_done(_: asyncio.Task) -> None: # internal engine abort). Surface it as an error rather than silently # returning a truncated result to the client. engine_terminal = True - watchdog.pending() - yield _engine_error_response( + response = _engine_error_response( request_id, "generation stream ended before all outputs finished", result_handle, ) + if request_metrics is not None: + request_metrics.observe_response(response) + watchdog.pending() + yield response except asyncio.CancelledError: raise # Broad by design: this is the RPC boundary, so any engine/processing # failure after acceptance must be reported to the client as an error # event and must still trigger abort + cleanup in the finally below. except Exception as error: + if request_metrics is not None: + request_metrics.engine_failed() logger.error(f"OpenEngine request {request_id} failed after acceptance: {error}") abort_request("response processing failed") engine_terminal = True diff --git a/tensorrt_llm/metrics/collector.py b/tensorrt_llm/metrics/collector.py index 50846d323599..35faa80b8c11 100644 --- a/tensorrt_llm/metrics/collector.py +++ b/tensorrt_llm/metrics/collector.py @@ -16,7 +16,13 @@ import math import time -from typing import Dict, List, Optional, Union +from functools import partial +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +if TYPE_CHECKING: + from prometheus_client import CollectorRegistry + + from tensorrt_llm.llmapi.llm_args import BaseLlmArgs from .enums import MetricNames @@ -116,6 +122,36 @@ class MetricsCollector: labelname_cache_level = "cache_level" labelname_cache_tier = "cache_tier" + @classmethod + def from_args( + cls, + args: "BaseLlmArgs", + model: str, + *, + registry: "Optional[CollectorRegistry]" = None, + gauge_multiprocess_mode: str = "all", + ) -> "MetricsCollector": + """Build a collector using the shared serving labels and bucket configuration.""" + config = getattr(args, "prometheus_metrics_config", None) + buckets = { + name: getattr(config, name, None) + for name in ("e2e_request_latency_buckets", + "time_to_first_token_buckets", + "time_per_output_token_buckets", + "request_queue_time_buckets", + "request_prefill_time_buckets", + "request_decode_time_buckets", + "request_inference_time_buckets") + } + return cls( + { + "model_name": model, + "engine_type": args.backend or "unknown" + }, + **buckets, + registry=registry, + gauge_multiprocess_mode=gauge_multiprocess_mode) + def __init__( self, labels: Dict[str, str], @@ -126,8 +162,17 @@ def __init__( request_prefill_time_buckets: Optional[List[float]] = None, request_decode_time_buckets: Optional[List[float]] = None, request_inference_time_buckets: Optional[List[float]] = None, + registry: "Optional[CollectorRegistry]" = None, + gauge_multiprocess_mode: str = "all", ) -> None: - from prometheus_client import Counter, Gauge, Histogram + from prometheus_client import REGISTRY, Counter, Gauge, Histogram + self._registry = REGISTRY if registry is None else registry + self._gauge_multiprocess_mode = gauge_multiprocess_mode + Counter = partial(Counter, registry=self._registry) + Histogram = partial(Histogram, registry=self._registry) + Gauge = partial(Gauge, + registry=self._registry, + multiprocess_mode=gauge_multiprocess_mode) _bucket_params = { "e2e_request_latency_buckets": e2e_request_latency_buckets, "time_to_first_token_buckets": time_to_first_token_buckets, @@ -530,6 +575,95 @@ def __init__( documentation="Total request errors, labeled by HTTP status code.", labelnames=self.labels_with_http_code.keys()) + def log_config_from_args(self, args: "BaseLlmArgs", model: str) -> None: + """Extract configuration from generator args and log as Prometheus info gauges.""" + from tensorrt_llm.logger import logger + + # Model config + model_config = { + "model": str(args.model), + "served_model_name": model, + "dtype": str(args.dtype), + } + quant_config = getattr(args, "quant_config", None) + if quant_config is not None: + quant_algo = getattr(quant_config, "quant_algo", None) + model_config["quantization"] = str( + quant_algo) if quant_algo else "none" + else: + model_config["quantization"] = "none" + max_seq_len = getattr(args, "max_seq_len", None) + if max_seq_len is not None: + model_config["max_model_len"] = str(max_seq_len) + try: + import torch + if torch.cuda.is_available(): + model_config["gpu_type"] = torch.cuda.get_device_name(0) + except (ImportError, RuntimeError) as e: + logger.debug("Could not detect GPU type for config metrics: %s", e) + + # Parallel config — prefer parallel_config from generator args + # for accurate values including cp_size and world_size. + par_cfg = getattr(args, "parallel_config", None) + if par_cfg is not None: + tp_size = getattr(par_cfg, "tp_size", 1) or 1 + pp_size = getattr(par_cfg, "pp_size", 1) or 1 + cp_size = getattr(par_cfg, "cp_size", 1) or 1 + world_size = getattr(par_cfg, "world_size", + tp_size * pp_size * cp_size) + else: + tp_size = getattr(args, "tensor_parallel_size", 1) or 1 + pp_size = getattr(args, "pipeline_parallel_size", 1) or 1 + cp_size = 1 + world_size = tp_size * pp_size * cp_size + parallel_config = { + "tensor_parallel_size": str(tp_size), + "pipeline_parallel_size": str(pp_size), + "context_parallel_size": str(cp_size), + "gpu_count": str(world_size), + } + ep_size = getattr(par_cfg, "moe_ep_size", None) if par_cfg else \ + getattr(args, "moe_expert_parallel_size", None) + if ep_size is not None and ep_size > 0: + parallel_config["expert_parallel_size"] = str(ep_size) + + # Speculative decoding config + spec_config_obj = getattr(args, "speculative_config", None) or getattr( + args, "decoding_config", None) + speculative_config = None + if spec_config_obj is not None: + speculative_config = {"spec_enabled": "true"} + decoding_type = getattr(spec_config_obj, "decoding_type", None) + if decoding_type is not None: + speculative_config["spec_method"] = str(decoding_type) + max_draft_len = getattr(spec_config_obj, "max_draft_len", None) + if max_draft_len is not None: + speculative_config["spec_num_tokens"] = str(max_draft_len) + draft_model = getattr(spec_config_obj, "speculative_model", None) + if draft_model is not None: + speculative_config["spec_draft_model"] = str(draft_model) + + # KV cache config + kv_cache_config_obj = getattr(args, "kv_cache_config", None) + kv_cache_config = None + if kv_cache_config_obj is not None: + kv_cache_config = {} + for field in ("page_size", "enable_block_reuse", + "enable_partial_reuse", "free_gpu_memory_fraction"): + val = getattr(kv_cache_config_obj, field, None) + if val is not None: + kv_cache_config[field] = str(val) + kv_dtype = getattr(kv_cache_config_obj, "dtype", None) + if kv_dtype is not None: + kv_cache_config["cache_dtype"] = str(kv_dtype) + + self.log_config_info( + model_config=model_config, + parallel_config=parallel_config, + speculative_config=speculative_config, + kv_cache_config=kv_cache_config if kv_cache_config else None, + ) + def log_config_info( self, model_config: Optional[Dict[str, str]] = None, @@ -548,6 +682,9 @@ def log_config_info( kv_cache_config: KV cache configuration labels (page_size, enable_block_reuse, etc.) """ from prometheus_client import Gauge + Gauge = partial(Gauge, + registry=self._registry, + multiprocess_mode=self._gauge_multiprocess_mode) if model_config: info_labels = {**self.labels, **model_config} diff --git a/tensorrt_llm/serve/_perf_metrics_config.py b/tensorrt_llm/serve/_perf_metrics_config.py new file mode 100644 index 000000000000..25d56fbe199d --- /dev/null +++ b/tensorrt_llm/serve/_perf_metrics_config.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared performance-metrics enablement for serving frontends.""" + +from dataclasses import dataclass + +from pydantic import TypeAdapter + +from tensorrt_llm.llmapi.llm_args import BaseLlmArgs + +_RETURN_PERF_METRICS = TypeAdapter(BaseLlmArgs.model_fields["return_perf_metrics"].annotation) + + +@dataclass(frozen=True) +class PerfMetricsConfig: + expose_response_metrics: bool + collect_request_metrics: bool + + +def resolve_perf_metrics_config( + return_perf_metrics: object = False, + perf_metrics_output_dir: str | None = None, +) -> PerfMetricsConfig: + """Validate the opt-in and distinguish response exposure from collection. + + File output enables request-performance collection and Prometheus, while + response headers still require the original return_perf_metrics opt-in. + No Prometheus imports occur here: directory preparation must come first. + """ + expose = _RETURN_PERF_METRICS.validate_python(return_perf_metrics) + return PerfMetricsConfig( + expose_response_metrics=expose, + collect_request_metrics=expose or perf_metrics_output_dir is not None, + ) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index ac186ca15ad3..9047548e6776 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -67,6 +67,7 @@ from tensorrt_llm.runtime.kv_cache_hash import \ get_effective_kv_cache_event_hash_algo from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams +from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from tensorrt_llm.serve.anthropic_adapter import ( AnthropicRequestError, AnthropicResponseError, anthropic_error_response, convert_anthropic_count_tokens_request, convert_anthropic_request, @@ -704,12 +705,13 @@ def __init__( self.model = model self.metrics_collector = None args = getattr(self.generator, "args", None) - self._expose_perf_metrics = bool( - args and getattr(args, "return_perf_metrics", False)) perf_metrics_output_dir = (getattr(args, "perf_metrics_output_dir", None) if args else None) - self._collect_perf_metrics = (self._expose_perf_metrics - or perf_metrics_output_dir is not None) + perf_metrics = resolve_perf_metrics_config( + getattr(args, "return_perf_metrics", False), + perf_metrics_output_dir) + self._expose_perf_metrics = perf_metrics.expose_response_metrics + self._collect_perf_metrics = perf_metrics.collect_request_metrics # Per-request spec-decode acceptance stats. Deliberately independent of # return_perf_metrics: coupling them would mean asking for acceptance # numbers silently mounts the Prometheus endpoint too. @@ -998,29 +1000,10 @@ def _init_llm(self, chat_template: Optional[str] = None): if self.generator.args.return_perf_metrics: set_prometheus_multiproc_dir() - args = self.generator.args - pmc = getattr(args, "prometheus_metrics_config", None) - self.metrics_collector = MetricsCollector( - { - "model_name": self.model, - "engine_type": args.backend or "unknown" - }, - e2e_request_latency_buckets=(pmc.e2e_request_latency_buckets - if pmc else None), - time_to_first_token_buckets=(pmc.time_to_first_token_buckets - if pmc else None), - time_per_output_token_buckets=(pmc.time_per_output_token_buckets - if pmc else None), - request_queue_time_buckets=(pmc.request_queue_time_buckets - if pmc else None), - request_prefill_time_buckets=(pmc.request_prefill_time_buckets - if pmc else None), - request_decode_time_buckets=(pmc.request_decode_time_buckets - if pmc else None), - request_inference_time_buckets=( - pmc.request_inference_time_buckets if pmc else None), - ) - self._log_config_info_metrics() + self.metrics_collector = MetricsCollector.from_args( + self.generator.args, self.model) + self.metrics_collector.log_config_from_args(self.generator.args, + self.model) @staticmethod def _ensure_post_processor_hook_supported( @@ -1075,95 +1058,6 @@ def _has_cache_transceiver_config(self) -> bool: return (cache_transceiver_config is not None and getattr( cache_transceiver_config, "backend", None) is not None) - def _log_config_info_metrics(self) -> None: - """Extract configuration from generator args and log as Prometheus info gauges.""" - args = self.generator.args - - # Model config - model_config = { - "model": str(args.model), - "served_model_name": self.model, - "dtype": str(args.dtype), - } - quant_config = getattr(args, "quant_config", None) - if quant_config is not None: - quant_algo = getattr(quant_config, "quant_algo", None) - model_config["quantization"] = str( - quant_algo) if quant_algo else "none" - else: - model_config["quantization"] = "none" - max_seq_len = getattr(args, "max_seq_len", None) - if max_seq_len is not None: - model_config["max_model_len"] = str(max_seq_len) - try: - import torch - if torch.cuda.is_available(): - model_config["gpu_type"] = torch.cuda.get_device_name(0) - except (ImportError, RuntimeError) as e: - logger.debug("Could not detect GPU type for config metrics: %s", e) - - # Parallel config — prefer parallel_config from generator args - # for accurate values including cp_size and world_size. - par_cfg = getattr(args, "parallel_config", None) - if par_cfg is not None: - tp_size = getattr(par_cfg, "tp_size", 1) or 1 - pp_size = getattr(par_cfg, "pp_size", 1) or 1 - cp_size = getattr(par_cfg, "cp_size", 1) or 1 - world_size = getattr(par_cfg, "world_size", - tp_size * pp_size * cp_size) - else: - tp_size = getattr(args, "tensor_parallel_size", 1) or 1 - pp_size = getattr(args, "pipeline_parallel_size", 1) or 1 - cp_size = 1 - world_size = tp_size * pp_size * cp_size - parallel_config = { - "tensor_parallel_size": str(tp_size), - "pipeline_parallel_size": str(pp_size), - "context_parallel_size": str(cp_size), - "gpu_count": str(world_size), - } - ep_size = getattr(par_cfg, "moe_ep_size", None) if par_cfg else \ - getattr(args, "moe_expert_parallel_size", None) - if ep_size is not None and ep_size > 0: - parallel_config["expert_parallel_size"] = str(ep_size) - - # Speculative decoding config - spec_config_obj = getattr(args, "speculative_config", None) or getattr( - args, "decoding_config", None) - speculative_config = None - if spec_config_obj is not None: - speculative_config = {"spec_enabled": "true"} - decoding_type = getattr(spec_config_obj, "decoding_type", None) - if decoding_type is not None: - speculative_config["spec_method"] = str(decoding_type) - max_draft_len = getattr(spec_config_obj, "max_draft_len", None) - if max_draft_len is not None: - speculative_config["spec_num_tokens"] = str(max_draft_len) - draft_model = getattr(spec_config_obj, "speculative_model", None) - if draft_model is not None: - speculative_config["spec_draft_model"] = str(draft_model) - - # KV cache config - kv_cache_config_obj = getattr(args, "kv_cache_config", None) - kv_cache_config = None - if kv_cache_config_obj is not None: - kv_cache_config = {} - for field in ("page_size", "enable_block_reuse", - "enable_partial_reuse", "free_gpu_memory_fraction"): - val = getattr(kv_cache_config_obj, field, None) - if val is not None: - kv_cache_config[field] = str(val) - kv_dtype = getattr(kv_cache_config_obj, "dtype", None) - if kv_dtype is not None: - kv_cache_config["cache_dtype"] = str(kv_dtype) - - self.metrics_collector.log_config_info( - model_config=model_config, - parallel_config=parallel_config, - speculative_config=speculative_config, - kv_cache_config=kv_cache_config if kv_cache_config else None, - ) - @staticmethod def _attach_multimodal_data( prompt: Dict[str, Any], diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 86d013d5461b..109abecc178b 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -192,6 +192,7 @@ l0_cpu: - unittest/grpc/openengine/test_control.py - unittest/grpc/openengine/test_node.py - unittest/grpc/openengine/test_openengine.py + - unittest/grpc/openengine/test_metrics.py - unittest/grpc/openengine/test_servicer.py - unittest/grpc/smg/test_smg.py::test_streaming_deltas_preserve_sequences_logprobs_and_final_tokens - unittest/grpc/smg/test_smg.py::test_streaming_empty_token_history diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 8c06a6d0ea57..e40f5bccd0a4 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -378,6 +378,24 @@ commands: is_flag: true flags: - "--openengine-enable-load-metrics" + openengine_metrics_host: + type: str + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--openengine-metrics-host" + openengine_metrics_port: + type: int + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--openengine-metrics-port" otlp_traces_endpoint: type: str default: null diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py new file mode 100644 index 000000000000..9a85fa0ccddc --- /dev/null +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -0,0 +1,468 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Observable exporter contracts in fresh processes with real Prometheus registries.""" + +import asyncio +import os +import socket +import subprocess +import sys +import time +from pathlib import Path +from types import SimpleNamespace +from urllib.error import HTTPError +from urllib.request import urlopen + +import pytest + +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + + +def _run_worker(mode: str) -> None: + result = subprocess.run( + [sys.executable, __file__, mode], capture_output=True, text=True, timeout=90, check=False + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize( + "mode", + [ + "serving", + "multiprocess", + "lifecycle", + "collection_requests", + "collection_iterations", + "http_enablement", + ], +) +def test_prometheus_contract(mode: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Protect counts, group isolation, and real listener cleanup across process boundaries.""" + # Existing operator-owned samples must neither be scraped nor deleted. + sentinel = tmp_path / "counter_999.db" + sentinel.write_bytes(b"operator-owned metrics") + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + _run_worker(mode) + assert list(tmp_path.iterdir()) == [sentinel] + assert sentinel.read_bytes() == b"operator-owned metrics" + + +def _sample(text: str, name: str, **labels: str) -> float: + from prometheus_client.parser import text_string_to_metric_families + + return sum( + sample.value + for family in text_string_to_metric_families(text) + for sample in family.samples + if sample.name == name + and all(sample.labels.get(key) == value for key, value in labels.items()) + ) + + +async def _scrape(port: int) -> str: + def read() -> str: + with urlopen(f"http://127.0.0.1:{port}/prometheus/metrics", timeout=5) as response: + assert response.status == 200 + assert "text/plain" in response.headers["Content-Type"] + return response.read().decode() + + return await asyncio.to_thread(read) + + +def _args() -> SimpleNamespace: + return SimpleNamespace( + model="test-model", + dtype="auto", + backend="pytorch", + return_perf_metrics=True, + enable_iter_perf_stats=True, + guided_decoding_backend=None, + reasoning_parser=None, + max_seq_len=100, + max_batch_size=8, + max_num_tokens=100, + tensor_parallel_size=1, + pipeline_parallel_size=1, + context_parallel_size=1, + kv_cache_config=None, + prometheus_metrics_config=SimpleNamespace(e2e_request_latency_buckets=[0.05, 0.15]), + ) + + +def _result_metrics(tokens: int, prompt_tokens: int = 0) -> dict: + from tensorrt_llm.metrics.enums import MetricNames + + return { + "finished_reason": "length", + MetricNames.PROMPT_TOKENS: prompt_tokens, + MetricNames.GENERATION_TOKENS: tokens, + MetricNames.E2E: 0.1, + MetricNames.TTFT: 0.02, + } + + +async def _serving() -> None: + import grpc + from conftest import FakeServicerContext + from test_servicer import _FakeLlm, _FakeResultHandle + + from tensorrt_llm.grpc.openengine.bindings import generation_pb2, openengine_pb2_grpc + from tensorrt_llm.grpc.openengine.metrics import OpenEngineMetrics + from tensorrt_llm.grpc.openengine.server import OpenEngineServer + from tensorrt_llm.grpc.openengine.servicer import OpenEngineInferenceServicer + + gate = asyncio.Event() + + def output(index: int, tokens: list[int], reason: str | None) -> SimpleNamespace: + return SimpleNamespace( + index=index, + token_ids=tokens, + text="A" * len(tokens), + logprobs=[], + prompt_logprobs=[], + finish_reason=reason, + stop_reason=None, + ) + + partial = SimpleNamespace( + prompt_token_ids=[1, 2, 3], + outputs=[output(0, [10], None), output(1, [11], None)], + cached_tokens=0, + error=None, + finished=False, + ) + final = SimpleNamespace( + prompt_token_ids=[1, 2, 3], + outputs=[output(0, [10, 12], "length"), output(1, [11], "length")], + cached_tokens=0, + error=None, + finished=True, + candidate_metrics=[_result_metrics(2, 3), _result_metrics(1)], + ) + + class Handle(_FakeResultHandle): + async def __aiter__(self): + async for result in super().__aiter__(): + yield result + await gate.wait() + + llm = _FakeLlm([]) + llm.args = _args() + llm.llm_id = "test-engine" + llm._check_health = lambda: True + reads = [] + + def get_stats(timeout): + del timeout + reads.append(threading.get_ident()) + return [{"numActiveRequests": 2, "iterLatencyMS": 3}] + + import threading + + llm.get_stats = get_stats + llm.result_handle = Handle([partial, final]) + metrics = OpenEngineMetrics(llm, "test-model", "127.0.0.1", 0, owner=True) + server = OpenEngineServer("127.0.0.1", 0, llm, "test-model", metrics=metrics) + await server.start() + try: + async with grpc.aio.insecure_channel(f"127.0.0.1:{server.port}") as channel: + stub = openengine_pb2_grpc.InferenceStub(channel) + request = generation_pb2.GenerateRequest( + request_id="two-outputs", + model="alias", + prompt="hello", + sampling=generation_pb2.SamplingParams(num_sequences=2), + ) + stream = stub.Generate(request, timeout=10) + assert (await stream.read()).HasField("token") + deadline = time.monotonic() + 5 + while True: + text = await _scrape(metrics.port) + if _sample(text, "trtllm_num_requests_running") == 2: + break + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + assert _sample(text, "trtllm_openengine_rpc_outcomes_total") == 0 + assert reads and all(tid != threading.get_ident() for tid in reads) + gate.set() + while await stream.read() is not grpc.aio.EOF: + pass + text = await _scrape(metrics.port) + assert _sample(text, "trtllm_request_success_total") == 2 + assert _sample(text, "trtllm_prompt_tokens_total") == 3 + assert _sample(text, "trtllm_generation_tokens_total") == 3 + assert _sample(text, "trtllm_e2e_request_latency_seconds_count") == 2 + assert _sample(text, "trtllm_e2e_request_latency_seconds_bucket", le="0.15") == 2 + assert _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="completed") == 1 + assert _sample(text, "trtllm_request_success_total", model_name="alias") == 0 + + # An in-band engine error may finish with OK gRPC transport status. + llm.result_handle = _FakeResultHandle( + [SimpleNamespace(error="engine failed", outputs=[], finished=True)] + ) + events = [event async for event in stub.Generate(request, timeout=5)] + assert events[-1].HasField("error") + text = await _scrape(metrics.port) + assert ( + _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="engine_error") == 1 + ) + + with pytest.raises(grpc.aio.AioRpcError): + async for _ in stub.Generate(generation_pb2.GenerateRequest(), timeout=5): + pass + text = await _scrape(metrics.port) + assert _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="rejected") == 1 + + # An exception inside engine submission remains an INTERNAL engine + # failure, even if its type is also used by frontend coordination. + from tensorrt_llm.grpc.openengine.coordinator import CoordinationError + + with pytest.MonkeyPatch.context() as patch: + + def fail_submit(**kwargs): + raise CoordinationError("submission failed") + + patch.setattr(llm, "generate_async", fail_submit) + with pytest.raises(grpc.aio.AioRpcError) as error: + async for _ in stub.Generate(request, timeout=5): + pass + assert error.value.code() == grpc.StatusCode.INTERNAL + text = await _scrape(metrics.port) + assert ( + _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="engine_error") == 2 + ) + assert _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="unavailable") == 0 + + # A recording failure must not change a successful inference result. + def fail_record(_): + raise OSError("metrics storage full") + + metrics.collector.log_request_metrics_dict = fail_record + llm.result_handle = _FakeResultHandle([final]) + events = [event async for event in stub.Generate(request, timeout=5)] + assert events[-1].HasField("finished") + text = await _scrape(metrics.port) + assert _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="completed") == 2 + assert _sample(text, "trtllm_openengine_metrics_errors_total", source="request") == 1 + + # The watchdog must finalize even if its suspended generator never closes, + # and closing it later must not record a second outcome. + import tensorrt_llm.grpc.openengine.servicer as servicer_module + + servicer_module.RESPONSE_STALL_TIMEOUT_SECONDS = 0 + llm.result_handle = _FakeResultHandle([partial]) + servicer = OpenEngineInferenceServicer(llm, "test-model", metrics=metrics) + responses = servicer.Generate(request, FakeServicerContext()) + await responses.__anext__() + for _ in range(3): + await asyncio.sleep(0) + assert ( + _sample( + await _scrape(metrics.port), + "trtllm_openengine_rpc_outcomes_total", + outcome="stalled", + ) + == 1 + ) + await responses.aclose() + assert ( + _sample( + await _scrape(metrics.port), + "trtllm_openengine_rpc_outcomes_total", + outcome="stalled", + ) + == 1 + ) + finally: + await server.stop(0) + + +async def _multiprocess(attached: bool) -> None: + from tensorrt_llm.grpc.openengine.metrics import OpenEngineMetrics + + def get_stats(timeout): + del timeout + assert not attached, "attached frontend consumed shared engine stats" + return [{"numActiveRequests": 7, "numCompletedRequests": 2}] + + llm = SimpleNamespace(args=_args(), get_stats=get_stats) + metrics = OpenEngineMetrics(llm, "test-model", "127.0.0.1", 0, owner=not attached) + metrics.record_result(SimpleNamespace(candidate_metrics=[_result_metrics(1, 2)])) + await metrics.start() + try: + if attached: + return + _run_worker("attached") + deadline = time.monotonic() + 5 + while True: + text = await _scrape(metrics.port) + if _sample(text, "trtllm_num_requests_running") == 7: + break + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + assert _sample(text, "trtllm_request_success_total") == 2 + assert _sample(text, "trtllm_prompt_tokens_total") == 4 + assert _sample(text, "trtllm_num_requests_completed_total") == 2 + assert "pid=" not in text + # Another live group using the same parent cannot see these counters. + _run_worker("isolated") + assert _sample(await _scrape(metrics.port), "trtllm_request_success_total") == 2 + finally: + await metrics.stop() + + +async def _lifecycle() -> None: + import threading + + from tensorrt_llm.grpc.openengine.metrics import OpenEngineMetrics + + llm = SimpleNamespace(args=_args(), get_stats=lambda timeout: []) + with socket.create_server(("127.0.0.1", 0)) as occupied: + service = OpenEngineMetrics(llm, "test-model", "127.0.0.1", occupied.getsockname()[1], True) + try: + with pytest.raises(OSError): + await service.start() + finally: + await service.stop() + service = OpenEngineMetrics(llm, "test-model", "127.0.0.1", 0, True) + await service.start() + port = service.port + try: + service.check_health() + await _scrape(port) + with pytest.raises(HTTPError) as error: + await asyncio.to_thread(urlopen, f"http://127.0.0.1:{port}/health", timeout=5) + assert error.value.code == 404 + service._http.should_exit = True + await asyncio.to_thread(service._thread.join, 5) + with pytest.raises(RuntimeError, match="stopped unexpectedly"): + service.check_health() + finally: + await service.stop() + with socket.create_server(("127.0.0.1", port)): + pass + + entered = threading.Event() + release = threading.Event() + + def blocked_stats(timeout): + del timeout + entered.set() + assert release.wait(10) + return [{"numActiveRequests": 999}] + + llm.get_stats = blocked_stats + service = OpenEngineMetrics(llm, "test-model", "127.0.0.1", 0, True) + await service.start() + try: + while not entered.is_set(): + await asyncio.sleep(0.01) + # The engine owner must reach engine shutdown before joining this RPC. + await asyncio.wait_for(service.stop(), timeout=3) + finally: + release.set() + await service.wait_for_collection() + await service.stop() + from prometheus_client import CollectorRegistry, generate_latest, multiprocess + + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + assert _sample(generate_latest(registry).decode(), "trtllm_num_requests_running") == 0 + + +async def _collection_mode(requests: bool) -> None: + from tensorrt_llm.grpc.openengine.metrics import OpenEngineMetrics + + args = _args() + args.return_perf_metrics = requests + args.enable_iter_perf_stats = not requests + reads = [] + + def get_stats(timeout): + del timeout + reads.append(True) + return [{"numActiveRequests": 7}] + + llm = SimpleNamespace(args=args, get_stats=get_stats) + metrics = OpenEngineMetrics(llm, "test-model", "127.0.0.1", 0, True) + metrics.record_result(SimpleNamespace(candidate_metrics=[_result_metrics(3, 2)])) + metrics.record_outcome("aggregated", "completed") + await metrics.start() + try: + if requests: + await asyncio.sleep(1.2) + else: + deadline = time.monotonic() + 5 + while _sample(await _scrape(metrics.port), "trtllm_num_requests_running") != 7: + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + metrics.check_health() + text = await _scrape(metrics.port) + assert bool(reads) is not requests + assert _sample(text, "trtllm_request_success_total") == int(requests) + assert _sample(text, "trtllm_generation_tokens_total") == (3 if requests else 0) + assert _sample(text, "trtllm_num_requests_running") == (0 if requests else 7) + assert _sample(text, "trtllm_openengine_rpc_outcomes_total", outcome="completed") == 1 + finally: + await metrics.stop() + await metrics.wait_for_collection() + + +def _http_enablement() -> None: + """File collection must expose Prometheus without opting HTTP headers in.""" + from tensorrt_llm.serve.openai_server import OpenAIServer + + with pytest.MonkeyPatch.context() as patch: + # Exercise the real constructor and route registration without loading a model. + patch.setattr( + OpenAIServer, "_init_llm", lambda self, _: setattr(self, "use_harmony", False) + ) + for requested, output_dir in [(False, None), (True, None), (False, "unused-jsonl")]: + args = _args() + args.return_perf_metrics = requested + args.perf_metrics_output_dir = output_dir + generator = SimpleNamespace( + args=args, _executor=SimpleNamespace(resource_governor_queue=None) + ) + server = OpenAIServer(generator, "test-model", None, None, None) + try: + enabled = requested or output_dir is not None + assert ( + any(route.path == "/prometheus/metrics" for route in server.app.routes) + == enabled + ) + assert server._expose_perf_metrics is requested + assert server._collect_perf_metrics is enabled + assert args.return_perf_metrics is enabled + finally: + server._input_proc_executor.shutdown(wait=True) + server._media_load_executor.shutdown(wait=True) + server._metrics_stats_executor.shutdown(wait=True) + + +if __name__ == "__main__": + from tensorrt_llm.grpc.openengine.metrics import MetricsDirectory + + mode = sys.argv[1] + previous = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + directory = MetricsDirectory(attached=mode == "attached") + owned_path = Path(os.environ["PROMETHEUS_MULTIPROC_DIR"]) + try: + if mode == "isolated": + assert not list(owned_path.glob("*.db")) + elif mode in ("attached", "multiprocess"): + asyncio.run(_multiprocess(attached=mode == "attached")) + elif mode == "serving": + asyncio.run(_serving()) + elif mode.startswith("collection_"): + asyncio.run(_collection_mode(requests=mode == "collection_requests")) + elif mode == "http_enablement": + _http_enablement() + else: + asyncio.run(_lifecycle()) + finally: + directory.close() + if mode != "attached": + assert not owned_path.exists() + assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") == previous diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index fdd464dbe70a..abb5df182fec 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -152,6 +152,117 @@ class _StopLaunch(Exception): """Ends launch_server once the server would start accepting requests.""" +@pytest.mark.parametrize( + "request_stats,iteration_stats,output_dir,enabled", + [ + (False, False, None, False), + (True, False, None, True), + (False, True, None, False), + (True, True, None, True), + ("false", False, None, False), + ("0", False, None, False), + ("off", True, None, False), + (False, False, "/unused/perf-metrics", True), + ], +) +@pytest.mark.parametrize("metrics_options", [{}, {"metrics_port": 0}]) +def test_metrics_listener_preserves_engine_collection_config( + monkeypatch: pytest.MonkeyPatch, + request_stats: bool | str, + iteration_stats: bool, + output_dir: str | None, + enabled: bool, + metrics_options: dict, +) -> None: + """Schema coercion and file-output opt-in must match HTTP; ports never enable export.""" + created = [] + import tensorrt_llm.grpc.openengine.server as oe_server + + class Model: + def __init__(self, **kwargs) -> None: + assert kwargs["return_perf_metrics"] == ( + True if output_dir is not None else request_stats + ) + assert kwargs["enable_iter_perf_stats"] is iteration_stats + self.llm_id = "test-instance" + + def shutdown(self) -> None: + pass + + class Metrics: + def __init__(self, llm, model, host, port, **kwargs) -> None: + created.append((host, port)) + + async def wait_for_collection(self) -> None: + pass + + class Server: + def __init__(self, **kwargs) -> None: + pass + + async def start(self) -> None: + raise _StopLaunch + + async def stop(self) -> None: + pass + + monkeypatch.setattr(oe_server, "PyTorchLLM", Model) + monkeypatch.setattr(oe_server, "OpenEngineMetrics", Metrics) + monkeypatch.setattr(oe_server, "OpenEngineServer", Server) + with pytest.raises(_StopLaunch): + oe_server.launch_server( + "127.0.0.1", + 0, + { + "model": "test-model", + "backend": "pytorch", + "return_perf_metrics": request_stats, + "enable_iter_perf_stats": iteration_stats, + "perf_metrics_output_dir": output_dir, + }, + **metrics_options, + ) + assert created == ( + [("127.0.0.1", metrics_options.get("metrics_port", 8001))] if enabled else [] + ) + + +@pytest.mark.parametrize( + ("arguments", "message"), + [ + (["--openengine-metrics-port", "8001"], "requires --grpc --grpc-protocol openengine"), + ( + ["--grpc", "--openengine-metrics-port", "8001"], + "requires --grpc --grpc-protocol openengine", + ), + (["--openengine-metrics-host", "0.0.0.0"], "requires --grpc --grpc-protocol openengine"), + ( + [ + "--grpc", + "--grpc-protocol", + "openengine", + "--port", + "8001", + "--openengine-metrics-port", + "8001", + ], + "must differ from --port", + ), + ], +) +def test_metrics_cli_rejects_incompatible_server_options( + arguments: list[str], message: str +) -> None: + """Reject invalid metrics deployment combinations before attempting model load.""" + from click.testing import CliRunner + + from tensorrt_llm.commands.serve import serve + + result = CliRunner().invoke(serve, ["nonexistent-model", *arguments]) + assert result.exit_code == 2 + assert message in result.output + + @pytest.mark.parametrize( ("value", "gc_enabled"), [("1", False), ("0", True), (None, True)], diff --git a/tests/unittest/grpc/openengine/test_servicer.py b/tests/unittest/grpc/openengine/test_servicer.py index 9d2de9d32821..841f9309e178 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -1234,7 +1234,11 @@ def test_generate_context_only_ends_at_prefill_ready() -> None: prompt_token_ids=[1, 2], outputs=[final_output], cached_tokens=0, error=None ), ] - servicer = OpenEngineInferenceServicer(_FakeLlm(results), model="test-model") + outcomes = [] + metrics = SimpleNamespace( + record_outcome=lambda *args: outcomes.append(args), record_result=lambda _: None + ) + servicer = OpenEngineInferenceServicer(_FakeLlm(results), model="test-model", metrics=metrics) request = generation_pb2.GenerateRequest( request_id="request-ctx", model="test-model", @@ -1244,7 +1248,11 @@ def test_generate_context_only_ends_at_prefill_ready() -> None: request.extra.update({"request_type": "context_only"}) async def collect_responses() -> list[generation_pb2.GenerateResponse]: - return [response async for response in servicer.Generate(request, FakeServicerContext())] + responses = [] + async for response in servicer.Generate(request, FakeServicerContext()): + assert outcomes == [] + responses.append(response) + return responses responses = asyncio.run(collect_responses()) events = [response.WhichOneof("event") for response in responses] @@ -1253,6 +1261,7 @@ async def collect_responses() -> list[generation_pb2.GenerateResponse]: assert "finished" not in events assert events[-1] == "prefill_ready" assert responses[-1].prefill_ready.kv_session.session_id == "99" + assert outcomes == [("prefill", "completed")] def _set_priority_metadata(request: generation_pb2.GenerateRequest) -> None: @@ -1489,7 +1498,8 @@ def decode(self, *a: Any, **k: Any) -> str: assert all(info.token == "" and not info.candidates for info in infos) -def test_oversized_prompt_logprobs_are_refused_not_built() -> None: +@pytest.mark.parametrize("error_on_final_result", [False, True]) +def test_oversized_prompt_logprobs_are_refused_not_built(error_on_final_result: bool) -> None: """A prompt_length * candidates blow-up is rejected, not computed. PromptOutput is built in one non-preemptible slice on the shared event loop, @@ -1524,6 +1534,41 @@ def test_oversized_prompt_logprobs_are_refused_not_built() -> None: assert responses[0].error.code == openengine_servicer.error_pb2.ERROR_CODE_INVALID_ARGUMENT assert "prompt logprobs" in responses[0].error.message + # Formatting can fail without result.error. Neither immediate completion + # nor a later normal engine result may turn that RPC into a success. + output.finish_reason = "length" if error_on_final_result else None + results = [result] + if not error_on_final_result: + results.append( + SimpleNamespace( + prompt_token_ids=[1], + outputs=[SimpleNamespace(**{**vars(output), "finish_reason": "length"})], + cached_tokens=0, + error=None, + ) + ) + outcomes = [] + metrics = SimpleNamespace( + record_outcome=lambda *args: outcomes.append(args), record_result=lambda _: None + ) + servicer = OpenEngineInferenceServicer(_FakeLlm(results), model="test-model", metrics=metrics) + request = generation_pb2.GenerateRequest( + request_id="oversized-logprobs", + model="test-model", + prompt="hello", + response=generation_pb2.ResponseOptions( + return_prompt_logprobs=True, + prompt_candidates=generation_pb2.CandidateTokenSelection(top_n=2), + ), + ) + + async def collect_responses() -> list[generation_pb2.GenerateResponse]: + return [response async for response in servicer.Generate(request, FakeServicerContext())] + + responses = asyncio.run(collect_responses()) + assert responses[0].HasField("error") + assert outcomes == [("aggregated", "engine_error")] + def test_generate_aborts_the_engine_when_the_rpc_is_cancelled() -> None: """A cancelled RPC must abort the engine request, not just stop streaming. @@ -1542,7 +1587,9 @@ def test_generate_aborts_the_engine_when_the_rpc_is_cancelled() -> None: ) result = SimpleNamespace(prompt_token_ids=[1], outputs=[output], cached_tokens=0, error=None) llm = _FakeLlm([result, result]) - servicer = OpenEngineInferenceServicer(llm, model="test-model") + outcomes = [] + metrics = SimpleNamespace(record_outcome=lambda *args: outcomes.append(args)) + servicer = OpenEngineInferenceServicer(llm, model="test-model", metrics=metrics) context = FakeServicerContext() request = generation_pb2.GenerateRequest( request_id="cancelled-1", model="test-model", prompt="hello" @@ -1559,6 +1606,7 @@ async def cancel_after_first() -> None: assert llm.result_handle.aborted assert servicer.active_request_count() == 0 + assert outcomes == [("aggregated", "cancelled")] @pytest.mark.parametrize("detokenize", [None, True, False]) From e2dd225f0038e77d45bd2e86e05bff1b6747fc5a Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 11:21:19 -0700 Subject: [PATCH 2/8] docs(openengine): simplify metrics setup guide Signed-off-by: Tanmay Verma --- tensorrt_llm/grpc/openengine/README.md | 78 ++++++++++---------------- 1 file changed, 31 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 101a81024b4e..a38bdc0cd358 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -93,79 +93,63 @@ Direct ZMQ sockets are unauthenticated. Keep them on a trusted network and restr ## Prometheus metrics -Select collection with the same engine configuration fields used by the HTTP server. For request timing and token metrics without iteration collection, save this as `metrics.yaml`: +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: false +enable_iter_perf_stats: true ``` -Then expose an HTTP scrape endpoint alongside the OpenEngine gRPC listener: +Start the server and check the metrics endpoint: ```bash trtllm-serve Qwen/Qwen3-0.6B \ --config metrics.yaml \ - --grpc --grpc-protocol openengine --port 50051 + --grpc --grpc-protocol openengine --host 127.0.0.1 --port 50051 curl --fail http://127.0.0.1:8001/prometheus/metrics ``` -Like the HTTP server, OpenEngine exposes Prometheus when `return_perf_metrics: true`. Configuring `perf_metrics_output_dir` also opts into request-performance collection and Prometheus, matching HTTP's existing enablement behavior. Otherwise, the endpoint stays disabled. Boolean values use the LLM schema's validation, so quoted values such as `"false"`, `"0"`, and `"off"` do not enable it. No metrics-specific enable flag is needed. The listener uses port `8001` and the existing `--host` bind address by default. Use `--openengine-metrics-port` or `--openengine-metrics-host` only to override the address. Specifying either override does not enable metrics. This HTTP endpoint is unauthenticated; restrict its network exposure. The metrics port must differ from the gRPC port. Only `/prometheus/metrics` is exposed on the HTTP listener. - -Iteration collection remains independently configured. Without `perf_metrics_output_dir`, the settings behave as follows: - -| `return_perf_metrics` | `enable_iter_perf_stats` | Collected measurements | -| --- | --- | --- | -| `false` | `false` | Prometheus endpoint disabled | -| `true` | `false` | Also request timing and token metrics; no iteration-statistics reader runs | -| `false` | `true` | Prometheus endpoint disabled; engine iteration collection remains enabled | -| `true` | `true` | Both request and iteration metrics | - -Both engine fields default to `false`. Set both to `true` for full collection. The exporter always collects inexpensive Generate RPC outcomes and static configuration while the listener is enabled. Disabled metric families receive no observations; some shared collector counters may still expose their initial zero value. Collection costs depend on the engine timing paths and workload: iteration-only collection is not necessarily faster than full collection. Benchmark the selected configuration on the target build. - -The separate Prometheus listener does not write HTTP per-request JSONL artifacts; matching endpoint enablement does not add that HTTP frontend feature to OpenEngine. - -Prometheus collection does not enable KV-event publishing or change `GetLoad`: `--openengine-enable-load-metrics` still independently controls routing-load snapshots. Health checks use `Control.Health`; a successful scrape only proves that the exporter is responding. The OpenEngine health RPC is not the standard Kubernetes gRPC health-check service. +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 full collection, the existing KV-cache sampling setting reduces how often KV statistics are read from the engine: +For Prometheus, add a separate scrape target: ```yaml -kv_cache_config: - iteration_stats_interval: 4 +scrape_configs: + - job_name: trtllm-openengine + scrape_interval: 15s + metrics_path: /prometheus/metrics + static_configs: + - targets: ["engine-host:8001"] ``` -The engine accumulates KV deltas between samples. Reports can therefore be delayed until another collection iteration occurs. This setting is independent of the one-second exporter polling cadence and the Prometheus scrape interval. Leave `iter_perf_stats_interval: 1` when exact counters from every executor iteration are required: increasing that broader sampling interval can undercount counters such as speculative-decoding draft/accepted totals, even though selected completion and KV counters are carried forward. Request latency/token accounting is collected separately at completion. +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. -The exporter reuses TensorRT-LLM's `trtllm_*` metrics, served-model labels, and `prometheus_metrics_config` histogram buckets: +### Collection settings -| Metric | Counting and timing contract | -| --- | --- | -| `trtllm_request_success_total` | Existing engine finish-reason accounting, per completed output candidate. A successful `num_sequences=2` RPC contributes two candidates. Use the RPC outcome metric below to distinguish cancellation and errors. | -| `trtllm_prompt_tokens_total` | Shared prompt counted once, rather than once per candidate. | -| `trtllm_generation_tokens_total` | Sum of generated tokens across candidates. | -| Request latency histograms | Engine-measured latency for this worker's phase; prefill/decode workers do not measure cross-worker client end-to-end latency. | -| Iteration counters and gauges | Engine statistics collected by one owner approximately once per second, including during active streams. All available iteration deltas are processed; gauges retain the last engine sample when no new sample is produced. Coverage follows the engine's statistics API; these are not a new per-rank reporting API. | -| `trtllm_openengine_rpc_outcomes_total{phase,outcome}` | One outcome per `Generate` RPC. Phases are `aggregated`, `prefill`, and `decode`; outcomes are `completed`, `rejected`, `engine_error`, `cancelled`, `timeout`, `stalled`, and `unavailable`. An in-band engine error is counted even when gRPC transport status is OK. | -| `trtllm_openengine_metrics_errors_total{source}` | Request or iteration metrics recording failures. These failures preserve inference responses. | +- `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`. -RPC completion is recorded after the response stream has been consumed by the gRPC serving layer; for prefill this includes the handoff response. It does not acknowledge application-level receipt. Engine performance metrics are recorded once when the engine result is final; rejected or interrupted requests may have no candidate timing metrics. Internal `Control.Health` inference probes are excluded from Generate outcomes and candidate metrics, but their work is included in engine-level iteration statistics. +### Reading the metrics -With `--num_serve_frontends`, frontend zero owns one HTTP listener and the sole iteration-statistics consumer. Frontends contribute their own request counters and histograms to a shared registry. Engine gauges have one writer and export without PID labels. Scrape serialization and engine-statistics reads run outside the gRPC event loop; HTTP concurrency is bounded. A listener startup failure prevents serving readiness, and unexpected listener termination fails serving-group supervision. +- `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. -Each serving group uses a fresh temporary metrics directory, inherited by its attached frontends. If `PROMETHEUS_MULTIPROC_DIR` is supplied, it must be an existing writable parent directory; the server creates its own child and never deletes the parent's existing contents. Storage is removed after the group exits normally. A forced process kill may leave a child directory behind, but a new group uses a fresh child and does not aggregate the old samples. Initialize this environment before importing `prometheus_client` when embedding the server; an already initialized single-process client is rejected. +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. -Example Prometheus configuration: +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. -```yaml -scrape_configs: - - job_name: trtllm-openengine - scrape_interval: 15s - metrics_path: /prometheus/metrics - static_configs: - - targets: ["engine-host:8001"] -``` +### Performance and deployment notes -For deployment acceptance, verify `up{job="trtllm-openengine"} == 1`, run a known generation workload, and query the RPC outcome, token-counter, and latency-histogram changes. Compare metrics disabled, collection enabled without scrapes, and collection enabled at the intended scrape cadence using matched workloads and frontend counts. Use repeated runs to distinguish measurement noise; target at most 2% throughput loss and 5% p99 inter-token latency growth, and record CPU and resident memory as well. GPU inference and performance qualification require a compatible GPU, model weights, and the matching TensorRT-LLM build. +- 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 From cd8e97b30b207a646e158dd2b6e33a8f3d018fc1 Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 11:51:23 -0700 Subject: [PATCH 3/8] fix(openengine): correct metrics shutdown and RPC outcomes Signed-off-by: Tanmay Verma --- tensorrt_llm/executor/proxy.py | 8 +- tensorrt_llm/executor/rpc/rpc_client.py | 103 ++++++++-------- tensorrt_llm/executor/rpc_proxy.py | 6 +- tensorrt_llm/grpc/openengine/metrics.py | 13 +- tensorrt_llm/grpc/openengine/servicer.py | 8 +- tests/unittest/grpc/openengine/conftest.py | 3 + .../unittest/grpc/openengine/test_metrics.py | 94 +++++++++++++++ .../unittest/grpc/openengine/test_servicer.py | 114 ++++++++++++++++++ 8 files changed, 293 insertions(+), 56 deletions(-) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index ca28c4960bbe..68692a0adddc 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/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 = [] diff --git a/tensorrt_llm/executor/rpc/rpc_client.py b/tensorrt_llm/executor/rpc/rpc_client.py index 7f9b4c545e5f..3fce73c3571b 100644 --- a/tensorrt_llm/executor/rpc/rpc_client.py +++ b/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 diff --git a/tensorrt_llm/executor/rpc_proxy.py b/tensorrt_llm/executor/rpc_proxy.py index bb9e37295f0d..2108628b7cb3 100644 --- a/tensorrt_llm/executor/rpc_proxy.py +++ b/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 = [] diff --git a/tensorrt_llm/grpc/openengine/metrics.py b/tensorrt_llm/grpc/openengine/metrics.py index 467649401ab6..7d49dbf3c1a9 100644 --- a/tensorrt_llm/grpc/openengine/metrics.py +++ b/tensorrt_llm/grpc/openengine/metrics.py @@ -12,6 +12,8 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Literal +import grpc + from tensorrt_llm.logger import logger from tensorrt_llm.metrics.collector import MetricsCollector @@ -135,10 +137,14 @@ def complete(self, result: "RequestOutput", *, handoff_ready: bool) -> None: # Engine candidate accounting is independent of the RPC's outcome. self._metrics.record_result(result) - def finish(self) -> None: + def finish(self, status: grpc.StatusCode | None = None) -> None: """Record at most once, including when a suspended stream resumes later.""" if self._finished: return + if status == grpc.StatusCode.CANCELLED: + self._outcome = RpcOutcome.CANCELLED + elif status == grpc.StatusCode.DEADLINE_EXCEEDED: + self._outcome = RpcOutcome.TIMEOUT self._finished = True self._metrics.record_outcome(self._phase, self._outcome.value) @@ -262,8 +268,9 @@ async def start(self) -> None: """Bind the HTTP listener and start periodic collection in the owner only.""" if not self.owner: return - family = socket.getaddrinfo(self.host, self.port, type=socket.SOCK_STREAM)[0][0] - self._socket = socket.create_server((self.host, self.port), family=family, backlog=8) + bind_host = self.host.strip("[]") + family = socket.getaddrinfo(bind_host, self.port, type=socket.SOCK_STREAM)[0][0] + self._socket = socket.create_server((bind_host, self.port), family=family, backlog=8) self.port = self._socket.getsockname()[1] self._thread = threading.Thread( target=self._http.run, diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index f6826401dfbc..e33756d437bd 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -153,16 +153,20 @@ async def Generate( try: if reservation is not None: await asyncio.shield(reservation.release()) - except CoordinationError as error: + except asyncio.CancelledError as error: if request_metrics is not None: request_metrics.observe_exception(error) + raise + except CoordinationError as error: # Cleanup still fails the serving group through FrontendClient. # Preserve cancellation or an already selected RPC error. if not context.cancelled() and context.code() in (None, grpc.StatusCode.OK): + if request_metrics is not None: + request_metrics.observe_exception(error) await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) finally: if request_metrics is not None: - request_metrics.finish() + request_metrics.finish(context.code()) async def _generate( self, diff --git a/tests/unittest/grpc/openengine/conftest.py b/tests/unittest/grpc/openengine/conftest.py index 5ead36e3ba31..d981dbd8d460 100644 --- a/tests/unittest/grpc/openengine/conftest.py +++ b/tests/unittest/grpc/openengine/conftest.py @@ -46,6 +46,9 @@ def invocation_metadata(self) -> list[SimpleNamespace]: def cancelled(self) -> bool: return self.is_cancelled + def code(self) -> grpc.StatusCode | None: + return self.abort_code + def add_done_callback(self, callback: Any) -> None: self.done_callbacks.append(callback) diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 9a85fa0ccddc..239bc410f093 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -35,6 +35,7 @@ def _run_worker(mode: str) -> None: "collection_requests", "collection_iterations", "http_enablement", + "rpc_shutdown", ], ) def test_prometheus_contract(mode: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: @@ -318,6 +319,22 @@ async def _lifecycle() -> None: from tensorrt_llm.grpc.openengine.metrics import OpenEngineMetrics llm = SimpleNamespace(args=_args(), get_stats=lambda timeout: []) + if socket.has_ipv6: + try: + ipv6_socket = socket.create_server(("::1", 0), family=socket.AF_INET6) + except OSError: + pass # IPv6 can be disabled even when the Python build supports it. + else: + ipv6_socket.close() + ipv6 = OpenEngineMetrics(llm, "test-model", "[::1]", 0, True) + try: + await ipv6.start() + await asyncio.to_thread( + lambda: socket.create_connection(("::1", ipv6.port), timeout=2).close() + ) + finally: + await ipv6.stop() + await ipv6.wait_for_collection() with socket.create_server(("127.0.0.1", 0)) as occupied: service = OpenEngineMetrics(llm, "test-model", "127.0.0.1", occupied.getsockname()[1], True) try: @@ -409,6 +426,81 @@ def get_stats(timeout): await metrics.wait_for_collection() +async def _rpc_shutdown() -> None: + """An unanswered stats RPC must time out or release its caller on close.""" + import threading + + from tensorrt_llm.executor.proxy import GenerationExecutorProxy + from tensorrt_llm.executor.rpc import RPCCancelled, RPCClient, RPCServer, RPCTimeout + from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr + from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy + + for mode in ("deadline", "rpc_deadline", "sync", "future", "async", "queued"): + entered, release = threading.Event(), threading.Event() + + class App: + async def fetch_stats_wait_async(self, timeout): + entered.set() + while not release.is_set(): + await asyncio.sleep(0.01) + return [] + + server = RPCServer(App(), async_run_task=True) + server.bind(get_unique_ipc_addr()) + server.start() + client = RPCClient(server.address, hmac_key=server.hmac_key, num_workers=1) + proxy = SimpleNamespace(rpc_client=client) + try: + if mode == "queued": + # Hold the sole worker until close has stopped the RPC loop. + blocker = client._executor.submit(release.wait) + read = asyncio.wrap_future( + client.fetch_stats_wait_async(timeout=0.1).remote_future() + ) + closing = asyncio.create_task(asyncio.to_thread(client.close)) + await asyncio.to_thread(client._loop_thread.join, 3) + assert not client._loop_thread.is_alive() + release.set() + await asyncio.wait_for(closing, 3) + assert blocker.result(timeout=1) + with pytest.raises(RPCCancelled): + await asyncio.wait_for(read, 3) + with pytest.raises(RPCCancelled): + GenerationExecutorProxy.get_stats(proxy, timeout=0.1) + continue + if mode == "async": + read = asyncio.create_task( + client.fetch_stats_wait_async(timeout=0.1).remote_async() + ) + elif mode == "future": + read = asyncio.wrap_future( + client.fetch_stats_wait_async(timeout=0.1).remote_future() + ) + else: + proxy_type = ( + GenerationExecutorRpcProxy + if mode == "rpc_deadline" + else GenerationExecutorProxy + ) + read = asyncio.create_task( + asyncio.to_thread(proxy_type.get_stats, proxy, timeout=0.1) + ) + assert await asyncio.to_thread(entered.wait, 3) + if mode == "deadline": + with pytest.raises(RPCTimeout): + await asyncio.wait_for(read, 3) + elif mode == "rpc_deadline": + assert await asyncio.wait_for(read, 3) == [] + else: + await asyncio.wait_for(asyncio.to_thread(client.close), 3) + with pytest.raises(RPCCancelled): + await asyncio.wait_for(read, 3) + finally: + release.set() + await asyncio.to_thread(client.close) + await asyncio.to_thread(server.shutdown) + + def _http_enablement() -> None: """File collection must expose Prometheus without opting HTTP headers in.""" from tensorrt_llm.serve.openai_server import OpenAIServer @@ -459,6 +551,8 @@ def _http_enablement() -> None: asyncio.run(_collection_mode(requests=mode == "collection_requests")) elif mode == "http_enablement": _http_enablement() + elif mode == "rpc_shutdown": + asyncio.run(_rpc_shutdown()) else: asyncio.run(_lifecycle()) finally: diff --git a/tests/unittest/grpc/openengine/test_servicer.py b/tests/unittest/grpc/openengine/test_servicer.py index 841f9309e178..916c538eb4eb 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -1570,6 +1570,120 @@ async def collect_responses() -> list[generation_pb2.GenerateResponse]: assert outcomes == [("aggregated", "engine_error")] +def test_generate_reserve_failure_is_unavailable() -> None: + """A coordination failure before submission is not an engine failure.""" + from tensorrt_llm.grpc.openengine.coordinator import CoordinationError + + class Frontend: + async def reserve(self, request_id): + raise CoordinationError("reservation failed") + + outcomes = [] + metrics = SimpleNamespace(record_outcome=lambda *args: outcomes.append(args)) + servicer = OpenEngineInferenceServicer( + _FakeLlm([]), "test-model", frontend=Frontend(), metrics=metrics + ) + context = FakeServicerContext() + request = generation_pb2.GenerateRequest( + request_id="reserve-failure", model="test-model", prompt="hello" + ) + + async def run(): + with pytest.raises(AbortError) as error: + async for _ in servicer.Generate(request, context): + pass + assert error.value.code == grpc.StatusCode.UNAVAILABLE + + asyncio.run(run()) + assert outcomes == [("aggregated", "unavailable")] + + +@pytest.mark.parametrize( + "failure,expected_status,expected_outcome", + [ + ("cancel_release", grpc.StatusCode.CANCELLED, "cancelled"), + ("reject_release", grpc.StatusCode.INVALID_ARGUMENT, "rejected"), + ("abort_submission", grpc.StatusCode.CANCELLED, "cancelled"), + ], +) +def test_generate_cleanup_preserves_rpc_outcome(failure, expected_status, expected_outcome): + """Real transport cancellation and cleanup errors must not mislabel RPCs.""" + from tensorrt_llm.grpc.openengine.coordinator import CoordinationError + + async def run(): + entered, release = asyncio.Event(), asyncio.Event() + outcomes = [] + + class Reservation: + aborted = failure == "abort_submission" + + def bind(self, handle): + pass + + async def release(self): + entered.set() + if failure == "cancel_release": + await release.wait() + elif failure == "reject_release": + raise CoordinationError("release failed") + + class Frontend: + async def reserve(self, request_id): + return Reservation() + + output = SimpleNamespace( + index=0, + token_ids=[10], + text="A", + logprobs=[], + prompt_logprobs=[], + finish_reason="length", + stop_reason=None, + ) + result = SimpleNamespace( + prompt_token_ids=[1], outputs=[output], cached_tokens=0, error=None, finished=True + ) + metrics = SimpleNamespace( + record_outcome=lambda *args: outcomes.append(args), record_result=lambda _: None + ) + servicer = OpenEngineInferenceServicer( + _FakeLlm([result]), "test-model", frontend=Frontend(), metrics=metrics + ) + server = grpc.aio.server() + openengine_pb2_grpc.add_InferenceServicer_to_server(servicer, server) + port = server.add_insecure_port("127.0.0.1:0") + await server.start() + try: + async with grpc.aio.insecure_channel(f"127.0.0.1:{port}") as channel: + request = generation_pb2.GenerateRequest( + request_id="cleanup", + model="" if failure == "reject_release" else "model", + prompt="hello", + ) + call = openengine_pb2_grpc.InferenceStub(channel).Generate(request, timeout=5) + + async def drain(): + try: + async for _ in call: + pass + except (grpc.aio.AioRpcError, asyncio.CancelledError): + pass + + task = asyncio.create_task(drain()) + if failure == "cancel_release": + await asyncio.wait_for(entered.wait(), 3) + call.cancel() + await task + assert await call.code() == expected_status + await asyncio.wait_for(server.stop(0), 3) + assert outcomes == [("aggregated", expected_outcome)] + finally: + release.set() + await server.stop(0) + + asyncio.run(run()) + + def test_generate_aborts_the_engine_when_the_rpc_is_cancelled() -> None: """A cancelled RPC must abort the engine request, not just stop streaming. From 67d8ff9e857fb5e971fa8003759fd2e13491d332 Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 12:16:53 -0700 Subject: [PATCH 4/8] fix(openengine): handle metrics binding and RPC close races Signed-off-by: Tanmay Verma --- tensorrt_llm/executor/rpc/rpc_client.py | 39 +++++++--- tensorrt_llm/grpc/openengine/metrics.py | 2 +- .../unittest/grpc/openengine/test_metrics.py | 77 ++++++++++++++++++- 3 files changed, 104 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/executor/rpc/rpc_client.py b/tensorrt_llm/executor/rpc/rpc_client.py index 3fce73c3571b..00ea55dd9983 100644 --- a/tensorrt_llm/executor/rpc/rpc_client.py +++ b/tensorrt_llm/executor/rpc/rpc_client.py @@ -22,6 +22,8 @@ from .rpc_common import (RPCCancelled, RPCParams, RPCRequest, RPCResponse, RPCStreamingError, RPCTimeout) +_CLOSE_CHECK_INTERVAL_SECONDS = 0.5 + class RemoteCall: """Helper class to enable chained remote call syntax like client.method().remote()""" @@ -548,8 +550,15 @@ def _call_sync(self, method_name, *args, **kwargs): raise RPCCancelled("RPC client closed") future = asyncio.run_coroutine_threadsafe( self._call_async(method_name, *args, **kwargs), self._loop) - result = future.result() - return result + while True: + try: + return future.result(timeout=_CLOSE_CHECK_INTERVAL_SECONDS) + except concurrent.futures.TimeoutError: + if future.done(): + return future.result() + if self._closed and not self._loop.is_running(): + future.cancel() + raise RPCCancelled("RPC client closed") def _call_future(self, name: str, *args, **kwargs) -> concurrent.futures.Future: @@ -592,6 +601,8 @@ async def _call_streaming(self, name: str, *args, # Register queue with lock to ensure thread-safe access with self._streaming_queues_lock: + if self._closed: + raise RPCCancelled("RPC client closed") self._streaming_queues[request_id] = queue #logger_debug(f"[{datetime.now().isoformat()}] Registered streaming queue for request_id={request_id}") @@ -617,15 +628,21 @@ async def _send_streaming_request(req: RPCRequest): f"[client][{datetime.now().isoformat()}] Streaming request sent successfully: {req.method_name}, request_id={req.request_id}" ) - send_future = asyncio.run_coroutine_threadsafe( - _send_streaming_request(request), self._loop) - - # Wait until the request is actually on the wire before entering the - # user-visible streaming loop. We wrap the concurrent.futures.Future so - # we can await it in the caller's asyncio context. - await asyncio.wrap_future(send_future) - + send_future = None try: + with self._call_lock: + if self._closed: + raise RPCCancelled("RPC client closed") + send_future = asyncio.run_coroutine_threadsafe( + _send_streaming_request(request), self._loop) + sent = asyncio.wrap_future(send_future) + # Closing can stop the RPC loop while the transport send is pending. + while not sent.done(): + await asyncio.wait((sent, ), + timeout=_CLOSE_CHECK_INTERVAL_SECONDS) + if self._closed: + raise RPCCancelled("RPC client closed") + await sent logger_debug( f"[client] [{datetime.now().isoformat()}] Starting to read streaming responses for request_id={request_id}" ) @@ -661,6 +678,8 @@ async def _send_streaming_request(req: RPCRequest): raise RPCTimeout( f"Streaming request '{name}' timed out after {timeout}s") finally: + if send_future is not None: + send_future.cancel() # Clean up with self._streaming_queues_lock: self._streaming_queues.pop(request_id, None) diff --git a/tensorrt_llm/grpc/openengine/metrics.py b/tensorrt_llm/grpc/openengine/metrics.py index 7d49dbf3c1a9..0f33c2bfdf69 100644 --- a/tensorrt_llm/grpc/openengine/metrics.py +++ b/tensorrt_llm/grpc/openengine/metrics.py @@ -269,7 +269,7 @@ async def start(self) -> None: if not self.owner: return bind_host = self.host.strip("[]") - family = socket.getaddrinfo(bind_host, self.port, type=socket.SOCK_STREAM)[0][0] + family = socket.AF_INET6 if ":" in bind_host else socket.AF_INET self._socket = socket.create_server((bind_host, self.port), family=family, backlog=8) self.port = self._socket.getsockname()[1] self._thread = threading.Thread( diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 239bc410f093..82ceb4821c4b 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -319,6 +319,23 @@ async def _lifecycle() -> None: from tensorrt_llm.grpc.openengine.metrics import OpenEngineMetrics llm = SimpleNamespace(args=_args(), get_stats=lambda timeout: []) + resolve = socket.getaddrinfo + + def ipv6_first(host, *args, **kwargs): + addresses = resolve(host, *args, **kwargs) + if host == "localhost": + addresses.sort(key=lambda address: address[0] != socket.AF_INET6) + return addresses + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(socket, "getaddrinfo", ipv6_first) + localhost = OpenEngineMetrics(llm, "test-model", "localhost", 0, True) + try: + await localhost.start() + await _scrape(localhost.port) # Scrape over IPv4 even with IPv6-first DNS. + finally: + await localhost.stop() + await localhost.wait_for_collection() if socket.has_ipv6: try: ipv6_socket = socket.create_server(("::1", 0), family=socket.AF_INET6) @@ -435,7 +452,18 @@ async def _rpc_shutdown() -> None: from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy - for mode in ("deadline", "rpc_deadline", "sync", "future", "async", "queued"): + for mode in ( + "deadline", + "rpc_deadline", + "sync", + "future", + "async", + "queued", + "send_sync", + "send_future", + "send_stream", + "stream_admission", + ): entered, release = threading.Event(), threading.Event() class App: @@ -451,6 +479,29 @@ async def fetch_stats_wait_async(self, timeout): client = RPCClient(server.address, hmac_key=server.hmac_key, num_workers=1) proxy = SimpleNamespace(rpc_client=client) try: + if mode == "stream_admission": + from tensorrt_llm.executor.rpc import rpc_client + + class ClosingQueue(rpc_client.AsyncQueue): + def __init__(self): + super().__init__() + client.close() + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(rpc_client, "AsyncQueue", ClosingQueue) + stream = client.never_stream().remote_streaming() + with pytest.raises(RPCCancelled): + await asyncio.wait_for(stream.__anext__(), 3) + assert not client._streaming_queues + continue + if mode.startswith("send_"): + + async def blocked_send(request): + entered.set() + while not release.is_set(): + await asyncio.sleep(0.01) + + client._client_socket.put_async = blocked_send if mode == "queued": # Hold the sole worker until close has stopped the RPC loop. blocker = client._executor.submit(release.wait) @@ -468,11 +519,14 @@ async def fetch_stats_wait_async(self, timeout): with pytest.raises(RPCCancelled): GenerationExecutorProxy.get_stats(proxy, timeout=0.1) continue - if mode == "async": + if mode == "send_stream": + stream = client.never_stream().remote_streaming() + read = asyncio.create_task(stream.__anext__()) + elif mode == "async": read = asyncio.create_task( client.fetch_stats_wait_async(timeout=0.1).remote_async() ) - elif mode == "future": + elif mode in ("future", "send_future"): read = asyncio.wrap_future( client.fetch_stats_wait_async(timeout=0.1).remote_future() ) @@ -495,9 +549,26 @@ async def fetch_stats_wait_async(self, timeout): await asyncio.wait_for(asyncio.to_thread(client.close), 3) with pytest.raises(RPCCancelled): await asyncio.wait_for(read, 3) + assert not client._streaming_queues finally: release.set() await asyncio.to_thread(client.close) + if mode.startswith("send_"): + # Drain cancellation callbacks queued after the deliberately stopped loop. + async def drain_cancelled_sends(): + pending = asyncio.all_tasks() - {asyncio.current_task()} + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + + drain_thread = threading.Thread(target=client._loop.run_forever) + drain_thread.start() + await asyncio.wrap_future( + asyncio.run_coroutine_threadsafe(drain_cancelled_sends(), client._loop) + ) + client._loop.call_soon_threadsafe(client._loop.stop) + await asyncio.to_thread(drain_thread.join, 3) + client._loop.close() await asyncio.to_thread(server.shutdown) From 0296d56939a82b9b9cb9cacee5077f0164545b6d Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 12:37:49 -0700 Subject: [PATCH 5/8] fix(deps): exclude vulnerable filelock releases Signed-off-by: Tanmay Verma --- requirements.txt | 3 +++ security_scanning/poetry.lock | 8 ++++---- security_scanning/pyproject.toml | 3 ++- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/requirements.txt b/requirements.txt index 31dc876b9bb8..7e0808af2bb1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,8 @@ --extra-index-url https://download.pytorch.org/whl/cu132 -c constraints.txt +# Patched filelock requires Python 3.11; 3.20.3 predates the affected lock implementations. +filelock==3.20.3; python_version < "3.11" +filelock>=4.2.1; python_version >= "3.11" accelerate>=1.7.0 build colored diff --git a/security_scanning/poetry.lock b/security_scanning/poetry.lock index e2f4b2388be1..feafd87c2d7c 100644 --- a/security_scanning/poetry.lock +++ b/security_scanning/poetry.lock @@ -1355,14 +1355,14 @@ standard-no-fastapi-cloud-cli = ["email-validator (>=2.0.0)", "fastapi-cli[stand [[package]] name = "filelock" -version = "4.1.0" +version = "3.20.3" description = "A platform independent file lock." optional = false python-versions = ">=3.10" groups = ["main"] files = [ - {file = "filelock-4.1.0-py3-none-any.whl", hash = "sha256:2ce9818e3e2d8f284c1a964414447ef148d42a5fd5e2a477a7118e574b293ec1"}, - {file = "filelock-4.1.0.tar.gz", hash = "sha256:ad7f724afef953e731b1cc39bcd3a09166d72ed7fcdf29e6e88b1c3235c6715d"}, + {file = "filelock-3.20.3-py3-none-any.whl", hash = "sha256:4b0dda527ee31078689fc205ec4f1c1bf7d56cf88b6dc9426c4f230e46c2dce1"}, + {file = "filelock-3.20.3.tar.gz", hash = "sha256:18c57ee915c7ec61cff0ecf7f0f869936c7c30191bb0cf406f1341778d0834e1"}, ] [[package]] @@ -7457,4 +7457,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "588023ad0e11bf5966c75a2ad04641e27948158adf6aa36a35875882146be71c" +content-hash = "51bc42164fc784eab2f12f065da39c6595490a1cb80642e245d2d1e928c27223" diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index 3164fb6e11c3..c0557cc7d4f4 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -90,7 +90,8 @@ dependencies = [ "cache-dit (>=1.3.5)", "msgpack (>=1.2.3,<2.0.0)", "uvloop (>=0.19.0,<0.23.0)", - "mooncake-transfer-engine-cuda13 (==0.3.13)" + "mooncake-transfer-engine-cuda13 (==0.3.13)", + "filelock (==3.20.3)" ] From a7f55ae1b292342e8a3f19ac029402e89614cba9 Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 13:13:04 -0700 Subject: [PATCH 6/8] fix(openengine): resolve metrics listener address family Signed-off-by: Tanmay Verma --- tensorrt_llm/grpc/openengine/metrics.py | 9 +++++-- .../unittest/grpc/openengine/test_metrics.py | 25 ++++++++++++------- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/metrics.py b/tensorrt_llm/grpc/openengine/metrics.py index 0f33c2bfdf69..f7b0e629ebf7 100644 --- a/tensorrt_llm/grpc/openengine/metrics.py +++ b/tensorrt_llm/grpc/openengine/metrics.py @@ -269,8 +269,13 @@ async def start(self) -> None: if not self.owner: return bind_host = self.host.strip("[]") - family = socket.AF_INET6 if ":" in bind_host else socket.AF_INET - self._socket = socket.create_server((bind_host, self.port), family=family, backlog=8) + addresses = socket.getaddrinfo( + bind_host or None, self.port, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE + ) + address = next( + (address for address in addresses if address[0] == socket.AF_INET), addresses[0] + ) + self._socket = socket.create_server(address[4], family=address[0], backlog=8) self.port = self._socket.getsockname()[1] self._thread = threading.Thread( target=self._http.run, diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 82ceb4821c4b..55cc48bbeaf4 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -343,15 +343,22 @@ def ipv6_first(host, *args, **kwargs): pass # IPv6 can be disabled even when the Python build supports it. else: ipv6_socket.close() - ipv6 = OpenEngineMetrics(llm, "test-model", "[::1]", 0, True) - try: - await ipv6.start() - await asyncio.to_thread( - lambda: socket.create_connection(("::1", ipv6.port), timeout=2).close() - ) - finally: - await ipv6.stop() - await ipv6.wait_for_collection() + + def ipv6_only(host, *args, **kwargs): + return resolve("::1" if host == "metrics-v6.test" else host, *args, **kwargs) + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(socket, "getaddrinfo", ipv6_only) + for host in ("[::1]", "metrics-v6.test"): + ipv6 = OpenEngineMetrics(llm, "test-model", host, 0, True) + try: + await ipv6.start() + await asyncio.to_thread( + lambda: socket.create_connection(("::1", ipv6.port), timeout=2).close() + ) + finally: + await ipv6.stop() + await ipv6.wait_for_collection() with socket.create_server(("127.0.0.1", 0)) as occupied: service = OpenEngineMetrics(llm, "test-model", "127.0.0.1", occupied.getsockname()[1], True) try: From 977fcc7f58a53cf3876739839818173ef3a94003 Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Mon, 26 Oct 2026 15:07:02 -0700 Subject: [PATCH 7/8] fix(openengine): classify expired RPC deadlines as timeouts Signed-off-by: Tanmay Verma --- tensorrt_llm/grpc/openengine/README.md | 2 ++ tensorrt_llm/grpc/openengine/servicer.py | 10 +++++++++- tests/unittest/grpc/openengine/conftest.py | 3 +++ tests/unittest/grpc/openengine/test_servicer.py | 14 +++++++++----- 4 files changed, 23 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/grpc/openengine/README.md b/tensorrt_llm/grpc/openengine/README.md index 17dc0963e788..95c7d193db64 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -157,6 +157,8 @@ 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. +Monitoring tools scrape each service's metrics endpoint directly. Configure the engine endpoint as its own target; another service does not need to re-export these metrics. Metrics endpoint discovery through `GetServerInfo` is deferred, so configure the metrics host, port, and path in your monitoring setup. + For Prometheus, add a separate scrape target: ```yaml diff --git a/tensorrt_llm/grpc/openengine/servicer.py b/tensorrt_llm/grpc/openengine/servicer.py index e33756d437bd..2a356dd2e656 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -166,7 +166,15 @@ async def Generate( await context.abort(grpc.StatusCode.UNAVAILABLE, str(error)) finally: if request_metrics is not None: - request_metrics.finish(context.code()) + status = context.code() + remaining = context.time_remaining() + if ( + status in (None, grpc.StatusCode.OK) + and remaining is not None + and remaining <= 0 + ): + status = grpc.StatusCode.DEADLINE_EXCEEDED + request_metrics.finish(status) async def _generate( self, diff --git a/tests/unittest/grpc/openengine/conftest.py b/tests/unittest/grpc/openengine/conftest.py index 2d47b8b1a12a..3e2861bfc55f 100644 --- a/tests/unittest/grpc/openengine/conftest.py +++ b/tests/unittest/grpc/openengine/conftest.py @@ -52,6 +52,9 @@ def cancelled(self) -> bool: def code(self) -> grpc.StatusCode | None: return self.abort_code + def time_remaining(self) -> float | None: + return None + def add_done_callback(self, callback: Any) -> None: self.done_callbacks.append(callback) diff --git a/tests/unittest/grpc/openengine/test_servicer.py b/tests/unittest/grpc/openengine/test_servicer.py index d835655102c5..c484ded1217b 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -1603,12 +1603,13 @@ async def run(): "failure,expected_status,expected_outcome", [ ("cancel_release", grpc.StatusCode.CANCELLED, "cancelled"), + ("deadline_release", grpc.StatusCode.DEADLINE_EXCEEDED, "timeout"), ("reject_release", grpc.StatusCode.INVALID_ARGUMENT, "rejected"), ("abort_submission", grpc.StatusCode.CANCELLED, "cancelled"), ], ) def test_generate_cleanup_preserves_rpc_outcome(failure, expected_status, expected_outcome): - """Real transport cancellation and cleanup errors must not mislabel RPCs.""" + """Real transport deadlines, cancellation, and cleanup errors retain their outcomes.""" from tensorrt_llm.grpc.openengine.coordinator import CoordinationError async def run(): @@ -1623,7 +1624,7 @@ def bind(self, handle): async def release(self): entered.set() - if failure == "cancel_release": + if failure in ("cancel_release", "deadline_release"): await release.wait() elif failure == "reject_release": raise CoordinationError("release failed") @@ -1661,7 +1662,9 @@ async def reserve(self, request_id): model="" if failure == "reject_release" else "model", prompt="hello", ) - call = openengine_pb2_grpc.InferenceStub(channel).Generate(request, timeout=5) + call = openengine_pb2_grpc.InferenceStub(channel).Generate( + request, timeout=0.5 if failure == "deadline_release" else 5 + ) async def drain(): try: @@ -1671,9 +1674,10 @@ async def drain(): pass task = asyncio.create_task(drain()) - if failure == "cancel_release": + if failure in ("cancel_release", "deadline_release"): await asyncio.wait_for(entered.wait(), 3) - call.cancel() + if failure == "cancel_release": + call.cancel() await task assert await call.code() == expected_status await asyncio.wait_for(server.stop(0), 3) From b04fac4b8550abc594c8b5f5b3d9b2f0f9e197dd Mon Sep 17 00:00:00 2001 From: Tanmay Verma Date: Tue, 27 Oct 2026 12:12:03 -0700 Subject: [PATCH 8/8] test(openengine): release gRPC transport cycles after cleanup checks Signed-off-by: Tanmay Verma --- tests/unittest/grpc/openengine/test_servicer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unittest/grpc/openengine/test_servicer.py b/tests/unittest/grpc/openengine/test_servicer.py index eebe0fe00539..8a9f989b8ca3 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -1719,7 +1719,11 @@ async def drain(): release.set() await server.stop(0) - asyncio.run(run()) + try: + asyncio.run(run()) + finally: + # gRPC transport cycles can retain its global poller after the loop closes. + gc.collect() def test_generate_aborts_the_engine_when_the_rpc_is_cancelled() -> None: