chore: import upstream snapshot with attribution

tools/README.md

@@ -0,0 +1,236 @@
+# DGL Utility Scripts
+
+This folder contains the utilities that do not belong to DGL core package as standalone executable
+scripts.
+
+## Graph Chunking
+
+`chunk_graph.py` provides an example of chunking an existing DGLGraph object into the on-disk
+[chunked graph format](http://13.231.216.217/guide/distributed-preprocessing.html#chunked-graph-format).
+
+<!-- TODO: change the link of documentation once it's merged to master -->
+
+An example of chunking the OGB MAG240M dataset:
+
+```python
+import ogb.lsc
+
+dataset = ogb.lsc.MAG240MDataset('.')
+etypes = [
+    ('paper', 'cites', 'paper'),
+    ('author', 'writes', 'paper'),
+    ('author', 'affiliated_with', 'institution')]
+g = dgl.heterograph({k: tuple(dataset.edge_index(*k)) for k in etypes})
+chunk_graph(
+    g,
+    'mag240m',
+    {'paper': {
+        'feat': 'mag240m_kddcup2021/processed/paper/node_feat.npy',
+        'label': 'mag240m_kddcup2021/processed/paper/node_label.npy',
+        'year': 'mag240m_kddcup2021/processed/paper/node_year.npy'}},
+    {},
+    4,
+    'output')
+```
+
+The output chunked graph metadata will go as follows (assuming the current directory as
+`/home/user`:
+
+```json
+{
+    "graph_name": "mag240m",
+    "node_type": [
+        "author",
+        "institution",
+        "paper"
+    ],
+    "num_nodes_per_chunk": [
+        [
+            30595778,
+            30595778,
+            30595778,
+            30595778
+        ],
+        [
+            6431,
+            6430,
+            6430,
+            6430
+        ],
+        [
+            30437917,
+            30437917,
+            30437916,
+            30437916
+        ]
+    ],
+    "edge_type": [
+        "author:affiliated_with:institution",
+        "author:writes:paper",
+        "paper:cites:paper"
+    ],
+    "num_edges_per_chunk": [
+        [
+            11148147,
+            11148147,
+            11148146,
+            11148146
+        ],
+        [
+            96505680,
+            96505680,
+            96505680,
+            96505680
+        ],
+        [
+            324437232,
+            324437232,
+            324437231,
+            324437231
+        ]
+    ],
+    "edges": {
+        "author:affiliated_with:institution": {
+            "format": {
+                "name": "csv",
+                "delimiter": " "
+            },
+            "data": [
+                "/home/user/output/edge_index/author:affiliated_with:institution0.txt",
+                "/home/user/output/edge_index/author:affiliated_with:institution1.txt",
+                "/home/user/output/edge_index/author:affiliated_with:institution2.txt",
+                "/home/user/output/edge_index/author:affiliated_with:institution3.txt"
+            ]
+        },
+        "author:writes:paper": {
+            "format": {
+                "name": "csv",
+                "delimiter": " "
+            },
+            "data": [
+                "/home/user/output/edge_index/author:writes:paper0.txt",
+                "/home/user/output/edge_index/author:writes:paper1.txt",
+                "/home/user/output/edge_index/author:writes:paper2.txt",
+                "/home/user/output/edge_index/author:writes:paper3.txt"
+            ]
+        },
+        "paper:cites:paper": {
+            "format": {
+                "name": "csv",
+                "delimiter": " "
+            },
+            "data": [
+                "/home/user/output/edge_index/paper:cites:paper0.txt",
+                "/home/user/output/edge_index/paper:cites:paper1.txt",
+                "/home/user/output/edge_index/paper:cites:paper2.txt",
+                "/home/user/output/edge_index/paper:cites:paper3.txt"
+            ]
+        }
+    },
+    "node_data": {
+        "paper": {
+            "feat": {
+                "format": {
+                    "name": "numpy"
+                },
+                "data": [
+                    "/home/user/output/node_data/paper/feat-0.npy",
+                    "/home/user/output/node_data/paper/feat-1.npy",
+                    "/home/user/output/node_data/paper/feat-2.npy",
+                    "/home/user/output/node_data/paper/feat-3.npy"
+                ]
+            },
+            "label": {
+                "format": {
+                    "name": "numpy"
+                },
+                "data": [
+                    "/home/user/output/node_data/paper/label-0.npy",
+                    "/home/user/output/node_data/paper/label-1.npy",
+                    "/home/user/output/node_data/paper/label-2.npy",
+                    "/home/user/output/node_data/paper/label-3.npy"
+                ]
+            },
+            "year": {
+                "format": {
+                    "name": "numpy"
+                },
+                "data": [
+                    "/home/user/output/node_data/paper/year-0.npy",
+                    "/home/user/output/node_data/paper/year-1.npy",
+                    "/home/user/output/node_data/paper/year-2.npy",
+                    "/home/user/output/node_data/paper/year-3.npy"
+                ]
+            }
+        }
+    },
+    "edge_data": {}
+}
+```
+
+## Change edge type to canonical edge type for partition configuration json
+
+In the upcoming DGL v1.0, we will require the partition configuration file to contain only canonical edge type. This tool is designed to help migrating existing configuration files from old style to new one.
+
+### Sample Usage
+
+```
+python tools/change_etype_to_canonical_etype.py --part_config "{configuration file path}"
+```
+
+### Requirement
+
+Partition algorithms produce one configuration file and multiple data folders, and each data folder corresponds to a partition. **This tool needs to read from the partition configuration file (specified by the commandline argument) *and* the graph structure data (stored in `graph.dgl` under the data folder) of the first partition.** They can be local files or shared files among network, if you follow this [official tutorial](https://docs.dgl.ai/en/latest/tutorials/dist/1_node_classification.html#sphx-glr-tutorials-dist-1-node-classification-py) for distributed training, you don't need to care about this as all files are shared by every participant through NFS.
+
+**For example, below is a typical data folder expected by this tool:**
+```
+data_root_dir/
+|-- graph_name.json    # specified by part_config
+|-- part0/
+    ...
+    |-- graph.dgl
+...
+```
+
+For more information about partition algorithm, see https://docs.dgl.ai/en/latest/generated/dgl.distributed.partition.partition_graph.html.
+
+### Input arguments
+
+1. *part_config*: The path of partition json file. < **Required**>
+
+### Result
+
+This tool changes the key of ``etypes`` and ``edge_map`` from format ``str`` to ``str:str:str`` and it overwrites the original file instead of creating a new one.
+
+E.g. **File content before running the script**
+```json
+{
+    "edge_map": {
+        "r1": [ [ 0, 6 ], [ 16, 20 ] ],
+        "r2": [ [ 6, 11 ], [ 20, 25 ] ],
+        "r3": [ [ 11, 16 ], [ 25, 30 ] ]
+    },
+    "etypes": {
+        "r1": 0,
+        "r2": 1,
+        "r3": 2
+    },
+    ...
+}
+```
+
+**After running**
+```json
+{
+    "edge_map": {
+        "n1:r1:n2": [ [ 0, 6 ], [ 16, 20 ] ],
+        "n1:r2:n3": [ [ 6, 11 ], [ 20, 25 ] ],
+        "n2:r3:n3": [ [ 11, 16 ], [ 25, 30 ] ] },
+    "etypes": {
+        "n1:r1:n2": 0,
+        "n1:r2:n3": 1,
+        "n2:r3:n3": 2
+    }
+    ...
+}
+```

tools/change_etype_to_canonical_etype.py

@@ -0,0 +1,147 @@
+import argparse
+import json
+import logging
+import os
+import time
+
+import dgl
+
+import torch
+from dgl._ffi.base import DGLError
+from dgl.data.utils import load_graphs
+from dgl.utils import toindex
+
+ETYPES_KEY = "etypes"
+EDGE_MAP_KEY = "edge_map"
+NTYPES_KEY = "ntypes"
+NUM_PARTS_KEY = "num_parts"
+CANONICAL_ETYPE_DELIMITER = ":"
+
+
+def convert_conf(part_config):
+    with open(part_config, "r+", encoding="utf-8") as f:
+        config = json.load(f)
+        logging.info("Checking if the provided json file need to be changed.")
+        if is_old_version(config):
+            logging.info("Changing the partition configuration file.")
+            canonical_etypes = {}
+            if len(config[NTYPES_KEY]) == 1:
+                ntype = list(config[NTYPES_KEY].keys())[0]
+                canonical_etypes = {
+                    CANONICAL_ETYPE_DELIMITER.join((ntype, etype, ntype)): eid
+                    for etype, eid in config[ETYPES_KEY].items()
+                }
+            else:
+                canonical_etypes = etype2canonical_etype(part_config, config)
+            reverse_c_etypes = {v: k for k, v in canonical_etypes.items()}
+            # Convert edge_map keys from etype -> c_etype.
+            new_edge_map = {}
+            for e_type, range in config[EDGE_MAP_KEY].items():
+                eid = config[ETYPES_KEY][e_type]
+                c_etype = reverse_c_etypes[eid]
+                new_edge_map[c_etype] = range
+            config[EDGE_MAP_KEY] = new_edge_map
+            config[ETYPES_KEY] = canonical_etypes
+            logging.info("Dumping the content to disk.")
+            f.seek(0)
+            json.dump(config, f, indent=4)
+            f.truncate()
+
+
+def etype2canonical_etype(part_config, config):
+    num_parts = config[NUM_PARTS_KEY]
+    edge_map = config[EDGE_MAP_KEY]
+    etypes = list(edge_map.keys())
+    # Get part id of each seed edge.
+    partition_ids = []
+    for _, bound in edge_map.items():
+        for i in range(num_parts):
+            if bound[i][1] > bound[i][0]:
+                partition_ids.append(i)
+                break
+    partition_ids = torch.tensor(partition_ids)
+
+    # Get starting index of each partition.
+    shifts = []
+    for i in range(num_parts):
+        shifts.append(edge_map[etypes[0]][i][0])
+    shifts = torch.tensor(shifts)
+
+    canonical_etypes = {}
+    part_ids = [
+        part_id for part_id in range(num_parts) if part_id in partition_ids
+    ]
+    for part_id in part_ids:
+        seed_etypes = [
+            etypes[i] for i in range(len(etypes)) if partition_ids[i] == part_id
+        ]
+        c_etype = _find_c_etypes_in_partition(
+            part_id,
+            seed_etypes,
+            config[ETYPES_KEY],
+            config[NTYPES_KEY],
+            edge_map,
+            shifts,
+            part_config,
+        )
+        canonical_etypes.update(c_etype)
+    return canonical_etypes
+
+
+def _find_c_etypes_in_partition(
+    part_id, seed_etypes, etypes, ntypes, edge_map, shifts, config_path
+):
+    try:
+        folder = os.path.dirname(os.path.realpath(config_path))
+        local_g = load_graphs(f"{folder}/part{part_id}/graph.dgl")[0][0]
+        local_eids = [
+            edge_map[etype][part_id][0] - shifts[part_id]
+            for etype in seed_etypes
+        ]
+        local_eids = toindex(torch.tensor(local_eids))
+        local_eids = local_eids.tousertensor()
+        local_src, local_dst = local_g.find_edges(local_eids)
+        src_ntids, dst_ntids = (
+            local_g.ndata[dgl.NTYPE][local_src],
+            local_g.ndata[dgl.NTYPE][local_dst],
+        )
+        ntypes = {v: k for k, v in ntypes.items()}
+        src_ntypes = [ntypes[ntid.item()] for ntid in src_ntids]
+        dst_ntypes = [ntypes[ntid.item()] for ntid in dst_ntids]
+        c_etypes = list(zip(src_ntypes, seed_etypes, dst_ntypes))
+        c_etypes = [
+            CANONICAL_ETYPE_DELIMITER.join(c_etype) for c_etype in c_etypes
+        ]
+        return {k: etypes[v] for (k, v) in zip(c_etypes, seed_etypes)}
+    except DGLError as e:
+        print(e)
+        logging.fatal(
+            f"Graph data of partition {part_id} is requested but not found."
+        )
+
+
+def is_old_version(config):
+    first_etype = list(config[ETYPES_KEY].keys())[0]
+    etype_tuple = first_etype.split(CANONICAL_ETYPE_DELIMITER)
+    return len(etype_tuple) == 1
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Change edge type in config file from format (str)"
+        " to (str,str,str), the original file will be overwritten",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--part_config", type=str, help="The file of the partition config"
+    )
+    args, _ = parser.parse_known_args()
+    assert (
+        args.part_config is not None
+    ), "A user has to specify a partition config file with --part_config."
+
+    start = time.time()
+    convert_conf(args.part_config)
+    end = time.time()
+    logging.info(f"elplased time in seconds: {end - start}")

tools/chunk_graph.py

@@ -0,0 +1,222 @@
+# See the __main__ block for usage of chunk_graph().
+import json
+import logging
+import os
+import pathlib
+from contextlib import contextmanager
+
+import dgl
+
+import torch
+from distpartitioning import array_readwriter
+from files import setdir
+
+
+def chunk_numpy_array(arr, fmt_meta, chunk_sizes, path_fmt):
+    paths = []
+    offset = 0
+
+    for j, n in enumerate(chunk_sizes):
+        path = os.path.abspath(path_fmt % j)
+        arr_chunk = arr[offset : offset + n]
+        logging.info("Chunking %d-%d" % (offset, offset + n))
+        array_readwriter.get_array_parser(**fmt_meta).write(path, arr_chunk)
+        offset += n
+        paths.append(path)
+
+    return paths
+
+
+def _chunk_graph(
+    g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt
+):
+    # First deal with ndata and edata that are homogeneous (i.e. not a dict-of-dict)
+    if len(g.ntypes) == 1 and not isinstance(
+        next(iter(ndata_paths.values())), dict
+    ):
+        ndata_paths = {g.ntypes[0]: ndata_paths}
+    if len(g.etypes) == 1 and not isinstance(
+        next(iter(edata_paths.values())), dict
+    ):
+        edata_paths = {g.etypes[0]: ndata_paths}
+    # Then convert all edge types to canonical edge types
+    etypestrs = {etype: ":".join(etype) for etype in g.canonical_etypes}
+    edata_paths = {
+        ":".join(g.to_canonical_etype(k)): v for k, v in edata_paths.items()
+    }
+
+    metadata = {}
+
+    metadata["graph_name"] = name
+    metadata["node_type"] = g.ntypes
+
+    # Compute the number of nodes per chunk per node type
+    metadata["num_nodes_per_chunk"] = num_nodes_per_chunk = []
+    for ntype in g.ntypes:
+        num_nodes = g.num_nodes(ntype)
+        num_nodes_list = []
+        for i in range(num_chunks):
+            n = num_nodes // num_chunks + (i < num_nodes % num_chunks)
+            num_nodes_list.append(n)
+        num_nodes_per_chunk.append(num_nodes_list)
+    num_nodes_per_chunk_dict = {
+        k: v for k, v in zip(g.ntypes, num_nodes_per_chunk)
+    }
+
+    metadata["edge_type"] = [etypestrs[etype] for etype in g.canonical_etypes]
+
+    # Compute the number of edges per chunk per edge type
+    metadata["num_edges_per_chunk"] = num_edges_per_chunk = []
+    for etype in g.canonical_etypes:
+        num_edges = g.num_edges(etype)
+        num_edges_list = []
+        for i in range(num_chunks):
+            n = num_edges // num_chunks + (i < num_edges % num_chunks)
+            num_edges_list.append(n)
+        num_edges_per_chunk.append(num_edges_list)
+    num_edges_per_chunk_dict = {
+        k: v for k, v in zip(g.canonical_etypes, num_edges_per_chunk)
+    }
+
+    # Split edge index
+    metadata["edges"] = {}
+    with setdir("edge_index"):
+        for etype in g.canonical_etypes:
+            etypestr = etypestrs[etype]
+            logging.info("Chunking edge index for %s" % etypestr)
+            edges_meta = {}
+            fmt_meta = {"name": "csv", "delimiter": " "}
+            edges_meta["format"] = fmt_meta
+
+            srcdst = torch.stack(g.edges(etype=etype), 1)
+            edges_meta["data"] = chunk_numpy_array(
+                srcdst.numpy(),
+                fmt_meta,
+                num_edges_per_chunk_dict[etype],
+                etypestr + "%d.txt",
+            )
+            metadata["edges"][etypestr] = edges_meta
+
+    # Chunk node data
+    reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {"name": data_fmt}
+    file_suffix = "npy" if data_fmt == "numpy" else "parquet"
+    metadata["node_data"] = {}
+    with setdir("node_data"):
+        for ntype, ndata_per_type in ndata_paths.items():
+            ndata_meta = {}
+            with setdir(ntype):
+                for key, path in ndata_per_type.items():
+                    logging.info(
+                        "Chunking node data for type %s key %s" % (ntype, key)
+                    )
+                    ndata_key_meta = {}
+                    arr = array_readwriter.get_array_parser(
+                        **reader_fmt_meta
+                    ).read(path)
+                    ndata_key_meta["format"] = writer_fmt_meta
+                    ndata_key_meta["data"] = chunk_numpy_array(
+                        arr,
+                        writer_fmt_meta,
+                        num_nodes_per_chunk_dict[ntype],
+                        key + "-%d." + file_suffix,
+                    )
+                    ndata_meta[key] = ndata_key_meta
+
+            metadata["node_data"][ntype] = ndata_meta
+
+    # Chunk edge data
+    metadata["edge_data"] = {}
+    with setdir("edge_data"):
+        for etypestr, edata_per_type in edata_paths.items():
+            edata_meta = {}
+            with setdir(etypestr):
+                for key, path in edata_per_type.items():
+                    logging.info(
+                        "Chunking edge data for type %s key %s"
+                        % (etypestr, key)
+                    )
+                    edata_key_meta = {}
+                    arr = array_readwriter.get_array_parser(
+                        **reader_fmt_meta
+                    ).read(path)
+                    edata_key_meta["format"] = writer_fmt_meta
+                    etype = tuple(etypestr.split(":"))
+                    edata_key_meta["data"] = chunk_numpy_array(
+                        arr,
+                        writer_fmt_meta,
+                        num_edges_per_chunk_dict[etype],
+                        key + "-%d." + file_suffix,
+                    )
+                    edata_meta[key] = edata_key_meta
+
+            metadata["edge_data"][etypestr] = edata_meta
+
+    metadata_path = "metadata.json"
+    with open(metadata_path, "w") as f:
+        json.dump(metadata, f, sort_keys=True, indent=4)
+    logging.info("Saved metadata in %s" % os.path.abspath(metadata_path))
+
+
+def chunk_graph(
+    g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt="numpy"
+):
+    """
+    Split the graph into multiple chunks.
+
+    A directory will be created at :attr:`output_path` with the metadata and chunked
+    edge list as well as the node/edge data.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        The graph.
+    name : str
+        The name of the graph, to be used later in DistDGL training.
+    ndata_paths : dict[str, pathlike] or dict[ntype, dict[str, pathlike]]
+        The dictionary of paths pointing to the corresponding numpy array file for each
+        node data key.
+    edata_paths : dict[etype, pathlike] or dict[etype, dict[str, pathlike]]
+        The dictionary of paths pointing to the corresponding numpy array file for each
+        edge data key. ``etype`` could be canonical or non-canonical.
+    num_chunks : int
+        The number of chunks
+    output_path : pathlike
+        The output directory saving the chunked graph.
+    """
+    for ntype, ndata in ndata_paths.items():
+        for key in ndata.keys():
+            ndata[key] = os.path.abspath(ndata[key])
+    for etype, edata in edata_paths.items():
+        for key in edata.keys():
+            edata[key] = os.path.abspath(edata[key])
+    with setdir(output_path):
+        _chunk_graph(
+            g, name, ndata_paths, edata_paths, num_chunks, output_path, data_fmt
+        )
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level="INFO")
+    input_dir = "/data"
+    output_dir = "/chunked-data"
+    (g,), _ = dgl.load_graphs(os.path.join(input_dir, "graph.dgl"))
+    chunk_graph(
+        g,
+        "mag240m",
+        {
+            "paper": {
+                "feat": os.path.join(input_dir, "paper/feat.npy"),
+                "label": os.path.join(input_dir, "paper/label.npy"),
+                "year": os.path.join(input_dir, "paper/year.npy"),
+            }
+        },
+        {
+            "cites": {"count": os.path.join(input_dir, "cites/count.npy")},
+            "writes": {"year": os.path.join(input_dir, "writes/year.npy")},
+            # you can put the same data file if they indeed share the features.
+            "rev_writes": {"year": os.path.join(input_dir, "writes/year.npy")},
+        },
+        4,
+        output_dir,
+    )
+# The generated metadata goes as in tools/sample-config/mag240m-metadata.json.

tools/copy_files.py

@@ -0,0 +1,149 @@
+"""Copy the partitions to a cluster of machines."""
+import argparse
+import copy
+import json
+import logging
+import os
+import signal
+import stat
+import subprocess
+import sys
+
+
+def copy_file(file_name, ip, workspace, param=""):
+    print("copy {} to {}".format(file_name, ip + ":" + workspace + "/"))
+    cmd = "scp " + param + " " + file_name + " " + ip + ":" + workspace + "/"
+    subprocess.check_call(cmd, shell=True)
+
+
+def exec_cmd(ip, cmd):
+    cmd = "ssh -o StrictHostKeyChecking=no " + ip + " '" + cmd + "'"
+    subprocess.check_call(cmd, shell=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Copy data to the servers.")
+    parser.add_argument(
+        "--workspace",
+        type=str,
+        required=True,
+        help="Path of user directory of distributed tasks. \
+                        This is used to specify a destination location where \
+                        data are copied to on remote machines.",
+    )
+    parser.add_argument(
+        "--rel_data_path",
+        type=str,
+        required=True,
+        help="Relative path in workspace to store the partition data.",
+    )
+    parser.add_argument(
+        "--part_config",
+        type=str,
+        required=True,
+        help="The partition config file. The path is on the local machine.",
+    )
+    parser.add_argument(
+        "--script_folder",
+        type=str,
+        required=True,
+        help="The folder contains all the user code scripts.",
+    )
+    parser.add_argument(
+        "--ip_config",
+        type=str,
+        required=True,
+        help="The file of IP configuration for servers. \
+                        The path is on the local machine.",
+    )
+    args = parser.parse_args()
+
+    hosts = []
+    with open(args.ip_config) as f:
+        for line in f:
+            res = line.strip().split(" ")
+            ip = res[0]
+            hosts.append(ip)
+
+    # We need to update the partition config file so that the paths are relative to
+    # the workspace in the remote machines.
+    with open(args.part_config) as conf_f:
+        part_metadata = json.load(conf_f)
+        tmp_part_metadata = copy.deepcopy(part_metadata)
+        num_parts = part_metadata["num_parts"]
+        assert num_parts == len(
+            hosts
+        ), "The number of partitions needs to be the same as the number of hosts."
+        graph_name = part_metadata["graph_name"]
+        node_map = part_metadata["node_map"]
+        edge_map = part_metadata["edge_map"]
+        if not isinstance(node_map, dict):
+            assert (
+                node_map[-4:] == ".npy"
+            ), "node map should be stored in a NumPy array."
+            tmp_part_metadata["node_map"] = "{}/{}/node_map.npy".format(
+                args.workspace, args.rel_data_path
+            )
+        if not isinstance(edge_map, dict):
+            assert (
+                edge_map[-4:] == ".npy"
+            ), "edge map should be stored in a NumPy array."
+            tmp_part_metadata["edge_map"] = "{}/{}/edge_map.npy".format(
+                args.workspace, args.rel_data_path
+            )
+
+        for part_id in range(num_parts):
+            part_files = tmp_part_metadata["part-{}".format(part_id)]
+            part_files["edge_feats"] = "{}/part{}/edge_feat.dgl".format(
+                args.rel_data_path, part_id
+            )
+            part_files["node_feats"] = "{}/part{}/node_feat.dgl".format(
+                args.rel_data_path, part_id
+            )
+            part_files["part_graph"] = "{}/part{}/graph.dgl".format(
+                args.rel_data_path, part_id
+            )
+    tmp_part_config = "/tmp/{}.json".format(graph_name)
+    with open(tmp_part_config, "w") as outfile:
+        json.dump(tmp_part_metadata, outfile, sort_keys=True, indent=4)
+
+    # Copy ip config.
+    for part_id, ip in enumerate(hosts):
+        remote_path = "{}/{}".format(args.workspace, args.rel_data_path)
+        exec_cmd(ip, "mkdir -p {}".format(remote_path))
+
+        copy_file(args.ip_config, ip, args.workspace)
+        copy_file(
+            tmp_part_config,
+            ip,
+            "{}/{}".format(args.workspace, args.rel_data_path),
+        )
+        node_map = part_metadata["node_map"]
+        edge_map = part_metadata["edge_map"]
+        if not isinstance(node_map, dict):
+            copy_file(node_map, ip, tmp_part_metadata["node_map"])
+        if not isinstance(edge_map, dict):
+            copy_file(edge_map, ip, tmp_part_metadata["edge_map"])
+        remote_path = "{}/{}/part{}".format(
+            args.workspace, args.rel_data_path, part_id
+        )
+        exec_cmd(ip, "mkdir -p {}".format(remote_path))
+
+        part_files = part_metadata["part-{}".format(part_id)]
+        copy_file(part_files["node_feats"], ip, remote_path)
+        copy_file(part_files["edge_feats"], ip, remote_path)
+        copy_file(part_files["part_graph"], ip, remote_path)
+        # copy script folder
+        copy_file(args.script_folder, ip, args.workspace, "-r")
+
+
+def signal_handler(signal, frame):
+    logging.info("Stop copying")
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    fmt = "%(asctime)s %(levelname)s %(message)s"
+    logging.basicConfig(format=fmt, level=logging.INFO)
+    signal.signal(signal.SIGINT, signal_handler)
+    main()

tools/dispatch_data.py

@@ -0,0 +1,217 @@
+"""Launching distributed graph partitioning pipeline """
+import argparse
+import json
+import logging
+import os
+import sys
+
+from partition_algo.base import load_partition_meta
+
+INSTALL_DIR = os.path.abspath(os.path.join(__file__, ".."))
+LAUNCH_SCRIPT = "distgraphlaunch.py"
+PIPELINE_SCRIPT = "distpartitioning/data_proc_pipeline.py"
+
+UDF_WORLD_SIZE = "world-size"
+UDF_PART_DIR = "partitions-dir"
+UDF_INPUT_DIR = "input-dir"
+UDF_GRAPH_NAME = "graph-name"
+UDF_SCHEMA = "schema"
+UDF_NUM_PARTS = "num-parts"
+UDF_OUT_DIR = "output"
+
+LARG_PROCS_MACHINE = "num_proc_per_machine"
+LARG_IPCONF = "ip_config"
+LARG_MASTER_PORT = "master_port"
+LARG_SSH_PORT = "ssh_port"
+
+
+def get_launch_cmd(args) -> str:
+    cmd = sys.executable + " " + os.path.join(INSTALL_DIR, LAUNCH_SCRIPT)
+    cmd = f"{cmd} --{LARG_SSH_PORT} {args.ssh_port} "
+    cmd = f"{cmd} --{LARG_PROCS_MACHINE} 1 "
+    cmd = f"{cmd} --{LARG_IPCONF} {args.ip_config} "
+    cmd = f"{cmd} --{LARG_MASTER_PORT} {args.master_port} "
+
+    return cmd
+
+
+def submit_jobs(args) -> str:
+    # read the json file and get the remaining argument here.
+    schema_path = args.metadata_filename
+    with open(os.path.join(args.in_dir, schema_path)) as schema:
+        schema_map = json.load(schema)
+
+    graph_name = schema_map["graph_name"]
+
+    # retrieve num_parts
+    num_parts = 0
+    partition_path = os.path.join(args.partitions_dir, "partition_meta.json")
+    if os.path.isfile(partition_path):
+        part_meta = load_partition_meta(partition_path)
+        num_parts = part_meta.num_parts
+
+    assert (
+        num_parts != 0
+    ), f"Invalid value for no. of partitions. Please check partition_meta.json file."
+
+    # verify ip_config
+    with open(args.ip_config, "r") as f:
+        num_ips = len(f.readlines())
+        assert (
+            num_parts % num_ips == 0
+        ), f"The num_parts[{args.num_parts}] should be a multiple of number of lines(ip addresses)[{args.ip_config}]."
+
+    argslist = ""
+    argslist += "--world-size {} ".format(num_ips)
+    argslist += "--partitions-dir {} ".format(
+        os.path.abspath(args.partitions_dir)
+    )
+    argslist += "--input-dir {} ".format(os.path.abspath(args.in_dir))
+    argslist += "--graph-name {} ".format(graph_name)
+    argslist += "--schema {} ".format(schema_path)
+    argslist += "--num-parts {} ".format(num_parts)
+    argslist += "--output {} ".format(os.path.abspath(args.out_dir))
+    argslist += "--process-group-timeout {} ".format(args.process_group_timeout)
+    argslist += "--log-level {} ".format(args.log_level)
+    argslist += "--save-orig-nids " if args.save_orig_nids else ""
+    argslist += "--save-orig-eids " if args.save_orig_eids else ""
+    argslist += "--use-graphbolt " if args.use_graphbolt else ""
+    argslist += "--store-eids " if args.store_eids else ""
+    argslist += "--store-inner-node " if args.store_inner_node else ""
+    argslist += "--store-inner-edge " if args.store_inner_edge else ""
+    argslist += (
+        f"--graph-formats {args.graph_formats} " if args.graph_formats else ""
+    )
+
+    # (BarclayII) Is it safe to assume all the workers have the Python executable at the same path?
+    pipeline_cmd = os.path.join(INSTALL_DIR, PIPELINE_SCRIPT)
+    udf_cmd = f"{args.python_path} {pipeline_cmd} {argslist}"
+
+    launch_cmd = get_launch_cmd(args)
+    launch_cmd += '"' + udf_cmd + '"'
+
+    print(launch_cmd)
+    os.system(launch_cmd)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Dispatch edge index and data to partitions",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    parser.add_argument(
+        "--in-dir",
+        type=str,
+        help="Location of the input directory where the dataset is located",
+    )
+    parser.add_argument(
+        "--metadata-filename",
+        type=str,
+        default="metadata.json",
+        help="Filename for the metadata JSON file that describes the dataset to be dispatched.",
+    )
+    parser.add_argument(
+        "--partitions-dir",
+        type=str,
+        help="Location of the partition-id mapping files which define node-ids and their respective partition-ids, relative to the input directory",
+    )
+    parser.add_argument(
+        "--out-dir",
+        type=str,
+        help="Location of the output directory where the graph partitions will be created by this pipeline",
+    )
+    parser.add_argument(
+        "--ip-config",
+        type=str,
+        help="File location of IP configuration for server processes",
+    )
+    parser.add_argument(
+        "--master-port",
+        type=int,
+        default=12345,
+        help="port used by gloo group to create randezvous point",
+    )
+    parser.add_argument(
+        "--log-level",
+        required=False,
+        type=str,
+        help="Log level to use for execution.",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+    )
+    parser.add_argument(
+        "--python-path",
+        type=str,
+        default=sys.executable,
+        help="Path to the Python executable on all workers",
+    )
+    parser.add_argument("--ssh-port", type=int, default=22, help="SSH Port.")
+    parser.add_argument(
+        "--process-group-timeout",
+        type=int,
+        default=1800,
+        help="timeout[seconds] for operations executed against the process group",
+    )
+    parser.add_argument(
+        "--save-orig-nids",
+        action="store_true",
+        help="Save original node IDs into files",
+    )
+    parser.add_argument(
+        "--save-orig-eids",
+        action="store_true",
+        help="Save original edge IDs into files",
+    )
+    parser.add_argument(
+        "--use-graphbolt",
+        action="store_true",
+        help="Use GraphBolt for distributed partition.",
+    )
+    parser.add_argument(
+        "--store-inner-node",
+        action="store_true",
+        default=False,
+        help="Store inner nodes.",
+    )
+
+    parser.add_argument(
+        "--store-inner-edge",
+        action="store_true",
+        default=False,
+        help="Store inner edges.",
+    )
+    parser.add_argument(
+        "--store-eids",
+        action="store_true",
+        default=False,
+        help="Store edge IDs.",
+    )
+    parser.add_argument(
+        "--graph-formats",
+        type=str,
+        default=None,
+        help="Save partitions in specified formats. It could be any combination(joined with ``,``) "
+        "of ``coo``, ``csc`` and ``csr``. If not specified, save one format only according to "
+        "what format is available. If multiple formats are available, selection priority "
+        "from high to low is ``coo``, ``csc``, ``csr``.",
+    )
+
+    args, _ = parser.parse_known_args()
+
+    fmt = "%(asctime)s %(levelname)s %(message)s"
+    logging.basicConfig(
+        format=fmt,
+        level=getattr(logging, args.log_level, None),
+    )
+
+    assert os.path.isdir(args.in_dir)
+    assert os.path.isdir(args.partitions_dir)
+    assert os.path.isfile(args.ip_config)
+    assert isinstance(args.master_port, int)
+
+    submit_jobs(args)
+
+
+if __name__ == "__main__":
+    main()

tools/distgraphlaunch.py

@@ -0,0 +1,492 @@
+"""Launching tool for DGL distributed training"""
+import argparse
+import json
+import logging
+import multiprocessing
+import os
+import re
+import signal
+import stat
+import subprocess
+import sys
+import time
+from functools import partial
+from threading import Thread
+from typing import Optional
+
+DEFAULT_PORT = 30050
+
+
+def cleanup_proc(get_all_remote_pids, conn):
+    """This process tries to clean up the remote training tasks."""
+    print("cleanupu process runs")
+    # This process should not handle SIGINT.
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    data = conn.recv()
+    # If the launch process exits normally, this process doesn't need to do anything.
+    if data == "exit":
+        sys.exit(0)
+    else:
+        remote_pids = get_all_remote_pids()
+        # Otherwise, we need to ssh to each machine and kill the training jobs.
+        for (ip, port), pids in remote_pids.items():
+            kill_process(ip, port, pids)
+    print("cleanup process exits")
+
+
+def kill_process(ip, port, pids):
+    """ssh to a remote machine and kill the specified processes."""
+    curr_pid = os.getpid()
+    killed_pids = []
+    # If we kill child processes first, the parent process may create more again. This happens
+    # to Python's process pool. After sorting, we always kill parent processes first.
+    pids.sort()
+    for pid in pids:
+        assert curr_pid != pid
+        print("kill process {} on {}:{}".format(pid, ip, port), flush=True)
+        kill_cmd = (
+            "ssh -o StrictHostKeyChecking=no -p "
+            + str(port)
+            + " "
+            + ip
+            + " 'kill {}'".format(pid)
+        )
+        subprocess.run(kill_cmd, shell=True)
+        killed_pids.append(pid)
+    # It's possible that some of the processes are not killed. Let's try again.
+    for i in range(3):
+        killed_pids = get_killed_pids(ip, port, killed_pids)
+        if len(killed_pids) == 0:
+            break
+        else:
+            killed_pids.sort()
+            for pid in killed_pids:
+                print(
+                    "kill process {} on {}:{}".format(pid, ip, port), flush=True
+                )
+                kill_cmd = (
+                    "ssh -o StrictHostKeyChecking=no -p "
+                    + str(port)
+                    + " "
+                    + ip
+                    + " 'kill -9 {}'".format(pid)
+                )
+                subprocess.run(kill_cmd, shell=True)
+
+
+def get_killed_pids(ip, port, killed_pids):
+    """Get the process IDs that we want to kill but are still alive."""
+    killed_pids = [str(pid) for pid in killed_pids]
+    killed_pids = ",".join(killed_pids)
+    ps_cmd = (
+        "ssh -o StrictHostKeyChecking=no -p "
+        + str(port)
+        + " "
+        + ip
+        + " 'ps -p {} -h'".format(killed_pids)
+    )
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    pids = []
+    for p in res.stdout.decode("utf-8").split("\n"):
+        l = p.split()
+        if len(l) > 0:
+            pids.append(int(l[0]))
+    return pids
+
+
+def execute_remote(
+    cmd: str, ip: str, port: int, username: Optional[str] = ""
+) -> Thread:
+    """Execute command line on remote machine via ssh.
+
+    Args:
+        cmd: User-defined command (udf) to execute on the remote host.
+        ip: The ip-address of the host to run the command on.
+        port: Port number that the host is listening on.
+        thread_list:
+        username: Optional. If given, this will specify a username to use when issuing commands over SSH.
+            Useful when your infra requires you to explicitly specify a username to avoid permission issues.
+
+    Returns:
+        thread: The Thread whose run() is to run the `cmd` on the remote host. Returns when the cmd completes
+            on the remote host.
+    """
+    ip_prefix = ""
+    if username:
+        ip_prefix += "{username}@".format(username=username)
+
+    # Construct ssh command that executes `cmd` on the remote host
+    ssh_cmd = "ssh -o StrictHostKeyChecking=no -p {port} {ip_prefix}{ip} '{cmd}'".format(
+        port=str(port),
+        ip_prefix=ip_prefix,
+        ip=ip,
+        cmd=cmd,
+    )
+
+    # thread func to run the job
+    def run(ssh_cmd):
+        subprocess.check_call(ssh_cmd, shell=True)
+
+    thread = Thread(target=run, args=(ssh_cmd,))
+    thread.setDaemon(True)
+    thread.start()
+    return thread
+
+
+def get_remote_pids(ip, port, cmd_regex):
+    """Get the process IDs that run the command in the remote machine."""
+    pids = []
+    curr_pid = os.getpid()
+    # Here we want to get the python processes. We may get some ssh processes, so we should filter them out.
+    ps_cmd = (
+        "ssh -o StrictHostKeyChecking=no -p "
+        + str(port)
+        + " "
+        + ip
+        + " 'ps -aux | grep python | grep -v StrictHostKeyChecking'"
+    )
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    for p in res.stdout.decode("utf-8").split("\n"):
+        l = p.split()
+        if len(l) < 2:
+            continue
+        # We only get the processes that run the specified command.
+        res = re.search(cmd_regex, p)
+        if res is not None and int(l[1]) != curr_pid:
+            pids.append(l[1])
+
+    pid_str = ",".join([str(pid) for pid in pids])
+    ps_cmd = (
+        "ssh -o StrictHostKeyChecking=no -p "
+        + str(port)
+        + " "
+        + ip
+        + " 'pgrep -P {}'".format(pid_str)
+    )
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    pids1 = res.stdout.decode("utf-8").split("\n")
+    all_pids = []
+    for pid in set(pids + pids1):
+        if pid == "" or int(pid) == curr_pid:
+            continue
+        all_pids.append(int(pid))
+    all_pids.sort()
+    return all_pids
+
+
+def get_all_remote_pids(hosts, ssh_port, udf_command):
+    """Get all remote processes."""
+    remote_pids = {}
+    for node_id, host in enumerate(hosts):
+        ip, _ = host
+        # When creating training processes in remote machines, we may insert some arguments
+        # in the commands. We need to use regular expressions to match the modified command.
+        cmds = udf_command.split()
+        new_udf_command = " .*".join(cmds)
+        pids = get_remote_pids(ip, ssh_port, new_udf_command)
+        remote_pids[(ip, ssh_port)] = pids
+    return remote_pids
+
+
+def construct_torch_dist_launcher_cmd(
+    num_trainers: int,
+    num_nodes: int,
+    node_rank: int,
+    master_addr: str,
+    master_port: int,
+) -> str:
+    """Constructs the torch distributed launcher command.
+    Helper function.
+
+    Args:
+        num_trainers:
+        num_nodes:
+        node_rank:
+        master_addr:
+        master_port:
+
+    Returns:
+        cmd_str.
+    """
+    torch_cmd_template = (
+        "-m torch.distributed.launch "
+        "--nproc_per_node={nproc_per_node} "
+        "--nnodes={nnodes} "
+        "--node_rank={node_rank} "
+        "--master_addr={master_addr} "
+        "--master_port={master_port}"
+    )
+    return torch_cmd_template.format(
+        nproc_per_node=num_trainers,
+        nnodes=num_nodes,
+        node_rank=node_rank,
+        master_addr=master_addr,
+        master_port=master_port,
+    )
+
+
+def wrap_udf_in_torch_dist_launcher(
+    udf_command: str,
+    num_trainers: int,
+    num_nodes: int,
+    node_rank: int,
+    master_addr: str,
+    master_port: int,
+) -> str:
+    """Wraps the user-defined function (udf_command) with the torch.distributed.launch module.
+
+     Example: if udf_command is "python3 run/some/trainer.py arg1 arg2", then new_df_command becomes:
+         "python3 -m torch.distributed.launch <TORCH DIST ARGS> run/some/trainer.py arg1 arg2
+
+    udf_command is assumed to consist of pre-commands (optional) followed by the python launcher script (required):
+    Examples:
+        # simple
+        python3.7 path/to/some/trainer.py arg1 arg2
+
+        # multi-commands
+        (cd some/dir && python3.7 path/to/some/trainer.py arg1 arg2)
+
+    IMPORTANT: If udf_command consists of multiple python commands, then this will result in undefined behavior.
+
+    Args:
+        udf_command:
+        num_trainers:
+        num_nodes:
+        node_rank:
+        master_addr:
+        master_port:
+
+    Returns:
+
+    """
+    torch_dist_cmd = construct_torch_dist_launcher_cmd(
+        num_trainers=num_trainers,
+        num_nodes=num_nodes,
+        node_rank=node_rank,
+        master_addr=master_addr,
+        master_port=master_port,
+    )
+    # Auto-detect the python binary that kicks off the distributed trainer code.
+    # Note: This allowlist order matters, this will match with the FIRST matching entry. Thus, please add names to this
+    #       from most-specific to least-specific order eg:
+    #           (python3.7, python3.8) -> (python3)
+    # The allowed python versions are from this: https://www.dgl.ai/pages/start.html
+    python_bin_allowlist = (
+        "python3.6",
+        "python3.7",
+        "python3.8",
+        "python3.9",
+        "python3",
+        # for backwards compatibility, accept python2 but technically DGL is a py3 library, so this is not recommended
+        "python2.7",
+        "python2",
+    )
+    # If none of the candidate python bins match, then we go with the default `python`
+    python_bin = "python"
+    for candidate_python_bin in python_bin_allowlist:
+        if candidate_python_bin in udf_command:
+            python_bin = candidate_python_bin
+            break
+
+    # transforms the udf_command from:
+    #     python path/to/dist_trainer.py arg0 arg1
+    # to:
+    #     python -m torch.distributed.launch [DIST TORCH ARGS] path/to/dist_trainer.py arg0 arg1
+    # Note: if there are multiple python commands in `udf_command`, this may do the Wrong Thing, eg launch each
+    #       python command within the torch distributed launcher.
+    new_udf_command = udf_command.replace(
+        python_bin, f"{python_bin} {torch_dist_cmd}"
+    )
+
+    return new_udf_command
+
+
+def construct_dgl_server_env_vars(
+    ip_config: str,
+    num_proc_per_machine: int,
+    pythonpath: Optional[str] = "",
+) -> str:
+    """Constructs the DGL server-specific env vars string that are required for DGL code to behave in the correct
+    server role.
+    Convenience function.
+
+    Args:
+        ip_config: IP config file containing IP addresses of cluster hosts.
+            Relative path to workspace.
+        num_proc_per_machine:
+        pythonpath: Optional. If given, this will pass this as PYTHONPATH.
+
+    Returns:
+        server_env_vars: The server-specific env-vars in a string format, friendly for CLI execution.
+
+    """
+    server_env_vars_template = (
+        "DGL_IP_CONFIG={DGL_IP_CONFIG} "
+        "DGL_NUM_SERVER={DGL_NUM_SERVER} "
+        "{suffix_optional_envvars}"
+    )
+    suffix_optional_envvars = ""
+    if pythonpath:
+        suffix_optional_envvars += f"PYTHONPATH={pythonpath} "
+    return server_env_vars_template.format(
+        DGL_IP_CONFIG=ip_config,
+        DGL_NUM_SERVER=num_proc_per_machine,
+        suffix_optional_envvars=suffix_optional_envvars,
+    )
+
+
+def wrap_cmd_with_local_envvars(cmd: str, env_vars: str) -> str:
+    """Wraps a CLI command with desired env vars with the following properties:
+    (1) env vars persist for the entire `cmd`, even if it consists of multiple "chained" commands like:
+        cmd = "ls && pwd && python run/something.py"
+    (2) env vars don't pollute the environment after `cmd` completes.
+
+    Example:
+        >>> cmd = "ls && pwd"
+        >>> env_vars = "VAR1=value1 VAR2=value2"
+        >>> wrap_cmd_with_local_envvars(cmd, env_vars)
+        "(export VAR1=value1 VAR2=value2; ls && pwd)"
+
+    Args:
+        cmd:
+        env_vars: A string containing env vars, eg "VAR1=val1 VAR2=val2"
+
+    Returns:
+        cmd_with_env_vars:
+
+    """
+    # use `export` to persist env vars for entire cmd block. required if udf_command is a chain of commands
+    # also: wrap in parens to not pollute env:
+    #     https://stackoverflow.com/a/45993803
+    return f"(export {env_vars}; {cmd})"
+
+
+def wrap_cmd_with_extra_envvars(cmd: str, env_vars: list) -> str:
+    """Wraps a CLI command with extra env vars
+
+    Example:
+        >>> cmd = "ls && pwd"
+        >>> env_vars = ["VAR1=value1", "VAR2=value2"]
+        >>> wrap_cmd_with_extra_envvars(cmd, env_vars)
+        "(export VAR1=value1 VAR2=value2; ls && pwd)"
+
+    Args:
+        cmd:
+        env_vars: A list of strings containing env vars, e.g., ["VAR1=value1", "VAR2=value2"]
+
+    Returns:
+        cmd_with_env_vars:
+    """
+    env_vars = " ".join(env_vars)
+    return wrap_cmd_with_local_envvars(cmd, env_vars)
+
+
+def submit_jobs(args, udf_command):
+    """Submit distributed jobs (server and client processes) via ssh"""
+    hosts = []
+    thread_list = []
+    server_count_per_machine = 0
+
+    # Get the IP addresses of the cluster.
+    # ip_config = os.path.join(args.workspace, args.ip_config)
+    ip_config = args.ip_config
+    with open(ip_config) as f:
+        for line in f:
+            result = line.strip().split()
+            if len(result) == 2:
+                ip = result[0]
+                port = int(result[1])
+                hosts.append((ip, port))
+            elif len(result) == 1:
+                ip = result[0]
+                port = DEFAULT_PORT
+                hosts.append((ip, port))
+            else:
+                raise RuntimeError("Format error of ip_config.")
+            server_count_per_machine = args.num_proc_per_machine
+
+    # launch server tasks
+    server_env_vars = construct_dgl_server_env_vars(
+        ip_config=args.ip_config,
+        num_proc_per_machine=args.num_proc_per_machine,
+        pythonpath=os.environ.get("PYTHONPATH", ""),
+    )
+    for i in range(len(hosts) * server_count_per_machine):
+        ip, _ = hosts[int(i / server_count_per_machine)]
+        server_env_vars_cur = f"{server_env_vars} RANK={i} MASTER_ADDR={hosts[0][0]} MASTER_PORT={args.master_port}"
+        cmd = wrap_cmd_with_local_envvars(udf_command, server_env_vars_cur)
+        print(cmd)
+        thread_list.append(
+            execute_remote(cmd, ip, args.ssh_port, username=args.ssh_username)
+        )
+
+    # Start a cleanup process dedicated for cleaning up remote training jobs.
+    conn1, conn2 = multiprocessing.Pipe()
+    func = partial(get_all_remote_pids, hosts, args.ssh_port, udf_command)
+    process = multiprocessing.Process(target=cleanup_proc, args=(func, conn1))
+    process.start()
+
+    def signal_handler(signal, frame):
+        logging.info("Stop launcher")
+        # We need to tell the cleanup process to kill remote training jobs.
+        conn2.send("cleanup")
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, signal_handler)
+
+    for thread in thread_list:
+        thread.join()
+    # The training processes complete. We should tell the cleanup process to exit.
+    conn2.send("exit")
+    process.join()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Launch a distributed job")
+    parser.add_argument("--ssh_port", type=int, default=22, help="SSH Port.")
+    parser.add_argument(
+        "--ssh_username",
+        default="",
+        help="Optional. When issuing commands (via ssh) to cluster, use the provided username in the ssh cmd. "
+        "Example: If you provide --ssh_username=bob, then the ssh command will be like: 'ssh bob@1.2.3.4 CMD' "
+        "instead of 'ssh 1.2.3.4 CMD'",
+    )
+    parser.add_argument(
+        "--num_proc_per_machine",
+        type=int,
+        help="The number of server processes per machine",
+    )
+    parser.add_argument(
+        "--master_port",
+        type=int,
+        help="This port is used to form gloo group (randevouz server)",
+    )
+    parser.add_argument(
+        "--ip_config",
+        type=str,
+        help="The file (in workspace) of IP configuration for server processes",
+    )
+
+    args, udf_command = parser.parse_known_args()
+    assert len(udf_command) == 1, "Please provide user command line."
+    assert (
+        args.num_proc_per_machine is not None and args.num_proc_per_machine > 0
+    ), "--num_proc_per_machine must be a positive number."
+    assert (
+        args.ip_config is not None
+    ), "A user has to specify an IP configuration file with --ip_config."
+
+    udf_command = str(udf_command[0])
+    if "python" not in udf_command:
+        raise RuntimeError(
+            "DGL launching script can only support Python executable file."
+        )
+
+    submit_jobs(args, udf_command)
+
+
+if __name__ == "__main__":
+    fmt = "%(asctime)s %(levelname)s %(message)s"
+    logging.basicConfig(format=fmt, level=logging.INFO)
+    main()

tools/distpartitioning/README.md

@@ -0,0 +1,40 @@
+### 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.
+##### Step 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/array_readwriter/__init__.py

@@ -0,0 +1,2 @@
+from . import csv, numpy_array, parquet
+from .registry import get_array_parser, register_array_parser

tools/distpartitioning/array_readwriter/csv.py

@@ -0,0 +1,39 @@
+import logging
+
+import pandas as pd
+import pyarrow
+import pyarrow.csv
+
+from .registry import register_array_parser
+
+
+@register_array_parser("csv")
+class CSVArrayParser(object):
+    def __init__(self, delimiter=","):
+        self.delimiter = delimiter
+
+    def read(self, path):
+        logging.debug(
+            "Reading from %s using CSV format with configuration %s"
+            % (path, self.__dict__)
+        )
+        # do not read the first line as header
+        read_options = pyarrow.csv.ReadOptions(autogenerate_column_names=True)
+        parse_options = pyarrow.csv.ParseOptions(delimiter=self.delimiter)
+        arr = pyarrow.csv.read_csv(
+            path, read_options=read_options, parse_options=parse_options
+        )
+        logging.debug("Done reading from %s" % path)
+        return arr.to_pandas().to_numpy()
+
+    def write(self, path, arr):
+        logging.debug(
+            "Writing to %s using CSV format with configuration %s"
+            % (path, self.__dict__)
+        )
+        write_options = pyarrow.csv.WriteOptions(
+            include_header=False, delimiter=self.delimiter
+        )
+        arr = pyarrow.Table.from_pandas(pd.DataFrame(arr))
+        pyarrow.csv.write_csv(arr, path, write_options=write_options)
+        logging.debug("Done writing to %s" % path)

tools/distpartitioning/array_readwriter/numpy_array.py

@@ -0,0 +1,26 @@
+import logging
+
+import numpy as np
+from numpy.lib.format import open_memmap
+
+from .registry import register_array_parser
+
+
+@register_array_parser("numpy")
+class NumpyArrayParser(object):
+    def __init__(self):
+        pass
+
+    def read(self, path):
+        logging.debug("Reading from %s using numpy format" % path)
+        arr = np.load(path, mmap_mode="r")
+        logging.debug("Done reading from %s" % path)
+        return arr
+
+    def write(self, path, arr):
+        logging.debug("Writing to %s using numpy format" % path)
+        # np.save would load the entire memmap array up into CPU.  So we manually open
+        # an empty npy file with memmap mode and manually flush it instead.
+        new_arr = open_memmap(path, mode="w+", dtype=arr.dtype, shape=arr.shape)
+        new_arr[:] = arr[:]
+        logging.debug("Done writing to %s" % path)

tools/distpartitioning/array_readwriter/parquet.py

@@ -0,0 +1,62 @@
+import logging
+
+import numpy as np
+import pandas as pd
+import pyarrow
+import pyarrow.parquet
+
+from .registry import register_array_parser
+
+
+@register_array_parser("parquet")
+class ParquetArrayParser(object):
+    def __init__(self):
+        pass
+
+    def read(self, path):
+        logging.debug("Reading from %s using parquet format" % path)
+        metadata = pyarrow.parquet.read_metadata(path)
+        metadata = metadata.schema.to_arrow_schema().metadata
+
+        # As parquet data are tabularized, we assume the dim of ndarray is 2.
+        # If not, it should be explictly specified in the file as metadata.
+        if metadata:
+            shape = metadata.get(b"shape", None)
+        else:
+            shape = None
+        table = pyarrow.parquet.read_table(path, memory_map=True)
+
+        data_types = table.schema.types
+        # Spark ML feature processing produces single-column parquet files where each row is a vector object
+        if len(data_types) == 1 and isinstance(data_types[0], pyarrow.ListType):
+            arr = np.array(table.to_pandas().iloc[:, 0].to_list())
+            logging.debug(
+                f"Parquet data under {path} converted from single vector per row to ndarray"
+            )
+        else:
+            arr = table.to_pandas().to_numpy()
+        if not shape:
+            logging.debug(
+                "Shape information not found in the metadata, read the data as "
+                "a 2 dim array."
+            )
+        logging.debug("Done reading from %s" % path)
+        shape = tuple(eval(shape.decode())) if shape else arr.shape
+        return arr.reshape(shape)
+
+    def write(self, path, array, vector_rows=False):
+        logging.debug("Writing to %s using parquet format" % path)
+        shape = array.shape
+        if len(shape) > 2:
+            array = array.reshape(shape[0], -1)
+        if vector_rows:
+            table = pyarrow.table(
+                [pyarrow.array(array.tolist())], names=["vector"]
+            )
+            logging.debug("Writing to %s using single-vector rows..." % path)
+        else:
+            table = pyarrow.Table.from_pandas(pd.DataFrame(array))
+            table = table.replace_schema_metadata({"shape": str(shape)})
+
+        pyarrow.parquet.write_table(table, path)
+        logging.debug("Done writing to %s" % path)

tools/distpartitioning/array_readwriter/registry.py

@@ -0,0 +1,14 @@
+REGISTRY = {}
+
+
+def register_array_parser(name):
+    def _deco(cls):
+        REGISTRY[name] = cls
+        return cls
+
+    return _deco
+
+
+def get_array_parser(**fmt_meta):
+    cls = REGISTRY[fmt_meta.pop("name")]
+    return cls(**fmt_meta)

tools/distpartitioning/constants.py

@@ -0,0 +1,43 @@
+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"
+
+STR_NODE_TYPE = "node_type"
+STR_EDGE_TYPE = "edge_type"
+STR_EDGES = "edges"
+STR_FORMAT = "format"
+STR_FORMAT_DELIMITER = "delimiter"
+STR_DATA = "data"
+STR_NODE_DATA = "node_data"
+STR_EDGE_DATA = "edge_data"
+
+STR_NUMPY = "numpy"
+STR_PARQUET = "parquet"
+STR_CSV = "csv"
+STR_NAME = "name"
+
+STR_GRAPH_NAME = "graph_name"
+STR_NODE_FEATURES = "node_features"
+STR_EDGE_FEATURES = "edge_features"
+
+STR_NUM_NODES_PER_TYPE = "num_nodes_per_type"
+STR_NUM_EDGES_PER_TYPE = "num_edges_per_type"
+
+STR_NTYPES = "ntypes"

tools/distpartitioning/convert_partition.py

