#!/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.


#   This text is cited from
#
#		http://www.mpi-forum.org/docs/mpi-20-html/node42.htm,
#
# section "4.1. Portable MPI Process Startup":
#
# A number of implementations of MPI-1 provide a startup command for MPI
# programs that is of the form
#
#     mpirun <mpirun arguments> <program> <program arguments>
#
# Separating the command to start the program from the program itself
# provides flexibility, particularly for network and heterogeneous
# implementations. For example, the startup script need not run on one of
# the machines that will be executing the MPI program itself.
#
# Having a standard startup mechanism also extends the portability of MPI
# programs one step further, to the command lines and scripts that manage
# them. For example, a validation suite script that runs hundreds of
# programs can be a portable script if it is written using such a standard
# starup mechanism. In order that the ``standard'' command not be confused
# with existing practice, which is not standard and not portable among
# implementations, instead of mpirun MPI specifies mpiexec.
#
# While a standardized startup mechanism improves the usability of MPI,
# the range of environments is so diverse (e.g., there may not even be a
# command line interface) that MPI cannot mandate such a mechanism.
# Instead, MPI specifies an mpiexec startup command and recommends but
# does not require it, as advice to implementors. However, if an
# implementation does provide a command called mpiexec, it must be of the
# form described below.
#
# It is suggested that
#
#     mpiexec -n <numprocs> <program>
#
# be at least one way to start <program> with an initial MPI_COMM_WORLD
# whose group contains <numprocs> processes. Other arguments to mpiexec
# may be implementation-dependent.
#
# This is advice to implementors, rather than a required part of MPI-2. It
# is not suggested that this be the only way to start MPI programs. If an
# implementation does provide a command called mpiexec, however, it must
# be of the form described here.
#
#
# / Advice to implementors./
#
# Implementors, if they do provide a special startup command for MPI
# programs, are advised to give it the following form. The syntax is
# chosen in order that mpiexec be able to be viewed as a command-line
# version of MPI_COMM_SPAWN (See Section Reserved Keys <node97.htm#Node97>).
#
# Analogous to MPI_COMM_SPAWN, we have
#
#
#     mpiexec -n    <maxprocs>
#            -soft  <        >
#            -host  <        >
#            -arch  <        >
#            -wdir  <        >
#            -path  <        >
#            -file  <        >
#             ...
#            <command line>
#
# for the case where a single command line for the application program and
# its arguments will suffice. See Section Reserved Keys
# <node97.htm#Node97> for the meanings of these arguments. For the case
# corresponding to MPI_COMM_SPAWN_MULTIPLE there are two possible formats:
#
# Form A:
#
#
#     mpiexec { <above arguments> } : { ... } : { ... } : ... : { ... }
#
# As with MPI_COMM_SPAWN, all the arguments are optional. (Even the -n x
# argument is optional; the default is implementation dependent. It might
# be 1, it might be taken from an environment variable, or it might be
# specified at compile time.) The names and meanings of the arguments are
# taken from the keys in the info argument to MPI_COMM_SPAWN. There may be
# other, implementation-dependent arguments as well.
#
# Note that Form A, though convenient to type, prevents colons from being
# program arguments. Therefore an alternate, file-based form is allowed:
#
# Form B:
#
#
#     mpiexec -configfile <filename>
#
# where the lines of /</filename/>/ are of the form separated by the
# colons in Form A. Lines beginning with ` #' are comments, and lines may
# be continued by terminating the partial line with `\'.
#
#
# * Example* Start 16 instances of myprog on the current or default machine:
#
#     mpiexec -n 16 myprog
#
#
# * Example* Start 10 processes on the machine called ferrari:
#
#     mpiexec -n 10 -host ferrari myprog
#
#
# * Example* Start three copies of the same program with different
# command-line arguments:
#
#     mpiexec myprog infile1 : myprog infile2 : myprog infile3
#
#
# * Example* Start the ocean program on five Suns and the atmos program on
# 10 RS/6000's:
#
#     mpiexec -n 5 -arch sun ocean : -n 10 -arch rs6000 atmos
#
# It is assumed that the implementation in this case has a method for
# choosing hosts of the appropriate type. Their ranks are in the order
# specified.
# * Example* Start the ocean program on five Suns and the atmos program on
# 10 RS/6000's (Form B):
#
#     mpiexec -configfile myfile
#
# where myfile contains
#
#     -n 5  -arch sun    ocean
#     -n 10 -arch rs6000 atmos
#
# (/ End of advice to implementors./)
# ...
# MPI-2.0 of July 18, 1997
# HTML Generated on September 10, 2001

