#!/bin/sh
#
# Copyright (C) 1994-2021 Altair Engineering, Inc.
# For more information, contact Altair at www.altair.com.
#
# This file is part of both the OpenPBS software ("OpenPBS")
# and the PBS Professional ("PBS Pro") software.
#
# Open Source License Information:
#
# OpenPBS is free software. You can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# OpenPBS is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public
# License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Commercial License Information:
#
# PBS Pro is commercially licensed software that shares a common core with
# the OpenPBS software.  For a copy of the commercial license terms and
# conditions, go to: (http://www.pbspro.com/agreement.html) or contact the
# Altair Legal Department.
#
# Altair's dual-license business model allows companies, individuals, and
# organizations to create proprietary derivative works of OpenPBS and
# distribute them - whether embedded or bundled with other software -
# under a commercial license agreement.
#
# Use of Altair's trademarks, including but not limited to "PBS™",
# "OpenPBS®", "PBS Professional®", and "PBS Pro™" and Altair's logos is
# subject to Altair's trademark licensing policies.

#

### FILE: pbsrun.intelmpi.init ###

####################################################################
# strict_pbs: set to 1 if you want the wrapper script to be executed
#             only if in a PBS environment
#
#####################################################################
strict_pbs=0

#####################################################################
# options_to_retain:  space-separted list of options and values that
# "pbsrun" will pass on to the actual mpirun call. options must begin
# with "-" or "--", option name can contain a wildcard (*) to match a
# number of characters (e.g. --totalnum=*), and option arguments must
# be specified by some arbitrary name bounded by angle brackets,
# as in "<val1>".
#####################################################################
options_to_retain="\
-h \
--help \
-help \
-nolocal \
-perhost <n> \
-tv \
-n <n> \
-np <n> \
-wdir <directory> \
-path <directory> \
-soft <spec> \
-arch <arch> \
-env <envvar> <value> \
-envall \
-envnone \
-envlist <list_of_env_var_names> \
-version \
-V \
-genv <envvar> <value> \
-genvnone \
-gn <n> \
-gnp <n> \
-gwdir <directory> \
-gpath <directory> \
-gsoft <spec> \
-garch <arch> \
-genv <envvar> <value> \
-genvall \
-genvnone \
-genvlist <list_of_env_var_names> \
-gversion \
-gV \
-l \
-if \
-kx \
-bnr \
-d \
-v \
--help \
--rsh=* \
--user=* \
--mpd=* \
--loccons \
--remcons \
--shell \
--verbose \
-1 \
--ncpus=* \
-r <rshcmd> \
-u <user> \
-m <mpdcmd> \
-file <XML_job_description> \
-s <spec> \
"

#####################################################################
# options_to_ignore:  space-separted list of options and values that
# "pbsrun" will NOT pass on to the actual mpirun call. options must begin
# with "-" or "--", option name can contain a wildcard (*) to match a
# number of characters (e.g. --totalnum=*), and option arguments must
# be specified by some arbitrary name bounded by angle brackets,
# as in "<val1>".
#####################################################################
options_to_ignore="\
-host <hostname> \
-ghost <hostname> \
--totalnum=* \
--file=* \
-f <mpd_hosts_file> \
-machinefile <file> \
"

#####################################################################
# options_to_transform:  space-separted list of options and values that
# "pbsrun" will modify before passing on to the actual mpirun call.
# options must begin with "-" or "--", option name can contain a
# wildcard (*) to match a number of characters (e.g. --totalnum=*),
# and option arguments must be specified by some arbitrary name
# bounded by angle brackets, as in "<val1>".
# NOTE: Adding values here require code to be added to
# transform_action() function appearing later in this file.
#####################################################################
options_to_transform="\
"

#####################################################################
# options_to_fail: space-separated list of options that will cause "pbsrun"
# to exit upon encountering a match.
# options must begin with "-" or "--", and option name can contain a
# wildcard (*) to match a number of characters (e.g. --totalnum=*).
#
#####################################################################
options_to_fail="\
"

