wrapper.sh

#!/bin/bash

## Set UCX variables
export OMPI_MCA_pml=ucx
export OMPI_MCA_coll=^hcoll
export UCX_PROTO_ENABLE=n
export UCX_MAX_RNDV_RAILS=1
export UCX_TLS=^gdr_copy
export UCX_CUDA_COPY_DMABUF=y

NICS=(mlx5_0:1 mlx5_3:1 mlx5_4:1 mlx5_5:1 mlx5_6:1 mlx5_9:1 mlx5_10:1 mlx5_11:1)

## Get global rank from SLURM or OpenMPI
export CUDA_VISIBLE_DEVICES=
export MINICOMBUST_RANK_ID="${OMPI_COMM_WORLD_RANK:-$SLURM_PROCID}"

export MINICOMBUST_FRANKS=$(( MINICOMBUST_NODES * MINICOMBUST_GPUS ))

## Set CUDA_VISIBLE_DEVICES for flow ranks, get first flow rank to confirm UCX variables.
if [ $MINICOMBUST_RANK_ID -lt $(( MINICOMBUST_FRANKS )) ]; then
  export CUDA_VISIBLE_DEVICES=$(( MINICOMBUST_RANK_ID % MINICOMBUST_GPUS ))
  echo "RANK $MINICOMBUST_RANK_ID VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"

  if [ $MINICOMBUST_RANK_ID -eq 0 ]; then
    echo "OMPI_MCA_pml $OMPI_MCA_pml UCX_TLS $UCX_TLS UCX_MEMTYPE_CACHE $UCX_MEMTYPE_CACHE UCX_MAX_RNDV_RAILS $UCX_MAX_RNDV_RAILS UCX_IB_PREFER_NEAREST_DEVICE $UCX_IB_PREFER_NEAREST_DEVICE OMPI_MCA_coll $OMPI_MCA_coll UCX_CUDA_COPY_DMABUF $UCX_CUDA_COPY_DMABUF "
  fi
fi

export MINICOMBUST_OUTPUT_ITER=-1
export MINICOMBUST_DIR=$PWD

## Only set LOCAL_PROF_CMD for set ranks.
export LOCAL_PROF_CMD=
if [[ ",$PROF_RANKS," = *",$MINICOMBUST_RANK_ID,"* ]]
then
  export LOCAL_PROF_CMD=$GLOBAL_PROF_CMD
fi


if [ $MINICOMBUST_RANK_ID -eq 0 ];
then
  echo "MINICOMBUST_RANK_ID=$MINICOMBUST_RANK_ID WRAPPER PROF_CMD: $GLOBAL_PROF_CMD"
  echo "MINICOMBUST_RANK_ID=$MINICOMBUST_RANK_ID WRAPPER CMD:      $@"
fi

## Run application
NSYS_CONFIG_DIRECTIVES='AgentLaunchTimeoutSec=120' $LOCAL_PROF_CMD "$@"