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