mpi4py是Python中用于MPI(Message Passing Interface)编程的模块,主要用于多节点的并行计算。以下是一个使用mpi4py进行跨节点并行计算的简单示例:
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
# 示例任务:每个进程生成一个随机数,然后所有进程打印出所有随机数
if rank == 0:
print("Total processes:", size)
# 每个进程生成一个随机数
local_number = np.random.rand()
# 收集所有进程的随机数
all_numbers = comm.gather(local_number, root=0)
if rank == 0:
print("Rank 0:", local_number)
print("All numbers:", all_numbers)
在这个例子中,我们首先导入了必要的模块,并获取了MPI的通信器、进程总数和当前进程的序号。然后,如果当前进程是根进程(rank为0),则打印出总进程数。接着,每个进程生成一个随机数。最后,通过comm.gather方法收集所有进程的随机数到根进程,根进程再打印出本地生成的随机数和收集到的所有随机数。
注意:在实际使用时,需要在具有多个节点和多个处理器的计算环境中运行此代码。此外,确保MPI环境已正确配置,以便mpi4py能够正确地在多个节点间通信。