多进程中使用初始化器传递函数的原因及共享内存替代方案咨询

嗨，我来帮你拆解这几个问题，结合Linux下多进程的运行机制给你讲明白：

先说说你两种方案的差异和初始化器的优势 #

第一种全局f的方案 #

你猜的没错但也不全对——在Linux下，multiprocessing.Pool默认用fork机制创建子进程：主进程创建全局f（包含大数组M）后，子进程会通过**写时复制（COW）**机制共享父进程的内存页。关键点是：只要主进程和子进程都只读取M、不修改它，那M就不会被复制到每个worker，大家共享同一块内存。你觉得会复制，可能是混淆了Windows的spawn机制（Windows下子进程会重新导入模块，必须显式传递数据）。

第二种初始化器的方案 #

为啥说这种风格更好？主要有两个核心原因：

1. 跨平台兼容性：如果你的代码要跑在Windows上，spawn机制下子进程不会继承父进程的全局变量，必须用初始化器在每个子进程启动时执行初始化逻辑，这时候init_worker就派上用场了。
2. 代码结构更清晰：初始化逻辑和任务执行逻辑分离，别人看代码的时候一眼就能知道worker进程启动时要做什么准备，比依赖全局变量更易维护、更易排查问题。

至于init_worker里要声明f为全局，是因为init_worker是在子进程的独立上下文里执行的，如果不声明全局，f只是这个函数的局部变量，后面process_task根本找不到它。

那为啥这个方案更慢？因为每个worker进程启动时都会调用make_func()重新创建一遍大数组M——相当于你创建了num_cores个完全一样的M，总初始化时间自然比第一种方案（主进程只创建一次M）长很多。

核心需求：避免大数组被复制的最优方案 #

既然你用的是Linux，其实有两种高效的方式实现：

方式1：利用Linux的COW机制（最简单、性能最优） #

只要调整第一种方案的顺序，让主进程先创建M，再创建Pool（也就是fork子进程），并且保证后续不会修改M，就能让所有worker共享同一块内存，完全不会复制：

from multiprocessing import Pool
from time import perf_counter as now
import numpy as np
import os


def make_func():
    n = 20000
    np.random.seed(7)
    M = np.random.rand(n, n)
    return lambda x, y: M[x, x] + M[y, y]


class ParallelProcessor:
    def __init__(self):
        pass
        
    def process_task(self, args):
        """Unpack arguments internally"""
        index, integer_arg = args
        print(f(index, integer_arg))

    def run_parallel(self, tasks, num_cores=None):
        """Simplified parallel execution without partial"""
        num_cores = num_cores or len(os.sched_getaffinity(0))
        task_args = [(idx, val) for idx, val in enumerate(tasks)]
        
        # 先创建f（也就是M），再创建Pool，确保fork时M已经存在
        start = now()
        global f
        f = make_func()
        print(f"************** {now() - start} seconds to make f")
        
        start = now()
        with Pool(num_cores) as pool:
            results = pool.map(self.process_task, task_args)
        print(f"************** {now() - start} seconds to run all jobs")
        return results


if __name__ == "__main__":
    processor = ParallelProcessor()
    processor.run_parallel(tasks=[1, 2, 3, 4, 5], num_cores=2)

方式2：显式使用共享内存（跨平台通用） #

如果你的代码需要跨平台，或者必须修改M（又不想复制），可以用Python3.8+提供的multiprocessing.shared_memory，直接在进程间共享numpy数组：

from multiprocessing import Pool, shared_memory
from time import perf_counter as now
import numpy as np
import os


def init_worker(shm_name, shm_shape, shm_dtype):
    global f, M
    # 子进程连接到主进程创建的共享内存
    existing_shm = shared_memory.SharedMemory(name=shm_name)
    # 创建numpy数组视图，不复制数据，直接映射共享内存
    M = np.ndarray(shm_shape, dtype=shm_dtype, buffer=existing_shm.buf)
    f = lambda x, y: M[x, x] + M[y, y]


class ParallelProcessor:
    def __init__(self):
        pass
        
    def process_task(self, args):
        index, integer_arg = args
        print(f(index, integer_arg))

    def run_parallel(self, tasks, num_cores=None):
        num_cores = num_cores or len(os.sched_getaffinity(0))
        task_args = [(idx, val) for idx, val in enumerate(tasks)]
        
        start = now()
        # 主进程创建大数组
        n = 20000
        np.random.seed(7)
        M = np.random.rand(n, n)
        # 创建共享内存，大小和数组一致
        shm = shared_memory.SharedMemory(create=True, size=M.nbytes)
        # 将数组数据复制到共享内存
        shm_M = np.ndarray(M.shape, dtype=M.dtype, buffer=shm.buf)
        shm_M[:] = M[:]
        print(f"************** {now() - start} seconds to create and share array")
        
        start = now()
        # 传递共享内存的元信息给子进程初始化
        with Pool(num_cores, initializer=init_worker, initargs=(shm.name, M.shape, M.dtype)) as pool:
            results = pool.map(self.process_task, task_args)
        # 主进程关闭并销毁共享内存
        shm.close()
        shm.unlink()
        print(f"************** {now() - start} seconds to run all jobs")
        return results


if __name__ == "__main__":
    processor = ParallelProcessor()
    processor.run_parallel(tasks=[1, 2, 3, 4, 5], num_cores=2)

这个方案里，所有worker共享同一块共享内存，不管是Linux还是Windows都能用，完全避免了数组复制。

备注：内容来源于stack exchange，提问作者Simd