arx

#!/usr/bin/env perl
use strict;
use Getopt::Std;
use File::Basename; #for basename, dirname
use POSIX "sys_wait_h"; # mostly for the handy uname() function
use Cwd qw(abs_path cwd);
use Fcntl qw(:DEFAULT :seek); #mostly for the SEEK constants
my $usage=q{arx : array job executor in current working directory
Usage:
 arx <command> [options] arguments...
 Array job executor for SLURM (or for GNU parallel, when used with -P).
 Usage example, running tasks for all lines in ../samples.manifest with
 a concurrency of 20 and a config file specified for environment variables:
   arx sub -a 1- -j20 -m 24G -t 16:00:00 --cfg ../fq2cram_config.cfg \
     task_fq2cram.sh ../samples.manifest

 Or use -C option to execute each line of <cmdfile> as a task in an array job.
 Example for local run (using GNU parallel) with a maximum concurrency of 6:
    arx sub -P -j6 -C <cmdfile>

 Use -D for a "dry run" showing the commands that would be run.
 
 The concurency (maximum number of tasks to be executed at the same time) 
 should be set with -j or using %concurency suffix for -a/--array (as in SLURM)
 If not given, the default concurrency is set to 6 for local (parallel) runs and 
 20 for SLURM jobs.
};

## double-dash recognized options for SLURM (anything else is just passed to sbatch)
# -c, --cpus-per-task=<ncpus>
# -a, --array=<indexes>
# -t, --time=hh:mm:ss
# -p, --partition=<partition_name>
# -J, --job-name=<jobname>
# --mem=<size>[units] (-m also recognized as a shortcut by arx, even though sbatch sees it as --distribution)
# -o, --output=<filename_pattern> #default file name is "slurm-%A_%a.out", "%A" is replaced by the job ID and "%a" with the array index
# -e, --error=<filename_pattern>
# --mail-type, 
# --mail-user (=> -M shortcut added by arx, also auto enables '--mail-type END,FAIL')
# --cfg,--config => arx config file, can define environment variables to be passed to the job script

my @cmds = qw/status st run runline sub submit qsub watch/;
my %cmds; 
@cmds{@cmds}=();

die("$usage\n") unless exists($cmds{$ARGV[0]});

my $arxcmd=join(' ',$0, @ARGV);
my $cmd=shift(@ARGV);

my @ddash; # double dash arguments
my $i = 0;
my $cfgfile;
while ($i < @ARGV) {
  my $iscfg;
  if ($ARGV[$i] =~ /^--([^=]+)=(.*)/) {# '=' found, split on it and add the key-value pair
      my ($k,$v)=($1,$2);
      if ($k=~m/(cfg|config)/) {
        $cfgfile=$v;
      } else { push @ddash, [$k, $v]; }
      splice(@ARGV, $i, 1); # remove the current item from @ARGV
  } elsif ($ARGV[$i] =~ /^--(.+)/) {# no '=' is found but the argument starts with '--'
      my $k=$1; 
      my $r=1;
      $iscfg=($k=~m/(cfg|config)/);
      if ($i<$#ARGV && $ARGV[$i + 1]!~m/^-/) {
         if ($iscfg) { $cfgfile=$ARGV[$i + 1] }
                else { push @ddash, [$k, $ARGV[$i + 1]]; } ## assume the next argument is the value
         $r++;
      } else {  push @ddash, [ $k, ''] unless $iscfg }
      splice(@ARGV, $i, $r); # remove this argument and the next one from @ARGV
  } else {  $i++ } # just move to the next argument
}

getopts('MPCDJ:j:t:l:a:m:p:o:e:c:') || die($usage."\n");

my $JGRID=$Getopt::Std::opt_P ? 'par' :'slurm';
my $dryrun=$Getopt::Std::opt_D;
my $jpar=$Getopt::Std::opt_j;
#my $ge=lc($Getopt::Std::opt_g) || 'slurm'; # grid engine 
my $PWD=abs_path(cwd()); #from Cwd module
my $HOSTNAME = (&POSIX::uname)[1];

chomp($HOSTNAME);
$HOSTNAME=lc($HOSTNAME);

my $HOST=$HOSTNAME;

($HOST)=($HOST=~m/^([\w\-]+)/);
$ENV{HOST}=$HOST;
$ENV{MACHINE}=$HOST;
$ENV{HOSTNAME}=$HOSTNAME;

my $HOME=$ENV{HOME}; 
my $JOB=$ENV{'SLURM_JOBID'} || '0';
#my $JOB=$ENV{'JOB_ID'} || '0';
my $TASK=$ENV{'SLURM_ARRAY_TASK_ID'} || '0';
#$TASK=$ENV{'SGE_TASK_ID'} unless $TASK; ## support SGE?

