MPI-3 中線程安全的 Mprobe

本文從本人簡書博客同步過來

在上一篇中我們介紹了 MPI 中多線程的使用，下面我們將介紹 MPI-3 中線程安全的 Mprobe。

Probe 和 Iprobe 操作的作用是在不實際接收消息的情況下檢查消息中包含的信息，據此決定接收消息的具體方式——如爲待接收消息申請空間等。

但是 Probe 和 Iprobe 操作不是線程安全的，如在多線程中使用 Probe 或 Iprobe 來確定一個消息的來源和大小，然後接收該消息的操作就不是線程安全的，有可能會導致程序無限阻塞或發生錯誤。因此在多線程環境下，必須通過同步/互斥機制使得同一時刻僅有一個線程執行與期望通信相匹配的操作，如某個線程的接收操作使用 Probe 或 Iprobe 檢測得到的 <源進程，tag> 參數時，不得再有其它線程使用同樣的 <源進程，tag> 二元組去匹配同一個發送操作來接收數據。

線程安全性在上一篇中已經作過介紹，具體來說，線程安全性是指多個線程可以同時執行消息傳遞的相關調用而不會相互影響。MPI-3 中引進了線程安全的 Mprobe 和 Improbe 操作，該操作與一個新引進的對象 MPI.Message 綁定，該對象標識被匹配到的特定消息，某個消息一旦被 Mprobe 或 Improbe 匹配到，就不會再被其它的 Mprobe 或 Improbe 匹配到，也不會被相應的 Mrecv 接收到，因此可以在多線程環境下安全地使用，避免 Probe 和 Iprobe 的以上限制，增加程序的安全性並簡化編程。被 Mprobe 或 Improbe 匹配到的消息必須由 MPI.Message.Recv 或 MPI.Message.Irecv 來接收，而不能以通常的 MPI.Comm.Recv 或 MPI.Comm.Irecv 來接收。

方法接口

非線程安全的 probe

MPI.Comm.probe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

檢測與指定參數相匹配的消息。source 爲消息源，tag 爲消息 tag，status 如果非 None，爲一個 MPI.Status 對象。可使用 MPI.ANY_SOURCE 和 MPI.ANY_TAG 來匹配任意消息。該函數爲阻塞操作，直到檢測到所匹配的消息纔會結束阻塞而返回，此時如果傳遞了 status 參數，則會爲其設置所檢測到的消息的狀態信息。

注意：此是以小寫字母開頭的方法，一般應該用來檢測被 pickle 系列化的消息。

MPI.Comm.iprobe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Comm.probe 的非阻塞版本，參數同。該函數被調用後會立即返回，並且可以對相同的消息檢測多次，如果檢測到所匹配的消息，會返回 True，否則返回 False。

注意：此是以小寫字母開頭的方法，一般應該用來檢測被 pickle 系列化的消息。

MPI.Comm.Probe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Comm.probe 的以大寫字母開頭的版本，爲阻塞操作，一般應該用來檢測以緩衝區形式發送的消息。

MPI.Comm.Iprobe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Comm.iprobe 的以大寫字母開頭的版本，爲非阻塞操作，一般應該用來檢測以緩衝區形式發送的消息。

線程安全的 mprobe

線程安全的 mprobe/improbe/Mprobe/Improbe 都有兩種等價的形式，一種爲 MPI.Comm 類的方法， 另一種爲 MPI.Message 類的類方法。這兩種形式的區別是，前者直接在一個通信子對象上調用，後者需將通信子作爲一個參數傳人。兩者都返回一個 MPI.Message 對象。注意：被 mprobe/improbe/Mprobe/Improbe 檢測到的消息應該用 MPI.Message 類的接收方法 recv/irecv/Recv/Irecv 來接收，阻塞的檢測方法可以用阻塞的接收方法也可以用非阻塞的接收方法，同樣非阻塞的檢測方法可以用阻塞的接收方法也可以用非阻塞的接收方法，但以小（大）寫字母開頭的檢測方法一般應該與以小（大）寫字母開頭的接收方法配對使用。

下面是這些方法的使用接口：

