Python多处理：如何限制等待进程的数量？

当使用Pool.apply\u async运行大量任务（具有大参数）时，进程将被分配并进入等待状态，等待进程的数量没有限制。这可能会消耗所有内存，如下例所示：

import multiprocessing
import numpy as np

def f(a,b):
    return np.linalg.solve(a,b)

def test():

    p = multiprocessing.Pool()
    for _ in range(1000):
        p.apply_async(f, (np.random.rand(1000,1000),np.random.rand(1000)))
    p.close()
    p.join()

if __name__ == '__main__':
    test()

---

我正在寻找一种限制等待队列的方法，即只有有限数量的等待进程，而Pool.apply\u async在等待队列已满时被阻止。

多处理。Pool

有一个

\u taskqueue

类型的成员

多处理.queue

，它采用可选的

maxsize

参数；不幸的是，它构造它时没有

maxsize

参数集

我建议使用

multiprocessing.Pool.\uuuu init\uuuu

的副本粘贴子类化

multiprocessing.Pool

，它将

maxsize

传递给

\u taskqueue

构造函数

猴子修补对象（池或队列）也可以，但您必须修补

pool.\u taskqueue.\u maxsize

和

pool.\u taskqueue.\u sem

因此它非常脆弱：

pool._taskqueue._maxsize = maxsize
pool._taskqueue._sem = BoundedSemaphore(maxsize)

---

您可以添加带有maxsize参数的显式队列，并使用

Queue.put（）

而不是

pool.apply\u async（）

。然后工作进程可以：

for a, b in iter(queue.get, sentinel):
    # process it

如果要将内存中创建的输入参数/结果的数量限制为活动工作进程的数量，则可以使用

pool.imap*（）

方法：

#!/usr/bin/env python
import multiprocessing
import numpy as np

def f(a_b):
    return np.linalg.solve(*a_b)

def main():
    args = ((np.random.rand(1000,1000), np.random.rand(1000))
            for _ in range(1000))
    p = multiprocessing.Pool()
    for result in p.imap_unordered(f, args, chunksize=1):
        pass
    p.close()
    p.join()

if __name__ == '__main__':
    main()

---

如果

池。\u taskqueue

超过所需大小，请稍候：

import multiprocessing
import time

import numpy as np


def f(a,b):
    return np.linalg.solve(a,b)

def test(max_apply_size=100):
    p = multiprocessing.Pool()
    for _ in range(1000):
        p.apply_async(f, (np.random.rand(1000,1000),np.random.rand(1000)))

        while p._taskqueue.qsize() > max_apply_size:
            time.sleep(1)

    p.close()
    p.join()

if __name__ == '__main__':
    test()

---

下面是一个猴子补丁替代顶部答案：

import queue
from multiprocessing.pool import ThreadPool as Pool


class PatchedQueue():
  """
  Wrap stdlib queue and return a Queue(maxsize=...)
  when queue.SimpleQueue is accessed
  """

  def __init__(self, simple_queue_max_size=5000):
    self.simple_max = simple_queue_max_size  

  def __getattr__(self, attr):
    if attr == "SimpleQueue":
      return lambda: queue.Queue(maxsize=self.simple_max)
    return getattr(queue, attr)


class BoundedPool(Pool):
  # Override queue in this scope to use the patcher above
  queue = PatchedQueue()

pool = BoundedPool()
pool.apply_async(print, ("something",))

---

这与Python3.8的预期一样，在Python3.8中，多处理池使用

queue.SimpleQueue

来设置任务队列。这听起来像是

多处理.Pool

的实现自2.7以来可能发生了变化，因为我使用的是Python 2.7.3，而_taskqueue的类型是Queue.Queue。这意味着它是一个简单的队列，而不是一个multiprocessing.Queue。子类化multiprocessing.Pool并重写init可以很好地工作，但是猴子修补对象并没有像预期的那样工作。不过，这就是我搜索的黑客，谢谢。使用

imap

没有任何区别。输入队列仍然是无限的，使用此解决方案将消耗所有内存。@Radim:answer中的

imap

代码可以工作，即使您给它一个无限的生成器。不幸的是，在Python 2中不是这样（没有查看py3中的代码）。有关一些解决方法，请参阅。我只想补充一点，我发现这是解决多处理内存问题的最简单方法。我使用了max_apply_size=10，这很好地解决了我的问题，这是一个缓慢的文件转换。按照@ecatmur的建议使用信号量似乎是一个更强大的解决方案，但对于简单的脚本来说可能会有些过火。TaylorMonacelli您的编辑被拒绝就是mods的一个例子。您的编辑修复了一个bug@greg-449是一个“驱动模式”，只批准了15%的编辑，并给出了一个不合理的拒绝理由。我没有测试ThreadPool one，但是如果我从multiprocessing.pool import pool将其修改为

，它将不起作用（限制没有更改，而且似乎
SimpleQueue
没有更改为
队列
）。你知道怎么解决这个问题吗？