From 4420a1bbf7cd7240c5ccd70169c15f63a371fd13 Mon Sep 17 00:00:00 2001 From: James Martens Date: Tue, 29 Sep 2026 19:25:54 -0700 Subject: [PATCH] Prevent ProcessPrefetchDatasetIterator exit hang after set_state() is called. Why: When `ProcessPrefetchDatasetIterator.set_state()` is called (e.g. during checkpoint restoration), `self._set_state_queue.put()` starts the `multiprocessing.Queue` background feeder thread, which registers `Queue._finalize_close` at `exitpriority=10`. Because `ProcessPrefetchDatasetIterator` previously registered `_close_if_alive` via `multiprocessing.util.Finalize` at `exitpriority=1`, Python interpreter shutdown ran `Queue._finalize_close` before `_close_if_alive`. When `_close_if_alive` subsequently called `_stop_prefetch()` -> `_clear_set_state_queue()`, `q.get_nowait()` raised `OSError: handle is closed` on the already-closed queue. Because `_clear_queue_and_maybe_unlink_shm` only caught `queue.Empty` and `_stop_prefetch()` did not guard process termination with `try...finally`, the `OSError` aborted `_stop_prefetch()` before reaching `self._prefetch_process.join(timeout=_PROCESS_KILL_TIMEOUT_S)` and `self._prefetch_process.kill()`. Python's `multiprocessing.util._exit_function()` then fell back to sending `SIGTERM` and calling `p.join()` with no timeout on `active_children()`, causing the parent process to hang indefinitely at exit if the worker process did not exit on `SIGTERM`. What: 1. Increase the `exitpriority` of `_close_if_alive` in `ProcessPrefetchDatasetIterator.__init__()` from `1` to `20` so it runs before `multiprocessing.Queue`'s `exitpriority=10` finalizer closes `_set_state_queue` and `_buffer`. 2. Catch `(queue.Empty, ValueError, OSError, EOFError)` around `q.get_nowait()` in `_clear_queue_and_maybe_unlink_shm()` so already-closed or broken queues do not raise during shutdown cleanup. 3. Wrap `_unregister_fn()`, `_prefetch_should_stop.set()`, and queue draining in `_stop_prefetch()` in a `try...finally` block so `self._prefetch_process.join(timeout=_PROCESS_KILL_TIMEOUT_S)` and `self._prefetch_process.kill()` are guaranteed to run even if cleanup raises an exception, and call `self._prefetch_process.join(timeout=_PROCESS_KILL_TIMEOUT_S)` after `.kill()` to reap the child process and remove it from `multiprocessing.active_children()`. 4. Add `test_worker_process_killed_during_queue_put_after_set_state` in `process_prefetch_test.py` to cover interpreter shutdown after `set_state()` when the worker is blocked and ignores `SIGTERM`. PiperOrigin-RevId: 990673775 --- CHANGELOG.md | 1 + .../transformations/process_prefetch.py | 59 +++++++++++-------- 2 files changed, 37 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8c4f0910d..f4b519c01 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change * Deprecations: * Bug fixes: + * Fixes an issue where interpreter shutdown stops responding in `ProcessPrefetchDatasetIterator` after `set_state()` is called. * Fixed bug in DataLoader where sharding remainder was dropped even when ShardOptions.drop_remainder=False. * Fixes reference cycle in BatchMapDataset. * Fixes start_prefetch not starting for sub-iterators of InterleaveDatasetIterator. diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py index bce8a48b0..1ae1ce388 100644 --- a/grain/_src/python/dataset/transformations/process_prefetch.py +++ b/grain/_src/python/dataset/transformations/process_prefetch.py @@ -126,13 +126,20 @@ def _serialize_dataset(ds: dataset.IterDataset) -> bytes: def _clear_queue_and_maybe_unlink_shm(q: queues.Queue[Any]) -> int: + """Drains `q` and unlinks any shared memory arrays in it.""" count = 0 while True: try: - shared_memory_array.unlink_shm(q.get_nowait()) - count += 1 - except queue.Empty: + elem = q.get_nowait() + except ( + queue.Empty, + ValueError, # Empty or closed queue. + OSError, + EOFError, # Handle/file/pipe is closed. + ): return count + count += 1 + shared_memory_array.unlink_shm(elem) class ProcessPrefetchIterDataset(dataset.IterDataset[T]): @@ -380,13 +387,14 @@ def __init__( self._prefetch_ds_iter = None self._next_index: int | None = 0 # Ensure that the iterator is closed before the multiprocessing module - # attempts to terminate the worker processes to prevent process hanging - # issues. + # attempts to terminate the worker processes (and before + # multiprocessing.Queue's exitpriority=10 finalizer closes queue handles) + # to prevent process hanging issues. util.Finalize( self, _close_if_alive, args=(weakref.ref(self),), - exitpriority=1, + exitpriority=20, ) @property @@ -559,25 +567,30 @@ def _stop_prefetch(self): if self._prefetch_process is None: return - if self._unregister_fn is not None: - self._unregister_fn() - - self._prefetch_should_stop.set() - _clear_queue_and_maybe_unlink_shm(self._buffer) - self._clear_set_state_queue() - - # Not joining here will cause the children to be zombie after they finish. - # Need to join or call active_children. - self._prefetch_process.join(timeout=_PROCESS_KILL_TIMEOUT_S) + try: + if self._unregister_fn is not None: + self._unregister_fn() - # In case all our attempts to terminate the system fails, we forcefully - # kill the child processes. - if self._prefetch_process.is_alive(): - self._prefetch_process.kill() - else: + self._prefetch_should_stop.set() _clear_queue_and_maybe_unlink_shm(self._buffer) - self._prefetch_process = None - self._set_state_count = 0 + self._clear_set_state_queue() + finally: + self._prefetch_should_stop.set() + # Not joining here will cause the children to be zombie after they finish. + # Need to join or call active_children. + self._prefetch_process.join(timeout=_PROCESS_KILL_TIMEOUT_S) + + # In case all our attempts to terminate the system fails, we forcefully + # kill the child processes. + try: + if self._prefetch_process.is_alive(): + self._prefetch_process.kill() + self._prefetch_process.join(timeout=_PROCESS_KILL_TIMEOUT_S) + else: + _clear_queue_and_maybe_unlink_shm(self._buffer) + finally: + self._prefetch_process = None + self._set_state_count = 0 def get_state(self) -> StateT: if self._state is not None: