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/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