被裝飾函數無法多進程原因分析以及解決辦法

被裝飾器裝飾的函數無法多進程原因

       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()