Merge pull request #152 from eurunuela/fix/n-jobs-parallel-scheduler

Wire n_jobs to Dask threaded scheduler in all voxel-wise solvers

Author: eurunuela

docs/conf.py

@@ -50,7 +50,6 @@
     "sphinx.ext.mathjax",
     "sphinx.ext.napoleon",
     "sphinx.ext.todo",
-    "sphinx_charts.charts",
     "sphinx_gallery.load_style",
     "sphinxcontrib.jquery",
 ]
@@ -75,8 +74,6 @@
 # nb_render_plugin = "plotly"
 
 
-
-
 # if LooseVersion(sphinx.__version__) < LooseVersion("1.4"):
 #     extensions.append("sphinx.ext.pngmath")
 # else:
@@ -192,9 +189,7 @@ def setup(app):
 # The following is used by sphinx.ext.linkcode to provide links to github
 linkcode_resolve = make_linkcode_resolve(
     "pySPFM",
-    "https://github.com/ParadigmFreeMapping/"
-    "pySPFM/blob/{revision}/"
-    "{package}/{path}#L{lineno}",
+    "https://github.com/ParadigmFreeMapping/pySPFM/blob/{revision}/{package}/{path}#L{lineno}",
 )
 
 # Example configuration for intersphinx: refer to the Python standard library.

pySPFM/decomposition.py

@@ -103,6 +103,33 @@ def _generate_hrf_matrix(self, n_scans):
         hrf_obj.generate_hrf(tr=self.tr, n_scans=n_scans)
         return hrf_obj.hrf_
 
