Fixed parallelisation of automated ICA.

Author: cgohil8

osl_dynamics/meeg/parallel.py

@@ -14,13 +14,36 @@
     "MKL_NUM_THREADS",
     "OPENBLAS_NUM_THREADS",
     "VECLIB_MAXIMUM_THREADS",
+    "NUMEXPR_MAX_THREADS",
 ]
 
 
+def _limit_onnx_threads():
+    """Patch ONNX Runtime to use 1 thread per session."""
+    try:
+        import onnxruntime as ort
+
+        _OriginalSession = ort.InferenceSession
+
+        class _SingleThreadSession(_OriginalSession):
+            def __init__(self, *args, **kwargs):
+                if "sess_options" not in kwargs or kwargs["sess_options"] is None:
+                    opts = ort.SessionOptions()
+                    opts.intra_op_num_threads = 1
+                    opts.inter_op_num_threads = 1
+                    kwargs["sess_options"] = opts
+                super().__init__(*args, **kwargs)
+
+        ort.InferenceSession = _SingleThreadSession
+    except ImportError:
+        pass
+
+
 def _worker(
     args: Tuple[Callable, str, Any, Path, dict],
 ) -> Tuple[str, bool]:
     """Wrapper that handles logging and error catching for a single item."""
+    _limit_onnx_threads()
     func, id, item, log_dir, kwargs = args
     with MEEGSessionLogger(id, log_dir) as logger:
         try: