diff --git a/src/spikeinterface/core/tests/test_waveform_tools.py b/src/spikeinterface/core/tests/test_waveform_tools.py index 5e0350f833..f02c5e8fc0 100644 --- a/src/spikeinterface/core/tests/test_waveform_tools.py +++ b/src/spikeinterface/core/tests/test_waveform_tools.py @@ -5,6 +5,7 @@ import numpy as np +import spikeinterface.core.waveform_tools as waveform_tools from spikeinterface.core import generate_recording, generate_sorting, generate_ground_truth_recording, ms_to_samples from spikeinterface.core.waveform_tools import ( extract_waveforms_to_buffers, @@ -212,6 +213,106 @@ def test_estimate_templates_with_accumulator(): # plt.show() +class _SharedArrayRecorder: + """ + Drop-in replacement for `make_shared_array` that records every segment it allocates and + whether `unlink()` was called on it. + """ + + def __init__(self): + self.segments = [] + + def __call__(self, shape, dtype): + from spikeinterface.core.core_tools import make_shared_array + + arr, shm = make_shared_array(shape, dtype) + record = {"shm": shm, "unlinked": False} + self.segments.append(record) + real_unlink = shm.unlink + + def tracked_unlink(): + record["unlinked"] = True + real_unlink() + + shm.unlink = tracked_unlink + return arr, shm + + def cleanup(self): + # do not let a failing test leave real segments behind + for record in self.segments: + if not record["unlinked"]: + record["shm"].unlink() + record["unlinked"] = True + + +@pytest.mark.parametrize("return_std", [False, True]) +def test_estimate_templates_with_accumulator_releases_shared_memory_on_error(monkeypatch, return_std): + # when the parallel run raises, the shared memory accumulator must still be unlinked. + # otherwise the segment survives the failure and python reports + # "There appear to be N leaked shared_memory objects to clean up at shutdown". See issue #4566. + recording, sorting = get_dataset() + + nbefore = ms_to_samples(1.0, recording.sampling_frequency) + nafter = ms_to_samples(1.5, recording.sampling_frequency) + spikes = sorting.to_spike_vector()[::10] + + recorder = _SharedArrayRecorder() + monkeypatch.setattr(waveform_tools, "make_shared_array", recorder) + + def failing_worker(*args, **kwargs): + raise RuntimeError("boom") + + monkeypatch.setattr(waveform_tools, "_worker_estimate_templates", failing_worker) + + try: + with pytest.raises(RuntimeError, match="boom"): + estimate_templates_with_accumulator( + recording, + spikes, + sorting.unit_ids, + nbefore, + nafter, + return_in_uV=True, + return_std=return_std, + n_jobs=1, + chunk_duration="1s", + progress_bar=False, + ) + + assert len(recorder.segments) == (2 if return_std else 1) + assert all(record["unlinked"] for record in recorder.segments) + finally: + recorder.cleanup() + + +def test_estimate_templates_with_accumulator_allocation_error_is_explicit(monkeypatch): + # a failed allocation of the shared accumulator is an opaque "OSError: [Errno 28] No space left + # on device". It should instead name the sizes that produced it. See issue #4566. + recording, sorting = get_dataset() + + nbefore = ms_to_samples(1.0, recording.sampling_frequency) + nafter = ms_to_samples(1.5, recording.sampling_frequency) + spikes = sorting.to_spike_vector()[::10] + + def no_space(shape, dtype): + raise OSError(28, "No space left on device") + + monkeypatch.setattr(waveform_tools, "make_shared_array", no_space) + + with pytest.raises(MemoryError, match="n_jobs"): + estimate_templates_with_accumulator( + recording, + spikes, + sorting.unit_ids, + nbefore, + nafter, + return_in_uV=True, + n_jobs=2, + chunk_duration="1s", + progress_bar=False, + ) + + def test_estimate_templates(): recording, sorting = get_dataset() diff --git a/src/spikeinterface/core/waveform_tools.py b/src/spikeinterface/core/waveform_tools.py index 70f29d94b0..8695589b39 100644 --- a/src/spikeinterface/core/waveform_tools.py +++ b/src/spikeinterface/core/waveform_tools.py @@ -566,7 +566,16 @@ def extract_waveforms_to_single_buffer( processor = TimeSeriesChunkExecutor( recording, func, init_func, init_args, job_name=job_name, verbose=verbose, **job_kwargs ) - processor.run() + try: + processor.run() + except Exception: + # the buffer is never handed over to the caller when the run fails, so nobody else can + # release it. Without this, a failed run leaves a shared memory segment behind. + if mode == "shared_memory" and shm is not None: + del all_waveforms + shm.unlink() + shm.close() + raise if mode == "memmap": return all_waveforms @@ -830,16 +839,21 @@ def estimate_templates( copy=False, **job_kwargs, ) - templates_array = np.zeros( - (len(unit_ids), all_waveforms.shape[1], all_waveforms.shape[2]), dtype=all_waveforms.dtype - ) - for unit_index, unit_id in enumerate(unit_ids): - wfs = all_waveforms[spikes["unit_index"] == unit_index] - templates_array[unit_index, :, :] = np.median(wfs, axis=0) - # release shared memory after the median - del all_waveforms - wf_array_info["shm"].close() - wf_array_info["shm"].unlink() + try: + templates_array = np.zeros( + (len(unit_ids), all_waveforms.shape[1], all_waveforms.shape[2]), dtype=all_waveforms.dtype + ) + for unit_index, unit_id in enumerate(unit_ids): + wfs = all_waveforms[spikes["unit_index"] == unit_index] + templates_array[unit_index, :, :] = np.median(wfs, axis=0) + finally: + # release shared memory after the median, also when the median raises + del all_waveforms + shm = wf_array_info["shm"] + if shm is not None: + # empty arrays have no shared memory + shm.unlink() + shm.close() else: raise ValueError(f"estimate_templates(..., operator={operator}) wrong operator must be average or median") @@ -847,6 +861,31 @@ def estimate_templates( return templates_array +def _allocate_accumulator_shared_array(shape, dtype): + """ + Allocate one shared memory accumulator for `estimate_templates_with_accumulator()`. + + The accumulator is allocated per worker, so its size scales with n_jobs and the allocation + can fail on machines with a small shared memory area (/dev/shm on Linux). The raw failure is + an opaque "OSError: [Errno 28] No space left on device", so it is re-raised with the sizes + that caused it. + """ + try: + return make_shared_array(shape, dtype) + except (OSError, MemoryError) as err: + dtype = np.dtype(dtype) + num_worker, num_units, num_samples, num_chans = shape + nbytes = int(np.prod(shape)) * dtype.itemsize + size = f"{nbytes / 1e9:.2f} GB" if nbytes >= 1e9 else f"{nbytes / 1e6:.2f} MB" + raise MemoryError( + f"estimate_templates_with_accumulator() could not allocate its shared memory accumulator of " + f"{size}. The accumulator is allocated per worker, so its size is " + f"n_jobs * num_units * num_samples * num_channels * itemsize = " + f"{num_worker} * {num_units} * {num_samples} * {num_chans} * {dtype.itemsize} bytes. " + f"Lower n_jobs, or pass an explicit sparsity so that num_channels is smaller." + ) from err + + def estimate_templates_with_accumulator( recording: BaseRecording, spikes: np.ndarray, @@ -921,68 +960,86 @@ def estimate_templates_with_accumulator( shape = (num_worker, num_units, nbefore + nafter, num_chans) dtype = np.dtype("float32") - waveform_accumulator_per_worker, shm = make_shared_array(shape, dtype) + waveform_accumulator_per_worker, shm = _allocate_accumulator_shared_array(shape, dtype) shm_name = shm.name - if return_std: - waveform_squared_accumulator_per_worker, shm_squared = make_shared_array(shape, dtype) - shm_squared_name = shm_squared.name - else: - waveform_squared_accumulator_per_worker = None - shm_squared_name = None - - func = _worker_estimate_templates - init_func = _init_worker_estimate_templates + waveform_squared_accumulator_per_worker = None + shm_squared = None + shm_squared_name = None - init_args = ( - recording, - spikes, - shm_name, - shm_squared_name, - shape, - dtype, - nbefore, - nafter, - return_in_uV, - sparsity_mask, - ) + # everything below must release the shared memory, including on error. A segment that is not + # unlinked outlives the failure, keeps consuming the resource that is already scarce here, and + # is reported at interpreter shutdown as "leaked shared_memory objects". + try: + if return_std: + waveform_squared_accumulator_per_worker, shm_squared = _allocate_accumulator_shared_array(shape, dtype) + shm_squared_name = shm_squared.name - if job_name is None: - job_name = "estimate_templates_with_accumulator" - processor = TimeSeriesChunkExecutor( - recording, func, init_func, init_args, job_name=job_name, verbose=verbose, need_worker_index=True, **job_kwargs - ) - processor.run() + func = _worker_estimate_templates + init_func = _init_worker_estimate_templates - # average - waveforms_sum = np.sum(waveform_accumulator_per_worker, axis=0) - if return_std: - # we need a copy here because we will use the means to compute the stds - template_means = waveforms_sum.copy() - else: - # waveforms_sum will also be changed in this case when acting on template_means - template_means = waveforms_sum + init_args = ( + recording, + spikes, + shm_name, + shm_squared_name, + shape, + dtype, + nbefore, + nafter, + return_in_uV, + sparsity_mask, + ) - unit_indices, spike_count = np.unique(spikes["unit_index"], return_counts=True) - template_means[unit_indices, :, :] /= spike_count[:, np.newaxis, np.newaxis] + if job_name is None: + job_name = "estimate_templates_with_accumulator" + processor = TimeSeriesChunkExecutor( + recording, + func, + init_func, + init_args, + job_name=job_name, + verbose=verbose, + need_worker_index=True, + **job_kwargs, + ) + processor.run() - if return_std: - waveforms_squared_sum = np.sum(waveform_squared_accumulator_per_worker, axis=0) - # standard deviation - template_stds = np.zeros_like(template_means) - for unit_index, count in zip(unit_indices, spike_count): - residuals = ( - waveforms_squared_sum[unit_index] - 2 * template_means[unit_index] * waveforms_sum[unit_index] - ) + count * template_means[unit_index] ** 2 - residuals[residuals < 0] = 0 - template_stds[unit_index] = np.sqrt(residuals / count) + # average + waveforms_sum = np.sum(waveform_accumulator_per_worker, axis=0) + if return_std: + # we need a copy here because we will use the means to compute the stds + template_means = waveforms_sum.copy() + else: + # waveforms_sum will also be changed in this case when acting on template_means + template_means = waveforms_sum + + unit_indices, spike_count = np.unique(spikes["unit_index"], return_counts=True) + template_means[unit_indices, :, :] /= spike_count[:, np.newaxis, np.newaxis] + + if return_std: + waveforms_squared_sum = np.sum(waveform_squared_accumulator_per_worker, axis=0) + # standard deviation + template_stds = np.zeros_like(template_means) + for unit_index, count in zip(unit_indices, spike_count): + residuals = ( + waveforms_squared_sum[unit_index] - 2 * template_means[unit_index] * waveforms_sum[unit_index] + ) + count * template_means[unit_index] ** 2 + residuals[residuals < 0] = 0 + template_stds[unit_index] = np.sqrt(residuals / count) + finally: + # important : release the sharedmem. + # the numpy views on the buffers must be dropped before closing, otherwise numpy raises + # "BufferError: cannot close exported pointers exist". unlink() is what actually frees the + # segment so it comes first, a failing close() must not leave the segment behind. + del waveform_accumulator_per_worker del waveform_squared_accumulator_per_worker - shm_squared.unlink() - shm_squared.close() + shm.unlink() + shm.close() + if shm_squared is not None: + shm_squared.unlink() + shm_squared.close() - # important : release the sharedmem - del waveform_accumulator_per_worker - shm.unlink() - shm.close() + # the returned arrays come from np.sum / np.zeros_like, so they do not view the shared buffers if return_std: return template_means, template_stds