text: import logging
  import datetime
  import time
  import os
  import shutil
  
  import numpy
  import ray
  from ray import ray_constants
  import pygloo
  
  from alpa.collective.collective_group import gloo_util
  from alpa.collective.collective_group.base_collective_group import BaseGroup
  from alpa.collective.types import (AllReduceOptions, BarrierOptions, Backend,
                                     ReduceOptions, BroadcastOptions,
                                     AllGatherOptions, ReduceScatterOptions,
                                     SendOptions, RecvOptions)
  from alpa.collective.const import get_store_name
  
  logger = logging.getLogger(__name__)
  
  
  class Rendezvous:
      """A rendezvous class for different actor/task processes to meet.
  
      To initialize an GLOO collective communication group, different
      actors/tasks spawned in Ray in a collective group needs to meet
      each other to synchronize the GLOOUniqueID. This class guarantees
      they meet via the GLOOUniqueIDStore, initialized on the rank=0
      process.
  
      Args:
          group_name (str): the unique user-specified group name.
      """
  
      def __init__(self, group_name, context, store_type, device_type):
          self._group_name = group_name
          self._context = context
          self._redis_ip_address, self._redis_port = (
              ray.worker._global_node.redis_address.split(":"))
          self._process_ip_address = (ray.util.get_node_ip_address())
          logger.debug(
              f"Redis address: {self._redis_ip_address}, port: {self._redis_port}, "
              f"this actor address: {self._process_ip_address}.")
          self._store_type = store_type
          self._device_type = device_type
          self._store = None
          self._device = None
          self.create_store(store_type)
          self.create_device(device_type)
  
      def create_store(self, store_type):
          if store_type == "redis":
              redis_store = pygloo.rendezvous.RedisStore(self._redis_ip_address,
                                                         int(self._redis_port))
              redis_password = ray_constants.REDIS_DEFAULT_PASSWORD
              redis_store.authorize(redis_password)
              self._store = redis_store
          elif store_type == "file":
              store_name = get_store_name(self._group_name)
              store_path = gloo_util.get_gloo_store_path(store_name)
              if self._context.rank == 0:
                  if not os.path.exists(store_path):
                      os.makedirs(store_path)
                  elif os.listdir(store_path) and os.listdir(store_path):
                      shutil.rmtree(store_path)
                      os.makedirs(store_path)
              else:
                  while not os.path.exists(store_path):
                      time.sleep(0.1)
              # Note: multi-machines needs a shared NFS.
              file_store = pygloo.rendezvous.FileStore(store_path)
              self._store = pygloo.rendezvous.PrefixStore(self._group_name,
                                                          file_store)
          elif store_type == "hash":
              raise NotImplementedError("No implementation for hash store.")
          else:
              raise RuntimeError(
                  f"Unrecognized store type: {store_type}.")
  
      def create_device(self, device_type):
          if device_type == "tcp":
              attr = pygloo.transport.tcp.attr(self._process_ip_address)
              self._device = pygloo.transport.tcp.CreateDevice(attr)
          elif device_type == "uv":
              raise NotImplementedError("No implementation for uv.")
  
      def meet(self, timeout_s=180):
          """Meet at the named actor store.
  
          Args:
              timeout_s (int): timeout in seconds.
  
          Return:
              None
          """
          if timeout_s <= 0:
              raise ValueError("The 'timeout' argument must be positive. "
                               f"Got '{timeout_s}'.")
  
          timeout_delta = datetime.timedelta(seconds=timeout_s)
          elapsed = datetime.timedelta(seconds=0)
          start_time = datetime.datetime.now()
          q, s = None, None
  
          if self._store_type == "redis":
              while elapsed < timeout_delta:
                  try:
                      q = ray.get_actor("gloo_queue")
                      s = ray.get_actor(f"gloo_{self._group_name}_signal")
                      break
                  except ValueError:
                      if self._context.rank == 0:
                          if not q:
                              ray.remote(gloo_util.GlooQueue).options(
                                  name="gloo_queue",
                                  lifetime="detached").remote(1000)
                          if not s:
                              gloo_util.SignalActor.options(
                                  name=f"gloo_{self._group_name}_signal",
                                  lifetime="detached").remote(self._context.size)
                      else:
                          time.sleep(0.1)
                  elapsed = datetime.datetime.now() - start_time
              if not q:
                  raise RuntimeError("Unable to get gloo_queue.")
              if self._context.rank == 0:
                  ray.get(q.put_nowait.remote(self._group_name))
              while (ray.get(q.index.remote(self._group_name))):
                  time.sleep(0.1)
              self._context.connectFullMesh(self._store, self._device)
              ray.get(s.send.remote(self._context.rank))
              if self._context.rank == 0:
                  ray.get(s.wait.remote())
                  keys = []
                  keys += [f"rank_{i}" for i in range(self._context.size)]
                  keys += [f"{i}" for i in range(self._context.size)]
                  self._store.delKeys(keys)
                  group_name = ray.get(q.get_nowait.remote())
                  assert group_name == self._group_name
                  ray.kill(s)
  
      @property
      def store_type(self):
          return self._store_type
  
      @property
      def store(self):
          return self._store
  
      @property
      def device_type(self):
          return self._device_type
  
      @property
      def device(self):
          return self._device
  
      def destroy(self):
          """GC the store and device used by this rendevzous."""
          self._device = None
  
  
  class GLOOGroup(BaseGroup):
  
      def __init__(self,
                   world_size,
                   rank,
                   group_name,
                   store_type="redis",
                   device_type="tcp"):
          """Init an GLOO collective group.
  
          Args:
              world_size (int): The number of processes.
              rank (int): The id of process
              group_name (str): The unique user-specified group name.
              store_type (str): The store type. Optional: "redis",
                                "file", "hash".
              device_type (str): The device type to transport.
                                 Optional: "tcp", "uv".
          """
          super().__init__(world_size, rank, group_name)
          self._gloo_context = gloo_util.create_gloo_context(
              self.rank, self.world_size)
          self._rendezvous = Rendezvous(self.group_name, self._gloo_context,
                                        store_type, device_type)
          self._rendezvous.meet()
  
      def destroy_group(self):
          """Destroy the group and release GLOO communicators."""
          self._rendezvous.destroy()
  
          if self._gloo_context is not None:
              pygloo.barrier(self._gloo_context)
              # destroy the communicator
              self._gloo_context = None
  
          if self.rank == 0 and self._rendezvous.store_type == "file":
              store_name = get_store_name(self._group_name)
              store_path = gloo_util.get_gloo_store_path(store_name)
              if os.path.exists(store_path):
                  shutil.rmtree(store_path)
  
      @classmethod
      def backend(cls):
          return Backend.GLOO
  
      def allreduce(self, tensors, allreduce_options=AllReduceOptions()):
          """AllReduce a list of tensors following options.
  
          Args:
              tensor: the tensor to be reduced, each tensor locates on CPU
              allreduce_options:
  
          Returns:
              None
          """
  
          def collective_fn(input_tensor, output_tensor, context):
              pygloo.allreduce(
                  context, gloo_util.get_tensor_ptr(input_tensor),
                  gloo_util.get_tensor_ptr(output_tensor),
                  gloo_util.get_tensor_n_elements(input_tensor),
                  gloo_util.get_gloo_tensor_dtype(input_tensor),
                  gloo_util.get_gloo_reduce_op(allreduce_options.reduce_op))
  
          self._collective(tensors, tensors, collective_fn)
  
      def barrier(self, barrier_options=BarrierOptions()):
          """Blocks until all processes reach this barrier.
  
          Args:
              barrier_options: barrier options.
  
          Returns:
              None
          """
          barrier_tensor = numpy.array([1])
          self.allreduce([barrier_tensor])
  
      def reduce(self, tensors, reduce_options=ReduceOptions()):
          """Reduce tensors following options.
  
          Args:
              tensors (List): the list of tensors to be reduced,
                              this list only have one tensor.
              reduce_options: reduce options.
  
          Returns:
              None
          """
          root_rank = reduce_options.root_rank
  
          def collective_fn(input_tensor, output_tensor, context):
              pygloo.reduce(context, gloo_util.get_tensor_ptr(input_tensor),
                            gloo_util.get_tensor_ptr(output_tensor),
                            gloo_util.get_tensor_n_elements(input_tensor),
                            gloo_util.get_gloo_tensor_dtype(input_tensor),
                            gloo_util.get_gloo_reduce_op(reduce_options.reduce_op),
                            root_rank)
  
          self._collective(tensors, tensors, collective_fn)
  
      def broadcast(self, tensors, broadcast_options=BroadcastOptions()):
          """Broadcast tensors to all other processes following options.
  
          Args:
              tensors (List): tensors to be broadcast or received.
              broadcast_options: broadcast options.
  
          Returns:
              None
          """
          root_rank = broadcast_options.root_rank
  
          def collective_fn(input_tensor, output_tensor, context):
              pygloo.broadcast(context, gloo_util.get_tensor_ptr(input_tensor),
                               gloo_util.get_tensor_ptr(output_tensor),
                               gloo_util.get_tensor_n_elements(input_tensor),
                               gloo_util.get_gloo_tensor_dtype(input_tensor),
                               root_rank)
  
          self._collective(tensors, tensors, collective_fn)
  
      def allgather(self,
                    tensor_lists,
                    tensors,
                    allgather_options=AllGatherOptions()):
          """Allgather tensors on CPU into a list of tensors.
  
          Args:
              tensor_lists (List[List[Tensor]]): allgathered tensors.
              tensors: the list of tensors to allgather across the group.
                       Each tensor must locate on CPU.
              allgather_options: allgather options.
  
          Returns:
              None
          """
  
          def collective_fn(input_tensor, output_tensor, context):
              pygloo.allgather(context, gloo_util.get_tensor_ptr(input_tensor),
                               gloo_util.get_tensor_ptr(output_tensor),
                               gloo_util.get_tensor_n_elements(input_tensor),
                               gloo_util.get_gloo_tensor_dtype(input_tensor))
  
          _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
          output_flattened = [
              _flatten_for_scatter_gather(tensor_list, copy=False)
              for tensor_list in tensor_lists
          ]
  
          def postprocess_fn():
              for i, tensor_list in enumerate(tensor_lists):
                  for j, tensor in enumerate(tensor_list):
                      gloo_util.copy_tensor(tensor, output_flattened[i][j])
  
          self._collective(tensors,
                           output_flattened,
                           collective_fn,
                           postprocess_fn=postprocess_fn)
  
      def reducescatter(self,
                        tensors,
                        tensor_lists,
                        reducescatter_options=ReduceScatterOptions()):
          """Reduce the scatter a list of tensors across the group.
  
          Args:
              tensors (List): the output tensors (could be unspecified), each
                              located on CPU.
              tensor_lists (List[List]): the list of tensors to be reduced then
                                         scattered.
              reducescatter_options: reduce-scatter options.
  
          Returns:
              None
          """
  
          def collective_fn(input_tensor, output_tensor, context):
              size = gloo_util.get_tensor_n_elements(input_tensor)
              world_size = self._gloo_context.size
              pygloo.reduce_scatter(
                  context, gloo_util.get_tensor_ptr(input_tensor),
                  gloo_util.get_tensor_ptr(output_tensor), size,
                  [size // world_size for _ in range(world_size)],
                  gloo_util.get_gloo_tensor_dtype(output_tensor),
                  gloo_util.get_gloo_reduce_op(reducescatter_options.reduce_op))
  
          _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists)
          input_flattened = [
              _flatten_for_scatter_gather(tensor_list, copy=False)
              for tensor_list in tensor_lists
          ]
  
          def preprocess_fn():
              for i, tensor_list in enumerate(tensor_lists):
                  for j, tensor in enumerate(tensor_list):
                      gloo_util.copy_tensor(input_flattened[i][j], tensor)
  
          self._collective(input_flattened,
                           tensors,
                           collective_fn,
                           preprocess_fn=preprocess_fn)
  
      def send(self, tensors, send_options=SendOptions()):
          """Send a tensor to a destination rank in the group.
  
          Args:
              tensors (List): the tensor to send.
              send_options: send options.
  
          Returns:
              None
          """
  
          def p2p_fn(tensor, context, peer):
              pygloo.send(context, gloo_util.get_tensor_ptr(tensor),
                          gloo_util.get_tensor_n_elements(tensor),
                          gloo_util.get_gloo_tensor_dtype(tensor), peer)
  
          self._point2point(tensors, p2p_fn, send_options.dst_rank)
  
      def recv(self, tensors, recv_options=RecvOptions()):
          """Receive a tensor from a source rank in the group.
  
          Args:
              tensors (List): the received tensor.
              recv_options: Receive options.
  
          Returns:
              None
          """
  
          def p2p_fn(tensor, context, peer):
              pygloo.recv(context, gloo_util.get_tensor_ptr(tensor),
                          gloo_util.get_tensor_n_elements(tensor),
                          gloo_util.get_gloo_tensor_dtype(tensor), peer)
  
          self._point2point(tensors, p2p_fn, recv_options.src_rank)
  
      def _collective(self,
                      input_tensors,
                      output_tensors,
                      collective_fn,
                      preprocess_fn=None,
                      postprocess_fn=None):
          """A method to encapsulate all collective calls.
  
          Args:
              input_tensors: the list of the input tensors.
              output_tensors: the list of the output tensors.
              collective_fn: the collective function call.
              preprocess_fn: preprocess procedures before collective calls.
              postprocess_fn: postprocess procedures after collective calls.
  
          Returns:
              None
          """
          _check_cpu_tensors(input_tensors)
          _check_cpu_tensors(output_tensors)
  
          if preprocess_fn:
              preprocess_fn()
          collective_fn(input_tensors[0], output_tensors[0], self._gloo_context)
          if postprocess_fn:
              postprocess_fn()
  
      def _point2point(self, tensors, p2p_fn, peer_rank: int):
          """A method to encapsulate all peer-to-peer calls (i.e., send/recv).
  
          Args:
              tensors: the tensor to send or receive.
              p2p_fn: the p2p function call.
              peer_rank (int): the rank of the peer process.
  
          Returns:
              None
          """
          _check_cpu_tensors(tensors)
  
          p2p_fn(tensors[0], self._gloo_context, peer_rank)
  
  
  def _check_cpu_tensors(tensors):
      """Check only have one tensor and located on CPU."""
      if not tensors or not isinstance(tensors, list):
          raise RuntimeError("'tensors' must be a nonempty list.")
      if len(tensors) != 1:
          raise RuntimeError("Gloo only accept one tensor in the tensor list."
                             f" Got {len(tensors)} != 1.")
      d = gloo_util.get_tensor_device(tensors[0])
      if d != "cpu":
          raise RuntimeError("Gloo only accept cpu tensor."
                             f" Got {d}.")
  
  
  def _flatten_for_scatter_gather(tensor_list, copy=False):
      """Flatten the tensor for gather/scatter operations.
  
      Args:
          tensor_list: the list of tensors to be scattered/gathered.
          copy: whether the copy the tensors in tensor_list into the buffer.
  
      Returns:
          The flattened tensor buffer.
      """
      if not tensor_list:
          raise RuntimeError("Received an empty list.")
  
      t = tensor_list[0]
      # note we need a numpy dtype here.
      dtype = gloo_util.get_numpy_tensor_dtype(t)
      buffer_shape = [len(tensor_list)] + gloo_util.get_tensor_shape(t)
  
      buffer = numpy.empty(buffer_shape, dtype=dtype)
      if copy:
          for i, tensor in enumerate(tensor_list):
              gloo_util.copy_tensor(buffer[i], tensor)
      return buffer
  
  
  def _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists):
      """Check the compatibility between tensor input and tensor list input."""
      if not tensors or not isinstance(tensors, list):
          raise RuntimeError(
              "The first argument 'tensors' expects a list of tensors.")
  
      if len(tensors) != 1:
          raise RuntimeError(
              "Gloo only accept one tensor in the first argument 'tensors'."
              f" Got {len(tensors)} != 1.")
  
      if not tensor_lists or not isinstance(tensor_lists, list):
          raise RuntimeError("The second argument 'tensor_lists' "
                             "expects a list of tensor list.")
  
      if len(tensor_lists) != 1:
          raise RuntimeError("Gloo only accept one tensor list "
                             "in the second argument 'tensor_lists'."
                             f" Got {len(tensor_lists)} != 1.")
  
      dtype = gloo_util.get_gloo_tensor_dtype(tensors[0])
      shape = gloo_util.get_tensor_shape(tensors[0])
  
      # check all tensors in `tensor_lists` match.
      for t in tensor_lists[0]:
          # check dtype
          dt = gloo_util.get_gloo_tensor_dtype(t)
          if dt != dtype:
              raise RuntimeError("All tensor operands to scatter/gather must "
                                 f"have the same dtype. Got '{dt}' and '{dtype}'.")
          s = gloo_util.get_tensor_shape(t)
          if s != shape:
              raise RuntimeError("All tensor operands to scatter/gather must "
                                 f"have the same shape. Got '{s}' and '{shape}'.")
  