Repository: dmlc/dgl

Issue:
Bus error(core dumped) in EdgeDataLoader of unsupervised graphsage example (20M edges)
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## To Reproduce

Steps to reproduce the behavior:

1. In the example code: https://github.com/dmlc/dgl/blob/master/examples/pytorch/graphsage/advanced/train_sampling_unsupervised.py
2. Feed a dataset with 10M nodes and 20M edges
3. Bus error(core dumped)  happened in EdgeDataLoader before training and DataLoader in inference.

<!-- 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): 0.8.1cu101
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.8.1+cu101
 - OS (e.g., Linux): Docker
 - How you installed DGL (`conda`, `pip`, source): pip
 - Build command you used (if compiling from source):
 - Python version: 3.6.13
 - CUDA/cuDNN version (if applicable): 101
 - GPU models and configuration (e.g. V100): V100
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
The num of workers is set to 0 in EdgeDataLoader
The bug didn't show when the number of train_seeds in EdgeDataLoader is small, like 1M
How many GPUs did you use? Have you changed the args such as `--graph-device`, `--data-device`?
> How many GPUs did you use? Have you changed the args such as `--graph-device`, `--data-device`?

one GPU.  Nope.  Have tried different args but got the same result.
Can you try adding `--shm-size=64g` (large enough to store the whole graph) to your `docker run` command?
I can not change the docker command. I set the num of workers to 0 so I think it doesn't use shared memory.
> Can you try adding `--shm-size=64g` (large enough to store the whole graph) to your `docker run` command?

The docker environment is created automatically. I am not able to change the shm size.
This line of code in dataloader will create a shared memory array for shuffling.
https://github.com/dmlc/dgl/blob/5ba5106acab6a642e9b790e5331ee519112a5623/python/dgl/dataloading/dataloader.py#L146-L149
When `len(train_seeds)` > 8M, the shared tensor will run out of Docker's default shm size 64MB.
@jermainewang @BarclayII 
Can we avoid that when users do not intend to use shared memory?
According to the comments, the shared tensor is used for `persistent_workers=True` (`or num_workers > 0` I think?). We can change the code to use shared tensors only when these conditions hold. @BarclayII may know more details.
> We can change the code to use shared tensors only when these conditions hold.

That sounds reasonable.  I'm not sure how PyTorch shares the tensor to forked subprocesses though: if PyTorch uses shared memory then we are technically still copying the ID tensor into shared memory implicitly.

Resolving patch:
diff --git a/python/dgl/dataloading/dataloader.py b/python/dgl/dataloading/dataloader.py
--- a/python/dgl/dataloading/dataloader.py
+++ b/python/dgl/dataloading/dataloader.py
@@ -134,7 +134,7 @@ class TensorizedDataset(torch.utils.data.IterableDataset):
     """Custom Dataset wrapper that returns a minibatch as tensors or dicts of tensors.
     When the dataset is on the GPU, this significantly reduces the overhead.
     """
-    def __init__(self, indices, batch_size, drop_last, shuffle):
+    def __init__(self, indices, batch_size, drop_last, shuffle, use_shared_memory):
         if isinstance(indices, Mapping):
             self._mapping_keys = list(indices.keys())
             self._device = next(iter(indices.values())).device
@@ -147,7 +147,9 @@ def __init__(self, indices, batch_size, drop_last, shuffle):
         # Use a shared memory array to permute indices for shuffling.  This is to make sure that
         # the worker processes can see it when persistent_workers=True, where self._indices
         # would not be duplicated every epoch.
-        self._indices = torch.arange(self._id_tensor.shape[0], dtype=torch.int64).share_memory_()
+        self._indices = torch.arange(self._id_tensor.shape[0], dtype=torch.int64)
+        if use_shared_memory:
+            self._indices.share_memory_()
         self.batch_size = batch_size
         self.drop_last = drop_last
         self._shuffle = shuffle
@@ -551,15 +553,16 @@ def __call__(self, worker_id):
 
 
 def create_tensorized_dataset(indices, batch_size, drop_last, use_ddp, ddp_seed,
-                              shuffle):
+                              shuffle, use_shared_memory):
     """Converts a given indices tensor to a TensorizedDataset, an IterableDataset
     that returns views of the original tensor, to reduce overhead from having
     a list of scalar tensors in default PyTorch DataLoader implementation.
     """
     if use_ddp:
+        # DDP always uses shared memory
         return DDPTensorizedDataset(indices, batch_size, drop_last, ddp_seed, shuffle)
     else:
-        return TensorizedDataset(indices, batch_size, drop_last, shuffle)
+        return TensorizedDataset(indices, batch_size, drop_last, shuffle, use_shared_memory)
 
 
 def _get_device(device):
@@ -825,7 +828,8 @@ def __init__(self, graph, indices, graph_sampler, device=None, use_ddp=False,
                 isinstance(indices, Mapping) and
                 all(torch.is_tensor(v) for v in indices.values()))):
             self.dataset = create_tensorized_dataset(
-                indices, batch_size, drop_last, use_ddp, ddp_seed, shuffle)
+                indices, batch_size, drop_last, use_ddp, ddp_seed, shuffle,
+                kwargs.get('persistent_workers', False))
         else:
             self.dataset = indices
 


Test patch: