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
@@ -340,6 +340,8 @@ 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/feature_fetcher.py b/python/dgl/graphbolt/feature_fetcher.py
--- a/python/dgl/graphbolt/feature_fetcher.py
+++ b/python/dgl/graphbolt/feature_fetcher.py
@@ -174,10 +174,5 @@ def _read(self, data):
         with torch.cuda.stream(self.stream):
             data = self._read_data(data, current_stream)
             if self.stream is not None:
-                event = torch.cuda.current_stream().record_event()
-
-                def _wait():
-                    event.wait()
-
-                data.wait = _wait
+                data.wait = torch.cuda.current_stream().record_event().wait
             return data


Test patch: