Handle terminate

pytorch · Feb 15, 2023 · d3863ed · d3863ed
1 parent e81a430
commit d3863ed
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