utils.py 19.9 KB
Newer Older
1
import json
2
3
import logging
import os
4

5
import constants
6
7

import dgl
8
9
import numpy as np
import psutil
10
import pyarrow
11
from dgl.distributed.partition import _dump_part_config
12
from pyarrow import csv
13

14

15
def read_ntype_partition_files(schema_map, input_dir):
16
    """
17
18
    Utility method to read the partition id mapping for each node.
    For each node type, there will be an file, in the input directory argument
19
    containing the partition id mapping for a given nodeid.
20
21
22

    Parameters:
    -----------
23
24
25
    schema_map : dictionary
        dictionary created by reading the input metadata json file
    input_dir : string
26
        directory in which the node-id to partition-id mappings files are
27
        located for each of the node types in the input graph
28
29
30

    Returns:
    --------
31
    numpy array :
32
        array of integers representing mapped partition-ids for a given node-id.
33
34
35
36
37
        The line number, in these files, are used as the type_node_id in each of
        the files. The index into this array will be the homogenized node-id and
        value will be the partition-id for that node-id (index). Please note that
        the partition-ids of each node-type are stacked together vertically and
        in this way heterogenous node-ids are converted to homogenous node-ids.
38
    """
39
40
    assert os.path.isdir(input_dir)

41
    # iterate over the node types and extract the partition id mappings
42
43
44
    part_ids = []
    ntype_names = schema_map[constants.STR_NODE_TYPE]
    for ntype in ntype_names:
45
46
47
48
49
50
51
52
        df = csv.read_csv(
            os.path.join(input_dir, "{}.txt".format(ntype)),
            read_options=pyarrow.csv.ReadOptions(
                autogenerate_column_names=True
            ),
            parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
        )
        ntype_partids = df["f0"].to_numpy()
53
54
        part_ids.append(ntype_partids)
    return np.concatenate(part_ids)
55

56

57
58
59
def read_json(json_file):
    """
    Utility method to read a json file schema
60

61
62
63
64
65
66
67
68
69
70
71
72
73
74
    Parameters:
    -----------
    json_file : string
        file name for the json schema

    Returns:
    --------
    dictionary, as serialized in the json_file
    """
    with open(json_file) as schema:
        val = json.load(schema)

    return val

75

76
77
78
79
80
81
82
83
84
85
86
87
88
89
def get_etype_featnames(etype_name, schema_map):
    """Retrieves edge feature names for a given edge_type

    Parameters:
    -----------
    eype_name : string
        a string specifying a edge_type name

    schema : dictionary
        metadata json object as a dictionary, which is read from the input
        metadata file from the input dataset

    Returns:
    --------
90
    list :
91
92
93
94
95
96
        a list of feature names for a given edge_type
    """
    edge_data = schema_map[constants.STR_EDGE_DATA]
    feats = edge_data.get(etype_name, {})
    return [feat for feat in feats]

97
98

def get_ntype_featnames(ntype_name, schema_map):
99
100
101
102
103
104
105
106
107
108
109
110
111
112
    """
    Retrieves node feature names for a given node_type

    Parameters:
    -----------
    ntype_name : string
        a string specifying a node_type name

    schema : dictionary
        metadata json object as a dictionary, which is read from the input
        metadata file from the input dataset

    Returns:
    --------
113
    list :
114
115
        a list of feature names for a given node_type
    """
116
117
118
    node_data = schema_map[constants.STR_NODE_DATA]
    feats = node_data.get(ntype_name, {})
    return [feat for feat in feats]
119

120

121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def get_edge_types(schema_map):
    """Utility method to extract edge_typename -> edge_type mappings
    as defined by the input schema

    Parameters:
    -----------
    schema_map : dictionary
        Input schema from which the edge_typename -> edge_typeid
        dictionary is created.

    Returns:
    --------
    dictionary
        with keys as edge type names and values as ids (integers)
    list
        list of etype name strings
    dictionary
        with keys as etype ids (integers) and values as edge type names
    """
    etypes = schema_map[constants.STR_EDGE_TYPE]
141
142
    etype_etypeid_map = {e: i for i, e in enumerate(etypes)}
    etypeid_etype_map = {i: e for i, e in enumerate(etypes)}
143
144
    return etype_etypeid_map, etypes, etypeid_etype_map

145

146
def get_node_types(schema_map):
147
    """
148
149
150
151
152
    Utility method to extract node_typename -> node_type mappings
    as defined by the input schema

    Parameters:
    -----------
153
    schema_map : dictionary
154
155
156
157
158
        Input schema from which the node_typename -> node_type
        dictionary is created.

    Returns:
    --------
159
160
161
162
163
164
    dictionary
        with keys as node type names and values as ids (integers)
    list
        list of ntype name strings
    dictionary
        with keys as ntype ids (integers) and values as node type names
165
    """
166
    ntypes = schema_map[constants.STR_NODE_TYPE]
167
168
    ntype_ntypeid_map = {e: i for i, e in enumerate(ntypes)}
    ntypeid_ntype_map = {i: e for i, e in enumerate(ntypes)}
169
170
    return ntype_ntypeid_map, ntypes, ntypeid_ntype_map

171
172

def get_gnid_range_map(node_tids):
173
    """
174
    Retrieves auxiliary dictionaries from the metadata json object
175
176
177

    Parameters:
    -----------
178
179
    node_tids: dictionary
        This dictionary contains the information about nodes for each node_type.
180
        Typically this information contains p-entries, where each entry has a file-name,
181
182
183
        starting and ending type_node_ids for the nodes in this file. Keys in this dictionary
        are the node_type and value is a list of lists. Each individual entry in this list has
        three items: file-name, starting type_nid and ending type_nid
184
185
186

    Returns:
    --------
187
    dictionary :
188
189
        a dictionary where keys are node_type names and values are global_nid range, which is a tuple.

190
    """
191
    ntypes_gid_range = {}
192
    offset = 0
193
    for k, v in node_tids.items():
194
195
196
197
        ntypes_gid_range[k] = [offset + int(v[0][0]), offset + int(v[-1][1])]
        offset += int(v[-1][1])

    return ntypes_gid_range
198

199
200
201
202

def write_metadata_json(
    input_list, output_dir, graph_name, world_size, num_parts
):
203
    """
204
    Merge json schema's from each of the rank's on rank-0.
205
206
207
208
209
210
211
212
213
214
215
    This utility function, to be used on rank-0, to create aggregated json file.

    Parameters:
    -----------
    metadata_list : list of json (dictionaries)
        a list of json dictionaries to merge on rank-0
    output_dir : string
        output directory path in which results are stored (as a json file)
    graph-name : string
        a string specifying the graph name
    """
216
    # Preprocess the input_list, a list of dictionaries
217
    # each dictionary will contain num_parts/world_size metadata json
218
219
    # which correspond to local partitions on the respective ranks.
    metadata_list = []
220
    for local_part_id in range(num_parts // world_size):
221
        for idx in range(world_size):
222
223
224
225
226
            metadata_list.append(
                input_list[idx][
                    "local-part-id-" + str(local_part_id * world_size + idx)
                ]
            )
227

228
    # Initialize global metadata
229
230
    graph_metadata = {}

231
    # Merge global_edge_ids from each json object in the input list
232
233
234
235
236
    edge_map = {}
    x = metadata_list[0]["edge_map"]
    for k in x:
        edge_map[k] = []
        for idx in range(len(metadata_list)):
237
238
239
240
241
242
            edge_map[k].append(
                [
                    int(metadata_list[idx]["edge_map"][k][0][0]),
                    int(metadata_list[idx]["edge_map"][k][0][1]),
                ]
            )
243
244
245
246
247
248
    graph_metadata["edge_map"] = edge_map

    graph_metadata["etypes"] = metadata_list[0]["etypes"]
    graph_metadata["graph_name"] = metadata_list[0]["graph_name"]
    graph_metadata["halo_hops"] = metadata_list[0]["halo_hops"]

249
    # Merge global_nodeids from each of json object in the input list
250
251
252
253
254
    node_map = {}
    x = metadata_list[0]["node_map"]
    for k in x:
        node_map[k] = []
        for idx in range(len(metadata_list)):
255
256
257
258
259
260
            node_map[k].append(
                [
                    int(metadata_list[idx]["node_map"][k][0][0]),
                    int(metadata_list[idx]["node_map"][k][0][1]),
                ]
            )
261
262
263
    graph_metadata["node_map"] = node_map

    graph_metadata["ntypes"] = metadata_list[0]["ntypes"]
264
265
266
267
268
269
    graph_metadata["num_edges"] = int(
        sum([metadata_list[i]["num_edges"] for i in range(len(metadata_list))])
    )
    graph_metadata["num_nodes"] = int(
        sum([metadata_list[i]["num_nodes"] for i in range(len(metadata_list))])
    )
270
271
272
273
    graph_metadata["num_parts"] = metadata_list[0]["num_parts"]
    graph_metadata["part_method"] = metadata_list[0]["part_method"]

    for i in range(len(metadata_list)):
274
275
276
277
278
        graph_metadata["part-{}".format(i)] = metadata_list[i][
            "part-{}".format(i)
        ]

    _dump_part_config(f"{output_dir}/metadata.json", graph_metadata)
279
280


281
282
283
def augment_edge_data(
    edge_data, lookup_service, edge_tids, rank, world_size, num_parts
):
284
285
    """
    Add partition-id (rank which owns an edge) column to the edge_data.
286

287
288
289
290
    Parameters:
    -----------
    edge_data : numpy ndarray
        Edge information as read from the xxx_edges.txt file
291
292
293
294
295
296
297
298
299
300
    lookup_service : instance of class DistLookupService
       Distributed lookup service used to map global-nids to respective partition-ids and▒
       shuffle-global-nids
    edge_tids: dictionary
        dictionary where keys are canonical edge types and values are list of tuples
        which indicate the range of edges assigned to each of the partitions
    rank : integer
        rank of the current process
    world_size : integer
        total no. of process participating in the communication primitives
301
302
    num_parts : integer
        total no. of partitions requested for the input graph
303
304
305

    Returns:
    --------
306
307
    dictionary :
        dictionary with keys as column names and values as numpy arrays and this information is
308
309
        loaded from input dataset files. In addition to this we include additional columns which
        aid this pipelines computation, like constants.OWNER_PROCESS
310
    """
311
    # add global_nids to the node_data
312
313
    etype_offset = {}
    offset = 0
314
    for etype_name, tid_range in edge_tids.items():
315
        assert int(tid_range[0][0]) == 0
316
        assert len(tid_range) == num_parts
317
318
319
320
        etype_offset[etype_name] = offset + int(tid_range[0][0])
        offset += int(tid_range[-1][1])

    global_eids = []
321
    for etype_name, tid_range in edge_tids.items():
322
323
324
325
326
327
        for idx in range(num_parts):
            if map_partid_rank(idx, world_size) == rank:
                global_eid_start = etype_offset[etype_name]
                begin = global_eid_start + int(tid_range[idx][0])
                end = global_eid_start + int(tid_range[idx][1])
                global_eids.append(np.arange(begin, end, dtype=np.int64))
328
329
    global_eids = np.concatenate(global_eids)
    assert global_eids.shape[0] == edge_data[constants.ETYPE_ID].shape[0]
330
331
    edge_data[constants.GLOBAL_EID] = global_eids

332
    return edge_data
333

334

335
def read_edges_file(edge_file, edge_data_dict):
336
    """
337
338
339
340
341
342
343
344
345
346
347
    Utility function to read xxx_edges.txt file

    Parameters:
    -----------
    edge_file : string
        Graph file for edges in the input graph

    Returns:
    --------
    dictionary
        edge data as read from xxx_edges.txt file and columns are stored
348
        in a dictionary with key-value pairs as column-names and column-data.
349
350
351
352
    """
    if edge_file == "" or edge_file == None:
        return None

353
354
    # Read the file from here.
    # <global_src_id> <global_dst_id> <type_eid> <etype> <attributes>
355
356
357
    # global_src_id -- global idx for the source node ... line # in the graph_nodes.txt
    # global_dst_id -- global idx for the destination id node ... line # in the graph_nodes.txt

358
359
360
361
362
    edge_data_df = csv.read_csv(
        edge_file,
        read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
        parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
    )
363
    edge_data_dict = {}
364
365
366
367
    edge_data_dict[constants.GLOBAL_SRC_ID] = edge_data_df["f0"].to_numpy()
    edge_data_dict[constants.GLOBAL_DST_ID] = edge_data_df["f1"].to_numpy()
    edge_data_dict[constants.GLOBAL_TYPE_EID] = edge_data_df["f2"].to_numpy()
    edge_data_dict[constants.ETYPE_ID] = edge_data_df["f3"].to_numpy()
368
369
    return edge_data_dict

370

371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
def read_node_features_file(nodes_features_file):
    """
    Utility function to load tensors from a file

    Parameters:
    -----------
    nodes_features_file : string
        Features file for nodes in the graph

    Returns:
    --------
    dictionary
        mappings between ntype and list of features
    """

    node_features = dgl.data.utils.load_tensors(nodes_features_file, False)
    return node_features

389

390
def read_edge_features_file(edge_features_file):
391
    """
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
    Utility function to load tensors from a file

    Parameters:
    -----------
    edge_features_file : string
        Features file for edges in the graph

    Returns:
    --------
    dictionary
        mappings between etype and list of features
    """
    edge_features = dgl.data.utils.load_tensors(edge_features_file, True)
    return edge_features

407

408
409
410
411
412
413
414
415
def write_node_features(node_features, node_file):
    """
    Utility function to serialize node_features in node_file file

    Parameters:
    -----------
    node_features : dictionary
        dictionary storing ntype <-> list of features
416
    node_file     : string
417
418
419
420
        File in which the node information is serialized
    """
    dgl.data.utils.save_tensors(node_file, node_features)

421
422

def write_edge_features(edge_features, edge_file):
423
424
425
426
427
428
429
    """
    Utility function to serialize edge_features in edge_file file

    Parameters:
    -----------
    edge_features : dictionary
        dictionary storing etype <-> list of features
430
    edge_file     : string
431
432
433
434
        File in which the edge information is serialized
    """
    dgl.data.utils.save_tensors(edge_file, edge_features)

435

436
def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes):
437
438
439
440
441
442
443
444
445
    """
    Utility function to serialize graph dgl objects

    Parameters:
    -----------
    graph_obj : dgl graph object
        graph dgl object, as created in convert_partition.py, which is to be serialized
    graph_file : string
        File name in which graph object is serialized
446
447
448
449
    formats : str or list[str]
        Save graph in specified formats.
    sort_etypes : bool
        Whether to sort etypes in csc/csr.
450
    """
