mpi4py 的一些使用技巧

本文从本人简书博客同步过来

在上一篇中我们介绍了 mpi4py 中的访问文件数据操作方法，至此 mpi4py 中最主要的内容已经基本介绍完毕，下面我们将介绍 mpi4py 的一些使用技巧。

兼容非 MPI 编程

从前面的介绍可知，使用 mpi4py 进行 Python 环境下的 MPI 编程是比较容易的，在不牺牲 Python 本身的灵活性和易用性的基础上，mpi4py 可以使我们轻松地利用多核甚至多计算节点进行并行甚至分布式的计算任务，以显著地提高计算效率。但是在有些情况下，我们却必须在非 MPI 环境下运行我们的程序，或者做相应的计算。为我们的并行计算程序再准备一个非并行的版本是一种解决方案，但是却要付出额外的劳动，对大型的或特别复杂的应用程序，维护两个版本的程序可能需要很高的成本，而且容易出错。一种更好的方案是让我们的并行计算程序也能兼容非 MPI 编程环境，也就是说，在 MPI 环境下，就利用多个进程以加速程序的计算，但是在非 MPI 环境下，就回归到单进程的串行程序，在可能花费更多时间的情况下完成所需的计算。使用 mpi4py 并行编程时怎么做到这一点呢？使用 mpi4py 做并行计算，一般需要导入 mpi4py 中的 MPI 模块，在非 MPI 环境下这一导入过程会出错（抛出 ImportError 异常），我们可以使用 Python 的异常处理机制捕获这一异常，使程序能够顺利执行。下面以一个简单的例子展示如何做到这一点。这个例子利用级数求和公式来近似得到圆周率 π: π/4 = 1 - 1/3 + 1/5 - 1/7 + 1/9 - 1/11 + …

例程

兼容非 MPI 编程。

# trick1.py

"""
Make the script work both with and without mpi4py.

Run this with 4 processes like:
$ mpiexec -n 4 python trick1.py
or
$ python trick1.py
"""

import warnings


rank = 0
size = 1
comm = None

## try to setup MPI and get the comm, rank and size
## if not they should end up as comm = None, rank=0, size=1
try:
    from mpi4py import MPI

    comm = MPI.COMM_WORLD

    rank = comm.Get_rank()
    size = comm.Get_size()

except ImportError:
    warnings.warn("Warning: mpi4py not installed.")


# compute pi/4 = 1 - 1/3 + 1/5 - 1/7 + 1/9 - 1/11 + ...
N_max = 10000
local_num = N_max / size
local_sum = 0.0
for i in range(rank*local_num, (rank + 1)*local_num):
    local_sum += (-1)**i * 1.0 / (2*i + 1)

# reduce
if comm is None:
    pi = 4.0 * local_sum
else:
    sum_ = comm.reduce(local_sum, root=0, op=MPI.SUM)
    if rank == 0:
        pi = 4.0 * sum_

if rank == 0:
    print 'pi =', pi

使用 mpi4py （4个进程）运行结果如下：

$ mpiexec -n 4 python trick1.py
pi = 3.14149265359

在非 MPI 环境下运行的结果如下：

$ python trick1.py
trick1.py:30: UserWarning: Warning: mpi4py not installed.
warnings.warn("Warning: mpi4py not installed.")
pi = 3.14149265359

可以看出无论有无 MPI 环境，均可得到相同的结果。

一个需要注意的地方是，在非 MPI 环境下运行这个程序和使用 mpi4py 单进程运行该程序还是有一些差别的（注意：在 mpi4py 可用的情况下命令 mpiexec -n 1 python trick1.py 和 python trick1.py 的效果相同），如上例，在使用 mpi4py 单进程运行时，依然会执行规约操作 reduce，而这是不需要的，因此会带来一些性能上的损失。可以对以上执行 reduce 这一部分作一点改进，使其只在进程数多于 1 个时才执行规约操作，如下

...

# reduce
if size