Fix shared memory closed by two threads after mp_prefetch finishes a shard

grain/_src/python/dataset/transformations/BUILD

@@ -66,6 +66,7 @@ py_test(
     shard_count = 50,
     srcs_version = "PY3",
     deps = [
+        "//grain/_src/core:config",
         "//grain/_src/core:transforms",
         "//grain/_src/python:options",
         "//grain/_src/python/dataset",
@@ -403,6 +404,7 @@ py_test(
         "//grain/_src/python:options",
         "//grain/_src/python/dataset",
         "//grain/_src/python/dataset:base",
+        "//grain/_src/python/ipc:shared_memory_array",
         "@abseil-py//absl/logging",
         "@abseil-py//absl/testing:absltest",
         "@abseil-py//absl/testing:parameterized",

grain/_src/python/dataset/transformations/prefetch.py

@@ -444,6 +444,10 @@ def _put_iterator_elements_in_buffer(
       element = stats.record_bytes_consumed(iterator.__next__())
       state = copy.deepcopy(iterator.get_state())
       buffer.put((element, state, None))
+      # The consumer owns the element now. Drop the reference so that the
+      # traceback of an exception raised by the next `__next__` call, which
+      # references this frame, does not keep the element alive.
+      del element, state
   except Exception as e:  # pylint: disable=broad-except
     buffer.put((None, None, e))  # pyrefly: ignore[bad-argument-type]
 
@@ -616,7 +620,14 @@ def __next__(self):
 
     if err is not None:
       self._stop_prefetch()
-      raise err
+      try:
+        raise err
+      finally:
+        # The traceback of `err` references this frame. Drop the local to break
+        # the reference cycle, which would otherwise keep the exception, the
+        # producer frame and the parent iterator alive until the cyclic garbage
+        # collector runs.
+        del err
     self._state = state
     if self._next_index is not None:
       self._next_index += 1

grain/_src/python/dataset/transformations/prefetch_test.py

@@ -13,16 +13,19 @@
 # limitations under the License.
 from concurrent import futures
 import dataclasses
+import gc
 import os
 import platform
 import sys
 import threading
 import time
 from typing import TypeVar, cast
 from unittest import mock
+import weakref
 
 from absl.testing import absltest
 from absl.testing import parameterized
+from grain._src.core import config
 from grain._src.core import transforms
 import multiprocessing as mp
 from grain._src.python import options
@@ -448,6 +451,10 @@ def test_set_next_index(self):
       self.assertEqual(next(ds_iter), i)
 
 
+class _WeakReferenceableElement:
+  pass
+
+
 @absltest.skipThisClass('Base class')
 class _ThreadPrefetchIterDatasetTestBase(parameterized.TestCase):
 
@@ -716,6 +723,91 @@ def __getitem__(self, index):
       time.sleep(1)
       self.assertGreater(count, 8)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='stop_iteration',
+          fail_at=None,
+          traceback_filtering_mode='auto',
+      ),
+      dict(
+          testcase_name='error',
+          fail_at=3,
+          traceback_filtering_mode='auto',
+      ),
+      dict(
+          testcase_name='error_without_traceback_filtering',
+          fail_at=3,
+          traceback_filtering_mode='off',
+      ),
+  )
+  def test_finished_iterator_is_freed_without_cyclic_gc(
+      self, fail_at, traceback_filtering_mode
+  ):
+    self.addCleanup(
+        config.config.update,
+        'py_traceback_filtering',
+        config.config.get_or_default('py_traceback_filtering'),
+    )
+    config.config.update('py_traceback_filtering', traceback_filtering_mode)
+
+    def _make_element(index):
+      if index == fail_at:
+        raise ValueError('Test error.')
+      return _WeakReferenceableElement()
+
+    ds = (
+        dataset.MapDataset.range(5)
+        .map(_make_element)
+        .to_iter_dataset(options.ReadOptions(prefetch_buffer_size=0))
+    )
+    ds = prefetch.ThreadPrefetchIterDataset(ds, prefetch_buffer_size=2)
+    gc.disable()
+    self.addCleanup(gc.enable)
+    ds_iter = ds.__iter__()
+    parent = weakref.ref(ds_iter._maybe_nonnative_parent)
+    last_element = None
+    try:
+      for element in ds_iter:
+        last_element = weakref.ref(element)
+        del element
+    except ValueError:
+      pass
+    ds_iter.close()
+    del ds_iter
+    self.assertIsNotNone(last_element)
+    self.assertIsNone(last_element())
+    self.assertIsNone(parent())
+
+  def test_raised_error_does_not_keep_last_element_alive(self):
+
+    def _make_element(index):
+      if index == 3:
+        raise ValueError('Test error.')
+      return _WeakReferenceableElement()
+
+    ds = (
+        dataset.MapDataset.range(5)
+        .map(_make_element)
+        .to_iter_dataset(options.ReadOptions(prefetch_buffer_size=0))
+    )
+    ds = prefetch.ThreadPrefetchIterDataset(ds, prefetch_buffer_size=2)
+    gc.disable()
+    self.addCleanup(gc.enable)
+    ds_iter = ds.__iter__()
+    last_element = None
+    error = None
+    try:
+      for element in ds_iter:
+        last_element = weakref.ref(element)
+        del element
+    except ValueError as e:
+      error = e
+    ds_iter.close()
+    del ds_iter
+    self.assertIsInstance(error, ValueError)
+    self.assertIsNotNone(last_element)
+    self.assertIsNone(last_element())
+
   def test_element_spec(self):
     ds = dataset.MapDataset.range(2).to_iter_dataset()
     ds = prefetch.ThreadPrefetchIterDataset(ds, prefetch_buffer_size=1)

grain/_src/python/dataset/transformations/process_prefetch.py

@@ -536,13 +536,17 @@ def __next__(self):
         ):
           self._stop_prefetch()
         self._exhausted = True
-        if isinstance(err, StopIteration):
-          raise err
-        if err_traceback is not None:
+        if not isinstance(err, StopIteration) and err_traceback is not None:
           err = err.with_traceback(
               traceback_util.reconstruct_traceback(err_traceback)
           )
-        raise err
+        try:
+          raise err
+        finally:
+          # The traceback of `err` references this frame. Drop the local to
+          # break the reference cycle, which would otherwise keep the exception
+          # and this iterator alive until the cyclic garbage collector runs.
+          del err
       self._state = state
       self._next_index = next_index
     with self._stats.record_self_time(offset_ns=timer.value()):

grain/_src/python/dataset/transformations/process_prefetch_test.py

@@ -14,6 +14,7 @@
 from collections.abc import Sequence
 from concurrent import futures
 import dataclasses
+import gc
 import logging as std_logging
 import os
 import sys
@@ -34,6 +35,7 @@
 from grain._src.python.dataset import dataset
 from grain._src.python.dataset.transformations import process_prefetch
 from grain._src.python.dataset.transformations import testing_util
+from grain._src.python.ipc import shared_memory_array
 import numpy as np
 
 
@@ -1087,6 +1089,37 @@ def map_fn(x):
         ],
     )
 
+  def test_finished_iteration_leaves_no_cyclic_garbage(self):
+    ds = (
+        dataset.MapDataset.range(8)
+        .map(lambda x: np.full((64, 1024), x, dtype=np.float32))
+        .to_iter_dataset()
+    )
+    ds = process_prefetch.multiprocess_prefetch(ds, num_workers=2)
+    gc.collect()
+    gc.disable()
+    self.addCleanup(gc.enable)
+    ds_iter = ds.__iter__()
+    self.assertEqual([int(x[0, 0]) for x in ds_iter], list(range(8)))
+    ds_iter.close()
+    del ds_iter
+    gc.set_debug(gc.DEBUG_SAVEALL)
+    try:
+      gc.collect()
+      garbage_types = {type(obj) for obj in gc.garbage}
+    finally:
+      gc.set_debug(0)
+      gc.garbage.clear()
+      gc.collect()
+    self.assertNotIn(shared_memory_array.SharedMemoryArray, garbage_types)
+    self.assertNotIn(
+        process_prefetch.ProcessPrefetchDatasetIterator, garbage_types
+    )
+    self.assertFalse(
+        any(issubclass(t, BaseException) for t in garbage_types),
+        garbage_types,
+    )
+
   def test_element_spec(self):
     ds = dataset.MapDataset.range(2).to_iter_dataset()
     ds = process_prefetch.multiprocess_prefetch(

grain/_src/python/ipc/shared_memory_array.py

@@ -48,6 +48,25 @@ def _del_shm(shm: shared_memory.SharedMemory, unlink: bool) -> None:
     shm.unlink()
 
 
+class _SharedMemory(shared_memory.SharedMemory):
+  """A `SharedMemory` whose close can be handed over to a delete thread.
+
+  `SharedMemory.__del__` closes the shared memory. When the cyclic garbage
+  collector frees a `SharedMemoryArray` together with its `SharedMemory`, it
+  calls both finalizers in the same pass. If `SharedMemoryArray.__del__` has
+  handed the close to a delete thread, `SharedMemory.close()` then runs in two
+  threads at once and both can call `os.close` on the same file descriptor;
+  the second call closes whatever file reused that descriptor number in
+  between. Once the close is handed over, the delete thread is its only owner.
+  """
+
+  close_handed_over: bool = False
+
+  def __del__(self) -> None:
+    if not self.close_handed_over:
+      super().__del__()
+
+
 class SharedMemoryArray(np.ndarray):
   """A NumPy array subclass which is backed by shared memory.
 
@@ -71,7 +90,7 @@ def __new__(
   ):
     # See https://numpy.org/doc/stable/user/basics.subclassing.html
     size = math.prod(shape) * np.dtype(dtype).itemsize
-    shm = shared_memory.SharedMemory(create=True, size=size)
+    shm = _SharedMemory(create=True, size=size)
     return cls.from_shared_memory(shm, shape, dtype)
 
   def __array_finalize__(self, obj):
@@ -101,7 +120,7 @@ def __array_wrap__(self, obj, context=None, return_scalar=False):  # pylint: dis
   def from_metadata(
       cls, metadata: SharedMemoryArrayMetadata
   ) -> SharedMemoryArray:
-    shm = shared_memory.SharedMemory(metadata.name)
+    shm = _SharedMemory(metadata.name)
     return cls.from_shared_memory(shm, metadata.shape, metadata.dtype)
 
   @property
@@ -180,11 +199,15 @@ def __del__(self) -> None:
     outstanding_del_requests = SharedMemoryArray._outstanding_del_requests
     shm = self.shm
     assert isinstance(shm, shared_memory.SharedMemory)
-    if thread_pool:
+    # Only a `_SharedMemory` can hand its close over to a delete thread. Any
+    # other `SharedMemory` closes itself in its own `__del__`, which may run
+    # concurrently with the delete thread, so it is closed here.
+    if thread_pool and isinstance(shm, _SharedMemory):
       assert outstanding_del_requests is not None
       # We use a semaphore to make sure that we don't accumulate too many
       # requests to close/unlink shared memory, which could lead to OOM errors.
       if outstanding_del_requests.acquire(blocking=False):
+        shm.close_handed_over = True
         thread_pool.apply_async(
             SharedMemoryArray.close_shm_async, args=(shm, self._unlink_on_del)
         )

grain/_src/python/ipc/shared_memory_array_test.py

@@ -13,7 +13,9 @@
 # limitations under the License.
 """Tests for shared memory array."""
 
+import gc
 from multiprocessing import shared_memory
+import os
 import platform
 import threading
 import time
@@ -53,6 +55,13 @@ def _wait_for_deletion(metadata: SharedMemoryArrayMetadata) -> None:
       break
 
 
+def _unlink_if_exists(metadata: SharedMemoryArrayMetadata) -> None:
+  try:
+    metadata.close_and_unlink_shm()
+  except FileNotFoundError:
+    pass
+
+
 @absltest.skipIf(platform.system() == "Windows", "Timeouts on Windows.")
 class SharedMemoryArrayTest(parameterized.TestCase):
 
@@ -185,6 +194,42 @@ def my_close_shm_async(shm, unlink_on_del):
             2 * max_outstanding_requests, mock_close_shm_async.call_count
         )
 
+  def test_cyclic_gc_closes_shm_once_with_async_del(self):
+    SharedMemoryArray._disable_async_del()
+    SharedMemoryArray.enable_async_del()
+    shm_array = SharedMemoryArray((1024,), np.uint8)
+    shm_array.unlink_on_del()
+    metadata = shm_array.metadata
+    fd = shm_array.shm._fd
+    # A collection while the array is alive moves its `SharedMemory` after the
+    # array in the garbage collector's list, so that the cyclic collection
+    # below finalizes the array first.
+    gc.collect()
+    # Make the array cyclic garbage.
+    shm_array.self_reference = shm_array
+    del shm_array
+
+    real_close = os.close
+    closes = []
+
+    def _close(fd_to_close: int) -> None:
+      closes.append(fd_to_close)
+      real_close(fd_to_close)
+      if fd_to_close == fd:
+        # Release the GIL between `os.close` and `SharedMemory._fd = -1`, as
+        # `os.close` itself does, so that a concurrent closer can run.
+        time.sleep(0.5)
+
+    self.addCleanup(_unlink_if_exists, metadata)
+    with mock.patch.object(os, "close", side_effect=_close):
+      gc.collect()
+      # The pool has a single thread, so this returns after the deletion
+      # request made by the collection above has finished.
+      SharedMemoryArray._del_thread_pool.apply(lambda: None)
+    self.assertEqual(closes.count(fd), 1)
+    with self.assertRaises(FileNotFoundError):
+      _ = shared_memory.SharedMemory(name=metadata.name, create=False)
+
   def test_copy_and_open_shm_single_array(self):
     arr = np.arange(10).astype(np.int32)
     shm_struct = copy_to_shm(arr)