3 # Copyright (c) 2004-2005 The Regents of The University of Michigan
6 # Redistribution and use in source and binary forms, with or without
7 # modification, are permitted provided that the following conditions are
8 # met: redistributions of source code must retain the above copyright
9 # notice, this list of conditions and the following disclaimer;
10 # redistributions in binary form must reproduce the above copyright
11 # notice, this list of conditions and the following disclaimer in the
12 # documentation and/or other materials provided with the distribution;
13 # neither the name of the copyright holders nor the names of its
14 # contributors may be used to endorse or promote products derived from
15 # this software without specific prior written permission.
17 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
18 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
19 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
20 # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
21 # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
22 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
23 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
24 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
25 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
26 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
27 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29 # Authors: Steve Reinhardt
39 progname
= os
.path
.basename(sys
.argv
[0])
41 usage
= "%prog [options] command [command arguments]"
42 optparser
= optparse
.OptionParser(usage
=usage
)
43 optparser
.allow_interspersed_args
=False
44 optparser
.add_option('-e', dest
='stderr_file',
45 help='command stderr output file')
46 optparser
.add_option('-o', dest
='stdout_file',
47 help='command stdout output file')
48 optparser
.add_option('-l', dest
='save_log', action
='store_true',
49 help='save qsub output log file')
50 optparser
.add_option('-N', dest
='job_name',
52 optparser
.add_option('-q', dest
='dest_queue',
53 help='qsub destination queue')
54 optparser
.add_option('--qwait', dest
='qsub_timeout', type='int',
55 help='qsub queue wait timeout', default
=30*60)
56 optparser
.add_option('-t', dest
='cmd_timeout', type='int',
57 help='command execution timeout', default
=600*60)
59 (options
, cmd
) = optparser
.parse_args()
62 print >>sys
.stderr
, "%s: missing command" % progname
65 # If we want to do this, need to add check here to make sure cmd[0] is
66 # a valid PBS job name, else qsub will die on us.
68 #if not options.job_name:
69 # options.job_name = cmd[0]
73 # Deal with systems where /n is a symlink to /.automount
74 if cwd
.startswith('/.automount/'):
75 cwd
= cwd
.replace('/.automount/', '/n/', 1)
77 if not cwd
.startswith('/n/poolfs/'):
78 print >>sys
.stderr
, "Error: current directory must be under /n/poolfs."
81 # The Shell class wraps pexpect.spawn with some handy functions that
82 # assume the thing on the other end is a Bourne/bash shell.
83 class Shell(pexpect
.spawn
):
84 # Regexp to match the shell prompt. We change the prompt to
85 # something fixed and distinctive to make it easier to match
87 prompt_re
= re
.compile('qdo\$ ')
89 def __init__(self
, cmd
):
90 # initialize base pexpect.spawn object
92 pexpect
.spawn
.__init
__(self
, cmd
)
93 except pexpect
.ExceptionPexpect
, exc
:
94 print "%s:" % progname
, exc
96 # full_output accumulates the full output of the session
98 self
.quick_timeout
= 15
99 # wait for a prompt, then change it
101 self
.expect('\$ ', options
.qsub_timeout
)
102 except pexpect
.TIMEOUT
:
103 print >>sys
.stderr
, "%s: qsub timed out." % progname
105 self
.close(wait
=True)
107 self
.do_command('unset PROMPT_COMMAND; PS1="qdo$ "')
109 # version of expect that updates full_output too
110 def expect(self
, regexp
, timeout
= -1):
111 pexpect
.spawn
.expect(self
, regexp
, timeout
)
112 self
.full_output
+= self
.before
+ self
.after
114 # Just issue a command and wait for the next prompt.
115 # Returns a string containing the output of the command.
116 def do_bare_command(self
, cmd
, timeout
= -1):
119 # read back the echo of the command
121 # wait for the next prompt
122 self
.expect(self
.prompt_re
, timeout
)
123 output
= self
.before
.rstrip()
126 # Issue a command, then query its exit status.
127 # Returns a (string, int) tuple with the command output and the status.
128 def do_command(self
, cmd
, timeout
= -1):
129 # do the command itself
130 output
= self
.do_bare_command(cmd
, timeout
)
132 status
= int(self
.do_bare_command("echo $?", self
.quick_timeout
))
133 return (output
, status
)
135 # Check to see if the given directory exists.
136 def dir_exists(self
, dirname
):
137 (output
, status
) = shell
.do_command('[ -d %s ]' % dirname
,
142 # Spawn the interactive pool job.
144 # Hack to do link on poolfs... disabled for now since
145 # compiler/linker/library versioning problems between poolfs and
146 # nodes. May never work since poolfs is x86-64 and nodes are 32-bit.
147 if False and len(cmd
) > 50:
148 shell_cmd
= 'ssh -t poolfs /bin/sh -l'
149 print "%s: running %s on poolfs" % (progname
, cmd
[0])
151 shell_cmd
= 'qsub -I -S /bin/sh'
153 shell_cmd
+= ' -N "%s"' % options
.job_name
154 if options
.dest_queue
:
155 shell_cmd
+= ' -q ' + options
.dest_queue
157 shell
= Shell(shell_cmd
)
161 (output
, status
) = shell
.do_command('cd ' + cwd
)
164 raise OSError, "Can't chdir to %s" % cwd
166 # wacky hack: sometimes scons will create an output directory then
167 # fork a job to generate files in that directory, and the job will
168 # get run before the directory creation propagates through NFS.
169 # This hack looks for a '-o' option indicating an output file and
170 # waits for the corresponding directory to appear if necessary.
172 if 'cc' in cmd
[0] or 'g++' in cmd
[0]:
173 output_dir
= os
.path
.dirname(cmd
[cmd
.index('-o')+1])
175 output_dir
= cmd
[cmd
.index('-d')+1]
178 except (ValueError, IndexError):
179 # no big deal if there's no '-o'/'-d' or if it's the final argument
184 while not shell
.dir_exists(output_dir
) and secs_waited
< 90:
188 print "waited", secs_waited
, "seconds for", output_dir
191 if options
.stdout_file
:
192 cmd
+= ['>', options
.stdout_file
]
193 if options
.stderr_file
:
194 cmd
+= ['2>', options
.stderr_file
]
196 (output
, status
) = shell
.do_command(' '.join(cmd
), options
.cmd_timeout
)
197 except pexpect
.TIMEOUT
:
198 print >>sys
.stderr
, "%s: command timed out after %d seconds." \
199 % (progname
, options
.cmd_timeout
)
200 shell
.sendline('~.') # qsub/ssh termination escape sequence
201 shell
.close(wait
=True)
209 shell
.sendline('exit')
210 shell
.expect('qsub: job .* completed\r\n')
211 shell
.close(wait
=True)
213 # if there was an error, log the output even if not requested
214 if status
!= 0 or options
.save_log
:
215 log
= file('qdo-log.' + str(os
.getpid()), 'w')
216 log
.write(shell
.full_output
)