在 IPython 中使用 mpi4py

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

在上一篇中我们介绍了 mpi4py 与 OpenMP 混合编程，下面我们将介绍在 IPython 中使用 mpi4py。

IPython 简介

Python 的一个重要特性是其提供的交互式解释器，在其中可以快速地做一些计算和检验自己的想法而不需要像其它很多编程语言一样必须首先创建一个完整的程序文件。但是 Python 自带的解释器的功能是非常有限的，在很多情况下使用并不是很方便。

IPython 就是为弥补和克服 Python 自带的解释器的不足而开发的。IPython 的目标是创造一个交互式和探索式计算的完全的开发环境，为此 IPython 有以下 3 个主要的成分：

• 一个增强的交互式 Python shell；
• 一个解藕的双过程通信模型，允许多个客户端连接到同一个计算内核上（最常用的内核为基于网络的 notebook）；
• 支持交互式并行计算的基础结构。

对 IPython 的更多介绍请参考其文档。

用 IPython 做并行计算

IPthon 拥有成熟的和功能强大的支持并行和分布式计算的基础结构，其基础结构以一种通用的方式抽象了并行计算机制，从而允许 IPython 支持多种不同的并行计算模型，包括：

• 单程序多数据（SPMD）并行；
• 多程序多数据（MPMD）并行；
• MPI 信息传递模型；
• 任务运算（task farming）；
• 数据并行；
• 以上方式的组合；
• 用户定制的并行方式。

IPython 最显著的优势在于其允许以上各种类型的并行应用以一种交互式的方式开发，执行，调试和监测。

IPython 并行基础结构包含四个主要成分：

• IPython 引擎(IPython engine)；
• IPython 中心（IPython hub）；
• IPython 调度器（IPython schedulers）；
• 控制客户端（controller client）。

这些成分都包含在 IPython.parallel 包中并会随 IPython 一起安装，这些成分的关系如下图所示：

要使用 IPython 进行并行计算，必须首先启动一个控制器和至少一个引擎。最简单的方式是使用 ipcluster 命令，比如说要启动一个控制器和 4 个引擎，可以使用如下命令：

$ ipcluster start -n 4

对启动 IPython 控制器和引擎及对 ipcluster 命令的更多介绍请参考这里的文档。

在 IPython 中使用 mpi4py

可以在 IPython 中交互使用 MPI 进行并行分布式计算，需要的条件是：

• 一个标准的 MPI 实现，如 OpenMPI 和 MPICH 等；
• mpi4py 包。

具体的使用方式如下，我们需要首先启动至少一个使用 MPI 的 IPython 引擎，这可以有多种方式，一种简单的启动方式如下（比如说启动 4 个引擎）：

$ ipcluster start -n 4 --engines=MPI

一些魔法命令

在使用 IPython 做并行计算之前需要了解一些魔法命令（magic command），这些魔法命令可以使我们更方便地在 IPython shell 中交互式地在启动的引擎上执行 Python 语句。这些魔法命令实际上是 DirectView.execute() 和 AsyncResult.display_outputs() 函数的简写使用方式。这些魔法命令会在创建一个客户端时自动启用，如：

In [1]: from IPython.parallel import Client
In [2]: rc = Client()

这样会初始化一个默认的活动视图（参数设置为 targets=‘all’, block=True，表示使用所有启动的引擎进行阻塞式同步运算，即得到计算结果后才会返回，如果 block = False，则会立即返回一个 AsyncResult 对象，计算结果可能随后才能得到）。

下面是几个常用的魔法命令：

• %px：执行单个 Python 命令在指定的引擎上（通过 target 属性来指定所使用的引擎，如果不指定会使用所有引擎）；
• %%px：可以使用此魔法命令执行一段 Python 程序；
• %pxresult：在非阻塞模式下使用 %px 魔法命令的执行结果为一个 AsyncResult 对象，而使用 %pxresult 才会显示真正的结果。%pxresult 相当于在阻塞模式下使用 %px。
• %autopx：切换执行模式，使用 ％autopx 后面的命令都会在引擎上自动执行，再一次使用 %autopx 关闭此项功能。

对并行相关的魔法命令的更多介绍参见这里的文档。

交互式使用

按照上面介绍的命令启动 4 个使用 MPI 的引擎后，就可以交互式地执行并行计算语句了，举例如下：

In [1]: from IPython.parallel import Client

In [2]: c = Client()

In [3]: c.block = True # use block mode

In [4]: %px from mpi4py import MPI

In [5]: %px comm = MPI.COMM_WORLD

In [6]: %px print comm.size
[stdout:0] 4
[stdout:1] 4
[stdout:2] 4
[stdout:3] 4

In [7]: %px import numpy as np

In [8]: %px a = np.arange(4, dtype='d') + 4 * comm.rank

In [9]: %px print a
[stdout:0] [ 8.  9. 10. 11.]
[stdout:1] [0. 1. 2. 3.]
[stdout:2] [12. 13. 14. 15.]
[stdout:3] [4. 5. 6. 7.]

In [10]: %px locsum = np.array(np.sum(a))

In [11]: %px rcvBuf = np.array(0.0, 'd')

In [12]: %px comm.Allreduce([locsum, MPI.DOUBLE], [rcvBuf, MPI.DOUBLE], op=MPI.SUM)

In [13]: %px print rcvBuf
[stdout:0] 120.0
[stdout:1] 120.0
[stdout:2] 120.0
[stdout:3] 120.0

也可以导入 Python 文件并调用文件中的函数，比如说我们有下面这个文件：

# psum.py

import numpy as np
from mpi4py import MPI


comm = MPI.COMM_WORLD

def psum(a):
    # locsum = np.sum(a)
    locsum = np.array(np.sum(a))
    rcvBuf = np.array(0.0, 'd')
    comm.Allreduce([locsum, MPI.DOUBLE], [rcvBuf, MPI.DOUBLE], op=MPI.SUM)

    return rcvBuf

则可以用类似下面的方式执行该文件并调用其中的函数 psum：

In [1]: import numpy as np

In [2]: from IPython.parallel import Client

In [3]: c = Client()

In [4]: view = c[:]

In [5]: view.activate() # enable magics

In [6]: view.run('psum.py')
Out[6]: <AsyncResult: execute>

In [7]: view.scatter('a', np.arange(16, dtype='float'))
Out[7]: <AsyncResult: scatter>

In [8]: view['a']
Out[8]:
[array([0., 1., 2., 3.]),
 array([4., 5., 6., 7.]),
 array([ 8.,  9., 10., 11.]),
 array([12., 13., 14., 15.])]

In [9]: %px totalsum = psum(a)
Out[9]: <AsyncResult: execute>

In [10]: view['totalsum']
Out[10]: [array(120.), array(120.), array(120.), array(120.)]

以上简单地介绍了在 IPython shell 中交互式使用 mpi4py 进行并行计算的方法，更多内容可以参见其文档。