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