multiprocessingで一つのtqdmを複数プロセスから更新する

複数プロセス（worker）を並列実行して，一つのプログレスバーを更新したい．

overview

• Pool + imap_unordered()という手もあるが，
  • 引数にリストなどのiterableを渡さなければならない
  • 一つのiterableしかimap_unordered()の引数に渡せない．
  • workerの引数も一つしか渡せない．（単純にやるからだけど）
  • pickle化するのでリストの要素が
• そこでProcess + queueを使ってみる．
  • リストの代わりにqueueに詰める．
  • tqdmをManagerで管理する．
  • workerにqueueとtqdmも渡す．
    • workerの中でtqdmをupdate

motivation

pool.imap_unordered()を使っていたけれども，各workerに渡すのが別々のものと共通のものとが混在していて，iterableを作る効率が悪い．

そこで別々のものはqueueに詰めてしまって，各workerはそこから取り出すことにする．共通のものは単にworkerの引数にわたすだけ．

code

ではまず必要なモジュールのインポート．

モジュールのインポート

from multiprocessing import Process, Manager
from multiprocessing.managers import BaseManager
from tqdm import tqdm
import queue

次にworkerの定義．

• 引数はキューq，ロックオブジェクト，tqdmのプログレスバー，worker番号
  • worker番号iは単なる例．なんでもよい．
• キューから取り出す．キューが空なら即終了．キューが空になるまでループ．
• 取り出したデータを処理
• lockしてtqdmのプログレスバーを更新
  • lockしないと更新表示がずれる

def worker(q, lock, pbar, i):

    while True:
        try:
            item = q.get(timeout=0)
        except queue.Empty:
            return

        # ここでitemを処理

        with lock:
            pbar.update(1)
            pbar.set_postfix_str('{0:3d} {1:3d}'.format(item, i))

tqdmをmanagerで管理するための定義．

class TqdmManager(BaseManager):
    # https://docs.python.org/ja/3/library/multiprocessing.html#customized-managers
    pass


TqdmManager.register('Tqdm', tqdm)

ではメイン．

• キューq，ロック，tqdmプログレスバーpbarのオブジェクトをmanagerで作成．
• qにデータを詰める
• Processで個数num_workerのプロセスを生成．
  • 引数はargsで指定．ここでiterable以外にいろいろ渡せる．
• workerをスタート．
  • joinで終了まで待つ．
  • closeでクリーンアップ．
• 最後にpbarをclose

メイン

if __name__ == '__main__':

    with TqdmManager() as tqdm_manager, Manager() as manager:

        num_workers = 13
        len_data = 1182

        q = manager.Queue()
        lock = manager.Lock()
        pbar = tqdm_manager.Tqdm(total=len_data)

        for item in range(len_data):
            q.put(item)

        p_all = [Process(target=worker,
                         args=(q, lock, pbar, i))
                 for i in range(num_workers)]
        [p.start() for p in p_all]
        [p.join() for p in p_all]
        [p.close() for p in p_all]

        pbar.close()

execute

$ python3 ./queue_test.py
100%|█████████████████████████████████| 1182/1182 [00:01<00:00, 601.64it/s, 1181   1]
$ 

limitation

managerを使うと，そのための別プロセスが起動するのでやや遅くなる．

misc

• tqdmをBaseManagerで派生して管理するという例は見かけない．なんでもmanagerにできるので便利．
• [p.start() for p in p_all]というインライン表記もあまり見かけない．なぜ？

Yet another version

queueにデータを入れるのが時間かかる場合，先にworkerをスタートしてからqueueにデータを投げ入れる．

• queueに入れればすぐにworkerに受け渡される
• workerではq.getのtimeoutを1程度にしておく（データ待ちの時間が発生するため）
• queueに入れるデータが尽きたら，queueに終了目印のNoneを入れておく．
  • workerの方はNoneが来たら即終了

queue_test.py

from multiprocessing import Process, Manager
from tqdm import tqdm
import queue
from multiprocessing.managers import BaseManager


def worker(q, lock, pbar, i):

    while True:
        try:
            item = q.get(timeout=1)
        except queue.Empty:
            return
        if item is None:
            return

        # ここでitemを処理

        with lock:
            pbar.update(1)
            pbar.set_postfix_str('{0:3d} {1:3d}'.format(item, i))


class TqdmManager(BaseManager):
    # https://docs.python.org/ja/3/library/multiprocessing.html#customized-managers
    pass


TqdmManager.register('Tqdm', tqdm)


if __name__ == '__main__':

    with TqdmManager() as tqdm_manager, Manager() as manager:

        num_workers = 13
        len_data = 1182

        q = manager.Queue()
        lock = manager.Lock()
        pbar = tqdm_manager.Tqdm(total=len_data)

        p_all = [Process(target=worker,
                         args=(q, lock, pbar, i))
                 for i in range(num_workers)]
        [p.start() for p in p_all]

        for item in range(len_data):
            q.put(item)
        for _ in range(num_workers):
            q.put(None)

        [p.join() for p in p_all]
        [p.close() for p in p_all]

        pbar.close()

追記：SyncManagerから派生

BaseManagerから派生させるのではなく，SyncManagerから派生させればqueueも使えた．managerが一つになるのでこのほうがよい．

SyncManagerから派生するバージョン

from multiprocessing import Process
from tqdm import tqdm
import queue
from multiprocessing.managers import SyncManager


def worker(q, lock, pbar, worker_id):

    while True:
        try:
            item = q.get(timeout=1)
        except queue.Empty:
            return
        if item is None:
            return

        # ここでitemを処理

        with lock:
            pbar.update(1)
            pbar.set_postfix_str('item {0:3d} id {1:3d}'.format(item, worker_id))


class TqdmManager(SyncManager):
    # https://docs.python.org/ja/3/library/multiprocessing.html#customized-managers
    pass


if __name__ == '__main__':

    TqdmManager.register('Tqdm', tqdm)

    with TqdmManager() as manager:

        num_workers = 13
        len_data = 1182

        q = manager.Queue()
        lock = manager.Lock()
        pbar = manager.Tqdm(total=len_data)

        p_all = [Process(target=worker,
                         args=(q, lock, pbar, worker_id))
                 for worker_id in range(num_workers)]
        [p.start() for p in p_all]

        for item in range(len_data):
            q.put(item)
        for _ in range(num_workers):
            q.put(None)

        [p.join() for p in p_all]
        [p.close() for p in p_all]

        pbar.close()

なおSyncManagerから派生せずにSyncManagerに直接registerもできるが，汚染してしまうのでやめたほうがよい．