[Distributed] update Metis (#1856)

* fix bugs.

* improve.

* more updates.

* remove print.

* change the heterograph version.

* fix.
Co-authored-by: Ubuntu <ubuntu@ip-172-31-19-1.us-west-2.compute.internal>
Co-authored-by: Ubuntu <ubuntu@ip-172-31-30-164.us-west-2.compute.internal>

python/dgl/partition.py

@@ -202,7 +202,14 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False):
 
     # 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))
+        if balance_ntypes is None:
+            vwgt.append(F.astype(g.in_degrees(), F.int64))
+        else:
+            for ntype in uniq_ntypes:
+                nids = F.asnumpy(F.nonzero_1d(balance_ntypes == ntype))
+                degs = np.zeros((g.number_of_nodes(),), np.int64)
+                degs[nids] = F.asnumpy(g.in_degrees(nids))
+                vwgt.append(F.tensor(degs))
 
     # The vertex weights have to be stored in a vector.
     if len(vwgt) > 0:

python/dgl/transform.py

@@ -1071,7 +1071,14 @@ def metis_partition_assignment(g, k, balance_ntypes=None, balance_edges=False):
 
     # 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))
+        if balance_ntypes is None:
+            vwgt.append(F.astype(g.in_degrees(), F.int64))
+        else:
+            for ntype in uniq_ntypes:
+                nids = F.asnumpy(F.nonzero_1d(balance_ntypes == ntype))
+                degs = np.zeros((g.number_of_nodes(),), np.int64)
+                degs[nids] = F.asnumpy(g.in_degrees(nids))
+                vwgt.append(F.tensor(degs))
 
     # The vertex weights have to be stored in a vector.
     if len(vwgt) > 0:

src/graph/metis_partition.cc

@@ -41,6 +41,10 @@ IdArray GraphOp::MetisPartition(GraphPtr g, int k, NDArray vwgt_arr) {
     vwgt = static_cast<idx_t*>(vwgt_arr->data);
   }
 
+  idx_t options[METIS_NOPTIONS];
+  METIS_SetDefaultOptions(options);
+  options[METIS_OPTION_ONDISK] = 1;
+
   int ret = METIS_PartGraphKway(&nvtxs,      // The number of vertices
                                 &ncon,       // The number of balancing constraints.
                                 xadj,        // indptr
@@ -52,7 +56,7 @@ IdArray GraphOp::MetisPartition(GraphPtr g, int k, NDArray vwgt_arr) {
                                 &nparts,     // The number of partitions.
                                 NULL,        // the desired weight for each partition and constraint
                                 NULL,        // the allowed load imbalance tolerance
-                                NULL,        // the array of options
+                                options,     // the array of options
                                 &objval,      // the edge-cut or the total communication volume of
                                 // the partitioning solution
                                 part);

src/graph/transform/metis_partition_hetero.cc

@@ -45,6 +45,10 @@ IdArray MetisPartition(UnitGraphPtr g, int k, NDArray vwgt_arr) {
     vwgt = static_cast<idx_t *>(vwgt_arr->data);
   }
 
+  idx_t options[METIS_NOPTIONS];
+  METIS_SetDefaultOptions(options);
+  options[METIS_OPTION_ONDISK] = 1;
+
   int ret = METIS_PartGraphKway(
     &nvtxs,  // The number of vertices
     &ncon,   // The number of balancing constraints.
@@ -57,7 +61,7 @@ IdArray MetisPartition(UnitGraphPtr g, int k, NDArray vwgt_arr) {
     &nparts,  // The number of partitions.
     NULL,     // the desired weight for each partition and constraint
     NULL,     // the allowed load imbalance tolerance
-    NULL,     // the array of options
+    options,  // the array of options
     &objval,  // the edge-cut or the total communication volume of
     // the partitioning solution
     part);