#!/usr/bin/perl

use File::Copy;
use strict;

# Execute a Vitess command pipe on a Grid Engine.

# This script is meant as a working template to submit a
# grid batch job using n_nodes of the grid, calling n_instances
# of parallel pipe instances on each node.
# The pipe commands are taken from a file generated by
# "Save Grid Command" within vitess.

# M. Fromme, HMI, jan 2008

###############################################################
### configuration
my $testmode = 0; 		# 0 for ordinary execution, 1 for tests

my $qsubcom = 'qsub';

my $batchhead =<<'EOS';
#!/bin/bash
#$ -S /bin/bash
#$ -cwd
EOS
### end configuration
###############################################################


# The input file named $filename is a text file
# with variable definitions of
# V the source module directory
# P the working directory
#  and command option intro strings for the
# Z random seed
# L logfile
# F output file
# and the command pipe.

# Grid execution has to deal with
# - different file systems, $V and $P may not be present on the grid
# - parallel execution of pipes

# Because execution files may be invisible on grid nodes, we copy them to the
# grid modules directory. We copy all files from $P to
# the $grid_work_dir which are accessed in the pipe, too.

# We assume n_instances to be greater 1, but == 1 works, too.
# We want to start parallel pipes.
# If the last command of the pipe specifies an output file (--F<a filename>)
# we have to replace <a filename> by temporary file names, perform the
# parallel execution, and we need to concatenate the temporary file
# to <a filename> afterwards.
# If the last pipe command specifies --Fno_file, the user probably wants to
# perform the pipe stages preceeding the last command in parallel,
# gather the output of those to temporary files, and feed results
# to the final command later as a concatenated input stream.
#
# Two modes of grid execution are possible, selected by n_nodes:
#    a) submit this script once to the batch system (n_nodes 1),
#       to be started on a single node,
#       and fire n_instances of pipes on that node
#    b) submit generated batch scripts n_nodes times to the batch system
# A is the easier one, because we have control on the work node
# and may gather the results after completion of all pipes.
# For B we need to know when the instances are finished.
# Because this is too complicated here we just submit n_nodes jobs
# then (only if not --Fno_file), replacing the final --F files
# and telling the user where and how to fetch them when all batch
# jobs are done.

# Example:
# I'm user JOE, my working dir is /net/home/JOE/vitess on the Linux
# server dirac at HMI.
# After proper building an instrument with Vitess I found my simulation
# does take so long I'd like to speed that up using the Grid Engine.
# So first I create the grd file /net/home/JOE/vitess/test.grd
# using the Vitess "Save Grid Command" menu. Contents of this file are
#   V=/net/home/JOE/vitess/MODULES
#   P=/net/home/JOE/vitess/FILES/SAN_DIR
#   Z=--Z
#   L=--L
#   F=--F
#   $V/source_Linux_x86_64 -S1 ...
# To file 4 pipes on a single node of the Grid Engine I call

# grid_command.pl test.grd 1 4 /home/JOE/vitmod /home/JOE/vitwork -l cpu=8 -l vf=1G

# I use /home/JOE because this filesystem is visible on the grid nodes, too.
# The script creates /home/JOE/vitmod/ and /home/JOE/vitwork/ if they do not exist.
# Because /net/home/JOE/vitess/MODULES/ is invisible for dirac compute nodes,
# all used *_Linux_x86_64 executables are copied to /home/JOE/vitmod, plus this
# script itself, and referenced files from /net/home/JOE/vitess/FILES/BEND_DIR
# plus test.grd are copied to /home/JOE/vitwork.
# The following batch command file /home/JOE/vitwork/jobxtttttttt is written, where
#    x for number of nodes, tttttttt unix time
#    -0 4 is special for single node execution, 4 pipes

#  #!/bin/bash
#  # Grid Engine Parameters
#  #$ -S /bin/bash
#  #$ -cwd
#  #$ -l cpu=8
#  #$ -l vf=1G
#  /home/JOE/vitmod/grid_command.pl /home/JOE/vitmod/test.grd -0 4 /home/JOE/vitmod /home/JOE/vitwork

# The job is submitted internally by
# qsub vitjob.bat

# When the Grid Engine starts grid_command.pl on a node, /home/JOE/vitmod/grid_command.pl fires 4 pipes
# on that node, and gathers the result when the pipes are done.

