8620f1f9856e1d41814f776296b293f1f8646b02
[uiuc_igb_scripts.git] / 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__