使用multiprocessing.Pool时输出长度与输入可迭代对象不一致问题

问题原因 #

出现该现象的核心原因是部分输入在子进程中执行get_dict方法时触发异常，未正常返回结果，常见触发场景如下：

1. Python层逻辑边界异常：get_dict中未处理边界场景，比如某原子的一阶邻居数为0时，next(Gn.neighbors(0))会直接抛出StopIteration异常；
2. 底层扩展崩溃：依赖的acat、ase、networkx底层调用C/C++扩展时触发段错误，子进程直接被系统终止，不会抛出Python级别的异常到主进程；
3. 系统资源不足：并行任务数过高时触发OOM（内存不足），系统主动杀死部分子进程，无返回结果。

排查与修复步骤 #

步骤1：先串行执行定位问题

临时修改__call__方法为串行逻辑，同时加入异常捕获，确认是否存在Python层逻辑错误：

def __call__(self, images):
    dicts = []
    for idx, atoms in enumerate(images):
        try:
            dicts.append(self.get_dict(atoms))
        except Exception as e:
            print(f"处理第{idx}个输入失败: {repr(e)}")
    return dicts

如果运行后有报错，针对性修复get_dict的边界逻辑即可，比如针对邻居为空的场景加判断：

# 在get_dict的for循环内加入边界判断
nbrs = np.where(A[i] == 1)[0]
if len(nbrs) < 2:
    # 自定义边界场景处理逻辑，比如跳过当前原子的统计
    continue
An = A[nbrs,:][:,nbrs]
Gn = nx.from_numpy_matrix(An)
if 0 not in Gn or len(list(Gn.neighbors(0))) == 0:
    continue
path = max(nx.all_simple_paths(Gn, source=0, target=next(Gn.neighbors(0))),
           key=lambda x: len(x))

步骤2：并行场景下捕获子进程异常

如果串行执行完全正常，说明是并行时子进程异常退出，修改并行逻辑逐个捕获任务结果：

def __call__(self, images):
    pool = Pool(self.n_jobs)
    tasks = [pool.apply_async(self.get_dict, (atoms,)) for atoms in images]
    pool.close()
    pool.join()
    dicts = []
    for idx, task in enumerate(tasks):
        try:
            dicts.append(task.get())
        except Exception as e:
            print(f"第{idx}个并行任务失败: {repr(e)}")
    return dicts

步骤3：排查系统资源问题

如果上述步骤没有捕获到Python异常，说明子进程被系统主动杀死，可查看系统日志确认是否存在OOM记录，适当调小n_jobs参数（比如从20改为10）后重新测试。

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