Internal

PiperOrigin-RevId: 878126445

Author: nithintatikonda2002-source

grain/_src/python/dataset/dataset.py

@@ -1740,16 +1740,23 @@ def __iter__(self) -> DatasetIterator[T]:
     """Performs any injection that needs to happen at the end of the pipeline."""
 
     # Loaded lazily due to a circular dependency (dataset <-> prefetch).
+    # Loaded lazily due to a circular dependency (dataset <-> process_prefetch).
     # pylint: disable=g-import-not-at-top
     from grain._src.python.dataset.transformations import prefetch
+    from grain._src.python.dataset.transformations import process_prefetch
     # pylint: enable=g-import-not-at-top
     iterator = self._parent.__iter__()
     if (
         is_thread_prefetch_injection_enabled()
         and not iterator._ctx.is_dataloader_pipeline  # pylint: disable=protected-access
+        and not process_prefetch.is_in_worker_process()
     ):
       if not prefetch.is_prefetch_iterator(iterator):
-        iterator = prefetch.ThreadPrefetchDatasetIterator(iterator, 1)
+        try:
+          iterator = prefetch.ThreadPrefetchDatasetIterator(iterator, 1)
+        except AttributeError:
+          # Some legacy iterators do not implement the `get_state` method.
+          pass
 
     filter_mode = traceback_filter_mode()
     if filter_mode != "off":

grain/_src/python/dataset/transformations/prefetch.py

@@ -454,7 +454,7 @@ def _put_iterator_elements_in_buffer(
   try:
     while not should_stop.is_set():
       element = stats.record_bytes_consumed(iterator.__next__())
-      state = iterator.get_state()
+      state = copy.deepcopy(iterator.get_state())
       buffer.put((element, state, None))
   except Exception as e:  # pylint: disable=broad-except
     buffer.put((None, None, e))
@@ -736,11 +736,17 @@ def multithread_prefetch(
 
 def is_prefetch_iterator(it: dataset.DatasetIterator) -> bool:
   """Returns whether the iterator is a prefetch iterator."""
+  # Loaded lazily due to a circular dependency (prefetch <-> process_prefetch).
+  # pylint: disable=g-import-not-at-top
+  from grain._src.python.dataset.transformations import process_prefetch
+  # pylint: enable=g-import-not-at-top
+
   return isinstance(
       it,
       (
           PrefetchDatasetIterator,
           ThreadPrefetchDatasetIterator,
           interleave.InterleaveDatasetIterator,
+          process_prefetch._ProcessPrefetchDatasetIterator,  # pylint: disable=protected-access
       ),
   )

grain/_src/python/dataset/transformations/process_prefetch.py

@@ -57,6 +57,11 @@
 _is_in_worker_process = False
 
 
+def is_in_worker_process() -> bool:
+  """Returns whether the current process is a worker process."""
+  return _is_in_worker_process
+
+
 def _run_all(fns: Sequence[Callable[[], None]]):
   for fn in fns:
     fn()