Replace launching explicit threads with a threadpoolexecutor

divyanshk · web-flow · commit db87e2baf586 · 2025-04-14T15:22:16.000-07:00
Differential Revision: D72825598 Pull Request resolved: pytorch#1469
diff --git a/test/nodes/test_map.py b/test/nodes/test_map.py
@@ -5,9 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import itertools
-
 import unittest
 from typing import List, Optional
+from unittest import mock
 
 from parameterized import parameterized
 from torch.testing._internal.common_utils import IS_WINDOWS, TEST_CUDA, TestCase
@@ -131,7 +131,11 @@ def test_out_of_order_process_prebatch(self):
         )
     )
     def test_save_load_state_thread(
-        self, midpoint: int, in_order: bool, snapshot_frequency: int, prebatch: Optional[int]
+        self,
+        midpoint: int,
+        in_order: bool,
+        snapshot_frequency: int,
+        prebatch: Optional[int],
     ):
         method = "thread"
         batch_size = 6
@@ -159,7 +163,11 @@ def test_save_load_state_thread(
         )
     )
     def test_save_load_state_process(
-        self, midpoint: int, in_order: bool, snapshot_frequency: int, prebatch: Optional[int]
+        self,
+        midpoint: int,
+        in_order: bool,
+        snapshot_frequency: int,
+        prebatch: Optional[int],
     ):
         method = "process"
         batch_size = 6
@@ -179,3 +187,56 @@ def test_save_load_state_process(
         )
         node = Prefetcher(node, prefetch_factor=2)
         run_test_save_load_state(self, node, midpoint)
+
+    def test_thread_pool_executor_shutdown_on_del(self):
+        """Test that the ThreadPoolExecutor is properly shut down when the iterator is deleted."""
+        # Create a ParallelMapper with method="thread"
+        src = MockSource(num_samples=10)
+        node = ParallelMapper(
+            src,
+            RandomSleepUdf(),
+            num_workers=2,
+            method="thread",
+        )
+
+        # Reset the node to create the iterator
+        node.reset()
+
+        # We need to consume some items to ensure the ThreadPoolExecutor is created
+        # and the worker threads are started
+        for _ in range(5):
+            next(node)
+
+        # Use mock.patch to intercept the ThreadPoolExecutor.shutdown method
+        with mock.patch("concurrent.futures.ThreadPoolExecutor.shutdown") as mock_shutdown:
+            # Delete the node, which should trigger the shutdown of the ThreadPoolExecutor
+            del node
+
+            # Verify that shutdown was called
+            mock_shutdown.assert_called()
+
+    def test_thread_pool_executor_shutdown_on_exception(self):
+        """Test that the ThreadPoolExecutor is properly shut down when the iterator is deleted."""
+        # Create a ParallelMapper with method="thread"
+        src = MockSource(num_samples=10)
+        node = ParallelMapper(
+            src,
+            udf_raises,
+            num_workers=2,
+            method="thread",
+        )
+
+        # Reset the node to create the iterator
+        node.reset()
+
+        # Use mock.patch to intercept the ThreadPoolExecutor.shutdown method
+        with mock.patch("concurrent.futures.ThreadPoolExecutor.shutdown") as mock_shutdown:
+            # Consumer the iterator to ensure the ThreadPoolExecutor is created
+            # and exception is raised
+            try:
+                next(node)
+            except ValueError:
+                pass
+
+            # Verify that shutdown was called
+            mock_shutdown.assert_called()
diff --git a/torchdata/nodes/map.py b/torchdata/nodes/map.py
@@ -7,6 +7,8 @@
 import queue
 import threading
 import time
+
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Callable, Dict, Generic, Iterator, List, Literal, Optional, Protocol, Sequence, TypeVar, Union
 
 import torch.multiprocessing as mp
@@ -65,7 +67,11 @@ def __call__(self, xlist: Sequence[X]) -> Sequence[T]:
         return [self.map_fn(x) for x in xlist]
 
 
-def _sort_worker(in_q: Union[queue.Queue, mp.Queue], out_q: queue.Queue, stop_event: threading.Event):
+def _sort_worker(
+    in_q: Union[queue.Queue, mp.Queue],
+    out_q: queue.Queue,
+    stop_event: threading.Event,
+):
     buffer: Dict[int, Any] = {}
     cur_idx = 0
     while not stop_event.is_set():
