timed_parmap.py

#!/usr/bin/env python

"""
    timed_parmap.py
"""

import sys
from joblib import delayed
from concurrent.futures import ProcessPoolExecutor, as_completed

def __delayed_runner(job):
    return job[0](*job[1], **job[2])

def timed_parmap(jobs, n_jobs=2, timeout=1):
    with ProcessPoolExecutor(max_workers=n_jobs) as runner:
        jobs = {runner.submit(__delayed_runner, job) for job in jobs}
        try:
            for job in as_completed(jobs, timeout=timeout):
                yield job.result()
        except:
            print('timed_parallel: timeout at %d' % timeout, file=sys.stderr)
            return