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     return @qo;
 302 }
 303
 304 sub read_array_options{
 305     my ($options) = @_;
 306     my $fh = \*STDIN;
 307     if (defined $options->{array_from}) {
 308         $fh = IO::File->new(defined $options->{array_from}) or
 309             die "Unable to open $options->{array_from} for reading: $!";
 310     }
 311     my @array;
 312     for (<$fh>) {
 313         chomp;
 314         push @array,$_;
 315     }
 316     return @array;
 317 }
 318
 319 sub call_qsub {
 320     my ($qsub_options,$script) = @_;
 321     my $qsub_fh;
 322     open $qsub_fh,'|-',$JOB_SUBMITTER,@{$qsub_options} or
 323         die "Unable to start $JOB_SUBMITTER: $!";
 324     print {$qsub_fh} $script or
 325         die "Unable to print to $JOB_SUBMITTER: $!";
 326     close($qsub_fh) or
 327         die "Unable to close $JOB_SUBMITTER filehandle: $!";
 328 }
 329
 330 sub write_qsub_script {
 331     my ($opt,$arg,$array) = @_;
 332
 333     my $script = "#!/bin/bash\n";
 334     my $command = join(' ',map {$_ =~ /\s/?qq('$_'):$_} @{$arg});
 335         $script .= <<EOF;
 336 # this script was written by dqsub
 337 EOF
 338     my $directory = getcwd;
 339     if (defined $opt->{dir}) {
 340         $directory = abs_path($opt->{dir});
 341     }
 342     # we really should be quoting this instead
 343     $script .=<<EOF;
 344 # change to the working directory
 345 cd "$directory";
 346 EOF
 347     if (defined $opt->{array}) {
 348         my @subshell = ('','');
 349         my $array_opt = join("\n",@{$array});
 350         my $max_array = scalar @{$array};
 351         my $apjm1 = $opt->{array_per_job} - 1;
 352         $script .= <<EOF;
 353 if [ -n "\$PBS_ARRAYID" ]; then
 354    MYARRAYID="\${PBS_ARRAYID:=1}"
 355 else
 356    MYARRAYID="\${SLURM_ARRAY_TASK_ID:=1}"
 357 fi;
 358 EOF
 359         if ($opt->{array_per_job} > 1) {
 360             # we will use subshells if there are more than one array
 361             # items per job
 362             @subshell = ('(',')');
 363             $script .= <<EOF;
 364 for i in \$(seq 1 $opt->{array_per_job}); do
 365 # in some cases, the jobs aren't going to come out evenly. Handle that.
 366 JOBNUM=\$(( \${MYARRAYID:=1} * $opt->{array_per_job} + \$i - $opt->{array_per_job} ))
 367 if [ \$JOBNUM -le $max_array ]; then
 368 OPT=\$(sed -n -e "\$JOBNUM p"<<'_HERE_DOC_END_'
 369 EOF
 370         } else {
 371             $script .= <<EOF;
 372 OPT=\$(sed -n -e "\${MYARRAYID:=1} p"<<'_HERE_DOC_END_'
 373 EOF
 374         }
 375         $script .= <<EOF;
 376 $array_opt
 377 _HERE_DOC_END_
 378 )
 379 EOF
 380         if ($opt->{array} eq 'chdir') {
 381             $script .= <<EOF;
 382 $subshell[0]
 383 cd "\$OPT";
 384 exec ${command};
 385 $subshell[1]
 386 EOF
 387         } else {
 388             $script .= <<EOF;
 389 $subshell[0]
 390 exec ${command} "\$OPT";
 391 $subshell[1]
 392 EOF
 393         }
 394         if ($opt->{array_per_job} > 1) {
 395             $script .= <<EOF
 396 fi;
 397 done;
 398 EOF
 399         }
 400     } else {
 401         $script .= <<EOF;
 402 # there's no array, so just executing the command with arguments
 403 exec $command;
 404 EOF
 405     }
 406     return $script;
 407 }
 408
 409
 410 __END__