如何在Python装饰器中使用多进程处理可迭代参数

问题 #

想要创建一个装饰器，使原函数在接受单个参数的同时，能够并行处理可迭代参数。示例代码如下：

import functools
import time

from multiprocessing import Pool

def parallel(func):

    def wrapper(iterable):
        with Pool() as pool:
            result = pool.map(func, iterable)

        return result

    return wrapper

@parallel
def test(i):
    time.sleep(1)
    print(f"{i}: {i * i}")


def main():
    test(range(10))


if __name__ == "__main__":
    main()

运行后出现如下报错：

Traceback (most recent call last):
  File "/home/user/projects/amdb/s2.py", line 29, in <module>
    main()
  File "/home/user/projects/amdb/s2.py", line 25, in main
    test(range(10))
  File "/home/user/projects/amdb/s2.py", line 10, in wrapper
    result = pool.map(func, iterable)
  File "/usr/lib/python3.10/multiprocessing/pool.py", line 367, in map
    return self._map_async(func, iterable, mapstar, chunksize).get()
  File "/usr/lib/python3.10/multiprocessing/pool.py", line 774, in get
    raise self._value
  File "/usr/lib/python3.10/multiprocessing/pool.py", line 540, in _handle_tasks
    put(task)
  File "/usr/lib/python3.10/multiprocessing/connection.py", line 206, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/usr/lib/python3.10/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
_pickle.PicklingError: Can't pickle <function test at 0x7fef63309120>: it's not the same object as __main__.test

已知通过test_multi = parallel(test)的方式可以解决，但希望保留装饰器语法实现。

解决方案 #

原因分析 #

报错核心是装饰器替换了原函数的引用：使用@parallel后，test变量指向的是装饰器返回的wrapper函数，原函数被隐藏。当multiprocessing.Pool尝试序列化原函数func时，pickle会检查该函数在主模块中的引用一致性，此时发现原函数标识与__main__.test不匹配，导致序列化失败。

修复代码 #

使用functools.wraps装饰wrapper函数，它会将原函数的元数据（名称、模块、文档字符串等）复制到wrapper上，同时让被装饰后的函数保持原函数的引用标识，解决pickle的问题：

import functools
import time

from multiprocessing import Pool

def parallel(func):
    @functools.wraps(func)  # 添加这一行
    def wrapper(iterable):
        with Pool() as pool:
            result = pool.map(func, iterable)
        return result
    return wrapper

@parallel
def test(i):
    time.sleep(1)
    print(f"{i}: {i * i}")


def main():
    test(range(10))


if __name__ == "__main__":
    main()

原理说明 #

functools.wraps通过update_wrapper函数修改wrapper的属性，让它在被pickle时能被识别为原函数的引用，子进程就能正确找到并加载原函数，避免序列化错误。

内容的提问来源于stack exchange，提问作者PaleNeutron