From 218b394b55216d19789d5517c1c81a8fb9afe1d6 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Tue, 6 Oct 2026 14:08:21 +0200 Subject: [PATCH 1/9] draft proof of concept of num_chunk_per_batch for the zarr shards --- src/spikeinterface/core/globals.py | 2 +- src/spikeinterface/core/job_tools.py | 86 ++++++++++++++++++++++------ 2 files changed, 68 insertions(+), 20 deletions(-) diff --git a/src/spikeinterface/core/globals.py b/src/spikeinterface/core/globals.py index 2712f4db3a..08fe6458c9 100644 --- a/src/spikeinterface/core/globals.py +++ b/src/spikeinterface/core/globals.py @@ -102,7 +102,7 @@ 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 diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index 59af6fd254..dc71b22f32 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/src/spikeinterface/core/job_tools.py @@ -225,12 +225,14 @@ 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): 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]) + if num_chunk_per_batch is not None: + slices = [slices[i : i + num_chunk_per_batch] for i in range(0, len(slices), num_chunk_per_batch)] return slices @@ -383,6 +385,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 reeive a batch of continuous chunks to process at once. + This is usefull 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 +423,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 +472,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 +527,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 = [] @@ -539,10 +549,17 @@ def run(self, slices=None): 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) + 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) else: n_jobs = min(self.n_jobs, len(slices)) @@ -583,10 +600,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 +648,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 +683,28 @@ 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 From d7a6e966e268a4c4fd6fcba85e57ebd60c54b1bc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 6 Oct 2026 12:09:23 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/spikeinterface/core/globals.py | 8 +++++++- src/spikeinterface/core/job_tools.py | 3 +-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/globals.py b/src/spikeinterface/core/globals.py index 08fe6458c9..941249b45b 100644 --- a/src/spikeinterface/core/globals.py +++ b/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, num_chunk_per_batch=None, + 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 diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index dc71b22f32..9d2c951e52 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/src/spikeinterface/core/job_tools.py @@ -232,7 +232,7 @@ def divide_time_series_into_chunks(recording, chunk_size, num_chunk_per_batch): chunks = divide_segment_into_chunks(num_frames, chunk_size) slices.extend([(segment_index, frame_start, frame_stop) for frame_start, frame_stop in chunks]) if num_chunk_per_batch is not None: - slices = [slices[i : i + num_chunk_per_batch] for i in range(0, len(slices), num_chunk_per_batch)] + slices = [slices[i : i + num_chunk_per_batch] for i in range(0, len(slices), num_chunk_per_batch)] return slices @@ -706,7 +706,6 @@ def __call__(self, args): raise ValueError("args must be a tuple or a list of tuples") - # see # https://stackoverflow.com/questions/10117073/how-to-use-initializer-to-set-up-my-multiprocess-pool # the trick is : this variable is global per worker (so not shared in the same process) From 9bba81be814290f6fa3654675cee3943edebe00d Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 12 Oct 2026 11:49:38 +0200 Subject: [PATCH 3/9] fix: add num_chunk_per_batch to job_keys --- src/spikeinterface/core/job_tools.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index 9d2c951e52..bbcc6d370a 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/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 From dafb095166bc4e4c8fb9b2841ea188baec56a1c8 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 12 Oct 2026 12:50:40 +0200 Subject: [PATCH 4/9] fix: failing tests --- conftest.py | 4 ++-- src/spikeinterface/core/job_tools.py | 2 +- src/spikeinterface/core/tests/test_globals.py | 3 +++ src/spikeinterface/core/tests/test_job_tools.py | 2 +- 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/conftest.py b/conftest.py index 9b1296bb0b..0bcb48c796 100644 --- a/conftest.py +++ b/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 @@ -51,7 +51,7 @@ def pytest_collection_modifyitems(config, items): rootdir = Path(config.rootdir) modules_location = rootdir / "src" / "spikeinterface" for item in items: - if config.getoption("--mp-context") is not None and item.name == "test_global_job_kwargs": + if config.getoption("--mp-context") is not None and config.getoption("--mp-context").lower() != "none" and item.name == "test_global_job_kwargs": item.add_marker(pytest.mark.skip(reason="--mp-context changes the default job kwargs")) try: diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index bbcc6d370a..4dae1b00d9 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/src/spikeinterface/core/job_tools.py @@ -226,7 +226,7 @@ def divide_segment_into_chunks(num_frames, chunk_size): return chunks -def divide_time_series_into_chunks(recording, chunk_size, num_chunk_per_batch): +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) diff --git a/src/spikeinterface/core/tests/test_globals.py b/src/spikeinterface/core/tests/test_globals.py index b8f54dfe0e..0a0de8b54c 100644 --- a/src/spikeinterface/core/tests/test_globals.py +++ b/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()) diff --git a/src/spikeinterface/core/tests/test_job_tools.py b/src/spikeinterface/core/tests/test_job_tools.py index 79cc76b85e..052443bce5 100644 --- a/src/spikeinterface/core/tests/test_job_tools.py +++ b/src/spikeinterface/core/tests/test_job_tools.py @@ -184,7 +184,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 From c07303b63810555ad76cf2b786c91c4bf89216ef Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 12 Oct 2026 13:01:44 +0200 Subject: [PATCH 5/9] fix: serialization tests --- .github/scripts/serialization/objects.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index 78fb3aed0b..b92d4d8f37 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -155,7 +155,8 @@ def _build_recording_with_interleaved_probes(): if parse(si_version) <= parse("0.105.0"): rec_with_probe = rec.set_probegroup(probegroup, in_place=False) # old API returns a new recording; portable across versions else: - rec_with_probe = rec.set_probegroup(probegroup) # new API returns a new recording; portable across versions + rec.set_probegroup(probegroup) # new API returns a new recording; portable across versions + rec_with_probe = rec return rec_with_probe From 29f083c63c539c3817c3ba1dfc33b99758b15b12 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 12 Oct 2026 13:08:44 +0200 Subject: [PATCH 6/9] fix: serialization tests 2 --- .github/scripts/serialization/objects.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index b92d4d8f37..6f760cdbff 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -151,13 +151,8 @@ def _build_recording_with_interleaved_probes(): probegroup.add_probe(probe0) probegroup.add_probe(probe1) probegroup.set_global_device_channel_indices([0, 2, 4, 6, 1, 3, 5, 7]) - # Interleave the two probes' channels: channel i alternates between probe0 and probe1. - if parse(si_version) <= parse("0.105.0"): - rec_with_probe = rec.set_probegroup(probegroup, in_place=False) # old API returns a new recording; portable across versions - else: - rec.set_probegroup(probegroup) # new API returns a new recording; portable across versions - rec_with_probe = rec - return rec_with_probe + rec_with_probe = rec.set_probegroup(probegroup) + return rec_with_probe if rec_with_probe is not None else rec def _check_recording_with_interleaved_probes(rec): From dcfeb47eda22e0934c9c568f4fa023394e8cd1a1 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 12 Oct 2026 14:49:59 +0200 Subject: [PATCH 7/9] fix: serialization tests 3 --- .github/scripts/serialization/objects.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/.github/scripts/serialization/objects.py b/.github/scripts/serialization/objects.py index 6f760cdbff..765d07ce8c 100644 --- a/.github/scripts/serialization/objects.py +++ b/.github/scripts/serialization/objects.py @@ -105,11 +105,8 @@ def _build_recording_with_probe(): rec = generate_recording(num_channels=8, durations=[DEFAULT_DURATION], sampling_frequency=30000.0, seed=0) probe = generate_linear_probe(num_elec=8) probe.set_device_channel_indices(np.arange(8)) - if parse(si_version) <= parse("0.105.0"): - rec_with_probe = rec.set_probe(probe, in_place=False) # old API returns a new recording; portable across versions - else: - rec_with_probe = rec.set_probe(probe) # new API returns a new recording; portable across versions - return rec_with_probe + rec_with_probe = rec.set_probe(probe) + return rec_with_probe if rec_with_probe is not None else rec def _build_recording_with_timestamps(): From b4e4bedc37bfe8953baf66ed61ab11ccc8b3e98c Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Fri, 30 Oct 2026 14:28:54 +0100 Subject: [PATCH 8/9] fix: batches should not cross segments! --- src/spikeinterface/core/job_tools.py | 9 +++++---- src/spikeinterface/core/tests/test_job_tools.py | 12 ++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index 0d857d6cdc..efdf1e5dfc 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/src/spikeinterface/core/job_tools.py @@ -230,10 +230,11 @@ def divide_time_series_into_chunks(recording, chunk_size, num_chunk_per_batch=No 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]) - if num_chunk_per_batch is not None: - slices = [slices[i : i + num_chunk_per_batch] for i in range(0, len(slices), num_chunk_per_batch)] + 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 diff --git a/src/spikeinterface/core/tests/test_job_tools.py b/src/spikeinterface/core/tests/test_job_tools.py index 05838707e5..c0241e2166 100644 --- a/src/spikeinterface/core/tests/test_job_tools.py +++ b/src/spikeinterface/core/tests/test_job_tools.py @@ -380,3 +380,15 @@ def test_num_chunk_per_batch_same_results(format, n_jobs, create_cache_folder): 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 From 8502e8a100d5b7caf3a0386a5e899842a676ff52 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Tue, 3 Nov 2026 17:32:50 +0100 Subject: [PATCH 9/9] typos from Chris Co-authored-by: Chris Halcrow <57948917+chrishalcrow@users.noreply.github.com> --- src/spikeinterface/core/job_tools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index efdf1e5dfc..7f0e73fcb8 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/src/spikeinterface/core/job_tools.py @@ -389,9 +389,9 @@ class TimeSeriesChunkExecutor: 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 reeive a batch of continuous chunks to process at once. - This is usefull in the context of zarr v3 shards concept. + 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().