如何在Python多进程imap_unordered()中为每个Worker添加独立Rich进度条

现有一段通过logging区分4个Worker的多线程代码，希望为每个Worker添加独立的rich进度条，同时理想情况下增加一个总进度条追踪所有任务的完成进度（已知tqdm的实现方式，但更倾向使用自定义性更强的rich）。

原代码 #

import time
import random
import logging
import logging.handlers
from multiprocessing.dummy import Pool

def do_something(number):
    logger.info(number)
    time.sleep(number/100)

logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)-7s] (%(threadName)-10s) %(message)s"))
logger.addHandler(handler)
logger.setLevel(logging.INFO)

numbers = random.sample(range(1, 101), 50)

pool = Pool(4)
pool.imap_unordered(do_something, numbers)
pool.close()
pool.join()

实现方案 #

利用rich.progress模块的线程安全特性，结合线程局部存储为每个Worker绑定独立进度条，同时添加总进度条追踪整体任务完成情况。

完整代码

import time
import random
import logging
import threading
from multiprocessing.dummy import Pool
from rich.progress import Progress
from threading import local

# 线程局部存储：保存每个线程对应的进度条任务ID，避免线程间干扰
thread_local = local()

def do_something(number):
    # 获取当前线程绑定的进度条任务ID
    task_id = thread_local.task_id
    # 初始化当前任务的进度条（单个任务总进度为1）
    progress.update(task_id, total=1, completed=0)
    
    logger.info(number)
    time.sleep(number/100)
    
    # 标记当前任务完成
    progress.update(task_id, completed=1)
    # 推进总进度条
    progress.update(total_task_id, advance=1)

# 配置logging
logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)-7s] (%(threadName)-10s) %(message)s"))
logger.addHandler(handler)
logger.setLevel(logging.INFO)

# 生成50个1-100的随机数作为任务
numbers = random.sample(range(1, 101), 50)

# 初始化Rich进度条上下文
with Progress() as progress:
    # 添加总进度条（跟踪50个任务的完成情况）
    total_task_id = progress.add_task("[green]总进度", total=len(numbers))
    # 为4个Worker预先创建进度条，线程名对应Pool默认的Thread-1~Thread-4
    worker_task_ids = {
        f"Thread-{i+1}": progress.add_task(f"[blue]Worker {i+1}", total=0)
        for i in range(4)
    }

    def init_worker():
        # 每个Worker线程启动时，根据线程名绑定对应的进度条任务ID
        thread_name = threading.current_thread().name
        thread_local.task_id = worker_task_ids[thread_name]

    # 创建线程池，指定初始化函数为每个线程绑定进度条
    pool = Pool(4, initializer=init_worker)
    # 提交任务并异步执行
    pool.imap_unordered(do_something, numbers)
    pool.close()
    pool.join()

关键逻辑说明

• 线程局部存储：用threading.local()存储每个线程对应的进度条任务ID，确保线程间的进度条更新互不干扰。
• 进度条初始化：在创建线程池时，通过initializer回调为每个Worker线程绑定对应的进度条任务ID，匹配Pool默认生成的线程名（Thread-1到Thread-4）。
• 进度更新：每个任务执行时，先重置Worker进度条的总任务数为1，任务完成后标记进度为100%，同时推进总进度条的完成计数。
• 视觉区分：总进度条用绿色标注，Worker进度条用蓝色标注，便于直观区分。

内容的提问来源于stack exchange，提问作者hobbes3