text: # Copyright 2020 Petuum. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Graph Transformer."""
  
  from tensorflow.core.framework.attr_value_pb2 import AttrValue as pb2_AttrValue
  from tensorflow.python.eager import context
  from tensorflow.python.framework import device_spec
  
  from autodist.graph_item import GraphItem
  from autodist.kernel.partitioner import VariablePartitioner
  from autodist.kernel.replicator import Replicator
  from autodist.kernel.synchronization.synchronizer import Synchronizer
  from autodist.utils import logging, visualization_util
  
  #import autodist.aggregators
  import aggregathor.aggregators 
  
  class GraphTransformer:
      """
      Graph Transformer.
  
      This is the bulk of the AutoDist backend logic, taking a single-node,
      single-GPU graph and transforming it into a distributed
      graph. This all happens based on the `Strategy` provided.
  
      The transformation occurs over several steps:
  
      1. Partitions the necessary variables
      2. Replicates the graph the desired number of times
      3. Within the graph, synchronizes gradients with in-graph logic
      4. Adds between-graph gradient synchronization logic
  
      """
  
      def __init__(self, compiled_strategy, cluster, graph_item):
          self._strategy = compiled_strategy
          self._cluster = cluster
          self.graph_item = graph_item
  
          # Set in _initialize_synchronizers
          self._num_local_replicas = 0
          self._num_workers = 0
          self._synchronizers = {}
  
      def transform(self):
          """Call graph transformer to transform a graph item based on strategy and cluster."""
          logging.info('Transforming the original graph to a distributed graph...')
          with context.graph_mode():
              graph_item = self.graph_item
              # Ensure the transformation happens under graph mode, no matter the outer mode is under eager or graph.
  
              visualization_util.log_graph(graph=graph_item.graph, name='0-original')
  
              graph_item, self._strategy.node_config = VariablePartitioner.apply(self._strategy.node_config, graph_item)
  
              visualization_util.log_graph(graph=graph_item.graph, name='1-after-partition')
  
              # Create Synchronizers for each node in the strategy
              self._initialize_synchronizers()
  
              # Replicate the graph (both in-graph and between-graph)
              new_graph_item = Replicator.apply(
                  config=self._strategy.graph_config.replicas,
                  cluster=self._cluster,
                  graph_item=graph_item
              )
  
              # Apply synchronizers
              if self._num_local_replicas >= 1:
                  new_graph_item = self._in_graph_apply(new_graph_item)
                  logging.debug('Successfully applied local in-graph replication')
                  visualization_util.log_graph(new_graph_item.graph, '2-after-in-graph')
  
              if self._num_workers >= 1:
                  new_graph_item = self._between_graph_apply(new_graph_item)
                  logging.debug('Successfully applied between-graph replication')
  
              final_item = new_graph_item
              logging.info('Successfully built the distributed graph.')
              visualization_util.log_graph(graph=final_item.graph, name='3-transformed')
  
          return final_item
  
      def _initialize_synchronizers(self):
          self._synchronizers = {}
          for node in self._strategy.node_config:
              partitioner = getattr(node, 'partitioner')
              if partitioner:
                  for part in node.part_config:
                      self._synchronizers[part.var_name] = \
                          Synchronizer.create(part.WhichOneof('synchronizer'),
                                              getattr(part, part.WhichOneof('synchronizer')))
              else:
                  self._synchronizers[node.var_name] = \
                      Synchronizer.create(node.WhichOneof('synchronizer'),
                                          getattr(node, node.WhichOneof('synchronizer')))
  
          config = self._strategy.graph_config.replicas
          replica_devices = {device_spec.DeviceSpecV2.from_string(s) for s in config}
          replica_hosts = {self._cluster.get_address_from_task(d.job, d.task) for d in replica_devices}
          self._num_workers = len(replica_hosts)
  
          local_canonical_replica_devices = sorted({
              d.to_string() for d in replica_devices
              if self._cluster.get_local_address() == self._cluster.get_address_from_task(d.job, d.task)
          })
          logging.debug('Local replica devices: {}'.format(local_canonical_replica_devices))
          self._num_local_replicas = len(local_canonical_replica_devices)
  
          local_worker_id = self._cluster.get_local_worker_task_index()
          local_worker_device = '/job:worker/task:{}'.format(local_worker_id)
  
          for synchronizer in self._synchronizers.values():
              synchronizer.assign_cluster_information(
                  num_workers=self._num_workers,
                  num_replicas=self._num_local_replicas,
                  worker_device=local_worker_device,
                  worker_id=local_worker_id,
                  canonical_replica_devices=sorted({d.to_string() for d in replica_devices}),
                  is_chief=self._cluster.is_chief())
  
      @property
      def num_local_replicas(self):
          """Return the number of local replicas."""
          assert self._num_local_replicas != 0  # ensure initialized
          return self._num_local_replicas
  
      def _in_graph_apply(self, graph_item: GraphItem):
          """
          Perform in-graph synchronization of the graph.
  
          Args:
              graph_item (GraphItem): The graph to replication.
  
          Returns:
              GraphItem
          """
          new_graph_item = graph_item
          BFTaggregator = aggregathor.aggregators.instantiate("average", 7, 0, []) # median, krum-co, krum-tf(2.1.0), average, bulyan-co
          for var_name, syncer in self._synchronizers.items():
              new_graph_item = syncer.in_graph_apply(new_graph_item, var_name, BFTaggregator)
          return new_graph_item
  
      def _between_graph_apply(self, multi_gpu_graph_item: GraphItem):
          """
          Perform between-graph replication of the graph.
  
          Args:
              multi_gpu_graph_item (GraphItem): The graph to replication.
  
          Returns:
              GraphItem
          """
          new_graph_item = multi_gpu_graph_item
          for var_name, syncer in self._synchronizers.items():
              new_graph_item = syncer.between_graph_apply(new_graph_item, var_name)
          self._prune_colocation_groups(new_graph_item)
          # TODO: make this work
          # update_shard_values_for_worker(num_workers, worker_id)
          return new_graph_item
  
      # TODO(Hao): this seems still problematic
      @staticmethod
      def _prune_colocation_groups(graph_item):
          for op in graph_item.graph.get_operations():
              # Now prune the graph to have the right colocation constraints
              colocation_groups = [(c, graph_item.get_colocation_op(c)) for c in op.colocation_groups()]
              # We don't want any colocation groups that are just this `op`
              colocation_groups = [(c, bind_op) for (c, bind_op) in colocation_groups if bind_op != op]
              if colocation_groups:
                  device_to_bind_to = colocation_groups[-1][1].device
                  new_colocation_groups = [c for (c, op) in colocation_groups if op.device == device_to_bind_to]
                  op._set_device(device_to_bind_to)
                  op._set_attr("_class", pb2_AttrValue(list=pb2_AttrValue.ListValue(s=new_colocation_groups)))
              else:
                  try:
                      if op.get_attr("_class"):
                          op._clear_attr("_class")
                  except ValueError:
                      pass
  