python - HPC上的mpi4py:comm.gather

我的目标是遍历大型2D数组（data）的每个元素，并对每个元素进行一些繁重的处理。因此，我想使用多个MPI来分别占用阵列的一部分进行处理。我遇到的问题是我不知道如何准确地编写代码以最后收集所有数据。这是一些示例代码：

import numpy as np
import math
from mpi4py import MPI

M = 400
N = 300
data = np.random.rand(M,N)
result_a = np.zeros((M,N))
result_b = np.zeros((M,N))

def process_function(data):
    a = data**2
    b = data**0.5
    return a,b

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
minimum = 0
maximum = int(M*N)
perrank = maximum//size

for index in range(minimum + rank*perrank, minimum + (rank+1)*perrank):
    i = int(math.floor(index/N))
    j = int(index % N)

    a,b = process_function(data[i,j])
    result_a[i,j] = a
    result_b[i,j] = b

a_gath = comm.gather(result_a, root=0)
b_gath = comm.gather(result_b, root=0)

print(np.shape(a_gath))
print('---')
print(np.shape(b_gath))

不幸的是，对于我真正的问题，当我将a_gath和b_gath都保存到磁盘（作为泡菜）时，当我重新加载它们时，它们仅包含一次出现的()（即键入None）。 comm.gather之前/之后，我还有其他事情要做吗？

这是我的提交脚本：

#!/bin/bash -l

#$ -S /bin/bash
#$ -l h_rt=00:05:00
#$ -l mem=2G
#$ -l tmpfs=10G
#$ -pe mpi 5
#$ -N stack_test
#$ -notify
#$ -wd /home/user/Scratch/

module load gcc-libs
module load python3/recommended
module unload compilers mpi
module load compilers/gnu/4.9.2
module load mpi/openmpi/3.1.1/gnu-4.9.2
module load mpi4py
module list

python_infile=test.py

echo ""
echo "Running python < $python_infile ..."
echo ""
gerun python $python_infile

我使用qsub js_test.sh提交此脚本

对于这个伪造的示例，返回的.o文件显示，在这种情况下，4/5 mpis包含类型None信息：在这种情况下，如果我随后将a_gath和b_gath保存到磁盘，它将保存最后的MPI？哪个是None类型？我希望在使用comm.gather之后，对于变量MxN和a_gath我将有一个大小为b_gath的单个数组

()
---
()
()
---
()
()
---
()
(5, 400, 300)
---
(5, 400, 300)
()
---
()

非常感谢。

最佳答案

为了澄清您在注释中为自己找到的答案：MPI_Gather是一个根操作：其结果在所有等级上都不相同，特别是在root参数中提供的等级上有所不同。

在Gather的情况下，您发现以数据结尾的等级为0，这与您使用root=0的方式完全正确。

虽然原则上MPI支持多个程序，多个数据执行，其中不同的行级运行不同的代码，但实际上大多数MPI代码是用单个程序编写的，具有多种数据样式，就像您编写的一样。因为所有等级都运行相同的代码体，所以由您决定从诸如MPI_Gather之类的根操作返回后，要运行的等级是否是根，并相应地执行不同的代码路径。如果不这样做，那么每个等级都将执行以下行：

print(np.shape(a_gath))
print('---')
print(np.shape(b_gath))

如前所述，除了等级0以外，它不会打印出您期望的a_gath和b_gath结果。

请尝试以下操作：

a_gath = comm.gather(result_a, root=0)
b_gath = comm.gather(result_b, root=0)

if rank == 0:
    print(np.shape(a_gath))
    print('---')
    print(np.shape(b_gath))

关于python - HPC上的mpi4py:comm.gather，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/57359727/