Repository: SpikeInterface/spikeinterface
Pull request #1207 (2023-01-19 09:07:56)

[feat-sorters] Set n_jobs to None for sorters to propagate global job kwargs
With this PR, the sorters use (by default) the globally set n_jobs


Patch:
diff --git a/spikeinterface/core/globals.py b/spikeinterface/core/globals.py
index 36643fbdc..21c8eae40 100644
--- a/spikeinterface/core/globals.py
+++ b/spikeinterface/core/globals.py
@@ -122,9 +122,7 @@ def set_global_job_kwargs(**job_kwargs):
     for k in job_kwargs:
         assert k in job_keys, (f"{k} is not a valid job keyword argument. "
                                f"Available keyword arguments are: {list(job_keys)}")
-    current_global_kwargs = get_global_job_kwargs()
-    current_global_kwargs.update(job_kwargs)
-    global_job_kwargs = current_global_kwargs
+    global_job_kwargs.update(job_kwargs)
     global global_job_kwargs_set
     global_job_kwargs_set = True
 
@@ -134,7 +132,7 @@ def reset_global_job_kwargs():
     Reset the global job kwargs.
     """
     global global_job_kwargs
-    global_job_kwargs = dict()
+    global_job_kwargs = dict(n_jobs=1, chunk_duration="1s", progress_bar=True)
 
 
 def is_set_global_job_kwargs_set():
diff --git a/spikeinterface/core/job_tools.py b/spikeinterface/core/job_tools.py
index 972a56f56..372bb8885 100644
--- a/spikeinterface/core/job_tools.py
+++ b/spikeinterface/core/job_tools.py
@@ -42,22 +42,26 @@ job_keys = ('n_jobs', 'total_memory', 'chunk_size', 'chunk_memory', 'chunk_durat
 def fix_job_kwargs(runtime_job_kwargs):
     from .globals import get_global_job_kwargs
     job_kwargs = get_global_job_kwargs()
-    
+
     for k in runtime_job_kwargs:
         assert k in job_keys, (f"{k} is not a valid job keyword argument. "
                                f"Available keyword arguments are: {list(job_keys)}")
-    job_kwargs.update(runtime_job_kwargs)
-
-    # if n_jobs is -1, set to os.cpu_count()
-    if "n_jobs" in job_kwargs:
-        assert isinstance(job_kwargs["n_jobs"], (float, np.integer, int))
-        if isinstance(job_kwargs["n_jobs"], float):
-            n_jobs = int(job_kwargs["n_jobs"] * os.cpu_count())
-        elif job_kwargs["n_jobs"] < 0:            
-            n_jobs = os.cpu_count() + 1 + job_kwargs["n_jobs"]
-        else:
-            n_jobs = job_kwargs["n_jobs"]
-        job_kwargs["n_jobs"] = max(n_jobs, 1)
+    # remove None
+    runtime_job_kwargs_exclude_none = runtime_job_kwargs.copy()
+    for job_key, job_value in runtime_job_kwargs.items():
+        if job_value is None:
+            del runtime_job_kwargs_exclude_none[job_key]
+    job_kwargs.update(runtime_job_kwargs_exclude_none)
+
+    # if n_jobs is -1, set to os.cpu_count() (n_jobs is always in global job_kwargs)
+    n_jobs = job_kwargs['n_jobs']
+    assert isinstance(n_jobs, (float, np.integer, int))
+    if isinstance(n_jobs, float):
+        n_jobs = int(n_jobs * os.cpu_count())
+    elif n_jobs < 0:
+        n_jobs = os.cpu_count() + 1 + n_jobs
+    job_kwargs["n_jobs"] = max(n_jobs, 1)
+
     return job_kwargs
 
 
diff --git a/spikeinterface/postprocessing/spike_amplitudes.py b/spikeinterface/postprocessing/spike_amplitudes.py
index 094fe492e..f0e4e5a5c 100644
--- a/spikeinterface/postprocessing/spike_amplitudes.py
+++ b/spikeinterface/postprocessing/spike_amplitudes.py
@@ -106,18 +106,17 @@ class SpikeAmplitudesCalculator(BaseWaveformExtractorExtension):
             as a dict with units as key and spike amplitudes as values.
         """
         we = self.waveform_extractor
-        recording = we.recording
         sorting = we.sorting
         all_spikes = sorting.get_all_spike_trains(outputs='unit_index')
 
         if outputs == 'concatenated':
             amplitudes = []
-            for segment_index in range(recording.get_num_segments()):
+            for segment_index in range(we.get_num_segments()):
                 amplitudes.append(self._extension_data[f'amplitude_segment_{segment_index}'])
             return amplitudes
         elif outputs == 'by_unit':
             amplitudes_by_unit = []
-            for segment_index in range(recording.get_num_segments()):
+            for segment_index in range(we.get_num_segments()):
                 amplitudes_by_unit.append({})
                 for unit_index, unit_id in enumerate(sorting.unit_ids):
                     _, spike_labels = all_spikes[segment_index]
