michael@0: #!/usr/bin/env python
michael@0: 
michael@0: import argparse
michael@0: import collections
michael@0: import ConfigParser
michael@0: import multiprocessing
michael@0: import time
michael@0: 
michael@0: ProcessNode = collections.namedtuple('ProcessNode', ['maxtime', 'children'])
michael@0: 
michael@0: class ProcessLauncher(object):
michael@0: 
michael@0:     """ Create and Launch process trees specified by a '.ini' file
michael@0: 
michael@0:         Typical .ini file accepted by this class :
michael@0: 
michael@0:         [main]
michael@0:         children=c1, 1*c2, 4*c3
michael@0:         maxtime=10
michael@0: 
michael@0:         [c1]
michael@0:         children= 2*c2, c3
michael@0:         maxtime=20
michael@0: 
michael@0:         [c2]
michael@0:         children=3*c3
michael@0:         maxtime=5
michael@0: 
michael@0:         [c3]
michael@0:         maxtime=3
michael@0: 
michael@0:         This generates a process tree of the form:
michael@0:             [main]
michael@0:                 |---[c1]
michael@0:                 |     |---[c2]
michael@0:                 |     |     |---[c3]
michael@0:                 |     |     |---[c3]
michael@0:                 |     |     |---[c3]
michael@0:                 |     |
michael@0:                 |     |---[c2]
michael@0:                 |     |     |---[c3]
michael@0:                 |     |     |---[c3]
michael@0:                 |     |     |---[c3]
michael@0:                 |     |
michael@0:                 |     |---[c3]
michael@0:                 |
michael@0:                 |---[c2]
michael@0:                 |     |---[c3]
michael@0:                 |     |---[c3]
michael@0:                 |     |---[c3]
michael@0:                 |
michael@0:                 |---[c3]
michael@0:                 |---[c3]
michael@0:                 |---[c3]
michael@0: 
michael@0:         Caveat: The section names cannot contain a '*'(asterisk) or a ','(comma)
michael@0:         character as these are used as delimiters for parsing.
michael@0:     """
michael@0: 
michael@0:     # Unit time for processes in seconds
michael@0:     UNIT_TIME = 1
michael@0: 
michael@0:     def __init__(self, manifest, verbose=False):
michael@0:         """
michael@0:         Parses the manifest and stores the information about the process tree
michael@0:         in a format usable by the class.
michael@0: 
michael@0:         Raises IOError if :
michael@0:             - The path does not exist
michael@0:             - The file cannot be read
michael@0:         Raises ConfigParser.*Error if:
michael@0:             - Files does not contain section headers
michael@0:             - File cannot be parsed because of incorrect specification
michael@0: 
michael@0:         :param manifest: Path to the manifest file that contains the
michael@0:         configuration for the process tree to be launched
michael@0:         :verbose: Print the process start and end information.
michael@0:         Genrates a lot of output. Disabled by default.
michael@0:         """
michael@0: 
michael@0:         self.verbose=verbose
michael@0: 
michael@0:         # Children is a dictionary used to store information from the,
michael@0:         # Configuration file in a more usable format.
michael@0:         # Key : string contain the name of child process
michael@0:         # Value : A Named tuple of the form (max_time, (list of child processes of Key))
michael@0:         #   Where each child process is a list of type: [count to run, name of child]
michael@0:         self.children = {}
michael@0: 
michael@0: 
michael@0:         cfgparser = ConfigParser.ConfigParser()
michael@0: 
michael@0:         if not cfgparser.read(manifest):
michael@0:             raise IOError('The manifest %s could not be found/opened', manifest)
michael@0: 
michael@0:         sections = cfgparser.sections()
michael@0:         for section in sections:
michael@0:             # Maxtime is a mandatory option
michael@0:             # ConfigParser.NoOptionError is raised if maxtime does not exist
michael@0:             if '*' in section or ',' in section:
michael@0:                 raise ConfigParser.ParsingError('%s is not a valid section name. Section names cannot contain a \'*\' or \',\'.' % section)
michael@0:             m_time = cfgparser.get(section, 'maxtime')
michael@0:             try:
michael@0:                 m_time = int(m_time)
michael@0:             except ValueError:
michael@0:                 raise ValueError('Expected maxtime to be an integer, specified %s' % m_time)
michael@0: 
michael@0:             # No children option implies there are no further children
michael@0:             # Leaving the children option blank is an error.
michael@0:             try:
michael@0:                 c = cfgparser.get(section, 'children')
michael@0:                 if not c:
michael@0:                     # If children is an empty field, assume no children
michael@0:                     children = None
michael@0: 
michael@0:                 else:
michael@0:                     # Tokenize chilren field, ignore empty strings
michael@0:                     children = [[y.strip() for y in x.strip().split('*', 1)]
michael@0:                                 for x in c.split(',') if x]
michael@0:                     try:
michael@0:                         for i, child in enumerate(children):
michael@0:                             # No multiplicate factor infront of a process implies 1
michael@0:                             if len(child) == 1:
michael@0:                                 children[i] = [1, child[0]]
michael@0:                             else:
michael@0:                                 children[i][0] = int(child[0])
michael@0: 
michael@0:                             if children[i][1] not in sections:
michael@0:                                 raise ConfigParser.ParsingError('No section corresponding to child %s' % child[1])
michael@0:                     except ValueError:
michael@0:                         raise ValueError('Expected process count to be an integer, specified %s' % child[0])
michael@0: 
michael@0:             except ConfigParser.NoOptionError:
michael@0:                 children = None
michael@0:             pn = ProcessNode(maxtime=m_time,
michael@0:                              children=children)
michael@0:             self.children[section] = pn
michael@0: 
michael@0:     def run(self):
michael@0:         """
michael@0:         This function launches the process tree.
michael@0:         """
michael@0:         self._run('main', 0)
michael@0: 
michael@0:     def _run(self, proc_name, level):
michael@0:         """
michael@0:         Runs the process specified by the section-name `proc_name` in the manifest file.
michael@0:         Then makes calls to launch the child processes of `proc_name`
michael@0: 
michael@0:         :param proc_name: File name of the manifest as a string.
michael@0:         :param level: Depth of the current process in the tree.
michael@0:         """
michael@0:         if proc_name not in self.children.keys():
michael@0:             raise IOError("%s is not a valid process" % proc_name)
michael@0: 
michael@0:         maxtime = self.children[proc_name].maxtime
michael@0:         if self.verbose:
michael@0:             print "%sLaunching %s for %d*%d seconds" % (" "*level, proc_name, maxtime, self.UNIT_TIME)
michael@0: 
michael@0:         while self.children[proc_name].children:
michael@0:             child = self.children[proc_name].children.pop()
michael@0: 
michael@0:             count, child_proc = child
michael@0:             for i in range(count):
michael@0:                 p = multiprocessing.Process(target=self._run, args=(child[1], level+1))
michael@0:                 p.start()
michael@0: 
michael@0:         self._launch(maxtime)
michael@0:         if self.verbose:
michael@0:             print "%sFinished %s" % (" "*level, proc_name)
michael@0: 
michael@0:     def _launch(self, running_time):
michael@0:         """
michael@0:         Create and launch a process and idles for the time specified by
michael@0:         `running_time`
michael@0: 
michael@0:         :param running_time: Running time of the process in seconds.
michael@0:         """
michael@0:         elapsed_time = 0
michael@0: 
michael@0:         while elapsed_time < running_time:
michael@0:             time.sleep(self.UNIT_TIME)
michael@0:             elapsed_time += self.UNIT_TIME
michael@0: 
michael@0: if __name__ == '__main__':
michael@0: 
michael@0:     parser = argparse.ArgumentParser()
michael@0:     parser.add_argument("manifest", help="Specify the configuration .ini file")
michael@0:     args = parser.parse_args()
michael@0: 
michael@0:     proclaunch = ProcessLauncher(args.manifest)
michael@0:     proclaunch.run()