code: def gen_dist_partitions(rank, world_size, params):
      """
      Function which will be executed by all Gloo processes to begin execution of the pipeline.
      This function expects the input dataset is split across multiple file format.
  
      Input dataset and its file structure is described in metadata json file which is also part of the
      input dataset. On a high-level, this metadata json file contains information about the following items
      a) Nodes metadata, It is assumed that nodes which belong to each node-type are split into p files
         (wherer `p` is no. of partitions).
      b) Similarly edge metadata contains information about edges which are split into p-files.
      c) Node and Edge features, it is also assumed that each node (and edge) feature, if present, is also
         split into `p` files.
  
      For example, a sample metadata json file might be as follows: :
      (In this toy example, we assume that we have "m" node-types, "k" edge types, and for node_type = ntype0-name
       we have two features namely feat0-name and feat1-name. Please note that the node-features are also split into
       `p` files. This will help in load-balancing during data-shuffling phase).
  
      Terminology used to identify any particular "id" assigned to nodes, edges or node features. Prefix "global" is
      used to indicate that this information is either read from the input dataset or autogenerated based on the information
      read from input dataset files. Prefix "type" is used to indicate a unique id assigned to either nodes or edges.
      For instance, type_node_id means that a unique id, with a given node type,  assigned to a node. And prefix "shuffle"
      will be used to indicate a unique id, across entire graph, assigned to either a node or an edge. For instance,
      SHUFFLE_GLOBAL_NID means a unique id which is assigned to a node after the data shuffle is completed.
  
      Some high-level notes on the structure of the metadata json file.
      1. path(s) mentioned in the entries for nodes, edges and node-features files can be either absolute or relative.
         if these paths are relative, then it is assumed that they are relative to the folder from which the execution is
         launched.
      2. The id_startx and id_endx represent the type_node_id and type_edge_id respectively for nodes and edge data. This
         means that these ids should match the no. of nodes/edges read from any given file. Since these are type_ids for
         the nodes and edges in any given file, their global_ids can be easily computed as well.
  
      {
          "graph_name" : xyz,
          "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
          ],
          "edge_type" : ["src_ntype:edge_type:dst_ntype", ....], #k edge types
          "num_edges_per_chunk" : [
              [a0, a1, ...a<p-1>], #p partitions
              [b0, b1, ... b<p-1>],
              ....
              [c0, c1, ..., c<p-1>] #no, of edge types
          ],
          "node_data" : {
              "ntype0-name" : {
                  "feat0-name" : {
                      "format" : {"name": "numpy"},
                      "data" :   [ #list of lists
                          ["<path>/feat-0.npy", 0, id_end0],
                          ["<path>/feat-1.npy", id_start1, id_end1],
                          ....
                          ["<path>/feat-<p-1>.npy", id_start<p-1>, id_end<p-1>]
                      ]
                  },
                  "feat1-name" : {
                      "format" : {"name": "numpy"},
                      "data" : [ #list of lists
                          ["<path>/feat-0.npy", 0, id_end0],
                          ["<path>/feat-1.npy", id_start1, id_end1],
                          ....
                          ["<path>/feat-<p-1>.npy", id_start<p-1>, id_end<p-1>]
                      ]
                  }
              }
          },
          "edges": { #k edge types
              "src_ntype:etype0-name:dst_ntype" : {
                  "format": {"name" : "csv", "delimiter" : " "},
                  "data" : [
                      ["<path>/etype0-name-0.txt", 0, id_end0], #These are type_edge_ids for edges of this type
                      ["<path>/etype0-name-1.txt", id_start1, id_end1],
                      ...,
                      ["<path>/etype0-name-<p-1>.txt", id_start<p-1>, id_end<p-1>]
                  ]
              },
              ...,
              "src_ntype:etype<k-1>-name:dst_ntype" : {
                  "format": {"name" : "csv", "delimiter" : " "},
                  "data" : [
                      ["<path>/etype<k-1>-name-0.txt", 0, id_end0],
                      ["<path>/etype<k-1>-name-1.txt", id_start1, id_end1],
                      ...,
                      ["<path>/etype<k-1>-name-<p-1>.txt", id_start<p-1>, id_end<p-1>]
                  ]
              },
          },
      }
  
      The function performs the following steps:
      1. Reads the metis partitions to identify the owner process of all the nodes in the entire graph.
      2. Reads the input data set, each partitipating process will map to a single file for the edges,
          node-features and edge-features for each node-type and edge-types respectively. Using nodes metadata
          information, nodes which are owned by a given process are generated to optimize communication to some
          extent.
      3. Now each process shuffles the data by identifying the respective owner processes using metis
          partitions.
          a. To identify owner processes for nodes, metis partitions will be used.
          b. For edges, the owner process of the destination node will be the owner of the edge as well.
          c. For node and edge features, identifying the owner process is a little bit involved.
              For this purpose, graph metadata json file is used to first map the locally read node features
              to their global_nids. Now owner process is identified using metis partitions for these global_nids
              to retrieve shuffle_global_nids. A similar process is used for edge_features as well.
          d. After all the data shuffling is done, the order of node-features may be different when compared to
              their global_type_nids. Node- and edge-data are ordered by node-type and edge-type respectively.
              And now node features and edge features are re-ordered to match the order of their node- and edge-types.
      4. Last step is to create the DGL objects with the data present on each of the processes.
          a. DGL objects for nodes, edges, node- and edge- features.
          b. Metadata is gathered from each process to create the global metadata json file, by process rank = 0.
  
      Parameters:
      ----------
      rank : int
          integer representing the rank of the current process in a typical distributed implementation
      world_size : int
          integer representing the total no. of participating processes in a typical distributed implementation
      params : argparser object
          this object, key value pairs, provides access to the command line arguments from the runtime environment
      """
      global_start = timer()
      logging.info(
          f"[Rank: {rank}] Starting distributed data processing pipeline..."
      )
      memory_snapshot("Pipeline Begin: ", rank)
  
      # init processing
      schema_map = read_json(
          os.path.join(params.input_dir, params.schema)
      )
  
      # The resources, which are node-id to partition-id mappings, are split
      # into `world_size` number of parts, where each part can be mapped to
      # each physical node.
      id_lookup = DistLookupService(
          os.path.join(params.input_dir, params.partitions_dir),
          schema_map[constants.STR_NODE_TYPE],
          rank,
          world_size,
          params.num_parts,
      )
  
      # get the id to name mappings here.
      ntypes_ntypeid_map, ntypes, ntypeid_ntypes_map = get_node_types(
          schema_map
      )
      etypes_etypeid_map, etypes, etypeid_etypes_map = get_edge_types(
          schema_map
      )
      logging.info(
          f"[Rank: {rank}] Initialized metis partitions and node_types map..."
      )
  
      # Initialize distributed lookup service for partition-id and shuffle-global-nids mappings
      # for global-nids
      _, global_nid_ranges = get_idranges(
          schema_map[constants.STR_NODE_TYPE],
          get_ntype_counts_map(
              schema_map[constants.STR_NODE_TYPE],
              schema_map[constants.STR_NUM_NODES_PER_TYPE],
          ),
      )
      id_map = dgl.distributed.id_map.IdMap(global_nid_ranges)
      id_lookup.set_idMap(id_map)
  
      # read input graph files and augment these datastructures with
      # appropriate information (global_nid and owner process) for node and edge data
      (
          node_features,
          node_feat_tids,
          edge_data,
          edge_typecounts,
          edge_features,
          edge_feat_tids,
      ) = read_dataset(
          rank,
          world_size,
          id_lookup,
          params,
          schema_map,
          get_ntype_counts_map(
              schema_map[constants.STR_NODE_TYPE],
              schema_map[constants.STR_NUM_NODES_PER_TYPE],
          ),
      )
      logging.info(
          f"[Rank: {rank}] Done augmenting file input data with auxilary columns"
      )
      memory_snapshot("DatasetReadComplete: ", rank)
  
      # send out node and edge data --- and appropriate features.
      # this function will also stitch the data recvd from other processes
      # and return the aggregated data
      # ntypes_gnid_range_map = get_gnid_range_map(node_tids)
      # etypes_geid_range_map = get_gnid_range_map(edge_tids)
      ntypes_gnid_range_map = get_gid_offsets(
          schema_map[constants.STR_NODE_TYPE],
          get_ntype_counts_map(
              schema_map[constants.STR_NODE_TYPE],
              schema_map[constants.STR_NUM_NODES_PER_TYPE],
          ),
      )
      etypes_geid_range_map = get_gid_offsets(
          schema_map[constants.STR_EDGE_TYPE], edge_typecounts
      )
  
      (
          node_data,
          rcvd_node_features,
          rcvd_global_nids,
          edge_data,
          rcvd_edge_features,
          rcvd_global_eids,
      ) = exchange_graph_data(
          rank,
          world_size,
          params.num_parts,
          node_features,
          edge_features,
          node_feat_tids,
          edge_feat_tids,
          edge_data,
          id_lookup,
          ntypes_ntypeid_map,
          ntypes_gnid_range_map,
          etypes_geid_range_map,
          ntypeid_ntypes_map,
          schema_map,
      )
      gc.collect()
      logging.debug(f"[Rank: {rank}] Done with data shuffling...")
      memory_snapshot("DataShuffleComplete: ", rank)
  
      # sort node_data by ntype
      node_data = reorder_data(
          params.num_parts, world_size, node_data, constants.NTYPE_ID
      )
      logging.debug(f"[Rank: {rank}] Sorted node_data by node_type")
      memory_snapshot("NodeDataSortComplete: ", rank)
  
      # resolve global_ids for nodes
      # Synchronize before assigning shuffle-global-ids to nodes
      dist.barrier()
      assign_shuffle_global_nids_nodes(
          rank, world_size, params.num_parts, node_data
      )
      logging.debug(
          f"[Rank: {rank}] Done assigning global-ids to nodes..."
      )
      memory_snapshot("ShuffleGlobalID_Nodes_Complete: ", rank)
  
      # shuffle node feature according to the node order on each rank.
      for ntype_name in ntypes:
          featnames = get_ntype_featnames(ntype_name, schema_map)
          for featname in featnames:
              # if a feature name exists for a node-type, then it should also have
              # feature data as well. Hence using the assert statement.
              for local_part_id in range(
                  params.num_parts // world_size
              ):
                  feature_key = (
                      ntype_name
                      + "/"
                      + featname
                      + "/"
                      + str(local_part_id)
                  )
                  assert feature_key in rcvd_global_nids
                  global_nids = rcvd_global_nids[feature_key]
  
                  _, idx1, _ = np.intersect1d(
                      node_data[
                          constants.GLOBAL_NID
                          + "/"
                          + str(local_part_id)
                      ],
                      global_nids,
                      return_indices=True,
                  )
                  shuffle_global_ids = node_data[
                      constants.SHUFFLE_GLOBAL_NID
                      + "/"
                      + str(local_part_id)
                  ][idx1]
                  feature_idx = shuffle_global_ids.argsort()
  
                  rcvd_node_features[feature_key] = rcvd_node_features[
                      feature_key
                  ][feature_idx]
      memory_snapshot("ReorderNodeFeaturesComplete: ", rank)
  
      # Sort edge_data by etype
      edge_data = reorder_data(
          params.num_parts, world_size, edge_data, constants.ETYPE_ID
      )
      logging.debug(f"[Rank: {rank}] Sorted edge_data by edge_type")
      memory_snapshot("EdgeDataSortComplete: ", rank)
  
      # Synchronize before assigning shuffle-global-nids for edges end points.
      dist.barrier()
      shuffle_global_eid_offsets = assign_shuffle_global_nids_edges(
          rank, world_size, params.num_parts, edge_data
      )
      logging.debug(
          f"[Rank: {rank}] Done assigning global_ids to edges ..."
      )
  
      memory_snapshot("ShuffleGlobalID_Edges_Complete: ", rank)
  
      # Shuffle edge features according to the edge order on each rank.
      for etype_name in etypes:
          featnames = get_etype_featnames(etype_name, schema_map)
          for featname in featnames:
              for local_part_id in range(
                  params.num_parts // world_size
              ):
                  feature_key = (
                      etype_name
                      + "/"
                      + featname
                      + "/"
                      + str(local_part_id)
                  )
                  assert feature_key in rcvd_global_eids
                  global_eids = rcvd_global_eids[feature_key]
  
                  _, idx1, _ = np.intersect1d(
                      edge_data[
                          constants.GLOBAL_EID
                          + "/"
                          + str(local_part_id)
                      ],
                      global_eids,
                      return_indices=True,
                  )
                  shuffle_global_ids = edge_data[
                      constants.SHUFFLE_GLOBAL_EID
                      + "/"
                      + str(local_part_id)
                  ][idx1]
                  feature_idx = shuffle_global_ids.argsort()
  
                  rcvd_edge_features[feature_key] = rcvd_edge_features[
                      feature_key
                  ][feature_idx]
  
      # determine global-ids for edge end-points
      # Synchronize before retrieving shuffle-global-nids for edges end points.
      dist.barrier()
      edge_data = lookup_shuffle_global_nids_edges(
          rank,
          world_size,
          params.num_parts,
          edge_data,
          id_lookup,
          node_data,
      )
      logging.debug(
          f"[Rank: {rank}] Done resolving orig_node_id for local node_ids..."
      )
      memory_snapshot("ShuffleGlobalID_Lookup_Complete: ", rank)
  
      def prepare_local_data(src_data, local_part_id):
          local_data = {}
          for k, v in src_data.items():
              tokens = k.split("/")
              if tokens[len(tokens) - 1] == str(local_part_id):
                  local_data["/".join(tokens[:-1])] = v
          return local_data
  
      # create dgl objects here
      output_meta_json = {}
      start = timer()
  
      graph_formats = None
      if params.graph_formats:
          graph_formats = params.graph_formats.split(",")
  
      for local_part_id in range(params.num_parts // world_size):
          # Synchronize for each local partition of the graph object.
          dist.barrier()
  
          num_edges = shuffle_global_eid_offsets[local_part_id]
          node_count = len(
              node_data[constants.NTYPE_ID + "/" + str(local_part_id)]
          )
          edge_count = len(
              edge_data[constants.ETYPE_ID + "/" + str(local_part_id)]
          )
          local_node_data = prepare_local_data(node_data, local_part_id)
          local_edge_data = prepare_local_data(edge_data, local_part_id)
          (
              graph_obj,
              ntypes_map_val,
              etypes_map_val,
              ntypes_map,
              etypes_map,
              orig_nids,
              orig_eids,
          ) = create_dgl_object(
              schema_map,
              rank + local_part_id * world_size,
              local_node_data,
              local_edge_data,
              num_edges,
              get_ntype_counts_map(
                  schema_map[constants.STR_NODE_TYPE],
                  schema_map[constants.STR_NUM_NODES_PER_TYPE],
              ),
              edge_typecounts,
              params.save_orig_nids,
              params.save_orig_eids,
          )
          sort_etypes = len(etypes_map) > 1
          local_node_features = prepare_local_data(
              rcvd_node_features, local_part_id
          )
          local_edge_features = prepare_local_data(
              rcvd_edge_features, local_part_id
          )
          write_dgl_objects(
              graph_obj,
              local_node_features,
              local_edge_features,
              params.output,
              rank + (local_part_id * world_size),
              orig_nids,
              orig_eids,
              graph_formats,
              sort_etypes,
          )
          memory_snapshot("DiskWriteDGLObjectsComplete: ", rank)
  
          # get the meta-data
          json_metadata = create_metadata_json(
              params.graph_name,
              node_count,
              edge_count,
              local_part_id * world_size + rank,
              params.num_parts,
              ntypes_map_val,
              etypes_map_val,
              ntypes_map,
              etypes_map,
              params.output,
          )
          output_meta_json[
              "local-part-id-" + str(local_part_id * world_size + rank)
          ] = json_metadata
          memory_snapshot("MetadataCreateComplete: ", rank)
  
      if rank == 0:
          # get meta-data from all partitions and merge them on rank-0
          metadata_list = gather_metadata_json(
              output_meta_json, rank, world_size
          )
          metadata_list[0] = output_meta_json
          write_metadata_json(
              metadata_list,
              params.output,
              params.graph_name,
              world_size,
              params.num_parts,
          )
      else:
          # send meta-data to Rank-0 process
          gather_metadata_json(output_meta_json, rank, world_size)
      end = timer()
      logging.info(
          f"[Rank: {rank}] Time to create dgl objects: {timedelta(seconds = end - start)}"
      )
      memory_snapshot("MetadataWriteComplete: ", rank)
  
      global_end = timer()
      logging.info(
          f"[Rank: {rank}] Total execution time of the program: {timedelta(seconds = global_end - global_start)}"
      )
      memory_snapshot("PipelineComplete: ", rank)
  
imports: import gc
  import logging
  import math
  import os
  import sys
  from datetime import timedelta
  from timeit import default_timer
  import constants
  import dgl
  import numpy
  import torch
  import torch.distributed
  import torch.multiprocessing
  from convert_partition import create_dgl_object, create_metadata_json
  from dataset_utils import get_dataset
  from dist_lookup import DistLookupService
  from globalids import assign_shuffle_global_nids_edges, assign_shuffle_global_nids_nodes, lookup_shuffle_global_nids_edges
  from gloo_wrapper import allgather_sizes, alltoallv_cpu, gather_metadata_json
  from utils import augment_edge_data, DATA_TYPE_ID, get_edge_types, get_etype_featnames, get_gid_offsets, get_gnid_range_map, get_idranges, get_node_types, get_ntype_counts_map, get_ntype_featnames, map_partid_rank, memory_snapshot, read_json, read_ntype_partition_files, REV_DATA_TYPE_ID, write_dgl_objects, write_metadata_json
label_desc: Use image node_id 8 for calling a global function with example usage: gen_dist_partitions(rank, world_size, params) and returns: local_data