qdo revision 1884
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. 281884Sstever@eecs.umich.edu 291884Sstever@eecs.umich.eduimport sys 301884Sstever@eecs.umich.eduimport os 311884Sstever@eecs.umich.eduimport re 321884Sstever@eecs.umich.eduimport time 331884Sstever@eecs.umich.eduimport optparse 341884Sstever@eecs.umich.edu 351884Sstever@eecs.umich.eduimport pexpect 361884Sstever@eecs.umich.edu 371884Sstever@eecs.umich.eduprogname = os.path.basename(sys.argv[0]) 381884Sstever@eecs.umich.edu 391884Sstever@eecs.umich.eduusage = "%prog [options] command [command arguments]" 401884Sstever@eecs.umich.eduoptparser = optparse.OptionParser(usage=usage) 411884Sstever@eecs.umich.eduoptparser.allow_interspersed_args=False 421884Sstever@eecs.umich.eduoptparser.add_option('-e', dest='stderr_file', 431884Sstever@eecs.umich.edu help='command stderr output file') 441884Sstever@eecs.umich.eduoptparser.add_option('-o', dest='stdout_file', 451884Sstever@eecs.umich.edu help='command stdout output file') 461884Sstever@eecs.umich.eduoptparser.add_option('-l', dest='save_log', action='store_true', 471884Sstever@eecs.umich.edu help='save qsub output log file') 481884Sstever@eecs.umich.eduoptparser.add_option('-q', dest='qsub_timeout', type='int', 491884Sstever@eecs.umich.edu help='qsub queue wait timeout', default=30*60) 501884Sstever@eecs.umich.eduoptparser.add_option('-t', dest='cmd_timeout', type='int', 511884Sstever@eecs.umich.edu help='command execution timeout', default=600*60) 521884Sstever@eecs.umich.edu 531884Sstever@eecs.umich.edu(options, cmd) = optparser.parse_args() 541884Sstever@eecs.umich.edu 551884Sstever@eecs.umich.eduif cmd == []: 561884Sstever@eecs.umich.edu print >>sys.stderr, "%s: missing command" % progname 571884Sstever@eecs.umich.edu sys.exit(1) 581884Sstever@eecs.umich.edu 591884Sstever@eecs.umich.educwd = os.getcwd() 601884Sstever@eecs.umich.edu 611884Sstever@eecs.umich.edu# Deal with systems where /n is a symlink to /.automount 621884Sstever@eecs.umich.eduif cwd.startswith('/.automount/'): 631884Sstever@eecs.umich.edu cwd = cwd.replace('/.automount/', '/n/', 1) 641884Sstever@eecs.umich.edu 651884Sstever@eecs.umich.eduif not cwd.startswith('/n/poolfs/'): 661884Sstever@eecs.umich.edu print >>sys.stderr, "Error: current directory must be under /n/poolfs." 671884Sstever@eecs.umich.edu sys.exit(1) 681884Sstever@eecs.umich.edu 691884Sstever@eecs.umich.edu# The Shell class wraps pexpect.spawn with some handy functions that 701884Sstever@eecs.umich.edu# assume the thing on the other end is a Bourne/bash shell. 711884Sstever@eecs.umich.educlass Shell(pexpect.spawn): 721884Sstever@eecs.umich.edu # Regexp to match the shell prompt. We change the prompt to 731884Sstever@eecs.umich.edu # something fixed and distinctive to make it easier to match 741884Sstever@eecs.umich.edu # reliably. 751884Sstever@eecs.umich.edu prompt_re = re.compile('qdo\$ ') 761884Sstever@eecs.umich.edu 771884Sstever@eecs.umich.edu def __init__(self, cmd): 781884Sstever@eecs.umich.edu # initialize base pexpect.spawn object 791884Sstever@eecs.umich.edu try: 801884Sstever@eecs.umich.edu pexpect.spawn.__init__(self, cmd) 811884Sstever@eecs.umich.edu except pexpect.ExceptionPexpect, exc: 821884Sstever@eecs.umich.edu print "%s:" % progname, exc 831884Sstever@eecs.umich.edu sys.exit(1) 841884Sstever@eecs.umich.edu # full_output accumulates the full output of the session 851884Sstever@eecs.umich.edu self.full_output = "" 861884Sstever@eecs.umich.edu self.quick_timeout = 15 871884Sstever@eecs.umich.edu # wait for a prompt, then change it 881884Sstever@eecs.umich.edu try: 891884Sstever@eecs.umich.edu self.expect('\$ ', options.qsub_timeout) 901884Sstever@eecs.umich.edu except pexpect.TIMEOUT: 911884Sstever@eecs.umich.edu print >>sys.stderr, "%s: qsub timed out." % progname 921884Sstever@eecs.umich.edu self.kill(15) 931884Sstever@eecs.umich.edu self.close(wait=True) 941884Sstever@eecs.umich.edu sys.exit(1) 951884Sstever@eecs.umich.edu self.do_command('PS1="qdo$ "') 961884Sstever@eecs.umich.edu 971884Sstever@eecs.umich.edu # version of expect that updates full_output too 981884Sstever@eecs.umich.edu def expect(self, regexp, timeout = -1): 991884Sstever@eecs.umich.edu pexpect.spawn.expect(self, regexp, timeout) 1001884Sstever@eecs.umich.edu self.full_output += self.before + self.after 1011884Sstever@eecs.umich.edu 1021884Sstever@eecs.umich.edu # Just issue a command and wait for the next prompt. 1031884Sstever@eecs.umich.edu # Returns a string containing the output of the command. 1041884Sstever@eecs.umich.edu def do_bare_command(self, cmd, timeout = -1): 1051884Sstever@eecs.umich.edu global full_output 1061884Sstever@eecs.umich.edu self.sendline(cmd) 1071884Sstever@eecs.umich.edu # read back the echo of the command 1081884Sstever@eecs.umich.edu self.readline() 1091884Sstever@eecs.umich.edu # wait for the next prompt 1101884Sstever@eecs.umich.edu self.expect(self.prompt_re, timeout) 1111884Sstever@eecs.umich.edu output = self.before.rstrip() 1121884Sstever@eecs.umich.edu return output 1131884Sstever@eecs.umich.edu 1141884Sstever@eecs.umich.edu # Issue a command, then query its exit status. 1151884Sstever@eecs.umich.edu # Returns a (string, int) tuple with the command output and the status. 1161884Sstever@eecs.umich.edu def do_command(self, cmd, timeout = -1): 1171884Sstever@eecs.umich.edu # do the command itself 1181884Sstever@eecs.umich.edu output = self.do_bare_command(cmd, timeout) 1191884Sstever@eecs.umich.edu # collect status 1201884Sstever@eecs.umich.edu status = int(self.do_bare_command("echo $?", self.quick_timeout)) 1211884Sstever@eecs.umich.edu return (output, status) 1221884Sstever@eecs.umich.edu 1231884Sstever@eecs.umich.edu # Check to see if the given directory exists. 1241884Sstever@eecs.umich.edu def dir_exists(self, dirname): 1251884Sstever@eecs.umich.edu (output, status) = shell.do_command('[ -d %s ]' % dirname, 1261884Sstever@eecs.umich.edu self.quick_timeout) 1271884Sstever@eecs.umich.edu return status == 0 1281884Sstever@eecs.umich.edu 1291884Sstever@eecs.umich.edu 1301884Sstever@eecs.umich.edu# Spawn the interactive pool job. 1311884Sstever@eecs.umich.edu 1321884Sstever@eecs.umich.edu# Hack to do link on poolfs... disabled for now since 1331884Sstever@eecs.umich.edu# compiler/linker/library versioning problems between poolfs and 1341884Sstever@eecs.umich.edu# nodes. May never work since poolfs is x86-64 and nodes are 32-bit. 1351884Sstever@eecs.umich.eduif False and len(cmd) > 50: 1361884Sstever@eecs.umich.edu shell_cmd = 'ssh -t poolfs /bin/sh -l' 1371884Sstever@eecs.umich.edu print "%s: running %s on poolfs" % (progname, cmd[0]) 1381884Sstever@eecs.umich.eduelse: 1391884Sstever@eecs.umich.edu shell_cmd = 'qsub -I -S /bin/sh' 1401884Sstever@eecs.umich.edu 1411884Sstever@eecs.umich.edushell = Shell(shell_cmd) 1421884Sstever@eecs.umich.edu 1431884Sstever@eecs.umich.edutry: 1441884Sstever@eecs.umich.edu # chdir to cwd 1451884Sstever@eecs.umich.edu (output, status) = shell.do_command('cd ' + cwd) 1461884Sstever@eecs.umich.edu 1471884Sstever@eecs.umich.edu if status != 0: 1481884Sstever@eecs.umich.edu raise OSError, "Can't chdir to %s" % cwd 1491884Sstever@eecs.umich.edu 1501884Sstever@eecs.umich.edu # wacky hack: sometimes scons will create an output directory then 1511884Sstever@eecs.umich.edu # fork a job to generate files in that directory, and the job will 1521884Sstever@eecs.umich.edu # get run before the directory creation propagates through NFS. 1531884Sstever@eecs.umich.edu # This hack looks for a '-o' option indicating an output file and 1541884Sstever@eecs.umich.edu # waits for the corresponding directory to appear if necessary. 1551884Sstever@eecs.umich.edu try: 1561884Sstever@eecs.umich.edu if 'cc' in cmd[0] or 'g++' in cmd[0]: 1571884Sstever@eecs.umich.edu output_dir = os.path.dirname(cmd[cmd.index('-o')+1]) 1581884Sstever@eecs.umich.edu elif 'm5' in cmd[0]: 1591884Sstever@eecs.umich.edu output_dir = cmd[cmd.index('-d')+1] 1601884Sstever@eecs.umich.edu else: 1611884Sstever@eecs.umich.edu output_dir = None 1621884Sstever@eecs.umich.edu except (ValueError, IndexError): 1631884Sstever@eecs.umich.edu # no big deal if there's no '-o'/'-d' or if it's the final argument 1641884Sstever@eecs.umich.edu output_dir = None 1651884Sstever@eecs.umich.edu 1661884Sstever@eecs.umich.edu if output_dir: 1671884Sstever@eecs.umich.edu secs_waited = 0 1681884Sstever@eecs.umich.edu while not shell.dir_exists(output_dir) and secs_waited < 45: 1691884Sstever@eecs.umich.edu time.sleep(5) 1701884Sstever@eecs.umich.edu secs_waited += 5 1711884Sstever@eecs.umich.edu if secs_waited > 10: 1721884Sstever@eecs.umich.edu print "waited", secs_waited, "seconds for", output_dir 1731884Sstever@eecs.umich.edu 1741884Sstever@eecs.umich.edu # run command 1751884Sstever@eecs.umich.edu if options.stdout_file: 1761884Sstever@eecs.umich.edu cmd += ['>', options.stdout_file] 1771884Sstever@eecs.umich.edu if options.stderr_file: 1781884Sstever@eecs.umich.edu cmd += ['2>', options.stderr_file] 1791884Sstever@eecs.umich.edu try: 1801884Sstever@eecs.umich.edu (output, status) = shell.do_command(' '.join(cmd), options.cmd_timeout) 1811884Sstever@eecs.umich.edu except pexpect.TIMEOUT: 1821884Sstever@eecs.umich.edu print >>sys.stderr, "%s: command timed out after %d seconds." \ 1831884Sstever@eecs.umich.edu % (progname, options.cmd_timeout) 1841884Sstever@eecs.umich.edu shell.sendline('~.') # qsub/ssh termination escape sequence 1851884Sstever@eecs.umich.edu shell.close(wait=True) 1861884Sstever@eecs.umich.edu status = 3 1871884Sstever@eecs.umich.edu if output: 1881884Sstever@eecs.umich.edu print output 1891884Sstever@eecs.umich.edu 1901884Sstever@eecs.umich.edufinally: 1911884Sstever@eecs.umich.edu # end job 1921884Sstever@eecs.umich.edu if shell.isalive(): 1931884Sstever@eecs.umich.edu shell.sendline('exit') 1941884Sstever@eecs.umich.edu shell.expect('qsub: job .* completed\r\n') 1951884Sstever@eecs.umich.edu shell.close(wait=True) 1961884Sstever@eecs.umich.edu 1971884Sstever@eecs.umich.edu # if there was an error, log the output even if not requested 1981884Sstever@eecs.umich.edu if status != 0 or options.save_log: 1991884Sstever@eecs.umich.edu log = file('qdo-log.' + str(os.getpid()), 'w') 2001884Sstever@eecs.umich.edu log.write(shell.full_output) 2011884Sstever@eecs.umich.edu log.close() 2021884Sstever@eecs.umich.edu 2031884Sstever@eecs.umich.edudel shell 2041884Sstever@eecs.umich.edu 2051884Sstever@eecs.umich.edusys.exit(status) 206