#####################################################################
# option_to_configfile: the SINGLE option and value that refers to the
# name of the "configfile" containing command line segments, found
# in certain versions of mpirun.
# option must begin with "-" or "--".
#
#####################################################################
option_to_configfile="\
-configfile <configfile> \
"
#####################################################################
# options_with_another_form: space-separated list of options and values
# that can be found in options_to_retain, options_to_ignore, or
# options_to_transform, whose syntax has an alternate, unsupported
# form.
# This usually occurs if a version of mpirun has different forms for
# the same option.
# For instance,
#       MPICH2's mpirun provides:
#               mpirun -localonly <x>
#               mpirun -localonly
# If options_to_retain lists "-localonly" as supported value, then
# set options_with_another_form="-localonly" as well.
# This would cause "pbsrun" to issue a
# warning about alternate forms upon encountering the  option.
#
# options must begin with "-" or "--", and option name can contain a
# wildcard (*) to match a number of characters (e.g. --totalnum=*).
#
#####################################################################

options_with_another_form="\
-s <spec> \
"

#####################################################################
# pbs_attach: full path or relative path to the pbs_attach executable.
#
#####################################################################
pbs_attach=$PBS_EXEC/bin/pbs_attach

#####################################################################
# options_to_pbs_attach:  are the special options to pass to the
# pbs_attach call. You may pass variable references (e.g. $PBS_JOBID)
# and they will be substituted  by pbsrun to actual values.
#####################################################################
options_to_pbs_attach="-j $PBS_JOBID -P"


################## transform_action() ###################################
# The action to be performed for each actual item and value matched in
# options_to_transform.
# RETURN: echo the replacement item and value for the matched arguments.
# NOTES:
# (1) "echo" produces the return value of this function;
#     do not arbitrarily invoke echo statements.
# (2) Please use 'printf "%s" "<value>"' (must quote the <value>) in place
#     of "echo" command here since we're returning an option, and 'echo' is
#     notorius for the following behavior:
#               echo "-n"               --> prints empty (-n is an echo opt)
#     Desired behavior:
#               printf "%s" "-n"        --> prints "-n"
#
########################################################################
transform_action () {
    args=$*
    printf "%s" "$args"
}

################## boot_action() ###################################
# The action to be performed BEFORE before calling mpirun.
# The location of actual mpirun is passed as first argument.
# RETURN: 0 for success; non-zero otherwise.
# NOTES:
# (1) 'return' produces the exit status of this function:
#     0 for success; non-zero for failure.
####################################################################
boot_action () {

	mpirun_location=$1

	return 0
}

# Check if version is older than "4.0.3"
# awk exit 0 means true (version is older)
# awk exit 1 means false (version is equal or newer)
ver_older () {
	file="/tmp/ckver$$"
	cat <<-EOF > $file
	{
		if (NF < 3)
			exit 0;
		if (\$1 < 4)
			exit 0;
		if (\$1 > 4)
			exit 1;
		if (\$2 > 0)
			exit 1;
		if (\$3 < 3)
			exit 0;
		exit 1;
	}
	EOF
	echo $1 | awk -F. -f $file
	ret=$?
	rm $file
	return $ret
}

# get mpirun version and see if it is at least 4.0.3
hydra_supported () {
	file="/tmp/getver$$"
	cat <<-EOF > $file
	BEGIN {
		ver = "0.0";
		update = "0";
	}

	{
		for(i=1;i<=NF;i++) {
			if (\$i == "Version" && i<NF)
				ver = \$(i+1);
			if (\$i == "Update" && i<NF)
				update = \$(i+1);
		}
	}

	END {
		printf("%s.%s\n", ver, update);
	}
	EOF

	ver=`$mpirun -V |awk -f $file`
	rm $file

	if ver_older $ver ;then
		return 1
	else
		return 0	# hydra supported
	fi
}

################## evaluate_options_action() #######################
# The action to be performed on the actual options and values matched
# in options_to_retain, not including those in options_to_ignore, and
# those changed arguments in options_to_transform, as well as any
# other transformation needed on the program name and program arguments.
#
# RETURN: echo the list of final arguments and program arguments
# to be passed on to mpirun command line.

