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 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)
docstr_tok: ["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",")"]
code: 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
  
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 3 for calling the DistLookupService obj's underlying member method code with example usage: obj.get_partition_ids(agg_global_nids) and returns: agg_partition_ids