parmetis_preprocess.py 14.6 KB
Newer Older
1
2
3
import argparse
import logging
import os
4
import platform
5
import sys
6
from datetime import timedelta
7
from pathlib import Path
8
from timeit import default_timer as timer
9

10
11
import array_readwriter

12
13
import constants

14
15
16
import numpy as np
import pyarrow
import pyarrow.csv as csv
17
import pyarrow.parquet as pq
18
import torch
19
from utils import generate_read_list, get_idranges, get_node_types, read_json
20
21
22
23
24
25
26


def get_proc_info():
    """Helper function to get the rank from the
    environment when `mpirun` is used to run this python program.

    Please note that for mpi(openmpi) installation the rank is retrieved from the
27
28
    environment using OMPI_COMM_WORLD_RANK. For mpich it is
    retrieved from the environment using PMI_RANK.
29

30
31
32
33
34
35
    Returns:
    --------
    integer :
        Rank of the current process.
    """
    env_variables = dict(os.environ)
36
37
38
    # mpich
    if "PMI_RANK" in env_variables:
        return int(env_variables["PMI_RANK"])
39
    # openmpi
40
41
42
43
    elif "OMPI_COMM_WORLD_RANK" in env_variables:
        return int(env_variables["OMPI_COMM_WORLD_RANK"])
    else:
        return 0
44

45

46
def gen_edge_files(schema_map, params):
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    """Function to create edges files to be consumed by ParMETIS
    for partitioning purposes.

    This function creates the edge files and each of these will have the
    following format (meaning each line of these file is of the following format)
    <global_src_id> <global_dst_id>

    Here ``global`` prefix means that globally unique identifier assigned each node
    in the input graph. In this context globally unique means unique across all the
    nodes in the input graph.

    Parameters:
    -----------
    schema_map : json dictionary
        Dictionary created by reading the metadata.json file for the input dataset.
    output : string
        Location of storing the node-weights and edge files for ParMETIS.
    """
    rank = get_proc_info()
    type_nid_dict, ntype_gnid_offset = get_idranges(
        schema_map[constants.STR_NODE_TYPE],
68
69
70
71
72
73
        dict(
            zip(
                schema_map[constants.STR_NODE_TYPE],
                schema_map[constants.STR_NUM_NODES_PER_TYPE],
            )
        ),
74
75
76
77
78
79
80
    )

    # Regenerate edge files here.
    edge_data = schema_map[constants.STR_EDGES]
    etype_names = schema_map[constants.STR_EDGE_TYPE]
    etype_name_idmap = {e: idx for idx, e in enumerate(etype_names)}

81
    outdir = Path(params.output_dir)
82
83
    os.makedirs(outdir, exist_ok=True)
    edge_files = []
84
    num_parts = params.num_parts
85
    for etype_name, etype_info in edge_data.items():
86
87
        edges_format = etype_info[constants.STR_FORMAT][constants.STR_NAME]
        edge_data_files = etype_info[constants.STR_DATA]
88
89
90
91
92
93
94
95
96

        # ``edgetype`` strings are in canonical format, src_node_type:edge_type:dst_node_type
        tokens = etype_name.split(":")
        assert len(tokens) == 3

        src_ntype_name = tokens[0]
        rel_name = tokens[1]
        dst_ntype_name = tokens[2]

97
98
99
        def process_and_write_back(data_df, idx):
            data_f0 = data_df[:, 0]
            data_f1 = data_df[:, 1]
100

101
102
103
104
            global_src_id = data_f0 + ntype_gnid_offset[src_ntype_name][0, 0]
            global_dst_id = data_f1 + ntype_gnid_offset[dst_ntype_name][0, 0]
            cols = [global_src_id, global_dst_id]
            col_names = ["global_src_id", "global_dst_id"]
105

106
            out_file = edge_data_files[idx].split("/")[-1]
