#!/bin/bash
# The line above this is the "shebang" line.  It must be first line in script
#-----------------------------------------------------
#	Example Submission Script: python version of Hello World using MPI
#	Runs an MPI-enabled python "Hello World" script
#-----------------------------------------------------
#
# Slurm sbatch parameters section:
#       These SBATCH lines should immediately follow the "shebang" line
#       They must occur before any non-empty/non-comment lines in the script
# 
#	Request 60 tasks with one CPU core each
#SBATCH -n  60
#SBATCH -c 1
#	Request 5 minutes of walltime
#SBATCH -t 5
#	Request 1 GB of memory per CPU core
#SBATCH --mem-per-cpu=1024
#	Allow other jobs to run on same node
#SBATCH --oversubscribe
#	Run on debug partition for rapid turnaround.  You will need
#	to change this (remove the line) if walltime > 15 minutes
#SBATCH --partition=debug

# This job will run the code in hello_mpi.py script from the submission dir
# We create a working directory on parallel file system and run the job
# from there

# Process our dot files (and in particular define the "module" command)
source /etc/profile

# Module load section
# First clear our module list 
module purge
# and reload the standard modules
module load hpcc/zaratan
# Load the desired compiler, MPI, and python modules
# NOTE: You need to use the same compiler and MPI module used
# when compiling the python (and its mpi4py package).  Module load-ing
# python first, followed by openmpi will ensure you get a compatible set
module load python
module load openmpi

# Section to make a scratch directory for this job
# Because different MPI tasks, which might be on different nodes, need
# access to it, we put it in scratch.  We include the SLURM jobid in the
# directory name to avoid interference if multiple jobs running at same time.
TMPWORKDIR="${HOME}/scratch/mpi4py-job.${SLURM_JOBID}"
mkdir $TMPWORKDIR
cd $TMPWORKDIR

# Section to output information identifying the job, etc.
echo "Slurm job submit-python-mpi.sh ${SLURM_JOBID} running on:"
hostname
echo "Ntasks=${SLURM_NTASKS}; CPU cores per task=${SLURM_CPUS_PER_TASK}"
echo "On ${SLURM_JOB_NUM_NODES} nodes: ${SLURM_JOB_NODELIST}"
date
pwd
echo "Loaded modules are:"
module list
echo "Running in $TMPWORKDIR"

# Setting this variable will suppress the warnings
# about lack of CUDA support on non-GPU enabled nodes.  We
# are not using CUDA, so warning is harmless.
export OMPI_MCA_mpi_cuda_support=0

# Get the full path to our python executable.  It is best
# to provide full path of our executable to , etc.
MYPYTHON=`which python`
MYPYSCRIPT="${SLURM_SUBMIT_DIR}/hello-mpi.py"
echo "Using python $MYPYTHON"
echo "Using python script $MYPYSCRIPT"

# Run our script using mpirun
# We do not specify the number of tasks here, and instead rely on
# it defaulting to the number of tasks requested of Slurm
mpirun  ${MYPYTHON} $MYPYSCRIPT > hello.out 2>&1
# Save the exit code from the previous command
ECODE=$?

echo "Job finished with exit code $ECODE.  Working dir is $TMPWORKDIR"
date

# Exit with the cached exit code
exit $ECODE