Use joblib for parallelism in regress_out (#1695)

* Use joblib for parallism in regress_out * release note * fix link in release notes * Add todo for resource test
scverse · Mar 3, 2021 · d69aa18 · d69aa18
1 parent ae3f8b0
commit d69aa18
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 8 deletions.
diff --git a/docs/release-notes/1.7.2.rst b/docs/release-notes/1.7.2.rst
@@ -7,6 +7,7 @@
 
 - Fix :func:`scanpy.pl.paga_path` `TypeError` with recent versions of anndata :pr:`1047` :smaller:`P Angerer`
 - :func:`scanpy.logging.print_versions` now works when `python<3.8` :pr:`1691` :smaller:`I Virshup`
+- :func:`scanpy.pp.regress_out` now uses `joblib` as the parallel backend, and should stop oversubscribing threads :pr:`1694` :smaller:`I Virshup`
 
 .. rubric:: Deprecations
 

diff --git a/scanpy/preprocessing/_simple.py b/scanpy/preprocessing/_simple.py
@@ -662,15 +662,10 @@ def regress_out(
             regres = regressors
         tasks.append(tuple((data_chunk, regres, variable_is_categorical)))
 
-    if n_jobs > 1 and n_chunks > 1:
-        import multiprocessing
+    from joblib import Parallel, delayed
 
-        pool = multiprocessing.Pool(n_jobs)
-        res = pool.map_async(_regress_out_chunk, tasks).get(9999999)
-        pool.close()
-
-    else:
-        res = list(map(_regress_out_chunk, tasks))
+    # TODO: figure out how to test that this doesn't oversubscribe resources
+    res = Parallel(n_jobs=n_jobs)(delayed(_regress_out_chunk)(task) for task in tasks)
 
     # res is a list of vectors (each corresponding to a regressed gene column).
     # The transpose is needed to get the matrix in the shape needed