#! /bin/sh 
#$Id: mpijob.tmpl,v 5.3 1997/01/24 17:28:40 chang Exp $
#  This sample script is a wrapper for running MPI jobs under lsbatch.
#  Submit job by saying "bsub [otheroptions] -n k mpijob command line",
#  where k is the number of hosts to use; mpijob is the name of this script.
#
#  Note the variable LSB_HOSTS is assigned by lsbatch system when this script
#  is started by lsbatch.

$LSB_TRAPSIGS 

dryrun=0
verbose=0
device_tcp=1
device_mpl=0
lsf_requeue=0
#
# runit.tmpl module will be included upon installation
#
# Do not remove the following two lines!
#RUNIT_MODULE_BEGIN#
#
# runit_unix()
#
# This func takes the global var $COMMANDLINE,
# executes the command and sets the $exstat with the exit
# status of the command
#
runit_unix() {
    eval $COMMANDLINE
    exstat=$?
}

#
# runit_mpl()
#
# This func takes the global var $COMMANDLINE
# and then sets the var $exstat to 
#    (1) first value of $REQUEUE_EXIT_VALUES   if (poe couln't allocate nodes)
#    (2) exit status of the command            otherwise
#
runit_mpl() {
    POE_ERRFILE="/tmp/poejob.$$"
    tmpfiles="$tmpfiles $POE_ERRFILE"

    eval poe $COMMANDLINE 2>$POE_ERRFILE
    exstat=$?

    if [ -n "`grep 't allocate nodes for parallel' $POE_ERRFILE`" ] ; then
        if [ x"$LSB_EXIT_REQUEUE" = "x" ] ; then
	    true
        else 
	    lsf_requeue=1
	    # take the first value of the REQUEUE_EXIT_VALUES
	    if [ x"REQUEUE_EXIT_VALUES" = x ] ; then
	        echo "poejob: fatal error REQUEUE_EXIT_VALUES not set!" 1>&2
	    else
		exstat=`echo $REQUEUE_EXIT_VALUES|awk {print $1}`
            fi
        fi
    fi
}
#RUNIT_MODULE_END#

#
# Temporary files to be removed before exit
#
tmpfiles=


#
# Parse Arguments
#
while [ 1 -le $# ] ; do
    arg=$1
    case $arg in
	-tcp)
          shift
	  ;;
	-mpl)
	  device_mpl=1
	  device_tcp=0
	  shift
	  ;;
        -dryrun)
	  dryrun=1
	  shift
	  ;;
        -v)
          verbose=1
	  shift
          ;;
         *)
	  break
	  ;;
    esac
done

CMD=$1

#
# if no programs specified 
# ignore it as mpirun will complain about this.
#
if [ $# -gt 1 ] ; then
    progname=$2
else
    progname=$@
fi

COMMANDLINE="$@"

#
# Build absolute path for the program instance
#
PWD_TRIAL=`pwd | sed -e 's%/tmp_mnt/%/%g'`
if [ ! -d $PWD_TRIAL ] ; then 
    echo "Warning: your default path uses the automounter; this may"
    echo "cause some problems if you use other NFS-connected systems."
    PWD_TRIAL=`pwd`
fi

#
# The test for PWD_TRIAL usable is within the p4/execer/tcp test (since that
# is where it matters)
#
tail=`echo $progname | sed 's/\/.*//'` 
if  [ "$tail" = "" ] ; then
    #echo Absolute path, don't need to change anything
    true
else
    #echo Relative path
    progname="$PWD_TRIAL/$progname"
fi

if [ $verbose = 1 ] ; then
    echo Absolute path $progname
fi

#
# Generate hostfile
#
nhosts=0

#
# sanity check
#
if [ x"$LSB_HOSTS" = "x" ] ; then
    echo "mpijob is only used to run batch jobs"
    echo "please use bsub to submit a job"
    exit 1
fi
 
HOSTFILE=$HOME/.lsbatch/host$$.`hostname`
rm -f $HOSTFILE
touch $HOSTFILE
tmpfiles="$tmpfiles $HOSTFILE"

firsthost=1
for host in $LSB_HOSTS 
do
#
# Create procgroup file
#
# for the MPICH implementation
#
    if [ $device_tcp = 1 ] ; then
      if [ $firsthost = 1 ] ; then
        echo $host 0 $progname>> $HOSTFILE
        firsthost=0
      else
	echo $host 1 $progname>> $HOSTFILE
      fi
    else
      echo $host >> $HOSTFILE
    fi
done

if [ $verbose = 1 ] ; then
    echo PROCGROUP
    cat $HOSTFILE
fi

if [ $device_tcp = 1 ] ; then
    COMMANDLINE="$COMMANDLINE -p4pg $HOSTFILE"
else
    COMMANDLINE="$COMMANDLINE -machinefile $HOSTFILE"
fi    

#
# run MPI job and save exit status.
#
if [ $dryrun != 1 ] ; then
     if [ $device_tcp = 1 ] ; then
         runit_unix
     else
         runit_mpl
     fi
else
    echo $COMMANDLINE
fi

# Kill runaway processes
if [ $lsf_requeue = 1 ] ; then
    true
else
    if [ $device_tcp = 1 ] ; then
	if [ $dryrun != 1 ] ; then
          lsgrun -p -f $HOSTFILE zapit -9 mpijob >/dev/null 2>&1
          lsgrun -p -f $HOSTFILE zapit -9 $CMD mpijob >/dev/null 2>&1
	fi
    fi
fi

if [ x"$tmpfiles" != "x" ] ; then
    rm -f $tmpfiles
fi

exit $exstat