my ($JARR, $JMEM, $JCPUS, $JTIME, $JPART, $JMAIL, $JMAILTYPE, $JNAME, $JLOG, $JERRLOG);
parseConfigFile($cfgfile) if $cfgfile;
my @ddpass; ## pass-through options
foreach my $dd (@ddash) {
  if ($$dd[0] eq 'array') { $JARR=$$dd[1] 
  } elsif ($$dd[0] eq 'cpus-per-task') {  $JCPUS=$$dd[1] 
  } elsif ($$dd[0] eq 'mem') {  $JMEM=$$dd[1] 
  } elsif ($$dd[0] eq 'time') {  $JTIME=$$dd[1] 
  } elsif ($$dd[0] eq 'partition') {  $JPART=$$dd[1] 
  } elsif ($$dd[0] eq 'mail-type') {  $JMAILTYPE=$$dd[1] 
  } elsif ($$dd[0] eq 'mail-user') {  $JMAIL=$$dd[1] 
  } elsif ($$dd[0] eq 'job-name') {  $JNAME=$$dd[1] 
  } elsif ($$dd[0] eq 'output') {  $JLOG=$$dd[1] 
  } elsif ($$dd[0] eq 'error') {  $JERRLOG=$$dd[1] }
  else { push @ddpass, "--$$dd[0]=$$dd[1]" } # pass-through unrecognized options
}

$JARR=$Getopt::Std::opt_a if $Getopt::Std::opt_a;
$JMEM=$Getopt::Std::opt_m if $Getopt::Std::opt_m;
$JPART=$Getopt::Std::opt_p if $Getopt::Std::opt_p;
$JCPUS=$Getopt::Std::opt_c if $Getopt::Std::opt_c;
my $mail=$Getopt::Std::opt_M;
if ($mail) {
  #$JMAIL=$Getopt::Std::opt_M; use the email from ~/.slurm/default
  $JMAILTYPE='END,FAIL' unless $JMAILTYPE;
}
$JNAME=$Getopt::Std::opt_J if $Getopt::Std::opt_J;
$JTIME=$Getopt::Std::opt_t if $Getopt::Std::opt_t;
$JLOG=$Getopt::Std::opt_o if $Getopt::Std::opt_o;
$JERRLOG=$Getopt::Std::opt_e if $Getopt::Std::opt_e;

my $cmdfile=$Getopt::Std::opt_C; # execute each line as a command

