Per-step chunk size and execution settings for preprocessing and saving #4905

I am coming at this after thinking about #4833 and also about other performance and memory issues like #4837, #4857 and #4859.

I think they hint at a common problem: each step of a processing and saving chain can need different execution settings, but right now we can't give phase_shift, bandpass_filter, common_reference and the writer their own chunk_duration.

We can illustrate this coupling with the current zarr problem. All the chunks of a shard have to be written by the same worker, otherwise two workers writing into the same shard overwrite each other. So with shards chunk_size is set to the shard length (adjust_chunks_shards_and_job_kwargs), and every step computes the whole shard at once. But avoiding the race condition only needs the shard to stay in one worker. Inside that worker the chain could compute small pieces, and the writer could write pieces of its own size, up to the whole shard in one call. #4833 keeps each shard in one worker with a new job keyword, num_chunk_per_batch, and computes the chain one zarr chunk at a time. The writer still shares that size with every step, so it also writes one zarr chunk at a time.

Outside zarr, steps can benefit from different chunk sizes. A highpass with a low cutoff benefits from longer chunks because it uses large margins. The automatic margin is five cutoff periods, so a 1 Hz filter uses 5 s on each side. Away from the recording boundaries, producing 1 s means filtering 11 s, while producing 30 s means filtering 40 s. Longer chunks reduce this repeated work. For phase_shift, shorter chunks reduce the memory used by its temporary arrays. Giving both steps the same chunk size forces a compromise between margin overhead and peak memory.

I propose giving each step its own execution settings, supplied separately from its algorithm parameters. The writer is a step too: its chunk_duration says how much data it writes per call, just as a preprocessor's chunk_duration says how much it computes.

I am envisioning something that looks like this, using the recording objects we already construct. The chain says what is computed and a separate execution plan says how each step runs. ZarrWriter and execute are proposed additions:

source = recording
shifted = phase_shift(source)
filtered = bandpass_filter(shifted, freq_min=1.0, freq_max=300.0, ignore_low_freq_error=True)
referenced = common_reference(filtered, operator="median")
writer = ZarrWriter(
    referenced,
    folder="preprocessed.zarr",
    chunks=(30_000, 384),
    shards=(1_800_000, 384),
)

execute(
    writer,
    n_jobs=4,
    step_options={
        shifted: {"chunk_duration": "1s"},
        filtered: {"chunk_duration": "30s"},
        referenced: {"chunk_duration": "30s"},
        writer: {"chunk_duration": "60s"},
    },
)

ZarrWriter describes the write without starting it. execute follows the parent links back to source, and the keys in step_options are the actual objects. The same chain can run with different step_options on a different machine without touching its parameters. Without overrides it would use the existing execution behavior. The preprocessing constructors stay unchanged. The runner supplies settings during execution, for example through adapters that split requests to a step's get_traces.

The advantages are:

• Each step runs at the size that suits it. The filter gets long chunks for its margins while phase_shift stays at 1 s and keeps its memory small.
• Shard ownership stays a zarr concern. The zarr writer decides that one worker owns each shard, and that no longer forces every step to compute a whole shard.
• It becomes possible to tune and benchmark one step at a time.

And in addition, once we decouple this and every part of the chain has its own configuration it opens the door to more goodies:

• Distribution across channels. How many channels a step processes at a time, the one we discussed on the PRs for the memory changes (#4840, #4903 and #4904, tracked in #4837). It is a module constant right now.
• Threads. The thread configuration inside a step that Galen discussed in #4857.
• Backend. Running specific steps on a GPU while the rest stays on the CPU.

For example the step_options entry for the phase shift could become shifted: {"chunk_duration": "1s", "num_channels_per_block": 16, "num_threads": 4}. These controls would need support inside the relevant steps to read their execution context.

execute would run full processing chains, with or without saving at the end. I think this separates the multiprocessing machinery from the preprocessing definitions. We could also expose the same per-step settings through the existing recording.save() entry point by adding a step_options argument:

referenced.save(
    format="zarr",
    folder="preprocessed.zarr",
    chunks=(30_000, 384),
    shards=(1_800_000, 384),
    n_jobs=4,
    chunk_duration="60s",
    step_options={
        shifted: {"chunk_duration": "1s"},
        filtered: {"chunk_duration": "30s"},
        referenced: {"chunk_duration": "30s"},
    },
)