[Distributed] add constraints for metis partitioning. (#1708)

* add constraint.

* fix bugs.

* temp fix.

* fix lint

* add balance_ntypes and balance_edges

* add comments.

* fix.

* fix

include/dgl/graph_op.h

@@ -172,9 +172,10 @@ class GraphOp {
    * The partitioning algorithm assigns each vertex to a partition.
    * \param graph The input graph
    * \param k The number of partitions.
+   * \param vwgt the vertex weight array.
    * \return The partition assignments of all vertices.
    */
-  static IdArray MetisPartition(GraphPtr graph, int32_t k);
+  static IdArray MetisPartition(GraphPtr graph, int32_t k, NDArray vwgt);
 };
 
 }  // namespace dgl

python/dgl/distributed/partition.py

@@ -178,7 +178,7 @@ def load_partition_book(conf_file, part_id, graph=None):
         return GraphPartitionBook(part_id, num_parts, node_map, edge_map, graph)
 
 def partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method="metis",
-                    reshuffle=True):
+                    reshuffle=True, balance_ntypes=None, balance_edges=False):
     ''' Partition a graph for distributed training and store the partitions on files.
 
     The partitioning occurs in three steps: 1) run a partition algorithm (e.g., Metis) to
@@ -213,6 +213,16 @@ def partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method=
     * "orig_id" exists when reshuffle=True. It indicates the original node Ids in the original
     graph before reshuffling.
 
+    When performing Metis partitioning, we can put some constraint on the partitioning.
+    Current, it supports two constrants to balance the partitioning. By default, Metis
+    always tries to balance the number of nodes in each partition.
+
+    * `balance_ntypes` balances the number of nodes of different types in each partition.
+    * `balance_edges` balances the number of edges in each partition.
+
+    To balance the node types, a user needs to pass a vector of N elements to indicate
+    the type of each node. N is the number of nodes in the input graph.
+
     Parameters
     ----------
     g : DGLGraph
@@ -230,6 +240,10 @@ def partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method=
     reshuffle : bool
         Reshuffle nodes and edges so that nodes and edges in a partition are in
         contiguous Id range.
+    balance_ntypes : tensor
+        Node type of each node
+    balance_edges : bool
+        Indicate whether to balance the edges.
     '''
     if num_parts == 1:
         client_parts = {0: g}
