[MNT] Add threads limit for numba and TF (#1604)

* Add threads limit for numba and TF

* pre-commit fix

* fix pre-commit

* Fix precommit

* Add env variable setting and move imports

* Update numba import

---------

Co-authored-by: Antoine GUILLAUME <[email protected]>

aeon/utils/numba/general.py

@@ -95,7 +95,7 @@ def generate_new_default_njit_func(
         base_func_py = base_func
     else:
         raise TypeError(
-            "Expected base_func to be of callable or CPUDispatcher type (numba "
+            "Expected base_func to be of type callable or CPUDispatcher type (numba "
             f"function), but got {type(base_func)}"
         )
     signature = inspect.signature(base_func_py)

conftest.py

@@ -10,8 +10,6 @@
 
 __maintainer__ = []
 
-from aeon.testing import test_config
-
 
 def pytest_addoption(parser):
     """Pytest command line parser options adder."""
@@ -28,5 +26,35 @@ def pytest_addoption(parser):
 
 def pytest_configure(config):
     """Pytest configuration preamble."""
+    import os
+
+    # Must be called before any numpy imports
+    os.environ["MKL_NUM_THREADS"] = "1"
+    os.environ["NUMEXPR_NUM_THREADS"] = "1"
+    os.environ["OMP_NUM_THREADS"] = "1"
+    os.environ["OPENBLAS_NUM_THREADS"] = "1"
+    os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
+
+    import numba
+
+    from aeon.testing import test_config
+    from aeon.utils.validation._dependencies import _check_soft_dependencies
+
+    numba.set_num_threads(1)
+
+    if _check_soft_dependencies("tensorflow", severity="none"):
+        from tensorflow.config.threading import (
+            set_inter_op_parallelism_threads,
+            set_intra_op_parallelism_threads,
+        )
+
+        set_inter_op_parallelism_threads(1)
+        set_intra_op_parallelism_threads(1)
+
+    if _check_soft_dependencies("torch", severity="none"):
+        import torch
+
+        torch.set_num_threads(1)
+
     if config.getoption("--prtesting") in [True, "True", "true"]:
         test_config.PR_TESTING = True