👻
python のマルチプロセスにおける tqdm-friendly な logging
名工大アドベントカレンダー 3 日目です。
モチベーション
- マルチプロセスの進捗状況を tqdm で管理していて、子プロセスでも tqdm と干渉しない log を行いたい
アイデア
マルチプロセスでも QueueHandler で全ログを共有 Queue に集約し、メインプロセスの QueueListener + TqdmLoggingHandler が tqdm.write で一括出力することで、複数プロセスのログが tqdm の表示と競合せずに安全に流れる
ロジック
-
TqdmLoggingHandlerでtqdm.writeを使ってログを出すハンドラを定義し、プログレスバーを壊さずにログを吐き出せる -
worker_initを各子プロセス起動時に呼び、ルートロガーにQueueHandlerをセットして全ログを共有Queueに流す - ワーカ関数は通常どおり
logging.getLogger(__name__)を使ってログを出すだけで、すべてのレコードがキュー経由でメインプロセスへ届く - メインプロセスでは
QueueListenerにTqdmLoggingHandlerを渡して起動し、mp.Pool生成時にinitializerとしてworker_initを指定してキューを渡す。最後にlistener.stop()でクリーンアップする
ソースコード
実装方法
import logging
import logging.handlers
import multiprocessing as mp
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm
import numpy as np
class TqdmLoggingHandler(logging.Handler):
def emit(self, record):
try:
msg = self.format(record)
tqdm.write(msg)
except Exception:
self.handleError(record)
def worker_init(log_queue):
"""
This runs once in each worker process.
It configures the root logger to send records to the shared Queue.
"""
queue_handler = logging.handlers.QueueHandler(log_queue)
root = logging.getLogger()
root.setLevel(logging.INFO)
root.handlers.clear()
root.addHandler(queue_handler)
def my_function(x):
logger = logging.getLogger(__name__)
logger.info(f"[worker] start x={x}")
time.sleep(np.random.uniform(5,10))
logger.info(f"[worker] end x={x}")
return x
def main():
log_queue = mp.Queue()
tqdm_handler = TqdmLoggingHandler()
listener = logging.handlers.QueueListener(
log_queue,
tqdm_handler,
respect_handler_level=True,
)
listener.start()
x_array = range(20)
try:
with ProcessPoolExecutor(
max_workers=4,
initializer=worker_init,
initargs=(log_queue,),
) as executor:
futures = [executor.submit(my_function, x_item) for x_item in x_array]
results = []
for future in tqdm(
as_completed(futures),
total=len(futures),
):
results.append(future.result())
print("Results:", results)
finally:
listener.stop()
if __name__ == "__main__":
main()
モジュール化
import contextlib
import logging
import logging.handlers
import multiprocessing as mp
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from collections.abc import Iterator
from tqdm import tqdm
import numpy as np
class TqdmLoggingHandler(logging.Handler):
def emit(self, record):
try:
msg = self.format(record)
tqdm.write(msg)
except Exception:
self.handleError(record)
def my_function(x):
logger = logging.getLogger(__name__)
logger.info(f"[worker] start x={x}")
time.sleep(np.random.uniform(5, 10))
logger.info(f"[worker] end x={x}")
return x
@contextlib.contextmanager
def queue_logging_context(
level: int = logging.INFO,
) -> Iterator[mp.Queue]:
log_queue: mp.Queue = mp.Queue()
listener_handler = TqdmLoggingHandler()
listener = logging.handlers.QueueListener(
log_queue,
listener_handler,
respect_handler_level=True,
)
root = logging.getLogger()
previous_handlers = root.handlers[:]
previous_level = root.level
queue_handler = logging.handlers.QueueHandler(log_queue)
root.handlers.clear()
root.addHandler(queue_handler)
root.setLevel(level)
listener.start()
try:
yield log_queue
finally:
listener.stop()
root.handlers = previous_handlers
root.setLevel(previous_level)
def worker_init(
log_queue: mp.Queue,
) -> None:
root = logging.getLogger()
queue_handler = logging.handlers.QueueHandler(log_queue)
root.handlers.clear()
root.addHandler(queue_handler)
root.setLevel(logging.INFO)
def main():
x_array = range(20)
with queue_logging_context() as log_queue:
with ProcessPoolExecutor(
max_workers=4,
initializer=worker_init,
initargs=(log_queue,),
) as executor:
futures = [executor.submit(my_function, x_item) for x_item in x_array]
results = []
for future in tqdm(
as_completed(futures),
total=len(futures),
):
results.append(future.result())
print("Results:", results)
if __name__ == "__main__":
main()
おまけ
子プロセスの待機時間が同じ場合
子プロセスの待機時間が同じ場合、たまにプログレスバーが更新されないことがあるのでそれに対処
import logging
import logging.handlers
import multiprocessing as mp
import time
from concurrent.futures import Future, ProcessPoolExecutor, as_completed
from tqdm import tqdm
class TqdmLoggingHandler(logging.Handler):
def emit(self, record):
try:
msg = self.format(record)
tqdm.write(msg)
except Exception:
self.handleError(record)
def worker_init(log_queue):
"""
This runs once in each worker process.
It configures the root logger to send records to the shared Queue.
"""
queue_handler = logging.handlers.QueueHandler(log_queue)
root = logging.getLogger()
root.setLevel(logging.INFO)
root.handlers.clear()
root.addHandler(queue_handler)
def my_function(x):
logger = logging.getLogger(__name__)
logger.info(f"[worker] start x={x}")
time.sleep(5)
logger.info(f"[worker] end x={x}")
return x
def main():
log_queue = mp.Queue()
tqdm_handler = TqdmLoggingHandler()
listener = logging.handlers.QueueListener(
log_queue,
tqdm_handler,
respect_handler_level=True,
)
listener.start()
x_array = range(20)
try:
with ProcessPoolExecutor(
max_workers=4,
initializer=worker_init,
initargs=(log_queue,),
) as executor:
results: list = []
futures: list[Future[int]] = []
progress = tqdm(total=len(x_array), desc="Processing")
try:
for x_item in x_array:
future = executor.submit(my_function, x_item)
future.add_done_callback(lambda _: progress.update(1))
futures.append(future)
for future in as_completed(futures):
results.append(future.result())
finally:
progress.close()
print("Results:", results)
finally:
listener.stop()
if __name__ == "__main__":
main()
Discussion