if [ $# -eq 1 ] && [ $1 = "--version" ]; then
   echo pbs_version = @PBS_VERSION@
   exit 0
fi

#	startup initializations
init()
{
	pbsconffile=${PBS_CONF_FILE:-"@PBS_CONF_FILE@"}
	if [ -r $pbsconffile ]
	then
		. $pbsconffile
		export PBS_TMPDIR="${PBS_TMPDIR:-${TMPDIR:-/var/tmp}}"
	else
		logerr "cannot read PBS configuration file \"$pbsconffile\""
		exit 1
	fi

	vendor_init ${1+"$@"}

	configfile=""
	runfile="`mktemp ${PBS_TMPDIR}/mpiexec_runfileXXXXXX`"
	tmpconfigfile="`mktemp ${PBS_TMPDIR}/mpiexec_configfileXXXXXX`"
	trap "rm -f $runfile $tmpconfigfile" 0 1 2 3 15

	ranknum=0
	reset_rank

	if [ -n "$PBS_MPI_DEBUG" ]
	then
		debug=1
	else
		debug=0
	fi
}

#	The purpose of this function is to find another mpiexec in the user's
#	path and hand control over to it.  If executing outside of PBS, the
#	user's normal PATH is used;  if within PBS, we consult the pre-PBS path
#	(available within PBS as $PBS_O_PATH).  In either case, we take special
#	precautions to avoid attempting to re-exec ourselves.
rempiexec()
{
	#	if PBS_ENVIRONMENT is not set, this function is called before
	#	the init() function.
	if [ -z "$PBS_ENVIRONMENT" ]
	then
	        testPATH=$PATH

		# make implicit "." in $testPATH explicit
		prepPATH=`echo $testPATH | sed	-e 's/^:/.:/'		\
						-e 's/:$/:./'		\
						-e 's/:::*/:.:/g'`

		for component in `echo $prepPATH | tr : " "`
		do
			if [ $component = `dirname $0` ]
			then
				continue
			fi
			if [ -x $component/mpiexec ]
			then
				exec $component/mpiexec ${1+"$@"}
			fi
		done
		logerr "unexpected error - no non-PBS mpiexec in PATH"
	else
		testPATH=$PBS_O_PATH
		pbsbindirID="`filestat $PBS_EXEC/bin`"

		# make implicit "." in $testPATH explicit
		prepPATH=`echo $testPATH | sed	-e 's/^:/.:/'		\
						-e 's/:$/:./'		\
						-e 's/:::*/:.:/g'`

		for component in `echo $prepPATH | tr : " "`
		do
			#	Check to see whether . is $PBS_EXEC/bin
			if [ $component = "." ]
			then
				if [ "`filestat .`" = $pbsbindirID ]
				then
					continue
				fi
			else
				if [ $component = "$PBS_EXEC/bin" ]
				then
					continue
				fi
			fi
			if [ -x $component/mpiexec ]
			then
				exec $component/mpiexec ${1+"$@"}
			fi
		done
		logerr "unexpected error - no non-PBS mpiexec in PBS_O_PATH"
	fi

	exit 1
}

filestat()
{
	if [ $# -ne 1 ]
	then
		logerr "filestat:  unexpected internal error ($#)"
		exit 1
	else
		#	Check for GNU stat(1) command, which allows us to
		#	return a tag (the file's serial number on a given
		#	device) more likely to be unique than the serial
		#	number alone.
		if type stat > /dev/null 2>&1
		then
			statformat="%d:%i"
			stat -c $statformat $1
		else
			#	Sigh - best we can do under the circumstances
			ls -id $1 | awk '{print $1}'
		fi
	fi
}

#	reset per-rank settings
reset_rank()
{
	maxprocs=`wc -l $PBS_NODEFILE | cut -f1 -d' '`
	arch=""
	file=""
	host=""
	path=""
	prog=""
	progargs=""
	set=""
	wdir=""
}

usage()
{
	printf "Usage:\n"
	printf "\t%s\n" "mpiexec -n    <maxprocs>"
	printf "\t%s\n" "-soft  <        >"
	printf "\t%s\n" "-host  <        >"
	printf "\t%s\n" "-arch  <        >"
	printf "\t%s\n" "-wdir  <        >"
	printf "\t%s\n" "-path  <        >"
	printf "\t%s\n" "-file  <        >"
	printf "\t%s\n" "..."
	printf "\t%s\n" "<command line>"
	printf "or\n"
	printf "\t%s\n" "mpiexec <args as above> [ : <args as above> ... ]"
	printf "or\n"
	printf "\t%s\n" "mpiexec -configfile <filename>"
	printf "\t%s\n" "mpiexec --version"

	exit 1
}

logerr()
{
	printf "%s:  %s\n" $MyName ${1+"$@"} > /dev/stderr
}

evalsimpleargs()
{
	if [ $# -le 1 ]
	then
		if [ $# -eq 1 ]
		then
			logerr "$1:  argument expected"
		fi
		usage
	else
		opt="$1"
		arg="$2"
	fi

	case "$opt" in
		"-n")		maxprocs=$arg	;;
		"-soft")	set=$arg	;;	# unimplemented?
		"-host")	host=$arg	;;
		"-arch")	arch=$arg	;;
		"-wdir")	wdir=$arg	;;
		"-path")	path=$arg	;;
		"-file")	file=$arg	;;
		*)		logerr "internal error - option \"$opt\""
				exit 1
				;;
	esac
}

#	debugging hook
printargs()
{
	printf "%s:\n" $MyName
	printf "\tmaxprocs:  %s\n" $maxprocs
	printf "\tsoft:  %s\n" $soft
	printf "\thost:  %s\n" $host
	printf "\tarch:  %s\n" $arch
	printf "\twdir:  %s\n" $wdir
	printf "\tpath:  %s\n" $path
	printf "\tfile:  %s\n" $file
	printf "\tprog:  %s\n" $prog
	printf "\targs:  %s\n" $progargs
	printf "\tconfigfile:  %s\n" $configfile
}

dorank()
{
	line=""
	[ -n "$maxprocs" ] && line="$line -n $maxprocs"
	[ -n "$soft" ] && line="$line -soft $soft"
	[ -n "$host" ] && line="$line -host $host"
	[ -n "$arch" ] && line="$line -arch $arch"
	[ -n "$wdir" ] && line="$line -wdir $wdir"
	[ -n "$path" ] && line="$line -path $path"
	[ -n "$file" ] && line="$line -file $file"
	[ -n "$prog" ] && line="$line $prog"
	[ -n "$progargs" ] && line="$line $progargs"

	echo "$line" >> $tmpconfigfile
	reset_rank
}

#	This function is passed the script's initial arguments (via init()).
#	It does any necessary vendor-specific initializations, including
#	determining whether this is a supported platform.
#
#	Currently the only supported platforms are Altix systems running
#	either the SGI MPI bundle of Performance Suite or older SGIs with
#	ProPack version 4 or greater.  On Altix systems with an earlier
#	version of ProPack, we complain about an unsupported version.  On
#	non-Altix systems, we assume we were invoked by mistake and use
#	the value of PBS_O_PATH to search for the correct version of mpiexec
#	to execute.
vendor_init()
{
	supported_platform=0

	sgi_release="/etc/sgi-release"
	sgi_compute_node="/etc/sgi-compute-node-release"
	sgi_service_node="/etc/sgi-service-node-release"

	if [ -f $sgi_release -o -f $sgi_compute_node -o -f $sgi_service_node ]
	then
		sgiinit
		supported_platform=1
	fi

	if [ $supported_platform -eq 0 ]
	then
		rempiexec ${1+"$@"}
	fi
}

