#!/usr/bin/bash
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

CONFIG=/etc/sysconfig/prun
MODE=native
LOGLEVEL=2 # 0=error,1=warn,2=info,3=debug

function usage() {
	echo " "
	echo "This OpenHPC utility is used to launch parallel (MPI) applications"
	echo "within a supported resource manager."
	echo " "
	echo "Usage: ${DELIM} <options> executable [arguments]"
	echo " "
	echo "where available options are as follows:"
	echo " "
	echo "  -h        generate help message and exit"
	echo "  -v        enable verbose output"
	echo " "
	exit 1
}

function parse_config_file() {
	if [ -e "${CONFIG}" ]; then
		# shellcheck disable=SC1090
		source "${CONFIG}" 2>/dev/null
	fi
}

function _error() {
	if [ "${LOGLEVEL}" -ge 0 ]; then
		echo -n "[${DELIM}] Error: "
		echo "$@"
	fi
}

function _warn() {
	if [ "${LOGLEVEL}" -ge 1 ]; then
		echo -n "[${DELIM}] "
		echo "$@"
	fi
}

function _info() {
	if [ "${LOGLEVEL}" -ge 2 ]; then
		echo -n "[${DELIM}] "
		echo "$@"
	fi
}

function _debug() {
	if [ "${LOGLEVEL}" -ge 3 ]; then
		echo -n "[${DELIM}] "
		echo "$@"
	fi
}

function launch_flux() {
	verify_launcher_avail flux
	# Unlike srun inside an salloc/sbatch job, "flux run" with no
	# sizing flags does not inherit the enclosing allocation's node
	# and task count; it defaults to a single task regardless of how
	# large the allocation actually is. FLUX_JOB_NNODES/FLUX_JOB_NTASKS
	# are set by run_mpi_binary()'s flux branch for exactly this
	# reason. Fall back to one task per node in this instance for
	# standalone use outside that harness.
	local nnodes="${FLUX_JOB_NNODES:-$(flux getattr size)}"
	local ntasks="${FLUX_JOB_NTASKS:-${nnodes}}"
	run_command flux run -N "${nnodes}" -n "${ntasks}" "$@"
}

function verify_launcher_avail() {
	if [ -z "$(type -P "${1}")" ]; then
		_error "Expected Job launcher $1 not found for ${LMOD_FAMILY_MPI}"
		exit 1
	fi
}

function set_variable() {
	local var=${1}

	if [[ -z ${!var} ]]; then
		export "${var}"="${2}"
		_debug "Setting env variable: ${var}=${2}"
	else
		_debug "Keeping existing env variable: ${var}=${!var}"
	fi
}

function run_command() {
	local launch_cmd
	printf -v launch_cmd '%q ' "$@"
	_info "Launch cmd = ${launch_cmd% } (family=${LMOD_FAMILY_MPI})"
	"$@"
}

function launch_impi() {
	local -a cmd
	if [ "${MODE}" == "native" ]; then
		if [ "${RM}" == "slurm" ]; then
			cmd=(srun --mpi=pmi2)
		elif [ "${RM}" == "openpbs" ]; then
			LAUNCHER="mpiexec.hydra"
			verify_launcher_avail "${LAUNCHER}"

			cmd=("${LAUNCHER}" -rmk pbs)
		elif [ "${RM}" == "flux" ]; then
			launch_flux "$@"
			return $?
		else
			_error "Unknown resource manager -> ${RM}"
			exit 1
		fi

		run_command "${cmd[@]}" "$@"
	else
		_error "Unsupported launch mode -> ${MODE} (${LMOD_FAMILY_MPI})"
		return 1
	fi
}

function launch_mpich() {
	local -a cmd
	if [[ "${RM}" == "slurm" ]]; then
		if [[ ${OHPC_MPI_LAUNCHERS} =~ "pmix" ]]; then
			verify_launcher_avail srun
			cmd=(srun --mpi=pmix)
		else
			verify_launcher_avail mpiexec.hydra
			cmd=(mpiexec.hydra -bootstrap slurm)
		fi
	elif [[ "${RM}" == "openpbs" ]]; then
		verify_launcher_avail mpiexec.hydra
		cmd=(mpiexec.hydra -rmk pbs)
	elif [[ "${RM}" == "flux" ]]; then
		# Flux's own default simple PMI-1/2 is sufficient for mpich
		# and its derivatives (mvapich2); no bootstrap flags needed.
		launch_flux "$@"
		return $?
	else
		_error "Unknown resource manager -> ${RM}"
		exit 1
	fi

	run_command "${cmd[@]}" "$@"

	return $?
}

function setupenv_mvapich() {
	# quiet job launch nonhomogeneous warning
	set_variable MV2_SUPPRESS_JOB_STARTUP_PERFORMANCE_WARNING 1
}

function launch_mvapich2() {
	local -a cmd
	if [ "${MODE}" == "native" ]; then

		setupenv_mvapich

		if [ "${RM}" == "flux" ]; then
			# Flux's own default simple PMI-1/2 is sufficient for
			# this mpich derivative; no mpiexec.hydra bootstrap
			# wrapper needed, same as the plain mpich case.
			launch_flux "$@"
			return $?
		else
			LAUNCHER="mpiexec.hydra"
			verify_launcher_avail "${LAUNCHER}"

			if [ "${RM}" == "slurm" ]; then
				cmd=("${LAUNCHER}" -bootstrap slurm)
			elif [ "${RM}" == "openpbs" ]; then
				cmd=("${LAUNCHER}" -rmk pbs)
			else
				_error "Unknown resource manager -> ${RM}"
				exit 1
			fi
		fi

		run_command "${cmd[@]}" "$@"

		local _status=$?

		return "${_status}"
	else
		_error "Unsupported launch mode -> ${MODE} (${LMOD_FAMILY_MPI})"
		return 1
	fi
}

function setupenv_openmpi() {
	# disable plugin loader warnings if not not available
	set_variable OMPI_MCA_mca_base_component_show_load_errors 0
	set_variable PMIX_MCA_mca_base_component_show_load_errors 0

	# no tm module avail when using slurm
	if [[ "${RM}" == "slurm" ]]; then
		set_variable OMPI_MCA_ras ^tm
		set_variable OMPI_MCA_ess ^tm
		set_variable OMPI_MCA_plm ^tm
	fi

	# use romio for openmpi3
	if [[ "${LMOD_FAMILY_MPI}" == "openmpi3" ]]; then
		set_variable OMPI_MCA_io romio314
	fi
}

function launch_openmpi() {
	local -a cmd

	setupenv_openmpi

	if [[ "${RM}" == "slurm" ]]; then
		if [[ "${OHPC_MPI_LAUNCHERS}" =~ "pmix" ]]; then
			verify_launcher_avail srun
			cmd=(srun --mpi=pmix)
		else
			verify_launcher_avail mpirun
			cmd=(mpirun)
		fi
	elif [[ "${RM}" == "openpbs" ]]; then
		verify_launcher_avail mpiexec
		cmd=(mpiexec -x LD_LIBRARY_PATH --prefix "${MPI_DIR}" --hostfile "${PBS_NODEFILE}")
	elif [[ "${RM}" == "flux" ]]; then
		# Unlike mpich and its derivatives, OpenMPI does not correctly
		# bootstrap over Flux's default simple PMI-1/2: it detects the
		# PMI environment but falls back to running disconnected
		# singleton processes instead of one real parallel job.
		# -o pmi=pmix requests flux-pmix's shell plugin instead.
		launch_flux -o pmi=pmix "$@"
		return $?
	else
		_error "Unknown resource manager -> ${RM}"
		exit 1
	fi

	run_command "${cmd[@]}" "$@"

	return $?
}

DELIM=prun

# Parse command-line args

while getopts "hv" opt; do
	case "${opt}" in
	h)
		usage
		;;
	v)
		LOGLEVEL=3
		;;
	\?)
		echo "Invalid option: -${OPTARG}" >&2
		exit 1
		;;
	esac
done
shift $((OPTIND - 1))

if [ $# -lt 1 ]; then
	usage
fi

parse_config_file

EXEC=$1
shift

# Resource manager detection. If not specified from the config file, determine
# from locally installed packages. Like Highlander, we assume there can be only one.

if [ -z "${RM}" ]; then
	[[ -s /etc/pbs.conf ]] && RM=openpbs
	[[ -s /etc/slurm/slurm.conf ]] && RM=slurm
	[[ -s /var/run/slurm/conf/slurm.conf ]] && RM=slurm
	# Unlike slurm/openpbs, flux has no single well-known system config
	# file to check for. prun always runs from inside an already
	# active flux job (submitted via "flux batch"/"flux run"), where
	# the flux shell has already set FLUX_URI for job tasks, so check
	# that instead.
	[[ -n "${FLUX_URI}" ]] && RM=flux
fi

if ! [ -x "$(command -v "${EXEC}")" ]; then
	if [ ! -x "${EXEC}" ]; then
		_error "Unable to access executable -> ${EXEC}"
		exit 1
	fi
fi

if [ -z "${LMOD_FAMILY_MPI}" ]; then
	_error "LMOD_FAMILY_MPI environment variable must be set to desired MPI runtime"
	exit 1
fi

PRIMARY_HOST=$(hostname -s)
_info "Master compute host = ${PRIMARY_HOST}"
_info "Resource manager = ${RM}"

if [ "${LMOD_FAMILY_MPI}" == "impi" ]; then
	launch_impi "${EXEC}" "$@"
elif [ "${LMOD_FAMILY_MPI}" == "mpich" ]; then
	launch_mpich "${EXEC}" "$@"
elif [ "${LMOD_FAMILY_MPI}" == "mvapich2" ]; then
	launch_mvapich2 "${EXEC}" "$@"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi" ]; then
	launch_openmpi "${EXEC}" "$@"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi3" ]; then
	launch_openmpi "${EXEC}" "$@"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi4" ]; then
	launch_openmpi "${EXEC}" "$@"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi5" ]; then
	# There have been errors where an application argument was interpreted by mpirun directly
	launch_openmpi "--" "${EXEC}" "$@"
else
	_error "Unsupported or unknown MPI family -> ${LMOD_FAMILY_MPI}"
	exit 1
fi