sub usage {
  print <<EOS;
usage:
grid_command.pl filename nodes instances mod_dir work_dir {batchopt batcharg}

  filename           name of a file generated by Vitess "Save Grid Command"
  nodes              number of grid nodes to be used,
                     if preceeded by - directly executes pipes
  instances          number of parallel pipes per node
  mod_dir            directory on compute nodes for vitess modules
                     will be created if necessary and parent dir exists
  work_dir           working directory on compute nodes
                     will be created like mod_dir
  batchopt batcharg  additional arguments for the batch system like -l cpu=8

Called with
  nodes >= 1   this script submits nodes batch jobs with $qsubcom
  nodes == -0  directly fires instances pipes
  nodes == -1  thus called from batch job 1
EOS
  exit;
}

my ($filename, $n_nodes, $n_instances,
    $grid_mod_dir, $grid_work_dir,
   %Batcharg) = @ARGV;

my $script = $0;

usage() unless -s $filename && $grid_work_dir ne '' && $grid_mod_dir ne '';

my ($fmodir, $acount, $fpdir, $seedopt, $logopt, $outopt,
    $pipe, $dir, $bdir, $nextseed, $logdir, $logi, @Logfiles, @Pipe);

my $now = time();

readGrdFile();

die "something's wrong with $filename" unless @Pipe;

if ($n_nodes =~ /^-/) {
  doPipes();
} else {
  generateBatchFile();
}

sub readGrdFile {
  open F, $filename;
  while (<F>) {
    next if /^\s*\#/;
    if (/^([VPLZF])=(.+)$/) {
      if ($1 eq 'V') {
	$fmodir = $2;
      } elsif ($1 eq 'P') {
	$fpdir = $2;
      } else {
	my $o = substr($2, 0, 3);
	if ($1 eq 'L') {$logopt = $o;} elsif ($1 eq 'F') {$outopt = $o;} else {$seedopt = $o;}
      }
    } else {
      chomp;
      next unless $_;
      @Pipe = split('\|', $_);
      last;
    }
  }
  close F;
}

sub generateBatchFile {

  my ($i, $ofn, $sfn, $gfn);

  die "unknown directory $fmodir in $filename:\n$_" unless -d $fmodir;
  die "unknown directory $fpdir in $filename:\n$_" unless -d $fpdir;

  # Create working directories for grid usage if necessary.
  foreach $dir ($grid_work_dir, $grid_mod_dir) {
    next if -d $dir;
    my $bdir = basedir($dir);
    die "parent directory of $dir does not exist" unless -d $bdir;
    mkdir $dir, 0765;
  }

  # Copy needed execution modules from $fmodir to $grid_mod_dir.
  # Extract names from pipe:
  foreach my $p (@Pipe) {
    $_ = $p;			# necessary, otherwise we change @Pipe
    $_ = (split)[0];
    die "unable to parse $_" unless /\/([^\/]+)$/;
    myCopy("$fmodir/$1", "$grid_mod_dir/$1", 0755);
  }

  # Copy this script to $fmodir to be sure it may be called on
  # a compute node of the grid.
  $script = "$ENV{PWD}/$script" unless $script =~ m@^/@;
  die "unable to parse script name $script" unless $script =~ /\/([^\/]+)$/;
  $sfn = "$grid_mod_dir/$1";
  myCopy($script, $sfn, 0755);

  # Copy the grd file to the working directory.
  $gfn = "$grid_work_dir/" . ($filename =~ /\/([^\/]+)$/ ? $1 : $filename);
  myCopy($filename, $gfn);

  # Copy all existing files from the working directory $fpdir,
  # which are acessed in the pipe, to $grid_work_dir.

  # Extract used files from pipe:
  foreach (@Pipe) {
    foreach (split) {
      myCopy("$fpdir/$1", "$grid_work_dir/$1") if /\$P\/(.+)$/;
    }
  }

  if ($n_nodes == 1) {
    submitBatchJob(0, $sfn, $gfn);
  } else {
    for ($i = 1; $i <= $n_nodes; $i++) {
      submitBatchJob($i, $sfn, $gfn);
    }
  }
}

sub submitBatchJob {
  my ($i, $sfn, $gfn) = @_;

  # Create the batch job.
  my $fn = "$grid_work_dir/job$i" . $now;
  open F, ">$fn";
  print F $batchhead;
  print F "\#\$ $_ $Batcharg{$_}\n" foreach keys %Batcharg;
  print F "$sfn $gfn -$i $n_instances $grid_mod_dir $grid_work_dir\n";
  close F;
  chmod 0755, $fn;

  # Submit the batch job.
  print "submitting $fn\n";
  system "cd $grid_work_dir; $qsubcom $fn";
}


