[TRTLLM-6295][test] Exit as early as possible and propagate exit status correctly for multi-node testing (#7739)

Signed-off-by: Yanchao Lu <[email protected]>

Author: chzblych

jenkins/scripts/slurm_run.sh

@@ -45,8 +45,7 @@ export LLM_ROOT=$llmSrcNode
 export LLM_MODELS_ROOT=$MODEL_CACHE_DIR
 export UCX_TLS=^gdr_copy
 
-# TODO: Move back to tensorrt_llm/llmapi/trtllm-llmapi-launch later
-llmapiLaunchScript="$llmSrcNode/jenkins/scripts/trtllm-llmapi-launch"
+llmapiLaunchScript="$llmSrcNode/tensorrt_llm/llmapi/trtllm-llmapi-launch"
 chmod +x $llmapiLaunchScript
 cd $llmSrcNode/tests/integration/defs
 testCmdLines=(

jenkins/scripts/trtllm-llmapi-launch

@@ -1,5 +1,5 @@
 #!/bin/bash
-set -e
+set -Eeo pipefail
 
 task_with_command=("$@")
 native_mpi_rank=$OMPI_COMM_WORLD_RANK
@@ -47,7 +47,7 @@ log_stderr "tllm_mpi_size: $tllm_mpi_size"
 export_free_tcp_addr_for_spawn_proxy_process
 
 if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
-    log_stderr "rank${mpi_rank} run ${task_with_command[@]} in background"
+    log_stderr "Rank${mpi_rank} run ${task_with_command[@]} in background"
 
     # MPI doesn't allow spawn a process sharing the MPI environment in a MPI
     # process, or duplicate MPI_Init in the child process will cause undefined
@@ -70,16 +70,60 @@ if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
             done
         done
 
+        # Turn off "exit on error" so the following lines always run
+        set +e
+
         # Execute the task with cleaned environment
-       "${task_with_command[@]}"
-        # stop the MPI Comm server
+        "${task_with_command[@]}"
+        task_exit_code=$?
+        log_stderr "Rank${mpi_rank} Task exit code: $task_exit_code"
+
+        # Stop the MPI Comm server
         python3 -m tensorrt_llm.llmapi.mgmn_leader_node --action stop
+        mpi_exit_code=$?
+        log_stderr "Rank${mpi_rank} MPI Comm server exit code: $mpi_exit_code"
+
+        # Propagate task exit status
+        if [ $task_exit_code -ne 0 ]; then
+            exit $task_exit_code
+        else
+            exit $mpi_exit_code
+        fi
     ) &
 
-    log_stderr "rank${mpi_rank} run mgmn leader node with mpi_world_size: $(mpi_world_size) ..."
-    log_stderr "rank0 host: $HOSTNAME"
+    # Turn off "exit on error" so the following lines always run
+    set +e
+
+    # Capture subshell PID
+    subshell_pid=$!
+    log_stderr "Rank${mpi_rank} Subshell PID: $subshell_pid"
+
+    log_stderr "Rank${mpi_rank} run mgmn leader node with mpi_world_size: $(mpi_world_size) ..."
+    log_stderr "Rank0 host: $HOSTNAME"
     python3 -m tensorrt_llm.llmapi.mgmn_leader_node
+    mgmn_leader_node_exit_code=$?
+    log_stderr "Rank${mpi_rank} MGMN leader node exit code: $mgmn_leader_node_exit_code"
+
+    # Wait for subshell
+    wait $subshell_pid
+    # This is subshell's exit code
+    subshell_exit_code=$?
+    log_stderr "Rank${mpi_rank} Subshell exit code: $subshell_exit_code"
+
+    # Propagate subshell exit status
+    if [ $subshell_exit_code -ne 0 ]; then
+        exit $subshell_exit_code
+    else
+        exit $mgmn_leader_node_exit_code
+    fi
 else
-    log_stderr "rank${mpi_rank} run mgmn worker node with mpi_world_size: $(mpi_world_size) ..."
+    # Turn off "exit on error" so the following lines always run
+    set +e
+
+    log_stderr "Rank${mpi_rank} run mgmn worker node with mpi_world_size: $(mpi_world_size) ..."
     python3 -m tensorrt_llm.llmapi.mgmn_worker_node
+    mgmn_worker_node_exit_code=$?
+    log_stderr "Rank${mpi_rank} MGMN worker node exit code: $mgmn_worker_node_exit_code"
+
+    exit $mgmn_worker_node_exit_code
 fi