用于字典迭代的python MPI

Question

我想使用MPI（mpi4.py）（消息传递接口）拆分字典的迭代

比如说,

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

tmp_list = []
for key, value in some_dict.items():
    tmp_values = some_function(key,value)
    tmp_list.append(tmp_values)

这里有一些简单的代码。 如何为迭代编写MPI代码

Answer 1

您首先需要将dict转换为一个列表，然后将该列表划分为您将要使用的多个进程。这是必要的，以便comm.scatter可以跨所有进程发送部分数据。然后可以使用comm.gather收集最终结果

script.py

#!/usr/bin/python

from mpi4py import MPI

comm = MPI.COMM_WORLD
size = comm.Get_size() # get number of processes
rank = comm.Get_rank() # get the current rank of process

def some_fun(x,y): # some random function
    return x+y

def chunkIt(seq, num): #function to chunk list into {size} parts for scatter to work
    avg = len(seq) / float(num)
    out = []
    last = 0.0

    while last < len(seq):
        out.append(seq[int(last):int(last + avg)])
        last += avg

    return out

some_dict = {i: i**2 for i in range(100)} # some random data to work on

if rank == 0:
    data = chunkIt(list(some_dict.items()), size) # convert dict into list first and then divide the data into {size} parts
    # print(f"rank: {rank} / data: {data}")
else:
    data = None
    
data = comm.scatter(data, root=0) # scatter the data accross given processes
print(f"rank: {rank} / data: {data}\n")

sub_dict = dict(data) # convert the list into dict

tmp_list = [] #local to each process
for key, value in sub_dict.items():
    tmp_values = some_fun(key, value)
    tmp_list.append(tmp_values)
  
tmp_list = comm.gather(tmp_list, root=0) # gathering data from all procs to root proc
if rank == 0:
    print(f"length of tmp_list on rank: {rank} is: {len(tmp_list)}")
    print(f"tmp_list: {tmp_list}") #tmp_list is list ot lists. make sure to convert it into required ds 
else:
    assert tmp_list is None

使用chmod使其可执行

chmod +x script.py

然后跑

mpiexec -n 4 script.py

-n是要运行的进程数

注意：我使用的是ubuntu 16.04、python 3.7.10和mpi4py==3.0.3