1#! /usr/bin/env python2.7
2
3# Copyright (c) 2004-2005, 2007 The Regents of The University of Michigan
4# All rights reserved.
5#
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.
16#
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.
28#
29# Authors: Steve Reinhardt
30#          Ali Saidi
31
32# Important!
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
37#   PS1='$ '
38
39
40import sys
41import os
42import re
43import time
44import optparse
45
46import pexpect
47
48progname = os.path.basename(sys.argv[0])
49
50usage = "%prog [options] command [command arguments]"
51optparser = optparse.OptionParser(usage=usage)
52optparser.allow_interspersed_args=False
53optparser.add_option('-e', dest='stderr_file',
54                     help='command stderr output file')
55optparser.add_option('-o', dest='stdout_file',
56                     help='command stdout output file')
57optparser.add_option('-l', dest='save_log', action='store_true',
58                     help='save oarsub output log file')
59optparser.add_option('-N', dest='job_name',
60                     help='oarsub job name')
61optparser.add_option('-q', dest='dest_queue',
62                     help='oarsub destination queue')
63optparser.add_option('--qwait', dest='oarsub_timeout', type='int',
64                     help='oarsub queue wait timeout', default=30*60)
65optparser.add_option('-t', dest='cmd_timeout', type='int',
66                     help='command execution timeout', default=600*60)
67
68(options, cmd) = optparser.parse_args()
69
70if cmd == []:
71    print >>sys.stderr, "%s: missing command" % progname
72    sys.exit(1)
73
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.
76#
77#if not options.job_name:
78#    options.job_name = cmd[0]
79
80cwd = os.getcwd()
81
82# Deal with systems where /n is a symlink to /.automount
83if cwd.startswith('/.automount/'):
84    cwd = cwd.replace('/.automount/', '/n/', 1)
85
86if not cwd.startswith('/n/poolfs/'):
87    print >>sys.stderr, "Error: current directory must be under /n/poolfs."
88    sys.exit(1)
89
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.
92class 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
95    # reliably.
96    prompt_re = re.compile('qdo\$ ')
97
98    def __init__(self, cmd):
99        # initialize base pexpect.spawn object
100        try:
101            pexpect.spawn.__init__(self, cmd)
102        except pexpect.ExceptionPexpect, exc:
103            print "%s:" % progname, exc
104            sys.exit(1)
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
109        try:
110            self.expect('\$ ', options.oarsub_timeout)
111        except pexpect.TIMEOUT:
112            print >>sys.stderr, "%s: oarsub timed out." % progname
113            self.kill(9)
114            self.safe_close()
115            sys.exit(1)
116        self.do_command('unset PROMPT_COMMAND; PS1="qdo$ "')
117
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
122
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):
126        global full_output
127        self.sendline(cmd)
128        # read back the echo of the command
129        self.readline()
130        # wait for the next prompt
131        self.expect(self.prompt_re, timeout)
132        output = self.before.rstrip()
133        return output
134
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)
140        # collect status
141        status = int(self.do_bare_command("echo $?", self.quick_timeout))
142        return (output, status)
143
144    # Check to see if the given directory exists.
145    def dir_exists(self, dirname):
146        (output, status) = shell.do_command('[ -d %s ]' % dirname,
147                                            self.quick_timeout)
148        return status == 0
149
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):
155        count = 0
156        while self.isalive() and count < 10:
157            time.sleep(1)
158        self.close(force=False)
159
160# Spawn the interactive pool job.
161
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.
165if False and len(cmd) > 50:
166    shell_cmd = 'ssh -t poolfs /bin/sh -l'
167    print "%s: running %s on poolfs" % (progname, cmd[0])
168else:
169    shell_cmd = 'oarsub -I'
170    if options.job_name:
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
175
176shell = Shell(shell_cmd)
177
178try:
179    # chdir to cwd
180    (output, status) = shell.do_command('cd ' + cwd)
181
182    if status != 0:
183        raise OSError, "Can't chdir to %s" % cwd
184
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.
190    try:
191        if 'cc' in cmd[0] or 'g++' in cmd[0]:
192            output_dir = os.path.dirname(cmd[cmd.index('-o')+1])
193        elif 'm5' in cmd[0]:
194            output_dir = cmd[cmd.index('-d')+1]
195        else:
196            output_dir = None
197    except (ValueError, IndexError):
198        # no big deal if there's no '-o'/'-d' or if it's the final argument
199        output_dir = None
200
201    if output_dir:
202        secs_waited = 0
203        while not shell.dir_exists(output_dir) and secs_waited < 90:
204            time.sleep(5)
205            secs_waited += 5
206        if secs_waited > 30:
207            print "waited", secs_waited, "seconds for", output_dir
208
209    # run command
210    if options.stdout_file:
211        cmd += ['>', options.stdout_file]
212    if options.stderr_file:
213        cmd += ['2>', options.stderr_file]
214    try:
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
220            shell.safe_close()
221            status = 3
222    if output:
223        print output
224finally:
225    # end job
226    if shell.isalive():
227        shell.sendline('exit')
228        shell.expect('Disconnected from OAR job .*')
229        shell.safe_close()
230
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)
235        log.close()
236del shell
237
238sys.exit(status)
239