diff --git a/spikeinterface/sorters/basesorter.py b/spikeinterface/sorters/basesorter.py
index 78b3c0ec3..7ec47549e 100644
--- a/spikeinterface/sorters/basesorter.py
+++ b/spikeinterface/sorters/basesorter.py
@@ -20,21 +20,16 @@ from spikeinterface.core.job_tools import fix_job_kwargs, split_job_kwargs
 from .utils import SpikeSortingError, ShellScript
 
 
-default_job_kwargs = {"n_jobs": 1, 
-                      "total_memory": None,
-                      "chunk_size": None,
-                      "chunk_memory": None,
-                      "chunk_duration": "1s",
-                      "progress_bar": True}
+default_job_kwargs = {"n_jobs": -1}
 
 default_job_kwargs_description = {
-    "n_jobs": "Number of jobs (when saving ti binary) - default 1", 
-    "chunk_size": "Number of samples per chunk (when saving ti binary) - default None",
-    "chunk_memory": "Memory usage for each job (e.g. '100M', '1G') (when saving to binary) - default None",
-    "total_memory": "Total memory usage (e.g. '500M', '2G') (when saving to binary) - default None",
+    "n_jobs": "Number of jobs (when saving ti binary) - default -1 (all cores)",
+    "chunk_size": "Number of samples per chunk (when saving ti binary) - default global",
+    "chunk_memory": "Memory usage for each job (e.g. '100M', '1G') (when saving to binary) - default global",
+    "total_memory": "Total memory usage (e.g. '500M', '2G') (when saving to binary) - default global",
     "chunk_duration": "Chunk duration in s if float or with units if str (e.g. '1s', '500ms') (when saving to binary)" \
-                      " - default '1s'",
-    "progress_bar": "If True, progress bar is shown (when saving to binary) - default True"}
+                      " - default global",
+    "progress_bar": "If True, progress bar is shown (when saving to binary) - default global"}
 
 
 class BaseSorter:
@@ -147,7 +142,8 @@ class BaseSorter:
     def default_params(cls):
         p = copy.deepcopy(cls._default_params)
         if cls.requires_binary_data:
-            p.update(default_job_kwargs)
+            job_kwargs = fix_job_kwargs(default_job_kwargs)
+            p.update(job_kwargs)
         return p
 
     @classmethod
diff --git a/spikeinterface/sorters/si_based_sorters/spyking_circus2.py b/spikeinterface/sorters/si_based_sorters/spyking_circus2.py
index 28acfc087..496291d1c 100644
--- a/spikeinterface/sorters/si_based_sorters/spyking_circus2.py
+++ b/spikeinterface/sorters/si_based_sorters/spyking_circus2.py
@@ -32,7 +32,7 @@ class Spykingcircus2Sorter(ComponentsBasedSorter):
         'registration' : {},
         'apply_preprocessing': True,
         'shared_memory' : False,
-        'job_kwargs' : {'n_jobs' : -1, 'chunk_duration' : '1s', 'verbose' : False}
+        'job_kwargs' : {}
     }
 
     @classmethod
diff --git a/spikeinterface/sorters/si_based_sorters/tridesclous2.py b/spikeinterface/sorters/si_based_sorters/tridesclous2.py
index 237803d6d..b6a72b453 100644
--- a/spikeinterface/sorters/si_based_sorters/tridesclous2.py
+++ b/spikeinterface/sorters/si_based_sorters/tridesclous2.py
@@ -23,7 +23,7 @@ class Tridesclous2Sorter(ComponentsBasedSorter):
         'selection' : {'n_peaks_per_channel' : 5000, 'min_n_peaks' : 20000},
         'localization' : {'max_distance_um':1000, 'optimizer': 'minimize_with_log_penality'},
         'matching':  {'peak_shift_ms':  0.2,},
-        'job_kwargs' : {'n_jobs' : -1, 'chunk_duration' : '1s', 'progress_bar': True}
+        'job_kwargs' : {}
     }
 
     @classmethod


Test patch:
diff --git a/spikeinterface/core/tests/test_globals.py b/spikeinterface/core/tests/test_globals.py
index 5a5f0e1c8..631583711 100644
--- a/spikeinterface/core/tests/test_globals.py
+++ b/spikeinterface/core/tests/test_globals.py
@@ -47,6 +47,11 @@ def test_global_job_kwargs():
     assert job_kwargs_split['n_jobs'] == new_job_kwargs['n_jobs']
     assert job_kwargs_split['chunk_duration'] == job_kwargs['chunk_duration']
     assert job_kwargs_split['progress_bar'] == job_kwargs['progress_bar']
+    # test that None values do not change existing global kwargs
+    none_job_kwargs = dict(n_jobs=None, progress_bar=None, chunk_duration=None)
+    job_kwargs_split = fix_job_kwargs(none_job_kwargs)
+    assert job_kwargs_split['chunk_duration'] == job_kwargs['chunk_duration']
+    assert job_kwargs_split['progress_bar'] == job_kwargs['progress_bar']
     reset_global_job_kwargs()