From a6bccf55a47a65ca55cf0a771c8dc304f1a50fb6 Mon Sep 17 00:00:00 2001
From: Don Armstrong <don@donarmstrong.com>
Date: Wed, 8 Feb 2017 10:52:13 -0800
Subject: [PATCH] support slurm

---
 dqsub | 116 +++++++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 102 insertions(+), 14 deletions(-)

diff --git a/dqsub b/dqsub
index 1682d05..5b4b845 100755
--- a/dqsub
+++ b/dqsub
@@ -71,6 +71,11 @@ Account name to use
 Whether to join STDOUT and STDERR. On by default; disable with
 C<--nojoin>.
 
+=item B<--batch>
+
+Which batch system to use. If sbatch exists, assume it's slurm,
+otherwise, PBS.
+
 =item B<--debug, -d>
 
 Debug verbosity. (Default 0)
@@ -110,6 +115,7 @@ my %options = (nodes           => 1,
 
 GetOptions(\%options,
            'queue|q=s',
+           'batch=s',
            'interactive|I!',
            'nodes=i',
            'array=s',
@@ -117,7 +123,7 @@ GetOptions(\%options,
            'array_per_job|array-per-job=i',
            'array_slot_limit|array-slot-limit=i',
            'array_all_in_one_job|array-all-in-one-job!',
-           'ppn|processors-per-node=i',
+           'ppn|cpus|processors-per-node=i',
            'account|A=s',
            'join|J!',
            'mem|memory=s',
@@ -149,14 +155,40 @@ if ($options{interactive} and @ARGV) {
     push @USAGE_ERRORS,"Don't provide commands when you're asking for an interactive shell";
 }
 
+if (not defined $options{batch}) {
+    qx/which sbatch/;
+    if ($?) {
+        $options{batch} = 'slurm'
+    } else {
+       $options{batch} = 'pbs'
+    }
+}
+
+if ($options{batch} !~ /^pbs|slurm$/) {
+    push @USAGE_ERRORS,"Unsupported batch system '$options{batch}'; ".
+        "supported systems are pbs or slurm";
+}
+
 # pod2usage(join("\n",@USAGE_ERRORS)) if @USAGE_ERRORS;
 if (@USAGE_ERRORS) {
     print STDERR map {"$_\n"} @USAGE_ERRORS;
     exit 1;
 }
 
+
+my $JOB_SUBMITTER = 'qsub';
 # OK. Generate the options to qsub which we'll be using
-my @qsub_options = generate_qsub_options(\%options,\@ARGV);
+my @qsub_options;
+if ($options{batch} eq 'pbs') {
+    @qsub_options = generate_qsub_options(\%options,\@ARGV);
+    $JOB_SUBMITTER = 'qsub';
+} elsif ($options{batch} eq 'slurm') {
+    @qsub_options = generate_slurm_options(\%options,\@ARGV);
+    $JOB_SUBMITTER = 'sbatch';
+} else {
+   die "Unsupported batch system '$options{batch}'";
+}
+
 
 if ($options{interactive}) {
     print STDERR 'running: qsub '.join(' ',@qsub_options) if $DEBUG;
@@ -184,11 +216,6 @@ sub generate_qsub_options{
     if (defined $options->{queue} and length $options->{queue}) {
         push @qo,'-q',$options->{queue};
     }
-    if (defined $options->{dir}) {
-        push @qo,'-d',abs_path($options->{dir});
-    } else {
-        push @qo,'-d',getcwd;
-    }
     ## handle the -l options
     my @l;
     push @l, 'nodes='.$options->{nodes};
@@ -224,7 +251,52 @@ sub generate_qsub_options{
                           @{$args}[0..min($#{$args},2)]);
     }
     # join error and output streams
-    if ($options->{join}) {
+   if ($options->{join}) {
+        push @qo,'-j','oe';
+    }
+    return @qo;
+}
+
+sub generate_slrum_options{
+    my ($options,$args) = @_;
+    my @qo;
+    if (defined $options->{queue} and length $options->{queue}) {
+        push @qo,'-p',$options->{queue};
+    }
+    ## handle the -l options
+    if (defined $options->{account}) {
+        push @qo,'-A',$options->{account};
+    }
+    my %options_map =
+        (mem => 'mem',
+         ppn => 'cpus-per-task',
+         time => 'time',
+         cputime => 'cput',
+         host    => 'host',
+         pmem => 'pmem',
+         pvmem => 'pvmem',
+         max_file => 'file',
+        );
+    for my $k (keys %options_map) {
+        if ($options->{$k}) {
+            push @qo,'--'.$options_map{$k}.'=',$options{$k};
+        }
+    }
+    if ($options{mem}) {
+        push @qo,'--mem=',$options{mem};
+    }
+    if ($options->{interactive}) {
+        push @qo,'-I';
+    }
+    if ($options->{name}) {
+        push @qo,'-J',$options->{name};
+    } else {
+        push @qo,'-J',join('_',
+                           map {my $a = $_; $a =~ s/[^a-zA-Z0-9]*//g; $a;}
+                          @{$args}[0..min($#{$args},2)]);
+    }
+    # join error and output streams
+   if ($options->{join}) {
         push @qo,'-j','oe';
     }
     return @qo;
@@ -248,12 +320,12 @@ sub read_array_options{
 sub call_qsub {
     my ($qsub_options,$script) = @_;
     my $qsub_fh;
-    open $qsub_fh,'|-','qsub',@{$qsub_options},'-' or
-        die "Unable to start qsub: $!";
+    open $qsub_fh,'|-',$JOB_SUBMITTER,@{$qsub_options},'-' or
+        die "Unable to start $JOB_SUBMITTER: $!";
     print {$qsub_fh} $script or
-        die "Unable to print to qsub: $!";
+        die "Unable to print to $JOB_SUBMITTER: $!";
     close($qsub_fh) or
-        die "Unable to close qsub filehandle: $!";
+        die "Unable to close $JOB_SUBMITTER filehandle: $!";
 }
 
 sub write_qsub_script {
@@ -263,12 +335,28 @@ sub write_qsub_script {
     my $command = join(' ',map {$_ =~ /\s/?qq('$_'):$_} @{$arg});
         $script .= <<EOF;
 # this script was written by dqsub
+EOF
+    my $directory = getcwd;
+    if (defined $opt->{dir}) {
+        $directory = abs_path($opt->{dir});
+    }
+    # we really should be quoting this instead
+    $script .=<<EOF;
+# change to the working directory
+cd "$directory";
 EOF
     if (defined $opt->{array}) {
         my @subshell = ('','');
         my $array_opt = join("\n",@{$array});
         my $max_array = scalar @{$array};
         my $apjm1 = $opt->{array_per_job} - 1;
+        $script .= <<EOF;
+if [ -n "\$PBS_ARRAYID" ]; then
+   MYARRAYID="\${PBS_ARRAYID:=1}"
+else
+   MYARRAYID="\${SLURM_ARRAY_TASK_ID:=1}"
+fi;
+EOF
         if ($opt->{array_per_job} > 1) {
             # we will use subshells if there are more than one array
             # items per job
@@ -276,13 +364,13 @@ EOF
             $script .= <<EOF;
 for i in \$(seq 1 $opt->{array_per_job}); do
 # in some cases, the jobs aren't going to come out evenly. Handle that.
-JOBNUM=\$(( \${PBS_ARRAYID:=1} * $opt->{array_per_job} + \$i - $opt->{array_per_job} ))
+JOBNUM=\$(( \${MYARRAYID:=1} * $opt->{array_per_job} + \$i - $opt->{array_per_job} ))
 if [ \$JOBNUM -le $max_array ]; then 
 OPT=\$(sed -n -e "\$JOBNUM p"<<'_HERE_DOC_END_'
 EOF
         } else {
             $script .= <<EOF;
-OPT=\$(sed -n -e "\${PBS_ARRAYID:=1} p"<<'_HERE_DOC_END_'
+OPT=\$(sed -n -e "\${MYARRAYID:=1} p"<<'_HERE_DOC_END_'
 EOF
         }
         $script .= <<EOF;
-- 
2.39.2