[TRTLLM-14510][feat] serve: multi-process HTTP frontends on the classic IPC executor path

tensorrt_llm/commands/serve.py

@@ -11,7 +11,7 @@
 import sys
 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 +26,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 (LlmLauncherEnvs, MAX_NUM_FRONTENDS)
 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
 from tensorrt_llm.llmapi import KvCacheConfig
 from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig,
@@ -37,7 +37,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 +197,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 +271,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 +361,87 @@ 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:
+    """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).
+    """
+    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 = []
+    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"
+        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,
@@ -372,17 +456,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)
+
     addr_info = socket.getaddrinfo(host, port, socket.AF_UNSPEC,
                                    socket.SOCK_STREAM)
     address_family = socket.AF_INET6 if all(
         [info[0] == socket.AF_INET6 for info in addr_info]) else socket.AF_INET
     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,6 +501,11 @@ def launch_server(
                 f"{backend} is not a known backend, check help for available options.",
                 param_hint="backend")
 
+        frontend_children = []
+        if multi_frontend.is_launcher:
+            frontend_children = _spawn_attached_frontends(
+                llm, multi_frontend.num_frontends)
+
         server = OpenAIServer(
             generator=llm,
             model=model,
@@ -427,7 +524,11 @@ def launch_server(
         if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1":
             gc.disable()
 
-        uvloop.run(server(host, port, sockets=[s]))
+        try:
+            uvloop.run(server(host, port, sockets=[s]))
+        finally:
+            if frontend_children:
+                _terminate_attached_frontends(frontend_children)
 
 
 def launch_grpc_server(host: str,
@@ -886,6 +987,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 +1175,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 +1280,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 +1908,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,

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: