Revert "[Distributed Training Pipeline] Initial implementation of Distributed...

Revert "[Distributed Training Pipeline] Initial implementation of Distributed data processing step in the Dis… (#3926)" (#4037)

This reverts commit 4b87e47f.

tools/convert_partition.py

+import os
+import json
+import time
+import argparse
+import numpy as np
+import dgl
+import torch as th
+import pyarrow
+import pandas as pd
+from pyarrow import csv
+
+
+parser = argparse.ArgumentParser(description='Construct graph partitions')
+parser.add_argument('--input-dir', required=True, type=str,
+                    help='The directory path that contains the partition results.')
+parser.add_argument('--graph-name', required=True, type=str,
+                    help='The graph name')
+parser.add_argument('--schema', required=True, type=str,
+                    help='The schema of the graph')
+parser.add_argument('--num-parts', required=True, type=int,
+                    help='The number of partitions')
+parser.add_argument('--num-node-weights', required=True, type=int,
+                    help='The number of node weights used by METIS.')
+parser.add_argument('--workspace', type=str, default='/tmp',
+                    help='The directory to store the intermediate results')
+parser.add_argument('--node-attr-dtype', type=str, default=None,
+                    help='The data type of the node attributes')
+parser.add_argument('--edge-attr-dtype', type=str, default=None,
+                    help='The data type of the edge attributes')
+parser.add_argument('--output', required=True, type=str,
+                    help='The output directory of the partitioned results')
+parser.add_argument('--removed-edges', help='a file that contains the removed self-loops and duplicated edges',
+                    default=None, type=str)
+
+args = parser.parse_args()
+
+input_dir = args.input_dir
+graph_name = args.graph_name
+num_parts = args.num_parts
+num_node_weights = args.num_node_weights
+node_attr_dtype = args.node_attr_dtype
+edge_attr_dtype = args.edge_attr_dtype
+workspace_dir = args.workspace
+output_dir = args.output
+
+self_loop_edges = None
+duplicate_edges = None
+if args.removed_edges is not None:
+    removed_file = '{}/{}'.format(input_dir, args.removed_edges)
+    removed_df = csv.read_csv(removed_file, read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
+                              parse_options=pyarrow.csv.ParseOptions(delimiter=' '))
+    assert removed_df.num_columns == 4
+    src_id = removed_df['f0'].to_numpy()
+    dst_id = removed_df['f1'].to_numpy()
+    orig_id = removed_df['f2'].to_numpy()
+    etype = removed_df['f3'].to_numpy()
+    self_loop_idx = src_id == dst_id
+    not_self_loop_idx = src_id != dst_id
+    self_loop_edges = [src_id[self_loop_idx], dst_id[self_loop_idx],
+                       orig_id[self_loop_idx], etype[self_loop_idx]]
+    duplicate_edges = [src_id[not_self_loop_idx], dst_id[not_self_loop_idx],
+                       orig_id[not_self_loop_idx], etype[not_self_loop_idx]]
+    print('There are {} self-loops and {} duplicated edges in the removed edges'.format(len(self_loop_edges[0]),
+                                                                                        len(duplicate_edges[0])))
+
+with open(args.schema) as json_file:
+    schema = json.load(json_file)
+nid_ranges = schema['nid']
+eid_ranges = schema['eid']
+nid_ranges = {key: np.array(nid_ranges[key]).reshape(
+    1, 2) for key in nid_ranges}
+eid_ranges = {key: np.array(eid_ranges[key]).reshape(
+    1, 2) for key in eid_ranges}
+id_map = dgl.distributed.id_map.IdMap(nid_ranges)
+
+ntypes = [(key, nid_ranges[key][0, 0]) for key in nid_ranges]
+ntypes.sort(key=lambda e: e[1])
+ntype_offset_np = np.array([e[1] for e in ntypes])
+ntypes = [e[0] for e in ntypes]
+ntypes_map = {e: i for i, e in enumerate(ntypes)}
+etypes = [(key, eid_ranges[key][0, 0]) for key in eid_ranges]
+etypes.sort(key=lambda e: e[1])
+etype_offset_np = np.array([e[1] for e in etypes])
+etypes = [e[0] for e in etypes]
+etypes_map = {e: i for i, e in enumerate(etypes)}
+
+
+def read_feats(file_name):
+    attrs = csv.read_csv(file_name, read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
+                         parse_options=pyarrow.csv.ParseOptions(delimiter=' '))
+    num_cols = len(attrs.columns)
+    return np.stack([attrs.columns[i].to_numpy() for i in range(num_cols)], 1)
+
+
+max_nid = np.iinfo(np.int32).max
+num_edges = 0
+num_nodes = 0
+node_map_val = {ntype: [] for ntype in ntypes}
+edge_map_val = {etype: [] for etype in etypes}
+for part_id in range(num_parts):
+    part_dir = output_dir + '/part' + str(part_id)
+    os.makedirs(part_dir, exist_ok=True)
+
+    node_file = 'p{:03}-{}_nodes.txt'.format(part_id, graph_name)
+    # The format of each line in the node file:
+    # <node_id> <node_type> <weight1> ... <orig_type_node_id> <attributes>
+    # The node file contains nodes that belong to a partition. It doesn't include HALO nodes.
+    orig_type_nid_col = 3 + num_node_weights
+    first_attr_col = 4 + num_node_weights
+
+    # Get the first two columns which is the node ID and node type.
+    tmp_output = workspace_dir + '/' + node_file + '.tmp'
+    os.system('awk \'{print $1, $2, $' + str(orig_type_nid_col) + '}\''
+              + ' {} > {}'.format(input_dir + '/' + node_file, tmp_output))
+    nodes = csv.read_csv(tmp_output, read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
+                         parse_options=pyarrow.csv.ParseOptions(delimiter=' '))
+    nids, ntype_ids, orig_type_nid = nodes.columns[0].to_numpy(), nodes.columns[1].to_numpy(), \
+        nodes.columns[2].to_numpy()
+    orig_homo_nid = ntype_offset_np[ntype_ids] + orig_type_nid
+    assert np.all(nids[1:] - nids[:-1] == 1)
+    nid_range = (nids[0], nids[-1])
+    num_nodes += len(nodes)
+
+    if node_attr_dtype is not None:
+        # Get node attributes
+        # Here we just assume all nodes have the same attributes.
+        # In practice, this is not the same, in which we need more complex solution to
+        # encode and decode node attributes.
+        os.system('cut -d\' \' -f {}- {} > {}'.format(first_attr_col,
+                                                      input_dir + '/' + node_file,
+                                                      tmp_output))
+        node_attrs = read_feats(tmp_output)
+        node_feats = {}
+        # nodes in a partition has been sorted based on node types.
+        for ntype_name in nid_ranges:
+            ntype_id = ntypes_map[ntype_name]
+            type_nids = nids[ntype_ids == ntype_id]
+            assert np.all(type_nids == np.arange(
+                type_nids[0], type_nids[-1] + 1))
+            node_feats[ntype_name +
+                       '/feat'] = th.as_tensor(node_attrs[ntype_ids == ntype_id])
+        dgl.data.utils.save_tensors(os.path.join(
+            part_dir, "node_feat.dgl"), node_feats)
+
+    # Determine the node ID ranges of different node types.
+    for ntype_name in nid_ranges:
+        ntype_id = ntypes_map[ntype_name]
+        type_nids = nids[ntype_ids == ntype_id]
+        node_map_val[ntype_name].append(
+            [int(type_nids[0]), int(type_nids[-1]) + 1])
+
+    edge_file = 'p{:03}-{}_edges.txt'.format(part_id, graph_name)
+    # The format of each line in the edge file:
+    # <src_id> <dst_id> <orig_src_id> <orig_dst_id> <orig_edge_id> <edge_type> <attributes>
+
+    tmp_output = workspace_dir + '/' + edge_file + '.tmp'
+    os.system('awk \'{print $1, $2, $3, $4, $5, $6}\'' + ' {} > {}'.format(input_dir + '/' + edge_file,
+                                                                           tmp_output))
+    edges = csv.read_csv(tmp_output, read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
+                         parse_options=pyarrow.csv.ParseOptions(delimiter=' '))
+    num_cols = len(edges.columns)
+    src_id, dst_id, orig_src_id, orig_dst_id, orig_edge_id, etype_ids = [
+        edges.columns[i].to_numpy() for i in range(num_cols)]
+
+    # Let's merge the self-loops and duplicated edges to the partition.
+    src_id_list, dst_id_list = [src_id], [dst_id]
+    orig_src_id_list, orig_dst_id_list = [orig_src_id], [orig_dst_id]
+    orig_edge_id_list, etype_id_list = [orig_edge_id], [etype_ids]
+    if self_loop_edges is not None and len(self_loop_edges[0]) > 0:
+        uniq_orig_nids, idx = np.unique(orig_dst_id, return_index=True)
+        common_nids, common_idx1, common_idx2 = np.intersect1d(
+            uniq_orig_nids, self_loop_edges[0], return_indices=True)
+        idx = idx[common_idx1]
+        # the IDs after ID assignment
+        src_id_list.append(dst_id[idx])
+        dst_id_list.append(dst_id[idx])
+        # homogeneous IDs in the input graph.
+        orig_src_id_list.append(self_loop_edges[0][common_idx2])
+        orig_dst_id_list.append(self_loop_edges[0][common_idx2])
+        # edge IDs and edge type.
+        orig_edge_id_list.append(self_loop_edges[2][common_idx2])
+        etype_id_list.append(self_loop_edges[3][common_idx2])
+        print('Add {} self-loops in partition {}'.format(len(idx), part_id))
+    if duplicate_edges is not None and len(duplicate_edges[0]) > 0:
+        part_ids = orig_src_id.astype(
+            np.int64) * max_nid + orig_dst_id.astype(np.int64)
+        uniq_orig_ids, idx = np.unique(part_ids, return_index=True)
+        duplicate_ids = duplicate_edges[0].astype(
+            np.int64) * max_nid + duplicate_edges[1].astype(np.int64)
+        common_nids, common_idx1, common_idx2 = np.intersect1d(
+            uniq_orig_ids, duplicate_ids, return_indices=True)
+        idx = idx[common_idx1]
+        # the IDs after ID assignment
+        src_id_list.append(src_id[idx])
+        dst_id_list.append(dst_id[idx])
+        # homogeneous IDs in the input graph.
+        orig_src_id_list.append(duplicate_edges[0][common_idx2])
+        orig_dst_id_list.append(duplicate_edges[1][common_idx2])
+        # edge IDs and edge type.
+        orig_edge_id_list.append(duplicate_edges[2][common_idx2])
+        etype_id_list.append(duplicate_edges[3][common_idx2])
+        print('Add {} duplicated edges in partition {}'.format(len(idx), part_id))
+    src_id = np.concatenate(src_id_list) if len(
+        src_id_list) > 1 else src_id_list[0]
+    dst_id = np.concatenate(dst_id_list) if len(
+        dst_id_list) > 1 else dst_id_list[0]
+    orig_src_id = np.concatenate(orig_src_id_list) if len(
+        orig_src_id_list) > 1 else orig_src_id_list[0]
+    orig_dst_id = np.concatenate(orig_dst_id_list) if len(
+        orig_dst_id_list) > 1 else orig_dst_id_list[0]
+    orig_edge_id = np.concatenate(orig_edge_id_list) if len(
+        orig_edge_id_list) > 1 else orig_edge_id_list[0]
+    etype_ids = np.concatenate(etype_id_list) if len(
+        etype_id_list) > 1 else etype_id_list[0]
+    print('There are {} edges in partition {}'.format(len(src_id), part_id))
+
+    # It's not guaranteed that the edges are sorted based on edge type.
+    # Let's sort edges and all attributes on the edges.
+    sort_idx = np.argsort(etype_ids)
+    src_id, dst_id, orig_src_id, orig_dst_id, orig_edge_id, etype_ids = src_id[sort_idx], dst_id[sort_idx], \
+        orig_src_id[sort_idx], orig_dst_id[sort_idx], orig_edge_id[sort_idx], etype_ids[sort_idx]
+    assert np.all(np.diff(etype_ids) >= 0)
+
+    if edge_attr_dtype is not None:
+        # Get edge attributes
+        # Here we just assume all edges have the same attributes.
+        # In practice, this is not the same, in which we need more complex solution to
+        # encode and decode edge attributes.
+        os.system('cut -d\' \' -f 7- {} > {}'.format(input_dir +
+                  '/' + edge_file, tmp_output))
+        edge_attrs = th.as_tensor(read_feats(tmp_output))[sort_idx]
+        edge_feats = {}
+        for etype_name in eid_ranges:
+            etype_id = etypes_map[etype_name]
+            edge_feats[etype_name +
+                       '/feat'] = th.as_tensor(edge_attrs[etype_ids == etype_id])
+        dgl.data.utils.save_tensors(os.path.join(
+            part_dir, "edge_feat.dgl"), edge_feats)
+
+    # Determine the edge ID range of different edge types.
+    edge_id_start = num_edges
+    for etype_name in eid_ranges:
+        etype_id = etypes_map[etype_name]
+        edge_map_val[etype_name].append([int(edge_id_start),
+                                         int(edge_id_start + np.sum(etype_ids == etype_id))])
+        edge_id_start += np.sum(etype_ids == etype_id)
+
+    # Here we want to compute the unique IDs in the edge list.
+    # It is possible that a node that belongs to the partition but it doesn't appear
+    # in the edge list. That is, the node is assigned to this partition, but its neighbor
+    # belongs to another partition so that the edge is assigned to another partition.
+    # This happens in a directed graph.
+    # To avoid this kind of nodes being removed from the graph, we add node IDs that
+    # belong to this partition.
+    ids = np.concatenate(
+        [src_id, dst_id, np.arange(nid_range[0], nid_range[1] + 1)])
+    uniq_ids, idx, inverse_idx = np.unique(
+        ids, return_index=True, return_inverse=True)
+    assert len(uniq_ids) == len(idx)
+    # We get the edge list with their node IDs mapped to a contiguous ID range.
+    local_src_id, local_dst_id = np.split(inverse_idx[:len(src_id) * 2], 2)
+    compact_g = dgl.graph((local_src_id, local_dst_id))
+    compact_g.edata['orig_id'] = th.as_tensor(orig_edge_id)
+    compact_g.edata[dgl.ETYPE] = th.as_tensor(etype_ids)
+    compact_g.edata['inner_edge'] = th.ones(
+        compact_g.number_of_edges(), dtype=th.bool)
+
+    # The original IDs are homogeneous IDs.
+    # Similarly, we need to add the original homogeneous node IDs
+    orig_ids = np.concatenate([orig_src_id, orig_dst_id, orig_homo_nid])
+    orig_homo_ids = orig_ids[idx]
+    ntype, per_type_ids = id_map(orig_homo_ids)
+    compact_g.ndata['orig_id'] = th.as_tensor(per_type_ids)
+    compact_g.ndata[dgl.NTYPE] = th.as_tensor(ntype)
+    compact_g.ndata[dgl.NID] = th.as_tensor(uniq_ids)
+    compact_g.ndata['inner_node'] = th.as_tensor(np.logical_and(
+        uniq_ids >= nid_range[0], uniq_ids <= nid_range[1]))
+    local_nids = compact_g.ndata[dgl.NID][compact_g.ndata['inner_node'].bool()]
+    assert np.all((local_nids == th.arange(
+        local_nids[0], local_nids[-1] + 1)).numpy())
+    print('|V|={}'.format(compact_g.number_of_nodes()))
+    print('|E|={}'.format(compact_g.number_of_edges()))
+
+    # We need to reshuffle nodes in a partition so that all local nodes are labelled starting from 0.
+    reshuffle_nodes = th.arange(compact_g.number_of_nodes())
+    reshuffle_nodes = th.cat([reshuffle_nodes[compact_g.ndata['inner_node'].bool()],
+                              reshuffle_nodes[compact_g.ndata['inner_node'] == 0]])
+    compact_g1 = dgl.node_subgraph(compact_g, reshuffle_nodes)
+    compact_g1.ndata['orig_id'] = compact_g.ndata['orig_id'][reshuffle_nodes]
+    compact_g1.ndata[dgl.NTYPE] = compact_g.ndata[dgl.NTYPE][reshuffle_nodes]
+    compact_g1.ndata[dgl.NID] = compact_g.ndata[dgl.NID][reshuffle_nodes]
+    compact_g1.ndata['inner_node'] = compact_g.ndata['inner_node'][reshuffle_nodes]
+    compact_g1.edata['orig_id'] = compact_g.edata['orig_id'][compact_g1.edata[dgl.EID]]
+    compact_g1.edata[dgl.ETYPE] = compact_g.edata[dgl.ETYPE][compact_g1.edata[dgl.EID]]
+    compact_g1.edata['inner_edge'] = compact_g.edata['inner_edge'][compact_g1.edata[dgl.EID]]
+
+    # reshuffle edges on ETYPE as node_subgraph relabels edges
+    idx = th.argsort(compact_g1.edata[dgl.ETYPE])
+    u, v = compact_g1.edges()
+    u = u[idx]
+    v = v[idx]
+    compact_g2 = dgl.graph((u, v))
+    compact_g2.ndata['orig_id'] = compact_g1.ndata['orig_id']
+    compact_g2.ndata[dgl.NTYPE] = compact_g1.ndata[dgl.NTYPE]
+    compact_g2.ndata[dgl.NID] = compact_g1.ndata[dgl.NID]
+    compact_g2.ndata['inner_node'] = compact_g1.ndata['inner_node']
+    compact_g2.edata['orig_id'] = compact_g1.edata['orig_id'][idx]
+    compact_g2.edata[dgl.ETYPE] = compact_g1.edata[dgl.ETYPE][idx]
+    compact_g2.edata['inner_edge'] = compact_g1.edata['inner_edge'][idx]
+    compact_g2.edata[dgl.EID] = th.arange(
+        num_edges, num_edges + compact_g2.number_of_edges())
+    num_edges += compact_g2.number_of_edges()
+
+    dgl.save_graphs(part_dir + '/graph.dgl', [compact_g2])
+
+part_metadata = {'graph_name': graph_name,
+                 'num_nodes': num_nodes,
+                 'num_edges': num_edges,
+                 'part_method': 'metis',
+                 'num_parts': num_parts,
+                 'halo_hops': 1,
+                 'node_map': node_map_val,
+                 'edge_map': edge_map_val,
+                 'ntypes': ntypes_map,
+                 'etypes': etypes_map}
+
+for part_id in range(num_parts):
+    part_dir = 'part' + str(part_id)
+    node_feat_file = os.path.join(part_dir, "node_feat.dgl")
+    edge_feat_file = os.path.join(part_dir, "edge_feat.dgl")
+    part_graph_file = os.path.join(part_dir, "graph.dgl")
+    part_metadata['part-{}'.format(part_id)] = {'node_feats': node_feat_file,
+                                                'edge_feats': edge_feat_file,
+                                                'part_graph': part_graph_file}
+with open('{}/{}.json'.format(output_dir, graph_name), 'w') as outfile:
+    json.dump(part_metadata, outfile, sort_keys=True, indent=4)

tools/distpartitioning/README.md

-### xxx_nodes.txt format
-This file is used to provide node information to this framework. Following is the format for each line in this file:
-```
-<node_type> <weight1> <weight2> <weight3> <weight4> <global_type_node_id> <attributes>
-```
-where node_type is the type id of this node, weights can be any number of columns as determined by the user, global_type_node_id are the contiguous ids starting from `0` for a particular node_type. And attributes can be any number of columns at the end of each line. 
-
-###xxx___edges.txt format
-This file is used to provide edge information to this framework. Following is the format for each line in this file:
-```
-<global_src_id> <global_dst_id> <global_type_edge_id> <edge_type> <attributes>
-```
-where global_src_id and global_dst_id are two end points of an edge, global_type_edge_id is the unique id assigned to each edge type and are contiguous, and starting from 0, for each edge_type. Attributes can be any number of columns at the end of each line. 
-
-###Naming convention 
-`global_` prefix (for any node or edge ids) indicate that these ids are read from graph input files. These ids are allocated to nodes and edges before `data shuffling`. These ids are globally unique across all partitions.
-`shuffle_global_` prefix (for any node or edge ids) indicate that these ids are assigned after the `data shuffling` is completed. These ids are globally unique across all partitions.
-`part_local_` prefix (for any node or edge ids) indicate that these ids are assigned after the `data shuffling` and are unique within a given partition.
-For instance, if a variable is named as `global_src_id` it means that this id is read from the graph input file and is assumed to be globally unique across all partitions. Similarly if a variable is named `part_local_node_id`  then it means that this node_id is assigned after the data shuffling is complete and is unique with a given partition.
-
-###High level description of the algorithm
-####Single file format for graph input files
-Here we assume that all the nodes' related data is present in one single file and similarly all the edges are in one single file. 
-In this case following steps are executed to write dgl objects for each partition, as assigned my any partitioning algorithm, for example METIS. 
-#####Step 1 (Data Loading):
-Rank-0 process reads in all the graph files which are xxx_nodes.txt, xxx_edges.txt, node_feats.dgl, edge_feats.dgl and xxx_removed_edges.txt.
-Rank-0 process determines the ownership of nodes by using the output of partitioning algorithm (here, we expect the output of partitioning step is a mapping between a node and its partition id for the entire graph). Edge ownership is determined by the `destination` node-id for that edge. Each edge belongs to the partition-id of the destination node-id of each edge. 
-#####Step 2 (Data Shuffling):
-Rank-0 process will send node-data, edge-data, node-features, edge-features to their respective processes by using the ownership rules described in Step-1. Non-Rank-0 processes will receive their own nodes, edges, node-features and edge-features and store them in local data-structures. Upon completion of sending information Rank-0 process will delete nodes, edges, node-features and edge-features which are not owned by rank-0. 
-#####Step 3 (ID assignment and resolution): 
-At this time all the ranks will have their own local information in their respective data structures. Then each process will perform the following steps: a) Assign shuffle_global_xxx (here xxx is node_ids and edge_ids) for nodes and edges by performing prefix sum on all ranks. b) Assign part_local_xxx (xxx means node_ids and edge_ids) to nodes and edges so that they can be used to index into the node and edge features, and c) Retrieve shuffle_global_node_ids by using global_node_ids to determine the ownership of any given node. This step is done for the node_ids (present locally on any given rank) for which shuffle_global_node_ids were assigned on a different rank'ed process.
-#####Ste 4 (Serialization): 
-After every rank has global-ids, shuffle_global-ids, part_local-ids for all the nodes and edges present locally, then it proceeds by DGL object creation. Finally Rank-0 process will aggregate graph-level metadata and create a json file with graph-level information. 
-
-###How to use this tool
-To run this code on a single machine using multiple processes, use the following command
-```
-python3 data_proc_pipeline.py --world-size 2 --nodes-file mag_nodes.txt --edges-file mag_edges.txt --node-feats-file node_feat.dgl --metis-partitions mag_part.2 --input-dir /home/ubuntu/data --graph-name mag --schema mag.json --num-parts 2 --num-node-weights 4 --workspace /home/ubuntu/data --node-attr-dtype float --output /home/ubuntu/data/outputs --removed-edges mag_removed_edges.txt
-```
-Above command, assumes that there are `2` partitions and number of node weights are `4`. All other command line arguments are self-explanatory.

tools/distpartitioning/constants.py

-GLOBAL_NID = "global_node_id"
-GLOBAL_EID = "global_edge_id"
-
-SHUFFLE_GLOBAL_NID = "shuffle_global_node_id"
-SHUFFLE_GLOBAL_EID = "shuffle_global_edge_id"
-
-NTYPE_ID = "node_type_id"
-ETYPE_ID = "edge_type_id"
-
-GLOBAL_TYPE_NID = "global_type_node_id"
-GLOBAL_TYPE_EID = "global_type_edge_id"
-
-GLOBAL_SRC_ID = "global_src_id"
-GLOBAL_DST_ID = "global_dst_id"
-SHUFFLE_GLOBAL_SRC_ID = "shuffle_global_src_id"
-SHUFFLE_GLOBAL_DST_ID = "shuffle_global_dst_id"
-
-OWNER_PROCESS = "owner_proc_id"
-
-PART_LOCAL_NID = "part_local_nid"

tools/distpartitioning/convert_partition.py

-import os
-import json
-import time
-import argparse
-import numpy as np
-import dgl
-import torch as th
-import pyarrow
-import pandas as pd
-import constants
-from pyarrow import csv
-
-def create_dgl_object(graph_name, num_parts, \
-                        schema, part_id, node_data, \
-                        edge_data, nodeid_offset, edgeid_offset):
-    """
-    This function creates dgl objects for a given graph partition, as in function
-    arguments. 
-
-    Parameters:
-    -----------
-    graph_name : string
-        name of the graph
-    num_parts : int
-        total no. of partitions (of the original graph)
-    schame : json object
-        json object created by reading the graph metadata json file
-    part_id : int
-        partition id of the graph partition for which dgl object is to be created
-    node_data : numpy ndarray
-        node_data, where each row is of the following format:
-        <global_nid> <ntype_id> <global_type_nid>
-    edge_data : numpy ndarray
-        edge_data, where each row is of the following format: 
-        <global_src_id> <global_dst_id> <etype_id> <global_type_eid>
-    nodeid_offset : int
-        offset to be used when assigning node global ids in the current partition
-    edgeid_offset : int
-        offset to be used when assigning edge global ids in the current partition
-
-    return compact_g2, node_map_val, edge_map_val, ntypes_map, etypes_map
-
-    Returns: 
-    --------
-    dgl object
-        dgl object created for the current graph partition
-    dictionary
-        map between node types and the range of global node ids used
-    dictionary
-        map between edge types and the range of global edge ids used
-    dictionary
-        map between node type(string)  and node_type_id(int)
-    dictionary
-        map between edge type(string)  and edge_type_id(int)
-    """
-
-    #create auxiliary data structures from the schema object
-    global_nid_ranges = schema['nid']
-    global_eid_ranges = schema['eid']
-    global_nid_ranges = {key: np.array(global_nid_ranges[key]).reshape(
-        1, 2) for key in global_nid_ranges}
-    global_eid_ranges = {key: np.array(global_eid_ranges[key]).reshape(
-        1, 2) for key in global_eid_ranges}
-    id_map = dgl.distributed.id_map.IdMap(global_nid_ranges)
-
-    ntypes = [(key, global_nid_ranges[key][0, 0]) for key in global_nid_ranges]
-    ntypes.sort(key=lambda e: e[1])
-    ntype_offset_np = np.array([e[1] for e in ntypes])
-    ntypes = [e[0] for e in ntypes]
-    ntypes_map = {e: i for i, e in enumerate(ntypes)}
-    etypes = [(key, global_eid_ranges[key][0, 0]) for key in global_eid_ranges]
-    etypes.sort(key=lambda e: e[1])
-    etype_offset_np = np.array([e[1] for e in etypes])
-    etypes = [e[0] for e in etypes]
-    etypes_map = {e: i for i, e in enumerate(etypes)}
-
-    node_map_val = {ntype: [] for ntype in ntypes}
-    edge_map_val = {etype: [] for etype in etypes}
-
-    shuffle_global_nids, ntype_ids, global_type_nid = node_data[constants.SHUFFLE_GLOBAL_NID], \
-            node_data[constants.NTYPE_ID], node_data[constants.GLOBAL_TYPE_NID]
-
-    global_homo_nid = ntype_offset_np[ntype_ids] + global_type_nid
-    assert np.all(shuffle_global_nids[1:] - shuffle_global_nids[:-1] == 1)
-    shuffle_global_nid_range = (shuffle_global_nids[0], shuffle_global_nids[-1])
-
-
-    # Determine the node ID ranges of different node types.
-    for ntype_name in global_nid_ranges:
-        ntype_id = ntypes_map[ntype_name]
-        type_nids = shuffle_global_nids[ntype_ids == ntype_id]
-        node_map_val[ntype_name].append(
-            [int(type_nids[0]), int(type_nids[-1]) + 1])
-
-    #process edges
-    shuffle_global_src_id, shuffle_global_dst_id, global_src_id, global_dst_id, global_edge_id, etype_ids = \
-                edge_data[constants.SHUFFLE_GLOBAL_SRC_ID], edge_data[constants.SHUFFLE_GLOBAL_DST_ID], \
-                edge_data[constants.GLOBAL_SRC_ID], edge_data[constants.GLOBAL_DST_ID], \
-                edge_data[constants.GLOBAL_TYPE_EID], edge_data[constants.ETYPE_ID]
-    print('There are {} edges in partition {}'.format(len(shuffle_global_src_id), part_id))
-
-    # It's not guaranteed that the edges are sorted based on edge type.
-    # Let's sort edges and all attributes on the edges.
-    sort_idx = np.argsort(etype_ids)
-    shuffle_global_src_id, shuffle_global_dst_id, global_src_id, global_dst_id, global_edge_id, etype_ids = \
-            shuffle_global_src_id[sort_idx], shuffle_global_dst_id[sort_idx], global_src_id[sort_idx], \
-            global_dst_id[sort_idx], global_edge_id[sort_idx], etype_ids[sort_idx]
-    assert np.all(np.diff(etype_ids) >= 0)
-
-    # Determine the edge ID range of different edge types.
-    edge_id_start = edgeid_offset 
-    for etype_name in global_eid_ranges:
-        etype_id = etypes_map[etype_name]
-        edge_map_val[etype_name].append([edge_id_start,
-                                         edge_id_start + np.sum(etype_ids == etype_id)])
-        edge_id_start += np.sum(etype_ids == etype_id)
-
-    # Here we want to compute the unique IDs in the edge list.
-    # It is possible that a node that belongs to the partition but it doesn't appear
-    # in the edge list. That is, the node is assigned to this partition, but its neighbor
-    # belongs to another partition so that the edge is assigned to another partition.
-    # This happens in a directed graph.
-    # To avoid this kind of nodes being removed from the graph, we add node IDs that
-    # belong to this partition.
-    ids = np.concatenate(
-        [shuffle_global_src_id, shuffle_global_dst_id, np.arange(shuffle_global_nid_range[0], shuffle_global_nid_range[1] + 1)])
-    uniq_ids, idx, inverse_idx = np.unique(
-        ids, return_index=True, return_inverse=True)
-    assert len(uniq_ids) == len(idx)
-    # We get the edge list with their node IDs mapped to a contiguous ID range.
-    part_local_src_id, part_local_dst_id = np.split(inverse_idx[:len(shuffle_global_src_id) * 2], 2)
-    compact_g = dgl.graph((part_local_src_id, part_local_dst_id))
-    compact_g.edata['orig_id'] = th.as_tensor(global_edge_id)
-    compact_g.edata[dgl.ETYPE] = th.as_tensor(etype_ids)
-    compact_g.edata['inner_edge'] = th.ones(
-        compact_g.number_of_edges(), dtype=th.bool)
-
-    # The original IDs are homogeneous IDs.
-    # Similarly, we need to add the original homogeneous node IDs
-    global_nids = np.concatenate([global_src_id, global_dst_id, global_homo_nid])
-    global_homo_ids = global_nids[idx]
-    ntype, per_type_ids = id_map(global_homo_ids)
-    compact_g.ndata['orig_id'] = th.as_tensor(per_type_ids)
-    compact_g.ndata[dgl.NTYPE] = th.as_tensor(ntype)
-    compact_g.ndata[dgl.NID] = th.as_tensor(uniq_ids)
-    compact_g.ndata['inner_node'] = th.as_tensor(np.logical_and(
-        uniq_ids >= shuffle_global_nid_range[0], uniq_ids <= shuffle_global_nid_range[1]))
-    part_local_nids = compact_g.ndata[dgl.NID][compact_g.ndata['inner_node'].bool()]
-    assert np.all((part_local_nids == th.arange(
-        part_local_nids[0], part_local_nids[-1] + 1)).numpy())
-    print('|V|={}'.format(compact_g.number_of_nodes()))
-    print('|E|={}'.format(compact_g.number_of_edges()))
-
-    # We need to reshuffle nodes in a partition so that all local nodes are labelled starting from 0.
-    reshuffle_nodes = th.arange(compact_g.number_of_nodes())
-    reshuffle_nodes = th.cat([reshuffle_nodes[compact_g.ndata['inner_node'].bool()],
-                              reshuffle_nodes[compact_g.ndata['inner_node'] == 0]])
-    compact_g1 = dgl.node_subgraph(compact_g, reshuffle_nodes)
-    compact_g1.ndata['orig_id'] = compact_g.ndata['orig_id'][reshuffle_nodes]
-    compact_g1.ndata[dgl.NTYPE] = compact_g.ndata[dgl.NTYPE][reshuffle_nodes]
-    compact_g1.ndata[dgl.NID] = compact_g.ndata[dgl.NID][reshuffle_nodes]
-    compact_g1.ndata['inner_node'] = compact_g.ndata['inner_node'][reshuffle_nodes]
-    compact_g1.edata['orig_id'] = compact_g.edata['orig_id'][compact_g1.edata[dgl.EID]]
-    compact_g1.edata[dgl.ETYPE] = compact_g.edata[dgl.ETYPE][compact_g1.edata[dgl.EID]]
-    compact_g1.edata['inner_edge'] = compact_g.edata['inner_edge'][compact_g1.edata[dgl.EID]]
-
-    # reshuffle edges on ETYPE as node_subgraph relabels edges
-    idx = th.argsort(compact_g1.edata[dgl.ETYPE])
-    u, v = compact_g1.edges()
-    u = u[idx]
-    v = v[idx]
-    compact_g2 = dgl.graph((u, v))
-    compact_g2.ndata['orig_id'] = compact_g1.ndata['orig_id']
-    compact_g2.ndata[dgl.NTYPE] = compact_g1.ndata[dgl.NTYPE]
-    compact_g2.ndata[dgl.NID] = compact_g1.ndata[dgl.NID]
-    compact_g2.ndata['inner_node'] = compact_g1.ndata['inner_node']
-    compact_g2.edata['orig_id'] = compact_g1.edata['orig_id'][idx]
-    compact_g2.edata[dgl.ETYPE] = compact_g1.edata[dgl.ETYPE][idx]
-    compact_g2.edata['inner_edge'] = compact_g1.edata['inner_edge'][idx]
-    compact_g2.edata[dgl.EID] = th.arange(
-        edgeid_offset, edgeid_offset + compact_g2.number_of_edges(), dtype=th.int64)
-    edgeid_offset += compact_g2.number_of_edges()
-
-    return compact_g2, node_map_val, edge_map_val, ntypes_map, etypes_map
-
-def create_metadata_json(graph_name, num_nodes, num_edges, num_parts, node_map_val, \
-                            edge_map_val, ntypes_map, etypes_map, output_dir ):
-    """
-    Auxiliary function to create json file for the graph partition metadata
-
-    Parameters:
-    -----------
-    graph_name : string
-        name of the graph
-    num_nodes : int
-        no. of nodes in the graph partition
-    num_edges : int
-        no. of edges in the graph partition
-    num_parts : int
-        total no. of partitions of the original graph
-    node_map_val : dictionary
-        map between node types and the range of global node ids used
-    edge_map_val : dictionary
-        map between edge types and the range of global edge ids used
-    ntypes_map : dictionary
-        map between node type(string)  and node_type_id(int)
-    etypes_map : dictionary
-        map between edge type(string)  and edge_type_id(int)
-    output_dir : string
-        directory where the output files are to be stored 
-
-    Returns:
-    --------
-    dictionary
-        map describing the graph information
-
-    """
-    part_metadata = {'graph_name': graph_name,
-                     'num_nodes': num_nodes,
-                     'num_edges': num_edges,
-                     'part_method': 'metis',
-                     'num_parts': num_parts,
-                     'halo_hops': 1,
-                     'node_map': node_map_val,
-                     'edge_map': edge_map_val,
-                     'ntypes': ntypes_map,
-                     'etypes': etypes_map}
-
-    for part_id in range(num_parts):
-        part_dir = 'part' + str(part_id)
-        node_feat_file = os.path.join(part_dir, "node_feat.dgl")
-        edge_feat_file = os.path.join(part_dir, "edge_feat.dgl")
-        part_graph_file = os.path.join(part_dir, "graph.dgl")
-        part_metadata['part-{}'.format(part_id)] = {'node_feats': node_feat_file,
-                                                'edge_feats': edge_feat_file,
-                                                'part_graph': part_graph_file}
-    return part_metadata

