238 lines
8.8 KiB
Python
Executable file
238 lines
8.8 KiB
Python
Executable file
#! /usr/bin/env python
|
|
|
|
# Copyright (c) 2004-2005, 2007 The Regents of The University of Michigan
|
|
# All rights reserved.
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are
|
|
# met: redistributions of source code must retain the above copyright
|
|
# notice, this list of conditions and the following disclaimer;
|
|
# redistributions in binary form must reproduce the above copyright
|
|
# notice, this list of conditions and the following disclaimer in the
|
|
# documentation and/or other materials provided with the distribution;
|
|
# neither the name of the copyright holders nor the names of its
|
|
# contributors may be used to endorse or promote products derived from
|
|
# this software without specific prior written permission.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
|
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
|
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
|
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
|
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
|
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
|
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
|
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
|
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
|
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
#
|
|
# Authors: Steve Reinhardt
|
|
# Ali Saidi
|
|
|
|
# Important!
|
|
# This script expects a simple $ prompt, if you are using a shell other than
|
|
# sh which defaults to this you'll need to add something like the following
|
|
# to your bashrc/bash_profile script:
|
|
#if [ "$OAR_USER" = "xxxx" ]; then
|
|
# PS1='$ '
|
|
|
|
|
|
import sys
|
|
import os
|
|
import re
|
|
import time
|
|
import optparse
|
|
|
|
import pexpect
|
|
|
|
progname = os.path.basename(sys.argv[0])
|
|
|
|
usage = "%prog [options] command [command arguments]"
|
|
optparser = optparse.OptionParser(usage=usage)
|
|
optparser.allow_interspersed_args=False
|
|
optparser.add_option('-e', dest='stderr_file',
|
|
help='command stderr output file')
|
|
optparser.add_option('-o', dest='stdout_file',
|
|
help='command stdout output file')
|
|
optparser.add_option('-l', dest='save_log', action='store_true',
|
|
help='save oarsub output log file')
|
|
optparser.add_option('-N', dest='job_name',
|
|
help='oarsub job name')
|
|
optparser.add_option('-q', dest='dest_queue',
|
|
help='oarsub destination queue')
|
|
optparser.add_option('--qwait', dest='oarsub_timeout', type='int',
|
|
help='oarsub queue wait timeout', default=30*60)
|
|
optparser.add_option('-t', dest='cmd_timeout', type='int',
|
|
help='command execution timeout', default=600*60)
|
|
|
|
(options, cmd) = optparser.parse_args()
|
|
|
|
if cmd == []:
|
|
print >>sys.stderr, "%s: missing command" % progname
|
|
sys.exit(1)
|
|
|
|
# If we want to do this, need to add check here to make sure cmd[0] is
|
|
# a valid PBS job name, else oarsub will die on us.
|
|
#
|
|
#if not options.job_name:
|
|
# options.job_name = cmd[0]
|
|
|
|
cwd = os.getcwd()
|
|
|
|
# Deal with systems where /n is a symlink to /.automount
|
|
if cwd.startswith('/.automount/'):
|
|
cwd = cwd.replace('/.automount/', '/n/', 1)
|
|
|
|
if not cwd.startswith('/n/poolfs/'):
|
|
print >>sys.stderr, "Error: current directory must be under /n/poolfs."
|
|
sys.exit(1)
|
|
|
|
# The Shell class wraps pexpect.spawn with some handy functions that
|
|
# assume the thing on the other end is a Bourne/bash shell.
|
|
class Shell(pexpect.spawn):
|
|
# Regexp to match the shell prompt. We change the prompt to
|
|
# something fixed and distinctive to make it easier to match
|
|
# reliably.
|
|
prompt_re = re.compile('qdo\$ ')
|
|
|
|
def __init__(self, cmd):
|
|
# initialize base pexpect.spawn object
|
|
try:
|
|
pexpect.spawn.__init__(self, cmd)
|
|
except pexpect.ExceptionPexpect, exc:
|
|
print "%s:" % progname, exc
|
|
sys.exit(1)
|
|
# full_output accumulates the full output of the session
|
|
self.full_output = ""
|
|
self.quick_timeout = 15
|
|
# wait for a prompt, then change it
|
|
try:
|
|
self.expect('\$ ', options.oarsub_timeout)
|
|
except pexpect.TIMEOUT:
|
|
print >>sys.stderr, "%s: oarsub timed out." % progname
|
|
self.kill(9)
|
|
self.safe_close()
|
|
sys.exit(1)
|
|
self.do_command('unset PROMPT_COMMAND; PS1="qdo$ "')
|
|
|
|
# version of expect that updates full_output too
|
|
def expect(self, regexp, timeout = -1):
|
|
pexpect.spawn.expect(self, regexp, timeout)
|
|
self.full_output += self.before + self.after
|
|
|
|
# Just issue a command and wait for the next prompt.
|
|
# Returns a string containing the output of the command.
|
|
def do_bare_command(self, cmd, timeout = -1):
|
|
global full_output
|
|
self.sendline(cmd)
|
|
# read back the echo of the command
|
|
self.readline()
|
|
# wait for the next prompt
|
|
self.expect(self.prompt_re, timeout)
|
|
output = self.before.rstrip()
|
|
return output
|
|
|
|
# Issue a command, then query its exit status.
|
|
# Returns a (string, int) tuple with the command output and the status.
|
|
def do_command(self, cmd, timeout = -1):
|
|
# do the command itself
|
|
output = self.do_bare_command(cmd, timeout)
|
|
# collect status
|
|
status = int(self.do_bare_command("echo $?", self.quick_timeout))
|
|
return (output, status)
|
|
|
|
# Check to see if the given directory exists.
|
|
def dir_exists(self, dirname):
|
|
(output, status) = shell.do_command('[ -d %s ]' % dirname,
|
|
self.quick_timeout)
|
|
return status == 0
|
|
|
|
# Don't actually try to close it.. just wait until it closes by itself
|
|
# We can't actually kill the pid which is what it's trying to do, and if
|
|
# we call wait we could be in an unfortunate situation of it printing input
|
|
# right as we call wait, so the input is never read and the process never ends
|
|
def safe_close(self):
|
|
count = 0
|
|
while self.isalive() and count < 10:
|
|
time.sleep(1)
|
|
self.close(force=False)
|
|
|
|
# Spawn the interactive pool job.
|
|
|
|
# Hack to do link on poolfs... disabled for now since
|
|
# compiler/linker/library versioning problems between poolfs and
|
|
# nodes. May never work since poolfs is x86-64 and nodes are 32-bit.
|
|
if False and len(cmd) > 50:
|
|
shell_cmd = 'ssh -t poolfs /bin/sh -l'
|
|
print "%s: running %s on poolfs" % (progname, cmd[0])
|
|
else:
|
|
shell_cmd = 'oarsub -I'
|
|
if options.job_name:
|
|
shell_cmd += ' -n "%s"' % options.job_name
|
|
if options.dest_queue:
|
|
shell_cmd += ' -q ' + options.dest_queue
|
|
shell_cmd += ' -d %s' % cwd
|
|
|
|
shell = Shell(shell_cmd)
|
|
|
|
try:
|
|
# chdir to cwd
|
|
(output, status) = shell.do_command('cd ' + cwd)
|
|
|
|
if status != 0:
|
|
raise OSError, "Can't chdir to %s" % cwd
|
|
|
|
# wacky hack: sometimes scons will create an output directory then
|
|
# fork a job to generate files in that directory, and the job will
|
|
# get run before the directory creation propagates through NFS.
|
|
# This hack looks for a '-o' option indicating an output file and
|
|
# waits for the corresponding directory to appear if necessary.
|
|
try:
|
|
if 'cc' in cmd[0] or 'g++' in cmd[0]:
|
|
output_dir = os.path.dirname(cmd[cmd.index('-o')+1])
|
|
elif 'm5' in cmd[0]:
|
|
output_dir = cmd[cmd.index('-d')+1]
|
|
else:
|
|
output_dir = None
|
|
except (ValueError, IndexError):
|
|
# no big deal if there's no '-o'/'-d' or if it's the final argument
|
|
output_dir = None
|
|
|
|
if output_dir:
|
|
secs_waited = 0
|
|
while not shell.dir_exists(output_dir) and secs_waited < 90:
|
|
time.sleep(5)
|
|
secs_waited += 5
|
|
if secs_waited > 30:
|
|
print "waited", secs_waited, "seconds for", output_dir
|
|
|
|
# run command
|
|
if options.stdout_file:
|
|
cmd += ['>', options.stdout_file]
|
|
if options.stderr_file:
|
|
cmd += ['2>', options.stderr_file]
|
|
try:
|
|
(output, status) = shell.do_command(' '.join(cmd), options.cmd_timeout)
|
|
except pexpect.TIMEOUT:
|
|
print >>sys.stderr, "%s: command timed out after %d seconds." \
|
|
% (progname, options.cmd_timeout)
|
|
shell.sendline('~.') # oarsub/ssh termination escape sequence
|
|
shell.safe_close()
|
|
status = 3
|
|
if output:
|
|
print output
|
|
finally:
|
|
# end job
|
|
if shell.isalive():
|
|
shell.sendline('exit')
|
|
shell.expect('Disconnected from OAR job .*')
|
|
shell.safe_close()
|
|
|
|
# if there was an error, log the output even if not requested
|
|
if status != 0 or options.save_log:
|
|
log = file('qdo-log.' + str(os.getpid()), 'w')
|
|
log.write(shell.full_output)
|
|
log.close()
|
|
del shell
|
|
|
|
sys.exit(status)
|