python - 多处理池和生成器

Question

先看下面的代码:

pool = multiprocessing.Pool(processes=N)
batch = []
for item in generator():
    batch.append(item)
    if len(batch) == 10:
        pool.apply_async(my_fun, args=(batch,))
        batch = []
# leftovers
pool.apply_async(my_fun, args=(batch,))

本质上，我是从生成器中检索数据，将其收集到一个列表中，然后生成一个使用这批数据的进程。

这可能看起来不错，但是当消费者(也称为池进程)比生产者(也称为生成器)慢时，主进程的内存使用量会增加，直到生成器停止或...系统内存不足。

如何避免这个问题？

Answer 1

在这种情况下，您可能希望使用有限大小的队列。

q = multiprocessing.Queue(maxSize).

当与最大值一起使用时。大小，这将为您提供必要的计数并在 q.put() 已满时阻塞调用 q.put() 的线程，因此您永远不会在其上发布超过一定数量的工作项，从而限制存储所需的内存待处理的项目。

或者，您可以使用计数信号量(例如，multiprocessing.BoundedSemaphore(maxSize))。每次从生成器获取工作项时获取它，并在处理该项目后在工作函数 (my_fun) 中释放它。这样，等待处理的最大工作项数永远不会超过信号量的初始值。