Python类中多进程并行化plateID匹配问题求助（M1 macOS+Jupyter）

问题背景 #

我编写了一个Python类用于系列计算，初始化时会在嵌套字典中存储多个pandas.DataFrame。类初始化的性能瓶颈在于：需要为覆盖全球的规则格点集匹配所属构造板块的ID，单轮耗时十数秒，远慢于预期的1秒以内。尝试用multiprocessing.Pool并行处理格点匹配逻辑时，每个worker进程都会重新初始化整个类，不符合仅并行目标函数的预期。

相关核心代码如下：

格点匹配函数 #

def process_plateIDs(
        geometries_data: list,
        lats_chunk: _numpy.array,
        lons_chunk: _numpy.array,
    ) -> list:
    plateIDs = _numpy.zeros(len(lats_chunk))
    
    for topology_geometry, topology_plateID in geometries_data:
        mask = topology_geometry.contains(_geopandas.points_from_xy(lons_chunk, lats_chunk))
        plateIDs[mask] = topology_plateID

    return plateIDs

几何数据提取函数 #

def extract_geometry_data(topology_geometries):
    return [(geom, plateID) for geom, plateID in zip(topology_geometries.geometry, topology_geometries.PLATEID1)]

并行尝试代码 #

# Use all available CPUs if num_workers is not specified
if num_workers is None:
   num_workers = _os.cpu_count()

# Split the data into chunks
chunk_size = len(lats) // num_workers
chunks = [(geometries_data.copy(), lats[i:i + chunk_size].copy(), lons[i:i + chunk_size].copy()) for i in range(0, len(lats), chunk_size)]

# Create a Pool of workers
with Pool(num_workers) as pool:
   # Map the process_chunk function to chunks
   results = pool.starmap(process_plateIDs, chunks)

# Concatenate results from all chunks
plateIDs = _numpy.concatenate(results)

运行环境为Jupyter Notebook，设备是搭载M1芯片的MacOS。

原因分析 #

1. M1 Mac的进程启动机制：M1芯片采用ARM架构，Python的multiprocessing默认使用spawn启动方式（而非Intel Mac的fork）。spawn会创建全新的Python进程，重新执行所有顶层代码，包括类的定义与初始化逻辑。
2. Jupyter环境特性：Jupyter的每个单元格代码都运行在__main__命名空间下，worker进程启动时会重复执行所有已运行的单元格代码，导致类被多次初始化。
3. 不必要的对象复制：代码中对geometries_data和格点数组的显式复制不仅无效，还会增加序列化开销，加重性能负担。

解决方案 #

1. 分离主进程与worker进程的代码逻辑 #

通过判断进程名称，确保只有主进程执行类的定义、初始化及顶层业务代码，worker进程仅加载必要的函数和依赖：

import numpy as _numpy
import geopandas as _geopandas
import multiprocessing as mp
import os as _os

# 仅定义worker进程需要的函数，放在顶层
def process_plateIDs(geometries_data, lats_chunk, lons_chunk):
    plateIDs = _numpy.zeros(len(lats_chunk))
    for topology_geometry, topology_plateID in geometries_data:
        mask = topology_geometry.contains(_geopandas.points_from_xy(lons_chunk, lats_chunk))
        plateIDs[mask] = topology_plateID
    return plateIDs

def extract_geometry_data(topology_geometries):
    return [(geom, plateID) for geom, plateID in zip(topology_geometries.geometry, topology_geometries.PLATEID1)]

# 仅主进程执行类定义与初始化
if mp.current_process().name == 'MainProcess':
    class PlateCalculator:
        def __init__(self):
            # 初始化嵌套字典中的DataFrame（加载板块几何数据等操作）
            self.topology_geometries = {}
            # 示例：self.topology_geometries[age][case] = _geopandas.read_file(...)
            pass

        def compute_plate_ids(self, lats, lons, age, case, num_workers=None):
            # 获取目标年代/场景的板块几何数据
            topology_df = self.topology_geometries[age][case]
            geometries_data = extract_geometry_data(topology_df)

            # 确定工作进程数
            if num_workers is None:
                num_workers = _os.cpu_count()

            # 针对M1 Mac设置进程启动方式为fork（避免spawn的重复初始化问题）
            try:
                mp.set_start_method('fork', force=True)
            except RuntimeError:
                # 若已设置过启动方式，忽略错误
                pass

            # 拆分格点数据为多个块（无需显式复制，fork模式下共享内存）
            chunk_size = len(lats) // num_workers
            chunks = [
                (geometries_data, lats[i:i+chunk_size], lons[i:i+chunk_size])
                for i in range(0, len(lats), chunk_size)
            ]

            # 并行处理
            with mp.Pool(num_workers) as pool:
                results = pool.starmap(process_plateIDs, chunks)

            # 合并结果
            return _numpy.concatenate(results)

    # 主进程实例化类并执行计算
    calculator = PlateCalculator()
    # 示例调用：plate_ids = calculator.compute_plate_ids(lats_array, lons_array, age=0, case='default')

2. 额外优化建议 #

• 几何数据序列化优化：若geometries_data中的几何对象序列化开销大，可将其转换为WKB（Well-Known Binary）格式传递，在worker进程中再解析回几何对象：

  # 主进程转换为WKB
  geometries_data = [(geom.wkb, plateID) for geom, plateID in zip(topology_df.geometry, topology_df.PLATEID1)]
  
  # worker进程中解析
  def process_plateIDs(geometries_data, lats_chunk, lons_chunk):
      plateIDs = _numpy.zeros(len(lats_chunk))
      for geom_wkb, plateID in geometries_data:
          geom = _geopandas.GeoSeries.from_wkb([geom_wkb])[0]
          mask = geom.contains(_geopandas.points_from_xy(lons_chunk, lats_chunk))
          plateIDs[mask] = plateID
      return plateIDs
  

• 空间连接替代循环：考虑用geopandas.sjoin替代逐多边形的contains循环，本身就有更高的效率，再结合并行可进一步提速。

内容的提问来源于stack exchange，提问作者thomas.s