[tools] tqdm joblib

function

import contextlib
import joblib
from tqdm import tqdm
from joblib import Parallel, delayed

#Code from this nice Stack Overflow discussion: https://stackoverflow.com/questions/37804279/how-can-we-use-tqdm-in-a-parallel-execution-with-joblib
@contextlib.contextmanager
def tqdm_joblib(tqdm_object):
    """Context manager to patch joblib to report into tqdm progress bar given as argument"""
    class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
        def __call__(self, *args, **kwargs):
            tqdm_object.update(n=self.batch_size)
            return super().__call__(*args, **kwargs)

    old_batch_callback = joblib.parallel.BatchCompletionCallBack
    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
    try:
        yield tqdm_object
    finally:
        joblib.parallel.BatchCompletionCallBack = old_batch_callback
        tqdm_object.close()

用法

with tqdm_joblib(tqdm(desc="progress", total=len(zipped_structures))) as progress_bar:
    values = Parallel(n_jobs=-1)(delayed(func)(a,b,c) for a,b,c in zipped_structures)