Set the chunk size of each preprocessing step with per_step_job_kwargs

src/spikeinterface/core/globals.py

@@ -133,6 +133,11 @@ def set_global_job_kwargs(**job_kwargs):
         assert k in job_keys, (
             f"{k} is not a valid job keyword argument. " f"Available keyword arguments are: {list(job_keys)}"
         )
+    if "per_step_job_kwargs" in job_kwargs:
+        raise ValueError(
+            "per_step_job_kwargs describes the steps of one processing chain, pass it to the call that runs "
+            "the chain instead of setting it globally"
+        )
     global_job_kwargs.update(job_kwargs)
     global global_job_kwargs_set
     global_job_kwargs_set = True

src/spikeinterface/core/job_tools.py

@@ -9,6 +9,7 @@
 import threading
 import warnings
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from contextvars import ContextVar
 
 import numpy as np
 from threadpoolctl import threadpool_limits
@@ -65,8 +66,49 @@
     "progress_bar",
     "mp_context",
     "max_threads_per_worker",
+    "per_step_job_kwargs",
 )
 
+# The `per_step_job_kwargs` of the running job. `TimeSeriesChunkExecutor` sets it in each worker and the preprocessing
+# steps read it when they are asked for traces.
+_per_step_job_kwargs = ContextVar("per_step_job_kwargs", default=None)
+
+
+def get_per_step_job_kwargs() -> dict | None:
+    return _per_step_job_kwargs.get()
+
+
+def get_step_chunk_size(options, time_series) -> int:
+    """
+    The number of samples a step computes at once, from its `chunk_size` or its `chunk_duration`.
+    """
+    if "chunk_size" in options:
+        return int(options["chunk_size"])
+    return chunk_duration_to_chunk_size(options["chunk_duration"], time_series)
+
+
+def check_per_step_job_kwargs(per_step_job_kwargs, time_series, chunk_size):
+    """
+    Check that `per_step_job_kwargs` names preprocessing functions, sets either `chunk_size` or `chunk_duration`,
+    and that no step is given more than the `chunk_size` of each job, which is the most a step can be asked for.
+    """
+    from spikeinterface.preprocessing.preprocessing_classes import _all_preprocesser_dict
+
+    step_names = [function.__name__ for function in _all_preprocesser_dict.values()]
+    for step_name, options in per_step_job_kwargs.items():
+        if step_name not in step_names:
+            raise ValueError(f"per_step_job_kwargs: '{step_name}' is not a preprocessing function")
+        if set(options) not in ({"chunk_size"}, {"chunk_duration"}):
+            raise ValueError(
+                f"per_step_job_kwargs: give either 'chunk_size' or 'chunk_duration', got {set(options)} for '{step_name}'"
+            )
+        if chunk_size is not None and get_step_chunk_size(options, time_series) > chunk_size:
+            raise ValueError(
+                f"per_step_job_kwargs: the chunk size of '{step_name}' is longer than the chunks of the job "
+                f"({chunk_size} samples)"
+            )
+
+
 # theses key are the same and should not be in th final dict
 _mutually_exclusive = (
     "total_memory",
@@ -419,6 +461,7 @@ def __init__(
         job_name="",
         max_threads_per_worker=1,
         need_worker_index=False,
+        per_step_job_kwargs=None,
     ):
         self.time_series = time_series
         self.func = func
@@ -470,6 +513,10 @@ def __init__(
 
         self.need_worker_index = need_worker_index
 
+        if per_step_job_kwargs:
+            check_per_step_job_kwargs(per_step_job_kwargs, self.time_series, self.chunk_size)
+        self.per_step_job_kwargs = per_step_job_kwargs
+
         if verbose:
             chunk_memory = self.get_chunk_memory()
             total_memory = chunk_memory * self.n_jobs
@@ -514,7 +561,13 @@ def run(self, slices=None):
         """
         Runs the defined jobs.
         """
+        token = _per_step_job_kwargs.set(self.per_step_job_kwargs)
+        try:
+            return self._run(slices)
+        finally:
+            _per_step_job_kwargs.reset(token)
 
+    def _run(self, slices):
         if slices is None:
             # TODO: rename
             slices = divide_time_series_into_chunks(self.time_series, self.chunk_size)
@@ -571,6 +624,7 @@ def run(self, slices=None):
                         self.need_worker_index,
                         lock,
                         array_pid,
+                        self.per_step_job_kwargs,
                     ),
                 ) as executor:
                     results = executor.map(process_function_wrapper, slices)
