如何實現多處理 python 包裝器？

Question

我想編寫一個包裝器，用於在 asyncio 中調用 CPU 要求高的函數。

我希望它像這樣使用：

@cpu_bound
def fact(x: int):
    res: int = 1
    while x != 1:
        res *= x
        x -= 1
    return res

async def foo(x: int):
    res = await fact(x)
    ...

起初，我寫道：

def cpu_bound(func: Callable[P, R]) -> Callable[P, Awaitable[R]]:
    @functools.wraps(func)
    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        executor =  get_executor() # This is a part where I implemented myself.
        return await loop.run_in_executor(
            executor, functools.partial(func, *args, **kwargs)
        )

    return wrapper

但是，我在腌制方面遇到了問題。

Traceback（最近一次調用最后一次）：文件“C:\Users\Lenovo\AppData\Local\Programs\Python\Python39\lib\multiprocessing\queues.py”，第 245 行，在 _feed obj = _ForkingPickler.dumps(obj) 文件中“C:\Users\Lenovo\AppData\Local\Programs\Python\Python39\lib\multiprocessing\reduction.py”，第 51 行，轉儲 cls(buf, protocol).dump(obj) _pickle.PicklingError: Can't pickle <function fact at 0x000001C2D7D40820>: 它與main .fact 不同的 object

也許原來的 function 和包裝的沒有相同的id是問題所在？

那么，有沒有辦法編寫這樣的包裝器？

我確實知道我可以使用loop.run_in_executor ，但是擁有這樣的包裝器會很有幫助。

非常感謝。

Answer 1

也許原來的 function 和包裝的沒有相同的 id 是問題所在？

在某種程度上，是的。 Before the function is sent to the target process it's pickled, which fails in your case, because the func object in the decorator's scope is different from the fact object in your main module, after rebinding by the decorator. 看看這個和這個問題的一些背景。

基於這些答案，我創建了一個關於如何實現您想要的示例。 訣竅是創建一個可腌制的“跑步者”function，目標進程可以使用它從某種注冊表（例如字典...）中查找您的原始 function 並運行它。 這當然只是一個例子。 您可能不想在裝飾器中創建ProcessPoolExecutor 。

import asyncio
from concurrent.futures import ProcessPoolExecutor
import functools

original_functions={}

def func_runner(name, *args):
    return original_functions[name](*args)

def cpu_bound(func):
    original_functions[func.__name__]=func
    @functools.wraps(func)
    async def wrapper(*args):
        with ProcessPoolExecutor(1) as pool:
            res =  await asyncio.get_running_loop().run_in_executor(
            pool, functools.partial(func_runner, func.__name__, *args)
        )
        return res

    return wrapper

@cpu_bound
def fact(arg):
    return arg

async def foo():
    res = await fact("ok")
    print(res)

if __name__ == "__main__":
    asyncio.run(foo())