Python UMAP模型在multiprocessing.Process中训练卡死问题咨询

问题根因 #

• Unix类系统下Python multiprocessing默认使用fork模式启动子进程，若父进程中已经运行过带多线程的计算逻辑（比如你代码中主进程先调用了UMAP拟合、或是UMAP底层依赖的OpenMP/BLAS库的多线程逻辑），fork出来的子进程只会保留当前调用线程的状态，其余线程的锁状态会出现异常，导致子进程内部的多线程运算直接触发死锁。
• UMAP默认启用多线程拟合，在fork模式的子进程中运行时就会触发上述死锁，所以你会看到代码停在reducer.fit_transform步骤永远不返回。

可落地的解决方法 #

方法1：修改multiprocessing启动模式为spawn #

这是最通用的解决方案，spawn模式会全新启动独立的Python进程，不会继承父进程的线程和锁状态，完全规避fork带来的多线程死锁问题。
只需要在导入multiprocessing后、所有多进程逻辑运行前新增一行设置代码即可：

import umap
import multiprocessing
multiprocessing.set_start_method('spawn') # 新增这行，位置不能错
import numpy as np
import sys
import time

其余代码保持不变即可正常运行。

方法2：强制UMAP在子进程中使用单线程运行 #

如果必须保留fork启动模式，可以在初始化UMAP的时候指定n_jobs=1，关闭UMAP内部的多线程逻辑，避免死锁：

reducer = umap.UMAP(n_jobs=1)

该方案不需要修改启动模式，缺点是UMAP拟合速度会有下降。

方法3：全局设置OpenMP线程数为1 #

如果是底层依赖的BLAS/OpenMP库导致的死锁，可以在代码最开头设置环境变量：

import os
os.environ['OMP_NUM_THREADS'] = '1'

该方案可以同时解决所有依赖OpenMP的库在fork子进程中的死锁问题，缺点是所有相关库的运算都会变成单线程，整体性能会受影响。

补充临时方案 #

如果你的业务逻辑允许主进程不提前运行UMAP拟合，也可以把主进程的单线程测试逻辑移到子进程运行逻辑之后，也能规避死锁，但通用性不如上述三个方案。

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