code: def exchange_features(
      rank,
      world_size,
      num_parts,
      feature_tids,
      type_id_map,
      id_lookup,
      feature_data,
      feat_type,
      data,
  ):
      """
      This function is used to shuffle node features so that each process will receive
      all the node features whose corresponding nodes are owned by the same process.
      The mapping procedure to identify the owner process is not straight forward. The
      following steps are used to identify the owner processes for the locally read node-
      features.
      a. Compute the global_nids for the locally read node features. Here metadata json file
          is used to identify the corresponding global_nids. Please note that initial graph input
          nodes.txt files are sorted based on node_types.
      b. Using global_nids and metis partitions owner processes can be easily identified.
      c. Now each process sends the global_nids for which shuffle_global_nids are needed to be
          retrieved.
      d. After receiving the corresponding shuffle_global_nids these ids are added to the
          node_data and edge_data dictionaries
  
      This pipeline assumes all the input data in numpy format, except node/edge features which
      are maintained as tensors throughout the various stages of the pipeline execution.
  
      Parameters:
      -----------
      rank : int
          rank of the current process
      world_size : int
          total no. of participating processes.
      feature_tids : dictionary
          dictionary with keys as node-type names with suffixes as feature names
          and value is a dictionary. This dictionary contains information about
          node-features associated with a given node-type and value is a list.
          This list contains a of indexes, like [starting-idx, ending-idx) which
          can be used to index into the node feature tensors read from
          corresponding input files.
      type_id_map : dictionary
          mapping between type names and global_ids, of either nodes or edges,
          which belong to the keys in this dictionary
      id_lookup : instance of class DistLookupService
         Distributed lookup service used to map global-nids to respective
         partition-ids and shuffle-global-nids
      feat_type : string
          this is used to distinguish which features are being exchanged. Please
          note that for nodes ownership is clearly defined and for edges it is
          always assumed that destination end point of the edge defines the
          ownership of that particular edge
      data: dicitonary
          dictionry in which node or edge features are stored and this information
          is read from the appropriate node features file which belongs to the
          current process
  
      Returns:
      --------
      dictionary :
          a dictionary is returned where keys are type names and
          feature data are the values
      list :
          a dictionary of global_ids either nodes or edges whose features are
          received during the data shuffle process
      """
      start = timer()
      own_features = {}
      own_global_ids = {}
  
      # To iterate over the node_types and associated node_features
      for feat_key, type_info in feature_tids.items():
          # To iterate over the feature data, of a given (node or edge )type
          # type_info is a list of 3 elements (as shown below):
          #   [feature-name, starting-idx, ending-idx]
          #       feature-name is the name given to the feature-data,
          #       read from the input metadata file
          #       [starting-idx, ending-idx) specifies the range of indexes
          #        associated with the features data
          # Determine the owner process for these features.
          # Note that the keys in the node features (and similarly edge features)
          # dictionary is of the following format:
          #   `node_type/feature_name/local_part_id`:
          #    where node_type and feature_name are self-explanatory and
          #    local_part_id denotes the partition-id, in the local process,
          #    which will be used a suffix to store all the information of a
          #    given partition which is processed by the current process. Its
          #    values start from 0 onwards, for instance 0, 1, 2 ... etc.
          #    local_part_id can be easily mapped to global partition id very
          #    easily, using cyclic ordering. All local_part_ids = 0 from all
          #    processes will form global partition-ids between 0 and world_size-1.
          #    Similarly all local_part_ids = 1 from all processes will form
          #    global partition ids in the range [world_size, 2*world_size-1] and
          #    so on.
          tokens = feat_key.split("/")
          assert len(tokens) == 3
          type_name = tokens[0]
          feat_name = tokens[1]
          logging.debug(
              f"[Rank: {rank}] processing feature: {feat_key}"
          )
  
          for feat_info in type_info:
              # Compute the global_id range for this feature data
              type_id_start = int(feat_info[0])
              type_id_end = int(feat_info[1])
              begin_global_id = type_id_map[type_name][0]
              gid_start = begin_global_id + type_id_start
              gid_end = begin_global_id + type_id_end
  
              # Check if features exist for this type_name + feat_name.
              # This check should always pass, because feature_tids are built
              # by reading the input metadata json file for existing features.
              assert feat_key in feature_data
  
              for local_part_id in range(num_parts // world_size):
                  featdata_key = feature_data[feat_key]
  
                  # Synchronize for each feature
                  dist.barrier()
                  own_features, own_global_ids = exchange_feature(
                      rank,
                      data,
                      id_lookup,
                      feat_type,
                      feat_key,
                      featdata_key,
                      gid_start,
                      gid_end,
                      type_id_start,
                      type_id_end,
                      local_part_id,
                      world_size,
                      num_parts,
                      own_features,
                      own_global_ids,
                  )
  
      end = timer()
      logging.info(
          f"[Rank: {rank}] Total time for feature exchange: {timedelta(seconds = end - start)}"
      )
      for k, v in own_features.items():
          logging.debug(f"Rank: {rank}] Key - {k} Value - {v.shape}")
      return own_features, own_global_ids
  
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 4 for calling a global function with example usage: exchange_features(rank, world_size, num_parts, feature_tids, type_id_map, id_lookup, feature_data, feat_type, data) and returns: own_features,  own_global_ids