pyrocko/src/parimap.py

# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------

try:
    import queue
except ImportError:
    import Queue as queue


import logging
import platform
import multiprocessing
import traceback
import errno


logger = logging.getLogger('pyrocko.parimap')


def worker(
        q_in, q_out, function, eprintignore, pshared,
        startup, startup_args, cleanup):

    kwargs = {}
    if pshared is not None:
        kwargs['pshared'] = pshared

    if startup is not None:
        startup(*startup_args)

    while True:
        i, args = q_in.get()
        if i is None:
            if cleanup is not None:
                cleanup()

            break

        res, exception = None, None
        try:
            res = function(*args, **kwargs)
        except Exception as e:
            if eprintignore is None or not isinstance(e, eprintignore):
                traceback.print_exc()
            exception = e
        q_out.put((i, res, exception))


def parimap(function, *iterables, **kwargs):
    assert all(
        k in (
            'nprocs', 'eprintignore', 'pshared', 'startup', 'startup_args',
            'cleanup')
        for k in kwargs.keys())

    nprocs = kwargs.get('nprocs', None)
    eprintignore = kwargs.get('eprintignore', 'all')
    pshared = kwargs.get('pshared', None)
    startup = kwargs.get('startup', None)
    startup_args = kwargs.get('startup_args', ())
    cleanup = kwargs.get('cleanup', None)

    if eprintignore == 'all':
        eprintignore = None

    if platform.system() == 'Windows' and nprocs != 1:
        logger.warn(
            'The parimap module relies on fork() for parallelism. This does '
            'not work on Windows. Using serial code.')
        nprocs = 1

    if nprocs == 1:
        iterables = list(map(iter, iterables))
        kwargs = {}
        if pshared is not None:
            kwargs['pshared'] = pshared

        if startup is not None:
            startup(*startup_args)

        try:
            while True:
                args = []
                for it in iterables:
                    try:
                        args.append(next(it))
                    except StopIteration:
                        return

                yield function(*args, **kwargs)

        finally:
            if cleanup is not None:
                cleanup()

        return

    if nprocs is None:
        nprocs = multiprocessing.cpu_count()

    q_in = multiprocessing.Queue(1)
    q_out = multiprocessing.Queue()

    procs = []

    results = []
    nrun = 0
    nwritten = 0
    iout = 0
    all_written = False
    error_ahead = False
    iterables = list(map(iter, iterables))
    while True:
        if nrun < nprocs and not all_written and not error_ahead:
            args = []
            for it in iterables:
                try:
                    args.append(next(it))
                except StopIteration:
                    pass

            if len(args) == len(iterables):
                if len(procs) < nrun + 1:
                    p = multiprocessing.Process(
                        target=worker,
                        args=(q_in, q_out, function, eprintignore, pshared,
                              startup, startup_args, cleanup))
                    p.daemon = True
                    p.start()
                    procs.append(p)

                q_in.put((nwritten, args))
                nwritten += 1
                nrun += 1
            else:
                all_written = True
                [q_in.put((None, None)) for p in procs]
                q_in.close()

        try:
            while nrun > 0:
                if nrun < nprocs and not all_written and not error_ahead:
                    results.append(q_out.get_nowait())
                else:
                    while True:
                        try:
                            results.append(q_out.get())
                            break
                        except IOError as e:
                            if e.errno != errno.EINTR:
                                raise

                nrun -= 1

        except queue.Empty:
            pass

        if results:
            results.sort()
            # check for error ahead to prevent further enqueuing
            if any(exc for (_, _, exc) in results):
                error_ahead = True

            while results:
                (i, r, exc) = results[0]
                if i == iout:
                    results.pop(0)
                    if exc is not None:
                        if not all_written:
                            [q_in.put((None, None)) for p in procs]
                            q_in.close()
                        raise exc
                    else:
                        yield r

                    iout += 1
                else:
                    break

        if all_written and nrun == 0:
            break

    [p.join() for p in procs]
    return