Merge pull request #56 from KempnerInstitute/mmap-init-failure-cleanup

Release partial mmaps when MemoryMappedDataset init fails

Author: mmshad

kempnerforge/data/dataset.py

@@ -13,6 +13,7 @@
 from __future__ import annotations
 
 import bisect
+import contextlib
 import logging
 from pathlib import Path
 
@@ -108,20 +109,27 @@ def __init__(
         self._cumulative_samples: list[int] = [0]
         total_tokens = 0
 
-        for f in self._files:
-            if self._is_bin:
-                # Raw binary: flat array of tokens. Infer dtype from file size
-                # or use uint32 (most common for modern tokenizers with vocab > 65535)
-                file_size = f.stat().st_size
-                dtype = np.uint32 if file_size % 4 == 0 else np.uint16
-                n_tokens = file_size // np.dtype(dtype).itemsize
-                mmap = np.memmap(str(f), dtype=dtype, mode="r", shape=(n_tokens,))
-            else:
-                mmap = np.load(str(f), mmap_mode="r")
-            n_samples = len(mmap) // seq_len
-            self._mmaps.append(mmap)
-            total_tokens += len(mmap)
-            self._cumulative_samples.append(self._cumulative_samples[-1] + n_samples)
+        # If any open fails partway, close the ones we already opened so they
+        # don't leak via the exception traceback (pytest, logger.exception,
+        # post-mortem debuggers all pin the partial `self` and its mmaps).
+        try:
+            for f in self._files:
+                if self._is_bin:
+                    # Raw binary: flat array of tokens. Infer dtype from file size
+                    # or use uint32 (most common for modern tokenizers with vocab > 65535)
+                    file_size = f.stat().st_size
+                    dtype = np.uint32 if file_size % 4 == 0 else np.uint16
+                    n_tokens = file_size // np.dtype(dtype).itemsize
+                    mmap = np.memmap(str(f), dtype=dtype, mode="r", shape=(n_tokens,))
+                else:
+                    mmap = np.load(str(f), mmap_mode="r")
+                n_samples = len(mmap) // seq_len
+                self._mmaps.append(mmap)
+                total_tokens += len(mmap)
+                self._cumulative_samples.append(self._cumulative_samples[-1] + n_samples)
+        except Exception:
+            self._close_mmaps()
+            raise
 
         self._total_samples = self._cumulative_samples[-1]
         logger.info(
@@ -177,6 +185,27 @@ def load_state_dict(self, state: dict) -> None:
         """Restore from checkpoint. Only ``epoch`` is restored; sample count is derived."""
         self._epoch = state.get("epoch", 0)
 
+    def _close_mmaps(self) -> None:
+        """Release the underlying mmap objects. Idempotent."""
+        for mm in self._mmaps:
+            inner = getattr(mm, "_mmap", None)
+            if inner is not None and not inner.closed:
+                # BufferError: live views into the mapping still exist — can't
+                # force-close safely; drop the ref and let GC finish it.
+                # ValueError: already closed by another code path.
+                with contextlib.suppress(BufferError, ValueError):
+                    inner.close()
+        self._mmaps.clear()
+
+    def close(self) -> None:
+        """Release the underlying mmaps. Preferred path; do not rely on ``__del__``."""
+        self._close_mmaps()
+
+    def __del__(self) -> None:
+        """GC safety net only. Prefer explicit :meth:`close`."""
+        with contextlib.suppress(Exception):
+            self._close_mmaps()
+
 
 class HuggingFaceDataset(Dataset):
     """HuggingFace dataset with on-the-fly tokenization and sequence packing.

tests/unit/test_dataset_mmap_cleanup.py

@@ -0,0 +1,107 @@
+"""Unit tests for MemoryMappedDataset mmap lifecycle.
+
+Covers the partial-open failure path (prior mmaps must be released so they
+don't leak through exception tracebacks) and the success-path invariant
+(mmaps stay open for the life of the dataset).
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from kempnerforge.data.dataset import MemoryMappedDataset
+
+
+def _write_npy_files(tmp_path, n_files: int, tokens_per_file: int = 1024) -> None:
+    for i in range(n_files):
+        arr = np.arange(tokens_per_file, dtype=np.uint32)
+        np.save(tmp_path / f"shard_{i:03d}.npy", arr)
+
+
+def test_partial_open_failure_closes_prior_mmaps(tmp_path):
+    """If np.load raises partway through, the mmaps already opened must be closed.
+
+    Without the fix, the prior mmaps stay live through any exception traceback
+    (pytest frames, logger.exception, post-mortem debuggers), accumulating
+    virtual-memory mappings on Lustre/NFS clusters under retry loops.
+    """
+    _write_npy_files(tmp_path, n_files=5)
+
+    original = np.load
+    calls = {"n": 0}
+    opened_mmaps: list = []
+
+    def flaky(*args, **kwargs):
+        calls["n"] += 1
+        if calls["n"] == 3:
+            raise RuntimeError("simulated Lustre hiccup")
+        mm = original(*args, **kwargs)
+        opened_mmaps.append(mm)
+        return mm
+
+    with (
+        patch("kempnerforge.data.dataset.np.load", side_effect=flaky),
+        pytest.raises(RuntimeError, match="Lustre hiccup"),
+    ):
+        MemoryMappedDataset(str(tmp_path), seq_len=128)
+
+    assert len(opened_mmaps) == 2, f"expected 2 opens before failure, got {len(opened_mmaps)}"
+    for mm in opened_mmaps:
+        inner = getattr(mm, "_mmap", None)
+        assert inner is not None
+        assert inner.closed, "mmap was not closed after __init__ raised"
+
+
+def test_close_is_idempotent_and_releases_mmaps(tmp_path):
+    _write_npy_files(tmp_path, n_files=3)
+    ds = MemoryMappedDataset(str(tmp_path), seq_len=128)
+
+    inners = [mm._mmap for mm in ds._mmaps]
+    assert all(not i.closed for i in inners)
+
+    ds.close()
+    assert all(i.closed for i in inners)
+
+    # Second close is a no-op, not a crash.
+    ds.close()
+
+
+def test_successful_init_keeps_mmaps_open(tmp_path):
+    """Regression guard: the fix must not close mmaps on the success path."""
+    _write_npy_files(tmp_path, n_files=2)
+    ds = MemoryMappedDataset(str(tmp_path), seq_len=128)
+    assert all(not mm._mmap.closed for mm in ds._mmaps)
+    sample = ds[0]
+    assert "input_ids" in sample
+    ds.close()
+
+
+def test_partial_open_failure_on_bin_files(tmp_path):
+    """Same leak guarantee applies to the .bin branch."""
+    for i in range(4):
+        (tmp_path / f"shard_{i:03d}.bin").write_bytes(np.arange(512, dtype=np.uint32).tobytes())
+
+    original = np.memmap
+    opened_mmaps: list = []
+    calls = {"n": 0}
+
+    def flaky_memmap(*args, **kwargs):
+        calls["n"] += 1
+        if calls["n"] == 2:
+            raise RuntimeError("simulated bin open failure")
+        mm = original(*args, **kwargs)
+        opened_mmaps.append(mm)
+        return mm
+
+    with (
+        patch("kempnerforge.data.dataset.np.memmap", side_effect=flaky_memmap),
+        pytest.raises(RuntimeError, match="bin open failure"),
+    ):
+        MemoryMappedDataset(str(tmp_path), seq_len=128, file_pattern="*.bin")
+
+    assert len(opened_mmaps) == 1
+    inner = getattr(opened_mmaps[0], "_mmap", None)
+    assert inner is not None and inner.closed