mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-08 20:21:48 +08:00
17 lines
473 B
Bash
Executable File
17 lines
473 B
Bash
Executable File
#!/bin/bash
|
|
set -e
|
|
|
|
task_with_command="${@:1}"
|
|
native_mpi_rank=$OMPI_COMM_WORLD_RANK
|
|
# Works with Slurm launching with `--mpi=pmix`
|
|
mpi_rank=${PMIX_RANK:-$native_mpi_rank}
|
|
echo "mpi_rank: $mpi_rank" >> /dev/stderr
|
|
|
|
if [ -z "$mpi_rank" ] || [ "$mpi_rank" -eq 0 ]; then
|
|
echo "${mpi_rank} run ${task_with_command} ..." >> /dev/stderr
|
|
$task_with_command
|
|
else
|
|
echo "${mpi_rank} launch worker ..." >> /dev/stderr
|
|
python3 -m tensorrt_llm.hlapi.mgmn_worker_node
|
|
fi
|