Python 多處理：如何限制等待進程的數量？

Question

使用Pool.apply_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_async 在等待隊列已滿時被阻塞。

Answer 1

multiprocessing.Pool具有_taskqueue類型的構件multiprocessing.Queue ，這需要一個可選maxsize參數; 不幸的是，它在沒有maxsize參數集的情況下構建它。

我建議你繼承multiprocessing.Pool用的復制粘貼multiprocessing.Pool.__init__是傳球maxsize到_taskqueue構造。

猴子修補對象（池或隊列）也可以，但您必須對pool._taskqueue._maxsize和pool._taskqueue._sem進行猴子pool._taskqueue._maxsize ，因此它會非常脆弱：

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

Answer 2

如果pool._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()

Answer 3

這是最佳答案的猴子修補替代方案：

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",))

這在 Python 3.8 中按預期工作，其中多處理池使用queue.SimpleQueue來設置任務隊列。 聽起來multiprocessing.Pool的實現可能自 2.7 以來發生了變化

Answer 4

在這種情況下，您可以使用 maxsize 參數添加顯式 Queue 並使用queue.put()而不是pool.apply_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()