From 7af5c6c213e5122858d0f1b383331ae0ef8b8af6 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:42:37 -0700 Subject: [PATCH 1/2] [None][fix] serve: make multi-frontend metrics, profiling and liveness launcher-owned With num_serve_frontends > 1 every frontend drained the engine's iteration stats and kept its own Prometheus registry, so /metrics and /prometheus/metrics reported a random frontend's share; /start_profile returned 500 on attached frontends; the Anthropic batch store was per-process; attached frontends never noticed a dead launcher or engine and kept answering /health 200; and hub-id models with tp > 1 crashed attached frontends in CachedModelLoader (no MPI session). - The launcher's uvicorn also listens on a Unix socket in the multi-frontend ipc dir; attached frontends forward /metrics, /kv_cache_events, /start_profile, /stop_profile and /v1/messages/batches* to it (tensorrt_llm/serve/multi_frontend.py). - Only the launcher runs the iteration-stats collector; it also polls on a 1 s cadence (off the event loop) since attached-frontend requests never wake it. - The launcher exports one PROMETHEUS_MULTIPROC_DIR before spawning the children so request counters/histograms aggregate across frontends; config-info gauges are logged by the launcher only. - An attached-frontend watchdog turns a vanished parent or failing launcher /health into a fatal executor error plus graceful shutdown. - CachedModelLoader runs node tasks locally when there is no MPI session. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/commands/serve.py | 99 ++++++++- tensorrt_llm/llmapi/llm_utils.py | 7 +- tensorrt_llm/serve/multi_frontend.py | 202 ++++++++++++++++++ tensorrt_llm/serve/openai_server.py | 122 +++++++++-- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_serve_multi_frontend_helpers.py | 172 +++++++++++++++ 6 files changed, 582 insertions(+), 21 deletions(-) create mode 100644 tensorrt_llm/serve/multi_frontend.py create mode 100644 tests/unittest/llmapi/test_serve_multi_frontend_helpers.py diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index a045382b6256..7533ba29601b 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -19,7 +19,7 @@ from pathlib import Path from types import FrameType from typing import (TYPE_CHECKING, Any, Dict, NamedTuple, NoReturn, Optional, - Sequence, Set) + Sequence, Set, Tuple) import click import torch @@ -52,6 +52,8 @@ from tensorrt_llm.logger import logger, severity_map from tensorrt_llm.mapping import CpType from tensorrt_llm.serve import OpenAIDisaggServer, OpenAIServer +from tensorrt_llm.serve.multi_frontend import (LAUNCHER_UDS_NAME, + MultiFrontendServing) from tensorrt_llm.serve.tool_parser import ToolParserFactory from tensorrt_llm.serve.tool_parser.tool_parser_factory import ( MODEL_TYPE_TO_TOOL_PARSER, resolve_auto_tool_parser) @@ -490,7 +492,68 @@ def _init_multi_frontend_mode(llm_args: dict, return mode -def _spawn_attached_frontends(llm, num_frontends: int) -> list: +def _attached_launcher_uds(multi_frontend: MultiFrontendMode) -> Optional[str]: + """The launcher's Unix socket path handed to this attached frontend. + + Only the path is read; the attach env var itself (which also carries the + executor HMAC keys) is consumed and deleted by GenerationExecutor.create. + """ + if not multi_frontend.is_attached_frontend: + return None + attach_env = os.getenv("TLLM_EXECUTOR_ATTACH_INFO") + if not attach_env: + return None + return json.loads(attach_env).get("launcher_uds") + + +def _bind_launcher_uds(llm) -> Tuple[socket.socket, str]: + """Bind the launcher's Unix socket inside the multi-frontend ipc dir. + + uvicorn serves it next to the TCP socket, so attached frontends can + forward launcher-owned routes (/metrics, /kv_cache_events, profiling, + message batches) to the one process that owns the engine. The ipc dir + is private to the launcher's user and removed by the executor proxy. + """ + ipc_dir = getattr(getattr(llm, "_executor", None), + "_multi_frontend_ipc_dir", None) + if not ipc_dir: + raise ValueError( + "num_serve_frontends > 1 requires the classic IPC executor " + "proxy in multi-frontend mode") + path = os.path.join(ipc_dir, LAUNCHER_UDS_NAME) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + sock.bind(path) + return sock, path + + +def _share_prometheus_multiproc_dir(llm, ipc_dir: str) -> None: + """Give every frontend one PROMETHEUS_MULTIPROC_DIR before spawning. + + prometheus_client's multiprocess mode aggregates the .db files under + that directory. Left to OpenAIServer, each process would create its own + after the children already exist, and every /prometheus/metrics scrape + would report a single frontend's share of the request counters. + + The directory lives inside the multi-frontend ipc dir, which the + executor proxy removes at shutdown. It is deliberately not created via + set_prometheus_multiproc_dir(): that helper keeps only its most recent + TemporaryDirectory alive, so its second call (OpenAIServer.__init__) + would garbage-collect, and thereby delete, a shared directory created + by the first while every process still writes into it. + """ + args = getattr(llm, "args", None) + if args is None or os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + return + if (getattr(args, "return_perf_metrics", False) + or getattr(args, "perf_metrics_output_dir", None)): + path = tempfile.mkdtemp(prefix="prometheus_", dir=ipc_dir) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = path + logger.info( + f"Shared PROMETHEUS_MULTIPROC_DIR for all frontends: {path}") + + +def _spawn_attached_frontends(llm, num_frontends: int, + launcher_uds: str) -> list: """Spawn num_frontends - 1 attached serving frontend processes. Each child re-execs this trtllm-serve command line with env vars @@ -513,6 +576,9 @@ def _spawn_attached_frontends(llm, num_frontends: int) -> list: raise ValueError( "num_serve_frontends > 1 requires the classic IPC executor " f"proxy in multi-frontend mode, got {type(executor).__name__}") + # Where the launcher's uvicorn also listens; attached frontends forward + # launcher-owned routes there (see serve/multi_frontend.py). + attach_info["launcher_uds"] = launcher_uds # Carries the executor HMAC keys; the child deletes it from its env # once consumed (GenerationExecutor.create). attach_env = json.dumps(attach_info) @@ -634,6 +700,8 @@ def launch_server( model = served_model_name or llm_args["model"] multi_frontend = _init_multi_frontend_mode(llm_args, multi_frontend_enabled) + # Read before the LLM consumes (and deletes) the attach env var. + attached_launcher_uds = _attached_launcher_uds(multi_frontend) # Same hazard the disaggregated fleet guard covers: _spawn_attached_frontends # re-execs this command line verbatim, so with port 0 every frontend binds # its own kernel-assigned port instead of sharing one, and every frontend @@ -711,10 +779,24 @@ def launch_server( # server construction, middleware registration, and runtime, or a # failure in between leaks the child processes. frontend_children = [] + sockets = [s] + launcher_uds_sock = None + multi_frontend_serving = None + if attached_launcher_uds is not None: + multi_frontend_serving = MultiFrontendServing( + launcher_uds=attached_launcher_uds, + is_launcher=False, + launcher_pid=os.getppid()) try: if multi_frontend.is_launcher: + launcher_uds_sock, launcher_uds = _bind_launcher_uds(llm) + _share_prometheus_multiproc_dir(llm, + os.path.dirname(launcher_uds)) frontend_children = _spawn_attached_frontends( - llm, multi_frontend.num_frontends) + llm, multi_frontend.num_frontends, launcher_uds) + multi_frontend_serving = MultiFrontendServing( + launcher_uds=launcher_uds, is_launcher=True) + sockets.append(launcher_uds_sock) server = OpenAIServer( generator=llm, @@ -728,7 +810,8 @@ def launch_server( allow_request_chat_template=allow_request_chat_template, input_processor_workers=num_input_processor_workers, media_load_workers=num_media_load_workers, - internal_disagg_auth_key=internal_disagg_auth_key) + internal_disagg_auth_key=internal_disagg_auth_key, + multi_frontend_serving=multi_frontend_serving) _apply_fastapi_middlewares(server.app, middleware) # Optionally disable GC (default: not disabled) @@ -736,10 +819,16 @@ def launch_server( gc.disable() _signal_frontend_ready(multi_frontend) - uvloop.run(server(host, port, sockets=[s])) + uvloop.run(server(host, port, sockets=sockets)) finally: if frontend_children: _terminate_attached_frontends(frontend_children) + if launcher_uds_sock is not None: + launcher_uds_sock.close() + try: + os.unlink(launcher_uds_sock.getsockname()) + except OSError: + pass def launch_mm_encoder_server( diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index e09f076f5919..66b54cbc3d1a 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -390,7 +390,12 @@ def _submit_to_all_workers( *args, **kwargs, ) -> List[Any]: - if self.llm_args.parallel_config.is_multi_gpu: + # An attached serving frontend (trtllm-serve num_serve_frontends > 1) + # is multi-GPU by configuration but owns no MPI session: the launcher + # already ran the task on every node, so running it locally is + # enough (e.g. a hub download resolves from the populated cache). + if (self.llm_args.parallel_config.is_multi_gpu + and self.mpi_session is not None): return self.mpi_session.submit_sync(task, *args, **kwargs) else: return [task(*args, **kwargs)] diff --git a/tensorrt_llm/serve/multi_frontend.py b/tensorrt_llm/serve/multi_frontend.py new file mode 100644 index 000000000000..9d3621e3e1e6 --- /dev/null +++ b/tensorrt_llm/serve/multi_frontend.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Helpers for trtllm-serve multi-frontend serving (``num_serve_frontends > 1``). + +Several frontend processes share one executor and one ``SO_REUSEPORT`` port, +and a client cannot choose which of them answers. Anything that needs exactly +one owner per server therefore lives in the launcher (frontend 0): draining +the engine's iteration statistics and KV-cache events, runtime profiling, and +the in-memory Anthropic Message Batches store. The launcher's uvicorn also +listens on a Unix domain socket in the multi-frontend ipc directory; attached +frontends forward those routes to it (:class:`LauncherForwarder`) and watch it +(:class:`AttachedFrontendWatchdog`) so that a dead launcher or engine does not +leave them answering ``/health`` 200 with nothing behind them. +""" + +import asyncio +import os +from dataclasses import dataclass +from typing import Callable, Optional + +import aiohttp +from fastapi import Request +from fastapi.responses import JSONResponse, Response + +from tensorrt_llm.logger import logger + +LAUNCHER_UDS_NAME = "launcher.sock" + +# Routes an attached frontend hands to the launcher, as FastAPI route +# templates. They are registered before the regular routes, and FastAPI +# matches in registration order, so they shadow the local handlers. The +# catch-all covers /v1/messages/batches/{batch_id}[/cancel|/results]. +FORWARDED_ROUTES = ( + ("GET", "/metrics"), + ("POST", "/kv_cache_events"), + ("POST", "/start_profile"), + ("POST", "/stop_profile"), + ("GET", "/v1/messages/batches"), + ("POST", "/v1/messages/batches"), + ("GET", "/v1/messages/batches/{rest:path}"), + ("POST", "/v1/messages/batches/{rest:path}"), + ("DELETE", "/v1/messages/batches/{rest:path}"), +) + +_FORWARDED_REQUEST_HEADERS = ("content-type", "accept", "authorization") + + +@dataclass(frozen=True) +class MultiFrontendServing: + """This process's role in a multi-frontend server. + + ``None`` in place of an instance means single-frontend serving. + """ + + launcher_uds: str + is_launcher: bool + # The launcher's pid as seen by an attached frontend (its parent). + launcher_pid: Optional[int] = None + + @property + def is_attached(self) -> bool: + return not self.is_launcher + + +class LauncherForwarder: + """Forward HTTP requests from an attached frontend to the launcher.""" + + def __init__(self, uds_path: str, timeout_s: float = 120.0): + self._uds_path = uds_path + self._timeout = aiohttp.ClientTimeout(total=timeout_s) + self._session: Optional[aiohttp.ClientSession] = None + + async def _get_session(self) -> aiohttp.ClientSession: + # Created lazily so it binds to the serving event loop. + if self._session is None or self._session.closed: + # force_close: no idle keep-alive connections to race against a + # launcher that is shutting down; these routes are low-rate. + self._session = aiohttp.ClientSession( + connector=aiohttp.UnixConnector(path=self._uds_path, force_close=True), + timeout=self._timeout, + ) + return self._session + + async def forward(self, request: Request) -> Response: + """FastAPI endpoint: replay ``request`` against the launcher verbatim.""" + target = request.url.path + if request.url.query: + target = f"{target}?{request.url.query}" + headers = { + k: v for k, v in request.headers.items() if k.lower() in _FORWARDED_REQUEST_HEADERS + } + body = await request.body() + try: + session = await self._get_session() + # The host is required by the URL syntax and ignored by the + # Unix connector. + async with session.request( + request.method, f"http://launcher{target}", data=body, headers=headers + ) as resp: + content = await resp.read() + return Response( + content=content, + status_code=resp.status, + media_type=resp.headers.get("Content-Type"), + ) + except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as e: + logger.warning( + f"Forwarding {request.method} {target} to the launcher frontend failed: {e!r}" + ) + return JSONResponse( + status_code=503, + content={ + "error": f"The launcher frontend that serves {request.url.path} " + f"is unavailable: {e!r}" + }, + ) + + async def get_status(self, path: str, timeout_s: float = 5.0) -> Optional[int]: + """HTTP status of ``GET path`` on the launcher, ``None`` if unreachable.""" + try: + session = await self._get_session() + async with session.get( + f"http://launcher{path}", timeout=aiohttp.ClientTimeout(total=timeout_s) + ) as resp: + await resp.read() + return resp.status + except (aiohttp.ClientError, asyncio.TimeoutError, OSError): + return None + + async def close(self) -> None: + if self._session is not None and not self._session.closed: + await self._session.close() + self._session = None + + +class AttachedFrontendWatchdog: + """Detect a lost launcher or engine from an attached frontend. + + An attached frontend owns no worker processes, so the only signals that + the engine is gone are its parent (the launcher) disappearing, or the + launcher's own ``/health`` failing, which happens once the launcher's + executor recorded a fatal engine error. Either way ``mark_dead`` is + called exactly once; the caller records the fatal error on its executor + (failing in-flight requests fast) and starts server shutdown. + """ + + def __init__( + self, + forwarder: LauncherForwarder, + launcher_pid: int, + mark_dead: Callable[[BaseException], None], + *, + ppid_interval_s: float = 1.0, + health_interval_s: float = 5.0, + health_failures: int = 3, + ): + self._forwarder = forwarder + self._launcher_pid = launcher_pid + self._mark_dead = mark_dead + self._ppid_interval = ppid_interval_s + self._health_interval = health_interval_s + self._health_failures = health_failures + + async def run(self) -> None: + loop = asyncio.get_running_loop() + failures = 0 + next_health = loop.time() + self._health_interval + while True: + await asyncio.sleep(self._ppid_interval) + if os.getppid() != self._launcher_pid: + self._die(f"launcher frontend (pid {self._launcher_pid}) exited") + return + if loop.time() < next_health: + continue + next_health = loop.time() + self._health_interval + status = await self._forwarder.get_status("/health") + if status == 200: + failures = 0 + continue + failures += 1 + if failures >= self._health_failures: + self._die( + f"launcher frontend /health failed {failures} " + f"consecutive times (last status: {status})" + ) + return + + def _die(self, why: str) -> None: + logger.error(f"Attached frontend lost its engine: {why}") + self._mark_dead(RuntimeError(f"attached frontend lost its engine: {why}")) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index ff98cb6c444d..7b39a3143b02 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -90,6 +90,10 @@ from tensorrt_llm.serve.encode_batcher import (EncodeBatcher, InputTooLongError, QueueFullError) from tensorrt_llm.serve.metadata_server import create_metadata_server +from tensorrt_llm.serve.multi_frontend import (FORWARDED_ROUTES, + AttachedFrontendWatchdog, + LauncherForwarder, + MultiFrontendServing) from tensorrt_llm.serve.openai_protocol import ( ChatCompletionMessageParam, ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, @@ -675,7 +679,8 @@ def __init__( media_load_workers: int = 8, internal_disagg_auth_key: Optional[str] = None, enable_rl_control_endpoints: bool = False, - rl_control_api_key: Optional[str] = None): + rl_control_api_key: Optional[str] = None, + multi_frontend_serving: Optional[MultiFrontendServing] = None): if enable_rl_control_endpoints and not rl_control_api_key: raise ValueError( "rl_control_api_key is required when RL control endpoints are enabled" @@ -750,6 +755,22 @@ def __init__( ) if server_role is not None else "server" self._perf_metrics_writer = PerfMetricsJsonlWriter( perf_metrics_output_dir, server_kind) + # Multi-frontend serving (num_serve_frontends > 1): the launcher is + # the sole consumer of the engine's iteration stats / KV events and + # the only owner of profiling and the batch store; attached + # frontends forward those routes to it over its Unix socket and + # watch it for engine or launcher death (see serve/multi_frontend.py). + self._multi_frontend = multi_frontend_serving + self._launcher_forwarder: Optional[LauncherForwarder] = None + self._launcher_watchdog_task: Optional[asyncio.Task] = None + if multi_frontend_serving is not None and multi_frontend_serving.is_attached: + self._launcher_forwarder = LauncherForwarder( + multi_frontend_serving.launcher_uds) + # Requests finishing on attached frontends never wake the launcher's + # collector, so the launcher also polls on a fixed cadence. + self._iteration_stats_poll_interval: Optional[float] = ( + 1.0 if multi_frontend_serving is not None + and multi_frontend_serving.is_launcher else None) self._iteration_stats_collector_task = None self._iteration_stats_wakeup_event = asyncio.Event() # Bounded snapshot of iteration stats for the GET /metrics handler. @@ -818,7 +839,18 @@ async def lifespan(app: FastAPI): # Start background iteration stats collector if metrics are enabled # The args for pytorch and autodeploy backend has attribute `enable_iter_perf_stats` while # tensorrt backend does not have this attribute but it always has iter stats enabled. - if self.metrics_collector and getattr( + # An attached frontend never drains the engine queue: the + # launcher owns it and /metrics is forwarded there. + if self._launcher_forwarder is not None: + self._launcher_watchdog_task = asyncio.create_task( + AttachedFrontendWatchdog( + self._launcher_forwarder, + self._multi_frontend.launcher_pid, + self._on_launcher_lost).run()) + logger.info( + "Attached frontend: forwarding launcher-owned routes " + f"to {self._multi_frontend.launcher_uds}") + elif self.metrics_collector and getattr( self.generator.args, "enable_iter_perf_stats", True): # The background loop becomes the sole consumer of the # engine stats queue; /metrics reads from a tee buffer @@ -851,6 +883,14 @@ async def lifespan(app: FastAPI): yield await self._perf_metrics_writer.close() + if self._launcher_watchdog_task is not None: + self._launcher_watchdog_task.cancel() + try: + await self._launcher_watchdog_task + except asyncio.CancelledError: + pass + if self._launcher_forwarder is not None: + await self._launcher_forwarder.close() if self.embedding_batcher is not None: await self.embedding_batcher.shutdown() logger.info("Stopped encode dynamic batcher") @@ -911,6 +951,14 @@ async def validation_exception_handler(request, exc): self._init_embedding_batcher() self.register_embedding_routes() else: + if self._launcher_forwarder is not None: + # Registered first so they shadow the local handlers below + # (FastAPI matches in registration order). + for method, path in FORWARDED_ROUTES: + self.app.add_api_route(path, + self._launcher_forwarder.forward, + methods=[method], + include_in_schema=False) self.register_routes() if self._collect_perf_metrics: @@ -1044,7 +1092,11 @@ def _init_llm(self, chat_template: Optional[str] = None): request_inference_time_buckets=( pmc.request_inference_time_buckets if pmc else None), ) - self._log_config_info_metrics() + # With a shared PROMETHEUS_MULTIPROC_DIR every process would + # export its own per-pid config-info series; the launcher's + # suffices (attached frontends share its engine and config). + if self._launcher_forwarder is None: + self._log_config_info_metrics() @staticmethod def _ensure_post_processor_hook_supported( @@ -1668,6 +1720,22 @@ async def data_transceiver_state(self) -> JSONResponse: base64.b64encode(state).decode("utf-8") }) + def _on_launcher_lost(self, error: BaseException) -> None: + """Watchdog callback of an attached frontend: fail fast, then shut down. + + Mirrors the fatal-error branch of ``health``: record the error on the + executor so in-flight and new requests raise EngineDeadError instead + of hanging, then raise SIGINT once for uvicorn's graceful shutdown. + """ + executor = getattr(self.generator, '_executor', None) + if executor is None: + return + if getattr(executor, '_fatal_error', None) is None: + executor._set_fatal_error(error) + if not getattr(executor, 'doing_shutdown', True): + _record_generator_termination(self.generator) + signal.raise_signal(signal.SIGINT) + async def health(self) -> Response: if self._check_health(): return Response(status_code=200) @@ -1849,8 +1917,10 @@ async def _extract_metrics(self, res: RequestOutput, raw_request: Request): self.metrics_collector.log_request_metrics_dict( res.metrics_dict) # Note: Iteration stats are collected by the background _iteration_stats_collector_loop task - # Wake up the stats collector to drain iteration stats - if getattr(self.generator.args, "enable_iter_perf_stats", True): + # Wake up the stats collector to drain iteration stats (an + # attached frontend has no collector: the launcher drains). + if self._iteration_stats_collector_task is not None and getattr( + self.generator.args, "enable_iter_perf_stats", True): self._iteration_stats_wakeup_event.set() async def _create_chat_response( @@ -1901,22 +1971,37 @@ async def _iteration_stats_collector_loop(self): try: logger.info("Iteration stats collector loop started") while True: - # Wait for signal that requests have completed and stats may be available - await self._iteration_stats_wakeup_event.wait() + # Wait for signal that requests have completed and stats may + # be available. A multi-frontend launcher also wakes on a + # fixed cadence: requests finishing on attached frontends + # never set this event, yet their iterations land in the + # engine queue this loop alone drains. + woken_by_timer = False + try: + await asyncio.wait_for( + self._iteration_stats_wakeup_event.wait(), + timeout=getattr(self, "_iteration_stats_poll_interval", + None)) + except asyncio.TimeoutError: + woken_by_timer = True # Clear the event for next wakeup self._iteration_stats_wakeup_event.clear() # Drain all available iteration stats and log each one to Prometheus. try: - async for llm_stat in self.generator.get_stats_async( - timeout=0.5): - self.metrics_collector.log_iteration_stats(llm_stat) - # Tee into the /metrics snapshot buffer so the HTTP - # handler can serve without competing for the engine - # queue (nvbug 6102381). - if self._iteration_stats_buffer is not None: - self._iteration_stats_buffer.append(llm_stat) + if woken_by_timer: + # Idle poll: the stats RPC would otherwise wait out + # its timeout inside a synchronous .remote() call and + # stall this event loop, so fetch once without + # waiting, off the loop. + for llm_stat in await asyncio.to_thread( + self.generator.get_stats, 0): + self._record_iteration_stat(llm_stat) + else: + async for llm_stat in self.generator.get_stats_async( + timeout=0.5): + self._record_iteration_stat(llm_stat) except Exception as e: # Log errors but continue collecting stats logger.error(f"Error collecting iteration stats: {e}", @@ -1927,6 +2012,13 @@ async def _iteration_stats_collector_loop(self): logger.info("Iteration stats collector loop cancelled") raise + def _record_iteration_stat(self, llm_stat: dict) -> None: + self.metrics_collector.log_iteration_stats(llm_stat) + # Tee into the /metrics snapshot buffer so the HTTP handler can serve + # without competing for the engine queue (nvbug 6102381). + if self._iteration_stats_buffer is not None: + self._iteration_stats_buffer.append(llm_stat) + async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0e7e5e66f7f1..2920bce5e75a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -146,6 +146,7 @@ l0_cpu: - unittest/llmapi/test_rl_control_auth.py - unittest/llmapi/test_sampling_params.py - unittest/llmapi/test_serialization.py + - unittest/llmapi/test_serve_multi_frontend_helpers.py - unittest/llmapi/test_serve_report_addr.py - unittest/llmapi/test_session_prefetcher.py - unittest/llmapi/test_tokenizer_aliases.py diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py new file mode 100644 index 000000000000..832011e28b11 --- /dev/null +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""CPU coverage for trtllm-serve's multi-frontend helpers. + +An attached frontend forwards launcher-owned routes to the launcher over a +Unix socket and watches the launcher for engine or process death. Both are +exercised here against a fake launcher; the GPU path is covered end to end by +the serving integration tests. +""" + +import asyncio +import os +import threading +from contextlib import asynccontextmanager + +import pytest +from aiohttp import web +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from tensorrt_llm.serve.multi_frontend import ( + FORWARDED_ROUTES, + AttachedFrontendWatchdog, + LauncherForwarder, +) + +pytestmark = pytest.mark.cpu_only + + +class _FakeLauncher: + """aiohttp app on a Unix socket standing in for the launcher's uvicorn.""" + + def __init__(self, path: str, health_status: int = 200) -> None: + self.path = path + self.health_status = health_status + self.seen = [] + self._loop = asyncio.new_event_loop() + self._ready = threading.Event() + self._thread = threading.Thread(target=self._serve, daemon=True) + self._thread.start() + assert self._ready.wait(10), "fake launcher did not start" + + def _serve(self) -> None: + asyncio.set_event_loop(self._loop) + + async def echo(request: web.Request) -> web.Response: + self.seen.append((request.method, request.path_qs, await request.text())) + return web.json_response( + {"method": request.method, "path": request.path_qs}, + status=201 if request.method == "POST" else 200, + ) + + async def health(_: web.Request) -> web.Response: + return web.Response(status=self.health_status) + + app = web.Application() + app.router.add_get("/health", health) + app.router.add_route("*", "/{tail:.*}", echo) + self._runner = web.AppRunner(app) + self._loop.run_until_complete(self._runner.setup()) + self._loop.run_until_complete(web.UnixSite(self._runner, self.path).start()) + self._ready.set() + self._loop.run_forever() + + def close(self) -> None: + self._loop.call_soon_threadsafe(self._loop.stop) + self._thread.join(10) + self._loop.run_until_complete(self._runner.cleanup()) + self._loop.close() + + +def _attached_app(forwarder: LauncherForwarder) -> FastAPI: + """The forwarded routes as OpenAIServer registers them, plus a local one.""" + + @asynccontextmanager + async def lifespan(_: FastAPI): + yield + await forwarder.close() + + app = FastAPI(lifespan=lifespan) + for method, path in FORWARDED_ROUTES: + app.add_api_route(path, forwarder.forward, methods=[method]) + + @app.get("/health") + async def local_health(): + return {"local": True} + + return app + + +def test_forwarder_replays_launcher_owned_routes(tmp_path) -> None: + launcher = _FakeLauncher(str(tmp_path / "launcher.sock")) + try: + with TestClient(_attached_app(LauncherForwarder(launcher.path))) as client: + r = client.get("/metrics?limit=3") + assert r.status_code == 200 + assert r.json() == {"method": "GET", "path": "/metrics?limit=3"} + + r = client.post("/v1/messages/batches/b1/cancel", json={"reason": "x"}) + assert r.status_code == 201 + assert r.json()["path"] == "/v1/messages/batches/b1/cancel" + method, path, body = launcher.seen[-1] + assert (method, path) == ("POST", "/v1/messages/batches/b1/cancel") + assert '"reason"' in body + + # Routes outside FORWARDED_ROUTES stay local. + assert client.get("/health").json() == {"local": True} + finally: + launcher.close() + + +def test_forwarder_reports_unreachable_launcher(tmp_path) -> None: + with TestClient(_attached_app(LauncherForwarder(str(tmp_path / "missing.sock")))) as client: + r = client.get("/metrics") + assert r.status_code == 503 + assert "unavailable" in r.json()["error"] + + +def test_watchdog_marks_dead_after_consecutive_health_failures(tmp_path) -> None: + launcher = _FakeLauncher(str(tmp_path / "launcher.sock"), health_status=503) + deaths = [] + + async def run() -> None: + forwarder = LauncherForwarder(launcher.path) + try: + watchdog = AttachedFrontendWatchdog( + forwarder, + os.getppid(), + deaths.append, + ppid_interval_s=0.01, + health_interval_s=0.02, + health_failures=2, + ) + await asyncio.wait_for(watchdog.run(), timeout=10) + finally: + await forwarder.close() + + try: + asyncio.run(run()) + finally: + launcher.close() + assert len(deaths) == 1 + assert "failed 2 consecutive" in str(deaths[0]) + + +def test_watchdog_marks_dead_when_launcher_process_is_gone(tmp_path) -> None: + deaths = [] + + async def run() -> None: + forwarder = LauncherForwarder(str(tmp_path / "none.sock")) + try: + # No process can have pid -1 as our parent: the launcher is gone. + watchdog = AttachedFrontendWatchdog(forwarder, -1, deaths.append, ppid_interval_s=0.01) + await asyncio.wait_for(watchdog.run(), timeout=5) + finally: + await forwarder.close() + + asyncio.run(run()) + assert len(deaths) == 1 + assert "exited" in str(deaths[0]) From 5660fa87120ab9819da5559967baa265af530776 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 20 Sep 2026 00:40:15 -0700 Subject: [PATCH 2/2] [None][fix] serve: gate multi-frontend route ownership with a CPU test Every route the LLM OpenAIServer registers must be either launcher-owned (FORWARDED_ROUTES, forwarded by attached frontends) or explicitly listed as safe to serve from whichever frontend accepted the connection, with the reason. A new endpoint backed by per-process state that should be server-wide now fails the test instead of silently shipping K copies. Also catches forwarding templates that no longer match a route and stale allowlist entries. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_serve_multi_frontend_route_gate.py | 168 ++++++++++++++++++ 2 files changed, 169 insertions(+) create mode 100644 tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 2920bce5e75a..3a99dfbb7627 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -147,6 +147,7 @@ l0_cpu: - unittest/llmapi/test_sampling_params.py - unittest/llmapi/test_serialization.py - unittest/llmapi/test_serve_multi_frontend_helpers.py + - unittest/llmapi/test_serve_multi_frontend_route_gate.py - unittest/llmapi/test_serve_report_addr.py - unittest/llmapi/test_session_prefetcher.py - unittest/llmapi/test_tokenizer_aliases.py diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py new file mode 100644 index 000000000000..bdf218211c79 --- /dev/null +++ b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Route-ownership gate for trtllm-serve multi-frontend serving. + +With ``num_serve_frontends > 1`` a client cannot choose which frontend process +answers, so every HTTP route of the LLM ``OpenAIServer`` must be classified: + +* **launcher-owned** -- backed by state that exists once per server (engine + stats queue, profiler, batch store). Attached frontends forward these to the + launcher; the list is ``tensorrt_llm.serve.multi_frontend.FORWARDED_ROUTES``. +* **per-frontend safe** -- correct (or deliberately per-process) when served by + whichever frontend accepted the connection; listed below with the reason. + +A new route that is in neither set fails this test, forcing the author to +decide instead of silently shipping a per-process copy of server-wide state. +""" + +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from starlette.routing import Match + +from tensorrt_llm.serve.multi_frontend import FORWARDED_ROUTES +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + +# (method, path template) -> why it is fine to answer from any frontend. +PER_FRONTEND_SAFE = { + ( + "GET", + "/health", + ): "liveness of the answering process; engine death propagates via the watchdog", + ("GET", "/health_generate"): "runs a real request through the shared executor", + ("GET", "/version"): "constant", + ("GET", "/v1/models"): "constant", + ("GET", "/server_info"): "static server/model configuration, identical in every frontend", + ("POST", "/v1/completions"): "request path: the whole point of multiple frontends", + ("POST", "/v1/chat/completions"): "request path", + ("POST", "/v1/messages"): "request path (Anthropic adapter)", + ("POST", "/v1/messages/count_tokens"): "tokenizer-only, identical in every frontend", + ( + "POST", + "/v1/responses", + ): "request path; the stateful store is disabled group-wide in multi-frontend mode", + ( + "GET", + "/v1/responses/{response_id}", + ): "always 404 in multi-frontend mode (store disabled group-wide)", + ( + "DELETE", + "/v1/responses/{response_id}", + ): "always 404 in multi-frontend mode (store disabled group-wide)", + ("POST", "/_internal/tokenize"): "tokenizer-only", + ("GET", "/energy_metrics"): "reads GPU energy counters; same devices from every process", + ("GET", "/v1/data_transceiver_state"): "fetched from the shared executor over RPC", + ("GET", "/steady_clock_offset"): "clock probe of the answering process", + ( + "POST", + "/steady_clock_offset", + ): "known limitation: calibrates the answering frontend only (documented)", + # RL control requires AsyncLLM, which is not the classic IPC executor path + # multi-frontend mode runs on; the two never coexist. + ("POST", "/release_memory"): "AsyncLLM-only; never registered together with attached frontends", + ("POST", "/resume_memory"): "AsyncLLM-only; never registered together with attached frontends", + ("POST", "/update_weights"): "AsyncLLM-only; never registered together with attached frontends", +} + + +def _llm_server_routes() -> list[tuple[str, str]]: + """(method, path) pairs the LLM server registers, without building an LLM.""" + server = object.__new__(OpenAIServer) + server.app = FastAPI() + server.generator = SimpleNamespace( + _executor=SimpleNamespace(resource_governor_queue=None), + args=SimpleNamespace(return_perf_metrics=False), + ) + server.use_harmony = False + server._enable_rl_control_endpoints = True + server.resource_governor = None + server.register_routes() + return sorted( + (method, route.path) + for route in server.app.routes + if hasattr(route, "methods") and route.methods + for method in route.methods + if method != "HEAD" + ) + + +def _forwarding_app() -> FastAPI: + app = FastAPI() + + async def _stub(): + return {} + + for method, path in FORWARDED_ROUTES: + app.add_api_route(path, _stub, methods=[method]) + return app + + +def _is_forwarded(app: FastAPI, method: str, path_template: str) -> bool: + # Concretise the registered template so the forwarding templates + # (including the {rest:path} catch-all) can be matched against it. + concrete = path_template + for name in ("{batch_id}", "{response_id}"): + concrete = concrete.replace(name, "x") + scope = {"type": "http", "method": method, "path": concrete, "path_params": {}} + return any(route.matches(scope)[0] == Match.FULL for route in app.routes) + + +def test_every_llm_route_is_classified() -> None: + fwd = _forwarding_app() + unclassified = [] + double = [] + for method, path in _llm_server_routes(): + forwarded = _is_forwarded(fwd, method, path) + safe = (method, path) in PER_FRONTEND_SAFE + if forwarded and safe: + double.append((method, path)) + elif not forwarded and not safe: + unclassified.append((method, path)) + assert not double, f"routes both forwarded and per-frontend-safe: {double}" + assert not unclassified, ( + "New OpenAIServer route(s) without a multi-frontend ownership decision: " + f"{unclassified}. Either add them to FORWARDED_ROUTES in " + "tensorrt_llm/serve/multi_frontend.py (state lives once per server, e.g. " + "an engine queue, the profiler or an in-memory store) or to " + "PER_FRONTEND_SAFE in this test with the reason they are correct when " + "served by whichever frontend accepted the connection." + ) + + +def test_every_forwarded_route_exists() -> None: + """A renamed route must not leave a dangling forwarding template.""" + registered = _llm_server_routes() + reg_app = FastAPI() + + async def _stub(): + return {} + + for method, path in registered: + reg_app.add_api_route(path, _stub, methods=[method]) + for method, template in FORWARDED_ROUTES: + probe = template.replace("{rest:path}", "x/cancel") + scope = {"type": "http", "method": method, "path": probe, "path_params": {}} + assert any(route.matches(scope)[0] == Match.FULL for route in reg_app.routes), ( + f"FORWARDED_ROUTES entry {method} {template} matches no registered route" + ) + + +def test_allowlist_has_no_stale_entries() -> None: + registered = set(_llm_server_routes()) + stale = [r for r in PER_FRONTEND_SAFE if r not in registered] + assert not stale, f"PER_FRONTEND_SAFE lists routes that no longer exist: {stale}"