Repository: dmlc/dgl

Issue:
[Bugfix] Fix corrupted memory when using multi-processing
## Description
<!-- Brief description. Refer to the related issues if existed.
It'll be great if relevant reviewers can be assigned as well.-->

When sampling with multi-processing, the sampler sometimes returns a corrupted original node ID array for destination nodes. The problem seems to be due to:
1. A PyTorch array is converted to a DGL array via DLPack.
1. It is then converted back to PyTorch array via DLPack. There is no data modification in the middle, so the conversion is actually redundant.
1. The array is then transmitted to another process via multi-processing mechanism.

The patch attempts to fix this by removing the redundant steps 1 and 2, though I haven't thoroughly verified whether the problem is completely solved.

## Checklist
Please feel free to remove inapplicable items for your PR.
- [x] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [x] Changes are complete (i.e. I finished coding on this PR)
- [x] All changes have test coverage
- [x] Code is well-documented
- [x] To the my best knowledge, examples are either not affected by this change,
      or have been fixed to be compatible with this change


Discussion:
Turns out to be a bug in PyTorch's DLPack implementation.  Submitted an issue in https://github.com/pytorch/pytorch/issues/36922.  The reproduction code goes as follows:
```python
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.dlpack import to_dlpack, from_dlpack

def collate(a):
    a = torch.LongTensor(a)
    acap = to_dlpack(a)
    b = from_dlpack(acap)

    return a, b

loader = DataLoader(dataset=[i for i in range(10000)],
                        batch_size = 2000,
                        collate_fn=collate,
                        shuffle=True,
                        num_workers=10)

for i, (a, b) in enumerate(loader):
    assert torch.equal(a, b)
```

Probably because `a` and `b` refer to the same content but somehow PyTorch & DLPack mess it up during tensor serialization.
Note that I cannot reproduce the error again if the tensor is created from DGL:
```python
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.dlpack import to_dlpack, from_dlpack
import dgl
import numpy as np

def collate(a):  # This causes the bug.
    a = torch.LongTensor(a)
    acap = to_dlpack(a)
    b = from_dlpack(acap)
    return a, b

def collate1(a):  # This works fine.
    a = np.array(a, dtype=np.int64)
    acap = dgl.ndarray.array(a).to_dlpack()
    b = from_dlpack(acap)
    return torch.tensor(a), b

loader = DataLoader(dataset=[i for i in range(10000)],
                    batch_size=200,
                    collate_fn=collate1,
                    shuffle=True,
                    num_workers=10)

for i, (a, b) in enumerate(loader):
    assert torch.equal(a, b)
```

Resolving patch:
diff --git a/python/dgl/transform.py b/python/dgl/transform.py
--- a/python/dgl/transform.py
+++ b/python/dgl/transform.py
@@ -1713,21 +1713,14 @@ def to_block(g, dst_nodes=None, include_dst_in_src=True):
             raise ValueError(
                 'Graph has more than one node type; please specify a dict for dst_nodes.')
         dst_nodes = {g.ntypes[0]: dst_nodes}
-    dst_nodes = {
-        ntype: utils.toindex(nodes, g._idtype_str).tousertensor()
-        for ntype, nodes in dst_nodes.items()}
 
-    # dst_nodes is now a dict
-    dst_nodes_nd = []
-    for ntype in g.ntypes:
-        nodes = dst_nodes.get(ntype, None)
-        if nodes is not None:
-            dst_nodes_nd.append(F.to_dgl_nd(nodes))
-        else:
-            dst_nodes_nd.append(nd.NULL[g._idtype_str])
+    dst_node_ids = [
+        utils.toindex(dst_nodes.get(ntype, []), g._idtype_str).tousertensor()
+        for ntype in g.ntypes]
+    dst_node_ids_nd = [F.to_dgl_nd(nodes) for nodes in dst_node_ids]
 
     new_graph_index, src_nodes_nd, induced_edges_nd = _CAPI_DGLToBlock(
-        g._graph, dst_nodes_nd, include_dst_in_src)
+        g._graph, dst_node_ids_nd, include_dst_in_src)
 
     # The new graph duplicates the original node types to SRC and DST sets.
     new_ntypes = (g.ntypes, g.ntypes)
@@ -1735,7 +1728,6 @@ def to_block(g, dst_nodes=None, include_dst_in_src=True):
     assert new_graph.is_unibipartite  # sanity check
 
     src_node_ids = [F.from_dgl_nd(src) for src in src_nodes_nd]
-    dst_node_ids = [F.from_dgl_nd(dst) for dst in dst_nodes_nd]
     edge_ids = [F.from_dgl_nd(eid) for eid in induced_edges_nd]
 
     node_frames = utils.extract_node_subframes_for_block(g, src_node_ids, dst_node_ids)


Test patch: