From 967dd9760b5ccbb797642afa6c8b65fbdc887188 Mon Sep 17 00:00:00 2001 From: Mohammad Sadegh Salehi Date: Sun, 2 Aug 2026 00:48:07 +0100 Subject: [PATCH] Fix ThreadPrefetch close deadlock when producer blocks in parent close() used to set a stop event, join the producer, and only then close the parent. The producer only checks the stop event at the top of its loop, so if it is blocked inside parent.__next__(), the join waits forever. That hang is deterministic for both single-level and nested ThreadPrefetch (device_put shape), 10 of 10 runs each. The fix splits cancellation from waiting. request_stop() marks the iterator closed, signals the local thread, wakes buffer waiters, and propagates non-blocking cancel down the parent chain without joining. close() calls request_stop(), then parent.close(), then joins (still skipping join when sys.is_finalizing()). __del__ calls request_stop() only, so finalizers stay best effort. next() after close() still raises ValueError on every buffer size, including 0, where a stop sentinel in an unbounded queue must not be read as end of stream. The healthy path still uses a blocking get. Measured median throughput versus base on this machine is within about 1 to 2 percent on the fast path and the slow-producer path. This is separate from google/grain#1196. The reporter device_put script does not hang at HEAD on macOS arm64 (the is_finalizing join guard from bc17ead covers that shutdown path). A known follow-up remains: parent StopIteration is a BaseException, so the producer except Exception path does not put end of stream on the buffer for a consumer blocked in get. --- .../dataset/transformations/prefetch.py | 186 +++++++++++++++- .../dataset/transformations/prefetch_test.py | 201 ++++++++++++++++++ 2 files changed, 378 insertions(+), 9 deletions(-) diff --git a/grain/_src/python/dataset/transformations/prefetch.py b/grain/_src/python/dataset/transformations/prefetch.py index 8209bc26d..c127473f3 100644 --- a/grain/_src/python/dataset/transformations/prefetch.py +++ b/grain/_src/python/dataset/transformations/prefetch.py @@ -413,6 +413,45 @@ def _element_spec(self) -> Any: BufferElementT = tuple[T, StateT, Exception | None] +class _PrefetchStopped(Exception): + """Internal signal that thread prefetch was stopped and the buffer was closed. + + Placed on the buffer to wake a consumer blocked in ``Queue.get`` when + cancellation is requested. Not meant to be raised to user code after an + explicit ``close()``. + """ + + +def _buffer_put( + buffer: queue.Queue[BufferElementT], + item: BufferElementT, + should_stop: threading.Event, +) -> None: + """Puts ``item`` on ``buffer``, never blocking after a stop request. + + A stop sentinel may already occupy the only free slot. Blocking put after + cancellation would deadlock with close(), which joins this thread. + """ + if should_stop.is_set(): + try: + buffer.put_nowait(item) + except queue.Full: + pass + return + while True: + if should_stop.is_set(): + try: + buffer.put_nowait(item) + except queue.Full: + pass + return + try: + buffer.put(item, timeout=0.05) + return + except queue.Full: + continue + + def _put_iterator_elements_in_buffer( iterator: dataset.DatasetIterator[T], buffer: queue.Queue[BufferElementT], @@ -424,9 +463,34 @@ def _put_iterator_elements_in_buffer( while not should_stop.is_set(): element = stats.record_bytes_consumed(iterator.__next__()) state = copy.deepcopy(iterator.get_state()) - buffer.put((element, state, None)) + # Re-check stop before a potentially blocking put so cancellation can + # discard the element rather than wait for buffer space. + if should_stop.is_set(): + return + _buffer_put(buffer, (element, state, None), should_stop) except Exception as e: # pylint: disable=broad-except - buffer.put((None, None, e)) # pyrefly: ignore[bad-argument-type] + _buffer_put( + buffer, (None, None, e), should_stop + ) # pyrefly: ignore[bad-argument-type] + + +def _request_stop_iterator_tree(iterator: dataset.DatasetIterator) -> None: + """Non-blocking cancel propagation through a parent iterator chain. + + ThreadPrefetch nodes get ``request_stop()``. Other nodes are traversed without + setting ``_closed`` on intermediate transforms, so a later ``close()`` can + still walk the chain and join nested prefetch threads. Leaf iterators (no + parents) are marked closed so a cooperative ``__next__`` can unblock. + """ + if isinstance(iterator, ThreadPrefetchDatasetIterator): + iterator.request_stop() + return + parents = iterator._parents # pylint: disable=protected-access + if not parents: + iterator._closed = True # pylint: disable=protected-access + return + for parent in parents: + _request_stop_iterator_tree(parent) class CheckpointableIterator(Iterator[T], Protocol[T]): @@ -534,6 +598,11 @@ def start_prefetch(self): stage_category=dataset_stats.IPL_CAT_PREFETCH ) def __next__(self): + # Check closed before any buffer read. A stop sentinel left in the queue + # (especially with unbounded maxsize=0) must not surface as StopIteration + # after close(); the documented contract is ValueError. + if self._closed: + raise ValueError("Attempting to use a closed iterator.") if self._state is None: self._state = self._maybe_nonnative_parent.get_state() @@ -542,7 +611,7 @@ def __next__(self): with timer: if self._target_prefetch_buffer_size > 0: self.start_prefetch() - element, state, err = self._buffer.get() + element, state, err = self._buffer_get() else: try: # In case of 0 prefetch buffer size, we still try to get from the @@ -555,6 +624,11 @@ def __next__(self): err = None if err is not None: + # A stop sentinel means cancellation was already requested (possibly by a + # parent close). Do not join here: the closing thread owns the join, and + # joining from a nested producer can deadlock with close(). + if isinstance(err, _PrefetchStopped): + raise StopIteration from err self._stop_prefetch() raise err self._state = state @@ -564,12 +638,55 @@ def __next__(self): element = self._stats.record_bytes_produced(element) return self._stats.record_output_spec(element) - def close(self): - """Stops the iterator. No further calls to the iterator are expected.""" + def request_stop(self) -> None: + """Requests cancellation without waiting for producer threads. + + Marks this iterator closed, signals the local prefetch thread, wakes any + reader blocked on the buffer, and propagates the same non-blocking request + down the parent chain. Does not join. Explicit ``close()`` joins after + requesting stop. Safe to call from ``__del__``. + """ + already_closed = self._closed self._closed = True - self._stop_prefetch() + self._request_stop_prefetch() + if already_closed: + return + parent = self._maybe_nonnative_parent + if isinstance(parent, dataset.DatasetIterator): + _request_stop_iterator_tree(parent) + + def close(self): + """Stops the iterator. No further calls to the iterator are expected. + + Cancellation is two-phase so a producer blocked in ``parent.__next__`` can + be woken by closing the parent before this iterator joins its thread: + + 1. ``request_stop()``: non-blocking cancel on this node and parents. + 2. ``parent.close()``: blocking cleanup of the parent chain (joins nested + prefetch threads). + 3. Join the local prefetch thread (skipped while the interpreter finalizes). + """ + self.request_stop() if isinstance(self._maybe_nonnative_parent, dataset.DatasetIterator): self._maybe_nonnative_parent.close() + # Parent close may have unblocked our producer mid-put; clear again so the + # join cannot stall on a full buffer. Re-issue the stop sentinel afterwards + # and do not clear after join: a nested reader may still be blocked in + # buffer.get(), and a post-join clear would steal the sentinel from it. + self._clear_buffer() + self._put_stop_sentinel() + self._join_prefetch_thread(clear_buffer=False) + + def __del__(self): + # Best-effort only: propagate non-blocking cancel, never join. Explicit + # close() is the deterministic path. Joining from a finalizer can hang if + # the producer is blocked in parent.__next__ or during interpreter + # shutdown. request_stop() still marks parents closed so a producer parked + # in parent.__next__ can observe cancel and drop its reference. + try: + self.request_stop() + except Exception: # pylint: disable=broad-except + pass def _clear_buffer(self): while True: @@ -578,20 +695,66 @@ def _clear_buffer(self): except queue.Empty: return - def _stop_prefetch(self, clear_buffer: bool = True): - """Stops the prefetching thread if it's currently running.""" + def _buffer_get(self) -> BufferElementT: + """Gets the next buffer item, honouring cancellation. + + The healthy path uses a blocking ``get`` (same as before this change) so a + busy producer is not paced by a poll interval. After stop is requested, + uses a short timeout so a nested reader can observe that this iterator was + closed or that its prefetch thread finished even if a stop sentinel was + drained by another closer. A stop sentinel still wakes a blocking get. + """ + while True: + if self._closed or self._prefetch_should_stop.is_set(): + try: + return self._buffer.get(timeout=0.05) + except queue.Empty: + if self._closed: + raise StopIteration + thread = self._prefetch_thread + if thread is None or not thread.is_alive(): + raise StopIteration + continue + # Running: park until an element or a stop sentinel arrives. + return self._buffer.get() + + def _put_stop_sentinel(self): + """Wakes a consumer blocked in ``buffer.get`` after a stop request.""" + sentinel: BufferElementT = (None, None, _PrefetchStopped()) + try: + self._buffer.put_nowait(sentinel) + except queue.Full: + try: + self._buffer.get_nowait() + except queue.Empty: + pass + try: + self._buffer.put_nowait(sentinel) + except queue.Full: + pass + + def _request_stop_prefetch(self, clear_buffer: bool = True): + """Non-blocking cancellation request for the prefetch thread.""" if self._prefetch_thread is None: return self._prefetch_should_stop.set() if clear_buffer: # Remove entries from the buffer to unblock the producer, so that it - # checks producer_running.is_set() and exits. + # checks should_stop and exits. self._clear_buffer() else: assert isinstance(self._buffer, variable_size_queue.VariableSizeQueue) # Increase the buffer size by 1 to unblock the producer. self._buffer.set_max_size(self._target_prefetch_buffer_size + 1) # pytype: disable=attribute-error + # Wake any reader blocked in get() (for example a parent producer in a + # nested ThreadPrefetch pipeline). + self._put_stop_sentinel() + + def _join_prefetch_thread(self, clear_buffer: bool = True): + """Waits for the prefetch thread to exit after a stop request.""" + if self._prefetch_thread is None: + return if not sys.is_finalizing(): # Joining the worker thread is not necessary when the Python interpreter @@ -607,6 +770,11 @@ def _stop_prefetch(self, clear_buffer: bool = True): # on exit. self._clear_buffer() + def _stop_prefetch(self, clear_buffer: bool = True): + """Stops the prefetching thread if it's currently running.""" + self._request_stop_prefetch(clear_buffer=clear_buffer) + self._join_prefetch_thread(clear_buffer=clear_buffer) + def get_state(self) -> StateT: if self._state is not None: return self._state diff --git a/grain/_src/python/dataset/transformations/prefetch_test.py b/grain/_src/python/dataset/transformations/prefetch_test.py index a1757bc74..68e3755a8 100644 --- a/grain/_src/python/dataset/transformations/prefetch_test.py +++ b/grain/_src/python/dataset/transformations/prefetch_test.py @@ -14,6 +14,7 @@ from concurrent import futures import dataclasses import platform +import queue import sys import threading import time @@ -71,6 +72,118 @@ def get_state(self): return {} +# Timeouts for the subprocess-isolated ThreadPrefetch close lifecycle test. +# Related to google/grain#1196 (join-before-cancel shape), not the reporter's +# exact device_put script on CPython 3.12. The child must finish well under the +# join budget; if close deadlocks, the parent kills the child and fails the test. +_PREFETCH_CLOSE_JOIN_TIMEOUT_SEC = 10.0 +_PREFETCH_CLOSE_MAX_SEC = 1.0 + + +class _BlockingAfterIterDataset(dataset.IterDataset[int]): + """Yields ``block_after`` elements, then blocks in ``__next__`` until close. + + Parks a thread-prefetch producer inside ``parent.__next__``. The parent only + unblocks when its own ``close()`` sets ``_closed``. That models real parents + that finish in-flight work only when cancelled, so ``ThreadPrefetch.close`` + must request stop, close the parent chain, then join (not join first). + """ + + def __init__(self, n: int = 100, block_after: int = 1): + super().__init__() + self._n = n + self._block_after = block_after + self.entered_block = threading.Event() + + def __iter__(self) -> dataset.DatasetIterator[int]: + return _BlockingAfterIterator( + self._n, self._block_after, self.entered_block + ) + + +class _BlockingAfterIterator(dataset.DatasetIterator[int]): + + def __init__( + self, n: int, block_after: int, entered_block: threading.Event + ): + super().__init__() + self._n = n + self._block_after = block_after + self._entered_block = entered_block + self._i = 0 + + def __next__(self) -> int: + if self._i >= self._n: + raise StopIteration + if self._i >= self._block_after: + self._entered_block.set() + # Block until close() sets _closed. A stop event on the child prefetch + # thread does not affect this wait. + while not self._closed: + time.sleep(0.01) + raise StopIteration + value = self._i + self._i += 1 + return value + + def get_state(self): + return {'i': self._i} + + def set_state(self, state): + self._i = state['i'] + + +def _close_when_producer_blocked_in_parent_worker( + result_queue: mp.Queue, nested: bool +): + """Child process: close() while a producer is blocked in parent.__next__. + + When ``nested`` is True, the pipeline matches the device_put shape + (ThreadPrefetch -> map -> ThreadPrefetch). When False, a single + ThreadPrefetch wraps the blocking source. The consumer raises after the + first element and keeps the iterator, then calls close once a producer is + known to be inside the blocking parent. + """ + try: + source_ds = _BlockingAfterIterDataset(n=100, block_after=1) + ds: dataset.IterDataset[int] = prefetch.ThreadPrefetchIterDataset( + source_ds, prefetch_buffer_size=1 + ) + if nested: + ds = ds.map(lambda x: x) + ds = prefetch.ThreadPrefetchIterDataset(ds, prefetch_buffer_size=1) + + data_iter = ds.__iter__() + try: + for _ in data_iter: + raise RuntimeError('consumer error') + except RuntimeError: + pass + + # Ensure a prefetch producer is blocked in parent.__next__ before close. + if not source_ds.entered_block.wait(timeout=2.0): + result_queue.put({ + 'error': 'producer did not enter blocking parent.__next__', + }) + return + + close_start = time.time() + data_iter.close() + close_elapsed = time.time() - close_start + + live_after = [ + t.name + for t in threading.enumerate() + if t.name.startswith('grain-thread-prefetch') and t.is_alive() + ] + result_queue.put({ + 'close_elapsed': close_elapsed, + 'live_prefetch_threads': live_after, + }) + except Exception as e: # pylint: disable=broad-exception-caught + result_queue.put({'error': repr(e)}) + + class PrefetchIterDatasetTest(parameterized.TestCase): def setUp(self): @@ -772,6 +885,94 @@ def new_get_state(self): class ThreadPrefetchIterDatasetTest(_ThreadPrefetchIterDatasetTestBase): """Runs tests without provided executor.""" + @parameterized.parameters(0, 1, 5) + def test_next_after_close_raises_value_error(self, prefetch_buffer_size: int): + """After close(), __next__ must raise ValueError for every buffer size. + + DatasetIterator.close documents that subsequent __next__ calls raise + ValueError. A stop sentinel left in an unbounded (maxsize=0) queue must not + turn that into StopIteration, which a training loop would treat as end of + epoch. + """ + ds = dataset.MapDataset.range(20).to_iter_dataset() + ds = prefetch.ThreadPrefetchIterDataset( + ds, prefetch_buffer_size=prefetch_buffer_size + ) + it = ds.__iter__() + _ = next(it) + it.close() + with self.assertRaisesRegex( + ValueError, 'Attempting to use a closed iterator' + ): + next(it) + + @parameterized.parameters(False, True) + def test_close_does_not_hang_when_producer_blocked_in_parent( + self, nested: bool + ): + """close() must not hang when a producer is blocked in parent.__next__. + + ThreadPrefetchDatasetIterator.close joins its producer before closing the + parent. A stop event does not interrupt a blocked parent.__next__, so close + can wait forever. This is the join-before-cancel shape related to + google/grain#1196. It is not the reporter's exact device_put script on + CPython 3.12 (that script exits there). The source blocks after the first + element to hold a producer in the parent. Parameterized over single-level + and nested ThreadPrefetch (device_put-like nesting). + + A deadlock must fail this test, never hang the suite: the child runs in a + subprocess; the parent joins with a hard timeout and asserts exit code, + wall time, and absence of live ``grain-thread-prefetch`` threads. + """ + ctx = mp.get_context('spawn') + result_queue = ctx.Queue() + process = ctx.Process( + target=_close_when_producer_blocked_in_parent_worker, + args=(result_queue, nested), + daemon=True, + ) + start = time.time() + process.start() + process.join(timeout=_PREFETCH_CLOSE_JOIN_TIMEOUT_SEC) + wall = time.time() - start + + if process.is_alive(): + process.terminate() + process.join(timeout=5.0) + nesting = 'nested' if nested else 'single-level' + self.fail( + 'Child process hung for' + f' >{_PREFETCH_CLOSE_JOIN_TIMEOUT_SEC}s during {nesting}' + ' ThreadPrefetch close while producer blocked in parent.__next__' + f' (related to google/grain#1196). wall={wall:.2f}s' + f' exitcode={process.exitcode}' + ) + + self.assertEqual( + process.exitcode, + 0, + f'Child exited with {process.exitcode}, wall={wall:.2f}s', + ) + self.assertLess( + wall, + _PREFETCH_CLOSE_JOIN_TIMEOUT_SEC, + f'Child wall time {wall:.2f}s exceeded join budget', + ) + try: + result = result_queue.get(timeout=1.0) + except queue.Empty: + self.fail('Child produced no result') + self.assertNotIn('error', result, msg=result) + self.assertLess( + result['close_elapsed'], + _PREFETCH_CLOSE_MAX_SEC, + f"close() took {result['close_elapsed']:.2f}s", + ) + self.assertEmpty( + result['live_prefetch_threads'], + f"live prefetch threads after close: {result['live_prefetch_threads']}", + ) + class _MpContextCheckIterDataset(dataset.IterDataset[_T]):