<commit_before># Copyright 2019, The TensorFlow Federated Authors.
#
# 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.
"""A strategy for composing federated types and intrinsics in disjoint scopes.

                +------------+
                | Federating |   +----------+
            +-->+ Executor   +-->+ unplaced |
            |   +--+---------+   | executor |
            |      |             +----------+
  +---------+--+   |
  | Federated  +<--+
  | Composing  |
  | Strategy   |
  +--+--+------+
     |  |
     |  |   +----------+
     |  +-->+ server   |
     |      | executor |
     |      +----------+
     |
     |   +-----------+
     +-->+ target    |
         | executors |
         +-----------+

  This strategy:

  * Delegates handling unplaced types, computations, and processing back to the
    `FederatingExecutor`.
  * Delegates handling server-side processing to the `server_executor`.
  * Delegates handling federating types and intrinsics to the target
    executor(s).
"""

import asyncio
from typing import Any, List

import tensorflow as tf

from tensorflow_federated.proto.v0 import computation_pb2 as pb
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.common_libs import structure
from tensorflow_federated.python.common_libs import tracing
from tensorflow_federated.python.core.api import computation_types
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
from tensorflow_federated.python.core.impl.executors import executor_base
from tensorflow_federated.python.core.impl.executors import executor_utils
from tensorflow_federated.python.core.impl.executors import executor_value_base
from tensorflow_federated.python.core.impl.executors import federating_executor
from tensorflow_federated.python.core.impl.types import placement_literals
from tensorflow_federated.python.core.impl.types import type_analysis


class FederatedComposingStrategyValue(executor_value_base.ExecutorValue):
  """A value embedded in a `FederatedExecutor`."""

  def __init__(self, value, type_signature):
    """Creates a `FederatedComposingStrategyValue` embedding the given `value`.

    Args:
      value: An object to embed in the executor, one of the supported types
        defined by the `FederatingExecutor`.
      type_signature: A `tff.Type` or something convertible to instance of
        `tff.Type` via `tff.to_type`, the type of `value`.
    """
    self._value = value
    self._type_signature = computation_types.to_type(type_signature)

  @property
  def internal_representation(self):
    return self._value

  @property
  def type_signature(self) -> computation_types.Type:
    return self._type_signature

  @tracing.trace
  async def compute(self):
    """Returns the result of computing the embedded value.

    Raises:
      TypeError: If the embedded value is composed of values that are not
        embedded in the executor.
      RuntimeError: If the embedded value is not a kind supported by the
        `FederatingExecutor`.
    """
    if isinstance(self._value, executor_value_base.ExecutorValue):
      return await self._value.compute()
    elif isinstance(self._value, structure.Struct):
      results = await asyncio.gather(*[
          FederatedComposingStrategyValue(v, t).compute()
          for v, t in zip(self._value, self._type_signature)
      ])
      element_types = structure.iter_elements(self._type_signature)
      return structure.Struct(
          (n, v) for (n, _), v in zip(element_types, results))
    elif isinstance(self._value, list):
      py_typecheck.check_type(self._type_signature,
                              computation_types.FederatedType)
      for value in self._value:
        py_typecheck.check_type(value, executor_value_base.ExecutorValue)
      if self._type_signature.all_equal:
        return await self._value[0].compute()
      else:
        result = []
        values = await asyncio.gather(*[v.compute() for v in self._value])
        for value in values:
          py_typecheck.check_type(value, list)
          result.extend(value)
        return result
    else:
      raise RuntimeError(
          'Computing values of type {} represented as {} is not supported in '
          'this executor.'.format(self._type_signature,
                                  py_typecheck.type_string(type(self._value))))


