Pythonのmultiprocessing.PoolでAttributeErrorが出た

Pythonのmultiprocessingモジュールを試そうとしたらエラーが発生してしまったという話。問題のコードは下の通りです。

#!/usr/bin/env python3

import multiprocessing

def foo(num):
    def bar(x):
        return x*x
    cpu = multiprocessing.cpu_count()
    args = range(num)
    with multiprocessing.Pool(cpu) as pool:
        result = pool.map(bar, args)
    return result

if __name__ == '__main__':
    print(foo(5))

これを実行したところ以下のエラーメッセージがでました。

AttributeError: Can't pickle local object 'foo.<locals>.bar'

こちらの方の記事やStack OverflowのこちらのQ&Aによると、関数内関数が問題のようでした。そこで、bar関数を外に出し、トップレベルで定義し直したところ、上記のエラーは発生しなくなりました。

オマケ Joblibを使う方法

一応考えた上で関数内関数を定義していたので、それを外に出してしまうのはちょっと悔しい。何か手はないかと思い、試しにJoblibを利用したところ、問題なく実行できました。

#!/usr/bin/env python3

from joblib import Parallel
from joblib import delayed

def foo(num):
    def bar(x):
        return x*x
    args = range(num)
    result = Parallel(n_jobs=-1)([delayed(bar)(arg) 
　　　　　　　　　　　　　　for arg in args])
    return result

if __name__ == '__main__':
    print(foo(5))

Joblibの方がお手軽かもしれませんね。最後までお読みいただきありがとうございました。