python中使用多进程执行计算密集型任务（完整代码示例）

问题

你有大量的数据，需要对这些大量的数据进行相关计算，并且计算的结果要和给定的数据顺序对应一致。

解决方案

创建两个被多个子进程共享的 Queue 对象，一个用于主进程将数据分块放入该 Queue 对象中，然后子进程从该 Queue 对象中取出数据进行计算；一个用于子进程计算的结果放入 Queue 对象中，然后主进程从该 Queue 对象中取出计算结果。

假设密集型计算任务（函数）如下：

def cal(nums: List[int]):
    ans = []
    for num in nums:
        res = 0
        for i in range(num):
            res += i + i ** i + i * i
        ans.append(res)
    return ans

假设（模拟）需要有如下大量的数据需要进行计算

nums = [10] * 10000000     # 7个0

创建用于进程间通信的 Queue 对象，然后将 Queue 对象作为子进程执行函数的参数，子进程从 Queue 对象中获取需要计算的数据；

import multiprocessing as mp

def start_multi_precess_pool(pool_nums: int):
    """
    pool_nums: 创建的进程数量
    """
    # 使用 spawn 启动进程，具体请参考：https://docs.python.org/3/library/multiprocessing.html?highlight=multiprocessing#contexts-and-start-methods
    ctx = mp.get_context("spawn")

    # 创建用于输入和输出的 Queue 对象
    input_queue = ctx.Queue()
    output_queue = ctx.Queue()  
    process = []
    
    # 创建 pool_nums 数量的进程
    for i in range(pool_nums):
        # 创建子进程，将 Queue 对象作为进程执行函数的参数，用于进程间的通信；
        # 并且给进程传递序号 i，用以保证后续计算结果的有序返回；
        # cal_multi_process_worker 函数用从 Queue 获取数据，并执行密集计算任务；
        p = ctx.Process(target=cal_multi_process_worker, args=(i, input_queue, output_queue))
        p.start()
        process.append(p)

    # 将 Queue 对象和进程对象集合返回，将 主进程将数据传递给 Queue 对象的行为 封装为另一个函数
    return {"input": input_queue, "output": output_queue, "process": process}

子进程从 Queue 对象中拿去数据并执行计算；

def cal_multi_process_worker(id: int, input_queue: Queue, output_queue: Queue):
    while True:
        try:
            id, nums = input_queue.get()
            ans = cal(nums)
            output_queue.put([id, ans])
        except queue.Empty:
            break

将数据分块，并写入 Queue 对象；然后从 Queue 对象取出计算后的结果；

def cal_multi_process(nums: List[int], pool: Dict[str, object], chunk_size: int = None):
    """
    对数据进行分块，将分块后的数据传入Queue对象
    """
    if chunk_size is None:
        chunk_size = min(math.ceil(len(nums) / len(pool["process"]) / 10), 500)

    input_queue = pool["input"]
    last_chunk_id = 0
    chunk = []

    for num in nums:
        chunk.append(num)
        if len(chunk) >= chunk_size:
            input_queue.put([last_chunk_id, chunk])
            last_chunk_id += 1
            chunk = []

    if len(chunk) > 0:
        input_queue.put([last_chunk_id, chunk])
        last_chunk_id += 1

    # 从 Queue 对象中的数据获取计算后的结果
    output_queue = pool["output"]
    
    # 根据数据分块号对返回的结果进行排序，以保证得到的计算结果和输入的数据顺序保持一致；
    result_list = sorted([output_queue.get() for _ in range(last_chunk_id)], key=lambda x : x[0])
    
    return result_list

关闭子进程，释放相关资源；

def stop_multi_process_pool(pool):
    # 终止子进程
    for p in pool["process"]:
        p.terminate()

    # 释放子进程的相关资源
    for p in pool["process"]:
        p.join()
        p.close()

    pool["input"].close()
    pool["output"].close()

完整代码

import multiprocessing as mp
from multiprocessing import Queue
import queue
import time
from typing import List, Dict
import math

def cal(nums: List[int]):
    """
    计算密集型函数
    """
    ans = []
    for num in nums:
        res = 0
        for i in range(num):
            res += i + i ** i + i * i
        ans.append(res)
    return ans

def cal_multi_process_worker(input_queue: Queue, output_queue: Queue):
    """
    子进程从Queue对象中获取数据，并执行计算
    """
    while True:
        try:
            id, nums = input_queue.get()
            ans = cal(nums)
            output_queue.put([id, ans])
        except queue.Empty:
            break
            
            
def cal_multi_process(nums: List[int], pool: Dict[str, object], chunk_size: int = None):
    """
    对数据进行分块，将分块后的数据传入Queue对象
    """
    if chunk_size is None:
        chunk_size = min(math.ceil(len(nums) / len(pool["process"]) / 10), 500)

    input_queue = pool["input"]
    last_chunk_id = 0
    chunk = []

    for num in nums:
        chunk.append(num)
        if len(chunk) >= chunk_size:
            input_queue.put([last_chunk_id, chunk])
            last_chunk_id += 1
            chunk = []

    if len(chunk) > 0:
        input_queue.put([last_chunk_id, chunk])
        last_chunk_id += 1

    # 从 Queue 对象中的数据获取计算后的结果
    output_queue = pool["output"]
    
    # 根据数据分块号对返回的结果进行排序，以保证得到的计算结果和输入的数据顺序保持一致；
    result_list = sorted([output_queue.get() for _ in range(last_chunk_id)], key=lambda x : x[0])
    
    return result_list

def start_multi_precess_pool(pool_nums: int):
    """
    pool_nums: 创建的进程数量
    """
    # 使用 spawn 启动进程，具体请参考：https://docs.python.org/3/library/multiprocessing.html?highlight=multiprocessing#contexts-and-start-methods
    ctx = mp.get_context("spawn")

    # 创建用于输入和输出的 Queue 对象
    input_queue = ctx.Queue()
    output_queue = ctx.Queue()  
    process = []
    
    # 创建 pool_nums 数量的进程
    for i in range(pool_nums):
        # 创建子进程，将 Queue 对象作为进程执行函数的参数，用于进程间的通信；
        # 并且给进程传递序号 i，用以保证后续计算结果的有序返回；
        # cal_multi_process_worker 函数用从 Queue 获取数据，并执行密集计算任务；
        p = ctx.Process(target=cal_multi_process_worker, args=(input_queue, output_queue))
        p.start()
        process.append(p)

    # 将 Queue 对象和进程对象集合返回，将主进程将数据传递给 Queue 对象的行为 封装为另一个函数
    return {"input": input_queue, "output": output_queue, "process": process}

def stop_multi_process_pool(pool):
    # 终止子进程
    for p in pool["process"]:
        p.terminate()

    # 释放子进程的相关资源
    for p in pool["process"]:
        p.join()
        p.close()

    pool["input"].close()
    pool["output"].close()

if __name__ == "__main__":
    start_time = time.time()

    nums = [10] * 10000000   # 7个0

    # pool = start_multi_precess_pool(6)
    # ans = cal_multi_process(nums, pool)
    # stop_multi_process_pool(pool)
		# 25.43005394935608

    ans = cal(nums)    # 6.784837245941162
    total_time = time.time() - start_time
    print(f"total_time = {total_time}")