greylag-mp.py

   1 #!/usr/bin/env python
   2
   3 '''Run a greylag job split across multiple local processes.  This is not for
   4 use on a cluster, but useful if you have just one multi-CPU machine, or for
   5 debugging.
   6 '''
   7
   8 __copyright__ = '''
   9     greylag, Copyright (C) 2006-2007, Stowers Institute for Medical Research
  10
  11     This program is free software; you can redistribute it and/or modify
  12     it under the terms of the GNU General Public License as published by
  13     the Free Software Foundation; either version 2 of the License, or
  14     (at your option) any later version.
  15
  16     This program is distributed in the hope that it will be useful,
  17     but WITHOUT ANY WARRANTY; without even the implied warranty of
  18     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  19     GNU General Public License for more details.
  20
  21     You should have received a copy of the GNU General Public License along
  22     with this program; if not, write to the Free Software Foundation, Inc.,
  23     51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  24 '''
  25
  26
  27 # This might better be implemented as a shell script under Unix.  It's done in
  28 # Python here as a demonstration, and so that greylag can be easily run on
  29 # multiple CPUs under Windows, as X!Tandem can.
  30
  31 # This is not intended to be used on cluster nodes.  In particular, like
  32 # X!Tandem, it just divides the spectra into N parts and then processes them
  33 # separately, making no further attempt at load balancing.
  34
  35
  36 import os
  37 import os.path
  38 from socket import gethostname
  39 import subprocess
  40 import sys
  41
  42
  43 GREYLAG_PROGRAM = 'greylag'
  44
  45
  46 def warn(s):
  47     print >> sys.stderr, 'warning:', s
  48 def error(s):
  49     sys.exit('error: ' + s)
  50
  51 def usage():
  52     print >> sys.stderr, ('Usage: %s <processes> <greylag-options-and-args>...'
  53                           '\n\n'
  54                           '%s\n(see "greylag.py --help" for more information)'
  55                           % (os.path.basename(sys.argv[0]), __doc__))
  56     sys.exit()
  57
  58
  59 def run_parts(processes, partprefix, args):
  60     ret = subprocess.call([GREYLAG_PROGRAM, '--part-split=%s' % processes,
  61                            '--part-prefix=%s' % partprefix] + args)
  62     if ret:
  63         error("part split failed")
  64
  65     # list of currently running subprocess.Popen objects
  66     subprocs = []
  67     try:
  68         for n in range(processes):
  69             p = subprocess.Popen([GREYLAG_PROGRAM,
  70                                   '--part=%sof%s' % (n+1, processes),
  71                                   '--part-prefix=%s' % partprefix] + args)
  72             subprocs.append(p)
  73         while subprocs:
  74             p = subprocs.pop()
  75             if p.wait():
  76                 raise EnvironmentError("error status")
  77     except EnvironmentError, e:
  78         error("part process failed [%s]" % e)
  79     finally:
  80         # upon error, try to kill any remaining processes
  81         try:
  82             import signal
  83             for p in subprocs:
  84                 os.kill(p.pid, signal.SIGINT)
  85         except Exception:
  86             pass
  87
  88     ret = subprocess.call([GREYLAG_PROGRAM, '--part-merge=%s' % processes,
  89                            '--part-prefix=%s' % partprefix] + args)
  90     if ret:
  91         error("part merge failed")
  92
  93
  94 def main():
  95     args = sys.argv[1:]
  96     if len(args) < 2 or '-h' in args or '--help' in args:
  97         usage()
  98     try:
  99         processes = int(args[0])
 100         if processes < 1:
 101             raise ValueError
 102     except ValueError:
 103         error("<processes> must be a positive integer")
 104     # This limit can be raised by setting environment variable.
 105     GREYLAGPROCESSES = os.environ.get('GREYLAGPROCESSES', 8)
 106     if processes > GREYLAGPROCESSES:
 107         error("current limit is %s processes (see source code for more)"
 108               % GREYLAGPROCESSES)
 109     if any(x for x in args if x.startswith('--part')):
 110         error("no --part options may be specified")
 111
 112     # try to generate a unique prefix, to avoid (unlikely) collision
 113     partprefix = '#greylag-part-%s-%s' % (gethostname(), os.getpid())
 114
 115     try:
 116         run_parts(processes, partprefix, args[1:])
 117     finally:
 118         # try to clean up temporary part files
 119         for fn in os.listdir('.'):
 120             if fn.startswith(partprefix):
 121                 try:
 122                     os.remove(fn)
 123                 except IOError, e:
 124                     warn("could not remove '%s' [%s]" % (fn, e))
 125
 126
 127 if __name__ == '__main__':
 128     main()