wehub-resource-sync
247 lines
9.3 KiB
Python
247 lines
9.3 KiB
Python
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)
|