qdo revision 5147:a7b91336a3fc
1#! /usr/bin/env python 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