MNT use new threadpoolctl API (global threadpool controller) (#21206)

jeremiedbb · web-flow · commit 1cab25c40d18 · 2021-10-08T15:14:00.000+02:00
diff --git a/sklearn/cluster/_kmeans.py b/sklearn/cluster/_kmeans.py
@@ -15,13 +15,13 @@
 
 import numpy as np
 import scipy.sparse as sp
-from threadpoolctl import threadpool_limits
-from threadpoolctl import threadpool_info
 
 from ..base import BaseEstimator, ClusterMixin, TransformerMixin
 from ..metrics.pairwise import euclidean_distances
 from ..metrics.pairwise import _euclidean_distances
 from ..utils.extmath import row_norms, stable_cumsum
+from ..utils.fixes import threadpool_limits
+from ..utils.fixes import threadpool_info
 from ..utils.sparsefuncs_fast import assign_rows_csr
 from ..utils.sparsefuncs import mean_variance_axis
 from ..utils import check_array
diff --git a/sklearn/cluster/tests/test_k_means.py b/sklearn/cluster/tests/test_k_means.py
@@ -4,13 +4,13 @@
 
 import numpy as np
 from scipy import sparse as sp
-from threadpoolctl import threadpool_limits
 
 import pytest
 
 from sklearn.utils._testing import assert_array_equal
 from sklearn.utils._testing import assert_allclose
 from sklearn.utils.fixes import _astype_copy_false
+from sklearn.utils.fixes import threadpool_limits
 from sklearn.base import clone
 from sklearn.exceptions import ConvergenceWarning
 
diff --git a/sklearn/utils/fixes.py b/sklearn/utils/fixes.py
@@ -13,11 +13,13 @@
 from functools import update_wrapper
 import functools
 
+import sklearn
 import numpy as np
 import scipy.sparse as sp
 import scipy
 import scipy.stats
 from scipy.sparse.linalg import lsqr as sparse_lsqr  # noqa
+import threadpoolctl
 from .._config import config_context, get_config
 from ..externals._packaging.version import parse as parse_version
 
@@ -271,3 +273,33 @@ def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis
             dtype=dtype,
             axis=axis,
         )
+
+
+# compatibility fix for threadpoolctl >= 3.0.0
+# since version 3 it's possible to setup a global threadpool controller to avoid
+# looping through all loaded shared libraries each time.
+# the global controller is created during the first call to threadpoolctl.
+def _get_threadpool_controller():
+    if not hasattr(threadpoolctl, "ThreadpoolController"):
+        return None
+
+    if not hasattr(sklearn, "_sklearn_threadpool_controller"):
+        sklearn._sklearn_threadpool_controller = threadpoolctl.ThreadpoolController()
+
+    return sklearn._sklearn_threadpool_controller
+
+
+def threadpool_limits(limits=None, user_api=None):
+    controller = _get_threadpool_controller()
+    if controller is not None:
+        return controller.limit(limits=limits, user_api=user_api)
+    else:
+        return threadpoolctl.threadpool_limits(limits=limits, user_api=user_api)
+
+
+def threadpool_info():
+    controller = _get_threadpool_controller()
+    if controller is not None:
+        return controller.info()
+    else:
+        return threadpoolctl.threadpool_info()