# NOTES:
# (1) "echo" produces the return value of this function;
#     do not arbitrarily invoke echo statements.
# (2) Please use 'printf "%s" "<value>"' (must quote <value>) in place of
#     "echo" command here since we're returning option string, and
#     'echo' is notorius for the following behavior:
#               echo "-n"               --> prints empty (-n is an echo opt)
#     Desired behavior:
#               printf "%s" "-n"        --> prints "-n"
#
########################################################################
evaluate_options_action () {

# 	Special handling of --rsh and -r because they have to
#	be added before other options.
	first=""
	args=""

	while [ $# -gt 0 ] ; do
		case "$1" in
		$pbs_attach)
			args="$args $*"
			break
			;;

		-r|-u|-m)
			first="$first $1 $2"
			shift
			;;

		--rsh=*|--user=*|--mpd=*)
			first="$first $1"
			;;

		*)
			args="$args $1"
			;;
		esac
		shift
	done

	pbs_hostsfile="${PBS_TMPDIR:-/var/tmp}/pbsrun_hosts$$"
	pbs_machinefile="${PBS_TMPDIR:-/var/tmp}/pbsrun_machfile$$"

	sort -u $PBS_NODEFILE > $pbs_hostsfile
	num_mpds=`wc -l $pbs_hostsfile | awk '{print $1}'`

	cat -n ${PBS_NODEFILE} | \
		sort -k2 | uniq -f1 -c | \
		awk '{if ($1 == 1) print $2, $3; else print $2, $3 ":" $1}'|\
		sort -n | awk '{print $2}' > $pbs_machinefile

	if [ "x$I_MPI_PROCESS_MANAGER" != "xmpd" -a \
				"x$I_MPI_PROCESS_MANAGER" != "xMPD" ] && hydra_supported ;then
		printf "%s" "-hostfile=$pbs_hostsfile $first -machinefile $pbs_machinefile $args"
	else
		printf "%s" "--totalnum=$num_mpds --file=$pbs_hostsfile $first -machinefile $pbs_machinefile $args"
	fi
}


################## configfile_cmdline_action() #######################
# If the option_to_configfile (e.g. -configfile) is specified in the
# mpirun command line, then this function gets passed any leading options
# and values found before option_to_configfile.
#
# RETURN: return the actual options and values to be put in before the
# option_to_configfile parameter. For instance,
#	returning "--totalnum=N --file=Y" would result in
# an mpirun command line of:
#	mpirun --totalnum=N --file=Y -configfile pbs_config
#
########################################################################
configfile_cmdline_action () {

# 	Special handling of --rsh and -r because they have to
#	be added before other options.
	first=""
	args=""

	while [ $# -gt 0 ] ; do
		case "$1" in
		$pbs_attach)
			args="$args $*"
			break
			;;

		-r|-u|-m)
			first="$first $1 $2"
			shift
			;;

		--rsh=*|--user=*|--mpd=*)
			first="$first $1"
			;;

		*)
			args="$args $1"
			;;
		esac
		shift
	done

	pbs_hostsfile="${PBS_TMPDIR:-/var/tmp}/pbsrun_hosts$$"

	sort -u $PBS_NODEFILE > $pbs_hostsfile
	num_mpds=`wc -l $pbs_hostsfile | awk '{print $1}'`

	if [ "x$I_MPI_PROCESS_MANAGER" != "xmpd" -a \
				"x$I_MPI_PROCESS_MANAGER" != "xMPD" ] && hydra_supported ;then
		printf "%s" "-hostfile=$pbs_hostsfile $first $args"
	else
		printf "%s" "--totalnum=$num_mpds --file=$pbs_hostsfile $first $args"
	fi
}

################## configfile_firstline_action () #######################
# If the option_to_configfile (e.g. -configfile) is specified in the
# mpirun command line, return here the item that will be put in the
# FIRST line of the configuration file.
# This is the place to put  the "-machinefile <filename>" parameter
# which determines the processes to hosts mappings. Some versions
# of mpirun (MPICH2, Intel MPI) require that the -machinefile parameter
# must appear inside the config file and not on the commmand line.
########################################################################
configfile_firstline_action () {
	pbs_machinefile="${PBS_TMPDIR:-/var/tmp}/pbsrun_machfile$$"

        cat -n ${PBS_NODEFILE} | \
        sort -k2 | uniq -f1 -c | \
        awk '{if ($1 == 1) print $2, $3; else print $2, $3 ":" $1}'|\
        sort -n | awk '{print $2}' > $pbs_machinefile

        printf "%s" "-machinefile $pbs_machinefile"
}

################## end_action() ########################################
# The action to be performed AFTER calling mpirun, and also when
# mpirun wrap script is prematurely interrupted.
# INPUT: The location of actual mpirun is passed as first argument.
# RETURN: none
########################################################################
end_action () {

    mpirun_location=$1

    pbs_hostsfile="${PBS_TMPDIR:-/var/tmp}/pbsrun_hosts$$"
    pbs_machinefile="${PBS_TMPDIR:-/var/tmp}/pbsrun_machfile$$"

    rm -f $pbs_hostsfile
    rm -f $pbs_machinefile
}
