diff --git a/security_scanning/poetry.lock b/security_scanning/poetry.lock index 193758861794..b7418562db42 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]] @@ -7690,4 +7690,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "cdf966b4679037304dfa5f63e42616e5f119d627b38ad566c6d58deae1417939" +content-hash = "c5bb6c4b97ae2a2e60a92bca494c9d30b37a43f6812e03a3e5b2434303f9d5fd" diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 08d018dcfb44..2ab54f6c1859 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1272,6 +1272,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, @@ -1380,6 +1393,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], @@ -1412,6 +1427,20 @@ 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( "--moe_cluster_parallel_size / --cluster_size is deprecated and " @@ -1671,6 +1700,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/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..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()""" @@ -125,6 +127,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 +160,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 +172,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 +306,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 +325,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 +468,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 +485,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,10 +545,20 @@ 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) - result = future.result() - return result + 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) + 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: @@ -552,12 +575,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 +588,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 @@ -583,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}") @@ -608,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}" ) @@ -652,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/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/README.md b/tensorrt_llm/grpc/openengine/README.md index 5ce68cf169c1..95c7d193db64 100644 --- a/tensorrt_llm/grpc/openengine/README.md +++ b/tensorrt_llm/grpc/openengine/README.md @@ -134,6 +134,68 @@ Consumers must track rank-local sequence numbers and resynchronize on gaps. Publ Direct ZMQ sockets are unauthenticated. Keep them on a trusted network and restrict access to intended subscribers. +## Prometheus metrics + +OpenEngine serves metrics on a separate HTTP port, using the same collection settings and `trtllm_*` metric families as the HTTP server. + +Save this as `metrics.yaml` to collect request timing, token counts, and engine statistics such as KV-cache usage: + +```yaml +return_perf_metrics: true +enable_iter_perf_stats: true +``` + +Start the server and check the metrics endpoint: + +```bash +trtllm-serve Qwen/Qwen3-0.6B \ + --config metrics.yaml \ + --grpc --grpc-protocol openengine --host 127.0.0.1 --port 50051 + +curl --fail http://127.0.0.1:8001/prometheus/metrics +``` + +Metrics use port `8001` and inherit `--host`. Override them with `--openengine-metrics-port` and `--openengine-metrics-host`; the metrics port must differ from the gRPC port. These options change the address, not whether collection is enabled. + +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 +scrape_configs: + - job_name: trtllm-openengine + scrape_interval: 15s + metrics_path: /prometheus/metrics + static_configs: + - targets: ["engine-host:8001"] +``` + +Replace `engine-host` with an address Prometheus can reach. For remote scraping, bind metrics to a reachable interface. The endpoint has no authentication; restrict access to a trusted network. Use `Control.Health` to check inference health; a successful scrape only checks the exporter. + +### Collection settings + +- `return_perf_metrics: true` enables the endpoint and request metrics. Both collection settings default to `false`. +- `enable_iter_perf_stats: true` adds engine statistics. Set it to `false` for request-only collection; it does not enable the endpoint by itself. +- `perf_metrics_output_dir` also enables collection and the endpoint, matching HTTP enablement, but OpenEngine does not write HTTP per-request JSONL files. +- KV events and routing load remain separate: metrics do not enable event publishing or `--openengine-enable-load-metrics`. + +### Reading the metrics + +- `trtllm_request_success_total` counts completed output candidates by finish reason. One successful request with `num_sequences=2` contributes two candidates. +- `trtllm_prompt_tokens_total` counts the prompt once per request; `trtllm_generation_tokens_total` sums tokens across outputs. +- `trtllm_openengine_rpc_outcomes_total{phase,outcome}` counts one outcome per `Generate` call, including rejection, cancellation, and engine errors. Use this to distinguish RPC outcomes from candidate completion. Completion means the server consumed the response stream, not that the client acknowledged receipt. +- `trtllm_openengine_metrics_errors_total{source}` reports collection failures without failing inference. + +Request latency histograms measure this worker's execution phase. For separate prefill and decode workers, they do not measure the client's full request latency. Rejected or interrupted requests may have no candidate timing observations. Health probes are excluded from request metrics, but their engine work can appear in iteration statistics. + +Engine statistics are collected approximately once per second. Gauges retain the last engine sample when no new sample is available. Served-model labels and `prometheus_metrics_config` histogram buckets match HTTP; OpenEngine's engine gauges omit `pid` labels. + +### Performance and deployment notes + +- Measure overhead with your workload and scrape interval. To sample KV statistics less often, increase `kv_cache_config.iteration_stats_interval`; this can delay updates. Keep `iter_perf_stats_interval: 1` for complete iteration counters, including speculative-decoding counts. +- With multiple frontends, one process serves metrics and collects engine statistics; request metrics are combined across all frontends. +- Metrics storage is managed automatically. If setting `PROMETHEUS_MULTIPROC_DIR`, use an existing writable parent directory; the server creates and cleans up a private child. Forced termination can leave that child behind. When embedding the server, initialize the environment before importing `prometheus_client`. + ## Transport security The listener is plaintext h2c with no authentication. Any client that can reach diff --git a/tensorrt_llm/grpc/openengine/metrics.py b/tensorrt_llm/grpc/openengine/metrics.py new file mode 100644 index 000000000000..f7b0e629ebf7 --- /dev/null +++ b/tensorrt_llm/grpc/openengine/metrics.py @@ -0,0 +1,340 @@ +# 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 + +import grpc + +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, 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) + + +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 + bind_host = self.host.strip("[]") + 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, + 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 1e764cb28a62..c9d810f65d01 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -31,12 +31,14 @@ _spawn_attached_frontends, _terminate_attached_frontends, ) +from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from ._health import OpenEngineHealthServicer 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"] @@ -111,10 +113,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 self._frontend = frontend self._stopping = False self._health_task: asyncio.Task | None = None @@ -141,6 +145,7 @@ def __init__( kv_transfer_backend=kv_transfer_backend, frontend=frontend, subagent_affinity_auth_key=subagent_affinity_auth_key, + metrics=metrics, input_executor=self._input_executor, media_executor=self._media_executor, ) @@ -187,6 +192,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() await self._health.set("liveness", health_pb2.HealthCheckResponse.SERVING) await self._refresh_inference_health() @@ -276,10 +283,18 @@ async def stop(self, grace: float = 5.0) -> None: await asyncio.gather(self._health_task, return_exceptions=True) await self._server.stop(grace=grace) finally: - await asyncio.gather( - asyncio.to_thread(self._input_executor.shutdown, wait=True, cancel_futures=True), - asyncio.to_thread(self._media_executor.shutdown, wait=True, cancel_futures=True), - ) + try: + if self._metrics is not None: + await self._metrics.stop() + finally: + await asyncio.gather( + asyncio.to_thread( + self._input_executor.shutdown, wait=True, cancel_futures=True + ), + asyncio.to_thread( + self._media_executor.shutdown, wait=True, cancel_futures=True + ), + ) logger.info("OpenEngine server stopped") async def wait_for_termination(self) -> None: @@ -307,6 +322,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. @@ -318,6 +335,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() @@ -332,6 +351,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") @@ -346,8 +375,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: @@ -402,6 +433,8 @@ async def supervise() -> None: if terminal_engine_failure(): fail("OpenEngine terminal engine failure") 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) @@ -420,6 +453,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. @@ -434,6 +469,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 = llm_args["_openengine_discovery"]["engine_id"] if mode.is_launcher: directory = tempfile.TemporaryDirectory(prefix="tllm-openengine-") @@ -491,9 +534,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) @@ -530,8 +575,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 71a8e5fd3963..edb3b18ccc07 100644 --- a/tensorrt_llm/grpc/openengine/servicer.py +++ b/tensorrt_llm/grpc/openengine/servicer.py @@ -32,6 +32,7 @@ from .disagg import disaggregated_params_from_request from .errors import AbortFailedError, MediaUnavailableError, UnsupportedFeatureError from .formatting import _engine_error_response, _stop_texts +from .metrics import OpenEngineMetrics, RequestMetrics from .request_mapping import ( _load_prepared_input, _metadata_from_context, @@ -131,6 +132,7 @@ def __init__( kv_transfer_backend: str = "", frontend: FrontendClient | None = None, subagent_affinity_auth_key: str | None = None, + metrics: OpenEngineMetrics | None = None, input_executor: Executor | None = None, media_executor: Executor | None = None, ) -> None: @@ -147,6 +149,7 @@ def __init__( self._media_capabilities = media_capabilities_for_model(llm) self.subagent_affinity_auth_key = subagent_affinity_auth_key self._active_requests: dict[str, Any] = {} + self._metrics = metrics @property def supported_media_modalities(self) -> frozenset[str]: @@ -211,6 +214,7 @@ async def Generate( local_request = None registration = None reservation = None + request_metrics = RequestMetrics(self._metrics) if self._metrics else None try: if not request.request_id: await context.abort( @@ -227,7 +231,9 @@ async def Generate( reservation = await self._frontend.reserve(request.request_id) reservation.bind(local_request) async with aclosing( - self._generate(request, context, local_request, registration, reservation) + self._generate( + request, context, local_request, registration, reservation, request_metrics + ) ) as responses: async for response in responses: yield response @@ -235,21 +241,44 @@ async def Generate( 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 local_request is not None and local_request.preprocessing is not None: - local_request.preprocessing.abandon() - if registration is not None: - registration.release() - if reservation is not None: - try: + try: + if local_request is not None and local_request.preprocessing is not None: + local_request.preprocessing.abandon() + if registration is not None: + registration.release() + 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 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: + 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, @@ -258,6 +287,7 @@ async def _generate( local_request: _LocalRequest, registration: ActiveRequest, 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 @@ -280,6 +310,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, @@ -369,6 +401,8 @@ async def prepare_inputs() -> Any: raise except Exception as error: # Unexpected preparation failures are worker errors, not bad inputs. + if request_metrics is not None: + request_metrics.engine_failed() logger.error(f"OpenEngine request {request_id} could not be prepared: {error}") await context.abort(grpc.StatusCode.INTERNAL, str(error)) return @@ -424,6 +458,8 @@ async def prepare_inputs() -> Any: # 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 @@ -431,11 +467,15 @@ async def prepare_inputs() -> Any: # 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 local_request.bind_engine(result_handle) + if request_metrics is not None: + request_metrics.accepted() engine_terminal = False abort_requested = False @@ -455,6 +495,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" @@ -490,6 +532,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() @@ -508,6 +552,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 @@ -516,18 +564,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 b01445d2a086..e3956e67c874 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, @@ -718,12 +719,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. @@ -1034,31 +1036,12 @@ 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.metrics_collector = MetricsCollector.from_args( + self.generator.args, self.model) # One config-info series per server, not per frontend process. if self._launcher_forwarder is None: - self._log_config_info_metrics() + self.metrics_collector.log_config_from_args( + self.generator.args, self.model) @staticmethod def _ensure_post_processor_hook_supported( @@ -1113,95 +1096,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 edf634c48fcf..40d5165e86b7 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -215,6 +215,7 @@ l0_cpu: - unittest/grpc/openengine/test_multimodal.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/conftest.py b/tests/unittest/grpc/openengine/conftest.py index d6e10501bf34..3e2861bfc55f 100644 --- a/tests/unittest/grpc/openengine/conftest.py +++ b/tests/unittest/grpc/openengine/conftest.py @@ -49,6 +49,12 @@ 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 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_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py new file mode 100644 index 000000000000..55cc48bbeaf4 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -0,0 +1,640 @@ +# 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", + "rpc_shutdown", + ], +) +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: []) + 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) + except OSError: + pass # IPv6 can be disabled even when the Python build supports it. + else: + ipv6_socket.close() + + 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: + 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() + + +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", + "send_sync", + "send_future", + "send_stream", + "stream_admission", + ): + 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 == "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) + 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 == "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 in ("future", "send_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) + 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) + + +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() + elif mode == "rpc_shutdown": + asyncio.run(_rpc_shutdown()) + 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_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index bc2cb9865266..d6daa3338491 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -1674,7 +1674,9 @@ def preprocess(*args): monkeypatch.setattr(request_mapping, "ImageMediaIO", _ImageIo) monkeypatch.setattr(llm, "preprocess", preprocess) - servicer = OpenEngineInferenceServicer(llm, model="image-model") + outcomes = [] + metrics = SimpleNamespace(record_outcome=lambda *args: outcomes.append(args)) + servicer = OpenEngineInferenceServicer(llm, model="image-model", metrics=metrics) request = generation_pb2.GenerateRequest( request_id="preprocess-error", model="image-model", @@ -1690,6 +1692,10 @@ def preprocess(*args): pass assert context.abort_code == expected_code + expected_outcome = ( + "rejected" if expected_code == grpc.StatusCode.INVALID_ARGUMENT else "engine_error" + ) + assert outcomes == [("aggregated", expected_outcome)] assert llm.generate_kwargs is None assert servicer.active_request_count() == 0 diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index 7856d31c95bc..241646f0fc06 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -63,6 +63,7 @@ async def exercise_stop(): ) server._health = SimpleNamespace(enter_graceful_shutdown=AsyncMock()) server._health_task = None + server._metrics = SimpleNamespace(stop=AsyncMock()) server._input_executor = ThreadPoolExecutor(max_workers=1) server._media_executor = ThreadPoolExecutor(max_workers=1) futures = [ @@ -74,6 +75,7 @@ async def exercise_stop(): else: await server.stop(grace=0) assert server._stopping + server._metrics.stop.assert_awaited_once() server._health.enter_graceful_shutdown.assert_awaited_once() assert all(future.result() for future in futures) for pool in (server._input_executor, server._media_executor): @@ -539,6 +541,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 ca7760c1e25d..a953bceba01f 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -1268,7 +1268,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", @@ -1278,7 +1282,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] @@ -1287,6 +1295,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: @@ -1523,7 +1532,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, @@ -1558,6 +1568,168 @@ 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_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"), + ("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 deadlines, cancellation, and cleanup errors retain their outcomes.""" + from tensorrt_llm.grpc.openengine.coordinator import CoordinationError + + async def run(): + entered, release = asyncio.Event(), asyncio.Event() + outcomes = [] + recorded = asyncio.Event() + + class Reservation: + aborted = failure == "abort_submission" + + def bind(self, handle): + pass + + async def release(self): + entered.set() + if failure in ("cancel_release", "deadline_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 + ) + + def record_outcome(*args): + outcomes.append(args) + recorded.set() + + metrics = SimpleNamespace(record_outcome=record_outcome, 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=0.5 if failure == "deadline_release" else 5 + ) + + async def drain(): + try: + async for _ in call: + pass + except (grpc.aio.AioRpcError, asyncio.CancelledError): + pass + + task = asyncio.create_task(drain()) + if failure in ("cancel_release", "deadline_release"): + await asyncio.wait_for(entered.wait(), 3) + if failure == "cancel_release": + call.cancel() + await task + assert await call.code() == expected_status + await asyncio.wait_for(recorded.wait(), 3) + await asyncio.wait_for(server.stop(0), 3) + assert outcomes == [("aggregated", expected_outcome)] + finally: + release.set() + await server.stop(0) + + 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: """A cancelled RPC must abort the engine request, not just stop streaming. @@ -1576,7 +1748,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" @@ -1593,6 +1767,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])