chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,4 @@
|
||||
"""
|
||||
This package contains DistGNN and Libra based graph partitioning tools.
|
||||
"""
|
||||
from . import partition, tools
|
||||
@@ -0,0 +1,4 @@
|
||||
"""
|
||||
This package contains Libra graph partitioner.
|
||||
"""
|
||||
from .libra_partition import partition_graph
|
||||
@@ -0,0 +1,353 @@
|
||||
r"""Libra partition functions.
|
||||
|
||||
Libra partition is a vertex-cut based partitioning algorithm from
|
||||
`Distributed Power-law Graph Computing:
|
||||
Theoretical and Empirical Analysis
|
||||
<https://proceedings.neurips.cc/paper/2014/file/67d16d00201083a2b118dd5128dd6f59-Paper.pdf>`__
|
||||
from Xie et al.
|
||||
"""
|
||||
|
||||
# Copyright (c) 2021 Intel Corporation
|
||||
# \file distgnn/partition/libra_partition.py
|
||||
# \brief Libra - Vertex-cut based graph partitioner for distributed training
|
||||
# \author Vasimuddin Md <vasimuddin.md@intel.com>,
|
||||
# Guixiang Ma <guixiang.ma@intel.com>
|
||||
# Sanchit Misra <sanchit.misra@intel.com>,
|
||||
# Ramanarayan Mohanty <ramanarayan.mohanty@intel.com>,
|
||||
# Sasikanth Avancha <sasikanth.avancha@intel.com>
|
||||
# Nesreen K. Ahmed <nesreen.k.ahmed@intel.com>
|
||||
# \cite Distributed Power-law Graph Computing: Theoretical and Empirical Analysis
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import torch as th
|
||||
|
||||
from dgl import DGLGraph
|
||||
from dgl._sparse_ops import (
|
||||
libra2dgl_build_adjlist,
|
||||
libra2dgl_build_dict,
|
||||
libra2dgl_set_lr,
|
||||
libra_vertex_cut,
|
||||
)
|
||||
from dgl.base import DGLError
|
||||
from dgl.data.utils import save_graphs, save_tensors
|
||||
|
||||
|
||||
def libra_partition(num_community, G, resultdir):
|
||||
"""
|
||||
Performs vertex-cut based graph partitioning and converts the partitioning
|
||||
output to DGL input format.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
num_community : Number of partitions to create
|
||||
G : Input graph to be partitioned
|
||||
resultdir : Output location for storing the partitioned graphs
|
||||
|
||||
Output
|
||||
------
|
||||
1. Creates X partition folder as XCommunities (say, X=2, so, 2Communities)
|
||||
XCommunities contains file name communityZ.txt per partition Z (Z <- 0 .. X-1);
|
||||
each such file contains a list of edges assigned to that partition.
|
||||
These files constitute the output of Libra graph partitioner
|
||||
(An intermediate result of this function).
|
||||
2. The folder also contains partZ folders, each of these folders stores
|
||||
DGL/DistGNN graphs for the Z partitions;
|
||||
these graph files are used as input to DistGNN.
|
||||
3. The folder also contains a json file which contains partitions' information.
|
||||
"""
|
||||
|
||||
num_nodes = G.num_nodes() # number of nodes
|
||||
num_edges = G.num_edges() # number of edges
|
||||
print("Number of nodes in the graph: ", num_nodes)
|
||||
print("Number of edges in the graph: ", num_edges)
|
||||
|
||||
in_d = G.in_degrees()
|
||||
out_d = G.out_degrees()
|
||||
node_degree = in_d + out_d
|
||||
edgenum_unassigned = node_degree.clone()
|
||||
|
||||
u_t, v_t = G.edges()
|
||||
weight_ = th.ones(u_t.shape[0], dtype=th.int64)
|
||||
community_weights = th.zeros(num_community, dtype=th.int64)
|
||||
|
||||
# self_loop = 0
|
||||
# for p, q in zip(u_t, v_t):
|
||||
# if p == q:
|
||||
# self_loop += 1
|
||||
# print("#self loops in the dataset: ", self_loop)
|
||||
|
||||
# del G
|
||||
|
||||
## call to C/C++ code
|
||||
out = th.zeros(u_t.shape[0], dtype=th.int32)
|
||||
libra_vertex_cut(
|
||||
num_community,
|
||||
node_degree,
|
||||
edgenum_unassigned,
|
||||
community_weights,
|
||||
u_t,
|
||||
v_t,
|
||||
weight_,
|
||||
out,
|
||||
num_nodes,
|
||||
num_edges,
|
||||
resultdir,
|
||||
)
|
||||
|
||||
print("Max partition size: ", int(community_weights.max()))
|
||||
print(" ** Converting libra partitions to dgl graphs **")
|
||||
fsize = int(community_weights.max()) + 1024 ## max edges in partition
|
||||
# print("fsize: ", fsize, flush=True)
|
||||
|
||||
node_map = th.zeros(num_community, dtype=th.int64)
|
||||
indices = th.zeros(num_nodes, dtype=th.int64)
|
||||
lrtensor = th.zeros(num_nodes, dtype=th.int64)
|
||||
gdt_key = th.zeros(num_nodes, dtype=th.int64)
|
||||
gdt_value = th.zeros([num_nodes, num_community], dtype=th.int64)
|
||||
offset = th.zeros(1, dtype=th.int64)
|
||||
ldt_ar = []
|
||||
|
||||
gg_ar = [DGLGraph() for i in range(num_community)]
|
||||
part_nodes = []
|
||||
|
||||
print(">>> ", "num_nodes ", " ", "num_edges")
|
||||
## Iterator over number of partitions
|
||||
for i in range(num_community):
|
||||
g = gg_ar[i]
|
||||
|
||||
a_t = th.zeros(fsize, dtype=th.int64)
|
||||
b_t = th.zeros(fsize, dtype=th.int64)
|
||||
ldt_key = th.zeros(fsize, dtype=th.int64)
|
||||
ldt_ar.append(ldt_key)
|
||||
|
||||
## building node, parition dictionary
|
||||
## Assign local node ids and mapping to global node ids
|
||||
ret = libra2dgl_build_dict(
|
||||
a_t,
|
||||
b_t,
|
||||
indices,
|
||||
ldt_key,
|
||||
gdt_key,
|
||||
gdt_value,
|
||||
node_map,
|
||||
offset,
|
||||
num_community,
|
||||
i,
|
||||
fsize,
|
||||
resultdir,
|
||||
)
|
||||
|
||||
num_nodes_partition = int(ret[0])
|
||||
num_edges_partition = int(ret[1])
|
||||
part_nodes.append(num_nodes_partition)
|
||||
print(">>> ", num_nodes_partition, " ", num_edges_partition)
|
||||
g.add_edges(a_t[0:num_edges_partition], b_t[0:num_edges_partition])
|
||||
|
||||
########################################################
|
||||
## fixing lr - 1-level tree for the split-nodes
|
||||
libra2dgl_set_lr(gdt_key, gdt_value, lrtensor, num_community, num_nodes)
|
||||
########################################################
|
||||
# graph_name = dataset
|
||||
graph_name = resultdir.split("_")[-1].split("/")[0]
|
||||
part_method = "Libra"
|
||||
num_parts = num_community ## number of paritions/communities
|
||||
num_hops = 0
|
||||
node_map_val = node_map.tolist()
|
||||
edge_map_val = 0
|
||||
out_path = resultdir
|
||||
|
||||
part_metadata = {
|
||||
"graph_name": graph_name,
|
||||
"num_nodes": G.num_nodes(),
|
||||
"num_edges": G.num_edges(),
|
||||
"part_method": part_method,
|
||||
"num_parts": num_parts,
|
||||
"halo_hops": num_hops,
|
||||
"node_map": node_map_val,
|
||||
"edge_map": edge_map_val,
|
||||
}
|
||||
############################################################
|
||||
|
||||
for i in range(num_community):
|
||||
g = gg_ar[0]
|
||||
num_nodes_partition = part_nodes[i]
|
||||
adj = th.zeros([num_nodes_partition, num_community - 1], dtype=th.int64)
|
||||
inner_node = th.zeros(num_nodes_partition, dtype=th.int32)
|
||||
lr_t = th.zeros(num_nodes_partition, dtype=th.int64)
|
||||
ldt = ldt_ar[0]
|
||||
|
||||
try:
|
||||
feat = G.ndata["feat"]
|
||||
except KeyError:
|
||||
feat = G.ndata["features"]
|
||||
|
||||
try:
|
||||
labels = G.ndata["label"]
|
||||
except KeyError:
|
||||
labels = G.ndata["labels"]
|
||||
|
||||
trainm = G.ndata["train_mask"].int()
|
||||
testm = G.ndata["test_mask"].int()
|
||||
valm = G.ndata["val_mask"].int()
|
||||
|
||||
feat_size = feat.shape[1]
|
||||
gfeat = th.zeros([num_nodes_partition, feat_size], dtype=feat.dtype)
|
||||
|
||||
glabels = th.zeros(num_nodes_partition, dtype=labels.dtype)
|
||||
gtrainm = th.zeros(num_nodes_partition, dtype=trainm.dtype)
|
||||
gtestm = th.zeros(num_nodes_partition, dtype=testm.dtype)
|
||||
gvalm = th.zeros(num_nodes_partition, dtype=valm.dtype)
|
||||
|
||||
## build remote node databse per local node
|
||||
## gather feats, train, test, val, and labels for each partition
|
||||
libra2dgl_build_adjlist(
|
||||
feat,
|
||||
gfeat,
|
||||
adj,
|
||||
inner_node,
|
||||
ldt,
|
||||
gdt_key,
|
||||
gdt_value,
|
||||
node_map,
|
||||
lr_t,
|
||||
lrtensor,
|
||||
num_nodes_partition,
|
||||
num_community,
|
||||
i,
|
||||
feat_size,
|
||||
labels,
|
||||
trainm,
|
||||
testm,
|
||||
valm,
|
||||
glabels,
|
||||
gtrainm,
|
||||
gtestm,
|
||||
gvalm,
|
||||
feat.shape[0],
|
||||
)
|
||||
|
||||
g.ndata["adj"] = adj ## database of remote clones
|
||||
g.ndata["inner_node"] = inner_node ## split node '0' else '1'
|
||||
g.ndata["feat"] = gfeat ## gathered features
|
||||
g.ndata["lf"] = lr_t ## 1-level tree among split nodes
|
||||
|
||||
g.ndata["label"] = glabels
|
||||
g.ndata["train_mask"] = gtrainm
|
||||
g.ndata["test_mask"] = gtestm
|
||||
g.ndata["val_mask"] = gvalm
|
||||
|
||||
# Validation code, run only small graphs
|
||||
# for l in range(num_nodes_partition):
|
||||
# index = int(ldt[l])
|
||||
# assert glabels[l] == labels[index]
|
||||
# assert gtrainm[l] == trainm[index]
|
||||
# assert gtestm[l] == testm[index]
|
||||
# for j in range(feat_size):
|
||||
# assert gfeat[l][j] == feat[index][j]
|
||||
|
||||
print("Writing partition {} to file".format(i), flush=True)
|
||||
|
||||
part = g
|
||||
part_id = i
|
||||
part_dir = os.path.join(out_path, "part" + str(part_id))
|
||||
node_feat_file = os.path.join(part_dir, "node_feat.dgl")
|
||||
edge_feat_file = os.path.join(part_dir, "edge_feat.dgl")
|
||||
part_graph_file = os.path.join(part_dir, "graph.dgl")
|
||||
part_metadata["part-{}".format(part_id)] = {
|
||||
"node_feats": node_feat_file,
|
||||
"edge_feats": edge_feat_file,
|
||||
"part_graph": part_graph_file,
|
||||
}
|
||||
os.makedirs(part_dir, mode=0o775, exist_ok=True)
|
||||
save_tensors(node_feat_file, part.ndata)
|
||||
save_graphs(part_graph_file, [part])
|
||||
|
||||
del g
|
||||
del gg_ar[0]
|
||||
del ldt
|
||||
del ldt_ar[0]
|
||||
|
||||
with open("{}/{}.json".format(out_path, graph_name), "w") as outfile:
|
||||
json.dump(part_metadata, outfile, sort_keys=True, indent=4)
|
||||
|
||||
print("Conversion libra2dgl completed !!!")
|
||||
|
||||
|
||||
def partition_graph(num_community, G, resultdir):
|
||||
"""
|
||||
Performs vertex-cut based graph partitioning and converts the partitioning
|
||||
output to DGL input format.
|
||||
|
||||
Given a graph, this function will create a folder named ``XCommunities`` where ``X``
|
||||
stands for the number of communities. It will contain ``X`` files named
|
||||
``communityZ.txt`` for each partition Z (from 0 to X-1);
|
||||
each such file contains a list of edges assigned to that partition.
|
||||
These files constitute the output of Libra graph partitioner.
|
||||
|
||||
The folder also contains X subfolders named ``partZ``, each of these folders stores
|
||||
DGL/DistGNN graphs for partition Z; these graph files are used as input to
|
||||
DistGNN.
|
||||
|
||||
The folder also contains a json file which contains partitions' information.
|
||||
|
||||
Currently we require the graph's node data to contain the following columns:
|
||||
|
||||
* ``features`` for node features.
|
||||
* ``label`` for node labels.
|
||||
* ``train_mask`` as a boolean mask of training node set.
|
||||
* ``val_mask`` as a boolean mask of validation node set.
|
||||
* ``test_mask`` as a boolean mask of test node set.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
num_community : int
|
||||
Number of partitions to create.
|
||||
G : DGLGraph
|
||||
Input graph to be partitioned.
|
||||
resultdir : str
|
||||
Output location for storing the partitioned graphs.
|
||||
"""
|
||||
|
||||
print("num partitions: ", num_community)
|
||||
print("output location: ", resultdir)
|
||||
|
||||
## create ouptut directory
|
||||
try:
|
||||
os.makedirs(resultdir, mode=0o775, exist_ok=True)
|
||||
except:
|
||||
raise DGLError("Error: Could not create directory: ", resultdir)
|
||||
|
||||
tic = time.time()
|
||||
print(
|
||||
"####################################################################"
|
||||
)
|
||||
print("Executing parititons: ", num_community)
|
||||
ltic = time.time()
|
||||
try:
|
||||
resultdir = os.path.join(resultdir, str(num_community) + "Communities")
|
||||
os.makedirs(resultdir, mode=0o775, exist_ok=True)
|
||||
except:
|
||||
raise DGLError("Error: Could not create sub-directory: ", resultdir)
|
||||
|
||||
## Libra partitioning
|
||||
libra_partition(num_community, G, resultdir)
|
||||
|
||||
ltoc = time.time()
|
||||
print(
|
||||
"Time taken by {} partitions {:0.4f} sec".format(
|
||||
num_community, ltoc - ltic
|
||||
)
|
||||
)
|
||||
print()
|
||||
|
||||
toc = time.time()
|
||||
print(
|
||||
"Generated ",
|
||||
num_community,
|
||||
" partitions in {:0.4f} sec".format(toc - tic),
|
||||
flush=True,
|
||||
)
|
||||
print("Partitioning completed successfully !!!")
|
||||
@@ -0,0 +1,4 @@
|
||||
"""
|
||||
This package contains extra routines related to Libra graph partitioner.
|
||||
"""
|
||||
from .tools import load_proteins
|
||||
@@ -0,0 +1,162 @@
|
||||
r"""
|
||||
Copyright (c) 2021 Intel Corporation
|
||||
\file distgnn/tools/tools.py
|
||||
\brief Tools for use in Libra graph partitioner.
|
||||
\author Vasimuddin Md <vasimuddin.md@intel.com>
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
|
||||
import requests
|
||||
import torch as th
|
||||
from scipy.io import mmread
|
||||
|
||||
import dgl
|
||||
from dgl.base import DGLError
|
||||
from dgl.data.utils import load_graphs, save_graphs, save_tensors
|
||||
|
||||
|
||||
def rep_per_node(prefix, num_community):
|
||||
"""
|
||||
Used on Libra partitioned data.
|
||||
This function reports number of split-copes per node (replication) of
|
||||
a partitioned graph
|
||||
Parameters
|
||||
----------
|
||||
prefix: Partition folder location (contains replicationlist.csv)
|
||||
num_community: number of partitions or communities
|
||||
"""
|
||||
ifile = os.path.join(prefix, "replicationlist.csv")
|
||||
fhandle = open(ifile, "r")
|
||||
r_dt = {}
|
||||
|
||||
fline = fhandle.readline() ## reading first line, contains the comment.
|
||||
print(fline)
|
||||
for line in fhandle:
|
||||
if line[0] == "#":
|
||||
raise DGLError("[Bug] Read Hash char in rep_per_node func.")
|
||||
|
||||
node = line.strip("\n")
|
||||
if r_dt.get(node, -100) == -100:
|
||||
r_dt[node] = 1
|
||||
else:
|
||||
r_dt[node] += 1
|
||||
|
||||
fhandle.close()
|
||||
## sanity checks
|
||||
for v in r_dt.values():
|
||||
if v >= num_community:
|
||||
raise DGLError(
|
||||
"[Bug] Unexpected event in rep_per_node() in tools.py."
|
||||
)
|
||||
|
||||
return r_dt
|
||||
|
||||
|
||||
def download_proteins():
|
||||
"""
|
||||
Downloads the proteins dataset
|
||||
"""
|
||||
print("Downloading dataset...")
|
||||
print("This might a take while..")
|
||||
url = "https://portal.nersc.gov/project/m1982/GNN/"
|
||||
file_name = "subgraph3_iso_vs_iso_30_70length_ALL.m100.propermm.mtx"
|
||||
url = url + file_name
|
||||
try:
|
||||
req = requests.get(url)
|
||||
except:
|
||||
raise DGLError(
|
||||
"Error: Failed to download Proteins dataset!! Aborting.."
|
||||
)
|
||||
|
||||
with open("proteins.mtx", "wb") as handle:
|
||||
handle.write(req.content)
|
||||
|
||||
|
||||
def proteins_mtx2dgl():
|
||||
"""
|
||||
This function converts Proteins dataset from mtx to dgl format.
|
||||
"""
|
||||
print("Converting mtx2dgl..")
|
||||
print("This might a take while..")
|
||||
a_mtx = mmread("proteins.mtx")
|
||||
coo = a_mtx.tocoo()
|
||||
u = th.tensor(coo.row, dtype=th.int64)
|
||||
v = th.tensor(coo.col, dtype=th.int64)
|
||||
g = dgl.DGLGraph()
|
||||
|
||||
g.add_edges(u, v)
|
||||
|
||||
n = g.num_nodes()
|
||||
feat_size = 128 ## arbitrary number
|
||||
feats = th.empty([n, feat_size], dtype=th.float32)
|
||||
|
||||
## arbitrary numbers
|
||||
train_size = 1000000
|
||||
test_size = 500000
|
||||
val_size = 5000
|
||||
nlabels = 256
|
||||
|
||||
train_mask = th.zeros(n, dtype=th.bool)
|
||||
test_mask = th.zeros(n, dtype=th.bool)
|
||||
val_mask = th.zeros(n, dtype=th.bool)
|
||||
label = th.zeros(n, dtype=th.int64)
|
||||
|
||||
for i in range(train_size):
|
||||
train_mask[i] = True
|
||||
|
||||
for i in range(test_size):
|
||||
test_mask[train_size + i] = True
|
||||
|
||||
for i in range(val_size):
|
||||
val_mask[train_size + test_size + i] = True
|
||||
|
||||
for i in range(n):
|
||||
label[i] = random.choice(range(nlabels))
|
||||
|
||||
g.ndata["feat"] = feats
|
||||
g.ndata["train_mask"] = train_mask
|
||||
g.ndata["test_mask"] = test_mask
|
||||
g.ndata["val_mask"] = val_mask
|
||||
g.ndata["label"] = label
|
||||
|
||||
return g
|
||||
|
||||
|
||||
def save(g, dataset):
|
||||
"""
|
||||
This function saves input dataset to dgl format
|
||||
Parameters
|
||||
----------
|
||||
g : graph to be saved
|
||||
dataset : output folder name
|
||||
"""
|
||||
print("Saving dataset..")
|
||||
part_dir = os.path.join("./" + dataset)
|
||||
node_feat_file = os.path.join(part_dir, "node_feat.dgl")
|
||||
part_graph_file = os.path.join(part_dir, "graph.dgl")
|
||||
os.makedirs(part_dir, mode=0o775, exist_ok=True)
|
||||
save_tensors(node_feat_file, g.ndata)
|
||||
save_graphs(part_graph_file, [g])
|
||||
print("Graph saved successfully !!")
|
||||
|
||||
|
||||
def load_proteins(dataset):
|
||||
"""
|
||||
This function downloads, converts, and load Proteins graph dataset
|
||||
Parameter
|
||||
---------
|
||||
dataset: output folder name
|
||||
"""
|
||||
part_dir = dataset
|
||||
graph_file = os.path.join(part_dir + "/graph.dgl")
|
||||
|
||||
if not os.path.exists("proteins.mtx"):
|
||||
download_proteins()
|
||||
if not os.path.exists(graph_file):
|
||||
g = proteins_mtx2dgl()
|
||||
save(g, dataset)
|
||||
## load
|
||||
graph = load_graphs(graph_file)[0][0]
|
||||
return graph
|
||||
Reference in New Issue
Block a user