[itamarst]

The expected joblib usage pattern is that it has a thread-local counter, which gets incremented on ThreadedBackend.configure(), and decremented on ThreadedBackend.finalize(), and if it hits zero then OpenBLASController.clear_controlling_thread() is called.

The result is that so long as you have Parallel() running in a particular thread, typically the main thread, other threads can't set OpenBLAS-pthreads thread pool size.

A separate companion PR #278 exposes an API to allow joblib to override thread limits so they're either 1 or all cores, since we want that too, I think.

[ogrisel]

I like this approach. If we ever have a second case of process wide shared threadpool, we could even introduce a mixin class for those controller methods. Or we could use ducktyping on the controllers from joblib.

[itamarst]

Here's what the corresponding joblib changes would look like, as a sketch:

2 files changed, 38 insertions(+)
joblib/_parallel_backends.py | 36 ++++++++++++++++++++++++++++++++++++
joblib/parallel.py           |  2 ++

modified   joblib/_parallel_backends.py
@@ -13,6 +13,8 @@ from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from typing import Any, Callable
 
+from threadpoolctl import OpenBLASController
+
 try:
     from threadpoolctl import get_cached_controller as _get_threadpool_controller
 except ImportError:
@@ -360,6 +362,24 @@ class ParallelBackendBase(metaclass=ABCMeta):
         return isinstance(threading.current_thread(), threading._MainThread)
 
 
+class _OpenBLASOwnership(threading.local):
+    """Count how many ``Parallel`` have modified OpenBLAS limits."""
+
+    count = 0
+
+    def increment(self):
+        """Increment ownership."""
+        self.count += 1
+
+    def decrement(self):
+        """Decrement ownership, return whether it hit zero."""
+        self.count -= 1
+        return self.count <= 0
+
+
+_OPENBLAS_OWNERSHIP = _OpenBLASOwnership()
+
+
 class SequentialBackend(ParallelBackendBase):
     """A ParallelBackend which will execute all batches sequentially.
 
@@ -373,6 +393,10 @@ class SequentialBackend(ParallelBackendBase):
     supports_retrieve_callback = False
     supports_sharedmem = True
 
+    def configure(self, *args, **kwargs):
+        _OPENBLAS_OWNERSHIP.increment()
+        return super().configure(*args, **kwargs)
+
     def effective_n_jobs(self, n_jobs):
         """Determine the number of jobs which are going to run in parallel"""
         if n_jobs == 0:
@@ -394,6 +418,13 @@ class SequentialBackend(ParallelBackendBase):
         # default backend or the number of jobs. Just return the current one.
         return get_active_backend()
 
+    def terminate(self):
+        if _OPENBLAS_OWNERSHIP.decrement() and hasattr(
+            OpenBLASController, "clear_controlling_thread"
+        ):
+            OpenBLASController.clear_controlling_thread()
+        return super().terminate()
+
 
 class PoolManagerMixin(object):
     """A helper class for managing pool of workers."""
@@ -582,6 +613,7 @@ class ThreadingBackend(PoolManagerMixin, ParallelBackendBase):
 
     def configure(self, n_jobs=1, parallel=None, **backend_kwargs):
         """Build a process or thread pool and return the number of workers"""
+        _OPENBLAS_OWNERSHIP.increment()
         self._threadpool_controller = _get_threadpool_controller()
         # Used to restore limits when when done:
         self._thread_limiter = self._threadpool_controller.limit()
@@ -611,6 +643,10 @@ class ThreadingBackend(PoolManagerMixin, ParallelBackendBase):
 
     def terminate(self):
         self._thread_limiter.restore_original_limits()
+        if _OPENBLAS_OWNERSHIP.decrement() and hasattr(
+            OpenBLASController, "clear_controlling_thread"
+        ):
+            OpenBLASController.clear_controlling_thread()
         super().terminate()
 
 
modified   joblib/parallel.py
@@ -1951,6 +1951,8 @@ class Parallel(Logger):
             self._iterating = False
             self._original_iterator = None
             self.print_progress()
+            if not self._managed_backend:
+                self._backend.terminate()
 
     def _reset_run_tracking(self):
         """Reset the counters and flags used to track the execution."""

[itamarst]

Once #280 is merged can rerun and see if CI will be green.

[ogrisel]

Thanks @itamarst. I think the approach implemented in this PR is the best way to deal with OpenBLAS threading semantics out of all the previously explored options.

Here are a few small improvements/fixes but otherwise LGTM.

[ogrisel]

BTW: I am concurrently attempting to fix the flexiblas CI build error in #281.

[itamarst]

I've addressed the review comments.

[ogrisel]

I really like the latest documentation/docstring updates. Thanks!

I think this PR is good to go the way it is but I wonder if the threadpool_limits class itself could implement a thread-local counter when used as a context manager so as to call clear_controlling_thread in __exit__ whenever applicable?

I don't think we would be able to leverage that in joblib but maybe this could be convenient for other projects that use the threadpoolctl public API directly?

[itamarst]

I guess if it's well-factored could also use it in joblib, either directly if we make it a public API, or just copy/pasting.

[itamarst]

Hm. One problem with clearing automatically is that if a user mixes with/wrap style usage with imperative usage, this could result in unexpected behavior.

In practice... that's probably not that likely.

[itamarst]

@ogrisel ok did the decorator/context manager thing.

[ogrisel]

Shall we make the OpenBLAS ownership singleton a public class attribute of the OpenBLAS controller class to make easier to reuse it directly from joblib without copy pasting code and without importing private threadpoolctl names? Or do you have a better idea?

[itamarst]

Public singleton would make it easier for joblib to use and interoperate with, yeah. Will get to that tomorrow.