<commit_before># Lint as: python3
# Copyright 2020 The TensorFlow Authors. 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.
# ==============================================================================
"""Parameter server strategy V2 class.

This is currently under development and the API is subject to change.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import distribute_utils
from tensorflow.python.distribute import parameter_server_strategy
from tensorflow.python.distribute import sharded_variable
from tensorflow.python.eager import remote
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
from tensorflow.python.training.tracking import base as trackable
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export


@tf_export("distribute.experimental.ParameterServerStrategy", v1=[])
class ParameterServerStrategyV2(distribute_lib.Strategy):
  """An multi-worker tf.distribute strategy with parameter servers.

  Parameter server training is a common data-parallel method to scale up a
  machine learning model on multiple machines. A parameter server training
  cluster consists of workers and parameter servers. Variables are created on
  parameter servers and they are read and updated by workers in each step.
  By default, workers read and update these variables independently without
  synchronizing with each other. Under this configuration, it is known as
  asynchronous training.

  In TensorFlow 2, we recommend an architecture based on central coordination
  for parameter server training. Each worker and parameter server runs a
  `tf.distribute.Server`, and on top of that, a coordinator task is responsible
  for creating resources on workers and parameter servers, dispatching
  functions, and coordinating the training. The coordinator uses a
  `tf.distribute.experimental.coordinator.ClusterCoordinator` to coordinate the
  cluster, and a `tf.distribute.experimental.ParameterServerStrategy` to define
  variables on parameter servers and computation on workers.

  For the training to work, the coordinator dispatches `tf.function`s to be
  executed on remote workers. Upon receiving requests from the coordinator, a
  worker executes the `tf.function` by reading the variables from parameter
  servers, executing the ops, and updating the variables on the parameter
  servers. Each of the worker only processes the requests from the coordinator,
  and communicates with parameter servers, without direct interactions with
  other workers in the cluster.

  As a result, failures of some workers do not prevent the cluster from
  continuing the work, and this allows the cluster to train with instances that
  can be occasionally unavailable (e.g. preemptible or spot instances). The
  coordinator and parameter servers though, must be available at all times for
  the cluster to make progress.

  Note that the coordinator is not one of the training workers. Instead, it
  creates resources such as variables and datasets, dispatchs `tf.function`s,
  saves checkpoints and so on. In addition to workers, parameter servers and
  the coordinator, an optional evaluator can be run on the side that
  periodically reads the checkpoints saved by the coordinator and runs
  evaluations against each checkpoint.

  `tf.distribute.experimental.ParameterServerStrategy` has to work in
  conjunction with a `tf.distribute.experimental.coordinator.ClusterCoordinator`
  object. Standalone usage of
  `tf.distribute.experimental.ParameterServerStrategy` without central
  coordination is not supported at this time.

  __Example code for coordinator__

  Here's an example usage of the API, with a custom training loop to train a
  model. This code snippet is intended to be run on (the only) one task that
  is designated as the coordinator. Note that `cluster_resolver`,
  `variable_partitioner`, and `dataset_fn` arguments are explained in the
  following "Cluster setup", "Variable partitioning", and "Dataset preparation"
  sections.

  ```python
  # Set the environment variable to allow reporting worker and ps failure to the
  # coordinator. This a short-term workaround.
  os.environ["GRPC_FAIL_FAST"] = "use_caller"

  # Prepare a strategy to use with the cluster and variable partitioning info.
  strategy = tf.distribute.experimental.ParameterServerStrategy(
      cluster_resolver=...,
      variable_partitioner=...)
  coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
      strategy=strategy)

  # Prepare a distribute dataset that will place datasets on the workers.
  distributed_dataset = coordinator.create_per_worker_dataset(dataset_fn=...)

  with strategy.scope():
    model = ...
    optimizer, metrics = ...  # Keras optimizer/metrics are great choices
    checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
    checkpoint_manager = tf.train.CheckpointManager(
        checkpoint, checkpoint_dir, max_to_keep=2)
    # `load_checkpoint` infers initial epoch from `optimizer.iterations`.
    initial_epoch = load_checkpoint(checkpoint_manager) or 0

  @tf.function
  def worker_fn(iterator):

    def replica_fn(inputs):
      batch_data, labels = inputs
      # calculate gradient, applying gradient, metrics update etc.

    strategy.run(replica_fn, args=(next(iterator),))

  for epoch in range(initial_epoch, num_epoch):
    distributed_iterator = iter(distributed_dataset)  # Reset iterator state.
    for step in range(steps_per_epoch):

      # Asynchronously schedule the `worker_fn` to be executed on an arbitrary
      # worker. This call returns immediately.
      coordinator.schedule(worker_fn, args=(distributed_iterator,))

    # `join` blocks until all scheduled `worker_fn`s finish execution. Once it
    # returns, we can read the metrics and save checkpoints as needed.
    coordinator.join()
    logging.info('Metric result: %r', metrics.result())
    train_accuracy.reset_states()
    checkpoint_manager.save()
  ```

  __Example code for worker and parameter servers__

  In addition to the coordinator, there should be tasks designated as
  "worker" or "ps". They should run the following code to start a TensorFlow
  server, waiting for coordinator's requests:

  ```python
  # Set the environment variable to allow reporting worker and ps failure to the
  # coordinator.
  os.environ["GRPC_FAIL_FAST"] = "use_caller"

  # Provide a `tf.distribute.cluster_resolver.ClusterResolver` that serves
  # the cluster information. See below "Cluster setup" section.
  cluster_resolver = ...

  server = tf.distribute.Server(
      cluster_resolver.cluster_spec(),
      job_name=cluster_resolver.task_type,
      task_index=cluster_resolver.task_id,
      protocol="grpc")

  # Blocking the process that starts a server from exiting.
  server.join()
  ```

  __Cluster setup__

  In order for the tasks in the cluster to know other tasks' addresses,
  a `tf.distribute.cluster_resolver.ClusterResolver` is required to be used
  in coordinator, worker, and ps. The
  `tf.distribute.cluster_resolver.ClusterResolver` is responsible for providing
  the cluster information, as well as the task type and id of the current task.
  See `tf.distribute.cluster_resolver.ClusterResolver` for more information.

  If `TF_CONFIG` environment variable is set, a
  `tf.distribute.cluster_resolver.TFConfigClusterResolver` should be used as
  well. Note that for legacy reason, on some platform, "chief" is used as the
  task type for the coordinator, as the following example demonstrates. Here we
  set `TF_CONFIG` for the task designated as a parameter server (task type "ps")
  and index 1 (the second task), in a cluster with 1 chief, 2 parameter servers,
  and 3 workers. Note that the it needs to be set before the use of
  `tf.distribute.cluster_resolver.TFConfigClusterResolver`.

  Example code for cluster setup:
  ```python
  os.environ['TF_CONFIG'] = '''
  {
    "cluster": {
      "chief": ["chief.example.com:2222"],
      "ps": ["ps0.example.com:2222", "ps1.example.com:2222"],
      "worker": ["worker0.example.com:2222", "worker1.example.com:2222",
                 "worker2.example.com:2222"]
    },
    "task": {
      "type": "ps",
      "index": 1<commit_msg>PSv2: Make it explicitly in the docstring that "chief", "worker", and "ps" task type naming should be used.<commit_after># Lint as: python3
# Copyright 2020 The TensorFlow Authors. 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.
# ==============================================================================
"""Parameter server strategy V2 class.

This is currently under development and the API is subject to change.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import distribute_utils
from tensorflow.python.distribute import parameter_server_strategy
from tensorflow.python.distribute import sharded_variable
from tensorflow.python.eager import remote
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import server_lib
from tensorflow.python.training.tracking import base as trackable
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export


@tf_export("distribute.experimental.ParameterServerStrategy", v1=[])
class ParameterServerStrategyV2(distribute_lib.Strategy):
  """An multi-worker tf.distribute strategy with parameter servers.

  Parameter server training is a common data-parallel method to scale up a
  machine learning model on multiple machines. A parameter server training
  cluster consists of workers and parameter servers. Variables are created on
  parameter servers and they are read and updated by workers in each step.
  By default, workers read and update these variables independently without
  synchronizing with each other. Under this configuration, it is known as
  asynchronous training.

  In TensorFlow 2, we recommend an architecture based on central coordination
  for parameter server training. Each worker and parameter server runs a
  `tf.distribute.Server`, and on top of that, a coordinator task is responsible
  for creating resources on workers and parameter servers, dispatching
  functions, and coordinating the training. The coordinator uses a
  `tf.distribute.experimental.coordinator.ClusterCoordinator` to coordinate the
  cluster, and a `tf.distribute.experimental.ParameterServerStrategy` to define
  variables on parameter servers and computation on workers.

  For the training to work, the coordinator dispatches `tf.function`s to be
  executed on remote workers. Upon receiving requests from the coordinator, a
  worker executes the `tf.function` by reading the variables from parameter
  servers, executing the ops, and updating the variables on the parameter
  servers. Each of the worker only processes the requests from the coordinator,
  and communicates with parameter servers, without direct interactions with
  other workers in the cluster.

  As a result, failures of some workers do not prevent the cluster from
  continuing the work, and this allows the cluster to train with instances that
  can be occasionally unavailable (e.g. preemptible or spot instances). The
  coordinator and parameter servers though, must be available at all times for
  the cluster to make progress.

  Note that the coordinator is not one of the training workers. Instead, it
  creates resources such as variables and datasets, dispatchs `tf.function`s,
  saves checkpoints and so on. In addition to workers, parameter servers and
  the coordinator, an optional evaluator can be run on the side that
  periodically reads the checkpoints saved by the coordinator and runs
  evaluations against each checkpoint.

  `tf.distribute.experimental.ParameterServerStrategy` has to work in
  conjunction with a `tf.distribute.experimental.coordinator.ClusterCoordinator`
  object. Standalone usage of
  `tf.distribute.experimental.ParameterServerStrategy` without central
  coordination is not supported at this time.

  __Example code for coordinator__

  Here's an example usage of the API, with a custom training loop to train a
  model. This code snippet is intended to be run on (the only) one task that
  is designated as the coordinator. Note that `cluster_resolver`,
  `variable_partitioner`, and `dataset_fn` arguments are explained in the
  following "Cluster setup", "Variable partitioning", and "Dataset preparation"
  sections.

  ```python
  # Set the environment variable to allow reporting worker and ps failure to the
  # coordinator. This a short-term workaround.
  os.environ["GRPC_FAIL_FAST"] = "use_caller"

  # Prepare a strategy to use with the cluster and variable partitioning info.
  strategy = tf.distribute.experimental.ParameterServerStrategy(
      cluster_resolver=...,
      variable_partitioner=...)
  coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
      strategy=strategy)

  # Prepare a distribute dataset that will place datasets on the workers.
  distributed_dataset = coordinator.create_per_worker_dataset(dataset_fn=...)

  with strategy.scope():
    model = ...
    optimizer, metrics = ...  # Keras optimizer/metrics are great choices
    checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer)
    checkpoint_manager = tf.train.CheckpointManager(
        checkpoint, checkpoint_dir, max_to_keep=2)
    # `load_checkpoint` infers initial epoch from `optimizer.iterations`.
    initial_epoch = load_checkpoint(checkpoint_manager) or 0

  @tf.function
  def worker_fn(iterator):

    def replica_fn(inputs):
      batch_data, labels = inputs
      # calculate gradient, applying gradient, metrics update etc.

    strategy.run(replica_fn, args=(next(iterator),))

  for epoch in range(initial_epoch, num_epoch):
    distributed_iterator = iter(distributed_dataset)  # Reset iterator state.
    for step in range(steps_per_epoch):

      # Asynchronously schedule the `worker_fn` to be executed on an arbitrary
      # worker. This call returns immediately.
      coordinator.schedule(worker_fn, args=(distributed_iterator,))

    # `join` blocks until all scheduled `worker_fn`s finish execution. Once it
    # returns, we can read the metrics and save checkpoints as needed.
    coordinator.join()
    logging.info('Metric result: %r', metrics.result())
    train_accuracy.reset_states()
    checkpoint_manager.save()
  ```

  __Example code for worker and parameter servers__

  In addition to the coordinator, there should be tasks designated as
  "worker" or "ps". They should run the following code to start a TensorFlow
  server, waiting for coordinator's requests:

  ```python
  # Set the environment variable to allow reporting worker and ps failure to the
  # coordinator.
  os.environ["GRPC_FAIL_FAST"] = "use_caller"

  # Provide a `tf.distribute.cluster_resolver.ClusterResolver` that serves
  # the cluster information. See below "Cluster setup" section.
  cluster_resolver = ...

  server = tf.distribute.Server(
      cluster_resolver.cluster_spec(),
      job_name=cluster_resolver.task_type,
      task_index=cluster_resolver.task_id,
      protocol="grpc")

  # Blocking the process that starts a server from exiting.
  server.join()
  ```

  __Cluster setup__

  In order for the tasks in the cluster to know other tasks' addresses,
  a `tf.distribute.cluster_resolver.ClusterResolver` is required to be used
  in coordinator, worker, and ps. The
  `tf.distribute.cluster_resolver.ClusterResolver` is responsible for providing
  the cluster information, as well as the task type and id of the current task.
  See `tf.distribute.cluster_resolver.ClusterResolver` for more information.

  If `TF_CONFIG` environment variable is set, a
  `tf.distribute.cluster_resolver.TFConfigClusterResolver` should be used as
  well.

  Since there are assumptions in
  `tf.distribute.experimental.ParameterServerStrategy` around the naming of the
  task types, "chief", "ps", and "worker" should be used in the
  `tf.distribute.cluster_resolver.ClusterResolver` to refer to the coordinator,
  parameter servers, and workers, respectively.

  The following example demonstrates setting `TF_CONFIG` for the task designated
  as a parameter server (task type "ps") and index 1 (the second task), in a
  cluster with 1 chief, 2 parameter servers, and 3 workers. Note that it needs
  to be set before the use of
  `tf.distribute.cluster_resolver.TFConfigClusterResolver`.

  Example code for cluster setup:
  ```python
  os.environ['TF_CONFIG'] = '''
  {
    "cluster": {
      "chief": ["chief.example.com:2222"],
      "ps": ["ps0.example.com:2222", "ps1.example.com:2222"],
      "worker": ["worker0.example.com:2222", "worker1.example.com:2222",
                 "worker2.example.com:2222"]
    },
    "task": {
      "type": "ps",
      "index": 1
{'fraction_non_alphanumeric': 0.0732271602281287, 'fraction_numerical': 0.005326589906381147, 'mean_word_length': 4.463550852439742, 'pattern_counts': {'":': 14, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2402370', 'n_tokens_mistral': 4860, 'n_tokens_neox': 4665, 'n_words': 2009}