Repository: dmlc/dgl
Pull request: [GraphBolt][CUDA] Dataloader `num_workers > 0` fix.
## Description
Fixes #6921. We defer calling `is_pinned` to avoid initializing torch cuda internals by forcing the user to provide the `nodes` tensor on the GPU to be able to use pinned graph sampling via UVA. If the `nodes` or `index` tensor is not on the GPU, then we don't need to check if the graph is pinned or not. If the `nodes` tensor is on the GPU, then that means torch cuda is already initialized, and calling `is_pinned` is fine.

## Checklist
Please feel free to remove inapplicable items for your PR.
- [ ] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [ ] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code.
- [ ] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
- [ ] All changes have test coverage
- [ ] Code is well-documented
- [ ] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
- [ ] Related issue is referred in this PR
- [ ] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md).

## Changes
<!-- You could use following template
- [ ] Feature1, tests, (and when applicable, API doc)
- [ ] Feature2, tests, (and when applicable, API doc)
-->


Problem statement:
[GraphBolt] `RuntimeError: CUDA error: initialization error` in `sample_neighbors()` when num_workers > 0
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. build CUDA version upon DGL 2.0.x branch or install latest nightly build DGL.
2. run `DGL_HOME=/home/ubuntu/workspace/dgl_0 DGL_LIBRARY_PATH=$DGL_HOME/build PYTHONPATH=tools:tools/distpartitioning:tests:$DGL_HOME/python:tests/python/pytorch/graphbolt:$PYTHONPATH python3 examples/sampling/graphbolt/node_classification.py --num-workers 5 --device cpu`

