File _service:extract_file:prun of Package prun

File _service:extract_file:prun of Package prun (Revision b5d463149c1c6eabe6606a6dd752d4ab)

Currently displaying revision b5d463149c1c6eabe6606a6dd752d4ab , Show latest

#!/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 flux_run_prefix() {
	# 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}}"
	echo "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
}

cmd="${LAUNCHER} -rmk pbs $*"
		elif [ "${RM}" == "flux" ]; then
			verify_launcher_avail flux
			cmd="$(flux_run_prefix) $*"
		fi

_info "Launch cmd = ${cmd} (family=${LMOD_FAMILY_MPI})"
		${cmd}
	fi
}

_info "Launch cmd = ${cmd} (family=${LMOD_FAMILY_MPI})"
	${cmd}

return $?
}

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

function launch_mvapich2() {
	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.
			verify_launcher_avail flux
			cmd="$(flux_run_prefix) $*"
		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}"
			fi
		fi

_info "Launch cmd = ${cmd} (family=${LMOD_FAMILY_MPI})"
		${cmd}

local _status=$?

return "${_status}"
	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() {

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.
		verify_launcher_avail flux
		cmd="$(flux_run_prefix) -o pmi=pmix $*"
	else
		_error "Unknown resource manager -> ${RM}"
	fi

_info "Launch cmd = ${cmd} (family=${LMOD_FAMILY_MPI})"
	${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
ARGS=$*

# 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}" "${ARGS}"
elif [ "${LMOD_FAMILY_MPI}" == "mpich" ]; then
	launch_mpich "${EXEC}" "${ARGS}"
elif [ "${LMOD_FAMILY_MPI}" == "mvapich2" ]; then
	launch_mvapich2 "${EXEC}" "${ARGS}"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi" ]; then
	launch_openmpi "${EXEC}" "${ARGS}"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi3" ]; then
	launch_openmpi "${EXEC}" "${ARGS}"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi4" ]; then
	launch_openmpi "${EXEC}" "${ARGS}"
elif [ "${LMOD_FAMILY_MPI}" == "openmpi5" ]; then
	# There have been errors where one of the ARGS was interpreted by mpirun directly
	launch_openmpi "--" "${EXEC}" "${ARGS}"
else
	_error "Unsupported or unknown MPI family -> ${LMOD_FAMILY_MPI}"
fi