text: import copy
  from argparse import Namespace
  from itertools import cycle
  from typing import Optional, Dict, List, Union, Set
  
  from .. import BasePod
  from .. import Pea
  from .. import Pod
  from ..networking import get_connect_host
  from ... import helper
  from ...enums import SocketType, SchedulerType
  from ...helper import random_identity
  
  
  class CompoundPod(BasePod):
      """A CompoundPod is a immutable set of pods, which run in parallel.
      A CompoundPod is an abstraction using a composable pattern to abstract the usage of parallel Pods that act as shards.
  
      CompoundPod will make sure to add a `HeadPea` and a `TailPea` to serve as routing/merging pattern for the different Pod shards
  
      :param args: pod arguments parsed from the CLI. These arguments will be used for each of the shards
      :param needs: pod names of preceding pods, the output of these pods are going into the input of this pod
      """
  
      head_args = None
      tail_args = None
  
      def __init__(
          self, args: Union['Namespace', Dict], needs: Optional[Set[str]] = None
      ):
          super().__init__()
          args.upload_files = BasePod._set_upload_files(args)
          self.args = args
          self.needs = (
              needs or set()
          )  #: used in the :class:`jina.flow.Flow` to build the graph
          # we will see how to have `CompoundPods` in remote later when we add tests for it
          self.is_head_router = True
          self.is_tail_router = True
          self.head_args = BasePod._copy_to_head_args(args, args.polling)
          self.tail_args = BasePod._copy_to_tail_args(self.args, self.args.polling)
          # uses before with shards apply to shards and not to replicas
          self.shards = []  # type: List['Pod']
          # BACKWARDS COMPATIBILITY:
          self.args.parallel = self.args.shards
          self.assign_shards()
  
      def assign_shards(self):
          """Assign shards to the CompoundPod"""
          self.shards.clear()
          cargs = copy.copy(self.args)
          cargs.uses_before = None
          cargs.uses_after = None
          self.shards_args = self._set_shard_args(cargs, self.head_args, self.tail_args)
  
          for _args in self.shards_args:
              if getattr(self.args, 'noblock_on_start', False):
                  _args.noblock_on_start = True
              self.shards.append(Pod(_args))
  
      def __exit__(self, exc_type, exc_val, exc_tb) -> None:
          super().__exit__(exc_type, exc_val, exc_tb)
          self.join()
  
      @property
      def port_jinad(self) -> int:
          """Get the JinaD remote port
  
          .. # noqa: DAR201
          """
          return self.head_args.port_jinad
  
      @property
      def host(self) -> str:
          """Get the host name of this Pod
  
          .. # noqa: DAR201
          """
          return self.head_args.host
  
      def _parse_pod_args(self, args: Namespace) -> List[Namespace]:
          return self._set_shard_args(
              args,
              head_args=self.head_args,
              tail_args=self.tail_args,
          )
  
      @property
      def num_peas(self) -> int:
          """
          Get the number of running :class:`Pod`
  
          :return: total number of peas including head and tail
          """
          return sum([shard.num_peas for shard in self.shards]) + 2
  
      def __eq__(self, other: 'CompoundPod'):
          return self.num_peas == other.num_peas and self.name == other.name
  
      def _enter_pea(self, pea: 'Pea') -> None:
          self.enter_context(pea)
  
      def start(self) -> 'CompoundPod':
          """
          Start to run all :class:`Pod` and :class:`Pea` in this CompoundPod.
  
          :return: started CompoundPod
  
          .. note::
              If one of the :class:`Pod` fails to start, make sure that all of them
              are properly closed.
          """
          if getattr(self.args, 'noblock_on_start', False):
              head_args = self.head_args
              head_args.noblock_on_start = True
              self.head_pea = Pea(head_args)
              self._enter_pea(self.head_pea)
              for shard in self.shards:
                  self._enter_shard(shard)
              tail_args = self.tail_args
              tail_args.noblock_on_start = True
              self.tail_pea = Pea(tail_args)
              self._enter_pea(self.tail_pea)
          else:
              try:
                  head_args = self.head_args
                  self.head_pea = Pea(head_args)
                  self._enter_pea(self.head_pea)
                  for shard in self.shards:
                      self._enter_shard(shard)
                      shard.activate()
                  tail_args = self.tail_args
                  self.tail_pea = Pea(tail_args)
                  self._enter_pea(self.tail_pea)
              except:
                  self.close()
                  raise
              return self
  
      def wait_start_success(self) -> None:
          """
          Block until all pods and peas start successfully.
          If not successful, it will raise an error hoping the outer function to catch it
          """
  
          if not self.args.noblock_on_start:
              raise ValueError(
                  f'{self.wait_start_success!r} should only be called when `noblock_on_start` is set to True'
              )
  
          try:
              self.head_pea.wait_start_success()
              self.tail_pea.wait_start_success()
              for p in self.shards:
                  p.wait_start_success()
          except:
              self.close()
              raise
  
      def _enter_shard(self, shard: 'Pod') -> None:
          self.enter_context(shard)
  
      def join(self):
          """Wait until all pods and peas exit."""
          try:
              if getattr(self, 'head_pea', None):
                  self.head_pea.join()
              if getattr(self, 'tail_pea', None):
                  self.tail_pea.join()
              for p in self.shards:
                  p.join()
          except KeyboardInterrupt:
              pass
  
      @property
      def is_ready(self) -> bool:
          """
          Checks if Pod is read.
          :return: true if the peas and pods are ready to serve requests
  
          .. note::
              A Pod is ready when all the Peas it contains are ready
          """
          return all(
              [p.is_ready.is_set() for p in [self.head_pea, self.tail_pea]]
              + [p.is_ready for p in self.shards]
          )
  
      @staticmethod
      def _set_shard_args(
          args: Namespace,
          head_args: Namespace,
          tail_args: Namespace,
      ) -> List[Namespace]:
          """
          Sets the arguments of the shards in the compound pod.
  
          :param args: arguments configured by the user for the shards
          :param head_args: head args from the compound pod
          :param tail_args: tail args from the compound pod
  
          :return: list of arguments for the shards
          """
          result = []
          _host_list = (
              args.peas_hosts
              if args.peas_hosts
              else [
                  args.host,
              ]
          )
          host_generator = cycle(_host_list)
          for idx in range(args.shards):
              _args = copy.deepcopy(args)
              pod_host_list = [
                  host for _, host in zip(range(args.replicas), host_generator)
              ]
              _args.peas_hosts = pod_host_list
              _args.shard_id = idx
              # BACKWARDS COMPATIBILITY:
              _args.pea_id = _args.shard_id
              _args.identity = random_identity()
              if _args.name:
                  _args.name += f'/shard-{idx}'
              else:
                  _args.name = f'{idx}'
  
              _args.port_in = head_args.port_out
              _args.port_out = tail_args.port_in
              _args.port_ctrl = helper.random_port()
  
              if args.polling.is_push:
                  if args.scheduling == SchedulerType.ROUND_ROBIN:
                      _args.socket_in = SocketType.PULL_CONNECT
                  elif args.scheduling == SchedulerType.LOAD_BALANCE:
                      _args.socket_in = SocketType.DEALER_CONNECT
                  else:
                      raise ValueError(
                          f'{args.scheduling} is not supported as a SchedulerType!'
                      )
              else:
                  _args.socket_in = SocketType.SUB_CONNECT
  
              _args.socket_out = SocketType.PUSH_CONNECT
  
              _args.dynamic_routing = False
              # ugly trick to avoid Head of shard to have wrong host in
              tmp_args = copy.deepcopy(_args)
              if _args.shards > 1:
                  tmp_args.runs_in_docker = False
                  tmp_args.uses = ''
  
              _args.host_in = get_connect_host(
                  bind_host=head_args.host,
                  bind_expose_public=head_args.expose_public,
                  connect_args=tmp_args,
              )
              _args.host_out = get_connect_host(
                  bind_host=tail_args.host,
                  bind_expose_public=tail_args.expose_public,
                  connect_args=tmp_args,
              )
              result.append(_args)
          return result
  
      async def rolling_update(
          self, dump_path: Optional[str] = None, *, uses_with: Optional[Dict] = None
      ):
          """Reload all Pods of this Compound Pod.
          :param dump_path: **backwards compatibility** This function was only accepting dump_path as the only potential arg to override
          :param uses_with: a Dictionary of arguments to restart the executor with
          """
          tasks = []
          try:
              import asyncio
  
              tasks = [
                  asyncio.create_task(
                      shard.rolling_update(dump_path=dump_path, uses_with=uses_with)
                  )
                  for shard in self.shards
              ]
              for future in asyncio.as_completed(tasks):
                  _ = await future
          except:
              for task in tasks:
                  if not task.done():
                      task.cancel()
  
      @property
      def _mermaid_str(self) -> List[str]:
          """String that will be used to represent the Pod graphically when `Flow.plot()` is invoked
  
  
          .. # noqa: DAR201
          """
          mermaid_graph = [f'subgraph {self.name};\n', f'direction LR;\n']
          head_name = self.head_args.name
          tail_name = self.tail_args.name
          pod_names = []
          for shard in self.shards:
              pod_names.append(shard.name)
              shard_mermaid_graph = shard._mermaid_str
              shard_mermaid_graph = [
                  node.replace(';', '\n') for node in shard_mermaid_graph
              ]
              mermaid_graph.extend(shard_mermaid_graph)
              mermaid_graph.append('\n')
  
          for name in pod_names:
              mermaid_graph.append(f'{head_name}:::HEADTAIL --> {name};')
              mermaid_graph.append(f'{name} --> {tail_name}:::HEADTAIL;')
          mermaid_graph.append('end;')
  
          return mermaid_graph
  