@@ -616,6 +670,7 @@ def run(self, slices=None):
                         thread_local_data,
                         self.need_worker_index,
                         lock,
+                        self.per_step_job_kwargs,
                     ),
                 ) as executor:
                     slices2 = [(thread_local_data,) + tuple(args) for args in slices]
@@ -665,9 +720,13 @@ def __call__(self, args):
 global _process_func_wrapper
 
 
-def process_worker_initializer(func, init_func, init_args, max_threads_per_worker, need_worker_index, lock, array_pid):
+def process_worker_initializer(
+    func, init_func, init_args, max_threads_per_worker, need_worker_index, lock, array_pid, per_step_job_kwargs
+):
     global _process_func_wrapper
 
+    _per_step_job_kwargs.set(per_step_job_kwargs)
+
     if need_worker_index:
         child_process = multiprocessing.current_process()
         lock.acquire()
@@ -703,8 +762,9 @@ def process_function_wrapper(args):
 
 
 def thread_worker_initializer(
-    func, init_func, init_args, max_threads_per_worker, thread_local_data, need_worker_index, lock
+    func, init_func, init_args, max_threads_per_worker, thread_local_data, need_worker_index, lock, per_step_job_kwargs
 ):
+    _per_step_job_kwargs.set(per_step_job_kwargs)
 
     if need_worker_index:
         lock.acquire()

src/spikeinterface/preprocessing/basepreprocessor.py

@@ -1,4 +1,9 @@
+import functools
+
+import numpy as np
+
 from spikeinterface.core import BaseRecording, BaseRecordingSegment
+from spikeinterface.core.job_tools import get_per_step_job_kwargs, get_step_chunk_size
 
 
 class BasePreprocessor(BaseRecording):
@@ -24,7 +29,58 @@ def __init__(self, recording, sampling_frequency=None, channel_ids=None, dtype=N
         # self._kwargs have to be handled in subclass
 
 
+def split_requests_by_per_step_job_kwargs(get_traces):
+    """
+    Wrap the `get_traces` of a preprocessing segment so that, inside a job with `per_step_job_kwargs`, requests longer
+    than the step's chunk size are computed in pieces and assembled.
+    """
+
+    @functools.wraps(get_traces)
+    def get_traces_in_pieces(self, start_frame=None, end_frame=None, channel_indices=None):
+        per_step_job_kwargs = get_per_step_job_kwargs()
+        if not per_step_job_kwargs:
+            return get_traces(self, start_frame, end_frame, channel_indices)
+
+        from .preprocessing_classes import _all_preprocesser_dict
+
+        recording = self.parent_extractor
+        function = _all_preprocesser_dict.get(type(recording))
+        options = per_step_job_kwargs.get(function.__name__) if function is not None else None
+        if options is None:
+            return get_traces(self, start_frame, end_frame, channel_indices)
+
+        if start_frame is None:
+            start_frame = 0
+        if end_frame is None:
+            end_frame = self.get_num_samples()
+
+        # Pieces of equal length instead of full pieces and a remainder. A request a bit longer than the
+        # chunk size, for example one that includes the margin of the step below, would otherwise leave
+        # a tiny last piece that the step computes with its own full margin.
+        chunk_size = get_step_chunk_size(options, recording)
+        num_samples = end_frame - start_frame
+        num_pieces = -(-num_samples // chunk_size)
+        if num_pieces <= 1:
+            return get_traces(self, start_frame, end_frame, channel_indices)
+
+        piece_boundaries = start_frame + np.arange(num_pieces + 1) * num_samples // num_pieces
+        traces = None
+        for piece_start, piece_end in zip(piece_boundaries[:-1], piece_boundaries[1:]):
+            piece = get_traces(self, int(piece_start), int(piece_end), channel_indices)
+            if traces is None:
+                traces = np.empty((num_samples, piece.shape[1]), dtype=piece.dtype)
+            traces[piece_start - start_frame : piece_end - start_frame] = piece
+        return traces
+
+    return get_traces_in_pieces
+
+
 class BasePreprocessorSegment(BaseRecordingSegment):
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if "get_traces" in cls.__dict__:
+            cls.get_traces = split_requests_by_per_step_job_kwargs(cls.__dict__["get_traces"])
+
     def __init__(self, parent_recording_segment):
         BaseRecordingSegment.__init__(self, **parent_recording_segment.get_times_kwargs())
         self.parent_recording_segment = parent_recording_segment