class_docstr: 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
docstr: 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.
docstr_tok: ["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","."]
code: 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
  
file_path: dgl/tools/distpartitioning/dist_lookup.py
imports: import copy
  import logging
  import os
  import numpy
  import pyarrow
  import torch
  from gloo_wrapper import allgather_sizes, alltoallv_cpu
  from pyarrow import csv
  from utils import map_partid_rank
label_desc: Use image node_id 4 for calling the DistLookupService obj's underlying member method code with example usage: obj.get_shuffle_nids(global_nids, my_global_nids, my_shuffle_global_nids, world_size) and returns: shuffle_global_nids