Skip to content

Set the chunk size of each preprocessing step with per_step_job_kwargs - #4930

Open
h-mayorquin wants to merge 2 commits into
SpikeInterface:mainfrom
h-mayorquin:chunk_per_step_prototype
Open

h-mayorquin wants to merge 2 commits into
SpikeInterface:mainfrom
h-mayorquin:chunk_per_step_prototype

Conversation

@h-mayorquin

Copy link
Copy Markdown
Contributor

@alejoe91 mentioned in #4833 that distributing the work across processing steps would be very hard to implement. This PR is a prototype of the simplest part of #4905 to show that this part is simple: each step of a chain gets its own chunk size inside the worker that owns the job. It adds a per_step_job_kwargs job keyword that maps the name of a preprocessing function to its own chunk_size or chunk_duration:

from spikeinterface.preprocessing import bandpass_filter, common_reference, phase_shift

shifted = phase_shift(recording)
filtered = bandpass_filter(shifted)
referenced = common_reference(filtered, operator="median")

referenced.save(
    folder="preprocessed",
    n_jobs=4,
    chunk_duration="10s",  # common_reference is not listed, so it computes the whole 10 s job
    per_step_job_kwargs={"phase_shift": {"chunk_duration": "1s"}, "bandpass_filter": {"chunk_size": 150_000}},
)

The top level job kwargs keep their meaning and set what each worker gets. per_step_job_kwargs only splits that job inside the worker: a step asked for more than its chunk size computes the request in equal pieces. It also works for detect_peaks and SortingAnalyzer.compute without changing them as it is a job keyword. It can't be set globally as it describes the steps of one chain.

The main limitation is that a step can't use a larger chunk size than the steps after it in the chain. Each step only computes what the next one asks for, so a larger setting has no effect, and the chunk size of the job is the largest of all.

I think this covers most cases and is the smallest change that:

  1. lets the steps before the writer compute small pieces inside the worker that owns a zarr shard. Giving the writer its own chunk size as proposed in Per-step chunk size and execution settings for preprocessing and saving #4905, instead of taking it from the flat job kwargs of recording.save, would then fix the problem in Add num_chunk_per_batch to job kwargs and use it for zarr v3 shard writes #4833.
  2. gives each step its own settings, so the channel blocks from Process channels in banks in preprocessors to reduce RAM usage #4837 and the threads from Threads inside preprocessing steps #4857 could go in the same dictionary
  3. stays simple, about 120 lines in three files

I have an idea to make the sizes of the steps independent of each other (inside the samples a worker gets), but I think this is enough progress for now and it would make the proposal harder to discuss. If we decide we want to go there, there is a path forward.

Also, this keeps n_jobs shared across the same pipeline.

@alejoe91 alejoe91 added core Changes to core module discussion General discussions and community feedback labels Nov 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Changes to core module discussion General discussions and community feedback

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants