qdo revision 2665
11884Sstever@eecs.umich.edu#! /usr/bin/env python
21884Sstever@eecs.umich.edu
31884Sstever@eecs.umich.edu# Copyright (c) 2004-2005 The Regents of The University of Michigan
41884Sstever@eecs.umich.edu# All rights reserved.
51884Sstever@eecs.umich.edu#
61884Sstever@eecs.umich.edu# Redistribution and use in source and binary forms, with or without
71884Sstever@eecs.umich.edu# modification, are permitted provided that the following conditions are
81884Sstever@eecs.umich.edu# met: redistributions of source code must retain the above copyright
91884Sstever@eecs.umich.edu# notice, this list of conditions and the following disclaimer;
101884Sstever@eecs.umich.edu# redistributions in binary form must reproduce the above copyright
111884Sstever@eecs.umich.edu# notice, this list of conditions and the following disclaimer in the
121884Sstever@eecs.umich.edu# documentation and/or other materials provided with the distribution;
131884Sstever@eecs.umich.edu# neither the name of the copyright holders nor the names of its
141884Sstever@eecs.umich.edu# contributors may be used to endorse or promote products derived from
151884Sstever@eecs.umich.edu# this software without specific prior written permission.
161884Sstever@eecs.umich.edu#
171884Sstever@eecs.umich.edu# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
181884Sstever@eecs.umich.edu# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
191884Sstever@eecs.umich.edu# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
201884Sstever@eecs.umich.edu# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
211884Sstever@eecs.umich.edu# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
221884Sstever@eecs.umich.edu# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
231884Sstever@eecs.umich.edu# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
241884Sstever@eecs.umich.edu# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
251884Sstever@eecs.umich.edu# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
261884Sstever@eecs.umich.edu# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
271884Sstever@eecs.umich.edu# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
282665Ssaidi@eecs.umich.edu#
292665Ssaidi@eecs.umich.edu# Authors: Steve Reinhardt
301884Sstever@eecs.umich.edu
311884Sstever@eecs.umich.eduimport sys
321884Sstever@eecs.umich.eduimport os
331884Sstever@eecs.umich.eduimport re
341884Sstever@eecs.umich.eduimport time
351884Sstever@eecs.umich.eduimport optparse
361884Sstever@eecs.umich.edu
371884Sstever@eecs.umich.eduimport pexpect
381884Sstever@eecs.umich.edu
391884Sstever@eecs.umich.eduprogname = os.path.basename(sys.argv[0])
401884Sstever@eecs.umich.edu
411884Sstever@eecs.umich.eduusage = "%prog [options] command [command arguments]"
421884Sstever@eecs.umich.eduoptparser = optparse.OptionParser(usage=usage)
431884Sstever@eecs.umich.eduoptparser.allow_interspersed_args=False
441884Sstever@eecs.umich.eduoptparser.add_option('-e', dest='stderr_file',
451884Sstever@eecs.umich.edu                     help='command stderr output file')
461884Sstever@eecs.umich.eduoptparser.add_option('-o', dest='stdout_file',
471884Sstever@eecs.umich.edu                     help='command stdout output file')
481884Sstever@eecs.umich.eduoptparser.add_option('-l', dest='save_log', action='store_true',
491884Sstever@eecs.umich.edu                     help='save qsub output log file')
501930Sstever@eecs.umich.eduoptparser.add_option('-N', dest='job_name',
511930Sstever@eecs.umich.edu                     help='qsub job name')
521930Sstever@eecs.umich.eduoptparser.add_option('-q', dest='dest_queue',
531930Sstever@eecs.umich.edu                     help='qsub destination queue')
541930Sstever@eecs.umich.eduoptparser.add_option('--qwait', dest='qsub_timeout', type='int',
551884Sstever@eecs.umich.edu                     help='qsub queue wait timeout', default=30*60)
561884Sstever@eecs.umich.eduoptparser.add_option('-t', dest='cmd_timeout', type='int',
571884Sstever@eecs.umich.edu                     help='command execution timeout', default=600*60)
581884Sstever@eecs.umich.edu
591884Sstever@eecs.umich.edu(options, cmd) = optparser.parse_args()
601884Sstever@eecs.umich.edu
611884Sstever@eecs.umich.eduif cmd == []:
621884Sstever@eecs.umich.edu    print >>sys.stderr, "%s: missing command" % progname
631884Sstever@eecs.umich.edu    sys.exit(1)
641884Sstever@eecs.umich.edu
651940Sstever@eecs.umich.edu# If we want to do this, need to add check here to make sure cmd[0] is
661940Sstever@eecs.umich.edu# a valid PBS job name, else qsub will die on us.
671940Sstever@eecs.umich.edu#
681940Sstever@eecs.umich.edu#if not options.job_name:
691940Sstever@eecs.umich.edu#    options.job_name = cmd[0]
701930Sstever@eecs.umich.edu
711884Sstever@eecs.umich.educwd = os.getcwd()
721884Sstever@eecs.umich.edu
731884Sstever@eecs.umich.edu# Deal with systems where /n is a symlink to /.automount
741884Sstever@eecs.umich.eduif cwd.startswith('/.automount/'):
751884Sstever@eecs.umich.edu    cwd = cwd.replace('/.automount/', '/n/', 1)
761884Sstever@eecs.umich.edu
771884Sstever@eecs.umich.eduif not cwd.startswith('/n/poolfs/'):
781884Sstever@eecs.umich.edu    print >>sys.stderr, "Error: current directory must be under /n/poolfs."
791884Sstever@eecs.umich.edu    sys.exit(1)
801884Sstever@eecs.umich.edu
811884Sstever@eecs.umich.edu# The Shell class wraps pexpect.spawn with some handy functions that
821884Sstever@eecs.umich.edu# assume the thing on the other end is a Bourne/bash shell.
831884Sstever@eecs.umich.educlass Shell(pexpect.spawn):
841884Sstever@eecs.umich.edu    # Regexp to match the shell prompt.  We change the prompt to
851884Sstever@eecs.umich.edu    # something fixed and distinctive to make it easier to match
861884Sstever@eecs.umich.edu    # reliably.
871884Sstever@eecs.umich.edu    prompt_re = re.compile('qdo\$ ')
881884Sstever@eecs.umich.edu
891884Sstever@eecs.umich.edu    def __init__(self, cmd):
901884Sstever@eecs.umich.edu        # initialize base pexpect.spawn object
911884Sstever@eecs.umich.edu	try:
921884Sstever@eecs.umich.edu            pexpect.spawn.__init__(self, cmd)
931884Sstever@eecs.umich.edu	except pexpect.ExceptionPexpect, exc:
941884Sstever@eecs.umich.edu	    print "%s:" % progname, exc
951884Sstever@eecs.umich.edu	    sys.exit(1)
961884Sstever@eecs.umich.edu        # full_output accumulates the full output of the session
971884Sstever@eecs.umich.edu        self.full_output = ""
981884Sstever@eecs.umich.edu        self.quick_timeout = 15
991884Sstever@eecs.umich.edu        # wait for a prompt, then change it
1001884Sstever@eecs.umich.edu        try:
1011884Sstever@eecs.umich.edu            self.expect('\$ ', options.qsub_timeout)
1021884Sstever@eecs.umich.edu        except pexpect.TIMEOUT:
1031884Sstever@eecs.umich.edu            print >>sys.stderr, "%s: qsub timed out." % progname
1041964Sstever@eecs.umich.edu            self.kill(9)
1051884Sstever@eecs.umich.edu            self.close(wait=True)
1061884Sstever@eecs.umich.edu            sys.exit(1)
1071891Sstever@eecs.umich.edu        self.do_command('unset PROMPT_COMMAND; PS1="qdo$ "')
1081884Sstever@eecs.umich.edu
1091884Sstever@eecs.umich.edu    # version of expect that updates full_output too
1101884Sstever@eecs.umich.edu    def expect(self, regexp, timeout = -1):
1111884Sstever@eecs.umich.edu        pexpect.spawn.expect(self, regexp, timeout)
1121884Sstever@eecs.umich.edu        self.full_output += self.before + self.after
1131884Sstever@eecs.umich.edu
1141884Sstever@eecs.umich.edu    # Just issue a command and wait for the next prompt.
1151884Sstever@eecs.umich.edu    # Returns a string containing the output of the command.
1161884Sstever@eecs.umich.edu    def do_bare_command(self, cmd, timeout = -1):
1171884Sstever@eecs.umich.edu        global full_output
1181884Sstever@eecs.umich.edu        self.sendline(cmd)
1191884Sstever@eecs.umich.edu        # read back the echo of the command
1201884Sstever@eecs.umich.edu        self.readline()
1211884Sstever@eecs.umich.edu        # wait for the next prompt
1221884Sstever@eecs.umich.edu        self.expect(self.prompt_re, timeout)
1231884Sstever@eecs.umich.edu        output = self.before.rstrip()
1241884Sstever@eecs.umich.edu        return output
1251884Sstever@eecs.umich.edu
1261884Sstever@eecs.umich.edu    # Issue a command, then query its exit status.
1271884Sstever@eecs.umich.edu    # Returns a (string, int) tuple with the command output and the status.
1281884Sstever@eecs.umich.edu    def do_command(self, cmd, timeout = -1):
1291884Sstever@eecs.umich.edu        # do the command itself
1301884Sstever@eecs.umich.edu        output = self.do_bare_command(cmd, timeout)
1311884Sstever@eecs.umich.edu        # collect status
1321884Sstever@eecs.umich.edu        status = int(self.do_bare_command("echo $?", self.quick_timeout))
1331884Sstever@eecs.umich.edu        return (output, status)
1341884Sstever@eecs.umich.edu
1351884Sstever@eecs.umich.edu    # Check to see if the given directory exists.
1361884Sstever@eecs.umich.edu    def dir_exists(self, dirname):
1371884Sstever@eecs.umich.edu        (output, status) = shell.do_command('[ -d %s ]' % dirname,
1381884Sstever@eecs.umich.edu                                            self.quick_timeout)
1391884Sstever@eecs.umich.edu        return status == 0
1401884Sstever@eecs.umich.edu
1411884Sstever@eecs.umich.edu
1421884Sstever@eecs.umich.edu# Spawn the interactive pool job.
1431884Sstever@eecs.umich.edu
1441884Sstever@eecs.umich.edu# Hack to do link on poolfs... disabled for now since
1451884Sstever@eecs.umich.edu# compiler/linker/library versioning problems between poolfs and
1461884Sstever@eecs.umich.edu# nodes.  May never work since poolfs is x86-64 and nodes are 32-bit.
1471884Sstever@eecs.umich.eduif False and len(cmd) > 50:
1481884Sstever@eecs.umich.edu    shell_cmd = 'ssh -t poolfs /bin/sh -l'
1491884Sstever@eecs.umich.edu    print "%s: running %s on poolfs" % (progname, cmd[0])
1501884Sstever@eecs.umich.eduelse:
1511884Sstever@eecs.umich.edu    shell_cmd = 'qsub -I -S /bin/sh'
1521940Sstever@eecs.umich.edu    if options.job_name:
1531940Sstever@eecs.umich.edu        shell_cmd += ' -N "%s"' % options.job_name
1541930Sstever@eecs.umich.edu    if options.dest_queue:
1551930Sstever@eecs.umich.edu        shell_cmd += ' -q ' + options.dest_queue
1561884Sstever@eecs.umich.edu
1571884Sstever@eecs.umich.edushell = Shell(shell_cmd)
1581884Sstever@eecs.umich.edu
1591884Sstever@eecs.umich.edutry:
1601884Sstever@eecs.umich.edu    # chdir to cwd
1611884Sstever@eecs.umich.edu    (output, status) = shell.do_command('cd ' + cwd)
1621884Sstever@eecs.umich.edu
1631884Sstever@eecs.umich.edu    if status != 0:
1641884Sstever@eecs.umich.edu        raise OSError, "Can't chdir to %s" % cwd
1651884Sstever@eecs.umich.edu
1661884Sstever@eecs.umich.edu    # wacky hack: sometimes scons will create an output directory then
1671884Sstever@eecs.umich.edu    # fork a job to generate files in that directory, and the job will
1681884Sstever@eecs.umich.edu    # get run before the directory creation propagates through NFS.
1691884Sstever@eecs.umich.edu    # This hack looks for a '-o' option indicating an output file and
1701884Sstever@eecs.umich.edu    # waits for the corresponding directory to appear if necessary.
1711884Sstever@eecs.umich.edu    try:
1721884Sstever@eecs.umich.edu        if 'cc' in cmd[0] or 'g++' in cmd[0]:
1731884Sstever@eecs.umich.edu            output_dir = os.path.dirname(cmd[cmd.index('-o')+1])
1741884Sstever@eecs.umich.edu        elif 'm5' in cmd[0]:
1751884Sstever@eecs.umich.edu            output_dir = cmd[cmd.index('-d')+1]
1761884Sstever@eecs.umich.edu        else:
1771884Sstever@eecs.umich.edu            output_dir = None
1781884Sstever@eecs.umich.edu    except (ValueError, IndexError):
1791884Sstever@eecs.umich.edu        # no big deal if there's no '-o'/'-d' or if it's the final argument
1801884Sstever@eecs.umich.edu        output_dir = None
1811884Sstever@eecs.umich.edu
1821884Sstever@eecs.umich.edu    if output_dir:
1831884Sstever@eecs.umich.edu        secs_waited = 0
1842441Sstever@eecs.umich.edu        while not shell.dir_exists(output_dir) and secs_waited < 90:
1851884Sstever@eecs.umich.edu            time.sleep(5)
1861884Sstever@eecs.umich.edu            secs_waited += 5
1872441Sstever@eecs.umich.edu        if secs_waited > 30:
1881884Sstever@eecs.umich.edu            print "waited", secs_waited, "seconds for", output_dir
1891884Sstever@eecs.umich.edu
1901884Sstever@eecs.umich.edu    # run command
1911884Sstever@eecs.umich.edu    if options.stdout_file:
1921884Sstever@eecs.umich.edu        cmd += ['>', options.stdout_file]
1931884Sstever@eecs.umich.edu    if options.stderr_file:
1941884Sstever@eecs.umich.edu        cmd += ['2>', options.stderr_file]
1951884Sstever@eecs.umich.edu    try:
1961884Sstever@eecs.umich.edu        (output, status) = shell.do_command(' '.join(cmd), options.cmd_timeout)
1971884Sstever@eecs.umich.edu    except pexpect.TIMEOUT:
1981884Sstever@eecs.umich.edu            print >>sys.stderr, "%s: command timed out after %d seconds." \
1991884Sstever@eecs.umich.edu                  % (progname, options.cmd_timeout)
2001884Sstever@eecs.umich.edu            shell.sendline('~.') # qsub/ssh termination escape sequence
2011884Sstever@eecs.umich.edu            shell.close(wait=True)
2021884Sstever@eecs.umich.edu            status = 3
2031884Sstever@eecs.umich.edu    if output:
2041884Sstever@eecs.umich.edu        print output
2051884Sstever@eecs.umich.edu
2061884Sstever@eecs.umich.edufinally:
2071884Sstever@eecs.umich.edu    # end job
2081884Sstever@eecs.umich.edu    if shell.isalive():
2091884Sstever@eecs.umich.edu        shell.sendline('exit')
2101884Sstever@eecs.umich.edu        shell.expect('qsub: job .* completed\r\n')
2111884Sstever@eecs.umich.edu        shell.close(wait=True)
2121884Sstever@eecs.umich.edu
2131884Sstever@eecs.umich.edu    # if there was an error, log the output even if not requested
2141884Sstever@eecs.umich.edu    if status != 0 or options.save_log:
2151884Sstever@eecs.umich.edu        log = file('qdo-log.' + str(os.getpid()), 'w')
2161884Sstever@eecs.umich.edu        log.write(shell.full_output)
2171884Sstever@eecs.umich.edu        log.close()
2181884Sstever@eecs.umich.edu
2191884Sstever@eecs.umich.edudel shell
2201884Sstever@eecs.umich.edu
2211884Sstever@eecs.umich.edusys.exit(status)
222