#!/usr/bin/perl -w
# +-======-+ 
#  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
#============================================================================
use strict;

# global variables
#-----------------
my ($node, $filteredflags, $xtraflags, $progx);
my (@mpiARGs, @progARGs, $mpicmd, $mpi_type);
my ($inherit_limits_flag);

# main program
#-------------
{
    my ($cmd, $flags);

    &init();
    &extract_prog();
    &filter_flags();   # must be called after &extract_prog for $mpi_type
    &get_xtraflags();  # must be called after &extract_prog for $mpi_type

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

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

#============================================================================
# name - init
# purpose - get runtime inputs; check for existence of mpirun command.
#
# note -
#============================================================================
sub init {
    use Getopt::Long;
    Getopt::Long::Configure("pass_through");

    my $status;

    # get node information
    #---------------------
    $node = `uname -n`;

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

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

    $status = system "which mpirun >& /dev/null";
    chomp($mpicmd = `which mpirun`) unless $status;

    unless (-e $mpicmd) {
        $status = system "which mpiexec >& /dev/null";
        chomp($mpicmd = `which mpiexec`) unless $status;
    }

    die ">>> Error <<< Cannot find mpirun or mpiexec command."
        unless (-e $mpicmd);
}

#============================================================================
# name - filter_flags
# purpose - check flags which have been identified as potentially needing
#           to be filtered.
#============================================================================
sub filter_flags {

    $filteredflags = "";
    if ($inherit_limits_flag) {
 	$filteredflags .= "-inherit_limits" if ($mpi_type eq "scali")
    }
}

#============================================================================
# name - get_xtraflags
# purpose - set other flags needed by the node where executing.
#
# notes (from D.Kokron):
# - borga      => Dempsey node ..... 4 core/node
# - borg[bcde] => Woodcrest node ... 4 core/node
# - borg[fghi] => Harpertown node .. 8 core/node 
#============================================================================
sub get_xtraflags {
    my ($s1, $pbs_nodefile, $numnodes);
    $xtraflags = "";

    # add -perhost flag for newer borg nodes
    #---------------------------------------
    $xtraflags .= " -perhost 8" if $node =~ "borg";

    # add flags needed on pleiades nodes
    #-----------------------------------
    if (&pleiades_node) {
        $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 8";
    }

    # add -inherit_limits flag for scali type mpi
    #--------------------------------------------
    unless ($filteredflags =~ "-inherit_limits") {
        $xtraflags .= " -inherit_limits" if ($mpi_type eq "scali");
    }
}

#============================================================================
# name - pleiades_node
# purpose - determine whether the node is a "pleiades" node
#
# return value
# =1 for pleiades nodes
# =0 if not a pleiades node
#============================================================================
sub pleiades_node {
    return 1 if $node =~ /^pfe\d*$/;
    return 1 if $node =~ /^p4fe\d*$/;
    return 1 if $node =~ /^r\d+i\d+n\d+$/;
    return 0;
}

#============================================================================
# name - extract_prog
# purpose - extract the executable program from the argument list.
#
# note -
# 1. unclaimed args which follow the prog are @progARGs
# 2. unclaimed args which precede the prog are @mpiARGs
#============================================================================
sub extract_prog {
    my ($found, $lastarg, $arg);

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

    $lastarg = $#ARGV;
    foreach (0..$lastarg) {
        $arg = shift @ARGV;
        unless ($found) {
            if (-x $arg) {
                $progx = $arg;
                $found = 1;
                &get_mpi_type();
                next;
            }
            push @mpiARGs, $arg;
            next;
        }
        push @progARGs, $arg
    }
    $progx = $ARGV[-1] unless $progx;
}

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

    $mpi_type = "unknown";
    $libs = `ldd $progx`;
    $mpi_type = "scali" if ($libs =~ "scali");
}