Python3 multiprocessing joinable queue 模板

最近需要在服務器上處理一批文件，每個文件的處理過程很簡單，基本就是讀入文件，計算一些統計值，然後把統計值彙總。一想這可以多線程啊老鐵！調試了一下Python3的multiprocessing，這裏留下一個模板以備之後使用。

程序運行的邏輯是這樣的

• 主進程掃描需要處理的文件，生成文件列表。
• 主進程創建job隊列和result隊列。此時隊列都爲空。
• 主進程創建所有子進程。子進程啓動。監聽來自job隊列的信息（blocked get()）。
• 主進程將文件列表的內容逐一put到job隊列內（JoinableQueue）。子進程get來自job隊列的信息，處理文件，將處理結果put到result隊列。
• 主進程開始從result隊列get數據並臨時存儲。
• 主進程將所有job發送完畢。主進程將所有result隊列的數據get完畢。主進程join job隊列。
• 主進程發送終止標誌給所有子進程。
• 子進程終止。
• 主進程join所有子進程。
• 主進程開始處理所有result（根據文件次序排序，等等）。
• 主進程退出。

調試過程還是比較順利的，之前也簡單使用過python自帶的多線程工具。期間遇到一個問題，就是主進程先發送了終止標誌給子進程，然後纔開始從result隊列獲取數據，導致主進程在從result隊裏get數據時形成無限block。其原理基本是這樣的：

• 子進程通過result隊列put信息，當result隊列已經有過多尚未get的數據時，子進程put的信息被一個pipe緩衝起來，等待result隊列有更多空間時再轉入隊列。
• 若主進程沒有及時從result隊列get數據，導致result隊列有尚未進入的緩衝數據，並且主進程發送終止標誌給子進程，子進程在未完成result隊列的put的情況下退出，導致數據丟失，數據沒有及時進入result隊列。

Lesson learned: 利用get() 處理隊列的進程要保持開啓直到所有需要put()的數據都已put並且get到隊列爲空，此時才能安全地終止調用put()的進程。

此外，Python3對Queue package的命名與Python2不同，處理異常時需要注意。

以下是模板源碼。注意必須爲Python3執行。

# Author: Yaoyu Hu <[email protected]>

import argparse
import multiprocessing
from queue import Empty
import time

def cprint(msg, flagSilent=False):
    if ( not flagSilent ):
        print(msg)

def process_single_file(name, jobStr, flagSilent=False):
    """
    name is the name of the process.
    """

    startTime = time.time()

    cprint("%s. " % (jobStr))

    endTime = time.time()
    
    s = "%s: %ds for processing." % (name, endTime - startTime )

    cprint(s, flagSilent)
    cprint("%s: " % (name), flagSilent)

    return s

def worker(name, q, p, rq, flagSilent=False):
    """
    name: String, the name of this worker process.
    q: A JoinableQueue.
    p: A pipe connection object. Only for receiving.
    """

    cprint("%s: Worker starts." % (name), flagSilent)

    while (True):
        if (p.poll()):
            command = p.recv()

            cprint("%s: %s command received." % (name, command), flagSilent)

            if ("exit" == command):
                break

        try:
            jobStrList = q.get(True, 1)
            # print("{}: {}.".format(name, jobStrList))

            s = process_single_file(name, jobStrList[0], flagSilent)

            rq.put([s], block=True)

            q.task_done()
        except Empty as exp:
            pass
    
    cprint("%s: Work done." % (name), flagSilent)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Filter the files.")

    parser.add_argument("--jobs", type=int, default=100, \
        help="The number of jobs for testing.")

    parser.add_argument("--np", type=int, default=2, \
        help="The number of processes.")

    args = parser.parse_args()

    assert args.jobs > 0
    assert args.np > 0

    startTime = time.time()

    print("Main: Main process.")

    jobQ    = multiprocessing.JoinableQueue()
    resultQ = multiprocessing.Queue()

    processes = []
    pipes     = []

    print("Main: Create %d processes." % (args.np))

    for i in range(int(args.np)):
        [conn1, conn2] = multiprocessing.Pipe(False)
        processes.append( multiprocessing.Process( \
            target=worker, args=["P%03d" % (i), jobQ, conn1, resultQ, False]) )
        pipes.append(conn2)

    for p in processes:
        p.start()

    print("Main: All processes started.")

    for dj in range(args.jobs):
        jobQ.put([ str(dj) ])

    print("Main: All jobs submitted.")

    resultList = []
    resultCount = 0

    while(resultCount < args.jobs):
        try:
            print("Main: Get index %d. " % (resultCount))
            r = resultQ.get(block=True, timeout=1)
            resultList.append(r)
            resultCount += 1
        except Empty as exp:
            if ( resultCount == args.jobs ):
                print("Main: Last element of the result queue is reached.")
                break

    jobQ.join()

    print("Main: Queue joined.")

    for p in pipes:
        p.send("exit")

    print("Main: Exit command sent to all processes.")

    for p in processes:
        p.join()

    print("Main: All processes joined.")

    print("Main: Starts process the result.")

    print(resultList)

    endTime = time.time()

    print("Main: Job done. Total time is %ds." % (endTime - startTime))