code: def all_gather(data):
      """
      Run all_gather on arbitrary picklable data (not necessarily tensors)
      Args:
          data: any picklable object
      Returns:
          list[data]: list of data gathered from each rank
      """
      world_size = get_world_size()
      if world_size == 1:
          return [data]
  
      # serialized to a Tensor
      buffer = pickle.dumps(data)
      storage = torch.ByteStorage.from_buffer(buffer)
      tensor = torch.ByteTensor(storage).to("cuda")
  
      # obtain Tensor size of each rank
      local_size = torch.tensor([tensor.numel()], device="cuda")
      size_list = [
          torch.tensor([0], device="cuda") for _ in range(world_size)
      ]
      dist.all_gather(size_list, local_size)
      size_list = [int(size.item()) for size in size_list]
      max_size = max(size_list)
  
      # receiving Tensor from all ranks
      # we pad the tensor because torch all_gather does not support
      # gathering tensors of different shapes
      tensor_list = []
      for _ in size_list:
          tensor_list.append(
              torch.empty((max_size,), dtype=torch.uint8, device="cuda")
          )
      if local_size != max_size:
          padding = torch.empty(
              size=(max_size - local_size,),
              dtype=torch.uint8,
              device="cuda",
          )
          tensor = torch.cat((tensor, padding), dim=0)
      dist.all_gather(tensor_list, tensor)
  
      data_list = []
      for size, tensor in zip(size_list, tensor_list):
          buffer = tensor.cpu().numpy().tobytes()[:size]
          data_list.append(pickle.loads(buffer))
  
      return data_list
  
imports: from collections import defaultdict, deque
  import datetime
  import pickle
  import time
  import errno
  import os
  from contextlib import contextmanager
  import torch
  import torch.distributed
label_desc: Use image node_id 1 for calling a global function with example usage: all_gather(data) and returns: data_list, list