#!/bin/bash
#SBATCH --account lcls_g
#SBATCH -C gpu
#SBATCH -q regular
#SBATCH -t 1:00:00
#SBATCH --exclusive
#SBATCH --gpus-per-node 4
#SBATCH --mail-type=ALL

source perlmutter_env.sh
mamba activate $1

NODES=$SLURM_JOB_NUM_NODES

SCRIPT=$2
echo "CONDA ENV: $CONDA_PREFIX"
echo "RUNNING: $SCRIPT on $NODES nodes"
# echo "GIT COMMIT: $(git log --pretty=format:'%h' -n 1)"

RUN_DIR=$(dirname "$SCRIPT")
SCRIPT=$(basename "$SCRIPT")

pushd $RUN_DIR

IFS="." read -ra OUTPREFIX <<< "$SCRIPT"
OUTPUT_DIR=output/"$OUTPREFIX"_nodes"$NODES"
mkdir -p $OUTPUT_DIR
scontrol write batch_script $SLURM_JOB_ID $OUTPUT_DIR/job.sh

#cp $SCRIPT $OUTPUT_DIR
cp *.py $OUTPUT_DIR
pushd $OUTPUT_DIR

export FI_MR_CACHE_MONITOR=memhooks
export FI_CXI_RX_MATCH_MODE=software
export FI_CXI_DEFAULT_TX_SIZE=4096
export FI_CXI_DEFAULT_CQ_SIZE=131072
export FI_CXI_DISABLE_HOST_REGISTER=1
# export FI_CXI_REQ_BUF_MIN_POSTED=10
# export FI_MR_CACHE_MAX_COUNT=0
# export FI_LOG_LEVEL=debug
# export FI_LOG_PROV=cxi
export GASNET_OFI_DEVICE=cxi0
export GASNET_OFI_DEVICE_1=cxi1
export GASNET_OFI_DEVICE_2=cxi2
export GASNET_OFI_DEVICE_3=cxi3
export GASNET_OFI_DEVICE_TYPE=Node
export GASNET_OFI_RECEIVE_BUFF_SIZE=8M
export GASNET_OFI_MAX_MEDIUM=32768
export GASNET_BACKTRACE=1
export GASNET_FREEZE_ON_ERROR=0

export LEGION_DEFAULT_ARGS="-gex:obcount $(( SLURM_JOB_NUM_NODES * 12 ))"
export LEGATE_LIMIT_STDOUT=1
export LEGATE_AUTO_CONFIG=0
export LEGATE_IO_USE_VFD_GDS=0
export XLA_PYTHON_CLIENT_PREALLOCATE=false

export OMP_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export MKL_NUM_THREADS=1
export VECLIB_MAXIMUM_THREADS=1
export NUMEXPR_NUM_THREADS=1

legate --ranks-per-node 4 --launcher srun --gpu-bind "0/1/2/3" --cpu-bind 48-63,112-127/32-47,96-111/16-31,80-95/0-15,64-79 --mem-bind 3/2/1/0 --gpus 1 --cpus 1 --omps 0 --fbmem 30000 --sysmem 32768 --nsys $SCRIPT legateio |& tee output.txt

popd
popd

if [ -f slurm-$SLURM_JOB_ID.out ]; then
    cp slurm-$SLURM_JOB_ID.out $RUN_DIR/$OUTPUT_DIR
fi
