Add num_chunk_per_batch to job kwargs and use it for zarr v3 shard writes

conftest.py

@@ -32,7 +32,7 @@ def pytest_addoption(parser):
 
 def pytest_configure(config):
     mp_context = config.getoption("--mp-context")
-    if mp_context is not None:
+    if mp_context is not None and mp_context.lower() != "none":
         import multiprocessing
         import spikeinterface.core.globals as si_globals
 
@@ -50,7 +50,6 @@ def pytest_collection_modifyitems(config, items):
     rootdir = Path(config.rootdir)
     modules_location = rootdir / "src" / "spikeinterface"
     for item in items:
-        # TODO: remove once writing to zarr is supported with zarr>=3
         if config.getoption("--mp-context") is not None and item.name == "test_global_job_kwargs":
             item.add_marker(pytest.mark.skip(reason="--mp-context changes the default job kwargs"))
 

src/spikeinterface/core/globals.py

@@ -102,7 +102,13 @@ def is_set_global_dataset_folder() -> bool:
 
 ########################################
 _default_job_kwargs = dict(
-    pool_engine="process", n_jobs=1, chunk_duration="1s", progress_bar=True, mp_context=None, max_threads_per_worker=1
+    pool_engine="process",
+    n_jobs=1,
+    chunk_duration="1s",
+    progress_bar=True,
+    mp_context=None,
+    max_threads_per_worker=1,
+    num_chunk_per_batch=None,
 )
 
 global global_job_kwargs

src/spikeinterface/core/job_tools.py

@@ -65,6 +65,7 @@
     "progress_bar",
     "mp_context",
     "max_threads_per_worker",
+    "num_chunk_per_batch",
 )
 
 # theses key are the same and should not be in th final dict
@@ -225,12 +226,15 @@ def divide_segment_into_chunks(num_frames, chunk_size):
     return chunks
 
 
-def divide_time_series_into_chunks(recording, chunk_size):
+def divide_time_series_into_chunks(recording, chunk_size, num_chunk_per_batch=None):
     slices = []
     for segment_index in range(recording.get_num_segments()):
         num_frames = recording.get_num_samples(segment_index)
-        chunks = divide_segment_into_chunks(num_frames, chunk_size)
-        slices.extend([(segment_index, frame_start, frame_stop) for frame_start, frame_stop in chunks])
+        chunks = [(segment_index, s, e) for s, e in divide_segment_into_chunks(num_frames, chunk_size)]
+        if num_chunk_per_batch is None:
+            slices.extend(chunks)
+        else:
+            slices.extend(chunks[i : i + num_chunk_per_batch] for i in range(0, len(chunks), num_chunk_per_batch))
     return slices
 
 
@@ -383,6 +387,12 @@ class TimeSeriesChunkExecutor:
         Size of each chunk in number of samples. If "total_memory" or "chunk_memory" are used, it is ignored.
     chunk_duration : str or float or None
         Chunk duration in s if float or with units if str (e.g. "1s", "500ms")
+    num_chunk_per_batch : int or None, default: None
+        Number of chunks to process per batch.
+        Chunks (slice of traces) are distributed over a pool of n_jobs workers one by one.
+        Optionally each worker can receive a batch of continuous chunks to process at once.
+        This is useful in the context of zarr v3 shards concept.
+        Internally, instead of  sending slices to each worker it sends a list of slices.
     mp_context : "fork" | "spawn" | None, default: None
         "fork" or "spawn". If None, the context is taken by the recording.get_preferred_mp_context().
         "fork" is only safely available on LINUX systems.
