#!/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.mvapich2.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="\
-l \
-bnr \
-s <spec> \
-1 \
-ifhn \
-tv \
-gdb \
-m \
-a \
-ecfn \
-n <n> \
-np <n> \
-wdir <dirname> \
-umask <umask> \
-path <dirname> \
-soft <spec> \
-arch <arch> \
-envall \
-envnone \
-envlist <list> \
-env <var> <val> \
-gn <x> \
-gnp <x> \
-gwdir <dirname> \
-gsoft <spec> \
-garch <arch> \
-genv <var> <val> \
-genvall \
-genvlist <env_var_names> \
-genvnone \
-gpath <search_path> \
-gumask <umask> \
"

#####################################################################
# 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> \
-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="\
"

#####################################################################
# pbs_attach: full path or relative path to the pbs_attach executable.
#
#####################################################################
pbs_attach=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"


################## 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

#	If user has done mpdboot then we don't need to do it.
	$mpirun_location/mpdtrace >/dev/null 2>/dev/null
	if [ $? -eq 0 ]; then
		return 0
	fi

	# Get the host.list from PBS and uniq it.
	pbs_hostsfile="${PBS_TMPDIR:-/var/tmp}/pbsrun_hosts$$"
        sort -u $PBS_NODEFILE > $pbs_hostsfile
	num_mpds=`wc -l $pbs_hostsfile | awk '{print $1}'`
	$mpirun_location/mpdboot -n $num_mpds --file=$pbs_hostsfile
	return 0
}


################## 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 () {
	args=$*

	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 $args"
}


################## 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 () {
	args=$*

        printf "%s" "$args"
}

################## 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$$"

	if [ -f $pbs_hostsfile ]; then
    		$mpirun_location/mpdallexit  >/dev/null 2>/dev/null
		rm -f $pbs_hostsfile
	fi
	rm -f $pbs_machinefile
}
