code: def 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,
      cur_features,
      cur_global_ids,
  ):
      """This function is used to send/receive one feature for either nodes or
      edges of the input graph dataset.
  
      Parameters:
      -----------
      rank : int
          integer, unique id assigned to the current process
      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
      id_lookup : instance of DistLookupService
          instance of an implementation of dist. lookup service to retrieve values
          for keys
      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
      feat_key : string
          this string is used as a key in the dictionary to store features, as
          tensors, in local dictionaries
      featdata_key : numpy array
          features associated with this feature key being processed
      gid_start : int
          starting global_id, of either node or edge, for the feature data
      gid_end : int
          ending global_if, of either node or edge, for the feature data
      type_id_start : int
          starting type_id for the feature data
      type_id_end : int
          ending type_id for the feature data
      local_part_id : int
          integers used to the identify the local partition id used to locate
          data belonging to this partition
      world_size : int
          total number of processes created
      num_parts : int
          total number of partitions
      cur_features : dictionary
          dictionary to store the feature data which belongs to the current
          process
      cur_global_ids : dictionary
          dictionary to store global ids, of either nodes or edges, for which
          the features stored in the cur_features dictionary
  
      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
      """
      # type_ids for this feature subset on the current rank
      gids_feat = np.arange(gid_start, gid_end)
      tids_feat = np.arange(type_id_start, type_id_end)
      local_idx = np.arange(0, type_id_end - type_id_start)
  
      feats_per_rank = []
      global_id_per_rank = []
  
      tokens = feat_key.split("/")
      assert len(tokens) == 3
      local_feat_key = "/".join(tokens[:-1]) + "/" + str(local_part_id)
  
      logging.debug(
          f"[Rank: {rank} feature: {feat_key}, gid_start - {gid_start} and gid_end - {gid_end}"
      )
  
      # Get the partition ids for the range of global nids.
      if feat_type == constants.STR_NODE_FEATURES:
          # Retrieve the partition ids for the node features.
          # Each partition id will be in the range [0, num_parts).
          partid_slice = id_lookup.get_partition_ids(
              np.arange(gid_start, gid_end, dtype=np.int64)
          )
      else:
          # Edge data case.
          # Ownership is determined by the destination node.
          assert data is not None
          global_eids = np.arange(gid_start, gid_end, dtype=np.int64)
          if data[constants.GLOBAL_EID].shape[0] > 0:
              logging.debug(
                  f"[Rank: {rank} disk read global eids - min - {np.amin(data[constants.GLOBAL_EID])}, max - {np.amax(data[constants.GLOBAL_EID])}, count - {data[constants.GLOBAL_EID].shape}"
              )
  
          # Now use `data` to extract destination nodes' global id
          # and use that to get the ownership
          common, idx1, idx2 = np.intersect1d(
              data[constants.GLOBAL_EID],
              global_eids,
              return_indices=True,
          )
          assert (
              common.shape[0] == idx2.shape[0]
          ), f"Rank {rank}: {common.shape[0]} != {idx2.shape[0]}"
          assert (
              common.shape[0] == global_eids.shape[0]
          ), f"Rank {rank}: {common.shape[0]} != {global_eids.shape[0]}"
  
          global_dst_nids = data[constants.GLOBAL_DST_ID][idx1]
          assert np.all(global_eids == data[constants.GLOBAL_EID][idx1])
          partid_slice = id_lookup.get_partition_ids(global_dst_nids)
  
      # determine the shape of the feature-data
      # this is needed to so that ranks where feature-data is not present
      # should use the correct shape for sending the padded vector.
      # exchange length here.
      feat_dim_len = 0
      if featdata_key is not None:
          feat_dim_len = len(featdata_key.shape)
      all_lens = allgather_sizes(
          [feat_dim_len], world_size, num_parts, return_sizes=True
      )
      if all_lens[0] <= 0:
          logging.debug(
              f"[Rank: {rank} No process has any feature data to shuffle for {local_feat_key}"
          )
          return cur_features, cur_global_ids
  
      rank0_shape_len = all_lens[0]
      for idx in range(1, world_size):
          assert (all_lens[idx] == 0) or (
              all_lens[idx] == rank0_shape_len
          ), (
              f"feature: {local_feat_key} shapes does not match "
              f"at rank - {idx} and rank - 0"
          )
  
      # exchange actual data here.
      if featdata_key != None:
          feat_dims_dtype = list(featdata_key.shape)
          feat_dims_dtype.append(DATA_TYPE_ID[featdata_key.dtype])
      else:
          feat_dims_dtype = list(
              np.zeros((rank0_shape_len), dtype=np.int64)
          )
          feat_dims_dtype.append(DATA_TYPE_ID[torch.float32])
  
      logging.debug(
          f"Sending the feature shape information - {feat_dims_dtype}"
      )
      all_dims_dtype = allgather_sizes(
          feat_dims_dtype, world_size, num_parts, return_sizes=True
      )
  
      for idx in range(world_size):
          cond = partid_slice == (idx + local_part_id * world_size)
          gids_per_partid = gids_feat[cond]
          tids_per_partid = tids_feat[cond]
          local_idx_partid = local_idx[cond]
  
          if gids_per_partid.shape[0] == 0:
              assert len(all_dims_dtype) % world_size == 0
              dim_len = int(len(all_dims_dtype) / world_size)
              rank0_shape = tuple(
                  list(np.zeros((dim_len - 1), dtype=np.int32))
              )
              rank0_dtype = REV_DATA_TYPE_ID[
                  all_dims_dtype[(dim_len - 1) : (dim_len)][0]
              ]
              data = torch.empty(rank0_shape, dtype=rank0_dtype)
              feats_per_rank.append(data)
              global_id_per_rank.append(
                  torch.empty((0,), dtype=torch.int64)
              )
          else:
              feats_per_rank.append(featdata_key[local_idx_partid])
              global_id_per_rank.append(
                  torch.from_numpy(gids_per_partid).type(torch.int64)
              )
      for idx, tt in enumerate(feats_per_rank):
          logging.debug(
              f"[Rank: {rank} features shape - {tt.shape} and ids - {global_id_per_rank[idx].shape}"
          )
  
      # features (and global nids) per rank to be sent out are ready
      # for transmission, perform alltoallv here.
      output_feat_list = alltoallv_cpu(
          rank, world_size, feats_per_rank, retain_nones=False
      )
      output_id_list = alltoallv_cpu(
          rank, world_size, global_id_per_rank, retain_nones=False
      )
      logging.debug(
          f"[Rank : {rank} feats - {output_feat_list}, ids - {output_id_list}"
      )
      assert len(output_feat_list) == len(output_id_list), (
          "Length of feature list and id list are expected to be equal while "
          f"got {len(output_feat_list)} and {len(output_id_list)}."
      )
  
      # stitch node_features together to form one large feature tensor
      if len(output_feat_list) > 0:
          output_feat_list = torch.cat(output_feat_list)
          output_id_list = torch.cat(output_id_list)
          if local_feat_key in cur_features:
              temp = cur_features[local_feat_key]
              cur_features[local_feat_key] = torch.cat(
                  [temp, output_feat_list]
              )
              temp = cur_global_ids[local_feat_key]
              cur_global_ids[local_feat_key] = torch.cat(
                  [temp, output_id_list]
              )
          else:
              cur_features[local_feat_key] = output_feat_list
              cur_global_ids[local_feat_key] = output_id_list
  
      return cur_features, cur_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 3 for calling a global function with example usage: 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, cur_features, cur_global_ids) and returns: cur_features,  cur_global_ids, cur_features,  cur_global_ids