#!/usr/bin/env perl
# +-======-+ 
#  Copyright (c) 2003-2007 United States Government as represented by 
#  the Admistrator of the National Aeronautics and Space Administration.  
#  All Rights Reserved.
#  
#  THIS OPEN  SOURCE  AGREEMENT  ("AGREEMENT") DEFINES  THE  RIGHTS  OF USE,
#  REPRODUCTION,  DISTRIBUTION,  MODIFICATION AND REDISTRIBUTION OF CERTAIN 
#  COMPUTER SOFTWARE ORIGINALLY RELEASED BY THE UNITED STATES GOVERNMENT AS 
#  REPRESENTED BY THE GOVERNMENT AGENCY LISTED BELOW ("GOVERNMENT AGENCY").  
#  THE UNITED STATES GOVERNMENT, AS REPRESENTED BY GOVERNMENT AGENCY, IS AN 
#  INTENDED  THIRD-PARTY  BENEFICIARY  OF  ALL  SUBSEQUENT DISTRIBUTIONS OR 
#  REDISTRIBUTIONS  OF THE  SUBJECT  SOFTWARE.  ANYONE WHO USES, REPRODUCES, 
#  DISTRIBUTES, MODIFIES  OR REDISTRIBUTES THE SUBJECT SOFTWARE, AS DEFINED 
#  HEREIN, OR ANY PART THEREOF,  IS,  BY THAT ACTION, ACCEPTING IN FULL THE 
#  RESPONSIBILITIES AND OBLIGATIONS CONTAINED IN THIS AGREEMENT.
#  
#  Government Agency: National Aeronautics and Space Administration
#  Government Agency Original Software Designation: GSC-15354-1
#  Government Agency Original Software Title:  GEOS-5 GCM Modeling Software
#  User Registration Requested.  Please Visit http://opensource.gsfc.nasa.gov
#  Government Agency Point of Contact for Original Software:  
#  			Dale Hithon, SRA Assistant, (301) 286-2691
#  
# +-======-+ 
#============================================================================
# name - esma_mpirun
# purpose - Wrapper for the mpirun command; filters out command-line flags
#           which are not applicable for the loaded version of mpi, and loads
#           other flags required for the operating environment.
#
# note -
# 1. Any flag which has not been identified as potentially needing 
#    to be filtered is passed directly through to the mpirun command.
# 2. Specifically, the -h (or -help) flag is passed through to the
#    mpirun command.
#
# revision history
# 02Oct2008  Stassi     Initial version of code.
# 23Jun2009  Stassi     Add $xtraflags for pleiades nodes
# 03May2010  Todling    Quick fix to work under mvapich2
# 13May2010  Stassi     Extended fix to work on NAS nodes as well as NCCS
# 27Jul2010  Stassi     Modified to work with either intel-mpi or mvapich
# 18May2011  Stassi     Capture status of cmd and die if not success
#============================================================================
use strict;
use warnings;

# global variables
#-----------------
my ($siteID, $node, $mpi_type, $perhost);
my ($progx, @mpiARGs, @progARGs);
my ($mpicmd, $xtraflags);
my ($scriptname);

# main program
#-------------
{
    init();
    get_siteID();
    get_mpi_type();

    parse_command_line();
    which_mpi_cmd();
    get_xtraflags();

    run_mpicmd();
}

#============================================================================
# name - init
# purpose - get runtime inputs
#
# notes:
# 1. -inherit_limits flag only valid for scali mpi
# 2. -perhost flag only valid for intel mpi
#============================================================================
sub init {
    use Getopt::Long;
    Getopt::Long::Configure("pass_through");
    my ($inherit_limits_flag);

    # get runtime flags
    #------------------
    GetOptions("inherit_limits" => \$inherit_limits_flag,
               "perhost=i"      => \$perhost );

    $scriptname = $0;
}

#============================================================================
# name - get_siteID
# purpose - set global variable, $siteID
#============================================================================
sub get_siteID {

    $siteID = "other";
    $node = `uname -n`;

    if    ( $node =~ /^pfe\d*$/       ) { $siteID = "nas" }
    elsif ( $node =~ /^p4fe\d*$/      ) { $siteID = "nas" }
    elsif ( $node =~ /^r\d+i\d+n\d+$/ ) { $siteID = "nas" }
        
    elsif ( $node =~ /^cfe/      )      { $siteID = "nccs" }
    elsif ( $node =~ /^dali/     )      { $siteID = "nccs" }
    elsif ( $node =~ /^discover/ )      { $siteID = "nccs" }
    elsif ( $node =~ /^borg/     )      { $siteID = "nccs" }

    print "$scriptname: node = $node\n";
    print "$scriptname: siteID = $siteID\n";
    return;
}

