| # Lint as: python2, python3 |
| from __future__ import absolute_import |
| from __future__ import division |
| from __future__ import print_function |
| import os, sys |
| import six |
| from six.moves import range |
| |
| |
| class ParallelError(Exception): |
| def __init__(self, str, errors): |
| self.str = str |
| self.errors = errors |
| Exception.__init__(self, str) |
| |
| |
| class ParallelExecute(object): |
| def __init__(self, functions, max_simultaneous_procs=20): |
| """\ |
| This takes in a dictionary of functions which map to a set of |
| functions that they depend on. |
| |
| functions: This is either a list of or dictionary of functions to |
| be run. If it's a dictionary, the value should be a set |
| of other functions this function is dependent on. If its |
| a list (or tuple or anything iterable that returns a |
| single element each iteration), then it's assumed that |
| there are no dependencies. |
| |
| max_simultaneous_procs: Throttle the number of processes we have |
| running at once. |
| """ |
| if not isinstance(functions, dict): |
| function_list = functions |
| functions = {} |
| for fn in function_list: |
| functions[fn] = set() |
| |
| dependents = {} |
| for fn, deps in six.iteritems(functions): |
| dependents[fn] = [] |
| for fn, deps in six.iteritems(functions): |
| for dep in deps: |
| dependents[dep].append(fn) |
| |
| self.max_procs = max_simultaneous_procs |
| self.functions = functions |
| self.dependents = dependents |
| self.pid_map = {} |
| self.ready_to_run = [] |
| |
| |
| def _run(self, function): |
| self.functions.pop(function) |
| pid = os.fork() |
| if pid: |
| self.pid_map[pid] = function |
| else: |
| function() |
| sys.exit(0) |
| |
| |
| def run_until_completion(self): |
| for fn, deps in six.iteritems(self.functions): |
| if len(deps) == 0: |
| self.ready_to_run.append(fn) |
| |
| errors = [] |
| while len(self.pid_map) > 0 or len(self.ready_to_run) > 0: |
| max_allowed = self.max_procs - len(self.pid_map) |
| max_able = len(self.ready_to_run) |
| for i in range(min(max_allowed, max_able)): |
| self._run(self.ready_to_run.pop()) |
| |
| # Handle one proc that's finished. |
| pid, status = os.wait() |
| fn = self.pid_map.pop(pid) |
| if status != 0: |
| errors.append("%s failed" % fn.__name__) |
| continue |
| |
| for dependent in self.dependents[fn]: |
| self.functions[dependent].remove(fn) |
| if len(self.functions[dependent]) == 0: |
| self.ready_to_run.append(dependent) |
| |
| if len(self.functions) > 0 and len(errors) == 0: |
| errors.append("Deadlock detected") |
| |
| if len(errors) > 0: |
| msg = "Errors occurred during execution:" |
| msg = '\n'.join([msg] + errors) |
| raise ParallelError(msg, errors) |
| |
| |
| def redirect_io(log_file='/dev/null'): |
| # Always redirect stdin. |
| in_fd = os.open('/dev/null', os.O_RDONLY) |
| try: |
| os.dup2(in_fd, 0) |
| finally: |
| os.close(in_fd) |
| |
| out_fd = os.open(log_file, os.O_WRONLY | os.O_CREAT) |
| try: |
| os.dup2(out_fd, 2) |
| os.dup2(out_fd, 1) |
| finally: |
| os.close(out_fd) |
| |
| sys.stdin = os.fdopen(0, 'r') |
| sys.stdout = os.fdopen(1, 'w') |
| sys.stderr = os.fdopen(2, 'w') |