+    def _dask_compute(self, futures):
+        """Run dask compute with the scheduler selected by ``n_jobs``.
+
+        Uses the synchronous scheduler when ``n_jobs=1`` (default, safe for
+        all environments) and the threaded scheduler otherwise.  Because the
+        per-voxel solvers (LARS, FISTA) call into numpy/scipy C extensions
+        that release the GIL, threads provide real parallelism here.
+
+        Parameters
+        ----------
+        futures : list of dask.delayed
+            Delayed objects to compute.
+
+        Returns
+        -------
+        results : list
+            Computed results in the same order as ``futures``.
+        """
+        if self.n_jobs == 0 or self.n_jobs < -1:
+            raise ValueError(
+                f"n_jobs must be -1 (use all CPUs) or a positive integer >= 1, got {self.n_jobs!r}"
+            )
+        if self.n_jobs == 1:
+            return compute(futures, scheduler="synchronous")[0]
+        num_workers = None if self.n_jobs == -1 else self.n_jobs
+        return compute(futures, scheduler="threads", num_workers=num_workers)[0]
+
     @abstractmethod
     def fit(self, X, y=None):
         """Fit the deconvolution model.
@@ -378,7 +405,7 @@ def _fit_lars(self, X, n_scans, n_voxels):
             )
             futures.append(fut)
 
-        results = compute(futures, scheduler="synchronous")[0]
+        results = self._dask_compute(futures)
 
         for vox_idx in range(n_voxels):
             self.coef_[:, vox_idx] = np.squeeze(results[vox_idx][0])
@@ -403,7 +430,7 @@ def _fit_fista(self, X, n_scans, n_voxels):
             )
             futures.append(fut)
 
-        results = compute(futures, scheduler="synchronous")[0]
+        results = self._dask_compute(futures)
 
         for vox_idx in range(n_voxels):
             self.coef_[:, vox_idx] = np.squeeze(results[vox_idx][0])
@@ -628,7 +655,7 @@ def fit(self, X, y=None):
                 )
                 futures.append(fut)
 
-            results = compute(futures, scheduler="synchronous")[0]
+            results = self._dask_compute(futures)
 
             for vox_idx in range(n_voxels):
                 self.coef_[:, vox_idx] = np.squeeze(results[vox_idx][0])
@@ -806,7 +833,7 @@ def fit(self, X, y=None):
             )
             futures.append(fut)
 
-        results = compute(futures, scheduler="synchronous")[0]
+        results = self._dask_compute(futures)
 
         for vox_idx in range(n_voxels):
             self.selection_frequency_[:, vox_idx] = np.squeeze(results[vox_idx])

pySPFM/tests/test_decomposition.py

@@ -484,3 +484,48 @@ def test_threshold_binary_selection(self, sample_data):
         # coef_ should only contain 0s and 1s
         unique_vals = np.unique(model.coef_)
         assert set(unique_vals).issubset({0.0, 1.0})
+
+
+class TestDaskCompute:
+    """Tests for _BaseDeconvolution._dask_compute scheduler selection."""
+
+    def test_n_jobs_1_uses_synchronous(self):
+        """n_jobs=1 must select the synchronous scheduler."""
+        from unittest.mock import patch
+
+        from pySPFM import SparseDeconvolution
+
+        model = SparseDeconvolution(tr=2.0, n_jobs=1)
+        with patch("pySPFM.decomposition.compute", return_value=([],)) as mock_compute:
+            model._dask_compute([])
+        mock_compute.assert_called_once_with([], scheduler="synchronous")
+
+    def test_n_jobs_gt1_uses_threaded(self):
+        """n_jobs=2 must select the threaded scheduler with num_workers=2."""
+        from unittest.mock import patch
+
+        from pySPFM import SparseDeconvolution
+
+        model = SparseDeconvolution(tr=2.0, n_jobs=2)
+        with patch("pySPFM.decomposition.compute", return_value=([],)) as mock_compute:
+            model._dask_compute([])
+        mock_compute.assert_called_once_with([], scheduler="threads", num_workers=2)
+
+    def test_n_jobs_minus1_uses_all_cpus(self):
+        """n_jobs=-1 must select the threaded scheduler with num_workers=None."""
+        from unittest.mock import patch
+
+        from pySPFM import SparseDeconvolution
+
+        model = SparseDeconvolution(tr=2.0, n_jobs=-1)
+        with patch("pySPFM.decomposition.compute", return_value=([],)) as mock_compute:
+            model._dask_compute([])
+        mock_compute.assert_called_once_with([], scheduler="threads", num_workers=None)
+
+    def test_invalid_n_jobs_raises(self):
+        """n_jobs=0 must raise ValueError with a message mentioning n_jobs."""
+        from pySPFM import SparseDeconvolution
+
+        model = SparseDeconvolution(tr=2.0, n_jobs=0)
+        with pytest.raises(ValueError, match="n_jobs"):
+            model._dask_compute([])

pySPFM/tests/test_integration.py

@@ -8,9 +8,9 @@
 import shutil
 import tarfile
 from gzip import GzipFile
+from importlib.resources import files
 
 import pytest
-from pkg_resources import resource_filename
 
 
 def extract_test_data(tarball_path, outpath):
@@ -292,7 +292,7 @@ def test_integration_auc_to_estimates(
     assert ret.success
 
     # compare the generated output files
-    fn = resource_filename("pySPFM", "tests/data/auc_to_estimates_outputs.txt")
+    fn = str(files("pySPFM") / "tests" / "data" / "auc_to_estimates_outputs.txt")
     check_integration_outputs(fn, out_dir, "auc_to_estimates")
 
     ############################
@@ -329,7 +329,7 @@ def test_integration_auc_to_estimates(
     assert ret.success
 
     # compare the generated output files
-    fn = resource_filename("pySPFM", "tests/data/auc_to_estimates_outputs.txt")
+    fn = str(files("pySPFM") / "tests" / "data" / "auc_to_estimates_outputs.txt")
     check_integration_outputs(fn, out_dir, "auc_to_estimates")
 
     ############################
@@ -362,7 +362,7 @@ def test_integration_auc_to_estimates(
     assert ret.success
 
     # compare the generated output files
-    fn = resource_filename("pySPFM", "tests/data/auc_to_estimates_outputs.txt")
+    fn = str(files("pySPFM") / "tests" / "data" / "auc_to_estimates_outputs.txt")
     check_integration_outputs(fn, out_dir, "auc_to_estimates")
 
     ############################
@@ -395,7 +395,7 @@ def test_integration_auc_to_estimates(
     assert ret.success
 
     # compare the generated output files
-    fn = resource_filename("pySPFM", "tests/data/auc_to_estimates_outputs.txt")
+    fn = str(files("pySPFM") / "tests" / "data" / "auc_to_estimates_outputs.txt")
     check_integration_outputs(fn, out_dir, "auc_to_estimates")
 
 

pyproject.toml

@@ -63,7 +63,6 @@ doc = [
     "sphinx_design",
     "sphinx_gallery",
     "sphinx-book-theme",
-    "sphinx_charts",
     "sphinxcontrib-jquery",
 ]