[None][fix] serve: make multi-frontend metrics, profiling and liveness launcher-owned

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,18 +810,25 @@ 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)
             if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1":
                 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(

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)]

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}"))