Repository: dmlc/dgl

Issue:
[GraphBolt][CUDA] Link prediction example crashing.
## 🔨Work Item

**IMPORTANT:**
* This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.
* DO NOT create a new work item if the purpose is to fix an existing issue or feature request. We will directly use the issue in the project tracker.

Project tracker: https://github.com/orgs/dmlc/projects/2

## Description
```
Training in cuda-cuda mode.
Loading data
The dataset is already preprocessed.
Training...
49it [00:05,  9.20it/s]
Epoch 00000 | Loss 0.4254 | Time 5.3299 s
Traceback (most recent call last):
  File "/localscratch/dgl-3/examples/sampling/graphbolt/lightning/../link_prediction.py", line 432, in <module>
    main(args)
  File "/localscratch/dgl-3/examples/sampling/graphbolt/lightning/../link_prediction.py", line 414, in main
    train(args, model, graph, features, train_set)
  File "/localscratch/dgl-3/examples/sampling/graphbolt/lightning/../link_prediction.py", line 322, in train
    for step, data in tqdm.tqdm(enumerate(dataloader)):
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 438, in __iter__
    return self._get_iterator()
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 383, in _get_iterator
    return _SingleProcessDataLoaderIter(self)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 659, in __init__
    super().__init__(loader)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 574, in __init__
    self._dataset = torch.utils.data.graph_settings.apply_random_seed(self._dataset, shared_rng)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph_settings.py", line 139, in apply_random_seed
    graph = traverse_dps(datapipe)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph.py", line 98, in traverse_dps
    return _traverse_helper(datapipe, only_datapipe=True, cache=cache)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph.py", line 145, in _traverse_helper
    d[dp_id][1].update(_traverse_helper(item, only_datapipe, cache.copy()))
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph.py", line 145, in _traverse_helper
    d[dp_id][1].update(_traverse_helper(item, only_datapipe, cache.copy()))
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph.py", line 145, in _traverse_helper
    d[dp_id][1].update(_traverse_helper(item, only_datapipe, cache.copy()))
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph.py", line 140, in _traverse_helper
    items = _list_connected_datapipes(datapipe, only_datapipe, cache.copy())
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/graph.py", line 67, in _list_connected_datapipes
    p.dump(scan_obj)
AttributeError: Can't pickle local object 'FeatureFetcher._read.<locals>._wait'
```



Discussion:
Seems to be related to: https://github.com/pytorch/data/issues/1161

Resolving patch:
diff --git a/examples/sampling/graphbolt/link_prediction.py b/examples/sampling/graphbolt/link_prediction.py
--- a/examples/sampling/graphbolt/link_prediction.py
+++ b/examples/sampling/graphbolt/link_prediction.py
@@ -328,8 +328,6 @@ def train(args, model, graph, features, train_set):
 
             total_loss += loss.item()
             if step + 1 == args.early_stop:
-                # Early stopping requires a new dataloader to reset its state.
-                dataloader = create_dataloader(args, graph, features, train_set)
                 break
 
         end_epoch_time = time.time()
diff --git a/python/dgl/graphbolt/base.py b/python/dgl/graphbolt/base.py
--- a/python/dgl/graphbolt/base.py
+++ b/python/dgl/graphbolt/base.py
@@ -313,6 +313,20 @@ def __iter__(self):
         while len(self.buffer) > 0:
             yield self.buffer.popleft()
 
+    def __getstate__(self):
+        state = (self.datapipe, self.buffer.maxlen)
+        if IterDataPipe.getstate_hook is not None:
+            return IterDataPipe.getstate_hook(state)
+        return state
+
+    def __setstate__(self, state):
+        self.datapipe, buffer_size = state
+        self.buffer = deque(maxlen=buffer_size)
+
+    def reset(self):
+        """Resets the state of the datapipe."""
+        self.buffer.clear()
+
 
 @functional_datapipe("wait")
 class Waiter(IterDataPipe):


Test patch:
diff --git a/tests/python/pytorch/graphbolt/test_dataloader.py b/tests/python/pytorch/graphbolt/test_dataloader.py
--- a/tests/python/pytorch/graphbolt/test_dataloader.py
+++ b/tests/python/pytorch/graphbolt/test_dataloader.py
@@ -111,3 +111,9 @@ def test_gpu_sampling_DataLoader(
     )
     assert len(bufferers) == bufferer_awaiter_cnt
     assert len(list(dataloader)) == N // B
+
+    for i, _ in enumerate(dataloader):
+        if i >= 1:
+            break
+
+    assert len(list(dataloader)) == N // B