Tracking progress of joblib.Parallel execution

Question

Is there a simple way to track the overall progress of a joblib.Parallel execution?

I have a long-running execution composed of thousands of jobs, which I want to track and record in a database. However, to do that, whenever Parallel finishes a task, I need it to execute a callback, reporting how many remaining jobs are left.

I've accomplished a similar task before with Python's stdlib multiprocessing.Pool, by launching a thread that records the number of pending jobs in Pool's job list.

Looking at the code, Parallel inherits Pool, so I thought I could pull off the same trick, but it doesn't seem to use these that list, and I haven't been able to figure out how else to "read" it's internal status any other way.

Solution 1

Why can't you simply use tqdm? The following worked for me

from joblib import Parallel, delayed
from datetime import datetime
from tqdm import tqdm

def myfun(x):
    return x**2

results = Parallel(n_jobs=8)(delayed(myfun)(i) for i in tqdm(range(1000))
100%|██████████| 1000/1000 [00:00<00:00, 10563.37it/s]

Solution 2

Yet another step ahead from dano's and Connor's answers is to wrap the whole thing as a context manager:

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

Then you can use it like this and don't leave monkey patched code once you're done:

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

which is awesome I think and it looks similar to tqdm pandas integration.

Solution 3

The documentation you linked to states that Parallel has an optional progress meter. It's implemented by using the callback keyword argument provided by multiprocessing.Pool.apply_async:

# This is inside a dispatch function
self._lock.acquire()
job = self._pool.apply_async(SafeFunction(func), args,
            kwargs, callback=CallBack(self.n_dispatched, self))
self._jobs.append(job)
self.n_dispatched += 1

...

class CallBack(object):
    """ Callback used by parallel: it is used for progress reporting, and
        to add data to be processed
    """
    def __init__(self, index, parallel):
        self.parallel = parallel
        self.index = index

    def __call__(self, out):
        self.parallel.print_progress(self.index)
        if self.parallel._original_iterable:
            self.parallel.dispatch_next()

And here's print_progress:

def print_progress(self, index):
    elapsed_time = time.time() - self._start_time

    # This is heuristic code to print only 'verbose' times a messages
    # The challenge is that we may not know the queue length
    if self._original_iterable:
        if _verbosity_filter(index, self.verbose):
            return
        self._print('Done %3i jobs       | elapsed: %s',
                    (index + 1,
                     short_format_time(elapsed_time),
                    ))
    else:
        # We are finished dispatching
        queue_length = self.n_dispatched
        # We always display the first loop
        if not index == 0:
            # Display depending on the number of remaining items
            # A message as soon as we finish dispatching, cursor is 0
            cursor = (queue_length - index + 1
                      - self._pre_dispatch_amount)
            frequency = (queue_length // self.verbose) + 1
            is_last_item = (index + 1 == queue_length)
            if (is_last_item or cursor % frequency):
                return
        remaining_time = (elapsed_time / (index + 1) *
                    (self.n_dispatched - index - 1.))
        self._print('Done %3i out of %3i | elapsed: %s remaining: %s',
                    (index + 1,
                     queue_length,
                     short_format_time(elapsed_time),
                     short_format_time(remaining_time),
                    ))

The way they implement this is kind of weird, to be honest - it seems to assume that tasks will always be completed in the order that they're started. The index variable that goes to print_progress is just the self.n_dispatched variable at the time the job was actually started. So the first job launched will always finish with an index of 0, even if say, the third job finished first. It also means they don't actually keep track of the number of completed jobs. So there's no instance variable for you to monitor.

I think your best best is to make your own CallBack class, and monkey patch Parallel:

from math import sqrt
from collections import defaultdict
from joblib import Parallel, delayed

class CallBack(object):
    completed = defaultdict(int)

    def __init__(self, index, parallel):
        self.index = index
        self.parallel = parallel

    def __call__(self, index):
        CallBack.completed[self.parallel] += 1
        print("done with {}".format(CallBack.completed[self.parallel]))
        if self.parallel._original_iterable:
            self.parallel.dispatch_next()

import joblib.parallel
joblib.parallel.CallBack = CallBack

if __name__ == "__main__":
    print(Parallel(n_jobs=2)(delayed(sqrt)(i**2) for i in range(10)))

Output:

done with 1
done with 2
done with 3
done with 4
done with 5
done with 6
done with 7
done with 8
done with 9
done with 10
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]

That way, your callback gets called whenever a job completes, rather than the default one.

Solution 4

Expanding on dano's answer for the newest version of the joblib library. There were a couple of changes to the internal implementation.

from joblib import Parallel, delayed
from collections import defaultdict

# patch joblib progress callback
class BatchCompletionCallBack(object):
  completed = defaultdict(int)

  def __init__(self, time, index, parallel):
    self.index = index
    self.parallel = parallel

  def __call__(self, index):
    BatchCompletionCallBack.completed[self.parallel] += 1
    print("done with {}".format(BatchCompletionCallBack.completed[self.parallel]))
    if self.parallel._original_iterator is not None:
      self.parallel.dispatch_next()

import joblib.parallel
joblib.parallel.BatchCompletionCallBack = BatchCompletionCallBack

Solution 5

TLDR solution:

Works with joblib 0.14.0 and tqdm 4.46.0 using python 3.5. Credits to frenzykryger for contextlib suggestions, dano and Connor for monkey patching idea.

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

@contextlib.contextmanager
def tqdm_joblib(tqdm_object):
    """Context manager to patch joblib to report into tqdm progress bar given as argument"""

    def tqdm_print_progress(self):
        if self.n_completed_tasks > tqdm_object.n:
            n_completed = self.n_completed_tasks - tqdm_object.n
            tqdm_object.update(n=n_completed)

    original_print_progress = joblib.parallel.Parallel.print_progress
    joblib.parallel.Parallel.print_progress = tqdm_print_progress

    try:
        yield tqdm_object
    finally:
        joblib.parallel.Parallel.print_progress = original_print_progress
        tqdm_object.close()

You can use this the same way as described by frenzykryger

import time
def some_method(wait_time):
    time.sleep(wait_time)

with tqdm_joblib(tqdm(desc="My method", total=10)) as progress_bar:
    Parallel(n_jobs=2)(delayed(some_method)(0.2) for i in range(10))

Longer explanation:

The solution by Jon is simple to implement, but it only measures the dispatched task. If the task takes a long time, the bar will be stuck at 100% while waiting for the last dispatched task to finish execution.

The context manager approach by frenzykryger, improved from dano and Connor, is better, but the BatchCompletionCallBack can also be called with ImmediateResult before the task completes (See Intermediate results from joblib). This is going to get us a count that is over 100%.

Instead of monkey patching the BatchCompletionCallBack, we can just patch the print_progress function in Parallel. The BatchCompletionCallBack already calls this print_progress anyway. If the verbose is set (i.e. Parallel(n_jobs=2, verbose=100)), the print_progress will be printing out completed tasks, though not as nice as tqdm. Looking at the code, the print_progress is a class method, so it already has self.n_completed_tasks that logs the number we want. All we have to do is just to compare this with the current state of joblib's progress and update only if there is a difference.

This was tested in joblib 0.14.0 and tqdm 4.46.0 using python 3.5.

View more solutions

Comments

[Cerin]

Great research, thanks. I didn't notice the callback attribute.

[Ziyuan]

I found the documentation of joblib is very limited. I have to dig into the source code for this CallBack class. My question: can I customize the arguments when __call__ is called? (sub-classing the whole Parallel class may be one way but it's heavy for me).

[Cerin]

Very neat. Thanks.

[Noah]

I don't think this is actually monitoring the completion of running jobs, just the queuing of jobs. If you were to insert a time.sleep(1) at the start of myfun you would find the tqdm progress finishes almost instantly, but results takes a few more seconds to populate.

[Jon]

Yes, that’s partly correct. It is tracking the job starts vs the completions, but the other issue is that there is also a delay caused by overhead after all jobs are completed. Once all tasks are completed results need to be collected and this can take quite a while.

[devforfu]

I believe this answer doesn't really answer the question. As it was mentioned, one will track queuing and not the execution itself with this approach. The approach with callback shown below seems to be more precise in relation to the question.

[Jon]

@devforfu yes, that was addressed in the former comment.

[lawrencegripper]

Works great. If you want to print a time estimate too you can adapt __call__ with the following: ``` time_remaining = (this_batch_duration / self.batch_size) * (total_n_jobs - self.parallel.n_completed_tasks) print( "\rProgress: [{0:50s}] {1:.1f}% est {2:1f}mins left".format('#' * int(progress * 50), progress*100, time_remaining/60) , end="", flush=True) ```

[TheDimLebowski]

This should be the best answer! Thanks

[dennisobrien]

Excellent solution. Tested with joblib 0.14.1 and tqdm 4.41.0 -- works great. This would be a great addition to tqdm!

[Andrew]

I can't edit it, but minor typo in solution where joblib.parallel.BatchCompletionCallback is actually BatchCompletionCallBack (note the camelcase on CallBack)

[Henry Henrinson]

This answer is incorrect, as it does not answer the question. This answer should be unaccepted.

[Christian Steinmeyer]

The provided answer by frenzykryger below contains a great solution to the problem of this answer.

[anilbey]

It's wrong. It only counts the job start times which happens immediately.

[anilbey]

Wrong, this only counts the job start times which will be immediate no matter what function you are wrapping.

[Дмитро Олександрович]

How can it be wrong if it's from the official documentation? joblib.readthedocs.io/en/latest Ctrl+F for "Parallel(n_jobs=1)" And my answer was about running tqdm in Jupyter notebook. It is almost the same as the accepted one. The only difference is that it is intended for use in Jupyter notebook.

[Дмитро Олександрович]

I think I got it. Looks like you're right.

[Дмитро Олександрович]

However, it is not instant in Jupyter notebook. For example, 14% 14/100 [00:05<00:31, 2.77it/s] It takes time to complete with random time sleeps.

[Marc Maxmeister]

This worked for me with a reasonably complex logistic regression function called on thousands of probes in parallel: stats = parallel(func(data, phenotype) for data in tqdm(meth_data, total=len(all_probes), desc='Probes') ) [meth_data is a dataframe and I'm passing each column through the function]

[Labo]

I just posted this code to PyPI: github.com/louisabraham/tqdm_joblib Now you can just pip install tqdm_joblib and from tqdm_joblib import tqdm_joblib