mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-31 16:21:07 +08:00
9bd15f1937
* TensorRT-LLM Release 0.10.0 --------- Co-authored-by: Loki <lokravi@amazon.com> Co-authored-by: meghagarwal <16129366+megha95@users.noreply.github.com>
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
|