class FederatedComposingStrategy(federating_executor.FederatingStrategy):
  """A strategy for composing federated types and intrinsics in disjoint scopes.

  This strategy implements the `FederatingStrategy` interface and provides
  logic for handling federated types and federated intrinsics on an underlying
  collection of target executors. This strategy can be used in hierarchical
  aggregation structures with federating executors handling disjoint scopes at
  the leafs.

  This strategy only supports the placements:

  * `tff.SERVER`
  * `tff.CLIENTS`
  """

  @classmethod
  def factory(cls, server_executor: executor_base.Executor,
              target_executors: List[executor_base.Executor]):
    return lambda executor: cls(executor, server_executor, target_executors)

  def __init__(self, executor: federating_executor.FederatingExecutor,
               server_executor: executor_base.Executor,
               target_executors: List[executor_base.Executor]):
    """Creates a `FederatedComposingStrategy`.

    Args:
      executor: A `federating_executor.FederatingExecutor` to use to handle
        unplaced types, computations, and processing.
      server_executor: The parent executor to use for all processing at the
        parent, such as combining values from child executors, unplaced and
        server-side processing, etc.
      target_executors: The list of executors that manage disjoint scopes to
        combine in this executor, delegate to and collect or aggregate from.

    Raises:
      TypeError: If `server_executor` is not an `executor_base.Executor` or if
        `target_executors` is not a `list` of `executor_base.Executor`s.
    """
    super().__init__(executor)
    py_typecheck.check_type(server_executor, executor_base.Executor)
    py_typecheck.check_type(target_executors, list)
    for e in target_executors:
      py_typecheck.check_type(e, executor_base.Executor)
    self._server_executor = server_executor
    self._target_executors = target_executors
    self._cardinalities_task = None

  def close(self):
    self._server_executor.close()
    for e in self._target_executors:
      e.close()

  def ingest_value(
      self, value: Any, type_signature: computation_types.Type
  ) -> executor_value_base.ExecutorValue:
    return FederatedComposingStrategyValue(value, type_signature)

  async def _get_cardinalities(self):
    """Returns information about the number of clients in the child executors.

    Returns:
      A `list` with one element for each element in `self._target_executors`;
      each of these elements is an integer representing the total number of
      clients located in the corresponding child executor.
    """

    async def _num_clients(executor):<commit_msg>Removes stale cardinalities_task internal attribute on composing strategy.<commit_after># Copyright 2019, The TensorFlow Federated Authors.
#
# 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.
"""A strategy for composing federated types and intrinsics in disjoint scopes.

                +------------+
                | Federating |   +----------+
            +-->+ Executor   +-->+ unplaced |
            |   +--+---------+   | executor |
            |      |             +----------+
  +---------+--+   |
  | Federated  +<--+
  | Composing  |
  | Strategy   |
  +--+--+------+
     |  |
     |  |   +----------+
     |  +-->+ server   |
     |      | executor |
     |      +----------+
     |
     |   +-----------+
     +-->+ target    |
         | executors |
         +-----------+

  This strategy:

  * Delegates handling unplaced types, computations, and processing back to the
    `FederatingExecutor`.
  * Delegates handling server-side processing to the `server_executor`.
  * Delegates handling federating types and intrinsics to the target
    executor(s).
"""

import asyncio
from typing import Any, List

import tensorflow as tf

from tensorflow_federated.proto.v0 import computation_pb2 as pb
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.common_libs import structure
from tensorflow_federated.python.common_libs import tracing
from tensorflow_federated.python.core.api import computation_types
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
from tensorflow_federated.python.core.impl.executors import executor_base
from tensorflow_federated.python.core.impl.executors import executor_utils
from tensorflow_federated.python.core.impl.executors import executor_value_base
from tensorflow_federated.python.core.impl.executors import federating_executor
from tensorflow_federated.python.core.impl.types import placement_literals
from tensorflow_federated.python.core.impl.types import type_analysis


