[DistDGL] enable exclude_edges for sample_etype_neighbors() (dmlc#7427)

Author: Rhett-Ying

examples/distributed/graphsage/node_classification_unsupervised.py

@@ -257,6 +257,21 @@ def run(args, device, data):
             for step, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(
                 dataloader
             ):
+                if args.debug:
+                    # Verify exclude_edges functionality.
+                    for block in blocks:
+                        current_eids = block.edata[dgl.EID]
+                        seed_eids = pos_graph.edata[dgl.EID]
+                        if exclude is None:
+                            assert th.any(th.isin(current_eids, seed_eids))
+                        elif exclude == "self":
+                            assert not th.any(th.isin(current_eids, seed_eids))
+                        elif exclude == "reverse_id":
+                            assert not th.any(th.isin(current_eids, seed_eids))
+                        else:
+                            raise ValueError(
+                                f"Unsupported exclude type: {exclude}"
+                            )
                 tic_step = time.time()
                 sample_t.append(tic_step - start)
 
@@ -449,6 +464,12 @@ def main(args):
         action="store_true",
         help="whether to remove edges during sampling",
     )
+    parser.add_argument(
+        "--debug",
+        default=False,
+        action="store_true",
+        help="whether to verify functionality of remove edges",
+    )
     args = parser.parse_args()
     print(args)
     main(args)

python/dgl/dataloading/dist_dataloader.py

@@ -1,5 +1,6 @@
 """Distributed dataloaders.
 """
+
 import inspect
 from abc import ABC, abstractmethod, abstractproperty
 from collections.abc import Mapping
@@ -83,6 +84,8 @@ def _find_exclude_eids(g, exclude_mode, eids, **kwargs):
     if exclude_mode is None:
         return None
     elif exclude_mode == "self":
