diff --git a/CHANGES.md b/CHANGES.md index 87bd115e..7317c063 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -13,6 +13,7 @@ thread controlls OpenBLAS in these cases, and decorator and context manager limit APIs will automatically call that API when they are done. https://github.com/joblib/threadpoolctl/pull/275 + https://github.com/joblib/threadpoolctl/pull/289 - Add a new API to get a cached `ThreadpoolController` instance (`get_cached_controller()`), as an optimization for cases where limits are set diff --git a/tests/test_threadpoolctl.py b/tests/test_threadpoolctl.py index f8760396..82c918f9 100644 --- a/tests/test_threadpoolctl.py +++ b/tests/test_threadpoolctl.py @@ -8,7 +8,7 @@ import subprocess import sys import time -from threading import Thread, current_thread +from threading import Lock, Thread, current_thread import pytest @@ -792,11 +792,41 @@ def clear_and_set(): run_in_thread(openblas_lib.set_num_threads, 1) assert openblas_lib.get_num_threads() == 1 + # Dead threads lose control: + openblas_lib.set_num_threads(orig_num_threads) + assert openblas_lib.get_num_threads() == orig_num_threads + + # Release control: + openblas_lib.clear_controlling_thread() + # clear_controlling_thread(force=True) allows any thread to clear: + was_set, finish = Lock(), Lock() + was_set.acquire() + finish.acquire() + + def set_and_exit_after_lock_released(): + openblas_lib.set_num_threads(1) + was_set.release() + finish.acquire() + + t = Thread(target=set_and_exit_after_lock_released) + t.start() + was_set.acquire() + assert openblas_lib.get_num_threads() == 1 + + # This thread can't set since it's not the controller: + openblas_lib.set_num_threads(orig_num_threads) + assert openblas_lib.get_num_threads() == 1 + + # But when forced clear happens, we can set: openblas_lib.clear_controlling_thread(force=True) openblas_lib.set_num_threads(orig_num_threads) assert openblas_lib.get_num_threads() == orig_num_threads + # Clean up + finish.release() + t.join() + def _openblas_controlling_thread(): openblas_controller = ThreadpoolController().select(internal_api="openblas") diff --git a/threadpoolctl.py b/threadpoolctl.py index 10ba21df..756168cb 100644 --- a/threadpoolctl.py +++ b/threadpoolctl.py @@ -557,24 +557,26 @@ def set_num_threads(self, num_threads): # A shared process-wide thread pool, so don't want different Python # threads fighting over limits: this_thread = current_thread() - - match self._SHARED_STATE.controlling_thread: - case None: - # No thread controls OpenBLAS, so try to take control: - with self._SHARED_STATE.lock: - # Another thread might have changed it, so ensure that's - # not the case: - if self._SHARED_STATE.controlling_thread is None: - self._SHARED_STATE.controlling_thread = this_thread - # Controlling thread should now be set, to either this thread - # or another, so try again: - return self.set_num_threads(num_threads) - case controlling_thread if controlling_thread is this_thread: - # This thread controls OpenBLAS: - return self._real_set_num_threads(num_threads) - case _: - # Some other thread controlls OpenBLAS: - return + controlling_thread = self._SHARED_STATE.controlling_thread + if controlling_thread is None or not controlling_thread.is_alive(): + # No thread controls OpenBLAS, so try to take control: + with self._SHARED_STATE.lock: + # Another thread might have changed it, so ensure that's not + # the case: + if ( + self._SHARED_STATE.controlling_thread is None + or not self._SHARED_STATE.controlling_thread.is_alive() + ): + self._SHARED_STATE.controlling_thread = this_thread + # Controlling thread should now be set, to either this thread or + # another, so try again: + return self.set_num_threads(num_threads) + elif controlling_thread is this_thread: + # This thread controls OpenBLAS: + return self._real_set_num_threads(num_threads) + else: + # Some other thread controlls OpenBLAS: do nothing. + return def _real_set_num_threads(self, num_threads): # The OpenBLAS limiting API is process-wide, and we want current thread