utils.py 21.3 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
12

import torch
13
from dgl.distributed.partition import _dump_part_config
14
from pyarrow import csv
15

16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
DATA_TYPE_ID = {
    data_type: id
    for id, data_type in enumerate(
        [
            torch.float32,
            torch.float64,
            torch.float16,
            torch.uint8,
            torch.int8,
            torch.int16,
            torch.int32,
            torch.int64,
            torch.bool,
        ]
    )
}

REV_DATA_TYPE_ID = {id: data_type for data_type, id in DATA_TYPE_ID.items()}

35

36
def read_ntype_partition_files(schema_map, input_dir):
37
    """
38
39
    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
40
    containing the partition id mapping for a given nodeid.
41
42
43

    Parameters:
    -----------
44
45
46
    schema_map : dictionary
        dictionary created by reading the input metadata json file
    input_dir : string
47
        directory in which the node-id to partition-id mappings files are
48
        located for each of the node types in the input graph
49
50
51

    Returns:
    --------
52
    numpy array :
53
        array of integers representing mapped partition-ids for a given node-id.
54
55
56
57
58
        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.
59
    """
60
61
    assert os.path.isdir(input_dir)

62
    # iterate over the node types and extract the partition id mappings
63
64
65
    part_ids = []
    ntype_names = schema_map[constants.STR_NODE_TYPE]
    for ntype in ntype_names:
66
67
68
69
70
71
72
73
        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()
74
75
        part_ids.append(ntype_partids)
    return np.concatenate(part_ids)
76

77

78
79
80
def read_json(json_file):
    """
    Utility method to read a json file schema
81

82
83
84
85
86
87
88
89
90
91
92
93
94
95
    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

96

97
98
99
100
101
102
103
104
105
106
107
108
109
110
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:
    --------
111
    list :
112
113
114
115
116
117
        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]

118
119

def get_ntype_featnames(ntype_name, schema_map):
120
121
122
123
124
125
126
127
128
129
130
131
132
133
    """
    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:
    --------
134
    list :
135
136
        a list of feature names for a given node_type
    """
137
138
139
    node_data = schema_map[constants.STR_NODE_DATA]
    feats = node_data.get(ntype_name, {})
    return [feat for feat in feats]
140

141

142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
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]
162
163
    etype_etypeid_map = {e: i for i, e in enumerate(etypes)}
    etypeid_etype_map = {i: e for i, e in enumerate(etypes)}
164
165
    return etype_etypeid_map, etypes, etypeid_etype_map

166

167
def get_node_types(schema_map):
168
    """
169
170
171
172
173
    Utility method to extract node_typename -> node_type mappings
    as defined by the input schema

    Parameters:
    -----------
174
    schema_map : dictionary
175
176
177
178
179
        Input schema from which the node_typename -> node_type
        dictionary is created.

    Returns:
    --------
180
181
182
183
184
185
    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
186
    """
187
    ntypes = schema_map[constants.STR_NODE_TYPE]
188
189
    ntype_ntypeid_map = {e: i for i, e in enumerate(ntypes)}
    ntypeid_ntype_map = {i: e for i, e in enumerate(ntypes)}
190
191
    return ntype_ntypeid_map, ntypes, ntypeid_ntype_map

192

193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
def get_gid_offsets(typenames, typecounts):
    """
    Builds a map where the key-value pairs are typnames and respective
    global-id offsets.

    Parameters:
    -----------
    typenames : list of strings
        a list of strings which can be either node typenames or edge typenames
    typecounts : list of integers
        a list of integers indicating the total number of nodes/edges for its
        typeid which is the index in this list

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

    """
    assert len(typenames) == len(
        typecounts
    ), f"No. of typenames does not match with its type counts names = {typenames}, counts = {typecounts}"

    counts = []
    for name in typenames:
        counts.append(typecounts[name])
    starts = np.cumsum([0] + counts[:-1])
    ends = np.cumsum(counts)

    gid_offsets = {}
    for idx, name in enumerate(typenames):
        gid_offsets[name] = [starts[idx], ends[idx]]
    return gid_offsets

    """
    starts = np.cumsum([0] + type_counts[:-1])
    ends = np.cumsum(type_counts)
    gid_offsets = {}
    for idx, name in enumerate(typenames):
        gid_offsets[name] = [start[idx], ends[idx]]

    return gid_offsets
    """


239
def get_gnid_range_map(node_tids):
240
    """
241
    Retrieves auxiliary dictionaries from the metadata json object
242
243
244

    Parameters:
    -----------
245
246
    node_tids: dictionary
        This dictionary contains the information about nodes for each node_type.
247
        Typically this information contains p-entries, where each entry has a file-name,
248
249
250
        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
251
252
253

    Returns:
    --------
254
    dictionary :
255
256
        a dictionary where keys are node_type names and values are global_nid range, which is a tuple.

