Skip to content

Release the shared memory accumulator when template estimation fails - #4732

Open
adityasingh2400 wants to merge 1 commit into
SpikeInterface:mainfrom
adityasingh2400:fix-4566-shm-leak
Open

Release the shared memory accumulator when template estimation fails#4732
adityasingh2400 wants to merge 1 commit into
SpikeInterface:mainfrom
adityasingh2400:fix-4566-shm-leak

Conversation

@adityasingh2400

Copy link
Copy Markdown

Related to #4566.

estimate_templates_with_accumulator() allocates a shared memory accumulator, runs the chunk executor over it, and then unlinks it. The unlink and close calls only sit on the success path, so any exception raised inside processor.run(), or in the averaging that follows, returns without ever releasing the segment. Python then prints "There appear to be 1 leaked shared_memory objects to clean up at shutdown", which is exactly the message in the report, and the orphaned segment keeps occupying the shared memory area that a retry needs.

This change wraps the whole body in try/finally so the segment is released either way. The numpy views are still dropped before close() so numpy does not raise "cannot close exported pointers exist", and unlink() now runs before close() so a failing close cannot leave the segment behind. The returned arrays come from np.sum and np.zeros_like, so they do not view the shared buffers and are safe to return after the release.

Two sibling functions in the same module had the same unprotected release and are fixed the same way. extract_waveforms_to_single_buffer() leaked its buffer when the run failed, in both copy modes, because the caller never receives the handle when the call raises. The median branch of estimate_templates() leaked its buffer if the median computation raised.

I checked the four other make_shared_array call sites and deliberately left them alone, since they hand ownership to the caller: allocate_waveforms_buffers, numpyextractors.py:527, time_series_tools.py:241, and the matching components.

The second part of this change is about the error message. The accumulator is allocated per worker, so its size is n_jobs * num_units * num_samples * num_channels * 4 bytes. When estimate_sparsity() is the caller there is no sparsity mask yet, so num_channels is the full channel count. For the reporter's numbers, roughly 1500 units on 384 channels with 105 samples, that is about 2.4 GB of shared memory at 10 jobs. On Linux /dev/shm defaults to half of RAM, so the allocation fails with a bare OSError: [Errno 28] No space left on device, which reads nothing like a job or sparsity problem. That failure is now re-raised as a MemoryError that prints the requested size, spells out the four factors that produced it, and points at lowering n_jobs or passing an explicit sparsity. No pre-emptive cap is added, so nothing that works today starts refusing to run.

To be clear about what this does not do. It does not change the fact that the accumulator scales with n_jobs * num_units, and that scaling is the actual root cause of the crash in #4566. It matches the reporter's own measurements well: 10 cores topped out around 800 units, 6 cores around 1000, and 4 cores handled all 1500, so the product of cores and units is roughly constant, which is what a per worker allocation predicts. Working through the three data points at 30 kHz with the default 1.0 and 2.5 ms window on 384 channels in float32 gives 1.29 GB for 10 by 800 which fails, and 0.97 GB for both 6 by 1000 and 4 by 1500 which work, so there is a fixed ceiling between 1.0 and 1.3 GB.

Reworking the accumulator so it does not scale with the worker count is a much larger design change and I did not want to bundle it into this PR, so I have left that decision to you. What this change does fix is that a failure no longer leaves shared memory behind to make the next attempt worse, and that the failure now says what went wrong.

Tests are added in src/spikeinterface/core/tests/test_waveform_tools.py. One records every segment created through make_shared_array, forces the per chunk worker to raise, and asserts every recorded segment was unlinked, covering both the return_std=False and return_std=True shapes. The other asserts the allocation failure surfaces as a MemoryError mentioning n_jobs. Both fail on current main.

Reverting the source to main makes all three new tests fail. They pass after, and the regression runs complete at 12 passed for waveform_tools plus sparsity, 14 passed adding template_tools, and 41 passed for sortinganalyzer plus analyzer_extension_core. black reports both files unchanged.

Disclosure: this change was prepared with AI assistance. I have reviewed and tested it.

estimate_templates_with_accumulator() only unlinked its shared memory
accumulator on the success path, so any error raised inside the parallel
run left the segment behind. Python then reports "There appear to be 1
leaked shared_memory objects to clean up at shutdown", and the orphaned
segment keeps consuming the shared memory that a retry needs.

The accumulator is allocated per worker, so its size is
n_jobs * num_units * num_samples * num_channels * 4 bytes. Allocating it
can therefore fail on a machine with a small /dev/shm, which surfaced as
an opaque OSError with errno 28. That failure now names the sizes that
caused it and points at n_jobs and sparsity.

The same unprotected release is fixed in extract_waveforms_to_single_buffer()
and in the median branch of estimate_templates().
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant