7.4 高级图分区

本章涵盖了一些关于图分区的高级主题。

METIS分区算法

METIS 是一种最先进的图分割算法，能够生成具有最少跨分区边数的分割，使其适用于分布式消息传递，其中网络通信量与跨分区边数成正比。DGL 已将 METIS 集成为其 dgl.distributed.partition_graph() API 中的默认分割算法。

输出格式

无论使用何种分区算法，分区结果都存储在按以下方式组织的数据文件中：

data_root_dir/
  |-- graph_name.json       # partition configuration file in JSON
  |-- part0/                # data for partition 0
  |  |-- node_feats.dgl     # node features stored in binary format
  |  |-- edge_feats.dgl     # edge features stored in binary format
  |  |-- graph.dgl          # graph structure of this partition stored in binary format
  |
  |-- part1/                # data for partition 1
  |  |-- node_feats.dgl
  |  |-- edge_feats.dgl
  |  |-- graph.dgl
  |
  |-- ...                   # data for other partitions

当分发到集群时，元数据 JSON 应该被复制到所有机器上，而 partX 文件夹应该相应地分发。

DGL 提供了一个 dgl.distributed.load_partition() 函数来加载一个分区进行检查。

>>> import dgl
>>> # load partition 0
>>> part_data = dgl.distributed.load_partition('data_root_dir/graph_name.json', 0)
>>> g, nfeat, efeat, partition_book, graph_name, ntypes, etypes = part_data  # unpack
>>> print(g)
Graph(num_nodes=966043, num_edges=34270118,
      ndata_schemes={'orig_id': Scheme(shape=(), dtype=torch.int64),
                     'part_id': Scheme(shape=(), dtype=torch.int64),
                     '_ID': Scheme(shape=(), dtype=torch.int64),
                     'inner_node': Scheme(shape=(), dtype=torch.int32)}
      edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64),
                     'inner_edge': Scheme(shape=(), dtype=torch.int8),
                     'orig_id': Scheme(shape=(), dtype=torch.int64)})

如`ID mapping`_部分所述，每个分区都携带保存为ndata或edata的辅助信息，例如原始节点/边ID、分区ID等。每个分区不仅保存其拥有的节点/边，还包括与分区相邻的节点/边（称为HALO节点/边）。inner_node和inner_edge指示节点/边是否真正属于该分区（值为True）或是HALO节点/边（值为False）。

load_partition() 函数一次性加载所有数据。用户可以使用 dgl.distributed.load_partition_feats() 和 dgl.distributed.load_partition_book() API 分别加载特征或分区书。

并行METIS分区

对于需要并行预处理的大规模图，DGL支持 ParMETIS 作为分区算法的选择之一。

注意

由于ParMETIS不支持异构图，用户需要在运行ParMETIS之前和之后进行ID转换。查看章节 7.5 异构图内部机制以获取解释。

注意

请确保输入到ParMETIS的图没有重复边（或平行边）和自环边。

ParMETIS 安装

ParMETIS 需要 METIS 和 GKLib。请按照这里的说明编译并安装 GKLib。对于编译和安装 METIS，请按照以下说明使用 GIT 克隆 METIS 并使用 int64 支持进行编译。

git clone https://github.com/KarypisLab/METIS.git
make config shared=1 cc=gcc prefix=~/local i64=1
make install

目前，我们需要手动编译和安装ParMETIS。我们按照以下方式克隆DGL分支的ParMETIS：

git clone --branch dgl https://github.com/KarypisLab/ParMETIS.git

然后编译并安装ParMETIS。

make config cc=mpicc prefix=~/local
make install

在运行ParMETIS之前，我们需要设置两个环境变量：PATH 和 LD_LIBRARY_PATH。

export PATH=$PATH:$HOME/local/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/local/lib/

输入格式

注意

作为先决条件，请阅读章节 guide-distributed-hetero 以了解 DGL 如何为分布式训练组织异构图。

ParMETIS的输入图存储在以下三个文件中： xxx_nodes.txt, xxx_edges.txt 和 xxx_stats.txt，其中 xxx 是图的名称。

xxx_nodes.txt 中的每一行存储一个节点的信息。行 ID 也是节点的同质 ID，例如，第 0 行对应节点 0；第 1 行对应节点 1，依此类推。每行的格式如下：

<node_type_id> <node_weight_list> <type_wise_node_id>

所有字段由空格分隔：

是一个从0开始的整数。每个节点类型都映射到一个整数。对于同质图，其值始终为0。
是整数（由空格分隔），表示 ParMETIS 用于平衡图分区的节点权重。对于同构图，列表中只有一个整数，而对于具有 \(T\) 种节点类型的异构图，列表中应有 \(T\) 个整数。如果节点属于节点类型 \(t\)，则除了第 \(t^{th}\) 个整数外，所有其他整数都为零；第 \(t^{th}\) 个整数是该节点的权重。ParMETIS 将尝试平衡每个分区的总节点权重。对于异构图，它将尝试将相同类型的节点分配到所有分区。推荐的节点权重为 1，用于平衡每个分区中的节点数量，或使用节点度数来平衡每个分区中的边数。
是一个整数，表示其自身类型中的节点ID。

下面展示了一个包含两种节点类型的异构图节点文件的示例。节点类型0有三个节点；节点类型1有四个节点。它使用两个节点权重来确保ParMETIS将为类型0生成大致相同数量的节点分区，并为类型1生成相同数量的节点分区。

同样，xxx_edges.txt中的每一行存储了一条边的信息。行ID也是边的同质ID，例如，第0行对应边0；第1行对应边1，依此类推。每行的格式如下：

<src_node_id> <dst_node_id> <type_wise_edge_id> <edge_type_id>

所有字段由空格分隔：

是源节点的同质 ID。
是目标节点的同质 ID。
是边类型的边ID。
是一个从0开始的整数。每种边类型都映射到一个整数。对于同构图，其值始终为0。

xxx_stats.txt 存储了图的一些基本统计信息。它只有一行，包含三个字段，由空格分隔：

<num_nodes> <num_edges> <total_node_weights>

num_nodes 存储了不考虑节点类型的节点总数。
num_edges 存储了不考虑边类型的边的总数。
total_node_weights 存储节点文件中节点权重的数量。

运行ParMETIS并输出格式

ParMETIS 包含一个名为 pm_dglpart 的命令，该命令从调用 pm_dglpart 的机器加载存储在三个文件中的图，将数据分发到集群中的所有机器，并调用 ParMETIS 对图进行分区。完成后，它会为每个分区生成三个文件：p-xxx_nodes.txt、p-xxx_edges.txt 和 p-xxx_stats.txt。

注意

ParMETIS 在分区过程中重新分配节点的ID。重新分配ID后，分区中的节点被分配连续的ID；此外，相同类型的节点也被分配连续的ID。

p-xxx_nodes.txt 存储分区的节点数据。每行代表一个节点，包含以下字段：

<node_id> <node_type_id> <node_weight_list> <type_wise_node_id>

是ID重新分配后的同质节点ID。
是节点类型 ID。
是 ParMETIS 使用的节点权重（从输入文件中复制）。
<type_wise_node_id> is an integer representing the node ID in its own type.

p-xxx_edges.txt 存储了分区的边数据。每一行代表一条边，包含以下字段：

<src_id> <dst_id> <orig_src_id> <orig_dst_id> <type_wise_edge_id> <edge_type_id>

是源节点在ID重新分配后的同质ID。
是ID重新分配后目标节点的同质ID。
是输入图中源节点的同质 ID。
是输入图中目标节点的同质 ID。
是其自身类型中的边ID。
是边的类型ID。

当调用pm_dglpart时，三个输入文件：xxx_nodes.txt、 xxx_edges.txt、xxx_stats.txt应位于pm_dglpart运行的目录中。以下命令运行四个ParMETIS进程，将名为xxx的图分割成八个分区（每个进程处理两个分区）。

mpirun -np 4 pm_dglpart xxx 2

ParMETIS 的输出文件需要转换为分区分配格式以便运行后续的预处理步骤。