#!/bin/bash
# The line above this is the "shebang" line.  It must be first line in script
#-----------------------------------------------------
#	Example Submission Script: Hello-UMD, Hybrid Ssh version
#	Runs a multithreaded hello-umd command using ssh to span nodes
#-----------------------------------------------------
#
# Slurm sbatch parameters section:
#	Request 3 "tasks" with 15 CPU cores each
#SBATCH -n 3
#SBATCH -c  15
#	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 hello-umd with multiple "tasks" each using 
# multiple threads.  The tasks are launched using ssh instead of MPI.
# We create a directory on parallel filesystem from  where we actually 
# will run the job.

# We start by defining a function to launch the program on every
# node allocated to the job (and for each node, as many times as
# the number of tasks allocated to the node).  
# The use of this function makes the code clearer, plus allows us
# to easily capture the output to a single file.
launch_ssh_tasks()
{	# First, get the full path to our hello-umd executable.  It is best
	# to provide the full path of our executable to ssh, etc.
	MYEXE=`which hello-umd`
	echo "Using executable $MYEXE"

	# Now we get a list of hostnames for nodes our job is running
	# on.  We want each hostname to appear in the list exactly as
	# many times as the number of tasks allocated to the node.
	# This information is in the Slurm environmental variables
	# SLURM_JOB_NODELIST and SLURM_TASKS_PER_NODE, but in an abbreviated
	# form.  We call a helper script to convert to the format we want.
	MYNODES=`/software/acigs-utilities/bin/slurm_hostnames_by_tasks \
		-r ' ' -m hosts_by_tasks`

	# Keep track of our task number and the pid for each task
	tasknum=0
	pidlist=()

	# We now use ssh to spawn a copy of hello-umd for each task 
	# for our, using $SLURM_CPUS_PER_TASK threads per task
	for node in $MYNODES
	do
		# We run in background to get parallelism
		ssh -q $node $MYEXE \
			-t ${SLURM_CPUS_PER_TASK} \
			-m "'Hello from task $tasknum'" &
		# Collect the process id (pid)
		pidlist[$tasknum]=$!
		# and increment tasknum
		tasknum=$(( tasknum + 1 ))
	done

	# Wait for all of our background jobs to complete and if any
	# had non-zero exit codes, warn us and set the main script's 
	# exit code to non-zero
	ECODE=0
	tasknum=0
	for pid in "${pidlist[@]}"
	do
		wait $pid
		tmpecode=$?
		if  [ $tmpecode -ne 0 ]; then
			ECODE=$tmpecode
			echo >&2 "Task $tasknum had non-zero exit code $ECODE"
		fi
		tasknum=$(( tasknum + 1 ))
	done
	return $ECODE
}

# 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 and package modules
module load gcc
module load hello-umd

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

# Section to output information identifying the job, etc.
echo "Slurm job submit-hello-umd-hybridssh.sh ${SLURM_JOB_ID} 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 "Job will be run in $TMPWORKDIR"

# Launch our code using the previously defined function, saving output
# to hello.out
launch_ssh_tasks > $TMPWORKDIR/hello.out 2>&1
# and save exit code
ECODE=$?

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

# Exit with the cached exit code
exit $ECODE