my $JMDIR="$HOME/_jobs";
my $logdir=$Getopt::Std::opt_l;
if ($cmd eq 'status') {
   print STDERR "Using grid engine: $JGRID\n";
   if ($JGRID eq 'slurm') {
      print STDERR "job $JOB.$TASK on $HOST\n";
   }
} elsif ($cmd =~ m/^sub/i) { # submit job
   if ($dryrun) {
     print STDERR "Using grid engine: $JGRID\n";
   }
   mkdir($JMDIR) unless $dryrun || -d $JMDIR;
   my @args=@ARGV; # default state
   my $nlines=0;
   if ($cmdfile) { # -C option : execute commands from file
      $cmdfile=shift(@ARGV);
      die("Error: command file $cmdfile not found!") unless $cmdfile && -f $cmdfile;
      unless($JARR) {
        $nlines=`wc -l < $cmdfile`;
        chomp($nlines);
        $JARR='1-'.$nlines;
      }
      @args=("$HOME/gscripts/task_linecmd.sh", $cmdfile);
   } else {
     die("Error: submit command requires a target script!\n") unless @ARGV>0;
     if ($JARR && -f $args[1]) { ## still an array job to be submitted
       ## index with linix
       my $nl=`linix $args[1]`;
       ($nlines)=($nl=~m/^(\d+)/);
     }
   }
   if ($JARR) {
     # default concurrency limit : how many tasks can run at the same time
     my $pdef=$JGRID eq 'par' ? 6: 20;
     my $plimit;
     if ($JARR=~s/%(\d+)$// ) {
       $plimit=$1;
     }
     $plimit=$jpar if $jpar && (!$plimit || $plimit>$jpar); # always prefer the lowest given value
     if ($JARR=~m/\-$/) { ## assume line-based job array is the argument
         die("Number of lines not determined.\n") unless $nlines;
         $JARR.=$nlines;
     }
     $plimit=$pdef unless $plimit;
     $JARR.='%'.$plimit if $plimit<$nlines;
   }
   ## whatever the case $args[0] must be a path to a script to submit
   if ( ! -f $args[0] ) {
      my $cpath=`which $args[0]`; 
      chomp($cpath);
      die("Error: script $args[0] not found.\n") unless $cpath && -f $cpath;
      $args[0]=$cpath;
   }
   unless ($JLOG) {
     $logdir='logs' unless $logdir;
     mkdir($logdir) unless $dryrun || -d $logdir; # all my array jobs write logs in a 'logs' subdir of cwd!
     $JLOG=$logdir.'/arx-%A_%a.tlog';
   }
   $JERRLOG = $JLOG if ($JLOG && !$JERRLOG);
   my ($job, $jobcmd);
   my ($sec, $min, $hour, $mday, $mon, $year) = localtime();
   my $jtime = sprintf("%02d%02d%02d_%02d%02d", $year % 100, $mon + 1, $mday, $hour, $min);
   my $parPID; # only populated in GNU parallel case with the parallel PID
   if ($JGRID eq 'slurm') {
     unshift(@args, '-a '.$JARR) if $JARR;
     unshift(@args, '--mem='.$JMEM) if $JMEM;
     unshift(@args, '-t '.$JTIME) if $JTIME;
     unshift(@args, '-p '.$JPART) if $JPART;
     unshift(@args, '-c '.$JCPUS) if $JCPUS;
     unshift(@args, '--mail-user='.$JMAIL) if $JMAIL; # or could be picked up from ~/.slurm/default
     unshift(@args, '--mail-type='.$JMAILTYPE) if $JMAILTYPE;
     unshift(@args, '-J '.$JNAME) if $JNAME;
     unshift(@args, '-o '.$JLOG) if $JLOG;
     unshift(@args, '-e '.$JERRLOG) if $JERRLOG;
     unshift(@args, @ddpass);
     $job=$dryrun ? 'dry-run' : submit(@args);
     $jobcmd="sbatch ".join(' ', @args);
     print STDERR "Job $job submitted:\n$jobcmd\n";
     $JMDIR.="/${jtime}.$job";
     mkdir($JMDIR) unless $dryrun;
   } else { # use GNU parallel (local parallelism)
     my $pID="$$";
     my $shell_script="parallel.arx$pID.sh";
     unlink($shell_script);
     my $pid_file="parallel.arx$pID.pid";
     unlink($pid_file);
     print STDERR "JARR=$JARR\n" if $dryrun;
     my ($ast, $aend)=($JARR=~m/^(\d+)\-(\d+)/);
     my ($apar)=($JARR=~m/%(\d+)$/);
     my $parprog=`which parallel`;chomp($parprog);
     $job=$pID;
     my $tpar='';
     $tpar="-j $apar" if $apar;
     $jobcmd="$parprog --delay .01 $tpar $args[0] $args[1] \{1\} '&>' $logdir/arx-${pID}_\{1\}.tlog ::: \{$ast..$aend\}";
     my $fparlog="parallel.arx$pID.log";
     if ($dryrun) {
       print STDERR "------ Dry-run ------:\n$arxcmd\n$jobcmd\n----------------\n";
     } else {
       open(PLOG, ">$fparlog") || die("Error creating $fparlog\n");
       print PLOG "$arxcmd\n";
       print PLOG "$jobcmd\n";
       print PLOG "## see $pid_file for GNU parallel's process ID\n";
       close(PLOG);
       #fork & exec
       defined(my $pid = fork) or die "Cannot fork: $!";
       if ($pid) { # parent process
          while (!-s $pid_file) {
            select(undef, undef, undef, 0.2); # wait 0.2s for the child process to write the PID to the file
          }
          open(my $fh, '<', $pid_file) or die "Cannot open PID file: $!";
          $parPID = <$fh>;
          close $fh;
          chomp($parPID); # PID of the parallel process
          print STDERR "Started parallel with PID=$parPID\n";
          $JMDIR.="/${jtime}.$pID";
          mkdir($JMDIR);
          open(PARID, ">$JMDIR/parallel.pid") || die("Error creating $JMDIR/parallel.pid\n");
          print PARID $parPID."\n";
          close(PARID);
       } else { # child process
          open(my $sh, '>', $shell_script) or die "Cannot open shell script file: $!";
          print $sh "$jobcmd >> $fparlog 2>&1 &\n";
          print $sh "echo \$! > $pid_file\n";
          close $sh;
          exec("env", "bash", $shell_script);
          exit 0;
       }
     }
   }
   if ($dryrun) {
     print STDERR "$JMDIR/j_cmd would contain this info:\n";
     print STDERR "$arxcmd\n";
     print STDERR "$jobcmd\n";
     print STDERR "#cwd: $PWD\n";
   } else {
     open(JINFO, ">$JMDIR/j_cmd") || die("Error creating $JMDIR/j_cmd\n");
     print JINFO "$arxcmd\n";
     print JINFO "$jobcmd\n";
     print JINFO "#cwd: $PWD\n";
     close(JINFO);
     $logdir=abs_path($logdir);
     symlink($PWD, "$JMDIR/j_cwd");
     if (!$cmdfile) {
       symlink("$logdir", "$JMDIR/j_logs");
     }
   }
   print STDERR "Submitted $JGRID job: $job (check $JMDIR)\n";
   if ($JGRID eq 'slurm') {
       #if ($mail) {
         my $jn=$JNAME ? "w_${JNAME}_$job" : "w_$job";
         @args=("-J $jn --mem 100 -c 1", "--dependency=afterany:$job -o j_$job.watcher.log --mail-type='END,FAIL'", __FILE__, "watch ${jtime}.$job");
         my $j=$dryrun ? 'w-dryrun' : submit(@args);
         print STDERR " ..submitted watcher job $j :\nsbatch ".join(' ', @args)."\n";
       #}
   }
} elsif ($cmd eq 'watch') {
  # get summary info for watched job
  my $tjob=$ARGV[0]; # YYMMDD_hhmm.job#  
  $JMDIR.="/$tjob";
  my ($job)=($tjob=~m/(\d+)$/);
  print STDERR "Job $tjob ended (see $JMDIR/j_cmd)\n";
  chdir($JMDIR) || die("Error chdir($JMDIR)\n");
  ##NOTE: task scripts SHOULD create logs: ~/_jobs/<jobID>/arx-%A_<taskID>.tlog
  my @tlogs=glob("j_logs/*[-_.]${job}_[0-9]*.*log");
  #print STDERR "tlogs array = ".join(', ',@tlogs)."\n";
  if (@tlogs>0) { #collect summary for all tasks
     my @terr; # tasks that failed
     foreach my $tlf (@tlogs) {
        my ($tid)=($tlf=~m/(\d+)\.t?log$/);
        open(TLOG, "<$tlf") || die("Error opening task log file $JMDIR/$tlf\n");
        seek(TLOG, -100, SEEK_END);
        my @last=<TLOG>;
        my $l=$last[-1]; #last line, should be shorter than 100 chars!
        push(@terr, $tid) unless ($l=~m/\b(done|ok|succeeded|finished)/i);
        close(TLOG);
     }
   open(TSUM, '>t_sum.txt') || die("Error creating $JMDIR/t_sum.txt\n");
   my $total_msg="total_tasks\t".scalar(@tlogs)."\n";
   print TSUM $total_msg;
   print STDERR $total_msg;
   if (@terr>0) {
      my $errtasks="error_tasks_list\t".join(',',@terr)."\n";
      print TSUM $errtasks;
      print STDERR $errtasks
   }
   close(TSUM);
  }
} elsif ($cmd eq 'run') { ## TODO : check implementation
  print "running task $TASK in ".$ENV{PWD}." with args: ".join(" ",@ARGV)."\n";
  while (<>) {
    $cmd=$_;
    last if $.==$TASK;
    $cmd='';
  }
  if (length($cmd)) {
   chomp($cmd);$cmd=~s/\s+$//;
   my $cl=length($cmd);
   print " executing command of len $cl:\n<$cmd>\n";
   system($cmd);
  } else {
    print " error: no cmd found for $TASK!\n";
  }
}

