n_workers kwarg for FilterRecording and CommonReferenceRecording

[galenlynch]

I was attempting to visualize my recordings as they would be seen by Kilosort4 in a spikeinterface pipeline, but the preprocessing performance was dreadfully slow for 90 minute recordings. I have since realized that this library primarily uses a batch processing model with TimeSeriesChunkExecutor, and parallelizes over chunks of data. However, TimeSeriesChunkExecutor is difficult to use if you just want to see a small stretch of a recording after preprocessing, which from my understanding requires preprocessing the entire recording, or at least saving to disk. Other workflows might also benefit from faster direct access to the data, instead of using the batch process API.

Because of this design choice, inner extractors are written to be single-threaded. This PR introduces extractor-level parallelism for FilterRecording and CommonReferenceRecording, which was intended to speed up get_traces calls that are not inside TimeSeriesChunkExecutor. The high-level takeaway is that on a 384x1M sample data segment, FilterRecording is now 2.8 times faster, and CommonReferenceRecording is 5.4 times faster. Those times reflect the entire extractor with overhead; the speedup for the core computation is 2.9x for scipy.sosfiltfilt and 10.6x for np.median. The core algorithm was not changed for either, and the multi-threading is opt-in by raising n_workers past its default of 1.

The companion PR #4563 was also intended to speed up get_traces calls by addressing the lion's share of the processing time: PhaseShiftRecording. However, the majority of the speedup in that PR is due to an algorithmic improvement and not simply extractor-level parallelism, which is why I separated it.

When these PRs are combined and tested on my data, the entire get_traces() call with phase shifting, CMR, and filtering is 13.1 times faster.

---

But how do these changes compose with TimeSeriesChunkExecutor? The good news is that adding extractor-level parallelism to TimeSeriesChunkExecutor can actually improve performance, and it might be a nice knob to have for RAM-conscious users that nonetheless have modern CPUs with many cores. See the table below for band pass + common reference filter chains, with 'outer' meaning TimeSeriesChunkExecutor n_jobs and 'inner' meaning extractor-level n_workers:

Setup: chunk_duration="1s" (SI default), different splits of a ~24-thread compute budget on the BP+CMR pipeline, per-caller-thread pools:

Budget	Config	Time	Notes
24 threads	outer=24, inner=1 each	1.54 s	clean, minimum thread count
192 threads	outer=24, inner=8 each	1.42 s	absolute peak, oversubscribed
24 threads	outer=8, inner=3 each	1.53 s	8 outer, tied with outer=24 inner=1; ~⅓ RAM
12 threads	outer=12, inner=1 each	1.59 s	~½ RAM of outer=24
12 threads	outer=6, inner=2 each	1.75 s	~¼ RAM of outer=24
12 threads	outer=4, inner=3 each	1.92 s	~⅙ RAM of outer=24
12 threads	outer=1, inner=12 each	4.31 s	single caller — sync overhead dominates

We can split this up for each extractor, with CRE meaning TimeSeriesChunkExecutor:

band pass specifically (inner pool = 8, matching CRE n_jobs=8):

Config	Time	Speedup	Parallelism axis
stock, CRE n=1 (baseline)	7.42 s	1.00×	—
stock, CRE n=8 thread	1.40 s	5.29×	outer only
n_workers=8, CRE n=1	2.39 s	3.04×	inner only
n_workers=8, CRE n=8 thread	1.24 s	6.00×	both

common mode reference specifically (inner pool = 16, exceeds CRE n_jobs=8):

Config	Time	Speedup	Parallelism axis
stock, CRE n=1 (baseline)	3.93 s	1.00×	—
stock, CRE n=8 thread	0.62 s	6.34×	outer only
n_workers=16, CRE n=1	0.86 s	4.57×	inner only
n_workers=16, CRE n=8 thread	0.33 s	11.83×	both

Compatibility

• No default behavior changes. n_workers=1 preserves existing semantics exactly.
• Round-trip dumpability. _kwargs dicts updated on both preprocessors; save() / load() round-trip the new kwargs correctly.
• No new deps. Uses stdlib concurrent.futures.ThreadPoolExecutor, threading, weakref.
• Propagates through the bandpass_filter, highpass_filter, filter, notch_filter, common_reference wrapper functions via **filter_kwargs.
• Long-running processes safe: per-caller pools are cleaned up when the calling thread is GC'd.

Companion PR

An independent companion PR #4563 adds a sinc-FIR alternative to PhaseShiftRecording with ~100× per-stage speedup while reducing memory. Combined, they give 13–20× on a typical PhaseShiftRecording → HighpassFilterRecording → CommonReferenceRecording chain for direct get_traces() callers, or ~3× on top of existing CRE parallelism.

[galenlynch]

Test failures are unrelated to this PR

[alejoe91]

Thanks @galenlynch

