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