#============================================================================
# name - get_mpi_type
# purpose - try to determine the type of mpi library
#============================================================================
sub get_mpi_type {
    my ($MPIHOME);

    $mpi_type = "unknown";

    # check MPI_ROOT env variable
    #----------------------------
    if    ( $ENV{"M_MPI_ROOT"} )             { $mpi_type = "mvapich" }
    elsif ( $ENV{"I_MPI_ROOT"} )             { $mpi_type = "intel"   }

    # check $MPIHOME env variable
    #----------------------------
    elsif ( $MPIHOME = $ENV{"MPIHOME"} ) {
        if    ( $MPIHOME =~ m[mvapich] )     { $mpi_type = "mvapich" }
        elsif ( $MPIHOME =~ m[/intel/mpi/] ) { $mpi_type = "intel"   }
    }

    # check $MPI_HOME env variable
    #-----------------------------
    elsif ( $MPIHOME = $ENV{"MPI_HOME"} ) {
        if    ( $MPIHOME =~ m[mvapich] )     { $mpi_type = "mvapich" }
        elsif ( $MPIHOME =~ m[/intel/mpi/] ) { $mpi_type = "intel"   }
    }

    print "$scriptname: mpi_type = $mpi_type\n";
    return;
}

#============================================================================
# name - parse_command_line
# purpose - extract the executable program from the argument list;
#           divide all other arguments between those which belong to the
#           mpi command and those which belong to the executable program
#
# note -
# 1. arguments which precede the prog are @mpiARGs
# 2. arguments which follow the prog are @progARGs
#============================================================================
sub parse_command_line {
    my ($found, $num, $arg);

    # extract executable program from argument list
    #----------------------------------------------
    $progx = undef;
    $found = 0;

    $num = scalar(@ARGV);
    foreach (1..$num) {
        $arg = shift @ARGV;
        unless ($found) {
            if (-x $arg) {
                $progx = $arg;
                $found = 1;
                next;
            }
            push @mpiARGs, $arg;
            next;
        }
        push @progARGs, $arg
    }
    die ">> Error << no executable program found in mpi command" unless $progx;
}

#============================================================================
# name - which_mpi_cmd
# purpose - determine which mpi command to use
#============================================================================
sub which_mpi_cmd {
    my ($progname, $status);

    # check for existence of mpirun or mpiexec command
    #-------------------------------------------------
    $mpicmd = "notfound";

    foreach $progname ( "mpirun_rsh", "mpirun", "mpiexec" ) {
        $status = system "which $progname >& /dev/null";
        next if $status;

        chomp($mpicmd = `which $progname`);
        last;
    }
    die ">>> Error <<< Cannot find command to run mpi" unless -e $mpicmd;

    return;
}

#============================================================================
# name - get_xtraflags
# purpose - set other flags needed by the node where executing.
#
# node types at NCCS (info from D.Kokron):
# - borga      => Dempsey node ..... 4 core/node
# - borg[bcde] => Woodcrest node ... 4 core/node
# - borg[fghi] => Harpertown node .. 8 core/node 
#
# notes
# 1. node-specific flags for different borg node types have been removed
# 2. have not updated node info with addition of nehalem nodes
#============================================================================
sub get_xtraflags {
    my ($s1, $pbs_nodefile, $numnodes);
    $xtraflags = "";

    # flags needed for mvapich mpi
    #-----------------------------
    $xtraflags .= " -ssh -hostfile \$PBS_NODEFILE" if $mpi_type eq "mvapich";

    # flags needed for intel mpi
    #---------------------------
    if ($mpi_type eq "intel") {

        # nccs site
        #----------
        if ($siteID eq "nccs") {
            unless ($perhost) {
                if ($node =~ "borg") {
                    $s1 = substr($node,4,1);
                    $perhost = 4 if $s1 ge "a" and $s1 le "e";
                }
            }
        }

        # default perhost value
        #----------------------
        $perhost = 8 unless $perhost;

        # nas site
        #---------
        if ($siteID eq "nas") {

            if ($mpi_type eq "intel") {
                $pbs_nodefile = $ENV{"PBS_NODEFILE"};
                chomp($numnodes = `sort -u $pbs_nodefile | wc -l`)
                    if $pbs_nodefile;
                $numnodes = 1 unless $numnodes;
    
                $xtraflags .= " --totalnum=$numnodes";
                $xtraflags .= " --file=$pbs_nodefile";
                $xtraflags .= " --rsh=ssh";
                $xtraflags .= " -perhost $perhost";
            }
        }
        $xtraflags .= " -perhost $perhost";
    }

    # flags needed for scali mpi
    #---------------------------
    $xtraflags .= " -inherit_limits" if $mpi_type eq "scali";
}

#============================================================================
# name - run_mpicmd
# purpose - run the mpi command
#============================================================================
sub run_mpicmd {
    my ($flags, $cmd, $status);

    $flags = "";
    $flags .= " $xtraflags"     if $xtraflags;
    $flags .= " @mpiARGs"       if @mpiARGs;

    $cmd = "$mpicmd $flags $progx @progARGs";
    print "$cmd\n";
    $status = system $cmd;

    if ($status) {
        $status = $status>>8;
        die ">> Error << $cmd: status = $status;";
    }
}