[Python] Overthreading in Multiprocessing for ML models

The problem


I am running machine learning models in parallel using multiprocessing. When using models with parameters stating the number of threads used- num_threads, num_jobs, etc. - the code works well. Recently, I have had to use models that do not have this parameter. Below is a minimal example of my code:

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)


How can I ensure that each model called only uses 1 thread?

Attempts to Fix the Problem

1. Manually setting libraries to 1 core I tried this but it still indicated full core usage by each model:

#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'

1. joblib. I used joblib but encountered the same problems:

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)
)

Continue reading...