tools/distpartitioning/data_proc_pipeline.py

-import argparse
-import numpy as np
-import torch.multiprocessing as mp
-from initialize import proc_exec, single_dev_init, multi_dev_init
-
-def log_params(params): 
-    """ Print all the command line arguments for debugging purposes.
-
-    Parameters:
-    -----------
-    params: argparse object
-        Argument Parser structure listing all the pre-defined parameters
-    """
-    print('Input Dir: ', params.input_dir)
-    print('Graph Name: ', params.graph_name)
-    print('Schema File: ', params.schema)
-    print('No. partitions: ', params.num_parts)
-    print('No. node weights: ', params.num_node_weights)
-    print('Workspace dir: ', params.workspace)
-    print('Node Attr Type: ', params.node_attr_dtype)
-    print('Edge Attr Dtype: ', params.edge_attr_dtype)
-    print('Output Dir: ', params.output)
-    print('Removed Edges File: ', params.removed_edges)
-    print('WorldSize: ', params.world_size)
-    print('Nodes File: ', params.nodes_file)
-    print('Edges File: ', params.edges_file)
-    print('Node feats: ', params.node_feats_file)
-    print('Edge feats: ', params.edge_feats_file)
-    print('Metis partitions: ', params.partitions_file)
-    print('Exec Type: ', params.exec_type)
-
-def start_local_run(params): 
-    """ Main function for distributed implementation on a single machine
-
-    Parameters:
-    -----------
-    params : argparser object
-        Argument Parser structure with pre-determined arguments as defined
-        at the bottom of this file.
-    """
-    log_params(params)
-    processes = []
-    mp.set_start_method("spawn")
-
-    #Invoke `target` function from each of the spawned process for distributed 
-    #implementation
-    for rank in range(params.world_size):
-        p = mp.Process(target=single_dev_init, args=(rank, params.world_size, proc_exec, params))
-        p.start()
-        processes.append(p)
-
-    for p in processes:
-        p.join()
-
-if __name__ == "__main__":
-    """ 
-    Start of execution from this point. 
-    Invoke the appropriate function to begin execution
-    """
-    #arguments which are already needed by the existing implementation of convert_partition.py
-    parser = argparse.ArgumentParser(description='Construct graph partitions')
-    parser.add_argument('--input-dir', required=True, type=str,
-                     help='The directory path that contains the partition results.')
-    parser.add_argument('--graph-name', required=True, type=str,
-                     help='The graph name')
-    parser.add_argument('--schema', required=True, type=str,
-                     help='The schema of the graph')
-    parser.add_argument('--num-parts', required=True, type=int,
-                     help='The number of partitions')
-    parser.add_argument('--num-node-weights', required=True, type=int,
-                     help='The number of node weights used by METIS.')
-    parser.add_argument('--workspace', type=str, default='/tmp',
-                    help='The directory to store the intermediate results')
-    parser.add_argument('--node-attr-dtype', type=str, default=None,
-                    help='The data type of the node attributes')
-    parser.add_argument('--edge-attr-dtype', type=str, default=None,
-                    help='The data type of the edge attributes')
-    parser.add_argument('--output', required=True, type=str,
-                    help='The output directory of the partitioned results')
-    parser.add_argument('--removed-edges', help='a file that contains the removed self-loops and duplicated edges',
-                    default=None, type=str)
-    parser.add_argument('--exec-type', type=int, default=0,
-                    help='Use 0 for single machine run and 1 for distributed execution')
-
-    #arguments needed for the distributed implementation
-    parser.add_argument('--world-size', help='no. of processes to spawn',
-                    default=1, type=int, required=True)
-    parser.add_argument('--nodes-file', help='filename of the nodes metadata', 
-                    default=None, type=str, required=True)
-    parser.add_argument('--edges-file', help='filename of the nodes metadata', 
-                    default=None, type=str, required=True)
-    parser.add_argument('--node-feats-file', help='filename of the nodes features', 
-                    default=None, type=str, required=True)
-    parser.add_argument('--edge-feats-file', help='filename of the nodes metadata', 
-                    default=None, type=str )
-    parser.add_argument('--partitions-file', help='filename of the output of dgl_part2 (metis partitions)',
-                    default=None, type=str)
-    params = parser.parse_args()
-
-    #invoke the starting function here.
-    if(params.exec_type == 0):
-        start_local_run(params)
-    else:
-        multi_dev_init(params)

tools/distpartitioning/globalids.py

-import numpy as np
-import torch
-import operator
-import itertools
-import constants
-from gloo_wrapper import allgather_sizes, alltoall_cpu, alltoallv_cpu
-
-def get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data):
-    """ 
-    For nodes which are not owned by the current rank, whose global_nid <-> shuffle_global-nid mapping
-    is not present at the current rank, this function retrieves their shuffle_global_ids from the owner rank
-
-    Parameters: 
-    -----------
-    rank : integer
-        rank of the process
-    world_size : integer
-        total no. of ranks configured
-    global_nids_ranks : list
-        list of numpy arrays (of global_nids), index of the list is the rank of the process
-                    where global_nid <-> shuffle_global_nid mapping is located. 
-    node_data : dictionary
-        node_data is a dictionary with keys as column names and values as numpy arrays
-
-    Returns:
-    --------
-    numpy ndarray
-        where the column-0 are global_nids and column-1 are shuffle_global_nids which are retrieved
-        from other processes. 
-    """
-    #build a list of sizes (lengths of lists)
-    sizes = [len(x) for x in global_nids_ranks]
-
-    #compute total_nodes whose mappings should be resolved, between orig-node-id <-> global-id
-    total_nodes = np.sum(sizes)
-    if (total_nodes == 0):
-        print('Rank: ', rank, ' -- All mappings are present locally... No need for to send any info.')
-        return None
-
-    #determine the no. of global_node_ids to send and receive and perform alltoall
-    send_counts = list(torch.Tensor(sizes).type(dtype=torch.int64).chunk(world_size))
-    recv_counts = list(torch.zeros([world_size], dtype=torch.int64).chunk(world_size))
-    alltoall_cpu(rank, world_size, recv_counts, send_counts)
-
-    #allocate buffers to receive node-ids
-    recv_nodes = []
-    for i in recv_counts:
-        recv_nodes.append(torch.zeros([i.item()], dtype=torch.int64))
-
-    #form the outgoing message
-    send_nodes = []
-    for i in range(world_size):
-        send_nodes.append(torch.Tensor(global_nids_ranks[i]).type(dtype=torch.int64))
-
-    #send-recieve messages
-    alltoallv_cpu(rank, world_size, recv_nodes, send_nodes)
-
-    # allocate buffers to receive global-ids
-    recv_shuffle_global_nids = []
-    for i in sizes: 
-        recv_shuffle_global_nids.append(torch.zeros((i), dtype=torch.int64))
-
-    # Use node_data to lookup global id to send over.
-    send_nodes = []
-    for proc_i_nodes in recv_nodes:
-        #list of node-ids to lookup
-        global_nids = proc_i_nodes.numpy()
-        if (len(global_nids) != 0):
-            common, ind1, ind2 = np.intersect1d(node_data[constants.GLOBAL_NID], global_nids, return_indices=True)
-            values = node_data[constants.SHUFFLE_GLOBAL_NID][ind1]
-            send_nodes.append(torch.Tensor(values).type(dtype=torch.int64))
-        else:
-            send_nodes.append(torch.empty((0,), dtype=torch.int64))
-
-    #send receive global-ids
-    alltoallv_cpu(rank, world_size, recv_shuffle_global_nids, send_nodes)
-
-    shuffle_global_nids = [x.numpy() for x in recv_shuffle_global_nids]
-    global_nids = [x for x in global_nids_ranks]
-    return np.column_stack((np.concatenate(global_nids), np.concatenate(shuffle_global_nids)))
-
-
-def get_shuffle_global_nids_edges(rank, world_size, edge_data, node_part_ids, node_data):
-    """
-    Edges which are owned by this rank, may have global_nids whose shuffle_global_nids are NOT present locally.
-    This function retrieves shuffle_global_nids for such global_nids.
-
-    Parameters: 
-    -----------
-    rank : integer
-        rank of the process
-    world_size : integer
-        total no. of processes used
-    edge_data : numpy ndarray
-        edge_data (augmented) as read from the xxx_edges.txt file
-    node_part_ids : numpy array 
-        list of partition ids indexed by global node ids.
-    node_data : dictionary
-        node_data, is a dictionary with keys as column_names and values as numpy arrays
-    """
-
-    #determine unique node-ids present locally
-    global_nids = np.sort(np.unique(np.concatenate((edge_data[constants.GLOBAL_SRC_ID], edge_data[constants.GLOBAL_DST_ID], node_data[constants.GLOBAL_NID]))))
-
-    #determine the rank which owns orig-node-id <-> partition/rank mappings
-    part_ids = node_part_ids[global_nids]
-
-    #form list of lists, each list includes global_nids whose mappings (shuffle_global_nids) needs to be retrieved.
-    #and rank will be the process which owns mappings of these global_nids
-    global_nids_ranks = []
-    for i in range(world_size):
-        if (i == rank):
-            global_nids_ranks.append(np.empty(shape=(0)))
-            continue
-
-        #not_owned_nodes = part_ids[:,0][part_ids[:,1] == i]
-        not_owned_node_ids = np.where(part_ids == i)[0] 
-        if not_owned_node_ids.shape[0] == 0: 
-            not_owned_nodes = np.empty(shape=(0))
-        else: 
-            not_owned_nodes = global_nids[not_owned_node_ids]
-        global_nids_ranks.append(not_owned_nodes)
-
-    #Retrieve Global-ids for respective node owners
-    resolved_global_nids = get_shuffle_global_nids(rank, world_size, global_nids_ranks, node_data)
-
-    #Add global_nid <-> shuffle_global_nid mappings to the received data
-    for i in range(world_size):
-        if (i == rank):
-            own_node_ids = np.where(part_ids == i)[0]
-            own_global_nids = global_nids[own_node_ids]
-            common, ind1, ind2 = np.intersect1d(node_data[constants.GLOBAL_NID], own_global_nids, return_indices=True)
-            my_shuffle_global_nids = node_data[constants.SHUFFLE_GLOBAL_NID][ind1]
-            local_mappings = np.column_stack((own_global_nids, my_shuffle_global_nids))
-            resolved_global_nids = np.concatenate((resolved_global_nids, local_mappings))
-
-    #form a dictionary of mappings between orig-node-ids and global-ids
-    resolved_mappings = dict(zip(resolved_global_nids[:,0], resolved_global_nids[:,1]))
-
-    #determine global-ids for the orig-src-id and orig-dst-id
-    shuffle_global_src_id = [resolved_mappings[x] for x in edge_data[constants.GLOBAL_SRC_ID]]
-    shuffle_global_dst_id = [resolved_mappings[x] for x in edge_data[constants.GLOBAL_DST_ID]]
-    edge_data[constants.SHUFFLE_GLOBAL_SRC_ID] = np.array(shuffle_global_src_id, dtype=np.int64)
-    edge_data[constants.SHUFFLE_GLOBAL_DST_ID] = np.array(shuffle_global_dst_id, dtype=np.int64)
-
-def assign_shuffle_global_nids_nodes(rank, world_size, node_data):
-    """
-    Utility function to assign shuffle global ids to nodes at a given rank
-    node_data gets converted from [ntype, global_type_nid, global_nid]
-    to [shuffle_global_nid, ntype, global_type_nid, global_nid, part_local_type_nid]
-    where shuffle_global_nid : global id of the node after data shuffle
-            ntype : node-type as read from xxx_nodes.txt
-            global_type_nid : node-type-id as read from xxx_nodes.txt
-            global_nid : node-id as read from xxx_nodes.txt, implicitly 
-                            this is the line no. in the file
-            part_local_type_nid : type_nid assigned by the current rank within its scope
-            
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    world_size : integer
-        total number of processes used in the process group
-    ntype_counts: list of tuples
-        list of tuples (x,y), where x=ntype and y=no. of nodes whose shuffle_global_nids are needed
-    node_data : dictionary
-        node_data is a dictionary with keys as column names and values as numpy arrays
-    """
-    # Compute prefix sum to determine node-id offsets
-    prefix_sum_nodes = allgather_sizes([node_data[constants.GLOBAL_NID].shape[0]], world_size)
-
-    # assigning node-ids from localNodeStartId to (localNodeEndId - 1)
-    # Assuming here that the nodeDataArr is sorted based on the nodeType.
-    shuffle_global_nid_start = prefix_sum_nodes[rank]
-    shuffle_global_nid_end = prefix_sum_nodes[rank + 1]
-
-    # add a column with global-ids (after data shuffle)
-    shuffle_global_nids = np.arange(shuffle_global_nid_start, shuffle_global_nid_end, dtype=np.int64)
-    node_data[constants.SHUFFLE_GLOBAL_NID] = shuffle_global_nids
-
-
-def assign_shuffle_global_nids_edges(rank, world_size, edge_data):
-    """
-    Utility function to assign shuffle_global_eids to edges
-    edge_data gets converted from [global_src_nid, global_dst_nid, global_type_eid, etype]
-    to [shuffle_global_src_nid, shuffle_global_dst_nid, global_src_nid, global_dst_nid, global_type_eid, etype]
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the current process
-    world_size : integer
-        total count of processes in execution
-    etype_counts : list of tuples
-        list of tuples (x,y), x = rank, y = no. of edges
-    edge_data : numpy ndarray
-        edge data as read from xxx_edges.txt file
-
-    Returns:
-    --------
-    integer
-        shuffle_global_eid_start, which indicates the starting value from which shuffle_global-ids are assigned to edges
-        on this rank
-    """
-    #get prefix sum of edge counts per rank to locate the starting point
-    #from which global-ids to edges are assigned in the current rank
-    prefix_sum_edges = allgather_sizes([edge_data[constants.GLOBAL_SRC_ID].shape[0]], world_size)
-    shuffle_global_eid_start = prefix_sum_edges[rank]
-    shuffle_global_eid_end = prefix_sum_edges[rank + 1]
-
-    # assigning edge-ids from localEdgeStart to (localEdgeEndId - 1)
-    # Assuming here that the edge_data is sorted by edge_type
-    shuffle_global_eids = np.arange(shuffle_global_eid_start, shuffle_global_eid_end, dtype=np.int64)
-    edge_data[constants.SHUFFLE_GLOBAL_EID] = shuffle_global_eids
-    return shuffle_global_eid_start

tools/distpartitioning/gloo_wrapper.py

-import numpy as np
-import torch
-import torch.distributed as dist
-
-def allgather_sizes(send_data, world_size):
-    """ 
-    Perform all gather on list lengths, used to compute prefix sums
-    to determine the offsets on each ranks. This is used to allocate
-    global ids for edges/nodes on each ranks.
-
-    Parameters
-    ----------
-    send_data : numpy array 
-        Data on which allgather is performed.
-    world_size : integer
-        No. of processes configured for execution
-
-    Returns : 
-    ---------
-        numpy array
-            array with the prefix sum
-    """
-
-    #compute the length of the local data
-    send_length = len(send_data)
-    out_tensor = torch.as_tensor(send_data, dtype=torch.int64)
-    in_tensor = [torch.zeros(send_length, dtype=torch.int64) 
-                    for _ in range(world_size)]
-
-    #all_gather message
-    dist.all_gather(in_tensor, out_tensor)
-
-    #gather sizes in on array to return to the invoking function
-    rank_sizes = np.zeros(world_size + 1)
-    count = rank_sizes[0]
-    for i, t in enumerate(in_tensor): 
-        count += t.item()
-        rank_sizes[i+1] = count
-
-    return rank_sizes
-
-def alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list):
-    """
-    Each process scatters list of input tensors to all processes in a cluster
-    and return gathered list of tensors in output list. The tensors should have the same shape.
-
-    Parameters
-    ----------
-    rank : int
-        The rank of current worker
-    world_size : int
-        The size of the entire
-    output_tensor_list : List of tensor
-        The received tensors
-    input_tensor_list : List of tensor
-        The tensors to exchange
-    """
-    input_tensor_list = [tensor.to(torch.device('cpu')) for tensor in input_tensor_list]
-    for i in range(world_size):
-        dist.scatter(output_tensor_list[i], input_tensor_list if i == rank else [], src=i)
-
-def alltoallv_cpu(rank, world_size, output_tensor_list, input_tensor_list):
-    """
-    Each process scatters list of input tensors to all processes in a cluster
-    and return gathered list of tensors in output list.
-
-    Parameters
-    ----------
-    rank : int
-        The rank of current worker
-    world_size : int
-        The size of the entire
-    output_tensor_list : List of tensor
-        The received tensors
-    input_tensor_list : List of tensor
-        The tensors to exchange
-    """
-    # send tensor to each target trainer using torch.distributed.isend
-    # isend is async
-    senders = []
-    for i in range(world_size):
-        if i == rank:
-            output_tensor_list[i] = input_tensor_list[i].to(torch.device('cpu'))
-        else:
-            sender = dist.isend(input_tensor_list[i].to(torch.device('cpu')), dst=i)
-            senders.append(sender)
-
-    for i in range(world_size):
-        if i != rank:
-            dist.recv(output_tensor_list[i], src=i)
-
-    torch.distributed.barrier()
-
-def gather_metadata_json(metadata, rank, world_size): 
-    """ 
-    Gather an object (json schema on `rank`)
-    Parameters:
-    -----------
-    metadata : json dictionary object
-        json schema formed on each rank with graph level data. 
-        This will be used as input to the distributed training in the later steps.
-    Returns:
-    --------
-    list : list of json dictionary objects
-        The result of the gather operation, which is the list of json dicitonary 
-        objects from each rank in the world
-    """
-
-    #Populate input obj and output obj list on rank-0 and non-rank-0 machines
-    input_obj = None if rank == 0 else metadata
-    output_objs = [None for _ in range(world_size)] if rank == 0 else None
-
-    #invoke the gloo method to perform gather on rank-0
-    dist.gather_object(input_obj, output_objs, dst=0)
-    return output_objs

tools/distpartitioning/initialize.py

-import os
-import json
-import numpy as np
-import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-import constants
-
-from timeit import default_timer as timer
-from datetime import timedelta
-from utils import augment_node_data, augment_edge_data,\
-                  read_nodes_file, read_edges_file,\
-                  read_node_features_file, read_edge_features_file,\
-                  read_partitions_file, read_json,\
-                  get_node_types, write_metadata_json, write_dgl_objects
-from globalids import assign_shuffle_global_nids_nodes, assign_shuffle_global_nids_edges,\
-                      get_shuffle_global_nids_edges
-from gloo_wrapper import gather_metadata_json
-from convert_partition import create_dgl_object, create_metadata_json
-
-def send_node_data(rank, node_data, part_ids):
-    """ 
-    Function to send node_data to non-rank-0 processes. 
-    
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    node_data : numpy ndarray
-        node_data, in the augmented form
-    part_ids : python list 
-        list of unique ranks/partition-ids
-    """
-
-    for part_id in part_ids: 
-        if part_id == rank: 
-            continue
-        
-        #extract <node_type>, <global_type_nid>, <global_nid>
-        #which belong to `part_id`
-        send_data_idx = (node_data[constants.OWNER_PROCESS] == part_id) 
-        idx = send_data_idx.reshape(node_data[constants.GLOBAL_NID].shape[0])
-        filt_data = np.column_stack((node_data[constants.NTYPE_ID], \
-                                    node_data[constants.GLOBAL_TYPE_NID], \
-                                    node_data[constants.GLOBAL_NID]))
-        filt_data = filt_data[idx == 1]
-
-        #prepare tensor to send
-        send_size = filt_data.shape
-        size_tensor = torch.tensor(filt_data.shape, dtype=torch.int64)
-
-        # Send size first, so that the part-id (rank)
-        # can create appropriately sized buffers
-        dist.send(size_tensor, dst=part_id)
-
-        #send actual node_data to part-id rank
-        start = timer()
-        send_tensor = torch.from_numpy(filt_data.astype(np.int64))
-        dist.send(send_tensor, dst=part_id)
-        end = timer()
-        print('Rank: ', rank, ' Sent data size: ', filt_data.shape, \
-                ', to Process: ', part_id, 'in: ', timedelta(seconds = end - start))
-
-def send_edge_data(rank, edge_data, part_ids): 
-    """ 
-    Function to send edge data to non-rank-0 processes
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    edge_data : numpy ndarray
-        edge_data, in the augmented form
-    part_ids : python list
-        list of unique ranks/partition-ids
-    """
-    for part_id in part_ids: 
-        if part_id == rank: 
-            continue
-
-        #extract global_sid, global_dit, global_type_eid, etype_id
-        send_data = (edge_data[constants.OWNER_PROCESS] == part_id) 
-        idx = send_data.reshape(edge_data[constants.GLOBAL_SRC_ID].shape[0])
-        filt_data = np.column_stack((edge_data[constants.GLOBAL_SRC_ID][idx == 1], \
-                                    edge_data[constants.GLOBAL_DST_ID][idx == 1], \
-                                    edge_data[constants.GLOBAL_TYPE_EID][idx == 1], \
-                                    edge_data[constants.ETYPE_ID][idx == 1]))
-
-        #send shape
-        send_size = filt_data.shape
-        size_tensor = torch.tensor(filt_data.shape, dtype=torch.int64)
-
-        # Send size first, so that the rProc can create appropriately sized tensor
-        dist.send(size_tensor, dst=part_id)
-
-        start = timer()
-        send_tensor = torch.from_numpy(filt_data)
-        dist.send(send_tensor, dst=part_id)
-        end = timer()
-
-        print('Rank: ', rank, ' Time to send Edges to proc: ', part_id, \
-                ' is : ', timedelta(seconds = end - start))
-
-def send_node_features(rank, node_data, node_features, part_ids, ntype_map):
-    """ 
-    Function to send node_features data to non-rank-0 processes.
-    
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    node_data : numpy ndarray of int64
-        node_data, read from the xxx_nodes.txt file
-    node_features : numpy ndarray of floats 
-        node_features, data from the node_feats.dgl
-    part_ids : list
-        list of unique ranks/partition-ids
-    ntype_map : dictionary 
-        mappings between ntype_name -> ntype
-    """
-
-    node_features_out = []
-    for part_id in part_ids: 
-        if part_id == rank: 
-            node_features_out.append(None)
-            continue
-
-        part_node_features = {}
-        for ntype_name, ntype in ntype_map.items():
-            
-            if (ntype_name +'/feat' in node_features) and (node_features[ntype_name+'/feat'].shape[0] > 0): 
-                #extract orig_type_node_id
-                idx = (node_data[constants.OWNER_PROCESS] == part_id) & (node_data[constants.NTYPE_ID] == ntype) 
-                filt_global_type_nids = node_data[constants.GLOBAL_TYPE_NID][idx] # extract global_ntype_id here
-                part_node_features[ntype_name+'/feat'] = node_features[ntype_name+'/feat'][filt_global_type_nids]
-
-        #accumulate subset of node_features targetted for part-id rank
-        node_features_out.append(part_node_features)
-
-    #send data
-    output_list = [None]
-    start = timer ()
-    dist.scatter_object_list(output_list, node_features_out, src=0)
-    end = timer ()
-    print('Rank: ', rank, ', Done sending Node Features to: ', part_id, \
-            ' in: ', timedelta(seconds = end - start))
-
-def send_data(rank, node_data, node_features, edge_data, node_part_ids, ntypes_map): 
-    """ 
-    Wrapper function to send graph data to non-rank-0 processes.
-    
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    node_data : numpy ndarray 
-        node_data, augmented, from xxx_nodes.txt file
-    node_features : numpy ndarray 
-        node_features, data from the node_feats.dgl
-    edge_data : numpy ndarray
-        edge_data, augmented, from xxx_edges.txt file
-    node_part_ids : ndarray  
-        array of part_ids indexed by global_nid
-    ntype_map : dictionary 
-        mappings between ntype_name -> ntype_id
-    """
-
-    part_ids = np.unique(node_part_ids)
-    part_ids.sort ()
-    print('Rank: ', rank, ', Unique partitions: ', part_ids)
-
-    send_node_data(rank, node_data, part_ids)
-    send_edge_data(rank, edge_data, part_ids) 
-    send_node_features(rank, node_data, node_features, part_ids, ntypes_map)
-
-def recv_data(rank, shape, dtype): 
-    """ 
-    Auxiliary function to receive a multi-dimensional tensor, used by the
-    non-rank-0 processes. 
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    shape : tuple of integers
-        shape of the received data
-    dtype : integer
-        type of the received data
-
-    Returns: 
-    --------
-    numpy array
-        received data after completing 'recv' gloo primitive.
-    """
-
-    #First receive the size of the data to be received from rank-0 process
-    recv_tensor_shape = torch.zeros(shape, dtype=torch.int64)
-    dist.recv(recv_tensor_shape, src=0)
-    recv_shape = list(map(lambda x: int(x), recv_tensor_shape))
-
-    #Receive the data message here for nodes here. 
-    recv_tensor_data = torch.zeros(recv_shape, dtype=dtype)
-    dist.recv(recv_tensor_data, src=0)
-    return recv_tensor_data.numpy()
-
-def recv_node_data(rank, shape, dtype): 
-    """ 
-    Function to receive node_data, used by non-rank-0 processes.
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    shape : tuple of integers 
-        shape of the received data
-    dtype : integer
-        type of the received data
-
-    Returns:
-    --------
-    numpy array
-        result of the 'recv' gloo primitive
-    """
-    return recv_data(rank, shape, dtype)
-
-def recv_edge_data(rank, shape, dtype): 
-    """
-    Function to receive edge_data, used by non-rank0 processes. 
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    shape : tuple of integers
-        shape of the received data
-    dtype : integer
-        type of the received data
-
-    Returns:
-    --------
-    numpy array
-        result of the 'recv' operation
-    """
-    return recv_data(rank, shape, dtype)
-
-def recv_node_features_obj(rank, world_size): 
-    """
-    Function to receive node_feautres as an object, as read from the node_feats.dgl file.
-    This is used by non-rank-0 processes. 
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    world_size : integer
-        no. of processes used
-
-    Returns:
-    --------
-    numpy ndarray
-        node_feature data, of floats, as received by the scatter_object_list function
-    """
-    send_objs = [None for _ in range(world_size)]
-    recv_obj = [None]
-    dist.scatter_object_list(recv_obj, send_objs, src=0)
-
-    node_features = recv_obj[0]
-    return node_features
-
-def read_graph_files(rank, params, node_part_ids): 
-    """
-    Read the files and return the data structures
-    Node data as read from files, which is in the following format: 
-        <node_type> <weight1> <weight2> <weight3> <weight4> <global_type_nid> <attributes>
-    is converted to 
-        <node_type> <weight1> <weight2> <weight3> <weight4> <global_type_nid> <nid> <recv_proc>
-    Edge data as read from files, which is in the following format: 
-        <global_src_id> <global_dst_id> <global_type_eid> <edge_type> <attributes>
-    is converted to the following format in this function:
-        <global_src_id> <global_dst_id> <global_type_eid> <edge_type> <recv_proc>
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    params : argparser object
-        argument parser data structure to access command line arguments
-    node_part_ids : numpy array 
-        array of part_ids indexed by global_nid
-
-    Returns:
-    --------
-    numpy ndarray
-        integer node_data with additional columns added
-    numpy ndarray
-        floats, node_features as read from the node features file
-    numpy ndarray
-        integer, edge_data with additional columns added
-    numpy ndarray
-        floats, edge_features are read from the edge feature file
-    """
-    node_data = read_nodes_file(params.input_dir+'/'+params.nodes_file)
-    augment_node_data(node_data, node_part_ids)
-    print('Rank: ', rank, ', Completed loading nodes data: ', node_data[constants.GLOBAL_TYPE_NID].shape)
-
-    edge_data = read_edges_file(params.input_dir+'/'+params.edges_file, None)
-    print('Rank: ', rank, ', Completed loading edge data: ', edge_data[constants.GLOBAL_SRC_ID].shape)
-    edge_data = read_edges_file(params.input_dir+'/'+params.removed_edges, edge_data)
-    augment_edge_data(edge_data, node_part_ids)
-    print('Rank: ', rank, ', Completed adding removed edges : ', edge_data[constants.GLOBAL_SRC_ID].shape)
-
-    node_features = {}
-    node_features = read_node_features_file( params.input_dir+'/'+params.node_feats_file )
-    print('Rank: ', rank, ', Completed loading node features reading from file ', len(node_features))
-
-    edge_features = {}
-    #edge_features = read_edge_features_file( params.input_dir+'/'+params.edge_feats_file )
-    #print( 'Rank: ', rank, ', Completed edge features reading from file ', len(edge_features) )
-
-    return node_data, node_features, edge_data, edge_features
-
-def proc_exec(rank, world_size, params):
-    """ 
-    `main` function for each rank in the distributed implementation.
-
-    This function is used when one-machine is used for executing the entire pipeline. 
-    In this case, all the gloo-processes will exist on the same machine. Also, this function
-    expects that the graph input files are in single file-format. Nodes, edges, node-features and 
-    edge features each will have their own file describing the appropriate parts of the input graph.
-
-    Parameters: 
-    -----------
-    rank : integer
-        rank of the current process
-    world_size : integer
-        total no. of ranks
-    params : argparser object
-        argument parser structure to access values passed from command line
-    """
-
-    #Read METIS partitions
-    node_part_ids = read_partitions_file(params.input_dir+'/'+params.partitions_file)
-    print('Rank: ', rank, ', Completed loading metis partitions: ', len(node_part_ids))
-
-    #read graph schema, get ntype_map(dict for ntype to ntype-id lookups) and ntypes list
-    schema_map = read_json(params.input_dir+'/'+params.schema)
-    ntypes_map, ntypes = get_node_types(schema_map)
-
-    # Rank-0 process will read the graph input files (nodes, edges, node-features and edge-features). 
-    # it will uses metis partitions, node-id to partition-id mappings, to determine the node and edge
-    # ownership and sends out data to all the other non rank-0 processes.
-    if rank == 0: 
-        #read input graph files
-        node_data, node_features, edge_data, edge_features = read_graph_files(rank, params, node_part_ids)
-
-        # order node_data by node_type before extracting node features. 
-        # once this is ordered, node_features are automatically ordered and 
-        # can be assigned contiguous ids starting from 0 for each type. 
-        #node_data = node_data[node_data[:, 0].argsort()]
-        sorted_idx = node_data[constants.NTYPE_ID].argsort()
-        for k, v in node_data.items(): 
-            node_data[k] = v[sorted_idx]
-
-        print('Rank: ', rank, ', node_data: ', len(node_data))
-        print('Rank: ', rank, ', node_features: ', len(node_features))
-        print('Rank: ', rank, ', edge_data: ', len(edge_data))
-        #print('Rank: ', rank, ', edge_features : ',len( edge_features))
-        print('Rank: ', rank, ', partitions : ', len(node_part_ids))
-
-        # shuffle data
-        send_data(rank, node_data, node_features, edge_data, node_part_ids, ntypes_map)
-
-        #extract features here for rank-0
-        for name, ntype_id in ntypes_map.items(): 
-            ntype = name + '/feat'
-            if(ntype in node_features): 
-                idx = node_data[constants.GLOBAL_TYPE_NID][(node_data[constants.NTYPE_ID] == ntype_id) & (node_data[constants.OWNER_PROCESS] == rank)]
-                node_features[ntype] = node_features[ntype][idx]
-
-        # Filter data owned by rank-0
-        #extract only ntype, global_type_nid, global_nid 
-        idx = np.where(node_data[constants.OWNER_PROCESS] == 0)[0]
-        for k, v in node_data.items(): 
-            node_data[k] = v[idx]
-
-        #extract only global_src_id, global_dst_id, global_type_eid etype
-        idx = np.where(edge_data[constants.OWNER_PROCESS] == 0)[0]
-        for k, v in edge_data.items(): 
-            edge_data[k] = v[idx]
-    else: 
-        #Non-rank-0 processes, receives nodes, edges, node-features and edge feautres from rank-0
-        # process and creates appropriate data structures. 
-        rcvd_node_data = recv_node_data(rank, 2, torch.int64)
-        node_data = {}
-        node_data[constants.NTYPE_ID] = rcvd_node_data[:,0]
-        node_data[constants.GLOBAL_TYPE_NID] = rcvd_node_data[:,1]
-        node_data[constants.GLOBAL_NID] = rcvd_node_data[:,2]
-
-        rcvd_edge_data = recv_edge_data(rank, 2, torch.int64)
-        edge_data = {}
-        edge_data[constants.GLOBAL_SRC_ID] = rcvd_edge_data[:,0]
-        edge_data[constants.GLOBAL_DST_ID] = rcvd_edge_data[:,1]
-        edge_data[constants.GLOBAL_TYPE_EID] = rcvd_edge_data[:,2]
-        edge_data[constants.ETYPE_ID] = rcvd_edge_data[:,3]
-
-        node_features = recv_node_features_obj(rank, world_size)
-        edge_features = {}
-
-    # From this point onwards, all the processes will follow the same execution logic and
-    # process the data which is owned by the current process. 
-    # At this time, all the processes will have all the data which it owns (nodes, edges, 
-    # node-features and edge-features).
-
-    #syncronize
-    dist.barrier()
-
-    # assign shuffle_global ids to nodes
-    assign_shuffle_global_nids_nodes(rank, world_size, node_data)
-    print('Rank: ', rank, ' Done assign Global ids to nodes...')
-
-    #sort edge_data by etype
-    sorted_idx = edge_data[constants.ETYPE_ID].argsort()
-    for k, v in edge_data.items(): 
-        edge_data[k] = v[sorted_idx]
-
-    # assign shuffle_global ids to edges
-    shuffle_global_eid_start = assign_shuffle_global_nids_edges(rank, world_size, edge_data)
-    print('Rank: ', rank, ' Done assign Global ids to edges...')
-
-    # resolve shuffle_global ids for nodes which are not locally owned
-    get_shuffle_global_nids_edges(rank, world_size, edge_data, node_part_ids, node_data)
-    print('Rank: ', rank, ' Done retrieving Global Node Ids for non-local nodes... ')
-
-    #create dgl objects
-    print('Rank: ', rank, ' Creating DGL objects for all partitions')
-    num_nodes = 0
-    num_edges = shuffle_global_eid_start
-    with open('{}/{}'.format(params.input_dir, params.schema)) as json_file: 
-        schema = json.load(json_file)
-    graph_obj, ntypes_map_val, etypes_map_val, ntypes_map, etypes_map = create_dgl_object(\
-            params.graph_name, params.num_parts, \
-            schema, rank, node_data, edge_data, num_nodes, num_edges)
-    write_dgl_objects(graph_obj, node_features, edge_features, params.output, rank)
-
-    #get the meta-data 
-    json_metadata = create_metadata_json(params.graph_name, num_nodes, num_edges, params.num_parts, ntypes_map_val, \
-                            etypes_map_val, ntypes_map, etypes_map, params.output)
-
-    if (rank == 0): 
-        #get meta-data from all partitions and merge them on rank-0
-        metadata_list = gather_metadata_json(json_metadata, rank, world_size)
-        metadata_list[0] = json_metadata
-        write_metadata_json(metadata_list, params.output, params.graph_name)
-    else: 
-        #send meta-data to Rank-0 process
-        gather_metadata_json(json_metadata, rank, world_size)
-
-def single_dev_init(rank, world_size, func_exec, params, backend="gloo"):
-    """
-    Init. function which is run by each process in the Gloo ProcessGroup
-
-    Parameters:
-    -----------
-    rank : integer
-        rank of the process
-    world_size : integer
-        number of processes configured in the Process Group
-    proc_exec : function name
-        function which will be invoked which has the logic for each process in the group
-    params : argparser object
-        argument parser object to access the command line arguments
-    backend : string
-        string specifying the type of backend to use for communication
-    """
-    os.environ["MASTER_ADDR"] = '127.0.0.1'
-    os.environ["MASTER_PORT"] = '29500'
-
-    #create Gloo Process Group
-    dist.init_process_group(backend, rank=rank, world_size=world_size)
-
-    #Invoke the main function to kick-off each process
-    func_exec(rank, world_size, params)
-
-def multi_dev_init(params):
-    """
-    Function to be invoked when executing data loading pipeline on multiple machines
-
-    Parameters:
-    -----------
-    params : argparser object
-        argparser object providing access to command line arguments.
-    """
-    #init the gloo process group here. 
-    dist.init_prcess_group("gloo", rank=params.rank, world_size=params.world_size)
-    print('[Rank: ', params.rank, '] Done with process group initialization...')
-
-    #invoke the main function here.
-    proc_exec(params.rank, params.world_size, params)
-    print('[Rank: ', params.rank, '] Done with Distributed data processing pipeline processing.')

tools/distpartitioning/utils.py

-import os
-import torch
-import numpy as np
-import json
-import dgl
-import constants
-
-def read_partitions_file(part_file):
-    """
-    Utility method to read metis partitions, which is the output of 
-    pm_dglpart2
-
-    Parameters:
-    -----------
-    part_file : string
-        file name which is the output of metis partitioning
-        algorithm (pm_dglpart2, in the METIS installation).
-        This function expects each line in `part_file` to be formatted as 
-        <global_nid> <part_id>
-        and the contents of this file are sorted by <global_nid>. 
-
-    Returns:
-    --------
-    numpy array
-        array of part_ids and the idx is the <global_nid>
-    """
-    partitions_map = np.loadtxt(part_file, delimiter=' ', dtype=np.int64)
-    #as a precaution sort the lines based on the <global_nid>
-    partitions_map = partitions_map[partitions_map[:,0].argsort()]
-    return partitions_map[:,1]
-
-def read_json(json_file):
-    """
-    Utility method to read a json file schema
-    
-    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
-
-def get_node_types(schema):
-    """ 
-    Utility method to extract node_typename -> node_type mappings
-    as defined by the input schema
-
-    Parameters:
-    -----------
-    schema : dictionary
-        Input schema from which the node_typename -> node_type
-        dictionary is created.
-
-    Returns:
-    --------
-    dictionary, list 
-        dictionary with ntype <-> type_nid mappings
-        list of ntype strings
-    """
-    #Get the node_id ranges from the schema
-    global_nid_ranges = schema['nid']
-    global_nid_ranges = {key: np.array(global_nid_ranges[key]).reshape(1,2)
-                    for key in global_nid_ranges}
-
-    #Create an array with the starting id for each node_type and sort
-    ntypes = [(key, global_nid_ranges[key][0,0]) for key in global_nid_ranges]
-    ntypes.sort(key=lambda e: e[1])
-
-    #Create node_typename -> node_type dictionary
-    ntypes = [e[0] for e in ntypes]
-    ntypes_map = {e: i for i, e in enumerate(ntypes)}
-
-    return ntypes_map, ntypes
-
-def write_metadata_json(metadata_list, output_dir, graph_name):
-    """
-    Merge json schema's from each of the rank's on rank-0. 
-    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
-    """
-    #Initialize global metadata
-    graph_metadata = {}
-
-    #Merge global_edge_ids from each json object in the input list
-    edge_map = {}
-    x = metadata_list[0]["edge_map"]
-    for k in x:
-        edge_map[k] = []
-        for idx in range(len(metadata_list)):
-            edge_map[k].append(metadata_list[idx]["edge_map"][k][0])
-    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"]
-
-    #Merge global_nodeids from each of json object in the input list
-    node_map = {}
-    x = metadata_list[0]["node_map"]
-    for k in x:
-        node_map[k] = []
-        for idx in range(len(metadata_list)):
-            node_map[k].append(metadata_list[idx]["node_map"][k][0])
-    graph_metadata["node_map"] = node_map
-
-    graph_metadata["ntypes"] = metadata_list[0]["ntypes"]
-    graph_metadata["num_edges"] = sum([metadata_list[i]["num_edges"] for i in range(len(metadata_list))])
-    graph_metadata["num_nodes"] = sum([metadata_list[i]["num_nodes"] for i in range(len(metadata_list))])
-    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)):
-        graph_metadata["part-{}".format(i)] = metadata_list[i]["part-{}".format(i)]
-
-    with open('{}/{}.json'.format(output_dir, graph_name), 'w') as outfile: 
-        json.dump(graph_metadata, outfile, sort_keys=True, indent=4)
-
-def augment_edge_data(edge_data, part_ids):
-    """
-    Add partition-id (rank which owns an edge) column to the edge_data.
-    
-    Parameters:
-    -----------
-    edge_data : numpy ndarray
-        Edge information as read from the xxx_edges.txt file
-    part_ids : numpy array
-        array of part_ids indexed by global_nid
-    """
-    edge_data[constants.OWNER_PROCESS] = part_ids[edge_data[constants.GLOBAL_DST_ID]]
-
-def augment_node_data(node_data, part_ids): 
-    """
-    Utility function to add auxilary columns to the node_data numpy ndarray.
-
-    Parameters:
-    -----------
-    node_data : numpy ndarray
-        Node information as read from xxx_nodes.txt file
-    part_ids : numpy array 
-        array of part_ids indexed by global_nid
-    """
-    #add global_nids to the node_data
-    global_nids = np.arange(len(node_data[constants.GLOBAL_TYPE_NID]), dtype=np.int64)
-    node_data[constants.GLOBAL_NID] = global_nids
-
-    #add owner proc_ids to the node_data
-    proc_ids = part_ids[node_data[constants.GLOBAL_NID]]
-    node_data[constants.OWNER_PROCESS] = proc_ids
-
-def read_nodes_file(nodes_file):
-    """
-    Utility function to read xxx_nodes.txt file
-    
-    Parameters:
-    -----------
-    nodesfile : string
-        Graph file for nodes in the input graph
-    
-    Returns:
-    --------
-    dictionary
-        Nodes data stored in dictionary where keys are column names
-        and values are the columns from the numpy ndarray as read from the
-        xxx_nodes.txt file
-    """
-    if nodes_file == "" or nodes_file == None:
-        return None
-
-    # Read the file from here.
-    # Assuming the nodes file is a numpy file
-    # nodes.txt file is of the following format
-    # <node_type> <weight1> <weight2> <weight3> <weight4> <global_type_nid> <attributes>
-    # For the ogb-mag dataset, nodes.txt is of the above format.
-    nodes_data = np.loadtxt(nodes_file, delimiter=' ', dtype='int64')
-    nodes_datadict = {}
-    nodes_datadict[constants.NTYPE_ID] = nodes_data[:,0]
-    nodes_datadict[constants.GLOBAL_TYPE_NID] = nodes_data[:,5]
-    return nodes_datadict
-
-def read_edges_file(edge_file, edge_data_dict):
-    """ 
-    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
-        in a dictionary with key-value pairs as column-names and column-data. 
-    """
-    if edge_file == "" or edge_file == None:
-        return None
-
-    #Read the file from here.
-    #<global_src_id> <global_dst_id> <type_eid> <etype> <attributes>
-    # 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
-
-    edge_data = np.loadtxt(edge_file , delimiter=' ', dtype = 'int64')
-
-    if (edge_data_dict == None): 
-        edge_data_dict = {}
-        edge_data_dict[constants.GLOBAL_SRC_ID] = edge_data[:,0]
-        edge_data_dict[constants.GLOBAL_DST_ID] = edge_data[:,1]
-        edge_data_dict[constants.GLOBAL_TYPE_EID] = edge_data[:,2]
-        edge_data_dict[constants.ETYPE_ID] = edge_data[:,3]
-    else: 
-        edge_data_dict[constants.GLOBAL_SRC_ID] = \
-            np.concatenate((edge_data_dict[constants.GLOBAL_SRC_ID], edge_data[:,0]))
-        edge_data_dict[constants.GLOBAL_DST_ID] = \
-            np.concatenate((edge_data_dict[constants.GLOBAL_DST_ID], edge_data[:,1]))
-        edge_data_dict[constants.GLOBAL_TYPE_EID] = \
-            np.concatenate((edge_data_dict[constants.GLOBAL_TYPE_EID], edge_data[:,2]))
-        edge_data_dict[constants.ETYPE_ID] = \
-            np.concatenate((edge_data_dict[constants.ETYPE_ID], edge_data[:,3]))
-    return edge_data_dict
-
-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
-
-def read_edge_features_file(edge_features_file):
-    """ 
-    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
-
-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
-    node_file     : string 
-        File in which the node information is serialized
-    """
-    dgl.data.utils.save_tensors(node_file, node_features)
-
-def write_edge_features(edge_features, edge_file): 
-    """
-    Utility function to serialize edge_features in edge_file file
-
-    Parameters:
-    -----------
-    edge_features : dictionary
-        dictionary storing etype <-> list of features
-    edge_file     : string 
-        File in which the edge information is serialized
-    """
-    dgl.data.utils.save_tensors(edge_file, edge_features)
-
-def write_graph_dgl(graph_file, graph_obj): 
-    """
-    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
-    """
-    dgl.save_graphs(graph_file, [graph_obj])
-
-def write_dgl_objects(graph_obj, node_features, edge_features, output_dir, part_id): 
-    """
-    Wrapper function to create dgl objects for graph, node-features and edge-features
-    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
-    """
-
-    part_dir = output_dir + '/part' + str(part_id)
-    os.makedirs(part_dir, exist_ok=True)
-    write_graph_dgl(os.path.join(part_dir ,'part'+str(part_id)), graph_obj)
-
-    if node_features != None:
-        write_node_features(node_features, os.path.join(part_dir, "node_feat.dgl"))
-
-    if (edge_features != None):
-        write_edge_features(edge_features, os.path.join(part_dir, "edge_feat.dgl"))