如何限制imap_unordered中的工人吞吐量？

Question

我正在使用multiprocessing库中的imap_unordered来并行化一些数据处理计算。问题在于，有时读取返回的迭代器的主进程处理计算结果的速度比工作者生成的慢（网络/磁盘速度限制等）。这导致程序消耗所有可用内存并崩溃。

我希望内部迭代器有一些内部大小限制，这样当返回的迭代器处理得太慢时，内部队列就会变满并阻塞生成器（异步工作器）。但显然事实并非如此。

如何实现这种行为最简单的方法是什么？

Answer 1

您可能需要考虑使用Queue：

import multiprocessing  # Don't use queue.Queue!

MAX_QUEUE_SIZE = 20

q = multiprocessing.Queue(MAX_QUEUE_SIZE)  # Inserts will block if the queue is full

然后，在您的主过程中：

while 1:
    do_something_with(q.get())

在你孩子的过程中：

while 1:
    q.put(create_something())

你必须重写一些机器（即你不能再使用imap_unordered），但使用Pool的低级方法应该是相当简单的。