107
            out_file = os.path.join(outdir, "edges_{}".format(out_file))
108

109
110
111
112
            # TODO(thvasilo): We should support writing to the same format as the input
            options = csv.WriteOptions(include_header=False, delimiter=" ")
            options.delimiter = " "
            csv.write_csv(
113
114
115
                pyarrow.Table.from_arrays(cols, names=col_names),
                out_file,
                options,
116
117
118
            )
            return out_file

119
120
121
122
123
124
125
126
        # handle any no. of files case here.
        file_idxes = generate_read_list(len(edge_data_files), params.num_parts)
        for idx in file_idxes[rank]:
            reader_fmt_meta = {
                "name": etype_info[constants.STR_FORMAT][constants.STR_NAME]
            }
            data_df = array_readwriter.get_array_parser(**reader_fmt_meta).read(
                edge_data_files[idx]
127
            )
128
129
            out_file = process_and_write_back(data_df, idx)
            edge_files.append(out_file)
130
131
132
133

    return edge_files


134
def gen_node_weights_files(schema_map, params):
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
    """Function to create node weight files for ParMETIS along with the edge files.

    This function generates node-data files, which will be read by the ParMETIS
    executable for partitioning purposes. Each line in these files will be of the
    following format:
        <node_type_id> <node_weight_list> <type_wise_node_id>
    node_type_id -  is id assigned to the node-type to which a given particular
        node belongs to
    weight_list - this is a one-hot vector in which the number in the location of
        the current nodes' node-type will be set to `1` and other will be `0`
    type_node_id - this is the id assigned to the node (in the context of the current
        nodes` node-type). Meaning this id is unique across all the nodes which belong to
        the current nodes` node-type.

    Parameters:
    -----------
    schema_map : json dictionary
        Dictionary created by reading the metadata.json file for the input dataset.
    output : string
        Location of storing the node-weights and edge files for ParMETIS.

    Returns:
    --------
    list :
        List of filenames for nodes of the input graph.
    list :
        List o ffilenames for edges of the input graph.
    """
    rank = get_proc_info()
    ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(schema_map)
    type_nid_dict, ntype_gnid_offset = get_idranges(
        schema_map[constants.STR_NODE_TYPE],
167
168
169
170
171
172
        dict(
            zip(
                schema_map[constants.STR_NODE_TYPE],
                schema_map[constants.STR_NUM_NODES_PER_TYPE],
            )
        ),
173
174
175
    )

    node_files = []
176
    outdir = Path(params.output_dir)
177
178
179
    os.makedirs(outdir, exist_ok=True)

    for ntype_id, ntype_name in ntid_ntype_map.items():
180
181
182
183
184
185

        # This ntype does not have any train/test/val masks...
        # Each rank will generate equal no. of rows for this node type.
        total_count = schema_map[constants.STR_NUM_NODES_PER_TYPE][ntype_id]
        per_rank_range = np.ones((params.num_parts,), dtype=np.int64) * (
            total_count // params.num_parts
186
        )
187
188
189
190
191
192
193
194
        for i in range(total_count % params.num_parts):
            per_rank_range[i] += 1

        tid_start = np.cumsum([0] + list(per_rank_range[:-1]))
        tid_end = np.cumsum(list(per_rank_range))
        local_tid_start = tid_start[rank]
        local_tid_end = tid_end[rank]
        sz = local_tid_end - local_tid_start
195
196
197

        cols = []
        col_names = []
198
199

        # ntype-id
200
201
202
203
204
        cols.append(
            pyarrow.array(np.ones(sz, dtype=np.int64) * np.int64(ntype_id))
        )
        col_names.append("ntype")

205
        # one-hot vector for ntype-id here.