@@ -415,6 +425,7 @@ def __init__(
         chunk_size=None,
         chunk_memory=None,
         chunk_duration=None,
+        num_chunk_per_batch=None,
         mp_context=None,
         job_name="",
         max_threads_per_worker=1,
@@ -463,6 +474,7 @@ def __init__(
             chunk_duration=chunk_duration,
             n_jobs=self.n_jobs,
         )
+        self.num_chunk_per_batch = num_chunk_per_batch
         self.job_name = job_name
         self.max_threads_per_worker = max_threads_per_worker
 
@@ -517,7 +529,7 @@ def run(self, slices=None):
 
         if slices is None:
             # TODO: rename
-            slices = divide_time_series_into_chunks(self.time_series, self.chunk_size)
+            slices = divide_time_series_into_chunks(self.time_series, self.chunk_size, self.num_chunk_per_batch)
 
         if self.handle_returns:
             returns = []
@@ -537,12 +549,14 @@ def run(self, slices=None):
             if self.need_worker_index:
                 worker_dict["worker_index"] = worker_index
 
-            for segment_index, frame_start, frame_stop in slices:
-                res = self.func(segment_index, frame_start, frame_stop, worker_dict)
-                if self.handle_returns:
-                    returns.append(res)
-                if self.gather_func is not None:
-                    self.gather_func(res)
+            for slice_ in slices:
+                # with batches, each item is a list of (segment_index, frame_start, frame_stop)
+                for segment_index, frame_start, frame_stop in [slice_] if self.num_chunk_per_batch is None else slice_:
+                    res = self.func(segment_index, frame_start, frame_stop, worker_dict)
+                    if self.handle_returns:
+                        returns.append(res)
+                    if self.gather_func is not None:
+                        self.gather_func(res)
 
         else:
             n_jobs = min(self.n_jobs, len(slices))
@@ -583,10 +597,17 @@ def run(self, slices=None):
                         )
 
                     for res in results:
-                        if self.handle_returns:
-                            returns.append(res)
-                        if self.gather_func is not None:
-                            self.gather_func(res)
+                        if self.num_chunk_per_batch is None:
+                            if self.handle_returns:
+                                returns.append(res)
+                            if self.gather_func is not None:
+                                self.gather_func(res)
+                        else:
+                            for res2 in res:
+                                if self.handle_returns:
+                                    returns.append(res2)
+                                if self.gather_func is not None:
+                                    self.gather_func(res2)
 
             elif self.pool_engine == "thread":
                 # this is need to create a per worker local dict where the initializer will push the func wrapper
@@ -624,10 +645,18 @@ def run(self, slices=None):
                     for res in results:
                         if self.progress_bar:
                             pbar.update(1)
-                        if self.handle_returns:
-                            returns.append(res)
-                        if self.gather_func is not None:
-                            self.gather_func(res)
+                        if self.num_chunk_per_batch is None:
+                            if self.handle_returns:
+                                returns.append(res)
+                            if self.gather_func is not None:
+                                self.gather_func(res)
+                        else:
+                            for res2 in res:
+                                if self.handle_returns:
+                                    returns.append(res2)
+                                if self.gather_func is not None:
+                                    self.gather_func(res2)
+
                 if self.progress_bar:
                     pbar.close()
                     del pbar
@@ -651,12 +680,27 @@ def __init__(self, func, worker_dict, max_threads_per_worker):
         self.max_threads_per_worker = max_threads_per_worker
 
     def __call__(self, args):
-        segment_index, start_frame, end_frame = args
-        if self.max_threads_per_worker is None:
-            return self.func(segment_index, start_frame, end_frame, self.worker_dict)
-        else:
-            with threadpool_limits(limits=self.max_threads_per_worker):
+        if isinstance(args[0], tuple):
+            # this is the case when num_chunk_per_batch is not None
+            res = []
+            for arg in args:
+                segment_index, start_frame, end_frame = arg
+                if self.max_threads_per_worker is None:
+                    res.append(self.func(segment_index, start_frame, end_frame, self.worker_dict))
+                else:
+                    with threadpool_limits(limits=self.max_threads_per_worker):
+                        res.append(self.func(segment_index, start_frame, end_frame, self.worker_dict))
+            return res
+        elif isinstance(args, tuple):
+            # this is the case when num_chunk_per_batch is None -> legacy mode
+            segment_index, start_frame, end_frame = args
+            if self.max_threads_per_worker is None:
                 return self.func(segment_index, start_frame, end_frame, self.worker_dict)
