[Python] Disable Numpy parallelization inside Numba JIT

The problem is illustrated by the following script, which works correctly if MKL is used for linear algebra operations:

from numba import njit, prange
from numpy import random, dot, empty
from threadpoolctl import ThreadpoolController
controller = ThreadpoolController()

numba_parallel=True # True False

@njit(parallel=False)
def internal_dot(a, b):
return dot(a, b)

@njit(parallel=numba_parallel)
def total_sum(b, c):
npoints=c.shape[0]
output=empty((npoints, c.shape[1], b.shape[1]))
for i in prange(npoints):
output=internal_dot(c, b)
return output

@controller.wrap(limits=1, user_api='blas')
def safe_total_sum(b, c):
return total_sum(b, c)

nvecs=256
dim1=256
dim2=256

vector=random.random((dim1, dim2))
matrix=random.random((nvecs, dim2, dim1))

_ = total_sum(vector, matrix)

_ = safe_total_sum(vector, matrix)


However, using it with OpenBLAS leads to warning OpenBLAS warning: precompiled NUM_THREADS exceeded, adding auxiliary array for thread metadata., indicating an oversubscription problem that in my (much more lengthy) usecase crashes the code. I am aware that setting export OPENBLAS_NUM_THREADS=1 solves the issue for this script, but it is not applicable in my usecase since my code calls other Numpy functions elsewhere and needs them parallelized. Using ThreadpoolController does not seem to help; I am also aware of the possibility to pack everything into a .pkl and unpack it in a subprocess environment with OPENBLAS_NUM_THREADS=1, but it'd really prefer to avoid this dirty trick.

Is there a proper Python solution for this problem?

Continue reading...