MAINT: Work around joblib bug (#11765)

mne-tools · Jun 30, 2023 · b92c88d · b92c88d
1 parent 3c4a6f6
commit b92c88d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/mne/parallel.py b/mne/parallel.py
@@ -105,8 +105,12 @@ def parallel_func(
             max_nbytes = None  # disable memmaping
         kwargs["temp_folder"] = cache_dir
         kwargs["max_nbytes"] = max_nbytes
-        parallel = Parallel(n_jobs, **kwargs)
+        n_jobs_orig = n_jobs
+        if n_jobs is not None:  # https://github.com/joblib/joblib/issues/1473
+            kwargs["n_jobs"] = n_jobs
+        parallel = Parallel(**kwargs)
         n_jobs = _check_n_jobs(parallel.n_jobs)
+        logger.debug(f"Got {n_jobs} parallel jobs after requesting {n_jobs_orig}")
         if max_jobs is not None:
             n_jobs = min(n_jobs, max(_ensure_int(max_jobs, "max_jobs"), 1))
         my_func = delayed(func)

diff --git a/mne/tests/test_parallel.py b/mne/tests/test_parallel.py
@@ -34,7 +34,12 @@ def fun(x):
     if isinstance(n_jobs, str):
         backend, n_jobs = n_jobs.split()
         n_jobs = want_jobs = int(n_jobs)
-        ctx = joblib.parallel_backend(backend, n_jobs)
+        try:
+            func = joblib.parallel_config
+        except AttributeError:
+            # joblib < 1.3
+            func = joblib.parallel_backend
+        ctx = func(backend, n_jobs=n_jobs)
         n_jobs = None
     else:
         ctx = nullcontext()
@@ -43,5 +48,5 @@ def fun(x):
         else:
             want_jobs = 1
     with ctx:
-        parallel, p_fun, got_jobs = parallel_func(fun, n_jobs)
+        parallel, p_fun, got_jobs = parallel_func(fun, n_jobs, verbose="debug")
     assert got_jobs == want_jobs