无外部依赖的局部函数多进程兼容方案（含不可序列化生成器场景）

你的问题核心在于标准库multiprocessing默认使用的pickle无法序列化局部作用域中的函数（因为pickle依赖全局命名空间的引用），同时还要避开fork的安全风险、处理生成器不可序列化的问题。下面是几个不依赖外部库的可行方案，按需选择：

方案一：基于源码重构的通用解法（推荐） #

这个方法通过提取局部函数的源码，在子进程中重新定义函数，彻底避开pickle序列化局部函数的限制，同时支持生成器和闭包变量，还能使用安全的spawn启动方式。

代码示例 #

import multiprocessing as mp
import inspect
import textwrap

def _execute_func(func_source, closure_vars, args, kwargs):
    # 在子进程中注入闭包变量并重构函数
    local_namespace = closure_vars.copy()
    exec(func_source, globals(), local_namespace)
    # 取出重构后的函数（这里假设源码里只有一个函数定义）
    func = next(v for v in local_namespace.values() if callable(v))
    return func(*args, **kwargs)

def main():
    # 带闭包变量和生成器的局部函数示例
    multiplier = 3
    def f(x):
        # 内部生成器，无需序列化，在子进程中执行时生成
        return sum(i * multiplier for i in range(x))
    
    # 提取函数源码并格式化（去除缩进）
    func_source = textwrap.dedent(inspect.getsource(f))
    # 获取函数的闭包变量（非局部变量）
    closure_vars = inspect.getclosurevars(f).nonlocals
    
    # 使用spawn启动方式，避免fork的安全风险
    with mp.Pool(5, mp.get_context('spawn')) as p:
        # 打包任务：函数源码、闭包变量、参数
        tasks = [(func_source, closure_vars, (x,), {}) for x in range(10)]
        results = p.starmap(_execute_func, tasks)
        print(results)

if __name__ == "__main__":
    main()

原理说明 #

1. 用inspect.getsource获取局部函数的源码，textwrap.dedent去除多余缩进，保证exec能正确执行。
2. 用inspect.getclosurevars提取函数依赖的闭包变量，一起传递到子进程。
3. 在子进程中通过exec重构函数并注入闭包变量，执行后返回结果。
4. 生成器是在子进程执行函数时创建的，不会被跨进程传递，自然不存在序列化问题。

方案二：临时绑定局部函数到全局命名空间 #

如果你的局部函数没有依赖闭包变量，这个方法实现更简单：临时把局部函数挂载到__main__模块的全局命名空间，让pickle能找到它，用完后清理。

代码示例 #

import multiprocessing as mp
import sys

def main():
    def f(x):
        return x * x
    
    # 生成唯一的临时函数名，避免冲突
    temp_func_name = f"_temp_local_func_{id(f)}"
    main_module = sys.modules['__main__']
    
    try:
        # 临时绑定到全局命名空间
        setattr(main_module, temp_func_name, f)
        with mp.Pool(5, mp.get_context('spawn')) as p:
            # 传递全局命名空间中的函数引用
            print(p.map(getattr(main_module, temp_func_name), range(10)))
    finally:
        # 清理临时函数，避免污染全局命名空间
        delattr(main_module, temp_func_name)

if __name__ == "__main__":
    main()

注意事项 #

• 仅适用于无闭包变量的局部函数，因为闭包变量不会被自动传递到子进程。
• 必须清理临时函数，避免全局命名空间污染。

方案三：自定义可序列化的函数包装器 #

如果需要更灵活的封装，可以把局部函数和它的依赖变量包装成一个可pickle的类，利用类的__call__方法实现调用。

代码示例 #

import multiprocessing as mp
import inspect
import textwrap

class PicklableFunc:
    def __init__(self, func):
        self.func_source = textwrap.dedent(inspect.getsource(func))
        self.closure_vars = inspect.getclosurevars(func).nonlocals
    
    def __call__(self, *args, **kwargs):
        local_namespace = self.closure_vars.copy()
        exec(self.func_source, globals(), local_namespace)
        func = next(v for v in local_namespace.values() if callable(v))
        return func(*args, **kwargs)

def main():
    multiplier = 2
    def f(x):
        return x * multiplier
    
    # 包装局部函数
    picklable_f = PicklableFunc(f)
    
    with mp.Pool(5, mp.get_context('spawn')) as p:
        print(p.map(picklable_f, range(10)))

if __name__ == "__main__":
    main()

原理说明 #

把函数的源码和闭包变量存储在类实例中，类本身是可pickle的（因为存储的都是可序列化的字符串和基本类型）。子进程中反序列化后，调用__call__方法重构函数并执行。

关键注意事项 #

• 强制使用spawn启动方式：mp.get_context('spawn')避免了fork在MacOS和Windows上的安全风险，这也是Python 3.8+在MacOS的默认启动方式。
• 生成器处理：所有方案都不需要序列化生成器对象，因为生成器是在子进程执行函数时创建的，完全在子进程内部运行。
• 闭包变量限制：如果闭包变量包含不可序列化的对象（比如打开的文件句柄），需要额外处理这些对象的传递逻辑，或者避免在闭包中使用它们。

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