class FederatedComposingStrategyValue(executor_value_base.ExecutorValue):
  """A value embedded in a `FederatedExecutor`."""

  def __init__(self, value, type_signature):
    """Creates a `FederatedComposingStrategyValue` embedding the given `value`.

    Args:
      value: An object to embed in the executor, one of the supported types
        defined by the `FederatingExecutor`.
      type_signature: A `tff.Type` or something convertible to instance of
        `tff.Type` via `tff.to_type`, the type of `value`.
    """
    self._value = value
    self._type_signature = computation_types.to_type(type_signature)

  @property
  def internal_representation(self):
    return self._value

  @property
  def type_signature(self) -> computation_types.Type:
    return self._type_signature

  @tracing.trace
  async def compute(self):
    """Returns the result of computing the embedded value.

    Raises:
      TypeError: If the embedded value is composed of values that are not
        embedded in the executor.
      RuntimeError: If the embedded value is not a kind supported by the
        `FederatingExecutor`.
    """
    if isinstance(self._value, executor_value_base.ExecutorValue):
      return await self._value.compute()
    elif isinstance(self._value, structure.Struct):
      results = await asyncio.gather(*[
          FederatedComposingStrategyValue(v, t).compute()
          for v, t in zip(self._value, self._type_signature)
      ])
      element_types = structure.iter_elements(self._type_signature)
      return structure.Struct(
          (n, v) for (n, _), v in zip(element_types, results))
    elif isinstance(self._value, list):
      py_typecheck.check_type(self._type_signature,
                              computation_types.FederatedType)
      for value in self._value:
        py_typecheck.check_type(value, executor_value_base.ExecutorValue)
      if self._type_signature.all_equal:
        return await self._value[0].compute()
      else:
        result = []
        values = await asyncio.gather(*[v.compute() for v in self._value])
        for value in values:
          py_typecheck.check_type(value, list)
          result.extend(value)
        return result
    else:
      raise RuntimeError(
          'Computing values of type {} represented as {} is not supported in '
          'this executor.'.format(self._type_signature,
                                  py_typecheck.type_string(type(self._value))))


class FederatedComposingStrategy(federating_executor.FederatingStrategy):
  """A strategy for composing federated types and intrinsics in disjoint scopes.

  This strategy implements the `FederatingStrategy` interface and provides
  logic for handling federated types and federated intrinsics on an underlying
  collection of target executors. This strategy can be used in hierarchical
  aggregation structures with federating executors handling disjoint scopes at
  the leafs.

  This strategy only supports the placements:

  * `tff.SERVER`
  * `tff.CLIENTS`
  """

  @classmethod
  def factory(cls, server_executor: executor_base.Executor,
              target_executors: List[executor_base.Executor]):
    return lambda executor: cls(executor, server_executor, target_executors)

  def __init__(self, executor: federating_executor.FederatingExecutor,
               server_executor: executor_base.Executor,
               target_executors: List[executor_base.Executor]):
    """Creates a `FederatedComposingStrategy`.

    Args:
      executor: A `federating_executor.FederatingExecutor` to use to handle
        unplaced types, computations, and processing.
      server_executor: The parent executor to use for all processing at the
        parent, such as combining values from child executors, unplaced and
        server-side processing, etc.
      target_executors: The list of executors that manage disjoint scopes to
        combine in this executor, delegate to and collect or aggregate from.

    Raises:
      TypeError: If `server_executor` is not an `executor_base.Executor` or if
        `target_executors` is not a `list` of `executor_base.Executor`s.
    """
    super().__init__(executor)
    py_typecheck.check_type(server_executor, executor_base.Executor)
    py_typecheck.check_type(target_executors, list)
    for e in target_executors:
      py_typecheck.check_type(e, executor_base.Executor)
    self._server_executor = server_executor
    self._target_executors = target_executors

  def close(self):
    self._server_executor.close()
    for e in self._target_executors:
      e.close()

  def ingest_value(
      self, value: Any, type_signature: computation_types.Type
  ) -> executor_value_base.ExecutorValue:
    return FederatedComposingStrategyValue(value, type_signature)

  async def _get_cardinalities(self):
    """Returns information about the number of clients in the child executors.

    Returns:
      A `list` with one element for each element in `self._target_executors`;
      each of these elements is an integer representing the total number of
      clients located in the corresponding child executor.
    """

    async def _num_clients(executor):
{'fraction_non_alphanumeric': 0.08750566746550942, 'fraction_numerical': 0.0014249627566552238, 'mean_word_length': 3.1650930671702184, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7593050', 'n_tokens_mistral': 4318, 'n_tokens_neox': 4170, 'n_words': 1280}