From 6b52b2e71767273ccdacf2aa9aa5e3434a054e6d Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Wed, 4 Nov 2026 12:34:47 -0600 Subject: [PATCH 1/2] Add step_options to recording.save to set chunk_duration per step --- src/spikeinterface/core/baserecording.py | 17 ++- src/spikeinterface/core/step_options.py | 136 +++++++++++++++++++++++ 2 files changed, 149 insertions(+), 4 deletions(-) create mode 100644 src/spikeinterface/core/step_options.py diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index d4bb36f2ab..9ead499948 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -908,6 +908,13 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): """ kwargs, job_kwargs = split_job_kwargs(save_kwargs) + recording = self + step_options = kwargs.pop("step_options", None) + if step_options is not None: + from .step_options import apply_step_options + + recording = apply_step_options(self, step_options) + if format == "binary": if "folder" not in kwargs: raise ValueError("Missing folder in recording.save(folder='...')") @@ -916,7 +923,7 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): folder = kwargs.pop("folder") cached = BinaryFolderRecording.write_recording( - self, folder_path=folder, verbose=verbose, **kwargs, **job_kwargs + recording, folder_path=folder, verbose=verbose, **kwargs, **job_kwargs ) elif format == "memory": @@ -924,12 +931,14 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): from .numpyextractors import SharedMemoryRecording cached = SharedMemoryRecording.from_recording( - self, with_metadata=True, with_time_vector=True, **job_kwargs + recording, with_metadata=True, with_time_vector=True, **job_kwargs ) else: from spikeinterface.core import NumpyRecording - cached = NumpyRecording.from_recording(self, with_metadata=True, with_time_vector=True, **job_kwargs) + cached = NumpyRecording.from_recording( + recording, with_metadata=True, with_time_vector=True, **job_kwargs + ) elif format == "zarr": if "folder" not in kwargs: @@ -939,7 +948,7 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): from .zarrextractors import ZarrRecordingExtractor cached = ZarrRecordingExtractor.write_recording( - self, folder_path=folder_path, verbose=verbose, **kwargs, **job_kwargs + recording, folder_path=folder_path, verbose=verbose, **kwargs, **job_kwargs ) else: diff --git a/src/spikeinterface/core/step_options.py b/src/spikeinterface/core/step_options.py new file mode 100644 index 0000000000..cee07d1d6d --- /dev/null +++ b/src/spikeinterface/core/step_options.py @@ -0,0 +1,136 @@ +import numpy as np + +from .base import BaseExtractor +from .baserecording import BaseRecording, BaseRecordingSegment +from .job_tools import chunk_duration_to_chunk_size + + +def apply_step_options(recording, step_options): + """ + Return a recording that computes the same chain as `recording`, with execution options for some of its steps. + + The only option for now is `chunk_duration`, the longest request a step computes in one call. Longer requests + to that step are split into pieces and the pieces are assembled, so a step's parents only see the pieces. + + The original chain is not modified. The steps between `recording` and each configured step are rebuilt + from their dictionaries so they read from the split requests. + + Parameters + ---------- + recording : BaseRecording + The last step of the chain. + step_options : dict + Maps steps of the chain, the recording objects themselves, to a dict of options, for example + `{shifted: {"chunk_duration": "1s"}, filtered: {"chunk_duration": "30s"}}`. + + Returns + ------- + recording_with_step_options : BaseRecording + The rebuilt chain. + """ + chunk_sizes = {} + for step, options in step_options.items(): + unsupported_options = set(options) - {"chunk_duration"} + if unsupported_options: + raise ValueError(f"Unsupported step options {unsupported_options}, only 'chunk_duration' is supported") + chunk_sizes[id(step)] = chunk_duration_to_chunk_size(options["chunk_duration"], step) + + # Keyed by id because recordings are not hashable. A recording used by several steps is rebuilt once. + rebuilt = {} + + def rebuild(step): + if id(step) in rebuilt: + return rebuilt[id(step)] + + changed = False + + def rebuild_parent(value): + nonlocal changed + if not isinstance(value, BaseRecording): + return value + new_value = rebuild(value) + changed = changed or new_value is not value + return new_value + + new_kwargs = {} + for name, value in step._kwargs.items(): + if isinstance(value, list): + new_kwargs[name] = [rebuild_parent(element) for element in value] + elif isinstance(value, dict): + new_kwargs[name] = {key: rebuild_parent(element) for key, element in value.items()} + else: + new_kwargs[name] = rebuild_parent(value) + + new_step = step + if changed: + step_dict = step.to_dict(include_annotations=True, include_properties=True) + step_dict["kwargs"] = new_kwargs + new_step = BaseExtractor.from_dict(step_dict) + if id(step) in chunk_sizes: + new_step = ChunkRequestsRecording(new_step, chunk_size=chunk_sizes[id(step)]) + + rebuilt[id(step)] = new_step + return new_step + + recording_with_step_options = rebuild(recording) + + steps_not_found = [step for step in step_options if id(step) not in rebuilt] + if steps_not_found: + raise ValueError(f"These steps are not part of the chain of the recording: {steps_not_found}") + + return recording_with_step_options + + +class ChunkRequestsRecording(BaseRecording): + """ + Split each request for traces into pieces of at most `chunk_size` samples before passing it to the parent, + so the parent computes at most `chunk_size` samples per call. Used by `apply_step_options`. + """ + + def __init__(self, recording, chunk_size): + BaseRecording.__init__( + self, + sampling_frequency=recording.get_sampling_frequency(), + channel_ids=recording.channel_ids, + dtype=recording.get_dtype(), + ) + recording.copy_metadata(self, only_main=False) + self._parent = recording + + for parent_segment in recording.segments: + self.add_recording_segment(ChunkRequestsRecordingSegment(parent_segment, chunk_size)) + + self._kwargs = dict(recording=recording, chunk_size=chunk_size) + + +class ChunkRequestsRecordingSegment(BaseRecordingSegment): + def __init__(self, parent_recording_segment, chunk_size): + BaseRecordingSegment.__init__(self, **parent_recording_segment.get_times_kwargs()) + self.parent_recording_segment = parent_recording_segment + self.chunk_size = chunk_size + + def get_num_samples(self): + return self.parent_recording_segment.get_num_samples() + + def 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 + # `chunk_size`, for example one that includes the margin of the step below, would otherwise leave a + # tiny last piece that the parent computes with its own full margin. + num_samples = end_frame - start_frame + num_pieces = -(-num_samples // self.chunk_size) + if num_pieces <= 1: + return self.parent_recording_segment.get_traces(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 = self.parent_recording_segment.get_traces(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 From c5c3cf6532dfe0986275ec047a3e7e15212fa140 Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Sun, 8 Nov 2026 11:10:19 -0600 Subject: [PATCH 2/2] Pass per_step_job_kwargs as a job keyword instead of wrapping the chain Each step accepts either chunk_size or chunk_duration. --- src/spikeinterface/core/baserecording.py | 17 +-- src/spikeinterface/core/globals.py | 5 + src/spikeinterface/core/job_tools.py | 64 ++++++++- src/spikeinterface/core/step_options.py | 136 ------------------ .../preprocessing/basepreprocessor.py | 56 ++++++++ 5 files changed, 127 insertions(+), 151 deletions(-) delete mode 100644 src/spikeinterface/core/step_options.py diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 9ead499948..d4bb36f2ab 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -908,13 +908,6 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): """ kwargs, job_kwargs = split_job_kwargs(save_kwargs) - recording = self - step_options = kwargs.pop("step_options", None) - if step_options is not None: - from .step_options import apply_step_options - - recording = apply_step_options(self, step_options) - if format == "binary": if "folder" not in kwargs: raise ValueError("Missing folder in recording.save(folder='...')") @@ -923,7 +916,7 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): folder = kwargs.pop("folder") cached = BinaryFolderRecording.write_recording( - recording, folder_path=folder, verbose=verbose, **kwargs, **job_kwargs + self, folder_path=folder, verbose=verbose, **kwargs, **job_kwargs ) elif format == "memory": @@ -931,14 +924,12 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): from .numpyextractors import SharedMemoryRecording cached = SharedMemoryRecording.from_recording( - recording, with_metadata=True, with_time_vector=True, **job_kwargs + self, with_metadata=True, with_time_vector=True, **job_kwargs ) else: from spikeinterface.core import NumpyRecording - cached = NumpyRecording.from_recording( - recording, with_metadata=True, with_time_vector=True, **job_kwargs - ) + cached = NumpyRecording.from_recording(self, with_metadata=True, with_time_vector=True, **job_kwargs) elif format == "zarr": if "folder" not in kwargs: @@ -948,7 +939,7 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): from .zarrextractors import ZarrRecordingExtractor cached = ZarrRecordingExtractor.write_recording( - recording, folder_path=folder_path, verbose=verbose, **kwargs, **job_kwargs + self, folder_path=folder_path, verbose=verbose, **kwargs, **job_kwargs ) else: diff --git a/src/spikeinterface/core/globals.py b/src/spikeinterface/core/globals.py index 2712f4db3a..d535d32b6f 100644 --- a/src/spikeinterface/core/globals.py +++ b/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 diff --git a/src/spikeinterface/core/job_tools.py b/src/spikeinterface/core/job_tools.py index 59af6fd254..e8205241b2 100644 --- a/src/spikeinterface/core/job_tools.py +++ b/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() diff --git a/src/spikeinterface/core/step_options.py b/src/spikeinterface/core/step_options.py deleted file mode 100644 index cee07d1d6d..0000000000 --- a/src/spikeinterface/core/step_options.py +++ /dev/null @@ -1,136 +0,0 @@ -import numpy as np - -from .base import BaseExtractor -from .baserecording import BaseRecording, BaseRecordingSegment -from .job_tools import chunk_duration_to_chunk_size - - -def apply_step_options(recording, step_options): - """ - Return a recording that computes the same chain as `recording`, with execution options for some of its steps. - - The only option for now is `chunk_duration`, the longest request a step computes in one call. Longer requests - to that step are split into pieces and the pieces are assembled, so a step's parents only see the pieces. - - The original chain is not modified. The steps between `recording` and each configured step are rebuilt - from their dictionaries so they read from the split requests. - - Parameters - ---------- - recording : BaseRecording - The last step of the chain. - step_options : dict - Maps steps of the chain, the recording objects themselves, to a dict of options, for example - `{shifted: {"chunk_duration": "1s"}, filtered: {"chunk_duration": "30s"}}`. - - Returns - ------- - recording_with_step_options : BaseRecording - The rebuilt chain. - """ - chunk_sizes = {} - for step, options in step_options.items(): - unsupported_options = set(options) - {"chunk_duration"} - if unsupported_options: - raise ValueError(f"Unsupported step options {unsupported_options}, only 'chunk_duration' is supported") - chunk_sizes[id(step)] = chunk_duration_to_chunk_size(options["chunk_duration"], step) - - # Keyed by id because recordings are not hashable. A recording used by several steps is rebuilt once. - rebuilt = {} - - def rebuild(step): - if id(step) in rebuilt: - return rebuilt[id(step)] - - changed = False - - def rebuild_parent(value): - nonlocal changed - if not isinstance(value, BaseRecording): - return value - new_value = rebuild(value) - changed = changed or new_value is not value - return new_value - - new_kwargs = {} - for name, value in step._kwargs.items(): - if isinstance(value, list): - new_kwargs[name] = [rebuild_parent(element) for element in value] - elif isinstance(value, dict): - new_kwargs[name] = {key: rebuild_parent(element) for key, element in value.items()} - else: - new_kwargs[name] = rebuild_parent(value) - - new_step = step - if changed: - step_dict = step.to_dict(include_annotations=True, include_properties=True) - step_dict["kwargs"] = new_kwargs - new_step = BaseExtractor.from_dict(step_dict) - if id(step) in chunk_sizes: - new_step = ChunkRequestsRecording(new_step, chunk_size=chunk_sizes[id(step)]) - - rebuilt[id(step)] = new_step - return new_step - - recording_with_step_options = rebuild(recording) - - steps_not_found = [step for step in step_options if id(step) not in rebuilt] - if steps_not_found: - raise ValueError(f"These steps are not part of the chain of the recording: {steps_not_found}") - - return recording_with_step_options - - -class ChunkRequestsRecording(BaseRecording): - """ - Split each request for traces into pieces of at most `chunk_size` samples before passing it to the parent, - so the parent computes at most `chunk_size` samples per call. Used by `apply_step_options`. - """ - - def __init__(self, recording, chunk_size): - BaseRecording.__init__( - self, - sampling_frequency=recording.get_sampling_frequency(), - channel_ids=recording.channel_ids, - dtype=recording.get_dtype(), - ) - recording.copy_metadata(self, only_main=False) - self._parent = recording - - for parent_segment in recording.segments: - self.add_recording_segment(ChunkRequestsRecordingSegment(parent_segment, chunk_size)) - - self._kwargs = dict(recording=recording, chunk_size=chunk_size) - - -class ChunkRequestsRecordingSegment(BaseRecordingSegment): - def __init__(self, parent_recording_segment, chunk_size): - BaseRecordingSegment.__init__(self, **parent_recording_segment.get_times_kwargs()) - self.parent_recording_segment = parent_recording_segment - self.chunk_size = chunk_size - - def get_num_samples(self): - return self.parent_recording_segment.get_num_samples() - - def 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 - # `chunk_size`, for example one that includes the margin of the step below, would otherwise leave a - # tiny last piece that the parent computes with its own full margin. - num_samples = end_frame - start_frame - num_pieces = -(-num_samples // self.chunk_size) - if num_pieces <= 1: - return self.parent_recording_segment.get_traces(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 = self.parent_recording_segment.get_traces(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 diff --git a/src/spikeinterface/preprocessing/basepreprocessor.py b/src/spikeinterface/preprocessing/basepreprocessor.py index d28de0b34f..60b3c62806 100644 --- a/src/spikeinterface/preprocessing/basepreprocessor.py +++ b/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