diff --git a/grain/_src/python/dataset/transformations/interleave.py b/grain/_src/python/dataset/transformations/interleave.py index 46a1d4730..865e01b3b 100644 --- a/grain/_src/python/dataset/transformations/interleave.py +++ b/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 diff --git a/grain/_src/python/dataset/transformations/interleave_test.py b/grain/_src/python/dataset/transformations/interleave_test.py index c80f02f9f..80323d652 100644 --- a/grain/_src/python/dataset/transformations/interleave_test.py +++ b/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):