如何从串行python脚本使用mpi4py应用程序

Question

我试图基于mpi4py创建一个库，但是我想在串行python代码中使用它。

$ python serial_source.py

但是在serial_source.py内部存在一些称为parallel_bar的函数

from foo import parallel_bar
# Can I to make this with mpi4py like a common python source code?
result = parallel_bar(num_proc = 5)

这个问题的动机是找到正确的方法来使用mpi4py来优化python中的程序，这些程序不一定设计为完全并行运行。

Answer 1

这确实是可行的，并且在mpi4py的文档“ 动态过程管理 ”部分中。您需要的是所谓的Spawn功能MSMPI（如果您使用的是Windows ，则不可用），另请参见未在MSMPI中实现的Spawn。

例

第一个文件为您的函数提供了一种包装器，以隐藏所有MPI内容，我想这是您的意图。在内部，它在4个新产生的进程中调用包含您的并行代码的“实际”脚本。

最后，您可以打开python终端并调用：

from my_prog import parallel_fun

parallel_fun()
# Hi from 0/4
# Hi from 3/4
# Hi from 1/4
# Hi from 2/4
# We got the magic number 6

my_prog.py

import sys
import numpy as np
from mpi4py import MPI

    def parallel_fun():
        comm = MPI.COMM_SELF.Spawn(
            sys.executable,
            args = ['child.py'],
            maxprocs=4)

        N = np.array(0, dtype='i')

        comm.Reduce(None, [N, MPI.INT], op=MPI.SUM, root=MPI.ROOT)

        print(f'We got the magic number {N}')

这是带有并行代码的子文件：

child.py

from mpi4py import MPI
import numpy as np


comm = MPI.Comm.Get_parent()

print(f'Hi from {comm.Get_rank()}/{comm.Get_size()}')
N = np.array(comm.Get_rank(), dtype='i')

comm.Reduce([N, MPI.INT], None, op=MPI.SUM, root=0)