python並行計算之mpi4py的安裝與基本使用
技術背景
在之前的部落格中我們介紹過concurrent等python多進程任務的方案，而之所以我們又在考慮MPI等方案來實現python並行計算的原因，其實是將python的計算任務與並行計算的任務調度分層實現。在concurrent和multiprocessing等方案中，我們的python計算任務和調度任務是一體化的，而且還有一個比較大的限制是沒辦法跨節點操作的，這對於任務與環境的訂製化程度要求是比較高的。而MPI的方案在設計初期就考慮到了多節點之間通訊的問題，而這種分層式的任務調度解決方案其實在架構上看也更加的合理。做計算的人只要考慮單個進程下的任務如何執行就可以了，至於任務如何並行如何調度，那就是上層的MPI該做的事情了。
mpi4py的安裝
這裡推薦使用conda直接安裝，如果採用pip安裝的話，可能會有些環境依賴的問題出現：
$ conda install mpi4py
Collecting package metadata (current_repodata.json): done
Solving environment: done
## Package Plan ##
environment location: /home/dechin/anaconda3
added / updated specs:
- mpi4py
The following packages will be downloaded:
package | build
---------------------------|-----------------
mpi-1.0 | mpich 13 KB defaults
mpi4py-3.0.3 | py38h028fd6f_0 572 KB defaults
mpich-3.3.2 | hc856adb_0 3.8 MB defaults
------------------------------------------------------------
Total: 4.4 MB
The following NEW packages will be INSTALLED:
mpi pkgs/main/linux-64::mpi-1.0-mpich
mpi4py pkgs/main/linux-64::mpi4py-3.0.3-py38h028fd6f_0
mpich pkgs/main/linux-64::mpich-3.3.2-hc856adb_0
The following packages will be UPDATED:
ca-certificates 2021.9.30-h06a4308_1 --> 2021.10.26-h06a4308_2
Proceed ([y]/n)? y
Downloading and Extracting Packages
mpi4py-3.0.3 | 572 KB | ############################################## | 100%
mpich-3.3.2 | 3.8 MB | ############################################## | 100%
mpi-1.0 | 13 KB | ############################################## | 100%
Preparing transaction: done
Verifying transaction: done
Executing transaction: done
安裝完成後可以通過python3 -c "from mpi4py import MPI"來檢查是否安裝成功，下面我們來看一些具體的使用案例。
使用案例
首先了解下mpi的基本使用方法，如果我們使用mpirun -n 3 python3 test.py這樣的指令去運行一個程式，那麼就會給每一個不同的test.py中發送一個互不相同的rank，這個rank的範圍是從0開始數的。比如如下案例我們使用Get_rank()方法就可以獲取到mpi所傳遞下來的rank id，這樣進程就知道了自己所處的進程編號，我們寫好每個編號下所需要執行的任務即可：
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
print('My rank is ',rank)
執行效果如下：
$ mpirun -n 4 python3 mpi_test.py
My rank is 2
My rank is 1
My rank is 0
My rank is 3
當然，因為每個任務都是平級的關係，因此mpi的rank id也是隨機發放的，這個沒辦法控制，但是我們可以用如下的方法在進程間通訊：
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if rank == 0:
idata = 1
comm.send(idata, dest=1)
print ('This is process {}'.format(rank), '\nData send to process 1 successfully!')
elif rank == 1:
idata = comm.recv(source=0)
print ('This is process {}, data is '.format(rank),idata)
在這個案例中，我們從rank id為0的進程向rank id為1的進程發送了一個整數變數的數據。因為我們並不知道這個腳本什麼時候會被分配到rank 0什麼時候會被分配到rank 1，因此在同一個腳本內我們就需要分別對這兩種可能發生的情況進行針對性的處理。運行結果如下：
$ mpirun -n 2 python3 mpi_test.py
This is process 0
Data send to process 1 successfully!
This is process 1, data is 1
整數型的變數被成功的傳遞，當然，這裡面其實還有個更加重要的意義是，rank為1的進程實際上是對rank為0的進程有時間序列上的依賴的，我們必須執行完rank 0中的任務，才能再執行rank 1中的任務，這個是有可能在實際的應用過程中被頻繁使用的功能，尤其是任務之間互相有依賴的情況下。當然，進程之間的通訊不僅僅可以傳遞整數型的變數，還可以傳遞其他類型，比如字典或者一個numpy的數組：
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if rank == 0:
numData = 10
comm.send(numData, dest=1)
data = np.linspace(0.0, 3.14, numData)
comm.Send(data, dest=1)
print ('This is process {}'.format(rank), '\nData send to process 1 successfully!')
elif rank == 1:
numData = comm.recv(source=0)
print('Number of data to receive: ', numData)
data = np.empty(numData, dtype='d')
comm.Recv(data, source=0)
print ('This is process {}, data is '.format(rank),data)
運行結果如下：
$ mpirun -n 2 python3 mpi_test.py
This is process 0
Data send to process 1 successfully!
Number of data to receive: 10
This is process 1, data is [0. 0.34888889 0.69777778 1.04666667 1.39555556 1.74444444
2.09333333 2.44222222 2.79111111 3.14 ]
可以看到結果是被成功的接收了的。
總結概要
在這篇文章中，我們並沒有介紹非常多的MPI的應用，其實在普通的並行或者分散式任務中，掌握了各個進程的任務處理過程以及進程間的通訊方法就足夠了。總體來說，MPI是一個非常通用也非常高效的並行計算軟體。有了這些專業的並行化任務調度軟體，我們就可以專註於專業任務的程式碼和演算法上，而不需要過多的去關注並行任務的調度和分配問題。
版權聲明
本文首發鏈接為：//www.cnblogs.com/dechinphy/p/mpi4py.html
作者ID：DechinPhy
更多原著文章請參考：//www.cnblogs.com/dechinphy/
打賞專用鏈接：//www.cnblogs.com/dechinphy/gallery/image/379634.html
騰訊雲專欄同步：//cloud.tencent.com/developer/column/91958
