使用带有修饰函数的多处理会导致 PicklingError

Question

我正在尝试基于multiprocessing库编写一个方便的函数，它接受任何函数和参数，并使用多个进程运行该函数。我正在导入以下文件“MultiProcFunctions.py”：

import multiprocessing
from multiprocessing import Manager

def MultiProcDecorator(f,*args):

    """
    Takes a function f, and formats it so that results are saved to a shared dict
    """

    def g(procnum,return_dict,*args):
        result = f(*args)
        return_dict[procnum] = result

    return g

def MultiProcFunction(f,n_procs,*args):
    """
    Takes a function f, and runs it in n_procs with given args
    """

    manager     = Manager()
    return_dict = manager.dict()

    jobs = []
    for i in range(n_procs):
        p = multiprocessing.Process( target = f, args = (i,return_dict) + args )
        jobs.append(p)
        p.start()

    for proc in jobs:
        proc.join()

    return dict(return_dict)

这是我运行的代码：

from MultiProcFunctions import *

def sq(x):
    return [i**2 for i in x]

g = MultiProcDecorator(sq)

if __name__ == '__main__':

    result = MultiProcFunction(g,2,[1,2,3])

我收到以下错误： PicklingError: Can't pickle <function g at 0x01BD83B0>: it's not found as MultiProcFunctions.g

如果我改用以下定义g，一切都很好：

def g(procnum,return_dict,x):
    result = [i**2 for i in x]
    return_dict[procnum] = result

为什么这两个定义g不同，我可以做些什么来使原始g定义“起作用”？

Answer 1

发生这种情况是因为g实际上被定义为 中的嵌套函数MultiProcFunctions，这意味着它实际上不能从该模块的顶层导入，这意味着它无法正确腌制。现在，当我们这样做时，我们实际上g在模块的顶层非常清楚地定义了__main__：

g = MultiProcDecorator(sq)

所以，它确实应该是可腌制的。__module__我们可以通过显式地将of设置g为"__main__"：

g = MultiProcDecorator(sq)
g.__module__ = "__main__"  # Fix the __module__

这将允许酸洗过程工作，因为它将查找gin的定义__main__，它是在顶层定义的，而不是MultiProcFunctions，它仅在嵌套范围中定义。

编辑：

请注意，您还可以在装饰器本身中进行更改：

def MultiProcDecorator(f,*args):

    """
    Takes a function f, and formats it so that results are saved to a shared dict
    """

    def g(procnum,return_dict,*args):
        result = f(*args)
        return_dict[procnum] = result
    g.__module__ = "__main__"

    return g

这对您来说可能更有意义，因为这个装饰器严格来说是为了特定multiprocessing目的而使用的。