util/batch/job.py

   1 #!/usr/bin/env python2.7
   2 # Copyright (c) 2006 The Regents of The University of Michigan
   3 # All rights reserved.
   4 #
   5 # Redistribution and use in source and binary forms, with or without
   6 # modification, are permitted provided that the following conditions are
   7 # met: redistributions of source code must retain the above copyright
   8 # notice, this list of conditions and the following disclaimer;
   9 # redistributions in binary form must reproduce the above copyright
  10 # notice, this list of conditions and the following disclaimer in the
  11 # documentation and/or other materials provided with the distribution;
  12 # neither the name of the copyright holders nor the names of its
  13 # contributors may be used to endorse or promote products derived from
  14 # this software without specific prior written permission.
  15 #
  16 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  17 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  18 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  19 # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  20 # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  21 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
  22 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  23 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  24 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  25 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  26 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27
  28 import os, os.path, shutil, signal, socket, sys
  29 from os import environ as env
  30 from os.path import join as joinpath, expanduser
  31
  32 def date():
  33     import time
  34     return time.strftime('%a %b %e %H:%M:%S %Z %Y', time.localtime())
  35
  36 def cleandir(dir):
  37     for root, dirs, files in os.walk(dir, False):
  38         for name in files:
  39             os.remove(joinpath(root, name))
  40         for name in dirs:
  41             os.rmdir(joinpath(root, name))
  42
  43 class rsync:
  44     def __init__(self):
  45         self.sudo = False
  46         self.rsync = 'rsync'
  47         self.compress = False
  48         self.archive = True
  49         self.delete = False
  50         self.options = ''
  51
  52     def do(self, src, dst):
  53         args = []
  54         if self.sudo:
  55             args.append('sudo')
  56
  57         args.append(self.rsync)
  58         if (self.archive):
  59             args.append('-a')
  60         if (self.compress):
  61             args.append('-z')
  62         if (self.delete):
  63             args.append('--delete')
  64         if len(self.options):
  65             args.append(self.options)
  66         args.append(src)
  67         args.append(dst)
  68
  69         return os.spawnvp(os.P_WAIT, args[0], args)
  70
  71 class JobDir(object):
  72     def __init__(self, dir):
  73         self.dir = dir
  74
  75     def file(self, filename):
  76         return joinpath(self.dir, filename)
  77
  78     def create(self):
  79         if os.path.exists(self.dir):
  80             if not os.path.isdir(self.dir):
  81                 sys.exit('%s is not a directory.  Cannot build job' % self.dir)
  82         else:
  83             os.mkdir(self.dir)
  84
  85     def exists(self):
  86         return os.path.isdir(self.dir)
  87
  88     def clean(self):
  89         cleandir(self.dir)
  90
  91     def hasfile(self, filename):
  92         return os.path.isfile(self.file(filename))
  93
  94     def echofile(self, filename, string):
  95         filename = self.file(filename)
  96         try:
  97             f = file(filename, 'w')
  98             print >>f, string
  99             f.flush()
 100             f.close()
 101         except IOError,e:
 102             sys.exit(e)
 103
 104     def rmfile(self, filename):
 105         filename = self.file(filename)
 106         if os.path.isfile(filename):
 107             os.unlink(filename)
 108
 109     def readval(self, filename):
 110         filename = self.file(filename)
 111         f = file(filename, 'r')
 112         value = f.readline().strip()
 113         f.close()
 114         return value
 115
 116     def setstatus(self, string):
 117         filename = self.file('.status')
 118         try:
 119             f = file(filename, 'a')
 120             print >>f, string
 121             f.flush()
 122             f.close()
 123         except IOError,e:
 124             sys.exit(e)
 125
 126     def getstatus(self):
 127         filename = self.file('.status')
 128         try:
 129             f = file(filename, 'r')
 130         except IOError, e:
 131             return 'none'
 132
 133         # fast forward to the end
 134         for line in f: pass
 135
 136         # the first word on the last line is the status
 137         return line.split(' ')[0]
 138
 139     def __str__(self):
 140         return self.dir
 141
 142 if __name__ == '__main__':
 143     import platform
 144     binaries = { 'i686' : 'm5.i386',
 145                  'x86_64' : 'm5.amd64' }
 146     binary = binaries[platform.machine()]
 147
 148     cwd = os.getcwd()
 149     rootdir = env.setdefault('ROOTDIR', os.path.dirname(cwd))
 150     oar_jobid = int(env['OAR_JOBID'])
 151     oar_jobname = os.path.basename(cwd)
 152     #pbs_jobname = env['PBS_JOBNAME']
 153     basedir = joinpath(rootdir, 'Base')
 154     jobname = env.setdefault('JOBNAME', oar_jobname)
 155     jobfile = env.setdefault('JOBFILE', joinpath(rootdir, 'Test.py'))
 156     outdir = env.setdefault('OUTPUT_DIR', cwd)
 157     env['POOLJOB'] = 'True'
 158
 159     if os.path.isdir("/work"):
 160         workbase = "/work"
 161     else:
 162         workbase = "/tmp/"
 163
 164     workdir = joinpath(workbase, '%s.%s' % (env['USER'], oar_jobid))
 165     host = socket.gethostname()
 166
 167     os.umask(0022)
 168
 169     jobdir = JobDir(outdir)
 170
 171     started = date()
 172     jobdir.echofile('.running', started)
 173     jobdir.rmfile('.queued')
 174     jobdir.echofile('.host', host)
 175
 176     jobdir.setstatus('running on %s on %s' % (host, started))
 177
 178     if os.path.isdir(workdir):
 179         cleandir(workdir)
 180     else:
 181         os.mkdir(workdir)
 182
 183     if False and os.path.isdir('/z/dist'):
 184         sync = rsync()
 185         sync.delete = True
 186         sync.sudo = True
 187         sync.do('poolfs::dist/m5/', '/z/dist/m5/')
 188
 189     try:
 190         os.chdir(workdir)
 191     except OSError,e:
 192         sys.exit(e)
 193
 194     os.symlink(jobdir.file('output'), 'status.out')
 195
 196     args = [ joinpath(basedir, binary), joinpath(basedir, 'run.py') ]
 197     if not len(args):
 198         sys.exit("no arguments")
 199
 200     print 'starting job... %s' % started
 201     print ' '.join(args)
 202     print
 203     sys.stdout.flush()
 204
 205     childpid = os.fork()
 206     if not childpid:
 207         # Execute command
 208         sys.stdin.close()
 209         fd = os.open(jobdir.file("output"),
 210                      os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
 211         os.dup2(fd, sys.stdout.fileno())
 212         os.dup2(fd, sys.stderr.fileno())
 213         os.execvp(args[0], args)
 214
 215     def handler(signum, frame):
 216         if childpid != 0:
 217             os.kill(childpid, signum)
 218
 219     signal.signal(signal.SIGHUP, handler)
 220     signal.signal(signal.SIGINT, handler)
 221     signal.signal(signal.SIGQUIT, handler)
 222     signal.signal(signal.SIGTERM, handler)
 223     signal.signal(signal.SIGCONT, handler)
 224     signal.signal(signal.SIGUSR1, handler)
 225     signal.signal(signal.SIGUSR2, handler)
 226
 227     done = 0
 228     while not done:
 229         try:
 230             thepid,ec = os.waitpid(childpid, 0)
 231             if ec:
 232                 print 'Exit code ', ec
 233                 status = 'failure'
 234             else:
 235                 status = 'success'
 236             done = 1
 237         except OSError:
 238             pass
 239
 240     complete = date()
 241     print '\njob complete... %s' % complete
 242     jobdir.echofile('.%s' % status, complete)
 243     jobdir.rmfile('.running')
 244     jobdir.setstatus('%s on %s' % (status, complete))