如何阻止Multiprocess.Pool在任务完成前继续执行后续代码？

使用Python 3.10处理大型数据库时，原本通过以下代码生成HOPS和GROUP列：

OGS_DB[['HOPS', 'GROUP']] = OGS_DB.apply(lambda x: hops_generator(const_id=const_id, num_hops=n_hops, sat_c=x['SAT'], ul_start=x['OGS_Start'], ul_end=x['OGS_End']), axis=1).apply(pd.Series)

为提升效率改用multiprocessing.Pool并行处理，主程序代码如下：

if __name__ == '__main__':
    with Manager() as manager:
        DB_HOP = manager.list([None] * len(OGS_DB))
        with Pool(processes=cpucores - 3) as pool:
            freeze_support()
            results = pool.map(process_row, [(row, row['SAT'], const_id, n_hops, row['OGS_Start'], row['OGS_End']) for index, row in OGS_DB.iterrows()])
            pool.close()
            pool.join()
        hops_list = [result['HOPS'] for result in results]
        group_list = [result['GROUP'] for result in results]
    OGS_DB['HOPS'] = pd.Series(hops_list)
    OGS_DB['GROUP'] = pd.Series(group_list)

执行后续代码时出现KeyError: 'HOPS'报错：

OGS_DB = OGS_DB.drop('SAT', axis=1)
OGS_DB = OGS_DB.explode('HOPS').reset_index(drop=True)

完整回溯信息：

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\multiprocessing\spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\multiprocessing\spawn.py", line 125, in _main
    prepare(preparation_data)
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\multiprocessing\spawn.py", line 236, in prepare
    _fixup_main_from_path(data['init_main_from_path'])
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\multiprocessing\spawn.py", line 287, in _fixup_main_from_path
    main_content = runpy.run_path(main_path,
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\runpy.py", line 289, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\runpy.py", line 96, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "C:\Users\pizz_lu\FL4NGSO\Development\Flexible_OISL\OISL_flexible.py", line 426, in <module>
    OISL_analysis(Loc)
  File "C:\Users\pizz_lu\FL4NGSO\Development\Flexible_OISL\OISL_flexible.py", line 99, in OISL_analysis
    OGS_DB = OGS_DB.explode('HOPS').reset_index(drop=True)
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\site-packages\pandas\core\frame.py", line 9540, in explode
    result = df[columns[0]].explode()
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\site-packages\pandas\core\frame.py", line 3893, in __getitem__
    indexer = self.columns.get_loc(key)
  File "C:\Users\pizz_lu\AppData\Local\miniforge3\lib\site-packages\pandas\core\indexes\base.py", line 3797, in get_loc
    raise KeyError(key) from err
KeyError: 'HOPS'

核心问题分析 #

Windows系统中multiprocessing默认使用spawn模式启动子进程，会重新导入整个脚本。如果OISL_analysis函数调用不在if __name__ == '__main__':块内，子进程会重复执行该函数，此时子进程未运行并行代码，导致HOPS列未生成，触发报错。此外原代码存在冗余逻辑和位置错误。

具体修复步骤 #

1. 严格控制主程序作用域
   将所有主逻辑（包括OISL_analysis调用、并行处理代码）全部放入if __name__ == '__main__':块内，避免子进程重复执行非必要代码。

2. 修正freeze_support()位置
   freeze_support()必须放在if __name__ == '__main__':的最开头，仅Windows环境需要，用于支持打包成可执行文件，原代码放在Pool内部无效。

3. 移除冗余代码
   DB_HOP = manager.list([None] * len(OGS_DB))未被使用，直接删除，pool.map会直接返回结果列表，无需借助Manager共享列表。

4. 验证process_row函数
   确保函数返回包含'HOPS'和'GROUP'键的字典，示例实现：

   def process_row(args):
       row, sat, const_id, n_hops, ogs_start, ogs_end = args
       hops, group = hops_generator(const_id=const_id, num_hops=n_hops, sat_c=sat, ul_start=ogs_start, ul_end=ogs_end)
       return {'HOPS': hops, 'GROUP': group}
   

5. 修正后的完整主程序代码

   import multiprocessing
   from multiprocessing import Pool
   import pandas as pd
   
   def process_row(args):
       row, sat, const_id, n_hops, ogs_start, ogs_end = args
       hops, group = hops_generator(const_id=const_id, num_hops=n_hops, sat_c=sat, ul_start=ogs_start, ul_end=ogs_end)
       return {'HOPS': hops, 'GROUP': group}
   
   def OISL_analysis(Loc):
       # 其他预处理逻辑...
       cpucores = multiprocessing.cpu_count()
       with Pool(processes=cpucores - 3) as pool:
           results = pool.map(
               process_row, 
               [(row, row['SAT'], const_id, n_hops, row['OGS_Start'], row['OGS_End']) 
                for index, row in OGS_DB.iterrows()]
           )
       # 生成新列
       OGS_DB['HOPS'] = [res['HOPS'] for res in results]
       OGS_DB['GROUP'] = [res['GROUP'] for res in results]
       # 后续处理
       OGS_DB = OGS_DB.drop('SAT', axis=1)
       OGS_DB = OGS_DB.explode('HOPS').reset_index(drop=True)
       # 其他后续逻辑...
   
   if __name__ == '__main__':
       multiprocessing.freeze_support()
       Loc = "your_location"  # 替换为实际参数
       OISL_analysis(Loc)
   

额外优化建议 #

• 替换iterrows()为itertuples()提升遍历效率，减少内存开销：

  results = pool.map(
      process_row, 
      [(row, row.SAT, const_id, n_hops, row.OGS_Start, row.OGS_End) 
       for row in OGS_DB.itertuples(index=False)]
  )
  

• 若DataFrame规模极大，改用pool.imap_unordered()分批获取结果，降低内存占用。

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