+        if isinstance(eids, Mapping):
+            eids = {g.to_canonical_etype(k): v for k, v in eids.items()}
         return eids
     elif exclude_mode == "reverse_id":
         return _find_exclude_eids_with_reverse_id(

python/dgl/distributed/dist_graph.py

@@ -10,12 +10,13 @@
 
 from .. import backend as F, graphbolt as gb, heterograph_index
 from .._ffi.ndarray import empty_shared_mem
-from ..base import ALL, dgl_warning, DGLError, EID, ETYPE, is_all, NID
+from ..base import ALL, DGLError, EID, ETYPE, is_all, NID
 from ..convert import graph as dgl_graph, heterograph as dgl_heterograph
 from ..frame import infer_scheme
 
 from ..heterograph import DGLGraph
 from ..ndarray import exist_shared_mem_array
+from ..sampling.utils import EidExcluder
 from ..transforms import compact_graphs
 from . import graph_services, role, rpc
 from .dist_tensor import DistTensor
@@ -1422,17 +1423,14 @@ def sample_neighbors(
         # pylint: disable=unused-argument
         """Sample neighbors from a distributed graph."""
         if len(self.etypes) > 1:
-            if exclude_edges is not None:
-                dgl_warning(
-                    "exclude_edges is not supported for a graph with multiple edge types."
-                )
             frontier = graph_services.sample_etype_neighbors(
                 self,
                 seed_nodes,
                 fanout,
                 replace=replace,
                 etype_sorted=etype_sorted,
                 prob=prob,
+                exclude_edges=None,
                 use_graphbolt=self._use_graphbolt,
             )
         else:
@@ -1442,9 +1440,18 @@ def sample_neighbors(
                 fanout,
                 replace=replace,
                 prob=prob,
-                exclude_edges=exclude_edges,
+                exclude_edges=None,
                 use_graphbolt=self._use_graphbolt,
             )
+        # [TODO][Rui]
+        # For now, exclude_edges is applied after sampling. Namely, we first sample
+        # the neighbors and then exclude the edges before returning frontier. This
+        # is probably not efficient. We could try to exclude the edges during
+        # sampling. Or we pass exclude_edges IDs to local and remote sampling
+        # functions and let them handle the exclusion.
+        if exclude_edges is not None:
+            eid_excluder = EidExcluder(exclude_edges)
+            frontier = eid_excluder(frontier)
         return frontier
 
     def _get_ndata_names(self, ntype=None):

python/dgl/distributed/graph_services.py

@@ -257,6 +257,7 @@ def _sample_etype_neighbors_dgl(
     fan_out,
     edge_dir="in",
     prob=None,
+    exclude_edges=None,
     replace=False,
     etype_offset=None,
     etype_sorted=False,
@@ -278,9 +279,10 @@ def _sample_etype_neighbors_dgl(
         local_ids,
         etype_offset,
         fan_out,
-        edge_dir,
-        prob,
-        replace,
+        edge_dir=edge_dir,
+        prob=prob,
+        exclude_edges=exclude_edges,
+        replace=replace,
         etype_sorted=etype_sorted,
         _dist_training=True,
     )
@@ -481,13 +483,15 @@ def __init__(
         fan_out,
         edge_dir="in",
         prob=None,
+        exclude_edges=None,
         replace=False,
         etype_sorted=True,
         use_graphbolt=False,
     ):
         self.seed_nodes = nodes
         self.edge_dir = edge_dir
         self.prob = prob
+        self.exclude_edges = exclude_edges
         self.replace = replace
         self.fan_out = fan_out
         self.etype_sorted = etype_sorted
@@ -498,6 +502,7 @@ def __setstate__(self, state):
             self.seed_nodes,
             self.edge_dir,
             self.prob,
+            self.exclude_edges,
             self.replace,
             self.fan_out,
             self.etype_sorted,
@@ -509,6 +514,7 @@ def __getstate__(self):
             self.seed_nodes,
             self.edge_dir,
             self.prob,
+            self.exclude_edges,
             self.replace,
             self.fan_out,
             self.etype_sorted,
@@ -536,6 +542,7 @@ def process_request(self, server_state):
             self.fan_out,
             edge_dir=self.edge_dir,
             prob=probs,
+            exclude_edges=self.exclude_edges,
             replace=self.replace,
             etype_offset=etype_offset,
             etype_sorted=self.etype_sorted,
@@ -801,20 +808,19 @@ def _frontier_to_heterogeneous_graph(g, frontier, gpb):
         src_ntype_id = g.get_ntype_id(src_ntype)
         dst_ntype_id = g.get_ntype_id(dst_ntype)
         type_idx = etype_ids == etid
-        if F.sum(type_idx, 0) > 0:
-            data_dict[etype] = (
-                F.boolean_mask(src, type_idx),
-                F.boolean_mask(dst, type_idx),
-            )
-            if "DGL_DIST_DEBUG" in os.environ:
-                assert torch.all(
-                    src_ntype_id == src_ntype_ids[type_idx]
-                ), "source ntype is is not expected."
-                assert torch.all(
-                    dst_ntype_id == dst_ntype_ids[type_idx]
-                ), "destination ntype is is not expected."
-            if type_wise_eids is not None:
-                edge_ids[etype] = F.boolean_mask(type_wise_eids, type_idx)
+        data_dict[etype] = (
+            F.boolean_mask(src, type_idx),
+            F.boolean_mask(dst, type_idx),
+        )
+        if "DGL_DIST_DEBUG" in os.environ:
+            assert torch.all(
+                src_ntype_id == src_ntype_ids[type_idx]
+            ), "source ntype is is not expected."
+            assert torch.all(
+                dst_ntype_id == dst_ntype_ids[type_idx]
+            ), "destination ntype is is not expected."
+        if type_wise_eids is not None:
+            edge_ids[etype] = F.boolean_mask(type_wise_eids, type_idx)
     hg = heterograph(
         data_dict,
         {ntype: g.num_nodes(ntype) for ntype in g.ntypes},
@@ -832,6 +838,7 @@ def sample_etype_neighbors(
     fanout,
     edge_dir="in",
     prob=None,
+    exclude_edges=None,
     replace=False,
     etype_sorted=True,
     use_graphbolt=False,
@@ -879,6 +886,8 @@ def sample_etype_neighbors(
         The features must be non-negative floats, and the sum of the features of
         inbound/outbound edges for every node must be positive (though they don't have
         to sum up to one).  Otherwise, the result will be undefined.
+    exclude_edges : tensor, optional
+        The edges to exclude when sampling.
     replace : bool, optional
         If True, sample with replacement.
 
@@ -947,6 +956,7 @@ def issue_remote_req(node_ids):
             fanout,
             edge_dir=edge_dir,
             prob=_prob,
+            exclude_edges=exclude_edges,
             replace=replace,
             etype_sorted=etype_sorted,
             use_graphbolt=use_graphbolt,
@@ -974,6 +984,7 @@ def local_access(local_g, partition_book, local_nids):
             fanout,
             edge_dir=edge_dir,
             prob=_prob,
+            exclude_edges=exclude_edges,
             replace=replace,
             etype_offset=etype_offset,
             etype_sorted=etype_sorted,

python/dgl/sampling/neighbor.py

@@ -49,9 +49,11 @@ def _prepare_edge_arrays(g, arg):
                     result.append(None)
 
             result = [
-                F.to_dgl_nd(F.copy_to(F.tensor([], dtype=dtype), ctx))
-                if x is None
-                else x
+                (
+                    F.to_dgl_nd(F.copy_to(F.tensor([], dtype=dtype), ctx))
+                    if x is None
+                    else x
+                )
                 for x in result
             ]
             return result
@@ -74,6 +76,7 @@ def sample_etype_neighbors(
     fanout,
     edge_dir="in",
     prob=None,
+    exclude_edges=None,
     replace=False,
     copy_ndata=True,
     copy_edata=True,
@@ -116,6 +119,11 @@ def sample_etype_neighbors(
 
         The features must be non-negative floats or boolean.  Otherwise, the
         result will be undefined.
+    exclude_edges: tensor or dict
+        Edge IDs to exclude during sampling neighbors for the seed nodes.
+
+        This argument can take a single ID tensor or a dictionary of edge types and ID tensors.
+        If a single tensor is given, the graph must only have one type of nodes.
     replace : bool, optional
         If True, sample with replacement.
     copy_ndata: bool, optional
@@ -154,6 +162,10 @@ def sample_etype_neighbors(
     As a result, users should avoid performing in-place operations
     on the node features of the new graph to avoid feature corruption.
     """
+    if exclude_edges is not None:
+        raise DGLError(
+            "exclude_edges is not supported for sample_etype_neighbors"
+        )
     if g.device != F.cpu():
         raise DGLError("The graph should be in cpu.")
     # (BarclayII) because the homogenized graph no longer contains the *name* of edge