sub doPipes {

  my ($i, $fn, $pid, @Pid, $lastfile, @Com);

  # If the pipes are done on a single node, we may gather results
  # of the pipes without the last command to files, and feed those
  # files to the last command in a final step.

  my $one_node_only = $n_nodes eq '-0';

  # Expand $P and $V in pipe commands.
  my @Pcopy;
  foreach (@Pipe) {
    @Com = ();
    foreach (split) {
      if (/^\$V\/(.+)/) {
	push @Com, "$grid_mod_dir/$1";   # expand $V
      } elsif (/^([^\$]+)\$P(.*)/) {
	push @Com, "$1$grid_work_dir$2"; # expand $P
      } else {
	push @Com, $_;
      }
    }
    push @Pcopy, join(' ', @Com);
  }

  # Does the last pipe command generate output?
  my $lastcom = $_ = pop @Pcopy;

  @Com = ();			# will be the last command without output option
  foreach (split) {
    if (/^$outopt(.+)/o) {
      $lastfile = $1;
    } else {
      push @Com, $_;
    }
  }

  if ($lastfile eq 'no_file' && $one_node_only) {
    # We should gather what comes before
    # erase $outopt (--F) from the 2. last command.
    $_ = pop @Pcopy;
    @Com = ();
    foreach (split) {
      push @Com, $_ unless /^$outopt/o;
    }
  } else {
    $lastcom = '';
  }

  # now @Com has the 2. last or last command, whatever is needed
  push @Pcopy, join(' ', @Com);

  my $resdir = "$grid_work_dir/res$n_nodes";
  my @Res;

  # Fire $n_instances pipes.

  for ($i = 1; $i <= $n_instances; $i++) {

    # Adopt pipe to
    # 1. use different logfiles
    # 2. use individual seeds

    @Pipe = ();
    foreach (@Pcopy) {
      @Com = ();
      foreach (split) {
	if (/^$seedopt.+/o) {
	  push @Com, $seedopt . nextSeed();
	} elsif (/^$logopt.+/o) {
	  push @Com, $logopt . nextLog();
	} else {
	  push @Com, $_;
	}
      }
      push @Pipe, join(' ', @Com);
    }

    # fork and exec
    my $resfile = "$resdir-$i.dat";
    $_ = join(' | ', @Pipe) . " $outopt$resfile";
    if ($testmode) {
      print "would do\n$_\n\n";
    } else {
      exec $_ unless ($pid = fork());
      push @Pid, $pid;
    }
    push @Res, $resfile;
  }

  # Wait for completion of pipes.
  my ($exitstat, $some_error);
  foreach (@Pid) {
    waitpid $_, 0 or die "Could not get status of created process: $!\n";
    $exitstat = $? / 256;
    $some_error = $exitstat if $exitstat != 0;
  }

  # Find part result files.
  my @R;
  if ($testmode) {
    @R = @Res;
  } else {
    foreach (@Res) {
      push @R, $_ if -s $_;
    }
  }

  if (@R) {
    if ($lastcom ne '') {
      $_ = 'cat ' . join(' ', @R) . " | $lastcom";
      if ($testmode) {
	print "would do\n$_\n";
      } else {
	system $_;
	# delete part results
	unlink $_ foreach @R;
      }
    } else {
      # concatenate parts
      $fn = $resdir . '-full.dat';
      concatenateFiles($fn, @R);
      print "result neutron data file $fn\n";
    }
  }

  # Concatenate log files.
  $fn = $resdir . '-log.dat';
  concatenateFiles($fn, @Logfiles);
  print "result log file $fn\n";
}

sub concatenateFiles {
  open OF, '>' . shift;
  foreach my $fn (@_) {
    open F, $fn or next;
    print OF $_ while <F>;
    close F;
    unlink $fn;
  }
  close OF;
}

sub nextSeed {
  if ($nextseed eq '') {
    # initialize distinct seed for all $n_nodes, $n_instances, and pipe commands.
    # $n_nodes comes coded as a negative value here, -0, -1, ..
    $nextseed = (1 - $n_nodes) * $n_instances * 1024;
  }
  $nextseed++;
}

sub nextLog {
  $logdir = "$grid_work_dir/pipelog$$-$now" if $logdir eq '';
  $logi++;
  $_ = "$logdir-$logi.log";
  push @Logfiles, $_;
  $_;
}

sub myCopy {
  my ($a, $b, $mode) = @_;
  return if $a eq $b;
  my @A = stat $a;
  return if 0 == @A;
  my @B = stat $b;
  copy($a, $b) if 0 == @B || $A[7] ne $B[7] || $A[9] > $B[9];
  chmod $mode, $b if $mode ne '';
}

sub basedir {
  $_[0] =~ /^(.+)\/([^\/]+)$/ ? $1 : '';
}

