Skip to content

Commit 8ca25bd

Browse files
authored
use threadpoolctl for detecting BLAS threadpools (#692)
We were using numpy.__config__.get_info to figure out what BLAS library was being used, in an effort to warn people if they hadn't disabled the internal threadpool for the BLAS they were using. This method is removed in numpy 1.26, causing errors. Fix by using the threadpoolctl library to detect both the BLAS library and the number of threads its configured for. While we could automatically configure the BLAS library to reduce the threadpool size to 1, this is process wide - and would have side effects for our users. Instead just warn here, and give instructions for people on how to configure themselves.
1 parent ab81242 commit 8ca25bd

File tree

3 files changed

+40
-14
lines changed

3 files changed

+40
-14
lines changed

implicit/utils.py

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
import os
21
import time
32
import warnings
43

54
import numpy as np
65
import scipy.sparse
6+
import threadpoolctl
77

88

99
def nonzeros(m, row):
@@ -22,19 +22,44 @@ def check_blas_config():
2222
global _checked_blas_config # pylint: disable=global-statement
2323
if _checked_blas_config:
2424
return
25-
_checked_blas_config = True
2625

27-
if np.__config__.get_info("openblas_info") and os.environ.get("OPENBLAS_NUM_THREADS") != "1":
28-
warnings.warn(
29-
"OpenBLAS detected. Its highly recommend to set the environment variable "
30-
"'export OPENBLAS_NUM_THREADS=1' to disable its internal multithreading"
31-
)
32-
if np.__config__.get_info("blas_mkl_info") and os.environ.get("MKL_NUM_THREADS") != "1":
33-
warnings.warn(
34-
"Intel MKL BLAS detected. Its highly recommend to set the environment "
35-
"variable 'export MKL_NUM_THREADS=1' to disable its internal "
36-
"multithreading"
37-
)
26+
for api in threadpoolctl.threadpool_info():
27+
num_threads = api["num_threads"]
28+
if api["user_api"] != "blas" or num_threads == 1:
29+
continue
30+
31+
internal_api = api["internal_api"]
32+
if internal_api == "openblas":
33+
warnings.warn(
34+
f"OpenBLAS is configured to use {num_threads} threads. It is highly recommended"
35+
" to disable its internal threadpool by setting the environment variable"
36+
" 'OPENBLAS_NUM_THREADS=1' or by calling 'threadpoolctl.threadpool_limits(1,"
37+
' "blas")\'. Having OpenBLAS use a threadpool can lead to severe performance'
38+
" issues here.",
39+
RuntimeWarning,
40+
stacklevel=2,
41+
)
42+
elif internal_api == "mkl":
43+
warnings.warn(
44+
f"Intel MKL BLAS is configured to use {num_threads} threads. It is highly"
45+
" recommended to disable its internal threadpool by setting the environment"
46+
" variable 'MKL_NUM_THREADS=1' or by callng 'threadpoolctl.threadpool_limits(1,"
47+
' "blas")\'. Having MKL use a threadpool can lead to severe performance issues',
48+
RuntimeWarning,
49+
stacklevel=2,
50+
)
51+
else:
52+
# probably using blis, which is by default single threaded. warn anyways
53+
# for when threadpoolctl gets support for VecLib/Accelerate etc
54+
warnings.warn(
55+
f"BLAS library {internal_api} is configured to use {num_threads} threads."
56+
" It is highly recommended to disable its internal threadpool by calling"
57+
" 'threadpoolctl.threadpool_limits(1, \"blas\")'.",
58+
RuntimeWarning,
59+
stacklevel=2,
60+
)
61+
62+
_checked_blas_config = True
3863

3964

4065
def check_random_state(random_state):

requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
scipy>=0.16.0
22
Cython>=0.24.0
33
tqdm>=4.27.0
4+
threadpoolctl

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,6 @@ def exclude_non_implicit_cmake_files(cmake_manifest):
4444
"Collaborative Filtering, Recommender Systems"
4545
),
4646
packages=find_packages(),
47-
install_requires=["numpy", "scipy>=0.16", "tqdm>=4.27"],
47+
install_requires=["numpy", "scipy>=0.16", "tqdm>=4.27", "threadpoolctl"],
4848
cmake_process_manifest_hook=exclude_non_implicit_cmake_files,
4949
)

0 commit comments

Comments
 (0)