From 65bac7d0a44d17f5ec57e6aeef95ac65c81c7204 Mon Sep 17 00:00:00 2001 From: Mahdi Shafiei Date: Wed, 7 Oct 2026 22:49:14 -0700 Subject: [PATCH] Fix shared memory closed by two threads after `mp_prefetch` finishes a shard. `ThreadPrefetchDatasetIterator.__next__` and `ProcessPrefetchDatasetIterator.__next__` re-raise a stored exception from a local variable. The exception's traceback references the raising frame and the frame references the exception, so every shard that ends leaves a reference cycle that holds the `ProcessPrefetchDatasetIterator` and, through the producer thread's frame, the last element. Only the cyclic garbage collector frees it. When the cyclic garbage collector frees a `SharedMemoryArray` together with its `SharedMemory`, it calls both finalizers in one pass: `SharedMemoryArray.__del__` hands the `SharedMemory` to the delete thread pool, which closes it, and `SharedMemory.__del__` closes it on the collecting thread. `SharedMemory.close()` releases the GIL in `os.close` before it sets `_fd = -1`, so both threads can close the same descriptor number, and the second call closes whatever file reused that number in between. Delete the exception locals after raising, drop the producer's reference to an element once it is in the buffer, and make the delete thread the only owner of a close that has been handed over to it. --- .../_src/python/dataset/transformations/BUILD | 2 + .../dataset/transformations/prefetch.py | 13 ++- .../dataset/transformations/prefetch_test.py | 92 +++++++++++++++++++ .../transformations/process_prefetch.py | 12 ++- .../transformations/process_prefetch_test.py | 33 +++++++ grain/_src/python/ipc/shared_memory_array.py | 29 +++++- .../python/ipc/shared_memory_array_test.py | 45 +++++++++ 7 files changed, 218 insertions(+), 8 deletions(-) diff --git a/grain/_src/python/dataset/transformations/BUILD b/grain/_src/python/dataset/transformations/BUILD index 01ff40058..049d3a695 100644 --- a/grain/_src/python/dataset/transformations/BUILD +++ b/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", diff --git a/grain/_src/python/dataset/transformations/prefetch.py b/grain/_src/python/dataset/transformations/prefetch.py index 013e68a95..4cae3cbe2 100644 --- a/grain/_src/python/dataset/transformations/prefetch.py +++ b/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 diff --git a/grain/_src/python/dataset/transformations/prefetch_test.py b/grain/_src/python/dataset/transformations/prefetch_test.py index 168d5871b..842539e3e 100644 --- a/grain/_src/python/dataset/transformations/prefetch_test.py +++ b/grain/_src/python/dataset/transformations/prefetch_test.py @@ -13,6 +13,7 @@ # limitations under the License. from concurrent import futures import dataclasses +import gc import os import platform import sys @@ -20,9 +21,11 @@ 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) diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py index 93e0f44e9..9660b995b 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch.py +++ b/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()): diff --git a/grain/_src/python/dataset/transformations/process_prefetch_test.py b/grain/_src/python/dataset/transformations/process_prefetch_test.py index f3ff30d65..523a177df 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch_test.py +++ b/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( diff --git a/grain/_src/python/ipc/shared_memory_array.py b/grain/_src/python/ipc/shared_memory_array.py index 118b606c7..8a1224cd1 100644 --- a/grain/_src/python/ipc/shared_memory_array.py +++ b/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) ) diff --git a/grain/_src/python/ipc/shared_memory_array_test.py b/grain/_src/python/ipc/shared_memory_array_test.py index bc73984ff..57e8d5ebe 100644 --- a/grain/_src/python/ipc/shared_memory_array_test.py +++ b/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)