wehub-resource-sync
247 lines
7.8 KiB
Python
247 lines
7.8 KiB
Python
import argparse
|
|
import logging
|
|
import os
|
|
import platform
|
|
|
|
import constants
|
|
|
|
import dgl
|
|
import numpy as np
|
|
|
|
import pyarrow
|
|
import pyarrow.parquet as pq
|
|
import torch as th
|
|
from dgl.data.utils import load_graphs, load_tensors
|
|
|
|
from dgl.distributed.partition import (
|
|
_etype_str_to_tuple,
|
|
_etype_tuple_to_str,
|
|
_get_inner_edge_mask,
|
|
_get_inner_node_mask,
|
|
load_partition,
|
|
RESERVED_FIELD_DTYPE,
|
|
)
|
|
from utils import get_idranges, read_json
|
|
from verification_utils import (
|
|
get_node_partids,
|
|
read_file,
|
|
read_orig_ids,
|
|
verify_graph_feats,
|
|
verify_metadata_counts,
|
|
verify_node_partitionids,
|
|
verify_partition_data_types,
|
|
verify_partition_formats,
|
|
)
|
|
|
|
|
|
def _read_graph(schema):
|
|
"""Read a DGL Graph object from storage using metadata schema, which is
|
|
a json object describing the DGL graph on disk.
|
|
|
|
Parameters:
|
|
-----------
|
|
schema : json object
|
|
json object describing the input graph to read from the disk
|
|
|
|
Returns:
|
|
--------
|
|
DGL Graph Object :
|
|
DGL Graph object is created which is read from the disk storage.
|
|
"""
|
|
edges = {}
|
|
edge_types = schema[constants.STR_EDGE_TYPE]
|
|
for etype in edge_types:
|
|
efiles = schema[constants.STR_EDGES][etype][constants.STR_DATA]
|
|
src = []
|
|
dst = []
|
|
for fname in efiles:
|
|
if (
|
|
schema[constants.STR_EDGES][etype][constants.STR_FORMAT][
|
|
constants.STR_NAME
|
|
]
|
|
== constants.STR_CSV
|
|
):
|
|
data = read_file(fname, constants.STR_CSV)
|
|
elif (
|
|
schema[constants.STR_EDGES][etype][constants.STR_FORMAT][
|
|
constants.STR_NAME
|
|
]
|
|
== constants.STR_PARQUET
|
|
):
|
|
data = read_file(fname)
|
|
else:
|
|
raise ValueError(
|
|
f"Unknown edge format for {etype} - {schema[constants.STR_EDGES][etype][constants.STR_FORMAT]}"
|
|
)
|
|
|
|
src.append(data[:, 0])
|
|
dst.append(data[:, 1])
|
|
src = np.concatenate(src)
|
|
dst = np.concatenate(dst)
|
|
edges[_etype_str_to_tuple(etype)] = (src, dst)
|
|
|
|
g = dgl.heterograph(edges)
|
|
# g = dgl.to_homogeneous(g)
|
|
|
|
g.ndata["orig_id"] = g.ndata[dgl.NID]
|
|
g.edata["orig_id"] = g.edata[dgl.EID]
|
|
|
|
# read features here.
|
|
for ntype in schema[constants.STR_NODE_TYPE]:
|
|
if ntype in schema[constants.STR_NODE_DATA]:
|
|
for featname, featdata in schema[constants.STR_NODE_DATA][
|
|
ntype
|
|
].items():
|
|
files = fdata[constants.STR_DATA]
|
|
feats = []
|
|
for fname in files:
|
|
feats.append(read_file(fname, constants.STR_NUMPY))
|
|
if len(feats) > 0:
|
|
g.nodes[ntype].data[featname] = th.from_numpy(
|
|
np.concatenate(feats)
|
|
)
|
|
|
|
# read edge features here.
|
|
for etype in schema[constants.STR_EDGE_TYPE]:
|
|
if etype in schema[constants.STR_EDGE_DATA]:
|
|
for featname, fdata in schema[constants.STR_EDGE_DATA][etype]:
|
|
files = fdata[constants.STR_DATA]
|
|
feats = []
|
|
for fname in files:
|
|
feats.append(read_file(fname))
|
|
if len(feats) > 0:
|
|
g.edges[etype].data[featname] = th.from_numpy(
|
|
np.concatenate(feats)
|
|
)
|
|
|
|
# print from graph
|
|
logging.info(f"|V|= {g.num_nodes()}")
|
|
logging.info(f"|E|= {g.num_edges()}")
|
|
for ntype in g.ntypes:
|
|
for name, data in g.nodes[ntype].data.items():
|
|
if isinstance(data, th.Tensor):
|
|
logging.info(
|
|
f"Input Graph: nfeat - {ntype}/{name} - data - {data.size()}"
|
|
)
|
|
|
|
for c_etype in g.canonical_etypes:
|
|
for name, data in g.edges[c_etype].data.items():
|
|
if isinstance(data, th.Tensor):
|
|
logging.info(
|
|
f"Input Graph: efeat - {etype}/{name} - data - {g.edges[etype].data[name].size()}"
|
|
)
|
|
|
|
return g
|
|
|
|
|
|
def _read_part_graphs(part_config, part_metafile):
|
|
"""Read partitioned graph objects from disk storage.
|
|
|
|
Parameters:
|
|
----------
|
|
part_config : json object
|
|
json object created using the metadata file for the partitioned graph.
|
|
part_metafile : string
|
|
absolute path of the metadata.json file for the partitioned graph.
|
|
|
|
Returns:
|
|
--------
|
|
list of tuples :
|
|
where each tuple contains 4 objects in the following order:
|
|
partitioned graph object
|
|
global partition book
|
|
node features
|
|
edge features
|
|
"""
|
|
part_graph_data = []
|
|
for i in range(part_config["num_parts"]):
|
|
part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
|
|
part_metafile, i
|
|
)
|
|
part_graph_data.append((part_g, node_feats, edge_feats, gpb))
|
|
return part_graph_data
|
|
|
|
|
|
def _validate_results(params):
|
|
"""Main function to verify the graph partitions
|
|
|
|
Parameters:
|
|
-----------
|
|
params : argparser object
|
|
to access the command line arguments
|
|
"""
|
|
logging.info(f"loading config files...")
|
|
part_config = os.path.join(params.part_graph_dir, "metadata.json")
|
|
part_schema = read_json(part_config)
|
|
num_parts = part_schema["num_parts"]
|
|
|
|
logging.info(f"loading config files of the original dataset...")
|
|
graph_config = os.path.join(params.orig_dataset_dir, "metadata.json")
|
|
graph_schema = read_json(graph_config)
|
|
|
|
logging.info(f"loading original ids from the dgl files...")
|
|
orig_nids = read_orig_ids(params.part_graph_dir, "orig_nids.dgl", num_parts)
|
|
orig_eids = read_orig_ids(params.part_graph_dir, "orig_eids.dgl", num_parts)
|
|
|
|
logging.info(f"loading node to partition-ids from files... ")
|
|
node_partids = get_node_partids(params.partitions_dir, graph_schema)
|
|
|
|
logging.info(f"loading the original dataset...")
|
|
g = _read_graph(graph_schema)
|
|
|
|
logging.info(f"Beginning the verification process...")
|
|
for i in range(num_parts):
|
|
part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
|
|
part_config, i
|
|
)
|
|
|
|
verify_partition_data_types(part_g)
|
|
verify_partition_formats(part_g, None)
|
|
verify_graph_feats(
|
|
g, gpb, part_g, node_feats, edge_feats, orig_nids, orig_eids
|
|
)
|
|
verify_metadata_counts(part_schema, part_g, graph_schema, g, i)
|
|
verify_node_partitionids(
|
|
node_partids, part_g, g, gpb, graph_schema, orig_nids, i
|
|
)
|
|
logging.info(f"Verification of partitioned graph - {i}... SUCCESS !!!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Construct graph partitions")
|
|
parser.add_argument(
|
|
"--orig-dataset-dir",
|
|
required=True,
|
|
type=str,
|
|
help="The directory path that contains the original graph input files.",
|
|
)
|
|
parser.add_argument(
|
|
"--part-graph-dir",
|
|
required=True,
|
|
type=str,
|
|
help="The directory path that contains the partitioned graph files.",
|
|
)
|
|
parser.add_argument(
|
|
"--partitions-dir",
|
|
required=True,
|
|
type=str,
|
|
help="The directory path that contains metis/random partitions results.",
|
|
)
|
|
parser.add_argument(
|
|
"--log-level",
|
|
type=str,
|
|
default="info",
|
|
help="To enable log level for debugging purposes. Available options: \
|
|
(Critical, Error, Warning, Info, Debug, Notset), default value \
|
|
is: Info",
|
|
)
|
|
params = parser.parse_args()
|
|
|
|
numeric_level = getattr(logging, params.log_level.upper(), None)
|
|
logging.basicConfig(
|
|
level=numeric_level,
|
|
format=f"[{platform.node()} %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
|
|
)
|
|
|
|
_validate_results(params)
|