[TRTLLM-16720][fix] Bound remote MPI failures and launcher shutdown

docs/source/llm-api/index.md

@@ -130,6 +130,84 @@ The following tips typically assist new LLM API users who are familiar with othe
 
   This limitation is applicable for multi-GPU inference only.
 
+### Launcher process ownership
+
+The rank-0 `trtllm-llmapi-launch` shell starts a small Python guard for the MPI
+communication server, task, and shutdown helper. Each guard owns a separate
+process group and monitors the shell, so killing only the shell with `SIGKILL`
+under plain `mpirun` still triggers cleanup of those groups. Startup registration
+finishes before a command can run. Guards also clean remaining group members
+when the command exits, then return its original exit status.
+
+Cleanup sends `SIGTERM`, followed by `SIGKILL` after
+`TLLM_LLMAPI_LAUNCH_TERM_GRACE_SECONDS` (default 5 seconds). This covers descendants
+that remain in their inherited process group. A command that deliberately
+creates a separate session or process group must manage those processes itself;
+killing a guard independently also requires the job's scheduler/container to
+reclaim its workload.
+
+### Remote MPI shutdown timeouts
+
+Configure these settings together when using `trtllm-llmapi-launch`:
+
+| Setting | Default | Scope |
+| --- | --- | --- |
+| `TLLM_MGMN_SHUTDOWN_GRACE_SECONDS` | 60 seconds | Server budget for draining a failed batch or stopping the MPI world. Final session and shared-executor teardown use the remaining budget. |
+| `TLLM_LLMAPI_LAUNCH_STOP_TIMEOUT` | 120 seconds | Launcher wait for the stop helper and server after the task exits. Also used separately for guard registration and for the task to exit naturally after a server failure. |
+| `TLLM_LLMAPI_LAUNCH_TERM_GRACE_SECONDS` | 5 seconds | Additional cleanup grace between TERM and KILL for owned process groups. |
+
+The server shutdown grace starts when it receives the stop request; an earlier
+worker-failure deadline can shorten the remaining grace. The launcher also waits
+for the stop helper's Python startup and `tensorrt_llm` imports before that request
+can be sent. Set the launcher stop timeout above the server grace by enough to
+cover measured cold startup/import time, connection/send and queue cleanup, and
+scheduling margin. The defaults leave 60 seconds of allowance; raising only the
+server grace does not raise the launcher timeout.
+
+For example, a 120-second server grace with a 180-second launcher timeout leaves
+60 seconds for this overhead. This is an illustrative allowance: increase it if
+cold or loaded nodes need more time.
+
+If task-first shutdown exhausts the launcher budget, the launcher terminates
+owned groups and returns 124 when the task succeeded; an existing nonzero task
+status takes precedence. After a server-first failure, the launcher instead
+preserves a naturally completed task's nonzero status, otherwise the server's
+failure status. TERM-to-KILL cleanup is a separate phase. Launcher waits count
+100 ms sleeps, so polling and scheduling overhead add time; these settings are
+not a strict end-to-end wall-clock limit.
+
+The launcher returns exit code 2 for invalid
+`TLLM_LLMAPI_LAUNCH_STOP_TIMEOUT` or
+`TLLM_LLMAPI_LAUNCH_TERM_GRACE_SECONDS` values; both require a positive integer
+of at most six digits. A process-guard registration timeout also returns 124.
+If registration of the stop helper fails after the task has exited, an existing
+nonzero task status still takes precedence.
+
+### Remote MPI batch ordering
+
+The remote MPI server dispatches the next batch only after every future in the
+previous batch has completed. The task wrapper retains its entry barrier to
+place one task on each rank, but has no trailing barrier: an individual rank can
+finish or report an error while its peers are still running. Cross-batch ordering
+is enforced by the server; directly invoking the wrapper does not provide a
+collective completion barrier.
+
+### Disaggregated MPI leader failures
+
+The disaggregated MPI leader uses the same remote MPI server. An asynchronous
+worker failure is fatal to that server: it reports the worker error, stops
+accepting new work, and keeps the error channel available until a stop request
+or the original failure deadline. Final teardown can abort the MPI communicator
+if its remaining grace expires. This does not restart the failed service.
+
+The leader launches its proxy through an autonomous process guard. Normal leader
+cleanup signals the guard; if the leader exits through `MPI_Abort` or is killed,
+the guard detects parent death and cleans the proxy's inherited process group.
+It allows 5 seconds between `SIGTERM` and `SIGKILL`. Descendants that create their
+own session or process group still require separate ownership, as described
+above. The Bash launcher's `TLLM_LLMAPI_LAUNCH_*` settings do not configure this
+Python launcher.
+
 ### FlashInfer JIT workspaces for MPI workers
 
 `trtllm-llmapi-launch` ranks and dynamically spawned `MpiPoolSession` workers