+            else:
+                with threadpool_limits(limits=self.max_threads_per_worker):
+                    return self.func(segment_index, start_frame, end_frame, self.worker_dict)
+        else:
+            raise ValueError("args must be a tuple or a list of tuples")
 
 
 # see

src/spikeinterface/core/tests/test_globals.py

@@ -38,6 +38,7 @@ def test_global_job_kwargs():
         progress_bar=True,
         mp_context=None,
         max_threads_per_worker=1,
+        num_chunk_per_batch=1,
     )
     global_job_kwargs = get_global_job_kwargs()
 
@@ -48,6 +49,7 @@ def test_global_job_kwargs():
         progress_bar=True,
         mp_context=None,
         max_threads_per_worker=1,
+        num_chunk_per_batch=None,
     )
     set_global_job_kwargs(**job_kwargs)
     assert get_global_job_kwargs() == job_kwargs
@@ -63,6 +65,7 @@ def test_global_job_kwargs():
         progress_bar=True,
         mp_context=None,
         max_threads_per_worker=1,
+        num_chunk_per_batch=1,
     )
     # test that fix_job_kwargs grabs global kwargs
     new_job_kwargs = dict(n_jobs=get_usable_cpu_count())

src/spikeinterface/core/tests/test_job_tools.py

@@ -1,3 +1,4 @@
+import numpy as np
 import os
 import time
 
@@ -184,7 +185,7 @@ def __call__(self, res):
         job_name="job_name",
     )
     processor.run()
-    num_chunks = len(divide_time_series_into_chunks(recording, processor.chunk_size))
+    num_chunks = len(divide_time_series_into_chunks(recording, processor.chunk_size, None))
 
     assert gathering_func2.pos == num_chunks
 
@@ -364,3 +365,30 @@ def test_get_best_job_kwargs():
     # test_get_best_job_kwargs()
 
     # quick_becnhmark()
+
+
+@pytest.mark.parametrize("format", ["binary", "zarr"])
+@pytest.mark.parametrize("n_jobs", [1, 2])
+def test_num_chunk_per_batch_same_results(format, n_jobs, create_cache_folder):
+    from spikeinterface.core import generate_recording
+
+    recording = generate_recording(durations=[2.0, 1.0], num_channels=4, sampling_frequency=10_000.0, seed=0)
+    kwargs = dict(format=format, chunk_size=1000, n_jobs=n_jobs, progress_bar=False)
+    rec_ref = recording.save(folder=create_cache_folder / f"ref_{format}_{n_jobs}", **kwargs)
+    rec_batch = recording.save(folder=create_cache_folder / f"batch_{format}_{n_jobs}", num_chunk_per_batch=4, **kwargs)
+    for segment_index in range(recording.get_num_segments()):
+        np.testing.assert_array_equal(
+            rec_ref.get_traces(segment_index=segment_index), rec_batch.get_traces(segment_index=segment_index)
+        )
+
+
+def test_divide_time_series_into_chunks_batches_do_not_cross_segments():
+    from spikeinterface.core import generate_recording
+
+    # 2.5 chunks in seg 0 and 1.5 chunks in seg 1, batches of 2 chunks
+    recording = generate_recording(durations=[0.25, 0.15], sampling_frequency=10_000.0, num_channels=2, seed=0)
+    batches = divide_time_series_into_chunks(recording, 1000, num_chunk_per_batch=2)
+    # seg 0: 3 chunks -> [2, 1]; seg 1: 2 chunks -> [2]
+    assert [len(b) for b in batches] == [2, 1, 2]
+    for batch in batches:
+        assert len({segment_index for segment_index, _, _ in batch}) == 1