@@ -0,0 +1,922 @@
+import copy
+import gc
+import logging
+import os
+
+import constants
+import dgl
+import dgl.backend as F
+import dgl.graphbolt as gb
+import numpy as np
+import torch as th
+import torch.distributed as dist
+from dgl import EID, ETYPE, NID, NTYPE
+
+from dgl.distributed.constants import DGL2GB_EID, GB_DST_ID
+from dgl.distributed.partition import (
+    _cast_to_minimum_dtype,
+    _etype_str_to_tuple,
+    _etype_tuple_to_str,
+    cast_various_to_minimum_dtype_gb,
+    RESERVED_FIELD_DTYPE,
+)
+from utils import get_idranges, memory_snapshot
+
+
+def _get_unique_invidx(srcids, dstids, nids, low_mem=True):
+    """This function is used to compute a list of unique elements,
+    and their indices in the input list, which is the concatenation
+    of srcids, dstids and uniq_nids. In addition, this function will also
+    compute inverse indices, in the list of unique elements, for the
+    elements in srcids, dstids and nids arrays. srcids, dstids will be
+    over-written to contain the inverse indices. Basically, this function
+    is mimicing the functionality of numpy's unique function call.
+    The problem with numpy's unique function call is its high memory
+    requirement. For an input list of 3 billion edges it consumes about
+    550GB of systems memory, which is limiting the capability of the
+    partitioning pipeline.
+
+    Note: This function is a workaround solution for the high memory requirement
+    of numpy's unique function call. This function is not a general purpose
+    function and is only used in the context of the partitioning pipeline.
+    What's more, this function does not behave exactly the same as numpy's
+    unique function call. Namely, this function does not return the exact same
+    inverse indices as numpy's unique function call. However, for the current
+    use case, this function is sufficient.
+
+    Current numpy uniques function returns 3 return parameters, which are
+        . list of unique elements
+        . list of indices, in the input argument list, which are first
+            occurance of the corresponding element in the uniques list
+        . list of inverse indices, which are indices from the uniques list
+            and can be used to rebuild the original input array
+    Compared to the above numpy's return parameters, this work around
+    solution returns 4 values
+        . list of unique elements,
+        . list of indices, which may not be the first occurance of the
+            corresponding element from the uniques
+        . list of inverse indices, here we only build the inverse indices
+            for srcids and dstids input arguments. For the current use case,
+            only these two inverse indices are needed.
+
+    Parameters:
+    -----------
+    srcids : numpy array
+        a list of numbers, which are the src-ids of the edges
+    dstids : numpy array
+        a list of numbers, which are the dst-ids of the edges
+    nids : numpy array
+        a list of numbers, a list of unique shuffle-global-nids.
+        This list is guaranteed to be a list of sorted consecutive unique
+        list of numbers. Also, this list will be a `super set` for the
+        list of dstids. Current implementation of the pipeline guarantees
+        this assumption and is used to simplify the current implementation
+        of the workaround solution.
+    low_mem : bool, optional
+        Indicates whether to use the low memory version of the function. If
+        ``False``, the function will use numpy's native ``unique`` function.
+        Otherwise, the function will use the low memory version of the
+        function.
+
+    Returns:
+    --------
+    numpy array :
+        a list of unique, sorted elements, computed from the input arguments
+    numpy array :
+        a list of integers. These are indices in the concatenated list
+        [srcids, dstids, uniq_nids], which are the input arguments to this function
+    numpy array :
+        a list of integers. These are inverse indices, which will be indices
+        from the unique elements list specifying the elements from the
+        input array, srcids
+    numpy array :
+        a list of integers. These are inverse indices, which will be indices
+        from the unique elements list specifying the elements from the
+        input array, dstids
+    """
+    assert len(srcids) == len(
+        dstids
+    ), f"Please provide the correct input parameters"
+    assert len(srcids) != 0, f"Please provide a non-empty edge-list."
+
+    if not low_mem:
+        logging.warning(
+            "Calling numpy's native function unique. This functions memory "
+            "overhead will limit size of the partitioned graph objects "
+            "processed by each node in the cluster."
+        )
+        uniques, idxes, inv_idxes = np.unique(
+            np.concatenate([srcids, dstids, nids]),
+            return_index=True,
+            return_inverse=True,
+        )
+        src_len = len(srcids)
+        dst_len = len(dstids)
+        return (
+            uniques,
+            idxes,
+            inv_idxes[:src_len],
+            inv_idxes[src_len : (src_len + dst_len)],
+        )
+
+    # find uniqes which appear only in the srcids list
+    mask = np.isin(srcids, nids, invert=True, kind="table")
+    srcids_only = srcids[mask]
+    srcids_idxes = np.where(mask == 1)[0]
+
+    # sort
+    uniques, unique_srcids_idx = np.unique(srcids_only, return_index=True)
+    idxes = srcids_idxes[unique_srcids_idx]
+
+    # build uniques and idxes, first and second return parameters
+    uniques = np.concatenate([uniques, nids])
+    idxes = np.concatenate(
+        [idxes, len(srcids) + len(dstids) + np.arange(len(nids))]
+    )
+
+    # sort and idxes
+    sort_idx = np.argsort(uniques)
+    uniques = uniques[sort_idx]
+    idxes = idxes[sort_idx]
+
+    # uniques and idxes are built
+    assert len(uniques) == len(idxes), f"Error building the idxes array."
+
+    srcids = np.searchsorted(uniques, srcids, side="left")
+
+    # process dstids now.
+    # dstids is guaranteed to be a subset of the `nids` list
+    # here we are computing index in the list of uniqes for
+    # each element in the list of dstids, in a two step process
+    # 1. locate the position of first element from nids in the
+    #       list of uniques - dstids cannot appear to the left
+    #       of this number, they are guaranteed to be on the right
+    #       side of this number.
+    # 2. dstids = dstids - nids[0]
+    #       By subtracting nids[0] from the list of dstids will make
+    #       the list of dstids to be in the range of [0, max(nids)-1]
+    # 3. dstids = dstids - nids[0] + offset
+    #       Now we move the list of dstids by `offset` which will be
+    #       the starting position of the nids[0] element. Note that
+    #       nids will ALWAYS be a SUPERSET of dstids.
+    offset = np.searchsorted(uniques, nids[0], side="left")
+    dstids = dstids - nids[0] + offset
+
+    # return the values
+    return uniques, idxes, srcids, dstids
+
+
+# Utility functions.
+def _is_homogeneous(ntypes, etypes):
+    """Checks if the provided ntypes and etypes form a homogeneous graph."""
+    return len(ntypes) == 1 and len(etypes) == 1
+
+
+def _coo2csc(src_ids, dst_ids):
+    src_ids, dst_ids = th.tensor(src_ids, dtype=th.int64), th.tensor(
+        dst_ids, dtype=th.int64
+    )
+    num_nodes = th.max(th.stack([src_ids, dst_ids], dim=0)).item() + 1
+    dst, idx = dst_ids.sort()
+    indptr = th.searchsorted(dst, th.arange(num_nodes + 1))
+    indices = src_ids[idx]
+    return indptr, indices, idx
+
+
+def _create_edge_data(edgeid_offset, etype_ids, num_edges):
+    eid = th.arange(
+        edgeid_offset,
+        edgeid_offset + num_edges,
+        dtype=RESERVED_FIELD_DTYPE[dgl.EID],
+    )
+    etype = th.as_tensor(etype_ids, dtype=RESERVED_FIELD_DTYPE[dgl.ETYPE])
+    inner_edge = th.ones(num_edges, dtype=RESERVED_FIELD_DTYPE["inner_edge"])
+    return eid, etype, inner_edge
+
+
+def _create_node_data(ntype, uniq_ids, reshuffle_nodes, inner_nodes):
+    node_type = th.as_tensor(ntype, dtype=RESERVED_FIELD_DTYPE[dgl.NTYPE])
+    node_id = th.as_tensor(uniq_ids[reshuffle_nodes])
+    inner_node = th.as_tensor(
+        inner_nodes[reshuffle_nodes],
+        dtype=RESERVED_FIELD_DTYPE["inner_node"],
+    )
+    return node_type, node_id, inner_node
+
+
+def _compute_node_ntype(
+    global_src_id, global_dst_id, global_homo_nid, idx, reshuffle_nodes, id_map
+):
+    global_ids = np.concatenate([global_src_id, global_dst_id, global_homo_nid])
+    part_global_ids = global_ids[idx]
+    part_global_ids = part_global_ids[reshuffle_nodes]
+    ntype, per_type_ids = id_map(part_global_ids)
+    return ntype, per_type_ids
+
+
+def _graph_orig_ids(
+    return_orig_nids,
+    return_orig_eids,
+    ntypes_map,
+    etypes_map,
+    node_attr,
+    edge_attr,
+    per_type_ids,
+    type_per_edge,
+    global_edge_id,
+):
+    orig_nids = None
+    orig_eids = None
+    if return_orig_nids:
+        orig_nids = {}
+        for ntype, ntype_id in ntypes_map.items():
+            mask = th.logical_and(
+                node_attr[dgl.NTYPE] == ntype_id,
+                node_attr["inner_node"],
+            )
+            orig_nids[ntype] = th.as_tensor(per_type_ids[mask])
+    if return_orig_eids:
+        orig_eids = {}
+        for etype, etype_id in etypes_map.items():
+            mask = th.logical_and(
+                type_per_edge == etype_id,
+                edge_attr["inner_edge"],
+            )
+            orig_eids[_etype_tuple_to_str(etype)] = th.as_tensor(
+                global_edge_id[mask]
+            )
+    return orig_nids, orig_eids
+
+
+def _create_edge_attr_gb(
+    part_local_dst_id, edgeid_offset, etype_ids, ntypes, etypes, etypes_map
+):
+    edge_attr = {}
+    # create edge data in graph.
+    num_edges = len(part_local_dst_id)
+    (
+        edge_attr[dgl.EID],
+        type_per_edge,
+        edge_attr["inner_edge"],
+    ) = _create_edge_data(edgeid_offset, etype_ids, num_edges)
+    assert "inner_edge" in edge_attr
+
+    is_homo = _is_homogeneous(ntypes, etypes)
+
+    edge_type_to_id = (
+        {gb.etype_tuple_to_str(("_N", "_E", "_N")): 0}
+        if is_homo
+        else {
+            gb.etype_tuple_to_str(etype): etid
+            for etype, etid in etypes_map.items()
+        }
+    )
+    return edge_attr, type_per_edge, edge_type_to_id
+
+
+def _create_node_attr(
+    idx,
+    global_src_id,
+    global_dst_id,
+    global_homo_nid,
+    uniq_ids,
+    reshuffle_nodes,
+    id_map,
+    inner_nodes,
+):
+    # compute per_type_ids and ntype for all the nodes in the graph.
+    ntype, per_type_ids = _compute_node_ntype(
+        global_src_id,
+        global_dst_id,
+        global_homo_nid,
+        idx,
+        reshuffle_nodes,
+        id_map,
+    )
+
+    # create node data in graph.
+    node_attr = {}
+    (
+        node_attr[dgl.NTYPE],
+        node_attr[dgl.NID],
+        node_attr["inner_node"],
+    ) = _create_node_data(ntype, uniq_ids, reshuffle_nodes, inner_nodes)
+    return node_attr, per_type_ids
+
+
+def remove_attr_gb(
+    edge_attr, node_attr, store_inner_node, store_inner_edge, store_eids
+):
+    edata, ndata = copy.deepcopy(edge_attr), copy.deepcopy(node_attr)
+    if not store_inner_edge:
+        assert "inner_edge" in edata
+        edata.pop("inner_edge")
+
+    if not store_eids:
+        assert dgl.EID in edata
+        edata.pop(dgl.EID)
+
+    if not store_inner_node:
+        assert "inner_node" in ndata
+        ndata.pop("inner_node")
+    return edata, ndata
+
+
+def _process_partition_gb(
+    node_attr,
+    edge_attr,
+    type_per_edge,
+    src_ids,
+    dst_ids,
+    sort_etypes,
+):
+    """Preprocess partitions before saving:
+    1. format data types.
+    2. sort csc/csr by tag.
+    """
+    for k, dtype in RESERVED_FIELD_DTYPE.items():
+        if k in node_attr:
+            node_attr[k] = F.astype(node_attr[k], dtype)
+        if k in edge_attr:
+            edge_attr[k] = F.astype(edge_attr[k], dtype)
+
+    indptr, indices, edge_ids = _coo2csc(src_ids, dst_ids)
+    if sort_etypes:
+        split_size = th.diff(indptr)
+        split_indices = th.split(type_per_edge, tuple(split_size), dim=0)
+        sorted_idxs = []
+        for split_indice in split_indices:
+            sorted_idxs.append(split_indice.sort()[1])
+
+        sorted_idx = th.cat(sorted_idxs, dim=0)
+        sorted_idx = (
+            th.repeat_interleave(indptr[:-1], split_size, dim=0) + sorted_idx
+        )
+
+    return indptr, indices[sorted_idx], edge_ids[sorted_idx]
+
+
+def _update_node_map(node_map_val, end_ids_per_rank, id_ntypes, prev_last_id):
+    """this function is modified from the function '_update_node_edge_map' in dgl.distributed.partition"""
+    # Update the node_map_val to be contiguous.
+    rank = dist.get_rank()
+    prev_end_id = (
+        end_ids_per_rank[rank - 1].item() if rank > 0 else prev_last_id
+    )
+    ntype_ids = {ntype: ntype_id for ntype_id, ntype in enumerate(id_ntypes)}
+    for ntype_id in list(ntype_ids.values()):
+        ntype = id_ntypes[ntype_id]
+        start_id = node_map_val[ntype][0][0]
+        end_id = node_map_val[ntype][0][1]
+        if not (start_id == -1 and end_id == -1):
+            continue
+        prev_ntype_id = (
+            ntype_ids[ntype] - 1
+            if ntype_ids[ntype] > 0
+            else max(ntype_ids.values())
+        )
+        prev_ntype = id_ntypes[prev_ntype_id]
+        if ntype_ids[ntype] == 0:
+            node_map_val[ntype][0][0] = prev_end_id
+        else:
+            node_map_val[ntype][0][0] = node_map_val[prev_ntype][0][1]
+        node_map_val[ntype][0][1] = node_map_val[ntype][0][0]
+    return node_map_val[ntype][0][-1]
+
+
+def create_graph_object(
+    tot_node_count,
+    tot_edge_count,
+    node_count,
+    edge_count,
+    num_parts,
+    schema,
+    part_id,
+    node_data,
+    edge_data,
+    edgeid_offset,
+    node_typecounts,
+    edge_typecounts,
+    last_ids={},
+    return_orig_nids=False,
+    return_orig_eids=False,
+    use_graphbolt=False,
+    **kwargs,
+):
+    """
+    This function creates dgl objects for a given graph partition, as in function
+    arguments.
+
+    The "schema" argument is a dictionary, which contains the metadata related to node ids
+    and edge ids. It contains two keys: "nid" and "eid", whose value is also a dictionary
+    with the following structure.
+
+    1. The key-value pairs in the "nid" dictionary has the following format.
+       "ntype-name" is the user assigned name to this node type. "format" describes the
+       format of the contents of the files. and "data" is a list of lists, each list has
+       3 elements: file-name, start_id and end_id. File-name can be either absolute or
+       relative path to this file and starting and ending ids are type ids of the nodes
+       which are contained in this file. These type ids are later used to compute global ids
+       of these nodes which are used throughout the processing of this pipeline.
+        "ntype-name" : {
+            "format" : "csv",
+            "data" : [
+                    [ <path-to-file>/ntype0-name-0.csv, start_id0, end_id0],
+                    [ <path-to-file>/ntype0-name-1.csv, start_id1, end_id1],
+                    ...
+                    [ <path-to-file>/ntype0-name-<p-1>.csv, start_id<p-1>, end_id<p-1>],
+            ]
+        }
+
+    2. The key-value pairs in the "eid" dictionary has the following format.
+       As described for the "nid" dictionary the "eid" dictionary is similarly structured
+       except that these entries are for edges.
+        "etype-name" : {
+            "format" : "csv",
+            "data" : [
+                    [ <path-to-file>/etype0-name-0, start_id0, end_id0],
+                    [ <path-to-file>/etype0-name-1 start_id1, end_id1],
+                    ...
+                    [ <path-to-file>/etype0-name-1 start_id<p-1>, end_id<p-1>]
+            ]
+        }
+
+    In "nid" dictionary, the type_nids are specified that
+    should be assigned to nodes which are read from the corresponding nodes file.
+    Along the same lines dictionary for the key "eid" is used for edges in the
+    input graph.
+
+    These type ids, for nodes and edges, are used to compute global ids for nodes
+    and edges which are stored in the graph object.
+
+    Parameters:
+    -----------
+    tot_node_count : int
+        the number of all nodes
+    tot_edge_count : int
+        the number of all edges
+    node_count : int
+        the number of nodes in partition
+    edge_count : int
+        the number of edges in partition
+    graph_formats : str
+        the format of graph
+    num_parts : int
+        the number of parts
+    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>
+    edgeid_offset : int
+        offset to be used when assigning edge global ids in the current partition
+    return_orig_ids : bool, optional
+        Indicates whether to return original node/edge IDs.
+
+    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)
+    dict of tensors
+        If `return_orig_nids=True`, return a dict of 1D tensors whose key is the node type
+        and value is a 1D tensor mapping between shuffled node IDs and the original node
+        IDs for each node type. Otherwise, ``None`` is returned.
+    dict of tensors
+        If `return_orig_eids=True`, return a dict of 1D tensors whose key is the edge type
+        and value is a 1D tensor mapping between shuffled edge IDs and the original edge
+        IDs for each edge type. Otherwise, ``None`` is returned.
+    """
+    # create auxiliary data structures from the schema object
+    memory_snapshot("CreateDGLObj_Begin", part_id)
+    _, global_nid_ranges = get_idranges(
+        schema[constants.STR_NODE_TYPE], node_typecounts
+    )
+    _, global_eid_ranges = get_idranges(
+        schema[constants.STR_EDGE_TYPE], edge_typecounts
+    )
+
+    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])
+    etypes = [e[0] for e in etypes]
+    etypes_map = {_etype_str_to_tuple(e): i for i, e in enumerate(etypes)}
+
+    node_map_val = {ntype: [] for ntype in ntypes}
+    edge_map_val = {_etype_str_to_tuple(etype): [] for etype in etypes}
+
+    memory_snapshot("CreateDGLObj_AssignNodeData", part_id)
+    shuffle_global_nids = node_data[constants.SHUFFLE_GLOBAL_NID]
+    node_data.pop(constants.SHUFFLE_GLOBAL_NID)
+    gc.collect()
+
+    ntype_ids = node_data[constants.NTYPE_ID]
+    node_data.pop(constants.NTYPE_ID)
+    gc.collect()
+
+    global_type_nid = node_data[constants.GLOBAL_TYPE_NID]
+    node_data.pop(constants.GLOBAL_TYPE_NID)
+    node_data = None
+    gc.collect()
+
+    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.
+    prev_last_id = last_ids.get(part_id - 1, 0)
+    for ntype_name in global_nid_ranges:
+        ntype_id = ntypes_map[ntype_name]
+        type_nids = shuffle_global_nids[ntype_ids == ntype_id]
+        if len(type_nids) == 0:
+            node_map_val[ntype_name].append([-1, -1])
+        else:
+            node_map_val[ntype_name].append(
+                [int(type_nids[0]), int(type_nids[-1]) + 1]
+            )
+            last_id = th.tensor(
+                [max(prev_last_id, int(type_nids[-1]) + 1)], dtype=th.int64
+            )
+    id_ntypes = list(global_nid_ranges.keys())
+
+    gather_last_ids = [
+        th.zeros(1, dtype=th.int64) for _ in range(dist.get_world_size())
+    ]
+
+    dist.all_gather(gather_last_ids, last_id)
+    prev_last_id = _update_node_map(
+        node_map_val, gather_last_ids, id_ntypes, prev_last_id
+    )
+    last_ids[part_id] = prev_last_id
+
+    # process edges
+    memory_snapshot("CreateDGLObj_AssignEdgeData: ", part_id)
+    shuffle_global_src_id = edge_data[constants.SHUFFLE_GLOBAL_SRC_ID]
+    edge_data.pop(constants.SHUFFLE_GLOBAL_SRC_ID)
+    gc.collect()
+
+    shuffle_global_dst_id = edge_data[constants.SHUFFLE_GLOBAL_DST_ID]
+    edge_data.pop(constants.SHUFFLE_GLOBAL_DST_ID)
+    gc.collect()
+
+    global_src_id = edge_data[constants.GLOBAL_SRC_ID]
+    edge_data.pop(constants.GLOBAL_SRC_ID)
+    gc.collect()
+
+    global_dst_id = edge_data[constants.GLOBAL_DST_ID]
+    edge_data.pop(constants.GLOBAL_DST_ID)
+    gc.collect()
+
+    global_edge_id = edge_data[constants.GLOBAL_TYPE_EID]
+    edge_data.pop(constants.GLOBAL_TYPE_EID)
+    gc.collect()
+
+    etype_ids = edge_data[constants.ETYPE_ID]
+    edge_data.pop(constants.ETYPE_ID)
+    edge_data = None
+    gc.collect()
+    logging.info(
+        f"There are {len(shuffle_global_src_id)} edges in partition {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.
+    if not np.all(np.diff(etype_ids) >= 0):
+        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)
+    else:
+        print(f"[Rank: {part_id} Edge data is already sorted !!!")
+
+    # Determine the edge ID range of different edge types.
+    edge_id_start = edgeid_offset
+    for etype_name in global_eid_ranges:
+        etype = _etype_str_to_tuple(etype_name)
+        assert len(etype) == 3
+        etype_id = etypes_map[etype]
+        edge_map_val[etype].append(
+            [edge_id_start, edge_id_start + np.sum(etype_ids == etype_id)]
+        )
+        edge_id_start += np.sum(etype_ids == etype_id)
+    memory_snapshot("CreateDGLObj_UniqueNodeIds: ", part_id)
+
+    # get the edge list in some order and then reshuffle.
+    # Here the order of nodes is defined by the sorted order.
+    uniq_ids, idx, part_local_src_id, part_local_dst_id = _get_unique_invidx(
+        shuffle_global_src_id,
+        shuffle_global_dst_id,
+        np.arange(shuffle_global_nid_range[0], shuffle_global_nid_range[1] + 1),
+    )
+
+    inner_nodes = th.as_tensor(
+        np.logical_and(
+            uniq_ids >= shuffle_global_nid_range[0],
+            uniq_ids <= shuffle_global_nid_range[1],
+        )
+    )
+
+    # get the list of indices, from inner_nodes, which will sort inner_nodes as [True, True, ...., False, False, ...]
+    # essentially local nodes will be placed before non-local nodes.
+    reshuffle_nodes = th.arange(len(uniq_ids))
+    reshuffle_nodes = th.cat(
+        [reshuffle_nodes[inner_nodes.bool()], reshuffle_nodes[inner_nodes == 0]]
+    )
+
+    """
+    Following procedure is used to map the part_local_src_id, part_local_dst_id to account for
+    reshuffling of nodes (to order localy owned nodes prior to non-local nodes in a partition)
+    1. Form a node_map, in this case a numpy array, which will be used to map old node-ids (pre-reshuffling)
+    to post-reshuffling ids.
+    2. Once the map is created, use this map to map all the node-ids in the part_local_src_id 
+    and part_local_dst_id list to their appropriate `new` node-ids (post-reshuffle order).
+    3. Since only the node's order is changed, we will have to re-order nodes related information when
+    creating dgl object: this includes dgl.NTYPE, dgl.NID and inner_node.
+    4. Edge's order is not changed. At this point in the execution path edges are still ordered by their etype-ids.
+    5. Create the dgl object appropriately and return the dgl object.
+    
+    Here is a  simple example to understand the above flow better.
+
+    part_local_nids = [0, 1, 2, 3, 4, 5]
+    part_local_src_ids = [0, 0, 0, 0, 2, 3, 4]
+    part_local_dst_ids = [1, 2, 3, 4, 4, 4, 5]
+
+    Assume that nodes {1, 5} are halo-nodes, which are not owned by this partition.
+
+    reshuffle_nodes = [0, 2, 3, 4, 1, 5]
+
+    A node_map, which maps node-ids from old to reshuffled order is as follows:
+    node_map = np.zeros((len(reshuffle_nodes,)))
+    node_map[reshuffle_nodes] = np.arange(len(reshuffle_nodes))
+
+    Using the above map, we have mapped part_local_src_ids and part_local_dst_ids as follows:
+    part_local_src_ids = [0, 0, 0, 0, 1, 2, 3]
+    part_local_dst_ids = [4, 1, 2, 3, 3, 3, 5]
+
+    In this graph above, note that nodes {0, 1, 2, 3} are inner_nodes and {4, 5} are NON-inner-nodes
+
+    Since the edge are re-ordered in any way, there is no reordering required for edge related data
+    during the DGL object creation.
+    """
+    # create the mappings to generate mapped part_local_src_id and part_local_dst_id
+    # This map will map from unshuffled node-ids to reshuffled-node-ids (which are ordered to prioritize
+    # locally owned nodes).
+    nid_map = np.zeros(
+        (
+            len(
+                reshuffle_nodes,
+            )
+        )
+    )
+    nid_map[reshuffle_nodes] = np.arange(len(reshuffle_nodes))
+
+    # Now map the edge end points to reshuffled_values.
+    part_local_src_id, part_local_dst_id = (
+        nid_map[part_local_src_id],
+        nid_map[part_local_dst_id],
+    )
+
+    """
+    Creating attributes for graphbolt and DGLGraph is as follows.
+
+    node attributes:    
+        this part is implemented in _create_node_attr.
+        compute the ntype and per type ids for each node with global node type id.
+        create ntype, nid and inner node with orig ntype and inner nodes
+    this part is shared by graphbolt and DGLGraph.
+
+    the attributes created for graphbolt are as follows:
+
+    edge attributes:
+        this part is implemented in _create_edge_attr_gb.
+        create eid, type per edge and inner edge with edgeid_offset.
+        create edge_type_to_id with etypes_map.
+    
+    The process to remove extra attribute is implemented in  remove_attr_gb.
+    the unused attributes like inner_node, inner_edge, eids will be removed following the arguments in kwargs.
+    edge_attr, node_attr are the variable that have removed extra attributes to construct csc_graph.
+    edata, ndata are the variable that reserve extra attributes to be used to generate orig_nid and orig_eid. 
+    
+    the src_ids and dst_ids will be transformed into indptr and indices in _coo2csc.
+
+    all variable mentioned above will be casted to minimum data type in cast_various_to_minimum_dtype_gb.
+
+    orig_nids and orig_eids will be generated in _graph_orig_ids with ndata and edata.
+    """
+    # create the graph here now.
+    ndata, per_type_ids = _create_node_attr(
+        idx,
+        global_src_id,
+        global_dst_id,
+        global_homo_nid,
+        uniq_ids,
+        reshuffle_nodes,
+        id_map,
+        inner_nodes,
+    )
+    if use_graphbolt:
+        edata, type_per_edge, edge_type_to_id = _create_edge_attr_gb(
+            part_local_dst_id,
+            edgeid_offset,
+            etype_ids,
+            ntypes,
+            etypes,
+            etypes_map,
+        )
+
+        assert edata is not None
+        assert ndata is not None
+
+        sort_etypes = len(etypes_map) > 1
+        indptr, indices, csc_edge_ids = _process_partition_gb(
+            ndata,
+            edata,
+            type_per_edge,
+            part_local_src_id,
+            part_local_dst_id,
+            sort_etypes,
+        )
+        edge_attr, node_attr = remove_attr_gb(
+            edge_attr=edata, node_attr=ndata, **kwargs
+        )
+        edge_attr = {
+            attr: edge_attr[attr][csc_edge_ids] for attr in edge_attr.keys()
+        }
+        cast_various_to_minimum_dtype_gb(
+            node_count=node_count,
+            edge_count=edge_count,
+            tot_node_count=tot_node_count,
+            tot_edge_count=tot_edge_count,
+            num_parts=num_parts,
+            indptr=indptr,
+            indices=indices,
+            type_per_edge=type_per_edge,
+            etypes=etypes,
+            ntypes=ntypes,
+            node_attributes=node_attr,
+            edge_attributes=edge_attr,
+        )
+        part_graph = gb.fused_csc_sampling_graph(
+            csc_indptr=indptr,
+            indices=indices,
+            node_type_offset=None,
+            type_per_edge=type_per_edge[csc_edge_ids],
+            node_attributes=node_attr,
+            edge_attributes=edge_attr,
+            node_type_to_id=ntypes_map,
+            edge_type_to_id=edge_type_to_id,
+        )
+    else:
+        num_edges = len(part_local_dst_id)
+        part_graph = dgl.graph(
+            data=(part_local_src_id, part_local_dst_id), num_nodes=len(uniq_ids)
+        )
+        # create edge data in graph.
+        (
+            part_graph.edata[dgl.EID],
+            part_graph.edata[dgl.ETYPE],
+            part_graph.edata["inner_edge"],
+        ) = _create_edge_data(edgeid_offset, etype_ids, num_edges)
+
+        ndata, per_type_ids = _create_node_attr(
+            idx,
+            global_src_id,
+            global_dst_id,
+            global_homo_nid,
+            uniq_ids,
+            reshuffle_nodes,
+            id_map,
+            inner_nodes,
+        )
+        for attr_name, node_attributes in ndata.items():
+            part_graph.ndata[attr_name] = node_attributes
+        type_per_edge = part_graph.edata[dgl.ETYPE]
+        ndata, edata = part_graph.ndata, part_graph.edata
+    # get the original node ids and edge ids from original graph.
+    orig_nids, orig_eids = _graph_orig_ids(
+        return_orig_nids,
+        return_orig_eids,
+        ntypes_map,
+        etypes_map,
+        ndata,
+        edata,
+        per_type_ids,
+        type_per_edge,
+        global_edge_id,
+    )
+    return (
+        part_graph,
+        node_map_val,
+        edge_map_val,
+        ntypes_map,
+        etypes_map,
+        orig_nids,
+        orig_eids,
+    )
+
+
+def create_metadata_json(
+    graph_name,
+    num_nodes,
+    num_edges,
+    part_id,
+    num_parts,
+    node_map_val,
+    edge_map_val,
+    ntypes_map,
+    etypes_map,
+    output_dir,
+    use_graphbolt,
+):
+    """
+    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
+    part_id : int
+       integer indicating the partition id
+    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
+    use_graphbolt : bool
+        whether to use graphbolt or not
+
+    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,
+    }
+
+    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")
+    if use_graphbolt:
+        part_graph_file = os.path.join(part_dir, "fused_csc_sampling_graph.pt")
+    else:
+        part_graph_file = os.path.join(part_dir, "graph.dgl")
+    part_graph_type = "part_graph_graphbolt" if use_graphbolt else "part_graph"
+    part_metadata["part-{}".format(part_id)] = {
+        "node_feats": node_feat_file,
+        "edge_feats": edge_feat_file,
+        part_graph_type: part_graph_file,
+    }
+    return part_metadata

tools/distpartitioning/data_proc_pipeline.py

@@ -0,0 +1,134 @@
+import argparse
+import logging
+import os
+import platform
+
+import numpy as np
+import torch.multiprocessing as mp
+
+from data_shuffle import multi_machine_run, single_machine_run
+
+
+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("Output Dir: ", params.output)
+    print("WorldSize: ", params.world_size)
+    print("Metis partitions: ", params.partitions_dir)
+
+
+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(
+        "--output",
+        required=True,
+        type=str,
+        help="The output directory of the partitioned results",
+    )
+    parser.add_argument(
+        "--partitions-dir",
+        help="directory of the partition-ids for each node type",
+        default=None,
+        type=str,
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="info",
+        help="To enable log level for debugging purposes. Available options: \
+			  (Critical, Error, Warning, Info, Debug, Notset), default value \
+			  is: Info",
+    )
+
+    # 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(
+        "--process-group-timeout",
+        required=True,
+        type=int,
+        help="timeout[seconds] for operations executed against the process group "
+        "(see torch.distributed.init_process_group)",
+    )
+    parser.add_argument(
+        "--save-orig-nids",
+        action="store_true",
+        help="Save original node IDs into files",
+    )
+    parser.add_argument(
+        "--save-orig-eids",
+        action="store_true",
+        help="Save original edge IDs into files",
+    )
+    parser.add_argument(
+        "--use-graphbolt",
+        action="store_true",
+        help="Use GraphBolt for distributed partition.",
+    )
+    parser.add_argument(
+        "--store-inner-node",
+        action="store_true",
+        default=False,
+        help="Store inner nodes.",
+    )
+
+    parser.add_argument(
+        "--store-inner-edge",
+        action="store_true",
+        default=False,
+        help="Store inner edges.",
+    )
+    parser.add_argument(
+        "--store-eids",
+        action="store_true",
+        default=False,
+        help="Store edge IDs.",
+    )
+    parser.add_argument(
+        "--graph-formats",
+        default=None,
+        type=str,
+        help="Save partitions in specified formats.",
+    )
+    params = parser.parse_args()
+    # invoke the pipeline function
+    numeric_level = getattr(logging, params.log_level.upper(), None)
+    logging.basicConfig(
+        level=numeric_level,
+        format=f"[{platform.node()} %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
+    )
+    multi_machine_run(params)

tools/distpartitioning/dataset_utils.py

@@ -0,0 +1,666 @@
+import gc
+import logging
+import os
+
+import array_readwriter
+import constants
+
+import numpy as np
+import pyarrow
+import pyarrow.parquet as pq
+import torch
+import torch.distributed as dist
+from gloo_wrapper import alltoallv_cpu
+from utils import (
+    DATA_TYPE_ID,
+    generate_read_list,
+    get_gid_offsets,
+    get_idranges,
+    map_partid_rank,
+    REV_DATA_TYPE_ID,
+)
+
+
+def _broadcast_shape(
+    data, rank, world_size, num_parts, is_feat_data, feat_name
+):
+    """Auxiliary function to broadcast the shape of a feature data.
+    This information is used to figure out the type-ids for the
+    local features.
+
+    Parameters:
+    -----------
+    data : numpy array
+        which is the feature data read from the disk
+    rank : integer
+        which represents the id of the process in the process group
+    world_size : integer
+        represents the total no. of process in the process group
+    num_parts : integer
+        specifying the no. of partitions
+    is_feat_data : bool
+        flag used to seperate feature data and edge data
+    feat_name : string
+        name of the feature
+
+    Returns:
+    -------
+    list of tuples :
+        which represents the range of type-ids for the data array.
+    """
+    assert len(data.shape) in [
+        1,
+        2,
+    ], f"Data is expected to be 1-D or 2-D but got {data.shape}."
+    data_shape = list(data.shape)
+
+    if len(data_shape) == 1:
+        data_shape.append(1)
+
+    if is_feat_data:
+        data_shape.append(DATA_TYPE_ID[data.dtype])
+
+    data_shape = torch.tensor(data_shape, dtype=torch.int64)
+    data_shape_output = [
+        torch.zeros_like(data_shape) for _ in range(world_size)
+    ]
+    dist.all_gather(data_shape_output, data_shape)
+    logging.debug(
+        f"[Rank: {rank} Received shapes from all ranks: {data_shape_output}"
+    )
+    shapes = [x.numpy() for x in data_shape_output if x[0] != 0]
+    shapes = np.vstack(shapes)
+
+    if is_feat_data:
+        logging.debug(
+            f"shapes: {shapes}, condition: {all(shapes[0,2] == s for s in shapes[:,2])}"
+        )
+        assert all(
+            shapes[0, 2] == s for s in shapes[:, 2]
+        ), f"dtypes for {feat_name} does not match on all ranks"
+
+    # compute tids here.
+    type_counts = list(shapes[:, 0])
+    tid_start = np.cumsum([0] + type_counts[:-1])
+    tid_end = np.cumsum(type_counts)
+    tid_ranges = list(zip(tid_start, tid_end))
+    logging.debug(f"starts -> {tid_start} ... end -> {tid_end}")
+
+    return tid_ranges
+
+
+def get_dataset(
+    input_dir, graph_name, rank, world_size, num_parts, schema_map, ntype_counts
+):
+    """
+    Function to read the multiple file formatted dataset.
+
+    Parameters:
+    -----------
+    input_dir : string
+        root directory where dataset is located.
+    graph_name : string
+        graph name string
+    rank : int
+        rank of the current process
+    world_size : int
+        total number of process in the current execution
+    num_parts : int
+        total number of output graph partitions
+    schema_map : dictionary
+        this is the dictionary created by reading the graph metadata json file
+        for the input graph dataset
+
+    Return:
+    -------
+    dictionary
+        where keys are node-type names and values are tuples. Each tuple represents the
+        range of type ids read from a file by the current process. Please note that node
+        data for each node type is split into "p" files and each one of these "p" files are
+        read a process in the distributed graph partitioning pipeline
+    dictionary
+        Data read from numpy files for all the node features in this dataset. Dictionary built
+        using this data has keys as node feature names and values as tensor data representing
+        node features
+    dictionary
+        in which keys are node-type and values are a triplet. This triplet has node-feature name,
+        and range of tids for the node feature data read from files by the current process. Each
+        node-type may have mutiple feature(s) and associated tensor data.
+    dictionary
+        Data read from edges.txt file and used to build a dictionary with keys as column names
+        and values as columns in the csv file.
+    dictionary
+        in which keys are edge-type names and values are triplets. This triplet has edge-feature name,
+        and range of tids for theedge feature data read from the files by the current process. Each
+        edge-type may have several edge features and associated tensor data.
+    dictionary
+        Data read from numpy files for all the edge features in this dataset. This dictionary's keys
+        are feature names and values are tensors data representing edge feature data.
+    dictionary
+        This dictionary is used for identifying the global-id range for the associated edge features
+        present in the previous return value. The keys are edge-type names and values are triplets.
+        Each triplet consists of edge-feature name and starting and ending points of the range of
+        tids representing the corresponding edge feautres.
+    """
+
+    # node features dictionary
+    # TODO: With the new file format, It is guaranteed that the input dataset will have
+    # no. of nodes with features (node-features) files and nodes metadata will always be the same.
+    # This means the dimension indicating the no. of nodes in any node-feature files and the no. of
+    # nodes in the corresponding nodes metadata file will always be the same. With this guarantee,
+    # we can eliminate the `node_feature_tids` dictionary since the same information is also populated
+    # in the `node_tids` dictionary. This will be remnoved in the next iteration of code changes.
+    node_features = {}
+    node_feature_tids = {}
+
+    """
+    The structure of the node_data is as follows, which is present in the input metadata json file.
+       "node_data" : {
+            "ntype0-name" : {
+                "feat0-name" : {
+                    "format" : {"name": "numpy"},
+                    "data" :   [ #list
+                        "<path>/feat-0.npy",
+                        "<path>/feat-1.npy",
+                        ....
+                        "<path>/feat-<p-1>.npy"
+                    ]
+                },
+                "feat1-name" : {
+                    "format" : {"name": "numpy"},
+                    "data" : [ #list
+                        "<path>/feat-0.npy",
+                        "<path>/feat-1.npy",
+                        ....
+                        "<path>/feat-<p-1>.npy"
+                    ]
+                }
+            }
+       }
+
+    As shown above, the value for the key "node_data" is a dictionary object, which is
+    used to describe the feature data for each of the node-type names. Keys in this top-level
+    dictionary are node-type names and value is a dictionary which captures all the features
+    for the current node-type. Feature data is captured with keys being the feature-names and
+    value is a dictionary object which has 2 keys namely format and data. Format entry is used
+    to mention the format of the storage used by the node features themselves and "data" is used
+    to mention all the files present for this given node feature.
+
+    Data read from each of the node features file is a multi-dimensional tensor data and is read
+    in numpy or parquet format, which is also the storage format of node features on the permanent storage.
+
+        "node_type" : ["ntype0-name", "ntype1-name", ....], #m node types
+        "num_nodes_per_chunk" : [
+            [a0, a1, ...a<p-1>], #p partitions
+            [b0, b1, ... b<p-1>],
+            ....
+            [c0, c1, ..., c<p-1>] #no, of node types
+        ],
+
+    The "node_type" points to a list of all the node names present in the graph
+    And "num_nodes_per_chunk" is used to mention no. of nodes present in each of the
+    input nodes files. These node counters are used to compute the type_node_ids as
+    well as global node-ids by using a simple cumulative summation and maitaining an
+    offset counter to store the end of the current.
+
+    Since nodes are NOT actually associated with any additional metadata, w.r.t to the processing
+    involved in this pipeline this information is not needed to be stored in files. This optimization
+    saves a considerable amount of time when loading massively large datasets for paritioning.
+    As opposed to reading from files and performing shuffling process each process/rank generates nodes
+    which are owned by that particular rank. And using the "num_nodes_per_chunk" information each
+    process can easily compute any nodes per-type node_id and global node_id.
+    The node-ids are treated as int64's in order to support billions of nodes in the input graph.
+    """
+
+    # read my nodes for each node type
+    """
+    node_tids, ntype_gnid_offset = get_idranges(
+        schema_map[constants.STR_NODE_TYPE],
+        schema_map[constants.STR_NUM_NODES_PER_CHUNK],
+        num_chunks=num_parts,
+    )
+    """
+    logging.debug(f"[Rank: {rank} ntype_counts: {ntype_counts}")
+    ntype_gnid_offset = get_gid_offsets(
+        schema_map[constants.STR_NODE_TYPE], ntype_counts
+    )
+    logging.debug(f"[Rank: {rank} - ntype_gnid_offset = {ntype_gnid_offset}")
+
+    # iterate over the "node_data" dictionary in the schema_map
+    # read the node features if exists
+    # also keep track of the type_nids for which the node_features are read.
+    dataset_features = schema_map[constants.STR_NODE_DATA]
+    if (dataset_features is not None) and (len(dataset_features) > 0):
+        for ntype_name, ntype_feature_data in dataset_features.items():
+            for feat_name, feat_data in ntype_feature_data.items():
+                assert feat_data[constants.STR_FORMAT][constants.STR_NAME] in [
+                    constants.STR_NUMPY,
+                    constants.STR_PARQUET,
+                ]
+
+                # It is guaranteed that num_chunks is always greater
+                # than num_partitions.
+                node_data = []
+                num_files = len(feat_data[constants.STR_DATA])
+                if num_files == 0:
+                    continue
+                reader_fmt_meta = {
+                    "name": feat_data[constants.STR_FORMAT][constants.STR_NAME]
+                }
+                read_list = generate_read_list(num_files, world_size)
+                for idx in read_list[rank]:
+                    data_file = feat_data[constants.STR_DATA][idx]
+                    if not os.path.isabs(data_file):
+                        data_file = os.path.join(input_dir, data_file)
+                    node_data.append(
+                        array_readwriter.get_array_parser(
+                            **reader_fmt_meta
+                        ).read(data_file)
+                    )
+                if len(node_data) > 0:
+                    node_data = np.concatenate(node_data)
+                else:
+                    node_data = np.array([])
+                node_data = torch.from_numpy(node_data)
+                cur_tids = _broadcast_shape(
+                    node_data,
+                    rank,
+                    world_size,
+                    num_parts,
+                    True,
+                    f"{ntype_name}/{feat_name}",
+                )
+                logging.debug(f"[Rank: {rank} - cur_tids: {cur_tids}")
+
+                # collect data on current rank.
+                for local_part_id in range(num_parts):
+                    data_key = (
+                        f"{ntype_name}/{feat_name}/{local_part_id//world_size}"
+                    )
+                    if map_partid_rank(local_part_id, world_size) == rank:
+                        if len(cur_tids) > local_part_id:
+                            start, end = cur_tids[local_part_id]
+                            assert node_data.shape[0] == (
+                                end - start
+                            ), f"Node feature data, {data_key}, shape = {node_data.shape} does not match with tids = ({start},{end})"
+                            node_features[data_key] = node_data
+                            node_feature_tids[data_key] = [(start, end)]
+                        else:
+                            node_features[data_key] = None
+                            node_feature_tids[data_key] = [(0, 0)]
+
+    # done building node_features locally.
+    if len(node_features) <= 0:
+        logging.debug(
+            f"[Rank: {rank}] This dataset does not have any node features"
+        )
+    else:
+        assert len(node_features) == len(node_feature_tids)
+
+        # Note that the keys in the node_features dictionary are as follows:
+        # `ntype_name/feat_name/local_part_id`.
+        #   where ntype_name and feat_name are self-explanatory, and
+        #   local_part_id indicates the partition-id, in the context of current
+        #   process which take the values 0, 1, 2, ....
+        for feat_name, feat_info in node_features.items():
+            if feat_info == None:
+                continue
+
+            logging.debug(
+                f"[Rank: {rank}] node feature name: {feat_name}, feature data shape: {feat_info.size()}"
+            )
+            tokens = feat_name.split("/")
+            assert len(tokens) == 3
+
+            # Get the range of type ids which are mapped to the current node.
+            tids = node_feature_tids[feat_name]
+
+            # Iterate over the range of type ids for the current node feature
+            # and count the number of features for this feature name.
+            count = tids[0][1] - tids[0][0]
+            assert (
+                count == feat_info.size()[0]
+            ), f"{feat_name}, {count} vs {feat_info.size()[0]}."
+
+    """
+    Reading edge features now.
+    The structure of the edge_data is as follows, which is present in the input metadata json file.
+       "edge_data" : {
+            "etype0-name" : {
+                "feat0-name" : {
+                    "format" : {"name": "numpy"},
+                    "data" :   [ #list
+                        "<path>/feat-0.npy",
+                        "<path>/feat-1.npy",
+                        ....
+                        "<path>/feat-<p-1>.npy"
+                    ]
+                },
+                "feat1-name" : {
+                    "format" : {"name": "numpy"},
+                    "data" : [ #list
+                        "<path>/feat-0.npy",
+                        "<path>/feat-1.npy",
+                        ....
+                        "<path>/feat-<p-1>.npy"
+                    ]
+                }
+            }
+       }
+
+    As shown above, the value for the key "edge_data" is a dictionary object, which is
+    used to describe the feature data for each of the edge-type names. Keys in this top-level
+    dictionary are edge-type names and value is a dictionary which captures all the features
+    for the current edge-type. Feature data is captured with keys being the feature-names and
+    value is a dictionary object which has 2 keys namely `format` and `data`. Format entry is used
+    to mention the format of the storage used by the node features themselves and "data" is used
+    to mention all the files present for this given node feature.
+
+    Data read from each of the node features file is a multi-dimensional tensor data and is read
+    in numpy format, which is also the storage format of node features on the permanent storage.
+    """
+    edge_features = {}
+    edge_feature_tids = {}
+
+    # Iterate over the "edge_data" dictionary in the schema_map.
+    # Read the edge features if exists.
+    # Also keep track of the type_eids for which the edge_features are read.
+    dataset_features = schema_map[constants.STR_EDGE_DATA]
+    if dataset_features and (len(dataset_features) > 0):
+        for etype_name, etype_feature_data in dataset_features.items():
+            for feat_name, feat_data in etype_feature_data.items():
+                assert feat_data[constants.STR_FORMAT][constants.STR_NAME] in [
+                    constants.STR_NUMPY,
+                    constants.STR_PARQUET,
+                ]
+
+                edge_data = []
+                num_files = len(feat_data[constants.STR_DATA])
+                if num_files == 0:
+                    continue
+                reader_fmt_meta = {
+                    "name": feat_data[constants.STR_FORMAT][constants.STR_NAME]
+                }
+                read_list = generate_read_list(num_files, world_size)
+                for idx in read_list[rank]:
+                    data_file = feat_data[constants.STR_DATA][idx]
+                    if not os.path.isabs(data_file):
+                        data_file = os.path.join(input_dir, data_file)
+                    logging.debug(
+                        f"[Rank: {rank}] Loading edges-feats of {etype_name}[{feat_name}] from {data_file}"
+                    )
+                    edge_data.append(
+                        array_readwriter.get_array_parser(
+                            **reader_fmt_meta
+                        ).read(data_file)
+                    )
+                if len(edge_data) > 0:
+                    edge_data = np.concatenate(edge_data)
+                else:
+                    edge_data = np.array([])
+                edge_data = torch.from_numpy(edge_data)
+
+                # exchange the amount of data read from the disk.
+                edge_tids = _broadcast_shape(
+                    edge_data,
+                    rank,
+                    world_size,
+                    num_parts,
+                    True,
+                    f"{etype_name}/{feat_name}",
+                )
+
+                # collect data on current rank.
+                for local_part_id in range(num_parts):
+                    data_key = (
+                        f"{etype_name}/{feat_name}/{local_part_id//world_size}"
+                    )
+                    if map_partid_rank(local_part_id, world_size) == rank:
+                        if len(edge_tids) > local_part_id:
+                            start, end = edge_tids[local_part_id]
+                            assert edge_data.shape[0] == (
+                                end - start
+                            ), f"Edge Feature data, for {data_key}, of shape = {edge_data.shape} does not match with tids = ({start}, {end})"
+                            edge_features[data_key] = edge_data
+                            edge_feature_tids[data_key] = [(start, end)]
+                        else:
+                            edge_features[data_key] = None
+                            edge_feature_tids[data_key] = [(0, 0)]
+
+    # Done with building node_features locally.
+    if len(edge_features) <= 0:
+        logging.debug(
+            f"[Rank: {rank}] This dataset does not have any edge features"
+        )
+    else:
+        assert len(edge_features) == len(edge_feature_tids)
+
+        for k, v in edge_features.items():
+            if v == None:
+                continue
+            logging.debug(
+                f"[Rank: {rank}] edge feature name: {k}, feature data shape: {v.shape}"
+            )
+            tids = edge_feature_tids[k]
+            count = tids[0][1] - tids[0][0]
+            assert count == v.size()[0]
+
+    """
+    Code below is used to read edges from the input dataset with the help of the metadata json file
+    for the input graph dataset.
+    In the metadata json file, we expect the following key-value pairs to help read the edges of the
+    input graph.
+
+    "edge_type" : [ # a total of n edge types
+        canonical_etype_0,
+        canonical_etype_1,
+        ...,
+        canonical_etype_n-1
+    ]
+
+    The value for the key is a list of strings, each string is associated with an edgetype in the input graph.
+    Note that these strings are in canonical edgetypes format. This means, these edge type strings follow the
+    following naming convention: src_ntype:etype:dst_ntype. src_ntype and dst_ntype are node type names of the
+    src and dst end points of this edge type, and etype is the relation name between src and dst ntypes.
+
+    The files in which edges are present and their storage format are present in the following key-value pair:
+
+    "edges" : {
+        "canonical_etype_0" : {
+            "format" : { "name" : "csv", "delimiter" : " " },
+            "data" : [
+                filename_0,
+                filename_1,
+                filename_2,
+                ....
+                filename_<p-1>
+            ]
+        },
+    }
+
+    As shown above the "edges" dictionary value has canonical edgetypes as keys and for each canonical edgetype
+    we have "format" and "data" which describe the storage format of the edge files and actual filenames respectively.
+    Please note that each edgetype data is split in to `p` files, where p is the no. of partitions to be made of
+    the input graph.
+
+    Each edge file contains two columns representing the source per-type node_ids and destination per-type node_ids
+    of any given edge. Since these are node-ids as well they are read in as int64's.
+    """
+
+    # read my edges for each edge type
+    etype_names = schema_map[constants.STR_EDGE_TYPE]
+    etype_name_idmap = {e: idx for idx, e in enumerate(etype_names)}
+
+    edge_tids = {}
+    edge_typecounts = {}
+    edge_datadict = {}
+    edge_data = schema_map[constants.STR_EDGES]
+
+    # read the edges files and store this data in memory.
+    for col in [
+        constants.GLOBAL_SRC_ID,
+        constants.GLOBAL_DST_ID,
+        constants.GLOBAL_TYPE_EID,
+        constants.ETYPE_ID,
+    ]:
+        edge_datadict[col] = []
+
+    for etype_name, etype_id in etype_name_idmap.items():
+        etype_info = edge_data[etype_name]
+        edge_info = etype_info[constants.STR_DATA]
+
+        # edgetype strings are in canonical format, src_node_type:edge_type:dst_node_type
+        tokens = etype_name.split(":")
+        assert len(tokens) == 3
+
+        src_ntype_name = tokens[0]
+        dst_ntype_name = tokens[2]
+
+        num_chunks = len(edge_info)
+        read_list = generate_read_list(num_chunks, world_size)
+        src_ids = []
+        dst_ids = []
+
+        """
+        curr_partids = []
+        for part_id in range(num_parts):
+            if map_partid_rank(part_id, world_size) == rank:
+                curr_partids.append(read_list[part_id])
+
+        for idx in np.concatenate(curr_partids):
+        """
+        for idx in read_list[rank]:
+            edge_file = edge_info[idx]
+            if not os.path.isabs(edge_file):
+                edge_file = os.path.join(input_dir, edge_file)
+            logging.debug(
+                f"[Rank: {rank}] Loading edges of etype[{etype_name}] from {edge_file}"
+            )
+
+            if (
+                etype_info[constants.STR_FORMAT][constants.STR_NAME]
+                == constants.STR_CSV
+            ):
+                read_options = pyarrow.csv.ReadOptions(
+                    use_threads=True,
+                    block_size=4096,
+                    autogenerate_column_names=True,
+                )
+                parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
+
+                if os.path.getsize(edge_file) == 0:
+                    # if getsize() == 0, the file is empty, indicating that the partition doesn't have this attribute.
+                    # The src_ids and dst_ids should remain empty.
+                    continue
+                with pyarrow.csv.open_csv(
+                    edge_file,
+                    read_options=read_options,
+                    parse_options=parse_options,
+                ) as reader:
+                    for next_chunk in reader:
+                        if next_chunk is None:
+                            break
+
+                        next_table = pyarrow.Table.from_batches([next_chunk])
+                        src_ids.append(next_table["f0"].to_numpy())
+                        dst_ids.append(next_table["f1"].to_numpy())
+            elif (
+                etype_info[constants.STR_FORMAT][constants.STR_NAME]
+                == constants.STR_PARQUET
+            ):
+                data_df = pq.read_table(edge_file)
+                data_df = data_df.rename_columns(["f0", "f1"])
+                src_ids.append(data_df["f0"].to_numpy())
+                dst_ids.append(data_df["f1"].to_numpy())
+            else:
+                raise ValueError(
+                    f"Unknown edge format {etype_info[constants.STR_FORMAT][constants.STR_NAME]} for edge type {etype_name}"
+                )
+
+        if len(src_ids) > 0:
+            src_ids = np.concatenate(src_ids)
+            dst_ids = np.concatenate(dst_ids)
+
+            # currently these are just type_edge_ids... which will be converted to global ids
+            edge_datadict[constants.GLOBAL_SRC_ID].append(
+                src_ids + ntype_gnid_offset[src_ntype_name][0]
+            )
+            edge_datadict[constants.GLOBAL_DST_ID].append(
+                dst_ids + ntype_gnid_offset[dst_ntype_name][0]
+            )
+            edge_datadict[constants.ETYPE_ID].append(
+                etype_name_idmap[etype_name]
+                * np.ones(shape=(src_ids.shape), dtype=np.int64)
+            )
+        else:
+            src_ids = np.array([])
+
+        # broadcast shape to compute the etype_id, and global_eid's later.
+        cur_tids = _broadcast_shape(
+            src_ids, rank, world_size, num_parts, False, None
+        )
+        edge_typecounts[etype_name] = cur_tids[-1][1]
+        edge_tids[etype_name] = cur_tids
+
+        for local_part_id in range(num_parts):
+            if map_partid_rank(local_part_id, world_size) == rank:
+                if len(cur_tids) > local_part_id:
+                    edge_datadict[constants.GLOBAL_TYPE_EID].append(
+                        np.arange(
+                            cur_tids[local_part_id][0],
+                            cur_tids[local_part_id][1],
+                            dtype=np.int64,
+                        )
+                    )
+                    # edge_tids[etype_name] = [(cur_tids[local_part_id][0], cur_tids[local_part_id][1])]
+                    assert len(edge_datadict[constants.GLOBAL_SRC_ID]) == len(
+                        edge_datadict[constants.GLOBAL_TYPE_EID]
+                    ), f"Error while reading edges from the disk, local_part_id = {local_part_id}, num_parts = {num_parts}, world_size = {world_size} cur_tids = {cur_tids}"
+
+    # stitch together to create the final data on the local machine
+    for col in [
+        constants.GLOBAL_SRC_ID,
+        constants.GLOBAL_DST_ID,
+        constants.GLOBAL_TYPE_EID,
+        constants.ETYPE_ID,
+    ]:
+        if len(edge_datadict[col]) > 0:
+            edge_datadict[col] = np.concatenate(edge_datadict[col])
+
+    if len(edge_datadict[constants.GLOBAL_SRC_ID]) > 0:
+        assert (
+            edge_datadict[constants.GLOBAL_SRC_ID].shape
+            == edge_datadict[constants.GLOBAL_DST_ID].shape
+        )
+        assert (
+            edge_datadict[constants.GLOBAL_DST_ID].shape
+            == edge_datadict[constants.GLOBAL_TYPE_EID].shape
+        )
+        assert (
+            edge_datadict[constants.GLOBAL_TYPE_EID].shape
+            == edge_datadict[constants.ETYPE_ID].shape
+        )
+        logging.debug(
+            f"[Rank: {rank}] Done reading edge_file: {len(edge_datadict)}, {edge_datadict[constants.GLOBAL_SRC_ID].shape}"
+        )
+    else:
+        assert edge_datadict[constants.GLOBAL_SRC_ID] == []
+        assert edge_datadict[constants.GLOBAL_DST_ID] == []
+        assert edge_datadict[constants.GLOBAL_TYPE_EID] == []
+
+        edge_datadict[constants.GLOBAL_SRC_ID] = np.array([], dtype=np.int64)
+        edge_datadict[constants.GLOBAL_DST_ID] = np.array([], dtype=np.int64)
+        edge_datadict[constants.GLOBAL_TYPE_EID] = np.array([], dtype=np.int64)
+        edge_datadict[constants.ETYPE_ID] = np.array([], dtype=np.int64)
+
+    logging.debug(f"Rank: {rank} edge_feat_tids: {edge_feature_tids}")
+
+    return (
+        node_features,
+        node_feature_tids,
+        edge_datadict,
+        edge_typecounts,
+        edge_tids,
+        edge_features,
+        edge_feature_tids,
+    )

tools/distpartitioning/dist_lookup.py

@@ -0,0 +1,451 @@
+import copy
+import logging
+import os
+
+import numpy as np
+import pyarrow
+import torch
+from gloo_wrapper import allgather_sizes, alltoallv_cpu
+from pyarrow import csv
+from utils import map_partid_rank
+
+
+class DistLookupService:
+    """
+    This is an implementation of a Distributed Lookup Service to provide the following
+    services to its users. Map 1) global node-ids to partition-ids, and 2) global node-ids
+    to shuffle global node-ids (contiguous, within each node for a give node_type and across
+    all the partitions)
+
+    This services initializes itself with the node-id to partition-id mappings, which are inputs
+    to this service. The node-id to partition-id  mappings are assumed to be in one file for each
+    node type. These node-id-to-partition-id mappings are split within the service processes so that
+    each process ends up with a contiguous chunk. It first divides the no of mappings (node-id to
+    partition-id) for each node type into equal chunks across all the service processes. So each
+    service process will be thse owner of a set of node-id-to-partition-id mappings. This class
+    has two functions which are as follows:
+
+    1) `get_partition_ids` function which returns the node-id to partition-id mappings to the user
+    2) `get_shuffle_nids` function which returns the node-id to shuffle-node-id mapping to the user
+
+    Parameters:
+    -----------
+    input_dir : string
+        string representing the input directory where the node-type partition-id
+        files are located
+    ntype_names : list of strings
+        list of strings which are used to read files located within the input_dir
+        directory and these files contents are partition-id's for the node-ids which
+        are of a particular node type
+    id_map : dgl.distributed.id_map instance
+        this id_map is used to retrieve ntype-ids, node type ids, and type_nids, per type
+        node ids, for any given global node id
+    rank : integer
+        integer indicating the rank of a given process
+    world_size : integer
+        integer indicating the total no. of processes
+    num_parts : integer
+        interger representing the no. of partitions
+    """
+
+    def __init__(self, input_dir, ntype_names, rank, world_size, num_parts):
+        assert os.path.isdir(input_dir)
+        assert ntype_names is not None
+        assert len(ntype_names) > 0
+
+        # These lists are indexed by ntype_ids.
+        type_nid_begin = []
+        type_nid_end = []
+        partid_list = []
+        ntype_count = []
+        ntypes = []
+
+        # Iterate over the node types and extract the partition id mappings.
+        for ntype in ntype_names:
+
+            filename = f"{ntype}.txt"
+            logging.debug(
+                f"[Rank: {rank}] Reading file: {os.path.join(input_dir, filename)}"
+            )
+
+            read_options = pyarrow.csv.ReadOptions(
+                use_threads=True,
+                block_size=4096,
+                autogenerate_column_names=True,
+            )
+            parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
+            ntype_partids = []
+            with pyarrow.csv.open_csv(
+                os.path.join(input_dir, "{}.txt".format(ntype)),
+                read_options=read_options,
+                parse_options=parse_options,
+            ) as reader:
+                for next_chunk in reader:
+                    if next_chunk is None:
+                        break
+                    next_table = pyarrow.Table.from_batches([next_chunk])
+                    ntype_partids.append(next_table["f0"].to_numpy())
+
+            ntype_partids = np.concatenate(ntype_partids)
+            count = len(ntype_partids)
+            ntype_count.append(count)
+            ntypes.append(ntype)
+
+            # Each rank assumes a contiguous set of partition-ids which are equally split
+            # across all the processes.
+            split_size = np.ceil(count / np.int64(world_size)).astype(np.int64)
+            start, end = (
+                np.int64(rank) * split_size,
+                np.int64(rank + 1) * split_size,
+            )
+            if rank == (world_size - 1):
+                end = count
+            type_nid_begin.append(start)
+            type_nid_end.append(end)
+
+            # Slice the partition-ids which belong to the current instance.
+            partid_list.append(copy.deepcopy(ntype_partids[start:end]))
+
+            # Explicitly release the array read from the file.
+            del ntype_partids
+
+        logging.debug(
+            f"[Rank: {rank}] ntypeid begin - {type_nid_begin} - {type_nid_end}"
+        )
+
+        # Store all the information in the object instance variable.
+        self.type_nid_begin = np.array(type_nid_begin, dtype=np.int64)
+        self.type_nid_end = np.array(type_nid_end, dtype=np.int64)
+        self.partid_list = partid_list
+        self.ntype_count = np.array(ntype_count, dtype=np.int64)
+        self.ntypes = ntypes
+        self.rank = rank
+        self.world_size = world_size
+        self.num_parts = num_parts
+
+    def set_idMap(self, id_map):
+        self.id_map = id_map
+
+    def get_partition_ids(self, agg_global_nids):
+        """
+        This function is used to get the partition-ids for a given set of global node ids
+
+        global_nids <-> partition-ids mappings are deterministically  distributed across
+        all the participating processes, within the service. A contiguous global-nids
+        (ntype-ids, per-type-nids) are stored within each process and this is determined
+        by the total no. of nodes of a given ntype-id and the rank of the process.
+
+        Process, where the global_nid <-> partition-id mapping is stored can be easily computed
+        as described above. Once this is determined we perform an alltoallv to send the request.
+        On the receiving side, each process receives a set of global_nids and retrieves corresponding
+        partition-ids using locally stored lookup tables. It builds responses to all the other
+        processes and performs alltoallv.
+
+        Once the response, partition-ids, is received, they are re-ordered corresponding to the
+        incoming global-nids order and returns to the caller.
+
+        Parameters:
+        -----------
+        self : instance of this class
+            instance of this class, which is passed by the runtime implicitly
+
+        agg_global_nids : numpy array
+            an array of aggregated global node-ids for which partition-ids are
+            to be retrieved by the distributed lookup service.
+
+        Returns:
+        --------
+        list of integers :
+            list of integers, which are the partition-ids of the global-node-ids (which is the
+            function argument)
+        """
+        CHUNK_SIZE = 200 * 1000 * 1000
+        # Determine the no. of times each process has to send alltoall messages.
+        local_rows = agg_global_nids.shape[0]
+        all_sizes = allgather_sizes(
+            [local_rows], self.world_size, self.num_parts, return_sizes=True
+        )
+        max_count = np.amax(all_sizes)
+
+        if max_count <= 0:
+            logging.debug(
+                f"[Rank: {self.rank}] No process has global_nids to process !!!"
+            )
+            return
+
+        num_splits = np.ceil(max_count / CHUNK_SIZE).astype(np.uint16)
+        LOCAL_CHUNK_SIZE = np.ceil(local_rows / num_splits).astype(np.int64)
+        agg_partition_ids = []
+
+        logging.debug(
+            f"[Rank: {self.rank}] BatchSize: {CHUNK_SIZE}, \
+                            max_count: {max_count}, \
+                            splits: {num_splits}, \
+                            rows: {agg_global_nids.shape}, \
+                            local batch_size: {LOCAL_CHUNK_SIZE}"
+        )
+
+        for split in range(num_splits):
+            # Compute the global_nids for this iteration
+            global_nids = agg_global_nids[
+                split * LOCAL_CHUNK_SIZE : (split + 1) * LOCAL_CHUNK_SIZE
+            ]
+
+            # Find the process where global_nid --> partition-id(owner) is stored.
+            if len(global_nids) > 0:
+                ntype_ids, type_nids = self.id_map(global_nids)
+                ntype_ids, type_nids = ntype_ids.numpy(), type_nids.numpy()
+            else:
+                ntype_ids = np.array([], dtype=np.int64)
+                type_nids = np.array([], dtype=np.int64)
+
+            assert len(ntype_ids) == len(global_nids)
+
+            # For each node-type, the per-type-node-id <-> partition-id mappings are
+            # stored as contiguous chunks by this lookup service.
+            # The no. of these mappings stored by each process, in the lookup service, are
+            # equally split among all the processes in the lookup service, deterministically.
+            typeid_counts = self.ntype_count[ntype_ids]
+            chunk_sizes = np.ceil(typeid_counts / self.world_size).astype(
+                np.int64
+            )
+            service_owners = np.floor_divide(type_nids, chunk_sizes).astype(
+                np.int64
+            )
+
+            # Now `service_owners` is a list of ranks (process-ids) which own the corresponding
+            # global-nid <-> partition-id mapping.
+
+            # Split the input global_nids into a list of lists where each list will be
+            # sent to the respective rank/process
+            # We also need to store the indices, in the indices_list, so that we can re-order
+            # the final result (partition-ids) in the same order as the global-nids (function argument)
+            send_list = []
+            indices_list = []
+            for idx in range(self.world_size):
+                idxes = np.where(service_owners == idx)
+                ll = global_nids[idxes[0]]
+                send_list.append(torch.from_numpy(ll))
+                indices_list.append(idxes[0])
+            assert len(np.concatenate(indices_list)) == len(global_nids)
+            assert np.all(
+                np.sort(np.concatenate(indices_list))
+                == np.arange(len(global_nids))
+            )
+
+            # Send the request to everyone else.
+            # As a result of this operation, the current process also receives a list of lists
+            # from all the other processes.
+            # These lists are global-node-ids whose global-node-ids <-> partition-id mappings
+            # are owned/stored by the current process
+            owner_req_list = alltoallv_cpu(
+                self.rank, self.world_size, send_list
+            )
+
+            # Create the response list here for each of the request list received in the previous
+            # step. Populate the respective partition-ids in this response lists appropriately
+            out_list = []
+            for idx in range(self.world_size):
+                if owner_req_list[idx] is None:
+                    out_list.append(torch.empty((0,), dtype=torch.int64))
+                    continue
+                # Get the node_type_ids and per_type_nids for the incoming global_nids.
+                ntype_ids, type_nids = self.id_map(owner_req_list[idx].numpy())
+                ntype_ids, type_nids = ntype_ids.numpy(), type_nids.numpy()
+
+                # Lists to store partition-ids for the incoming global-nids.
+                type_id_lookups = []
+                local_order_idx = []
+
+                # Now iterate over all the node_types and acculumulate all the partition-ids
+                # since all the partition-ids are based on the node_type order... they
+                # must be re-ordered as per the order of the input, which may be different.
+                for tid in range(len(self.partid_list)):
+                    cond = ntype_ids == tid
+                    local_order_idx.append(np.where(cond)[0])
+                    global_type_nids = type_nids[cond]
+                    if len(global_type_nids) <= 0:
+                        continue
+
+                    local_type_nids = (
+                        global_type_nids - self.type_nid_begin[tid]
+                    )
+
+                    assert np.all(local_type_nids >= 0)
+                    assert np.all(
+                        local_type_nids
+                        <= (
+                            self.type_nid_end[tid]
+                            + 1
+                            - self.type_nid_begin[tid]
+                        )
+                    )
+
+                    cur_owners = self.partid_list[tid][local_type_nids]
+                    type_id_lookups.append(cur_owners)
+
+                # Reorder the partition-ids, so that it agrees with the input order --
+                # which is the order in which the incoming message is received.
+                if len(type_id_lookups) <= 0:
+                    out_list.append(torch.empty((0,), dtype=torch.int64))
+                else:
+                    # Now reorder results for each request.
+                    sort_order_idx = np.argsort(np.concatenate(local_order_idx))
+                    lookups = np.concatenate(type_id_lookups)[sort_order_idx]
+                    out_list.append(torch.from_numpy(lookups))
+
+            # Send the partition-ids to their respective requesting processes.
+            owner_resp_list = alltoallv_cpu(
+                self.rank, self.world_size, out_list
+            )
+
+            # Owner_resp_list, is a list of lists of numpy arrays where each list
+            # is a list of partition-ids which the current process requested
+            # Now we need to re-order so that the parition-ids correspond to the
+            # global_nids which are passed into this function.
+
+            # Order according to the requesting order.
+            # Owner_resp_list is the list of owner-ids for global_nids (function argument).
+            owner_ids = [x for x in owner_resp_list if x is not None]
+            if len(owner_ids) > 0:
+                owner_ids = torch.cat(owner_ids).numpy()
+            else:
+                owner_ids = np.array([], dtype=np.int64)
+            assert len(owner_ids) == len(global_nids)
+
+            global_nids_order = np.concatenate(indices_list)
+            sort_order_idx = np.argsort(global_nids_order)
+            owner_ids = owner_ids[sort_order_idx]
+            global_nids_order = global_nids_order[sort_order_idx]
+            assert np.all(np.arange(len(global_nids)) == global_nids_order)
+
+            if len(owner_ids) > 0:
+                # Store the partition-ids for the current split
+                agg_partition_ids.append(owner_ids)
+
+        # Stitch the list of partition-ids and return to the caller
+        if len(agg_partition_ids) > 0:
+            agg_partition_ids = np.concatenate(agg_partition_ids)
+        else:
+            agg_partition_ids = np.array([], dtype=np.int64)
+        assert agg_global_nids.shape[0] == agg_partition_ids.shape[0]
+
+        # Now the owner_ids (partition-ids) which corresponding to the  global_nids.
+        return agg_partition_ids
+
+    def get_shuffle_nids(
+        self, global_nids, my_global_nids, my_shuffle_global_nids, world_size
+    ):
+        """
+        This function is used to retrieve shuffle_global_nids for a given set of incoming
+        global_nids. Note that global_nids are of random order and will contain duplicates
+
+        This function first retrieves the partition-ids of the incoming global_nids.
+        These partition-ids which are also the ranks of processes which own the respective
+        global-nids as well as shuffle-global-nids. alltoallv is performed to send the
+        global-nids to respective ranks/partition-ids where the mapping
+        global-nids <-> shuffle-global-nid is located.
+
+        On the receiving side, once the global-nids are received associated shuffle-global-nids
+        are retrieved and an alltoallv is performed to send the responses to all the other
+        processes.
+
+        Once the responses, shuffle-global-nids, are received, they are re-ordered according
+        to the incoming global-nids order and returns to the caller.
+
+        Parameters:
+        -----------
+        self : instance of this class
+            instance of this class, which is passed by the runtime implicitly
+        global_nids : numpy array
+            an array of global node-ids for which partition-ids are to be retrieved by
+            the distributed lookup service.
+        my_global_nids: numpy ndarray
+            array of global_nids which are owned by the current partition/rank/process
+            This process has the node <-> partition id mapping
+        my_shuffle_global_nids : numpy ndarray
+            array of shuffle_global_nids which are assigned by the current process/rank
+        world_size : int
+            total no. of processes in the MPI_WORLD
+
+        Returns:
+        --------
+        list of integers:
+            list of shuffle_global_nids which correspond to the incoming node-ids in the
+            global_nids.
+        """
+
+        # Get the owner_ids (partition-ids or rank).
+        owner_ids = self.get_partition_ids(global_nids)
+
+        # These owner_ids, which are also partition ids of the nodes in the
+        # input graph, are in the range 0 - (num_partitions - 1).
+        # These ids are generated using some kind of graph partitioning method.
+        # Distribuged lookup service, as used by the graph partitioning
+        # pipeline, is used to store ntype-ids (also type_nids) and their
+        # mapping to the associated partition-id.
+        # These ids are split into `num_process` chunks and processes in the
+        # dist. lookup service are assigned the owernship of these chunks.
+        # The pipeline also enforeces the following constraint among the
+        # pipeline input parameters: num_partitions, num_processes
+        #   num_partitions is an integer multiple of num_processes
+        #   which means each individual node in the cluster will be running
+        #   equal number of processes.
+        owner_ids = map_partid_rank(owner_ids, world_size)
+
+        # Ask these owners to supply for the shuffle_global_nids.
+        send_list = []
+        id_list = []
+        for idx in range(self.world_size):
+            cond = owner_ids == idx
+            idxes = np.where(cond)
+            ll = global_nids[idxes[0]]
+            send_list.append(torch.from_numpy(ll))
+            id_list.append(idxes[0])
+
+        assert len(np.concatenate(id_list)) == len(global_nids)
+        cur_global_nids = alltoallv_cpu(self.rank, self.world_size, send_list)
+
+        # At this point, current process received a list of lists each containing
+        # a list of global-nids whose corresponding shuffle_global_nids are located
+        # in the current process.
+        shuffle_nids_list = []
+        for idx in range(self.world_size):
+            if cur_global_nids[idx] is None:
+                shuffle_nids_list.append(torch.empty((0,), dtype=torch.int64))
+                continue
+
+            uniq_ids, inverse_idx = np.unique(
+                cur_global_nids[idx], return_inverse=True
+            )
+            common, idx1, idx2 = np.intersect1d(
+                uniq_ids,
+                my_global_nids,
+                assume_unique=True,
+                return_indices=True,
+            )
+            assert len(common) == len(uniq_ids)
+
+            req_shuffle_global_nids = my_shuffle_global_nids[idx2][inverse_idx]
+            assert len(req_shuffle_global_nids) == len(cur_global_nids[idx])
+            shuffle_nids_list.append(torch.from_numpy(req_shuffle_global_nids))
+
+        # Send the shuffle-global-nids to their respective ranks.
+        mapped_global_nids = alltoallv_cpu(
+            self.rank, self.world_size, shuffle_nids_list
+        )
+        for idx in range(len(mapped_global_nids)):
+            if mapped_global_nids[idx] == None:
+                mapped_global_nids[idx] = torch.empty((0,), dtype=torch.int64)
+
+        # Reorder to match global_nids (function parameter).
+        global_nids_order = np.concatenate(id_list)
+        shuffle_global_nids = torch.cat(mapped_global_nids).numpy()
+        assert len(shuffle_global_nids) == len(global_nids)
+
+        sorted_idx = np.argsort(global_nids_order)
+        shuffle_global_nids = shuffle_global_nids[sorted_idx]
+        global_nids_ordered = global_nids_order[sorted_idx]
+        assert np.all(global_nids_ordered == np.arange(len(global_nids)))
+
+        return shuffle_global_nids

tools/distpartitioning/globalids.py

@@ -0,0 +1,284 @@
+import itertools
+import operator
+
+import constants
+
+import numpy as np
+import torch
+from dist_lookup import DistLookupService
+from gloo_wrapper import allgather_sizes, alltoallv_cpu
+from utils import memory_snapshot
+
+
+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)
+    global_nids_ranks = [torch.from_numpy(x) for x in global_nids_ranks]
+    recv_nodes = alltoallv_cpu(rank, world_size, global_nids_ranks)
+
+    # 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
+        if proc_i_nodes is not None:
+            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,
+                )
+                shuffle_global_nids = node_data[constants.SHUFFLE_GLOBAL_NID][
+                    ind1
+                ]
+                send_nodes.append(
+                    torch.from_numpy(shuffle_global_nids).type(
+                        dtype=torch.int64
+                    )
+                )
+            else:
+                send_nodes.append(torch.empty((0), dtype=torch.int64))
+        else:
+            send_nodes.append(torch.empty((0), dtype=torch.int64))
+
+    # send receive global-ids
+    recv_shuffle_global_nids = alltoallv_cpu(rank, world_size, send_nodes)
+    shuffle_global_nids = np.concatenate(
+        [x.numpy() if x is not None else [] for x in recv_shuffle_global_nids]
+    )
+    global_nids = np.concatenate([x for x in global_nids_ranks])
+    ret_val = np.column_stack([global_nids, shuffle_global_nids])
+    return ret_val
+
+
+def lookup_shuffle_global_nids_edges(
+    rank, world_size, num_parts, edge_data, id_lookup, node_data
+):
+    """
+    This function is a helper function used to lookup shuffle-global-nids for a given set of
+    global-nids using a distributed lookup service.
+
+    Parameters:
+    -----------
+    rank : integer
+        rank of the process
+    world_size : integer
+        total number of processes used in the process group
+    num_parts : integer
+        total number of output graph partitions
+    edge_data : dictionary
+        edge_data is a dicitonary with keys as column names and values as numpy arrays representing
+        all the edges present in the current graph partition
+    id_lookup : instance of DistLookupService class
+        instance of a distributed lookup service class which is used to retrieve partition-ids and
+        shuffle-global-nids for any given set of global-nids
+    node_data : dictionary
+        node_data is a dictionary with keys as column names and values as numpy arrays representing
+        all the nodes owned by the current process
+
+    Returns:
+    --------
+    dictionary :
+        dictionary where keys are column names and values are numpy arrays representing all the
+        edges present in the current graph partition
+    """
+    # Make sure that the outgoing message size does not exceed 2GB in size.
+    # Even though gloo can handle upto 10GB size of data in the outgoing messages,
+    # it needs additional memory to store temporary information into the buffers which will increase
+    # the memory needs of the process.
+    MILLION = 1000 * 1000
+    BATCH_SIZE = 250 * MILLION
+    memory_snapshot("GlobalToShuffleIDMapBegin: ", rank)
+
+    local_nids = []
+    local_shuffle_nids = []
+    for local_part_id in range(num_parts // world_size):
+        local_nids.append(
+            node_data[constants.GLOBAL_NID + "/" + str(local_part_id)]
+        )
+        local_shuffle_nids.append(
+            node_data[constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)]
+        )
+
+    local_nids = np.concatenate(local_nids)
+    local_shuffle_nids = np.concatenate(local_shuffle_nids)
+
+    for local_part_id in range(num_parts // world_size):
+        node_list = edge_data[
+            constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
+        ]
+
+        # Determine the no. of times each process has to send alltoall messages.
+        all_sizes = allgather_sizes(
+            [node_list.shape[0]], world_size, num_parts, return_sizes=True
+        )
+        max_count = np.amax(all_sizes)
+        num_splits = max_count // BATCH_SIZE + 1
+
+        # Split the message into batches and send.
+        splits = np.array_split(node_list, num_splits)
+        shuffle_mappings = []
+        for item in splits:
+            shuffle_ids = id_lookup.get_shuffle_nids(
+                item, local_nids, local_shuffle_nids, world_size
+            )
+            shuffle_mappings.append(shuffle_ids)
+
+        shuffle_ids = np.concatenate(shuffle_mappings)
+        assert shuffle_ids.shape[0] == node_list.shape[0]
+        edge_data[
+            constants.SHUFFLE_GLOBAL_SRC_ID + "/" + str(local_part_id)
+        ] = shuffle_ids
+
+        # Destination end points of edges are owned by the current node and therefore
+        # should have corresponding SHUFFLE_GLOBAL_NODE_IDs.
+        # Here retrieve SHUFFLE_GLOBAL_NODE_IDs for the destination end points of local edges.
+        uniq_ids, inverse_idx = np.unique(
+            edge_data[constants.GLOBAL_DST_ID + "/" + str(local_part_id)],
+            return_inverse=True,
+        )
+        common, idx1, idx2 = np.intersect1d(
+            uniq_ids,
+            node_data[constants.GLOBAL_NID + "/" + str(local_part_id)],
+            assume_unique=True,
+            return_indices=True,
+        )
+        assert len(common) == len(uniq_ids)
+
+        edge_data[
+            constants.SHUFFLE_GLOBAL_DST_ID + "/" + str(local_part_id)
+        ] = node_data[constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)][
+            idx2
+        ][
+            inverse_idx
+        ]
+        assert len(
+            edge_data[
+                constants.SHUFFLE_GLOBAL_DST_ID + "/" + str(local_part_id)
+            ]
+        ) == len(edge_data[constants.GLOBAL_DST_ID + "/" + str(local_part_id)])
+
+    memory_snapshot("GlobalToShuffleIDMap_AfterLookupServiceCalls: ", rank)
+    return edge_data
+
+
+def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, 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
+    num_parts : integer
+        total number of output graph partitions
+    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
+    local_row_counts = []
+    for local_part_id in range(num_parts // world_size):
+        local_row_counts.append(
+            node_data[constants.GLOBAL_NID + "/" + str(local_part_id)].shape[0]
+        )
+
+    # Perform allgather to compute the local offsets.
+    prefix_sum_nodes = allgather_sizes(local_row_counts, world_size, num_parts)
+
+    for local_part_id in range(num_parts // world_size):
+        shuffle_global_nid_start = prefix_sum_nodes[
+            rank + (local_part_id * world_size)
+        ]
+        shuffle_global_nid_end = prefix_sum_nodes[
+            rank + 1 + (local_part_id * world_size)
+        ]
+        shuffle_global_nids = np.arange(
+            shuffle_global_nid_start, shuffle_global_nid_end, dtype=np.int64
+        )
+        node_data[
+            constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)
+        ] = shuffle_global_nids
+
+
+def assign_shuffle_global_nids_edges(rank, world_size, num_parts, 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
+    num_parts : integer
+        total number of output graph partitions
+    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
+    local_row_counts = []
+    for local_part_id in range(num_parts // world_size):
+        local_row_counts.append(
+            edge_data[constants.GLOBAL_SRC_ID + "/" + str(local_part_id)].shape[
+                0
+            ]
+        )
+
+    shuffle_global_eid_offset = []
+    prefix_sum_edges = allgather_sizes(local_row_counts, world_size, num_parts)
+    for local_part_id in range(num_parts // world_size):
+        shuffle_global_eid_start = prefix_sum_edges[
+            rank + (local_part_id * world_size)
+        ]
+        shuffle_global_eid_end = prefix_sum_edges[
+            rank + 1 + (local_part_id * world_size)
+        ]
+        shuffle_global_eids = np.arange(
+            shuffle_global_eid_start, shuffle_global_eid_end, dtype=np.int64
+        )
+        edge_data[
+            constants.SHUFFLE_GLOBAL_EID + "/" + str(local_part_id)
+        ] = shuffle_global_eids
+        shuffle_global_eid_offset.append(shuffle_global_eid_start)
+
+    return shuffle_global_eid_offset

tools/distpartitioning/gloo_wrapper.py

@@ -0,0 +1,222 @@
+import numpy as np
+import torch
+import torch.distributed as dist
+
+
+def allgather_sizes(send_data, world_size, num_parts, return_sizes=False):
+    """
+    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
+    num_parts : integer
+        No. of output graph partitions
+    return_sizes : bool
+        Boolean flag to indicate whether to return raw sizes from each process
+        or perform prefix sum on the raw sizes.
+
+    Returns :
+    ---------
+        numpy array
+            array with the prefix sum
+    """
+
+    # Assert on the world_size, num_parts
+    assert (num_parts % world_size) == 0
+
+    # 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)
+
+    # Return on the raw sizes from each process
+    if return_sizes:
+        return torch.cat(in_tensor).numpy()
+
+    # gather sizes in on array to return to the invoking function
+    rank_sizes = np.zeros(num_parts + 1, dtype=np.int64)
+    part_counts = torch.cat(in_tensor).numpy()
+
+    count = rank_sizes[0]
+    idx = 1
+    for local_part_id in range(num_parts // world_size):
+        for r in range(world_size):
+            count += part_counts[r * (num_parts // world_size) + local_part_id]
+            rank_sizes[idx] = count
+            idx += 1
+
+    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
+    ]
+    # TODO(#5002): As Boolean data is not supported in
+    # ``torch.distributed.scatter()``, we convert boolean into uint8 before
+    # scatter and convert it back afterwards.
+    dtypes = [t.dtype for t in input_tensor_list]
+    for i, dtype in enumerate(dtypes):
+        if dtype == torch.bool:
+            input_tensor_list[i] = input_tensor_list[i].to(torch.int8)
+            output_tensor_list[i] = output_tensor_list[i].to(torch.int8)
+    for i in range(world_size):
+        dist.scatter(
+            output_tensor_list[i], input_tensor_list if i == rank else [], src=i
+        )
+    # Convert back to original dtype
+    for i, dtype in enumerate(dtypes):
+        if dtype == torch.bool:
+            input_tensor_list[i] = input_tensor_list[i].to(dtype)
+            output_tensor_list[i] = output_tensor_list[i].to(dtype)
+
+
+def alltoallv_cpu(rank, world_size, input_tensor_list, retain_nones=True):
+    """
+    Wrapper function to providing the alltoallv functionality by using underlying alltoall
+    messaging primitive. This function, in its current implementation, supports exchanging
+    messages of arbitrary dimensions and is not tied to the user of this function.
+
+    This function pads all input tensors, except one, so that all the messages are of the same
+    size. Once the messages are padded, It first sends a vector whose first two elements are
+    1) actual message size along first dimension, and 2) Message size along first dimension
+    which is used for communication. The rest of the dimensions are assumed to be same across
+    all the input tensors. After receiving the message sizes, the receiving end will create buffers
+    of appropriate sizes. And then slices the received messages to remove the added padding, if any,
+    and returns to the caller.
+
+    Parameters:
+    -----------
+    rank : int
+        The rank of current worker
+    world_size : int
+        The size of the entire
+    input_tensor_list : List of tensor
+        The tensors to exchange
+    retain_nones : bool
+        Indicates whether to retain ``None`` data in returned value.
+
+    Returns:
+    --------
+    list :
+        list of tensors received from other processes during alltoall message
+
+    """
+    # ensure len of input_tensor_list is same as the world_size.
+    assert input_tensor_list != None
+    assert len(input_tensor_list) == world_size
+
+    # ensure that all the tensors in the input_tensor_list are of same size.
+    sizes = [list(x.size()) for x in input_tensor_list]
+    for idx in range(1, len(sizes)):
+        assert len(sizes[idx - 1]) == len(
+            sizes[idx]
+        )  # no. of dimensions should be same
+        assert (
+            input_tensor_list[idx - 1].dtype == input_tensor_list[idx].dtype
+        )  # dtype should be same
+        assert (
+            sizes[idx - 1][1:] == sizes[idx][1:]
+        )  # except first dimension remaining dimensions should all be the same
+
+    # decide how much to pad.
+    # always use the first-dimension for padding.
+    ll = [x[0] for x in sizes]
+
+    # dims of the padding needed, if any
+    # these dims are used for padding purposes.
+    diff_dims = [[np.amax(ll) - l[0]] + l[1:] for l in sizes]
+
+    # pad the actual message
+    input_tensor_list = [
+        torch.cat((x, torch.zeros(diff_dims[idx]).type(x.dtype)))
+        for idx, x in enumerate(input_tensor_list)
+    ]
+
+    # send useful message sizes to all
+    send_counts = []
+    recv_counts = []
+    for idx in range(world_size):
+        # send a vector, of atleast 3 elements, [a, b, ....] where
+        # a = useful message dim, b = actual message outgoing message size along the first dimension
+        # and remaining elements are the remaining dimensions of the tensor
+        send_counts.append(
+            torch.from_numpy(
+                np.array([sizes[idx][0]] + [np.amax(ll)] + sizes[idx][1:])
+            ).type(torch.int64)
+        )
+        recv_counts.append(
+            torch.zeros((1 + len(sizes[idx])), dtype=torch.int64)
+        )
+    __alltoall_cpu(rank, world_size, recv_counts, send_counts)
+
+    # allocate buffers for receiving message
+    output_tensor_list = []
+    recv_counts = [tsize.numpy() for tsize in recv_counts]
+    for idx, tsize in enumerate(recv_counts):
+        output_tensor_list.append(
+            torch.zeros(tuple(tsize[1:])).type(input_tensor_list[idx].dtype)
+        )
+
+    # send actual message itself.
+    __alltoall_cpu(rank, world_size, output_tensor_list, input_tensor_list)
+
+    # extract un-padded message from the output_tensor_list and return it
+    return_vals = []
+    for s, t in zip(recv_counts, output_tensor_list):
+        if s[0] == 0:
+            if retain_nones:
+                return_vals.append(None)
+        else:
+            return_vals.append(t[0 : s[0]])
+    return return_vals
+
+
+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/parmetis_postprocess.py

@@ -0,0 +1,150 @@
+import argparse
+import logging
+import os
+import platform
+import sys
+from pathlib import Path
+
+import constants
+
+import numpy as np
+import pyarrow
+import pyarrow.csv as csv
+from partition_algo.base import dump_partition_meta, PartitionMeta
+from utils import get_idranges, get_node_types, read_json
+
+
+def post_process(params):
+    """Auxiliary function to read the parmetis output file and generate
+    metis partition-id files, sorted, per node-type. These files are used
+    by the dist. graph partitioning pipeline for further processing.
+
+    Parameters:
+    -----------
+    params : argparser object
+        argparser object to capture command line options passed to the
+        executable
+    """
+    logging.info("Starting to process parmetis output.")
+
+    logging.info(params.postproc_input_dir)
+    logging.info(params.schema_file)
+    logging.info(params.parmetis_output_file)
+    assert os.path.isfile(
+        os.path.join(params.postproc_input_dir, params.schema_file)
+    )
+    assert os.path.isfile(params.parmetis_output_file)
+    schema = read_json(
+        os.path.join(params.postproc_input_dir, params.schema_file)
+    )
+
+    metis_df = csv.read_csv(
+        params.parmetis_output_file,
+        read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
+        parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
+    )
+    global_nids = metis_df["f0"].to_numpy()
+    partition_ids = metis_df["f1"].to_numpy()
+    num_parts = np.unique(partition_ids).size
+
+    sort_idx = np.argsort(global_nids)
+    global_nids = global_nids[sort_idx]
+    partition_ids = partition_ids[sort_idx]
+
+    ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(schema)
+    type_nid_dict, ntype_gnid_offset = get_idranges(
+        schema[constants.STR_NODE_TYPE],
+        dict(
+            zip(
+                schema[constants.STR_NODE_TYPE],
+                schema[constants.STR_NUM_NODES_PER_TYPE],
+            )
+        ),
+    )
+
+    outdir = Path(params.partitions_dir)
+    os.makedirs(outdir, exist_ok=True)
+    for ntype_id, ntype_name in ntid_ntype_map.items():
+        start = ntype_gnid_offset[ntype_name][0, 0]
+        end = ntype_gnid_offset[ntype_name][0, 1]
+        out_data = partition_ids[start:end]
+
+        out_file = os.path.join(outdir, f"{ntype_name}.txt")
+        options = csv.WriteOptions(include_header=False, delimiter=" ")
+
+        csv.write_csv(
+            pyarrow.Table.from_arrays([out_data], names=["partition-ids"]),
+            out_file,
+            options,
+        )
+        logging.info(f"Generated {out_file}")
+
+    # generate partition meta file.
+    part_meta = PartitionMeta(
+        version="1.0.0", num_parts=num_parts, algo_name="metis"
+    )
+    dump_partition_meta(part_meta, os.path.join(outdir, "partition_meta.json"))
+
+    logging.info("Done processing parmetis output")
+
+
+if __name__ == "__main__":
+    """Main function to convert the output of parmetis into metis partitions
+    which are accepted by graph partitioning pipeline.
+
+    ParMETIS currently generates one output file, which is in the following format:
+    <global-node-id> <partition-id>
+
+    Graph partitioing pipeline, per the new dataset file format rules expects the
+    metis partitions to be in the following format:
+    No. of files will be equal to the no. of node-types in the graph
+    Each file will have one-number/line which is <partition-id>.
+
+    Example usage:
+    --------------
+    python parmetis_postprocess.py
+        --input_file <metis-partitions-file>
+        --output-dir <directory where the output files are stored>
+        --schema <schema-file-path>
+    """
+    parser = argparse.ArgumentParser(
+        description="PostProcessing the ParMETIS\
+        output for partitioning pipeline"
+    )
+    parser.add_argument(
+        "--postproc_input_dir",
+        required=True,
+        type=str,
+        help="Base directory for post processing step.",
+    )
+    parser.add_argument(
+        "--schema_file",
+        required=True,
+        type=str,
+        help="The schema of the input graph",
+    )
+    parser.add_argument(
+        "--parmetis_output_file",
+        required=True,
+        type=str,
+        help="ParMETIS output file",
+    )
+    parser.add_argument(
+        "--partitions_dir",
+        required=True,
+        type=str,
+        help="The output\
+        will be files (with metis partition ids) and each file corresponds to\
+        a node-type in the input graph dataset.",
+    )
+    params = parser.parse_args()
+
+    # Configure logging.
+    logging.basicConfig(
+        level="INFO",
+        format=f"[{platform.node()} \
+        %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
+    )
+
+    # Invoke the function for post processing
+    post_process(params)

tools/distpartitioning/parmetis_preprocess.py

@@ -0,0 +1,461 @@
+import argparse
+import logging
+import os
+import platform
+from pathlib import Path
+
+import array_readwriter
+
+import constants
+
+import numpy as np
+import pyarrow
+import pyarrow.csv as csv
+from utils import (
+    generate_read_list,
+    generate_roundrobin_read_list,
+    get_idranges,
+    get_node_types,
+    read_json,
+)
+
+
+def get_proc_info():
+    """Helper function to get the rank from the
+    environment when `mpirun` is used to run this python program.
+
+    Please note that for mpi(openmpi) installation the rank is retrieved from the
+    environment using OMPI_COMM_WORLD_RANK. For mpich it is
+    retrieved from the environment using PMI_RANK.
+
+    Returns:
+    --------
+    integer :
+        Rank of the current process.
+    """
+    env_variables = dict(os.environ)
+    # mpich
+    if "PMI_RANK" in env_variables:
+        return int(env_variables["PMI_RANK"])
+    # openmpi
+    elif "OMPI_COMM_WORLD_RANK" in env_variables:
+        return int(env_variables["OMPI_COMM_WORLD_RANK"])
+    else:
+        return 0
+
+
+def get_world_size():
+    """Helper function to get the world size from the
+    environment when `mpirun` is used to run this python program.
+
+    Returns:
+    --------
+    integer :
+        Numer of processes created by the executor that created this process.
+    """
+    env_variables = dict(os.environ)
+    # mpich
+    if "PMI_SIZE" in env_variables:
+        return int(env_variables["PMI_SIZE"])
+    # openmpi
+    elif "OMPI_COMM_WORLD_SIZE" in env_variables:
+        return int(env_variables["OMPI_COMM_WORLD_SIZE"])
+    else:
+        return 1
+
+
+def gen_edge_files(rank, schema_map, params):
+    """Function to create edges files to be consumed by ParMETIS
+    for partitioning purposes.
+
+    This function creates the edge files and each of these will have the
+    following format (meaning each line of these file is of the following format)
+    <global_src_id> <global_dst_id>
+
+    Here ``global`` prefix means that globally unique identifier assigned each node
+    in the input graph. In this context globally unique means unique across all the
+    nodes in the input graph.
+
+    Parameters:
+    -----------
+    rank : int
+        rank of the current process
+    schema_map : json dictionary
+        Dictionary created by reading the metadata.json file for the input dataset.
+    output : string
+        Location of storing the node-weights and edge files for ParMETIS.
+    """
+    _, ntype_gnid_offset = get_idranges(
+        schema_map[constants.STR_NODE_TYPE],
+        dict(
+            zip(
+                schema_map[constants.STR_NODE_TYPE],
+                schema_map[constants.STR_NUM_NODES_PER_TYPE],
+            )
+        ),
+    )
+
+    # Regenerate edge files here.
+    edge_data = schema_map[constants.STR_EDGES]
+
+    outdir = Path(params.output_dir)
+    os.makedirs(outdir, exist_ok=True)
+
+    def process_and_write_back(data_df, idx):
+        data_f0 = data_df[:, 0]
+        data_f1 = data_df[:, 1]
+
+        global_src_id = data_f0 + ntype_gnid_offset[src_ntype_name][0, 0]
+        global_dst_id = data_f1 + ntype_gnid_offset[dst_ntype_name][0, 0]
+        cols = [global_src_id, global_dst_id]
+        col_names = ["global_src_id", "global_dst_id"]
+
+        out_file_name = Path(edge_data_files[idx]).stem.split(".")[0]
+        out_file = os.path.join(
+            outdir, etype_name, f"edges_{out_file_name}.csv"
+        )
+        os.makedirs(os.path.dirname(out_file), exist_ok=True)
+
+        options = csv.WriteOptions(include_header=False, delimiter=" ")
+        csv.write_csv(
+            pyarrow.Table.from_arrays(cols, names=col_names),
+            out_file,
+            options,
+        )
+        return out_file
+
+    edge_files = []
+    for etype_name, etype_info in edge_data.items():
+        edge_data_files = etype_info[constants.STR_DATA]
+
+        # ``edgetype`` strings are in canonical format, src_node_type:edge_type:dst_node_type
+        tokens = etype_name.split(":")
+        assert len(tokens) == 3
+
+        src_ntype_name = tokens[0]
+
+        dst_ntype_name = tokens[2]
+
+        rank_assignments = generate_roundrobin_read_list(
+            len(edge_data_files), params.num_parts
+        )
+        for file_idx in rank_assignments[rank]:
+            reader_fmt_meta = {
+                "name": etype_info[constants.STR_FORMAT][constants.STR_NAME],
+            }
+            if reader_fmt_meta["name"] == constants.STR_CSV:
+                reader_fmt_meta["delimiter"] = etype_info[constants.STR_FORMAT][
+                    constants.STR_FORMAT_DELIMITER
+                ]
+            data_df = array_readwriter.get_array_parser(**reader_fmt_meta).read(
+                os.path.join(params.input_dir, edge_data_files[file_idx])
+            )
+            out_file = process_and_write_back(data_df, file_idx)
+            edge_files.append(out_file)
+
+    return edge_files
+
+
+def gen_node_weights_files(schema_map, params):
+    """Function to create node weight files for ParMETIS along with the edge files.
+
+    This function generates node-data files, which will be read by the ParMETIS
+    executable for partitioning purposes. Each line in these files will be of the
+    following format:
+        <node_type_id> <node_weight_list> <type_wise_node_id>
+    node_type_id -  is id assigned to the node-type to which a given particular
+        node belongs to
+    weight_list - this is a one-hot vector in which the number in the location of
+        the current nodes' node-type will be set to `1` and other will be `0`
+    type_node_id - this is the id assigned to the node (in the context of the current
+        nodes` node-type). Meaning this id is unique across all the nodes which belong to
+        the current nodes` node-type.
+
+    Parameters:
+    -----------
+    schema_map : json dictionary
+        Dictionary created by reading the metadata.json file for the input dataset.
+    output : string
+        Location of storing the node-weights and edge files for ParMETIS.
+
+    Returns:
+    --------
+    list :
+        List of filenames for nodes of the input graph.
+    list :
+        List o ffilenames for edges of the input graph.
+    """
+    rank = get_proc_info()
+    ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(schema_map)
+    type_nid_dict, ntype_gnid_offset = get_idranges(
+        schema_map[constants.STR_NODE_TYPE],
+        dict(
+            zip(
+                schema_map[constants.STR_NODE_TYPE],
+                schema_map[constants.STR_NUM_NODES_PER_TYPE],
+            )
+        ),
+    )
+
+    node_files = []
+    outdir = Path(params.output_dir)
+    os.makedirs(outdir, exist_ok=True)
+
+    for ntype_id, ntype_name in ntid_ntype_map.items():
+
+        # This ntype does not have any train/test/val masks...
+        # Each rank will generate equal no. of rows for this node type.
+        total_count = schema_map[constants.STR_NUM_NODES_PER_TYPE][ntype_id]
+        per_rank_range = np.ones((params.num_parts,), dtype=np.int64) * (
+            total_count // params.num_parts
+        )
+        for i in range(total_count % params.num_parts):
+            per_rank_range[i] += 1
+
+        tid_start = np.cumsum([0] + list(per_rank_range[:-1]))
+        tid_end = np.cumsum(list(per_rank_range))
+        local_tid_start = tid_start[rank]
+        local_tid_end = tid_end[rank]
+        sz = local_tid_end - local_tid_start
+
+        cols = []
+        col_names = []
+
+        # ntype-id
+        cols.append(
+            pyarrow.array(np.ones(sz, dtype=np.int64) * np.int64(ntype_id))
+        )
+        col_names.append("ntype")
+
+        # one-hot vector for ntype-id here.
+        for i in range(len(ntypes)):
+            if i == ntype_id:
+                cols.append(pyarrow.array(np.ones(sz, dtype=np.int64)))
+            else:
+                cols.append(pyarrow.array(np.zeros(sz, dtype=np.int64)))
+            col_names.append("w{}".format(i))
+
+        # `type_nid` should be the very last column in the node weights files.
+        cols.append(
+            pyarrow.array(
+                np.arange(local_tid_start, local_tid_end, dtype=np.int64)
+            )
+        )
+        col_names.append("type_nid")
+
+        out_file = os.path.join(
+            outdir, "node_weights_{}_{}.txt".format(ntype_name, rank)
+        )
+        options = csv.WriteOptions(include_header=False, delimiter=" ")
+        options.delimiter = " "
+
+        csv.write_csv(
+            pyarrow.Table.from_arrays(cols, names=col_names), out_file, options
+        )
+        node_files.append(
+            (
+                ntype_gnid_offset[ntype_name][0, 0] + local_tid_start,
+                ntype_gnid_offset[ntype_name][0, 0] + local_tid_end,
+                out_file,
+            )
+        )
+
+    return node_files
+
+
+def gen_parmetis_input_args(params, schema_map):
+    """Function to create two input arguments which will be passed to the parmetis.
+    first argument is a text file which has a list of node-weights files,
+    namely parmetis-nfiles.txt, and second argument is a text file which has a
+    list of edge files, namely parmetis_efiles.txt.
+    ParMETIS uses these two files to read/load the graph and partition the graph
+    With regards to the file format, parmetis_nfiles.txt uses the following format
+    for each line in that file:
+        <filename> <global_node_id_start> <global_node_id_end>(exclusive)
+    While parmetis_efiles.txt just has <filename> in each line.
+
+    Parameters:
+    -----------
+    params : argparser instance
+        Instance of ArgParser class, which has all the input arguments passed to
+        run this program.
+    schema_map : json dictionary
+        Dictionary object created after reading the graph metadata.json file.
+    """
+
+    # TODO: This makes the assumption that all node files have the same number of chunks
+    ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(schema_map)
+    type_nid_dict, ntype_gnid_offset = get_idranges(
+        schema_map[constants.STR_NODE_TYPE],
+        dict(
+            zip(
+                schema_map[constants.STR_NODE_TYPE],
+                schema_map[constants.STR_NUM_NODES_PER_TYPE],
+            )
+        ),
+    )
+
+    # Check if <graph-name>_stats.txt exists, if not create one using metadata.
+    # Here stats file will be created in the current directory.
+    # No. of constraints, third column in the stats file is computed as follows:
+    #   num_constraints = no. of node types + train_mask + test_mask + val_mask
+    #   Here, (train/test/val) masks will be set to 1 if these masks exist for
+    #   all the node types in the graph, otherwise these flags will be set to 0
+    assert (
+        constants.STR_GRAPH_NAME in schema_map
+    ), "Graph name is not present in the json file"
+    graph_name = schema_map[constants.STR_GRAPH_NAME]
+    if not os.path.isfile(
+        os.path.join(params.input_dir, f"{graph_name}_stats.txt")
+    ):
+        num_nodes = np.sum(schema_map[constants.STR_NUM_NODES_PER_TYPE])
+        num_edges = np.sum(schema_map[constants.STR_NUM_EDGES_PER_TYPE])
+        num_ntypes = len(schema_map[constants.STR_NODE_TYPE])
+
+        num_constraints = num_ntypes
+
+        with open(
+            os.path.join(params.input_dir, f"{graph_name}_stats.txt"), "w"
+        ) as sf:
+            sf.write(f"{num_nodes} {num_edges} {num_constraints}")
+
+    node_files = []
+    outdir = Path(params.output_dir)
+    os.makedirs(outdir, exist_ok=True)
+    for ntype_id, ntype_name in ntid_ntype_map.items():
+        global_nid_offset = ntype_gnid_offset[ntype_name][0, 0]
+        total_count = schema_map[constants.STR_NUM_NODES_PER_TYPE][ntype_id]
+        per_rank_range = np.ones((params.num_parts,), dtype=np.int64) * (
+            total_count // params.num_parts
+        )
+        for i in range(total_count % params.num_parts):
+            per_rank_range[i] += 1
+        tid_start = np.cumsum([0] + list(per_rank_range[:-1]))
+        tid_end = np.cumsum(per_rank_range)
+        logging.info(f" tid-start = {tid_start}, tid-end = {tid_end}")
+        logging.info(f" per_rank_range - {per_rank_range}")
+
+        for part_idx in range(params.num_parts):
+            local_tid_start = tid_start[part_idx]
+            local_tid_end = tid_end[part_idx]
+            out_file = os.path.join(
+                outdir, "node_weights_{}_{}.txt".format(ntype_name, part_idx)
+            )
+            node_files.append(
+                (
+                    out_file,
+                    global_nid_offset + local_tid_start,
+                    global_nid_offset + local_tid_end,
+                )
+            )
+
+    with open(
+        os.path.join(params.output_dir, "parmetis_nfiles.txt"), "w"
+    ) as parmetis_nf:
+        for node_file in node_files:
+            # format: filename global_node_id_start global_node_id_end(exclusive)
+            parmetis_nf.write(
+                "{} {} {}\n".format(node_file[0], node_file[1], node_file[2])
+            )
+
+    # Regenerate edge files here.
+    # NOTE: The file names need to match the ones generated by gen_edge_files function
+    edge_data = schema_map[constants.STR_EDGES]
+    edge_files = []
+    for etype_name, etype_info in edge_data.items():
+        edge_data_files = etype_info[constants.STR_DATA]
+        for edge_file_path in edge_data_files:
+            out_file_name = Path(edge_file_path).stem.split(".")[0]
+            out_file = os.path.join(
+                outdir, etype_name, "edges_{}.csv".format(out_file_name)
+            )
+            edge_files.append(out_file)
+
+    with open(
+        os.path.join(params.output_dir, "parmetis_efiles.txt"), "w"
+    ) as parmetis_efile:
+        for edge_file in edge_files:
+            parmetis_efile.write("{}\n".format(edge_file))
+
+
+def run_preprocess_data(params):
+    """Main function which will help create graph files for ParMETIS processing
+
+    Parameters:
+    -----------
+    params : argparser object
+        An instance of argparser class which stores command line arguments.
+    """
+    logging.info("Starting to generate ParMETIS files...")
+    rank = get_proc_info()
+
+    assert os.path.isdir(
+        params.input_dir
+    ), f"Please check `input_dir` argument: {params.input_dit}."
+
+    schema_map = read_json(os.path.join(params.input_dir, params.schema_file))
+    gen_node_weights_files(schema_map, params)
+    logging.info("Done with node weights....")
+
+    gen_edge_files(rank, schema_map, params)
+    logging.info("Done with edge weights...")
+
+    if rank == 0:
+        gen_parmetis_input_args(params, schema_map)
+    logging.info("Done generating files for ParMETIS run ..")
+
+
+if __name__ == "__main__":
+    """Main function used to generate temporary files needed for ParMETIS execution.
+    This function generates node-weight files and edges files which are consumed by ParMETIS.
+
+    Example usage:
+    --------------
+    mpirun -np 4 python3 parmetis_preprocess.py --schema <file> --output <target-output-dir>
+    """
+    parser = argparse.ArgumentParser(
+        description="Generate ParMETIS files for input dataset"
+    )
+    parser.add_argument(
+        "--schema_file",
+        required=True,
+        type=str,
+        help="The schema of the input graph",
+    )
+    parser.add_argument(
+        "--input_dir",
+        required=True,
+        type=str,
+        help="This directory will be used as the relative directory to locate files, if absolute paths are not used",
+    )
+    parser.add_argument(
+        "--output_dir",
+        required=True,
+        type=str,
+        help="The output directory for the node weights files and auxiliary files for ParMETIS.",
+    )
+    parser.add_argument(
+        "--num_parts",
+        required=True,
+        type=int,
+        help="Total no. of output graph partitions.",
+    )
+    parser.add_argument(
+        "--log_level",
+        required=False,
+        type=str,
+        help="Log level to use for execution.",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+    )
+    params = parser.parse_args()
+
+    # Configure logging.
+    logging.basicConfig(
+        level=getattr(logging, params.log_level, None),
+        format=f"[{platform.node()} \
+        %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
+    )
+
+    # Invoke the function to generate files for parmetis
+    run_preprocess_data(params)

tools/distpartitioning/parmetis_wrapper.py

@@ -0,0 +1,174 @@
+import argparse
+import logging
+import os
+import platform
+import sys
+from pathlib import Path
+
+import constants
+from utils import read_json
+
+
+def check_dependencies():
+    """Check if all the dependencies needed for the execution of this file
+    are installed.
+    """
+
+    exec_path = os.get_exec_path()
+    mpi_install = False
+    for x in exec_path:
+        if os.path.isfile(os.path.join(x, "mpirun")):
+            mpi_install = True
+            break
+    assert (
+        mpi_install
+    ), "Could not locate the following dependency: MPI. Please install it and try again."
+
+    dgl_path = os.environ.get("DGL_HOME", "")
+    assert os.path.isdir(
+        dgl_path
+    ), "Environment variable DGL_HOME not found. Please define the DGL installation path"
+
+
+def run_parmetis_wrapper(params):
+    """Function to execute all the steps needed to run ParMETIS
+
+    Parameters:
+    -----------
+    params : argparser object
+        an instance of argparser class to capture command-line arguments
+    """
+    schema = read_json(
+        os.path.join(params.preproc_input_dir, params.schema_file)
+    )
+    graph_name = schema[constants.STR_GRAPH_NAME]
+    num_partitions = params.num_parts
+
+    # Check if parmetis_preprocess.py exists.
+    assert os.path.isfile(
+        os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "parmetis_preprocess.py"
+        )
+    ), "Please check DGL Installation, parmetis_preprocess.py file does not exist."
+
+    # Trigger pre-processing step to generate input files for ParMETIS.
+    preproc_cmd = (
+        f"mpirun -np {num_partitions} -hostfile {params.hostfile} "
+        f"python3 $DGL_HOME/tools/distpartitioning/parmetis_preprocess.py "
+        f"--schema_file {params.schema_file} "
+        f"--input_dir {params.preproc_input_dir} "
+        f"--output_dir {params.preproc_output_dir} "
+        f"--num_parts {num_partitions}"
+    )
+    logging.info(f"Executing Preprocessing Step: {preproc_cmd}")
+    os.system(preproc_cmd)
+    logging.info(f"Done Preprocessing Step")
+
+    # Trigger ParMETIS for creating metis partitions for the input graph.
+    parmetis_install_path = "pm_dglpart3"
+    if params.parmetis_install_path is not None:
+        parmetis_install_path = os.path.join(
+            params.parmetis_install_path, parmetis_install_path
+        )
+    parmetis_nfiles = os.path.join(
+        params.preproc_output_dir, "parmetis_nfiles.txt"
+    )
+    parmetis_efiles = os.path.join(
+        params.preproc_output_dir, "parmetis_efiles.txt"
+    )
+    parmetis_cmd = (
+        f"mpirun -np {num_partitions} -hostfile {params.hostfile} "
+        f"{parmetis_install_path} {graph_name} {num_partitions} "
+        f"{parmetis_nfiles} {parmetis_efiles}"
+    )
+    logging.info(f"Executing ParMETIS: {parmetis_cmd}")
+    os.system(parmetis_cmd)
+    logging.info(f"Done ParMETIS execution step")
+
+    # Trigger post-processing step to convert parmetis output to the form
+    # acceptable by dist. graph partitioning pipeline.
+    parmetis_output_file = os.path.join(
+        os.getcwd(), f"{graph_name}_part.{num_partitions}"
+    )
+    postproc_cmd = (
+        f"python3 $DGL_HOME/tools/distpartitioning/parmetis_postprocess.py "
+        f"--postproc_input_dir {params.preproc_input_dir} "
+        f"--schema_file {params.schema_file} "
+        f"--parmetis_output_file {parmetis_output_file} "
+        f"--partitions_dir {params.partitions_dir}"
+    )
+    logging.info(f"Executing PostProcessing: {postproc_cmd}")
+    os.system(postproc_cmd)
+    logging.info("Done Executing ParMETIS...")
+
+
+if __name__ == "__main__":
+    """Main function to invoke the parmetis wrapper function"""
+    parser = argparse.ArgumentParser(
+        description="Run ParMETIS as part of the graph partitioning pipeline"
+    )
+    # Preprocessing step.
+    parser.add_argument(
+        "--schema_file",
+        required=True,
+        type=str,
+        help="The schema of the input graph",
+    )
+    parser.add_argument(
+        "--preproc_input_dir",
+        type=str,
+        help="The input directory for preprocess where the dataset is located",
+    )
+    parser.add_argument(
+        "--preproc_output_dir",
+        required=True,
+        type=str,
+        help="The output directory for the node weights files and auxiliary\
+              files for ParMETIS.",
+    )
+    parser.add_argument(
+        "--hostfile",
+        required=True,
+        type=str,
+        help="A text file with a list of ip addresses.",
+    )
+    parser.add_argument(
+        "--num_parts",
+        required=True,
+        type=int,
+        help="integer representing no. of partitions.",
+    )
+
+    # ParMETIS step.
+    parser.add_argument(
+        "--parmetis_install_path",
+        required=False,
+        type=str,
+        help="The directory where ParMETIS is installed",
+    )
+
+    # Postprocessing step.
+    parser.add_argument(
+        "--parmetis_output_file",
+        required=True,
+        type=str,
+        help="ParMETIS output file (global_node_id to partition_id mappings)",
+    )
+    parser.add_argument(
+        "--partitions_dir",
+        required=True,
+        type=str,
+        help="The directory where the files (with metis partition ids) grouped \
+              by node_types",
+    )
+    params = parser.parse_args()
+
+    # Configure logging.
+    logging.basicConfig(
+        level="INFO",
+        format=f"[{platform.node()} \
+        %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
+    )
+
+    check_dependencies()
+    run_parmetis_wrapper(params)

tools/distpartitioning/utils.py

@@ -0,0 +1,785 @@
+import json
+import logging
+import os
+from itertools import cycle
+
+import constants
+
+import dgl
+import numpy as np
+import psutil
+import pyarrow
+
+import torch
+from dgl.distributed.partition import _dump_part_config
+from pyarrow import csv
+
+DATA_TYPE_ID = {
+    data_type: id
+    for id, data_type in enumerate(
+        [
+            torch.float32,
+            torch.float64,
+            torch.float16,
+            torch.uint8,
+            torch.int8,
+            torch.int16,
+            torch.int32,
+            torch.int64,
+            torch.bool,
+        ]
+    )
+}
+
+REV_DATA_TYPE_ID = {id: data_type for data_type, id in DATA_TYPE_ID.items()}
+
+
+def read_ntype_partition_files(schema_map, input_dir):
+    """
+    Utility method to read the partition id mapping for each node.
+    For each node type, there will be an file, in the input directory argument
+    containing the partition id mapping for a given nodeid.
+
+    Parameters:
+    -----------
+    schema_map : dictionary
+        dictionary created by reading the input metadata json file
+    input_dir : string
+        directory in which the node-id to partition-id mappings files are
+        located for each of the node types in the input graph
+
+    Returns:
+    --------
+    numpy array :
+        array of integers representing mapped partition-ids for a given node-id.
+        The line number, in these files, are used as the type_node_id in each of
+        the files. The index into this array will be the homogenized node-id and
+        value will be the partition-id for that node-id (index). Please note that
+        the partition-ids of each node-type are stacked together vertically and
+        in this way heterogenous node-ids are converted to homogenous node-ids.
+    """
+    assert os.path.isdir(input_dir)
+
+    # iterate over the node types and extract the partition id mappings
+    part_ids = []
+    ntype_names = schema_map[constants.STR_NODE_TYPE]
+    for ntype in ntype_names:
+        df = csv.read_csv(
+            os.path.join(input_dir, "{}.txt".format(ntype)),
+            read_options=pyarrow.csv.ReadOptions(
+                autogenerate_column_names=True
+            ),
+            parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
+        )
+        ntype_partids = df["f0"].to_numpy()
+        part_ids.append(ntype_partids)
+    return np.concatenate(part_ids)
+
+
+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_etype_featnames(etype_name, schema_map):
+    """Retrieves edge feature names for a given edge_type
+
+    Parameters:
+    -----------
+    eype_name : string
+        a string specifying a edge_type name
+
+    schema : dictionary
+        metadata json object as a dictionary, which is read from the input
+        metadata file from the input dataset
+
+    Returns:
+    --------
+    list :
+        a list of feature names for a given edge_type
+    """
+    edge_data = schema_map[constants.STR_EDGE_DATA]
+    feats = edge_data.get(etype_name, {})
+    return [feat for feat in feats]
+
+
+def get_ntype_featnames(ntype_name, schema_map):
+    """
+    Retrieves node feature names for a given node_type
+
+    Parameters:
+    -----------
+    ntype_name : string
+        a string specifying a node_type name
+
+    schema : dictionary
+        metadata json object as a dictionary, which is read from the input
+        metadata file from the input dataset
+
+    Returns:
+    --------
+    list :
+        a list of feature names for a given node_type
+    """
+    node_data = schema_map[constants.STR_NODE_DATA]
+    feats = node_data.get(ntype_name, {})
+    return [feat for feat in feats]
+
+
+def get_edge_types(schema_map):
+    """Utility method to extract edge_typename -> edge_type mappings
+    as defined by the input schema
+
+    Parameters:
+    -----------
+    schema_map : dictionary
+        Input schema from which the edge_typename -> edge_typeid
+        dictionary is created.
+
+    Returns:
+    --------
+    dictionary
+        with keys as edge type names and values as ids (integers)
+    list
+        list of etype name strings
+    dictionary
+        with keys as etype ids (integers) and values as edge type names
+    """
+    etypes = schema_map[constants.STR_EDGE_TYPE]
+    etype_etypeid_map = {e: i for i, e in enumerate(etypes)}
+    etypeid_etype_map = {i: e for i, e in enumerate(etypes)}
+    return etype_etypeid_map, etypes, etypeid_etype_map
+
+
+def get_node_types(schema_map):
+    """
+    Utility method to extract node_typename -> node_type mappings
+    as defined by the input schema
+
+    Parameters:
+    -----------
+    schema_map : dictionary
+        Input schema from which the node_typename -> node_type
+        dictionary is created.
+
+    Returns:
+    --------
+    dictionary
+        with keys as node type names and values as ids (integers)
+    list
+        list of ntype name strings
+    dictionary
+        with keys as ntype ids (integers) and values as node type names
+    """
+    ntypes = schema_map[constants.STR_NODE_TYPE]
+    ntype_ntypeid_map = {e: i for i, e in enumerate(ntypes)}
+    ntypeid_ntype_map = {i: e for i, e in enumerate(ntypes)}
+    return ntype_ntypeid_map, ntypes, ntypeid_ntype_map
+
+
+def get_gid_offsets(typenames, typecounts):
+    """
+    Builds a map where the key-value pairs are typnames and respective
+    global-id offsets.
+
+    Parameters:
+    -----------
+    typenames : list of strings
+        a list of strings which can be either node typenames or edge typenames
+    typecounts : list of integers
+        a list of integers indicating the total number of nodes/edges for its
+        typeid which is the index in this list
+
+    Returns:
+    --------
+    dictionary :
+        a dictionary where keys are node_type names and values are
+        global_nid range, which is a tuple.
+
+    """
+    assert len(typenames) == len(
+        typecounts
+    ), f"No. of typenames does not match with its type counts names = {typenames}, counts = {typecounts}"
+
+    counts = []
+    for name in typenames:
+        counts.append(typecounts[name])
+    starts = np.cumsum([0] + counts[:-1])
+    ends = np.cumsum(counts)
+
+    gid_offsets = {}
+    for idx, name in enumerate(typenames):
+        gid_offsets[name] = [starts[idx], ends[idx]]
+    return gid_offsets
+
+    """
+    starts = np.cumsum([0] + type_counts[:-1])
+    ends = np.cumsum(type_counts)
+    gid_offsets = {}
+    for idx, name in enumerate(typenames):
+        gid_offsets[name] = [start[idx], ends[idx]]
+
+    return gid_offsets
+    """
+
+
+def get_gnid_range_map(node_tids):
+    """
+    Retrieves auxiliary dictionaries from the metadata json object
+
+    Parameters:
+    -----------
+    node_tids: dictionary
+        This dictionary contains the information about nodes for each node_type.
+        Typically this information contains p-entries, where each entry has a file-name,
+        starting and ending type_node_ids for the nodes in this file. Keys in this dictionary
+        are the node_type and value is a list of lists. Each individual entry in this list has
+        three items: file-name, starting type_nid and ending type_nid
+
+    Returns:
+    --------
+    dictionary :
+        a dictionary where keys are node_type names and values are global_nid range, which is a tuple.
+
+    """
+    ntypes_gid_range = {}
+    offset = 0
+    for k, v in node_tids.items():
+        ntypes_gid_range[k] = [offset + int(v[0][0]), offset + int(v[-1][1])]
+        offset += int(v[-1][1])
+
+    return ntypes_gid_range
+
+
+def write_metadata_json(
+    input_list, output_dir, graph_name, world_size, num_parts
+):
+    """
+    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
+    """
+    # Preprocess the input_list, a list of dictionaries
+    # each dictionary will contain num_parts/world_size metadata json
+    # which correspond to local partitions on the respective ranks.
+    metadata_list = []
+    for local_part_id in range(num_parts // world_size):
+        for idx in range(world_size):
+            metadata_list.append(
+                input_list[idx][
+                    "local-part-id-" + str(local_part_id * world_size + idx)
+                ]
+            )
+
+    # 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(
+                [
+                    int(metadata_list[idx]["edge_map"][k][0][0]),
+                    int(metadata_list[idx]["edge_map"][k][0][1]),
+                ]
+            )
+    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(
+                [
+                    int(metadata_list[idx]["node_map"][k][0][0]),
+                    int(metadata_list[idx]["node_map"][k][0][1]),
+                ]
+            )
+    graph_metadata["node_map"] = node_map
+
+    graph_metadata["ntypes"] = metadata_list[0]["ntypes"]
+    graph_metadata["num_edges"] = int(
+        sum([metadata_list[i]["num_edges"] for i in range(len(metadata_list))])
+    )
+    graph_metadata["num_nodes"] = int(
+        sum([metadata_list[i]["num_nodes"] for i in range(len(metadata_list))])
+    )
+    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)
+        ]
+
+    _dump_part_config(f"{output_dir}/metadata.json", graph_metadata)
+
+
+def augment_edge_data(
+    edge_data, lookup_service, edge_tids, rank, world_size, num_parts
+):
+    """
+    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
+    lookup_service : instance of class DistLookupService
+       Distributed lookup service used to map global-nids to respective partition-ids and▒
+       shuffle-global-nids
+    edge_tids: dictionary
+        dictionary where keys are canonical edge types and values are list of tuples
+        which indicate the range of edges assigned to each of the partitions
+    rank : integer
+        rank of the current process
+    world_size : integer
+        total no. of process participating in the communication primitives
+    num_parts : integer
+        total no. of partitions requested for the input graph
+
+    Returns:
+    --------
+    dictionary :
+        dictionary with keys as column names and values as numpy arrays and this information is
+        loaded from input dataset files. In addition to this we include additional columns which
+        aid this pipelines computation, like constants.OWNER_PROCESS
+    """
+    # add global_nids to the node_data
+    etype_offset = {}
+    offset = 0
+    for etype_name, tid_range in edge_tids.items():
+        etype_offset[etype_name] = offset + int(tid_range[0][0])
+        offset += int(tid_range[-1][1])
+
+    global_eids = []
+    for etype_name, tid_range in edge_tids.items():
+        for idx in range(num_parts):
+            if map_partid_rank(idx, world_size) == rank:
+                if len(tid_range) > idx:
+                    global_eid_start = etype_offset[etype_name]
+                    begin = global_eid_start + int(tid_range[idx][0])
+                    end = global_eid_start + int(tid_range[idx][1])
+                    global_eids.append(np.arange(begin, end, dtype=np.int64))
+
+    global_eids = (
+        np.concatenate(global_eids)
+        if len(global_eids) > 0
+        else np.array([], dtype=np.int64)
+    )
+    assert global_eids.shape[0] == edge_data[constants.ETYPE_ID].shape[0]
+    edge_data[constants.GLOBAL_EID] = global_eids
+    return edge_data
+
+
+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_df = csv.read_csv(
+        edge_file,
+        read_options=pyarrow.csv.ReadOptions(autogenerate_column_names=True),
+        parse_options=pyarrow.csv.ParseOptions(delimiter=" "),
+    )
+    edge_data_dict = {}
+    edge_data_dict[constants.GLOBAL_SRC_ID] = edge_data_df["f0"].to_numpy()
+    edge_data_dict[constants.GLOBAL_DST_ID] = edge_data_df["f1"].to_numpy()
+    edge_data_dict[constants.GLOBAL_TYPE_EID] = edge_data_df["f2"].to_numpy()
+    edge_data_dict[constants.ETYPE_ID] = edge_data_df["f3"].to_numpy()
+    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_graghbolt(graph_file, graph_obj):
+    """
+    Utility function to serialize FusedCSCSamplingGraph
+
+    Parameters:
+    -----------
+    graph_obj : FusedCSCSamplingGraph
+        FusedCSCSamplingGraph, as created in convert_partition.py, which is to be serialized
+    graph_file : string
+        File name in which graph object is serialized
+    """
+    torch.save(graph_obj, graph_file)
+
+
+def write_graph_dgl(graph_file, graph_obj, formats, sort_etypes):
+    """
+    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
+    formats : str or list[str]
+        Save graph in specified formats.
+    sort_etypes : bool
+        Whether to sort etypes in csc/csr.
+    """
+    dgl.distributed.partition.process_partitions(
+        graph_obj, formats, sort_etypes
+    )
+    dgl.save_graphs(graph_file, [graph_obj], formats=formats)
+
+
+def _write_graph(
+    part_dir, graph_obj, formats=None, sort_etypes=None, use_graphbolt=False
+):
+    if use_graphbolt:
+        write_graph_graghbolt(
+            os.path.join(part_dir, "fused_csc_sampling_graph.pt"), graph_obj
+        )
+    else:
+        write_graph_dgl(
+            os.path.join(part_dir, "graph.dgl"), graph_obj, formats, sort_etypes
+        )
+
+
+def write_dgl_objects(
+    graph_obj,
+    node_features,
+    edge_features,
+    output_dir,
+    part_id,
+    orig_nids,
+    orig_eids,
+    formats,
+    sort_etypes,
+    use_graphbolt,
+):
+    """
+    Wrapper function to write graph, node/edge feature, original node/edge IDs.
+
+    Parameters:
+    -----------
+    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
+    output_dir : string
+        location where the output files will be located
+    part_id : int
+        integer indicating the partition-id
+    orig_nids : dict
+        original node IDs
+    orig_eids : dict
+        original edge IDs
+    formats : str or list[str]
+        Save graph in formats.
+    sort_etypes : bool
+        Whether to sort etypes in csc/csr.
+    use_graphbolt : bool
+        Whether to use graphbolt or not.
+    """
+    part_dir = output_dir + "/part" + str(part_id)
+    os.makedirs(part_dir, exist_ok=True)
+    _write_graph(
+        part_dir,
+        graph_obj,
+        formats=formats,
+        sort_etypes=sort_etypes,
+        use_graphbolt=use_graphbolt,
+    )
+    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")
+        )
+
+    if orig_nids is not None:
+        orig_nids_file = os.path.join(part_dir, "orig_nids.dgl")
+        dgl.data.utils.save_tensors(orig_nids_file, orig_nids)
+    if orig_eids is not None:
+        orig_eids_file = os.path.join(part_dir, "orig_eids.dgl")
+        dgl.data.utils.save_tensors(orig_eids_file, orig_eids)
+
+
+def get_idranges(names, counts, num_chunks=None):
+    """
+    counts will be a list of numbers of a dictionary.
+    Length is less than or equal to the num_parts variable.
+
+    Parameters:
+    -----------
+    names : list of strings
+        which are either node-types or edge-types
+    counts : list of integers
+        which are total no. of nodes or edges for a give node
+        or edge type
+    num_chunks : int, optional
+        specifying the no. of chunks
+
+    Returns:
+    --------
+    dictionary
+        dictionary where the keys are node-/edge-type names and values are
+        list of tuples where each tuple indicates the range of values for
+        corresponding type-ids.
+    dictionary
+        dictionary where the keys are node-/edge-type names and value is a tuple.
+        This tuple indicates the global-ids for the associated node-/edge-type.
+    """
+    gnid_start = 0
+    gnid_end = gnid_start
+    tid_dict = {}
+    gid_dict = {}
+
+    for idx, typename in enumerate(names):
+        gnid_end += counts[typename]
+        tid_dict[typename] = [[0, counts[typename]]]
+        gid_dict[typename] = np.array([gnid_start, gnid_end]).reshape([1, 2])
+        gnid_start = gnid_end
+
+    return tid_dict, gid_dict
+
+
+def get_ntype_counts_map(ntypes, ntype_counts):
+    """
+    Return a dictionary with key, value pairs as node type names and no. of
+    nodes of a particular type in the input graph.
+
+    Parameters:
+    -----------
+    ntypes : list of strings
+        where each string is a node-type name
+    ntype_counts : list of integers
+        where each integer is the total no. of nodes for that, idx, node type
+
+    Returns:
+    --------
+    dictinary :
+        a dictionary where node-type names are keys and values are total no.
+        of nodes for a given node-type name (which is also the key)
+    """
+    return dict(zip(ntypes, ntype_counts))
+
+
+def memory_snapshot(tag, rank):
+    """
+    Utility function to take a snapshot of the usage of system resources
+    at a given point of time.
+
+    Parameters:
+    -----------
+    tag : string
+        string provided by the user for bookmarking purposes
+    rank : integer
+        process id of the participating process
+    """
+    GB = 1024 * 1024 * 1024
+    MB = 1024 * 1024
+    KB = 1024
+
+    peak = dgl.partition.get_peak_mem() * KB
+    mem = psutil.virtual_memory()
+    avail = mem.available / MB
+    used = mem.used / MB
+    total = mem.total / MB
+
+    mem_string = f"{total:.0f} (MB) total, {peak:.0f} (MB) peak, {used:.0f} (MB) used, {avail:.0f} (MB) avail"
+    logging.debug(f"[Rank: {rank} MEMORY_SNAPSHOT] {mem_string} - {tag}")
+
+
+def map_partid_rank(partid, world_size):
+    """Auxiliary function to map a given partition id to one of the rank in the
+    MPI_WORLD processes. The range of partition ids is assumed to equal or a
+    multiple of the total size of MPI_WORLD. In this implementation, we use
+    a cyclical mapping procedure to convert partition ids to ranks.
+
+    Parameters:
+    -----------
+    partid : int
+        partition id, as read from node id to partition id mappings.
+
+    Returns:
+    --------
+    int :
+        rank of the process, which will be responsible for the given partition
+        id.
+    """
+    return partid % world_size
+
+
+def generate_read_list(num_files, world_size):
+    """
+    Generate the file IDs to read for each rank
+    using sequential assignment.
+
+
+    Parameters:
+    -----------
+    num_files : int
+        Total number of files.
+    world_size : int
+        World size of group.
+
+    Returns:
+    --------
+    read_list : np.array
+        Array of target file IDs to read. Each worker is expected
+        to read the list of file indexes in its rank's index in the list.
+        e.g. rank 0 reads the file indexed in read_list[0], rank 1 the
+        ones in read_list[1] etc.
+
+
+    Examples
+    --------
+    >>> tools.distpartitionning.utils.generate_read_list(10, 4)
+    [array([0, 1, 2]), array([3, 4, 5]), array([6, 7]), array([8, 9])]
+    """
+    return np.array_split(np.arange(num_files), world_size)
+
+
+def generate_roundrobin_read_list(num_files, world_size):
+    """
+    Generate the file IDs to read for each rank
+    using round robin assignment.
+
+    Parameters:
+    -----------
+    num_files : int
+        Total number of files.
+    world_size : int
+        World size of group.
+
+    Returns:
+    --------
+    read_list : np.array
+        Array of target file IDs to read. Each worker is expected
+        to read the list of file indexes in its rank's index in the list.
+        e.g. rank 0 reads the indexed in read_list[0], rank 1 the
+        ones in read_list[1] etc.
+
+    Examples
+    --------
+    >>> tools.distpartitionning.utils.generate_roundrobin_read_list(10, 4)
+    [[0, 4, 8], [1, 5, 9], [2, 6], [3, 7]]
+    """
+    assignment_lists = [[] for _ in range(world_size)]
+    for rank, part_idx in zip(cycle(range(world_size)), range(num_files)):
+        assignment_lists[rank].append(part_idx)
+
+    return assignment_lists

tools/files.py

@@ -0,0 +1,19 @@
+import logging
+import os
+from contextlib import contextmanager
+
+from numpy.lib.format import open_memmap
+
+
+@contextmanager
+def setdir(path):
+    try:
+        os.makedirs(path, exist_ok=True)
+        cwd = os.getcwd()
+        logging.info("Changing directory to %s" % path)
+        logging.info("Previously: %s" % cwd)
+        os.chdir(path)
+        yield
+    finally:
+        logging.info("Restoring directory to %s" % cwd)
+        os.chdir(cwd)

tools/launch.py

@@ -0,0 +1,775 @@
+"""Launching tool for DGL distributed training"""
+import argparse
+import json
+import logging
+import multiprocessing
+import os
+import queue
+import re
+import signal
+import subprocess
+import sys
+import time
+from functools import partial
+from threading import Thread
+from typing import Optional
+
+
+def cleanup_proc(get_all_remote_pids, conn):
+    """This process tries to clean up the remote training tasks."""
+    print("cleanup process runs")
+    # This process should not handle SIGINT.
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+    data = conn.recv()
+    # If the launch process exits normally, this process doesn't need to do anything.
+    if data == "exit":
+        sys.exit(0)
+    else:
+        remote_pids = get_all_remote_pids()
+        # Otherwise, we need to ssh to each machine and kill the training jobs.
+        for (ip, port), pids in remote_pids.items():
+            kill_process(ip, port, pids)
+    print("cleanup process exits")
+
+
+def kill_process(ip, port, pids):
+    """ssh to a remote machine and kill the specified processes."""
+    curr_pid = os.getpid()
+    killed_pids = []
+    # If we kill child processes first, the parent process may create more again. This happens
+    # to Python's process pool. After sorting, we always kill parent processes first.
+    pids.sort()
+    for pid in pids:
+        assert curr_pid != pid
+        print("kill process {} on {}:{}".format(pid, ip, port), flush=True)
+        kill_cmd = (
+            "ssh -o StrictHostKeyChecking=no -p "
+            + str(port)
+            + " "
+            + ip
+            + " 'kill {}'".format(pid)
+        )
+        subprocess.run(kill_cmd, shell=True)
+        killed_pids.append(pid)
+    # It's possible that some of the processes are not killed. Let's try again.
+    for i in range(3):
+        killed_pids = get_killed_pids(ip, port, killed_pids)
+        if len(killed_pids) == 0:
+            break
+        else:
+            killed_pids.sort()
+            for pid in killed_pids:
+                print(
+                    "kill process {} on {}:{}".format(pid, ip, port), flush=True
+                )
+                kill_cmd = (
+                    "ssh -o StrictHostKeyChecking=no -p "
+                    + str(port)
+                    + " "
+                    + ip
+                    + " 'kill -9 {}'".format(pid)
+                )
+                subprocess.run(kill_cmd, shell=True)
+
+
+def get_killed_pids(ip, port, killed_pids):
+    """Get the process IDs that we want to kill but are still alive."""
+    killed_pids = [str(pid) for pid in killed_pids]
+    killed_pids = ",".join(killed_pids)
+    ps_cmd = (
+        "ssh -o StrictHostKeyChecking=no -p "
+        + str(port)
+        + " "
+        + ip
+        + " 'ps -p {} -h'".format(killed_pids)
+    )
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    pids = []
+    for p in res.stdout.decode("utf-8").split("\n"):
+        l = p.split()
+        if len(l) > 0:
+            pids.append(int(l[0]))
+    return pids
+
+
+def execute_remote(
+    cmd: str,
+    state_q: queue.Queue,
+    ip: str,
+    port: int,
+    username: Optional[str] = "",
+) -> Thread:
+    """Execute command line on remote machine via ssh.
+
+    Args:
+        cmd: User-defined command (udf) to execute on the remote host.
+        state_q: A queue collecting Thread exit states.
+        ip: The ip-address of the host to run the command on.
+        port: Port number that the host is listening on.
+        thread_list:
+        username: Optional. If given, this will specify a username to use when issuing commands over SSH.
+            Useful when your infra requires you to explicitly specify a username to avoid permission issues.
+
+    Returns:
+        thread: The Thread whose run() is to run the `cmd` on the remote host. Returns when the cmd completes
+            on the remote host.
+    """
+    ip_prefix = ""
+    if username:
+        ip_prefix += "{username}@".format(username=username)
+
+    # Construct ssh command that executes `cmd` on the remote host
+    ssh_cmd = "ssh -o StrictHostKeyChecking=no -p {port} {ip_prefix}{ip} '{cmd}'".format(
+        port=str(port),
+        ip_prefix=ip_prefix,
+        ip=ip,
+        cmd=cmd,
+    )
+
+    # thread func to run the job
+    def run(ssh_cmd, state_q):
+        try:
+            subprocess.check_call(ssh_cmd, shell=True)
+            state_q.put(0)
+        except subprocess.CalledProcessError as err:
+            print(f"Called process error {err}")
+            state_q.put(err.returncode)
+        except Exception:
+            state_q.put(-1)
+
+    thread = Thread(
+        target=run,
+        args=(
+            ssh_cmd,
+            state_q,
+        ),
+    )
+    thread.setDaemon(True)
+    thread.start()
+    # sleep for a while in case of ssh is rejected by peer due to busy connection
+    time.sleep(0.2)
+    return thread
+
+
+def get_remote_pids(ip, port, cmd_regex):
+    """Get the process IDs that run the command in the remote machine."""
+    pids = []
+    curr_pid = os.getpid()
+    # Here we want to get the python processes. We may get some ssh processes, so we should filter them out.
+    ps_cmd = (
+        "ssh -o StrictHostKeyChecking=no -p "
+        + str(port)
+        + " "
+        + ip
+        + " 'ps -aux | grep python | grep -v StrictHostKeyChecking'"
+    )
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    for p in res.stdout.decode("utf-8").split("\n"):
+        l = p.split()
+        if len(l) < 2:
+            continue
+        # We only get the processes that run the specified command.
+        res = re.search(cmd_regex, p)
+        if res is not None and int(l[1]) != curr_pid:
+            pids.append(l[1])
+
+    pid_str = ",".join([str(pid) for pid in pids])
+    ps_cmd = (
+        "ssh -o StrictHostKeyChecking=no -p "
+        + str(port)
+        + " "
+        + ip
+        + " 'pgrep -P {}'".format(pid_str)
+    )
+    res = subprocess.run(ps_cmd, shell=True, stdout=subprocess.PIPE)
+    pids1 = res.stdout.decode("utf-8").split("\n")
+    all_pids = []
+    for pid in set(pids + pids1):
+        if pid == "" or int(pid) == curr_pid:
+            continue
+        all_pids.append(int(pid))
+    all_pids.sort()
+    return all_pids
+
+
+def get_all_remote_pids(hosts, ssh_port, udf_command):
+    """Get all remote processes."""
+    remote_pids = {}
+    for node_id, host in enumerate(hosts):
+        ip, _ = host
+        # When creating training processes in remote machines, we may insert some arguments
+        # in the commands. We need to use regular expressions to match the modified command.
+        cmds = udf_command.split()
+        new_udf_command = " .*".join(cmds)
+        pids = get_remote_pids(ip, ssh_port, new_udf_command)
+        remote_pids[(ip, ssh_port)] = pids
+    return remote_pids
+
+
+def construct_torch_dist_launcher_cmd(
+    num_trainers: int,
+    num_nodes: int,
+    node_rank: int,
+    master_addr: str,
+    master_port: int,
+) -> str:
+    """Constructs the torch distributed launcher command.
+    Helper function.
+
+    Args:
+        num_trainers:
+        num_nodes:
+        node_rank:
+        master_addr:
+        master_port:
+
+    Returns:
+        cmd_str.
+    """
+    torch_cmd_template = (
+        "-m torch.distributed.run "
+        "--nproc_per_node={nproc_per_node} "
+        "--nnodes={nnodes} "
+        "--node_rank={node_rank} "
+        "--master_addr={master_addr} "
+        "--master_port={master_port}"
+    )
+    return torch_cmd_template.format(
+        nproc_per_node=num_trainers,
+        nnodes=num_nodes,
+        node_rank=node_rank,
+        master_addr=master_addr,
+        master_port=master_port,
+    )
+
+
+def wrap_udf_in_torch_dist_launcher(
+    udf_command: str,
+    num_trainers: int,
+    num_nodes: int,
+    node_rank: int,
+    master_addr: str,
+    master_port: int,
+) -> str:
+    """Wraps the user-defined function (udf_command) with the torch.distributed.run module.
+
+     Example: if udf_command is "python3 run/some/trainer.py arg1 arg2", then new_df_command becomes:
+         "python3 -m torch.distributed.run <TORCH DIST ARGS> run/some/trainer.py arg1 arg2
+
+    udf_command is assumed to consist of pre-commands (optional) followed by the python launcher script (required):
+    Examples:
+        # simple
+        python3.7 path/to/some/trainer.py arg1 arg2
+
+        # multi-commands
+        (cd some/dir && python3.7 path/to/some/trainer.py arg1 arg2)
+
+    IMPORTANT: If udf_command consists of multiple python commands, then this will result in undefined behavior.
+
+    Args:
+        udf_command:
+        num_trainers:
+        num_nodes:
+        node_rank:
+        master_addr:
+        master_port:
+
+    Returns:
+
+    """
+    torch_dist_cmd = construct_torch_dist_launcher_cmd(
+        num_trainers=num_trainers,
+        num_nodes=num_nodes,
+        node_rank=node_rank,
+        master_addr=master_addr,
+        master_port=master_port,
+    )
+    # Auto-detect the python binary that kicks off the distributed trainer code.
+    # Note: This allowlist order matters, this will match with the FIRST matching entry. Thus, please add names to this
+    #       from most-specific to least-specific order eg:
+    #           (python3.7, python3.8) -> (python3)
+    # The allowed python versions are from this: https://www.dgl.ai/pages/start.html
+    python_bin_allowlist = (
+        "python3.6",
+        "python3.7",
+        "python3.8",
+        "python3.9",
+        "python3",
+        # for backwards compatibility, accept python2 but technically DGL is a py3 library, so this is not recommended
+        "python2.7",
+        "python2",
+    )
+    # If none of the candidate python bins match, then we go with the default `python`
+    python_bin = "python"
+    for candidate_python_bin in python_bin_allowlist:
+        if candidate_python_bin in udf_command:
+            python_bin = candidate_python_bin
+            break
+
+    # transforms the udf_command from:
+    #     python path/to/dist_trainer.py arg0 arg1
+    # to:
+    #     python -m torch.distributed.run [DIST TORCH ARGS] path/to/dist_trainer.py arg0 arg1
+    # Note: if there are multiple python commands in `udf_command`, this may do the Wrong Thing, eg launch each
+    #       python command within the torch distributed launcher.
+    new_udf_command = udf_command.replace(
+        python_bin, f"{python_bin} {torch_dist_cmd}"
+    )
+
+    return new_udf_command
+
+
+def construct_dgl_server_env_vars(
+    num_samplers: int,
+    num_server_threads: int,
+    tot_num_clients: int,
+    part_config: str,
+    ip_config: str,
+    num_servers: int,
+    graph_format: str,
+    pythonpath: Optional[str] = "",
+) -> str:
+    """Constructs the DGL server-specific env vars string that are required for DGL code to behave in the correct
+    server role.
+    Convenience function.
+
+    Args:
+        num_samplers:
+        num_server_threads:
+        tot_num_clients:
+        part_config: Partition config.
+            Relative path to workspace.
+        ip_config: IP config file containing IP addresses of cluster hosts.
+            Relative path to workspace.
+        num_servers:
+        graph_format:
+        pythonpath: Optional. If given, this will pass this as PYTHONPATH.
+
+    Returns:
+        server_env_vars: The server-specific env-vars in a string format, friendly for CLI execution.
+
+    """
+    server_env_vars_template = (
+        "DGL_ROLE={DGL_ROLE} "
+        "DGL_NUM_SAMPLER={DGL_NUM_SAMPLER} "
+        "OMP_NUM_THREADS={OMP_NUM_THREADS} "
+        "DGL_NUM_CLIENT={DGL_NUM_CLIENT} "
+        "DGL_CONF_PATH={DGL_CONF_PATH} "
+        "DGL_IP_CONFIG={DGL_IP_CONFIG} "
+        "DGL_NUM_SERVER={DGL_NUM_SERVER} "
+        "DGL_GRAPH_FORMAT={DGL_GRAPH_FORMAT} "
+        "{suffix_optional_envvars}"
+    )
+    suffix_optional_envvars = ""
+    if pythonpath:
+        suffix_optional_envvars += f"PYTHONPATH={pythonpath} "
+    return server_env_vars_template.format(
+        DGL_ROLE="server",
+        DGL_NUM_SAMPLER=num_samplers,
+        OMP_NUM_THREADS=num_server_threads,
+        DGL_NUM_CLIENT=tot_num_clients,
+        DGL_CONF_PATH=part_config,
+        DGL_IP_CONFIG=ip_config,
+        DGL_NUM_SERVER=num_servers,
+        DGL_GRAPH_FORMAT=graph_format,
+        suffix_optional_envvars=suffix_optional_envvars,
+    )
+
+
+def construct_dgl_client_env_vars(
+    num_samplers: int,
+    tot_num_clients: int,
+    part_config: str,
+    ip_config: str,
+    num_servers: int,
+    graph_format: str,
+    num_omp_threads: int,
+    group_id: int,
+    pythonpath: Optional[str] = "",
+) -> str:
+    """Constructs the DGL client-specific env vars string that are required for DGL code to behave in the correct
+    client role.
+    Convenience function.
+
+    Args:
+        num_samplers:
+        tot_num_clients:
+        part_config: Partition config.
+            Relative path to workspace.
+        ip_config: IP config file containing IP addresses of cluster hosts.
+            Relative path to workspace.
+        num_servers:
+        graph_format:
+        num_omp_threads:
+        group_id:
+            Used in client processes to indicate which group it belongs to.
+        pythonpath: Optional. If given, this will pass this as PYTHONPATH.
+
+    Returns:
+        client_env_vars: The client-specific env-vars in a string format, friendly for CLI execution.
+
+    """
+    client_env_vars_template = (
+        "DGL_DIST_MODE={DGL_DIST_MODE} "
+        "DGL_ROLE={DGL_ROLE} "
+        "DGL_NUM_SAMPLER={DGL_NUM_SAMPLER} "
+        "DGL_NUM_CLIENT={DGL_NUM_CLIENT} "
+        "DGL_CONF_PATH={DGL_CONF_PATH} "
+        "DGL_IP_CONFIG={DGL_IP_CONFIG} "
+        "DGL_NUM_SERVER={DGL_NUM_SERVER} "
+        "DGL_GRAPH_FORMAT={DGL_GRAPH_FORMAT} "
+        "OMP_NUM_THREADS={OMP_NUM_THREADS} "
+        "DGL_GROUP_ID={DGL_GROUP_ID} "
+        "{suffix_optional_envvars}"
+    )
+    # append optional additional env-vars
+    suffix_optional_envvars = ""
+    if pythonpath:
+        suffix_optional_envvars += f"PYTHONPATH={pythonpath} "
+    return client_env_vars_template.format(
+        DGL_DIST_MODE="distributed",
+        DGL_ROLE="client",
+        DGL_NUM_SAMPLER=num_samplers,
+        DGL_NUM_CLIENT=tot_num_clients,
+        DGL_CONF_PATH=part_config,
+        DGL_IP_CONFIG=ip_config,
+        DGL_NUM_SERVER=num_servers,
+        DGL_GRAPH_FORMAT=graph_format,
+        OMP_NUM_THREADS=num_omp_threads,
+        DGL_GROUP_ID=group_id,
+        suffix_optional_envvars=suffix_optional_envvars,
+    )
+
+
+def wrap_cmd_with_local_envvars(cmd: str, env_vars: str) -> str:
+    """Wraps a CLI command with desired env vars with the following properties:
+    (1) env vars persist for the entire `cmd`, even if it consists of multiple "chained" commands like:
+        cmd = "ls && pwd && python run/something.py"
+    (2) env vars don't pollute the environment after `cmd` completes.
+
+    Example:
+        >>> cmd = "ls && pwd"
+        >>> env_vars = "VAR1=value1 VAR2=value2"
+        >>> wrap_cmd_with_local_envvars(cmd, env_vars)
+        "(export VAR1=value1 VAR2=value2; ls && pwd)"
+
+    Args:
+        cmd:
+        env_vars: A string containing env vars, eg "VAR1=val1 VAR2=val2"
+
+    Returns:
+        cmd_with_env_vars:
+
+    """
+    # use `export` to persist env vars for entire cmd block. required if udf_command is a chain of commands
+    # also: wrap in parens to not pollute env:
+    #     https://stackoverflow.com/a/45993803
+    return f"(export {env_vars}; {cmd})"
+
+
+def wrap_cmd_with_extra_envvars(cmd: str, env_vars: list) -> str:
+    """Wraps a CLI command with extra env vars
+
+    Example:
+        >>> cmd = "ls && pwd"
+        >>> env_vars = ["VAR1=value1", "VAR2=value2"]
+        >>> wrap_cmd_with_extra_envvars(cmd, env_vars)
+        "(export VAR1=value1 VAR2=value2; ls && pwd)"
+
+    Args:
+        cmd:
+        env_vars: A list of strings containing env vars, e.g., ["VAR1=value1", "VAR2=value2"]
+
+    Returns:
+        cmd_with_env_vars:
+    """
+    env_vars = " ".join(env_vars)
+    return wrap_cmd_with_local_envvars(cmd, env_vars)
+
+
+def get_available_port(ip):
+    """Get available port with specified ip."""
+    import socket
+
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    for port in range(1234, 65535):
+        try:
+            sock.connect((ip, port))
+        except:
+            return port
+    raise RuntimeError("Failed to get available port for ip~{}".format(ip))
+
+
+def submit_jobs(args, udf_command, dry_run=False):
+    """Submit distributed jobs (server and client processes) via ssh"""
+    if dry_run:
+        print(
+            "Currently it's in dry run mode which means no jobs will be launched."
+        )
+    servers_cmd = []
+    clients_cmd = []
+    hosts = []
+    thread_list = []
+    server_count_per_machine = 0
+
+    # Get the IP addresses of the cluster.
+    ip_config = os.path.join(args.workspace, args.ip_config)
+    with open(ip_config) as f:
+        for line in f:
+            result = line.strip().split()
+            if len(result) == 2:
+                ip = result[0]
+                port = int(result[1])
+                hosts.append((ip, port))
+            elif len(result) == 1:
+                ip = result[0]
+                port = get_available_port(ip)
+                hosts.append((ip, port))
+            else:
+                raise RuntimeError("Format error of ip_config.")
+            server_count_per_machine = args.num_servers
+    # Get partition info of the graph data
+    part_config = os.path.join(args.workspace, args.part_config)
+    with open(part_config) as conf_f:
+        part_metadata = json.load(conf_f)
+    assert "num_parts" in part_metadata, "num_parts does not exist."
+    # The number of partitions must match the number of machines in the cluster.
+    assert part_metadata["num_parts"] == len(
+        hosts
+    ), "The number of graph partitions has to match the number of machines in the cluster."
+
+    state_q = queue.Queue()
+    tot_num_clients = args.num_trainers * (1 + args.num_samplers) * len(hosts)
+    # launch server tasks
+    server_env_vars = construct_dgl_server_env_vars(
+        num_samplers=args.num_samplers,
+        num_server_threads=args.num_server_threads,
+        tot_num_clients=tot_num_clients,
+        part_config=args.part_config,
+        ip_config=args.ip_config,
+        num_servers=args.num_servers,
+        graph_format=args.graph_format,
+        pythonpath=os.environ.get("PYTHONPATH", ""),
+    )
+    for i in range(len(hosts) * server_count_per_machine):
+        ip, _ = hosts[int(i / server_count_per_machine)]
+        server_env_vars_cur = f"{server_env_vars} DGL_SERVER_ID={i}"
+        cmd = wrap_cmd_with_local_envvars(udf_command, server_env_vars_cur)
+        cmd = (
+            wrap_cmd_with_extra_envvars(cmd, args.extra_envs)
+            if len(args.extra_envs) > 0
+            else cmd
+        )
+        cmd = "cd " + str(args.workspace) + "; " + cmd
+        servers_cmd.append(cmd)
+        if not dry_run:
+            thread_list.append(
+                execute_remote(
+                    cmd,
+                    state_q,
+                    ip,
+                    args.ssh_port,
+                    username=args.ssh_username,
+                )
+            )
+
+    # launch client tasks
+    client_env_vars = construct_dgl_client_env_vars(
+        num_samplers=args.num_samplers,
+        tot_num_clients=tot_num_clients,
+        part_config=args.part_config,
+        ip_config=args.ip_config,
+        num_servers=args.num_servers,
+        graph_format=args.graph_format,
+        num_omp_threads=os.environ.get(
+            "OMP_NUM_THREADS", str(args.num_omp_threads)
+        ),
+        group_id=0,
+        pythonpath=os.environ.get("PYTHONPATH", ""),
+    )
+
+    master_addr = hosts[0][0]
+    master_port = get_available_port(master_addr)
+    for node_id, host in enumerate(hosts):
+        ip, _ = host
+        # Transform udf_command to follow torch's dist launcher format: `PYTHON_BIN -m torch.distributed.run ... UDF`
+        torch_dist_udf_command = wrap_udf_in_torch_dist_launcher(
+            udf_command=udf_command,
+            num_trainers=args.num_trainers,
+            num_nodes=len(hosts),
+            node_rank=node_id,
+            master_addr=master_addr,
+            master_port=master_port,
+        )
+        cmd = wrap_cmd_with_local_envvars(
+            torch_dist_udf_command, client_env_vars
+        )
+        cmd = (
+            wrap_cmd_with_extra_envvars(cmd, args.extra_envs)
+            if len(args.extra_envs) > 0
+            else cmd
+        )
+        cmd = "cd " + str(args.workspace) + "; " + cmd
+        clients_cmd.append(cmd)
+        if not dry_run:
+            thread_list.append(
+                execute_remote(
+                    cmd, state_q, ip, args.ssh_port, username=args.ssh_username
+                )
+            )
+
+    # return commands of clients/servers directly if in dry run mode
+    if dry_run:
+        return clients_cmd, servers_cmd
+
+    # Start a cleanup process dedicated for cleaning up remote training jobs.
+    conn1, conn2 = multiprocessing.Pipe()
+    func = partial(get_all_remote_pids, hosts, args.ssh_port, udf_command)
+    process = multiprocessing.Process(target=cleanup_proc, args=(func, conn1))
+    process.start()
+
+    def signal_handler(signal, frame):
+        logging.info("Stop launcher")
+        # We need to tell the cleanup process to kill remote training jobs.
+        conn2.send("cleanup")
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, signal_handler)
+
+    err = 0
+    for thread in thread_list:
+        thread.join()
+        err_code = state_q.get()
+        if err_code != 0:
+            # Record err_code
+            # We record one of the error if there are multiple
+            err = err_code
+
+    # The training processes complete. We should tell the cleanup process to exit.
+    conn2.send("exit")
+    process.join()
+    if err != 0:
+        print("Task failed")
+        sys.exit(-1)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Launch a distributed job")
+    parser.add_argument("--ssh_port", type=int, default=22, help="SSH Port.")
+    parser.add_argument(
+        "--ssh_username",
+        default="",
+        help="Optional. When issuing commands (via ssh) to cluster, use the provided username in the ssh cmd. "
+        "Example: If you provide --ssh_username=bob, then the ssh command will be like: 'ssh bob@1.2.3.4 CMD' "
+        "instead of 'ssh 1.2.3.4 CMD'",
+    )
+    parser.add_argument(
+        "--workspace",
+        type=str,
+        help="Path of user directory of distributed tasks. \
+                        This is used to specify a destination location where \
+                        the contents of current directory will be rsyncd",
+    )
+    parser.add_argument(
+        "--num_trainers",
+        type=int,
+        help="The number of trainer processes per machine",
+    )
+    parser.add_argument(
+        "--num_omp_threads",
+        type=int,
+        help="The number of OMP threads per trainer",
+    )
+    parser.add_argument(
+        "--num_samplers",
+        type=int,
+        default=0,
+        help="The number of sampler processes per trainer process",
+    )
+    parser.add_argument(
+        "--num_servers",
+        type=int,
+        help="The number of server processes per machine",
+    )
+    parser.add_argument(
+        "--part_config",
+        type=str,
+        help="The file (in workspace) of the partition config",
+    )
+    parser.add_argument(
+        "--ip_config",
+        type=str,
+        help="The file (in workspace) of IP configuration for server processes",
+    )
+    parser.add_argument(
+        "--num_server_threads",
+        type=int,
+        default=1,
+        help="The number of OMP threads in the server process. \
+                        It should be small if server processes and trainer processes run on \
+                        the same machine. By default, it is 1.",
+    )
+    parser.add_argument(
+        "--graph_format",
+        type=str,
+        default="csc",
+        help='The format of the graph structure of each partition. \
+                        The allowed formats are csr, csc and coo. A user can specify multiple \
+                        formats, separated by ",". For example, the graph format is "csr,csc".',
+    )
+    parser.add_argument(
+        "--extra_envs",
+        nargs="+",
+        type=str,
+        default=[],
+        help="Extra environment parameters need to be set. For example, \
+                        you can set the LD_LIBRARY_PATH and NCCL_DEBUG by adding: \
+                        --extra_envs LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH NCCL_DEBUG=INFO ",
+    )
+    args, udf_command = parser.parse_known_args()
+    assert len(udf_command) == 1, "Please provide user command line."
+    assert (
+        args.num_trainers is not None and args.num_trainers > 0
+    ), "--num_trainers must be a positive number."
+    assert (
+        args.num_samplers is not None and args.num_samplers >= 0
+    ), "--num_samplers must be a non-negative number."
+    assert (
+        args.num_servers is not None and args.num_servers > 0
+    ), "--num_servers must be a positive number."
+    assert (
+        args.num_server_threads > 0
+    ), "--num_server_threads must be a positive number."
+    assert (
+        args.workspace is not None
+    ), "A user has to specify a workspace with --workspace."
+    assert (
+        args.part_config is not None
+    ), "A user has to specify a partition configuration file with --part_config."
+    assert (
+        args.ip_config is not None
+    ), "A user has to specify an IP configuration file with --ip_config."
+    if args.num_omp_threads is None:
+        # Here we assume all machines have the same number of CPU cores as the machine
+        # where the launch script runs.
+        args.num_omp_threads = max(
+            multiprocessing.cpu_count() // 2 // args.num_trainers, 1
+        )
+        print(
+            "The number of OMP threads per trainer is set to",
+            args.num_omp_threads,
+        )
+
+    udf_command = str(udf_command[0])
+    if "python" not in udf_command:
+        raise RuntimeError(
+            "DGL launching script can only support Python executable file."
+        )
+    submit_jobs(args, udf_command)
+
+
+if __name__ == "__main__":
+    fmt = "%(asctime)s %(levelname)s %(message)s"
+    logging.basicConfig(format=fmt, level=logging.INFO)
+    main()

tools/partition_algo/base.py

@@ -0,0 +1,82 @@
+import json
+from typing import Optional
+
+import pydantic as dt
+from dgl import DGLError
+
+
+class PartitionMeta(dt.BaseModel):
+    """Metadata that describes the partition assignment results.
+
+    Regardless of the choice of partitioning algorithm, a metadata JSON file
+    will be created in the output directory which includes the meta information
+    of the partition algorithm.
+
+    To generate a metadata JSON:
+
+    >>> part_meta = PartitionMeta(version='1.0.0', num_parts=4, algo_name='random')
+    >>> with open('metadata.json', 'w') as f:
+    ...     json.dump(part_meta.dict(), f)
+
+    To read a metadata JSON:
+
+    >>> with open('metadata.json') as f:
+    ...     part_meta = PartitionMeta(**(json.load(f)))
+
+    """
+
+    # version of metadata JSON.
+    version: Optional[str] = "1.0.0"
+    # number of partitions.
+    num_parts: int
+    # name of partition algorithm.
+    algo_name: str
+
+
+def dump_partition_meta(part_meta, meta_file):
+    """Dump partition metadata into json file.
+
+    Parameters
+    ----------
+    part_meta : PartitionMeta
+        The partition metadata.
+    meta_file : str
+        The target file to save data.
+    """
+    with open(meta_file, "w") as f:
+        json.dump(part_meta.dict(), f, sort_keys=True, indent=4)
+
+
+def load_partition_meta(meta_file):
+    """Load partition metadata and do sanity check.
+
+    Parameters
+    ----------
+    meta_file : str
+        The path of the partition metadata file.
+
+    Returns
+    -------
+    PartitionMeta
+        The partition metadata.
+    """
+    with open(meta_file) as f:
+        try:
+            part_meta = PartitionMeta(**(json.load(f)))
+        except dt.ValidationError as e:
+            raise DGLError(
+                f"Invalid partition metadata JSON. Error details: {e.json()}"
+            )
+        if part_meta.version != "1.0.0":
+            raise DGLError(
+                f"Invalid version[{part_meta.version}]. Supported versions: '1.0.0'"
+            )
+        if part_meta.num_parts <= 0:
+            raise DGLError(
+                f"num_parts[{part_meta.num_parts}] should be greater than 0."
+            )
+        if part_meta.algo_name not in ["random", "metis"]:
+            raise DGLError(
+                f"algo_name[{part_meta.num_parts}] is not supported."
+            )
+        return part_meta

tools/partition_algo/random_partition.py

@@ -0,0 +1,61 @@
+# Requires setting PYTHONPATH=${GITROOT}/tools
+import argparse
+import json
+import logging
+import os
+
+import numpy as np
+from base import dump_partition_meta, PartitionMeta
+from distpartitioning import array_readwriter
+from files import setdir
+
+
+def _random_partition(metadata, num_parts):
+    num_nodes_per_type = metadata["num_nodes_per_type"]
+    ntypes = metadata["node_type"]
+    for ntype, n in zip(ntypes, num_nodes_per_type):
+        logging.info("Generating partition for node type %s" % ntype)
+        parts = np.random.randint(0, num_parts, (n,))
+        array_readwriter.get_array_parser(name="csv").write(
+            ntype + ".txt", parts
+        )
+
+
+def random_partition(metadata, num_parts, output_path):
+    """
+    Randomly partition the graph described in metadata and generate partition ID mapping
+    in :attr:`output_path`.
+
+    A directory will be created at :attr:`output_path` containing the partition ID
+    mapping files named "<node-type>.txt" (e.g. "author.txt", "paper.txt" and
+    "institution.txt" for OGB-MAG240M).  Each file contains one line per node representing
+    the partition ID the node belongs to.
+    In addition, metadata which includes version, number of partitions is dumped.
+    """
+    with setdir(output_path):
+        _random_partition(metadata, num_parts)
+        part_meta = PartitionMeta(
+            version="1.0.0", num_parts=num_parts, algo_name="random"
+        )
+        dump_partition_meta(part_meta, "partition_meta.json")
+
+
+# Run with PYTHONPATH=${GIT_ROOT_DIR}/tools
+# where ${GIT_ROOT_DIR} is the directory to the DGL git repository.
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--in_dir",
+        type=str,
+        help="input directory that contains the metadata file",
+    )
+    parser.add_argument("--out_dir", type=str, help="output directory")
+    parser.add_argument(
+        "--num_partitions", type=int, help="number of partitions"
+    )
+    logging.basicConfig(level="INFO")
+    args = parser.parse_args()
+    with open(os.path.join(args.in_dir, "metadata.json")) as f:
+        metadata = json.load(f)
+    num_parts = args.num_partitions
+    random_partition(metadata, num_parts, args.out_dir)

tools/verification_utils.py

@@ -0,0 +1,293 @@
+import json
+import os
+
+import constants
+
+import dgl
+
+import numpy as np
+import pyarrow
+import pyarrow.parquet as pq
+import pytest
+import torch
+from dgl.data.utils import load_tensors
+from dgl.distributed.partition import (
+    _etype_str_to_tuple,
+    _etype_tuple_to_str,
+    _get_inner_edge_mask,
+    _get_inner_node_mask,
+    RESERVED_FIELD_DTYPE,
+)
+from distpartitioning.utils import get_idranges
+
+
+def read_file(fname, ftype):
+    """Read a file from disk
+    Parameters:
+    -----------
+    fname : string
+        specifying the absolute path to the file to read
+    ftype : string
+        supported formats are `numpy`, `parquet', `csv`
+
+    Returns:
+    --------
+    numpy ndarray :
+        file contents are returned as numpy array
+    """
+    reader_fmt_meta = {"name": ftype}
+    array_readwriter.get_array_parser(**reader_fmt_meta).read(fname)
+
+    return data
+
+
+def verify_partition_data_types(part_g):
+    """Validate the dtypes in the partitioned graphs are valid
+
+    Parameters:
+    -----------
+    part_g : DGL Graph object
+        created for the partitioned graphs
+    """
+    for k, dtype in RESERVED_FIELD_DTYPE.items():
+        if k in part_g.ndata:
+            assert part_g.ndata[k].dtype == dtype
+        if k in part_g.edata:
+            assert part_g.edata[k].dtype == dtype
+
+
+def verify_partition_formats(part_g, formats):
+    """Validate the partitioned graphs with supported formats
+
+    Parameters:
+    -----------
+    part_g : DGL Graph object
+        created for the partitioned graphs
+    formats : string
+        formats(csc, coo, csr) supported formats and multiple
+        values can be seperated by comma
+    """
+    # Verify saved graph formats
+    if formats is None:
+        assert "coo" in part_g.formats()["created"]
+    else:
+        formats = formats.split(",")
+        for format in formats:
+            assert format in part_g.formats()["created"]
+
+
+def verify_graph_feats(
+    g, gpb, part, node_feats, edge_feats, orig_nids, orig_eids
+):
+    """Verify the node/edge features of the partitioned graph with
+    the original graph
+
+    Parameters:
+    -----------
+    g : DGL Graph Object
+        of the original graph
+    gpb : global partition book
+        created for the partitioned graph object
+    node_feats : dictionary
+        with key, value pairs as node-types and features as numpy arrays
+    edge_feats : dictionary
+        with key, value pairs as edge-types and features as numpy arrays
+    orig_nids : dictionary
+        with key, value pairs as node-types and (global) nids from the
+        original graph
+    orig_eids : dictionary
+        with key, value pairs as edge-types and (global) eids from the
+        original graph
+    """
+    for ntype in g.ntypes:
+        ntype_id = g.get_ntype_id(ntype)
+        inner_node_mask = _get_inner_node_mask(part, ntype_id)
+        inner_nids = part.ndata[dgl.NID][inner_node_mask]
+        ntype_ids, inner_type_nids = gpb.map_to_per_ntype(inner_nids)
+        partid = gpb.nid2partid(inner_type_nids, ntype)
+        assert np.all(ntype_ids.numpy() == ntype_id)
+        assert np.all(partid.numpy() == gpb.partid)
+
+        orig_id = orig_nids[ntype][inner_type_nids]
+        local_nids = gpb.nid2localnid(inner_type_nids, gpb.partid, ntype)
+
+        for name in g.nodes[ntype].data:
+            if name in [dgl.NID, "inner_node"]:
+                continue
+            true_feats = g.nodes[ntype].data[name][orig_id]
+            ndata = node_feats[ntype + "/" + name][local_nids]
+            assert np.array_equal(ndata.numpy(), true_feats.numpy())
+
+    for etype in g.canonical_etypes:
+        etype_id = g.get_etype_id(etype)
+        inner_edge_mask = _get_inner_edge_mask(part, etype_id)
+        inner_eids = part.edata[dgl.EID][inner_edge_mask]
+        etype_ids, inner_type_eids = gpb.map_to_per_etype(inner_eids)
+        partid = gpb.eid2partid(inner_type_eids, etype)
+        assert np.all(etype_ids.numpy() == etype_id)
+        assert np.all(partid.numpy() == gpb.partid)
+
+        orig_id = orig_eids[_etype_tuple_to_str(etype)][inner_type_eids]
+        local_eids = gpb.eid2localeid(inner_type_eids, gpb.partid, etype)
+
+        for name in g.edges[etype].data:
+            if name in [dgl.EID, "inner_edge"]:
+                continue
+            true_feats = g.edges[etype].data[name][orig_id]
+            edata = edge_feats[_etype_tuple_to_str(etype) + "/" + name][
+                local_eids
+            ]
+            assert np.array_equal(edata.numpy(), true_feats.numpy())
+
+
+def verify_metadata_counts(part_schema, part_g, graph_schema, g, partid):
+    """Verify the partitioned graph objects with the metadata
+
+    Parameters:
+    -----------
+    part_schema : json object
+        which is created by reading the metadata.json file for the
+        partitioned graph
+    part_g : DGL graph object
+        of a graph partition
+    graph_schema : json object
+        which is created by reading the metadata.json file for the
+        original graph
+    g : DGL Graph object
+        created by reading the original graph from the disk.
+    partid : integer
+        specifying the partition id of the graph object, part_g
+    """
+    for ntype in part_schema[constants.STR_NTYPES]:
+        ntype_data = part_schema[constants.STR_NODE_MAP][ntype]
+        meta_ntype_count = ntype_data[partid][1] - ntype_data[partid][0]
+        inner_node_mask = _get_inner_node_mask(part_g, g.get_ntype_id(ntype))
+        graph_ntype_count = len(part_g.ndata[dgl.NID][inner_node_mask])
+        assert (
+            meta_ntype_count == graph_ntype_count
+        ), f"Metadata ntypecount = {meta_ntype_count} and graph_ntype_count = {graph_ntype_count}"
+
+    for etype in part_schema[constants.STR_ETYPES]:
+        etype_data = part_schema[constants.STR_EDGE_MAP][etype]
+        meta_etype_count = etype_data[partid][1] - etype_data[partid][0]
+        mask = _get_inner_edge_mask(
+            part_g, g.get_etype_id(_etype_str_to_tuple(etype))
+        )
+        graph_etype_count = len(part_g.edata[dgl.EID][mask])
+        assert (
+            meta_etype_count == graph_etype_count
+        ), f"Metadata etypecount = {meta_etype_count} does not match part graph etypecount = {graph_etype_count}"
+
+
+def get_node_partids(partitions_dir, graph_schema):
+    """load the node partition ids from the disk
+
+    Parameters:
+    ----------
+    partitions_dir : string
+        directory path where metis/random partitions are located
+    graph_schema : json object
+        which is created by reading the metadata.json file for the
+        original graph
+
+    Returns:
+    --------
+    dictionary :
+        where keys are node-types and value is a list of partition-ids for all the
+        nodes of that particular node-type.
+    """
+    assert os.path.isdir(
+        partitions_dir
+    ), f"Please provide a valid directory to read nodes to partition-id mappings."
+    _, gid_dict = get_idranges(
+        graph_schema[constants.STR_NODE_TYPE],
+        dict(
+            zip(
+                graph_schema[constants.STR_NODE_TYPE],
+                graph_schema[constants.STR_NODE_TYPE_COUNTS],
+            )
+        ),
+    )
+    node_partids = {}
+    for ntype_id, ntype in enumerate(graph_schema[constants.STR_NODE_TYPE]):
+        node_partids[ntype] = read_file(
+            os.path.join(partitions_dir, f"{ntype}.txt"), constants.STR_CSV
+        )
+        assert (
+            len(node_partids[ntype])
+            == graph_schema[constants.STR_NODE_TYPE_COUNTS][ntype_id]
+        ), f"Node count for {ntype} = {len(node_partids[ntype])} in the partitions_dir while it should be {graph_schema[constants.STR_NTYPE_COUNTS][ntype_id]} (from graph schema)."
+
+    return node_partids
+
+
+def verify_node_partitionids(
+    node_partids, part_g, g, gpb, graph_schema, orig_nids, partition_id
+):
+    """Verify partitioned graph objects node counts with the original graph
+
+    Parameters:
+    -----------
+    params : argparser object
+        to access command line arguments for this python script
+    part_data : list of tuples
+        partitioned graph objects read from the disk
+    g : DGL Graph object
+        created by reading the original graph from disk
+    graph_schema : json object
+        created by reading the metadata.json file for the original graph
+    orig_nids : dictionary
+        which contains the origial(global) node-ids
+    partition_id : integer
+        partition id of the partitioned graph, part_g
+    """
+    # read part graphs and verify the counts
+    # inner node masks, should give the node counts in each part-g and get the corresponding orig-ids to map to the original graph node-ids
+    for ntype_id, ntype in enumerate(graph_schema[constants.STR_NODE_TYPE]):
+        mask = _get_inner_node_mask(part_g, g.get_ntype_id(ntype))
+
+        # map these to orig-nids.
+        inner_nids = part_g.ndata[dgl.NID][mask]
+        ntype_ids, inner_type_nids = gpb.map_to_per_ntype(inner_nids)
+        partid = gpb.nid2partid(inner_type_nids, ntype)
+
+        assert np.all(ntype_ids.numpy() == ntype_id)
+        assert np.all(partid.numpy() == gpb.partid)
+
+        idxes = orig_nids[ntype][inner_type_nids]
+        assert np.all(idxes >= 0)
+
+        # get the partition-ids for these nodes.
+        assert np.all(
+            node_partids[ntype][idxes] == partition_id
+        ), f"All the nodes in the partition = {partid} does not their nodeid to partition-id maps are defined by the partitioning algorithm. Node-type = {ntype}"
+
+
+def read_orig_ids(out_dir, fname, num_parts):
+    """Read original id files for the partitioned graph objects
+
+    Parameters:
+    -----------
+    out_dir : string
+        specifying the directory where the files are located
+    fname : string
+        file name to read from
+    num_parts : integer
+        no. of partitions
+
+    Returns:
+    --------
+    dictionary :
+        where keys are node/edge types and values are original node
+        or edge ids from the original graph
+    """
+    orig_ids = {}
+    for i in range(num_parts):
+        ids_path = os.path.join(out_dir, f"part{i}", fname)
+        part_ids = load_tensors(ids_path)
+        for type, data in part_ids.items():
+            if type not in orig_ids:
+                orig_ids[type] = data.numpy()
+            else:
+                orig_ids[type] = np.concatenate((orig_ids[type], data))
+    return orig_ids

tools/verify_partitions.py

@@ -0,0 +1,246 @@
+import argparse
+import logging
+import os
+import platform
+
+import constants
+
+import dgl
+import numpy as np
+
+import pyarrow
+import pyarrow.parquet as pq
+import torch as th
+from dgl.data.utils import load_graphs, load_tensors
+
+from dgl.distributed.partition import (
+    _etype_str_to_tuple,
+    _etype_tuple_to_str,
+    _get_inner_edge_mask,
+    _get_inner_node_mask,
+    load_partition,
+    RESERVED_FIELD_DTYPE,
+)
+from utils import get_idranges, read_json
+from verification_utils import (
+    get_node_partids,
+    read_file,
+    read_orig_ids,
+    verify_graph_feats,
+    verify_metadata_counts,
+    verify_node_partitionids,
+    verify_partition_data_types,
+    verify_partition_formats,
+)
+
+
+def _read_graph(schema):
+    """Read a DGL Graph object from storage using metadata schema, which is
+    a json object describing the DGL graph on disk.
+
+    Parameters:
+    -----------
+    schema : json object
+        json object describing the input graph to read from the disk
+
+    Returns:
+    --------
+    DGL Graph Object :
+        DGL Graph object is created which is read from the disk storage.
+    """
+    edges = {}
+    edge_types = schema[constants.STR_EDGE_TYPE]
+    for etype in edge_types:
+        efiles = schema[constants.STR_EDGES][etype][constants.STR_DATA]
+        src = []
+        dst = []
+        for fname in efiles:
+            if (
+                schema[constants.STR_EDGES][etype][constants.STR_FORMAT][
+                    constants.STR_NAME
+                ]
+                == constants.STR_CSV
+            ):
+                data = read_file(fname, constants.STR_CSV)
+            elif (
+                schema[constants.STR_EDGES][etype][constants.STR_FORMAT][
+                    constants.STR_NAME
+                ]
+                == constants.STR_PARQUET
+            ):
+                data = read_file(fname)
+            else:
+                raise ValueError(
+                    f"Unknown edge format for {etype} - {schema[constants.STR_EDGES][etype][constants.STR_FORMAT]}"
+                )
+
+                src.append(data[:, 0])
+            dst.append(data[:, 1])
+        src = np.concatenate(src)
+        dst = np.concatenate(dst)
+        edges[_etype_str_to_tuple(etype)] = (src, dst)
+
+    g = dgl.heterograph(edges)
+    # g = dgl.to_homogeneous(g)
+
+    g.ndata["orig_id"] = g.ndata[dgl.NID]
+    g.edata["orig_id"] = g.edata[dgl.EID]
+
+    # read features here.
+    for ntype in schema[constants.STR_NODE_TYPE]:
+        if ntype in schema[constants.STR_NODE_DATA]:
+            for featname, featdata in schema[constants.STR_NODE_DATA][
+                ntype
+            ].items():
+                files = fdata[constants.STR_DATA]
+                feats = []
+                for fname in files:
+                    feats.append(read_file(fname, constants.STR_NUMPY))
+                if len(feats) > 0:
+                    g.nodes[ntype].data[featname] = th.from_numpy(
+                        np.concatenate(feats)
+                    )
+
+    # read edge features here.
+    for etype in schema[constants.STR_EDGE_TYPE]:
+        if etype in schema[constants.STR_EDGE_DATA]:
+            for featname, fdata in schema[constants.STR_EDGE_DATA][etype]:
+                files = fdata[constants.STR_DATA]
+                feats = []
+                for fname in files:
+                    feats.append(read_file(fname))
+                if len(feats) > 0:
+                    g.edges[etype].data[featname] = th.from_numpy(
+                        np.concatenate(feats)
+                    )
+
+    # print from graph
+    logging.info(f"|V|= {g.num_nodes()}")
+    logging.info(f"|E|= {g.num_edges()}")
+    for ntype in g.ntypes:
+        for name, data in g.nodes[ntype].data.items():
+            if isinstance(data, th.Tensor):
+                logging.info(
+                    f"Input Graph: nfeat - {ntype}/{name} - data - {data.size()}"
+                )
+
+    for c_etype in g.canonical_etypes:
+        for name, data in g.edges[c_etype].data.items():
+            if isinstance(data, th.Tensor):
+                logging.info(
+                    f"Input Graph: efeat - {etype}/{name} - data - {g.edges[etype].data[name].size()}"
+                )
+
+    return g
+
+
+def _read_part_graphs(part_config, part_metafile):
+    """Read partitioned graph objects from disk storage.
+
+    Parameters:
+    ----------
+    part_config : json object
+        json object created using the metadata file for the partitioned graph.
+    part_metafile : string
+        absolute path of the metadata.json file for the partitioned graph.
+
+    Returns:
+    --------
+    list of tuples :
+        where each tuple contains 4 objects in the following order:
+            partitioned graph object
+            global partition book
+            node features
+            edge features
+    """
+    part_graph_data = []
+    for i in range(part_config["num_parts"]):
+        part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
+            part_metafile, i
+        )
+        part_graph_data.append((part_g, node_feats, edge_feats, gpb))
+    return part_graph_data
+
+
+def _validate_results(params):
+    """Main function to verify the graph partitions
+
+    Parameters:
+    -----------
+    params : argparser object
+        to access the command line arguments
+    """
+    logging.info(f"loading config files...")
+    part_config = os.path.join(params.part_graph_dir, "metadata.json")
+    part_schema = read_json(part_config)
+    num_parts = part_schema["num_parts"]
+
+    logging.info(f"loading config files of the original dataset...")
+    graph_config = os.path.join(params.orig_dataset_dir, "metadata.json")
+    graph_schema = read_json(graph_config)
+
+    logging.info(f"loading original ids from the dgl files...")
+    orig_nids = read_orig_ids(params.part_graph_dir, "orig_nids.dgl", num_parts)
+    orig_eids = read_orig_ids(params.part_graph_dir, "orig_eids.dgl", num_parts)
+
+    logging.info(f"loading node to partition-ids from files... ")
+    node_partids = get_node_partids(params.partitions_dir, graph_schema)
+
+    logging.info(f"loading the original dataset...")
+    g = _read_graph(graph_schema)
+
+    logging.info(f"Beginning the verification process...")
+    for i in range(num_parts):
+        part_g, node_feats, edge_feats, gpb, _, _, _ = load_partition(
+            part_config, i
+        )
+
+        verify_partition_data_types(part_g)
+        verify_partition_formats(part_g, None)
+        verify_graph_feats(
+            g, gpb, part_g, node_feats, edge_feats, orig_nids, orig_eids
+        )
+        verify_metadata_counts(part_schema, part_g, graph_schema, g, i)
+        verify_node_partitionids(
+            node_partids, part_g, g, gpb, graph_schema, orig_nids, i
+        )
+        logging.info(f"Verification of partitioned graph - {i}... SUCCESS !!!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Construct graph partitions")
+    parser.add_argument(
+        "--orig-dataset-dir",
+        required=True,
+        type=str,
+        help="The directory path that contains the original graph input files.",
+    )
+    parser.add_argument(
+        "--part-graph-dir",
+        required=True,
+        type=str,
+        help="The directory path that contains the partitioned graph files.",
+    )
+    parser.add_argument(
+        "--partitions-dir",
+        required=True,
+        type=str,
+        help="The directory path that contains metis/random partitions results.",
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="info",
+        help="To enable log level for debugging purposes. Available options: \
+                          (Critical, Error, Warning, Info, Debug, Notset), default value \
+                          is: Info",
+    )
+    params = parser.parse_args()
+
+    numeric_level = getattr(logging, params.log_level.upper(), None)
+    logging.basicConfig(
+        level=numeric_level,
+        format=f"[{platform.node()} %(levelname)s %(asctime)s PID:%(process)d] %(message)s",
+    )
+
+    _validate_results(params)