#!/usr/bin/perl
#
# This script is a component of Warewulf,
# http://www.runlevelzero.net/greg/warewulf
#
#########################################################################
#
# Copyright (c) 2003, The Regents of the University of California, through
# Lawrence Berkeley National Laboratory (subject to receipt of any
# required approvals from the U.S. Dept. of Energy).  All rights reserved.
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# The GNU GPL Document can be found at:
# http://www.gnu.org/copyleft/gpl.html
#
#########################################################################
#
# Written and maintained by:
#       Greg Kurtzer, <gmkurtzer@lbl.gov>

use lib "/usr/lib/warewulf/", "/usr/lib64/warewulf/";
use Warewulf::Config;
use Warewulf::Status;
use Warewulf::Util;
use Getopt::Long;
use Sys::Hostname;
#use strict;

#Getopt::Long::Configure ("bundling");


$usage = "USAGE: $0 [options] [hostname]
  About:
    wwmpirun 
  Options:
   -np              Number of processes to run
   -m, --master     Include the number of jobs to run on master (no jobs
                    on master if option not used unless --lam defined)
   -l, --lam        Use LAM-MPI style lamboot/mpirun sequence
   -h, --help       Show this banner
  
  Enviornment Variables
   MPIRUN           The mpirun command to use (default `which mpirun`)
  
  This tool is part of the Warewulf cluster distribution
     http://warewulf-cluster.org/
";

GetOptions(
   'master|m=s'  => \$include_master,
   'np=s'        => \$mpirun_np,
   'lam|l'       => \$lam,
   'help|h'      => \$help,
);

if ( ! $mpirun_np ) {
   warn "ERROR: You must specify how many processes you wish to run!\n";
   $help = 1;
}

if ( ! @ARGV ) {
   warn "ERROR: You must specify a command to be run!\n";
   $help = 1;
}

if ( $help ) {
   print $usage;
   exit;
}

if ( $ENV{MPIRUN} ) {
   $mpirun = $ENV{MPIRUN};
} else {
   $mpirun = 'mpirun';
}

$CWD = $ENV{PWD};
$cpus = 0;

%config = &client_config();
%nodestatus = &node_status($config{'warewulf master'});
%filter = &users_filter();

if ( -d "/etc/warewulf/nodes" ) {
   %nodes = &node_config();
}

foreach ( sort keys %nodestatus ) {
   if ( ! $filter{$_} and defined %filter ) {
      next;
   }
   if ( $nodestatus{$_}{USERPROC} > 0 or $nodestatus{$_}{CPUUTIL} >= 5 ) {
     next;
   }
   if ( $nodestatus{$_}{NODESTATUS} eq 'READY' ) {
      push(@nodes_ready, $_);
   }
   $node_count++;
}

if ( $lam ) {
   if ( ! $include_master ) {
      $include_master = 1;
   }
   $machine_file = "$ENV{HOSTNAME} cpu=$include_master\n";
   $nodes_selected = "$ENV{HOSTNAME} ";
} elsif ( $include_master ) {
   for ($i=0;$i<$include_master;$i++) {
      $machine_file .= "$ENV{HOSTNAME}\n";
   }
   $nodes_selected = "$ENV{HOSTNAME} ";
}

$cpus += $include_master;

if ( @nodes_ready ) {
   foreach (sort @nodes_ready ) {
      if ( $lam ) {
         $machine_file .= "$nodestatus{$_}{NODENAME} cpu=$nodestatus{$_}{CPUCOUNT}\n";
         $cpus += $nodestatus{$_}{CPUCOUNT};
      } else {
         for ($i=0; $i<$nodestatus{$_}{CPUCOUNT}; $i++) {
            $machine_file .= "$nodestatus{$_}{NODENAME}\n";
            $cpus++;
         }
      }
      $nodes_selected .= "$nodestatus{$_}{NODENAME} ";
      if ( $cpus >= $mpirun_np ) {
         last;
      }
   }
}

$tmp_hostfile = "$ENV{HOME}/hostfile.";
$tmp_hostfile .= &generate_random_string(8);

if ( $mpirun_np > $cpus ) {
   die "Error: $cpus processors available and you want $mpirun_np!\n";
}

warn "Running job on:\n$nodes_selected\n\n";

open(HOSTFILE, "> $tmp_hostfile");
print HOSTFILE $machine_file;
close HOSTFILE;

chdir("$CWD");
if ( $lam ) {
   warn "CMD: lamboot $tmp_hostfile\n";
   system("lamboot $tmp_hostfile");
   warn "CMD: $mpirun -np $mpirun_np @ARGV\n";
   system("$mpirun -np $mpirun_np @ARGV");
   warn "CMD: lamhalt\n";
   system("lamhalt");
} else {
   warn "CMD: $mpirun -hostfile $tmp_hostfile -np $mpirun_np @ARGV\n";
   system("$mpirun -hostfile $tmp_hostfile -np $mpirun_np @ARGV");
}
unlink("$tmp_hostfile");

exit;
