mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
73 lines
2.4 KiB
Bash
Executable File
73 lines
2.4 KiB
Bash
Executable File
#!/bin/bash
|
|
set -e
|
|
|
|
task_with_command="${@:1}"
|
|
native_mpi_rank=$OMPI_COMM_WORLD_RANK
|
|
mpi_rank=${SLURM_PROCID:-${OMPI_COMM_WORLD_RANK:-${PMI_RANK:-${PMI_ID:-0}}}}
|
|
|
|
log_stderr() { echo -e "\033[33m$@\033[0m" >&2; }
|
|
log_stderr "mpi_rank: $mpi_rank"
|
|
|
|
pid=$(ps -o pid= -p $$ | tr -d ' ')
|
|
|
|
# Tell TRTLLM to spawn a additional process for the Proxy
|
|
export TLLM_SPAWN_PROXY_PROCESS=1
|
|
|
|
function mpi_world_size {
|
|
if [ -n "$SLURM_NTASKS" ]; then
|
|
echo "$SLURM_NTASKS"
|
|
elif [ -n "$OMPI_COMM_WORLD_SIZE" ]; then
|
|
echo "$OMPI_COMM_WORLD_SIZE"
|
|
else
|
|
echo "1"
|
|
fi
|
|
}
|
|
|
|
function export_free_tcp_addr {
|
|
# find free port with python -c, the port should start from 10000
|
|
local free_port=$(python -c 'import socket; s=socket.socket(); s.bind(("", 10000)); print(s.getsockname()[1]); s.close()')
|
|
export TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR="tcp://127.0.0.1:${free_port}"
|
|
log_stderr "TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR: $TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR"
|
|
}
|
|
|
|
export tllm_mpi_size=$(mpi_world_size)
|
|
log_stderr "tllm_mpi_size: $tllm_mpi_size"
|
|
|
|
export_free_tcp_addr
|
|
|
|
if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
|
|
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
|
|
# behavior. Thus we need to clean the MPI environment in the parent process
|
|
# before spawning the child process, and restore the MPI environment later
|
|
# before running MPI operations in the parent process.
|
|
mpi_blacklist=(
|
|
OMPI_ PMIX_ PMI_ SLURM_ MPI_ UCX_
|
|
I_MPI_ HYDRA_ KMP_ MPICH_ MV2_ CRAY_
|
|
)
|
|
|
|
(
|
|
# Remove MPI-related variables only in the subshell context
|
|
for var in $(compgen -e); do
|
|
for prefix in "${mpi_blacklist[@]}"; do
|
|
if [[ "$var" == "$prefix"* ]]; then
|
|
unset "$var"
|
|
break
|
|
fi
|
|
done
|
|
done
|
|
|
|
# Execute the task with cleaned environment
|
|
$task_with_command
|
|
) &
|
|
|
|
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
|
|
else
|
|
log_stderr "rank${mpi_rank} run mgmn worker node with mpi_world_size: $(mpi_world_size) ..."
|
|
python3 -m tensorrt_llm.llmapi.mgmn_worker_node
|
|
fi
|