```
Training in cpu mode.
Loading data...
The dataset is already preprocessed.
Training...
Traceback (most recent call last):
  File "/home/ubuntu/workspace/dgl_0/examples/sampling/graphbolt/node_classification.py", line 410, in <module>
    main(args)
  File "/home/ubuntu/workspace/dgl_0/examples/sampling/graphbolt/node_classification.py", line 392, in main
    train(args, graph, features, train_set, valid_set, num_classes, model)
  File "/home/ubuntu/workspace/dgl_0/examples/sampling/graphbolt/node_classification.py", line 292, in train
    for step, data in enumerate(dataloader):
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 674, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 41, in fetch
    data = next(self.dataset_iter)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 154, in __next__
    return self._get_next()
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 142, in _get_next
    result = next(self.iterator)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 226, in wrap_next
    result = next_func(*args, **kwargs)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/datapipe.py", line 381, in __next__
    return next(self._datapipe_iter)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 183, in wrap_generator
    response = gen.send(None)
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/base.py", line 162, in __iter__
    for data in self.datapipe:
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 183, in wrap_generator
    response = gen.send(None)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torchdata/datapipes/iter/util/prefetcher.py", line 103, in __iter__
    raise data
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torchdata/datapipes/iter/util/prefetcher.py", line 73, in thread_worker
    item = next(itr)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 183, in wrap_generator
    response = gen.send(None)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 122, in __iter__
    for data in self.datapipe:
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 183, in wrap_generator
    response = gen.send(None)
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/dataloader.py", line 64, in __iter__
    yield from self.dataloader
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1345, in _next_data
    return self._process_data(data)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1371, in _process_data
    data.reraise()
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 41, in fetch
    data = next(self.dataset_iter)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 154, in __next__
    return self._get_next()
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 142, in _get_next
    result = next(self.iterator)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 226, in wrap_next
    result = next_func(*args, **kwargs)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/datapipe.py", line 381, in __next__
    return next(self._datapipe_iter)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/_hook_iterator.py", line 183, in wrap_generator
    response = gen.send(None)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 123, in __iter__
    yield self._apply_fn(data)
  File "/opt/conda/envs/dgl-dev-gpu-dgl-0/lib/python3.10/site-packages/torch/utils/data/datapipes/iter/callable.py", line 88, in _apply_fn
    return self.fn(data)
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/minibatch_transformer.py", line 38, in _transformer
    minibatch = self.transformer(minibatch)
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/subgraph_sampler.py", line 57, in _sample
    ) = self.sample_subgraphs(seeds)
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/impl/neighbor_sampler.py", line 138, in sample_subgraphs
    subgraph = self.sampler(
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/impl/fused_csc_sampling_graph.py", line 557, in sample_neighbors
    C_sampled_subgraph = self._sample_neighbors(
  File "/home/ubuntu/workspace/dgl_0/python/dgl/graphbolt/impl/fused_csc_sampling_graph.py", line 657, in _sample_neighbors
    return self._c_csc_graph.sample_neighbors(
RuntimeError: CUDA error: initialization error
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

This exception is thrown by __iter__ of NeighborSampler(datapipe=ShardingFilterIterDataPipe, deduplicate=True, fanouts=[tensor([10]), tensor([10]), tensor([10])], graph=FusedCSCSamplingGraph(csc_indptr=tensor([        0,       174,       263,  ..., 123718233, 123718242,
                                         123718280]),
                      indices=tensor([152857,  32104,  23158,  ...,  18965, 162836,  53324]),
                      num_nodes=2449029, num_edges=123718280, edge_attributes={}), prob_name=None, replace=False)
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

 - DGL Version (e.g., 1.0): 2.0.0, nightly
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 2.1
 - OS (e.g., Linux): Linux
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable): 12.1
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->

Hints from issue discussion:
cpu only DGL works well with `num_workers>0, --device cpu`.
adding `mp.set_start_method("spawn")` at the top of `main()` resolves this crash
Probably related: https://pytorch.org/docs/master/notes/multiprocessing.html#cuda-in-multiprocessing
I don't know if there is a way to stop CUDA from being initialized when not using it. I will do some research on it.
Also related:
https://github.com/dmlc/dgl/blob/master/tests/python/pytorch/graphbolt/test_dataloader.py#L13
I identified the issue. When we are dispatching to the GPU, we check two things, whether the tensors are pinned or whether they are on CUDA. I verified that changing this function:
https://github.com/dmlc/dgl/blob/4323986b25ce5c6d0ced0a6db5550402740aeb99/graphbolt/src/utils.h#L18-L20
To
```python
return tensor.device().type() == c10::DeviceType::CUDA;
```
resolves the issue.
However, we need to figure out how to do the dispatch when it is one of the dataloader workers. Maybe we can mark each of the launched processes by a global variable which can be checked inside this function so that we don't check if tensors are pinned if they are dataloader workers.
If `is_dataloader_worker` indicates whether the process is a dataloader process, then the is_pinned check can be avoided as follows:
```python
 inline bool is_accessible_from_gpu(torch::Tensor tensor) { 
   return (!is_dataloader_worker && tensor.is_pinned()) || tensor.device().type() == c10::DeviceType::CUDA; 
 } 
 ```
That is if we don't want to add `mp.set_start_method("spawn")` to main for all examples or into the dataloader for when cuda is available in the system.
Found a better solution, opening a PR.

Commit message: fix UVA index select as well.

Patch to review:
diff --git a/graphbolt/src/fused_csc_sampling_graph.cc b/graphbolt/src/fused_csc_sampling_graph.cc
index 36a3993ecff4..8432d8a2da01 100644
--- a/graphbolt/src/fused_csc_sampling_graph.cc
+++ b/graphbolt/src/fused_csc_sampling_graph.cc
@@ -274,9 +274,8 @@ FusedCSCSamplingGraph::GetState() const {
 
 c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::InSubgraph(
     const torch::Tensor& nodes) const {
-  if (utils::is_accessible_from_gpu(indptr_) &&
+  if (utils::is_on_gpu(nodes) && utils::is_accessible_from_gpu(indptr_) &&
       utils::is_accessible_from_gpu(indices_) &&
-      utils::is_accessible_from_gpu(nodes) &&
       (!type_per_edge_.has_value() ||
        utils::is_accessible_from_gpu(type_per_edge_.value()))) {
     GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(c10::DeviceType::CUDA, "InSubgraph", {
@@ -616,9 +615,9 @@ c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::SampleNeighbors(
     probs_or_mask = this->EdgeAttribute(probs_name);
   }
 
-  if (!replace && utils::is_accessible_from_gpu(indptr_) &&
+  if (!replace && utils::is_on_gpu(nodes) &&
+      utils::is_accessible_from_gpu(indptr_) &&
       utils::is_accessible_from_gpu(indices_) &&
-      utils::is_accessible_from_gpu(nodes) &&
       (!probs_or_mask.has_value() ||
        utils::is_accessible_from_gpu(probs_or_mask.value())) &&
       (!type_per_edge_.has_value() ||
diff --git a/graphbolt/src/index_select.cc b/graphbolt/src/index_select.cc
index 34191a8f8257..00257061c675 100644
--- a/graphbolt/src/index_select.cc
+++ b/graphbolt/src/index_select.cc
@@ -13,8 +13,7 @@ namespace graphbolt {
 namespace ops {
 
 torch::Tensor IndexSelect(torch::Tensor input, torch::Tensor index) {
-  if (input.is_pinned() &&
-      (index.is_pinned() || index.device().type() == c10::DeviceType::CUDA)) {
+  if (utils::is_on_gpu(index) && input.is_pinned()) {
     GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
         c10::DeviceType::CUDA, "UVAIndexSelect",
         { return UVAIndexSelectImpl(input, index); });
@@ -26,9 +25,8 @@ std::tuple<torch::Tensor, torch::Tensor> IndexSelectCSC(
     torch::Tensor indptr, torch::Tensor indices, torch::Tensor nodes) {
   TORCH_CHECK(
       indices.sizes().size() == 1, "IndexSelectCSC only supports 1d tensors");
-  if (utils::is_accessible_from_gpu(indptr) &&
-      utils::is_accessible_from_gpu(indices) &&
-      utils::is_accessible_from_gpu(nodes)) {
+  if (utils::is_on_gpu(nodes) && utils::is_accessible_from_gpu(indptr) &&
+      utils::is_accessible_from_gpu(indices)) {
     GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
         c10::DeviceType::CUDA, "IndexSelectCSCImpl",
         { return IndexSelectCSCImpl(indptr, indices, nodes); });
diff --git a/graphbolt/src/isin.cc b/graphbolt/src/isin.cc
index ff9a976b03b9..c41b839b1651 100644
--- a/graphbolt/src/isin.cc
+++ b/graphbolt/src/isin.cc
@@ -48,8 +48,7 @@ torch::Tensor IsInCPU(
 
 torch::Tensor IsIn(
     const torch::Tensor& elements, const torch::Tensor& test_elements) {
-  if (utils::is_accessible_from_gpu(elements) &&
-      utils::is_accessible_from_gpu(test_elements)) {
+  if (utils::is_on_gpu(elements) && utils::is_on_gpu(test_elements)) {
     GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
         c10::DeviceType::CUDA, "IsInOperation",
         { return ops::IsIn(elements, test_elements); });
diff --git a/graphbolt/src/unique_and_compact.cc b/graphbolt/src/unique_and_compact.cc
index fd0d23e84e63..3a7e4963bbdd 100644
--- a/graphbolt/src/unique_and_compact.cc
+++ b/graphbolt/src/unique_and_compact.cc
@@ -19,9 +19,8 @@ namespace sampling {
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> UniqueAndCompact(
     const torch::Tensor& src_ids, const torch::Tensor& dst_ids,
     const torch::Tensor unique_dst_ids) {
-  if (utils::is_accessible_from_gpu(src_ids) &&
-      utils::is_accessible_from_gpu(dst_ids) &&
-      utils::is_accessible_from_gpu(unique_dst_ids)) {
+  if (utils::is_on_gpu(src_ids) && utils::is_on_gpu(dst_ids) &&
+      utils::is_on_gpu(unique_dst_ids)) {
     GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
         c10::DeviceType::CUDA, "unique_and_compact",
         { return ops::UniqueAndCompact(src_ids, dst_ids, unique_dst_ids); });
diff --git a/graphbolt/src/utils.h b/graphbolt/src/utils.h
index 093e920af017..5203a5f1048c 100644
--- a/graphbolt/src/utils.h
+++ b/graphbolt/src/utils.h
@@ -19,6 +19,13 @@ inline bool is_accessible_from_gpu(torch::Tensor tensor) {
   return tensor.is_pinned() || tensor.device().type() == c10::DeviceType::CUDA;
 }
 
+/**
+ * @brief Checks whether the tensor is stored on the GPU.
+ */
+inline bool is_on_gpu(torch::Tensor tensor) {
+  return tensor.device().type() == c10::DeviceType::CUDA;
+}
+
 /**
  * @brief Retrieves the value of the tensor at the given index.
  *
diff --git a/tests/python/pytorch/graphbolt/test_dataloader.py b/tests/python/pytorch/graphbolt/test_dataloader.py
index 04cef3958356..b59260cf654e 100644
--- a/tests/python/pytorch/graphbolt/test_dataloader.py
+++ b/tests/python/pytorch/graphbolt/test_dataloader.py
@@ -9,9 +9,6 @@
 
 
 def test_DataLoader():
-    # https://pytorch.org/docs/master/notes/multiprocessing.html#cuda-in-multiprocessing
-    mp.set_start_method("spawn", force=True)
-
     N = 40
     B = 4
     itemset = dgl.graphbolt.ItemSet(torch.arange(N), names="seed_nodes")


Code review comments:
diff_hunk: @@ -19,6 +19,13 @@ inline bool is_accessible_from_gpu(torch::Tensor tensor) {
   return tensor.is_pinned() || tensor.device().type() == c10::DeviceType::CUDA;
 }
 
+/**
+ * @brief Checks whether the tensor is stored on the GPU.
+ */
+inline bool is_on_gpu(torch::Tensor tensor) {
+  return tensor.device().type() == c10::DeviceType::CUDA;
path: graphbolt/src/utils.h
text: @user1:
How about this?
```
template<typename... Args>
inline bool is_cuda(const Args&... args) {
    return (... && (args.is_cuda()));
}