451
452
453
454
    dgl.distributed.partition._save_graphs(
        graph_file, [graph_obj], formats, sort_etypes
    )

455

456
457
458
459
460
461
462
463
464
465
466
def write_dgl_objects(
    graph_obj,
    node_features,
    edge_features,
    output_dir,
    part_id,
    orig_nids,
    orig_eids,
    formats,
    sort_etypes,
):
467
    """
468
    Wrapper function to write graph, node/edge feature, original node/edge IDs.
469
470
471

    Parameters:
    -----------
472
473
474
475
476
477
    graph_obj : dgl object
        graph dgl object as created in convert_partition.py file
    node_features : dgl object
        Tensor data for node features
    edge_features : dgl object
        Tensor data for edge features
478
479
480
481
    output_dir : string
        location where the output files will be located
    part_id : int
        integer indicating the partition-id
482
483
484
485
    orig_nids : dict
        original node IDs
    orig_eids : dict
        original edge IDs
486
487
488
489
    formats : str or list[str]
        Save graph in formats.
    sort_etypes : bool
        Whether to sort etypes in csc/csr.
490
    """
491
    part_dir = output_dir + "/part" + str(part_id)
492
    os.makedirs(part_dir, exist_ok=True)
493
494
495
    write_graph_dgl(
        os.path.join(part_dir, "graph.dgl"), graph_obj, formats, sort_etypes
    )
496
497

    if node_features != None:
498
499
500
        write_node_features(
            node_features, os.path.join(part_dir, "node_feat.dgl")
        )
501

502
503
504
505
    if edge_features != None:
        write_edge_features(
            edge_features, os.path.join(part_dir, "edge_feat.dgl")
        )
506

507
    if orig_nids is not None:
508
        orig_nids_file = os.path.join(part_dir, "orig_nids.dgl")
509
510
        dgl.data.utils.save_tensors(orig_nids_file, orig_nids)
    if orig_eids is not None:
511
        orig_eids_file = os.path.join(part_dir, "orig_eids.dgl")
512
513
        dgl.data.utils.save_tensors(orig_eids_file, orig_eids)

514
515

def get_idranges(names, counts, num_chunks=None):
516
    """
517
    Utility function to compute typd_id/global_id ranges for both nodes and edges.
518
519
520
521
522
523
524

    Parameters:
    -----------
    names : list of strings
        list of node/edge types as strings
    counts : list of lists
        each list contains no. of nodes/edges in a given chunk
525
526
527
528
529
    num_chunks : int, optional
        In distributed partition pipeline, ID ranges are grouped into chunks.
        In some scenarios, we'd like to merge ID ranges into specific number
        of chunks. This parameter indicates the expected number of chunks.
        If not specified, no merge is applied.
530
531
532
533
534

    Returns:
    --------
    dictionary
        dictionary where the keys are node-/edge-type names and values are
535
536
        list of tuples where each tuple indicates the range of values for
        corresponding type-ids.
537
538
539
540
541
542
543
544
    dictionary
        dictionary where the keys are node-/edge-type names and value is a tuple.
        This tuple indicates the global-ids for the associated node-/edge-type.
    """
    gnid_start = 0
    gnid_end = gnid_start
    tid_dict = {}
    gid_dict = {}
545
    orig_num_chunks = 0
546
    for idx, typename in enumerate(names):
547
548
549
550
551
552
553
554
        type_counts = counts[idx]
        tid_start = np.cumsum([0] + type_counts[:-1])
        tid_end = np.cumsum(type_counts)
        tid_ranges = list(zip(tid_start, tid_end))

        gnid_end += tid_ranges[-1][1]

        tid_dict[typename] = tid_ranges
555
        gid_dict[typename] = np.array([gnid_start, gnid_end]).reshape([1, 2])
556
557

        gnid_start = gnid_end
558
559
560
561
562
        orig_num_chunks = len(tid_start)

    if num_chunks is None:
        return tid_dict, gid_dict

563
564
565
    assert (
        num_chunks <= orig_num_chunks
    ), "Specified number of chunks should be less/euqual than original numbers of ID ranges."
566
567
568
569
570
    chunk_list = np.array_split(np.arange(orig_num_chunks), num_chunks)
    for typename in tid_dict:
        orig_tid_ranges = tid_dict[typename]
        tid_ranges = []
        for idx in chunk_list:
571
572
573
            tid_ranges.append(
                (orig_tid_ranges[idx[0]][0], orig_tid_ranges[idx[-1]][-1])
            )
574
        tid_dict[typename] = tid_ranges
575
576
577

    return tid_dict, gid_dict

578
579
580
581
582

def memory_snapshot(tag, rank):
    """
    Utility function to take a snapshot of the usage of system resources
    at a given point of time.
583
584

    Parameters:
585
586
587
588
589
590
591
592
593
594
    -----------
    tag : string
        string provided by the user for bookmarking purposes
    rank : integer
        process id of the participating process
    """
    GB = 1024 * 1024 * 1024
    MB = 1024 * 1024
    KB = 1024

595
    peak = dgl.partition.get_peak_mem() * KB
596
597
598
599
600
    mem = psutil.virtual_memory()
    avail = mem.available / MB
    used = mem.used / MB
    total = mem.total / MB

601
602
    mem_string = f"{total:.0f} (MB) total, {peak:.0f} (MB) peak, {used:.0f} (MB) used, {avail:.0f} (MB) avail"
    logging.debug(f"[Rank: {rank} MEMORY_SNAPSHOT] {mem_string} - {tag}")
603

604
605
606

def map_partid_rank(partid, world_size):
    """Auxiliary function to map a given partition id to one of the rank in the
607
    MPI_WORLD processes. The range of partition ids is assumed to equal or a
608
609
610
611
612
613
614
615
616
617
    multiple of the total size of MPI_WORLD. In this implementation, we use
    a cyclical mapping procedure to convert partition ids to ranks.

    Parameters:
    -----------
    partid : int
        partition id, as read from node id to partition id mappings.

    Returns:
    --------
618
    int :
619
620
621
622
        rank of the process, which will be responsible for the given partition
        id.
    """
    return partid % world_size
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645


def generate_read_list(num_files, world_size):
    """Generate the file IDs to read for each rank.

    Parameters:
    -----------
    num_files : int
        Total number of files.
    world_size : int
        World size of group.

    Returns:
    --------
    read_list : np.array
        Array of target file IDs to read.

    Examples
    --------
    >>> tools.distpartitionning.utils.generate_read_list(10, 4)
    [array([0, 1, 2]), array([3, 4, 5]), array([6, 7]), array([8, 9])]
    """
    return np.array_split(np.arange(num_files), world_size)