mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-02-07 11:41:47 +08:00
75057cd036
* Update TensorRT-LLM --------- Co-authored-by: Puneesh Khanna <puneesh.khanna@tii.ae> Co-authored-by: Ethan Zhang <26497102+ethnzhng@users.noreply.github.com>
17 lines
474 B
Bash
Executable File
17 lines
474 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.llmapi.mgmn_worker_node
|
|
fi
|