如何在并行化中使用非顶级函数？

Question

我想在我编写的代码中在资源繁重的计算中使用多处理，如以下淡化示例所示：

import numpy as np
import multiprocessing as multiproc

def function(r, phi, z, params):
    """returns an array of the timepoints and the corresponding values 
       (demanding computation in actual code, with iFFT and stuff)"""
    times = np.array([1.,2.,3.])
    tdependent_vals = r + z * times + phi
    return np.array([times, tdependent_vals])

def calculate_func(rmax, zmax, phi, param):
    rvals = np.linspace(0,rmax,5)
    zvals = np.linspace(0,zmax,5)
    for r in rvals:
        func_at_r = lambda z: function(r, phi, z, param)[1]
        with multiproc.Pool(2) as pool:
             fieldvals = np.array([*pool.map(func_at_r, zvals)])
             print(fieldvals) #for test, it's actually saved in a numpy array

calculate_func(3.,4.,5.,6.)

如果我运行它，它会失败并显示

AttributeError: Can't pickle local object 'calculate_func.<locals>.<lambda>'

我认为原因是，根据documentation，只能腌制顶级定义的函数，而我的函数内定义lambda 不能。但我看不出有什么办法可以让它成为一个独立的函数，至少不会用一堆顶级变量污染模块：在调用 calculate_func 之前参数是未知的，并且它们在每次迭代时都会发生变化超过rvals。这整个多处理的事情对我来说是非常新的，我想不出一个替代方案。在rvals 和zvals 上并行化循环的最简单工作方法是什么？

注意：我以answer 为起点。

Answer 1

这可能不是最好的答案，但它是一个答案，所以请不要讨厌:)

您可以编写一个可以序列化的顶级包装函数并让它执行函数...这有点像函数初始，但我在我的代码中解决了类似的问题。

这是一个简单的例子

def wrapper(arg_list, *args):
    func_str = arg_list[0]
    args = arg_list[1]
    code = marshal.loads(base64.b64decode(func_str.data))
    func = types.FunctionType(code, globals(), "wrapped_func")
    return func(*args)

def run_func(func, *args):
    func_str = base64.b64encode(marshal.dumps(func.__code__, 0))
    arg_list = [func_str, args]
    with mp.Pool(2) as pool:
        results = pool.map(wrapper, arg_list)
    return results