From 66117b3cf8f2ea811f4a902d48f86a4ecf7416f0 Mon Sep 17 00:00:00 2001 From: Tim Jenness Date: Wed, 15 Jul 2026 15:30:27 -0700 Subject: [PATCH 1/2] Improve coverage of disable_implicit_threading Add NUMBA_NUM_THREADS, ARROW_IO_THREADS, VECLIB_MAXIMUM_THREADS, RAYON_NUM_THREADS, POLARS_MAX_THREADS, and BLIS_NUM_THREADS to the environment variables set by set_thread_envvars. Resize existing pyarrow thread pools explicitly since they are sized from the environment only when each pool is first used and do not react to environment variables after creation. Log a warning when threadpoolctl is not installed since thread pools of already-loaded libraries can then no longer be limited. Co-Authored-By: Claude Fable 5 --- doc/changes/DM-55482.bugfix.rst | 4 +++ python/lsst/utils/threads.py | 27 +++++++++++++++++++-- tests/test_threads.py | 43 +++++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) create mode 100644 doc/changes/DM-55482.bugfix.rst diff --git a/doc/changes/DM-55482.bugfix.rst b/doc/changes/DM-55482.bugfix.rst new file mode 100644 index 00000000..33dcad4f --- /dev/null +++ b/doc/changes/DM-55482.bugfix.rst @@ -0,0 +1,4 @@ +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. +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..e4591cd0 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,19 @@ 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) + # 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..fba9ea90 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,10 @@ import threadpoolctl except ImportError: threadpoolctl = None +try: + import pyarrow +except ImportError: + pyarrow = None class ThreadsTestCase(unittest.TestCase): @@ -55,6 +60,44 @@ 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(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() From ca70172e816e8f04a3b707d6787607f81500f8ee Mon Sep 17 00:00:00 2001 From: Tim Jenness Date: Wed, 15 Jul 2026 17:16:05 -0700 Subject: [PATCH 2/2] Limit already-imported numba at runtime numba reads NUMBA_NUM_THREADS only at import time, so setting the environment variable in disable_implicit_threading has no effect on a numba that is already imported. Call numba.set_num_threads(1) as well. Co-Authored-By: Claude Fable 5 --- doc/changes/DM-55482.bugfix.rst | 1 + python/lsst/utils/threads.py | 5 +++++ tests/test_threads.py | 14 ++++++++++++++ 3 files changed, 20 insertions(+) diff --git a/doc/changes/DM-55482.bugfix.rst b/doc/changes/DM-55482.bugfix.rst index 33dcad4f..4b1c6be3 100644 --- a/doc/changes/DM-55482.bugfix.rst +++ b/doc/changes/DM-55482.bugfix.rst @@ -1,4 +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 e4591cd0..993c2e4a 100644 --- a/python/lsst/utils/threads.py +++ b/python/lsst/utils/threads.py @@ -97,6 +97,11 @@ def disable_implicit_threading() -> 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) diff --git a/tests/test_threads.py b/tests/test_threads.py index fba9ea90..9ef38a71 100644 --- a/tests/test_threads.py +++ b/tests/test_threads.py @@ -37,6 +37,10 @@ import pyarrow except ImportError: pyarrow = None +try: + import numba +except ImportError: + numba = None class ThreadsTestCase(unittest.TestCase): @@ -86,6 +90,16 @@ def testDisablePyarrow(self): 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