在 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 進行並行計算的方法，更多內容可以參見其文檔。