From e4cbe62522b29f62fb894838a624b0901690adb7 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:18:52 -0700 Subject: [PATCH] [None][feat] serve: multi-process HTTP frontends on the classic IPC executor path At high concurrency a trtllm-serve worker is host-bound on its single serving process: one asyncio event loop (one GIL) performs every request json.loads+pydantic validate and every SSE chunk write, and queuing on that loop dominates first-token latency while the GPUs idle (measured on DSv4 disagg GEN: gen_preprocessing p50 54ms@c512 -> 1520ms@c2048). vLLM/SGLang address the same limit with multiple API-server processes. TLLM_SERVE_NUM_FRONTENDS=K (env-gated, default off) runs K HTTP frontend processes against ONE executor on the default (classic IPC) orchestrator: - Launcher/attach split: frontend 0 builds the LLM and launches workers as usual, then spawns K-1 children that re-exec the command line with TLLM_EXECUTOR_ATTACH_INFO set; their executor attaches via the new GenerationExecutorFrontendProxy (no MPI session, no worker launch, and shutdown never emits the worker's None engine-shutdown sentinel -- the launcher alone owns the engine lifecycle). - Deterministic ipc endpoints, pre-generated by the launcher (TLLM_MULTI_FRONTEND_IPC_DIR/_HMAC): the rank0 worker BINDS the request ingress (PULL) so every frontend PUSH-connects; each frontend binds its own result lane (PULL) that the worker (non-postproc) or every postprocess worker (one PUSH pipe per frontend) sends to. - client-id namespacing: the top 16 bits of the uint64 client id carry the frontend id; responses are routed by client_id>>48. Responses without a usable client id (e.g. attention-DP dummy requests carry client_id=None) route to the launcher, preserving today's silent discard semantics. Frontend 0 keeps ids bit-identical to today. - All frontends bind the public port with SO_REUSEPORT, so the kernel load-balances accepted connections across their independent processes (and GILs); clients and the disagg orchestrator still see one URL. Known limitations (documented): /metrics and /perf_metrics are served per-frontend (SO_REUSEPORT samples one frontend per request); enable_resource_governor is rejected in multi-frontend mode. Validation: 12 CPU-only unit tests (id namespacing, response-lane bucketing incl. the ADP-dummy None guard, attached-proxy submit/cancel/ never-sends-sentinel over real ipc sockets). E2E A/B on DSv4-Pro disagg (11xCTX-DEP4 + 1xGEN-DEP16, c3120, back-to-back on identical nodes): gen_preprocessing p50 1075ms -> 37.6ms (-96.5%), p95 3145ms -> 75.7ms; per-request SSE speed p50 +30%; 16 frontends, 75 min, zero tracebacks; engine-side phases (gen_queue, kv_transfer) byte-identical to baseline. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/commands/serve.py | 100 +++++++- tensorrt_llm/executor/base_worker.py | 49 +++- tensorrt_llm/executor/executor.py | 24 ++ tensorrt_llm/executor/postproc_worker.py | 45 +++- tensorrt_llm/executor/proxy.py | 201 +++++++++++++-- tensorrt_llm/executor/utils.py | 76 ++++++ tensorrt_llm/executor/worker.py | 35 ++- tensorrt_llm/llmapi/llm.py | 8 +- .../executor/test_multi_frontend_routing.py | 235 ++++++++++++++++++ 9 files changed, 732 insertions(+), 41 deletions(-) create mode 100644 tests/unittest/executor/test_multi_frontend_routing.py diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 30c728d0b78e..6141962442a7 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -9,6 +9,7 @@ import socket import subprocess # nosec B404 import sys +import tempfile import uuid from pathlib import Path from typing import Any, Dict, Literal, Mapping, Optional, Sequence @@ -294,6 +295,64 @@ def get_llm_args( return llm_args, llm_args_extra_dict +def _spawn_attached_frontends(llm, num_frontends: int) -> list: + """Spawn num_frontends - 1 attached serving frontend processes. + + Multi-frontend serving (TLLM_SERVE_NUM_FRONTENDS > 1): each child + re-execs this trtllm-serve command line with env vars pointing at the + launcher executor's attach endpoints — the multi-frontend request + ingress plus per-frontend result lanes (see + GenerationExecutorProxy._setup_queues); the child's executor attaches to + the already-running worker instead of launching a new one (see + executor.py GenerationExecutor.create). All frontends bind the serving + port with SO_REUSEPORT, so the kernel load-balances accepted + connections across their independent processes (and GILs). + """ + from tensorrt_llm.executor.proxy import GenerationExecutorProxy + + executor = getattr(llm, "_executor", None) + if not isinstance(executor, GenerationExecutorProxy) or ( + attach_info := executor.multi_frontend_attach_info()) is None: + raise ValueError( + "TLLM_SERVE_NUM_FRONTENDS > 1 requires the classic IPC executor " + f"proxy in multi-frontend mode, got {type(executor).__name__}") + # mkstemp creates the file 0600: it carries the executor HMAC keys. + fd, attach_info_path = tempfile.mkstemp(prefix="trtllm_frontend_", + suffix=".json") + with os.fdopen(fd, "w") as f: + json.dump(attach_info, f) + + children = [] + for frontend_id in range(1, num_frontends): + env = os.environ.copy() + env["TLLM_EXECUTOR_ATTACH_INFO"] = attach_info_path + env["TLLM_EXECUTOR_FRONTEND_ID"] = str(frontend_id) + # Attached frontends are plain RPC clients: keep them out of the + # launcher's MPI job (an inherited PMI rank identity would make + # mpi4py try to (re-)join it at import time). + env["TLLM_DISABLE_MPI"] = "1" + for key in list(env): + if key.startswith(("OMPI_", "PMIX_", "PMI_")): + del env[key] + child = subprocess.Popen([sys.executable] + sys.argv, + env=env) # nosec B603 + children.append(child) + logger.info( + f"Launched attached serving frontend {frontend_id} (pid {child.pid})" + ) + return children + + +def _terminate_attached_frontends(children: list) -> None: + for child in children: + child.terminate() + for child in children: + try: + child.wait(timeout=10) + except subprocess.TimeoutExpired: + child.kill() + + def launch_server( host: str, port: int, @@ -309,6 +368,33 @@ def launch_server( backend = llm_args["backend"] model = served_model_name or llm_args["model"] + + # Multi-frontend serving (prototype): TLLM_SERVE_NUM_FRONTENDS=K runs K + # HTTP frontend processes against ONE executor. The launcher (this + # process, frontend 0) launches the worker as usual and spawns K-1 + # attached frontends; attached frontends (TLLM_EXECUTOR_ATTACH_INFO set) + # skip the spawning. Supported only on the classic IPC executor path + # (the default orchestrator). + num_frontends = int(os.getenv("TLLM_SERVE_NUM_FRONTENDS", "1") or "1") + if not 0 < num_frontends <= (1 << 16): + raise ValueError( + f"TLLM_SERVE_NUM_FRONTENDS out of range: {num_frontends}") + is_attached_frontend = os.getenv("TLLM_EXECUTOR_ATTACH_INFO") is not None + multi_frontend = num_frontends > 1 or is_attached_frontend + if multi_frontend and not is_attached_frontend: + if llm_args.get("orchestrator_type") is not None: + raise ValueError( + "TLLM_SERVE_NUM_FRONTENDS > 1 currently supports only the " + "default (classic IPC) executor path, not orchestrator_type=" + f"{llm_args.get('orchestrator_type')!r}") + # Opt the launcher executor into multi-frontend mode BEFORE it is + # created: pre-generate the shared ipc directory and HMAC key so the + # launcher proxy, the rank0 worker and the attached frontends agree + # on deterministic endpoints (GenerationExecutorProxy._setup_queues). + os.environ["TLLM_MULTI_FRONTEND_IPC_DIR"] = tempfile.mkdtemp( + prefix="trtllm_frontends_") + os.environ["TLLM_MULTI_FRONTEND_HMAC"] = os.urandom(32).hex() + addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM) address_family = socket.AF_INET6 if all( @@ -316,6 +402,10 @@ def launch_server( with socket.socket(address_family, socket.SOCK_STREAM) as s: # If disagg cluster config is provided and port is not specified, try to find a free port, otherwise try to bind to the specified port assert port > 0 or disagg_cluster_config is not None, "Port must be specified if disagg cluster config is not provided" + if multi_frontend: + # Every frontend process binds its own listening socket on the + # same port; the kernel load-balances accepts across them. + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) try: s.bind((host, port)) if port == 0: @@ -340,6 +430,10 @@ def launch_server( f"{backend} is not a known backend, check help for available options.", param_hint="backend") + frontend_children = [] + if multi_frontend and not is_attached_frontend: + frontend_children = _spawn_attached_frontends(llm, num_frontends) + server = OpenAIServer(generator=llm, model=model, tool_parser=tool_parser, @@ -354,7 +448,11 @@ def launch_server( if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": gc.disable() - asyncio.run(server(host, port, sockets=[s])) + try: + asyncio.run(server(host, port, sockets=[s])) + finally: + if frontend_children: + _terminate_attached_frontends(frontend_children) def launch_grpc_server(host: str, diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 0fc8ba8289f5..e6ec203f87a0 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -36,6 +36,7 @@ from .result import (GenerationResult, LogProbsResult, ResponseWrapper, compute_logprobs) from .utils import (ErrorResponse, IntraProcessQueue, RequestError, + bucket_responses_by_frontend, get_frontend_id, is_llm_response) if TYPE_CHECKING: @@ -105,6 +106,9 @@ def __init__( self.engine = None self.result_queue: Optional[IpcQueue] = None self.postproc_queues: Optional[List[IpcQueue]] = None + # Multi-frontend serving: one result lane per frontend process, + # selected by the frontend id in client_id's top bits. + self.frontend_result_queues: Optional[List[IpcQueue]] = None self.rank = mpi_rank() self.global_rank = global_mpi_rank() # mapping: client_id -> GenerationResult @@ -329,13 +333,24 @@ def fetch_kv_cache_events(self) -> list: def set_result_queue(self, queue): """In multi-gpu mode, result_queue will be set here to communicate between the proxy and the worker 0 process.""" assert self.postproc_queues is None + assert self.frontend_result_queues is None self.result_queue = queue def set_postproc_queues(self, queues: List["IpcQueue"]): """ Set the IPC queues for feeding post-processing processes. """ assert self.result_queue is None + assert self.frontend_result_queues is None self.postproc_queues = queues + def set_frontend_result_queues(self, queues: List["IpcQueue"]): + """Multi-frontend serving: one result lane per frontend process. + + The lane is selected by the frontend id in client_id's top bits. + """ + assert self.result_queue is None + assert self.postproc_queues is None + self.frontend_result_queues = queues + def _set_iteration_result_queue(self, it_result_queue: IterationResultQueue, queue: Union[Queue, FusedIpcQueue, IntraProcessQueue]): @@ -769,20 +784,20 @@ def responses_handler(self, responses: List[tllm.Response]): HandlerKind = AwaitResponseHelper.HandlerKind if self.handler_kind is HandlerKind.unknown: - if not (self.worker.result_queue is not None - or self.worker.postproc_queues is not None): + has_ipc_queues = (self.worker.result_queue is not None + or self.worker.postproc_queues is not None + or self.worker.frontend_result_queues is not None) + if not has_ipc_queues: logger_debug(f"creating await_response helper for Worker\n", color="yellow") # When ExecutorBindingWorker is used in the main process # aka the single process mode self.handler_kind = HandlerKind.single_process_worker - elif self.worker.result_queue is not None or self.worker.postproc_queues is not None: + else: # The ExecutorBindingProxy is used logger_debug(f"creating await_response helper for IPC\n", color="yellow") self.handler_kind = HandlerKind.ipc_batched - else: - raise NotImplementedError match self.handler_kind: case HandlerKind.single_process_worker: @@ -877,7 +892,15 @@ def handle_for_ipc_batched(self, responses: List[tllm.Response]) -> None: self.worker.postproc_queues[wid].put(batch) if rsp_batch: - self.worker.result_queue.put(rsp_batch) + if (lanes := self.worker.frontend_result_queues) is not None: + # Multi-frontend serving: route each response to its origin + # frontend's lane by the id in client_id's top bits. + for frontend_id, sub_batch in enumerate( + bucket_responses_by_frontend(rsp_batch, len(lanes))): + if sub_batch: + lanes[frontend_id].put(sub_batch) + else: + self.worker.result_queue.put(rsp_batch) def _get_params_for_first_rsp( @@ -981,7 +1004,19 @@ def _send_rsp( rsp_batch: Optional[List[tllm.Response]] = None): # if postproc_batches is set, append to batch instead of putting to IpcQueue - if worker.result_queue is not None: + if worker.frontend_result_queues is not None: + # Multi-frontend serving: route to the origin frontend's result lane + # (client_id may be None for e.g. ADP dummy requests -- those go to + # lane 0, the launcher, which silently discards them like today). + if rsp_batch is not None: + rsp_batch.append(response) + else: + lanes = worker.frontend_result_queues + frontend_id = get_frontend_id(getattr(response, "client_id", None)) + if frontend_id >= len(lanes): + frontend_id = 0 + lanes[frontend_id].put(response) + elif worker.result_queue is not None: if rsp_batch is not None: rsp_batch.append(response) else: diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index ed011ed82b57..254c02696c85 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -1,6 +1,8 @@ import atexit import faulthandler +import json import multiprocessing +import os import platform import signal import traceback @@ -543,6 +545,28 @@ def create( f"Using {postproc_worker_config.num_postprocess_workers} postprocess parallel processes.\n", "green") + # Multi-frontend serving: attach to an already-running executor + # instead of launching one. Set by trtllm-serve for the attached + # frontend processes (see commands/serve.py); the frontend id is + # picked up from TLLM_EXECUTOR_FRONTEND_ID. + attach_info_path = os.getenv("TLLM_EXECUTOR_ATTACH_INFO") + if attach_info_path: + with open(attach_info_path) as f: + attach_info = json.load(f) + if attach_info.get("mode") != "classic": + raise ValueError( + "TLLM_EXECUTOR_ATTACH_INFO only supports the classic IPC " + f"executor path, got mode={attach_info.get('mode')!r}") + from .proxy import GenerationExecutorFrontendProxy + frontend_id = int(os.environ["TLLM_EXECUTOR_FRONTEND_ID"]) + logger.info(f"Attaching executor frontend {frontend_id} to the " + f"running classic IPC worker via {attach_info_path}") + return GenerationExecutorFrontendProxy( + attach_info, + frontend_id=frontend_id, + postproc_worker_config=postproc_worker_config, + is_llm_executor=is_llm_executor) + worker_kwargs = { "engine": engine, "executor_config": executor_config, diff --git a/tensorrt_llm/executor/postproc_worker.py b/tensorrt_llm/executor/postproc_worker.py index cbaeb33bab2f..c58f0d534e1b 100644 --- a/tensorrt_llm/executor/postproc_worker.py +++ b/tensorrt_llm/executor/postproc_worker.py @@ -14,7 +14,7 @@ from ..logger import logger from ..sampling_params import SamplingParams from .ipc import ZeroMqQueue -from .utils import ErrorResponse, is_llm_response +from .utils import ErrorResponse, bucket_responses_by_frontend, is_llm_response if TYPE_CHECKING: from ..disaggregated_params import DisaggregatedParams @@ -78,7 +78,8 @@ class Output(NamedTuple): def __init__( self, pull_pipe_addr: tuple[str, Optional[bytes]], - push_pipe_addr: tuple[str, Optional[bytes]], + push_pipe_addr: Union[tuple[str, Optional[bytes]], + List[tuple[str, Optional[bytes]]]], tokenizer_dir: str, record_creator: Callable[ ["PostprocWorker.Input", TransformersTokenizer], Any], @@ -86,7 +87,10 @@ def __init__( ''' Args: pull_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the input IPC. - push_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the output IPC. + push_pipe_addr: The address and HMAC key of the output IPC, or a + list of them with multi-frontend serving -- one result lane + per frontend, selected by the frontend id in the output's + client_id top bits. tokenizer_dir (str): The directory to load tokenizer. record_creator (Callable[["ResponsePostprocessWorker.Input"], Any]): A creator for creating a record for a request. result_handler (Optional[Callable[[GenerationResultBase], Any]]): A callback handles the final result. @@ -98,11 +102,16 @@ def __init__( is_async=True, is_server=False, name="postprocess_pull_pipe") - self._push_pipe = ZeroMqQueue(address=push_pipe_addr, - is_async=True, - is_server=False, - socket_type=zmq.PUSH, - name="postprocess_push_pipe") + push_pipe_addrs = (push_pipe_addr if isinstance(push_pipe_addr, list) + else [push_pipe_addr]) + self._push_pipes = [ + ZeroMqQueue(address=addr, + is_async=True, + is_server=False, + socket_type=zmq.PUSH, + name=f"postprocess_push_pipe_{i}") + for i, addr in enumerate(push_pipe_addrs) + ] self._to_stop = asyncio.Event() self._q = deque() @@ -172,11 +181,21 @@ async def _batched_put(self): ''' Batched IPC send. ''' async for batch in self._mainloop(): if batch is None: - # notify dispatch_result corountine to quit - await self._push_pipe.put_async(None) + # notify the dispatch_result coroutine in every frontend to + # quit + for pipe in self._push_pipes: + await pipe.put_async(None) break assert isinstance(batch, list) - await self._push_pipe.put_async(batch) + if len(self._push_pipes) == 1: + await self._push_pipes[0].put_async(batch) + continue + # Multi-frontend serving: route each output to its origin + # frontend's result lane by the id in client_id's top bits. + for frontend_id, sub_batch in enumerate( + bucket_responses_by_frontend(batch, len(self._push_pipes))): + if sub_batch: + await self._push_pipes[frontend_id].put_async(sub_batch) async def _mainloop(self): ''' The loop for handle_response and keep producing outputs. ''' @@ -257,7 +276,9 @@ async def main(): @print_traceback_on_error def postproc_worker_main(feedin_ipc_addr: tuple[str, Optional[bytes]], - feedout_ipc_addr: tuple[str, Optional[bytes]], + feedout_ipc_addr: Union[tuple[str, Optional[bytes]], + List[tuple[str, + Optional[bytes]]]], tokenizer_dir: str, record_creator: Callable): worker = PostprocWorker(feedin_ipc_addr, feedout_ipc_addr, diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 68189351e6aa..9e21b4ce90fb 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -27,12 +27,15 @@ from .rpc import RPCClient from .rpc.rpc_common import RPCError, get_unique_ipc_addr from .utils import (ErrorResponse, RequestError, WorkerCommIpcAddrs, - create_mpi_comm_session, get_spawn_proxy_process_env, - is_llm_response, print_alive_threads) + create_mpi_comm_session, get_multi_frontend_ipc_info, + get_spawn_proxy_process_env, is_llm_response, + multi_frontend_request_addr, multi_frontend_result_addr, + namespace_client_id, print_alive_threads) from .worker import GenerationExecutorWorker, worker_main __all__ = [ "GenerationExecutorProxy", + "GenerationExecutorFrontendProxy", ] @@ -96,6 +99,22 @@ def __init__( self._enable_resource_governor = bool( getattr(_llm_args, "enable_resource_governor", False)) + # Multi-frontend serving (TLLM_SERVE_NUM_FRONTENDS > 1) on the classic + # IPC path: trtllm-serve (the launcher, frontend 0) pre-generates a + # shared ipc directory + HMAC key so this proxy, the rank0 worker and + # the attached frontends agree on deterministic endpoints (see + # _setup_queues). Attached frontends never reach this class -- they + # construct GenerationExecutorFrontendProxy instead. + self._multi_frontend_info = get_multi_frontend_ipc_info() + self._num_frontends = int( + os.getenv("TLLM_SERVE_NUM_FRONTENDS", "1") + or "1") if self._multi_frontend_info is not None else 1 + if self._num_frontends > 1 and self._enable_resource_governor: + raise ValueError( + "Multi-frontend serving does not support " + "enable_resource_governor: the resource-governor signal only " + "reaches the launcher frontend.") + # Generate RPC address and key for stats RPC self.rpc_addr = get_unique_ipc_addr() self.hmac_key = os.urandom(32) @@ -235,34 +254,88 @@ def _error_monitor_loop(self) -> None: self._shutdown_event.wait(timeout=5.0) def _setup_queues(self) -> WorkerCommIpcAddrs: - - self.request_queue = IpcQueue(is_server=True, - name="proxy_request_queue") + frontend_result_addrs = None + if self._num_frontends > 1: + # Multi-frontend serving: deterministic endpoints shared with the + # attached frontends. The rank0 worker BINDS the request ingress + # (PULL) so every frontend can PUSH-connect; each frontend + # (including this launcher, frontend 0) binds its own result lane + # (PULL) that the worker / postproc processes PUSH-connect to, + # selected by the frontend id in client_id's top bits. + ipc_dir, hmac_key = self._multi_frontend_info + request_addr = (multi_frontend_request_addr(ipc_dir), hmac_key) + frontend_result_addrs = [(multi_frontend_result_addr(ipc_dir, + i), hmac_key) + for i in range(self._num_frontends)] + self.request_queue = IpcQueue(request_addr, + is_server=False, + socket_type=zmq.PUSH, + name="proxy_request_queue") + self.result_queue = FusedIpcQueue(frontend_result_addrs[0], + is_server=True, + fuse_message=False, + socket_type=zmq.PULL, + name="proxy_result_queue") + else: + request_addr = None + self.request_queue = IpcQueue(is_server=True, + name="proxy_request_queue") + # TODO[chunweiy]: Unify IpcQueue and FusedIpcQueue + # Use PULL mode when enable_postprocess_parallel as there are + # multiple senders from multiple processes. + self.result_queue = FusedIpcQueue( + is_server=True, + fuse_message=False, + socket_type=zmq.PULL + if self.enable_postprocess_parallel else zmq.PAIR, + name="proxy_result_queue") self.worker_init_status_queue = IpcQueue( is_server=True, socket_type=zmq.ROUTER, name="worker_init_status_queue") - # TODO[chunweiy]: Unify IpcQueue and FusedIpcQueue - # Use PULL mode when enable_postprocess_parallel as there are - # multiple senders from multiple processes. - self.result_queue = FusedIpcQueue( - is_server=True, - fuse_message=False, - socket_type=zmq.PULL - if self.enable_postprocess_parallel else zmq.PAIR, - name="proxy_result_queue") self._resource_governor_queue = IpcQueue( is_server=True, name="proxy_resource_governor_queue" ) if self._enable_resource_governor else None # Stats and KV events are now fetched via RPC, not IPC queues. return WorkerCommIpcAddrs( - request_queue_addr=self.request_queue.address, + # A connect-mode queue has no bound .address; use the preset one. + request_queue_addr=request_addr + if request_addr is not None else self.request_queue.address, worker_init_status_queue_addr=self.worker_init_status_queue.address, result_queue_addr=self.result_queue.address, resource_governor_queue_addr=self._resource_governor_queue.address if self._resource_governor_queue is not None else None, + frontend_result_queue_addrs=frontend_result_addrs, ) + def multi_frontend_attach_info(self) -> Optional[dict]: + """The attach payload consumed by attached serving frontends. + + See GenerationExecutorFrontendProxy and commands/serve.py. Returns + None unless multi-frontend mode is active. + """ + if self._num_frontends <= 1: + return None + ipc_dir, hmac_key = self._multi_frontend_info + return { + "mode": + "classic", + "request_addr": + multi_frontend_request_addr(ipc_dir), + "result_addrs": [ + multi_frontend_result_addr(ipc_dir, i) + for i in range(self._num_frontends) + ], + "hmac_key": + hmac_key.hex(), + # Stats / KV events / disagg params RPC endpoint on the rank0 + # worker (ROUTER socket, natively multi-client). + "rpc_addr": + self.rpc_addr, + "rpc_hmac_key": + self.hmac_key.hex(), + } + @property def resource_governor_queue(self): return self._resource_governor_queue @@ -643,3 +716,101 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): self.shutdown() return False # propagate the exception + + +class GenerationExecutorFrontendProxy(GenerationExecutorProxy): + """An attached serving frontend for the classic IPC executor path. + + Used for multi-frontend serving (TLLM_SERVE_NUM_FRONTENDS > 1). + PUSH-connects to the request ingress bound by the rank0 worker and binds + its own per-frontend result lane (PULL); the worker routes responses to + this lane by the frontend id embedded in the top bits of client_id (see + rpc_common.namespace_client_id). It never owns the engine: no MPI + session, no worker launch, and shutdown never emits the worker's None + shutdown sentinel -- that right is the launcher frontend's alone. + """ + + def __init__( + self, + attach_info: dict, + *, + frontend_id: int, + postproc_worker_config: Optional[PostprocWorkerConfig] = None, + is_llm_executor: Optional[bool] = None, + ) -> None: + if not 0 < frontend_id < (1 << 16): + raise ValueError(f"frontend_id out of range: {frontend_id}") + if frontend_id >= len(attach_info["result_addrs"]): + raise ValueError( + f"frontend_id {frontend_id} has no result lane: only " + f"{len(attach_info['result_addrs'])} lanes were provisioned") + postproc_worker_config = postproc_worker_config or PostprocWorkerConfig( + ) + # Deliberately skip GenerationExecutorProxy.__init__: it creates an + # MPI session, launches workers, and registers the pre_shutdown + # atexit hook that emits the engine shutdown sentinel. + GenerationExecutor.__init__( + self, + num_postprocess_workers=postproc_worker_config. + num_postprocess_workers, + postprocess_tokenizer_dir=postproc_worker_config. + postprocess_tokenizer_dir, + is_llm_executor=is_llm_executor) + + self._frontend_id = frontend_id + self._results: Dict[int, GenerationResult] = {} + self.garbage_collection_gen0_threshold = None + self.workers_started = False + self.dispatch_result_thread: Optional[ManagedThread] = None + # The resource governor lives with the launcher frontend only; the + # inherited resource_governor_queue property must return None here so + # OpenAIServer takes its governor-disabled path (openai_server.py). + self._resource_governor_queue = None + + hmac_key = bytes.fromhex(attach_info["hmac_key"]) + self.request_queue = IpcQueue( + (attach_info["request_addr"], hmac_key), + is_server=False, + socket_type=zmq.PUSH, + name=f"frontend_{frontend_id}_request_queue") + self.result_queue = FusedIpcQueue( + (attach_info["result_addrs"][frontend_id], hmac_key), + is_server=True, + fuse_message=False, + socket_type=zmq.PULL, + name=f"frontend_{frontend_id}_result_queue") + + # Stats / KV events / disagg params share the rank0 worker's stats + # RPC server with the launcher (ROUTER socket, natively + # multi-client; per-frontend sampling). + self.rpc_client: Optional[RPCClient] = None + if attach_info.get("rpc_addr"): + self.rpc_client = RPCClient(attach_info["rpc_addr"], + hmac_key=bytes.fromhex( + attach_info["rpc_hmac_key"])) + + def _get_next_client_id(self) -> int: + # Embed the frontend id in the top bits so the worker routes the + # responses back to this frontend's result lane. + return namespace_client_id(self._frontend_id, + super()._get_next_client_id()) + + def pre_shutdown(self): + if self.doing_shutdown: + return + self.doing_shutdown = True + # Abort this frontend's in-flight requests so the engine frees their + # slots. Never send the None engine-shutdown sentinel: the launcher + # frontend owns the engine lifecycle (see the class docstring). + self._abort_all_requests() + + def shutdown(self): + self.pre_shutdown() + if self.rpc_client is not None: + self.rpc_client.close() + self.rpc_client = None + # The dispatch thread blocks on result_queue.get(); it is a daemon + # ManagedThread that exits with the process or on the worker's + # per-lane None sentinel at engine teardown. Closing its socket from + # another thread is not ZMQ-safe, so leave the queues to process + # teardown. diff --git a/tensorrt_llm/executor/utils.py b/tensorrt_llm/executor/utils.py index 12eaa9afefa9..6a3cfb5382d8 100644 --- a/tensorrt_llm/executor/utils.py +++ b/tensorrt_llm/executor/utils.py @@ -154,6 +154,82 @@ class WorkerCommIpcAddrs(NamedTuple): worker_init_status_queue_addr: tuple[str, Optional[bytes]] result_queue_addr: tuple[str, Optional[bytes]] resource_governor_queue_addr: Optional[tuple[str, Optional[bytes]]] = None + # Multi-frontend serving (classic IPC path): one result lane per frontend + # process, selected by the frontend id in client_id's top bits. When set, + # the rank0 worker BINDS the request queue (PULL) so every frontend can + # PUSH-connect, and routes responses to these lanes instead of + # result_queue_addr (which then aliases lane 0, the launcher's). + frontend_result_queue_addrs: Optional[list[tuple[str, + Optional[bytes]]]] = None + + +# Multi-frontend client_id namespacing: the top bits of the uint64 client id +# carry the frontend id, the low FRONTEND_ID_SHIFT bits carry the per-frontend +# request counter. Frontend id 0 keeps client ids bit-identical to the legacy +# single-frontend scheme. +FRONTEND_ID_SHIFT = 48 +FRONTEND_COUNTER_MASK = (1 << FRONTEND_ID_SHIFT) - 1 + + +def get_frontend_id(client_id: Optional[int]) -> int: + """Extract the originating frontend id from a namespaced client id.""" + if not isinstance(client_id, int): + return 0 + return client_id >> FRONTEND_ID_SHIFT + + +def namespace_client_id(frontend_id: int, client_id: int) -> int: + """Embed frontend_id in the top bits of a per-frontend client id.""" + return (frontend_id << FRONTEND_ID_SHIFT) | (client_id + & FRONTEND_COUNTER_MASK) + + +def bucket_responses_by_frontend(responses: list, + num_frontends: int) -> list[list]: + """Bucket responses by their originating frontend id (client_id top bits). + + Responses without a usable client_id (e.g. ADP dummy requests carry + client_id=None) and ids with an out-of-range frontend go to bucket 0 + (the launcher), matching legacy single-client visibility where such + responses are silently discarded by the launcher's dispatcher. + """ + buckets = [[] for _ in range(num_frontends)] + for rsp in responses: + frontend_id = get_frontend_id(getattr(rsp, "client_id", None)) + if frontend_id >= num_frontends: + frontend_id = 0 + buckets[frontend_id].append(rsp) + return buckets + + +def get_multi_frontend_ipc_info() -> Optional[tuple[str, bytes]]: + """The shared ipc dir and HMAC key for multi-frontend serving. + + Pre-generated by trtllm-serve (the launcher) on the classic IPC executor + path before the executor is created. None outside that mode. + """ + ipc_dir = os.getenv("TLLM_MULTI_FRONTEND_IPC_DIR") + hmac_hex = os.getenv("TLLM_MULTI_FRONTEND_HMAC") + if ipc_dir and hmac_hex: + return ipc_dir, bytes.fromhex(hmac_hex) + return None + + +def multi_frontend_request_addr(ipc_dir: str) -> str: + """The request ingress endpoint bound by the rank0 worker (PULL). + + Every frontend PUSH-connects to it. + """ + return f"ipc://{os.path.join(ipc_dir, 'request.sock')}" + + +def multi_frontend_result_addr(ipc_dir: str, frontend_id: int) -> str: + """The result lane endpoint bound by frontend ``frontend_id`` (PULL). + + The worker/postproc processes PUSH-connect to it. Deterministic so a + respawned process can rebind the same lane. + """ + return f"ipc://{os.path.join(ipc_dir, f'result_{frontend_id}.sock')}" def is_llm_response(instance): diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index 5243c3a0bd11..0a0697a44aaa 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -198,6 +198,10 @@ def _print_stacks(): postproc_worker_config = postproc_worker_config or PostprocWorkerConfig() is_leader: bool = mpi_rank() == 0 + # Multi-frontend serving (classic IPC path): per-frontend result lanes; + # the request queue is bound here (PULL) so every frontend PUSH-connects. + multi_frontend_addrs = worker_queues.frontend_result_queue_addrs + frontend_result_queues: Optional[List[FusedIpcQueue]] = None if tracer_init_kwargs is not None and is_leader: tracer = VizTracer(**tracer_init_kwargs) tracer.register_exit() @@ -215,7 +219,9 @@ def _print_stacks(): # inherit the log level from "TLLM_LOG_LEVEL" environment variable logger.set_level(log_level) request_queue = IpcQueue(worker_queues.request_queue_addr, - is_server=False, + is_server=multi_frontend_addrs is not None, + socket_type=zmq.PULL if multi_frontend_addrs + is not None else zmq.PAIR, name="worker_request_queue") worker_init_status_queue = IpcQueue( worker_queues.worker_init_status_queue_addr, @@ -237,6 +243,18 @@ def _print_stacks(): name=f"postprocess_{i}_feedin_queue") for i in range(postproc_worker_config.num_postprocess_workers) ] + elif multi_frontend_addrs is not None: + # Multi-frontend serving: one PUSH lane per frontend, selected by + # the frontend id in client_id's top bits (see base_worker + # _send_rsp). + frontend_result_queues = [ + FusedIpcQueue(addr, + is_server=False, + fuse_message=False, + socket_type=zmq.PUSH, + name=f"worker_result_queue_{i}") + for i, addr in enumerate(multi_frontend_addrs) + ] else: # IPC queue for sending results back to the proxy, and let the # Proxy process to handle the postprocess @@ -246,9 +264,12 @@ def _print_stacks(): name="worker_result_queue") def notify_proxy_threads_to_quit(): - # Signal the dispatcher thread in the proxy to quit + # Signal the dispatcher thread in every frontend proxy to quit if result_queue is not None: result_queue.put(None) + elif frontend_result_queues is not None: + for q in frontend_result_queues: + q.put(None) else: assert result_queues is not None for q in result_queues: @@ -258,13 +279,15 @@ def notify_proxy_threads_to_quit(): if is_leader and postproc_worker_config.enabled: logger_debug(f"initiate postprocess workers...", "yellow") - proxy_result_queue: tuple[ - str, Optional[bytes]] = worker_queues.result_queue_addr + # With multi-frontend serving each postproc worker gets every + # frontend's result lane and routes outputs by client_id's top bits. + proxy_result_queue = (multi_frontend_addrs if multi_frontend_addrs + is not None else worker_queues.result_queue_addr) assert result_queues is not None postproc_worker_pool = ProcessPoolExecutor( max_workers=postproc_worker_config.num_postprocess_workers) - assert isinstance(proxy_result_queue, tuple) + assert isinstance(proxy_result_queue, (tuple, list)) for i in range(postproc_worker_config.num_postprocess_workers): fut = postproc_worker_pool.submit( postproc_worker_main, @@ -323,6 +346,8 @@ def notify_proxy_threads_to_quit(): if is_leader: if postproc_worker_config.enabled: worker.set_postproc_queues(result_queues) + elif frontend_result_queues is not None: + worker.set_frontend_result_queues(frontend_result_queues) else: worker.set_result_queue(result_queue) diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 995298b3060b..a4c90fe0a560 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -231,7 +231,13 @@ def __init__(self, "yellow") self.mpi_session = self.args.mpi_session - if self.args.parallel_config.is_multi_gpu: + # Attached serving frontends (TLLM_EXECUTOR_ATTACH_INFO) connect to + # an already-running executor worker: they need no MPI session of + # their own and must not spawn one (see executor.py + # GenerationExecutor.create). + is_attached_frontend = os.getenv( + "TLLM_EXECUTOR_ATTACH_INFO") is not None + if self.args.parallel_config.is_multi_gpu and not is_attached_frontend: if os.getenv("RAY_LOCAL_WORLD_SIZE") is None and get_device_count( ) < self.args.parallel_config.world_size_per_node: raise RuntimeError( diff --git a/tests/unittest/executor/test_multi_frontend_routing.py b/tests/unittest/executor/test_multi_frontend_routing.py new file mode 100644 index 000000000000..ec03224de0ad --- /dev/null +++ b/tests/unittest/executor/test_multi_frontend_routing.py @@ -0,0 +1,235 @@ +# 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-only tests for classic-path multi-frontend serving. + +Multi-frontend serving (TLLM_SERVE_NUM_FRONTENDS) on the classic IPC +executor path. Covers client-id namespacing, worker-side response-lane routing, and the +attached-frontend proxy lifecycle over real ipc:// sockets. +""" + +import os +import tempfile +import time +from types import SimpleNamespace + +from tensorrt_llm.executor.utils import ( + FRONTEND_COUNTER_MASK, + bucket_responses_by_frontend, + get_frontend_id, + namespace_client_id, +) + + +class TestClientIdNamespacing: + def test_frontend_zero_keeps_legacy_ids(self): + for client_id in (1, 42, FRONTEND_COUNTER_MASK): + assert namespace_client_id(0, client_id) == client_id + + def test_roundtrip(self): + for frontend_id in (0, 1, 7, (1 << 16) - 1): + for counter in (1, 12345, FRONTEND_COUNTER_MASK): + client_id = namespace_client_id(frontend_id, counter) + assert get_frontend_id(client_id) == frontend_id + assert client_id & FRONTEND_COUNTER_MASK == counter + assert client_id < (1 << 64) + + def test_counter_wraparound_stays_in_namespace(self): + # A counter larger than 48 bits must not leak into the frontend bits. + client_id = namespace_client_id(3, FRONTEND_COUNTER_MASK + 5) + assert get_frontend_id(client_id) == 3 + assert client_id & FRONTEND_COUNTER_MASK == 4 + + def test_non_int_client_id_routes_to_launcher(self): + assert get_frontend_id(None) == 0 + + +def _response(client_id): + return SimpleNamespace(client_id=client_id) + + +class TestClassicResponseBucketing: + """The classic IPC path routes batches with bucket_responses_by_frontend.""" + + def test_buckets_by_namespace(self): + responses = [ + _response(namespace_client_id(0, 1)), + _response(namespace_client_id(1, 1)), + _response(namespace_client_id(1, 2)), + _response(namespace_client_id(2, 1)), + ] + buckets = bucket_responses_by_frontend(responses, 3) + assert [len(b) for b in buckets] == [1, 2, 1] + assert all(get_frontend_id(r.client_id) == 1 for r in buckets[1]) + + def test_none_client_id_routes_to_launcher(self): + # ADP dummy responses carry client_id=None: they must land in the + # launcher's lane (which silently discards them), not raise. + buckets = bucket_responses_by_frontend([_response(None)], 4) + assert len(buckets[0]) == 1 + assert all(not b for b in buckets[1:]) + + def test_out_of_range_frontend_routes_to_launcher(self): + buckets = bucket_responses_by_frontend([_response(namespace_client_id(7, 1))], 2) + assert len(buckets[0]) == 1 and not buckets[1] + + +class _LaneStub: + def __init__(self): + self.items = [] + + def put(self, obj): + self.items.append(obj) + + +class TestClassicSendRspLaneRouting: + """_send_rsp selects the origin frontend's lane on the non-postproc path.""" + + @staticmethod + def _fake_worker(num_lanes): + pops = [] + return SimpleNamespace( + result_queue=None, + postproc_queues=None, + frontend_result_queues=[_LaneStub() for _ in range(num_lanes)], + _pop_result=pops.append, + ), pops + + def test_error_response_routes_to_origin_lane(self): + from tensorrt_llm.executor.base_worker import _send_rsp + from tensorrt_llm.executor.utils import ErrorResponse + + worker, pops = self._fake_worker(3) + client_id = namespace_client_id(2, 7) + _send_rsp(worker, ErrorResponse(client_id, "boom", 1)) + assert [len(q.items) for q in worker.frontend_result_queues] == [0, 0, 1] + assert pops == [client_id] + + def test_none_client_id_routes_to_launcher_lane(self): + from tensorrt_llm.executor.base_worker import _send_rsp + from tensorrt_llm.executor.utils import ErrorResponse + + worker, _ = self._fake_worker(2) + _send_rsp(worker, ErrorResponse(None, "adp dummy", 1)) + assert len(worker.frontend_result_queues[0].items) == 1 + assert not worker.frontend_result_queues[1].items + + def test_rsp_batch_defers_lane_selection(self): + from tensorrt_llm.executor.base_worker import _send_rsp + from tensorrt_llm.executor.utils import ErrorResponse + + worker, _ = self._fake_worker(2) + rsp_batch = [] + _send_rsp(worker, ErrorResponse(namespace_client_id(1, 3), "x", 1), rsp_batch=rsp_batch) + assert len(rsp_batch) == 1 + assert all(not q.items for q in worker.frontend_result_queues) + + +class TestClassicFrontendProxyEndToEnd: + """GenerationExecutorFrontendProxy against a fake rank0 worker. + + Real ipc:// sockets: namespaced submit, cancel-on-shutdown, and -- + critically -- that an attached frontend NEVER emits the None + engine-shutdown sentinel. + """ + + @staticmethod + def _make_proxy_and_fake_worker(tmpdir, frontend_id=1, num_frontends=2): + import zmq + + from tensorrt_llm.executor.ipc import IpcQueue + from tensorrt_llm.executor.proxy import GenerationExecutorFrontendProxy + + hmac_key = os.urandom(32) + request_addr = f"ipc://{os.path.join(tmpdir, 'request.sock')}" + result_addrs = [ + f"ipc://{os.path.join(tmpdir, f'result_{i}.sock')}" for i in range(num_frontends) + ] + # The fake rank0 worker binds the request ingress (PULL), exactly + # like worker_main does in multi-frontend mode. + worker_ingress = IpcQueue( + (request_addr, hmac_key), + is_server=True, + socket_type=zmq.PULL, + name="fake_worker_request_queue", + ) + proxy = GenerationExecutorFrontendProxy( + { + "mode": "classic", + "request_addr": request_addr, + "result_addrs": result_addrs, + "hmac_key": hmac_key.hex(), + }, + frontend_id=frontend_id, + ) + return proxy, worker_ingress, hmac_key, result_addrs + + def test_submit_namespaces_and_shutdown_never_sends_sentinel(self): + from tensorrt_llm.executor.request import CancellingRequest, GenerationRequest + from tensorrt_llm.sampling_params import SamplingParams + + with tempfile.TemporaryDirectory() as tmpdir: + proxy, worker_ingress, _, _ = self._make_proxy_and_fake_worker(tmpdir) + + # Attributes read by OpenAIServer at init must exist (a missing + # _resource_governor_queue crashed all siblings in the first e2e). + assert proxy.resource_governor_queue is None + + result = proxy.submit(GenerationRequest([1, 2, 3], SamplingParams())) + assert get_frontend_id(result.request_id) == 1 + assert worker_ingress.poll(5) + received = worker_ingress.get() + assert isinstance(received, GenerationRequest) + assert received.id == result.request_id + + # Frontend shutdown aborts its in-flight requests (cancel) but + # must NOT emit the None engine-shutdown sentinel. + proxy.shutdown() + assert worker_ingress.poll(5) + cancel = worker_ingress.get() + assert isinstance(cancel, CancellingRequest) + assert cancel.id == result.request_id + assert not worker_ingress.poll(1), ( + "an attached frontend must never send the engine-shutdown sentinel" + ) + + def test_dispatch_routes_own_lane_responses(self): + import zmq + + from tensorrt_llm.executor.ipc import FusedIpcQueue + from tensorrt_llm.executor.request import GenerationRequest + from tensorrt_llm.executor.utils import ErrorResponse + from tensorrt_llm.sampling_params import SamplingParams + + with tempfile.TemporaryDirectory() as tmpdir: + proxy, _, hmac_key, result_addrs = self._make_proxy_and_fake_worker(tmpdir) + result = proxy.submit(GenerationRequest([1, 2, 3], SamplingParams())) + client_id = result.request_id + assert client_id in proxy._results + + # The fake worker pushes an ErrorResponse down this frontend's + # result lane; the dispatcher must deliver it and retire the + # request. + worker_lane = FusedIpcQueue( + (result_addrs[1], hmac_key), + is_server=False, + fuse_message=False, + socket_type=zmq.PUSH, + name="fake_worker_result_lane", + ) + worker_lane.put(ErrorResponse(client_id, "boom", 1)) + deadline = time.time() + 5 + while client_id in proxy._results and time.time() < deadline: + time.sleep(0.01) + assert client_id not in proxy._results