Python多進(jìn)程并行編程實(shí)踐: mpi4py的使用

在高性能計(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，即消息傳遞接口。

它并不是一門(mén)語(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等。

它是一種消息傳遞編程模型，顧名思義，它就是專(zhuān)門(mén)服務(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ù)共享，但是在編寫(xiě)起來(lái)仍然具有很大的不靈活性。而這一方面正是MPI所擅長(zhǎng)的領(lǐng)域，因此如果能夠在Python中調(diào)用MPI的接口那真是太完美了不是么。

MPI與mpi4py

mpi4py是一個(gè)構(gòu)建在MPI之上的Python庫(kù)，主要使用Cython編寫(xiě)。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

這里我開(kāi)始對(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中寫(xiě)入了初始化的操作，因此在我們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類(lèi)，其中Comm是通信域的基類(lèi)，Intracomm和Intercomm是其派生類(lèi)，這根MPI的C++實(shí)現(xiàn)中是相同的。

同時(shí)它也提供了兩個(gè)預(yù)定義的通信域?qū)ο?

包含所有進(jìn)程的COMM_WORLD

只包含調(diào)用進(jìn)程本身的COMM_SELF

In[1]:frommpi4pyimportMPI                  

                                                

In[2]:MPI.COMM_SELF                          

Out[2]: <mpi4py.MPI.Intracommat0x7f2fa2fd59d0>

                                                

In[3]:MPI.COMM_WORLD                          

Out[3]: <mpi4py.MPI.Intracommat0x7f2fa2fd59f0>

通信域?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.Groupat0x7f2fa40fec30>    

                                                

In[9]:comm.Split(0,0)                        

Out[9]: <mpi4py.MPI.Intracommat0x7f2fa2fd5bd0>

關(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ū)ο蟮姆椒ㄖ行?xiě)的接口，例如send(),recv(),isend()等。

如果需要直接傳遞數(shù)據(jù)對(duì)象，則需要調(diào)用大寫(xiě)的接口，例如Send(),Recv(),Isend()等，這與C++接口中的拼寫(xiě)是一樣的。

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ì)象。

frommpi4pyimportMPI

importnumpyasnp

comm=MPI.COMM_WORLD

rank=comm.Get_rank()

size=comm.Get_size()

ifrank==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-np2pythontemp.py

process0send[0,1,2,3,4,5,6,7,8,9]...

process1recv[0,1,2,3,4,5,6,7,8,9]...

非阻塞標(biāo)準(zhǔn)通信

所有的阻塞通信mpi都提供了一個(gè)非阻塞的版本，類(lèi)似與我們編寫(xiě)異步程序不阻塞在耗時(shí)的IO上是一樣的，MPI的非阻塞通信也不會(huì)阻塞消息的傳遞過(guò)程中，這樣能夠充分利用處理器資源提升整個(gè)程序的效率。

來(lái)張圖看看阻塞通信與非阻塞通信的對(duì)比：

非阻塞通信的消息發(fā)送和接受：

同樣的，我們也可以寫(xiě)一個(gè)上面例子的非阻塞版本。

frommpi4pyimportMPI                                        

importnumpyasnp                                            

                                                              

comm=MPI.COMM_WORLD                                          

rank=comm.Get_rank()                                        

size=comm.Get_size()                                        

                                                              

ifrank==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-np2pythontemp.py

process0immediatesend[0,1,2,3,4,5,6,7,8,9]...

process1recv[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的接口中也及其類(lèi)似, Comm.Send()中需要接收一個(gè)Python list作為參數(shù)，其中包含所傳數(shù)據(jù)的地址，長(zhǎng)度和類(lèi)型。

來(lái)個(gè)阻塞標(biāo)準(zhǔn)通信的例子：

frommpi4pyimportMPI                                                

importnumpyasnp                                                    

                                                                      

comm=MPI.COMM_WORLD                                                  

rank=comm.Get_rank()                                                

size=comm.Get_size()                                                

                                                                      

ifrank==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-np2pythontemp.py

process0Sendbuffer-likearray[0123456789]...

process1recvbuffer-likearray[0123456789]...

組通信

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)程中：

frommpi4pyimportMPI                                                    

                                                                          

comm=MPI.COMM_WORLD                                                      

rank=comm.Get_rank()                                                    

size=comm.Get_size()                                                    

                                                                          

ifrank==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-np5pythontemp.py

process0bcastdata[0,1,2,3,4,5,6,7,8,9]toother processes

process0recvdata[0,1,2,3,4,5,6,7,8,9]...

process1recvdata[0,1,2,3,4,5,6,7,8,9]...

process3recvdata[0,1,2,3,4,5,6,7,8,9]...

process2recvdata[0,1,2,3,4,5,6,7,8,9]...

process4recvdata[0,1,2,3,4,5,6,7,8,9]...

發(fā)散

與廣播不同，發(fā)散可以向不同的進(jìn)程發(fā)送不同的數(shù)據(jù)，而不是完全復(fù)制。

例如我想將0-9發(fā)送到不同的進(jìn)程中：

frommpi4pyimportMPI                                                            

importnumpyasnp                                                                

                                                                                  

comm=MPI.COMM_WORLD                                                            

rank=comm.Get_rank()                                                            

size=comm.Get_size()                                                            

                                                                                  

recv_data=None                                                                  

                                                                                  

ifrank==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-np10pythontemp.py

process0scatterdata[0,1,2,3,4,5,6,7,8,9]toother processes

process0recvdata0...

process3recvdata3...

process5recvdata5...

process8recvdata8...

process2recvdata2...

process7recvdata7...

process4recvdata4...

process1recvdata1...

process9recvdata9...

process6recvdata6..

收集

收集過(guò)程是發(fā)散過(guò)程的逆過(guò)程，每個(gè)進(jìn)程將發(fā)送緩沖區(qū)的消息發(fā)送給根進(jìn)程，根進(jìn)程根據(jù)發(fā)送進(jìn)程的進(jìn)程號(hào)將各自的消息存放到自己的消息緩沖區(qū)中。

收集結(jié)果：

zjshao@vaio:~/temp_codes/mpipy$/usr/bin/mpiexec-np5pythontemp.py

process2senddata2toroot...

process3senddata3toroot...

process0senddata0toroot...

process4senddata4toroot...

process1senddata1toroot...

process0gather alldata[0,1,2,3,4]...

其他的組內(nèi)通信還有歸約操作等等由于篇幅限制就不多講了，有興趣的可以去看看MPI的官方文檔和相應(yīng)的教材。

mpi4py并行編程實(shí)踐

這里我就上篇《Python 多進(jìn)程并行編程實(shí)踐: multiprocessing 模塊》中的二重循環(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)程中:

之后我需要在每個(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)

由于代碼都是涉及的專(zhuān)業(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的封裝文件和類(lèi)型映射，能夠幫助我們將Python同真正的C/C++以及Fortran程序在消息傳遞上實(shí)現(xiàn)統(tǒng)一。