From 6cce4218ec6ad3e28a3122ab59d2a71cd929f5c4 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 13:18:57 +0200 Subject: [PATCH 1/4] Read FIF simple numeric tags through a memory map FIF raw segments are read as byte-offset views into a PID-keyed memory map of the file instead of open/seek/read per call; buffer entries are selected with searchsorted on the sorted bounds. gzip, file-like objects, and non-simple tag types keep the legacy path. The generic memory-map cache in _read_segments_file also serves the other binary readers. --- mne/_fiff/utils.py | 78 ++++++++++++++++++++++++++++++++++-- mne/io/fiff/raw.py | 98 +++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 163 insertions(+), 13 deletions(-) diff --git a/mne/_fiff/utils.py b/mne/_fiff/utils.py index 2d9c0b0d53c..968bb9f8fd0 100644 --- a/mne/_fiff/utils.py +++ b/mne/_fiff/utils.py @@ -84,13 +84,14 @@ def _mult_cal_one(data_view, one, idx, cals, mult): else: assert cals is not None if isinstance(idx, slice): - # Hot path: gather + type-cast + calibration in a single pass - # (was three passes plus a full float64 temporary). - # Benchmark (128 ch x 1024 samples): ~85 -> ~30 us per call - # on BrainVision/FIF window reads. + # Hot path: gather + type-cast + calibration in a single pass, + # without materializing an intermediate float64 copy of `one` + # (`one[idx]` is a view for basic slices). Numerically identical + # to cast-then-scale because both are elementwise. np.multiply(one[idx], cals.reshape(-1, 1), out=data_view, casting="unsafe") else: one = np.asarray(one, dtype=data_view.dtype) + # faster than doing one = one[idx] np.take(one, idx, axis=0, out=data_view) data_view *= cals @@ -219,6 +220,8 @@ def _read_segments_file( if n_channels is None: n_channels = raw._raw_extras[fi]["orig_nchan"] + import os as _os + n_bytes = np.dtype(dtype).itemsize # data_offset and data_left count data samples (channels x time points), # not bytes. @@ -228,6 +231,51 @@ def _read_segments_file( # Read up to 100 MB of data at a time, block_size is in data samples block_size = ((int(100e6) // n_bytes) // n_channels) * n_channels block_size = min(data_left, block_size) + + # Reuse a memory map across calls (keyed by PID so forked processes -- + # e.g., PyTorch DataLoader workers -- create their own mapping instead of + # sharing one). This removes the per-call open/seek/syscall overhead. + ex = raw._raw_extras[fi] if fi < len(raw._raw_extras) else {} + mm = ex.get("_mm") if isinstance(ex, dict) else None + if mm is not None and ex.get("_mm_pid") != _os.getpid(): + mm = None + if mm is not None and ( + mm.dtype != np.dtype(dtype) + or mm.size * n_bytes < data_offset + data_left * n_bytes + ): + mm = None + if mm is None and isinstance(ex, dict): + try: + mm = np.memmap(raw.filenames[fi], dtype=dtype, mode="r") + ex["_mm"] = mm + ex["_mm_pid"] = _os.getpid() + except Exception: + mm = None + + if mm is not None: + base_idx = data_offset // n_bytes + for sample_start in np.arange(0, data_left, block_size) // n_channels: + count = min(block_size, data_left - sample_start * n_channels) + block = mm[ + base_idx + sample_start * n_channels : base_idx + + sample_start * n_channels + + count + ] + if block.size != count: + raise RuntimeError( + f"Incorrect number of samples ({block.size} != {count}), " + "please report this error to MNE-Python developers" + ) + block = block.reshape(n_channels, -1, order="F") + n_samples = block.shape[1] + sample_stop = sample_start + n_samples + if trigger_ch is not None: + stim_ch = trigger_ch[start:stop][sample_start:sample_stop] + block = np.vstack((block, stim_ch)) + data_view = data[:, sample_start:sample_stop] + _mult_cal_one(data_view, block, idx, cals, mult) + return + with open(raw.filenames[fi], "rb", buffering=0) as fid: fid.seek(data_offset) # extract data in chunks @@ -333,3 +381,25 @@ def _make_split_fnames(fname, n_splits, split_naming): path = Path(_construct_bids_filename(base, ext, i)) res.append(path) return res + + +def _memmap_for(extras, fname): + """Return a PID-keyed read-only uint8 memmap of *fname* from *extras*. + + The mapping is created lazily on first call and cached in *extras* (a + per-instance dict) together with the PID that created it, so forked worker + processes build their own mapping instead of sharing inherited state. + Returns None if the file cannot be mapped. There is deliberately no + staleness check: callers index the mapping through tables read at open + time (bounds, entries), which are invalid if the file changes anyway, + so a per-call stat would only add overhead. + """ + mm = extras.get("_mm") + if mm is not None and extras.get("_mm_pid") == os.getpid(): + return mm + try: + mm = np.memmap(str(fname), dtype=np.uint8, mode="r") + except Exception: + return None + extras["_mm"], extras["_mm_pid"] = mm, os.getpid() + return mm diff --git a/mne/io/fiff/raw.py b/mne/io/fiff/raw.py index 95c6db5dbec..07d58c77a09 100644 --- a/mne/io/fiff/raw.py +++ b/mne/io/fiff/raw.py @@ -12,9 +12,9 @@ from ..._fiff.constants import FIFF from ..._fiff.meas_info import read_meas_info from ..._fiff.open import _fiff_get_fid, _get_next_fname, fiff_open -from ..._fiff.tag import _call_dict, read_tag +from ..._fiff.tag import _call_dict, _simple_dict, read_tag from ..._fiff.tree import dir_tree_find -from ..._fiff.utils import _mult_cal_one +from ..._fiff.utils import _memmap_for, _mult_cal_one from ...annotations import Annotations, _read_annotations_fif from ...channels import fix_mag_coil_types from ...event import AcqParserFIF @@ -39,6 +39,40 @@ @fill_doc +def _fif_mm_plan(bounds, ents, entry_span, nchan, file_mapping, start, stop): + """Plan byte-offset reads covering [start, stop) from a FIF memory map. + + Returns a list of ``(dtype, byte_offset, n_bytes, n_samples)`` read tuples, + one per touched buffer entry, or None when the fast path does not apply: a + touched entry is missing (gaps are zero-filled by the legacy loop), holds a + non-simple tag type, or has a declared size inconsistent with ``nchan`` + samples. + """ + read_plan = [] + for ei in entry_span: + entry = ents[ei] + if entry is None or entry.type not in _simple_dict: + return None + dtype = np.dtype(_simple_dict[entry.type]) + n_entry_samples = bounds[ei + 1] - bounds[ei] + n_bytes_per_sample = dtype.itemsize * nchan + if getattr(entry, "size", None) != (n_entry_samples * n_bytes_per_sample): + return None + # samples of this entry that fall inside [start, stop) + first_pick = max(start - bounds[ei], 0) + last_pick = min(n_entry_samples, stop - bounds[ei]) + n_samples = last_pick - first_pick + if n_samples <= 0: + continue + # payload starts 16 bytes into the tag (header is four uint32: kind, + # type, size, next -- see _read_tag_header and Tag.next_pos) + byte_offset = entry.pos + 16 + first_pick * n_bytes_per_sample + read_plan.append( + (dtype, byte_offset, n_samples * n_bytes_per_sample, n_samples) + ) + return read_plan + + class Raw(BaseRaw): """Raw data in FIF format. @@ -401,15 +435,61 @@ def _dtype(self): return dtype def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): - """Read a segment of data from a file.""" + """Read a segment of data from a file. + + Writes ``data`` (shape ``(len(idx), stop - start)``) for samples + ``[start, stop)``, rows following ``idx`` order. Two equivalent + implementations: a memory-map fast path for simple uncompressed tags + (see the comment block below), and the legacy buffered loop. + """ n_bad = 0 - with _fiff_get_fid(self._raw_extras[fi]["filename"]) as fid: - bounds = self._raw_extras[fi]["bounds"] - ents = self._raw_extras[fi]["ent"] - nchan = self._raw_extras[fi]["orig_nchan"] - use = (stop > bounds[:-1]) & (start < bounds[1:]) + bounds = self._raw_extras[fi]["bounds"] + ents = self._raw_extras[fi]["ent"] + nchan = self._raw_extras[fi]["orig_nchan"] + fname = self._raw_extras[fi]["filename"] + # Entries overlapping [start, stop) via binary search on sorted bounds + # (O(log n) instead of a mask over every entry; matters for long + # recordings with thousands of buffer entries). + # indices of buffer entries overlapping [start, stop) + entry_span = range( + max(np.searchsorted(bounds, start, side="right") - 1, 0), + min(np.searchsorted(bounds, stop, side="left"), len(bounds) - 1), + ) + + # Fast path: serve [start, stop) from one persistent memory map of + # the file; see _fif_mm_plan() for eligibility and fallbacks. + fname = self._raw_extras[fi]["filename"] + file_mapping = ( + _memmap_for(self._raw_extras[fi], fname) + if isinstance(fname, Path) and fname.suffixes[-1] != ".gz" + else None + ) + read_plan = ( + _fif_mm_plan(bounds, ents, entry_span, nchan, file_mapping, start, stop) + if file_mapping is not None + else None + ) + if read_plan is not None: + col_start = 0 + for dtype, byte_offset, n_bytes, n_samples in read_plan: + values = np.frombuffer( + file_mapping[byte_offset : byte_offset + n_bytes], + dtype=dtype, + count=n_samples * nchan, + ).reshape(n_samples, nchan) + _mult_cal_one( + data[:, col_start : col_start + n_samples], + values.T, + idx, + cals, + mult, + ) + col_start += n_samples + return + + with _fiff_get_fid(fname) as fid: offset = 0 - for ei in np.where(use)[0]: + for ei in entry_span: first = bounds[ei] last = bounds[ei + 1] nsamp = last - first From 9c46efe684d58037feee1068f935242f8dca2703 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 15:21:37 +0200 Subject: [PATCH 2/4] Fix memory-map reader fallbacks and lifecycle --- doc/changes/dev/14213.newfeature.rst | 1 + mne/_fiff/tests/test_utils.py | 82 ++++++++++++++++++++++++++- mne/_fiff/utils.py | 85 ++++++++++++++++------------ mne/io/base.py | 11 ++-- mne/io/fiff/raw.py | 56 +++++++++--------- mne/io/tests/test_raw.py | 16 +++++- 6 files changed, 177 insertions(+), 74 deletions(-) create mode 100644 doc/changes/dev/14213.newfeature.rst diff --git a/doc/changes/dev/14213.newfeature.rst b/doc/changes/dev/14213.newfeature.rst new file mode 100644 index 00000000000..8d3f77e9ecd --- /dev/null +++ b/doc/changes/dev/14213.newfeature.rst @@ -0,0 +1 @@ +Speed up FIF reading by reading tag payloads through a memory map instead of open/seek/read per call, by `Bruno Aristimunha`_ diff --git a/mne/_fiff/tests/test_utils.py b/mne/_fiff/tests/test_utils.py index ba6748826a3..48757ce019a 100644 --- a/mne/_fiff/tests/test_utils.py +++ b/mne/_fiff/tests/test_utils.py @@ -4,7 +4,15 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -from mne._fiff.utils import _check_orig_units +import pickle +from copy import deepcopy +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +from numpy.testing import assert_array_equal + +from mne._fiff.utils import _check_orig_units, _memmap_for, _read_segments_file def test_check_orig_units(): @@ -16,3 +24,75 @@ def test_check_orig_units(): assert orig_units["Pz"] == "µV" assert orig_units["greekMu"] == "µV" assert orig_units["microSign"] == "µV" + + +def test_read_segments_file_unaligned_offset(tmp_path): + """Test reading multi-byte data at an unaligned byte offset.""" + fname = tmp_path / "test.bin" + values = np.array([[1, 2, 3], [10, 20, 30]], dtype=" None: - """Clean up the object. - - Does nothing for objects that close their file descriptors. - Things like Raw will override this method. - """ - pass # noqa + """Clean up any resources used by the object.""" + for extras in self._raw_extras: + cache = extras.get("_memmap_cache") + if cache is not None: + cache.close() def copy(self) -> Self: """Return copy of the instance. diff --git a/mne/io/fiff/raw.py b/mne/io/fiff/raw.py index 07d58c77a09..ffa080f0804 100644 --- a/mne/io/fiff/raw.py +++ b/mne/io/fiff/raw.py @@ -38,8 +38,7 @@ ) -@fill_doc -def _fif_mm_plan(bounds, ents, entry_span, nchan, file_mapping, start, stop): +def _fif_mm_plan(bounds, ents, entry_span, nchan, start, stop): """Plan byte-offset reads covering [start, stop) from a FIF memory map. Returns a list of ``(dtype, byte_offset, n_bytes, n_samples)`` read tuples, @@ -73,6 +72,7 @@ def _fif_mm_plan(bounds, ents, entry_span, nchan, file_mapping, start, stop): return read_plan +@fill_doc class Raw(BaseRaw): """Raw data in FIF format. @@ -458,34 +458,30 @@ def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): # Fast path: serve [start, stop) from one persistent memory map of # the file; see _fif_mm_plan() for eligibility and fallbacks. - fname = self._raw_extras[fi]["filename"] - file_mapping = ( - _memmap_for(self._raw_extras[fi], fname) - if isinstance(fname, Path) and fname.suffixes[-1] != ".gz" - else None - ) - read_plan = ( - _fif_mm_plan(bounds, ents, entry_span, nchan, file_mapping, start, stop) - if file_mapping is not None - else None - ) - if read_plan is not None: - col_start = 0 - for dtype, byte_offset, n_bytes, n_samples in read_plan: - values = np.frombuffer( - file_mapping[byte_offset : byte_offset + n_bytes], - dtype=dtype, - count=n_samples * nchan, - ).reshape(n_samples, nchan) - _mult_cal_one( - data[:, col_start : col_start + n_samples], - values.T, - idx, - cals, - mult, - ) - col_start += n_samples - return + read_plan = _fif_mm_plan(bounds, ents, entry_span, nchan, start, stop) + if ( + read_plan is not None + and isinstance(fname, Path) + and fname.suffixes[-1] != ".gz" + ): + file_mapping = _memmap_for(self._raw_extras[fi], fname) + if file_mapping is not None: + col_start = 0 + for dtype, byte_offset, n_bytes, n_samples in read_plan: + values = np.frombuffer( + file_mapping[byte_offset : byte_offset + n_bytes], + dtype=dtype, + count=n_samples * nchan, + ).reshape(n_samples, nchan) + _mult_cal_one( + data[:, col_start : col_start + n_samples], + values.T, + idx, + cals, + mult, + ) + col_start += n_samples + return with _fiff_get_fid(fname) as fid: offset = 0 diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py index 909ff9504ae..1218fe1c3b1 100644 --- a/mne/io/tests/test_raw.py +++ b/mne/io/tests/test_raw.py @@ -27,7 +27,7 @@ from mne._fiff.meas_info import Info, _get_valid_units, _writing_info_hdf5 from mne._fiff.pick import _ELECTRODE_CH_TYPES, _FNIRS_CH_TYPES_SPLIT from mne._fiff.proj import Projection -from mne._fiff.utils import _mult_cal_one +from mne._fiff.utils import _memmap_for, _mult_cal_one from mne.io import BaseRaw, RawArray, read_raw_fif from mne.io.base import _get_scaling from mne.transforms import Transform @@ -838,6 +838,20 @@ def _read_raw_arange(preload=False, verbose=None): return _RawArange(preload, verbose) +def test_raw_close_memmap_cache(tmp_path): + """Test that closing Raw releases cached memory maps.""" + fname = tmp_path / "test.bin" + fname.write_bytes(b"test") + with _read_raw_arange() as raw: + cache = raw._raw_extras[0] + mapping = _memmap_for(cache, fname) + assert not mapping._mmap.closed + + assert mapping._mmap.closed + assert cache["_memmap_cache"].mapping is None + assert cache["_memmap_cache"].pid is None + + def test_load_data_memmap(tmp_path): """Test loading raw data into a memmap via load_data.""" raw = _read_raw_arange(preload=False) From fb27de946659af9323d84aaf51555d926377517c Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 18:53:15 +0200 Subject: [PATCH 3/4] Withdraw persistent source-file memory mapping --- doc/changes/dev/14213.newfeature.rst | 1 - mne/_fiff/tests/test_utils.py | 82 +----------------------- mne/_fiff/utils.py | 91 ++------------------------- mne/io/base.py | 11 ++-- mne/io/fiff/raw.py | 94 +++------------------------- mne/io/tests/test_raw.py | 16 +---- 6 files changed, 21 insertions(+), 274 deletions(-) delete mode 100644 doc/changes/dev/14213.newfeature.rst diff --git a/doc/changes/dev/14213.newfeature.rst b/doc/changes/dev/14213.newfeature.rst deleted file mode 100644 index 8d3f77e9ecd..00000000000 --- a/doc/changes/dev/14213.newfeature.rst +++ /dev/null @@ -1 +0,0 @@ -Speed up FIF reading by reading tag payloads through a memory map instead of open/seek/read per call, by `Bruno Aristimunha`_ diff --git a/mne/_fiff/tests/test_utils.py b/mne/_fiff/tests/test_utils.py index 48757ce019a..ba6748826a3 100644 --- a/mne/_fiff/tests/test_utils.py +++ b/mne/_fiff/tests/test_utils.py @@ -4,15 +4,7 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -import pickle -from copy import deepcopy -from types import SimpleNamespace -from unittest.mock import Mock - -import numpy as np -from numpy.testing import assert_array_equal - -from mne._fiff.utils import _check_orig_units, _memmap_for, _read_segments_file +from mne._fiff.utils import _check_orig_units def test_check_orig_units(): @@ -24,75 +16,3 @@ def test_check_orig_units(): assert orig_units["Pz"] == "µV" assert orig_units["greekMu"] == "µV" assert orig_units["microSign"] == "µV" - - -def test_read_segments_file_unaligned_offset(tmp_path): - """Test reading multi-byte data at an unaligned byte offset.""" - fname = tmp_path / "test.bin" - values = np.array([[1, 2, 3], [10, 20, 30]], dtype=" ~30 us per call + # on BrainVision/FIF window reads. np.multiply(one[idx], cals.reshape(-1, 1), out=data_view, casting="unsafe") else: one = np.asarray(one, dtype=data_view.dtype) - # faster than doing one = one[idx] np.take(one, idx, axis=0, out=data_view) data_view *= cals @@ -229,37 +228,6 @@ def _read_segments_file( # Read up to 100 MB of data at a time, block_size is in data samples block_size = ((int(100e6) // n_bytes) // n_channels) * n_channels block_size = min(data_left, block_size) - - # Reuse a memory map across calls (keyed by PID so forked processes -- - # e.g., PyTorch DataLoader workers -- create their own mapping instead of - # sharing one). This removes the per-call open/seek/syscall overhead. - extras = raw._raw_extras[fi] if fi < len(raw._raw_extras) else {} - mm = _memmap_for(extras, raw.filenames[fi]) if isinstance(extras, dict) else None - if mm is not None and mm.size < data_offset + data_left * n_bytes: - mm = None - - if mm is not None: - for sample_start in np.arange(0, data_left, block_size) // n_channels: - count = min(block_size, data_left - sample_start * n_channels) - byte_start = data_offset + sample_start * n_channels * n_bytes - block = np.frombuffer( - mm[byte_start : byte_start + count * n_bytes], dtype=dtype, count=count - ) - if block.size != count: - raise RuntimeError( - f"Incorrect number of samples ({block.size} != {count}), " - "please report this error to MNE-Python developers" - ) - block = block.reshape(n_channels, -1, order="F") - n_samples = block.shape[1] - sample_stop = sample_start + n_samples - if trigger_ch is not None: - stim_ch = trigger_ch[start:stop][sample_start:sample_stop] - block = np.vstack((block, stim_ch)) - data_view = data[:, sample_start:sample_stop] - _mult_cal_one(data_view, block, idx, cals, mult) - return - with open(raw.filenames[fi], "rb", buffering=0) as fid: fid.seek(data_offset) # extract data in chunks @@ -365,54 +333,3 @@ def _make_split_fnames(fname, n_splits, split_naming): path = Path(_construct_bids_filename(base, ext, i)) res.append(path) return res - - -class _MemmapCache: - """Hold a memory map without copying or pickling its contents.""" - - def __init__(self): - self.mapping = None - self.pid = None - - def __deepcopy__(self, memodict): - """Create an empty cache when its owner is copied.""" - return type(self)() - - def __reduce__(self): - """Create an empty cache when its owner is pickled.""" - return type(self), () - - def close(self): - """Close the mapping and reset the cache.""" - mapping = self.mapping - self.mapping = self.pid = None - if mapping is not None: - mapping._mmap.close() - - -def _memmap_for(extras, fname): - """Return a process-local read-only uint8 memmap of *fname* from *extras*. - - The mapping is created lazily on first call and held by a cache that is - reset when *extras* is copied or pickled. The creating PID is recorded so - forked worker processes build their own mapping instead of sharing - inherited state. Returns None if the file cannot be mapped. There is - deliberately no staleness check: callers index the mapping through tables - read at open time (bounds, entries), which are invalid if the file changes - anyway, so a per-call stat would only add overhead. - """ - cache = extras.get("_memmap_cache") - if cache is None: - cache = extras["_memmap_cache"] = _MemmapCache() - mm = cache.mapping - pid = os.getpid() - if mm is not None: - if cache.pid == pid: - return mm - cache.close() - try: - mm = np.memmap(str(fname), dtype=np.uint8, mode="r") - except Exception: - return None - cache.mapping, cache.pid = mm, pid - return mm diff --git a/mne/io/base.py b/mne/io/base.py index b4bb5dd8529..0019a3cdf6b 100644 --- a/mne/io/base.py +++ b/mne/io/base.py @@ -2321,11 +2321,12 @@ def append( raise RuntimeError("Append error") # should never happen def close(self) -> None: - """Clean up any resources used by the object.""" - for extras in self._raw_extras: - cache = extras.get("_memmap_cache") - if cache is not None: - cache.close() + """Clean up the object. + + Does nothing for objects that close their file descriptors. + Things like Raw will override this method. + """ + pass # noqa def copy(self) -> Self: """Return copy of the instance. diff --git a/mne/io/fiff/raw.py b/mne/io/fiff/raw.py index ffa080f0804..95c6db5dbec 100644 --- a/mne/io/fiff/raw.py +++ b/mne/io/fiff/raw.py @@ -12,9 +12,9 @@ from ..._fiff.constants import FIFF from ..._fiff.meas_info import read_meas_info from ..._fiff.open import _fiff_get_fid, _get_next_fname, fiff_open -from ..._fiff.tag import _call_dict, _simple_dict, read_tag +from ..._fiff.tag import _call_dict, read_tag from ..._fiff.tree import dir_tree_find -from ..._fiff.utils import _memmap_for, _mult_cal_one +from ..._fiff.utils import _mult_cal_one from ...annotations import Annotations, _read_annotations_fif from ...channels import fix_mag_coil_types from ...event import AcqParserFIF @@ -38,40 +38,6 @@ ) -def _fif_mm_plan(bounds, ents, entry_span, nchan, start, stop): - """Plan byte-offset reads covering [start, stop) from a FIF memory map. - - Returns a list of ``(dtype, byte_offset, n_bytes, n_samples)`` read tuples, - one per touched buffer entry, or None when the fast path does not apply: a - touched entry is missing (gaps are zero-filled by the legacy loop), holds a - non-simple tag type, or has a declared size inconsistent with ``nchan`` - samples. - """ - read_plan = [] - for ei in entry_span: - entry = ents[ei] - if entry is None or entry.type not in _simple_dict: - return None - dtype = np.dtype(_simple_dict[entry.type]) - n_entry_samples = bounds[ei + 1] - bounds[ei] - n_bytes_per_sample = dtype.itemsize * nchan - if getattr(entry, "size", None) != (n_entry_samples * n_bytes_per_sample): - return None - # samples of this entry that fall inside [start, stop) - first_pick = max(start - bounds[ei], 0) - last_pick = min(n_entry_samples, stop - bounds[ei]) - n_samples = last_pick - first_pick - if n_samples <= 0: - continue - # payload starts 16 bytes into the tag (header is four uint32: kind, - # type, size, next -- see _read_tag_header and Tag.next_pos) - byte_offset = entry.pos + 16 + first_pick * n_bytes_per_sample - read_plan.append( - (dtype, byte_offset, n_samples * n_bytes_per_sample, n_samples) - ) - return read_plan - - @fill_doc class Raw(BaseRaw): """Raw data in FIF format. @@ -435,57 +401,15 @@ def _dtype(self): return dtype def _read_segment_file(self, data, idx, fi, start, stop, cals, mult): - """Read a segment of data from a file. - - Writes ``data`` (shape ``(len(idx), stop - start)``) for samples - ``[start, stop)``, rows following ``idx`` order. Two equivalent - implementations: a memory-map fast path for simple uncompressed tags - (see the comment block below), and the legacy buffered loop. - """ + """Read a segment of data from a file.""" n_bad = 0 - bounds = self._raw_extras[fi]["bounds"] - ents = self._raw_extras[fi]["ent"] - nchan = self._raw_extras[fi]["orig_nchan"] - fname = self._raw_extras[fi]["filename"] - # Entries overlapping [start, stop) via binary search on sorted bounds - # (O(log n) instead of a mask over every entry; matters for long - # recordings with thousands of buffer entries). - # indices of buffer entries overlapping [start, stop) - entry_span = range( - max(np.searchsorted(bounds, start, side="right") - 1, 0), - min(np.searchsorted(bounds, stop, side="left"), len(bounds) - 1), - ) - - # Fast path: serve [start, stop) from one persistent memory map of - # the file; see _fif_mm_plan() for eligibility and fallbacks. - read_plan = _fif_mm_plan(bounds, ents, entry_span, nchan, start, stop) - if ( - read_plan is not None - and isinstance(fname, Path) - and fname.suffixes[-1] != ".gz" - ): - file_mapping = _memmap_for(self._raw_extras[fi], fname) - if file_mapping is not None: - col_start = 0 - for dtype, byte_offset, n_bytes, n_samples in read_plan: - values = np.frombuffer( - file_mapping[byte_offset : byte_offset + n_bytes], - dtype=dtype, - count=n_samples * nchan, - ).reshape(n_samples, nchan) - _mult_cal_one( - data[:, col_start : col_start + n_samples], - values.T, - idx, - cals, - mult, - ) - col_start += n_samples - return - - with _fiff_get_fid(fname) as fid: + with _fiff_get_fid(self._raw_extras[fi]["filename"]) as fid: + bounds = self._raw_extras[fi]["bounds"] + ents = self._raw_extras[fi]["ent"] + nchan = self._raw_extras[fi]["orig_nchan"] + use = (stop > bounds[:-1]) & (start < bounds[1:]) offset = 0 - for ei in entry_span: + for ei in np.where(use)[0]: first = bounds[ei] last = bounds[ei + 1] nsamp = last - first diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py index 1218fe1c3b1..909ff9504ae 100644 --- a/mne/io/tests/test_raw.py +++ b/mne/io/tests/test_raw.py @@ -27,7 +27,7 @@ from mne._fiff.meas_info import Info, _get_valid_units, _writing_info_hdf5 from mne._fiff.pick import _ELECTRODE_CH_TYPES, _FNIRS_CH_TYPES_SPLIT from mne._fiff.proj import Projection -from mne._fiff.utils import _memmap_for, _mult_cal_one +from mne._fiff.utils import _mult_cal_one from mne.io import BaseRaw, RawArray, read_raw_fif from mne.io.base import _get_scaling from mne.transforms import Transform @@ -838,20 +838,6 @@ def _read_raw_arange(preload=False, verbose=None): return _RawArange(preload, verbose) -def test_raw_close_memmap_cache(tmp_path): - """Test that closing Raw releases cached memory maps.""" - fname = tmp_path / "test.bin" - fname.write_bytes(b"test") - with _read_raw_arange() as raw: - cache = raw._raw_extras[0] - mapping = _memmap_for(cache, fname) - assert not mapping._mmap.closed - - assert mapping._mmap.closed - assert cache["_memmap_cache"].mapping is None - assert cache["_memmap_cache"].pid is None - - def test_load_data_memmap(tmp_path): """Test loading raw data into a memmap via load_data.""" raw = _read_raw_arange(preload=False) From 1dc4c8b1c6a9c8a5818008fbb91042047bf36a55 Mon Sep 17 00:00:00 2001 From: Bru Date: Tue, 25 Aug 2026 19:57:59 +0200 Subject: [PATCH 4/4] Preserve caller-owned preload files --- doc/changes/dev/14213.bugfix.rst | 1 + mne/io/array/_array.py | 3 ++ mne/io/base.py | 30 +++++---------- mne/io/edf/tests/test_edf.py | 23 ++++++++++++ mne/io/fiff/tests/test_raw_fiff.py | 21 +++++++---- mne/io/tests/test_raw.py | 59 +++++++++++++++++++++++++++--- mne/utils/docs.py | 19 ++++++---- 7 files changed, 115 insertions(+), 41 deletions(-) create mode 100644 doc/changes/dev/14213.bugfix.rst diff --git a/doc/changes/dev/14213.bugfix.rst b/doc/changes/dev/14213.bugfix.rst new file mode 100644 index 00000000000..c3920c2c003 --- /dev/null +++ b/doc/changes/dev/14213.bugfix.rst @@ -0,0 +1 @@ +Preserve caller-owned memory-mapped preload files after Raw objects are destroyed, by `Bruno Aristimunha`_. diff --git a/mne/io/array/_array.py b/mne/io/array/_array.py index ccb3062ca2b..c41853133db 100644 --- a/mne/io/array/_array.py +++ b/mne/io/array/_array.py @@ -31,6 +31,9 @@ class RawArray(BaseRaw): Determines what gets copied on instantiation. "auto" (default) will copy info, and copy "data" only if necessary to get to double floating point precision. + If ``data`` is a memory-mapped array and is not copied, the caller + retains ownership of its backing file and is responsible for removing + it after the Raw object is no longer in use. .. versionadded:: 0.18 %(verbose)s diff --git a/mne/io/base.py b/mne/io/base.py index 0019a3cdf6b..7d41bb20be6 100644 --- a/mne/io/base.py +++ b/mne/io/base.py @@ -2,7 +2,6 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. -import os import shutil from collections import defaultdict from collections.abc import Callable, Sequence @@ -140,11 +139,12 @@ class BaseRaw( preload : bool | str | ndarray Preload data into memory for data manipulation and faster indexing. If True, the data will be preloaded into memory (fast, requires - large amount of memory). If preload is a string, preload is the - file name of a memory-mapped file which is used to store the data - on the hard drive (slower, requires less memory). If preload is an - ndarray, the data are taken from that array. If False, data are not - read until save. + large amount of memory). If preload is a string, it is the name of a + freshly created memory-mapped file used to store the data on the hard + drive (slower, requires less memory). An existing file is overwritten. + The caller owns the file and is responsible for removing it after the + Raw object is no longer in use. If preload is an ndarray, the data are + taken from that array. If False, data are not read until save. first_samps : sequence Sequence of the first sample number from each raw file. For unsplit raw files this should be a length-one list or tuple. @@ -602,8 +602,10 @@ def load_data( Parameters ---------- memmap : path-like | None - If not ``None``, preload data into a memory-mapped file at this - path. If ``None`` (default), preload data into RAM. + If not ``None``, preload data into a freshly created memory-mapped file + at this path. An existing file is overwritten. The caller owns the file + and is responsible for removing it after the Raw object is no longer in + use. If ``None`` (default), preload data into RAM. .. versionadded:: 1.13 %(verbose)s @@ -812,18 +814,6 @@ def set_annotations( return self - def __del__(self): # noqa: D105 - # remove file for memmap - fname = getattr(getattr(self, "_data", None), "filename", None) - if fname is not None: - # First, close the file out; happens automatically on del - del self._data - # Now file can be removed - try: - os.remove(fname) - except OSError: - pass # ignore file that no longer exists - def __enter__(self): """Entering with block.""" return self diff --git a/mne/io/edf/tests/test_edf.py b/mne/io/edf/tests/test_edf.py index 211641dc726..1672a131fc0 100644 --- a/mne/io/edf/tests/test_edf.py +++ b/mne/io/edf/tests/test_edf.py @@ -3,6 +3,7 @@ # Copyright the MNE-Python contributors. import datetime +import gc from contextlib import nullcontext from functools import partial from io import BytesIO @@ -317,6 +318,28 @@ def test_edf_data_broken(tmp_path): assert_allclose(data, data_new) +@pytest.mark.parametrize("method", ("constructor", "load_data")) +def test_edf_preload_memmap_ownership(method, tmp_path): + """Test ownership of a populated EDF preload memmap.""" + memmap_fname = tmp_path / f"edf-{method}-memmap.dat" + if method == "constructor": + raw = read_raw_edf(edf_stim_channel_path, preload=memmap_fname) + else: + raw = read_raw_edf(edf_stim_channel_path, preload=False) + raw.load_data(memmap=memmap_fname) + shape = raw._data.shape + expected = raw.get_data(picks=[0], start=0, stop=10)[0] + assert expected.any() + assert memmap_fname.stat().st_size == raw._data.nbytes + del raw + gc.collect() + + assert memmap_fname.is_file() + data = np.memmap(memmap_fname, dtype=np.float64, mode="r", shape=shape) + assert_array_equal(data[0, :10], expected) + del data + + def test_duplicate_channel_labels_edf(): """Test reading edf file with duplicate channel names.""" EXPECTED_CHANNEL_NAMES = ["EEG F1-Ref-0", "EEG F2-Ref", "EEG F1-Ref-1"] diff --git a/mne/io/fiff/tests/test_raw_fiff.py b/mne/io/fiff/tests/test_raw_fiff.py index 23b47812ed3..c996f187015 100644 --- a/mne/io/fiff/tests/test_raw_fiff.py +++ b/mne/io/fiff/tests/test_raw_fiff.py @@ -3,6 +3,7 @@ # Copyright the MNE-Python contributors. import datetime +import gc import os import pathlib import pickle @@ -315,7 +316,8 @@ def test_multiple_files(tmp_path): ) _compare_combo(raw, raw_combo, times, n_times) raw_combo = concatenate_raws( - [read_raw_fif(f) for f in [fif_fname, fif_fname]], preload="memmap8.dat" + [read_raw_fif(f) for f in [fif_fname, fif_fname]], + preload=tmp_path / "memmap8.dat", ) _compare_combo(raw, raw_combo, times, n_times) assert raw[:, :][0].shape[1] * 2 == raw_combo0[:, :][0].shape[1] @@ -349,13 +351,13 @@ def test_multiple_files(tmp_path): raw_combo = concatenate_raws( [read_raw_fif(fif_fname, preload=False), read_raw_fif(fif_fname, preload=True)], - preload="memmap3.dat", + preload=tmp_path / "memmap3.dat", ) _compare_combo(raw, raw_combo, times, n_times) raw_combo = concatenate_raws( [read_raw_fif(fif_fname, preload=True), read_raw_fif(fif_fname, preload=True)], - preload="memmap4.dat", + preload=tmp_path / "memmap4.dat", ) _compare_combo(raw, raw_combo, times, n_times) @@ -364,7 +366,7 @@ def test_multiple_files(tmp_path): read_raw_fif(fif_fname, preload=False), read_raw_fif(fif_fname, preload=False), ], - preload="memmap5.dat", + preload=tmp_path / "memmap5.dat", ) _compare_combo(raw, raw_combo, times, n_times) @@ -957,9 +959,9 @@ def test_io_complex(tmp_path, dtype): @testing.requires_testing_data -def test_getitem(): +def test_getitem(tmp_path): """Test getitem/indexing of Raw.""" - for preload in [False, True, "memmap1.dat"]: + for preload in [False, True, tmp_path / "memmap1.dat"]: raw = read_raw_fif(fif_fname, preload=preload) data, times = raw[0, :] data1, times1 = raw[0] @@ -1078,9 +1080,11 @@ def test_proj(tmp_path): @testing.requires_testing_data -@pytest.mark.parametrize("preload", [False, True, "memmap2.dat"]) +@pytest.mark.parametrize("preload", [False, True, "memmap"]) def test_preload_modify(preload, tmp_path): """Test preloading and modifying data.""" + if preload == "memmap": + preload = tmp_path / "memmap2.dat" rng = np.random.default_rng(0) raw = read_raw_fif(fif_fname, preload=preload) @@ -2030,6 +2034,9 @@ def test_memmap(tmp_path): raw_0._data[:] = 0.0 assert not raw_0._data.any() assert raw_1._data[:1, 3:5].all() + del raw_0, raw_1 + gc.collect() + assert Path(memmaps[3]).is_file() # other things like drop_channels and crop work but do not use memmapping, # eventually we might want to add support for some of these as users # require them. diff --git a/mne/io/tests/test_raw.py b/mne/io/tests/test_raw.py index 909ff9504ae..0f4b47d6e80 100644 --- a/mne/io/tests/test_raw.py +++ b/mne/io/tests/test_raw.py @@ -4,6 +4,7 @@ # License: BSD-3-Clause # Copyright the MNE-Python contributors. +import gc import math import re from contextlib import chdir, redirect_stdout @@ -838,16 +839,62 @@ def _read_raw_arange(preload=False, verbose=None): return _RawArange(preload, verbose) -def test_load_data_memmap(tmp_path): - """Test loading raw data into a memmap via load_data.""" - raw = _read_raw_arange(preload=False) - memmap_fname = tmp_path / "raw-load-data-memmap.dat" - raw.load_data(memmap=memmap_fname) +@pytest.mark.parametrize("method", ("constructor", "load_data")) +def test_preload_memmap_ownership(method, tmp_path): + """Test that a caller owns an explicit preload memmap path.""" + memmap_fname = tmp_path / f"raw-{method}-memmap.dat" + memmap_fname.write_bytes(b"stale" * 20_000) + if method == "constructor": + raw = _read_raw_arange(preload=memmap_fname) + else: + raw = _read_raw_arange(preload=False) + raw.load_data(memmap=memmap_fname) assert raw.preload assert isinstance(raw._data, np.memmap) assert Path(raw._data.filename) == memmap_fname - assert_array_equal(raw._data[:, 0], np.arange(1, 9)) + assert memmap_fname.stat().st_size == raw._data.nbytes + assert_array_equal(raw.get_data()[:, 0], np.arange(1, 9)) + raw.close() + assert_array_equal(raw.get_data()[:, 0], np.arange(1, 9)) + del raw + gc.collect() + + assert memmap_fname.is_file() + data = np.memmap(memmap_fname, dtype=np.float64, mode="r", shape=(8, 1000)) + assert_array_equal(data[:, 0], np.arange(1, 9)) + del data + replacement = tmp_path / "replacement.dat" + replacement.write_bytes(b"replacement") + replacement.replace(memmap_fname) + assert memmap_fname.read_bytes() == b"replacement" + + +def test_append_memmap_ownership(tmp_path): + """Test ownership of a caller-named concatenation memmap.""" + memmap_fname = tmp_path / "raw-append-memmap.dat" + raw = _read_raw_arange() + other = _read_raw_arange() + raw.append(other, preload=memmap_fname) + assert isinstance(raw._data, np.memmap) + expected = np.repeat(np.arange(1, 9)[:, np.newaxis], 2, axis=1) + assert_array_equal(raw.get_data()[:, [0, -1]], expected) + copied = raw.copy() + assert_array_equal(copied.get_data(), raw.get_data()) + del copied, other, raw + gc.collect() + assert memmap_fname.is_file() + + +def test_raw_array_memmap_ownership(tmp_path): + """Test ownership of a memmap supplied directly to RawArray.""" + memmap_fname = tmp_path / "raw-array-memmap.dat" + source = np.memmap(memmap_fname, dtype=np.float64, mode="w+", shape=(2, 10)) + source[:] = np.arange(20).reshape(2, 10) + raw = RawArray(source, create_info(2, 10.0), copy=None) + del source, raw + gc.collect() + assert memmap_fname.is_file() def test_test_raw_reader(): diff --git a/mne/utils/docs.py b/mne/utils/docs.py index 851952400b7..111ff21a112 100644 --- a/mne/utils/docs.py +++ b/mne/utils/docs.py @@ -3697,19 +3697,22 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75): preload : bool | str Preload data into memory for data manipulation and faster indexing. If True, the data will be preloaded into memory (fast, requires - large amount of memory). If preload is a string, preload is the - file name of a memory-mapped file which is used to store the data - on the hard drive (slower, requires less memory).""" + large amount of memory). If preload is a string, it is the name of a + freshly created memory-mapped file used to store the data on the hard + drive (slower, requires less memory). An existing file is overwritten. + The caller owns the file and is responsible for removing it after the + Raw object is no longer in use.""" docdict["preload_concatenate"] = """ preload : bool | str | None Preload data into memory for data manipulation and faster indexing. If True, the data will be preloaded into memory (fast, requires - large amount of memory). If preload is a string, preload is the - file name of a memory-mapped file which is used to store the data - on the hard drive (slower, requires less memory). If preload is - None, preload=True or False is inferred using the preload status - of the instances passed in. + large amount of memory). If preload is a string, it is the name of a + freshly created memory-mapped file used to store the data on the hard + drive (slower, requires less memory). An existing file is overwritten. + The caller owns the file and is responsible for removing it after the + Raw object is no longer in use. If preload is None, preload=True or False + is inferred using the preload status of the instances passed in. """ docdict["proj_epochs"] = """