@@ -91,6 +97,25 @@ def _sort_worker(in_q: Union[queue.Queue, mp.Queue], out_q: queue.Queue, stop_ev
             cur_idx += 1
 
 
+def _transformation_pool(
+    pool: ThreadPoolExecutor,
+    num_workers: int,
+    in_q: queue.Queue,
+    out_q: queue.Queue,
+    map_fn: Callable[[X], T],
+    stop_event: threading.Event,
+):
+    for worker_id in range(num_workers):
+        args = (
+            worker_id,
+            in_q,
+            out_q,
+            map_fn,
+            stop_event,
+        )
+        pool.submit(_apply_udf, *args)
+
+
 class _InlineMapperIter(Iterator[T]):
     """Non-Parallel implementation of Mapper"""
 
@@ -186,40 +211,49 @@ def __init__(
             name="read_thread(target=_populate_queue)",
             daemon=self.daemonic_reading,
         )
-        self._workers: List[Union[threading.Thread, mp.Process]] = []
-        for worker_id in range(self.num_workers):
-            args = (
-                worker_id,
+        self._read_thread.start()
+
+        if self.method == "thread":
+            self.pool = ThreadPoolExecutor(max_workers=self.num_workers)
+
+            _transformation_pool(
+                self.pool,
+                self.num_workers,
                 self._in_q,
                 self._intermed_q,
                 self.map_fn,
-                self._stop if self.method == "thread" else self._mp_stop,
+                self._stop,
             )
-            self._workers.append(
-                threading.Thread(
-                    target=_apply_udf,
-                    args=args,
-                    daemon=True,
-                    name=f"worker_thread_{worker_id}(target=_apply_udf)",
+
+        elif self.method == "process":
+            self._workers: List[mp.Process] = []
+            for worker_id in range(self.num_workers):
+                args = (
+                    worker_id,
+                    self._in_q,
+                    self._intermed_q,
+                    self.map_fn,
+                    self._mp_stop,
                 )
-                if self.method == "thread"
-                else mp_context.Process(target=_apply_udf, args=args, daemon=True)
-            )
+                self._workers.append(mp_context.Process(target=_apply_udf, args=args, daemon=True))
+            for t in self._workers:
+                t.start()
 
         self._out_q = self._intermed_q
         if self.in_order:
             self._sort_q: queue.Queue = queue.Queue()
             self._sort_thread = threading.Thread(
                 target=_sort_worker,
-                args=(self._intermed_q, self._sort_q, self._stop),
+                args=(
+                    self._intermed_q,
+                    self._sort_q,
+                    self._stop,
+                ),
                 daemon=True,
                 name="sort_thread(target=_sort_worker)",
             )
             self._out_q = self._sort_q
 
-        self._read_thread.start()
-        for t in self._workers:
-            t.start()
         if self.in_order:
             self._sort_thread.start()
 
@@ -260,6 +294,7 @@ def __next__(self) -> T:
             elif isinstance(item, ExceptionWrapper):
                 if not isinstance(item, StartupExceptionWrapper):
                     self._sem.release()
+                    self._shutdown()
                 item.reraise()
 
             self._steps_since_snapshot += 1
@@ -286,12 +321,14 @@ def _shutdown(self):
         self._mp_stop.set()
         if hasattr(self, "_read_thread") and self._read_thread.is_alive():
             self._read_thread.join(timeout=QUEUE_TIMEOUT * 5)
-        if hasattr(self, "_sort_thread") and self._sort_thread.is_alive():
-            self._sort_thread.join(timeout=QUEUE_TIMEOUT * 5)
+        if hasattr(self, "pool"):
+            self.pool.shutdown(wait=True)
         if hasattr(self, "_workers"):
             for t in self._workers:
                 if t.is_alive():
                     t.join(timeout=QUEUE_TIMEOUT * 5)
+        if hasattr(self, "_sort_thread") and self._sort_thread.is_alive():
+            self._sort_thread.join(timeout=QUEUE_TIMEOUT * 5)
 
 
 class _ParallelMapperImpl(BaseNode[T]):
