diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 0873ddcf5251..79f7ee8e0465 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -32,7 +32,7 @@ from pathlib import Path from types import FrameType from typing import (TYPE_CHECKING, Any, Dict, Iterator, NoReturn, Optional, - Sequence, Set) + Sequence, Set, Tuple) import click import torch @@ -72,7 +72,11 @@ from tensorrt_llm.serve import OpenAIDisaggServer, OpenAIServer from tensorrt_llm.serve._frontend_processes import ( _init_multi_frontend_mode, _signal_frontend_ready, - _spawn_attached_frontends, _terminate_attached_frontends) + _spawn_attached_frontends, _terminate_attached_frontends, + _watch_attached_frontends) +from tensorrt_llm.serve.multi_frontend import (LAUNCHER_UDS_ENV, + 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) @@ -517,6 +521,43 @@ def _provision_kv_cache_pool(llm_args: dict, yield +def _bind_launcher_uds(llm) -> Tuple[socket.socket, str]: + """Bind the Unix socket attached frontends forward launcher-owned routes to. + + It lives in the multi-frontend ipc dir, which is private to the user and + removed, socket file included, when the executor proxy shuts down. + """ + # Only the classic IPC proxy has the attribute. + ipc_dir = getattr(llm._executor, "_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) + try: + sock.bind(path) + except OSError: + sock.close() + raise + return sock, path + + +def _share_prometheus_multiproc_dir(llm, ipc_dir: str) -> None: + """Give every frontend one PROMETHEUS_MULTIPROC_DIR before spawning them. + + Not via set_prometheus_multiproc_dir(): it keeps only its latest + TemporaryDirectory alive, so OpenAIServer's later call would delete this one. + """ + if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + return + if llm.args.return_perf_metrics or llm.args.perf_metrics_output_dir: + 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 launch_server( host: str, port: int, @@ -548,6 +589,16 @@ def launch_server( model = served_model_name or llm_args["model"] multi_frontend = _init_multi_frontend_mode(llm_args, multi_frontend_enabled) + # An attached frontend gets the launcher's socket from its spawn env; the + # launcher is its parent. + multi_frontend_serving = None + if multi_frontend.is_attached_frontend: + attached_launcher_uds = os.environ.pop(LAUNCHER_UDS_ENV, None) + if attached_launcher_uds is not None: + multi_frontend_serving = MultiFrontendServing( + launcher_uds=attached_launcher_uds, + is_launcher=False, + launcher_pid=os.getppid()) if enable_runtime_control_endpoints and multi_frontend.num_frontends > 1: raise click.BadParameter( "Runtime control endpoints require num_serve_frontends=1 because " @@ -625,12 +676,25 @@ def launch_server( # server construction, middleware registration, and runtime, or a # failure in between leaks the child processes. frontend_children = [] + stop_watching_frontends = None + sockets = [s] + launcher_uds_sock = None try: if multi_frontend.is_launcher: + launcher_uds_sock, launcher_uds_path = _bind_launcher_uds( + llm) + _share_prometheus_multiproc_dir( + llm, os.path.dirname(launcher_uds_path)) frontend_children = _spawn_attached_frontends( llm, multi_frontend.num_frontends, + extra_env={LAUNCHER_UDS_ENV: launcher_uds_path}, report_failure=_report_observed_child_failure) + stop_watching_frontends = _watch_attached_frontends( + frontend_children) + multi_frontend_serving = MultiFrontendServing( + launcher_uds=launcher_uds_path, is_launcher=True) + sockets.append(launcher_uds_sock) server = OpenAIServer( generator=llm, @@ -647,7 +711,8 @@ def launch_server( internal_disagg_auth_key=internal_disagg_auth_key, enable_runtime_control_endpoints= enable_runtime_control_endpoints, - runtime_control_api_key=runtime_control_api_key) + runtime_control_api_key=runtime_control_api_key, + multi_frontend_serving=multi_frontend_serving) _apply_fastapi_middlewares(server.app, middleware) # Optionally disable GC (default: not disabled) @@ -655,10 +720,15 @@ 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 stop_watching_frontends is not None: + stop_watching_frontends.set() if frontend_children: _terminate_attached_frontends(frontend_children) + if launcher_uds_sock is not None: + # Its file goes with the ipc dir (see _bind_launcher_uds). + launcher_uds_sock.close() def launch_mm_encoder_server( diff --git a/tensorrt_llm/serve/_frontend_processes.py b/tensorrt_llm/serve/_frontend_processes.py index 9ea351799f87..799e4e8eeab9 100644 --- a/tensorrt_llm/serve/_frontend_processes.py +++ b/tensorrt_llm/serve/_frontend_processes.py @@ -196,6 +196,38 @@ def _signal_frontend_ready(multi_frontend: MultiFrontendMode) -> None: os.close(fd) +def _watch_attached_frontends(children: list, interval_s: float = 1.0) -> threading.Event: + """Log an attached frontend that exits while the launcher serves. + + The kernel drops a dead frontend's socket from the shared port, so the + others keep serving and only its in-flight requests are lost; nothing else + reports it. Set the returned event before terminating the group, so the + expected exits stay quiet. + """ + stopping = threading.Event() + + def watch() -> None: + alive = dict(enumerate(children, start=1)) + while alive and not stopping.wait(interval_s): + for frontend_id, child in list(alive.items()): + code = child.poll() + if code is None or stopping.is_set(): + continue + del alive[frontend_id] + if code == 0: # e.g. Ctrl-C reaching the whole process group + logger.info(f"Attached serving frontend {frontend_id} (pid {child.pid}) exited") + continue + how = f"exit code {code}" if code > 0 else f"signal {-code}" + logger.error( + f"Attached serving frontend {frontend_id} (pid {child.pid}) exited with " + f"{how}; {len(alive) + 1} of {len(children) + 1} frontends keep serving, " + "and the requests it was handling are lost." + ) + + threading.Thread(target=watch, name="attached-frontend-watch", daemon=True).start() + return stopping + + def _terminate_attached_frontends(children: list) -> None: """Stop and reap the group with one shared grace deadline.""" for child in children: diff --git a/tensorrt_llm/serve/multi_frontend.py b/tensorrt_llm/serve/multi_frontend.py new file mode 100644 index 000000000000..ea7aa6bd4cea --- /dev/null +++ b/tensorrt_llm/serve/multi_frontend.py @@ -0,0 +1,269 @@ +# 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``). + +The frontends share one executor and one ``SO_REUSEPORT`` port, so a client +cannot pick the process that answers. State that exists once per server lives +in the launcher (frontend 0), which also serves a Unix socket in the +multi-frontend ipc dir: attached frontends forward the routes backed by that +state to it, follow the clock offset it publishes, and watch it for death. +""" + +import asyncio +import contextlib +import math +import os +import tempfile +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" +LAUNCHER_UDS_ENV = "TLLM_SERVE_LAUNCHER_UDS" +CLOCK_OFFSET_NAME = "clock_offset" + +# Launcher-owned routes. Attached frontends register them ahead of the local +# handlers, which they shadow since FastAPI matches in registration order. +FORWARDED_ROUTES = ( + ("GET", "/metrics"), + ("POST", "/kv_cache_events"), + ("POST", "/start_profile"), + ("POST", "/stop_profile"), + ("POST", "/steady_clock_offset"), + ("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}"), +) + +# Hop-by-hop and framing headers, which aiohttp sets for the replayed request. +# The rest is forwarded: x-trtllm-msgpack selects the body codec, and auth +# headers feed the --middleware the launcher runs again. +_NOT_FORWARDED_REQUEST_HEADERS = frozenset( + ( + "accept-encoding", + "connection", + "content-length", + "expect", + "host", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + ) +) + + +@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 + launcher_pid: Optional[int] = None # attached frontends: their parent + + @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: + # Lazy, to bind to the serving loop. These routes are low-rate, so no + # keep-alive connections to race against a launcher shutting down. + if self._session is None or self._session.closed: + 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() not in _NOT_FORWARDED_REQUEST_HEADERS + ] + body = await request.body() + try: + session = await self._get_session() + # The Unix connector ignores the host. + 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 SharedClockOffset: + """The disaggregated server's clock offset, published by the launcher. + + The disaggregated server POSTs it once, but every frontend stamps its own + requests' perf metrics with it. + """ + + def __init__(self, ipc_dir: str): + self._path = os.path.join(ipc_dir, CLOCK_OFFSET_NAME) + self._seen = None + + def publish(self, offset: float) -> None: + """Launcher side: atomically replace the published offset.""" + fd, tmp = tempfile.mkstemp(prefix=f".{CLOCK_OFFSET_NAME}.", dir=os.path.dirname(self._path)) + try: + with os.fdopen(fd, "w") as f: + f.write(repr(float(offset))) + os.replace(tmp, self._path) + except BaseException: + with contextlib.suppress(OSError): + os.unlink(tmp) + raise + + def poll(self) -> Optional[float]: + """The published offset if it changed since the last poll, else None.""" + try: + with open(self._path) as f: + offset = float(f.read()) + except (OSError, ValueError): + return None + if not math.isfinite(offset) or offset == self._seen: + return None + self._seen = offset + return offset + + +async def follow_clock_offset( + shared: SharedClockOffset, apply: Callable[[float], None], interval_s: float = 0.25 +) -> None: + """Attached-frontend loop: apply every offset the launcher publishes.""" + while True: + offset = shared.poll() + if offset is not None: + apply(offset) + await asyncio.sleep(interval_s) + + +class AttachedFrontendWatchdog: + """Detect a lost launcher or engine from an attached frontend. + + An attached frontend owns no workers: its only signals are the launcher + process exiting and the launcher's ``/health`` failing. ``mark_dead`` is + called once. + """ + + 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 + # The launcher starts listening only after its attached frontends are + # ready, so failures count once it has answered. + launcher_answered = False + 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: + launcher_answered = True + failures = 0 + continue + if not launcher_answered: + 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 4d0a9c22b315..f9fd027d4f88 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -93,6 +93,12 @@ 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, + SharedClockOffset, + follow_clock_offset) from tensorrt_llm.serve.openai_protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, @@ -597,7 +603,8 @@ def __init__( enable_rl_control_endpoints: bool = False, rl_control_api_key: Optional[str] = None, enable_runtime_control_endpoints: bool = False, - runtime_control_api_key: Optional[str] = None): + runtime_control_api_key: Optional[str] = None, + multi_frontend_serving: Optional[MultiFrontendServing] = None): if enable_rl_control_endpoints: logger.warning( "RL control endpoints are deprecated; use generic runtime " @@ -641,7 +648,11 @@ def __init__( self._embedding_max_queue_size = embedding_max_queue_size self.embedding_batcher: Optional[EncodeBatcher] = None self.tool_parser = tool_parser - self.metadata_server = create_metadata_server(metadata_server_cfg) + # Attached frontends serve the launcher's port: only it registers. + self._registers_server = (multi_frontend_serving is None + or multi_frontend_serving.is_launcher) + self.metadata_server = (create_metadata_server(metadata_server_cfg) + if self._registers_server else None) self.disagg_cluster_config = disagg_cluster_config self.multimodal_server_config = multimodal_server_config # Seed media-IO decode defaults from the model's input processor, then @@ -725,6 +736,23 @@ def __init__( ) if server_role is not None else "server" self._perf_metrics_writer = PerfMetricsJsonlWriter( perf_metrics_output_dir, server_kind) + # Multi-frontend serving: 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) + self._shared_clock_offset: Optional[SharedClockOffset] = ( + SharedClockOffset( + os.path.dirname(multi_frontend_serving.launcher_uds)) + if multi_frontend_serving is not None else None) + self._clock_offset_task: Optional[asyncio.Task] = None + # Requests finishing on attached frontends never wake the launcher's + # collector, so it also polls. + 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. @@ -755,27 +783,7 @@ def __init__( @asynccontextmanager async def lifespan(app: FastAPI): await self._perf_metrics_writer.start() - if self.metadata_server is not None: - metadata = { - "model": self.model, - "version": VERSION, - "timestamp": datetime.now().isoformat(), - "server_role": server_role.name, - "url": self.binding_addr - } - # TODO: add more metadata - # Register with ETCD using the existing key format - self.metadata_server.put(f"trtllm/{self.generator.llm_id}", - metadata) - logger.info(f"trtllm/{self.generator.llm_id} is registered") - - if self.disagg_cluster_config: - self.disagg_cluster_storage = create_cluster_storage_client( - self.disagg_cluster_config.cluster_uri, - self.disagg_cluster_config.cluster_name) - self.disagg_cluster_worker = DisaggClusterWorker( - self.server_role, self.host, self.port, - self.disagg_cluster_config, self.disagg_cluster_storage) + self._register_server() # VisualGen has no args if not self._is_visual_gen: @@ -792,7 +800,20 @@ async def lifespan(app: FastAPI): # Start background iteration stats collector if metrics are enabled # The PyTorch backend args include `enable_iter_perf_stats`. - if self.metrics_collector and getattr( + # Attached frontends leave the engine queue to the launcher. + 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()) + self._clock_offset_task = asyncio.create_task( + follow_clock_offset(self._shared_clock_offset, + self._apply_shared_clock_offset)) + 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 @@ -825,6 +846,15 @@ async def lifespan(app: FastAPI): yield await self._perf_metrics_writer.close() + for task in (self._launcher_watchdog_task, self._clock_offset_task): + if task is not None: + task.cancel() + try: + await 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") @@ -838,11 +868,7 @@ async def lifespan(app: FastAPI): pass logger.info("Stopped background iteration stats collector task") - if self.metadata_server is not None: - self.metadata_server.remove(f"trtllm/{self.generator.llm_id}") - logger.info(f"trtllm/{self.generator.llm_id} is unregistered") - if self.disagg_cluster_worker: - await self.disagg_cluster_worker.deregister_worker() + await self._deregister_server() if self.resource_governor is not None: self.resource_governor.close() self.generator.shutdown() @@ -889,6 +915,7 @@ async def validation_exception_handler(request, exc): self._init_embedding_batcher() self.register_embedding_routes() else: + self._register_launcher_forwarding_routes() self.register_routes() if self._collect_perf_metrics: @@ -1020,7 +1047,9 @@ 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() + # One config-info series per server, not per frontend process. + if self._launcher_forwarder is None: + self._log_config_info_metrics() @staticmethod def _ensure_post_processor_hook_supported( @@ -1723,6 +1752,67 @@ async def data_transceiver_state(self) -> JSONResponse: base64.b64encode(state).decode("utf-8") }) + def _register_launcher_forwarding_routes(self) -> None: + """Attached frontends: forward launcher-owned routes (before register_routes()).""" + if self._launcher_forwarder is None: + return + for method, path in FORWARDED_ROUTES: + self.app.add_api_route(path, + self._launcher_forwarder.forward, + methods=[method], + include_in_schema=False) + + def _register_server(self) -> None: + """Register with the metadata server and/or disagg cluster (launcher only).""" + if not self._registers_server: + return + if self.metadata_server is not None: + metadata = { + "model": self.model, + "version": VERSION, + "timestamp": datetime.now().isoformat(), + "server_role": self.server_role.name, + "url": self.binding_addr + } + # TODO: add more metadata + # Register with ETCD using the existing key format + self.metadata_server.put(f"trtllm/{self.generator.llm_id}", + metadata) + logger.info(f"trtllm/{self.generator.llm_id} is registered") + + if self.disagg_cluster_config: + self.disagg_cluster_storage = create_cluster_storage_client( + self.disagg_cluster_config.cluster_uri, + self.disagg_cluster_config.cluster_name) + self.disagg_cluster_worker = DisaggClusterWorker( + self.server_role, self.host, self.port, + self.disagg_cluster_config, self.disagg_cluster_storage) + + async def _deregister_server(self) -> None: + if self.metadata_server is not None: + self.metadata_server.remove(f"trtllm/{self.generator.llm_id}") + logger.info(f"trtllm/{self.generator.llm_id} is unregistered") + if self.disagg_cluster_worker: + await self.disagg_cluster_worker.deregister_worker() + + def _apply_shared_clock_offset(self, offset: float) -> None: + self._adjusted_steady_clock.set_reference_offset(offset) + logger.info(f"Following the launcher's steady clock offset between " + f"local and disagg server: {offset} second") + + def _on_launcher_lost(self, error: BaseException) -> None: + """Watchdog callback: fail requests with EngineDeadError, then shut down. + + Mirrors the fatal-error branch of ``health``. + """ + executor = self.generator._executor + if executor is None: # the LLM already shut down + return + executor._set_fatal_error(error) # keeps only the first error + if not executor.doing_shutdown: + _record_generator_termination(self.generator) + signal.raise_signal(signal.SIGINT) + async def health(self) -> Response: if self._check_health(): return Response(status_code=200) @@ -1897,6 +1987,15 @@ async def set_steady_clock_offset( logger.info( f"The steady clock offset between local and disagg server: {offset} second" ) + # Multi-frontend launcher: every frontend stamps perf metrics with it. + if self._shared_clock_offset is not None: + try: + self._shared_clock_offset.publish(offset) + except OSError as e: + logger.error( + f"Failed to share the steady clock offset with the " + f"attached frontends: {e}") + return Response(status_code=500, content=str(e)) return Response(status_code=200) async def get_steady_clock_offset(self) -> JSONResponse: @@ -1944,7 +2043,8 @@ async def _extract_metrics(self, res: RequestOutput, raw_request: Request): 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): + # The collector only runs with enable_iter_perf_stats. + if self._iteration_stats_collector_task is not None: self._iteration_stats_wakeup_event.set() async def _create_chat_response( @@ -1995,22 +2095,32 @@ 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 polls). + woken_by_timer = False + try: + await asyncio.wait_for( + self._iteration_stats_wakeup_event.wait(), + timeout=self._iteration_stats_poll_interval) + 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: a non-blocking read on the metrics + # thread, idle while this loop owns the queue. + for llm_stat in await asyncio.get_running_loop( + ).run_in_executor(self._metrics_stats_executor, + 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}", @@ -2045,6 +2155,13 @@ def _apply_spec_decode_stats_opt_in(self, postproc_args.spec_decode_num_spec_tokens = getattr( self, "_spec_decode_num_spec_tokens", None) + 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 3ed9f4085675..4e1246422e81 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -252,6 +252,8 @@ l0_cpu: - unittest/llmapi/test_runtime_memory_api.py - 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_serve_sleep_mode.py - unittest/llmapi/test_session_prefetcher.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..0a6f3b9055bb --- /dev/null +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -0,0 +1,325 @@ +# 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 tests for trtllm-serve multi-frontend serving, against a fake launcher.""" + +import asyncio +import os +import signal +import subprocess +import sys +import threading +import time +from contextlib import asynccontextmanager +from types import SimpleNamespace +from unittest import mock + +import msgspec +import pytest +from aiohttp import web +from fastapi import FastAPI +from fastapi.testclient import TestClient +from starlette.routing import Match + +from tensorrt_llm._utils import AdjustedSteadyClock +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve import _frontend_processes, openai_server +from tensorrt_llm.serve._frontend_processes import _watch_attached_frontends +from tensorrt_llm.serve.multi_frontend import ( + FORWARDED_ROUTES, + AttachedFrontendWatchdog, + LauncherForwarder, + SharedClockOffset, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +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.health_hits = 0 + 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.read(), dict(request.headers)) + ) + return web.json_response( + {"path": request.path_qs}, status=201 if request.method == "POST" else 200 + ) + + async def health(_: web.Request) -> web.Response: + self.health_hits += 1 + 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_server(forwarder: LauncherForwarder) -> OpenAIServer: + """An attached-frontend OpenAIServer with just enough state to route.""" + + @asynccontextmanager + async def lifespan(_: FastAPI): + yield + await forwarder.close() + + server = object.__new__(OpenAIServer) + server.app = FastAPI(lifespan=lifespan) + server.generator = SimpleNamespace( + _executor=SimpleNamespace(resource_governor_queue=None), + args=SimpleNamespace(return_perf_metrics=False), + ) + server.use_harmony = False + server._enable_rl_control_endpoints = False + server._enable_runtime_control_endpoints = False + server.resource_governor = None + server._launcher_forwarder = forwarder + # The order OpenAIServer.__init__ registers them in. + server._register_launcher_forwarding_routes() + server.register_routes() + return server + + +def test_attached_server_forwards_launcher_owned_routes(tmp_path) -> None: + launcher = _FakeLauncher(str(tmp_path / "launcher.sock")) + forwarder = LauncherForwarder(launcher.path) + server = _attached_server(forwarder) + + # Every launcher-owned route resolves to the forwarder before its local handler. + for method, template in FORWARDED_ROUTES: + path = template.replace("{rest:path}", "b1/cancel") + scope = {"type": "http", "method": method, "path": path, "path_params": {}} + first = next(r for r in server.app.routes if r.matches(scope)[0] == Match.FULL) + assert first.endpoint == forwarder.forward, f"{method} {path} is answered locally" + + body = msgspec.msgpack.encode({"requests": []}) + try: + with TestClient(server.app) as client: + r = client.get("/metrics?limit=3") + assert (r.status_code, r.json()) == (200, {"path": "/metrics?limit=3"}) + # The body is replayed byte for byte, so its codec header must travel + # with it, and so must auth headers a launcher --middleware checks. + r = client.post( + "/v1/messages/batches", + content=body, + headers={"x-trtllm-msgpack": "1", "x-api-key": "secret"}, + ) + assert r.status_code == 201 + assert client.get("/version").status_code == 200 # stays local + finally: + launcher.close() + assert len(launcher.seen) == 2 + _, path, seen_body, headers = launcher.seen[-1] + headers = {k.lower(): v for k, v in headers.items()} + assert (path, seen_body) == ("/v1/messages/batches", body) + assert headers["x-trtllm-msgpack"] == "1" and headers["x-api-key"] == "secret" + assert headers["host"] != "testserver" # hop-by-hop headers are not replayed + + with TestClient(server.app) as client: + r = client.get("/metrics") + assert r.status_code == 503 and "unavailable" in r.json()["error"] + + +async def _wait_for(predicate, timeout: float = 10.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline, "condition not reached in time" + await asyncio.sleep(0.01) + + +def test_watchdog_counts_health_failures_only_after_the_launcher_answered(tmp_path) -> None: + path = str(tmp_path / "launcher.sock") + deaths = [] + + async def run() -> None: + forwarder = LauncherForwarder(path) + watchdog = AttachedFrontendWatchdog( + forwarder, + os.getppid(), + deaths.append, + ppid_interval_s=0.01, + health_interval_s=0.02, + health_failures=2, + ) + task = asyncio.create_task(watchdog.run()) + launcher = None + try: + # The launcher starts listening after its attached frontends: ~25 + # failed probes against a socket nobody serves yet are not fatal. + await asyncio.sleep(0.5) + launcher = await asyncio.to_thread(_FakeLauncher, path, 503) + await _wait_for(lambda: launcher.health_hits >= 3) + assert not deaths and not task.done() + launcher.health_status = 200 + await _wait_for(lambda: launcher.health_hits >= 4) + launcher.health_status = 503 + await asyncio.wait_for(task, timeout=10) + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await forwarder.close() + if launcher is not None: + await asyncio.to_thread(launcher.close) + + asyncio.run(run()) + assert len(deaths) == 1 and "failed 2 consecutive" in str(deaths[0]) + + +def test_watchdog_notices_a_gone_launcher_process(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 and "exited" in str(deaths[0]) + + +def test_launcher_logs_an_attached_frontend_that_dies(monkeypatch) -> None: + errors = [] + monkeypatch.setattr(_frontend_processes.logger, "error", errors.append) + sleeper = [sys.executable, "-c", "import time; time.sleep(60)"] + children = [subprocess.Popen(sleeper) for _ in range(2)] + try: + stopping = _watch_attached_frontends(children, interval_s=0.01) + children[0].kill() + deadline = time.monotonic() + 10 + while not errors and time.monotonic() < deadline: + time.sleep(0.01) + # Exits once the launcher terminates the group are expected. + stopping.set() + children[1].kill() + time.sleep(0.1) + finally: + for child in children: + child.kill() + child.wait() + assert len(errors) == 1 + assert "frontend 1 " in errors[0] and "signal 9" in errors[0] and "2 of 3" in errors[0] + + +@pytest.mark.parametrize("doing_shutdown", [False, True]) +def test_on_launcher_lost_fails_fast_then_shuts_down(monkeypatch, doing_shutdown: bool) -> None: + fatal_errors = [] + executor = SimpleNamespace(doing_shutdown=doing_shutdown, _set_fatal_error=fatal_errors.append) + server = object.__new__(OpenAIServer) + server.generator = SimpleNamespace(_executor=executor) + signals, terminations = [], [] + monkeypatch.setattr(openai_server.signal, "raise_signal", signals.append) + monkeypatch.setattr(openai_server, "_record_generator_termination", terminations.append) + + error = RuntimeError("attached frontend lost its engine: test") + server._on_launcher_lost(error) + + assert fatal_errors == [error] + # A shutdown already under way gets no second SIGINT (which would force-quit it). + assert signals == ([] if doing_shutdown else [signal.SIGINT]) + assert terminations == ([] if doing_shutdown else [server.generator]) + + +def test_shared_clock_offset_publish_and_poll(tmp_path) -> None: + launcher, follower = SharedClockOffset(str(tmp_path)), SharedClockOffset(str(tmp_path)) + assert follower.poll() is None + launcher.publish(-12.5) + assert follower.poll() == -12.5 + assert follower.poll() is None # unchanged + launcher.publish(0.125) + launcher.publish(0.25) + assert follower.poll() == 0.25 # the latest value wins + for junk in ("not a number", "nan"): + (tmp_path / "clock_offset").write_text(junk) + assert follower.poll() is None + assert [p.name for p in tmp_path.iterdir()] == ["clock_offset"] + + +@pytest.mark.parametrize( + "role,status", + [("launcher", 200), ("launcher_without_ipc_dir", 500), ("single_frontend", 200)], +) +def test_set_steady_clock_offset(tmp_path, role: str, status: int) -> None: + shared = { + "launcher": SharedClockOffset(str(tmp_path)), + "launcher_without_ipc_dir": SharedClockOffset(str(tmp_path / "gone")), + "single_frontend": None, + }[role] + server = object.__new__(OpenAIServer) + server._adjusted_steady_clock = AdjustedSteadyClock(time_source=lambda: 100.0) + server._shared_clock_offset = shared + + response = asyncio.run(server.set_steady_clock_offset(offset=-7.25)) + + # A non-200 makes the disaggregated server warn about misaligned timestamps. + assert response.status_code == status + assert server._adjusted_steady_clock.now() == 92.75 + if role == "launcher": + assert SharedClockOffset(str(tmp_path)).poll() == -7.25 + + +@pytest.mark.parametrize("is_launcher", [True, False]) +def test_only_the_launcher_registers_the_server(monkeypatch, is_launcher: bool) -> None: + storage, worker_cls = mock.Mock(), mock.Mock() + monkeypatch.setattr(openai_server, "create_cluster_storage_client", storage) + monkeypatch.setattr(openai_server, "DisaggClusterWorker", worker_cls) + server = object.__new__(OpenAIServer) + server._registers_server = is_launcher + server.metadata_server = mock.Mock() + server.disagg_cluster_config = SimpleNamespace(cluster_uri="etcd://h:2379", cluster_name="c") + server.disagg_cluster_worker = None + server.server_role = ServerRole.CONTEXT + server.model, server.binding_addr = "m", "http://h:8000" + server.host, server.port = "h", 8000 + server.generator = SimpleNamespace(llm_id="h-1-0") + + server._register_server() + + if is_launcher: + key, metadata = server.metadata_server.put.call_args.args + assert key == "trtllm/h-1-0" and metadata["url"] == "http://h:8000" + assert server.disagg_cluster_worker is worker_cls.return_value + else: + server.metadata_server.put.assert_not_called() + storage.assert_not_called() + worker_cls.assert_not_called() 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..e477c020d314 --- /dev/null +++ b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py @@ -0,0 +1,121 @@ +# 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. +"""Every OpenAIServer route must have a multi-frontend owner. + +A route either is forwarded to the launcher (state that exists once per server, +``FORWARDED_ROUTES``) or is safe to answer from whichever frontend accepted the +connection (``PER_FRONTEND_SAFE`` below, with the reason). A new route in +neither set fails here, so its author has to decide. +""" + +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 + +PER_FRONTEND_SAFE = { + ("GET", "/health"): "liveness of this process; engine death reaches it via the watchdog", + ("GET", "/health_generate"): "a real request through the shared executor", + ("GET", "/version"): "constant", + ("GET", "/v1/models"): "constant", + ("GET", "/server_info"): "static configuration", + ("POST", "/v1/completions"): "request path", + ("POST", "/v1/chat/completions"): "request path", + ("POST", "/v1/messages"): "request path", + ("POST", "/v1/messages/count_tokens"): "tokenizer only", + ("POST", "/v1/responses"): "request path; the stateful store is off with several frontends", + ("GET", "/v1/responses/{response_id}"): "store is off with several frontends", + ("DELETE", "/v1/responses/{response_id}"): "store is off with several frontends", + ("POST", "/_internal/tokenize"): "tokenizer only", + ("GET", "/energy_metrics"): "the same GPUs from every process", + ("GET", "/v1/data_transceiver_state"): "read from the shared executor", + ("GET", "/steady_clock_offset"): "every frontend reads the same host clock", + # RL control needs AsyncLLM and runtime control needs one frontend. + ("POST", "/release_memory"): "never served by attached frontends", + ("POST", "/resume_memory"): "never served by attached frontends", + ("POST", "/update_weights"): "never served by attached frontends", + ("GET", "/memory_status"): "never served by attached frontends", + ("GET", "/openapi.json"): "generated, identical everywhere", + ("GET", "/docs"): "static page", + ("GET", "/docs/oauth2-redirect"): "static page", + ("GET", "/redoc"): "static page", +} + + +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 + # Mutually exclusive in practice; both on to enumerate every route. + server._enable_rl_control_endpoints = True + server._enable_runtime_control_endpoints = True + server.resource_governor = None + server.register_routes() + return sorted( + (method, route.path) + for route in server.app.routes + if getattr(route, "methods", None) + for method in route.methods + if method != "HEAD" + ) + + +async def _stub(): + return {} + + +def _forwarding_app(routes=FORWARDED_ROUTES) -> FastAPI: + # No generated docs routes: they would "match" the server's own. + app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None) + for method, path in routes: + app.add_api_route(path, _stub, methods=[method]) + return app + + +def _is_forwarded(app: FastAPI, method: str, path_template: str) -> bool: + concrete = path_template.replace("{batch_id}", "x").replace("{response_id}", "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: + forwarding = _forwarding_app() + routes = _llm_server_routes() + forwarded = {r for r in routes if _is_forwarded(forwarding, *r)} + assert not forwarded & PER_FRONTEND_SAFE.keys(), "both forwarded and per-frontend safe" + unclassified = [r for r in routes if r not in forwarded and r not in PER_FRONTEND_SAFE] + assert not unclassified, ( + f"Routes without a multi-frontend owner: {unclassified}. Add them to " + "FORWARDED_ROUTES (server-wide state) or to PER_FRONTEND_SAFE with the reason." + ) + + +def test_no_stale_entries() -> None: + routes = _llm_server_routes() + for method, template in FORWARDED_ROUTES: + app = _forwarding_app([(method, template)]) + assert any(_is_forwarded(app, *r) for r in routes), f"{method} {template} matches nothing" + assert not PER_FRONTEND_SAFE.keys() - set(routes), "PER_FRONTEND_SAFE lists removed routes"