dqsub

   1 #!/usr/bin/perl
   2 # dqsub submits jobs using qsub with better options
   3 # and is released under the terms of the GNU GPL version 3, or any
   4 # later version, at your option. See the file README and COPYING for
   5 # more information.
   6 # Copyright 2014 by Don Armstrong <don@donarmstrong.com>.
   7
   8
   9 use warnings;
  10 use strict;
  11
  12 use Getopt::Long;
  13 # use Pod::Usage;
  14
  15 =head1 NAME
  16
  17 dqsub - submits jobs using qsub with better options
  18
  19 =head1 SYNOPSIS
  20
  21 dqsub [options]
  22
  23  Options:
  24    --queue, -q Queue to use
  25    --interactive, -I call qsub interactively
  26    --nodes nodes to use
  27    --array array mode (one of 'chdir' or 'xargs' or '')
  28    --array-from file to read arrays from (default STDIN)
  29    --array-per-job number of array items to handle in each job (default 1)
  30    --array-all-in-one-job Run all of the array items in one job
  31    --ppn processors per node to use
  32    --mem memory to request
  33    --dir Directory to run the script in (default current directory)
  34    --account, -A Account name to use
  35    --join, -J join error and output streams (default)
  36    --name, -N Name of the job
  37    --precommand Optional command to run before each command
  38    --debug, -d debugging level (Default 0)
  39    --help, -h display this help
  40    --man, -m display manual
  41
  42 =head1 OPTIONS
  43
  44 =over
  45
  46 =item B<--array>
  47
  48 This describes how dqsub will generate array jobs.
  49
  50 If no B<--array> is given, then the command and any additional
  51 arguments given will be run using qsub.
  52
  53 If B<--array> is C<chdir>, then each line of the input given in
  54 B<--array-from> will be used as a directory and the command and any
  55 additional arguments given will run in each directory.
  56
  57 IF B<--array> is C<xargs>, then each line of the input given will be
  58 considered to be an additional argument which will be given to the
  59 command run in the current directory.
  60
  61 =item B<--array-from>
  62
  63 File to read array arguments from. If not provided, and B<--array> is
  64 given, arguments will be read from STDIN.
  65
  66 =item B<--account, -A>
  67
  68 Account name to use
  69
  70 =item B<--join, J>
  71
  72 Whether to join STDOUT and STDERR. On by default; disable with
  73 C<--nojoin>.
  74
  75 =item B<--batch>
  76
  77 Which batch system to use. If sbatch exists, assume it's slurm,
  78 otherwise, PBS.
  79
  80 =item B<--debug, -d>
  81
  82 Debug verbosity. (Default 0)
  83
  84 =item B<--help, -h>
  85
  86 Display brief usage information.
  87
  88 =item B<--man, -m>
  89
  90 Display this manual.
  91
  92 =back
  93
  94 =head1 EXAMPLES
  95
  96 dqsub
  97
  98 =cut
  99
 100 use IO::File;
 101 use Cwd qw(getcwd abs_path);
 102 use POSIX qw(ceil);
 103 use List::Util qw(min);
 104 use vars qw($DEBUG);
 105
 106 my %options = (nodes           => 1,
 107                ppn             => 2,
 108                mem             => '2G',
 109                debug           => 0,
 110                help            => 0,
 111                man             => 0,
 112                interactive     => 0,
 113                array_per_job   => 1,
 114                join            => 1,
 115               );
 116
 117 GetOptions(\%options,
 118            'queue|q=s',
 119            'batch=s',
 120            'interactive|I!',
 121            'nodes=i',
 122            'array=s',
 123            'array_from|array-from=s',
 124            'array_per_job|array-per-job=i',
 125            'array_slot_limit|array-slot-limit=i',
 126            'array_all_in_one_job|array-all-in-one-job!',
 127            'ppn|cpus|processors-per-node=i',
 128            'account|A=s',
 129            'join|J!',
 130            'mem|memory=s',
 131            'time|walltime=s','cputime|cput=s','host=s',
 132            'pmem|process_mem|process-mem=s',
 133            'pvmem|process_virtual_mem|process-virtiual-mem=s',
 134            'max_file|max-file|file=s',
 135            'precommand|pre-command|pre_command=s',
 136            'dir=s',
 137            'name=s',
 138            'debug|d+','help|h|?','man|m');
 139
 140 # pod2usage() if $options{help};
 141 # pod2usage({verbose=>2}) if $options{man};
 142
 143 $DEBUG = $options{debug};
 144
 145 my @USAGE_ERRORS;
 146 if (not @ARGV and not $options{interactive}) {
 147     push @USAGE_ERRORS,"You must provide a command to run";
 148 }
 149 if (defined $options{array} and $options{array} !~ /^(?:|chdir|xargs)$/i) {
 150     push @USAGE_ERRORS,"--array must be one of chdir, xargs or '' if provided";
 151     $options{array} = lc($options{array});
 152     if ($options{array} eq '') {
 153         $options{array} = undef;
 154     }
 155 }
 156 if ($options{interactive} and @ARGV) {
 157     push @USAGE_ERRORS,"Don't provide commands when you're asking for an interactive shell";
 158 }
 159
 160 if (not defined $options{batch}) {
 161     qx/which sbatch/;
 162     if ($? == 0) {
 163         $options{batch} = 'slurm'
 164     } else {
 165        $options{batch} = 'pbs'
 166     }
 167 }
 168
 169 if ($options{batch} !~ /^pbs|slurm$/) {
 170     push @USAGE_ERRORS,"Unsupported batch system '$options{batch}'; ".
 171         "supported systems are pbs or slurm";
 172 }
 173
 174 # pod2usage(join("\n",@USAGE_ERRORS)) if @USAGE_ERRORS;
 175 if (@USAGE_ERRORS) {
 176     print STDERR map {"$_\n"} @USAGE_ERRORS;
 177     exit 1;
 178 }
 179
 180
 181 my $JOB_SUBMITTER = 'qsub';
 182 # OK. Generate the options to qsub which we'll be using
 183 my @qsub_options;
 184 if ($options{batch} eq 'pbs') {
 185     @qsub_options = generate_qsub_options(\%options,\@ARGV);
 186     $JOB_SUBMITTER = 'qsub';
 187 } elsif ($options{batch} eq 'slurm') {
 188     @qsub_options = generate_slurm_options(\%options,\@ARGV);
 189     $JOB_SUBMITTER = 'sbatch';
 190 } else {
 191    die "Unsupported batch system '$options{batch}'";
 192 }
 193
 194
 195 if ($options{interactive}) {
 196     print STDERR 'running: '.$JOB_SUBMITTER.' '.join(' ',@qsub_options) if $DEBUG;
 197     exec($JOB_SUBMITTER,@qsub_options);
 198 } else {
 199     my @array = ();
 200     if ($options{array}) {
 201         @array = read_array_options(\%options) if $options{array};
 202         # the -t option gives the range of elements for an array job
 203         if ($options{array_all_in_one_job}) {
 204             $options{array_per_job} = scalar @array;
 205         } else {
 206             if ($options{batch} eq 'pbs') {
 207                 push @qsub_options,'-t';
 208             } else {
 209                 push @qsub_options,'-a';
 210             }
 211             push @qsub_options,'1-'. ceil(scalar @array / $options{array_per_job});
 212             if ($options{array_slot_limit}) {
 213                 $qsub_options[$#qsub_options] .= '%'.$options{array_slot_limit};
 214             }
 215         }
 216     }
 217     if ($options{batch} eq 'pbs') {
 218         push @qsub_options,'-';
 219     }
 220     call_qsub(\@qsub_options,write_qsub_script(\%options,\@ARGV,\@array));
 221 }
 222
 223 sub generate_qsub_options{
 224     my ($options,$args) = @_;
 225     my @qo;
 226     if (defined $options->{queue} and length $options->{queue}) {
 227         push @qo,'-q',$options->{queue};
 228     }
 229     ## handle the -l options
 230     my @l;
 231     push @l, 'nodes='.$options->{nodes};
 232     if (defined $options->{ppn}) {
 233         $l[$#l] .= ':ppn='.$options->{ppn};
 234     }
 235     if (defined $options->{account}) {
 236         push @qo,'-A',$options->{account};
 237     }
 238     my %l_options =
 239         (mem => 'vmem',
 240          time => 'walltime',
 241          cputime => 'cput',
 242          host    => 'host',
 243          pmem => 'pmem',
 244          pvmem => 'pvmem',
 245          max_file => 'file',
 246         );
 247     for my $k (keys %l_options) {
 248         if ($options->{$k}) {
 249             push @l,$l_options{$k}.'='.$options{$k};
 250         }
 251     }
 252     push @qo,'-l',join(',',@l) if @l;
 253     if ($options->{interactive}) {
 254         push @qo,'-I';
 255     }
 256     if ($options->{name}) {
 257         push @qo,'-N',$options->{name};
 258     } else {
 259         push @qo,'-N',join('_',
 260                            map {my $a = $_; $a =~ s/[^a-zA-Z0-9]*//g; $a;}
 261                           @{$args}[0..min($#{$args},2)]);
 262     }
 263     # join error and output streams
 264    if ($options->{join}) {
 265         push @qo,'-j','oe';
 266     }
 267     return @qo;
 268 }
 269
 270 sub generate_slurm_options{
 271     my ($options,$args) = @_;
 272     my @qo;
 273     if (defined $options->{queue} and length $options->{queue}) {
 274         push @qo,'-p',$options->{queue};
 275     }
 276     ## handle the -l options
 277     if (defined $options->{account}) {
 278         push @qo,'-A',$options->{account};
 279     }
 280     my %options_map =
 281         (mem => 'mem',
 282          ppn => 'cpus-per-task',
 283          time => 'time',
 284          cputime => 'cput',
 285          host    => 'host',
 286          pmem => 'pmem',
 287          pvmem => 'pvmem',
 288          max_file => 'file',
 289         );
 290     for my $k (keys %options_map) {
 291         if ($options->{$k}) {
 292             push @qo,'--'.$options_map{$k}.'='.$options{$k};
 293         }
 294     }
 295     if ($options{mem}) {
 296         push @qo,'--mem='.$options{mem};
 297     }
 298     if ($options->{interactive}) {
 299         push @qo,'-I';
 300     }
 301     if ($options->{name}) {
 302         push @qo,'-J',$options->{name};
 303     } else {
 304         push @qo,'-J',join('_',
 305                            map {my $a = $_; $a =~ s/[^a-zA-Z0-9]*//g; $a;}
 306                           @{$args}[0..min($#{$args},2)]);
 307     }
 308     return @qo;
 309 }
 310
 311 sub read_array_options{
 312     my ($options) = @_;
 313     my $fh = \*STDIN;
 314     if (defined $options->{array_from}) {
 315         $fh = IO::File->new(defined $options->{array_from}) or
 316             die "Unable to open $options->{array_from} for reading: $!";
 317     }
 318     my @array;
 319     for (<$fh>) {
 320         chomp;
 321         push @array,$_;
 322     }
 323     return @array;
 324 }
 325
 326 sub call_qsub {
 327     my ($qsub_options,$script) = @_;
 328     my $qsub_fh;
 329     open $qsub_fh,'|-',$JOB_SUBMITTER,@{$qsub_options} or
 330         die "Unable to start $JOB_SUBMITTER: $!";
 331     print {$qsub_fh} $script or
 332         die "Unable to print to $JOB_SUBMITTER: $!";
 333     close($qsub_fh) or
 334         die "Unable to close $JOB_SUBMITTER filehandle: $!";
 335 }
 336
 337 sub write_qsub_script {
 338     my ($opt,$arg,$array) = @_;
 339
 340     my $script = "#!/bin/bash\n";
 341     my $command = join(' ',map {$_ =~ /\s/?qq('$_'):$_} @{$arg});
 342         $script .= <<EOF;
 343 # this script was written by dqsub
 344 EOF
 345     # if there is a precommand, write it out
 346     if ($opt->{precommand}) {
 347         $script .=<<EOF;
 348 # this is the precommand _BEGIN_
 349 $opt->{precommand}
 350 # precommand _END_
 351 EOF
 352     }
 353     my $directory = getcwd;
 354     if (defined $opt->{dir}) {
 355         $directory = abs_path($opt->{dir});
 356     }
 357     # we really should be quoting this instead
 358     $script .=<<EOF;
 359 # change to the working directory
 360 cd "$directory";
 361 EOF
 362     if (defined $opt->{array}) {
 363         my @subshell = ('','');
 364         my $array_opt = join("\n",@{$array});
 365         my $max_array = scalar @{$array};
 366         my $apjm1 = $opt->{array_per_job} - 1;
 367         $script .= <<EOF;
 368 if [ -n "\$PBS_ARRAYID" ]; then
 369    MYARRAYID="\${PBS_ARRAYID:=1}"
 370 else
 371    MYARRAYID="\${SLURM_ARRAY_TASK_ID:=1}"
 372 fi;
 373 EOF
 374         if ($opt->{array_per_job} > 1) {
 375             # we will use subshells if there are more than one array
 376             # items per job
 377             @subshell = ('(',')');
 378             $script .= <<EOF;
 379 for i in \$(seq 1 $opt->{array_per_job}); do
 380 # in some cases, the jobs aren't going to come out evenly. Handle that.
 381 JOBNUM=\$(( \${MYARRAYID:=1} * $opt->{array_per_job} + \$i - $opt->{array_per_job} ))
 382 if [ \$JOBNUM -le $max_array ]; then
 383 OPT=\$(sed -n -e "\$JOBNUM p"<<'_HERE_DOC_END_'
 384 EOF
 385         } else {
 386             $script .= <<EOF;
 387 OPT=\$(sed -n -e "\${MYARRAYID:=1} p"<<'_HERE_DOC_END_'
 388 EOF
 389         }
 390         $script .= <<EOF;
 391 $array_opt
 392 _HERE_DOC_END_
 393 )
 394 EOF
 395         if ($opt->{array} eq 'chdir') {
 396             $script .= <<EOF;
 397 $subshell[0]
 398 cd "\$OPT";
 399 exec ${command};
 400 $subshell[1]
 401 EOF
 402         } else {
 403             $script .= <<EOF;
 404 $subshell[0]
 405 exec ${command} "\$OPT";
 406 $subshell[1]
 407 EOF
 408         }
 409         if ($opt->{array_per_job} > 1) {
 410             $script .= <<EOF
 411 fi;
 412 done;
 413 EOF
 414         }
 415     } else {
 416         $script .= <<EOF;
 417 # there's no array, so just executing the command with arguments
 418 exec $command;
 419 EOF
 420     }
 421     return $script;
 422 }
 423
 424
 425 __END__