Repository: dmlc/dgl

Issue:
I can't run dgl/examples/pytorch/han/train_sampling.py with cuda
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. I just run dgl/examples/pytorch/han/train_sampling.py with cuda, it has an error

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->
<img width="1103" alt="image" src="https://user-images.githubusercontent.com/53086386/156488968-32cde64a-1c14-4de7-93de-211ca4dd306e.png">


## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

 - DGL Version: 0.8.0+cu11.1
 - Backend Library & Version:  PyTorch0.9.1
 - OS: Linux
 - How you installed DGL: pip
 - Build command you used (if compiling from source):
 - Python version: 3.7.4
 - CUDA/cuDNN version: cuda11.1
 - GPU models and configuration: A100
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Able to reproduce the issue with Python 3.9 + Pytorch 1.10.2 + DGL 0.8 + CUDA 11.3 with a different error message:
```
Downloading /home/somepath/.dgl/ACM.mat from https://data.dgl.ai/dataset/ACM.mat...
total_params: 252611
total trainable params: 252611
Traceback (most recent call last):
  File "/home/somepath/dgl/examples/pytorch/han/train_sampling.py", line 271, in <module>
    main(args)
  File "/home/somepath/dgl/examples/pytorch/han/train_sampling.py", line 219, in main
    blocks = [block.to(args['device']) for block in blocks]
  File "/home/somepath/dgl/examples/pytorch/han/train_sampling.py", line 219, in <listcomp>
    blocks = [block.to(args['device']) for block in blocks]
  File "/home/somepath/miniconda3/lib/python3.9/site-packages/dgl/heterograph.py", line 5429, in to
    ret._graph = self._graph.copy_to(utils.to_dgl_context(device))
  File "/home/somepath/miniconda3/lib/python3.9/site-packages/dgl/heterograph_index.py", line 236, in copy_to
    return _CAPI_DGLHeteroCopyTo(self, ctx.device_type, ctx.device_id)
  File "dgl/_ffi/_cython/./function.pxi", line 287, in dgl._ffi._cy3.core.FunctionBase.__call__
  File "dgl/_ffi/_cython/./function.pxi", line 222, in dgl._ffi._cy3.core.FuncCall
  File "dgl/_ffi/_cython/./function.pxi", line 211, in dgl._ffi._cy3.core.FuncCall3
  File "dgl/_ffi/_cython/./base.pxi", line 155, in dgl._ffi._cy3.core.CALL
dgl._ffi.base.DGLError: [02:06:03] /opt/dgl/src/runtime/cuda/cuda_device_api.cc:246: Check failed: e == cudaSuccess || e == cudaErrorCudartUnloading: CUDA: unspecified launch failure
Stack trace:
  [bt] (0) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(dmlc::LogMessageFatal::~LogMessageFatal()+0x4f) [0x7f34520de44f]
  [bt] (1) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(dgl::runtime::CUDADeviceAPI::CopyDataFromTo(void const*, unsigned long, void*, unsigned long, unsigned long, DLContext, DLContext, DLDataType, void*)+0x79) [0x7f34525b6a29]
  [bt] (2) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(dgl::runtime::NDArray::CopyFromTo(DLTensor*, DLTensor*, void*)+0x26c) [0x7f3452424d0c]
  [bt] (3) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(dgl::runtime::NDArray::CopyTo(DLContext const&, void* const&) const+0xfb) [0x7f345245ec8b]
  [bt] (4) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(dgl::UnitGraph::CopyTo(std::shared_ptr<dgl::BaseHeteroGraph>, DLContext const&, void* const&)+0x2ff) [0x7f345257964f]
  [bt] (5) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(dgl::HeteroGraph::CopyTo(std::shared_ptr<dgl::BaseHeteroGraph>, DLContext const&, void* const&)+0x109) [0x7f3452470849]
  [bt] (6) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(+0x6aeb59) [0x7f345247db59]
  [bt] (7) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/libdgl.so(DGLFuncCall+0x48) [0x7f3452402d78]
  [bt] (8) /home/somepath/miniconda3/lib/python3.9/site-packages/dgl/_ffi/_cy3/core.cpython-39-x86_64-linux-gnu.so(+0x16ce2) [0x7f3451b02ce2]
```
@yaox12 Also confirmed.  Though I'm not sure what `unspecified launch failure` means.  A `cuda-memcheck` reports the error in `torch`:
```
(base) ubuntu@ip-172-31-34-161:~/dgl/examples/pytorch/han$ CUDA_LAUNCH_BLOCKING=1 cuda-memcheck python train_sampling.py --device cuda
========= CUDA-MEMCHECK
Downloading /home/ubuntu/.dgl/ACM.mat from https://data.dgl.ai/dataset/ACM.mat...
========= Internal Memcheck Error: Initialization failed
=========     Saved host backtrace up to driver entry point at error
=========     Host Frame:/usr/lib/x86_64-linux-gnu/libcuda.so [0x28746b]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 [0x3127c]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 [0x1ff4e]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 [0x37974]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 [0x395aa]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 [0x2f32e]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 [0x1d788]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libcudart-a7b20f20.so.11.0 (cudaStreamIsCapturing + 0x11e) [0x631ee]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libc10_cuda.so [0x217be]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libc10_cuda.so [0x250ba]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libc10_cuda.so [0x264b2]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libc10_cuda.so [0x268e2]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cuda_cpp.so (_Z22THCStorage_resizeBytesP8THCStatePN3c1011StorageImplEl + 0xca) [0x2dd614a]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cuda_cpp.so (_ZN2at6native18empty_strided_cudaEN3c108ArrayRefIlEES3_NS1_8optionalINS1_10ScalarTypeEEENS4_INS1_6LayoutEEENS4_INS1_6DeviceEEENS4_IbEE + 0x303) [0x2d63993]     =========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cuda_cu.so [0x25abd14]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cuda_cu.so [0x25abd9f]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so [0x1d152c0]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (_ZN2at4_ops13empty_strided4callEN3c108ArrayRefIlEES4_NS2_8optionalINS2_10ScalarTypeEEENS5_INS2_6LayoutEEENS5_INS2_6DeviceEEENS5_IbEE + 0x1ca) [0x1affeda]            =========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (_ZN2at6native8_to_copyERKNS_6TensorEN3c108optionalINS4_10ScalarTypeEEENS5_INS4_6LayoutEEENS5_INS4_6DeviceEEENS5_IbEEbNS5_INS4_12MemoryFormatEEE + 0x136d) [0x164eced]=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so [0x1e2e51a]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so [0x1d1115e]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (_ZN2at4_ops8_to_copy10redispatchEN3c1014DispatchKeySetERKNS_6TensorENS2_8optionalINS2_10ScalarTypeEEENS7_INS2_6LayoutEEENS7_INS2_6DeviceEEENS7_IbEEbNS7_INS2_12MemoryFormatEEE + 0x137) [0x18d6767]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so [0x3239d84]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so [0x323a24d]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (_ZN2at4_ops8_to_copy4callERKNS_6TensorEN3c108optionalINS5_10ScalarTypeEEENS6_INS5_6LayoutEEENS6_INS5_6DeviceEEENS6_IbEEbNS6_INS5_12MemoryFormatEEE + 0x1d8) [0x194c138]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (_ZN2at6native2toERKNS_6TensorEN3c108optionalINS4_10ScalarTypeEEENS5_INS4_6LayoutEEENS5_INS4_6DeviceEEENS5_IbEEbbNS5_INS4_12MemoryFormatEEE + 0x155) [0x164ce25]      =========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so [0x1eb53c0]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_cpu.so (_ZN2at4_ops15to_dtype_layout4callERKNS_6TensorEN3c108optionalINS5_10ScalarTypeEEENS6_INS5_6LayoutEEENS6_INS5_6DeviceEEENS6_IbEEbbNS6_INS5_12MemoryFormatEEE + 0x1f9)
[0x1c01ab9]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_python.so [0x587163]
=========     Host Frame:/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/lib/libtorch_python.so [0x6d0ddb]
=========     Host Frame:python [0x1a4fd7]
=========     Host Frame:python [0xffb08]
=========     Host Frame:python (_PyEval_EvalCodeWithName + 0x7df) [0x18abaf]
=========     Host Frame:python (_PyFunction_Vectorcall + 0x1e3) [0x18b643]
...
```
Build DGL TOT from source with nightly PyTorch (I use an NGC container with PyTorch `1.11.0a0+2c916ef`) seems to solve the issue.
@BarclayII That specific error is due to `cuda-memcheck` failing to setup its hooks. Set `CUDA_MEMCHECK_PATCH_MODULE=1` to solve this (https://docs.nvidia.com/cuda/cuda-memcheck/index.html#environment-variables). Also, note that `cuda-memcheck` is being deprecated in favor of `compute-sanitizer` (https://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/index.html).
Upon further investigation, I believe this is related to the `IsPinned()` function being added in 0.8. The torch.utils.dataloader's child process are created with `fork()`, and when a cuda context is manipulated in one of these processes it corrupts the cuda state of the spawning process.

I was able to reproduce the issue, and when I either set the num_workers=0 or `mp.set_spawn_method("spawn")` it worked.
spawn should be fine since our graph structure now goes through the shared memory
This is a very specific fix for this issue (workaround), but it seems just calling `g.create_formats_()` before passing the graph to the dataloader avoids the cuda context being referenced by the worker processes. The IsPinned() was being called by the `SetCOO`/`SetCSR`/`SetCSC` methods on Heterograph in the worker processes. When all formats exist, the worker processes do not call `IsPinned()`.

@yijicheng Can you see if this solves your issue:
```
diff --git a/examples/pytorch/han/train_sampling.py b/examples/pytorch/han/train_sampling.py
index c84e50b6..ffabf29a 100644
--- a/examples/pytorch/han/train_sampling.py
+++ b/examples/pytorch/han/train_sampling.py
@@ -183,6 +183,8 @@ def main(args):
     else:
         raise NotImplementedError('Unsupported dataset {}'.format(args['dataset']))

+    g.create_formats_()
+
     # Is it need to set different neighbors numbers for different meta-path based graph?
     num_neighbors = args['num_neighbors']
     han_sampler = HANSampler(g, metapath_list, num_neighbors)
```
@nv-dlasalle It's really interesting, because I think `is_pinned()` would be called in 
```python
if not g.is_pinned() and (F.dtype(data) != g.idtype or F.context(data) != g.device):
```
in `dgl.utils.prepare_tensor()`.
@nv-dlasalle So far `SetCOOMatrix` is only called when a graph is forked from one process to another via `ForkingPickler` (in particular in `HeteroForkingUnpickle` C++ function).  Retrospectively I'm not sure if having a `SetCOOMatrix` or alike is a good idea since it replaces the underlying storage directly (which will cause problems if the COOMatrix is pinned and the graph is unpinned, or vice versa, hence the check).

That being said, @yaox12 is right on `is_pinned` being called elsewhere, e.g. in `sample_neighbors`:
```python
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 287, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 40, in fetch
    return self.collate_fn(data)
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/dataloader.py", line 531, in __call__
    batch = self.sample_func(self.g, items)
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/base.py", line 252, in sample
    result = self.sample_blocks(g, seed_nodes, exclude_eids=exclude_eids)
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/dataloading/neighbor_sampler.py", line 110, in sample_blocks
    frontier = g.sample_neighbors(
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/utils/internal.py", line 909, in _fn
    return func(*args, **kwargs)
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/sampling/neighbor.py", line 290, in sample_neighbors
    if F.device_type(g.device) == 'cpu' and not g.is_pinned():
  File "/home/ubuntu/miniconda3/lib/python3.8/site-packages/dgl-0.9-py3.8-linux-x86_64.egg/dgl/heterograph.py", line 5561, in is_pinned
    return self._graph.is_pinned()
```
@BarclayII This example is using randomwalks, not neighbor sampling, which is why that isn't an issue here. However I think we could change that specific instance to avoid the is_pinned() call:
```
if F.device_type(g.device) == 'cpu' and not g.is_pinned():
```
to something like
```
if F.device_type(g.device) == 'cpu' and F.device_type(F.context(nodes)) == 'cpu':
```
since UVA sampling requires the nodes to be on the GPU.

@yaox12 I had assumed prepare_tensor would be causing an issue, but in my testing its not in this example. I think its because this example uses the PinSage sampler which does not call prepare_tensor as far as I can tell. That being said, changing the order of that if statement will be helpful:
```
if not g.is_pinned() and (F.dtype(data) != g.idtype or F.context(data) != g.device):
```
to
```
if (F.dtype(data) != g.idtype or F.context(data) != g.device) and not g.is_pinned():
```
Since we restrict UVA sampling to num_workers=0, we wouldn't be calling `g.is_pinned()` in a forked process here unless sampling would fail with an error anyways.

These may be reasonable workarounds for the short term, but I think we need a solution like in PyTorch where we can avoid calling into the CUDA context in forked workers.


> @BarclayII This example is using randomwalks, not neighbor sampling, which is why that isn't an issue here

That particular stacktrace comes from me running the GraphSAGE example.

> These may be reasonable workarounds for the short term, but I think we need a solution like in PyTorch where we can avoid calling into the CUDA context in forked workers.

I think PyTorch uses `pthread_atfork` to prevent CUDA calls from running in forked childs with the `poison_fork` function:
https://github.com/pytorch/pytorch/blob/198d727d01335f58f66943610eb023eabb94e98e/torch/csrc/cuda/Module.cpp#L36-L54

A question though: when `is_pinned` is called in a forked child, should we raise an error or return False?  Looks like `is_pinned` is  very likely to be called in worker processes.

Another question is when we should register the `poison_fork` hook.  In PyTorch this is called before any CUDA call, so in DGL we should also run a similar hook after one of PyTorch and DGL has called any CUDA function.  The former can be done with `torch.cuda.is_initialized` but I'm not sure how we are supposed to check for the latter.
@nv-dlasalle Hmm I still see the error even if I remove the `is_pinned` Python call in `prepare_tensor`.  If I directly return False in `cuda_device_api.cc` then the code worked fine.

Resolving patch:
diff --git a/python/dgl/utils/checks.py b/python/dgl/utils/checks.py
--- a/python/dgl/utils/checks.py
+++ b/python/dgl/utils/checks.py
@@ -30,7 +30,7 @@ def prepare_tensor(g, data, name):
         Data in tensor object.
     """
     if F.is_tensor(data):
-        if not g.is_pinned() and (F.dtype(data) != g.idtype or F.context(data) != g.device):
+        if (F.dtype(data) != g.idtype or F.context(data) != g.device) and not g.is_pinned():
             raise DGLError('Expect argument "{}" to have data type {} and device '
                            'context {}. But got {} and {}.'.format(
                                name, g.idtype, g.device, F.dtype(data), F.context(data)))


Test patch: