#!/bin/bash
#!
#! Dask job script for Darwin (Sandy Bridge, ConnectX3)
#! Tim Cornwell
#!

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

#! sbatch directives begin here ###############################
#! Name of the job:
#SBATCH -J SDP_ARL
#! Which project should be charged:
#SBATCH -A SKA-SDP
#! How many whole nodes should be allocated?
#SBATCH --nodes=4
#! How many (MPI) tasks will there be in total? (<= nodes*16)
#SBATCH --ntasks=32
#! Memory limit
#SBATCH --mem 63900
#! How much wallclock time will be required?
#SBATCH --time=00:20:00
#! What types of email messages do you wish to receive?
#SBATCH --mail-type=FAIL
#! 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 sandybridge

#! Modify the settings below to specify the application's environment, location
#! and launch method:

#! Optionally modify the environment seen by the application
#! (note that SLURM reproduces the environment at submission irrespective of ~/.bashrc):
. /etc/profile.d/modules.sh                # Leave this line (enables the module command)
module purge                               # Removes all modules still loaded
module load default-impi                   # REQUIRED - loads the basic environment

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

#! Work directory (i.e. where the job will run):
workdir="$SLURM_SUBMIT_DIR"  # The value of SLURM_SUBMIT_DIR sets workdir to the directory
                             # in which sbatch is run.

#! Are you using OpenMP (NB this is unrelated to OpenMPI)? If so increase this
#! safe value to no more than 16:
export OMP_NUM_THREADS=1

cd $workdir
echo -e "Changed directory to `pwd`.\n"

JOBID=$SLURM_JOB_ID

if [ "$SLURM_JOB_NODELIST" ]; then
        #! Create a hostfile:
        export NODEFILE=`generate_pbs_nodefile`
        cat $NODEFILE | uniq > hostfile.$JOBID
        echo -e "\nNodes allocated:\n================"
        echo `cat hostfile.$JOBID | sed -e 's/\..*$//g'`
fi

echo -e "JobID: $JOBID\n======"
echo "Time: `date`"
echo "Master node: `hostname`"
echo "Current directory: `pwd`"

scheduler="`hostname`:8786"
echo "About to dask-ssh on:"
cat hostfile.$JOBID

#! dask-ssh related options:
#!  --nthreads INTEGER        Number of threads per worker process. Defaults to
#!                            number of cores divided by the number of processes
#!                            per host.
#!  --nprocs INTEGER          Number of worker processes per host.  Defaults to
#!                            one.
#!  --hostfile PATH           Textfile with hostnames/IP addresses
#!
dask-ssh --nprocs 8 --nthreads 1 --scheduler-port 8786 --log-directory `pwd` --hostfile hostfile.$JOBID &
sleep 60

#! We need to tell dask Client (inside python) where the scheduler is running
scheduler="`hostname`:8786"
echo "Scheduler is running at ${scheduler}"
CMD="python ./losing_workers-loop.py ${scheduler}"
echo "About to execute $CMD"
eval $CMD

# Archive the results
archive="output_${JOBID}"
echo "Moving results to ${archive}"
mkdir ${archive}
mv "slurm-${JOBID}".out ${archive}
mv hostfile.${JOBID} ${archive}
mv dask-ssh* ${archive}
cp *.py ${archive}
cp ${0}  ${archive}