code: def exchange_edge_data(
      rank, world_size, num_parts, edge_data, id_lookup
  ):
      """
      Exchange edge_data among processes in the world.
      Prepare list of sliced data targeting each process and trigger
      alltoallv_cpu to trigger messaging api
  
      Parameters:
      -----------
      rank : int
          rank of the process
      world_size : int
          total no. of processes
      edge_data : dictionary
          edge information, as a dicitonary which stores column names as keys and values
          as column data. This information is read from the edges.txt file.
      id_lookup : DistLookupService instance
          this object will be used to retrieve ownership information of nodes
  
      Returns:
      --------
      dictionary :
          the input argument, edge_data, is updated with the edge data received by other processes
          in the world.
      """
  
      # Synchronize at the beginning of this function
      dist.barrier()
  
      # Prepare data for each rank in the cluster.
      timer_start = timer()
  
      CHUNK_SIZE = (
          100 * 1000 * 1000
      )  # 100 * 8 * 5 = 1 * 4 = 8 GB/message/node
      num_edges = edge_data[constants.GLOBAL_SRC_ID].shape[0]
      all_counts = allgather_sizes(
          [num_edges], world_size, num_parts, return_sizes=True
      )
      max_edges = np.amax(all_counts)
      all_edges = np.sum(all_counts)
      num_chunks = (max_edges // CHUNK_SIZE) + (
          0 if (max_edges % CHUNK_SIZE == 0) else 1
      )
      LOCAL_CHUNK_SIZE = (num_edges // num_chunks) + (
          0 if (num_edges % num_chunks == 0) else 1
      )
      logging.debug(
          f"[Rank: {rank} Edge Data Shuffle - max_edges: {max_edges}, \
                          local_edges: {num_edges} and num_chunks: {num_chunks} \
                          Total edges: {all_edges} Local_CHUNK_SIZE: {LOCAL_CHUNK_SIZE}"
      )
  
      for local_part_id in range(num_parts // world_size):
          local_src_ids = []
          local_dst_ids = []
          local_type_eids = []
          local_etype_ids = []
          local_eids = []
  
          for chunk in range(num_chunks):
              chunk_start = chunk * LOCAL_CHUNK_SIZE
              chunk_end = (chunk + 1) * LOCAL_CHUNK_SIZE
  
              logging.debug(
                  f"[Rank: {rank}] EdgeData Shuffle: processing \
                      local_part_id: {local_part_id} and chunkid: {chunk}"
              )
              cur_src_id = edge_data[constants.GLOBAL_SRC_ID][
                  chunk_start:chunk_end
              ]
              cur_dst_id = edge_data[constants.GLOBAL_DST_ID][
                  chunk_start:chunk_end
              ]
              cur_type_eid = edge_data[constants.GLOBAL_TYPE_EID][
                  chunk_start:chunk_end
              ]
              cur_etype_id = edge_data[constants.ETYPE_ID][
                  chunk_start:chunk_end
              ]
              cur_eid = edge_data[constants.GLOBAL_EID][
                  chunk_start:chunk_end
              ]
  
              input_list = []
              owner_ids = id_lookup.get_partition_ids(cur_dst_id)
              for idx in range(world_size):
                  send_idx = owner_ids == (
                      idx + local_part_id * world_size
                  )
                  send_idx = send_idx.reshape(cur_src_id.shape[0])
                  filt_data = np.column_stack(
                      (
                          cur_src_id[send_idx == 1],
                          cur_dst_id[send_idx == 1],
                          cur_type_eid[send_idx == 1],
                          cur_etype_id[send_idx == 1],
                          cur_eid[send_idx == 1],
                      )
                  )
                  if filt_data.shape[0] <= 0:
                      input_list.append(
                          torch.empty((0, 5), dtype=torch.int64)
                      )
                  else:
                      input_list.append(torch.from_numpy(filt_data))
  
              # Now send newly formed chunk to others.
              dist.barrier()
              output_list = alltoallv_cpu(
                  rank, world_size, input_list, retain_nones=False
              )
  
              # Replace the values of the edge_data, with the received data from all the other processes.
              rcvd_edge_data = torch.cat(output_list).numpy()
              local_src_ids.append(rcvd_edge_data[:, 0])
              local_dst_ids.append(rcvd_edge_data[:, 1])
              local_type_eids.append(rcvd_edge_data[:, 2])
              local_etype_ids.append(rcvd_edge_data[:, 3])
              local_eids.append(rcvd_edge_data[:, 4])
  
          edge_data[
              constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
          ] = np.concatenate(local_src_ids)
          edge_data[
              constants.GLOBAL_DST_ID + "/" + str(local_part_id)
          ] = np.concatenate(local_dst_ids)
          edge_data[
              constants.GLOBAL_TYPE_EID + "/" + str(local_part_id)
          ] = np.concatenate(local_type_eids)
          edge_data[
              constants.ETYPE_ID + "/" + str(local_part_id)
          ] = np.concatenate(local_etype_ids)
          edge_data[
              constants.GLOBAL_EID + "/" + str(local_part_id)
          ] = np.concatenate(local_eids)
  
      # Check if the data was exchanged correctly
      local_edge_count = 0
      for local_part_id in range(num_parts // world_size):
          local_edge_count += edge_data[
              constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
          ].shape[0]
      shuffle_edge_counts = allgather_sizes(
          [local_edge_count], world_size, num_parts, return_sizes=True
      )
      shuffle_edge_total = np.sum(shuffle_edge_counts)
      assert shuffle_edge_total == all_edges
  
      timer_end = timer()
      logging.info(
          f"[Rank: {rank}] Time to send/rcv edge data: {timedelta(seconds=timer_end-timer_start)}"
      )
  
      # Clean up.
      edge_data.pop(constants.GLOBAL_SRC_ID)
      edge_data.pop(constants.GLOBAL_DST_ID)
      edge_data.pop(constants.GLOBAL_TYPE_EID)
      edge_data.pop(constants.ETYPE_ID)
      edge_data.pop(constants.GLOBAL_EID)
  
      return edge_data
  
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 2 for calling a global function with example usage: exchange_edge_data(rank, world_size, num_parts, edge_data, id_lookup) and returns: edge_data