#	This function is passed as its argument an mpiexec-style configuration
#	file and is expected to reformat it into a format suitable for native
#	consumption by the vendor's MPI infrastructure.
vendor_config()
{
	if [ $# -ne 1 ]
	then
		usage
	fi

	sgiconfig ${1+"$@"}
}

#	This function takes no arguments.  It causes the native-format MPI job
#	constructed by vendor_config() to be executed.
vendor_run()
{
	sgimpirun
}

# Handle initialization of an SGI system based on either the chkfeature
# utility or, as a fallback, the presence of various text files on an
# older ProPack system.
sgiinit()
{
	PATH=$PATH:/usr/sbin	# ensure access to chkfeature CLI if present
	if type chkfeature > /dev/null 2>&1
	then
	    #	If MPT is present, make sure it's loaded so we can exec mpirun
	    if chkfeature -p sgi-mpt > /dev/null 2>&1
	    then
		module load mpt > /dev/null 2>&1
	    else
		#	chkfeature says MPT is not present.  If there is an
		#	mpirun in our path, we assume (as we did before) that
		#	it's the SGI one.
		if type mpirun > /dev/null 2>&1
		then
		    :
		else
		    logerr "unexpected error - MPT mpirun unavailable"
		    exit 1
		fi
	    fi
	else
	    #	no chkfeature - older ProPack
	    sgippinit
	fi
}

#	For the release into which this change is targeted, only the Altix
#	should use the PBS version of mpiexec.  Therefore, our mpiexec will
#	determine whether it is executing on an SGI system running ProPack 4
#	or greater.  This is accomplished by examining the /etc/sgi-release
#	file to look for a string of the form
#
#		"SGI ProPack N ..."
#
#	where N is an integer greater than or equal to 4.  There are three
#	cases to consider:
#
#		-  if the file does not exist, 	PBS's mpiexec assumes that it
#		   was the unintentional recipient of control and searches the
#		   user's pre-PBS path, whose value is found in the PBS_O_PATH
#		   environment variable, for mpiexec.  If one is found, we exec
#		   it;  otherwise, an appropriate error message is displayed
#		   and we exit with an error
#
#		-  if the file does exist but its format is not what we expect
#		   to find, or N is less than 4, an appropriate error message
#		   is displayed and we exit with an error
#
#		-  otherwise, we proceed with normal execution
sgippinit()
{
	if ! sgicheckppversion
	then
		logerr "unexpected error - sgicheckppversion returned $?"
		exit 1
	fi
}

sgicheckppversion()
{
	sgi_release="/etc/sgi-release"
	sgi_compute_node="/etc/sgi-compute-node-release"
	sgi_ppversmin=4
	if [ -r $sgi_release ]
	then
		read sgi propack propackvers rest < $sgi_release
		if [ "$sgi" != "SGI" -o "$propack" != "ProPack" ]
		then
			logerr "$sgi_release:  unexpected contents"
			exit 1
		fi
		if [ `expr substr $propackvers 1 1` -lt $sgi_ppversmin ]
		then
			logerr "ProPack version $propackvers is unsupported"
			exit 1
		else
			return 0
		fi
	elif [ -r $sgi_compute_node ]
	then
		if grep "Build 5" $sgi_compute_node > /dev/null
		then
			return 0
		fi
	else
		return 1
	fi
}

#	Translate the mpiexec-style configuration file into an mpirun-style
#	one ...
sgiconfig()
{
	if [ $debug -eq 1 ]
	then
		report_config "$1"
	fi

	PBS_LIB_PATH=${PBS_EXEC}/lib
	if [ ! -d ${PBS_LIB_PATH} -a -d ${PBS_EXEC}/lib64 ] ; then
		PBS_LIB_PATH=${PBS_EXEC}/lib64
	fi

	awk -f ${PBS_LIB_PATH}/MPI/sgiMPI.awk	-v configfile="$1"	\
						-v runfile="$runfile"	\
						-v pbs_exec="$PBS_EXEC"	\
						-v debug=$debug
}

#	... and execute it.
sgimpirun()
{
	if [ $debug -eq 1 ]
	then
		report_run
	fi

	# The Performance Suite version of mpirun needs to be told
	# that it ought not complain when we exec it via this script.
	# Don't be confused by the name - it's simply a directive to
	# the SGI mpirun command to assert that mpirun need not worry
	# its pretty little head about the absence of a pbs_attach
	# command among its command line arguments.
	export MPI_IGNORE_PBS=1

	a_opt=''
	if [ -n "$PBS_MPI_SGIARRAY" ]
	then
		a_opt="-a $PBS_MPI_SGIARRAY"
	fi

	mpirun $a_opt -f $runfile
}

#	debugging hooks
report_config()
{
	echo "generated mpiexec configuration file ($1) contains"
	cat "$1" | sed -e 's/^/\t/'
}
report_run()
{
	if [ -n "$PBS_MPI_SGIARRAY" ]
	then
		echo "mpirun -a $PBS_MPI_SGIARRAY -f $runfile"
	else
		echo "mpirun -f $runfile"
	fi
	echo "where $runfile contains:"
	cat $runfile | sed -e 's/^/\t/'
}

MyName="`basename $0`"					# must occur first

[ -z "$PBS_ENVIRONMENT" ] && rempiexec ${1+"$@"}	# no work for us to do

init ${1+"$@"}

in_rankdef=0
while [ $# -gt 0 ]
do
	case "$1" in
		"-n"|"-soft"|"-host"|"-arch"|"-wdir"|"-path"|"-file")
			in_rankdef=1
			evalsimpleargs ${1+"$@"}
			shift 2
			;;
		"-configfile")
			if [ $in_rankdef -eq 1 ]
			then
				logerr "-configfile in rank definition"
				exit 1
			fi
			# first "-configfile" option terminates argument parsing
			shift
			configfile="$1"
			vendor_config $configfile && vendor_run
			exit $?
			;;
		":")
			in_rankdef=0
			shift
			evalsimpleargs ${1+"$@"}
			while [ $# -gt 0 -a "$1" != ":" ]
			do
				shift
			done
			;;
		*)
			prog="$1"
			shift
			while [ $# -gt 0 ]
			do
				if [ "$1" = ":" ]
				then
					in_rankdef=0
					dorank
					break
				else
					progargs="$progargs $1"
				fi
				shift
			done
			if [ $# -gt 0 ]
			then
				if [ `expr substr "$1" 1 1` = ":" ]
				then
					shift
				fi
			else
				ranknum=`expr $ranknum + 1`
				in_rankdef=0
				dorank
			fi
			;;
	esac
done

if [ $in_rankdef -eq 1 -a -z "$prog" ]
then
	logerr "rank $ranknum has no executable"
	exit 1
else
	vendor_config $tmpconfigfile && vendor_run
fi
