如何实现对 python multiprocessing.Pool 的自定义控制？

Question

通常我使用以下代码，当您不关心函数process_func将处理某些参数时，它可以正常工作：

params = [1,2,3,4,5 ... ]

def process_func():
    ...

pool = new Pool(40)
pool.map(process_func, params)
pool.close()
pool.join()

在上面的示例中，我们有一种类型的进程，最大同时数为 40。但是.. 假设我们有不同类型的进程（参数），应该同时执行。例如，在我的 selenium 网格中，我有 40 个 firefox，40 个 chrome。而且我有 5000 个测试用例，其中一些更喜欢 chrome，其中一些更喜欢 firefox，其中一些无关紧要。

例如，假设我们有以下类型：

• type firefox：最大同时数：40
• type chrome：最大同时数：40

在这种情况下，我们的池将最多同时有 80 个进程，但有严格的规定：其中 40 个必须是 firefox，其中 40 个必须是 chrome。

这意味着参数不会一个接一个地被取走。池必须从参数列表中选择值，以使每个进程类型具有最大值。

如何做到这一点？

Answer 1

我会修改您的process_func() 以再采用一个参数来告诉它要使用哪个“类型”并使用两个单独的池。添加functools.partial 将允许我们仍然使用pool.map()：

from functools import partial
from multiprocessing import Pool

params = [1,2,3,4,5 ... ]

def process_func(type, param):
    if type == 'Firefox':
        # do Firefox stuff
    else:
        # do Chrome stuff

chrome_pool = Pool(40)
fox_pool = Pool(40)

chrome_function = partial(process_func, 'Chrome')
fox_function = partial(process_func, 'Firefox')

chrome_pool.map(chrome_func, params)
fox_pool.map(fox_func, params)

chrome_pool.close()
fox_pool.close()
chrome_pool.join()
fox_pool.join()

functools.partial() 函数允许我们将参数绑定到特定值，方法是返回一个始终提供该参数的新函数对象。这种方法允许您将每个“类型”（因为没有更好的术语）限制为 40 个工作进程。