util/qdo

   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 # Important!
  30 # This script expects a simple $ prompt, if you are using a shell other than
  31 # sh which defaults to this you'll need to add something like the following
  32 # to your bashrc/bash_profile script:
  33 #if [ "$OAR_USER" = "xxxx" ]; then
  34 #   PS1='$ '
  35
  36
  37 import sys
  38 import os
  39 import re
  40 import time
  41 import optparse
  42
  43 import pexpect
  44
  45 progname = os.path.basename(sys.argv[0])
  46
  47 usage = "%prog [options] command [command arguments]"
  48 optparser = optparse.OptionParser(usage=usage)
  49 optparser.allow_interspersed_args=False
  50 optparser.add_option('-e', dest='stderr_file',
  51                      help='command stderr output file')
  52 optparser.add_option('-o', dest='stdout_file',
  53                      help='command stdout output file')
  54 optparser.add_option('-l', dest='save_log', action='store_true',
  55                      help='save oarsub output log file')
  56 optparser.add_option('-N', dest='job_name',
  57                      help='oarsub job name')
  58 optparser.add_option('-q', dest='dest_queue',
  59                      help='oarsub destination queue')
  60 optparser.add_option('--qwait', dest='oarsub_timeout', type='int',
  61                      help='oarsub queue wait timeout', default=30*60)
  62 optparser.add_option('-t', dest='cmd_timeout', type='int',
  63                      help='command execution timeout', default=600*60)
  64
  65 (options, cmd) = optparser.parse_args()
  66
  67 if cmd == []:
  68     print >>sys.stderr, "%s: missing command" % progname
  69     sys.exit(1)
  70
  71 # If we want to do this, need to add check here to make sure cmd[0] is
  72 # a valid PBS job name, else oarsub will die on us.
  73 #
  74 #if not options.job_name:
  75 #    options.job_name = cmd[0]
  76
  77 cwd = os.getcwd()
  78
  79 # Deal with systems where /n is a symlink to /.automount
  80 if cwd.startswith('/.automount/'):
  81     cwd = cwd.replace('/.automount/', '/n/', 1)
  82
  83 if not cwd.startswith('/n/poolfs/'):
  84     print >>sys.stderr, "Error: current directory must be under /n/poolfs."
  85     sys.exit(1)
  86
  87 # The Shell class wraps pexpect.spawn with some handy functions that
  88 # assume the thing on the other end is a Bourne/bash shell.
  89 class Shell(pexpect.spawn):
  90     # Regexp to match the shell prompt.  We change the prompt to
  91     # something fixed and distinctive to make it easier to match
  92     # reliably.
  93     prompt_re = re.compile('qdo\$ ')
  94
  95     def __init__(self, cmd):
  96         # initialize base pexpect.spawn object
  97         try:
  98             pexpect.spawn.__init__(self, cmd)
  99         except pexpect.ExceptionPexpect, exc:
 100             print "%s:" % progname, exc
 101             sys.exit(1)
 102         # full_output accumulates the full output of the session
 103         self.full_output = ""
 104         self.quick_timeout = 15
 105         # wait for a prompt, then change it
 106         try:
 107             self.expect('\$ ', options.oarsub_timeout)
 108         except pexpect.TIMEOUT:
 109             print >>sys.stderr, "%s: oarsub timed out." % progname
 110             self.kill(9)
 111             self.safe_close()
 112             sys.exit(1)
 113         self.do_command('unset PROMPT_COMMAND; PS1="qdo$ "')
 114
 115     # version of expect that updates full_output too
 116     def expect(self, regexp, timeout = -1):
 117         pexpect.spawn.expect(self, regexp, timeout)
 118         self.full_output += self.before + self.after
 119
 120     # Just issue a command and wait for the next prompt.
 121     # Returns a string containing the output of the command.
 122     def do_bare_command(self, cmd, timeout = -1):
 123         global full_output
 124         self.sendline(cmd)
 125         # read back the echo of the command
 126         self.readline()
 127         # wait for the next prompt
 128         self.expect(self.prompt_re, timeout)
 129         output = self.before.rstrip()
 130         return output
 131
 132     # Issue a command, then query its exit status.
 133     # Returns a (string, int) tuple with the command output and the status.
 134     def do_command(self, cmd, timeout = -1):
 135         # do the command itself
 136         output = self.do_bare_command(cmd, timeout)
 137         # collect status
 138         status = int(self.do_bare_command("echo $?", self.quick_timeout))
 139         return (output, status)
 140
 141     # Check to see if the given directory exists.
 142     def dir_exists(self, dirname):
 143         (output, status) = shell.do_command('[ -d %s ]' % dirname,
 144                                             self.quick_timeout)
 145         return status == 0
 146
 147     # Don't actually try to close it.. just wait until it closes by itself
 148     # We can't actually kill the pid which is what it's trying to do, and if
 149     # we call wait we could be in an unfortunate situation of it printing input
 150     # right as we call wait, so the input is never read and the process never ends
 151     def safe_close(self):
 152         count = 0
 153         while self.isalive() and count < 10:
 154             time.sleep(1)
 155         self.close(force=False)
 156
 157 # Spawn the interactive pool job.
 158
 159 # Hack to do link on poolfs... disabled for now since
 160 # compiler/linker/library versioning problems between poolfs and
 161 # nodes.  May never work since poolfs is x86-64 and nodes are 32-bit.
 162 if False and len(cmd) > 50:
 163     shell_cmd = 'ssh -t poolfs /bin/sh -l'
 164     print "%s: running %s on poolfs" % (progname, cmd[0])
 165 else:
 166     shell_cmd = 'oarsub -I'
 167     if options.job_name:
 168         shell_cmd += ' -n "%s"' % options.job_name
 169     if options.dest_queue:
 170         shell_cmd += ' -q ' + options.dest_queue
 171     shell_cmd += ' -d %s' % cwd
 172
 173 shell = Shell(shell_cmd)
 174
 175 try:
 176     # chdir to cwd
 177     (output, status) = shell.do_command('cd ' + cwd)
 178
 179     if status != 0:
 180         raise OSError, "Can't chdir to %s" % cwd
 181
 182     # wacky hack: sometimes scons will create an output directory then
 183     # fork a job to generate files in that directory, and the job will
 184     # get run before the directory creation propagates through NFS.
 185     # This hack looks for a '-o' option indicating an output file and
 186     # waits for the corresponding directory to appear if necessary.
 187     try:
 188         if 'cc' in cmd[0] or 'g++' in cmd[0]:
 189             output_dir = os.path.dirname(cmd[cmd.index('-o')+1])
 190         elif 'm5' in cmd[0]:
 191             output_dir = cmd[cmd.index('-d')+1]
 192         else:
 193             output_dir = None
 194     except (ValueError, IndexError):
 195         # no big deal if there's no '-o'/'-d' or if it's the final argument
 196         output_dir = None
 197
 198     if output_dir:
 199         secs_waited = 0
 200         while not shell.dir_exists(output_dir) and secs_waited < 90:
 201             time.sleep(5)
 202             secs_waited += 5
 203         if secs_waited > 30:
 204             print "waited", secs_waited, "seconds for", output_dir
 205
 206     # run command
 207     if options.stdout_file:
 208         cmd += ['>', options.stdout_file]
 209     if options.stderr_file:
 210         cmd += ['2>', options.stderr_file]
 211     try:
 212         (output, status) = shell.do_command(' '.join(cmd), options.cmd_timeout)
 213     except pexpect.TIMEOUT:
 214             print >>sys.stderr, "%s: command timed out after %d seconds." \
 215                   % (progname, options.cmd_timeout)
 216             shell.sendline('~.') # oarsub/ssh termination escape sequence
 217             shell.safe_close()
 218             status = 3
 219     if output:
 220         print output
 221 finally:
 222     # end job
 223     if shell.isalive():
 224         shell.sendline('exit')
 225         shell.expect('Disconnected from OAR job .*')
 226         shell.safe_close()
 227
 228     # if there was an error, log the output even if not requested
 229     if status != 0 or options.save_log:
 230         log = file('qdo-log.' + str(os.getpid()), 'w')
 231         log.write(shell.full_output)
 232         log.close()
 233 del shell
 234
 235 sys.exit(status)