Python递归场景下的并行计算优化方案问询

问题背景 #

我想要并行化Python程序中频繁调用compute_heavy()的func()函数，且并行化不能作为对外不可见的实现细节。最初用multiprocessing.Pool实现了fork-join式并行：

import time
from multiprocessing import Pool

def compute_heavy(x):
    time.sleep(1)
    return x*x

def func():
    with Pool(5) as p:
        result = p.map(compute_heavy, [1, 2, 3])
    return sum(result)

但实际函数逻辑包含递归，示例如下：

def func(x):
    if x == 0:
        return 0
    my_result = sum([compute_heavy(i) for i in range(x)])
    other_result = func(x - 1) # 递归调用
    
    return my_result + other_result

当前核心问题：调用func(10)时，my_result部分有10个可并行的compute_heavy任务，但Pool.map会阻塞到所有任务完成才执行递归的func(x-1)。哪怕有无限多处理器，func(10)也要10秒，无法达到理想的1秒执行时间。尝试将递归调用纳入pool.map参数，但可能引发死锁。

需求：在不耗尽资源的前提下实现最优运行时间，同时想了解async/await框架是否适用于这种计算密集型场景？

优化方案：进程池+异步提交+非阻塞递归 #

核心思路 #

放弃map这类阻塞式方法，改用apply_async异步提交compute_heavy任务，让递归调用不等待当前批次任务完成就提前执行，通过跨进程容器收集结果，避免串行等待。同时复用全局进程池，避免重复创建进程导致资源浪费。

代码实现 #

import time
from multiprocessing import Pool, Manager

# 全局复用进程池，避免重复创建进程
global_pool = Pool(5)
# 用Manager字典跨进程安全收集任务结果
result_store = Manager().dict()

def compute_heavy(x, task_key):
    time.sleep(1)
    res = x * x
    result_store[task_key] = res
    return res

def func(x):
    if x == 0:
        return 0
    
    # 异步提交当前x对应的所有compute_heavy任务
    pending_tasks = []
    for i in range(x):
        unique_key = f"level_{x}_task_{i}"
        task = global_pool.apply_async(compute_heavy, args=(i, unique_key))
        pending_tasks.append((task, unique_key))
    
    # 不等待当前层任务完成，直接递归执行下一层
    other_result = func(x - 1)
    
    # 等待当前层所有任务完成，汇总结果
    for task, key in pending_tasks:
        task.wait()
    my_result = sum(result_store[key] for _, key in pending_tasks)
    
    return my_result + other_result

if __name__ == "__main__":
    start_time = time.time()
    print(f"最终结果: {func(10)}")
    print(f"总耗时: {time.time() - start_time:.2f}秒")
    global_pool.close()
    global_pool.join()

效果说明 #

所有compute_heavy任务都会被异步提交到进程池，递归调用不会被当前层任务阻塞，所有任务可并行执行。在无限多处理器的情况下，总耗时接近1秒；若进程池大小为5，耗时约2秒（10个任务分两批执行），完全符合最优运行时间的要求。

关于async/await在计算密集型场景的适用性 #

async/await本质是单线程异步调度，依赖事件循环在I/O等待时让出CPU，仅适合I/O密集型任务。对于compute_heavy这类计算密集型任务，单纯的async/await无法利用多CPU核心——Python的GIL会阻止同一进程内的线程并行执行CPU密集型代码，任务会串行执行，反而比进程池方案更慢。

如果要结合async/await，需搭配concurrent.futures.ProcessPoolExecutor实现异步+多进程的组合，但逻辑和直接用multiprocessing池差异不大，没有明显优势，因此不推荐单独用async/await处理这类计算密集型递归任务。

注意事项 #

• 全局进程池需在if __name__ == "__main__"代码块外创建（Windows系统需避免进程初始化时重复创建池），或在主函数中创建后传入func。
• 任务结果的唯一键要确保不重复，避免不同递归层级的任务结果冲突。
• 程序结束时必须调用close()和join()关闭进程池，避免资源泄漏。

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