File size: 473 Bytes
5000658 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
#!/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
|