Python中并行计算与存储解耦的实现方案咨询

问题描述 #

给定大量计算任务，每次取一个任务块在CPU/GPU上并行计算，因结果过大无法全部存入内存，需在块计算完成后统一收集结果写入文件。存储机制复杂，必须统一收集后再存储，且仅允许单个进程执行存储操作。当前代码实现了并行计算，但计算与存储串行执行，存储耗时较长时效率低下，需实现计算→收集→存储的并行流程，即在存储时启动下一轮计算。

原示例代码：

import numpy as np
from multiprocessing import Pool
import time

def crunch(n):
    print(f"crunch dummy things for input: {n}")
    results = np.random.random(100)
    time.sleep(np.random.randint(0, 3))
    return results

def store(results_npz, index):
    print(f"storing iteration {index}")
    np.savetxt(f'test_{str(index).zfill(2)}.out', results_npz)

# all tasks
all_tasks = list(range(10))

# iterate over tasks in chunks
for i in range(5):
    print(f"start iteration {i}")
    input_chunk = [all_tasks.pop(0), all_tasks.pop(0)]
    with Pool(2) as mp:
        results = mp.map(crunch, input_chunk)

    print("storing results ...")
    # ideally, this should start and then the result computation can start again
    results_all = np.vstack(results)
    store(results, i)

解决方案思路 #

核心是通过进程间队列实现计算与存储的解耦，用单独的存储进程消费队列中的结果，计算进程池持续生产结果，两者并行执行：

• 用multiprocessing.Queue传递计算完成的块结果及索引，作为计算与存储的通信桥梁。
• 启动一个单独的守护进程负责存储，确保唯一的存储执行单元。
• 计算进程池批量处理任务块，完成后立即将结果放入队列，无需等待存储完成即可启动下一轮计算。
• 所有计算任务完成后，向队列发送终止信号，让存储进程自动退出。

优化后代码 #

import numpy as np
from multiprocessing import Pool, Queue, Process
import time
import signal

def crunch(n):
    print(f"crunch dummy things for input: {n}")
    results = np.random.random(100)
    time.sleep(np.random.randint(0, 3))
    return results

def store_worker(queue):
    # 忽略终止信号，确保存储完成再退出
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    while True:
        item = queue.get()
        # 收到终止信号则退出
        if item is None:
            break
        results, index = item
        print(f"storing iteration {index}")
        results_all = np.vstack(results)
        np.savetxt(f'test_{str(index).zfill(2)}.out', results_all)
        print(f"finished storing iteration {index}")

if __name__ == "__main__":
    # 初始化队列
    result_queue = Queue()
    # 启动存储进程
    store_process = Process(target=store_worker, args=(result_queue,))
    store_process.daemon = False
    store_process.start()

    all_tasks = list(range(10))
    # 保持计算进程池长期存在，避免重复创建开销
    with Pool(2) as mp:
        for i in range(5):
            print(f"start iteration {i}")
            input_chunk = [all_tasks.pop(0), all_tasks.pop(0)]
            # 同步计算并获取结果，放入队列
            results = mp.map(crunch, input_chunk)
            # 将结果和索引放入队列，存储进程会自动处理
            result_queue.put((results, i))
            print(f"iteration {i} computation done, result sent to queue")

    # 所有计算完成，发送终止信号
    result_queue.put(None)
    # 等待存储进程完成所有剩余存储任务
    store_process.join()
    print("All tasks completed")

关键优化点说明 #

• 队列通信：通过Queue实现计算结果的异步传递，计算进程无需等待存储完成即可继续下一轮任务。
• 单一存储进程：单独的store_worker进程确保只有一个单元执行存储操作，符合需求。
• 进程池复用：将Pool的创建移到循环外，避免每次迭代重新创建进程池的开销。
• 优雅退出：计算完成后发送None作为终止信号，等待存储进程处理完所有队列中的结果再退出，避免数据丢失。

内容的提问来源于stack exchange，提问作者Kam Sen