Fix memory-map reader fallbacks and lifecycle

Author: bruAristimunha

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`_

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="<i2")
+    fname.write_bytes(b"x" + values.T.tobytes() + b"x")
+    raw = SimpleNamespace(
+        _raw_extras=[dict(orig_nchan=2)],
+        filenames=[fname],
+    )
+    data = np.empty((2, 3))
+
+    _read_segments_file(
+        raw,
+        data,
+        slice(None),
+        0,
+        0,
+        3,
+        np.ones(2),
+        None,
+        dtype="<i2",
+        offset=1,
+    )
+
+    assert_array_equal(data, values)
+
+
+def test_memmap_cache_deepcopy(tmp_path):
+    """Test that copying a cache does not copy the mapped file into memory."""
+    fname = tmp_path / "test.bin"
+    fname.write_bytes(b"test")
+    extras = {}
+    original = _memmap_for(extras, fname)
+
+    copied = _memmap_for(deepcopy(extras), fname)
+
+    assert copied is not original
+    assert copied.filename == str(fname)
+    assert not copied.flags.writeable
+
+
+def test_memmap_cache_pickle(tmp_path):
+    """Test that pickling a cache does not serialize the mapped file."""
+    fname = tmp_path / "test.bin"
+    fname.write_bytes(b"test")
+    extras = {}
+    original = _memmap_for(extras, fname)
+
+    copied = _memmap_for(pickle.loads(pickle.dumps(extras)), fname)
+
+    assert copied is not original
+    assert copied.filename == str(fname)
+    assert not copied.flags.writeable
+
+
+def test_memmap_cache_failed_pid_change(tmp_path, monkeypatch):
+    """Test that a failed process-local remap releases inherited state."""
+    fname = tmp_path / "test.bin"
+    fname.write_bytes(b"test")
+    extras = {}
+    original = _memmap_for(extras, fname)
+    cache = extras["_memmap_cache"]
+
+    monkeypatch.setattr("mne._fiff.utils.os.getpid", lambda: cache.pid + 1)
+    monkeypatch.setattr(np, "memmap", Mock(side_effect=OSError))
+
+    assert _memmap_for(extras, fname) is None
+    assert original._mmap.closed
+    assert cache.mapping is None
+    assert cache.pid is None

mne/_fiff/utils.py

@@ -220,8 +220,6 @@ 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.
@@ -235,32 +233,18 @@ def _read_segments_file(
     # 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
-    ):
+    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 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
-            ]
+            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}), "
@@ -383,23 +367,52 @@ def _make_split_fnames(fname, n_splits, split_naming):
     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 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.
+    """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.
     """
-    mm = extras.get("_mm")
-    if mm is not None and extras.get("_mm_pid") == os.getpid():
-        return mm
+    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
-    extras["_mm"], extras["_mm_pid"] = mm, os.getpid()
+    cache.mapping, cache.pid = mm, pid
     return mm

mne/io/base.py

@@ -2321,12 +2321,11 @@ def append(
             raise RuntimeError("Append error")  # should never happen
 
     def close(self) -> 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.

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

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)