We have discussed this internally and we think that having per processor thread parallelization could be a common and useful use case.

Implementation-wise, we think we could move the pool creation to core/job_tools.py and make it more general on what functions it applies. Then each extractor could implement a get_traces_multi_thread, which by default would fall back to get_traces, but could be overridden by individual preprocessors (like filter and common_ref). We already have a parameters in the job_kwargs called max_threads_per_worker. This could be use to automatically enable multi-threaded trace retrieval, avoiding to expose the num_workers at the init level. I think this is a better option, since the same class could support different paradigms (e.g., for visualization: 1 process + multi-threaded; for batch processing: N processes + 1 thread) without the need for reinstantiation.

What do you think? We can help with the implementation side :)

[h-mayorquin]

The initial thread is really hard to read. I am fine with LLM contributions but I think they should be curated somehow.

[galenlynch]

@h-mayorquin I can clean up the initial thread. fwiw I did a lot of testing of this PR and in my hands found a massive speedup, particularly when using get_traces for viz which didn't expose parallelism vs TimeSeriesChunkExecutor which seems to be the library-preferred model of parallelism but doesn't work for viz workflows. The performance story is complicated because of TimeSeriesChunkExecutor, and requires explanation as a result.

Edit: original message rewritten.

[galenlynch]

@alejoe91 centralizing threading control seems like a great idea! Adding get_traces_multi_thread per extractor seems like the right approach to me.

One thing worth pointing out: as far as I understand max_threads_per_workers currently controls threadpoolctl.threadpool_limits, which controls parallelism in BLAS/OpenMP:

spikeinterface/src/spikeinterface/core/job_tools.py

Lines 631 to 632 in ec55b00

	with threadpool_limits(limits=self.max_threads_per_worker):
	return self.func(segment_index, start_frame, end_frame, self.worker_dict)

However, scipy.sosfiltfilt and np.median are both single-threaded C code which do not use BLAS, which is why this PR uses explicit Python-level fan-out using ThreadPoolExecutor. That's a different axis of threading than what is controlled by max_threads_per_worker in TimeSeriesChunkExecutor via WorkerFuncWrapper. There is yet another axis of parallelism used by the sister PR #4563 which is numba prange. Having the same number control all of these could open the door to thread oversubscription. Maybe each BaseRecordingSegment's get_traces_multi_thread could decide how to spend that budget?

Maybe one way to do this would be to add building blocks in core/job_tools.py:

# Per-caller-thread pool, what the current PR builds inline in filter.py
def get_inner_pool(max_threads: int) -> ThreadPoolExecutor:
    """Per-caller-thread ThreadPoolExecutor (WeakKeyDictionary keyed by Thread)."""
    ...

@contextmanager
def thread_budget(max_threads: int, *, blas: bool = False, numba: bool = False):
    """Cap the named runtimes for the duration of the context."""
    ...

Dispatch in BaseRecording:

class BaseRecording:
    def get_traces_multi_thread(
        self, *, segment_index=None, start_frame=None, end_frame=None,
        channel_ids=None, max_threads=None,
    ):
        """Like get_traces, but allowed to use up to max_threads
        threads internally. Default falls back to get_traces."""
        if max_threads is None:
            max_threads = get_global_job_kwargs()["max_threads_per_worker"]
        if max_threads <= 1:
            return self.get_traces(...)
        rs = self.segments[segment_index]
        return rs.get_traces_multi_thread(start_frame, end_frame,
                                          channel_indices, max_threads)

And different extractors could choose how to spend the budget:

# Filter (this PR, wants ThreadPoolExecutor):
class FilterRecordingSegment:
    def get_traces_multi_thread(self, start, end, channel_indices, max_threads):
        traces = self._fetch_with_margin(start, end, channel_indices)
        pool = get_inner_pool(max_threads)
        return self._apply_sos_parallel(traces, pool)  # GIL-released scipy

# CMR (this PR, wants ThreadPoolExecutor):
class CommonReferenceRecordingSegment:
    def get_traces_multi_thread(self, start, end, channel_indices, max_threads):
        traces = self._fetch(start, end, channel_indices)
        pool = get_inner_pool(max_threads)
        return self._parallel_reduce_axis1(traces, pool)

# PhaseShift FIR (sister PR, numba prange only):
class PhaseShiftRecordingSegment:
    def get_traces_multi_thread(self, start, end, channel_indices, max_threads):
        traces = self._fetch(...)
        with thread_budget(max_threads, numba=True):
            return _sinc_fir_kernel_tc(traces, self._kernels)

# Hypothetical BLAS-heavy segment (example, BLAS threading only):
class WhitenRecordingSegment:
    def get_traces_multi_thread(self, start, end, channel_indices, max_threads):
        traces = self._fetch(...)
        with thread_budget(max_threads, blas=True):
            return self._W @ traces.T