Handle terminate

pytorch · Feb 16, 2023 · 8042a1a · 8042a1a
1 parent a6df023
commit 8042a1a
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/torchdata/datapipes/iter/util/prefetcher.py b/torchdata/datapipes/iter/util/prefetcher.py
@@ -88,11 +88,14 @@ def __iter__(self):
             thread.start()
             self.thread = thread
 
+            # Lazily import to prevent circular import
+            from torchdata.dataloader2 import communication
+
             while not prefetch_data.stop_iteration or len(prefetch_data.prefetch_buffer) > 0:
                 if len(prefetch_data.prefetch_buffer) > 0:
                     data = prefetch_data.prefetch_buffer.popleft()
                     if isinstance(data, Exception):
-                        if isinstance(data, StopIteration):
+                        if isinstance(data, (StopIteration, communication.iter.TerminateRequired)):
                             break
                         raise data
                     yield data
@@ -227,11 +230,14 @@ def __iter__(self):
             thread.start()
             self.thread = thread
 
+            # Lazily import to prevent circular import
+            from torchdata.dataloader2 import communication
+
             while not prefetch_data.stop_iteration or len(prefetch_data.prefetch_buffer) > 0:
                 if len(prefetch_data.prefetch_buffer) > 0:
                     data = prefetch_data.prefetch_buffer.popleft()
                     if isinstance(data, Exception):
-                        if isinstance(data, StopIteration):
+                        if isinstance(data, (StopIteration, communication.iter.TerminateRequired)):
                             break
                         raise data
                     yield data