#!/bin/bash
#Need bash for arrays
#
#---------------------------------------------------------
# This script is under version control in git repository
#git@gitlab.umd.edu:dit-perl-packages/hpcc-pdb-packages/assorted-slurm-scripts.git 
#
# Please make all edits to git repo
#---------------------------------------------------------
VERSION="1.0"

#Set defaults
FLDDELIM=':'
RECDELIM="\n"
#No default for MODE
MODE=


usage()
{	cat - >&2 <<EOF
$0: Expands Slurm SLURM_JOB_NODELIST and SLURM_TASKS_PER_NODE

Usage:

$0 [ -V ] [ -h ] [ -M ] [ -R ] \\
	[ -f FLDDELIM ] [ -r RECDELIM ] -m MODE 

This script expands the Slurm variables SLURM_JOB_NODELIST
and SLURM_TASKS_PER_NODE for the job to more usable formats, 
depending on the MODE flag.

Flags:
-h prints this help text, then exits
-M prints a "man"-like page, then exits
-R prints a nroff source for man page, then exits

-f FLDDELIM sets the intrafield delimiter, defaults to '$FLDDELIM'
-r RECDELIM sets the intrarecord delimiter, defaults to '\\n'
-m MODE sets the mode, described below.

The following values are acceptable for MODE:

	hosts_by_tasks: This will print the hostname for each 
		node allocated to the job, possibly multiple 
		times.  The number of times any hostname occurs 
		is equal to the number of tasks on that host.
		Hostnames are separated by RECDELIM.

	counts: This will print a list of the task count by
		node, in the same order as in SLURM_JOB_NODELIST.
		Basically, just expands SLURM_TASKS_PER_NODE.
		Task counts are separated by RECDELIM.

	host_count: This will print hostnames followed by
		the task count for that host.  Each hostname
		will appear only once.  The hostname and task 
		count will be separated by FLDDELIM.  The
		hostname/taskcount pairs will be separated by
		RECDELIM.

	hostnames: This will print the hostnames for each node
		allocated to the job exactly once.  Basically,
		just expands SLURM_JOB_NODELIST, i.e. the 
		equivalent of "scontrol show hostnames".
		Hostnames are separated by RECDELIM.

EOF
}

rawman_usage()
{	progname=`basename $0`
	PROGNAME=`echo $progname | tr 'a-z' 'A-Z'`

	cat - << EOF
.TH $PROGNAME 1
.SH NAME
$progname - Expand Slurm variables SLURM_JOB_NODELIST and SLURM_TASKS_BY_NODE
.SH SYNOPSIS
.B $progname
.nh
[\fB\-h\fR]
[\fB\-M\fR]
[\fB\-R\fR]
[\fB\-V\fR]
[\fB\-f\fR\ \%\fNFLDDELIM\fR]
[\fB\-r\fR\ \%\fNRECDELIM\fR]
\fB\-m\fR\ \%\fNMODE\fR
.SH DESCRIPTION
.B $progname
is a script to expand the Slurm variables \fISLURM_JOB_NODELIST\fR and \fISLURM_TASKS_BY_NODE\fR
from the compact, abbreviated form Slurm uses to more usable formats.  The exact format
returned depends on the \fIMODE\fR parameter, as modified by \fIFLDDELIM\fR and
\fIRECDELIM\fR parameters.
.SH OPTIONS
.TP
.BR \-h
Prints a simple usage/help text and exits
.PP
.BR \-M
Prints this man page and exits
.PP
.BR \-R
Prints the raw, nroff for this man page, and exits
.PP
.BR \-V
Prints the version number and exits
.PP
.BR \-f\ \fIFLDDELIM\fR
Sets the delimiter to use between fields in a record.  Defaults to ':'.
.PP
.BR \-r\ \fIRECDELIM\fR
Sets the delimiter to use between records.  Defaults to a newline character.
.PP
.BR \-m\ \fIMODE\fR
Sets the mode of operation. This parameter is \fBREQUIRED\fR.  Valid values
are \fIhosts_by_tasks\fR, \fIcounts\fR, \fIhost_count\fR,  or
\fIhostnames\fR.  These are discussed in more detail below.
.PP
The format of the output is determined by the \fIMODE\fR parameter.
The following:
.PP
.BR hosts_by_tasks
This will print the hostname for each node allocated to the job,
possibly multiple times.  The number of times a hostname gets
printed corresponds to the number of tasks allocated to that node.
The hostnames will be separated by \fIRECDELIM\fR
.PP
.BR host_count
This will print hostnames followed by the number of tasks allocated
to that host.  The hostname for each node allocated to the job
will appear exactly once.  The hostname and task count will be
separated by \fIFLDDELIM\fR, and the hostname/task count pairs
will be separated by \fIRECDELIM\fR.
.PP
.BR hostnames
This will print the hostnames for the nodes allocated to the
job.  Each hostname will occur exactly once.  Basically just
does a scontrol show hostnames.  Hostnames will be separated
by \fIRECDELIM\fR.
.PP
.BR counts
This will just expand \fISLURM_TASKS_BY_NODE\fR.  It will print
a list of the number of tasks per host, with the hosts being
in the same order as in the \fISLURM_JOB_NODELIST\fR variable.
These numbers will be separated by \fIRECDELIM\fR.
.SH EXAMPLES
Consider the case where the job is allocated to 4 nodes,
compute\-0, compute\-1, compute\-2, and compute\-3, with 2 tasks
allocated to the first three nodes, and one task allocated to 
compute\-3.
.PP
In the above scenario, we would likely have the environment variables
set to
.PP
.nf
.RS
SLURM_JOB_NODELIST="compute\-[0\-3]"
SLURM_TASKS_PER_NODE="2(x3),1"
.RE
.fi
.PP
Running
.PP
.nf
.RS
$progname\ \-m\ hosts_by_tasks
.RE
.fi
.PP
would output
.PP
.nf
.RS
compute\-0
compute\-0
compute\-1
compute\-1
compute\-2
compute\-2
compute\-3
.RE
.fi
.PP
I.e., hosts compute\-0, compute\-1, and compute\-2 get printed
twice because they each had 2 tasks allocated to them, and 
compute\-3 gets printed once as only one task was allocated to it.
.PP
If one sets \fIRECDELIM\fR to ", ", e.g.
.PP
.nf
.RS
$progname\ \-r\ ",\ "\ \-m\ hosts_by_tasks
.RE
.fi
.PP
we would get
.PP
.nf
.RS
compute\-0,\ compute\-0,\ compute\-1,\ compute\-1,\ compute\-2,\ compute\-2,\ compute\-3
.RE
.fi
.PP
We can use a \fIMODE\fR of host_count, as follows
.PP
.nf
.RS
$progname\ \-m\ host_count
.RE
.fi
.PP
to get the output in the form
.PP
.nf
.RS
compute\-0:2
compute\-1:2
compute\-2:2
compute\-3:1
.RE
.fi
.PP
We can adjust the delimiters like
.PP
.nf
.RS
$progname\ \-r\ ";"\ \f\ ","\ \-m\ host_count
.RE
.fi
.PP
to get
.PP
.nf
.RS
compute\-0,2;compute\-1,2;compute\-2,2;compute\-3,1
.RE
.fi
.PP
The other two modes work similarly, e.g.
.PP
.nf
.RS
$progname\ \-r\ ",\ "\ \-m\ hostnames
.RE
.fi
.PP
will return
compute\-0, compute\-1, compute\-2, compute\-3
.RE
.fi
.PP
and 
.PP
.nf
.RS
$progname\ \-r\ ",\ "\ \-m\ counts
.RE
.fi
.PP
would print
.PP
.nf
.RS
2, 2, 2, 1
.RE
.fi
EOF
}

man_usage()
{	pager="$PAGER"
	if [ "x$pager" == "x" -o ! -x "$pager" ]; then
		#Default pager to "more"
		pager="more"
	fi
	nroff="nroff"

	rawman_usage | $nroff -man | $pager
}

#This defines a bash function to expand SLURM_TASKS_PER_NODE from its
#compact form (using '(xN)' for repeated nodes with same task count) to
# a simple space delimited list of tasks per node.
#This is a helper for expand_tasks_hosts() function below
expand_tasklist()
{	#Split SLURM_TASKS_PER_NODE at the commas
	IFS="," read -ra taskarry <<< $SLURM_TASKS_PER_NODE

	#Loop over each comma-delimited record
	for rec in "${taskarry[@]}"
	do
		#See if there is a '(x' pattern in rec
		echo $rec | grep "(x" > /dev/null
		if [ $? -eq 0 ]; then
			#There is a (x pattern in rec
			#Split into ntasks and count
			ntasks=`echo $rec | sed -e 's/(x.*$//'`
			count=`echo $rec | sed -e 's/^.*(x//' -e 's/)$//'`
		else
			#No (x pattern.  
			#ntasks = rec 
			#and count = 1
			ntasks="$rec"
			count=1
		fi

		#Print out ntasks count times
		for i in `seq 1 $count`
		do
			echo $ntasks
		done
	done
}