#------------------
sub submit {
  my $args=join(" ", @_);
  #print STDERR ">> submitting:\nsbatch $args\n";
  my $s=`sbatch $args`;
  my ($jid)=($s=~m/ubmitted batch job (\d+)/);
  die("Error: errror at sbatch $args ") unless $jid;
  return $jid;
}

sub read_line { # use the lidx
    my ($reqlno, $ftxt, $fidx) = @_;
    open my $idx_fh, '<', $fidx or die "Could not open '$fidx': $!";
    my $index_offset = ($reqlno - 1) * 8;
    seek $idx_fh, $index_offset, SEEK_SET;
    read $idx_fh, my $record, 8;
    close $idx_fh;
    my ($lno, $offset, $clen) = unpack('S I S', $record);
    die "Error: line number mismatch in index file!\n" if !defined($lno) || $lno != $reqlno;
    open my $txt_fh, '<', $ftxt or die "Could not open '$ftxt': $!";
    seek($txt_fh, $offset, SEEK_SET) || die "Error seeking to $offset in $ftxt : $!\n";
    my $line='------------------------------------------------------------------';
    my $brd=read($txt_fh, $line, $clen); 
    die "Error reading line $reqlno from $ftxt : $!\n" if !defined($brd);
    close $txt_fh;
    print $line."\n";
}

sub parseConfigFile {
 my ($fcfg)=@_; # consider sections? [fq2cram]
 open(F, $fcfg) || die("Error opening config file $fcfg\n");
 while(<F>) {
   chomp;
   my ($k, $v)=split(/\s?=\s?/, $_, 2);
   $ENV{$k}=$v if $k; ## --export=ALL should pick these up?
 }
 close(F)
}