wehub-resource-sync
260 lines
7.4 KiB
Python
260 lines
7.4 KiB
Python
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
|
|
),
|
|
)
|