#!/bin/bash -e

# Make sure GPUs are up
if [ $SLURM_LOCALID -eq 0 ] ; then
    rocm-smi
fi
sleep 2

export MIOPEN_USER_DB_PATH="/tmp/samantao-miopen-cache-$SLURM_NODEID"
export MIOPEN_CUSTOM_CACHE_DIR=$MIOPEN_USER_DB_PATH

# Set MIOpen cache to a temporary folder.
if [ $SLURM_LOCALID -eq 0 ] ; then
    rm -rf $MIOPEN_USER_DB_PATH
    mkdir -p $MIOPEN_USER_DB_PATH
fi
sleep 2
  
# Report affinity
echo "Rank $SLURM_PROCID --> $(taskset -p $$)"

# Set interfaces to be used by RCCL.
export NCCL_SOCKET_IFNAME=hsn0,hsn1,hsn2,hsn3

# Set environment for the app
export MASTER_PORT=29500
export WORLD_SIZE=$SLURM_NPROCS
export RANK=$SLURM_PROCID
export HIP_VISIBLE_DEVICES=$SLURM_LOCALID

# Run app
cd /workdir/mnist

source /opt/venv-extra/bin/activate 

if [ $RANK -eq 2 ] ; then

    ROCPROF=rocprofiler-sdk     $ROCM_PATH/bin/rocprof-compute profile -n myprof --roof-only --kernel-names --basenames -- python -u mnist_DDP.py --gpu --modelpath /workdir/mnist/model

    echo "Completed profiling!"
    exit -100 # Signal fail for slurm to end all processes.

else
    for i in {0..100} ; do
        python -u mnist_DDP.py --gpu --modelpath /workdir/mnist/model
    done
fi

