parallel processing with timeout

joblib_parallel_with_timeout.py

# Source: https://github.com/joblib/joblib/pull/366#issuecomment-267603530
# Source2: https://github.com/joblib/joblib/issues/522
# Source3: https://github.com/mapillary/OpenSfM/blob/199f2dc4a535397afc8fdeb76d67053e76015b61/opensfm/reconstruction.py#L470
import time
import multiprocessing
import functools

from joblib import Parallel, delayed


def with_timeout(timeout):
    def decorator(decorated):
        @functools.wraps(decorated)
        def inner(*args, **kwargs):
            pool = multiprocessing.pool.ThreadPool(1)
            async_result = pool.apply_async(decorated, args, kwargs)
            try:
                return async_result.get(timeout)
            except multiprocessing.TimeoutError:
                return
        return inner
    return decorator


@with_timeout(0.5)
def func(to_return, sleep_time):
    time.sleep(sleep_time)
    return to_return


print(Parallel(n_jobs=2)(delayed(func)(to_return, sleep_time)
                         for to_return, sleep_time in enumerate([
                                 0.1, 0.7, 0.2, 0.3])))