作者：邵正將
PytLab，Python 中文社區(qū)專欄作者。主要從事科學(xué)計(jì)算與高性能計(jì)算領(lǐng)域的應(yīng)用，主要語(yǔ)言為Python，C，C++。熟悉數(shù)值算法(最優(yōu)化方法，蒙特卡洛算法等）與并行化 算法（MPI,OpenMP等多線程以及多進(jìn)程并行化）以及python優(yōu)化方法，經(jīng)常使用C++給python寫擴(kuò)展。
blog：http://ipytlab.com
github：https://github.com/PytLab
前言
在高性能計(jì)算的項(xiàng)目中我們通常都會(huì)使用效率更高的編譯型的語(yǔ)言例如C、C++、Fortran等，但是由于Python的靈活性和易用性使得它在發(fā)展和驗(yàn)證算法方面?zhèn)涫苋藗兊那嗖A于是在高性能計(jì)算領(lǐng)域也經(jīng)常能看到Python的身影了。本文簡(jiǎn)單介紹在Python環(huán)境下使用MPI接口在集群上進(jìn)行多進(jìn)程并行計(jì)算的方法。
MPI(Message Passing Interface)
這里我先對(duì)MPI進(jìn)行一下簡(jiǎn)單的介紹，MPI的全稱是Message Passing Interface，即消息傳遞接口。
它并不是一門語(yǔ)言，而是一個(gè)庫(kù)，我們可以用Fortran、C、C++結(jié)合MPI提供的接口來(lái)將串行的程序進(jìn)行并行化處理，也可以認(rèn)為Fortran+MPI或者C+MPI是一種再原來(lái)串行語(yǔ)言的基礎(chǔ)上擴(kuò)展出來(lái)的并行語(yǔ)言。
它是一種標(biāo)準(zhǔn)而不是特定的實(shí)現(xiàn)，具體的可以有很多不同的實(shí)現(xiàn)，例如MPICH、OpenMPI等。
它是一種消息傳遞編程模型，顧名思義，它就是專門服務(wù)于進(jìn)程間通信的。
MPI的工作方式很好理解，我們可以同時(shí)啟動(dòng)一組進(jìn)程，在同一個(gè)通信域中不同的進(jìn)程都有不同的編號(hào)，程序員可以利用MPI提供的接口來(lái)給不同編號(hào)的進(jìn)程分配不同的任務(wù)和幫助進(jìn)程相互交流最終完成同一個(gè)任務(wù)。就好比包工頭給工人們編上了工號(hào)然后指定一個(gè)方案來(lái)給不同編號(hào)的工人分配任務(wù)并讓工人相互溝通完成任務(wù)。
Python中的并行
由于CPython中的GIL的存在我們可以暫時(shí)不奢望能在CPython中使用多線程利用多核資源進(jìn)行并行計(jì)算了，因此我們?cè)赑ython中可以利用多進(jìn)程的方式充分利用多核資源。
Python中我們可以使用很多方式進(jìn)行多進(jìn)程編程，例如os.fork()來(lái)創(chuàng)建進(jìn)程或者通過(guò)multiprocessing模塊來(lái)更方便的創(chuàng)建進(jìn)程和進(jìn)程池等。在上一篇《Python多進(jìn)程并行編程實(shí)踐-multiprocessing模塊》中我們使用進(jìn)程池來(lái)方便的管理Python進(jìn)程并且通過(guò)multiprocessing模塊中的Manager管理分布式進(jìn)程實(shí)現(xiàn)了計(jì)算的多機(jī)分布式計(jì)算。
與多線程的共享式內(nèi)存不同，由于各個(gè)進(jìn)程都是相互獨(dú)立的，因此進(jìn)程間通信再多進(jìn)程中扮演這非常重要的角色，Python中我們可以使用multiprocessing模塊中的pipe、queue、Array、Value等等工具來(lái)實(shí)現(xiàn)進(jìn)程間通訊和數(shù)據(jù)共享，但是在編寫起來(lái)仍然具有很大的不靈活性。而這一方面正是MPI所擅長(zhǎng)的領(lǐng)域，因此如果能夠在Python中調(diào)用MPI的接口那真是太完美了不是么。
MPI與mpi4py
mpi4py是一個(gè)構(gòu)建在MPI之上的Python庫(kù)，主要使用Cython編寫。mpi4py使得Python的數(shù)據(jù)結(jié)構(gòu)可以方便的在多進(jìn)程中傳遞。
mpi4py是一個(gè)很強(qiáng)大的庫(kù)，它實(shí)現(xiàn)了很多MPI標(biāo)準(zhǔn)中的接口，包括點(diǎn)對(duì)點(diǎn)通信，組內(nèi)集合通信、非阻塞通信、重復(fù)非阻塞通信、組間通信等，基本上我能想到用到的MPI接口mpi4py中都有相應(yīng)的實(shí)現(xiàn)。不僅是Python對(duì)象，mpi4py對(duì)numpy也有很好的支持并且傳遞效率很高。同時(shí)它還提供了SWIG和F2PY的接口能夠讓我們將自己的Fortran或者C/C++程序在封裝成Python后仍然能夠使用mpi4py的對(duì)象和接口來(lái)進(jìn)行并行處理。可見(jiàn)mpi4py的作者的功力的確是非常了得。
mpi4py
這里我開始對(duì)在Python環(huán)境中使用mpi4py的接口進(jìn)行并行編程進(jìn)行介紹。
MPI環(huán)境管理
mpi4py提供了相應(yīng)的接口Init()和Finalize()來(lái)初始化和結(jié)束mpi環(huán)境。但是mpi4py通過(guò)在__init__.py中寫入了初始化的操作，因此在我們from mpi4py import MPI的時(shí)候就已經(jīng)自動(dòng)初始化mpi環(huán)境。
MPI_Finalize()被注冊(cè)到了Python的C接口Py_AtExit()，這樣在Python進(jìn)程結(jié)束時(shí)候就會(huì)自動(dòng)調(diào)用MPI_Finalize()， 因此不再需要我們顯式的去掉用Finalize()。
通信域(Communicator)
mpi4py直接提供了相應(yīng)的通信域的Python類，其中Comm是通信域的基類，Intracomm和Intercomm是其派生類，這根MPI的C++實(shí)現(xiàn)中是相同的。
同時(shí)它也提供了兩個(gè)預(yù)定義的通信域?qū)ο?
- 包含所有進(jìn)程的
COMM_WORLD - 只包含調(diào)用進(jìn)程本身的
COMM_SELF
In [1]: from mpi4py import MPI
In [2]: MPI.COMM_SELF
Out[2]: <mpi4py.MPI.Intracomm at 0x7f2fa2fd59d0>
In [3]: MPI.COMM_WORLD
Out[3]: <mpi4py.MPI.Intracomm at 0x7f2fa2fd59f0>
通信域?qū)ο髣t提供了與通信域相關(guān)的接口，例如獲取當(dāng)前進(jìn)程號(hào)、獲取通信域內(nèi)的進(jìn)程數(shù)、獲取進(jìn)程組、對(duì)進(jìn)程組進(jìn)行集合運(yùn)算、分割合并等等。
In [4]: comm = MPI.COMM_WORLD
In [5]: comm.Get_rank()
Out[5]: 0
In [6]: comm.Get_size()
Out[6]: 1
In [7]: comm.Get_group()
Out[7]: <mpi4py.MPI.Group at 0x7f2fa40fec30>
In [9]: comm.Split(0, 0)
Out[9]: <mpi4py.MPI.Intracomm at 0x7f2fa2fd5bd0>
關(guān)于通信域與進(jìn)程組的操作這里就不細(xì)講了，可以參考Introduction to Groups and Communicators
點(diǎn)對(duì)點(diǎn)通信
mpi4py提供了點(diǎn)對(duì)點(diǎn)通信的接口使得多個(gè)進(jìn)程間能夠互相傳遞Python的內(nèi)置對(duì)象（基于pickle序列化），同時(shí)也提供了直接的數(shù)組傳遞（numpy數(shù)組，接近C語(yǔ)言的效率）。
如果我們需要傳遞通用的Python對(duì)象，則需要使用通信域?qū)ο蟮姆椒ㄖ行懙慕涌?，例?code>send(),recv(),isend()等。
如果需要直接傳遞數(shù)據(jù)對(duì)象，則需要調(diào)用大寫的接口，例如Send(),Recv(),Isend()等，這與C++接口中的拼寫是一樣的。
MPI中的點(diǎn)到點(diǎn)通信有很多中，其中包括標(biāo)準(zhǔn)通信，緩存通信，同步通信和就緒通信，同時(shí)上面這些通信又有非阻塞的異步版本等等。這些在mpi4py中都有相應(yīng)的Python版本的接口來(lái)讓我們更靈活的處理進(jìn)程間通信。這里我只用標(biāo)準(zhǔn)通信的阻塞和非阻塞版本來(lái)做個(gè)舉例：
阻塞標(biāo)準(zhǔn)通信
這里我嘗試使用mpi4py的接口在兩個(gè)進(jìn)程中傳遞Python list對(duì)象。
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
if rank == 0:
data = range(10)
comm.send(data, dest=1, tag=11)
print("process {} send {}...".format(rank, data))
else:
data = comm.recv(source=0, tag=11)
print("process {} recv {}...".format(rank, data))
執(zhí)行效果：
zjshao@vaio:~/temp_codes/mpipy$ mpiexec -np 2 python temp.py
process 0 send [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
process 1 recv [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
非阻塞標(biāo)準(zhǔn)通信
所有的阻塞通信mpi都提供了一個(gè)非阻塞的版本，類似與我們編寫異步程序不阻塞在耗時(shí)的IO上是一樣的，MPI的非阻塞通信也不會(huì)阻塞消息的傳遞過(guò)程中，這樣能夠充分利用處理器資源提升整個(gè)程序的效率。
來(lái)張圖看看阻塞通信與非阻塞通信的對(duì)比：
非阻塞通信的消息發(fā)送和接受：
同樣的，我們也可以寫一個(gè)上面例子的非阻塞版本。
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
if rank == 0:
data = range(10)
comm.isend(data, dest=1, tag=11)
print("process {} immediate send {}...".format(rank, data))
else:
data = comm.recv(source=0, tag=11)
print("process {} recv {}...".format(rank, data))
執(zhí)行結(jié)果,注意非阻塞發(fā)送也可以用阻塞接收來(lái)接收消息：
zjshao@vaio:~/temp_codes/mpipy$ mpiexec -np 2 python temp.py
process 0 immediate send [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
process 1 recv [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
支持Numpy數(shù)組
mpi4py的一個(gè)很好的特點(diǎn)就是他對(duì)Numpy數(shù)組有很好的支持，我們可以通過(guò)其提供的接口來(lái)直接傳遞數(shù)據(jù)對(duì)象，這種方式具有很高的效率，基本上和C/Fortran直接調(diào)用MPI接口差不多（方式和效果）
例如我想傳遞長(zhǎng)度為10的int數(shù)組，MPI的C++接口是：
void Comm::Send(const void * buf, int count, const Datatype & datatype, int dest, int tag) const
在mpi4py的接口中也及其類似, Comm.Send()中需要接收一個(gè)Python list作為參數(shù)，其中包含所傳數(shù)據(jù)的地址，長(zhǎng)度和類型。
來(lái)個(gè)阻塞標(biāo)準(zhǔn)通信的例子：
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
if rank == 0:
data = np.arange(10, dtype='i')
comm.Send([data, MPI.INT], dest=1, tag=11)
print("process {} Send buffer-like array {}...".format(rank, data))
else:
data = np.empty(10, dtype='i')
comm.Recv([data, MPI.INT], source=0, tag=11)
print("process {} recv buffer-like array {}...".format(rank, data))
執(zhí)行效果：
zjshao@vaio:~/temp_codes/mpipy$ /usr/bin/mpiexec -np 2 python temp.py
process 0 Send buffer-like array [0 1 2 3 4 5 6 7 8 9]...
process 1 recv buffer-like array [0 1 2 3 4 5 6 7 8 9]...
組通信
MPI組通信和點(diǎn)到點(diǎn)通信的一個(gè)重要區(qū)別就是，在某個(gè)進(jìn)程組內(nèi)所有的進(jìn)程同時(shí)參加通信，mpi4py提供了方便的接口讓我們完成Python中的組內(nèi)集合通信，方便編程同時(shí)提高程序的可讀性和可移植性。
下面就幾個(gè)常用的集合通信來(lái)小試牛刀吧。
廣播
廣播操作是典型的一對(duì)多通信，將跟進(jìn)程的數(shù)據(jù)復(fù)制到同組內(nèi)其他所有進(jìn)程中。
在Python中我想將一個(gè)列表廣播到其他進(jìn)程中：
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
if rank == 0:
data = range(10)
print("process {} bcast data {} to other processes".format(rank, data))
else:
data = None
data = comm.bcast(data, root=0)
print("process {} recv data {}...".format(rank, data))
執(zhí)行結(jié)果：
zjshao@vaio:~/temp_codes/mpipy$ /usr/bin/mpiexec -np 5 python temp.py
process 0 bcast data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] to other processes
process 0 recv data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
process 1 recv data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
process 3 recv data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
process 2 recv data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
process 4 recv data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]...
發(fā)散
與廣播不同，發(fā)散可以向不同的進(jìn)程發(fā)送不同的數(shù)據(jù)，而不是完全復(fù)制。
例如我想將0-9發(fā)送到不同的進(jìn)程中：
m mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
recv_data = None
if rank == 0:
send_data = range(10)
print("process {} scatter data {} to other processes".format(rank, send_data))
else:
send_data = None
recv_data = comm.scatter(send_data, root=0)
print("process {} recv data {}...".format(rank, recv_data))
發(fā)散結(jié)果：
zjshao@vaio:~/temp_codes/mpipy$ /usr/bin/mpiexec -np 10 python temp.py
process 0 scatter data [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] to other processes
process 0 recv data 0...
process 3 recv data 3...
process 5 recv data 5...
process 8 recv data 8...
process 2 recv data 2...
process 7 recv data 7...
process 4 recv data 4...
process 1 recv data 1...
process 9 recv data 9...
process 6 recv data 6...
收集
收集過(guò)程是發(fā)散過(guò)程的逆過(guò)程，每個(gè)進(jìn)程將發(fā)送緩沖區(qū)的消息發(fā)送給根進(jìn)程，根進(jìn)程根據(jù)發(fā)送進(jìn)程的進(jìn)程號(hào)將各自的消息存放到自己的消息緩沖區(qū)中。
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
send_data = rank
print "process {} send data {} to root...".format(rank, send_data)
recv_data = comm.gather(send_data, root=0)
if rank == 0:
print "process {} gather all data {}...".format(rank, recv_data)
收集結(jié)果：
zjshao@vaio:~/temp_codes/mpipy$ /usr/bin/mpiexec -np 5 python temp.py
process 2 send data 2 to root...
process 3 send data 3 to root...
process 0 send data 0 to root...
process 4 send data 4 to root...
process 1 send data 1 to root...
process 0 gather all data [0, 1, 2, 3, 4]...
其他的組內(nèi)通信還有歸約操作等等由于篇幅限制就不多講了，有興趣的可以去看看MPI的官方文檔和相應(yīng)的教材。
mpi4py并行編程實(shí)踐
這里我就上篇中的二重循環(huán)繪制map的例子來(lái)使用mpi4py進(jìn)行并行加速處理。
我打算同時(shí)啟動(dòng)10個(gè)進(jìn)程來(lái)將每個(gè)0軸需要計(jì)算和繪制的數(shù)據(jù)發(fā)送到不同的進(jìn)程進(jìn)行并行計(jì)算。
因此我需要將pO2s數(shù)組發(fā)散到10個(gè)進(jìn)程中:
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
if rank == 0:
pO2 = np.linspace(1e-5, 0.5, 10)
else:
pO2 = None
pO2 = comm.scatter(pO2, root=0)
pCOs = np.linspace(1e-5, 0.5, 10)
之后我需要在每個(gè)進(jìn)程中根據(jù)接受到的pO2s的數(shù)據(jù)再進(jìn)行一次pCOs循環(huán)來(lái)進(jìn)行計(jì)算。
最終將每個(gè)進(jìn)程計(jì)算的結(jié)果(TOF)進(jìn)行收集操作:
comm.gather(tofs_1d, root=0)
由于代碼都是涉及的專業(yè)相關(guān)的東西我就不全列出來(lái)了，將mpi4py改過(guò)的并行版本放到10個(gè)進(jìn)程中執(zhí)行可見(jiàn)：
效率提升了10倍左右。
總結(jié)
本文簡(jiǎn)單介紹了mpi4py的接口在python中進(jìn)行多進(jìn)程編程的方法，MPI的接口非常龐大，相應(yīng)的mpi4py也非常龐大，mpi4py還有實(shí)現(xiàn)了相應(yīng)的SWIG和F2PY的封裝文件和類型映射，能夠幫助我們將Python同真正的C/C++以及Fortran程序在消息傳遞上實(shí)現(xiàn)統(tǒng)一。有興趣的同學(xué)可以進(jìn)一步研究一下，歡迎交流。
