如何用tqdm为MPI并行任务生成单个进度条？

问题描述 #

当前代码会遍历n_mod个模型，任务分配给n_cpu个CPU执行，运行mpirun -np 4 python test_mpi.py时会生成4个独立进度条，希望通过tqdm实现单个进度条展示已完成的模型总数。

原代码如下：

from tqdm import tqdm
from mpi4py import MPI
import time

comm = MPI.COMM_WORLD
cpu_ind = comm.Get_rank()
n_cpu = comm.Get_size()

n_mod=100
for i in tqdm(range(n_mod)):
    if (cpu_ind == int(i/int(n_mod/n_cpu))%n_cpu):
        # 这里是依赖i的任务逻辑
        time.sleep(0.02)

解决方案 #

核心思路是让主进程（rank=0）单独负责进度条更新，其他工作进程完成任务后向主进程发送完成信号，主进程累计完成数并更新进度条。

修改后的代码如下：

from tqdm import tqdm
from mpi4py import MPI
import time

comm = MPI.COMM_WORLD
cpu_ind = comm.Get_rank()
n_cpu = comm.Get_size()

n_mod = 100
# 按CPU数量均分任务（确保每个CPU分配到的任务数尽量均衡）
tasks_per_cpu = n_mod // n_cpu
remainder = n_mod % n_cpu

# 给每个CPU分配对应的任务索引范围
if cpu_ind < remainder:
    start = cpu_ind * (tasks_per_cpu + 1)
    end = start + tasks_per_cpu + 1
else:
    start = remainder * (tasks_per_cpu + 1) + (cpu_ind - remainder) * tasks_per_cpu
    end = start + tasks_per_cpu

# 工作进程执行任务并发送完成信号
if cpu_ind != 0:
    for i in range(start, end):
        # 这里是你的任务逻辑
        time.sleep(0.02)
        # 向主进程发送任务完成的信号
        comm.send(1, dest=0)
else:
    # 主进程初始化进度条，总任务数是n_mod
    with tqdm(total=n_mod) as pbar:
        completed = 0
        while completed < n_mod:
            # 接收任意工作进程的完成信号
            comm.recv(source=MPI.ANY_SOURCE)
            completed += 1
            pbar.update(1)

代码说明 #

• 任务分配优化：不再用原代码的条件判断分配任务，而是直接给每个CPU划分连续的任务索引范围，避免不必要的循环判断，提升运行效率。
• 主从进程分工：
  • 非主进程（rank≠0）专注执行分配到的任务，每完成一个任务就向主进程发送完成信号。
  • 主进程（rank=0）不参与任务计算，仅负责监听所有工作进程的完成信号，累计完成数量并更新进度条。
• 进度条更新：主进程通过tqdm初始化总长度为n_mod的进度条，每收到一个完成信号就调用update(1)更新进度，最终实现单个进度条展示全局任务完成情况。

运行mpirun -np 4 python test_mpi.py时，只会在主进程输出一个进度条，清晰展示所有模型的完成进度。

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