@@ -242,7 +256,8 @@ def partition_graph(g, graph_name, num_parts, out_path, num_hops=1, part_method=
             g.ndata['orig_id'] = F.arange(0, g.number_of_nodes())
             g.edata['orig_id'] = F.arange(0, g.number_of_edges())
     elif part_method == 'metis':
-        node_parts = metis_partition_assignment(g, num_parts)
+        node_parts = metis_partition_assignment(g, num_parts, balance_ntypes=balance_ntypes,
+                                                balance_edges=balance_edges)
         client_parts = partition_graph_with_halo(g, node_parts, num_hops, reshuffle=reshuffle)
     elif part_method == 'random':
         node_parts = random_choice(num_parts, g.number_of_nodes())

python/dgl/transform.py

@@ -597,15 +597,12 @@ def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
     ------------
     g: DGLGraph
         The graph to be partitioned
-
     node_part: 1D tensor
         Specify which partition a node is assigned to. The length of this tensor
         needs to be the same as the number of nodes of the graph. Each element
         indicates the partition Id of a node.
-
     extra_cached_hops: int
         The number of hops a HALO node can be accessed.
-
     reshuffle : bool
         Resuffle nodes so that nodes in the same partition are in the same Id range.
 
@@ -656,9 +653,19 @@ def partition_graph_with_halo(g, node_part, extra_cached_hops, reshuffle=False):
         subg_dict[i] = subg
     return subg_dict
 
-def metis_partition_assignment(g, k):
+def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False):
     ''' This assigns nodes to different partitions with Metis partitioning algorithm.
 
+    When performing Metis partitioning, we can put some constraint on the partitioning.
+    Current, it supports two constrants to balance the partitioning. By default, Metis
+    always tries to balance the number of nodes in each partition.
+
+    * `balance_ntypes` balances the number of nodes of different types in each partition.
+    * `balance_edges` balances the number of edges in each partition.
+
+    To balance the node types, a user needs to pass a vector of N elements to indicate
+    the type of each node. N is the number of nodes in the input graph.
+
     After the partition assignment, we construct partitions.
 
     Parameters
@@ -667,6 +674,10 @@ def metis_partition_assignment(g, k):
         The graph to be partitioned
     k : int
         The number of partitions.
+    balance_ntypes : tensor
+        Node type of each node
+    balance_edges : bool
+        Indicate whether to balance the edges.
 
     Returns
     -------
@@ -676,14 +687,49 @@ def metis_partition_assignment(g, k):
     # METIS works only on symmetric graphs.
     # The METIS runs on the symmetric graph to generate the node assignment to partitions.
     sym_g = to_bidirected(g, readonly=True)
-    node_part = _CAPI_DGLMetisPartition(sym_g._graph, k)
+    vwgt = []
+    # To balance the node types in each partition, we can take advantage of the vertex weights
+    # in Metis. When vertex weights are provided, Metis will tries to generate partitions with
+    # balanced vertex weights. A vertex can be assigned with multiple weights. The vertex weights
+    # are stored in a vector of N * w elements, where N is the number of vertices and w
+    # is the number of weights per vertex. Metis tries to balance the first weight, and then
+    # the second weight, and so on.
+    # When balancing node types, we use the first weight to indicate the first node type.
+    # if a node belongs to the first node type, its weight is set to 1; otherwise, 0.
+    # Similary, we set the second weight for the second node type and so on. The number
+    # of weights is the same as the number of node types.
+    if balance_ntypes is not None:
+        assert len(balance_ntypes) == g.number_of_nodes(), \
+                "The length of balance_ntypes should be equal to #nodes in the graph"
+        balance_ntypes = utils.toindex(balance_ntypes)
+        balance_ntypes = balance_ntypes.tousertensor()
+        uniq_ntypes = F.unique(balance_ntypes)
+        for ntype in uniq_ntypes:
+            vwgt.append(F.astype(balance_ntypes == ntype, F.int64))
+
+    # When balancing edges in partitions, we use in-degree as one of the weights.
+    if balance_edges:
+        vwgt.append(F.astype(g.in_degrees(), F.int64))
+
+    # The vertex weights have to be stored in a vector.
+    if len(vwgt) > 0:
+        vwgt = F.stack(vwgt, 1)
+        shape = (np.prod(F.shape(vwgt),),)
+        vwgt = F.reshape(vwgt, shape)
+        vwgt = F.zerocopy_to_dgl_ndarray(vwgt)
+    else:
+        vwgt = F.zeros((0,), F.int64, F.cpu())
+        vwgt = F.zerocopy_to_dgl_ndarray(vwgt)
+
+    node_part = _CAPI_DGLMetisPartition(sym_g._graph, k, vwgt)
     if len(node_part) == 0:
         return None
     else:
         node_part = utils.toindex(node_part)
         return node_part.tousertensor()
 
-def metis_partition(g, k, extra_cached_hops=0, reshuffle=False):
+def metis_partition(g, k, extra_cached_hops=0, reshuffle=False,
+                    balance_ntypes=None, balance_edges=False):
     ''' This is to partition a graph with Metis partitioning.
 
     Metis assigns vertices to partitions. This API constructs subgraphs with the vertices assigned
@@ -691,6 +737,16 @@ def metis_partition(g, k, extra_cached_hops=0, reshuffle=False):
     not belong to the partition of a subgraph but are connected to the nodes
     in the partition within a fixed number of hops.
 
+    When performing Metis partitioning, we can put some constraint on the partitioning.
+    Current, it supports two constrants to balance the partitioning. By default, Metis
+    always tries to balance the number of nodes in each partition.
+
+    * `balance_ntypes` balances the number of nodes of different types in each partition.
+    * `balance_edges` balances the number of edges in each partition.
+
+    To balance the node types, a user needs to pass a vector of N elements to indicate
+    the type of each node. N is the number of nodes in the input graph.
+
     If `reshuffle` is turned on, the function reshuffles node Ids and edge Ids
     of the input graph before partitioning. After reshuffling, all nodes and edges
     in a partition fall in a contiguous Id range in the input graph.
@@ -705,26 +761,24 @@ def metis_partition(g, k, extra_cached_hops=0, reshuffle=False):
     ------------
     g: DGLGraph
         The graph to be partitioned
-
     k: int
         The number of partitions.
-
     extra_cached_hops: int
         The number of hops a HALO node can be accessed.
-
     reshuffle : bool
         Resuffle nodes so that nodes in the same partition are in the same Id range.
