Repository: SpikeInterface/spikeinterface
Pull request #2829 (2024-05-10 12:48:53)

Add extra check in fix_job_kwargs
For functions that take in an `n_jobs` argument they all run `fix_job_kwargs`, but if the user gives something like 4.0 to mean 4 jobs they end up with (if they have 4 cores)

int(4.0 * 4.0) = 16

So instead we should cast 4.0 -> 4 and just make it 4 jobs no? 

What do you think @alejoe91 ?


Patch:
diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py
index db8d8f633..f1454d902 100644
--- a/src/spikeinterface/core/job_tools.py
+++ b/src/spikeinterface/core/job_tools.py
@@ -85,11 +85,18 @@ def fix_job_kwargs(runtime_job_kwargs):
 
     # 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):
+    assert isinstance(n_jobs, (float, np.integer, int)) and n_jobs != 0, "n_jobs must be a non-zero int or float"
+
+    # for a fraction we do fraction of total cores
+    if isinstance(n_jobs, float) and 0 < n_jobs <= 1:
         n_jobs = int(n_jobs * os.cpu_count())
+    # for negative numbers we count down from total cores (with -1 being all)
     elif n_jobs < 0:
-        n_jobs = os.cpu_count() + 1 + n_jobs
+        n_jobs = int(os.cpu_count() + 1 + n_jobs)
+    # otherwise we just take the value given
+    else:
+        n_jobs = int(n_jobs)
+
     job_kwargs["n_jobs"] = max(n_jobs, 1)
 
     return job_kwargs
diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py
index 1ba96defa..4a0c5f8ee 100644
--- a/src/spikeinterface/core/zarrextractors.py
+++ b/src/spikeinterface/core/zarrextractors.py
@@ -72,7 +72,7 @@ class ZarrRecordingExtractor(BaseRecording):
             time_kwargs = {}
             time_vector = self._root.get(f"times_seg{segment_index}", None)
             if time_vector is not None:
-                time_kwargs["time_vector"] = time_vector[:]
+                time_kwargs["time_vector"] = time_vector
             else:
                 if t_starts is None:
                     t_start = None
diff --git a/src/spikeinterface/generation/template_database.py b/src/spikeinterface/generation/template_database.py
index a9b5ef030..2bfcc3224 100644
--- a/src/spikeinterface/generation/template_database.py
+++ b/src/spikeinterface/generation/template_database.py
@@ -2,7 +2,7 @@ from spikeinterface.core.template import Templates
 import zarr
 
 
-def fetch_templates_from_database(dataset="test_templates.zarr") -> Templates:
+def fetch_templates_from_database(dataset="test_templates") -> Templates:
 
     s3_path = f"s3://spikeinterface-template-database/{dataset}/"
     zarr_group = zarr.open_consolidated(s3_path, storage_options={"anon": True})


Test patch:
diff --git a/src/spikeinterface/core/tests/test_job_tools.py b/src/spikeinterface/core/tests/test_job_tools.py
index 1bfe3a5e7..26ca8e7e7 100644
--- a/src/spikeinterface/core/tests/test_job_tools.py
+++ b/src/spikeinterface/core/tests/test_job_tools.py
@@ -180,10 +180,10 @@ def test_fix_job_kwargs():
     else:
         assert fixed_job_kwargs["n_jobs"] == 1
 
-    # test minimum n_jobs
-    job_kwargs = dict(n_jobs=0, progress_bar=False, chunk_duration="1s")
+    # test float value > 1 is cast to correct int
+    job_kwargs = dict(n_jobs=4.0, progress_bar=False, chunk_duration="1s")
     fixed_job_kwargs = fix_job_kwargs(job_kwargs)
-    assert fixed_job_kwargs["n_jobs"] == 1
+    assert fixed_job_kwargs["n_jobs"] == 4
 
     # test wrong keys
     with pytest.raises(AssertionError):