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 == 1: # changes here
    pi = 4.0 * local_sum
else:
    sum_ = comm.reduce(local_sum, root=0, op=MPI.SUM)
    if rank == 0:
        pi = 4.0 * sum_

...

独立模块

每次写程序时都采取如上操作未免有些麻烦，我们完全可以将其独立出来作为一个模块放到一个单独的文件中，命名为 mpiutil.py，（为了区分下面的例程中会加一个序号），然后在程序中导入该模块，如下：

# mpiutil1.py

"""Utilities for making MPI usage transparent."""

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.")

# trick2.py

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

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

import mpiutil1

rank = mpiutil1.rank
size = mpiutil1.size

# 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 size == 1:
    pi = 4.0 * local_sum
else:
    sum_ = mpiutil1.comm.reduce(local_sum, root=0)
    if rank == 0:
        pi = 4.0 * sum_

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

我们甚至可以将 reduce 方法作一下包装放到 mpiutil.py 中，以供其它程序调用，如下：

# mpiutil2.py

"""Utilities for making MPI usage transparent."""

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.")


def reduce(sendobj, root=0, op=None, comm=comm):
    if comm is not None and comm.size > 1:
        return comm.reduce(sendobj, root=root, op=(op or MPI.SUM))
    else:
        return sendobj

# trick3.py

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

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

import mpiutil2

rank = mpiutil2.rank
size = mpiutil2.size

# 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
sum_ = mpiutil2.reduce(local_sum, root=0)

if rank == 0:
    pi = 4.0 * sum_
    print 'pi =', pi

使用其它通信子

在以上的例程 mpiutil2.py 中，我们包装的函数 reduce 有一个额外的可选参数 comm=comm，这个参数看似是并不需要的（你可以试试去掉这个可选参数，trick3.py 仍然可以正常执行，因为会使用默认的通信子 MPI.COMM_WORLD)，但实际上这个参数还是另有用途的，它允许我们传递一个另外的通信子对象，并在这个另外的通信子对象上执行规约操作，如下所示：

# trick4.py

"""
Use a new communicator.

Run this with 5 processes like:
$ mpiexec -n 5 python trick4.py
"""

import mpiutil2

rank = mpiutil2.rank
size = mpiutil2.size

# 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)

if size >= 5:
    comm = mpiutil2.comm
    # create a new communicator by including rank 0, 1, 2, 3 only
    new_comm = comm.Create(comm.Get_group().Incl([0, 1, 2, 3]))
    if rank <= 3:
        print 'use new comm with size: %d' % new_comm.size

        # reduce
        sum_ = mpiutil2.reduce(local_sum, root=0, comm=new_comm)

        if rank == 0:
            pi = 4.0 * sum_
            print 'pi =', pi

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

$ mpiexec -n 5 python trick4.py
use new comm with size: 4
use new comm with size: 4
use new comm with size: 4
use new comm with size: 4
pi = 3.14146765359

自包装

在以上的例程中，通过将兼容非 MPI 的代码独立出来放到 mpiutil.py 中，同时将我们需要用到的 reduce 方法包装后放进 mpiutil.py，我们只需要在自己的程序中导入 mpiutil 模块并使用其中定义的属性和方法，就可以让我们的程序同时支持 MPI 和非 MPI 环境，无需我们在自己的程序中添加一些判断语句以针对这两种情况作特殊处理。

但是上面的解决方案是不完善的，比如在上面的例子 trick2.py，trick3.py 和 trick4.py 中，我们没有设置规约算符的值（会使用默认的 MPI.SUM），如果显式地设置 op = MPI.SUM， 或者使用其它预定义的规约算符，我们就不得不再显式地导入 mpi4py 中的 MPI 模块，从而破坏已有的兼容性。这可以通过另外一个技巧来解决，构建一个 mpiutil 模块的自包装对象 SelfWrapper，然后在程序中传递 op = mpiutil.SUM，在 mpi4py 可用的情况下，SelfWrapper 会将 mpiutil.SUM 换成 MPI.SUM，但是在 mpi4py 不可用的时则会替换成 None。具体如下：

# mpiutil3.py

"""Utilities for making MPI usage transparent."""

import sys
import warnings
from types import ModuleType


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.")


def reduce(sendobj, root=0, op=None, comm=comm):
    if comm is not None and comm.size > 1:
        return comm.reduce(sendobj, root=root, op=(op or MPI.SUM))
    else:
        return sendobj


# this is a thin wrapper around THIS module (we patch sys.modules[__name__])
class SelfWrapper(ModuleType):
    def __init__(self, self_module, baked_args={}):
        for attr in ["__file__", "__hash__", "__buildins__", "__doc__", "__name__", "__package__"]:
            setattr(self, attr, getattr(self_module, attr, None))

        self.self_module = self_module

    def __getattr__(self, name):
        if name in globals():
            return globals()[name]
        elif comm is not None and name in MPI.__dict__:
            return MPI.__dict__[name]

    def __call__(self, **kwargs):
        # print 'here'
        return SelfWrapper(self.self_module, kwargs)


self = sys.modules[__name__]
sys.modules[__name__] = SelfWrapper(self)

# trick5.py

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

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

import mpiutil3

rank = mpiutil3.rank
size = mpiutil3.size

# 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
sum_ = mpiutil3.reduce(local_sum, root=0, op=mpiutil3.SUM)

if rank == 0:
    pi = 4.0 * sum_
    print 'pi =', pi

安装和使用 caput

以上所介绍的若干技巧都已经放在了 caput 软件包的 mpiutil.py 中，其中还包括很多像上面介绍的 reduce 这样的方便和常用的函数（在下一篇中将会介绍），caput 软件包还包括其它一些构建在 mpiutil.py 之上的功能强大的并行计算方法，这会在后面作相应的介绍，如果读者感兴趣，可以前往 https://github.com/zuoshifan/caput/tree/zuo/develop 下载并安装（注意：最好下载 zuo/develop 分支而非 master 分支，因为 zuo/develop 分支中包含一些尚未被 master 分支采纳的函数和方法）。下载并解压后的安装步骤为在软件包的顶层目录下执行：

$ python setup.py install

或者

$ python setup.py install --user

使用 caput 中的 mpiutil 模块的例程如下：

# trick6.py

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

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

from caput import mpiutil

rank = mpiutil.rank
size = mpiutil.size

# 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
sum_ = mpiutil.reduce(local_sum, root=0, op=mpiutil.SUM)

if rank == 0:
    pi = 4.0 * sum_
    print 'pi =', pi

以上我们介绍了 mpi4py 的若干使用技巧，并且简要介绍了 caput 及其 mpiutil 模块，在下一篇中我们将介绍 mpiutil 中提供的若干方便和易用的函数，这些函数可以使我们更加方便地进行 Python 并行编程，并且使我们的程序很容易地做到兼容非 MPI 编程环境。