```

@user1:
or directly call `nodes.is_cuda()` in user code?

@author:
I am not sure if there is a `is_cuda()` method.

@author:
But I will try getting a variable number of argument function working.

@user1:
https://github.com/pytorch/pytorch/blob/4a10e9eed4f8faf3e68106e5e86559158d83a200/aten/src/ATen/core/TensorBase.h#L430 ?

@author:
I wonder why they have this comment there:
https://github.com/pytorch/pytorch/blob/4a10e9eed4f8faf3e68106e5e86559158d83a200/aten/src/ATen/core/TensorBase.h#L431
```c++
// NB: this is not a native function to avoid dispatching overhead.
```

@user1:
https://github.com/pytorch/pytorch/blob/4a10e9eed4f8faf3e68106e5e86559158d83a200/c10/core/TensorImpl.h#L1097

@user1:
or use this one https://github.com/pytorch/pytorch/blob/4a10e9eed4f8faf3e68106e5e86559158d83a200/c10/core/Device.h#L81?

@author:
The current code already uses the last one you just pointed out.

@author:
Oh, you mean calling is_cuda on the device, Ok let me do that as the implementation is exactly the same as ours.

@user1:
I think the guideline is always use available APIs(and c++ APIs are usually consistent with python in pytorch) except for performance or other issue is identified.

diff_hunk: @@ -19,6 +19,13 @@ inline bool is_accessible_from_gpu(torch::Tensor tensor) {
   return tensor.is_pinned() || tensor.device().type() == c10::DeviceType::CUDA;
 }
 
+/**
+ * @brief Checks whether the tensor is stored on the GPU.
+ */
+inline bool is_on_gpu(torch::Tensor tensor) {
path: graphbolt/src/utils.h
text: @user1:
this raises a question: how to use GPU for sampling automatically? we need to make sure `graph` is pinned or on gpu and `nodes` should be on gpu?

@user1:
this is related to what will the `gb.DataLoader` interface be like, namely, how to choose sampling on GPU or CPU by user.

@author:
If graph is pinned or is on the device and if nodes is on device, we use gpu sampling.

@author:
It is already automatic, see the cuda examples PR #6861.

@author:
To dispatch to CUDA, at least one input tensor needs to be on GPU, and the rest either pinned or on GPU.

@user1:
I don't think it's automatic as user needs to call `graph.pin_memory_()` and add `copy_to()` datapipe to copy `nodes` to GPU manually in the examples you showed. This requires extra coding.

@author:
To use the GPU, we need to insert a CopyTo at the end of the data pipes anyway. What is the difference? It is just that we move the CopyTo to the beginning instead of having it at the end.

@author:
Do you think we can do better?

@user1:
we could add such copy_to datapipe within `gb.DataLoader`(including pinning graph) and user just need to specify a target device argument for dataloader.

@author:
We can discuss this in an issue, on Slack or in a meeting.