Python多進(jìn)程并行編程實踐: mpi4py 的使用
前言
在高性能計算的項目中我們通常都會使用效率更高的編譯型的語言例如C、C++、Fortran等，但是由于Python的靈活性和易用性使得它在發(fā)展和驗證算法方面?zhèn)涫苋藗兊那嗖A于是在高性能計算領(lǐng)域也經(jīng)常能看到Python的身影了。本文簡單介紹在Python環(huán)境下使用MPI接口在集群上進(jìn)行多進(jìn)程并行計算的方法。
MPI(Message Passing Interface)
這里我先對MPI進(jìn)行一下簡單的介紹，MPI的全稱是Message Passing Interface，即消息傳遞接口。
- 它并不是一門語言，而是一個庫，我們可以用Fortran、C、C++結(jié)合MPI提供的接口來將串行的程序進(jìn)行并行化處理，也可以認(rèn)為Fortran+MPI或者C+MPI是一種再原來串行語言的基礎(chǔ)上擴展出來的并行語言。
- 它是一種標(biāo)準(zhǔn)而不是特定的實現(xiàn)，具體的可以有很多不同的實現(xiàn)，例如MPICH、OpenMPI等。
- 它是一種消息傳遞編程模型，顧名思義，它就是專門服務(wù)于進(jìn)程間通信的。
MPI的工作方式很好理解，我們可以同時啟動一組進(jìn)程，在同一個通信域中不同的進(jìn)程都有不同的編號，程序員可以利用MPI提供的接口來給不同編號的進(jìn)程分配不同的任務(wù)和幫助進(jìn)程相互交流最終完成同一個任務(wù)。就好比包工頭給工人們編上了工號然后指定一個方案來給不同編號的工人分配任務(wù)并讓工人相互溝通完成任務(wù)。
Python中的并行
由于CPython中的GIL的存在我們可以暫時不奢望能在CPython中使用多線程利用多核資源進(jìn)行并行計算了，因此我們在Python中可以利用多進(jìn)程的方式充分利用多核資源。
Python中我們可以使用很多方式進(jìn)行多進(jìn)程編程，例如os.fork()來創(chuàng)建進(jìn)程或者通過multiprocessing模塊來更方便的創(chuàng)建進(jìn)程和進(jìn)程池等。在上一篇《Python多進(jìn)程并行編程實踐-multiprocessing模塊》中我們使用進(jìn)程池來方便的管理Python進(jìn)程并且通過multiprocessing模塊中的Manager管理分布式進(jìn)程實現(xiàn)了計算的多機分布式計算。
與多線程的共享式內(nèi)存不同，由于各個進(jìn)程都是相互獨立的，因此進(jìn)程間通信再多進(jìn)程中扮演這非常重要的角色，Python中我們可以使用multiprocessing模塊中的pipe、queue、Array、Value等等工具來實現(xiàn)進(jìn)程間通訊和數(shù)據(jù)共享，但是在編寫起來仍然具有很大的不靈活性。而這一方面正是MPI所擅長的領(lǐng)域，因此如果能夠在Python中調(diào)用MPI的接口那真是太***了不是么。
MPI與mpi4py
mpi4py是一個構(gòu)建在MPI之上的Python庫，主要使用Cython編寫。mpi4py使得Python的數(shù)據(jù)結(jié)構(gòu)可以方便的在多進(jìn)程中傳遞。
mpi4py是一個很強大的庫，它實現(xiàn)了很多MPI標(biāo)準(zhǔn)中的接口，包括點對點通信，組內(nèi)集合通信、非阻塞通信、重復(fù)非阻塞通信、組間通信等，基本上我能想到用到的MPI接口mpi4py中都有相應(yīng)的實現(xiàn)。不僅是Python對象，mpi4py對numpy也有很好的支持并且傳遞效率很高。同時它還提供了SWIG和F2PY的接口能夠讓我們將自己的Fortran或者C/C++程序在封裝成Python后仍然能夠使用mpi4py的對象和接口來進(jìn)行并行處理?？梢妋pi4py的作者的功力的確是非常了得。
mpi4py
這里我開始對在Python環(huán)境中使用mpi4py的接口進(jìn)行并行編程進(jìn)行介紹。
MPI環(huán)境管理
mpi4py提供了相應(yīng)的接口Init()和Finalize()來初始化和結(jié)束mpi環(huán)境。但是mpi4py通過在__init__.py中寫入了初始化的操作，因此在我們from mpi4py import MPI的時候就已經(jīng)自動初始化mpi環(huán)境。
MPI_Finalize()被注冊到了Python的C接口Py_AtExit()，這樣在Python進(jìn)程結(jié)束時候就會自動調(diào)用MPI_Finalize()， 因此不再需要我們顯式的去掉用Finalize()。
通信域(Communicator)
mpi4py直接提供了相應(yīng)的通信域的Python類，其中Comm是通信域的基類，Intracomm和Intercomm是其派生類，這根MPI的C++實現(xiàn)中是相同的。
同時它也提供了兩個預(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)程號、獲取通信域內(nèi)的進(jìn)程數(shù)、獲取進(jìn)程組、對進(jìn)程組進(jì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
點對點通信
mpi4py提供了點對點通信的接口使得多個進(jìn)程間能夠互相傳遞Python的內(nèi)置對象(基于pickle序列化)，同時也提供了直接的數(shù)組傳遞(numpy數(shù)組，接近C語言的效率)。
如果我們需要傳遞通用的Python對象，則需要使用通信域?qū)ο蟮姆椒ㄖ行懙慕涌?，例如send(),recv(),isend()等。
如果需要直接傳遞數(shù)據(jù)對象，則需要調(diào)用大寫的接口，例如Send(),Recv(),Isend()等，這與C++接口中的拼寫是一樣的。
MPI中的點到點通信有很多中，其中包括標(biāo)準(zhǔn)通信，緩存通信，同步通信和就緒通信，同時上面這些通信又有非阻塞的異步版本等等。這些在mpi4py中都有相應(yīng)的Python版本的接口來讓我們更靈活的處理進(jìn)程間通信。這里我只用標(biāo)準(zhǔn)通信的阻塞和非阻塞版本來做個舉例：
阻塞標(biāo)準(zhǔn)通信
這里我嘗試使用mpi4py的接口在兩個進(jìn)程中傳遞Python list對象。
- 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都提供了一個非阻塞的版本，類似與我們編寫異步程序不阻塞在耗時的IO上是一樣的，MPI的非阻塞通信也不會阻塞消息的傳遞過程中，這樣能夠充分利用處理器資源提升整個程序的效率。
來張圖看看阻塞通信與非阻塞通信的對比：
非阻塞通信的消息發(fā)送和接受：
同樣的，我們也可以寫一個上面例子的非阻塞版本。
- 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ā)送也可以用阻塞接收來接收消息：
- 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的一個很好的特點就是他對Numpy數(shù)組有很好的支持，我們可以通過其提供的接口來直接傳遞數(shù)據(jù)對象，這種方式具有很高的效率，基本上和C/Fortran直接調(diào)用MPI接口差不多(方式和效果)
例如我想傳遞長度為10的int數(shù)組，MPI的C++接口是：
- void Comm::Send(const void * buf, int count, const Datatype & datatype, int dest, int tag) const
在mpi4py的接口中也及其類似, Comm.Send()中需要接收一個Python list作為參數(shù)，其中包含所傳數(shù)據(jù)的地址，長度和類型。
來個阻塞標(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組通信和點到點通信的一個重要區(qū)別就是，在某個進(jìn)程組內(nèi)所有的進(jìn)程同時參加通信，mpi4py提供了方便的接口讓我們完成Python中的組內(nèi)集合通信，方便編程同時提高程序的可讀性和可移植性。
下面就幾個常用的集合通信來小試牛刀吧。
廣播
廣播操作是典型的一對多通信，將跟進(jìn)程的數(shù)據(jù)復(fù)制到同組內(nèi)其他所有進(jìn)程中。
在Python中我想將一個列表廣播到其他進(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)程中：
- from 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...
收集
收集過程是發(fā)散過程的逆過程，每個進(jìn)程將發(fā)送緩沖區(qū)的消息發(fā)送給根進(jìn)程，根進(jìn)程根據(jù)發(fā)送進(jìn)程的進(jìn)程號將各自的消息存放到自己的消息緩沖區(qū)中。
收集結(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并行編程實踐
這里我就上篇《Python 多進(jìn)程并行編程實踐: multiprocessing 模塊》中的二重循環(huán)繪制map的例子來使用mpi4py進(jìn)行并行加速處理。
我打算同時啟動10個進(jìn)程來將每個0軸需要計算和繪制的數(shù)據(jù)發(fā)送到不同的進(jìn)程進(jìn)行并行計算。
因此我需要將pO2s數(shù)組發(fā)散到10個進(jìn)程中:
之后我需要在每個進(jìn)程中根據(jù)接受到的pO2s的數(shù)據(jù)再進(jìn)行一次pCOs循環(huán)來進(jìn)行計算。
最終將每個進(jìn)程計算的結(jié)果(TOF)進(jìn)行收集操作:
- comm.gather(tofs_1d, root=0)
由于代碼都是涉及的專業(yè)相關(guān)的東西我就不全列出來了，將mpi4py改過的并行版本放到10個進(jìn)程中執(zhí)行可見：
效率提升了10倍左右。
總結(jié)
本文簡單介紹了mpi4py的接口在python中進(jìn)行多進(jìn)程編程的方法，MPI的接口非常龐大，相應(yīng)的mpi4py也非常龐大，mpi4py還有實現(xiàn)了相應(yīng)的SWIG和F2PY的封裝文件和類型映射，能夠幫助我們將Python同真正的C/C++以及Fortran程序在消息傳遞上實現(xiàn)統(tǒng)一。有興趣的同學(xué)可以進(jìn)一步研究一下，歡迎交流。
參考
- MPI for Python 2.0.0 documentation
- MPI Tutorial
- A Python Introduction to Parallel Programming with MPI
- 《高性能計算并行編程技術(shù)-MPI并行程序設(shè)計》
- 《MPI并行程序設(shè)計實例教程》