#This defines a bash function to expand SLURM_JOB_NODELIST from its
#compact form (using ranges, etc) to a simple list.
#Basically, just invokes scontrol show hostnames
expand_hostlist()
{	HLIST="${SLURM_JOB_NODELIST}"
	if [ "x$HLIST" = "x" ]; then
		#Try deprecated SLURM_NODELIST
		HLIST="${SLURM_NODELIST}"
	fi
	if [ "x$HLIST" = "x" ]; then
		echo >&2 "ERROR: SLURM_JOB_NODELIST not found"
		echo >&2 "This script should only be run within"
		echo >&2 "a Slurm job"
		exit 1
	fi

	hostnames=( `scontrol show hostnames $HLIST` )
	echo "${hostnames[@]}"
}

#This defines a bash function which uses both SLURM_TASKS_PER_NODE
#and SLURM_JOB_NODELIST and returns a list of  hostnames for the
#nodes allocated to this job, with each hostname occuring once for
#each task assigned to the node.  
#Usage: hostlist=`expand_tasks_hosts`
expand_tasks_hosts()
{	#First convert SLURM_JOB_NODELIST into an array of hostnames
	hostnames=( `expand_hostlist` )
	#And convert SLURM_TASKS_PER_NODE into an array of task counts
	taskcounts=( `expand_tasklist` )

	#Get the count of hostnames, less one
	n="${#hostnames[*]}"
	n=$(( n - 1 ))

	#Loop over the hostnames and taskcounts arrays
	for i in `seq 0 $n`
	do
		#Get the i-th hostname/taskcount
		host="${hostnames[$i]}"
		count="${taskcounts[$i]}"

		#Output each hostname as many times as it has tasks
		for j in `seq 1 $count`
		do
			echo "$host"
		done
	done
}

#-----------------	Main -----------------
#Can't use getopts because we want to ignore some options
while getopts :hMRf:r:m: arg
do
    case $arg in
	h)
		usage
		exit 0
		;;

	M)
		man_usage
		exit 0
		;;

	R)
		rawman_usage
		exit 0
		;;
	V)
		echo "$VERSION"
		exit 0
		;;

	r)
		RECDELIM="$OPTARG"
		;;

	f)
		FLDDELIM="$OPTARG"
		;;

	m)
		MODE="$OPTARG"
		;;

	*)
		usage
		echo >&2
		echo "Illegal option -$OPTARG, aborting." >&2
		exit 1
		;;
    esac
done
COUNT=$(( OPTIND - 1 ))
shift $COUNT

if [ "x$MODE" = "x" ]; then
	usage
	echo >&2
	echo "You must supply a value for MODE, aborting." >&2
	exit 1
fi

case $MODE in
   hostnames)
	hostnames=( `expand_hostlist` )
	delim=""
	for hostname in "${hostnames[@]}"
	do
		echo -n -e "$delim"
		delim="$RECDELIM"
		echo -n -e "${hostname}"
	done
	echo -n -e "\n"
	exit 0
	;;

    counts)
	taskcounts=( `expand_tasklist` )
	delim=""
	for taskcount in "${taskcounts[@]}"
	do
		echo -n -e "$delim"
		delim="$RECDELIM"
		echo -n -e "${taskcount}"
	done
	echo -n -e "\n"
	exit 0
	;;

    hosts_by_tasks)
	hostlist=( `expand_tasks_hosts` )
	delim=""
	for hostname in "${hostlist[@]}"
	do
		echo -n -e "$delim"
		delim="$RECDELIM"
		echo -n -e "${hostname}"
	done
	echo -n -e "\n"
	exit 0
	;;

    host_count)
	hostlist=( `expand_hostlist` )
	taskcounts=( `expand_tasklist` )
	n="${#hostlist[*]}"
	n=$(( n - 1 ))
	delim=""
	for i in `seq 0 $n`
	do
		#Get the i-th hostname/taskcount
		host="${hostlist[$i]}"
		count="${taskcounts[$i]}"

		echo -n -e "$delim"
		delim="$RECDELIM"
		echo -n -e "${host}${FLDDELIM}${count}"
	done
	echo -n -e "\n"
	exit 0
	;;

    *)
	usage
	echo >&2
	echo >&2 "ERROR: Unrecognized value '$MODE' for MODE"
	echo >&2 "Aborting"
	exit 1
	;;
esac