[DataPipe] Ensures Prefetcher shuts down properly

ghstack-source-id: ad503d7719137f964c705978574cbc78cd7e5598
Pull Request resolved: #1166

test/test_distributed.py

@@ -157,6 +157,13 @@ def _test_fullsync(rank, world_size, backend, q):
         it2 = iter(dp3)  # Reset
         next(it2)
 
+        dp4 = dp.prefetch(2)
+        it = iter(dp4)
+        next(it)
+        dp4.pause()
+        it2 = iter(dp4)  # Reset
+        next(it2)
+
         _finalize_distributed_queue(rank, q)
 
     @world_size_parametrize

torchdata/datapipes/iter/util/prefetcher.py

@@ -60,6 +60,7 @@ def __init__(self, source_datapipe, buffer_size: int = 10):
             raise ValueError("'buffer_size' is required to be a positive integer.")
         self.buffer_size = buffer_size
         self.thread: Optional[threading.Thread] = None
+        self.prefetch_data: Optional[_PrefetchData] = None
 
     @staticmethod
     def thread_worker(prefetch_data: _PrefetchData):
@@ -83,6 +84,8 @@ def thread_worker(prefetch_data: _PrefetchData):
             time.sleep(PRODUCER_SLEEP_INTERVAL * 10)
 
     def __iter__(self):
+        if self.thread is not None:
+            self.shutdown()
         try:
             prefetch_data = _PrefetchData(self.source_datapipe, self.buffer_size)
             self.prefetch_data = prefetch_data
@@ -104,9 +107,7 @@ def __iter__(self):
                 else:
                     time.sleep(CONSUMER_SLEEP_INTERVAL)
         finally:
-            prefetch_data.run_prefetcher = False
-            prefetch_data.stop_iteration = True
-            thread.join()
+            self.shutdown()
 
     def __getstate__(self):
         """
@@ -127,10 +128,12 @@ def __setstate__(self, state):
 
     @final
     def reset(self):
-        if self.thread is not None:
+        if self.prefetch_data is not None:
             self.prefetch_data.run_prefetcher = False
             self.prefetch_data.stop_iteration = True
             self.prefetch_data.paused = False
+            self.prefetch_data = None
+        if self.thread is not None:
             self.thread.join()
             self.thread = None
 
@@ -145,13 +148,26 @@ def pause(self):
 
     @final
     def resume(self):
-        if self.thread is not None and (
-            not self.prefetch_data.stop_iteration or len(self.prefetch_data.prefetch_buffer) > 0
+        if (
+            self.thread is not None
+            and self.prefetch_data is not None
+            and (not self.prefetch_data.stop_iteration or len(self.prefetch_data.prefetch_buffer) > 0)
         ):
-            assert self.prefetch_data is not None
             self.prefetch_data.run_prefetcher = True
             self.prefetch_data.paused = False
 
+    @final
+    def shutdown(self):
+        if self.prefetch_data:
+            self.prefetch_data.run_prefetcher = False
+            self.prefetch_data.stop_iteration = True
+            self.prefetch_data = None
+        if self.thread:
+            self.thread.join()
+
+    def __del__(self):
+        self.shutdown()
+
     def __len__(self) -> int:
         if isinstance(self.source_datapipe, Sized):
             return len(self.source_datapipe)
@@ -210,6 +226,8 @@ def thread_worker(prefetch_data: _PrefetchData, pin_memory_fn, device):  # type:
             time.sleep(PRODUCER_SLEEP_INTERVAL * 10)
 
     def __iter__(self):
+        if self.thread is not None:
+            self.shutdown()
         try:
             prefetch_data = _PrefetchData(self.source_datapipe, self.buffer_size)
             self.prefetch_data = prefetch_data
@@ -235,9 +253,7 @@ def __iter__(self):
                 else:
                     time.sleep(CONSUMER_SLEEP_INTERVAL)
         finally:
-            prefetch_data.run_prefetcher = False
-            prefetch_data.stop_iteration = True
-            thread.join()
+            self.shutdown()
 
     def __getstate__(self):
         state = super().__getstate__()