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

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

This reverts commit 7c598aac.

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"))