257
    """
258
    ntypes_gid_range = {}
259
    offset = 0
260
    for k, v in node_tids.items():
261
262
263
264
        ntypes_gid_range[k] = [offset + int(v[0][0]), offset + int(v[-1][1])]
        offset += int(v[-1][1])

    return ntypes_gid_range
265

266
267
268
269

def write_metadata_json(
    input_list, output_dir, graph_name, world_size, num_parts
):
270
    """
271
    Merge json schema's from each of the rank's on rank-0.
272
273
274
275
276
277
278
279
280
281
282
    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
    """
283
    # Preprocess the input_list, a list of dictionaries
284
    # each dictionary will contain num_parts/world_size metadata json
285
286
    # which correspond to local partitions on the respective ranks.
    metadata_list = []
287
    for local_part_id in range(num_parts // world_size):
288
        for idx in range(world_size):
289
290
291
292
293
            metadata_list.append(
                input_list[idx][
                    "local-part-id-" + str(local_part_id * world_size + idx)
                ]
            )
294

295
    # Initialize global metadata
296
297
    graph_metadata = {}

298
    # Merge global_edge_ids from each json object in the input list
299
300
301
302
303
    edge_map = {}
    x = metadata_list[0]["edge_map"]
    for k in x:
        edge_map[k] = []
        for idx in range(len(metadata_list)):
304
305
306
307
308
309
            edge_map[k].append(
                [
                    int(metadata_list[idx]["edge_map"][k][0][0]),
                    int(metadata_list[idx]["edge_map"][k][0][1]),
                ]
            )
310
311
312
313
314
315
    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"]

316
    # Merge global_nodeids from each of json object in the input list
317
318
319
320
321
    node_map = {}
    x = metadata_list[0]["node_map"]
    for k in x:
        node_map[k] = []
        for idx in range(len(metadata_list)):
322
323
324
325
326
327
            node_map[k].append(
                [
                    int(metadata_list[idx]["node_map"][k][0][0]),
                    int(metadata_list[idx]["node_map"][k][0][1]),
                ]
            )
328
329
330
    graph_metadata["node_map"] = node_map

    graph_metadata["ntypes"] = metadata_list[0]["ntypes"]
331
332
333
334
335
336
    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))])
    )
337
338
339
340
    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)):
341
342
343
344
345
        graph_metadata["part-{}".format(i)] = metadata_list[i][
            "part-{}".format(i)
        ]

    _dump_part_config(f"{output_dir}/metadata.json", graph_metadata)
346
347


348
349
350
def augment_edge_data(
    edge_data, lookup_service, edge_tids, rank, world_size, num_parts
):
351
352
    """
    Add partition-id (rank which owns an edge) column to the edge_data.
353

354
355
356
357
    Parameters:
    -----------
    edge_data : numpy ndarray
        Edge information as read from the xxx_edges.txt file
358
359
360
361
362
363
364
365
366
367
    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
368
369
    num_parts : integer
        total no. of partitions requested for the input graph
370
371
372

    Returns:
    --------
373
374
    dictionary :
        dictionary with keys as column names and values as numpy arrays and this information is
375
376
        loaded from input dataset files. In addition to this we include additional columns which
        aid this pipelines computation, like constants.OWNER_PROCESS
377
    """
378
    # add global_nids to the node_data
379
380
    etype_offset = {}
    offset = 0
381
    for etype_name, tid_range in edge_tids.items():
382
383
384
385
        etype_offset[etype_name] = offset + int(tid_range[0][0])
        offset += int(tid_range[-1][1])

    global_eids = []
386
    for etype_name, tid_range in edge_tids.items():
387
388
        for idx in range(num_parts):
            if map_partid_rank(idx, world_size) == rank:
389
390
391
392
393
394
395
396
397
398
399
                if len(tid_range) > idx:
                    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))

    global_eids = (
        np.concatenate(global_eids)
        if len(global_eids) > 0
        else np.array([], dtype=np.int64)
    )
400
    assert global_eids.shape[0] == edge_data[constants.ETYPE_ID].shape[0]
401
    edge_data[constants.GLOBAL_EID] = global_eids
402
    return edge_data
403

404

405
def read_edges_file(edge_file, edge_data_dict):
406
    """
407
408
409
410
411
412
413
414
415
416
417
    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
418
        in a dictionary with key-value pairs as column-names and column-data.
419
420
421
422
    """
    if edge_file == "" or edge_file == None:
        return None

423
424
    # Read the file from here.
    # <global_src_id> <global_dst_id> <type_eid> <etype> <attributes>
425
426
427
    # 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

428
429
430
431
432
    edge_data_df = csv.read_csv(
        edge_file,
        read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
        parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
    )
433
    edge_data_dict = {}
434
435
436
437
    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()
438
439
    return edge_data_dict

440

441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
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

459

460
def read_edge_features_file(edge_features_file):
461
    """
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
    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

477

478
479
480
481
482
483
484
485
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
486
    node_file     : string
487
488
489
490
        File in which the node information is serialized
    """
    dgl.data.utils.save_tensors(node_file, node_features)

491
492

def write_edge_features(edge_features, edge_file):
493
494
495
496
497
498
499
    """
    Utility function to serialize edge_features in edge_file file

    Parameters:
    -----------
    edge_features : dictionary
        dictionary storing etype <-> list of features