setup.py

@@ -740,7 +740,10 @@ def get_build_state_options():
             'trtllm-eval=tensorrt_llm.commands.eval:main'
         ],
     },
-    scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch'],
+    scripts=[
+        'tensorrt_llm/llmapi/trtllm-llmapi-launch',
+        'tensorrt_llm/llmapi/_llmapi_process_guard.py',
+    ],
     extras_require={
         "devel": devel_deps + grpc_smg_deps,
         "openengine": openengine_deps,

tensorrt_llm/commands/serve.py

@@ -81,7 +81,7 @@
     # code paths so plain LLM serving never pays its import cost.
     from tensorrt_llm.visual_gen.args import VisualGenArgs
 
-# Global variable to store the Popen object of the child process
+# Popen for the disaggregated proxy's process guard.
 _child_p_global: Optional[subprocess.Popen] = None
 
 
@@ -2648,9 +2648,14 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
     # 1. Run the LLM-API Proxy in a separate process for streaming performance.
     #      The Proxy will create a RemoteMpiSessionClient as mpi_session in LLM
     #      class.
+    guard_path = Path(
+        __file__).resolve().parents[1] / "llmapi" / "_llmapi_process_guard.py"
     command = [
-        "python3", sys.argv[0], "disaggregated_mpi_worker", "-c", config_file,
-        "--log_level", log_level
+        sys.executable, "-S",
+        str(guard_path), "--parent-pid",
+        str(os.getpid()), "--term-grace", "5", "--autonomous", "--role",
+        "proxy", "--", "python3", sys.argv[0], "disaggregated_mpi_worker", "-c",
+        config_file, "--log_level", log_level
     ]
     logger.info(
         f"rank {mpi_rank()} step1: preparing to launch command: {command}")
@@ -2663,7 +2668,11 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
     signal.signal(signal.SIGTERM, _command_telemetry.raise_signal_exit)
     signal.signal(signal.SIGINT, _command_telemetry.raise_signal_exit)
 
+    _child_p_global = None
+    failure_pending = False
     try:
+        # MPI_Abort and SIGKILL bypass finally. The guard remains the owner of
+        # the proxy group and reclaims it when this leader disappears.
         _child_p_global = subprocess.Popen(
             command,
             env=non_mpi_env,
@@ -2672,7 +2681,7 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
             start_new_session=True)
 
         logger.info(
-            f"Parent process (PID {os.getpid()}) launched child process (PID {_child_p_global.pid})."
+            f"Parent process (PID {os.getpid()}) launched proxy guard (PID {_child_p_global.pid})."
         )
 
         logger.info(f"rank {mpi_rank()} step2: start the mpi session server")
@@ -2682,33 +2691,33 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str,
         # This is a blocking call
         launch_remote_mpi_session_server(sub_comm)
 
+    except BaseException:
+        failure_pending = True
+        raise
     finally:
         # Restore original signal handlers
         signal.signal(signal.SIGTERM, original_sigterm_handler)
         signal.signal(signal.SIGINT, original_sigint_handler)
 
         if _child_p_global:  # If Popen was successful and object exists
             logger.info(
-                f"Parent process (PID {os.getpid()}) in finally block. Cleaning up child process (PID: {_child_p_global.pid})."
+                f"Parent process (PID {os.getpid()}) in finally block. Stopping proxy guard (PID: {_child_p_global.pid})."
             )
             # Check if child is still running
             if _child_p_global.poll() is None:
                 _child_p_global.terminate()
                 try:
                     _child_p_global.wait(timeout=30)
-                except subprocess.TimeoutExpired:
-                    logger.warning(
-                        f"Child process {_child_p_global.pid} timed out on terminate (30s), killing."
+                except subprocess.TimeoutExpired as error:
+                    message = (
+                        f"Proxy guard {_child_p_global.pid} cleanup is "
+                        "still pending after 30s; retaining its process-group ownership"
                     )
-                    _child_p_global.kill()
-                    try:
-                        _child_p_global.wait(timeout=30)
-                    except subprocess.TimeoutExpired:
-                        logger.error(
-                            f"Child process {_child_p_global.pid} failed to be killed even after 30s."
-                        )
-            assert _child_p_global.poll(
-            ) is not None, f"the subprocess should be terminated"
+                    logger.error(message)
+                    # Killing the guard could strand its payload. Leave it
+                    # responsible for cleanup and preserve an existing failure.
+                    if not failure_pending:
+                        raise RuntimeError(message) from error
 
     # Check if the process was launched and assert it's terminated
     if _child_p_global and hasattr(_child_p_global,