mirrors/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/tools/pytest_utils.py
+593
View File
@@ -0,0 +1,593 @@
import json
import logging
import os
import dgl
import numpy as np
import torch
from distpartitioning import array_readwriter
from distpartitioning.array_readwriter.parquet import ParquetArrayParser
from files import setdir
def _chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt, vector_rows=False):
paths = []
offset = 0
for j, n in enumerate(chunk_sizes):
path = os.path.abspath(path_fmt % j)
arr_chunk = arr[offset : offset + n]
shape = arr_chunk.shape
logging.info("Chunking %d-%d" % (offset, offset + n))
# If requested we write multi-column arrays as single-column vector Parquet files
array_parser = array_readwriter.get_array_parser(**fmt_meta)
if (
isinstance(array_parser, ParquetArrayParser)
and len(shape) > 1
and shape[1] > 1
):
array_parser.write(path, arr_chunk, vector_rows=vector_rows)
else:
array_parser.write(path, arr_chunk)
offset += n
paths.append(path)
return paths
def _initialize_num_chunks(g, num_chunks, kwargs=None):
"""Initialize num_chunks for each node/edge.
Parameters
----------
g: DGLGraph
Graph to be chunked.
num_chunks: int
Default number of chunks to be applied onto node/edge data.
kwargs: dict
Key word arguments to specify details for each node/edge data.
Returns
-------
num_chunks_data: dict
Detailed number of chunks for each node/edge.
"""
def _init(g, num_chunks, key, kwargs=None):
chunks_data = kwargs.get(key, None)
is_node = "_node" in key
data_types = g.ntypes if is_node else g.canonical_etypes
if isinstance(chunks_data, int):
chunks_data = {data_type: chunks_data for data_type in data_types}
elif isinstance(chunks_data, dict):
for data_type in data_types:
if data_type not in chunks_data:
chunks_data[data_type] = num_chunks
else:
chunks_data = {data_type: num_chunks for data_type in data_types}
for _, data in chunks_data.items():
if isinstance(data, dict):
n_chunks = list(data.values())
else:
n_chunks = [data]
assert all(
isinstance(v, int) for v in n_chunks
), "num_chunks for each data type should be int."
return chunks_data
num_chunks_data = {}
for key in [
"num_chunks_nodes",
"num_chunks_edges",
"num_chunks_node_data",
"num_chunks_edge_data",
]:
num_chunks_data[key] = _init(g, num_chunks, key, kwargs=kwargs)
return num_chunks_data
def _chunk_graph(
g,
name,
ndata_paths,
edata_paths,
num_chunks,
data_fmt,
edges_format,
vector_rows=False,
**kwargs,
):
# First deal with ndata and edata that are homogeneous
# (i.e. not a dict-of-dict)
if len(g.ntypes) == 1 and not isinstance(
next(iter(ndata_paths.values())), dict
):
ndata_paths = {g.ntypes[0]: ndata_paths}
if len(g.etypes) == 1 and not isinstance(
next(iter(edata_paths.values())), dict
):
edata_paths = {g.etypes[0]: ndata_paths}
# Then convert all edge types to canonical edge types
etypestrs = {etype: ":".join(etype) for etype in g.canonical_etypes}
edata_paths = {
":".join(g.to_canonical_etype(k)): v for k, v in edata_paths.items()
}
metadata = {}
metadata["graph_name"] = name
metadata["node_type"] = g.ntypes
# add node_type_counts
metadata["num_nodes_per_type"] = [g.num_nodes(ntype) for ntype in g.ntypes]
# Initialize num_chunks for each node/edge.
num_chunks_details = _initialize_num_chunks(g, num_chunks, kwargs=kwargs)
# Compute the number of nodes per chunk per node type
metadata["num_nodes_per_chunk"] = num_nodes_per_chunk = []
num_chunks_nodes = num_chunks_details["num_chunks_nodes"]
for ntype in g.ntypes:
num_nodes = g.num_nodes(ntype)
num_nodes_list = []
n_chunks = num_chunks_nodes[ntype]
for i in range(n_chunks):
n = num_nodes // n_chunks + (i < num_nodes % n_chunks)
num_nodes_list.append(n)
num_nodes_per_chunk.append(num_nodes_list)
metadata["edge_type"] = [etypestrs[etype] for etype in g.canonical_etypes]
metadata["num_edges_per_type"] = [
g.num_edges(etype) for etype in g.canonical_etypes
]
# Compute the number of edges per chunk per edge type
metadata["num_edges_per_chunk"] = num_edges_per_chunk = []
num_chunks_edges = num_chunks_details["num_chunks_edges"]
for etype in g.canonical_etypes:
num_edges = g.num_edges(etype)
num_edges_list = []
n_chunks = num_chunks_edges[etype]
for i in range(n_chunks):
n = num_edges // n_chunks + (i < num_edges % n_chunks)
num_edges_list.append(n)
num_edges_per_chunk.append(num_edges_list)
num_edges_per_chunk_dict = {
k: v for k, v in zip(g.canonical_etypes, num_edges_per_chunk)
}
idxes_etypestr = {
idx: (etype, etypestrs[etype])
for idx, etype in enumerate(g.canonical_etypes)
}
idxes = np.arange(len(idxes_etypestr))
# Split edge index
metadata["edges"] = {}
with setdir("edge_index"):
np.random.shuffle(idxes)
for idx in idxes:
etype = idxes_etypestr[idx][0]
etypestr = idxes_etypestr[idx][1]
logging.info("Chunking edge index for %s" % etypestr)
edges_meta = {}
if edges_format == "csv":
fmt_meta = {"name": edges_format, "delimiter": " "}
elif edges_format == "parquet":
fmt_meta = {"name": edges_format}
else:
raise RuntimeError(f"Invalid edges_fmt: {edges_format}")
edges_meta["format"] = fmt_meta
srcdst = torch.stack(g.edges(etype=etype), 1)
edges_meta["data"] = _chunk_numpy_array(
srcdst.numpy(),
fmt_meta,
num_edges_per_chunk_dict[etype],
etypestr + "%d.txt",
)
metadata["edges"][etypestr] = edges_meta
# Chunk node data
reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {"name": data_fmt}
file_suffix = "npy" if data_fmt == "numpy" else "parquet"
metadata["node_data"] = {}
num_chunks_node_data = num_chunks_details["num_chunks_node_data"]
with setdir("node_data"):
for ntype, ndata_per_type in ndata_paths.items():
ndata_meta = {}
with setdir(ntype):
for key, path in ndata_per_type.items():
logging.info(
"Chunking node data for type %s key %s" % (ntype, key)
)
chunk_sizes = []
num_nodes = g.num_nodes(ntype)
n_chunks = num_chunks_node_data[ntype]
if isinstance(n_chunks, dict):
n_chunks = n_chunks.get(key, num_chunks)
assert isinstance(n_chunks, int), (
f"num_chunks for {ntype}/{key} should be int while "
f"{type(n_chunks)} is got."
)
for i in range(n_chunks):
n = num_nodes // n_chunks + (i < num_nodes % n_chunks)
chunk_sizes.append(n)
ndata_key_meta = {}
arr = array_readwriter.get_array_parser(
**reader_fmt_meta
).read(path)
ndata_key_meta["format"] = writer_fmt_meta
ndata_key_meta["data"] = _chunk_numpy_array(
arr,
writer_fmt_meta,
chunk_sizes,
key + "-%d." + file_suffix,
vector_rows=vector_rows,
)
ndata_meta[key] = ndata_key_meta
metadata["node_data"][ntype] = ndata_meta
# Chunk edge data
metadata["edge_data"] = {}
num_chunks_edge_data = num_chunks_details["num_chunks_edge_data"]
with setdir("edge_data"):
for etypestr, edata_per_type in edata_paths.items():
edata_meta = {}
etype = tuple(etypestr.split(":"))
with setdir(etypestr):
for key, path in edata_per_type.items():
logging.info(
"Chunking edge data for type %s key %s"
% (etypestr, key)
)
chunk_sizes = []
num_edges = g.num_edges(etype)
n_chunks = num_chunks_edge_data[etype]
if isinstance(n_chunks, dict):
n_chunks = n_chunks.get(key, num_chunks)
assert isinstance(n_chunks, int), (
f"num_chunks for {etype}/{key} should be int while "
f"{type(n_chunks)} is got."
)
for i in range(n_chunks):
n = num_edges // n_chunks + (i < num_edges % n_chunks)
chunk_sizes.append(n)
edata_key_meta = {}
arr = array_readwriter.get_array_parser(
**reader_fmt_meta
).read(path)
edata_key_meta["format"] = writer_fmt_meta
edata_key_meta["data"] = _chunk_numpy_array(
arr,
writer_fmt_meta,
chunk_sizes,
key + "-%d." + file_suffix,
vector_rows=vector_rows,
)
edata_meta[key] = edata_key_meta
metadata["edge_data"][etypestr] = edata_meta
metadata_path = "metadata.json"
with open(metadata_path, "w") as f:
json.dump(metadata, f, sort_keys=True, indent=4)
logging.info("Saved metadata in %s" % os.path.abspath(metadata_path))
def chunk_graph(
g,
name,
ndata_paths,
edata_paths,
num_chunks,
output_path,
data_fmt="numpy",
edges_fmt="csv",
vector_rows=False,
**kwargs,
):
"""
Split the graph into multiple chunks.
A directory will be created at :attr:`output_path` with the metadata and
chunked edge list as well as the node/edge data.
Parameters
----------
g : DGLGraph
The graph.
name : str
The name of the graph, to be used later in DistDGL training.
ndata_paths : dict[str, pathlike] or dict[ntype, dict[str, pathlike]]
The dictionary of paths pointing to the corresponding numpy array file
for each node data key.
edata_paths : dict[etype, pathlike] or dict[etype, dict[str, pathlike]]
The dictionary of paths pointing to the corresponding numpy array file
for each edge data key. ``etype`` could be canonical or non-canonical.
num_chunks : int
The number of chunks
output_path : pathlike
The output directory saving the chunked graph.
data_fmt : str
Format of node/edge data: 'numpy' or 'parquet'.
edges_fmt : str
Format of edges files: 'csv' or 'parquet'.
vector_rows : str
When true will write parquet files as single-column vector row files.
kwargs : dict
Key word arguments to control chunk details.
"""
for ntype, ndata in ndata_paths.items():
for key in ndata.keys():
ndata[key] = os.path.abspath(ndata[key])
for etype, edata in edata_paths.items():
for key in edata.keys():
edata[key] = os.path.abspath(edata[key])
with setdir(output_path):
_chunk_graph(
g,
name,
ndata_paths,
edata_paths,
num_chunks,
data_fmt,
edges_fmt,
vector_rows,
**kwargs,
)
def create_chunked_dataset(
root_dir,
num_chunks,
data_fmt="numpy",
edges_fmt="csv",
vector_rows=False,
**kwargs,
):
"""
This function creates a sample dataset, based on MAG240 dataset.
Parameters:
-----------
root_dir : string
directory in which all the files for the chunked dataset will be stored.
"""
# Step0: prepare chunked graph data format.
# A synthetic mini MAG240.
num_institutions = 1200
num_authors = 1200
num_papers = 1200
def rand_edges(num_src, num_dst, num_edges):
eids = np.random.choice(num_src * num_dst, num_edges, replace=False)
src = torch.from_numpy(eids // num_dst)
dst = torch.from_numpy(eids % num_dst)
return src, dst
num_cite_edges = 24 * 1000
num_write_edges = 12 * 1000
num_affiliate_edges = 2400
# Structure.
data_dict = {
("paper", "cites", "paper"): rand_edges(
num_papers, num_papers, num_cite_edges
),
("author", "writes", "paper"): rand_edges(
num_authors, num_papers, num_write_edges
),
("author", "affiliated_with", "institution"): rand_edges(
num_authors, num_institutions, num_affiliate_edges
),
("institution", "writes", "paper"): rand_edges(
num_institutions, num_papers, num_write_edges
),
}
src, dst = data_dict[("author", "writes", "paper")]
data_dict[("paper", "rev_writes", "author")] = (dst, src)
g = dgl.heterograph(data_dict)
# paper feat, label, year
num_paper_feats = 3
paper_feat = np.random.randn(num_papers, num_paper_feats)
num_classes = 4
paper_label = np.random.choice(num_classes, num_papers)
paper_year = np.random.choice(2022, num_papers)
paper_orig_ids = np.arange(0, num_papers)
writes_orig_ids = np.arange(0, num_write_edges)
# masks.
paper_train_mask = np.random.choice([True, False], num_papers)
paper_test_mask = np.random.choice([True, False], num_papers)
paper_val_mask = np.random.choice([True, False], num_papers)
author_train_mask = np.random.choice([True, False], num_authors)
author_test_mask = np.random.choice([True, False], num_authors)
author_val_mask = np.random.choice([True, False], num_authors)
inst_train_mask = np.random.choice([True, False], num_institutions)
inst_test_mask = np.random.choice([True, False], num_institutions)
inst_val_mask = np.random.choice([True, False], num_institutions)
write_train_mask = np.random.choice([True, False], num_write_edges)
write_test_mask = np.random.choice([True, False], num_write_edges)
write_val_mask = np.random.choice([True, False], num_write_edges)
# Edge features.
cite_count = np.random.choice(10, num_cite_edges)
write_year = np.random.choice(2022, num_write_edges)
write2_year = np.random.choice(2022, num_write_edges)
# Save features.
input_dir = os.path.join(root_dir, "data_test")
os.makedirs(input_dir)
for sub_d in ["paper", "cites", "writes", "writes2"]:
os.makedirs(os.path.join(input_dir, sub_d))
paper_feat_path = os.path.join(input_dir, "paper/feat.npy")
with open(paper_feat_path, "wb") as f:
np.save(f, paper_feat)
g.nodes["paper"].data["feat"] = torch.from_numpy(paper_feat)
paper_label_path = os.path.join(input_dir, "paper/label.npy")
with open(paper_label_path, "wb") as f:
np.save(f, paper_label)
g.nodes["paper"].data["label"] = torch.from_numpy(paper_label)
paper_year_path = os.path.join(input_dir, "paper/year.npy")
with open(paper_year_path, "wb") as f:
np.save(f, paper_year)
g.nodes["paper"].data["year"] = torch.from_numpy(paper_year)
paper_orig_ids_path = os.path.join(input_dir, "paper/orig_ids.npy")
with open(paper_orig_ids_path, "wb") as f:
np.save(f, paper_orig_ids)
g.nodes["paper"].data["orig_ids"] = torch.from_numpy(paper_orig_ids)
cite_count_path = os.path.join(input_dir, "cites/count.npy")
with open(cite_count_path, "wb") as f:
np.save(f, cite_count)
g.edges["cites"].data["count"] = torch.from_numpy(cite_count)
write_year_path = os.path.join(input_dir, "writes/year.npy")
with open(write_year_path, "wb") as f:
np.save(f, write_year)
g.edges[("author", "writes", "paper")].data["year"] = torch.from_numpy(
write_year
)
g.edges["rev_writes"].data["year"] = torch.from_numpy(write_year)
writes_orig_ids_path = os.path.join(input_dir, "writes/orig_ids.npy")
with open(writes_orig_ids_path, "wb") as f:
np.save(f, writes_orig_ids)
g.edges[("author", "writes", "paper")].data["orig_ids"] = torch.from_numpy(
writes_orig_ids
)
write2_year_path = os.path.join(input_dir, "writes2/year.npy")
with open(write2_year_path, "wb") as f:
np.save(f, write2_year)
g.edges[("institution", "writes", "paper")].data["year"] = torch.from_numpy(
write2_year
)
etype = ("author", "writes", "paper")
write_train_mask_path = os.path.join(input_dir, "writes/train_mask.npy")
with open(write_train_mask_path, "wb") as f:
np.save(f, write_train_mask)
g.edges[etype].data["train_mask"] = torch.from_numpy(write_train_mask)
write_test_mask_path = os.path.join(input_dir, "writes/test_mask.npy")
with open(write_test_mask_path, "wb") as f:
np.save(f, write_test_mask)
g.edges[etype].data["test_mask"] = torch.from_numpy(write_test_mask)
write_val_mask_path = os.path.join(input_dir, "writes/val_mask.npy")
with open(write_val_mask_path, "wb") as f:
np.save(f, write_val_mask)
g.edges[etype].data["val_mask"] = torch.from_numpy(write_val_mask)
for sub_d in ["author", "institution"]:
os.makedirs(os.path.join(input_dir, sub_d))
paper_train_mask_path = os.path.join(input_dir, "paper/train_mask.npy")
with open(paper_train_mask_path, "wb") as f:
np.save(f, paper_train_mask)
g.nodes["paper"].data["train_mask"] = torch.from_numpy(paper_train_mask)
paper_test_mask_path = os.path.join(input_dir, "paper/test_mask.npy")
with open(paper_test_mask_path, "wb") as f:
np.save(f, paper_test_mask)
g.nodes["paper"].data["test_mask"] = torch.from_numpy(paper_test_mask)
paper_val_mask_path = os.path.join(input_dir, "paper/val_mask.npy")
with open(paper_val_mask_path, "wb") as f:
np.save(f, paper_val_mask)
g.nodes["paper"].data["val_mask"] = torch.from_numpy(paper_val_mask)
author_train_mask_path = os.path.join(input_dir, "author/train_mask.npy")
with open(author_train_mask_path, "wb") as f:
np.save(f, author_train_mask)
g.nodes["author"].data["train_mask"] = torch.from_numpy(author_train_mask)
author_test_mask_path = os.path.join(input_dir, "author/test_mask.npy")
with open(author_test_mask_path, "wb") as f:
np.save(f, author_test_mask)
g.nodes["author"].data["test_mask"] = torch.from_numpy(author_test_mask)
author_val_mask_path = os.path.join(input_dir, "author/val_mask.npy")
with open(author_val_mask_path, "wb") as f:
np.save(f, author_val_mask)
g.nodes["author"].data["val_mask"] = torch.from_numpy(author_val_mask)
inst_train_mask_path = os.path.join(input_dir, "institution/train_mask.npy")
with open(inst_train_mask_path, "wb") as f:
np.save(f, inst_train_mask)
g.nodes["institution"].data["train_mask"] = torch.from_numpy(
inst_train_mask
)
inst_test_mask_path = os.path.join(input_dir, "institution/test_mask.npy")
with open(inst_test_mask_path, "wb") as f:
np.save(f, inst_test_mask)
g.nodes["institution"].data["test_mask"] = torch.from_numpy(inst_test_mask)
inst_val_mask_path = os.path.join(input_dir, "institution/val_mask.npy")
with open(inst_val_mask_path, "wb") as f:
np.save(f, inst_val_mask)
g.nodes["institution"].data["val_mask"] = torch.from_numpy(inst_val_mask)
node_data = {
"paper": {
"feat": paper_feat_path,
"train_mask": paper_train_mask_path,
"test_mask": paper_test_mask_path,
"val_mask": paper_val_mask_path,
"label": paper_label_path,
"year": paper_year_path,
"orig_ids": paper_orig_ids_path,
},
"author": {
"train_mask": author_train_mask_path,
"test_mask": author_test_mask_path,
"val_mask": author_val_mask_path,
},
"institution": {
"train_mask": inst_train_mask_path,
"test_mask": inst_test_mask_path,
"val_mask": inst_val_mask_path,
},
}
edge_data = {
"cites": {"count": cite_count_path},
("author", "writes", "paper"): {
"year": write_year_path,
"orig_ids": writes_orig_ids_path,
"train_mask": write_train_mask_path,
"test_mask": write_test_mask_path,
"val_mask": write_val_mask_path,
},
"rev_writes": {"year": write_year_path},
("institution", "writes", "paper"): {"year": write2_year_path},
}
output_dir = os.path.join(root_dir, "chunked-data")
chunk_graph(
g,
"mag240m",
node_data,
edge_data,
num_chunks=num_chunks,
output_path=output_dir,
data_fmt=data_fmt,
edges_fmt=edges_fmt,
vector_rows=vector_rows,
**kwargs,
)
logging.debug("Done with creating chunked graph")
return g
tests/tools/test_array_readwriter.py
+25
View File
@@ -0,0 +1,25 @@
import os
import tempfile
import numpy as np
import pytest
from distpartitioning import array_readwriter
@pytest.mark.parametrize(
"shape", [[500], [300, 10], [200, 5, 5], [100, 5, 5, 5]]
)
@pytest.mark.parametrize("format", ["numpy", "parquet"])
def test_array_readwriter(format, shape):
original_array = np.random.rand(*shape)
fmt_meta = {"name": format}
with tempfile.TemporaryDirectory() as test_dir:
path = os.path.join(test_dir, f"nodes.{format}")
array_readwriter.get_array_parser(**fmt_meta).write(
path, original_array
)
array = array_readwriter.get_array_parser(**fmt_meta).read(path)
assert original_array.shape == array.shape
assert np.array_equal(original_array, array)
tests/tools/test_change_etype_to_canonical_etype.py
+99
View File
@@ -0,0 +1,99 @@
import json
import os
import tempfile
import unittest
from collections import Counter
import dgl
import pytest
from change_etype_to_canonical_etype import convert_conf, is_old_version
from dgl.distributed import partition_graph
from scipy import sparse as spsp
def create_random_hetero(type_n, node_n):
num_nodes = {}
for i in range(1, type_n + 1):
num_nodes[f"n{i}"] = node_n
c_etypes = []
count = 0
for i in range(1, type_n):
for j in range(i + 1, type_n + 1):
count += 1
c_etypes.append((f"n{i}", f"r{count}", f"n{j}"))
edges = {}
for etype in c_etypes:
src_ntype, _, dst_ntype = etype
arr = spsp.random(
num_nodes[src_ntype],
num_nodes[dst_ntype],
density=0.001,
format="coo",
random_state=100,
)
edges[etype] = (arr.row, arr.col)
return dgl.heterograph(edges, num_nodes), [
":".join(c_etype) for c_etype in c_etypes
]
@unittest.skip(reason="Skip due to glitch in CI")
@pytest.mark.parametrize(
"type_n, node_n, num_parts", [[3, 100, 2], [10, 500, 4], [10, 1000, 8]]
)
def test_hetero_graph(type_n, node_n, num_parts):
g, expected_c_etypes = create_random_hetero(type_n, node_n)
do_convert_and_check(g, "convert_conf_test", num_parts, expected_c_etypes)
@unittest.skip(reason="Skip due to glitch in CI")
@pytest.mark.parametrize("node_n, num_parts", [[100, 2], [500, 4]])
def test_homo_graph(node_n, num_parts):
g = dgl.rand_graph(node_n, node_n // 10)
do_convert_and_check(g, "convert_conf_test", num_parts, ["_N:_E:_N"])
def do_convert_and_check(g, graph_name, num_parts, expected_c_etypes):
with tempfile.TemporaryDirectory() as root_dir:
partition_graph(g, graph_name, num_parts, root_dir)
part_config = os.path.join(root_dir, graph_name + ".json")
old_config = _get_old_config(part_config)
# Call convert function
convert_conf(part_config)
with open(part_config, "r") as config_f:
config = json.load(config_f)
# Check we get all canonical etypes
assert Counter(expected_c_etypes) == Counter(
config["etypes"].keys()
)
# Check the id is match after transform from etypes -> canonical
assert old_config["etypes"] == _extract_etypes(config["etypes"])
def _get_old_config(part_config):
with open(part_config, "r+") as config_f:
config = json.load(config_f)
if not is_old_version(config):
config["etypes"] = _extract_etypes(config["etypes"])
config["edge_map"] = _extract_edge_map(config["edge_map"])
config_f.seek(0)
json.dump(config, config_f, indent=4)
config_f.truncate()
return config
def _extract_etypes(c_etypes):
etypes = {}
for c_etype, eid in c_etypes.items():
etype = c_etype.split(":")[1]
etypes[etype] = eid
return etypes
def _extract_edge_map(c_edge_map):
edge_map = {}
for c_etype, emap in c_edge_map.items():
etype = c_etype.split(":")[1]
edge_map[etype] = emap
return edge_map
tests/tools/test_convert_partition.py
+251
View File
@@ -0,0 +1,251 @@
import os
import tempfile
import numpy as np
import pytest
import utils
from convert_partition import _get_unique_invidx
@pytest.mark.parametrize(
"num_nodes, num_edges, nid_begin, nid_end",
[
[4000, 40000, 0, 1000],
[4000, 40000, 1000, 2000],
[4000, 40000, 2000, 3000],
[4000, 40000, 3000, 4000],
[4000, 100, 0, 1000],
[4000, 100, 1000, 2000],
[4000, 100, 2000, 3000],
[4000, 100, 3000, 4000],
[1, 1, 0, 1],
],
)
def test_get_unique_invidx_with_numpy(num_nodes, num_edges, nid_begin, nid_end):
# prepare data for the function
# generate synthetic edges
if num_edges > 0:
srcids = np.random.randint(0, num_nodes, (num_edges,)) # exclusive
dstids = np.random.randint(
nid_begin, nid_end, (num_edges,)
) # exclusive
else:
srcids = np.array([])
dstids = np.array([])
assert nid_begin <= nid_end
# generate unique node-ids for any
# partition. This list should be sorted.
# This is equivilant to shuffle_nids in a partition
unique_nids = np.arange(nid_begin, nid_end) # exclusive
# test with numpy unique here
orig_srcids = srcids.copy()
orig_dstids = dstids.copy()
input_arr = np.concatenate([srcids, dstids, unique_nids])
# test
uniques, idxes, srcids, dstids = _get_unique_invidx(
srcids, dstids, unique_nids
)
assert len(uniques) == len(idxes)
assert np.all(srcids < len(uniques))
assert np.all(dstids < len(uniques))
assert np.all(uniques[srcids].sort() == orig_srcids.sort())
assert np.all(uniques[dstids] == orig_dstids)
assert np.all(uniques == input_arr[idxes])
# numpy
np_uniques, np_idxes, np_inv_idxes = np.unique(
np.concatenate([orig_srcids, orig_dstids, unique_nids]),
return_index=True,
return_inverse=True,
)
# test uniques
assert np.all(np_uniques == uniques)
# test idxes array
assert np.all(input_arr[idxes].sort() == input_arr[np_idxes].sort())
# test srcids, inv_indices
assert np.all(
uniques[srcids].sort()
== np_uniques[np_inv_idxes[0 : len(srcids)]].sort()
)
# test dstids, inv_indices
assert np.all(
uniques[dstids].sort() == np_uniques[np_inv_idxes[len(srcids) :]].sort()
)
@pytest.mark.parametrize(
"num_nodes, num_edges, nid_begin, nid_end",
[
# dense networks, no. of edges more than no. of nodes
[4000, 40000, 0, 1000],
[4000, 40000, 1000, 2000],
[4000, 40000, 2000, 3000],
[4000, 40000, 3000, 4000],
# sparse networks, no. of edges smaller than no. of nodes
[4000, 100, 0, 1000],
[4000, 100, 1000, 2000],
[4000, 100, 2000, 3000],
[4000, 100, 3000, 4000],
# corner case
[1, 1, 0, 1],
],
)
def test_get_unique_invidx(num_nodes, num_edges, nid_begin, nid_end):
# prepare data for the function
# generate synthetic edges
if num_edges > 0:
srcids = np.random.randint(0, num_nodes, (num_edges,))
dstids = np.random.randint(nid_begin, nid_end, (num_edges,))
else:
srcids = np.array([])
dstids = np.array([])
assert nid_begin <= nid_end
# generate unique node-ids for any
# partition. This list should be sorted.
# This is equivilant to shuffle_nids in a partition
unique_nids = np.arange(nid_begin, nid_end)
# invoke the test target
uniques, idxes, src_ids, dst_ids = _get_unique_invidx(
srcids, dstids, unique_nids
)
# validate the outputs of this function
# array uniques should be sorted list of integers.
assert np.all(
np.diff(uniques) >= 0
), f"Output parameter uniques assert failing."
# idxes are list of integers
# these are indices in the concatenated list (srcids, dstids, unique_nids)
max_idx = len(src_ids) + len(dst_ids) + len(unique_nids)
assert np.all(idxes >= 0), f"Output parameter idxes has negative values."
assert np.all(
idxes < max_idx
), f"Output parameter idxes has invalid maximum value."
# srcids and dstids will be inverse indices in the uniques list
min_src = np.amin(src_ids)
max_src = np.amax(src_ids)
min_dst = np.amin(dst_ids)
max_dst = np.amax(dst_ids)
assert (
len(uniques) > max_src
), f"Inverse idx, src_ids, has invalid max value."
assert min_src >= 0, f"Inverse idx, src_ids has negative values."
assert len(uniques) > max_dst, f"Inverse idx, dst_ids, invalid max value."
assert max_dst >= 0, f"Inverse idx, dst_ids has negative values."
def test_get_unique_invidx_low_mem():
srcids = np.array([14, 0, 3, 3, 0, 3, 9, 5, 14, 12])
dstids = np.array([10, 16, 12, 13, 10, 17, 16, 13, 14, 16])
unique_nids = np.array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
uniques, idxes, srcids, dstids = _get_unique_invidx(
srcids,
dstids,
unique_nids,
low_mem=True,
)
expected_unqiues = np.array(
[0, 3, 5, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
)
expected_idxes = np.array(
[1, 2, 7, 6, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
)
expected_srcids = np.array([8, 0, 1, 1, 0, 1, 3, 2, 8, 6])
expected_dstids = np.array([4, 10, 6, 7, 4, 11, 10, 7, 8, 10])
assert np.all(
uniques == expected_unqiues
), f"unique is not expected. {uniques} != {expected_unqiues}"
assert np.all(
idxes == expected_idxes
), f"indices is not expected. {idxes} != {expected_idxes}"
assert np.all(
srcids == expected_srcids
), f"srcids is not expected. {srcids} != {expected_srcids}"
assert np.all(
dstids == expected_dstids
), f"dstdis is not expected. {dstids} != {expected_dstids}"
def test_get_unique_invidx_high_mem():
srcids = np.array([14, 0, 3, 3, 0, 3, 9, 5, 14, 12])
dstids = np.array([10, 16, 12, 13, 10, 17, 16, 13, 14, 16])
unique_nids = np.array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
uniques, idxes, srcids, dstids = _get_unique_invidx(
srcids,
dstids,
unique_nids,
low_mem=False,
)
expected_unqiues = np.array(
[0, 3, 5, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]
)
expected_idxes = np.array(
[1, 2, 7, 6, 10, 21, 9, 13, 0, 25, 11, 15, 28, 29]
)
expected_srcids = np.array([8, 0, 1, 1, 0, 1, 3, 2, 8, 6])
expected_dstids = np.array([4, 10, 6, 7, 4, 11, 10, 7, 8, 10])
assert np.all(
uniques == expected_unqiues
), f"unique is not expected. {uniques} != {expected_unqiues}"
assert np.all(
idxes == expected_idxes
), f"indices is not expected. {idxes} != {expected_idxes}"
assert np.all(
srcids == expected_srcids
), f"srcids is not expected. {srcids} != {expected_srcids}"
assert np.all(
dstids == expected_dstids
), f"dstdis is not expected. {dstids} != {expected_dstids}"
def test_get_unique_invidx_low_high_mem():
srcids = np.array([14, 0, 3, 3, 0, 3, 9, 5, 14, 12])
dstids = np.array([10, 16, 12, 13, 10, 17, 16, 13, 14, 16])
unique_nids = np.array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
uniques_low, idxes_low, srcids_low, dstids_low = _get_unique_invidx(
srcids,
dstids,
unique_nids,
low_mem=True,
)
uniques_high, idxes_high, srcids_high, dstids_high = _get_unique_invidx(
srcids,
dstids,
unique_nids,
low_mem=False,
)
assert np.all(
uniques_low == uniques_high
), f"unique is not expected. {uniques_low} != {uniques_high}"
assert not np.all(
idxes_low == idxes_high
), f"indices is not expected. {idxes_low} == {idxes_high}"
assert np.all(
srcids_low == srcids_high
), f"srcids is not expected. {srcids_low} != {srcids_high}"
assert np.all(
dstids_low == dstids_high
), f"dstdis is not expected. {dstids_low} != {dstids_high}"
tests/tools/test_dist_lookup.py
+259
View File
@@ -0,0 +1,259 @@
import logging
import os
import platform
import tempfile
from datetime import timedelta
import dgl
import numpy as np
import pyarrow
import pytest
import torch.distributed as dist
import torch.multiprocessing as mp
from pytest_utils import create_chunked_dataset
from tools.distpartitioning import constants, dist_lookup
from tools.distpartitioning.gloo_wrapper import allgather_sizes
from tools.distpartitioning.utils import (
get_idranges,
get_ntype_counts_map,
read_json,
)
try:
mp.set_start_method("spawn", force=True)
except RuntimeError:
pass
def _init_process_group(rank, world_size):
# init the gloo process group here.
dist.init_process_group(
backend="gloo",
rank=rank,
world_size=world_size,
timeout=timedelta(seconds=180),
)
print(f"[Rank: {rank}] Done with process group initialization...")
def _create_lookup_service(
partitions_dir, ntypes, id_map, rank, world_size, num_parts
):
id_lookup = dist_lookup.DistLookupService(
partitions_dir, ntypes, rank, world_size, num_parts
)
id_lookup.set_idMap(id_map)
# invoke the main function here.
print(f"[Rank: {rank}] Done with Dist Lookup Service initialization...")
return id_lookup
def _run(
port_num,
rank,
num_parts,
world_size,
partitions_dir,
ntypes,
id_map,
test_data,
):
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port_num)
_init_process_group(rank, world_size)
lookup = _create_lookup_service(
partitions_dir, ntypes, id_map, rank, world_size, num_parts
)
tests_exec = 0
for worker, data in test_data.items():
if f"rank-{rank}" == worker:
for item in data:
method = item[0]
request = item[1]
response = item[2]
if method == "getpartitionids":
ret_val = lookup.get_partition_ids(request)
tests_exec += 1
assert np.all(ret_val == response)
else:
assert False
# ensure all the tests are executed.
rank_counts = allgather_sizes([tests_exec], world_size, num_parts, True)
assert np.sum(rank_counts) == len(test_data)
def _single_machine_run(
num_parts, world_size, partitions_dir, ntypes, id_map, test_data
):
port_num = np.random.randint(10000, 20000, size=(1,), dtype=int)[0]
ctx = mp.get_context("spawn")
processes = []
for rank in range(world_size):
p = ctx.Process(
target=_run,
args=(
port_num,
rank,
num_parts,
world_size,
partitions_dir,
ntypes,
id_map,
test_data,
),
)
p.start()
processes.append(p)
for p in processes:
p.join()
p.close()
def _prepare_test_data(partitions_dir, ntypes, gid_ranges, world_size):
# read node-id to partition-id mappings from disk
ntype_partids = []
for ntype_id, ntype in enumerate(ntypes):
filename = f"{ntype}.txt"
assert os.path.isfile(os.path.join(partitions_dir, filename))
read_options = pyarrow.csv.ReadOptions(
use_threads=True,
block_size=4096,
autogenerate_column_names=True,
)
parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
with pyarrow.csv.open_csv(
os.path.join(partitions_dir, "{}.txt".format(ntype)),
read_options=read_options,
parse_options=parse_options,
) as reader:
for next_chunk in reader:
if next_chunk is None:
break
next_table = pyarrow.Table.from_batches([next_chunk])
ntype_partids.append(next_table["f0"].to_numpy())
# prepare test data for each rank here
# key = f'rank-{rank}'
# value is a list of tuple [(method-name, request, response)]
test_data = {}
for rank in range(world_size):
ntype_id = np.random.randint(0, len(ntypes) - 1)
ntype = ntypes[ntype_id]
request = (
np.arange(len(ntype_partids[ntype_id]))
+ gid_ranges[ntypes[ntype_id]][0, 0]
)
response = ntype_partids[ntype_id]
test_data[f"rank-{rank}"] = [("getpartitionids", request, response)]
# randomly shuffle the global-nids and retrieve their partition-ids.
for rank in range(world_size):
ntype_id = np.random.randint(0, len(ntypes) - 1)
ntype = ntypes[ntype_id]
idx = np.arange(len(ntype_partids[ntype_id]))
request = idx + gid_ranges[ntypes[ntype_id]][0, 0]
np.random.shuffle(idx)
request = request[idx]
response = ntype_partids[ntype_id][idx]
test_data[f"rank-{rank}"] = [("getpartitionids", request, response)]
# one final test
# mix all the ntypes and shuffle randomly
request = []
response = []
for idx in range(len(ntype_partids)):
request.append(
np.arange(len(ntype_partids[idx])) + gid_ranges[ntypes[idx]][0, 0]
)
response.append(ntype_partids[idx])
request = np.concatenate(request)
response = np.concatenate(response)
idx = np.arange(len(request))
np.random.shuffle(idx)
request = request[idx]
response = response[idx]
for idx in range(world_size):
test_data[f"rank-{idx}"] = [("getpartitionids", request, response)]
return test_data
@pytest.mark.parametrize(
"num_chunks, num_parts, world_size",
[[4, 4, 4], [8, 4, 2], [8, 4, 4], [9, 6, 3], [11, 11, 1], [11, 4, 1]],
)
def test_lookup_service(
num_chunks,
num_parts,
world_size,
num_chunks_nodes=None,
num_chunks_edges=None,
num_chunks_node_data=None,
num_chunks_edge_data=None,
):
with tempfile.TemporaryDirectory() as root_dir:
g = create_chunked_dataset(
root_dir,
num_chunks,
data_fmt="numpy",
num_chunks_nodes=num_chunks_nodes,
num_chunks_edges=num_chunks_edges,
num_chunks_node_data=num_chunks_node_data,
num_chunks_edge_data=num_chunks_edge_data,
)
# Step1: graph partition
in_dir = os.path.join(root_dir, "chunked-data")
output_dir = os.path.join(root_dir, "parted_data")
os.system(
"python3 tools/partition_algo/random_partition.py "
"--in_dir {} --out_dir {} --num_partitions {}".format(
in_dir, output_dir, num_parts
)
)
# metadata for original graph
orig_config = os.path.join(in_dir, "metadata.json")
orig_schema = read_json(orig_config)
ntypes = orig_schema[constants.STR_NODE_TYPE]
_, global_nid_ranges = get_idranges(
orig_schema[constants.STR_NODE_TYPE],
get_ntype_counts_map(
orig_schema[constants.STR_NODE_TYPE],
orig_schema[constants.STR_NUM_NODES_PER_TYPE],
),
num_chunks=num_parts,
)
id_map = dgl.distributed.id_map.IdMap(global_nid_ranges)
# run the test
_single_machine_run(
num_parts,
world_size,
output_dir,
ntypes,
id_map,
_prepare_test_data(
output_dir, ntypes, global_nid_ranges, world_size
),
)
tests/tools/test_dist_part.py
+572
View File
@@ -0,0 +1,572 @@
import json
import os
import tempfile
import dgl
import dgl.backend as F
import numpy as np
import pyarrow.parquet as pq
import pytest
import torch
from dgl.data.utils import load_graphs, load_tensors
from dgl.distributed.partition import (
_etype_tuple_to_str,
_get_inner_edge_mask,
_get_inner_node_mask,
load_partition,
RESERVED_FIELD_DTYPE,
)
from distpartitioning import array_readwriter
from distpartitioning.utils import generate_read_list
from pytest_utils import chunk_graph, create_chunked_dataset
from scipy import sparse as spsp
from tools.verification_utils import (
verify_graph_feats,
verify_partition_data_types,
verify_partition_formats,
)
def _test_chunk_graph(
num_chunks,
data_fmt="numpy",
edges_fmt="csv",
vector_rows=False,
num_chunks_nodes=None,
num_chunks_edges=None,
num_chunks_node_data=None,
num_chunks_edge_data=None,
):
with tempfile.TemporaryDirectory() as root_dir:
g = create_chunked_dataset(
root_dir,
num_chunks,
data_fmt=data_fmt,
edges_fmt=edges_fmt,
vector_rows=vector_rows,
num_chunks_nodes=num_chunks_nodes,
num_chunks_edges=num_chunks_edges,
num_chunks_node_data=num_chunks_node_data,
num_chunks_edge_data=num_chunks_edge_data,
)
# check metadata.json
output_dir = os.path.join(root_dir, "chunked-data")
json_file = os.path.join(output_dir, "metadata.json")
assert os.path.isfile(json_file)
with open(json_file, "rb") as f:
meta_data = json.load(f)
assert meta_data["graph_name"] == "mag240m"
assert len(meta_data["num_nodes_per_chunk"][0]) == num_chunks
# check edge_index
output_edge_index_dir = os.path.join(output_dir, "edge_index")
for c_etype in g.canonical_etypes:
c_etype_str = _etype_tuple_to_str(c_etype)
if num_chunks_edges is None:
n_chunks = num_chunks
else:
n_chunks = num_chunks_edges
for i in range(n_chunks):
fname = os.path.join(
output_edge_index_dir, f"{c_etype_str}{i}.txt"
)
assert os.path.isfile(fname)
if edges_fmt == "csv":
with open(fname, "r") as f:
header = f.readline()
num1, num2 = header.rstrip().split(" ")
assert isinstance(int(num1), int)
assert isinstance(int(num2), int)
elif edges_fmt == "parquet":
metadata = pq.read_metadata(fname)
assert metadata.num_columns == 2
else:
assert False, f"Invalid edges_fmt: {edges_fmt}"
# check node/edge_data
suffix = "npy" if data_fmt == "numpy" else "parquet"
reader_fmt_meta = {"name": data_fmt}
def test_data(sub_dir, feat, expected_data, expected_shape, num_chunks):
data = []
for i in range(num_chunks):
fname = os.path.join(sub_dir, f"{feat}-{i}.{suffix}")
assert os.path.isfile(fname), f"{fname} cannot be found."
feat_array = array_readwriter.get_array_parser(
**reader_fmt_meta
).read(fname)
assert feat_array.shape[0] == expected_shape
data.append(feat_array)
data = np.concatenate(data, 0)
assert torch.equal(torch.from_numpy(data), expected_data)
output_node_data_dir = os.path.join(output_dir, "node_data")
for ntype in g.ntypes:
sub_dir = os.path.join(output_node_data_dir, ntype)
if isinstance(num_chunks_node_data, int):
chunks_data = num_chunks_node_data
elif isinstance(num_chunks_node_data, dict):
chunks_data = num_chunks_node_data.get(ntype, num_chunks)
else:
chunks_data = num_chunks
for feat, data in g.nodes[ntype].data.items():
if isinstance(chunks_data, dict):
n_chunks = chunks_data.get(feat, num_chunks)
else:
n_chunks = chunks_data
test_data(
sub_dir,
feat,
data,
g.num_nodes(ntype) // n_chunks,
n_chunks,
)
output_edge_data_dir = os.path.join(output_dir, "edge_data")
for c_etype in g.canonical_etypes:
c_etype_str = _etype_tuple_to_str(c_etype)
sub_dir = os.path.join(output_edge_data_dir, c_etype_str)
if isinstance(num_chunks_edge_data, int):
chunks_data = num_chunks_edge_data
elif isinstance(num_chunks_edge_data, dict):
chunks_data = num_chunks_edge_data.get(c_etype, num_chunks)
else:
chunks_data = num_chunks
for feat, data in g.edges[c_etype].data.items():
if isinstance(chunks_data, dict):
n_chunks = chunks_data.get(feat, num_chunks)
else:
n_chunks = chunks_data
test_data(
sub_dir,
feat,
data,
g.num_edges(c_etype) // n_chunks,
n_chunks,
)
@pytest.mark.parametrize("num_chunks", [1, 8])
@pytest.mark.parametrize("data_fmt", ["numpy", "parquet"])
@pytest.mark.parametrize("edges_fmt", ["csv", "parquet"])
def test_chunk_graph_basics(num_chunks, data_fmt, edges_fmt):
_test_chunk_graph(num_chunks, data_fmt=data_fmt, edges_fmt=edges_fmt)
@pytest.mark.parametrize("num_chunks", [1, 8])
@pytest.mark.parametrize("vector_rows", [True, False])
def test_chunk_graph_vector_rows(num_chunks, vector_rows):
_test_chunk_graph(
num_chunks,
data_fmt="parquet",
edges_fmt="parquet",
vector_rows=vector_rows,
)
@pytest.mark.parametrize(
"num_chunks, "
"num_chunks_nodes, "
"num_chunks_edges, "
"num_chunks_node_data, "
"num_chunks_edge_data",
[
[1, None, None, None, None],
[8, None, None, None, None],
[4, 4, 4, 8, 12],
[4, 4, 4, {"paper": 10}, {("author", "writes", "paper"): 24}],
[
4,
4,
4,
{"paper": {"feat": 10}},
{("author", "writes", "paper"): {"year": 24}},
],
],
)
def test_chunk_graph_arbitrary_chunks(
num_chunks,
num_chunks_nodes,
num_chunks_edges,
num_chunks_node_data,
num_chunks_edge_data,
):
_test_chunk_graph(
num_chunks,
num_chunks_nodes=num_chunks_nodes,
num_chunks_edges=num_chunks_edges,
num_chunks_node_data=num_chunks_node_data,
num_chunks_edge_data=num_chunks_edge_data,
)
def create_mini_chunked_dataset(
root_dir,
num_chunks,
data_fmt,
edges_fmt,
vector_rows,
few_entity="node",
**kwargs,
):
num_nodes = {"n1": 1000, "n2": 1010, "n3": 1020}
etypes = [
("n1", "r1", "n2"),
("n2", "r1", "n1"),
("n1", "r2", "n3"),
("n2", "r3", "n3"),
]
node_items = ["n1", "n2", "n3"]
edges_coo = {}
for etype in etypes:
src_ntype, _, dst_ntype = etype
arr = spsp.random(
num_nodes[src_ntype],
num_nodes[dst_ntype],
density=0.001,
format="coo",
random_state=100,
)
edges_coo[etype] = (arr.row, arr.col)
edge_items = []
if few_entity == "edge":
edges_coo[("n1", "a0", "n2")] = (
torch.tensor([0, 1]),
torch.tensor([1, 0]),
)
edges_coo[("n1", "a1", "n3")] = (
torch.tensor([0, 1]),
torch.tensor([1, 0]),
)
edge_items.append(("n1", "a0", "n2"))
edge_items.append(("n1", "a1", "n3"))
elif few_entity == "node":
edges_coo[("n1", "r_few", "n_few")] = (
torch.tensor([0, 1]),
torch.tensor([1, 0]),
)
edges_coo[("a0", "a01", "n_1")] = (
torch.tensor([0, 1]),
torch.tensor([1, 0]),
)
edge_items.append(("n1", "r_few", "n_few"))
edge_items.append(("a0", "a01", "n_1"))
node_items.append("n_few")
node_items.append("n_1")
num_nodes["n_few"] = 2
num_nodes["n_1"] = 2
g = dgl.heterograph(edges_coo)
node_data = {}
edge_data = {}
# save feature
input_dir = os.path.join(root_dir, "data_test")
for ntype in node_items:
os.makedirs(os.path.join(input_dir, ntype))
feat = np.random.randn(num_nodes[ntype], 3)
feat_path = os.path.join(input_dir, f"{ntype}/feat.npy")
with open(feat_path, "wb") as f:
np.save(f, feat)
g.nodes[ntype].data["feat"] = torch.from_numpy(feat)
node_data[ntype] = {"feat": feat_path}
for etype in set(edge_items):
os.makedirs(os.path.join(input_dir, etype[1]))
num_edge = len(edges_coo[etype][0])
feat = np.random.randn(num_edge, 4)
feat_path = os.path.join(input_dir, f"{etype[1]}/feat.npy")
with open(feat_path, "wb") as f:
np.save(f, feat)
g.edges[etype].data["feat"] = torch.from_numpy(feat)
edge_data[etype] = {"feat": feat_path}
output_dir = os.path.join(root_dir, "chunked-data")
chunk_graph(
g,
"mag240m",
node_data,
edge_data,
num_chunks=num_chunks,
output_path=output_dir,
data_fmt=data_fmt,
edges_fmt=edges_fmt,
vector_rows=vector_rows,
**kwargs,
)
return g
def _test_pipeline(
num_chunks,
num_parts,
world_size,
graph_formats=None,
data_fmt="numpy",
num_chunks_nodes=None,
num_chunks_edges=None,
num_chunks_node_data=None,
num_chunks_edge_data=None,
use_verify_partitions=False,
):
if num_parts % world_size != 0:
# num_parts should be a multiple of world_size
return
with tempfile.TemporaryDirectory() as root_dir:
g = create_chunked_dataset(
root_dir,
num_chunks,
data_fmt=data_fmt,
num_chunks_nodes=num_chunks_nodes,
num_chunks_edges=num_chunks_edges,
num_chunks_node_data=num_chunks_node_data,
num_chunks_edge_data=num_chunks_edge_data,
)
# Step1: graph partition
in_dir = os.path.join(root_dir, "chunked-data")
output_dir = os.path.join(root_dir, "parted_data")
os.system(
"python3 tools/partition_algo/random_partition.py "
"--in_dir {} --out_dir {} --num_partitions {}".format(
in_dir, output_dir, num_parts
)
)
for ntype in ["author", "institution", "paper"]:
fname = os.path.join(output_dir, "{}.txt".format(ntype))
with open(fname, "r") as f:
header = f.readline().rstrip()
assert isinstance(int(header), int)
# Step2: data dispatch
partition_dir = os.path.join(root_dir, "parted_data")
out_dir = os.path.join(root_dir, "partitioned")
ip_config = os.path.join(root_dir, "ip_config.txt")
with open(ip_config, "w") as f:
for i in range(world_size):
f.write(f"127.0.0.{i + 1}\n")
cmd = "python3 tools/dispatch_data.py"
cmd += f" --in-dir {in_dir}"
cmd += f" --partitions-dir {partition_dir}"
cmd += f" --out-dir {out_dir}"
cmd += f" --ip-config {ip_config}"
cmd += " --ssh-port 22"
cmd += " --process-group-timeout 60"
cmd += " --save-orig-nids"
cmd += " --save-orig-eids"
cmd += f" --graph-formats {graph_formats}" if graph_formats else ""
os.system(cmd)
# check if verify_partitions.py is used for validation.
if use_verify_partitions:
cmd = "python3 tools/verify_partitions.py "
cmd += f" --orig-dataset-dir {in_dir}"
cmd += f" --part-graph {out_dir}"
cmd += f" --partitions-dir {output_dir}"
os.system(cmd)
return
# read original node/edge IDs
def read_orig_ids(fname):
orig_ids = {}
for i in range(num_parts):
ids_path = os.path.join(out_dir, f"part{i}", fname)
part_ids = load_tensors(ids_path)
for type, data in part_ids.items():
if type not in orig_ids:
orig_ids[type] = data
else:
orig_ids[type] = torch.cat((orig_ids[type], data))
return orig_ids
orig_nids = read_orig_ids("orig_nids.dgl")
orig_eids = read_orig_ids("orig_eids.dgl")
# load partitions and verify
part_config = os.path.join(out_dir, "metadata.json")
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, graph_formats)
verify_graph_feats(
g, gpb, part_g, node_feats, edge_feats, orig_nids, orig_eids
)
@pytest.mark.parametrize(
"num_chunks, num_parts, world_size",
[[4, 4, 4], [8, 4, 2], [8, 4, 4], [9, 6, 3], [11, 11, 1], [11, 4, 1]],
)
def test_pipeline_basics(num_chunks, num_parts, world_size):
_test_pipeline(num_chunks, num_parts, world_size)
_test_pipeline(
num_chunks, num_parts, world_size, use_verify_partitions=False
)
@pytest.mark.parametrize(
"graph_formats", [None, "csc", "coo,csc", "coo,csc,csr"]
)
def test_pipeline_formats(graph_formats):
_test_pipeline(4, 4, 4, graph_formats)
@pytest.mark.parametrize(
"num_chunks, "
"num_parts, "
"world_size, "
"num_chunks_node_data, "
"num_chunks_edge_data",
[
# Test cases where no. of chunks more than
# no. of partitions
[8, 4, 4, 8, 8],
[8, 4, 2, 8, 8],
[9, 7, 5, 9, 9],
[8, 8, 4, 8, 8],
# Test cases where no. of chunks smaller
# than no. of partitions
[7, 8, 4, 7, 7],
[1, 8, 4, 1, 1],
[1, 4, 4, 1, 1],
[3, 4, 4, 3, 3],
[1, 4, 2, 1, 1],
[3, 4, 2, 3, 3],
[1, 5, 3, 1, 1],
],
)
def test_pipeline_arbitrary_chunks(
num_chunks,
num_parts,
world_size,
num_chunks_node_data,
num_chunks_edge_data,
):
_test_pipeline(
num_chunks,
num_parts,
world_size,
num_chunks_node_data=num_chunks_node_data,
num_chunks_edge_data=num_chunks_edge_data,
)
@pytest.mark.parametrize(
"graph_formats", [None, "csc", "coo,csc", "coo,csc,csr"]
)
def test_pipeline_formats(graph_formats):
_test_pipeline(4, 4, 4, graph_formats)
@pytest.mark.parametrize("data_fmt", ["numpy", "parquet"])
def test_pipeline_feature_format(data_fmt):
_test_pipeline(4, 4, 4, data_fmt=data_fmt)
@pytest.mark.parametrize(
"num_chunks, num_parts, world_size",
[[4, 4, 4], [8, 4, 2], [8, 4, 4], [9, 6, 3], [11, 11, 1], [11, 4, 1]],
)
@pytest.mark.parametrize("few_entity", ["node", "edge"])
def test_partition_hetero_few_entity(
num_chunks,
num_parts,
world_size,
few_entity,
graph_formats=None,
data_fmt="numpy",
edges_fmt="csv",
vector_rows=False,
num_chunks_nodes=None,
num_chunks_edges=None,
num_chunks_node_data=None,
num_chunks_edge_data=None,
):
with tempfile.TemporaryDirectory() as root_dir:
g = create_mini_chunked_dataset(
root_dir,
num_chunks,
few_entity=few_entity,
data_fmt=data_fmt,
edges_fmt=edges_fmt,
vector_rows=vector_rows,
num_chunks_nodes=num_chunks_nodes,
num_chunks_edges=num_chunks_edges,
num_chunks_node_data=num_chunks_node_data,
num_chunks_edge_data=num_chunks_edge_data,
)
# Step1: graph partition
in_dir = os.path.join(root_dir, "chunked-data")
output_dir = os.path.join(root_dir, "parted_data")
os.system(
"python3 tools/partition_algo/random_partition.py "
"--in_dir {} --out_dir {} --num_partitions {}".format(
in_dir, output_dir, num_parts
)
)
# Step2: data dispatch
partition_dir = os.path.join(root_dir, "parted_data")
out_dir = os.path.join(root_dir, "partitioned")
ip_config = os.path.join(root_dir, "ip_config.txt")
with open(ip_config, "w") as f:
for i in range(world_size):
f.write(f"127.0.0.{i + 1}\n")
cmd = "python3 tools/dispatch_data.py"
cmd += f" --in-dir {in_dir}"
cmd += f" --partitions-dir {partition_dir}"
cmd += f" --out-dir {out_dir}"
cmd += f" --ip-config {ip_config}"
cmd += " --ssh-port 22"
cmd += " --process-group-timeout 60"
cmd += " --save-orig-nids"
cmd += " --save-orig-eids"
cmd += f" --graph-formats {graph_formats}" if graph_formats else ""
os.system(cmd)
# read original node/edge IDs
def read_orig_ids(fname):
orig_ids = {}
for i in range(num_parts):
ids_path = os.path.join(out_dir, f"part{i}", fname)
part_ids = load_tensors(ids_path)
for type, data in part_ids.items():
if type not in orig_ids:
orig_ids[type] = data
else:
orig_ids[type] = torch.cat((orig_ids[type], data))
return orig_ids
orig_nids = read_orig_ids("orig_nids.dgl")
orig_eids = read_orig_ids("orig_eids.dgl")
# load partitions and verify
part_config = os.path.join(out_dir, "metadata.json")
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, graph_formats)
verify_graph_feats(
g, gpb, part_g, node_feats, edge_feats, orig_nids, orig_eids
)
def test_utils_generate_read_list():
read_list = generate_read_list(10, 4)
assert np.array_equal(read_list[0], np.array([0, 1, 2]))
assert np.array_equal(read_list[1], np.array([3, 4, 5]))
assert np.array_equal(read_list[2], np.array([6, 7]))
assert np.array_equal(read_list[3], np.array([8, 9]))
tests/tools/test_dist_partition_graphbolt.py
+1023
View File
File diff suppressed because it is too large Load Diff
tests/tools/test_launch.py
+236
View File
@@ -0,0 +1,236 @@
import json
import os
import tempfile
import unittest
from launch import *
class TestWrapUdfInTorchDistLauncher(unittest.TestCase):
"""wrap_udf_in_torch_dist_launcher()"""
def test_simple(self):
# test that a simple udf_command is correctly wrapped
udf_command = "python3.7 path/to/some/trainer.py arg1 arg2"
wrapped_udf_command = wrap_udf_in_torch_dist_launcher(
udf_command=udf_command,
num_trainers=2,
num_nodes=2,
node_rank=1,
master_addr="127.0.0.1",
master_port=1234,
)
expected = (
"python3.7 -m torch.distributed.run "
"--nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr=127.0.0.1 "
"--master_port=1234 path/to/some/trainer.py arg1 arg2"
)
self.assertEqual(wrapped_udf_command, expected)
def test_chained_udf(self):
# test that a chained udf_command is properly handled
udf_command = (
"cd path/to && python3.7 path/to/some/trainer.py arg1 arg2"
)
wrapped_udf_command = wrap_udf_in_torch_dist_launcher(
udf_command=udf_command,
num_trainers=2,
num_nodes=2,
node_rank=1,
master_addr="127.0.0.1",
master_port=1234,
)
expected = (
"cd path/to && python3.7 -m torch.distributed.run "
"--nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr=127.0.0.1 "
"--master_port=1234 path/to/some/trainer.py arg1 arg2"
)
self.assertEqual(wrapped_udf_command, expected)
def test_py_versions(self):
# test that this correctly handles different py versions/binaries
py_binaries = (
"python3.7",
"python3.8",
"python3.9",
"python3",
"python",
)
udf_command = "{python_bin} path/to/some/trainer.py arg1 arg2"
for py_bin in py_binaries:
wrapped_udf_command = wrap_udf_in_torch_dist_launcher(
udf_command=udf_command.format(python_bin=py_bin),
num_trainers=2,
num_nodes=2,
node_rank=1,
master_addr="127.0.0.1",
master_port=1234,
)
expected = (
"{python_bin} -m torch.distributed.run ".format(
python_bin=py_bin
)
+ "--nproc_per_node=2 --nnodes=2 --node_rank=1 --master_addr=127.0.0.1 "
"--master_port=1234 path/to/some/trainer.py arg1 arg2"
)
self.assertEqual(wrapped_udf_command, expected)
class TestWrapCmdWithLocalEnvvars(unittest.TestCase):
"""wrap_cmd_with_local_envvars()"""
def test_simple(self):
self.assertEqual(
wrap_cmd_with_local_envvars("ls && pwd", "VAR1=value1 VAR2=value2"),
"(export VAR1=value1 VAR2=value2; ls && pwd)",
)
class TestConstructDglServerEnvVars(unittest.TestCase):
"""construct_dgl_server_env_vars()"""
def test_simple(self):
self.assertEqual(
construct_dgl_server_env_vars(
num_samplers=2,
num_server_threads=3,
tot_num_clients=4,
part_config="path/to/part.config",
ip_config="path/to/ip.config",
num_servers=5,
graph_format="csc",
),
(
"DGL_ROLE=server "
"DGL_NUM_SAMPLER=2 "
"OMP_NUM_THREADS=3 "
"DGL_NUM_CLIENT=4 "
"DGL_CONF_PATH=path/to/part.config "
"DGL_IP_CONFIG=path/to/ip.config "
"DGL_NUM_SERVER=5 "
"DGL_GRAPH_FORMAT=csc "
),
)
class TestConstructDglClientEnvVars(unittest.TestCase):
"""construct_dgl_client_env_vars()"""
def test_simple(self):
# with pythonpath
self.assertEqual(
construct_dgl_client_env_vars(
num_samplers=1,
tot_num_clients=2,
part_config="path/to/part.config",
ip_config="path/to/ip.config",
num_servers=3,
graph_format="csc",
num_omp_threads=4,
group_id=0,
pythonpath="some/pythonpath/",
),
(
"DGL_DIST_MODE=distributed "
"DGL_ROLE=client "
"DGL_NUM_SAMPLER=1 "
"DGL_NUM_CLIENT=2 "
"DGL_CONF_PATH=path/to/part.config "
"DGL_IP_CONFIG=path/to/ip.config "
"DGL_NUM_SERVER=3 "
"DGL_GRAPH_FORMAT=csc "
"OMP_NUM_THREADS=4 "
"DGL_GROUP_ID=0 "
"PYTHONPATH=some/pythonpath/ "
),
)
# without pythonpath
self.assertEqual(
construct_dgl_client_env_vars(
num_samplers=1,
tot_num_clients=2,
part_config="path/to/part.config",
ip_config="path/to/ip.config",
num_servers=3,
graph_format="csc",
num_omp_threads=4,
group_id=0,
),
(
"DGL_DIST_MODE=distributed "
"DGL_ROLE=client "
"DGL_NUM_SAMPLER=1 "
"DGL_NUM_CLIENT=2 "
"DGL_CONF_PATH=path/to/part.config "
"DGL_IP_CONFIG=path/to/ip.config "
"DGL_NUM_SERVER=3 "
"DGL_GRAPH_FORMAT=csc "
"OMP_NUM_THREADS=4 "
"DGL_GROUP_ID=0 "
),
)
def test_submit_jobs():
class Args:
pass
args = Args()
with tempfile.TemporaryDirectory() as test_dir:
num_machines = 8
ip_config = os.path.join(test_dir, "ip_config.txt")
with open(ip_config, "w") as f:
for i in range(num_machines):
f.write("{} {}\n".format("127.0.0." + str(i), 30050))
part_config = os.path.join(test_dir, "ogb-products.json")
with open(part_config, "w") as f:
json.dump({"num_parts": num_machines}, f)
args.num_trainers = 8
args.num_samplers = 1
args.num_servers = 4
args.workspace = test_dir
args.part_config = "ogb-products.json"
args.ip_config = "ip_config.txt"
args.num_server_threads = 1
args.graph_format = "csc"
args.extra_envs = ["NCCL_DEBUG=INFO"]
args.num_omp_threads = 1
udf_command = "python3 train_dist.py --num_epochs 10"
clients_cmd, servers_cmd = submit_jobs(args, udf_command, dry_run=True)
def common_checks():
assert "cd " + test_dir in cmd
assert "export " + args.extra_envs[0] in cmd
assert f"DGL_NUM_SAMPLER={args.num_samplers}" in cmd
assert (
f"DGL_NUM_CLIENT={args.num_trainers*(args.num_samplers+1)*num_machines}"
in cmd
)
assert f"DGL_CONF_PATH={args.part_config}" in cmd
assert f"DGL_IP_CONFIG={args.ip_config}" in cmd
assert f"DGL_NUM_SERVER={args.num_servers}" in cmd
assert f"DGL_GRAPH_FORMAT={args.graph_format}" in cmd
assert f"OMP_NUM_THREADS={args.num_omp_threads}" in cmd
assert udf_command[len("python3 ") :] in cmd
for cmd in clients_cmd:
common_checks()
assert "DGL_DIST_MODE=distributed" in cmd
assert "DGL_ROLE=client" in cmd
assert "DGL_GROUP_ID=0" in cmd
assert (
f"python3 -m torch.distributed.run --nproc_per_node={args.num_trainers} --nnodes={num_machines}"
in cmd
)
assert "--master_addr=127.0.0" in cmd
assert "--master_port=1234" in cmd
for cmd in servers_cmd:
common_checks()
assert "DGL_ROLE=server" in cmd
assert "DGL_SERVER_ID=" in cmd
if __name__ == "__main__":
unittest.main()
tests/tools/test_parmetis.py
+246
View File
@@ -0,0 +1,246 @@
import argparse
import json
import os
import sys
import tempfile
import unittest
import dgl
import numpy as np
import torch
from dgl.data.utils import load_graphs, load_tensors
from partition_algo.base import load_partition_meta
from pytest_utils import create_chunked_dataset
"""
TODO: skipping this test case since the dependency, mpirun, is
not yet configured in the CI framework.
"""
@unittest.skipIf(True, reason="mpi is not available in CI test framework.")
def test_parmetis_preprocessing():
with tempfile.TemporaryDirectory() as root_dir:
num_chunks = 2
g = create_chunked_dataset(root_dir, num_chunks)
# Trigger ParMETIS pre-processing here.
input_dir = os.path.join(root_dir, "chunked-data")
results_dir = os.path.join(root_dir, "parmetis-data")
os.system(
f"mpirun -np {num_chunks} python3 tools/distpartitioning/parmetis_preprocess.py "
f"--schema {metadata.json} "
f"--input_dir {input_dir} "
f"--output_dir {results_dir} "
f"--num_parts {num_chunks}"
)
# Now add all the tests and check whether the test has passed or failed.
# Read parmetis_nfiles and ensure all files are present.
parmetis_data_dir = os.path.join(root_dir, "parmetis-data")
assert os.path.isdir(parmetis_data_dir)
parmetis_nodes_file = os.path.join(
parmetis_data_dir, "parmetis_nfiles.txt"
)
assert os.path.isfile(parmetis_nodes_file)
# `parmetis_nfiles.txt` should have each line in the following format.
# <filename> <global_id_start> <global_id_end>
with open(parmetis_nodes_file, "r") as nodes_metafile:
lines = nodes_metafile.readlines()
total_node_count = 0
for line in lines:
tokens = line.split(" ")
assert len(tokens) == 3
assert os.path.isfile(tokens[0])
assert int(tokens[1]) == total_node_count
# check contents of each of the nodes files here
with open(tokens[0], "r") as nodes_file:
node_lines = nodes_file.readlines()
for line in node_lines:
val = line.split(" ")
# <ntype_id> <weight_list> <mask_list> <type_node_id>
assert len(val) == 8
node_count = len(node_lines)
total_node_count += node_count
assert int(tokens[2]) == total_node_count
# Meta_data object.
output_dir = os.path.join(root_dir, "chunked-data")
json_file = os.path.join(output_dir, "metadata.json")
assert os.path.isfile(json_file)
with open(json_file, "rb") as f:
meta_data = json.load(f)
# Count the total no. of nodes.
true_node_count = 0
num_nodes_per_chunk = meta_data["num_nodes_per_chunk"]
for i in range(len(num_nodes_per_chunk)):
node_per_part = num_nodes_per_chunk[i]
for j in range(len(node_per_part)):
true_node_count += node_per_part[j]
assert total_node_count == true_node_count
# Read parmetis_efiles and ensure all files are present.
# This file contains a list of filenames.
parmetis_edges_file = os.path.join(
parmetis_data_dir, "parmetis_efiles.txt"
)
assert os.path.isfile(parmetis_edges_file)
with open(parmetis_edges_file, "r") as edges_metafile:
lines = edges_metafile.readlines()
total_edge_count = 0
for line in lines:
edges_filename = line.strip()
assert os.path.isfile(edges_filename)
with open(edges_filename, "r") as edges_file:
edge_lines = edges_file.readlines()
total_edge_count += len(edge_lines)
for line in edge_lines:
val = line.split(" ")
assert len(val) == 2
# Count the total no. of edges
true_edge_count = 0
num_edges_per_chunk = meta_data["num_edges_per_chunk"]
for i in range(len(num_edges_per_chunk)):
edges_per_part = num_edges_per_chunk[i]
for j in range(len(edges_per_part)):
true_edge_count += edges_per_part[j]
assert true_edge_count == total_edge_count
def test_parmetis_postprocessing():
with tempfile.TemporaryDirectory() as root_dir:
num_chunks = 2
g = create_chunked_dataset(root_dir, num_chunks)
num_nodes = g.num_nodes()
num_institutions = g.num_nodes("institution")
num_authors = g.num_nodes("author")
num_papers = g.num_nodes("paper")
# Generate random parmetis partition ids for the nodes in the graph.
# Replace this code with actual ParMETIS executable when it is ready
output_dir = os.path.join(root_dir, "chunked-data")
assert os.path.isdir(output_dir)
parmetis_file = os.path.join(output_dir, "parmetis_output.txt")
node_ids = np.arange(num_nodes)
partition_ids = np.random.randint(0, 2, (num_nodes,))
parmetis_output = np.column_stack([node_ids, partition_ids])
# Create parmetis output, this is mimicking running actual parmetis.
with open(parmetis_file, "w") as f:
np.savetxt(f, parmetis_output)
assert os.path.isfile(parmetis_file)
# Check the post processing script here.
results_dir = os.path.join(output_dir, "partitions_dir")
json_file = os.path.join(output_dir, "metadata.json")
print(json_file)
print(results_dir)
print(parmetis_file)
os.system(
f"python3 tools/distpartitioning/parmetis_postprocess.py "
f"--postproc_input_dir {output_dir} "
f"--schema_file metadata.json "
f"--parmetis_output_file {parmetis_file} "
f"--partitions_dir {results_dir}"
)
ntype_count = {
"author": num_authors,
"paper": num_papers,
"institution": num_institutions,
}
for ntype_name in ["author", "paper", "institution"]:
fname = os.path.join(results_dir, f"{ntype_name}.txt")
print(fname)
assert os.path.isfile(fname)
# Load and check the partition ids in this file.
part_ids = np.loadtxt(fname)
assert part_ids.shape[0] == ntype_count[ntype_name]
assert np.min(part_ids) == 0
assert np.max(part_ids) == 1
# check partition meta file
part_meta_file = os.path.join(results_dir, "partition_meta.json")
assert os.path.isfile(part_meta_file)
part_meta = load_partition_meta(part_meta_file)
assert part_meta.num_parts == 2
assert part_meta.algo_name == "metis"
"""
TODO: skipping this test case since it depends on the dependency, mpi,
which is not yet configured in the CI framework.
"""
@unittest.skipIf(True, reason="mpi is not available in CI test framework.")
def test_parmetis_wrapper():
with tempfile.TemporaryDirectory() as root_dir:
num_chunks = 2
graph_name = "mag240m"
g = create_chunked_dataset(root_dir, num_chunks)
all_ntypes = g.ntypes
all_etypes = g.etypes
num_constraints = len(all_ntypes) + 3
num_institutions = g.num_nodes("institution")
num_authors = g.num_nodes("author")
num_papers = g.num_nodes("paper")
# Trigger ParMETIS.
schema_file = os.path.join(root_dir, "chunked-data/metadata.json")
preproc_input_dir = os.path.join(root_dir, "chunked-data")
preproc_output_dir = os.path.join(
root_dir, "chunked-data/preproc_output_dir"
)
parmetis_output_file = os.path.join(
os.getcwd(), f"{graph_name}_part.{num_chunks}"
)
partitions_dir = os.path.join(root_dir, "chunked-data/partitions_dir")
hostfile = os.path.join(root_dir, "ip_config.txt")
with open(hostfile, "w") as f:
f.write("127.0.0.1\n")
f.write("127.0.0.1\n")
num_nodes = g.num_nodes()
num_edges = g.num_edges()
stats_file = f"{graph_name}_stats.txt"
with open(stats_file, "w") as f:
f.write(f"{num_nodes} {num_edges} {num_constraints}")
os.system(
f"python3 tools/distpartitioning/parmetis_wrapper.py "
f"--schema_file {schema_file} "
f"--preproc_input_dir {preproc_input_dir} "
f"--preproc_output_dir {preproc_output_dir} "
f"--hostfile {hostfile} "
f"--num_parts {num_chunks} "
f"--parmetis_output_file {parmetis_output_file} "
f"--partitions_dir {partitions_dir} "
)
print("Executing Done.")
ntype_count = {
"author": num_authors,
"paper": num_papers,
"institution": num_institutions,
}
for ntype_name in ["author", "paper", "institution"]:
fname = os.path.join(partitions_dir, f"{ntype_name}.txt")
print(fname)
assert os.path.isfile(fname)
# Load and check the partition ids in this file.
part_ids = np.loadtxt(fname)
assert part_ids.shape[0] == ntype_count[ntype_name]
assert np.min(part_ids) == 0
assert np.max(part_ids) == (num_chunks - 1)
tests/tools/test_parmetis_preproc.py
+177
View File
@@ -0,0 +1,177 @@
import os
import tempfile
from collections import namedtuple
import numpy as np
import pytest
from distpartitioning import array_readwriter, constants
from distpartitioning.parmetis_preprocess import gen_edge_files
from distpartitioning.utils import generate_roundrobin_read_list
from numpy.testing import assert_array_equal
NODE_TYPE = "n1"
EDGE_TYPE = f"{NODE_TYPE}:e1:{NODE_TYPE}"
def _read_file(fname, fmt_name, fmt_delimiter):
"""Read a file
Parameters:
-----------
fname : string
filename of the input file to read
fmt_name : string
specifying whether it is a csv or a parquet file
fmt_delimiter : string
string specifying the delimiter used in the input file
"""
reader_fmt_meta = {
"name": fmt_name,
}
if fmt_name == constants.STR_CSV:
reader_fmt_meta["delimiter"] = fmt_delimiter
data_df = array_readwriter.get_array_parser(**reader_fmt_meta).read(fname)
return data_df
def _get_test_data(edges_dir, num_chunks, edge_fmt, edge_fmt_del):
"""Creates unit test input which are a set of edge files
in the following format "src_node_id<delimiter>dst_node_id"
Parameters:
-----------
edges_dir : str
folder where edge files are stored
num_chunks : int
no. of files to create for each edge type
edge_fmt : str, optional
to specify whether this file is csv or parquet
edge_fmt_del : str optional
delimiter to use in the edges file
Returns:
--------
dict :
dictionary created which represents the schema used for
creating the input dataset
"""
schema = {}
schema["num_nodes_per_type"] = [10]
schema["edge_type"] = [EDGE_TYPE]
schema["node_type"] = [NODE_TYPE]
edges = {}
edges[EDGE_TYPE] = {}
edges[EDGE_TYPE]["format"] = {}
edges[EDGE_TYPE]["format"]["name"] = edge_fmt
edges[EDGE_TYPE]["format"]["delimiter"] = edge_fmt_del
os.makedirs(edges_dir, exist_ok=True)
fmt_meta = {"name": edge_fmt}
if edge_fmt == "csv":
fmt_meta["delimiter"] = edge_fmt_del
edge_files = []
for idx in range(num_chunks):
path = os.path.join(edges_dir, f"test_file_{idx}.{fmt_meta['name']}")
array_parser = array_readwriter.get_array_parser(**fmt_meta)
edge_data = (
np.array([np.arange(10), np.arange(10)]).reshape(10, 2) + 10 * idx
)
array_parser.write(path, edge_data)
edge_files.append(path)
edges[EDGE_TYPE]["data"] = edge_files
schema["edges"] = edges
return schema
@pytest.mark.parametrize("num_chunks, num_parts", [[4, 1], [4, 2], [4, 4]])
@pytest.mark.parametrize("edges_fmt", ["csv", "parquet"])
@pytest.mark.parametrize("edges_delimiter", [" ", ","])
def test_gen_edge_files(num_chunks, num_parts, edges_fmt, edges_delimiter):
"""Unit test case for the function
tools/distpartitioning/parmetis_preprocess.py::gen_edge_files
Parameters:
-----------
num_chunks : int
no. of chunks the input graph needs to be partititioned into
num_parts : int
no. of partitions
edges_fmt : string
specifying the storage format for the edge files
edges_delimiter : string
specifying the delimiter used in the edge files
"""
# Create the input dataset
with tempfile.TemporaryDirectory() as root_dir:
# Create expected environment for test
input_dir = os.path.join(root_dir, "chunked-data")
output_dir = os.path.join(root_dir, "preproc_dir")
# Mock a parser object
fn_params = namedtuple("fn_params", "input_dir output_dir num_parts")
fn_params.input_dir = input_dir
fn_params.output_dir = output_dir
fn_params.num_parts = num_parts
# Create test files and get corresponding file schema
schema_map = _get_test_data(
input_dir, num_chunks, edges_fmt, edges_delimiter
)
edges_file_list = schema_map["edges"][EDGE_TYPE]["data"]
# This is breaking encapsulation, but no other good way to get file list
rank_assignments = generate_roundrobin_read_list(
len(edges_file_list), num_parts
)
# Get the global node id offsets for each node type
# There is only one node-type in the test graph
# which range from 0 thru 9.
ntype_gnid_offset = {}
ntype_gnid_offset[NODE_TYPE] = np.array([0, 10 * num_chunks]).reshape(
1, 2
)
# Iterate over no. of partitions
for rank in range(num_parts):
actual_results = gen_edge_files(rank, schema_map, fn_params)
# Get the original files
original_files = [
edges_file_list[file_idx] for file_idx in rank_assignments[rank]
]
# Validate the results with the baseline results
# Test 1. no. of files should have the same count per rank
assert len(original_files) == len(actual_results)
assert len(actual_results) > 0
# Test 2. Check the contents of each file and verify the
# file contents match with the expected results.
for actual_fname, original_fname in zip(
actual_results, original_files
):
# Check the actual file exists
assert os.path.isfile(actual_fname)
# Read both files and compare the edges
# Here note that the src and dst end points are global_node_ids
actual_data = _read_file(actual_fname, "csv", " ")
expected_data = _read_file(
original_fname, edges_fmt, edges_delimiter
)
# Subtract the global node id offsets, so that we get type node ids
# In the current unit test case, the graph has only one node-type.
# and this means that type-node-ids are same as the global-node-ids.
# Below two lines will take take into effect when the graphs have
# more than one node type.
actual_data[:, 0] -= ntype_gnid_offset[NODE_TYPE][0, 0]
actual_data[:, 1] -= ntype_gnid_offset[NODE_TYPE][0, 0]
# Verify that the contents are equal
assert_array_equal(expected_data, actual_data)