500
    edge_file     : string
501
502
503
504
        File in which the edge information is serialized
    """
    dgl.data.utils.save_tensors(edge_file, edge_features)

505

506
def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes):
507
508
509
510
511
512
513
514
515
    """
    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
516
517
518
519
    formats : str or list[str]
        Save graph in specified formats.
    sort_etypes : bool
        Whether to sort etypes in csc/csr.
520
    """
521
522
523
524
    dgl.distributed.partition._save_graphs(
        graph_file, [graph_obj], formats, sort_etypes
    )

525

526
527
528
529
530
531
532
533
534
535
536
def write_dgl_objects(
    graph_obj,
    node_features,
    edge_features,
    output_dir,
    part_id,
    orig_nids,
    orig_eids,
    formats,
    sort_etypes,
):
537
    """
538
    Wrapper function to write graph, node/edge feature, original node/edge IDs.
539
540
541

    Parameters:
    -----------
542
543
544
545
546
547
    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
548
549
550
551
    output_dir : string
        location where the output files will be located
    part_id : int
        integer indicating the partition-id
552
553
554
555
    orig_nids : dict
        original node IDs
    orig_eids : dict
        original edge IDs
556
557
558
559
    formats : str or list[str]
        Save graph in formats.
    sort_etypes : bool
        Whether to sort etypes in csc/csr.
560
    """
561
    part_dir = output_dir + "/part" + str(part_id)
562
    os.makedirs(part_dir, exist_ok=True)
563
564
565
    write_graph_dgl(
        os.path.join(part_dir, "graph.dgl"), graph_obj, formats, sort_etypes
    )
566
567

    if node_features != None:
568
569
570
        write_node_features(
            node_features, os.path.join(part_dir, "node_feat.dgl")
        )
571

572
573
574
575
    if edge_features != None:
        write_edge_features(
            edge_features, os.path.join(part_dir, "edge_feat.dgl")
        )
576

577
    if orig_nids is not None:
578
        orig_nids_file = os.path.join(part_dir, "orig_nids.dgl")
579
580
        dgl.data.utils.save_tensors(orig_nids_file, orig_nids)
    if orig_eids is not None:
581
        orig_eids_file = os.path.join(part_dir, "orig_eids.dgl")
582
583
        dgl.data.utils.save_tensors(orig_eids_file, orig_eids)

584
585

def get_idranges(names, counts, num_chunks=None):
586
    """
587
588
    counts will be a list of numbers of a dictionary.
    Length is less than or equal to the num_parts variable.
589
590
591
592

    Parameters:
    -----------
    names : list of strings
593
594
595
596
        which are either node-types or edge-types
    counts : list of integers
        which are total no. of nodes or edges for a give node
        or edge type
597
    num_chunks : int, optional
598
        specifying the no. of chunks
599
600
601
602
603

    Returns:
    --------
    dictionary
        dictionary where the keys are node-/edge-type names and values are
604
605
        list of tuples where each tuple indicates the range of values for
        corresponding type-ids.
606
607
608
609
610
611
612
613
614
    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 = {}

615
616
617
    for idx, typename in enumerate(names):
        gnid_end += counts[typename]
        tid_dict[typename] = [[0, counts[typename]]]
618
        gid_dict[typename] = np.array([gnid_start, gnid_end]).reshape([1, 2])
619
620
621
622
        gnid_start = gnid_end

    return tid_dict, gid_dict

623

624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
def get_ntype_counts_map(ntypes, ntype_counts):
    """
    Return a dictionary with key, value pairs as node type names and no. of
    nodes of a particular type in the input graph.

    Parameters:
    -----------
    ntypes : list of strings
        where each string is a node-type name
    ntype_counts : list of integers
        where each integer is the total no. of nodes for that, idx, node type

    Returns:
    --------
    dictinary :
        a dictionary where node-type names are keys and values are total no.
        of nodes for a given node-type name (which is also the key)
    """
    return dict(zip(ntypes, ntype_counts))


645
646
647
648
def memory_snapshot(tag, rank):
    """
    Utility function to take a snapshot of the usage of system resources
    at a given point of time.
649
650

    Parameters:
651
652
653
654
655
656
657
658
659
660
    -----------
    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

661
    peak = dgl.partition.get_peak_mem() * KB
662
663
664
665
666
    mem = psutil.virtual_memory()
    avail = mem.available / MB
    used = mem.used / MB
    total = mem.total / MB

667
668
    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}")
669

670
671
672

def map_partid_rank(partid, world_size):
    """Auxiliary function to map a given partition id to one of the rank in the
673
    MPI_WORLD processes. The range of partition ids is assumed to equal or a
674
675
676
677
678
679
680
681
682
683
    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:
    --------
684
    int :
685
686
687
688
        rank of the process, which will be responsible for the given partition
        id.
    """
    return partid % world_size
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711


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)