test_dist_lookup.py

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
            ),
        )