被装饰函数无法多进程原因分析以及解决办法

被装饰器装饰的函数无法多进程原因

       python多进程的原理是通过pickle多进程函数名，然后新建一个子进程并在子进程中导入模块后unpickle，通过访问模块的该函数来实现函数在子进程中的运行的，关于pickle更详细的说明可看笔者的这篇文章。关于多进程编程中的pickle的一个重点是，对于多进程函数的pickle，只会pickle其函数名，也即f.__name__属性，然后通过模块的点号访问，因此该函数必须是可以通过模块点号访问到的，这也就是为什么要求被pickle必须被定义在模块的顶层。

       问题在于，当一个函数被装饰器修饰过后，根据装饰器的语法规则，实际上是对被装饰器函数复制了一个新的wrapper函数的引用，这时被装饰函数的名称属性就发生了改变，变成了wrapper的名称，更具体的关于装饰器语法的说明可看笔者这篇文章和这篇文章。并且由于wrapper函数是定义在装饰器函数中的，即非模块顶层函数，无法通过点号访问到，因此这时如果直接将被装饰函数用以多进程运行，会报错类似这样的错误"AttributeError: Can't pickle local object 'decorator.<locals>.wrapper"。一个典型的错误实例代码如下所示。

import multiprocessing
import time


def decorator(f):
   
    def wrapper(*args):
        t1=time.time()
        r=f(*args)
        t2=time.time()
        print(t2-t1)
        return r
    return wrapper


@decorator
def f(x):
    l=[]
    for i in range(x):
        l.append(i)
    r = sum(l)
    return r

def run():
    print(f.__name__)
    
    pool = multiprocessing.Pool(2)
    res = [pool.apply_async(f,args=(x,)) for x in [1000000,2000000]]
    pool.close()
    pool.join()
    for r in res:
        print(r.get())

if __name__=='__main__':
    run()
 

解决办法

       直接运行上面的代码是会报错的，解决办法很简单，只要依然保持被装饰函数的__name__属性不变即可，这一点可以利用python的@wraps语法来实现，@wraps语法可以保持被装饰函数的__name__属性不变，从而使得多进程可以顺利运行。具体的用法如下。

import multiprocessing
import time
from functools import wraps


def decorator(f):
    @wraps(f)
    def wrapper(*args):
        t1=time.time()
        r=f(*args)
        t2=time.time()
        print(t2-t1)
        return r
    return wrapper


@decorator
def f(x):
    l=[]
    for i in range(x):
        l.append(i)
    r = sum(l)
    return r

def run():
    print(f.__name__)
    
    pool = multiprocessing.Pool(2)
    res = [pool.apply_async(f,args=(x,)) for x in [1000000,2000000]]
    pool.close()
    pool.join()
    for r in res:
        print(r.get())

if __name__=='__main__':
    run()