206
207
208
209
210
211
212
213
214
215
        for i in range(len(ntypes)):
            if i == ntype_id:
                cols.append(pyarrow.array(np.ones(sz, dtype=np.int64)))
            else:
                cols.append(pyarrow.array(np.zeros(sz, dtype=np.int64)))
            col_names.append("w{}".format(i))

        # `type_nid` should be the very last column in the node weights files.
        cols.append(
            pyarrow.array(
216
                np.arange(local_tid_start, local_tid_end, dtype=np.int64)
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
            )
        )
        col_names.append("type_nid")

        out_file = os.path.join(
            outdir, "node_weights_{}_{}.txt".format(ntype_name, rank)
        )
        options = csv.WriteOptions(include_header=False, delimiter=" ")
        options.delimiter = " "

        csv.write_csv(
            pyarrow.Table.from_arrays(cols, names=col_names), out_file, options
        )
        node_files.append(
            (
232
233
                ntype_gnid_offset[ntype_name][0, 0] + local_tid_start,
                ntype_gnid_offset[ntype_name][0, 0] + local_tid_end,
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
                out_file,
            )
        )

    return node_files


def gen_parmetis_input_args(params, schema_map):
    """Function to create two input arguments which will be passed to the parmetis.
    first argument is a text file which has a list of node-weights files,
    namely parmetis-nfiles.txt, and second argument is a text file which has a
    list of edge files, namely parmetis_efiles.txt.
    ParMETIS uses these two files to read/load the graph and partition the graph
    With regards to the file format, parmetis_nfiles.txt uses the following format
    for each line in that file:
        <filename> <global_node_id_start> <global_node_id_end>(exclusive)
    While parmetis_efiles.txt just has <filename> in each line.

    Parameters:
    -----------
    params : argparser instance
        Instance of ArgParser class, which has all the input arguments passed to
        run this program.
    schema_map : json dictionary
        Dictionary object created after reading the graph metadata.json file.
    """

261
    # TODO: This makes the assumption that all node files have the same number of chunks
262
263
264
    ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(schema_map)
    type_nid_dict, ntype_gnid_offset = get_idranges(
        schema_map[constants.STR_NODE_TYPE],
265
266
267
268
269
270
        dict(
            zip(
                schema_map[constants.STR_NODE_TYPE],
                schema_map[constants.STR_NUM_NODES_PER_TYPE],
            )
        ),
271
272
    )

273
    # Check if <graph-name>_stats.txt exists, if not create one using metadata.
274
    # Here stats file will be created in the current directory.
275
276
277
278
    # No. of constraints, third column in the stats file is computed as follows:
    #   num_constraints = no. of node types + train_mask + test_mask + val_mask
    #   Here, (train/test/val) masks will be set to 1 if these masks exist for
    #   all the node types in the graph, otherwise these flags will be set to 0
279
280
281
    assert (
        constants.STR_GRAPH_NAME in schema_map
    ), "Graph name is not present in the json file"
282
    graph_name = schema_map[constants.STR_GRAPH_NAME]
283
    if not os.path.isfile(f"{graph_name}_stats.txt"):
284
285
        num_nodes = np.sum(schema_map[constants.STR_NUM_NODES_PER_TYPE])
        num_edges = np.sum(schema_map[constants.STR_NUM_EDGES_PER_TYPE])
286
287
        num_ntypes = len(schema_map[constants.STR_NODE_TYPE])

288
        num_constraints = num_ntypes
289

290
291
        with open(f"{graph_name}_stats.txt", "w") as sf:
            sf.write(f"{num_nodes} {num_edges} {num_constraints}")
292

293
294
295
296
297
    node_files = []
    outdir = Path(params.output_dir)
    os.makedirs(outdir, exist_ok=True)
    for ntype_id, ntype_name in ntid_ntype_map.items():
        global_nid_offset = ntype_gnid_offset[ntype_name][0, 0]
298
299
300
301
302
303
304
305
306
307
308
309
310
311
        total_count = schema_map[constants.STR_NUM_NODES_PER_TYPE][ntype_id]
        per_rank_range = np.ones((params.num_parts,), dtype=np.int64) * (
            total_count // params.num_parts
        )
        for i in range(total_count % params.num_parts):
            per_rank_range[i] += 1
        tid_start = np.cumsum([0] + list(per_rank_range[:-1]))
        tid_end = np.cumsum(per_rank_range)
        logging.info(f" tid-start = {tid_start}, tid-end = {tid_end}")
        logging.info(f" per_rank_range - {per_rank_range}")

        for rank in range(params.num_parts):
            local_tid_start = tid_start[rank]
            local_tid_end = tid_end[rank]
312
            out_file = os.path.join(
313
                outdir, "node_weights_{}_{}.txt".format(ntype_name, rank)
314
315
316
317
            )
            node_files.append(
                (
                    out_file,
318
319
                    global_nid_offset + local_tid_start,
                    global_nid_offset + local_tid_end,
320
321
322
323
324
325
326
327
328
329
330
331
332
                )
            )

    nfile = open(os.path.join(params.output_dir, "parmetis_nfiles.txt"), "w")
    for f in node_files:
        # format: filename global_node_id_start global_node_id_end(exclusive)
        nfile.write("{} {} {}\n".format(f[0], f[1], f[2]))
    nfile.close()

    # Regenerate edge files here.
    edge_data = schema_map[constants.STR_EDGES]
    edge_files = []
    for etype_name, etype_info in edge_data.items():
333
334
335
        edge_data_files = etype_info[constants.STR_DATA]
        for edge_file_path in edge_data_files:
            out_file = os.path.basename(edge_file_path)
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
            out_file = os.path.join(outdir, "edges_{}".format(out_file))
            edge_files.append(out_file)

    with open(
        os.path.join(params.output_dir, "parmetis_efiles.txt"), "w"
    ) as efile:
        for f in edge_files:
            efile.write("{}\n".format(f))


def run_preprocess_data(params):
    """Main function which will help create graph files for ParMETIS processing

    Parameters:
    -----------
    params : argparser object
        An instance of argparser class which stores command line arguments.
    """
    logging.info(f"Starting to generate ParMETIS files...")
    rank = get_proc_info()
    schema_map = read_json(params.schema_file)
357
    gen_node_weights_files(schema_map, params)
358
359
    logging.info(f"Done with node weights....")

360
    gen_edge_files(schema_map, params)
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
    logging.info(f"Done with edge weights...")

    if rank == 0:
        gen_parmetis_input_args(params, schema_map)
    logging.info(f"Done generating files for ParMETIS run ..")


if __name__ == "__main__":
    """Main function used to generate temporary files needed for ParMETIS execution.
    This function generates node-weight files and edges files which are consumed by ParMETIS.

    Example usage:
    --------------
    mpirun -np 4 python3 parmetis_preprocess.py --schema <file> --output <target-output-dir>
    """
    parser = argparse.ArgumentParser(
        description="Generate ParMETIS files for input dataset"
    )
    parser.add_argument(
        "--schema_file",
        required=True,
        type=str,
        help="The schema of the input graph",
    )
385
386
    parser.add_argument(
        "--input_dir",
387
        required=True,
388
        type=str,
389
        help="This directory will be used as the relative directory to locate files, if absolute paths are not used",
390
    )
391
392
393
394
395
396
    parser.add_argument(
        "--output_dir",
        required=True,
        type=str,
        help="The output directory for the node weights files and auxiliary files for ParMETIS.",
    )
397
398
399
400
401
402
    parser.add_argument(
        "--num_parts",
        required=True,
        type=int,
        help="Total no. of output graph partitions.",
    )
403
404
    params = parser.parse_args()

405
406
407
408
409
410
411
    # Configure logging.
    logging.basicConfig(
        level="INFO",
        format=f"[{platform.node()} \
        %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
    )

412
413
    # Invoke the function to generate files for parmetis
    run_preprocess_data(params)