Close interleaved prefetch workers concurrently

grain/_src/python/dataset/transformations/interleave.py

@@ -17,6 +17,8 @@
 from collections.abc import Sequence
 import copy
 import functools
+import sys
+import threading
 from typing import Any, TypeVar, cast
 import weakref
 
@@ -381,9 +383,45 @@ def close(self) -> None:
       return
     self._closed = True
     self._prefetch_ds_iter.close()
-    for iterator in self._iterators_in_use:
-      if iterator is not None:
+    iterators = [it for it in self._iterators_in_use if it is not None]
+    # Python 3.11 can hang in Thread.start() during final garbage collection.
+    if len(iterators) < 2 or sys.is_finalizing():
+      for iterator in iterators:
         iterator.close()
+      return
+
+    def close_iterator(iterator, result):
+      try:
+        iterator.close()
+      except BaseException as error:
+        result.set_exception(error)
+      else:
+        result.set_result(None)
+
+    threads = []
+    results = []
+    for iterator in iterators:
+      result = futures.Future()
+      thread = threading.Thread(
+          target=close_iterator,
+          args=(iterator, result),
+          name="grain-interleave-close",
+          daemon=False,
+      )
+      try:
+        thread.start()
+      except RuntimeError:
+        # Interpreter shutdown may no longer allow starting threads. Unlike
+        # executor submission, a failed start cannot leave queued cleanup work.
+        close_iterator(iterator, result)
+      else:
+        threads.append(thread)
+      results.append(result)
+    # Start every close before waiting, so per-worker exit delays overlap.
+    for thread in threads:
+      thread.join()
+    for result in results:
+      result.result()
 
   def start_prefetch(self):
     self._started = True

grain/_src/python/dataset/transformations/interleave_test.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
+import sys
+import textwrap
 import threading
 from typing import cast
 from unittest import mock
@@ -101,6 +105,157 @@ def __iter__(self) -> dataset.DatasetIterator:
     return _IteratorIdDatasetIterator(self._parent.__iter__())
 
 
+class InterleaveCloseTest(parameterized.TestCase):
+
+  def _make_iterator(self, children):
+    iterator = interleave.InterleaveDatasetIterator([], cycle_length=0)
+    # pylint: disable=protected-access
+    iterator._prefetch_ds_iter = mock.Mock()
+    iterator._iterators_in_use = children
+    # pylint: enable=protected-access
+    return iterator
+
+  @parameterized.parameters(2, 4, 8)
+  def test_close_starts_all_children_before_waiting(self, num_children):
+    barrier = threading.Barrier(num_children, timeout=5)
+    children = [mock.Mock() for _ in range(num_children)]
+    for child in children:
+      child.close.side_effect = barrier.wait
+    iterator = self._make_iterator(children)
+
+    iterator.close()
+
+    for child in children:
+      child.close.assert_called_once_with()
+
+  @parameterized.named_parameters(
+      dict(testcase_name="empty", slots=[]),
+      dict(testcase_name="uninitialized", slots=[None]),
+      dict(testcase_name="partial", slots=[mock.sentinel.child, None]),
+  )
+  def test_close_handles_empty_slots_and_is_idempotent(self, slots):
+    child = mock.Mock()
+    children = [child if slot is not None else None for slot in slots]
+    iterator = self._make_iterator(children)
+
+    iterator.close()
+    iterator.close()
+
+    iterator._prefetch_ds_iter.close.assert_called_once_with()  # pylint: disable=protected-access
+    if mock.sentinel.child in slots:
+      child.close.assert_called_once_with()
+    else:
+      child.close.assert_not_called()
+    with self.assertRaisesRegex(ValueError, "closed iterator"):
+      next(iterator)
+
+  @parameterized.parameters(ValueError, RuntimeError)
+  def test_close_waits_for_all_children_before_propagating_error(self, error):
+    children = [mock.Mock() for _ in range(3)]
+    children[0].close.side_effect = error("close failed")
+    release = threading.Event()
+    finished = threading.Event()
+
+    def slow_close():
+      self.assertTrue(release.wait(timeout=5))
+      finished.set()
+
+    children[1].close.side_effect = slow_close
+    children[2].close.side_effect = release.set
+    iterator = self._make_iterator(children)
+
+    with self.assertRaisesRegex(error, "close failed"):
+      iterator.close()
+
+    for child in children:
+      child.close.assert_called_once_with()
+    self.assertTrue(finished.is_set())
+
+  def test_close_when_threads_cannot_start(self):
+    children = [mock.Mock() for _ in range(2)]
+    iterator = self._make_iterator(children)
+    with mock.patch.object(
+        threading.Thread,
+        "start",
+        side_effect=RuntimeError("cannot start thread"),
+    ):
+      iterator.close()
+    for child in children:
+      child.close.assert_called_once_with()
+
+  def test_close_does_not_start_threads_during_finalization(self):
+    children = [mock.Mock() for _ in range(2)]
+    iterator = self._make_iterator(children)
+    with (
+        mock.patch.object(sys, "is_finalizing", return_value=True),
+        mock.patch.object(threading.Thread, "start") as start,
+    ):
+      iterator.close()
+    start.assert_not_called()
+    for child in children:
+      child.close.assert_called_once_with()
+
+  @parameterized.named_parameters(
+      dict(testcase_name="atexit", final_gc=False),
+      dict(testcase_name="final_gc", final_gc=True),
+  )
+  def test_close_at_interpreter_exit(self, final_gc):
+    code = textwrap.dedent("""\
+        import atexit
+        import gc
+        import os
+        import sys
+        from grain._src.python.dataset.transformations import interleave
+
+        class Child:
+          def __init__(self):
+            self.count = 0
+
+          def close(self):
+            self.count += 1
+
+        class Holder:
+          def __init__(self):
+            self.final_gc = USE_FINAL_GC
+            self.iterator = interleave.InterleaveDatasetIterator([], 0)
+            self.iterator._prefetch_ds_iter = Child()
+            self.children = [Child(), Child()]
+            self.iterator._iterators_in_use = self.children
+            self.write = os.write
+            self.is_finalizing = sys.is_finalizing
+
+          def close(self):
+            self.iterator.close()
+            counts = " ".join(str(child.count) for child in self.children)
+            self.write(1, ("CLOSED " + counts + "\\n").encode())
+
+          def __del__(self):
+            if self.final_gc:
+              phase = str(self.is_finalizing()).encode()
+              self.write(1, b"FINAL_GC=" + phase + b"\\n")
+              self.close()
+
+        holder = Holder()
+        if USE_FINAL_GC:
+          holder.cycle = holder
+          gc.disable()
+          del holder
+        else:
+          atexit.register(holder.close)
+        """).replace("USE_FINAL_GC", repr(final_gc))
+    completed = subprocess.run(
+        [sys.executable, "-c", code],
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=30,
+        env={**os.environ, "PYTHONPATH": os.pathsep.join(sys.path)},
+    )
+    expected = "FINAL_GC=True\nCLOSED 1 1" if final_gc else "CLOSED 1 1"
+    self.assertEqual(completed.stdout.strip(), expected, msg=completed.stderr)
+    self.assertNotIn("Exception ignored", completed.stderr)
+
+
 @absltest.skipThisClass("Base class")
 class _InterleaveIterDatasetTestBase(parameterized.TestCase):