diff --git a/doc/changes/DM-55482.bugfix.rst b/doc/changes/DM-55482.bugfix.rst new file mode 100644 index 00000000..4b1c6be3 --- /dev/null +++ b/doc/changes/DM-55482.bugfix.rst @@ -0,0 +1,5 @@ +Improved the coverage of ``disable_implicit_threading``. +The thread-related environment variables now also include ``NUMBA_NUM_THREADS``, ``ARROW_IO_THREADS``, ``VECLIB_MAXIMUM_THREADS``, ``RAYON_NUM_THREADS``, ``POLARS_MAX_THREADS``, and ``BLIS_NUM_THREADS``. +Existing ``pyarrow`` thread pools are now resized explicitly since they do not react to environment variables after creation. +An already-imported ``numba`` is now limited at runtime since it reads its environment variable only at import time. +A warning is now logged if ``threadpoolctl`` is not installed since thread pools of already-loaded libraries can then no longer be limited. diff --git a/python/lsst/utils/threads.py b/python/lsst/utils/threads.py index ab0f1acd..993c2e4a 100644 --- a/python/lsst/utils/threads.py +++ b/python/lsst/utils/threads.py @@ -15,13 +15,17 @@ __all__ = ["disable_implicit_threading", "set_thread_envvars"] +import logging import os +import sys try: from threadpoolctl import threadpool_limits except ImportError: threadpool_limits = None +_LOG = logging.getLogger(__name__) + def set_thread_envvars(num_threads: int = 1, override: bool = False) -> None: """Set common threading environment variables to the given value. @@ -44,6 +48,12 @@ def set_thread_envvars(num_threads: int = 1, override: bool = False) -> None: "MPI_NUM_THREADS", "NUMEXPR_NUM_THREADS", "NUMEXPR_MAX_THREADS", + "NUMBA_NUM_THREADS", + "ARROW_IO_THREADS", + "VECLIB_MAXIMUM_THREADS", + "RAYON_NUM_THREADS", + "POLARS_MAX_THREADS", + "BLIS_NUM_THREADS", ) for var in envvars: @@ -63,8 +73,8 @@ def disable_implicit_threading() -> None: Notes ----- Explicitly limits the number of threads allowed to be used by ``numexpr`` - and attempts to limit the number of threads in all APIs supported by - the ``threadpoolctl`` package. + and ``pyarrow`` (if already imported) and attempts to limit the number of + threads in all APIs supported by the ``threadpoolctl`` package. """ # Force one thread and force override. set_thread_envvars(1, True) @@ -79,6 +89,24 @@ def disable_implicit_threading() -> None: else: numexpr.utils.set_num_threads(1) + # pyarrow sizes its thread pools from the environment only when each pool + # is first used, so pools that may already exist must be resized + # explicitly. If pyarrow has not been imported the environment variables + # set above are sufficient and there is no need to pay for an import here. + if (pyarrow := sys.modules.get("pyarrow")) is not None: + pyarrow.set_cpu_count(1) + pyarrow.set_io_thread_count(1) + + # numba reads its environment variable only at import time so an + # already-imported numba must be limited at runtime. + if (numba := sys.modules.get("numba")) is not None: + numba.set_num_threads(1) + # Try to set threads for openblas and openmp if threadpool_limits is not None: threadpool_limits(limits=1) + else: + _LOG.warning( + "threadpoolctl is not installed: thread pools of already-loaded libraries cannot be " + "limited and implicit threading may remain enabled." + ) diff --git a/tests/test_threads.py b/tests/test_threads.py index c3d5f983..9ef38a71 100644 --- a/tests/test_threads.py +++ b/tests/test_threads.py @@ -21,6 +21,7 @@ import os import unittest +import unittest.mock from lsst.utils.threads import disable_implicit_threading, set_thread_envvars @@ -32,6 +33,14 @@ import threadpoolctl except ImportError: threadpoolctl = None +try: + import pyarrow +except ImportError: + pyarrow = None +try: + import numba +except ImportError: + numba = None class ThreadsTestCase(unittest.TestCase): @@ -55,6 +64,54 @@ def testDisable(self): for api in info: self.assertEqual(api["num_threads"], 1, f"API: {api}") + def testEnvVarsForEnvOnlyLibraries(self): + """Libraries that are only controllable via environment variables + must be included in the list of variables that are set. + """ + set_thread_envvars(2, override=True) + for var in ( + "NUMBA_NUM_THREADS", + "ARROW_IO_THREADS", + "VECLIB_MAXIMUM_THREADS", + "RAYON_NUM_THREADS", + "POLARS_MAX_THREADS", + "BLIS_NUM_THREADS", + ): + self.assertEqual(os.environ.get(var), "2", f"Variable: {var}") + + @unittest.skipIf(pyarrow is None, "pyarrow is not available") + def testDisablePyarrow(self): + """Already-created pyarrow thread pools must be resized since they + do not react to environment variables after creation. + """ + pyarrow.set_cpu_count(4) + pyarrow.set_io_thread_count(4) + disable_implicit_threading() + self.assertEqual(pyarrow.cpu_count(), 1) + self.assertEqual(pyarrow.io_thread_count(), 1) + + @unittest.skipIf(numba is None, "numba is not available") + def testDisableNumba(self): + """An already-imported numba must be limited at runtime since it + reads its environment variable only at import time. + """ + if numba.config.NUMBA_NUM_THREADS > 1: + numba.set_num_threads(2) + disable_implicit_threading() + self.assertEqual(numba.get_num_threads(), 1) + + @unittest.skipIf(threadpoolctl is None, "threadpoolctl is not available") + def testMissingThreadpoolctlWarning(self): + """Absence of threadpoolctl silently weakens the thread disabling + so it must be reported. + """ + with ( + unittest.mock.patch("lsst.utils.threads.threadpool_limits", None), + self.assertLogs("lsst.utils.threads", "WARNING") as cm, + ): + disable_implicit_threading() + self.assertIn("threadpoolctl", "\n".join(cm.output)) + if __name__ == "__main__": unittest.main()