Add flag for non-daemonic _read_thread in ParallelMapper

Differential Revision: D71839829

Pull Request resolved: #1468

Author: syleshfb

torchdata/nodes/map.py

@@ -141,6 +141,7 @@ def __init__(
         max_concurrent: Optional[int],
         snapshot_frequency: int,
         initial_state: Optional[Dict[str, Any]],
+        daemonic_reading: bool,
     ):
         self.source = source
         self.map_fn = map_fn
@@ -149,6 +150,7 @@ def __init__(
         self.method = method
         self.mp_context = mp_context
         self.snapshot_frequency = snapshot_frequency
+        self.daemonic_reading = daemonic_reading
 
         self._in_q: Union[queue.Queue, mp.Queue] = queue.Queue() if method == "thread" else mp_context.Queue()
         self._intermed_q: Union[queue.Queue, mp.Queue] = queue.Queue() if method == "thread" else mp_context.Queue()
@@ -182,7 +184,7 @@ def __init__(
                 self._stop,
             ),
             name="read_thread(target=_populate_queue)",
-            daemon=True,
+            daemon=self.daemonic_reading,
         )
         self._workers: List[Union[threading.Thread, mp.Process]] = []
         for worker_id in range(self.num_workers):
@@ -311,6 +313,7 @@ def __init__(
         multiprocessing_context: Optional[str] = None,
         max_concurrent: Optional[int] = None,
         snapshot_frequency: int = 1,
+        daemonic_reading: bool = True,
     ):
         super().__init__()
         assert method in ["thread", "process"]
@@ -329,6 +332,7 @@ def __init__(
                 raise ValueError(f"{max_concurrent=} should be <= {num_workers=}!")
         self.max_concurrent = max_concurrent
         self.snapshot_frequency = snapshot_frequency
+        self.daemonic_reading = daemonic_reading
         self._it: Optional[Union[_InlineMapperIter[T], _ParallelMapperIter[T]]] = None
 
     def reset(self, initial_state: Optional[Dict[str, Any]] = None):
@@ -355,6 +359,7 @@ def _parallel_reset(self, initial_state: Optional[Dict[str, Any]]):
             max_concurrent=self.max_concurrent,
             snapshot_frequency=self.snapshot_frequency,
             initial_state=initial_state,
+            daemonic_reading=self.daemonic_reading,
         )
 
     def next(self) -> T:
@@ -403,6 +408,7 @@ def __init__(
         max_concurrent: Optional[int] = None,
         snapshot_frequency: int = 1,
         prebatch: Optional[int] = None,
+        daemonic_reading: bool = True,
     ):
         super().__init__()
         assert method in ["thread", "process"]
@@ -416,6 +422,7 @@ def __init__(
         self.max_concurrent = max_concurrent
         self.snapshot_frequency = snapshot_frequency
         self.prebatch = prebatch
+        self.daemonic_reading = daemonic_reading
         if prebatch is None:
             self.map_fn = map_fn
             self.source = source
@@ -434,6 +441,7 @@ def __init__(
             multiprocessing_context=self.multiprocessing_context,
             max_concurrent=self.max_concurrent,
             snapshot_frequency=self.snapshot_frequency,
+            daemonic_reading=self.daemonic_reading,
         )
 
         if self.prebatch is None: