diff --git a/conftest.py b/conftest.py index 6c4bfbef3a..d19cb5e6f9 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 @@ -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")) diff --git a/src/spikeinterface/core/globals.py b/src/spikeinterface/core/globals.py index 2712f4db3a..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 + 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..7f0e73fcb8 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 @@ -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 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..c0241e2166 100644 --- a/src/spikeinterface/core/tests/test_job_tools.py +++ b/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 diff --git a/src/spikeinterface/core/tests/test_zarrextractors.py b/src/spikeinterface/core/tests/test_zarrextractors.py index e8c76018de..622f1b47ce 100644 --- a/src/spikeinterface/core/tests/test_zarrextractors.py +++ b/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 diff --git a/src/spikeinterface/core/zarr_tools.py b/src/spikeinterface/core/zarr_tools.py index 41ff57a358..634d6de39e 100644 --- a/src/spikeinterface/core/zarr_tools.py +++ b/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 diff --git a/src/spikeinterface/sorters/internal/tests/test_lupin.py b/src/spikeinterface/sorters/internal/tests/test_lupin.py index f838f9f4eb..1c5ecdf1bd 100644 --- a/src/spikeinterface/sorters/internal/tests/test_lupin.py +++ b/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() diff --git a/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py b/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py index 6431532b87..0c3bfd68a7 100644 --- a/src/spikeinterface/sorters/internal/tests/test_tridesclous2.py +++ b/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()