Skip to content

Commit

Permalink
use threadpoolctl for detecting BLAS threadpools (#692)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
benfred authored Sep 29, 2023
1 parent ab81242 commit 8ca25bd
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 14 deletions.
51 changes: 38 additions & 13 deletions implicit/utils.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import os
import time
import warnings

import numpy as np
import scipy.sparse
import threadpoolctl


def nonzeros(m, row):
Expand All @@ -22,19 +22,44 @@ def check_blas_config():
global _checked_blas_config # pylint: disable=global-statement
if _checked_blas_config:
return
_checked_blas_config = True

if np.__config__.get_info("openblas_info") and os.environ.get("OPENBLAS_NUM_THREADS") != "1":
warnings.warn(
"OpenBLAS detected. Its highly recommend to set the environment variable "
"'export OPENBLAS_NUM_THREADS=1' to disable its internal multithreading"
)
if np.__config__.get_info("blas_mkl_info") and os.environ.get("MKL_NUM_THREADS") != "1":
warnings.warn(
"Intel MKL BLAS detected. Its highly recommend to set the environment "
"variable 'export MKL_NUM_THREADS=1' to disable its internal "
"multithreading"
)
for api in threadpoolctl.threadpool_info():
num_threads = api["num_threads"]
if api["user_api"] != "blas" or num_threads == 1:
continue

internal_api = api["internal_api"]
if internal_api == "openblas":
warnings.warn(
f"OpenBLAS is configured to use {num_threads} threads. It is highly recommended"
" to disable its internal threadpool by setting the environment variable"
" 'OPENBLAS_NUM_THREADS=1' or by calling 'threadpoolctl.threadpool_limits(1,"
' "blas")\'. Having OpenBLAS use a threadpool can lead to severe performance'
" issues here.",
RuntimeWarning,
stacklevel=2,
)
elif internal_api == "mkl":
warnings.warn(
f"Intel MKL BLAS is configured to use {num_threads} threads. It is highly"
" recommended to disable its internal threadpool by setting the environment"
" variable 'MKL_NUM_THREADS=1' or by callng 'threadpoolctl.threadpool_limits(1,"
' "blas")\'. Having MKL use a threadpool can lead to severe performance issues',
RuntimeWarning,
stacklevel=2,
)
else:
# probably using blis, which is by default single threaded. warn anyways
# for when threadpoolctl gets support for VecLib/Accelerate etc
warnings.warn(
f"BLAS library {internal_api} is configured to use {num_threads} threads."
" It is highly recommended to disable its internal threadpool by calling"
" 'threadpoolctl.threadpool_limits(1, \"blas\")'.",
RuntimeWarning,
stacklevel=2,
)

_checked_blas_config = True


def check_random_state(random_state):
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
scipy>=0.16.0
Cython>=0.24.0
tqdm>=4.27.0
threadpoolctl
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,6 @@ def exclude_non_implicit_cmake_files(cmake_manifest):
"Collaborative Filtering, Recommender Systems"
),
packages=find_packages(),
install_requires=["numpy", "scipy>=0.16", "tqdm>=4.27"],
install_requires=["numpy", "scipy>=0.16", "tqdm>=4.27", "threadpoolctl"],
cmake_process_manifest_hook=exclude_non_implicit_cmake_files,
)

0 comments on commit 8ca25bd

Please sign in to comment.