Python多进程Map中共享大型numpy只读数组的高效方案问询

核心问题分析 #

你的痛点在于并行化时的对象序列化开销——无论是将大数组作为类成员还是全局变量，跨进程传递时都会触发Pickle序列化，对于100GB+的数组来说，这种复制/序列化完全不可行。下面是几种直接解决问题的高效方案：

1. 优先选择：内存映射文件（Memory-Mapped Files） #

这是处理超大型只读数组最简洁高效的方案，完全规避序列化和内存复制：

• 原理：利用操作系统的文件缓存，让所有进程直接从磁盘映射数组，无需将整个数组加载到内存，所有进程共享缓存页，不会重复占用内存。
• 实现方式：
  主进程和所有子进程都通过mmap_mode="r"加载数组，不需要传递数组对象，只需要确保所有进程能访问到同一个.npy文件：

  import numpy as np
  
  def get_large_array():
      # 只读模式映射磁盘文件，仅加载元数据，实际数据按需从磁盘读取
      return np.load("/path/to/your/large_array.npy", mmap_mode="r")
  
  # 类A的处理逻辑直接使用该数组
  class A:
      def process_chunk(self, chunk_indices):
          arr = get_large_array()
          return arr[chunk_indices].sum()  # 示例处理逻辑
  
  # 类B的并行调用逻辑
  class B:
      def run_parallel(self, chunk_indices_list):
          from concurrent.futures import ProcessPoolExecutor
          with ProcessPoolExecutor() as executor:
              results = executor.map(A().process_chunk, chunk_indices_list)
          return list(results)
  

• 优势：零序列化开销、内存占用极低（仅加载数组元数据）、实现简单，无需额外处理进程间通信。

2. 共享内存（Shared Memory）方案 #

如果数组需要频繁访问，且磁盘IO成为瓶颈，可以用共享内存将数组一次性加载到内存，供所有进程直接访问：

• 原理：主进程将数组加载到共享内存区域，子进程通过内存地址直接访问，完全避免数据复制。
• 实现方式（Python 3.8+）：

  import numpy as np
  from multiprocessing import shared_memory, ProcessPoolExecutor
  
  # 主进程初始化共享内存
  def init_shared_memory():
      global shm, shared_arr
      # 先以内存映射方式加载数组，避免主进程占用过多内存
      original_arr = np.load("/path/to/large_array.npy", mmap_mode="r")
      # 创建共享内存块
      shm = shared_memory.SharedMemory(create=True, size=original_arr.nbytes)
      # 将数组数据复制到共享内存（仅主进程执行一次）
      shared_arr = np.ndarray(original_arr.shape, dtype=original_arr.dtype, buffer=shm.buf)
      shared_arr[:] = original_arr[:]
      return shm.name, original_arr.shape, original_arr.dtype
  
  # 子进程初始化函数：通过共享内存名获取数组
  def init_child_process(shm_name, arr_shape, arr_dtype):
      global big_array
      existing_shm = shared_memory.SharedMemory(name=shm_name)
      big_array = np.ndarray(arr_shape, dtype=arr_dtype, buffer=existing_shm.buf)
  
  # 子进程任务函数
  def task(chunk_indices):
      return big_array[chunk_indices].mean()
  
  class B:
      def run_parallel(self, chunk_indices_list):
          shm_name, arr_shape, arr_dtype = init_shared_memory()
          with ProcessPoolExecutor(initializer=init_child_process, 
                                  initargs=(shm_name, arr_shape, arr_dtype)) as executor:
              results = executor.map(task, chunk_indices_list)
          # 主进程最后关闭共享内存
          shm.close()
          shm.unlink()
          return list(results)
  

• 优势：访问速度比内存映射更快（直接从内存读取），适合高频访问场景。

3. 优化进程池初始化逻辑 #

如果必须在子进程中复用数组，可以利用进程池的initializer参数，让每个子进程仅加载一次数组：

• 原理：每个子进程启动时执行初始化函数，将数组加载为子进程的全局变量，避免每次任务重复加载。
• 实现方式：

  import numpy as np
  from concurrent.futures import ProcessPoolExecutor
  
  # 子进程全局变量，仅初始化一次
  child_global_arr = None
  
  def init_child():
      global child_global_arr
      # 内存映射加载数组，子进程仅执行一次
      child_global_arr = np.load("/path/to/large_array.npy", mmap_mode="r")
  
  class A:
      def process(self, indices):
          return child_global_arr[indices].max()
  
  class B:
      def run_parallel(self, indices_list):
          with ProcessPoolExecutor(initializer=init_child) as executor:
              a_instance = A()
              results = executor.map(a_instance.process, indices_list)
          return list(results)
  

• 优势：避免子进程重复加载数组，减少磁盘IO开销。

方案选择建议 #

1. 若数组规模超过物理内存：优先用内存映射文件，避免OOM。
2. 若数组能完全放入物理内存且访问频繁：用共享内存，获得最高访问速度。
3. 若进程池任务较多：配合initializer优化子进程初始化逻辑，减少重复加载。

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