3 # Copyright (c) 2004-2005, 2007 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
33 # This script expects a simple $ prompt, if you are using a shell other than
34 # sh which defaults to this you'll need to add something like the following
35 # to your bashrc/bash_profile script:
36 #if [ "$OAR_USER" = "xxxx" ]; then
48 progname
= os
.path
.basename(sys
.argv
[0])
50 usage
= "%prog [options] command [command arguments]"
51 optparser
= optparse
.OptionParser(usage
=usage
)
52 optparser
.allow_interspersed_args
=False
53 optparser
.add_option('-e', dest
='stderr_file',
54 help='command stderr output file')
55 optparser
.add_option('-o', dest
='stdout_file',
56 help='command stdout output file')
57 optparser
.add_option('-l', dest
='save_log', action
='store_true',
58 help='save oarsub output log file')
59 optparser
.add_option('-N', dest
='job_name',
60 help='oarsub job name')
61 optparser
.add_option('-q', dest
='dest_queue',
62 help='oarsub destination queue')
63 optparser
.add_option('--qwait', dest
='oarsub_timeout', type='int',
64 help='oarsub queue wait timeout', default
=30*60)
65 optparser
.add_option('-t', dest
='cmd_timeout', type='int',
66 help='command execution timeout', default
=600*60)
68 (options
, cmd
) = optparser
.parse_args()
71 print >>sys
.stderr
, "%s: missing command" % progname
74 # If we want to do this, need to add check here to make sure cmd[0] is
75 # a valid PBS job name, else oarsub will die on us.
77 #if not options.job_name:
78 # options.job_name = cmd[0]
82 # Deal with systems where /n is a symlink to /.automount
83 if cwd
.startswith('/.automount/'):
84 cwd
= cwd
.replace('/.automount/', '/n/', 1)
86 if not cwd
.startswith('/n/poolfs/'):
87 print >>sys
.stderr
, "Error: current directory must be under /n/poolfs."
90 # The Shell class wraps pexpect.spawn with some handy functions that
91 # assume the thing on the other end is a Bourne/bash shell.
92 class Shell(pexpect
.spawn
):
93 # Regexp to match the shell prompt. We change the prompt to
94 # something fixed and distinctive to make it easier to match
96 prompt_re
= re
.compile('qdo\$ ')
98 def __init__(self
, cmd
):
99 # initialize base pexpect.spawn object
101 pexpect
.spawn
.__init
__(self
, cmd
)
102 except pexpect
.ExceptionPexpect
, exc
:
103 print "%s:" % progname
, exc
105 # full_output accumulates the full output of the session
106 self
.full_output
= ""
107 self
.quick_timeout
= 15
108 # wait for a prompt, then change it
110 self
.expect('\$ ', options
.oarsub_timeout
)
111 except pexpect
.TIMEOUT
:
112 print >>sys
.stderr
, "%s: oarsub timed out." % progname
116 self
.do_command('unset PROMPT_COMMAND; PS1="qdo$ "')
118 # version of expect that updates full_output too
119 def expect(self
, regexp
, timeout
= -1):
120 pexpect
.spawn
.expect(self
, regexp
, timeout
)
121 self
.full_output
+= self
.before
+ self
.after
123 # Just issue a command and wait for the next prompt.
124 # Returns a string containing the output of the command.
125 def do_bare_command(self
, cmd
, timeout
= -1):
128 # read back the echo of the command
130 # wait for the next prompt
131 self
.expect(self
.prompt_re
, timeout
)
132 output
= self
.before
.rstrip()
135 # Issue a command, then query its exit status.
136 # Returns a (string, int) tuple with the command output and the status.
137 def do_command(self
, cmd
, timeout
= -1):
138 # do the command itself
139 output
= self
.do_bare_command(cmd
, timeout
)
141 status
= int(self
.do_bare_command("echo $?", self
.quick_timeout
))
142 return (output
, status
)
144 # Check to see if the given directory exists.
145 def dir_exists(self
, dirname
):
146 (output
, status
) = shell
.do_command('[ -d %s ]' % dirname
,
150 # Don't actually try to close it.. just wait until it closes by itself
151 # We can't actually kill the pid which is what it's trying to do, and if
152 # we call wait we could be in an unfortunate situation of it printing input
153 # right as we call wait, so the input is never read and the process never ends
154 def safe_close(self
):
156 while self
.isalive() and count
< 10:
158 self
.close(force
=False)
160 # Spawn the interactive pool job.
162 # Hack to do link on poolfs... disabled for now since
163 # compiler/linker/library versioning problems between poolfs and
164 # nodes. May never work since poolfs is x86-64 and nodes are 32-bit.
165 if False and len(cmd
) > 50:
166 shell_cmd
= 'ssh -t poolfs /bin/sh -l'
167 print "%s: running %s on poolfs" % (progname
, cmd
[0])
169 shell_cmd
= 'oarsub -I'
171 shell_cmd
+= ' -n "%s"' % options
.job_name
172 if options
.dest_queue
:
173 shell_cmd
+= ' -q ' + options
.dest_queue
174 shell_cmd
+= ' -d %s' % cwd
176 shell
= Shell(shell_cmd
)
180 (output
, status
) = shell
.do_command('cd ' + cwd
)
183 raise OSError, "Can't chdir to %s" % cwd
185 # wacky hack: sometimes scons will create an output directory then
186 # fork a job to generate files in that directory, and the job will
187 # get run before the directory creation propagates through NFS.
188 # This hack looks for a '-o' option indicating an output file and
189 # waits for the corresponding directory to appear if necessary.
191 if 'cc' in cmd
[0] or 'g++' in cmd
[0]:
192 output_dir
= os
.path
.dirname(cmd
[cmd
.index('-o')+1])
194 output_dir
= cmd
[cmd
.index('-d')+1]
197 except (ValueError, IndexError):
198 # no big deal if there's no '-o'/'-d' or if it's the final argument
203 while not shell
.dir_exists(output_dir
) and secs_waited
< 90:
207 print "waited", secs_waited
, "seconds for", output_dir
210 if options
.stdout_file
:
211 cmd
+= ['>', options
.stdout_file
]
212 if options
.stderr_file
:
213 cmd
+= ['2>', options
.stderr_file
]
215 (output
, status
) = shell
.do_command(' '.join(cmd
), options
.cmd_timeout
)
216 except pexpect
.TIMEOUT
:
217 print >>sys
.stderr
, "%s: command timed out after %d seconds." \
218 % (progname
, options
.cmd_timeout
)
219 shell
.sendline('~.') # oarsub/ssh termination escape sequence
227 shell
.sendline('exit')
228 shell
.expect('Disconnected from OAR job .*')
231 # if there was an error, log the output even if not requested
232 if status
!= 0 or options
.save_log
:
233 log
= file('qdo-log.' + str(os
.getpid()), 'w')
234 log
.write(shell
.full_output
)