Function test_lookup_service

tests/tools/test_dist_lookup.py:202–259

(
    num_chunks,
    num_parts,
    world_size,
    num_chunks_nodes=None,
    num_chunks_edges=None,
    num_chunks_node_data=None,
    num_chunks_edge_data=None,
)

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