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: