utils/parallel.py

# 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')