+    balance_ntypes : tensor
+        Node type of each node
+    balance_edges : bool
+        Indicate whether to balance the edges.
 
     Returns
     --------
     a dict of DGLGraphs
         The key is the partition Id and the value is the DGLGraph of the partition.
     '''
-    # METIS works only on symmetric graphs.
-    # The METIS runs on the symmetric graph to generate the node assignment to partitions.
-    sym_g = to_bidirected(g, readonly=True)
-    node_part = _CAPI_DGLMetisPartition(sym_g._graph, k)
-    if len(node_part) == 0:
+    node_part = metis_partition_assignment(g, k, balance_ntypes, balance_edges)
+    if node_part is None:
         return None
 
     # Then we split the original graph into parts based on the METIS partitioning results.

src/graph/metis_partition.cc

@@ -16,7 +16,7 @@ using namespace dgl::runtime;
 
 namespace dgl {
 
-IdArray GraphOp::MetisPartition(GraphPtr g, int k) {
+IdArray GraphOp::MetisPartition(GraphPtr g, int k, NDArray vwgt_arr) {
   // The index type of Metis needs to be compatible with DGL index type.
   CHECK_EQ(sizeof(idx_t), sizeof(dgl_id_t));
   ImmutableGraphPtr ig = std::dynamic_pointer_cast<ImmutableGraph>(g);
@@ -32,11 +32,23 @@ IdArray GraphOp::MetisPartition(GraphPtr g, int k) {
   IdArray part_arr = aten::NewIdArray(nvtxs);
   idx_t objval = 0;
   idx_t *part = static_cast<idx_t*>(part_arr->data);
+
+  int64_t vwgt_len = vwgt_arr->shape[0];
+  CHECK_EQ(sizeof(idx_t), vwgt_arr->dtype.bits / 8)
+      << "The vertex weight array doesn't have right type";
+  CHECK(vwgt_len % g->NumVertices() == 0)
+      << "The vertex weight array doesn't have right number of elements";
+  idx_t *vwgt = NULL;
+  if (vwgt_len > 0) {
+    ncon = vwgt_len / g->NumVertices();
+    vwgt = static_cast<idx_t*>(vwgt_arr->data);
+  }
+
   int ret = METIS_PartGraphKway(&nvtxs,      // The number of vertices
                                 &ncon,       // The number of balancing constraints.
                                 xadj,        // indptr
                                 adjncy,      // indices
-                                NULL,        // the weights of the vertices
+                                vwgt,        // the weights of the vertices
                                 NULL,        // The size of the vertices for computing
                                 // the total communication volume
                                 NULL,        // The weights of the edges
@@ -69,7 +81,8 @@ DGL_REGISTER_GLOBAL("transform._CAPI_DGLMetisPartition")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
     GraphRef g = args[0];
     int k = args[1];
-    *rv = GraphOp::MetisPartition(g.sptr(), k);
+    NDArray vwgt = args[2];
+    *rv = GraphOp::MetisPartition(g.sptr(), k, vwgt);
   });
 
 }   // namespace dgl

tests/compute/test_transform.py

@@ -255,9 +255,33 @@ def test_metis_partition():
     check_metis_partition(g, 0)
     check_metis_partition(g, 1)
     check_metis_partition(g, 2)
+    check_metis_partition_with_constraint(g)
+
+def check_metis_partition_with_constraint(g):
+    ntypes = np.zeros((g.number_of_nodes(),), dtype=np.int32)
+    ntypes[0:int(g.number_of_nodes()/4)] = 1
+    ntypes[int(g.number_of_nodes()*3/4):] = 2
+    subgs = dgl.transform.metis_partition(g, 4, extra_cached_hops=1, balance_ntypes=ntypes)
+    if subgs is not None:
+        for i in subgs:
+            subg = subgs[i]
+            parent_nids = F.asnumpy(subg.ndata[dgl.NID])
+            sub_ntypes = ntypes[parent_nids]
+            print('type0:', np.sum(sub_ntypes == 0))
+            print('type1:', np.sum(sub_ntypes == 1))
+            print('type2:', np.sum(sub_ntypes == 2))
+    subgs = dgl.transform.metis_partition(g, 4, extra_cached_hops=1,
+                                          balance_ntypes=ntypes, balance_edges=True)
+    if subgs is not None:
+        for i in subgs:
+            subg = subgs[i]
+            parent_nids = F.asnumpy(subg.ndata[dgl.NID])
+            sub_ntypes = ntypes[parent_nids]
+            print('type0:', np.sum(sub_ntypes == 0))
+            print('type1:', np.sum(sub_ntypes == 1))
+            print('type2:', np.sum(sub_ntypes == 2))
 
 def check_metis_partition(g, extra_hops):
-    # partitions with 1-hop HALO nodes
     subgs = dgl.transform.metis_partition(g, 4, extra_cached_hops=extra_hops)
     num_inner_nodes = 0
     num_inner_edges = 0
@@ -274,7 +298,7 @@ def check_metis_partition(g, extra_hops):
     if extra_hops == 0:
         return
 
-    # partitions with 1-hop HALO nodes and reshuffling nodes
+    # partitions with node reshuffling
     subgs = dgl.transform.metis_partition(g, 4, extra_cached_hops=extra_hops, reshuffle=True)
     num_inner_nodes = 0
     num_inner_edges = 0