#!/bin/sh
#
# Copyright (C) 2003-2017 Intel Corporation.  All Rights Reserved.
# 
# The source code contained or described herein and all documents
# related to the source code ("Material") are owned by Intel Corporation
# or its suppliers or licensors.  Title to the Material remains with
# Intel Corporation or its suppliers and licensors.  The Material is
# protected by worldwide copyright and trade secret laws and treaty
# provisions.  No part of the Material may be used, copied, reproduced,
# modified, published, uploaded, posted, transmitted, distributed, or
# disclosed in any way without Intel's prior express written permission.
# 
# No license under any patent, copyright, trade secret or other
# intellectual property right is granted to or conferred upon you by
# disclosure or delivery of the Materials, either expressly, by
# implication, inducement, estoppel or otherwise.  Any license under
# such intellectual property rights must be express and approved by
# Intel in writing.
#

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
    # LSF
    elif [ -n "$LSB_JOBID" ]; then
        export I_MPI_HYDRA_BOOTSTRAP=lsf
    # 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
