python - 跟踪joblib.Parallel执行的进度

Question

有没有一种简单的方法来跟踪 joblib.Parallel 的整体进度？执行？

我有一个由数千个作业组成的长时间运行的执行，我想在数据库中对其进行跟踪和记录。但是，要做到这一点，每当 Parallel 完成任务时，我需要它执行回调，报告剩余的作业数。

我之前使用 Python 的 stdlib multiprocessing.Pool 完成了类似的任务，方法是启动一个线程来记录 Pool 的作业列表中待处理作业的数量。

看代码，Parallel继承了Pool，所以我想我可以实现同样的技巧，但它似乎没有使用这些列表，我一直无法弄清楚如何“阅读” "任何其他方式都是内部状态。

Answer 1

比 dano 和 Connor 的答案更进一步的是将整个事情包装为上下文管理器:

import contextlib
import joblib
from tqdm import tqdm

@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()

然后你可以像这样使用它，完成后不要留下猴子修补的代码:

from math import sqrt
from joblib import Parallel, delayed

with tqdm_joblib(tqdm(desc="My calculation", total=10)) as progress_bar:
    Parallel(n_jobs=16)(delayed(sqrt)(i**2) for i in range(10))

我认为这太棒了，它看起来类似于 tqdm pandas 集成。