#!/bin/bash -e

wd=$(pwd)
jobid=$(squeue --me | head -2 | tail -n1 | awk '{print $1}')

#
# Example assume allocation was created, e.g.:
# N=1 ; salloc -p standard-g  --threads-per-core 1 --exclusive -N $N --gpus $((N*8)) -t 4:00:00 --mem 0
#

set -x

SIF=/pfs/lustref1/appl/local/laifs/containers/lumi-multitorch-u24r70f21m50t210-20260415_130625/lumi-multitorch-full-u24r70f21m50t210-20260415_130625.sif 

# Create configuration file.

rm -rf $wd/rocprof-sys-config.cfg

srun -N 1 -n 1 --gpus 8 --jobid=$jobid \
  singularity exec \
    --overlay /pfs/lustrep3/projappl/project_465002770/Exercises/AMD/Pytorch/deps.sqsh:ro \
    -B $wd:/workdir \
    $SIF \
    bash -c -eux '$ROCM_PATH/bin/rocprof-sys-avail -a -G /workdir/rocprof-sys-config.cfg'

sed -i 's/ROCPROFSYS_USE_SAMPLING.*/ROCPROFSYS_USE_SAMPLING = true/g' $wd/rocprof-sys-config.cfg
sed -i 's/ROCPROFSYS_USE_AMD_SMI.*/ROCPROFSYS_USE_AMD_SMI = false/g' $wd/rocprof-sys-config.cfg
sed -i 's/ROCPROFSYS_SAMPLING_CPUS.*/ROCPROFSYS_SAMPLING_CPUS = none/g' $wd/rocprof-sys-config.cfg
sed -i 's/ROCPROFSYS_SAMPLING_GPUS.*/ROCPROFSYS_SAMPLING_GPUS = 2/g' $wd/rocprof-sys-config.cfg

#
# Run application.
#

rm -rf $wd/run-me.sh 
cat > $wd/run-me.sh << EOF
#!/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/$(whoami)-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

pcmd='python -u'

if [ \$RANK -eq 2 ] ; then

  export ROCPROFSYS_CONFIG_FILE=/workdir/rocprof-sys-config.cfg

  # Set the environment for the newer rocprofiler-systems with Python support.
  source /opt/rocm-7.2.3-extra/share/rocprofiler-systems/setup-env.sh 

  pcmd='rocprof-sys-python -- ' 

  export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/opt/venv/lib/python3.12/site-packages/torch/lib

fi

\$pcmd mnist_DDP.py --gpu --modelpath /workdir/mnist/model

EOF
chmod +x $wd/run-me.sh

c=fe
MYMASKS="0x${c}000000000000,0x${c}00000000000000,0x${c}0000,0x${c}000000,0x${c},0x${c}00,0x${c}00000000,0x${c}0000000000"

Nodes=2

export MASTER_ADDR=$(scontrol show hostname "$SLURM_NODELIST" | head -n1)

srun --jobid=$jobid -N $((Nodes)) -n $((Nodes*8)) --gpus $((Nodes*8)) --cpu-bind=mask_cpu:$MYMASKS \
  singularity exec \
    --overlay /pfs/lustrep3/projappl/project_465002770/Exercises/AMD/Pytorch/deps.sqsh:ro \
    -B $wd:/workdir \
    $SIF /workdir/run-me.sh
