Add serialization logic

torchdata/datapipes/iter/util/prefetch.py

@@ -187,3 +187,25 @@ def reset(self):
             self._error = None
             self._sync_counter = torch.tensor([0], dtype=torch.int32)
             self._done_callback = False
+
+    def __getstate__(self):
+        if IterDataPipe.getstate_hook is not None:
+            return IterDataPipe.getstate_hook(self)
+        if self._executor is not None:
+            self._executor.shutdown()
+        state = (
+            self.datapipe,
+            self.timeout,
+        )
+        return state
+
+    def __setstate__(self, state):
+        self.datapipe, self.timeout = state
+        self._process_group = dist.new_group(backend="gloo")
+        self._world_size = dist.get_world_size()
+        self._lock = threading.RLock()
+        self._cv = threading.Condition(lock=self._lock)
+        self._executor = None
+        self._error = None
+        self._sync_counter = torch.tensor([0], dtype=torch.int32)
+        self._done_callback = False