Experimental `joblib.register_parallel_backend` to render progress bar.

logging_tqdm_multiprocessing.py

import logging
import threading
import time

import joblib
import tqdm


class TqdmLoggingHandler(logging.StreamHandler):
    def emit(self, record):
        try:
            msg = self.format(record)
            tqdm.tqdm.write(msg)
            self.flush()
        except (KeyboardInterrupt, SystemExit):
            raise
        except:
            self.handleError(record)


#handler = logging.StreamHandler(tqdm.tqdm)
#handler.terminator = ''
#handler = logging.FileHandler('out.txt')
handler = TqdmLoggingHandler()
logger = logging.getLogger(__name__)
logger.addHandler(handler)
logger.setLevel(logging.INFO)


class Callbacks:
    def __init__(self, *callbacks):
        self.callbacks = [cb for cb in callbacks if cb]

    def __call__(self, out):
        for cb in self.callbacks:
            cb(out)


class MultiprocessingCallbackBackend(
        joblib._parallel_backends.MultiprocessingBackend):

    def _init_progress_bar(self, n):
        self._progress_bar = tqdm.tqdm(range(n))

    def callback(self, results):
        for result in results:
            logging.getLogger(__name__).warning('Bar {}.'.format(result))
            self._progress_bar.update(1)
            self._progress_bar.set_description_str('Foo: {}.'.format(result))

    def apply_async(self, func, callback=None):
        return super().apply_async(func, Callbacks(callback, self.callback))



joblib.register_parallel_backend(
    'multiprocessing_callback', MultiprocessingCallbackBackend)


def func(i):
    time.sleep(i * .1)
    if i == 3:
        logging.getLogger(__name__).warning('I am not rendered correctly.')
    return i


with joblib.parallel_backend('multiprocessing_callback'):
    with joblib.Parallel(n_jobs=2, prefer='processes') as parallel:
        n = 10
        parallel._backend._init_progress_bar(n=n)
        parallel(joblib.delayed(func)(i) for i in range(n))
        parallel._backend._progress_bar.close()