资讯动态

python中mpi4py的所有基础使用案例详解

发布时间:2026/10/15 6:58:44 来源:尧图企业网站定制
前言mpi4py是 MPIMessage Passing Interface消息传递接口的 Python 绑定。MPI 是一套面向并行计算的消息传递标准它的典型使用场景是集群几十上百台机器每台多个核心用一个mpiexec命令把同一份程序启动成很多份各自算一块再通过显式的发送/接收把结果拼起来。这里要先纠正一个常见预期。很多人把mpi4py当成「multiprocessing的另一种写法」指望它能在单机上更方便地加速。两者的定位不一样multiprocessing面向单机多核进程靠fork或spawn在同一台机器上起用共享内存或管道通信mpi4py面向集群/多机进程之间没有共享内存一切数据都要靠显式的消息传递通信可能跨网络。单机上跑mpi4py当然也可以mpiexec -n 4就在本机起 4 个进程但如果你的目标只是「用满本机 CPU 核」multiprocessing通常更直接。还有一层安装上的现实mpi4py是第三方库而且必须先有一个 MPI 实现Linux 上常见 OpenMPI / MPICHWindows 上常见 MS-MPI再安装mpi4py去调用它。没有底层 MPI只pip install mpi4py编译不过去。本文按mpi4py4.x 的接口书写具体签名以该库官方文档为准示例需要 Python 3 环境Python 2.7 已于 2020-01-01 停止维护本文不适用。一、安装与运行方式# 1) 先装一个 MPI 实现示例为类 Unix 环境# Linux: apt install libopenmpi-dev openmpi-bin# macOS: brew install open-mpi# Windows: 安装 MS-MPI并确保 mpiexec 在 PATH 里# 2) 再安装 mpi4pypip install mpi4py# 3) 用 mpiexec 启动-n 指定进程数mpiexec -n 4 python hello_mpi.py关键在第三步必须用mpiexec或mpirun启动。直接python hello_mpi.py只会跑一个进程comm.size恒为 1你会以为程序「没报错就是对的」实际上完全没有并行起来。二、rank 与 size每个进程的身份# 适用于 Python 3.8需先安装 MPI 实现与 mpi4py用 mpiexec 启动from mpi4py import MPIcomm MPI.COMM_WORLD # 预定义的通信子包含本次启动的所有进程rank comm.Get_rank() # 本进程的编号从 0 开始size comm.Get_size() # 总进程数print(f我是第 {rank} 号进程一共 {size} 个)# 也可以直接用属性写法效果等价print(f属性写法rank{comm.rank}, size{comm.size})Get_rank()返回本进程的编号0 到size - 1Get_size()返回总进程数。comm.rank与comm.size是这两个方法的属性简写官方文档里分别写作「Rank of this process」和「Number of processes」。MPI.COMM_WORLD是预定义的两个通信子之一另一个是MPI.COMM_SELF只含自己。同一个程序文件被mpiexec -n 4启动成 4 份每份都会从头执行一遍rank是它区分「我是谁」的唯一依据。这就是 MPI 的SPMDsingle program, multiple data模式——一份程序、多份数据。概念含义常见写法MPI.COMM_WORLD全体进程组成的通信子所有进程都拿它通信Get_rank()/rank本进程编号0 起用if rank 0:区分主进程Get_size()/size进程总数决定数据怎么分片root某个集合操作的「主角」进程通常是 0 号tag点对点消息的标签用来区分不同消息三、点对点通信send / recv点对点point-to-point是最基础的通信一个进程发、另一个进程收。send与recv是阻塞版本——send要等到消息发出、recv要等到消息到达才返回。# 适用于 Python 3.8用 mpiexec -n 2 python p2p.py 启动from mpi4py import MPIcomm MPI.COMM_WORLDrank comm.Get_rank()if rank 0:data {a: 7, b: 3.14} # 任意可 pickle 的 Python 对象comm.send(data, dest1, tag11) # 发往 1 号进程elif rank 1:data comm.recv(source0, tag11) # 只收 0 号进程的消息print(1 号收到, data)签名上send(obj, dest[, tag])里obj是要发送的对象dest是目标编号recv([buf, source, tag, status])里source与tag可选给了就只接收匹配的消息。默认参数下recv可以接收任意来源、任意标签的消息。要记住的两点一是send/recv传 Python 对象时底层用的是 pickle因此对象必须可序列化二是recv的source若写成不存在的编号进程会一直阻塞而不是报错——这是 MPI 调试时最常见的「程序卡住」原因。非阻塞版本isend/irecv会立刻返回一个请求对象之后再等它完成# 适用于 Python 3.8用 mpiexec -n 2 python p2p_nb.py 启动from mpi4py import MPIcomm MPI.COMM_WORLDrank comm.Get_rank()if rank 0:req comm.isend({x: 1}, dest1, tag11)req.wait() # 等到发送真正完成elif rank 1:req comm.irecv(source0, tag11)data req.wait() # 等到数据到达返回收到的对象print(收到, data)非阻塞的价值在于「发起之后可以先干别的再等」适合让通信和计算重叠起来而不是干等。四、集合通信bcast / scatter / gather / reduce集合通信collective communication是通信子里的所有进程一起参与的操作。它们最大的特点是每个进程都必须调用同一个集合操作少一个就会挂起。广播bcast把 root 的数据发给所有人。# 适用于 Python 3.8mpiexec -n 4 python bcast_demo.pyfrom mpi4py import MPIcomm MPI.COMM_WORLDrank comm.Get_rank()if rank 0:config {epochs: 10, lr: 0.01} # 只有 root 准备数据else:config None # 其他进程随便给一个占位值config comm.bcast(config, root0) # 所有进程都调用返回值才是最终的print(rank, config)非 root 进程传进去的那个None没有意义它的值会被 root 的数据覆盖。所以正确写法是「root 准备数据、其他进程给占位、所有人都接收返回值」。分发scatter把 root 的一个序列拆开一人一份。# 适用于 Python 3.8mpiexec -n 4 python scatter_demo.pyfrom mpi4py import MPIcomm MPI.COMM_WORLDrank comm.ranksize comm.sizeif rank 0:# 序列长度必须等于进程数root 负责准备全部data [(i 1) ** 2 for i in range(size)]else:data Nonedata comm.scatter(data, root0) # 每个进程拿到自己那一份print(f进程 {rank} 分到 {data})汇集gather把每个人的结果收回到 root组成一个序列。# 适用于 Python 3.8mpiexec -n 4 python gather_demo.pyfrom mpi4py import MPIcomm MPI.COMM_WORLDrank comm.ranksize comm.sizelocal (rank 1) ** 2 # 每个进程各自算一份collected comm.gather(local, root0)if rank 0:assert len(collected) size # 只有 root 拿到列表print(汇总结果, collected)else:assert collected is None # 非 root 拿到的是 None这里有个必须记住的规则gather只在 root 上返回列表其他进程拿到的是None。所以非 root 去索引collected[0]会直接TypeError。归约reduce把所有人的值按某个运算合成一个。# 适用于 Python 3.8mpiexec -n 4 python reduce_demo.pyfrom mpi4py import MPIcomm MPI.COMM_WORLDrank comm.ranklocal_sum sum(range(rank 1)) # 每个进程算自己那部分total comm.reduce(local_sum, opMPI.SUM, root0)if rank 0:print(总和, total) # 只有 root 拿到结果op是归约运算MPI.SUM是求和此外还有MPI.MAX、MPI.MIN、MPI.PROD等用哪个以官方文档的列表为准。reduce与gather一样非 root 得到None。想让每个进程都拿到归约结果用allreduce(sendobj[, op])想让每个进程都拿到完整的收集结果用allgather。集合操作数据流向root 拿到什么其他进程拿到什么bcast(obj, root0)root → 全体原对象同一个对象scatter(data, root0)root → 每人一份自己那一份自己那一份gather(local, root0)每人 → root收齐的列表Nonereduce(local, op, root0)每人 → root归约结果Noneallgather(local)每人 → 全体收齐的列表收齐的列表allreduce(local, op)每人 → 全体归约结果归约结果五、一个把它们串起来的小例子# 适用于 Python 3.8用 mpiexec -n 4 python pipeline.py 启动from mpi4py import MPIcomm MPI.COMM_WORLDrank comm.ranksize comm.size# 1) root 准备全部数据广播给所有人一份配置config {scale: 10} if rank 0 else Noneconfig comm.bcast(config, root0)# 2) root 把数据切成 size 份scatter 分发chunk comm.scatter([list(range(i * 5, (i 1) * 5)) for i in range(size)] if rank 0 else None,root0,)# 3) 每个进程各自计算local_total sum(x * config[scale] for x in chunk)# 4) 把各自的结果 gather 回 rootparts comm.gather(local_total, root0)# 5) root 把结果归约成一个总数并打印if rank 0:total comm.reduce(local_total, opMPI.SUM, root0)print(分片, parts)print(总计, total)else:comm.reduce(local_total, opMPI.SUM, root0)注意最后那段if/elsereduce是集合操作所有进程都必须调用所以不能只写在if rank 0:里面否则非 root 进程不参与root 会永远挂起等它——这就是「集合操作漏调用」的典型后果。上面的写法让所有人都调用reduce只是只有 root 使用返回值。同理第 4 步的gather与第 2 步的scatter也都在所有进程上执行。常见坑点坑点 1没装 MPI 实现就pip install mpi4py。❌ 空机器上直接pip install mpi4py—— 找不到 MPI 头文件或库编译阶段就失败。 ✅ 先装一个 MPI 实现OpenMPI / MPICH / MS-MPI确认mpiexec可用再装mpi4py。坑点 2用python script.py直接跑。❌ 忘记mpiexec -n 4—— 只起一个进程size是 1看着「没报错」其实没并行。 ✅ 一律用mpiexec -n N python script.py启动用comm.size打印出来确认。坑点 3只在 root 分支里调用集合操作。❌if rank 0: comm.bcast(data, root0)—— 其他进程没参加root 永远等不到他们。 ✅ 集合操作写在公共路径上所有进程都调用只有返回值的使用按 rank 区分。坑点 4对非 root 拿到None的返回值做索引。❌gather/reduce之后所有进程都去访问collected[0]—— 非 root 拿到的是None直接TypeError。 ✅ 先if rank 0:再使用收集/归约结果想让所有人都拿到改用allgather/allreduce。坑点 5recv的source写错导致卡死。❌ 发送方dest2、接收方source3—— 谁也匹配不上双方静默阻塞。 ✅ 发送端与接收端的编号、tag要对齐调试时先用小规模-n 2核对。坑点 6把 MPI 当成单机多核的加速开关。❌ 在本机为了「用满 CPU」而上mpi4py—— 它面向集群/多机单机场景multiprocessing更轻、更简单。 ✅ 先判断规模单机多核用multiprocessing跨机器、跨节点才上 MPI。坑点 7以为进程间有共享内存。❌ 在 rank 0 里改一个全局变量指望 rank 1 能看到。 ✅ MPI 进程之间没有共享内存数据必须显式send/bcast/scatter过去。坑点 8所有进程都往同一个文件写、都打印同一段。❌ 不加 rank 判断地print或写文件 —— 输出重复、互相覆盖。 ✅ 输出集中在 root或按 rank 分文件/加前缀需要互斥时用集合操作把数据先收回 root。总结你想做的事用哪个关键点区分进程comm.rank/comm.size0 起编号SPMD 同程序多数据一对一传对象send/recv阻塞底层 picklesource写错会卡死一对一不阻塞isend/irecvwait()先发起后收结果把配置发给所有人bcast(obj, root0)非 root 的入参无意义把数据切给每人scatter(data, root0)root 准备长度为size的序列把结果收回 rootgather/reduce非 root 拿到None让所有人都拿到allgather/allreduce每个进程都有完整结果用好mpi4py的关键是两条纪律每个进程用 rank 认清自己该算哪一块以及集合操作所有进程都要调用返回值只在 root 用。前者决定「算得对不对」后者决定「会不会卡死」。把bcast → scatter → 本地计算 → reduce这条主干记住再按需替换成点对点的send/recvMPI 的基础用法就掌握了大半。至于它和multiprocessing的分工——单机多核选后者跨节点并行选前者别把两者混为一谈。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