#!/bin/sh
#
# Copyright 2003-2019 Intel Corporation.
# 
# This software and the related documents are Intel copyrighted materials, and
# your use of them is governed by the express license under which they were
# provided to you (License). Unless the License provides otherwise, you may
# not use, modify, copy, publish, distribute, disclose or transmit this
# software or the related documents without Intel's prior written permission.
# 
# This software and the related documents are provided as is, with no express
# or implied warranties, other than those that are expressly stated in the
# License.
#

tempdir="/tmp"
if [ -n "$TMPDIR" ]; then
    tempdir="$TMPDIR"
fi
if [ -n "$I_MPI_TMPDIR" ]; then
    tempdir="$I_MPI_TMPDIR"
fi

np_boot=
username=`whoami`
rc=0

if [ -z "$I_MPI_ROOT" -a -z "`uname -m | grep 1om`" ] ; then
    . ${0%/*}/mpivars.sh ""
fi

##### mpirun detection #####
export I_MPI_MPIRUN="mpirun"
############################

#############################################
#######  Job Scheduler autodetection  #######
#############################################
# PBS
if [ -n "$PBS_ENVIRONMENT" -a -z "$I_MPI_HYDRA_RMK" ] ; then
    export I_MPI_HYDRA_RMK=pbs
elif [ -n "$LSB_JOBID" -a -z "$I_MPI_HYDRA_RMK" ]; then
    export I_MPI_HYDRA_RMK=lsf
fi
if [ -z "$I_MPI_HYDRA_BOOTSTRAP" -a -z "$I_MPI_HYDRA_BOOTSTRAP_EXEC" ]; then 
    # SLURM
    if [ -n "$SLURM_JOBID" ]; then
        export I_MPI_HYDRA_BOOTSTRAP=slurm
    # LoadLeveler
    elif [ -n "$LOADL_HOSTFILE" ]; then
        export I_MPI_HYDRA_BOOTSTRAP=ll
    # SGE
    elif [ -n "$PE_HOSTFILE" ]; then
        export I_MPI_HYDRA_BOOTSTRAP=sge
    # Fujitsu NQS (Network Queuing System)
    elif [ -n "$ENVIRONMENT" -a -n "$QSUB_REQID" -a -n "$QSUB_NODEINF" ] ; then
        if [ -z "$I_MPI_HYDRA_HOST_FILE" ]; then
            export I_MPI_HYDRA_HOST_FILE=$QSUB_NODEINF
        fi
        if [ -x /usr/bin/plesh ] ; then
            export I_MPI_HYDRA_BOOTSTRAP=rsh
            export I_MPI_HYDRA_BOOTSTRAP_EXEC=/usr/bin/plesh
        fi
    fi
fi

# Slurm, LoadLeveler, LSF, SGE
if [ -n "$LOADL_HOSTFILE" -o \
     -n "$PE_HOSTFILE" ]; then
    # Create a host file
    if [ -n "$LOADL_HOSTFILE" ]; then
        mpiexec.hydra "$@" <&0
        rc=$?
    elif [ -n "$PE_HOSTFILE" ]; then
        machinefile="${tempdir}/sge_machinefile_$username.$$"
        > $machinefile
            while read line; do
            if [ -n "$line" ]; then
                host_name=`echo $line | sed -e "s/ .*//"`
                num_of_processes=`expr match "$line" '.* \([0-9]\+\) .*'`
                echo "$host_name:$num_of_processes" >> $machinefile
                np_boot=$(( np_boot + 1 ))
            fi
        done < $PE_HOSTFILE
        mpiexec.hydra -machinefile $machinefile "$@" <&0
        rc=$?
        rm -rf $machinefile
    fi
# YARN cluster
elif [ "x$I_MPI_YARN" = "xyes" -o "x$I_MPI_YARN" = "xenable" -o "x$I_MPI_YARN" = "xon" -o "x$I_MPI_YARN" = "x1" ]; then
    llamaMPIClient.py "$@" <&0
    rc=$?
# Netbatch
elif [ -n "$NB_PARALLEL_JOB_HOSTS" ]; then
    hosts_opt=`echo $NB_PARALLEL_JOB_HOSTS | tr ' ' ','`
    hosts_opt="-hosts `hostname -s`,$hosts_opt"
    mpiexec.hydra $hosts_opt "$@" <&0
    rc=$?
# PBS or ordinary job
else
    mpiexec.hydra "$@" <&0
    rc=$?
fi

#######################################
#######         Cleanup         #######
#######################################
cleanup=0
echo "$@" | grep "\-cleanup" >/dev/null 2>&1
if [ $? -eq 0 ]; then
    cleanup=1
else
    if [ -n "$I_MPI_HYDRA_CLEANUP" ]; then
        if [ "$I_MPI_HYDRA_CLEANUP" = "yes" -o "$I_MPI_HYDRA_CLEANUP" = "enable" -o "$I_MPI_HYDRA_CLEANUP" = "on" -o "$I_MPI_HYDRA_CLEANUP" = "1" ]; then
            cleanup=1
        fi
    fi
fi
if [ $cleanup -eq 1 ]; then
    mpicleanup -s
fi
exit $rc
