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)