diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 2f39353c521c..c40d7e43be16 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -5,13 +5,15 @@ import json import os import secrets +import select import signal import socket import subprocess # nosec B404 import sys +import time import uuid from pathlib import Path -from typing import Any, Dict, Optional, Sequence, Set +from typing import Any, Dict, NamedTuple, Optional, Sequence, Set import click import torch @@ -26,7 +28,7 @@ from tensorrt_llm.commands._serve_stability import stability_option from tensorrt_llm.commands.utils import (collect_explicit_cli_keys, get_is_diffusion_only_model) -from tensorrt_llm.executor.utils import LlmLauncherEnvs +from tensorrt_llm.executor.utils import MAX_NUM_FRONTENDS, LlmLauncherEnvs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, @@ -37,7 +39,7 @@ validate_config_bool) from tensorrt_llm.llmapi.llm_args import MultimodalConfig, TorchLlmArgs from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict -from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr +from tensorrt_llm.llmapi.mpi_session import find_free_ipc_addr, split_mpi_env from tensorrt_llm.llmapi.reasoning_parser import (ReasoningParserFactory, resolve_auto_reasoning_parser) from tensorrt_llm.logger import logger, severity_map @@ -197,6 +199,7 @@ def get_llm_args( free_gpu_memory_fraction: float = 0.9, kv_cache_dtype: str = "auto", num_postprocess_workers: int = 0, + num_serve_frontends: int = 1, trust_remote_code: bool = False, revision: Optional[str] = None, reasoning_parser: Optional[str] = None, @@ -270,6 +273,8 @@ def get_llm_args( max_seq_len, "num_postprocess_workers": num_postprocess_workers, + "num_serve_frontends": + num_serve_frontends, "enable_chunked_prefill": enable_chunked_prefill, "enable_attention_dp": @@ -358,6 +363,156 @@ def _diagnose_port_in_use(port: int) -> str: return "; ".join(details) +class MultiFrontendMode(NamedTuple): + """This process's role under multi-frontend serving (prototype).""" + num_frontends: int + is_attached_frontend: bool + + @property + def is_launcher(self) -> bool: + """Owns the engine and spawns/cleans the attached frontends.""" + return self.num_frontends > 1 and not self.is_attached_frontend + + +def _init_multi_frontend_mode(llm_args: dict, + enabled: bool) -> MultiFrontendMode: + """Resolve this process's multi-frontend serving role. + + num_serve_frontends=K runs K HTTP frontend processes against ONE + executor: the launcher (frontend 0) owns the engine and spawns K-1 + attached frontends (classic IPC executor path only). enabled=False + entry points (e.g. disaggregated MPI workers) never honor the knob. + """ + if not enabled: + if llm_args.pop("num_serve_frontends", 1) > 1: + logger.warning("num_serve_frontends is only supported on plain " + "trtllm-serve; ignored on this entry point.") + return MultiFrontendMode(1, False) + + mode = MultiFrontendMode(llm_args.get("num_serve_frontends", 1), + os.getenv("TLLM_EXECUTOR_ATTACH_INFO") is not None) + if mode.is_launcher and llm_args.get("orchestrator_type") is not None: + raise ValueError( + "num_serve_frontends > 1 requires the default (classic IPC) " + "executor path, not orchestrator_type=" + f"{llm_args.get('orchestrator_type')!r}") + return mode + + +def _spawn_attached_frontends(llm, num_frontends: int) -> list: + """Spawn num_frontends - 1 attached serving frontend processes. + + Each child re-execs this trtllm-serve command line with env vars + carrying the launcher executor's attach endpoints; its executor + attaches to the already-running worker instead of launching one (see + GenerationExecutor.create / GenerationExecutorFrontendProxy). + + Blocks until every child signals READY over its inherited pipe: a + successful Popen only proves the process exists, while the frontend + can still fail during executor attach or server setup. Any child + failure (or a missed deadline) fails the whole group, terminating + the children already started, so num_serve_frontends=K never + silently degrades to fewer frontends. + """ + 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( + "num_serve_frontends > 1 requires the classic IPC executor " + f"proxy in multi-frontend mode, got {type(executor).__name__}") + # Carries the executor HMAC keys; the child deletes it from its env + # once consumed (GenerationExecutor.create). + attach_env = json.dumps(attach_info) + + children, ready_fds = [], [] + try: + for frontend_id in range(1, num_frontends): + # Strip MPI/SLURM identity vars: an inherited rank identity would + # make the child's mpi4py try to (re-)join the launcher's job. + env, _ = split_mpi_env() + env["TLLM_EXECUTOR_ATTACH_INFO"] = attach_env + env["TLLM_EXECUTOR_FRONTEND_ID"] = str(frontend_id) + env["TLLM_DISABLE_MPI"] = "1" + read_fd, write_fd = os.pipe() + ready_fds.append(read_fd) + env["TLLM_FRONTEND_READY_FD"] = str(write_fd) + try: + child = subprocess.Popen([sys.executable] + sys.argv, + env=env, + pass_fds=(write_fd, )) # nosec B603 + finally: + # The child now holds the only write end; its exit before + # READY surfaces as EOF on read_fd. + os.close(write_fd) + children.append(child) + logger.info( + f"Launched attached serving frontend {frontend_id} (pid {child.pid})" + ) + _wait_attached_frontends_ready(children, ready_fds) + except BaseException: + _terminate_attached_frontends(children) + raise + finally: + for fd in ready_fds: + os.close(fd) + return children + + +def _wait_attached_frontends_ready(children: list, ready_fds: list) -> None: + """Block until every attached frontend writes its READY byte.""" + timeout = float(os.getenv("TLLM_FRONTEND_READY_TIMEOUT", "300")) + deadline = time.monotonic() + timeout + pending = dict(zip(ready_fds, children)) + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise RuntimeError( + f"{len(pending)} attached frontend(s) not ready within " + f"{timeout:.0f}s (TLLM_FRONTEND_READY_TIMEOUT)") + readable, _, _ = select.select(list(pending), [], [], + min(remaining, 1.0)) + for fd in readable: + child = pending.pop(fd) + if os.read(fd, 1) != b"R": # EOF: pipe closed without READY + raise RuntimeError( + f"Attached frontend (pid {child.pid}) exited before " + "signaling READY") + logger.info(f"Attached frontend (pid {child.pid}) is ready") + for fd, child in list(pending.items()): + if child.poll() is not None: + raise RuntimeError( + f"Attached frontend (pid {child.pid}) exited with code " + f"{child.returncode} before signaling READY") + + +def _signal_frontend_ready(multi_frontend: MultiFrontendMode) -> None: + """Report READY to the launcher over the inherited pipe. + + Called once everything fallible in an attached frontend's startup + (port bind, executor attach, LLM and OpenAIServer construction, + middleware registration) has succeeded; the launcher blocks group + startup on this byte (see _wait_attached_frontends_ready). + """ + ready_fd = os.environ.pop("TLLM_FRONTEND_READY_FD", None) + if not (multi_frontend.is_attached_frontend and ready_fd): + return + fd = int(ready_fd) + os.write(fd, b"R") + os.close(fd) + + +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, @@ -372,10 +527,25 @@ def launch_server( served_model_name: Optional[str] = None, allow_request_chat_template: bool = False, num_input_processor_workers: int = 8, - num_media_load_workers: int = 8): + num_media_load_workers: int = 8, + multi_frontend_enabled: bool = True): backend = llm_args["backend"] model = served_model_name or llm_args["model"] + + multi_frontend = _init_multi_frontend_mode(llm_args, multi_frontend_enabled) + if multi_frontend.is_launcher or multi_frontend.is_attached_frontend: + # The Responses API store is per-process in-memory: with several + # frontends behind one SO_REUSEPORT port, a follow-up request may + # land on a sibling that has no record of the previous response. + # Disable storage group-wide (OpenAIServer.enable_store). + if not os.getenv("TRTLLM_RESPONSES_API_DISABLE_STORE"): + logger.warning( + "num_serve_frontends > 1: stateful Responses API storage " + "(store/previous_response_id) is disabled; the per-frontend " + "in-memory store cannot be shared across frontends.") + os.environ["TRTLLM_RESPONSES_API_DISABLE_STORE"] = "1" + addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM) address_family = socket.AF_INET6 if all( @@ -383,6 +553,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.is_launcher or multi_frontend.is_attached_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: @@ -409,25 +583,39 @@ def launch_server( f"{backend} is not a known backend, check help for available options.", param_hint="backend") - server = OpenAIServer( - generator=llm, - model=model, - tool_parser=tool_parser, - server_role=server_role, - metadata_server_cfg=metadata_server_cfg, - disagg_cluster_config=disagg_cluster_config, - multimodal_server_config=multimodal_server_config, - chat_template=chat_template, - allow_request_chat_template=allow_request_chat_template, - input_processor_workers=num_input_processor_workers, - media_load_workers=num_media_load_workers) - _apply_fastapi_middlewares(server.app, middleware) + # The finally below is the cleanup boundary for the attached + # frontends: it must cover everything from their spawn through + # server construction, middleware registration, and runtime, or a + # failure in between leaks the child processes. + frontend_children = [] + try: + if multi_frontend.is_launcher: + frontend_children = _spawn_attached_frontends( + llm, multi_frontend.num_frontends) + + server = OpenAIServer( + generator=llm, + model=model, + tool_parser=tool_parser, + server_role=server_role, + metadata_server_cfg=metadata_server_cfg, + disagg_cluster_config=disagg_cluster_config, + multimodal_server_config=multimodal_server_config, + chat_template=chat_template, + allow_request_chat_template=allow_request_chat_template, + input_processor_workers=num_input_processor_workers, + media_load_workers=num_media_load_workers) + _apply_fastapi_middlewares(server.app, middleware) - # Optionally disable GC (default: not disabled) - if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": - gc.disable() + # Optionally disable GC (default: not disabled) + if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": + gc.disable() - uvloop.run(server(host, port, sockets=[s])) + _signal_frontend_ready(multi_frontend) + uvloop.run(server(host, port, sockets=[s])) + finally: + if frontend_children: + _terminate_attached_frontends(frontend_children) def launch_grpc_server(host: str, @@ -886,6 +1074,13 @@ def launch_visual_gen_server( help="Number of workers to postprocess raw responses " "to comply with OpenAI protocol.", status="prototype") +@stability_option("--num_serve_frontends", + type=click.IntRange(min=1, max=MAX_NUM_FRONTENDS), + default=1, + help="Number of HTTP frontend processes serving one " + "executor; values > 1 share the serving port via " + "SO_REUSEPORT (classic IPC executor path only).", + status="prototype") @stability_option("--num_input_processor_workers", type=click.IntRange(min=1), default=8, @@ -1067,29 +1262,30 @@ def launch_visual_gen_server( help= "Types of agents to schedule. Now Only Support Open Deep Research agent.", status="prototype") -def serve( - model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str], - post_processor_hook: Optional[str], host: str, port: int, - log_level: str, backend: str, max_beam_width: int, max_batch_size: int, - max_num_tokens: int, max_seq_len: int, tensor_parallel_size: int, - pipeline_parallel_size: int, context_parallel_size: int, - moe_expert_parallel_size: Optional[int], - moe_cluster_parallel_size: Optional[int], gpus_per_node: Optional[int], - free_gpu_memory_fraction: float, kv_cache_dtype: str, - num_postprocess_workers: int, num_input_processor_workers: int, - num_media_load_workers: int, trust_remote_code: bool, - revision: Optional[str], extra_llm_api_options: Optional[str], - reasoning_parser: Optional[str], tool_parser: Optional[str], - metadata_server_config_file: Optional[str], server_role: Optional[str], - fail_fast_on_attention_window_too_large: bool, - otlp_traces_endpoint: Optional[str], enable_chunked_prefill: bool, - enable_attention_dp: bool, disagg_cluster_uri: Optional[str], - media_io_kwargs: Optional[str], agent_percentage: float, - agent_types: Optional[str], video_pruning_rate: Optional[float], - telemetry: bool, custom_module_dirs: list[Path], - chat_template: Optional[str], allow_request_chat_template: bool, - middleware: tuple[str, ...], grpc: bool, enable_visual_gen: bool, - served_model_name: Optional[str], visual_gen_args: Optional[str]): +def serve(model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str], + post_processor_hook: Optional[str], host: str, port: int, + log_level: str, backend: str, max_beam_width: int, + max_batch_size: int, max_num_tokens: int, max_seq_len: int, + tensor_parallel_size: int, pipeline_parallel_size: int, + context_parallel_size: int, moe_expert_parallel_size: Optional[int], + moe_cluster_parallel_size: Optional[int], + gpus_per_node: Optional[int], free_gpu_memory_fraction: float, + kv_cache_dtype: str, num_postprocess_workers: int, + num_serve_frontends: int, num_input_processor_workers: int, + num_media_load_workers: int, trust_remote_code: bool, + revision: Optional[str], extra_llm_api_options: Optional[str], + reasoning_parser: Optional[str], tool_parser: Optional[str], + metadata_server_config_file: Optional[str], + server_role: Optional[str], + fail_fast_on_attention_window_too_large: bool, + otlp_traces_endpoint: Optional[str], enable_chunked_prefill: bool, + enable_attention_dp: bool, disagg_cluster_uri: Optional[str], + media_io_kwargs: Optional[str], agent_percentage: float, + agent_types: Optional[str], video_pruning_rate: Optional[float], + telemetry: bool, custom_module_dirs: list[Path], + chat_template: Optional[str], allow_request_chat_template: bool, + middleware: tuple[str, ...], grpc: bool, enable_visual_gen: bool, + served_model_name: Optional[str], visual_gen_args: Optional[str]): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path @@ -1171,6 +1367,7 @@ def _serve_llm(): free_gpu_memory_fraction=free_gpu_memory_fraction, kv_cache_dtype=kv_cache_dtype, num_postprocess_workers=num_postprocess_workers, + num_serve_frontends=num_serve_frontends, trust_remote_code=trust_remote_code, revision=revision, reasoning_parser=reasoning_parser, @@ -1798,7 +1995,9 @@ def _launch_disaggregated_server(disagg_config_file: str, llm_args: dict): host=server_cfg.hostname, port=server_cfg.port, llm_args=llm_args, - allow_request_chat_template=disagg_config.allow_request_chat_template) + allow_request_chat_template=disagg_config.allow_request_chat_template, + # Disagg ctx/gen MPI workers must not enter multi-frontend mode. + multi_frontend_enabled=False) def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str, diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 670d6f30508f..ec0403666c42 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -49,6 +49,7 @@ from .result import (GenerationResult, LogProbsResult, ResponseWrapper, compute_logprobs, get_metrics_dict) from .utils import (ErrorResponse, IntraProcessQueue, RequestError, + bucket_responses_by_frontend, frontend_lane_index, is_llm_response) if TYPE_CHECKING: @@ -118,6 +119,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 @@ -272,13 +276,21 @@ 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.""" + 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]): @@ -1087,20 +1099,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: @@ -1279,7 +1291,13 @@ 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: + 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( @@ -1393,7 +1411,16 @@ 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: + # Route to the origin frontend's result lane; None/out-of-range ids + # fall back to lane 0 (see frontend_lane_index). + if rsp_batch is not None: + rsp_batch.append(response) + else: + lanes = worker.frontend_result_queues + lanes[frontend_lane_index(response.client_id, + len(lanes))].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 ba5b843ff347..43adcbba93e6 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 @@ -560,6 +562,35 @@ 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 attached + # frontend processes). + attach_env = os.getenv("TLLM_EXECUTOR_ATTACH_INFO") + if attach_env: + attach_info = json.loads(attach_env) + # Consumed: it carries the HMAC keys and must not leak into + # descendant processes. + os.environ.pop("TLLM_EXECUTOR_ATTACH_INFO", None) + 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_env = os.getenv("TLLM_EXECUTOR_FRONTEND_ID") + if frontend_id_env is None: + raise ValueError( + "TLLM_EXECUTOR_ATTACH_INFO is set but " + "TLLM_EXECUTOR_FRONTEND_ID is not; both are set together " + "by trtllm-serve when spawning attached frontends.") + frontend_id = int(frontend_id_env) + logger.info(f"Attaching executor frontend {frontend_id} to the " + "running classic IPC worker") + 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 ece60b2374b1..552c5c8f2bd3 100644 --- a/tensorrt_llm/executor/postproc_worker.py +++ b/tensorrt_llm/executor/postproc_worker.py @@ -15,7 +15,7 @@ from ..sampling_params import SamplingParams from .ipc import ZeroMqQueue from .postprocessor_hook import load_post_processor_hook -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 @@ -86,7 +86,7 @@ class Output(NamedTuple): def __init__( self, pull_pipe_addr: tuple[str, Optional[bytes]], - push_pipe_addr: tuple[str, Optional[bytes]], + push_pipe_addrs: List[tuple[str, Optional[bytes]]], tokenizer_dir: str, record_creator: Callable[ ["PostprocWorker.Input", TransformersTokenizer], Any], @@ -95,7 +95,9 @@ 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_addrs: The addresses and HMAC keys of the output IPC + lanes, one per frontend (a single-element list in + single-frontend mode). 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. @@ -108,11 +110,14 @@ 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") + 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() @@ -192,11 +197,19 @@ 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 + 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. ''' @@ -289,13 +302,13 @@ 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_addrs: List[tuple[str, Optional[bytes]]], tokenizer_dir: str, record_creator: Callable, post_processor_hook: Optional[str] = None): # Pass the hook import path; PostprocWorker builds it once. worker = PostprocWorker(feedin_ipc_addr, - feedout_ipc_addr, + feedout_ipc_addrs, tokenizer_dir=tokenizer_dir, record_creator=record_creator, post_processor_hook=post_processor_hook) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 4850bfcf2f5d..aa661ab02789 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -16,6 +16,8 @@ import concurrent.futures import json import os +import shutil +import tempfile import threading import weakref from queue import Empty @@ -44,12 +46,14 @@ from .utils import (EngineDeadError, ErrorResponse, RequestError, WorkerCommIpcAddrs, create_mpi_comm_session, get_spawn_proxy_process_env, is_llm_response, - print_alive_threads) + multi_frontend_request_addr, multi_frontend_result_addr, + namespace_client_id, print_alive_threads) from .worker import GenerationExecutorWorker, worker_main from .worker_process_monitor import WorkerProcessIdentity, WorkerProcessMonitor __all__ = [ "GenerationExecutorProxy", + "GenerationExecutorFrontendProxy", ] # Methods that are explicitly implemented for multi-rank MPI/IPC executor @@ -161,6 +165,24 @@ def __init__( self._enable_resource_governor = bool( getattr(_llm_args, "enable_resource_governor", False)) + # Multi-frontend serving: this launcher proxy owns the shared ipc + # dir + HMAC key for the per-frontend endpoints (_setup_queues); + # trtllm-serve hands them to the attached frontends via + # multi_frontend_attach_info(). + self._num_frontends = (_llm_args.num_serve_frontends + if _llm_args is not None else 1) + self._multi_frontend_ipc_dir: Optional[str] = None + self._multi_frontend_hmac: Optional[bytes] = None + if self._num_frontends > 1: + if self._enable_resource_governor: + raise ValueError( + "Multi-frontend serving does not support " + "enable_resource_governor: the resource-governor signal " + "only reaches the launcher frontend.") + self._multi_frontend_ipc_dir = tempfile.mkdtemp( + prefix="trtllm_frontends_") + self._multi_frontend_hmac = os.urandom(32) + # Generate RPC address and key for stats RPC self.rpc_addr = get_unique_ipc_addr() self.hmac_key = os.urandom(32) @@ -395,34 +417,91 @@ 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: + # The rank0 worker BINDS the request ingress (PULL) so every + # frontend can PUSH-connect; each frontend (incl. this launcher, + # frontend 0) binds its own result lane (PULL). + ipc_dir = self._multi_frontend_ipc_dir + hmac_key = self._multi_frontend_hmac + 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 = self._multi_frontend_ipc_dir + hmac_key = self._multi_frontend_hmac + 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(), + # Attached frontends must apply the same collective_rpc guard + # as the launcher (see _check_collective_rpc_guard). + "model_world_size": + self.model_world_size, + # 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 @@ -630,8 +709,28 @@ def pre_shutdown(self): if not self.mpi_futures or any(not f.done() for f in self.mpi_futures): self.request_queue.put_noblock(None, retry=4) + def _get_next_client_id(self) -> int: + client_id = super()._get_next_client_id() + if self._num_frontends > 1: + # Lane 0 follows the same namespace rule as attached frontends + # so a long-lived counter can never bleed into the frontend-id + # bits (a no-op re-encode until the counter wraps). + client_id = namespace_client_id(0, client_id) + return client_id + + def _cleanup_multi_frontend_ipc_dir(self): + """Remove the launcher-owned multi-frontend ipc directory. + + Unlinking ipc socket paths does not disturb established zmq + connections; it only prevents new connects. + """ + if self._multi_frontend_ipc_dir is not None: + shutil.rmtree(self._multi_frontend_ipc_dir, ignore_errors=True) + self._multi_frontend_ipc_dir = None + def shutdown(self): if not self.workers_started: + self._cleanup_multi_frontend_ipc_dir() return if not self.doing_shutdown: @@ -671,6 +770,7 @@ def shutdown(self): self.result_queue.close() if self._resource_governor_queue is not None: self._resource_governor_queue.close() + self._cleanup_multi_frontend_ipc_dir() self.workers_started = False if self._owns_mpi_session: @@ -911,3 +1011,126 @@ 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 (num_serve_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 + utils.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: + num_lanes = len(attach_info["result_addrs"]) + if not 0 < frontend_id < num_lanes: + raise ValueError( + f"frontend_id {frontend_id} out of range: attached frontends " + f"use ids 1..{num_lanes - 1} (id 0 is the launcher)") + 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) + + # State consumed by methods inherited from GenerationExecutorProxy + # (submit / check_health / collective_rpc). The engine lives with + # the launcher: there are no local MPI workers, so the monitor + # stays empty and worker death reaches this frontend through its + # result lane / error queue instead. + self._engine_dead = False + self.model_world_size = attach_info.get("model_world_size", 1) + self._worker_process_monitor = WorkerProcessMonitor() + + self._frontend_id = frontend_id + self._num_frontends = num_lanes + self._results: Dict[int, GenerationResult] = {} + self.garbage_collection_gen0_threshold = None + self.workers_started = False + self.dispatch_result_thread: Optional[ManagedThread] = None + # Must be None: OpenAIServer reads the resource_governor_queue + # property at init; the governor lives with the launcher only. + 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, multi-client). + 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 check_health(self) -> bool: + """Health contract of an attached frontend. + + An attached frontend owns no workers, so there is no process or + MPI-future liveness to poll: it is healthy while no fatal error + has been recorded and shutdown has not begun. Engine death + reaches it through the per-lane result socket / dispatch-thread + error path, which records the fatal error checked here. + """ + if self.doing_shutdown or self._fatal_error is not None: + return False + + if self._drain_error_queue(): + return self._fatal_error is None and not self.doing_shutdown + + return True + + 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 + self._worker_process_monitor.close() + # 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 3bd0106241f8..e4a9f3333f4b 100644 --- a/tensorrt_llm/executor/utils.py +++ b/tensorrt_llm/executor/utils.py @@ -229,6 +229,66 @@ 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: one result lane per frontend. When set, the + # rank0 worker BINDS the request queue (PULL) and routes responses to + # these lanes; result_queue_addr then aliases lane 0 (the launcher). + frontend_result_queue_addrs: Optional[list[tuple[str, + Optional[bytes]]]] = None + + +# Multi-frontend client_id namespacing: a FRONTEND_ID_BITS-wide frontend id +# sits just below the sign bit -- bit 63 stays clear so ids remain positive +# in signed-int64 contexts -- and the low bits carry the per-frontend +# request counter. A stray un-namespaced (small) id reads as frontend 0, +# the launcher. +FRONTEND_ID_BITS = 6 +# Keep llm_args.num_serve_frontends le= in sync (it cannot import this +# module; test_multi_frontend_routing pins the two together). +MAX_NUM_FRONTENDS = 1 << FRONTEND_ID_BITS +FRONTEND_ID_SHIFT = 63 - FRONTEND_ID_BITS +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 frontend_lane_index(client_id: Optional[int], num_lanes: int) -> int: + """The result-lane index for a response's originating frontend. + + None (e.g. ADP dummy requests) and out-of-range frontend ids go to + lane 0, the launcher, which silently discards them like today. + """ + frontend_id = get_frontend_id(client_id) + return frontend_id if frontend_id < num_lanes else 0 + + +def bucket_responses_by_frontend(responses: list, + num_frontends: int) -> list[list]: + """Bucket responses by frontend_lane_index of their client_id.""" + buckets = [[] for _ in range(num_frontends)] + for rsp in responses: + buckets[frontend_lane_index(rsp.client_id, num_frontends)].append(rsp) + return buckets + + +def multi_frontend_request_addr(ipc_dir: str) -> str: + """The request ingress bound by the rank0 worker; frontends PUSH-connect.""" + 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 bound by a frontend; worker/postproc PUSH-connect.""" + 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 5958d296e1c4..233e394e7c77 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -221,6 +221,10 @@ def _print_stacks(): postproc_worker_config = postproc_worker_config or PostprocWorkerConfig() is_leader: bool = mpi_rank() == 0 + # Multi-frontend serving: the worker binds the request ingress (PULL) + # and pushes responses to per-frontend result lanes. + 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() @@ -238,7 +242,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, @@ -260,6 +266,16 @@ 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: + # One PUSH lane per frontend (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 @@ -269,9 +285,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: @@ -281,18 +300,20 @@ 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 + # Each postproc worker pushes to every frontend result lane (a + # single lane in single-frontend mode). + proxy_result_addrs = (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) for i in range(postproc_worker_config.num_postprocess_workers): fut = postproc_worker_pool.submit( postproc_worker_main, result_queues[i].address, - proxy_result_queue, + proxy_result_addrs, postproc_worker_config.postprocess_tokenizer_dir, PostprocWorker.default_record_creator, postproc_worker_config.post_processor_hook, @@ -349,6 +370,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 5d19670b0991..b02c5eab23b0 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -313,7 +313,11 @@ def __init__(self, load_post_processor_hook(_post_processor_path) if _post_processor_path else None) - if self.args.parallel_config.is_multi_gpu: + # Attached serving frontends connect to an already-running worker: + # they must not spawn an MPI session (see 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/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 57c3de25dcfe..f0baa32afaf1 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4121,6 +4121,19 @@ class BaseLlmArgs(StrictBaseModel): description="The path to the tokenizer directory for postprocessing.", status="prototype") + num_serve_frontends: int = Field( + default=1, + ge=1, + # = executor.utils.MAX_NUM_FRONTENDS (cannot be imported here); + # test_multi_frontend_routing pins the two together. + le=64, + description= + "The number of HTTP frontend processes serving one executor. Used by " + "trtllm-serve: values > 1 run additional attached frontend processes " + "that share the serving port via SO_REUSEPORT (classic IPC executor " + "path only).", + status="prototype") + reasoning_parser: Optional[str] = Field( default=None, description="The parser to separate reasoning content from output.", diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 53ab14af4234..1d7c471f550c 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2180,6 +2180,17 @@ async def create_streaming_generator(promise: RequestOutput, "Request.background is not supported yet, will fallback to foreground processing." ) + # Reject rather than silently ignore: with storage disabled + # (TRTLLM_RESPONSES_API_DISABLE_STORE, postproc workers, or + # multi-frontend serving) the previous response can never be + # resolved. + if request.previous_response_id is not None and not self.enable_store: + return self.create_error_response( + err_type="InvalidRequestError", + message=("'previous_response_id' requires response " + "storage, which is disabled on this server."), + ) + # Get prev response prev_response = None if self.enable_store: diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index ca6e762490cf..4d0261bf0131 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1016,6 +1016,13 @@ "kind": "value", "path": "num_postprocess_workers" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "num_serve_frontends" + }, { "allowed_values": [ "cutlass", diff --git a/tests/integration/test_lists/test-db/l0_cpu_x86.yml b/tests/integration/test_lists/test-db/l0_cpu_x86.yml index 9a39993347b8..a23f8b9a8f46 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_x86.yml @@ -14,3 +14,4 @@ l0_cpu_x86: orchestrator: mpi tests: - unittest/executor/test_rpc.py + - unittest/executor/test_multi_frontend_routing.py diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 7ff208c41ce3..959f249183e8 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -60,6 +60,10 @@ methods: annotation: Optional[str] default: null status: prototype + num_serve_frontends: + annotation: int + default: 1 + status: prototype custom_tokenizer: annotation: Optional[str] default: null diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 5d681cdeb9d8..c3430712d780 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -303,6 +303,15 @@ commands: is_flag: false flags: - "--num_postprocess_workers" + num_serve_frontends: + type: int + default: 1 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--num_serve_frontends" num_input_processor_workers: type: int default: 8 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..d34aa62d8dde --- /dev/null +++ b/tests/unittest/executor/test_multi_frontend_routing.py @@ -0,0 +1,302 @@ +# 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 (num_serve_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 + +import pytest + +from tensorrt_llm.executor.utils import ( + FRONTEND_COUNTER_MASK, + MAX_NUM_FRONTENDS, + bucket_responses_by_frontend, + frontend_lane_index, + get_frontend_id, + namespace_client_id, +) + +# The CI CPU stages collect with -m "cpu_only and not disabled" and skip +# files that don't mention pytest.mark.cpu_only (see unittest/conftest.py). +pytestmark = pytest.mark.cpu_only + + +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, MAX_NUM_FRONTENDS - 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 + # Bit 63 stays clear: ids remain positive as signed int64. + assert 0 < client_id < (1 << 63) + + def test_knob_cap_matches_encoding(self): + from tensorrt_llm.llmapi.llm_args import BaseLlmArgs + + meta = BaseLlmArgs.model_fields["num_serve_frontends"].metadata + assert any(getattr(m, "le", None) == MAX_NUM_FRONTENDS for m in meta) + + def test_counter_wraparound_stays_in_namespace(self): + # A counter overflowing its field 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 test_lane_index_clamps_to_launcher(self): + assert frontend_lane_index(namespace_client_id(2, 7), 3) == 2 + assert frontend_lane_index(namespace_client_id(7, 1), 2) == 0 + assert frontend_lane_index(None, 4) == 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 + + @staticmethod + def _stop_dispatch_thread(proxy, result_addrs, hmac_key): + """End the dispatch thread the way the worker does at engine teardown. + + The worker fans a per-lane None sentinel to every frontend (see + notify_proxy_threads_to_quit). Frontend shutdown deliberately leaves + the thread to process teardown, which the threadleak checker would + report as a leak. + """ + import zmq + + from tensorrt_llm.executor.ipc import FusedIpcQueue + + worker_lane = FusedIpcQueue( + (result_addrs[proxy._frontend_id], hmac_key), + is_server=False, + fuse_message=False, + socket_type=zmq.PUSH, + name="fake_worker_shutdown_lane", + ) + worker_lane.put(None) + proxy.dispatch_result_thread.join(timeout=10) + assert not proxy.dispatch_result_thread.is_alive() + + 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, hmac_key, result_addrs = 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" + ) + self._stop_dispatch_thread(proxy, result_addrs, hmac_key) + + def test_check_health_and_submit_reflect_fatal_error(self): + from tensorrt_llm.executor.request import GenerationRequest + from tensorrt_llm.executor.utils import EngineDeadError + from tensorrt_llm.sampling_params import SamplingParams + + with tempfile.TemporaryDirectory() as tmpdir: + proxy, _, _, _ = self._make_proxy_and_fake_worker(tmpdir) + + # Regression: this state lives in GenerationExecutorProxy.__init__, + # which an attached frontend deliberately skips, so it must be + # initialized explicitly (a missing _engine_dead broke the first + # submit(); a missing _worker_process_monitor broke /health). + assert proxy.check_health() + assert proxy.model_world_size == 1 + + proxy._set_fatal_error(RuntimeError("rank0 worker died")) + assert not proxy.check_health() + with pytest.raises(EngineDeadError): + proxy.submit(GenerationRequest([1], SamplingParams())) + + 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 + + self._stop_dispatch_thread(proxy, result_addrs, hmac_key) diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py index 85bd42b50f63..dce923c6211c 100644 --- a/tests/unittest/llmapi/test_executor.py +++ b/tests/unittest/llmapi/test_executor.py @@ -389,7 +389,7 @@ def ResponsePostprocessWorker_worker_task(pull_pipe_addr, push_pipe_addr, tokenizer_dir): worker = PostprocWorker( pull_pipe_addr=pull_pipe_addr, - push_pipe_addr=push_pipe_addr, + push_pipe_addrs=[push_pipe_addr], tokenizer_dir=tokenizer_dir, record_creator=ResponsePostprocessWorker_record_creator) worker.start()