Release the shared memory accumulator when template estimation fails - #4732
Open
adityasingh2400 wants to merge 1 commit into
Open
Release the shared memory accumulator when template estimation fails#4732adityasingh2400 wants to merge 1 commit into
adityasingh2400 wants to merge 1 commit into
Conversation
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().
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 insideprocessor.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", andunlink()now runs beforeclose()so a failing close cannot leave the segment behind. The returned arrays come fromnp.sumandnp.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 ofestimate_templates()leaked its buffer if the median computation raised.I checked the four other
make_shared_arraycall 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 * 4bytes. Whenestimate_sparsity()is the caller there is no sparsity mask yet, sonum_channelsis 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/shmdefaults to half of RAM, so the allocation fails with a bareOSError: [Errno 28] No space left on device, which reads nothing like a job or sparsity problem. That failure is now re-raised as aMemoryErrorthat prints the requested size, spells out the four factors that produced it, and points at loweringn_jobsor 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 throughmake_shared_array, forces the per chunk worker to raise, and asserts every recorded segment was unlinked, covering both thereturn_std=Falseandreturn_std=Trueshapes. The other asserts the allocation failure surfaces as aMemoryErrormentioningn_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.
blackreports both files unchanged.Disclosure: this change was prepared with AI assistance. I have reviewed and tested it.