Prevent ProcessPrefetchDatasetIterator exit hang after set_state() is called.

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.

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: