#!/usr/bin/env bash
# Start with srun --cpu-bind=none ./affinity-wrapper ./application --param1 …

lrank=$SLURM_LOCALID

case ${lrank} in
  0)
       export CUDA_VISIBLE_DEVICES=0
       export UCX_NET_DEVICES=mlx5_1:1
       CPU_BIND=18-23
       ;;
   1)
       export CUDA_VISIBLE_DEVICES=1
       export UCX_NET_DEVICES=mlx5_0:1
       CPU_BIND=6-11
       ;;
   2)
       export CUDA_VISIBLE_DEVICES=2
       export UCX_NET_DEVICES=mlx5_3:1
       CPU_BIND=42-47
       ;;
   3)
       export CUDA_VISIBLE_DEVICES=3
       export UCX_NET_DEVICES=mlx5_2:1
       CPU_BIND=30-35
       ;;
esac

numactl --physcpubind=${CPU_BIND} $*