dgl.metis_partition

dgl.metis_partition(g, k, extra_cached_hops=0, reshuffle=False, balance_ntypes=None, balance_edges=False, mode='k-way')[source]

This is to partition a graph with Metis partitioning.

Metis assigns vertices to partitions. This API constructs subgraphs with the vertices assigned to the partitions and their incoming edges. A subgraph may contain HALO nodes which does 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. The partitioend subgraphs have node data ‘orig_id’, which stores the node IDs in the original input graph.

The partitioned subgraph is stored in DGLGraph. The DGLGraph has the part_id node data that indicates the partition a node belongs to. The subgraphs do not contain the node/edge data in the input graph.

Parameters
  • 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.

  • mode (str, "k-way" or "recursive") – Whether use multilevel recursive bisection or multilevel k-way paritioning.

Returns

The key is the partition ID and the value is the DGLGraph of the partition.

Return type

a dict of DGLGraphs