如何从串行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

很遗憾，我觉得这可能不行，因为你必须用mpirun来专门运行MPI代码。

你能做的最好的是写一些通用的代码块，这些代码块可以被MPI进程或普通的Python进程调用。

另外一个解决办法是把你代码中的MPI部分封装成一个外部调用，然后在非MPI代码中用子进程来调用它。不过，这种方法会很依赖你的系统配置，实际上并不是特别灵活。

关于子进程的详细信息可以参考这个帖子 使用Python的子进程Popen，值得一看。这里的复杂之处在于首先要正确地进行调用，也就是

command = "/your/instance/of/mpirun /your/instance/of/python your_script.py -arguments"

然后把结果返回到你的单线程代码中。根据数据的大小，有很多方法可以做到这一点，但如果你需要传回大数组数据，像parallel hdf5这样的工具会是个不错的选择。

抱歉，我无法给你一个简单的解决方案。

Answer 2

这确实是可能的，而且在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)