MPI.Comm.mprobe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)
MPI.Message.probe(type cls, Comm comm, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Message.recv(self, Status status=None)

阻塞的 mprobe 和阻塞的 recv，用於 pickle 系列化的消息。

MPI.Comm.improbe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)
MPI.Message.iprobe(type cls, Comm comm, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Message.irecv(self)

非阻塞的 improbe 和非阻塞的 irecv，用於 pickle 系列化的消息。

MPI.Comm.Mprobe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)
MPI.Message.Probe(type cls, Comm comm, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Message.Recv(self, buf, Status status=None)

阻塞的 Mprobe 和阻塞的 Recv，用於緩衝區消息。

MPI.Comm.Improbe(self, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)
MPI.Message.Iprobe(type cls, Comm comm, int source=ANY_SOURCE, int tag=ANY_TAG, Status status=None)

MPI.Message.Irecv(self, buf)

非阻塞的 Improbe 和非阻塞的 Irecv，用於緩衝區消息。

例程

下面首先展示非線程安全的 probe 操作在多線程環境下可能導致的錯誤。在下面這個例程中，進程 1 的兩個線程都使用 probe 來檢測消息，然後接收檢測到的消息。如果這兩個線程在接收之前檢測到同一個消息，則先執行接收操作的線程會成功接收到消息，而另一個線程再執行接收動作時，因爲該消息已經被前一個線程接收，該線程可能會無限阻塞，或者發生錯誤。

# probe_error.py

"""
This is an example which shows the incorrect use of probe in
multi-threading environment.

Run this with 2 processes like:
$ mpiexec -n 2 python probe_error.py
"""

import sys
import time
import numpy as np
import threading
from mpi4py import MPI


if MPI.Query_thread() < MPI.THREAD_MULTIPLE:
    sys.stderr.write("MPI does not provide enough thread support\n")
    sys.exit(0)

comm = MPI.COMM_WORLD
rank = comm.rank

# -----------------------------------------------------------------------------------
# mprobe and recv
if rank == 0:
    comm.send(11, dest=1, tag=11)
    comm.send(22, dest=1, tag=22)
elif rank == 1:

    def recv():

        status = MPI.Status()
        comm.probe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        print '%s probed source = %d, tag = %d' % (threading.currentThread().getName(), status.source, status.tag)

        time.sleep(0.5)

        recv_obj = comm.recv(source=status.source, tag=status.tag)
        print '%s receives %d from rank 0' % (threading.currentThread().getName(), recv_obj)

    # create thread by using a function
    recv_thread1 = threading.Thread(target=recv, name='[rank-%d recv_thread 1]' % rank)
    recv_thread2 = threading.Thread(target=recv, name='[rank-%d recv_thread 2]' % rank)

    # start the threads
    recv_thread1.start()
    recv_thread2.start()

    # wait for complete
    recv_thread1.join()
    recv_thread2.join()

運行結果如下：

$ mpiexec -n 2 python probe_error.py
[rank-1 recv_thread 1] probed source = 0, tag = 11
[rank-1 recv_thread 2] probed source = 0, tag = 11
[rank-1 recv_thread 1] receives 11 from rank 0

以上例程中，線程 1 和 線程 2 都檢測到 source 爲 0，tag 爲 11 的同一個消息，線程 1 成功接收該消息，線程 2 則無限阻塞。

下面給出線程安全的 mprobe 使用例程。

# mprobe.py

"""
Demonstrates the usage of mprobe.

Run this with 2 processes like:
$ mpiexec -n 2 python mprobe.py
"""

import sys
import numpy as np
import threading
from mpi4py import MPI


if MPI.Query_thread() < MPI.THREAD_MULTIPLE:
    sys.stderr.write("MPI does not provide enough thread support\n")
    sys.exit(0)

comm = MPI.COMM_WORLD
rank = comm.rank

# -----------------------------------------------------------------------------------
# mprobe and recv
if rank == 0:
    comm.send(11, dest=1, tag=11)
    comm.send(22, dest=1, tag=22)
elif rank == 1:

    def recv():

        status = MPI.Status()
        msg = comm.mprobe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        # msg = MPI.Message.probe(comm, source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)

        recv_obj = msg.recv()
        print '%s receives %d from rank 0' % (threading.currentThread().getName(), recv_obj)

    # create thread by using a function
    recv_thread1 = threading.Thread(target=recv, name='[rank-%d recv_thread 1]' % rank)
    recv_thread2 = threading.Thread(target=recv, name='[rank-%d recv_thread 2]' % rank)

    # start the threads
    recv_thread1.start()
    recv_thread2.start()

    # wait for complete
    recv_thread1.join()
    recv_thread2.join()


comm.barrier()

# ------------------------------------------------------------------------------------
# Mprobe and Recv
if rank == 0:
    comm.Send(np.array([33, 33]), dest=1, tag=33)
    comm.Send(np.array([44, 44]), dest=1, tag=44)
elif rank == 1:

    def recv():

        status = MPI.Status()
        msg = comm.Mprobe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
        # msg = MPI.Message.Probe(comm, source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)

        recv_buf = np.array([-1, -1])
        msg.Recv(recv_buf)
        print '%s receives %s from rank 0' % (threading.currentThread().getName(), recv_buf)

    # create thread by using a function
    recv_thread1 = threading.Thread(target=recv, name='[rank-%d recv_thread 1]' % rank)
    recv_thread2 = threading.Thread(target=recv, name='[rank-%d recv_thread 2]' % rank)

    # start the threads
    recv_thread1.start()
    recv_thread2.start()

    # wait for complete
    recv_thread1.join()
    recv_thread2.join()


comm.barrier()

# ------------------------------------------------------------------------------------
# improbe and irecv
if rank == 0:
    comm.send(55, dest=1, tag=55)
    comm.send(66, dest=1, tag=66)
elif rank == 1:

    def recv():

        status = MPI.Status()
        msg = None
        while not msg:
            msg = comm.improbe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            # msg = MPI.Message.iprobe(comm, source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            if msg is None:
                print '%s improbe is not completed...' % threading.currentThread().getName()

        req = msg.irecv()
        recv_obj = req.wait()
        print '%s receives %d from rank 0' % (threading.currentThread().getName(), recv_obj)

    # create thread by using a function
    recv_thread1 = threading.Thread(target=recv, name='[rank-%d recv_thread 1]' % rank)
    recv_thread2 = threading.Thread(target=recv, name='[rank-%d recv_thread 2]' % rank)

    # start the threads
    recv_thread1.start()
    recv_thread2.start()

    # wait for complete
    recv_thread1.join()
    recv_thread2.join()


comm.barrier()

# ------------------------------------------------------------------------------------
# Improbe and Irecv
if rank == 0:
    comm.Send(np.array([77, 77]), dest=1, tag=77)
    comm.Send(np.array([88, 88]), dest=1, tag=88)
elif rank == 1:

    def recv():

        status = MPI.Status()
        msg = None
        while not msg:
            msg = comm.Improbe(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            # msg = MPI.Message.Iprobe(comm, source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
            if msg is None:
                print '%s Improbe is not completed...' % threading.currentThread().getName()

        recv_buf = np.array([-1, -1])
        req = msg.Irecv(recv_buf)
        req.Wait()
        print '%s receives %s from rank 0' % (threading.currentThread().getName(), recv_buf)

    # create thread by using a function
    recv_thread1 = threading.Thread(target=recv, name='[rank-%d recv_thread 1]' % rank)
    recv_thread2 = threading.Thread(target=recv, name='[rank-%d recv_thread 2]' % rank)

    # start the threads
    recv_thread1.start()
    recv_thread2.start()

    # wait for complete
    recv_thread1.join()
    recv_thread2.join()

運行結果如下：

$ mpiexec -n 2 python mprobe.py
[rank-1 recv_thread 1] receives 11 from rank 0
[rank-1 recv_thread 2] receives 22 from rank 0
[rank-1 recv_thread 1] receives [33 33] from rank 0
[rank-1 recv_thread 2] receives [44 44] from rank 0
[rank-1 recv_thread 1] improbe is not completed...
[rank-1 recv_thread 1] receives 55 from rank 0
[rank-1 recv_thread 2] improbe is not completed...
[rank-1 recv_thread 2] receives 66 from rank 0
[rank-1 recv_thread 1] Improbe is not completed...
[rank-1 recv_thread 1] receives [77 77] from rank 0
[rank-1 recv_thread 2] Improbe is not completed...
[rank-1 recv_thread 2] receives [88 88] from rank 0

以上介紹了非線程安全的 Probe 和 MPI-3 中線程安全的 Mprobe，在下一篇中我們將介紹 MPI-3 中大的計數及相關函數。