load_partition() fails for heterographs whose edge relations are not unique

🐛 Bug

load_partition() fails to partition a heterogeneous graph whose edges have the same relation but are different.

To Reproduce

Steps to reproduce the behavior:

1. Create a heterograph (Creating a Heterogeneous Graph section): https://docs.dgl.ai/en/0.6.x/guide/graph-heterogeneous.html

import dgl
import torch as th
# Create a heterograph with 3 node types and 3 edges types.
graph_data = {
   ('drug', 'interacts', 'drug'): (th.tensor([0, 1]), th.tensor([1, 2])),
   ('drug', 'interacts', 'gene'): (th.tensor([0, 1]), th.tensor([2, 3])),
   ('drug', 'treats', 'disease'): (th.tensor([1]), th.tensor([2]))
}
g = dgl.heterograph(graph_data)

1. Run dgl.distributed.partition_graph(g, "example", out_path="output/", num_parts=2) and get the following error:

>>> dgl.distributed.partition_graph(g, "example", out_path="output/", num_parts=2)
Convert a graph into a bidirected graph: 0.000 seconds
Construct multi-constraint weights: 0.000 seconds
[18:35:34] /home/tiger/usr_name/dgl/src/graph/transform/metis_partition_hetero.cc:78: Partition a graph with 10 nodes and 10 edges into 2 parts and get 1 edge cuts
Metis partitioning: 0.000 seconds
Reshuffle nodes and edges: 0.001 seconds
Split the graph: 0.000 seconds
Construct subgraphs: 0.004 seconds
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/tiger/anaconda3/envs/usr_name/lib/python3.8/site-packages/dgl-0.8-py3.8-linux-x86_64.egg/dgl/distributed/partition.py", line 631, in partition_graph
    inner_etype_mask = inner_etype == g.get_etype_id(etype)
  File "/home/tiger/anaconda3/envs/usr_name/lib/python3.8/site-packages/dgl-0.8-py3.8-linux-x86_64.egg/dgl/heterograph.py", line 1261, in get_etype_id
    etid = self._etypes_invmap.get(self.to_canonical_etype(etype), None)
  File "/home/tiger/anaconda3/envs/usr_name/lib/python3.8/site-packages/dgl-0.8-py3.8-linux-x86_64.egg/dgl/heterograph.py", line 1153, in to_canonical_etype
    raise DGLError('Edge type "%s" is ambiguous. Please use canonical edge type '
dgl._ffi.base.DGLError: Edge type "interacts" is ambiguous. Please use canonical edge type in the form of (srctype, etype, dsttype)

Expected behavior

Partition should work for graphs whose canonical edges are distinct, and it shouldn’t depend on edge relation.

Environment

• DGL Version (e.g., 1.0): 0.8
• Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): PyTorch 1.9.0
• OS (e.g., Linux): Debian
• How you installed DGL (conda, pip, source): source
• Build command you used (if compiling from source): successful
• Python version: 3.8.11
• CUDA/cuDNN version (if applicable): NA
• GPU models and configuration (e.g. V100): NA
• Any other relevant information:

1. The simplest solution is to use different edge relation names. However, this may not be a good solution, because this makes the concept of canonical edge types the same as edge types.

The problem seems to happen whenever (in partition.py)

for etype in g.etypes:
    etype_id = g.get_etype_id(etype)

It is related to to_canonical_etype(). One can change g.etypes to g.canonical_etypes. But this has other repercussions when saving meta graph data into json and reloading the graph. Essentially the edge relation has to be distinct.

[VoVAllen]

The pr merge is postponed now because the distributed heterograph didn’t take duplicate etype name into account when development. More modifications are needed than #3249. Currently we suggest to user to avoid duplicate types now

[VoVAllen]

Keep this issue open for track. Will close it once PR merged