机器学习模型多进程运行中的过度线程问题：如何限制单线程？

问题背景 #

使用多进程并行运行机器学习模型时，带num_threads、num_jobs等线程控制参数的模型能正常限制线程数，但对于无此类参数的模型（如sklearn.linear_model.PoissonRegressor），无法确保每个模型仅使用1个线程，极简复现代码如下：

import multiprocessing
import numpy as np
from sklearn.linear_model import PoissonRegressor

class MPModelWrapper:
    """ Wraps model for use in multiprocessing """
    def __init__(self, model):
        self.model = model
        
    def fit_predict(self, train_target_val_pair):
        X = train_target_val_pair[0]
        y = train_target_val_pair[1]
        val = train_target_val_pair[2]
        
        self.model.fit(X,y)
        return self.model.predict(val)


# Fake data to train models on
num_models_to_run = 100
train_target_val_pairs = [ 
    (np.random.random((1000, 10)), np.random.randint(0, 2, size=(1000,) ), np.random.random((10, 10)) )  
    for i in range(num_models_to_run) 
]

# Multiprocessor
num_concurrent_models = 1
model = PoissonRegressor()
with multiprocessing.Pool( num_concurrent_models ) as p:
    results = p.map(MPModelWrapper(model).fit_predict, train_target_val_pairs)

核心问题：如何确保每个调用的模型仅使用1个线程？

已尝试但无效的方案 #

• 手动设置全局环境变量：在代码开头设置线程数相关环境变量，但每个模型仍占用全部核心：

# At the top of the pgm above all the previous code
import os
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['BLAS_NUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'

• 使用joblib并行：替换为joblib的Parallel接口，问题依旧：

from joblib import Parallel, delayed
import os

# Set global threading limits
os.environ.update({
    'OMP_NUM_THREADS': '1',
    'MKL_NUM_THREADS': '1',
    'OPENBLAS_NUM_THREADS': '1',
    'BLAS_NUM_THREADS': '1'
})


preds = Parallel(n_jobs=os.cpu_count(), backend='threading')(
    delayed(fit_predict_wrapper)() for _ in range(1)
)

有效解决方案 #

1. 在子进程内部设置线程限制 #

部分科学计算库（如MKL、OpenBLAS）会在导入阶段读取环境变量，主进程提前设置的变量可能无法被子进程的库实例重新读取。因此需在子进程执行任务时，实时设置环境变量并调用库的API强制生效：

修改MPModelWrapper的fit_predict方法：

def fit_predict(self, train_target_val_pair):
    # 子进程内设置线程数环境变量
    import os
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    os.environ['OPENBLAS_NUM_THREADS'] = '1'
    os.environ['BLAS_NUM_THREADS'] = '1'
    os.environ['NUMEXPR_NUM_THREADS'] = '1'
    
    # 调用库API强制设置线程数（针对MKL/OpenBLAS）
    try:
        import mkl
        mkl.set_num_threads(1)
    except ImportError:
        pass
    
    try:
        import openblas
        openblas.set_num_threads(1)
    except ImportError:
        pass
    
    X = train_target_val_pair[0]
    y = train_target_val_pair[1]
    val = train_target_val_pair[2]
    
    self.model.fit(X,y)
    return self.model.predict(val)

2. 在子进程内创建模型实例 #

主进程创建的模型实例在传递给子进程时可能存在序列化/资源继承问题，导致线程限制失效。改为在子进程内动态创建模型：

class MPModelWrapper:
    """ Wraps model for use in multiprocessing """
    def __init__(self, model_cls, **model_kwargs):
        self.model_cls = model_cls
        self.model_kwargs = model_kwargs
        
    def fit_predict(self, train_target_val_pair):
        # 子进程内创建全新模型实例
        model = self.model_cls(**self.model_kwargs)
        
        # 设置线程限制（同方案1）
        import os
        os.environ['OMP_NUM_THREADS'] = '1'
        os.environ['MKL_NUM_THREADS'] = '1'
        os.environ['OPENBLAS_NUM_THREADS'] = '1'
        os.environ['BLAS_NUM_THREADS'] = '1'
        os.environ['NUMEXPR_NUM_THREADS'] = '1'
        
        try:
            import mkl
            mkl.set_num_threads(1)
        except ImportError:
            pass
        
        X = train_target_val_pair[0]
        y = train_target_val_pair[1]
        val = train_target_val_pair[2]
        
        model.fit(X,y)
        return model.predict(val)

调用方式修改为：

with multiprocessing.Pool(num_concurrent_models) as p:
    results = p.map(MPModelWrapper(PoissonRegressor).fit_predict, train_target_val_pairs)

3. 强制使用spawn启动子进程（Linux/macOS） #

Linux/macOS默认使用fork方式启动子进程，会继承主进程的所有资源（包括已初始化的线程池），导致线程限制无法重置。强制使用spawn方式，每个子进程会全新启动，确保环境变量生效：

import multiprocessing
# 在创建Pool前设置启动方式
multiprocessing.set_start_method('spawn')

# 后续创建Pool的代码不变
model = PoissonRegressor()
with multiprocessing.Pool(num_concurrent_models) as p:
    results = p.map(MPModelWrapper(model).fit_predict, train_target_val_pairs)

总结 #

最可靠的组合方案是：在子进程内创建模型实例 + 实时设置环境变量并调用库API强制限制线程数 + （Linux/macOS下）使用spawn启动方式，可确保每个模型仅使用1个线程运行。

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