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: qsub '.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             push @qsub_options,'-t','1-'. ceil(scalar @array / $options{array_per_job});
 207             if ($options{array_slot_limit}) {
 208                 $qsub_options[$#qsub_options] .= '%'.$options{array_slot_limit};
 209             }
 210         }
 211     }
 212     if ($options{batch} eq 'pbs') {
 213         push @qsub_options,'-';
 214     }
 215     call_qsub(\@qsub_options,write_qsub_script(\%options,\@ARGV,\@array));
 216 }
 217
 218 sub generate_qsub_options{
 219     my ($options,$args) = @_;
 220     my @qo;
 221     if (defined $options->{queue} and length $options->{queue}) {
 222         push @qo,'-q',$options->{queue};
 223     }
 224     ## handle the -l options
 225     my @l;
 226     push @l, 'nodes='.$options->{nodes};
 227     if (defined $options->{ppn}) {
 228         $l[$#l] .= ':ppn='.$options->{ppn};
 229     }
 230     if (defined $options->{account}) {
 231         push @qo,'-A',$options->{account};
 232     }
 233     my %l_options =
 234         (mem => 'vmem',
 235          time => 'walltime',
 236          cputime => 'cput',
 237          host    => 'host',
 238          pmem => 'pmem',
 239          pvmem => 'pvmem',
 240          max_file => 'file',
 241         );
 242     for my $k (keys %l_options) {
 243         if ($options->{$k}) {
 244             push @l,$l_options{$k}.'='.$options{$k};
 245         }
 246     }
 247     push @qo,'-l',join(',',@l) if @l;
 248     if ($options->{interactive}) {
 249         push @qo,'-I';
 250     }
 251     if ($options->{name}) {
 252         push @qo,'-N',$options->{name};
 253     } else {
 254         push @qo,'-N',join('_',
 255                            map {my $a = $_; $a =~ s/[^a-zA-Z0-9]*//g; $a;}
 256                           @{$args}[0..min($#{$args},2)]);
 257     }
 258     # join error and output streams
 259    if ($options->{join}) {
 260         push @qo,'-j','oe';
 261     }
 262     return @qo;
 263 }
 264
 265 sub generate_slurm_options{
 266     my ($options,$args) = @_;
 267     my @qo;
 268     if (defined $options->{queue} and length $options->{queue}) {
 269         push @qo,'-p',$options->{queue};
 270     }
 271     ## handle the -l options
 272     if (defined $options->{account}) {
 273         push @qo,'-A',$options->{account};
 274     }
 275     my %options_map =
 276         (mem => 'mem',
 277          ppn => 'cpus-per-task',
 278          time => 'time',
 279          cputime => 'cput',
 280          host    => 'host',
 281          pmem => 'pmem',
 282          pvmem => 'pvmem',
 283          max_file => 'file',
 284         );
 285     for my $k (keys %options_map) {
 286         if ($options->{$k}) {
 287             push @qo,'--'.$options_map{$k}.'='.$options{$k};
 288         }
 289     }
 290     if ($options{mem}) {
 291         push @qo,'--mem='.$options{mem};
 292     }
 293     if ($options->{interactive}) {
 294         push @qo,'-I';
 295     }
 296     if ($options->{name}) {
 297         push @qo,'-J',$options->{name};
 298     } else {
 299         push @qo,'-J',join('_',
 300                            map {my $a = $_; $a =~ s/[^a-zA-Z0-9]*//g; $a;}
 301                           @{$args}[0..min($#{$args},2)]);
 302     }
 303     return @qo;
 304 }
 305
 306 sub read_array_options{
 307     my ($options) = @_;
 308     my $fh = \*STDIN;
 309     if (defined $options->{array_from}) {
 310         $fh = IO::File->new(defined $options->{array_from}) or
 311             die "Unable to open $options->{array_from} for reading: $!";
 312     }
 313     my @array;
 314     for (<$fh>) {
 315         chomp;
 316         push @array,$_;
 317     }
 318     return @array;
 319 }
 320
 321 sub call_qsub {
 322     my ($qsub_options,$script) = @_;
 323     my $qsub_fh;
 324     open $qsub_fh,'|-',$JOB_SUBMITTER,@{$qsub_options} or
 325         die "Unable to start $JOB_SUBMITTER: $!";
 326     print {$qsub_fh} $script or
 327         die "Unable to print to $JOB_SUBMITTER: $!";
 328     close($qsub_fh) or
 329         die "Unable to close $JOB_SUBMITTER filehandle: $!";
 330 }
 331
 332 sub write_qsub_script {
 333     my ($opt,$arg,$array) = @_;
 334
 335     my $script = "#!/bin/bash\n";
 336     my $command = join(' ',map {$_ =~ /\s/?qq('$_'):$_} @{$arg});
 337         $script .= <<EOF;
 338 # this script was written by dqsub
 339 EOF
 340     # if there is a precommand, write it out
 341     if ($opt->{precommand}) {
 342         $script .=<<EOF;
 343 # this is the precommand _BEGIN_
 344 $opt->{precommand}
 345 # precommand _END_
 346 EOF
 347     }
 348     my $directory = getcwd;
 349     if (defined $opt->{dir}) {
 350         $directory = abs_path($opt->{dir});
 351     }
 352     # we really should be quoting this instead
 353     $script .=<<EOF;
 354 # change to the working directory
 355 cd "$directory";
 356 EOF
 357     if (defined $opt->{array}) {
 358         my @subshell = ('','');
 359         my $array_opt = join("\n",@{$array});
 360         my $max_array = scalar @{$array};
 361         my $apjm1 = $opt->{array_per_job} - 1;
 362         $script .= <<EOF;
 363 if [ -n "\$PBS_ARRAYID" ]; then
 364    MYARRAYID="\${PBS_ARRAYID:=1}"
 365 else
 366    MYARRAYID="\${SLURM_ARRAY_TASK_ID:=1}"
 367 fi;
 368 EOF
 369         if ($opt->{array_per_job} > 1) {
 370             # we will use subshells if there are more than one array
 371             # items per job
 372             @subshell = ('(',')');
 373             $script .= <<EOF;
 374 for i in \$(seq 1 $opt->{array_per_job}); do
 375 # in some cases, the jobs aren't going to come out evenly. Handle that.
 376 JOBNUM=\$(( \${MYARRAYID:=1} * $opt->{array_per_job} + \$i - $opt->{array_per_job} ))
 377 if [ \$JOBNUM -le $max_array ]; then
 378 OPT=\$(sed -n -e "\$JOBNUM p"<<'_HERE_DOC_END_'
 379 EOF
 380         } else {
 381             $script .= <<EOF;
 382 OPT=\$(sed -n -e "\${MYARRAYID:=1} p"<<'_HERE_DOC_END_'
 383 EOF
 384         }
 385         $script .= <<EOF;
 386 $array_opt
 387 _HERE_DOC_END_
 388 )
 389 EOF
 390         if ($opt->{array} eq 'chdir') {
 391             $script .= <<EOF;
 392 $subshell[0]
 393 cd "\$OPT";
 394 exec ${command};
 395 $subshell[1]
 396 EOF
 397         } else {
 398             $script .= <<EOF;
 399 $subshell[0]
 400 exec ${command} "\$OPT";
 401 $subshell[1]
 402 EOF
 403         }
 404         if ($opt->{array_per_job} > 1) {
 405             $script .= <<EOF
 406 fi;
 407 done;
 408 EOF
 409         }
 410     } else {
 411         $script .= <<EOF;
 412 # there's no array, so just executing the command with arguments
 413 exec $command;
 414 EOF
 415     }
 416     return $script;
 417 }
 418
 419
 420 __END__