python - 将 multiprocessing 与修饰函数一起使用会导致 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

尝试 dano's trick似乎只适用于 Python 2。在 Python 3 中尝试时，出现以下错误:

pickle.PicklingError: Can't pickle <function serialize at 0x7f7a1ac1fd08>: it's not the same object as __main__.orig_fn

我通过从 worker 的 init 中“装饰”函数解决了这个问题:

from functools import wraps
import sys

def worker_init(fn, *args):
    @wraps(fn)
    def wrapper(x):
        # wrapper logic
        pass

    setattr(sys.modules[fn.__module__], fn.__name__, wrapper)

pool = mp.Pool(initializer=worker_init, initargs=[orig_fn, *args])
# ...