src/spikeinterface/core/tests/test_zarrextractors.py

@@ -125,6 +125,26 @@ def test_sharding_options(tmp_path):
         )
 
 
+def test_sharding_num_chunk_per_batch(tmp_path):
+    from spikeinterface.core.zarr_tools import adjust_chunks_shards_and_job_kwargs
+
+    recording = generate_recording(durations=[10], num_channels=20)
+    _, _, job_kwargs = adjust_chunks_shards_and_job_kwargs(
+        chunks=(1000, 5), shards=(5000, 10), job_kwargs=dict(n_jobs=1), time_series=recording
+    )
+    assert job_kwargs["num_chunk_per_batch"] == 5
+    _, _, job_kwargs = adjust_chunks_shards_and_job_kwargs(
+        chunks=(1000, 2), shard_factor=(3, 2), job_kwargs=dict(n_jobs=1), time_series=recording
+    )
+    assert job_kwargs["num_chunk_per_batch"] == 3
+
+    # serial and parallel writing give the same result
+    for n_jobs in (1, 2):
+        folder = tmp_path / f"zarr_shards_{n_jobs}.zarr"
+        ZarrRecordingExtractor.write_recording(recording, folder, chunks=(1000, 5), shards=(5000, 10), n_jobs=n_jobs)
+        check_recordings_equal(recording, ZarrRecordingExtractor(folder))
+
+
 def test_ZarrSampleIndexSearch(tmp_path):
     rng = np.random.default_rng(0)
     # two "segments", each sorted, with long runs of equal values so that runs cross

src/spikeinterface/core/zarr_tools.py

@@ -64,9 +64,9 @@ def adjust_chunks_shards_and_job_kwargs(
             assert (
                 shards[dim] >= chunks[dim] and shards[dim] % chunks[dim] == 0
             ), "Shard size must be a multiple of chunk size"
-        # When sharding is used, chunk_size in job_kwargs is used to determine the number of samples per chunk to
-        # write in each job. Each process will write all chunks in a shard.
-        job_kwargs["chunk_size"] = shards[0]
+        # When sharding is used, we set the `num_chunk_per_batch` to match the number of chunks in
+        # the shard, so that each job writes one shard
+        job_kwargs["num_chunk_per_batch"] = int(shards[0] // chunks[0])
     elif shard_factor is not None:
         # If shard_factor is an integer, we only apply it to the first dimension. If it's an iterable,
         # we apply it to all dimensions.
@@ -76,7 +76,7 @@ def adjust_chunks_shards_and_job_kwargs(
             if len(shard_factor) != len(chunks):
                 raise ValueError("shard_factor must have the same length as chunks when it is an iterable")
             shards = tuple(chunks[dim] * shard_factor[dim] for dim in range(len(chunks)))
-        job_kwargs["chunk_size"] = shards[0]
+        job_kwargs["num_chunk_per_batch"] = int(shards[0] // chunks[0])
 
     return chunks, shards, job_kwargs
 

src/spikeinterface/sorters/internal/tests/test_lupin.py

@@ -12,8 +12,6 @@
 class LupinSorterCommonTestSuite(SorterCommonTestSuite, unittest.TestCase):
     SorterClass = LupinSorter
 
-    # TODO: remove once writing to zarr is supported with zarr>=3 (save_array writes the templates to zarr)
-
     def test_with_run(self):
         super().test_with_run()
 

src/spikeinterface/sorters/internal/tests/test_tridesclous2.py

@@ -12,8 +12,6 @@
 class Tridesclous2SorterCommonTestSuite(SorterCommonTestSuite, unittest.TestCase):
     SorterClass = Tridesclous2Sorter
 
-    # TODO: remove once writing to zarr is supported with zarr>=3 (save_array writes the templates to zarr)
-
     def test_with_run(self):
         super().test_with_run()