#!/bin/bash
#!
#! Dask job script for CSD3
#! Tim Cornwell
#!

#!#############################################################
#!#### Modify the options in this section as appropriate ######
#!#############################################################

#! sbatch directives begin here ###############################
#! Name of the job:
#SBATCH -J NNUMBER_NODES_WNUMBER_TASKS
#! Which project should be charged:
#SBATCH -A SKA-SDP-SL2-CPU
#! How many whole nodes should be allocated?
#SBATCH --nodes=NUMBER_NODES
#! How many (MPI) tasks will there be in total? (<= nodes*32)
#SBATCH --ntasks=NUMBER_TASKS
#! Memory limit
#SBATCH --mem 384000
#! How much wallclock time will be required?
#SBATCH --time=EXECUTION_TIME
#! What types of email messages do you wish to receive?
#SBATCH --mail-type=ALL
#! Uncomment this to prevent the job from being requeued (e.g. if
#! interrupted by node failure or system downtime):
##SBATCH --no-requeue
#! Do not change:
#SBATCH -p skylake-himem
#! What types of email messages do you wish to receive?
#SBATCH --mail-type=FAIL
#! Ask for exclusive access
#SBATCH --exclusive
#! Same switch
#SBATCH --switches=1
#! Uncomment this to prevent the job from being requeued (e.g. if
#! interrupted by node failure or system downtime):
##SBATCH --no-requeue
#! Optionally modify the environment seen by the application
#! (note that SLURM reproduces the environment at submission irrespective of ~/.bashrc):
#module purge                               # Removes all modules still loaded

#! Set up python
. $HOME/arlenv/bin/activate
export PYTHONPATH=$PYTHONPATH:$ARL
echo "PYTHONPATH is ${PYTHONPATH}"
echo -e "Running python: `which python`"
echo -e "Running dask-scheduler: `which dask-scheduler`"

module load slurm

JOBID=${SLURM_JOB_ID}
echo ${SLURM_JOB_NODELIST}

cd $SLURM_SUBMIT_DIR
echo -e "Changed directory to `pwd`.\n"
# Archive the results
archive="output_NNUMBER_NODES_WNUMBER_TASKS_${JOBID}_INCOMPLETE"
mkdir ${archive}
cd ${archive}

#! Create a hostfile:
scontrol show hostnames $SLURM_JOB_NODELIST > hostfile.$JOBID

hostIndex=0
for host in `cat hostfile.$JOBID`; do
    ibhost=${host}
    echo "Working on ${ibhost} ...."
    if [ "$hostIndex" = "0" ]; then
        scheduler=${ibhost}
        echo "run dask-scheduler"
        echo "ssh ${ibhost} dask-scheduler --host ${scheduler} --port=8786 &"
        ssh ${host} dask-scheduler --host ${scheduler} --port=8786 &
        sleep 5
    fi
    echo "run dask-worker"
    ssh $host dask-worker \
    --host ${ibhost} \
    --nprocs NUMBER_PROCS_PER_NODE \
    --nthreads NUMBER_THREADS  \
    --memory-limit 0.75 \
    --local-directory /tmp \
    ${scheduler}:8786 &
    sleep 1
    hostIndex="1"
done
echo "Scheduler and workers now running"

#! We need to tell dask Client (inside python) where the scheduler is running
export ARL_DASK_SCHEDULER=${scheduler}:8786
echo "Scheduler is running at ${scheduler}"

root=${ARL}/workflows/scripts/pipelines/performance/pipelines-timings
CMD="python ${root}/pipelines-timings.py --nnodes NUMBER_NODES --nworkers NUMBER_TASKS --memory MEMORY --nfreqwin NUMBER_FREQUENCY_WINDOWS --nthreads NUMBER_THREADS --jobid $JOBID --use_serial_imaging False --rmax 3000 --context wstack > pipelines-timings.out"
echo "About to execute $CMD"
eval $CMD

cp ${root}/pipelines-timings.py .
cp ${0}  ./
mv ../slurm-${JOBID}.out .
mv ../${archive} ../output_NNUMBER_NODES_WNUMBER_TASKS_${JOBID}_COMPLETE
exit

