DM-55482: Improve coverage of disable_implicit_threading

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.

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."
+        )

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()