Python中ProcessPoolExecutor惰性处理：限制生成器结果缓存的方法？

嘿，这个问题我之前也踩过坑——当任务量特别大的时候，ProcessPoolExecutor会一股脑把所有任务都提交进去，哪怕max_workers设得小，剩下的任务也会堆在Executor的内部队列里，加上as_completed一有结果就返回，要是结果没及时处理，分分钟内存就爆了。

当然有办法解决！核心思路就是限制内存中待处理的结果数量，同时控制并发任务数，让生成器在缓存达到指定大小后暂停提交新任务，直到你取出一些结果再继续。下面给你两种实用的方案：

方案一：手动控制任务提交与结果缓存 #

这种方法不需要额外依赖，直接通过控制提交的任务数和结果缓存队列来实现限流。逻辑很直观：保持max_workers个任务在后台运行，同时内存里缓存的已完成结果不超过设定的buffer_size，缓存满了就暂停提交新任务。

from concurrent.futures import ProcessPoolExecutor, wait, FIRST_COMPLETED
import itertools

def task_func(param):
    # 替换成你的实际任务逻辑，这里模拟耗时计算
    return param * 2

def limited_task_generator(tasks, max_workers=4, buffer_size=8):
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = set()
        completed_results = []
        task_iter = iter(tasks)
        
        # 先提交第一批任务，填满worker池
        for _ in range(min(max_workers, len(tasks))):
            try:
                param = next(task_iter)
                futures.add(executor.submit(task_func, param))
            except StopIteration:
                break
        
        while futures:
            # 等待至少一个任务完成
            done, futures = wait(futures, return_when=FIRST_COMPLETED)
            
            # 收集完成的结果（含异常处理）
            for future in done:
                try:
                    completed_results.append(future.result())
                except Exception as e:
                    # 可自定义异常处理，比如记录日志或继续抛出
                    completed_results.append(e)
            
            # 缓存满时先返回结果，直到缓存低于阈值
            while len(completed_results) >= buffer_size:
                yield completed_results.pop(0)
            
            # 继续提交新任务，保持worker池满负荷
            while len(futures) < max_workers:
                try:
                    param = next(task_iter)
                    futures.add(executor.submit(task_func, param))
                except StopIteration:
                    break
        
        # 返回最后剩余的结果
        while completed_results:
            yield completed_results.pop(0)

使用示例 #

比如你有1000个任务参数，调用方式如下：

tasks = range(1000)
for result in limited_task_generator(tasks, max_workers=4, buffer_size=8):
    # 处理每个结果，比如写入文件、数据库等
    print(result)

这个生成器会自动在缓存满8个结果时暂停提交新任务，直到你处理完几个结果、缓存空间释放后再继续，完美控制内存占用。

方案二：用有界队列实现自动阻塞 #

这种方法借助queue.Queue的有界特性，让任务完成后的回调函数在队列满时自动阻塞，从而间接限制内存中的结果数量，适合喜欢更“自动化”逻辑的场景。

from concurrent.futures import ProcessPoolExecutor
from queue import Queue
import threading

def task_func(param):
    # 替换成你的实际任务逻辑
    return param * 2

def queue_based_generator(tasks, max_workers=4, buffer_size=8):
    result_queue = Queue(maxsize=buffer_size)
    stop_flag = object()  # 用唯一对象作为结束标记
    
    def task_callback(future):
        try:
            result_queue.put(future.result())
        except Exception as e:
            # 把异常也放进队列，方便后续统一处理
            result_queue.put(e)
    
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        # 启动单独线程提交所有任务，避免阻塞生成器
        def submit_all_tasks():
            for param in tasks:
                future = executor.submit(task_func, param)
                future.add_done_callback(task_callback)
            # 所有任务提交完成后，放入结束标记
            result_queue.put(stop_flag)
        
        submit_thread = threading.Thread(target=submit_all_tasks)
        submit_thread.start()
        
        # 从队列取结果，直到拿到结束标记
        while True:
            item = result_queue.get()
            if item is stop_flag:
                break
            if isinstance(item, Exception):
                # 可选择抛出异常或记录日志
                raise item
            yield item
        
        # 等待提交线程完成
        submit_thread.join()

为什么这个能生效？ #

Queue(maxsize=buffer_size)是有界队列，当队列里的结果数量达到buffer_size时，task_callback里的put()方法会自动阻塞，直到队列中有空位。这意味着后续完成的任务无法把结果写入队列，相当于暂停了结果的“流入”，直到你取出一些结果、队列有空位后才会继续写入，从根源上控制了内存中的结果数量。

注意事项 #

• buffer_size的选择：根据你的结果大小和内存容量调整——如果结果很大，就把buffer_size设小一点；如果结果很小，可以适当调大，平衡内存占用和执行效率。
• 异常处理：两种方案都加入了基础的异常处理逻辑，实际使用时可以根据需求调整，比如记录错误日志而不是直接抛出。
• 任务取消：如果需要中途停止任务，可以在方案二中加入stop_event来中断提交线程，或者在方案一中检查外部标志位。

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