chore: import upstream snapshot with attribution

tensorflow/python/distribute/README.md

@@ -0,0 +1,81 @@
+# Tensorflow Distribute Libraries
+
+## Overview
+
+tf.distribute.Strategy is a TensorFlow API to distribute training across
+multiple GPUs, multiple machines or TPUs. Using this API, users can distribute
+their existing models and training code with minimal code changes.
+
+It can be used with TensorFlow's high level APIs, like tf.keras,
+with just a couple of lines of code change. It does so by changing the
+underlying components of TensorFlow to become strategy-aware.
+This includes variables, layers, models, optimizers, metrics, summaries,
+and checkpoints.
+
+## Documentation
+
+[Distributed Training Guide](https://www.tensorflow.org/guide/distributed_training)
+
+[Distributed Training With Keras Tutorial](https://www.tensorflow.org/tutorials/distribute/keras)
+
+[Distributed Training With Custom Training Loops Tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training)
+
+[Multiworker Training With Keras Tutorial](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras)
+
+[Save and Load with Distribution Strategy](https://www.tensorflow.org/tutorials/distribute/save_and_load)
+
+## Simple Examples
+
+### Using compile fit with GPUs.
+
+```python
+# Create the strategy instance. It will automatically detect all the GPUs.
+mirrored_strategy = tf.distribute.MirroredStrategy()
+
+# Create and compile the keras model under strategy.scope()
+with mirrored_strategy.scope():
+  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
+  model.compile(loss='mse', optimizer='sgd')
+
+# Call model.fit and model.evaluate as before.
+dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10)
+model.fit(dataset, epochs=2)
+model.evaluate(dataset)
+```
+
+### Custom training loop with TPUs.
+
+```python
+# Create the strategy instance.
+tpu_strategy = tf.distribute.TPUStrategy(resolver)
+
+
+# Create the keras model under strategy.scope()
+with tpu_strategy.scope():
+  model = keras.layers.Dense(1, name="dense")
+
+# Create custom training loop body as tf.function.
+@tf.function
+def train_step(iterator):
+  def step_fn(inputs):
+    images, targets = inputs
+    with tf.GradientTape() as tape:
+      outputs = model(images)
+      loss = tf.reduce_sum(outputs - targets)
+    grads = tape.gradient(loss, model.variables)
+    return grads
+
+  return tpu_strategy.run(
+      step_fn, args=(next(iterator),))
+
+# Run the loop body once on at dataset.
+dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10
+input_iterator = iter(tpu_strategy.experimental_distribute_dataset(dataset))
+train_step(input_iterator)
+```
+
+## Testing
+
+Tests here should cover all distribution strategies to ensure feature parity.
+This can be done using the test decorators in `strategy_combinations.py`.
+

tensorflow/python/distribute/__init__.py

@@ -0,0 +1,190 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Library for running a computation across multiple devices.
+
+The intent of this library is that you can write an algorithm in a stylized way
+and it will be usable with a variety of different `tf.distribute.Strategy`
+implementations. Each descendant will implement a different strategy for
+distributing the algorithm across multiple devices/machines.  Furthermore, these
+changes can be hidden inside the specific layers and other library classes that
+need special treatment to run in a distributed setting, so that most users'
+model definition code can run unchanged. The `tf.distribute.Strategy` API works
+the same way with eager and graph execution.
+
+*Guides*
+
+* [TensorFlow v2.x](https://www.tensorflow.org/guide/distributed_training)
+* [TensorFlow
+v1.x](https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/distribute_strategy.ipynb)
+
+*Tutorials*
+
+* [Distributed Training
+Tutorials](https://www.tensorflow.org/tutorials/distribute/)
+
+  The tutorials cover how to use `tf.distribute.Strategy` to do distributed
+  training with native Keras APIs, and custom training loops.
+  They also cover how to save/load model when using `tf.distribute.Strategy`.
+
+*Glossary*
+
+* _Data parallelism_ is where we run multiple copies of the model
+  on different slices of the input data. This is in contrast to
+  _model parallelism_ where we divide up a single copy of a model
+  across multiple devices.
+  Note: we only support data parallelism for now, but
+  hope to add support for model parallelism in the future.
+* A _device_ is a CPU or accelerator (e.g. GPUs, TPUs) on some machine that
+  TensorFlow can run operations on (see e.g. `tf.device`). You may have multiple
+  devices on a single machine, or be connected to devices on multiple
+  machines. Devices used to run computations are called _worker devices_.
+  Devices used to store variables are _parameter devices_. For some strategies,
+  such as `tf.distribute.MirroredStrategy`, the worker and parameter devices
+  will be the same (see mirrored variables below). For others they will be
+  different. For example, `tf.distribute.experimental.CentralStorageStrategy`
+  puts the variables on a single device (which may be a worker device or may be
+  the CPU), and `tf.distribute.experimental.ParameterServerStrategy` puts the
+  variables on separate machines called _parameter servers_ (see below).
+* A _replica_ is one copy of the model, running on one slice of the
+  input data. Right now each replica is executed on its own
+  worker device, but once we add support for model parallelism
+  a replica may span multiple worker devices.
+* A _host_ is the CPU device on a machine with worker devices, typically
+  used for running input pipelines.
+* A _worker_ is defined to be the physical machine(s) containing the physical
+  devices (e.g. GPUs, TPUs) on which the replicated computation is executed. A
+  worker may contain one or more replicas, but contains at least one
+  replica. Typically one worker will correspond to one machine, but in the case
+  of very large models with model parallelism, one worker may span multiple
+  machines. We typically run one input pipeline per worker, feeding all the
+  replicas on that worker.
+* _Synchronous_, or more commonly _sync_, training is where the updates from
+  each replica are aggregated together before updating the model variables. This
+  is in contrast to _asynchronous_, or _async_ training, where each replica
+  updates the model variables independently. You may also have replicas
+  partitioned into groups which are in sync within each group but async between
+  groups.
+* _Parameter servers_: These are machines that hold a single copy of
+  parameters/variables, used by some strategies (right now just
+  `tf.distribute.experimental.ParameterServerStrategy`). All replicas that want
+  to operate on a variable retrieve it at the beginning of a step and send an
+  update to be applied at the end of the step. These can in principle support
+  either sync or async training, but right now we only have support for async
+  training with parameter servers. Compare to
+  `tf.distribute.experimental.CentralStorageStrategy`, which puts all variables
+  on a single device on the same machine (and does sync training), and
+  `tf.distribute.MirroredStrategy`, which mirrors variables to multiple devices
+  (see below).
+
+* _Replica context_ vs. _Cross-replica context_ vs _Update context_
+
+  A _replica context_ applies
+  when you execute the computation function that was called with `strategy.run`.
+  Conceptually, you're in replica context when executing the computation
+  function that is being replicated.
+
+  An _update context_ is entered in a `tf.distribute.StrategyExtended.update`
+  call.
+
+  An _cross-replica context_ is entered when you enter a `strategy.scope`. This
+  is useful for calling `tf.distribute.Strategy` methods which operate across
+  the replicas (like `reduce_to()`). By default you start in a _replica context_
+  (the "default single _replica context_") and then some methods can switch you
+  back and forth.
+
+* _Distributed value_: Distributed value is represented by the base class
+  `tf.distribute.DistributedValues`. `tf.distribute.DistributedValues` is useful
+  to represent values on multiple devices, and it contains a map from replica id
+  to values. Two representative types of `tf.distribute.DistributedValues`
+  are `tf.types.experimental.PerReplica` and `tf.types.experimental.Mirrored`
+  values.
+
+  `PerReplica` values exist on the worker devices, with a different value for
+  each replica. They are produced by iterating through a distributed dataset
+  returned by `tf.distribute.Strategy.experimental_distribute_dataset` and
+  `tf.distribute.Strategy.distribute_datasets_from_function`. They are also the
+  typical result returned by `tf.distribute.Strategy.run`.
+
+  `Mirrored` values are like `PerReplica` values, except we know that the value
+  on all replicas are the same. `Mirrored` values are kept synchronized by the
+  distribution strategy in use, while `PerReplica` values are left
+  unsynchronized. `Mirrored` values typically represent model weights. We can
+  safely read a `Mirrored` value in a cross-replica context by using the value
+  on any replica, while PerReplica values can only be read within a replica
+  context.
+
+* _Unwrapping_ and _merging_: Consider calling a function `fn` on multiple
+  replicas, like `strategy.run(fn, args=[w])` with an
+  argument `w` that is a `tf.distribute.DistributedValues`. This means `w` will
+  have a map taking replica id `0` to `w0`, replica id `1` to `w1`, etc.
+  `strategy.run()` unwraps `w` before calling `fn`, so it calls `fn(w0)` on
+  device `d0`, `fn(w1)` on device `d1`, etc.  It then merges the return
+  values from `fn()`, which leads to one common object if the returned values
+  are the same object from every replica, or a `DistributedValues` object
+  otherwise.
+
+* _Reductions_ and _all-reduce_: A _reduction_ is a method of aggregating
+  multiple values into one value, like "sum" or "mean". If a strategy is doing
+  sync training, we will perform a reduction on the gradients to a parameter
+  from all replicas before applying the update. _All-reduce_ is an algorithm for
+  performing a reduction on values from multiple devices and making the result
+  available on all of those devices.
+
+* _Mirrored variables_: These are variables that are created on multiple
+  devices, where we keep the variables in sync by applying the same
+  updates to every copy. Mirrored variables are created with
+  `tf.Variable(...synchronization=tf.VariableSynchronization.ON_WRITE...)`.
+  Normally they are only used in synchronous training.
+
+* _SyncOnRead variables_
+
+  _SyncOnRead variables_ are created by
+  `tf.Variable(...synchronization=tf.VariableSynchronization.ON_READ...)`, and
+  they are created on multiple devices. In replica context, each
+  component variable on the local replica can perform reads and writes without
+  synchronization with each other. When the
+  _SyncOnRead variable_ is read in cross-replica context, the values from
+  component variables are aggregated and returned.
+
+  _SyncOnRead variables_ bring a lot of custom configuration difficulty to the
+  underlying logic, so we do not encourage users to instantiate and use
+  _SyncOnRead variable_ on their own. We have mainly used _SyncOnRead
+  variables_ for use cases such as batch norm and metrics. For performance
+  reasons, we often don't need to keep these statistics in sync every step and
+  they can be accumulated on each replica independently. The only time we want
+  to sync them is reporting or checkpointing, which typically happens in
+  cross-replica context. _SyncOnRead variables_ are also often used by advanced
+  users who want to control when variable values are aggregated. For example,
+  users sometimes want to maintain gradients independently on each replica for a
+  couple of steps without aggregation.
+
+* _Distribute-aware layers_
+
+  Layers are generally called in a replica context, except when defining a
+  Keras functional model. `tf.distribute.in_cross_replica_context` will let you
+  determine which case you are in. If in a replica context,
+  the `tf.distribute.get_replica_context` function will return the default
+  replica context outside a strategy scope, `None` within a strategy scope, and
+  a `tf.distribute.ReplicaContext` object inside a strategy scope and within a
+  `tf.distribute.Strategy.run` function. The `ReplicaContext` object has an
+  `all_reduce` method for aggregating across all replicas.
+
+
+Note that we provide a default version of `tf.distribute.Strategy` that is
+used when no other strategy is in scope, that provides the same API with
+reasonable default behavior.
+
+API docstring: tensorflow.distribute
+"""

tensorflow/python/distribute/central_storage_strategy.py

@@ -0,0 +1,226 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Class implementing a single machine parameter server strategy."""
+
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import parameter_server_strategy
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('distribute.experimental.CentralStorageStrategy', v1=[])
+class CentralStorageStrategy(distribute_lib.Strategy):
+  """A one-machine strategy that puts all variables on a single device.
+
+  Variables are assigned to local CPU or the only GPU. If there is more
+  than one GPU, compute operations (other than variable update operations)
+  will be replicated across all GPUs.
+
+  For Example:
+  ```
+  strategy = tf.distribute.experimental.CentralStorageStrategy()
+  # Create a dataset
+  ds = tf.data.Dataset.range(5).batch(2)
+  # Distribute that dataset
+  dist_dataset = strategy.experimental_distribute_dataset(ds)
+
+  with strategy.scope():
+    @tf.function
+    def train_step(val):
+      return val + 1
+
+    # Iterate over the distributed dataset
+    for x in dist_dataset:
+      # process dataset elements
+      strategy.run(train_step, args=(x,))
+  ```
+  """
+
+  def __init__(self, compute_devices=None, parameter_device=None):
+    extended = parameter_server_strategy.ParameterServerStrategyExtended(
+        self,
+        compute_devices=compute_devices,
+        parameter_device=parameter_device)
+    """Initializes the strategy with optional device strings.
+
+    Args:
+    compute_devices: an optional list of strings for device to replicate models
+      on. If this is not provided, all local GPUs will be used; if there is no
+      GPU, local CPU will be used.
+    parameter_device: an optional device string for which device to put
+      variables on. The default one is CPU or GPU if there is only one.
+    """
+    super(CentralStorageStrategy, self).__init__(extended)
+    distribute_lib.distribution_strategy_gauge.get_cell('V2').set(
+        'CentralStorageStrategy')
+
+  @classmethod
+  def _from_num_gpus(cls, num_gpus):
+    return cls(device_util.local_devices_from_num_gpus(num_gpus))
+
+  def experimental_distribute_dataset(self, dataset, options=None):  # pylint: disable=useless-super-delegation
+    """Distributes a tf.data.Dataset instance provided via dataset.
+
+    The returned dataset is a wrapped strategy dataset which creates a
+    multidevice iterator under the hood. It prefetches the input data to the
+    specified devices on the worker. The returned distributed dataset can be
+    iterated over similar to how regular datasets can.
+
+    NOTE: Currently, the user cannot add any more transformations to a
+    distributed dataset.
+
+    For Example:
+    ```
+    strategy = tf.distribute.CentralStorageStrategy()  # with 1 CPU and 1 GPU
+    dataset = tf.data.Dataset.range(10).batch(2)
+    dist_dataset = strategy.experimental_distribute_dataset(dataset)
+    for x in dist_dataset:
+      print(x)  # Prints PerReplica values [0, 1], [2, 3],...
+
+    ```
+    Args:
+      dataset: `tf.data.Dataset` to be prefetched to device.
+      options: `tf.distribute.InputOptions` used to control options on how this
+        dataset is distributed.
+
+    Returns:
+      A "distributed `Dataset`" that the caller can iterate over.
+    """
+    if (options and options.experimental_replication_moden ==
+        distribute_lib.InputReplicationMode.PER_REPLICA):
+      raise NotImplementedError(
+          'InputReplicationMode.PER_REPLICA '
+          'is only supported in '
+          '`experimental_distribute_datasets_from_function`.'
+      )
+    return super(CentralStorageStrategy, self).experimental_distribute_dataset(
+        dataset, options)
+
+  def experimental_local_results(self, value):  # pylint: disable=useless-super-delegation
+    """Returns the list of all local per-replica values contained in `value`.
+
+    In `CentralStorageStrategy` there is a single worker so the value returned
+    will be all the values on that worker.
+
+    Args:
+      value: A value returned by `run()`, `extended.call_for_each_replica()`,
+      or a variable created in `scope`.
+
+    Returns:
+      A tuple of values contained in `value`. If `value` represents a single
+      value, this returns `(value,).`
+    """
+    return super(CentralStorageStrategy, self).experimental_local_results(value)
+
+  def run(self, fn, args=(), kwargs=None, options=None):  # pylint: disable=useless-super-delegation
+    """Run `fn` on each replica, with the given arguments.
+
+    In `CentralStorageStrategy`, `fn` is  called on each of the compute
+    replicas, with the provided "per replica" arguments specific to that device.
+
+    Args:
+      fn: The function to run. The output must be a `tf.nest` of `Tensor`s.
+      args: (Optional) Positional arguments to `fn`.
+      kwargs: (Optional) Keyword arguments to `fn`.
+      options: (Optional) An instance of `tf.distribute.RunOptions` specifying
+        the options to run `fn`.
+
+    Returns:
+      Return value from running `fn`.
+    """
+    return super(CentralStorageStrategy, self).run(fn, args, kwargs, options)
+
+  def reduce(self, reduce_op, value, axis):  # pylint: disable=useless-super-delegation
+    """Reduce `value` across replicas.
+
+    Given a per-replica value returned by `run`, say a
+    per-example loss, the batch will be divided across all the replicas. This
+    function allows you to aggregate across replicas and optionally also across
+    batch elements.  For example, if you have a global batch size of 8 and 2
+    replicas, values for examples `[0, 1, 2, 3]` will be on replica 0 and
+    `[4, 5, 6, 7]` will be on replica 1. By default, `reduce` will just
+    aggregate across replicas, returning `[0+4, 1+5, 2+6, 3+7]`. This is useful
+    when each replica is computing a scalar or some other value that doesn't
+    have a "batch" dimension (like a gradient). More often you will want to
+    aggregate across the global batch, which you can get by specifying the batch
+    dimension as the `axis`, typically `axis=0`. In this case it would return a
+    scalar `0+1+2+3+4+5+6+7`.
+
+    If there is a last partial batch, you will need to specify an axis so
+    that the resulting shape is consistent across replicas. So if the last
+    batch has size 6 and it is divided into [0, 1, 2, 3] and [4, 5], you
+    would get a shape mismatch unless you specify `axis=0`. If you specify
+    `tf.distribute.ReduceOp.MEAN`, using `axis=0` will use the correct
+    denominator of 6. Contrast this with computing `reduce_mean` to get a
+    scalar value on each replica and this function to average those means,
+    which will weigh some values `1/8` and others `1/4`.
+
+    For Example:
+    ```
+    strategy = tf.distribute.experimental.CentralStorageStrategy(
+        compute_devices=['CPU:0', 'GPU:0'], parameter_device='CPU:0')
+    ds = tf.data.Dataset.range(10)
+    # Distribute that dataset
+    dist_dataset = strategy.experimental_distribute_dataset(ds)
+
+    with strategy.scope():
+      @tf.function
+      def train_step(val):
+        # pass through
+        return val
+
+      # Iterate over the distributed dataset
+      for x in dist_dataset:
+        result = strategy.run(train_step, args=(x,))
+
+    result = strategy.reduce(tf.distribute.ReduceOp.SUM, result,
+                             axis=None).numpy()
+    # result: array([ 4,  6,  8, 10])
+
+    result = strategy.reduce(tf.distribute.ReduceOp.SUM, result, axis=0).numpy()
+    # result: 28
+    ```
+
+    Args:
+      reduce_op: A `tf.distribute.ReduceOp` value specifying how values should
+        be combined.
+      value: A "per replica" value, e.g. returned by `run` to
+        be combined into a single tensor.
+      axis: Specifies the dimension to reduce along within each
+        replica's tensor. Should typically be set to the batch dimension, or
+        `None` to only reduce across replicas (e.g. if the tensor has no batch
+        dimension).
+
+    Returns:
+      A `Tensor`.
+    """
+    return super(CentralStorageStrategy, self).reduce(reduce_op, value, axis)
+
+
+@tf_export(v1=['distribute.experimental.CentralStorageStrategy'])  # pylint: disable=missing-docstring
+class CentralStorageStrategyV1(distribute_lib.StrategyV1):
+
+  __doc__ = CentralStorageStrategy.__doc__
+
+  def __init__(self, compute_devices=None, parameter_device=None):
+    super(CentralStorageStrategyV1, self).__init__(
+        parameter_server_strategy.ParameterServerStrategyExtended(
+            self,
+            compute_devices=compute_devices,
+            parameter_device=parameter_device))
+    distribute_lib.distribution_strategy_gauge.get_cell('V1').set(
+        'CentralStorageStrategy')
+
+  __init__.__doc__ = CentralStorageStrategy.__init__.__doc__

tensorflow/python/distribute/checkpoint_utils_test.py

@@ -0,0 +1,135 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for checkpoint_utils.init_from_checkpoint with Distribution Strategy.
+
+These tests are located here instead of as part of
+`python.training.CheckpointsTest` because they need access to distribution
+strategies which are only present in contrib right now.
+TODO(priyag): Move the tests to core `python.training.CheckpointsTest` when
+distribution strategy moves out of contrib.
+"""
+
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import checkpoint_utils
+from tensorflow.python.training import saver as saver_lib
+
+
+def _create_checkpoints(sess, checkpoint_dir):
+  checkpoint_prefix = os.path.join(checkpoint_dir, "model")
+  checkpoint_state_name = "checkpoint"
+  v1 = variable_scope.get_variable("var1", [1, 10])
+  v2 = variable_scope.get_variable("var2", [10, 10])
+  sess.run(variables.global_variables_initializer())
+  v1_value, v2_value = sess.run([v1, v2])
+  saver = saver_lib.Saver()
+  saver.save(
+      sess,
+      checkpoint_prefix,
+      global_step=0,
+      latest_filename=checkpoint_state_name)
+  return v1_value, v2_value
+
+
+class CheckpointUtilsWithDistributionStrategyTest(
+    test.TestCase, parameterized.TestCase):
+
+  def _get_test_object(self):
+    checkpoint_dir = self.get_temp_dir()
+    with self.cached_session() as session:
+      v1, v2 = _create_checkpoints(session, checkpoint_dir)
+    return checkpoint_dir, v1, v2
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=[
+              strategy_combinations.default_strategy,
+              strategy_combinations.one_device_strategy,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations
+              .mirrored_strategy_with_two_gpus_no_merge_call,
+          ],
+          in_replica_mode=[True, False],
+          mode=["graph"]))
+  def testInitFromCheckpoint(self, distribution, in_replica_mode):
+    checkpoint_dir, v1_value, v2_value = self._get_test_object()
+
+    def init_and_verify(g):
+      v1 = variable_scope.get_variable("new_var1", [1, 10])
+      v2 = variable_scope.get_variable(
+          "new_var2", [10, 10],
+          synchronization=variable_scope.VariableSynchronization.ON_READ,
+          aggregation=variable_scope.VariableAggregation.MEAN)
+      checkpoint_utils.init_from_checkpoint(checkpoint_dir, {
+          "var1": "new_var1",
+          "var2": "new_var2"
+      })
+      with self.session(graph=g) as session:
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(v1_value, self.evaluate(v1))
+        self.assertAllEqual(v2_value, self.evaluate(v2))
+
+    with ops.Graph().as_default() as g, distribution.scope():
+      if in_replica_mode:
+        distribution.extended.call_for_each_replica(init_and_verify, args=[g])
+      else:
+        init_and_verify(g)
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=[
+              strategy_combinations.default_strategy,
+              strategy_combinations.one_device_strategy,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations
+              .mirrored_strategy_with_two_gpus_no_merge_call,
+          ],
+          in_replica_mode=[True, False],
+          mode=["graph"]))
+  def testInitFromDifferentNameObject(self, distribution, in_replica_mode):
+    checkpoint_dir, v1_value, _ = self._get_test_object()
+
+    def init_and_verify(g):
+      v1 = variable_scope.get_variable("new_var1", [1, 10])
+      # Use string add to create new object in each replica
+      prefix = "new_"
+      suffix = "var1"
+      new_var1 = prefix + suffix
+      checkpoint_utils.init_from_checkpoint(checkpoint_dir, {
+          "var1": new_var1,
+      })
+      with self.test_session(graph=g) as session:
+        session.run(variables.global_variables_initializer())
+        self.assertAllEqual(v1_value, self.evaluate(v1))
+
+    with ops.Graph().as_default() as g, distribution.scope():
+      if in_replica_mode:
+        distribution.extended.call_for_each_replica(init_and_verify, [g])
+      else:
+        init_and_verify(g)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/checkpointing_test.py

@@ -0,0 +1,67 @@
+# Copyright 2018 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.
+# ==============================================================================
+
+import os
+
+from absl.testing import parameterized
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.eager import test
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables as variables_lib
+
+
+class TrainingCheckpointTests(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_one_cpu,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+              strategy_combinations.tpu_strategy_packed_var,
+              strategy_combinations.central_storage_strategy_with_two_gpus,
+          ],
+          mode=["eager"]))
+  def testInitializeFromCheckpoint(self, distribution):
+    variable_shape = [5]
+    save_checkpoint = trackable_utils.Checkpoint(v=variables_lib.Variable(
+        array_ops.ones(variable_shape)))
+    save_path = save_checkpoint.save(
+        os.path.join(self.get_temp_dir(), "checkpoint"))
+    with distribution.scope():
+      restore_checkpoint = trackable_utils.Checkpoint()
+      restore_checkpoint.restore(save_path)
+      initial_value = restore_checkpoint._preload_simple_restoration(
+          "v")
+      v = variables_lib.Variable(initial_value)
+      # Check that the variable is now tagged as restored. `Checkpoint` then
+      # knows it doesn't have to restore `v`'s value when it's assigned to an
+      # object.
+      self.assertGreater(v._update_uid, 0)
+      self.assertAllClose(array_ops.ones(variable_shape), v)
+      v.assign(array_ops.zeros(variable_shape))
+      # Assignment to an object should not trigger restoration, since we already
+      # restored the object through an initializer. This wouldn't be a
+      # correctness issue, but it would mean that models would use twice as much
+      # memory when loading (the buffer already assigned to the variable, and
+      # the new restoration).
+      restore_checkpoint.v = v
+      self.assertAllClose(array_ops.zeros(variable_shape), v)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/cluster_resolver/BUILD

@@ -0,0 +1,209 @@
+# Description: Operations defined for Cluster Resolvers
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "cluster_resolver_lib",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":base_cluster_resolver_py",
+        ":gce_cluster_resolver_py",
+        ":kubernetes_cluster_resolver_py",
+        ":sagemaker_cluster_resolver_py",
+        ":slurm_cluster_resolver_py",
+        ":tfconfig_cluster_resolver_py",
+        ":tpu_cluster_resolver_py",
+    ],  # placeholder for google-internal dependencies,,
+)
+
+py_library(
+    name = "base_cluster_resolver_py",
+    srcs = ["cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:tf_export",
+        "@six_archive//:six",
+    ],
+)
+
+py_library(
+    name = "gce_cluster_resolver_py",
+    srcs = ["gce_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        ":base_cluster_resolver_py",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "tfconfig_cluster_resolver_py",
+    srcs = ["tfconfig_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        ":base_cluster_resolver_py",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "sagemaker_cluster_resolver_py",
+    srcs = ["sagemaker_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        ":base_cluster_resolver_py",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "tpu_cluster_resolver_py",
+    srcs = ["tpu_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "slurm_cluster_resolver_py",
+    srcs = ["slurm_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        ":base_cluster_resolver_py",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "kubernetes_cluster_resolver_py",
+    srcs = ["kubernetes_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        ":base_cluster_resolver_py",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "base_cluster_resolver_py_test",
+    srcs = ["cluster_resolver_test.py"],
+    main = "cluster_resolver_test.py",
+    deps = [
+        ":base_cluster_resolver_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "gce_cluster_resolver_py_test",
+    size = "small",
+    srcs = ["gce_cluster_resolver_test.py"],
+    main = "gce_cluster_resolver_test.py",
+    deps = [
+        ":base_cluster_resolver_py",
+        ":gce_cluster_resolver_py",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "tfconfig_cluster_resolver_py_test",
+    size = "small",
+    srcs = ["tfconfig_cluster_resolver_test.py"],
+    grpc_enabled = True,
+    main = "tfconfig_cluster_resolver_test.py",
+    deps = [
+        ":tfconfig_cluster_resolver_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "sagemaker_cluster_resolver_py_test",
+    size = "small",
+    srcs = ["sagemaker_cluster_resolver_test.py"],
+    grpc_enabled = True,
+    main = "sagemaker_cluster_resolver_test.py",
+    deps = [
+        ":sagemaker_cluster_resolver_py",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "slurm_cluster_resolver_py_test",
+    size = "small",
+    srcs = ["slurm_cluster_resolver_test.py"],
+    main = "slurm_cluster_resolver_test.py",
+    tags = [],
+    deps = [
+        ":slurm_cluster_resolver_py",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "kubernetes_cluster_resolver_py_test",
+    size = "small",
+    srcs = ["kubernetes_cluster_resolver_test.py"],
+    main = "kubernetes_cluster_resolver_test.py",
+    deps = [
+        ":kubernetes_cluster_resolver_py",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+# copybara:uncomment_begin(google-only)
+# tf_py_strict_test(
+#     name = "brain_jobs_cluster_resolver_test",
+#     size = "small",
+#     srcs = ["brain_jobs_cluster_resolver_test.py"],
+#     grpc_enabled = True,
+#     main = "brain_jobs_cluster_resolver_test.py",
+#     deps = [
+#         ":brain_jobs_cluster_resolver_py",
+#         "//tensorflow/python/framework:test_lib",
+#         "//tensorflow/python/platform:client_testlib",
+#         "//tensorflow/python/training:server_lib",
+#     ],
+# )
+# copybara:uncomment_end

tensorflow/python/distribute/cluster_resolver/README.md

@@ -0,0 +1,5 @@
+# Cluster Resolvers
+
+Cluster Resolvers are a new way of specifying cluster information for distributed execution. Built on top of existing `ClusterSpec` framework, Cluster Resolvers allow users to simply specify a configuration and a cluster management service and a `ClusterResolver` will automatically fetch the relevant information from the service and populate `ClusterSpec`s.
+
+`ClusterResolvers` are designed to work well with `ManagedTrainingSession` and `ClusterSpec` propagation so that distributed training sessions remain robust in the face of node and network failures.

tensorflow/python/distribute/cluster_resolver/README_Slurm.md

@@ -0,0 +1,97 @@
+# Slurm Cluster Resolver
+
+The Slurm Cluster Resolver resolves cluster specification for distributing
+TensorFlow work launched on HPC systems running on Slurm. This implementation is
+able to handle homogeneous and heterogeneous tasks as long as the number of GPUs
+per node and task are the same. This means on nodes with 4 GPUs each it will be
+possible to allocate 4 processes on node A and only 2 on node B. The resolution
+is done by determining job configuration through a number of Slurm variables and
+can be overwritten by user input. By default everything is determined from the
+Slurm environment, hence for most uses case no manual setting of parameters is
+required.
+
+## How it works
+
+`SlurmClusterResolver` reads the environment variables that are set inside a job
+step launched by Slurm. This means it will only work correctly for applications
+launched via `srun`.
+
+The process ID/rank is extracted from environment variable `SLURM_PROCID` and
+the total number of tasks launched is extracted from `SLURM_STEP_NUM_TASKS`. The
+hostnames are resolved by inspection `SLURM_STEP_NODELIST`. The number of tasks
+per node is extracted from `SLURM_STEP_TASKS_PER_NODE`, unless a value is
+specified by user. By using this variable heterogeneous task distributions are
+possible. The user can set `tasks_per_node` to a single integer for homogeneous
+tasks or a dictionary mapping node names to number of tasks for heterogeneous
+distributions. However setting this is **NOT** recommended as there is a chance
+it makes `SLURM_PROCID` be wrong.
+
+A base port can be specified by user and in case there are more than one task
+launched per node the port number will be incremented for each additional tasks
+on that node. However a reasonable default is used.
+
+The number of GPUs present on each node and number of GPUs for each tasks are
+automatically detected. This is done by checking for `CUDA_VISIBLE_DEVICES`
+first (which is set by Slurm to a list of GPUs for the current node) and has a
+fallback to using `nvidia-smi`. If this doesn't work or non-NVIDIA GPUs are used
+those 2 values have to be specified by the user. By default allocated GPUs will
+be automatically exposed to processes according to specification by setting
+`CUDA_VISIBLE_DEVICES`.
+
+## Basic example
+
+-   Slurm allocation in shell `salloc --nodes=2 -t 01:30:00 --ntasks-per-node=2
+    --gres=gpu:k80:4 --exclusive`
+-   Run the example `srun python tf_example.py`
+-   Creating cluster in Python `import tensorflow as tf cluster_resolver =
+    tf.distribute.cluster_resolver.SlurmClusterResolver() strategy =
+    tf.distribute.experimental.MultiWorkerMirroredStrategy(cluster_resolver=cluster_resolver)
+    with strategy.scope(): # Load and compile model and data`
+
+The above example will allocate 4 jobs on 2 nodes with each node having 2 jobs
+and 4 GPUs. `cluster_resolver.cluster_spec()` will return a cluster
+specification object in protobuf format with the following value (host names may
+vary): `job { name: "worker" tasks { key: 0 value: "t02n13:8888" } tasks { key:
+1 value: "t02n13:8889" } tasks { key: 2 value: "t02n41:8888" } tasks { key: 3
+value: "t02n41:8889" } }`
+
+The `job_name` will be `worker` for all nodes and `task_index` will be `0` to
+`3`. Also GPUs will be allocated automatically, so the first job on each node
+will see GPU 0 and 1, and the second GPU 2 and 3.
+
+## Advanced example
+
+-   Assuming the same job parameters (`salloc` & `srun`) as above
+-   Creating cluster in Python ``` cluster_resolver =
+    tf.contrib.cluster_resolver.SlurmClusterResolver( {'ps': 1, 'worker': 3},
+    port_base=1337, tasks_per_node=2, gpus_per_node=2, gpus_per_task=1,
+    auto_set_gpu=False)
+
+cluster = cluster_resolver.cluster_spec() job_name, task_index =
+cluster_resolver.get_task_info() ```
+
+In this case 1 parameter server job and 3 worker jobs are used. The resulting
+protobuf specification will look similar to this: `job { name: "ps" tasks { key:
+0 value: "t02n13:1337" } } job { name: "worker" tasks { key: 0 value:
+"t02n13:1338" } tasks { key: 1 value: "t02n41:1337" } tasks { key: 2 value:
+"t02n41:1338" } }`
+
+The value of `job_name` will be `ps` for `t02n13:1337` and `worker` for all
+others. There will be no GPU allocation done by the cluster resolver, so this
+has to be done manually which is useful if e.g. GPUs 0 should go to the first
+process and GPU 3 to the second process on each node. Also note that only 1 GPU
+will be used per task.
+
+## Extension points
+
+The class `SlurmClusterResolver` provides some methods that are meant to be
+overwritten by deriving classes:
+
+-   `_resolve_own_rank`
+-   `_resolve_num_tasks`
+-   `_resolve_hostlist`
+-   `_resolve_task_configuration`
+
+    Those can be used to implement a cluster resolver that gets information from
+    a different source, e.g. via MPI, a file or other environment variables. See
+    the documentation of these methods on what to return.

tensorflow/python/distribute/cluster_resolver/__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Library imports for ClusterResolvers.
+
+  This library contains all implementations of ClusterResolvers.
+  ClusterResolvers are a way of specifying cluster information for distributed
+  execution. Built on top of existing `ClusterSpec` framework, ClusterResolvers
+  are a way for TensorFlow to communicate with various cluster management
+  systems (e.g. GCE, AWS, etc...).
+"""
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import UnionClusterResolver
+from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver
+from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import  ExecutableLocation
+from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver
+from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver
+from tensorflow.python.distribute.cluster_resolver.tfconfig_cluster_resolver import TFConfigClusterResolver
+from tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver import TPUClusterResolver

tensorflow/python/distribute/cluster_resolver/cluster_resolver.py

@@ -0,0 +1,624 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Cluster Resolvers are used for dynamic cluster IP/hostname resolution."""
+
+import abc
+
+import collections
+
+import six
+
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import ops
+from tensorflow.python.training.server_lib import ClusterSpec
+from tensorflow.python.util.tf_export import tf_export
+
+
+def format_master_url(master, rpc_layer=None):
+  if rpc_layer:
+    return '%s://%s' % (rpc_layer, master)
+  else:
+    return master
+
+
+def get_accelerator_devices(master, config_proto):
+  """Returns accelerator devices given a master and a configuration."""
+  if context.executing_eagerly():
+    logical_devices = config.list_logical_devices()
+    devices = []
+    for d in logical_devices:
+      if d.device_type == 'CPU' or d.device_type == 'XLA_CPU':  # Filter CPUs
+        continue
+      devices.append(session._DeviceAttributes(d.name, d.device_type, 0, 0))  # pylint: disable=protected-access
+    return devices
+  else:
+    with ops.Graph().as_default():
+      with session.Session(master, config=config_proto) as s:
+        devices = s.list_devices()
+    return devices
+
+
+@tf_export('distribute.cluster_resolver.ClusterResolver')
+@six.add_metaclass(abc.ABCMeta)
+class ClusterResolver(object):
+  """Abstract class for all implementations of ClusterResolvers.
+
+  This defines the skeleton for all implementations of ClusterResolvers.
+  ClusterResolvers are a way for TensorFlow to communicate with various cluster
+  management systems (e.g. GCE, AWS, etc...) and gives TensorFlow necessary
+  information to set up distributed training.
+
+  By letting TensorFlow communicate with these systems, we will be able to
+  automatically discover and resolve IP addresses for various TensorFlow
+  workers. This will eventually allow us to automatically recover from
+  underlying machine failures and scale TensorFlow worker clusters up and down.
+
+  Note to Implementors of `tf.distribute.cluster_resolver.ClusterResolver`
+  subclass: In addition to these abstract methods, when task_type, task_id, and
+  rpc_layer attributes are applicable, you should also implement them either as
+  properties with getters or setters, or directly set the attributes
+  `self._task_type`, `self._task_id`, or `self._rpc_layer` so the base class'
+  getters and setters are used. See
+  `tf.distribute.cluster_resolver.SimpleClusterResolver.__init__` for an
+  example.
+
+  In general, multi-client tf.distribute strategies such as
+  `tf.distribute.experimental.MultiWorkerMirroredStrategy` require task_type and
+  task_id properties to be available in the `ClusterResolver` they are using. On
+  the other hand, these concepts are not applicable in single-client strategies,
+  such as `tf.distribute.experimental.TPUStrategy`, because the program is only
+  expected to be run on one task, so there should not be a need to have code
+  branches according to task type and task id.
+
+  - task_type is the name of the server's current named job (e.g. 'worker',
+     'ps' in a distributed parameterized training job).
+  - task_id is the ordinal index of the server within the task type.
+  - rpc_layer is the protocol used by TensorFlow to communicate with other
+      TensorFlow servers in a distributed environment.
+  """
+
+  @abc.abstractmethod
+  def cluster_spec(self):
+    """Retrieve the current state of the cluster and return a `tf.train.ClusterSpec`.
+
+    Returns:
+      A `tf.train.ClusterSpec` representing the state of the cluster at the
+      moment this function is called.
+
+    Implementors of this function must take care in ensuring that the
+    ClusterSpec returned is up-to-date at the time of calling this function.
+    This usually means retrieving the information from the underlying cluster
+    management system every time this function is invoked and reconstructing
+    a cluster_spec, rather than attempting to cache anything.
+    """
+    raise NotImplementedError()
+
+  @abc.abstractmethod
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Retrieves the name or URL of the session master.
+
+    Note: this is only useful for TensorFlow 1.x.
+
+    Args:
+      task_type: (Optional) The type of the TensorFlow task of the master.
+      task_id: (Optional) The index of the TensorFlow task of the master.
+      rpc_layer: (Optional) The RPC protocol for the given cluster.
+
+    Returns:
+      The name or URL of the session master.
+
+    Implementors of this function must take care in ensuring that the master
+    returned is up-to-date at the time to calling this function. This usually
+    means retrieving the master every time this function is invoked.
+    """
+    raise NotImplementedError()
+
+  def num_accelerators(self,
+                       task_type=None,
+                       task_id=None,
+                       config_proto=None):
+    """Returns the number of accelerator cores per worker.
+
+    This returns the number of accelerator cores (such as GPUs and TPUs)
+    available per worker.
+
+    Optionally, we allow callers to specify the task_type, and task_id, for
+    if they want to target a specific TensorFlow task to query
+    the number of accelerators. This is to support heterogenous environments,
+    where the number of accelerators cores per host is different.
+
+    Args:
+      task_type: (Optional) The type of the TensorFlow task of the machine we
+        want to query.
+      task_id: (Optional) The index of the TensorFlow task of the machine we
+        want to query.
+      config_proto: (Optional) Configuration for starting a new session to
+        query how many accelerator cores it has.
+
+    Returns:
+      A map of accelerator types to number of cores.
+    """
+    master = self.master(task_type, task_id)
+    # TODO(b/126786766): in eager mode, we should check whether
+    # `tf.config.experimental_connect_to_cluster` is called or not.
+    devices = get_accelerator_devices(master, config_proto)
+    mapping = collections.defaultdict(int)
+    for device in devices:
+      if task_type is not None and task_id is not None:
+        job_path = '/job:%s' % task_type
+        task_path = '/task:%s' % task_id
+        if job_path not in device.name or task_path not in device.name:
+          continue
+      mapping[device.device_type] += 1
+    return mapping
+
+  @property
+  def environment(self):
+    """Returns the current environment which TensorFlow is running in.
+
+    There are two possible return values, "google" (when TensorFlow is running
+    in a Google-internal environment) or an empty string (when TensorFlow is
+    running elsewhere).
+
+    If you are implementing a ClusterResolver that works in both the Google
+    environment and the open-source world (for instance, a TPU ClusterResolver
+    or similar), you will have to return the appropriate string depending on the
+    environment, which you will have to detect.
+
+    Otherwise, if you are implementing a ClusterResolver that will only work
+    in open-source TensorFlow, you do not need to implement this property.
+    """
+    return ''
+
+  @property
+  def task_type(self):
+    """Returns the task type this `ClusterResolver` indicates.
+
+    In TensorFlow distributed environment, each job may have an applicable
+    task type. Valid task types in TensorFlow include
+    'chief': a worker that is designated with more responsibility,
+    'worker': a regular worker for training/evaluation,
+    'ps': a parameter server, or
+    'evaluator': an evaluator that evaluates the checkpoints for metrics.
+
+    See [Multi-worker configuration](
+    https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#multi-worker_configuration)
+    for more information about 'chief' and 'worker' task type, which are most
+    commonly used.
+
+    Having access to such information is useful when user needs to run specific
+    code according to task types. For example,
+
+    ```python
+    cluster_spec = tf.train.ClusterSpec({
+        "ps": ["localhost:2222", "localhost:2223"],
+        "worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
+    })
+
+    # SimpleClusterResolver is used here for illustration; other cluster
+    # resolvers may be used for other source of task type/id.
+    simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
+                                            task_id=1)
+
+    ...
+
+    if cluster_resolver.task_type == 'worker':
+      # Perform something that's only applicable on workers. This block
+      # will run on this particular instance since we've specified this task to
+      # be a worker in above cluster resolver.
+    elif cluster_resolver.task_type == 'ps':
+      # Perform something that's only applicable on parameter servers. This
+      # block will not run on this particular instance.
+    ```
+
+    Returns `None` if such information is not available or is not applicable
+    in the current distributed environment, such as training with
+    `tf.distribute.experimental.TPUStrategy`.
+
+    For more information, please see
+    `tf.distribute.cluster_resolver.ClusterResolver`'s class doc.
+    """
+    return getattr(self, '_task_type', None)
+
+  @property
+  def task_id(self):
+    """Returns the task id this `ClusterResolver` indicates.
+
+    In TensorFlow distributed environment, each job may have an applicable
+    task id, which is the index of the instance within its task type. This is
+    useful when user needs to run specific code according to task index. For
+    example,
+
+    ```python
+    cluster_spec = tf.train.ClusterSpec({
+        "ps": ["localhost:2222", "localhost:2223"],
+        "worker": ["localhost:2224", "localhost:2225", "localhost:2226"]
+    })
+
+    # SimpleClusterResolver is used here for illustration; other cluster
+    # resolvers may be used for other source of task type/id.
+    simple_resolver = SimpleClusterResolver(cluster_spec, task_type="worker",
+                                            task_id=0)
+
+    ...
+
+    if cluster_resolver.task_type == 'worker' and cluster_resolver.task_id == 0:
+      # Perform something that's only applicable on 'worker' type, id 0. This
+      # block will run on this particular instance since we've specified this
+      # task to be a 'worker', id 0 in above cluster resolver.
+    else:
+      # Perform something that's only applicable on other ids. This block will
+      # not run on this particular instance.
+    ```
+
+    Returns `None` if such information is not available or is not applicable
+    in the current distributed environment, such as training with
+    `tf.distribute.cluster_resolver.TPUClusterResolver`.
+
+    For more information, please see
+    `tf.distribute.cluster_resolver.ClusterResolver`'s class docstring.
+    """
+    return getattr(self, '_task_id', None)
+
+  @task_type.setter
+  def task_type(self, task_type):
+    """Setter of `task_type` property. See `task_type` property doc."""
+    self._task_type = task_type
+
+  @task_id.setter
+  def task_id(self, task_id):
+    """Setter of `task_id` property. See `task_type` property doc."""
+    self._task_id = task_id
+
+
+@tf_export('distribute.cluster_resolver.SimpleClusterResolver')
+class SimpleClusterResolver(ClusterResolver):
+  """Simple implementation of ClusterResolver that accepts all attributes.
+
+  Please see the base class for documentation of arguments of its constructor.
+
+  It is useful if you want to specify some or all attributes.
+
+  Usage example with `tf.distribute.Strategy`:
+
+    ```Python
+    cluster = tf.train.ClusterSpec({"worker": ["worker0.example.com:2222",
+                                               "worker1.example.com:2222"]})
+
+    # On worker 0
+    cluster_resolver = SimpleClusterResolver(cluster, task_type="worker",
+                                             task_id=0,
+                                             num_accelerators={"GPU": 8},
+                                             rpc_layer="grpc")
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=cluster_resolver)
+
+    # On worker 1
+    cluster_resolver = SimpleClusterResolver(cluster, task_type="worker",
+                                             task_id=1,
+                                             num_accelerators={"GPU": 8},
+                                             rpc_layer="grpc")
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=cluster_resolver)
+    ```
+  """
+
+  def __init__(self, cluster_spec, master='', task_type=None, task_id=None,
+               environment='', num_accelerators=None,
+               rpc_layer=None):
+    """Creates a SimpleClusterResolver from a ClusterSpec."""
+    super(SimpleClusterResolver, self).__init__()
+
+    self._task_type = task_type
+    self._task_id = task_id
+    self._environment = environment
+
+    self._num_accelerators = num_accelerators
+    self._rpc_layer = rpc_layer
+
+    if not isinstance(cluster_spec, ClusterSpec):
+      raise TypeError('cluster_spec must be a `tf.train.ClusterSpec`.')
+    self._cluster_spec = cluster_spec
+
+    if not isinstance(master, str):
+      raise TypeError('master must be a string.')
+    self._master = master
+
+  def cluster_spec(self):
+    """Returns the ClusterSpec passed into the constructor."""
+    return self._cluster_spec
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Returns the master address to use when creating a session.
+
+    Note: this is only useful for TensorFlow 1.x.
+
+    Args:
+      task_type: (Optional) The type of the TensorFlow task of the master.
+      task_id: (Optional) The index of the TensorFlow task of the master.
+      rpc_layer: (Optional) The RPC used by distributed TensorFlow.
+
+    Returns:
+      The name or URL of the session master.
+
+    If a task_type and task_id is given, this will override the `master`
+    string passed into the initialization function.
+    """
+    if task_type is not None and task_id is not None:
+      master = self.cluster_spec().task_address(task_type, task_id)
+    else:
+      master = self._master
+
+    return format_master_url(master, rpc_layer=rpc_layer or self._rpc_layer)
+
+  @property
+  def task_type(self):
+    return self._task_type
+
+  @property
+  def task_id(self):
+    return self._task_id
+
+  @task_type.setter
+  def task_type(self, task_type):
+    self._task_type = task_type
+
+  @task_id.setter
+  def task_id(self, task_id):
+    self._task_id = task_id
+
+  @property
+  def environment(self):
+    return self._environment
+
+  def num_accelerators(self,
+                       task_type=None,
+                       task_id=None,
+                       config_proto=None):
+    """Returns the number of accelerator cores per worker.
+
+    The SimpleClusterResolver does not do automatic detection of accelerators,
+    and thus all arguments are unused and we simply return the value provided
+    in the constructor.
+
+    Args:
+      task_type: Unused.
+      task_id: Unused.
+      config_proto: Unused.
+    """
+    # Unused
+    del task_type, task_id, config_proto
+    if self._num_accelerators is None:
+      return {}
+    return self._num_accelerators
+
+  @property
+  def rpc_layer(self):
+    return self._rpc_layer
+
+  @rpc_layer.setter
+  def rpc_layer(self, rpc_layer):
+    self._rpc_layer = rpc_layer
+
+
+@tf_export('distribute.cluster_resolver.UnionResolver')
+class UnionClusterResolver(ClusterResolver):
+  """Performs a union on underlying ClusterResolvers.
+
+  This class performs a union given two or more existing ClusterResolvers. It
+  merges the underlying ClusterResolvers, and returns one unified ClusterSpec
+  when cluster_spec is called. The details of the merge function is
+  documented in the cluster_spec function.
+
+  For additional ClusterResolver properties such as task type, task index,
+  rpc layer, environment, etc..., we will return the value from the first
+  ClusterResolver in the union.
+
+  An example to combine two cluster resolvers:
+
+    ```Python
+    cluster_0 = tf.train.ClusterSpec({"worker": ["worker0.example.com:2222",
+                                                 "worker1.example.com:2222"]})
+    cluster_resolver_0 = SimpleClusterResolver(cluster, task_type="worker",
+                                               task_id=0,
+                                               rpc_layer="grpc")
+
+    cluster_1 = tf.train.ClusterSpec({"ps": ["ps0.example.com:2222",
+                                             "ps1.example.com:2222"]})
+    cluster_resolver_1 = SimpleClusterResolver(cluster, task_type="ps",
+                                               task_id=0,
+                                               rpc_layer="grpc")
+
+    # Its task type would be "worker".
+    cluster_resolver = UnionClusterResolver(cluster_resolver_0,
+                                            cluster_resolver_1)
+    ```
+
+  An example to override the number of GPUs in a TFConfigClusterResolver
+  instance:
+
+    ```Python
+    tf_config = TFConfigClusterResolver()
+    gpu_override = SimpleClusterResolver(tf_config.cluster_spec(),
+                                         num_accelerators={"GPU": 1})
+    cluster_resolver = UnionResolver(gpu_override, tf_config)
+    ```
+  """
+
+  def __init__(self, *args, **kwargs):
+    """Initializes a UnionClusterResolver with other ClusterResolvers.
+
+    Args:
+      *args: `ClusterResolver` objects to be unionized.
+      **kwargs:
+        rpc_layer - (Optional) Override value for the RPC layer used by
+          TensorFlow.
+        task_type - (Optional) Override value for the current task type.
+        task_id - (Optional) Override value for the current task index.
+
+    Raises:
+      TypeError: If any argument is not a subclass of `ClusterResolvers`.
+      ValueError: If there are no arguments passed.
+    """
+    super(UnionClusterResolver, self).__init__()
+
+    self._rpc_layer = kwargs.pop('rpc_layer', None)
+    self._task_type = kwargs.pop('task_type', None)
+    self._task_id = kwargs.pop('task_id', None)
+
+    if kwargs:
+      raise ValueError('Unexpected kwargs provided {!r}'.format(kwargs))
+
+    if not args:
+      raise ValueError('At least one ClusterResolver is required.')
+
+    for cluster_resolver in args:
+      if not isinstance(cluster_resolver, ClusterResolver):
+        raise TypeError('All arguments must be a sub-class of '
+                        '`ClusterResolver.`')
+    self._cluster_resolvers = args
+
+  def cluster_spec(self):
+    """Returns a union of all the ClusterSpecs from the ClusterResolvers.
+
+    Returns:
+      A ClusterSpec containing host information merged from all the underlying
+      ClusterResolvers.
+
+    Raises:
+      KeyError: If there are conflicting keys detected when merging two or
+      more dictionaries, this exception is raised.
+
+    Note: If there are multiple ClusterResolvers exposing ClusterSpecs with the
+    same job name, we will merge the list/dict of workers.
+
+    If *all* underlying ClusterSpecs expose the set of workers as lists, we will
+    concatenate the lists of workers, starting with the list of workers from
+    the first ClusterResolver passed into the constructor.
+
+    If *any* of the ClusterSpecs expose the set of workers as a dict, we will
+    treat all the sets of workers as dicts (even if they are returned as lists)
+    and will only merge them into a dict if there is no conflicting keys. If
+    there is a conflicting key, we will raise a `KeyError`.
+    """
+
+    merged_cluster = {}
+
+    # We figure out whether it is all lists for a particular job, or whether
+    # there are dicts inside.
+    for cluster_resolver in self._cluster_resolvers:
+      cluster_spec = cluster_resolver.cluster_spec()
+      cluster_dict = cluster_spec.as_dict()
+
+      for job_name, tasks in cluster_dict.items():
+        if job_name in merged_cluster:
+          # If we see a dict, then we write a dict out regardless.
+          if isinstance(tasks, dict):
+            merged_cluster[job_name] = {}
+        else:
+          # We take whichever type is present.
+          if isinstance(tasks, list):
+            merged_cluster[job_name] = []
+          else:
+            merged_cluster[job_name] = {}
+
+    # We then do the merge as appropriate in merged_cluster[job].
+    for cluster_resolver in self._cluster_resolvers:
+      cluster_spec = cluster_resolver.cluster_spec()
+      cluster_dict = cluster_spec.as_dict()
+
+      for job_name, tasks in cluster_dict.items():
+        if isinstance(merged_cluster[job_name], list):
+          # We all have lists, we can just concatenate and be done.
+          merged_cluster[job_name].extend(tasks)
+        else:
+          if isinstance(tasks, list):
+            # We convert to a dictionary if the type is a list.
+            task_dict = dict(zip(range(0, len(tasks)), tasks))
+          else:
+            # We can simply make a copy (for update) and be done.
+            task_dict = tasks.copy()
+
+          # We detect if there are duplicates, and raise an error if so.
+          task_keys = set(task_dict)
+          merged_keys = set(merged_cluster[job_name].keys())
+          intersected_keys = task_keys.intersection(merged_keys)
+          if intersected_keys:
+            raise KeyError('Duplicate keys detected when merging two '
+                           'ClusterSpecs: %s' % repr(intersected_keys))
+
+          # We do the merge after all the processing.
+          merged_cluster[job_name].update(task_dict)
+
+    return ClusterSpec(merged_cluster)
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Returns the master address to use when creating a session.
+
+    This usually returns the master from the first ClusterResolver passed in,
+    but you can override this by specifying the task_type and task_id.
+
+    Note: this is only useful for TensorFlow 1.x.
+
+    Args:
+      task_type: (Optional) The type of the TensorFlow task of the master.
+      task_id: (Optional) The index of the TensorFlow task of the master.
+      rpc_layer: (Optional) The RPC protocol for the given cluster.
+
+    Returns:
+      The name or URL of the session master.
+    """
+    if task_type is not None and task_id is not None:
+      master = self.cluster_spec().task_address(task_type, task_id)
+      return format_master_url(master, rpc_layer or self._rpc_layer)
+
+    return self._cluster_resolvers[0].master(rpc_layer=rpc_layer)
+
+  @property
+  def task_type(self):
+    return self._task_type or self._cluster_resolvers[0].task_type
+
+  @property
+  def task_id(self):
+    return self._task_id or self._cluster_resolvers[0].task_id
+
+  @task_type.setter
+  def task_type(self, task_type):
+    self._task_type = task_type
+
+  @task_id.setter
+  def task_id(self, task_id):
+    self._task_id = task_id
+
+  @property
+  def environment(self):
+    return self._cluster_resolvers[0].environment
+
+  def num_accelerators(self,
+                       task_type=None,
+                       task_id=None,
+                       config_proto=None):
+    return self._cluster_resolvers[0].num_accelerators(
+        task_type, task_id, config_proto)
+
+  @property
+  def rpc_layer(self):
+    return self._rpc_layer or self._cluster_resolvers[0].rpc_layer
+
+  @rpc_layer.setter
+  def rpc_layer(self, rpc_layer):
+    self._rpc_layer = rpc_layer

tensorflow/python/distribute/cluster_resolver/cluster_resolver_test.py

@@ -0,0 +1,475 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for Cluster Resolvers."""
+
+from tensorflow.python.client import session
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+mock = test.mock
+
+
+class MockBaseClusterResolver(cluster_resolver.ClusterResolver):
+
+  def cluster_spec(self):
+    return None
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    return ""
+
+  def environment(self):
+    return ""
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class BaseClusterResolverTest(test.TestCase):
+
+  @mock.patch.object(config, "list_logical_devices")
+  @mock.patch.object(session.BaseSession, "list_devices")
+  def testNumAcceleratorsSuccess(self, mock_list_devices,
+                                 mock_eager_list_devices):
+    devices = [
+        context.LogicalDevice("/job:worker/task:0/device:GPU:0", "GPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:1", "GPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:2", "GPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:3", "GPU"),
+    ]
+    device_list = [
+        session._DeviceAttributes(d.name, d.device_type, 1024, 0)
+        for d in devices
+    ]
+    mock_eager_list_devices.return_value = devices
+    mock_list_devices.return_value = device_list
+
+    resolver = MockBaseClusterResolver()
+    self.assertEqual(resolver.num_accelerators(), {"GPU": 4})
+
+  @mock.patch.object(config, "list_logical_devices")
+  @mock.patch.object(session.BaseSession, "list_devices")
+  def testNumAcceleratorsMultiDeviceSuccess(self, mock_list_devices,
+                                            mock_eager_list_devices):
+    devices = [
+        context.LogicalDevice("/job:worker/task:0/device:TPU:0", "TPU"),
+        context.LogicalDevice("/job:worker/task:0/device:TPU:1", "TPU"),
+        context.LogicalDevice("/job:worker/task:0/device:TPU:2", "TPU"),
+        context.LogicalDevice("/job:worker/task:0/device:TPU:3", "TPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:0", "GPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:1", "GPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:2", "GPU"),
+        context.LogicalDevice("/job:worker/task:0/device:GPU:3", "GPU"),
+    ]
+    device_list = [
+        session._DeviceAttributes(d.name, d.device_type, 1024, 0)
+        for d in devices
+    ]
+    mock_eager_list_devices.return_value = devices
+    mock_list_devices.return_value = device_list
+
+    resolver = MockBaseClusterResolver()
+    self.assertEqual(resolver.num_accelerators(), {"TPU": 4, "GPU": 4})
+
+  @mock.patch.object(config, "list_logical_devices")
+  @mock.patch.object(session.BaseSession, "list_devices")
+  def testNumAcceleratorsFilterTasks(self, mock_list_devices,
+                                     mock_eager_list_devices):
+    devices = [
+        context.LogicalDevice("/job:worker1/task:0/device:TPU:0", "TPU"),
+        context.LogicalDevice("/job:worker1/task:0/device:TPU:1", "TPU"),
+        context.LogicalDevice("/job:worker1/task:0/device:GPU:0", "GPU"),
+        context.LogicalDevice("/job:worker1/task:0/device:GPU:1", "GPU"),
+        context.LogicalDevice("/job:worker2/task:1/device:TPU:2", "TPU"),
+        context.LogicalDevice("/job:worker2/task:2/device:TPU:3", "TPU"),
+        context.LogicalDevice("/job:worker2/task:3/device:GPU:2", "GPU"),
+        context.LogicalDevice("/job:worker2/task:4/device:GPU:3", "GPU"),
+    ]
+    device_list = [
+        session._DeviceAttributes(d.name, d.device_type, 1024, 0)
+        for d in devices
+    ]
+    mock_eager_list_devices.return_value = devices
+    mock_list_devices.return_value = device_list
+
+    resolver = MockBaseClusterResolver()
+    self.assertEqual(resolver.num_accelerators(task_type="worker1", task_id=0),
+                     {"TPU": 2, "GPU": 2})
+    self.assertEqual(resolver.num_accelerators(task_type="worker2", task_id=3),
+                     {"GPU": 1})
+    self.assertEqual(resolver.num_accelerators(task_type="worker2", task_id=4),
+                     {"GPU": 1})
+
+
+class UnionClusterResolverTest(test.TestCase):
+  # TODO(frankchn): Transform to parameterized test after it is included in the
+  # TF open source codebase.
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
+  def testSingleClusterResolver(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+    simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
+    union_resolver = cluster_resolver.UnionClusterResolver(simple_resolver)
+
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                         tasks { key: 1 value: 'worker1:2222' }
+                         tasks { key: 2 value: 'worker2:2222' } }
+    """
+    actual_cluster_spec = union_resolver.cluster_spec()
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testInitSimpleClusterResolver(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+
+    simple_resolver = cluster_resolver.SimpleClusterResolver(
+        base_cluster_spec, task_type="ps",
+        task_id=1, environment="cloud",
+        num_accelerators={"GPU": 8},
+        rpc_layer="grpc",
+    )
+
+    self.assertEqual(simple_resolver.task_type, "ps")
+    self.assertEqual(simple_resolver.task_id, 1)
+    self.assertEqual(simple_resolver.environment, "cloud")
+    self.assertEqual(simple_resolver.num_accelerators(), {"GPU": 8})
+    self.assertEqual(simple_resolver.rpc_layer, "grpc")
+
+  def testOverrideSimpleClusterResolver(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+
+    simple_resolver = cluster_resolver.SimpleClusterResolver(
+        base_cluster_spec, task_type="ps",
+        task_id=1, environment="cloud",
+        num_accelerators={"GPU": 8},
+        rpc_layer="grpc",
+    )
+
+    simple_resolver.task_type = "worker"
+    simple_resolver.task_id = 2
+    simple_resolver.rpc_layer = "http"
+
+    self.assertEqual(simple_resolver.task_type, "worker")
+    self.assertEqual(simple_resolver.task_id, 2)
+    self.assertEqual(simple_resolver.rpc_layer, "http")
+
+  def testSimpleOverrideMasterWithTaskIndexZero(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+
+    simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
+    actual_master = simple_resolver.master("worker", 0, rpc_layer="grpc")
+    self.assertEqual(actual_master, "grpc://worker0:2222")
+
+  def testSimpleOverrideMasterWithRpcLayer(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+
+    simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
+    actual_master = simple_resolver.master("worker", 2, rpc_layer="grpc")
+    self.assertEqual(actual_master, "grpc://worker2:2222")
+
+  def testSimpleOverrideMaster(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+
+    simple_resolver = cluster_resolver.SimpleClusterResolver(base_cluster_spec)
+    actual_master = simple_resolver.master("worker", 2)
+    self.assertEqual(actual_master, "worker2:2222")
+
+  def testUnionClusterResolverGetProperties(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+    resolver1 = cluster_resolver.SimpleClusterResolver(
+        cluster_spec_1, task_type="ps",
+        task_id=1, environment="cloud",
+        num_accelerators={"GPU": 8},
+        rpc_layer="grpc",
+    )
+
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "ps": ["ps2:2222", "ps3:2222"],
+        "worker": ["worker3:2222", "worker4:2222", "worker5:2222"]
+    })
+    resolver2 = cluster_resolver.SimpleClusterResolver(
+        cluster_spec_2, task_type="worker",
+        task_id=2, environment="local",
+        num_accelerators={"GPU": 16},
+        rpc_layer="http",
+    )
+
+    union_resolver = cluster_resolver.UnionClusterResolver(resolver1, resolver2)
+
+    self.assertEqual(union_resolver.task_type, "ps")
+    self.assertEqual(union_resolver.task_id, 1)
+    self.assertEqual(union_resolver.environment, "cloud")
+    self.assertEqual(union_resolver.num_accelerators(), {"GPU": 8})
+    self.assertEqual(union_resolver.rpc_layer, "grpc")
+
+    union_resolver.task_type = "worker"
+    union_resolver.task_id = 2
+    union_resolver.rpc_layer = "http"
+
+    self.assertEqual(union_resolver.task_type, "worker")
+    self.assertEqual(union_resolver.task_id, 2)
+    self.assertEqual(union_resolver.rpc_layer, "http")
+
+  def testTwoNonOverlappingJobMergedClusterResolver(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "ps": [
+            "ps0:2222",
+            "ps1:2222"
+        ]
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": [
+            "worker0:2222",
+            "worker1:2222",
+            "worker2:2222"
+        ]
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+    cluster_spec = union_cluster.cluster_spec()
+
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                         tasks { key: 1 value: 'worker1:2222' }
+                         tasks { key: 2 value: 'worker2:2222' } }
+    """
+    self._verifyClusterSpecEquality(cluster_spec, expected_proto)
+
+  def testMergedClusterResolverMaster(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "ps": [
+            "ps0:2222",
+            "ps1:2222"
+        ]
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": [
+            "worker0:2222",
+            "worker1:2222",
+            "worker2:2222"
+        ]
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+
+    unspecified_master = union_cluster.master()
+    self.assertEqual(unspecified_master, "")
+
+    specified_master = union_cluster.master("worker", 1)
+    self.assertEqual(specified_master, "worker1:2222")
+
+    rpc_master = union_cluster.master("worker", 1, rpc_layer="grpc")
+    self.assertEqual(rpc_master, "grpc://worker1:2222")
+
+  def testOverlappingJobMergedClusterResolver(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "worker": [
+            "worker4:2222",
+            "worker5:2222"
+        ]
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": [
+            "worker0:2222",
+            "worker1:2222",
+            "worker2:2222"
+        ]
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+    cluster_spec = union_cluster.cluster_spec()
+
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: 'worker4:2222' }
+                         tasks { key: 1 value: 'worker5:2222' }
+                         tasks { key: 2 value: 'worker0:2222' }
+                         tasks { key: 3 value: 'worker1:2222' }
+                         tasks { key: 4 value: 'worker2:2222' } }
+    """
+    self._verifyClusterSpecEquality(cluster_spec, expected_proto)
+
+  def testOverlappingSparseJobMergedClusterResolverThrowError(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "worker": {
+            7: "worker4:2222",
+            9: "worker5:2222"
+        }
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": {
+            3: "worker0:2222",
+            6: "worker1:2222",
+            7: "worker2:2222"
+        }
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+    self.assertRaises(KeyError, union_cluster.cluster_spec)
+
+  def testOverlappingDictAndListThrowError(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "worker": [
+            "worker4:2222",
+            "worker5:2222"
+        ]
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": {
+            1: "worker0:2222",
+            2: "worker1:2222",
+            3: "worker2:2222"
+        }
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+    self.assertRaises(KeyError, union_cluster.cluster_spec)
+
+  def testOverlappingJobNonOverlappingKey(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "worker": {
+            5: "worker4:2222",
+            9: "worker5:2222"
+        }
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": {
+            3: "worker0:2222",
+            6: "worker1:2222",
+            7: "worker2:2222"
+        }
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+    cluster_spec = union_cluster.cluster_spec()
+
+    expected_proto = """
+    job { name: 'worker' tasks { key: 3 value: 'worker0:2222' }
+                         tasks { key: 5 value: 'worker4:2222' }
+                         tasks { key: 6 value: 'worker1:2222' }
+                         tasks { key: 7 value: 'worker2:2222' }
+                         tasks { key: 9 value: 'worker5:2222' }}
+    """
+    self._verifyClusterSpecEquality(cluster_spec, expected_proto)
+
+  def testMixedModeNonOverlappingKey(self):
+    cluster_spec_1 = server_lib.ClusterSpec({
+        "worker": [
+            "worker4:2222",
+            "worker5:2222"
+        ]
+    })
+    cluster_spec_2 = server_lib.ClusterSpec({
+        "worker": {
+            3: "worker0:2222",
+            6: "worker1:2222",
+            7: "worker2:2222"
+        }
+    })
+    cluster_resolver_1 = cluster_resolver.SimpleClusterResolver(cluster_spec_1)
+    cluster_resolver_2 = cluster_resolver.SimpleClusterResolver(cluster_spec_2)
+
+    union_cluster = cluster_resolver.UnionClusterResolver(
+        cluster_resolver_1, cluster_resolver_2)
+    cluster_spec = union_cluster.cluster_spec()
+
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: 'worker4:2222' }
+                         tasks { key: 1 value: 'worker5:2222' }
+                         tasks { key: 3 value: 'worker0:2222' }
+                         tasks { key: 6 value: 'worker1:2222' }
+                         tasks { key: 7 value: 'worker2:2222' }}
+    """
+    self._verifyClusterSpecEquality(cluster_spec, expected_proto)
+
+  def testRetainSparseJobWithNoMerging(self):
+    base_cluster_spec = server_lib.ClusterSpec({
+        "worker": {
+            1: "worker0:2222",
+            3: "worker1:2222",
+            5: "worker2:2222"
+        }
+    })
+
+    base_cluster_resolver = cluster_resolver.SimpleClusterResolver(
+        base_cluster_spec)
+    union_cluster = cluster_resolver.UnionClusterResolver(base_cluster_resolver)
+    cluster_spec = union_cluster.cluster_spec()
+
+    expected_proto = """
+    job { name: 'worker' tasks { key: 1 value: 'worker0:2222' }
+                         tasks { key: 3 value: 'worker1:2222' }
+                         tasks { key: 5 value: 'worker2:2222' } }
+    """
+    self._verifyClusterSpecEquality(cluster_spec, expected_proto)
+
+
+# TODO(saeta): Include tests for master resolution
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver.py

@@ -0,0 +1,207 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Implementation of ClusterResolvers for GCE instance groups."""
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.training.server_lib import ClusterSpec
+from tensorflow.python.util.tf_export import tf_export
+
+
+_GOOGLE_API_CLIENT_INSTALLED = True
+try:
+  from googleapiclient import discovery  # pylint: disable=g-import-not-at-top
+  from oauth2client.client import GoogleCredentials  # pylint: disable=g-import-not-at-top
+except ImportError:
+  _GOOGLE_API_CLIENT_INSTALLED = False
+
+
+@tf_export('distribute.cluster_resolver.GCEClusterResolver')
+class GCEClusterResolver(ClusterResolver):
+  """ClusterResolver for Google Compute Engine.
+
+  This is an implementation of cluster resolvers for the Google Compute Engine
+  instance group platform. By specifying a project, zone, and instance group,
+  this will retrieve the IP address of all the instances within the instance
+  group and return a ClusterResolver object suitable for use for distributed
+  TensorFlow.
+
+  Note: this cluster resolver cannot retrieve `task_type`, `task_id` or
+  `rpc_layer`. To use it with some distribution strategies like
+  `tf.distribute.experimental.MultiWorkerMirroredStrategy`, you will need to
+  specify `task_type` and `task_id` in the constructor.
+
+  Usage example with tf.distribute.Strategy:
+
+    ```Python
+    # On worker 0
+    cluster_resolver = GCEClusterResolver("my-project", "us-west1",
+                                          "my-instance-group",
+                                          task_type="worker", task_id=0)
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=cluster_resolver)
+
+    # On worker 1
+    cluster_resolver = GCEClusterResolver("my-project", "us-west1",
+                                          "my-instance-group",
+                                          task_type="worker", task_id=1)
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=cluster_resolver)
+    ```
+  """
+
+  def __init__(self,
+               project,
+               zone,
+               instance_group,
+               port,
+               task_type='worker',
+               task_id=0,
+               rpc_layer='grpc',
+               credentials='default',
+               service=None):
+    """Creates a new GCEClusterResolver object.
+
+    This takes in a few parameters and creates a GCEClusterResolver project. It
+    will then use these parameters to query the GCE API for the IP addresses of
+    each instance in the instance group.
+
+    Args:
+      project: Name of the GCE project.
+      zone: Zone of the GCE instance group.
+      instance_group: Name of the GCE instance group.
+      port: Port of the listening TensorFlow server (default: 8470)
+      task_type: Name of the TensorFlow job this GCE instance group of VM
+        instances belong to.
+      task_id: The task index for this particular VM, within the GCE
+        instance group. In particular, every single instance should be assigned
+        a unique ordinal index within an instance group manually so that they
+        can be distinguished from each other.
+      rpc_layer: The RPC layer TensorFlow should use to communicate across
+        instances.
+      credentials: GCE Credentials. If nothing is specified, this defaults to
+        GoogleCredentials.get_application_default().
+      service: The GCE API object returned by the googleapiclient.discovery
+        function. (Default: discovery.build('compute', 'v1')). If you specify a
+        custom service object, then the credentials parameter will be ignored.
+
+    Raises:
+      ImportError: If the googleapiclient is not installed.
+    """
+    self._project = project
+    self._zone = zone
+    self._instance_group = instance_group
+    self._task_type = task_type
+    self._task_id = task_id
+    self._rpc_layer = rpc_layer
+    self._port = port
+    self._credentials = credentials
+
+    if credentials == 'default':
+      if _GOOGLE_API_CLIENT_INSTALLED:
+        self._credentials = GoogleCredentials.get_application_default()
+
+    if service is None:
+      if not _GOOGLE_API_CLIENT_INSTALLED:
+        raise ImportError('googleapiclient must be installed before using the '
+                          'GCE cluster resolver')
+      self._service = discovery.build(
+          'compute', 'v1',
+          credentials=self._credentials)
+    else:
+      self._service = service
+
+  def cluster_spec(self):
+    """Returns a ClusterSpec object based on the latest instance group info.
+
+    This returns a ClusterSpec object for use based on information from the
+    specified instance group. We will retrieve the information from the GCE APIs
+    every time this method is called.
+
+    Returns:
+      A ClusterSpec containing host information retrieved from GCE.
+    """
+    request_body = {'instanceState': 'RUNNING'}
+    request = self._service.instanceGroups().listInstances(
+        project=self._project,
+        zone=self._zone,
+        instanceGroups=self._instance_group,
+        body=request_body,
+        orderBy='name')
+
+    worker_list = []
+
+    while request is not None:
+      response = request.execute()
+
+      items = response['items']
+      for instance in items:
+        instance_name = instance['instance'].split('/')[-1]
+
+        instance_request = self._service.instances().get(
+            project=self._project,
+            zone=self._zone,
+            instance=instance_name)
+
+        if instance_request is not None:
+          instance_details = instance_request.execute()
+          ip_address = instance_details['networkInterfaces'][0]['networkIP']
+          instance_url = '%s:%s' % (ip_address, self._port)
+          worker_list.append(instance_url)
+
+      request = self._service.instanceGroups().listInstances_next(
+          previous_request=request,
+          previous_response=response)
+
+    worker_list.sort()
+    return ClusterSpec({self._task_type: worker_list})
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    task_type = task_type if task_type is not None else self._task_type
+    task_id = task_id if task_id is not None else self._task_id
+
+    if task_type is not None and task_id is not None:
+      master = self.cluster_spec().task_address(task_type, task_id)
+      if rpc_layer or self._rpc_layer:
+        return '%s://%s' % (rpc_layer or self._rpc_layer, master)
+      else:
+        return master
+
+    return ''
+
+  @property
+  def task_type(self):
+    return self._task_type
+
+  @property
+  def task_id(self):
+    return self._task_id
+
+  @task_type.setter
+  def task_type(self, task_type):
+    raise RuntimeError(
+        'You cannot reset the task_type of the GCEClusterResolver after it has '
+        'been created.')
+
+  @task_id.setter
+  def task_id(self, task_id):
+    self._task_id = task_id
+
+  @property
+  def rpc_layer(self):
+    return self._rpc_layer
+
+  @rpc_layer.setter
+  def rpc_layer(self, rpc_layer):
+    self._rpc_layer = rpc_layer

tensorflow/python/distribute/cluster_resolver/gce_cluster_resolver_test.py

@@ -0,0 +1,342 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for GCEClusterResolver."""
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import UnionClusterResolver
+from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GCEClusterResolver
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+mock = test.mock
+
+
+class GCEClusterResolverTest(test.TestCase):
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
+  def standard_mock_instance_groups(self, instance_map=None):
+    if instance_map is None:
+      instance_map = [
+          {'instance': 'https://gce.example.com/res/gce-instance-1'}
+      ]
+
+    mock_instance_group_request = mock.MagicMock()
+    mock_instance_group_request.execute.return_value = {
+        'items': instance_map
+    }
+
+    service_attrs = {
+        'listInstances.return_value': mock_instance_group_request,
+        'listInstances_next.return_value': None,
+    }
+    mock_instance_groups = mock.Mock(**service_attrs)
+    return mock_instance_groups
+
+  def standard_mock_instances(self, instance_to_ip_map=None):
+    if instance_to_ip_map is None:
+      instance_to_ip_map = {
+          'gce-instance-1': '10.123.45.67'
+      }
+
+    mock_get_request = mock.MagicMock()
+    mock_get_request.execute.return_value = {
+        'networkInterfaces': [
+            {'networkIP': '10.123.45.67'}
+        ]
+    }
+
+    def get_side_effect(project, zone, instance):
+      del project, zone  # Unused
+
+      if instance in instance_to_ip_map:
+        mock_get_request = mock.MagicMock()
+        mock_get_request.execute.return_value = {
+            'networkInterfaces': [
+                {'networkIP': instance_to_ip_map[instance]}
+            ]
+        }
+        return mock_get_request
+      else:
+        raise RuntimeError('Instance %s not found!' % instance)
+
+    service_attrs = {
+        'get.side_effect': get_side_effect,
+    }
+    mock_instances = mock.MagicMock(**service_attrs)
+    return mock_instances
+
+  def standard_mock_service_client(
+      self,
+      mock_instance_groups=None,
+      mock_instances=None):
+
+    if mock_instance_groups is None:
+      mock_instance_groups = self.standard_mock_instance_groups()
+    if mock_instances is None:
+      mock_instances = self.standard_mock_instances()
+
+    mock_client = mock.MagicMock()
+    mock_client.instanceGroups.return_value = mock_instance_groups
+    mock_client.instances.return_value = mock_instances
+    return mock_client
+
+  def gen_standard_mock_service_client(self, instances=None):
+    name_to_ip = {}
+    instance_list = []
+    for instance in instances:
+      name_to_ip[instance['name']] = instance['ip']
+      instance_list.append({
+          'instance': 'https://gce.example.com/gce/res/' + instance['name']
+      })
+
+    mock_instance = self.standard_mock_instances(name_to_ip)
+    mock_instance_group = self.standard_mock_instance_groups(instance_list)
+
+    return self.standard_mock_service_client(mock_instance_group, mock_instance)
+
+  def testSimpleSuccessfulRetrieval(self):
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        port=8470,
+        credentials=None,
+        service=self.standard_mock_service_client())
+
+    actual_cluster_spec = gce_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: '10.123.45.67:8470' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testMasterRetrieval(self):
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_id=0,
+        port=8470,
+        credentials=None,
+        service=self.standard_mock_service_client())
+    self.assertEqual(gce_cluster_resolver.master(), 'grpc://10.123.45.67:8470')
+
+  def testMasterRetrievalWithCustomTasks(self):
+    name_to_ip = [
+        {'name': 'instance1', 'ip': '10.1.2.3'},
+        {'name': 'instance2', 'ip': '10.2.3.4'},
+        {'name': 'instance3', 'ip': '10.3.4.5'},
+    ]
+
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(name_to_ip))
+
+    self.assertEqual(
+        gce_cluster_resolver.master('worker', 2, 'test'),
+        'test://10.3.4.5:8470')
+
+  def testOverrideParameters(self):
+    name_to_ip = [
+        {'name': 'instance1', 'ip': '10.1.2.3'},
+        {'name': 'instance2', 'ip': '10.2.3.4'},
+        {'name': 'instance3', 'ip': '10.3.4.5'},
+    ]
+
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='testworker',
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(name_to_ip))
+
+    gce_cluster_resolver.task_id = 1
+    gce_cluster_resolver.rpc_layer = 'test'
+
+    self.assertEqual(gce_cluster_resolver.task_type, 'testworker')
+    self.assertEqual(gce_cluster_resolver.task_id, 1)
+    self.assertEqual(gce_cluster_resolver.rpc_layer, 'test')
+    self.assertEqual(gce_cluster_resolver.master(), 'test://10.2.3.4:8470')
+
+  def testOverrideParametersWithZeroOrEmpty(self):
+    name_to_ip = [
+        {'name': 'instance1', 'ip': '10.1.2.3'},
+        {'name': 'instance2', 'ip': '10.2.3.4'},
+        {'name': 'instance3', 'ip': '10.3.4.5'},
+    ]
+
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='',
+        task_id=1,
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(name_to_ip))
+
+    self.assertEqual(gce_cluster_resolver.master(
+        task_type='', task_id=0), 'grpc://10.1.2.3:8470')
+
+  def testCustomJobNameAndPortRetrieval(self):
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='custom',
+        port=2222,
+        credentials=None,
+        service=self.standard_mock_service_client())
+
+    actual_cluster_spec = gce_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'custom' tasks { key: 0 value: '10.123.45.67:2222' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testMultipleInstancesRetrieval(self):
+    name_to_ip = [
+        {'name': 'instance1', 'ip': '10.1.2.3'},
+        {'name': 'instance2', 'ip': '10.2.3.4'},
+        {'name': 'instance3', 'ip': '10.3.4.5'},
+    ]
+
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(name_to_ip))
+
+    actual_cluster_spec = gce_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' }
+                         tasks { key: 1 value: '10.2.3.4:8470' }
+                         tasks { key: 2 value: '10.3.4.5:8470' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testUnionMultipleInstanceRetrieval(self):
+    worker1_name_to_ip = [
+        {'name': 'instance1', 'ip': '10.1.2.3'},
+        {'name': 'instance2', 'ip': '10.2.3.4'},
+        {'name': 'instance3', 'ip': '10.3.4.5'},
+    ]
+
+    worker2_name_to_ip = [
+        {'name': 'instance4', 'ip': '10.4.5.6'},
+        {'name': 'instance5', 'ip': '10.5.6.7'},
+        {'name': 'instance6', 'ip': '10.6.7.8'},
+    ]
+
+    ps_name_to_ip = [
+        {'name': 'ps1', 'ip': '10.100.1.2'},
+        {'name': 'ps2', 'ip': '10.100.2.3'},
+    ]
+
+    worker1_gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='worker',
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(worker1_name_to_ip))
+
+    worker2_gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='worker',
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(worker2_name_to_ip))
+
+    ps_gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='ps',
+        port=2222,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(ps_name_to_ip))
+
+    union_cluster_resolver = UnionClusterResolver(worker1_gce_cluster_resolver,
+                                                  worker2_gce_cluster_resolver,
+                                                  ps_gce_cluster_resolver)
+
+    actual_cluster_spec = union_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: '10.100.1.2:2222' }
+                     tasks { key: 1 value: '10.100.2.3:2222' } }
+    job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' }
+                         tasks { key: 1 value: '10.2.3.4:8470' }
+                         tasks { key: 2 value: '10.3.4.5:8470' }
+                         tasks { key: 3 value: '10.4.5.6:8470' }
+                         tasks { key: 4 value: '10.5.6.7:8470' }
+                         tasks { key: 5 value: '10.6.7.8:8470' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testSettingTaskTypeRaiseError(self):
+    name_to_ip = [
+        {
+            'name': 'instance1',
+            'ip': '10.1.2.3'
+        },
+        {
+            'name': 'instance2',
+            'ip': '10.2.3.4'
+        },
+        {
+            'name': 'instance3',
+            'ip': '10.3.4.5'
+        },
+    ]
+
+    gce_cluster_resolver = GCEClusterResolver(
+        project='test-project',
+        zone='us-east1-d',
+        instance_group='test-instance-group',
+        task_type='testworker',
+        port=8470,
+        credentials=None,
+        service=self.gen_standard_mock_service_client(name_to_ip))
+
+    with self.assertRaisesRegex(
+        RuntimeError, 'You cannot reset the task_type '
+        'of the GCEClusterResolver after it has '
+        'been created.'):
+      gce_cluster_resolver.task_type = 'foobar'
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver.py

@@ -0,0 +1,214 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Implementation of Cluster Resolvers for Kubernetes."""
+
+import enum
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
+from tensorflow.python.training import server_lib
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('distribute.cluster_resolver.KubernetesExecutableLocation')
+class ExecutableLocation(enum.Enum):
+  """Defines where the executable runs on.
+
+  This is used to determine how to resolve the configuration
+  to talk with the kube api server.
+
+  `WITHIN_CLUSTER` means that the TensorFlow code you are running is running
+  in a pod within the cluster itself.
+  `OFF_CLUSTER` means any other enviroment outside the cluster.
+  """
+
+  WITHIN_CLUSTER = 0
+  OFF_CLUSTER = 1
+
+
+@tf_export('distribute.cluster_resolver.KubernetesClusterResolver')
+class KubernetesClusterResolver(ClusterResolver):
+  """ClusterResolver for Kubernetes.
+
+  This is an implementation of cluster resolvers for Kubernetes. When given the
+  the Kubernetes namespace and label selector for pods, we will retrieve the
+  pod IP addresses of all running pods matching the selector, and return a
+  ClusterSpec based on that information.
+
+  Note: it cannot retrieve `task_type`, `task_id` or `rpc_layer`. To use it
+  with some distribution strategies like
+  `tf.distribute.experimental.MultiWorkerMirroredStrategy`, you will need to
+  specify `task_type` and `task_id` by setting these attributes.
+
+  Usage example with tf.distribute.Strategy:
+
+    ```Python
+    # On worker 0
+    cluster_resolver = KubernetesClusterResolver(
+        {"worker": ["job-name=worker-cluster-a", "job-name=worker-cluster-b"]})
+    cluster_resolver.task_type = "worker"
+    cluster_resolver.task_id = 0
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=cluster_resolver)
+
+    # On worker 1
+    cluster_resolver = KubernetesClusterResolver(
+        {"worker": ["job-name=worker-cluster-a", "job-name=worker-cluster-b"]})
+    cluster_resolver.task_type = "worker"
+    cluster_resolver.task_id = 1
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=cluster_resolver)
+    ```
+  """
+
+  def __init__(
+      self,
+      job_to_label_mapping=None,
+      tf_server_port=8470,
+      rpc_layer='grpc',
+      override_client=None,
+      executable_location=ExecutableLocation.WITHIN_CLUSTER,
+  ):
+    """Initializes a new KubernetesClusterResolver.
+
+    This initializes a new Kubernetes ClusterResolver. The ClusterResolver
+    will attempt to talk to the Kubernetes master to retrieve all the instances
+    of pods matching a label selector.
+
+    Args:
+      job_to_label_mapping: A mapping of TensorFlow jobs to label selectors.
+        This allows users to specify many TensorFlow jobs in one Cluster
+        Resolver, and each job can have pods belong with different label
+        selectors. For example, a sample mapping might be
+        ```
+        {'worker': ['job-name=worker-cluster-a', 'job-name=worker-cluster-b'],
+         'ps': ['job-name=ps-1', 'job-name=ps-2']}
+        ```
+      tf_server_port: The port the TensorFlow server is listening on.
+      rpc_layer: (Optional) The RPC layer TensorFlow should use to communicate
+        between tasks in Kubernetes. Defaults to 'grpc'.
+      override_client: The Kubernetes client (usually automatically retrieved
+        using `from kubernetes import client as k8sclient`). If you pass this
+        in, you are responsible for setting Kubernetes credentials manually and
+        calling `k8sconfig.load_kube_config()` or
+        `k8sconfig.load_incluster_config()` before using this ClusterResolver.
+      executable_location: Parameter that specifies whether or not this
+        TensorFlow code is running from within a K8S cluster or not.
+
+    Raises:
+      ImportError: If the Kubernetes Python client is not installed and no
+        `override_client` is passed in.
+      RuntimeError: If autoresolve_task is not a boolean or a callable.
+      ValueError: If `executable_location` is not a valid value.
+    """
+    try:
+      from kubernetes import config as k8sconfig  # pylint: disable=g-import-not-at-top
+
+      if not override_client:
+        if executable_location == ExecutableLocation.OFF_CLUSTER:
+          k8sconfig.load_kube_config()
+        elif executable_location == ExecutableLocation.WITHIN_CLUSTER:
+          k8sconfig.load_incluster_config()
+        else:
+          raise ValueError('The executable location provided is invalid.')
+
+    except ImportError:
+      if not override_client:
+        raise ImportError('The Kubernetes Python client must be installed '
+                          'before using the Kubernetes Cluster Resolver. '
+                          'To install the Kubernetes Python client, run '
+                          '`pip install kubernetes` on your command line.')
+
+    if not job_to_label_mapping:
+      job_to_label_mapping = {'worker': ['job-name=tensorflow']}
+
+    self._job_to_label_mapping = job_to_label_mapping
+    self._tf_server_port = tf_server_port
+    self._override_client = override_client
+
+    self.task_type = None
+    self.task_id = None
+    self.rpc_layer = rpc_layer
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Returns the master address to use when creating a session.
+
+    You must have set the task_type and task_id object properties before
+    calling this function, or pass in the `task_type` and `task_id`
+    parameters when using this function. If you do both, the function parameters
+    will override the object properties.
+
+    Note: this is only useful for TensorFlow 1.x.
+
+    Args:
+      task_type: (Optional) The type of the TensorFlow task of the master.
+      task_id: (Optional) The index of the TensorFlow task of the master.
+      rpc_layer: (Optional) The RPC protocol for the given cluster.
+
+    Returns:
+      The name or URL of the session master.
+    """
+    task_type = task_type if task_type is not None else self.task_type
+    task_id = task_id if task_id is not None else self.task_id
+
+    if task_type is not None and task_id is not None:
+      return format_master_url(
+          self.cluster_spec().task_address(task_type, task_id),
+          rpc_layer or self.rpc_layer)
+
+    return ''
+
+  def cluster_spec(self):
+    """Returns a ClusterSpec object based on the latest info from Kubernetes.
+
+    We retrieve the information from the Kubernetes master every time this
+    method is called.
+
+    Returns:
+      A ClusterSpec containing host information returned from Kubernetes.
+
+    Raises:
+      RuntimeError: If any of the pods returned by the master is not in the
+        `Running` phase.
+    """
+    if self._override_client:
+      client = self._override_client
+    else:
+      from kubernetes import config as k8sconfig  # pylint: disable=g-import-not-at-top
+      from kubernetes import client as k8sclient  # pylint: disable=g-import-not-at-top
+
+      k8sconfig.load_kube_config()
+      client = k8sclient.CoreV1Api()
+
+    cluster_map = {}
+
+    for tf_job in self._job_to_label_mapping:
+      all_pods = []
+      for selector in self._job_to_label_mapping[tf_job]:
+        ret = client.list_pod_for_all_namespaces(label_selector=selector)
+        selected_pods = []
+
+        # Sort the list by the name to make sure it doesn't change call to call.
+        for pod in sorted(ret.items, key=lambda x: x.metadata.name):
+          if pod.status.phase == 'Running':
+            selected_pods.append(
+                '%s:%s' % (pod.status.host_ip, self._tf_server_port))
+          else:
+            raise RuntimeError('Pod "%s" is not running; phase: "%s"' %
+                               (pod.metadata.name, pod.status.phase))
+        all_pods.extend(selected_pods)
+      cluster_map[tf_job] = all_pods
+
+    return server_lib.ClusterSpec(cluster_map)

tensorflow/python/distribute/cluster_resolver/kubernetes_cluster_resolver_test.py

@@ -0,0 +1,214 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for K8sClusterResolver."""
+
+import sys
+
+from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import ExecutableLocation
+from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+mock = test.mock
+
+
+def _mock_kubernetes_module():
+  sys.modules['kubernetes'] = mock.MagicMock()
+
+
+def _mock_kubernetes_client(ret):
+  mock_client = mock.MagicMock()
+  mock_client.list_pod_for_all_namespaces.side_effect = (
+      lambda *args, **kwargs: ret[kwargs['label_selector']])
+  return mock_client
+
+
+def _get_mock_pod_item(name, phase, host_ip):
+  mock_status = mock.Mock()
+  mock_status.configure_mock(phase=phase, host_ip=host_ip)
+
+  mock_metadata = mock.Mock()
+  mock_metadata.configure_mock(name=name)
+
+  mock_item = mock.Mock()
+  mock_item.configure_mock(status=mock_status, metadata=mock_metadata)
+  return mock_item
+
+
+def _create_pod_list(*args):
+  return mock.MagicMock(items=[_get_mock_pod_item(*x) for x in args])
+
+
+class KubernetesClusterResolverTest(test.TestCase):
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    """Verifies that the ClusterSpec generates the correct proto.
+
+    We are testing this four different ways to ensure that the ClusterSpec
+    returned by the TPUClusterResolver behaves identically to a normal
+    ClusterSpec when passed into the generic ClusterSpec libraries.
+
+    Args:
+      cluster_spec: ClusterSpec returned by the TPUClusterResolver
+      expected_proto: Expected protobuf
+    """
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(expected_proto,
+                           server_lib.ClusterSpec(
+                               cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(expected_proto,
+                           server_lib.ClusterSpec(
+                               cluster_spec.as_dict()).as_cluster_def())
+
+  def testSingleItemSuccessfulRetrievalInCluster(self):
+    ret = _create_pod_list(
+        ('tensorflow-abc123', 'Running', '10.1.2.3'),
+    )
+
+    cluster_resolver = KubernetesClusterResolver(
+        override_client=_mock_kubernetes_client({'job-name=tensorflow': ret}),
+        executable_location=ExecutableLocation.WITHIN_CLUSTER,
+    )
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.1.2.3:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
+
+  def testValueErrorRaisedOnInvalidExecutableLocation(self):
+
+    _mock_kubernetes_module()
+
+    with self.assertRaisesRegexp(ValueError, '.*'):
+      KubernetesClusterResolver(executable_location=None)
+
+  def testSingleItemSuccessfulRetrieval(self):
+    ret = _create_pod_list(('tensorflow-abc123', 'Running', '10.1.2.3'),)
+
+    cluster_resolver = KubernetesClusterResolver(
+        override_client=_mock_kubernetes_client(
+            {'job-name=tensorflow': ret}))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.1.2.3:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
+
+  def testSuccessfulRetrievalWithSort(self):
+    ret = _create_pod_list(
+        ('tensorflow-abc123', 'Running', '10.1.2.3'),
+        ('tensorflow-def456', 'Running', '10.1.2.4'),
+        ('tensorflow-999999', 'Running', '10.1.2.5'))
+
+    cluster_resolver = KubernetesClusterResolver(
+        override_client=_mock_kubernetes_client(
+            {'job-name=tensorflow': ret}))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.1.2.5:8470' }
+      tasks { key: 1 value: '10.1.2.3:8470' }
+      tasks { key: 2 value: '10.1.2.4:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
+
+  def testGetMasterWithOverrideParameters(self):
+    ret = _create_pod_list(
+        ('worker-0', 'Running', '10.1.2.3'),
+        ('worker-1', 'Running', '10.1.2.4'),
+        ('worker-2', 'Running', '10.1.2.5'))
+
+    cluster_resolver = KubernetesClusterResolver(
+        override_client=_mock_kubernetes_client(
+            {'job-name=tensorflow': ret}))
+    cluster_resolver.task_type = 'worker'
+    cluster_resolver.task_id = 0
+    self.assertEqual(cluster_resolver.task_type, 'worker')
+    self.assertEqual(cluster_resolver.task_id, 0)
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+    self.assertEqual(cluster_resolver.master('worker', 2),
+                     'grpc://10.1.2.5:8470')
+
+  def testNonRunningPod(self):
+    ret = _create_pod_list(('tensorflow-abc123', 'Failed', '10.1.2.3'),)
+
+    cluster_resolver = KubernetesClusterResolver(
+        override_client=_mock_kubernetes_client(
+            {'job-name=tensorflow': ret}))
+
+    error_msg = 'Pod "tensorflow-abc123" is not running; phase: "Failed"'
+    with self.assertRaisesRegex(RuntimeError, error_msg):
+      cluster_resolver.cluster_spec()
+
+  def testMultiplePodSelectorsAndWorkers(self):
+    worker1 = _create_pod_list(
+        ('tensorflow-abc123', 'Running', '10.1.2.3'),
+        ('tensorflow-def456', 'Running', '10.1.2.4'),
+        ('tensorflow-999999', 'Running', '10.1.2.5'))
+    worker2 = _create_pod_list(
+        ('tensorflow-abc124', 'Running', '10.1.2.6'),
+        ('tensorflow-def457', 'Running', '10.1.2.7'),
+        ('tensorflow-999990', 'Running', '10.1.2.8'))
+    ps = _create_pod_list(
+        ('tensorflow-ps-1', 'Running', '10.1.2.1'),
+        ('tensorflow-ps-2', 'Running', '10.1.2.2'))
+
+    cluster_resolver = KubernetesClusterResolver(
+        job_to_label_mapping={
+            'worker': ['job-name=worker1', 'job-name=worker2'],
+            'ps': ['job-name=ps']
+        },
+        override_client=_mock_kubernetes_client({
+            'job-name=worker1': worker1,
+            'job-name=worker2': worker2,
+            'job-name=ps': ps
+        }))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'ps'
+      tasks { key: 0 value: '10.1.2.1:8470' }
+      tasks { key: 1 value: '10.1.2.2:8470' }
+    }
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.1.2.5:8470' }
+      tasks { key: 1 value: '10.1.2.3:8470' }
+      tasks { key: 2 value: '10.1.2.4:8470' }
+      tasks { key: 3 value: '10.1.2.8:8470' }
+      tasks { key: 4 value: '10.1.2.6:8470' }
+      tasks { key: 5 value: '10.1.2.7:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/cluster_resolver/sagemaker_cluster_resolver.py

@@ -0,0 +1,204 @@
+# 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.
+# ==============================================================================
+"""Implementation of Cluster Resolvers for SageMaker Environment."""
+
+import json
+import os
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.training.server_lib import ClusterSpec
+
+# List of envs
+# https://github.com/aws/sagemaker-training-toolkit/blob/master/ENVIRONMENT_VARIABLES.md
+# Only support Multi-Worker Mirrored Strategy
+
+_SESSION_MASTER_KEY = 'session_master'
+_RPC_LAYER_KEY = 'rpc_layer'
+_TASK_KEY = 'task'
+_CLUSTER_KEY = 'cluster'
+_WORKER_KEY = 'worker'
+_INDEX_KEY = 'index'
+_TYPE_KEY = 'type'
+
+_SM_CURRENT_HOST = 'SM_CURRENT_HOST'
+_SM_HOSTS = 'SM_HOSTS'
+
+
+def format_master_url(master, rpc_layer=None):
+  if rpc_layer:
+    return '%s://%s' % (rpc_layer, master)
+  else:
+    return master
+
+
+def _load_tf_config(port):
+  # Create a tf_config from SM Variables
+  assert all([x in os.environ for x in [_SM_CURRENT_HOST, _SM_HOSTS]
+             ]), 'Not a SageMaker Environment'
+  hosts = sorted(json.loads(
+      os.environ[_SM_HOSTS])) if os.environ[_SM_HOSTS] != '' else []
+  current_host = os.environ[_SM_CURRENT_HOST]
+
+  if current_host not in hosts:
+    return {}
+
+  host_index = hosts.index(current_host)
+  # Assign ports
+  hosts = ['%s:%s' % (host, port) for host in hosts]
+
+  tf_config = {
+      _CLUSTER_KEY: {
+          _WORKER_KEY: hosts
+      },
+      _TASK_KEY: {
+          _TYPE_KEY: _WORKER_KEY,
+          _INDEX_KEY: host_index
+      }
+  }
+  return tf_config
+
+
+def _get_value_in_tfconfig(key, port, default=None):
+  tf_config = _load_tf_config(port)
+  return tf_config[key] if key in tf_config else default
+
+
+class SageMakerClusterResolver(ClusterResolver):
+  """Implementation of a ClusterResolver which reads the Sagemaker EnvVars. This is an implementation of cluster resolvers when running in a SageMaker environment to set information about the cluster.
+
+  The cluster spec returned will be initialized from the SageMaker
+  environment variables.
+  Currently this Cluster Resolver only supports Multi-Worker Mirrored Strategy.
+  It assumes all nodes in a SageMaker Cluster are workers.
+  """
+
+  def __init__(self,
+               port=2223,
+               task_type=None,
+               task_id=None,
+               rpc_layer=None,
+               environment=None):
+    """Creates a new SageMakerClusterResolver.
+
+    Args:
+      port: (integer, optional) Override default port usage of 2223
+      task_type: (String, optional) Overrides the task type.
+      task_id: (Integer, optional) Overrides the task index.
+      rpc_layer: (String, optional) Overrides the rpc layer TensorFlow uses.
+      environment: (String, optional) Overrides the environment TensorFlow
+        operates in.
+    """
+    self._task_type = task_type
+    self._task_id = task_id
+    self._rpc_layer = rpc_layer
+    self._environment = environment
+    self._port = str(port)
+
+  @property
+  def task_type(self):
+    if self._task_type is None:
+      task_info = _get_value_in_tfconfig(_TASK_KEY, self._port, {})
+      return str(task_info['type']) if 'type' in task_info else None
+    else:
+      return str(self._task_type)
+
+  @property
+  def task_id(self):
+    if self._task_id is None:
+      task_info = _get_value_in_tfconfig(_TASK_KEY, self._port, {})
+      return int(task_info['index']) if 'index' in task_info else None
+    else:
+      return int(self._task_id)
+
+  @task_type.setter
+  def task_type(self, task_type):
+    self._task_type = task_type
+
+  @task_id.setter
+  def task_id(self, task_id):
+    self._task_id = task_id
+
+  @property
+  def environment(self):
+    return self._environment
+
+  @property
+  def rpc_layer(self):
+    if self._rpc_layer is None:
+      return _get_value_in_tfconfig(_RPC_LAYER_KEY, self._port)
+    else:
+      return self._rpc_layer
+
+  @rpc_layer.setter
+  def rpc_layer(self, rpc_layer):
+    self._rpc_layer = rpc_layer
+
+  def num_accelerators(self, task_type=None, task_id=None, config_proto=None):
+    task_type = self.task_type if task_type is None else task_type
+    task_id = self.task_id if task_id is None else task_id
+    return super(SageMakerClusterResolver,
+                 self).num_accelerators(task_type, task_id, config_proto)
+
+  def cluster_spec(self):
+    """Returns a ClusterSpec based on the SageMaker environment variables.
+
+    Returns:
+      A ClusterSpec with information from the SageMaker environment variables.
+    """
+    tf_config = _load_tf_config(self._port)
+    if 'cluster' not in tf_config:
+      return ClusterSpec({})
+    return ClusterSpec(tf_config['cluster'])
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Returns the master address to use when creating a TensorFlow session.
+
+    Note: this is only useful for TensorFlow 1.x.
+
+    Args:
+      task_type: (String, optional) Overrides and sets the task_type of the
+        master.
+      task_id: (Integer, optional) Overrides and sets the task id of the master.
+      rpc_layer: (String, optional) Overrides and sets the protocol over which
+        TensorFlow nodes communicate with each other.
+
+    Returns:
+      The address of the master.
+
+    Raises:
+      RuntimeError: If the task_type or task_id is not specified and the
+        SageMaker environment variables does not contain a task section.
+    """
+
+    # If `session_master` is set, just use that.
+    session_master = _get_value_in_tfconfig(_SESSION_MASTER_KEY, self._port)
+    if session_master is not None:
+      return session_master
+
+    # Return an empty string if we are the only job in the ClusterSpec.
+    cluster_spec = self.cluster_spec()
+    if (not cluster_spec.jobs or
+        (len(cluster_spec.jobs) == 1 and
+         len(cluster_spec.job_tasks(cluster_spec.jobs[0])) == 1)):
+      return ''
+
+    # We try to auto-detect the task type and id, but uses the user-supplied one
+    # where available
+    task_type = task_type if task_type is not None else self.task_type
+    task_id = task_id if task_id is not None else self.task_id
+    rpc_layer = rpc_layer if rpc_layer is not None else self.rpc_layer
+
+    return format_master_url(
+        cluster_spec.task_address(task_type, task_id), rpc_layer)

tensorflow/python/distribute/cluster_resolver/sagemaker_cluster_resolver_test.py

@@ -0,0 +1,117 @@
+# 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.
+# ==============================================================================
+"""Tests for SageMakerClusterResolver."""
+
+import os
+
+from tensorflow.python.distribute.cluster_resolver.sagemaker_cluster_resolver import SageMakerClusterResolver
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+mock = test.mock
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class SageMakerClusterResolverTest(test.TestCase):
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
+  def testNormalClusterSpecRead(self):
+    os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
+    os.environ['SM_CURRENT_HOST'] = 'algo-2'
+
+    cluster_resolver = SageMakerClusterResolver()
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: 'algo-1:2223' }
+                         tasks { key: 1 value: 'algo-2:2223' } }
+    """
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testAutomaticMasterRead(self):
+    os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
+    os.environ['SM_CURRENT_HOST'] = 'algo-1'
+
+    cluster_resolver = SageMakerClusterResolver()
+    self.assertEqual('algo-1:2223', cluster_resolver.master())
+
+  def testSpecifiedTaskTypeAndIndexMasterRead(self):
+    os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
+    os.environ['SM_CURRENT_HOST'] = 'algo-2'
+
+    cluster_resolver = SageMakerClusterResolver()
+    self.assertEqual('algo-2:2223', cluster_resolver.master('worker', 1))
+
+  def testRpcLayerRead(self):
+    os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
+    os.environ['SM_CURRENT_HOST'] = 'algo-1'
+
+    cluster_resolver = SageMakerClusterResolver(rpc_layer='grpc')
+    self.assertEqual('grpc://algo-1:2223', cluster_resolver.master())
+
+  def testParameterOverrides(self):
+    os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
+    os.environ['SM_CURRENT_HOST'] = 'algo-1'
+
+    cluster_resolver = SageMakerClusterResolver(task_type='worker', task_id=0)
+
+    self.assertEqual('algo-1:2223', cluster_resolver.master())
+    self.assertEqual('worker', cluster_resolver.task_type)
+    self.assertEqual(0, cluster_resolver.task_id)
+
+    cluster_resolver.task_type = 'worker'
+    cluster_resolver.task_id = 1
+    cluster_resolver.rpc_layer = 'test'
+
+    self.assertEqual('test://algo-2:2223', cluster_resolver.master())
+    self.assertEqual('worker', cluster_resolver.task_type)
+    self.assertEqual(1, cluster_resolver.task_id)
+    self.assertEqual('test', cluster_resolver.rpc_layer)
+
+  def testTaskIndexOverride(self):
+    os.environ['SM_HOSTS'] = '["algo-1","algo-2"]'
+    os.environ['SM_CURRENT_HOST'] = 'algo-2'
+
+    cluster_resolver = SageMakerClusterResolver(task_id=1)
+    self.assertEqual(1, cluster_resolver.task_id)
+
+  def testZeroItemsInClusterSpecMasterRead(self):
+    os.environ['SM_HOSTS'] = ''
+    os.environ['SM_CURRENT_HOST'] = ''
+
+    cluster_resolver = SageMakerClusterResolver()
+    self.assertEqual('', cluster_resolver.master())
+
+  def testOneItemInClusterSpecMasterRead(self):
+    os.environ['SM_HOSTS'] = '["algo-1"]'
+    os.environ['SM_CURRENT_HOST'] = ''
+
+    cluster_resolver = SageMakerClusterResolver()
+    self.assertEqual('', cluster_resolver.master())
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver.py

@@ -0,0 +1,397 @@
+# Copyright 2018-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.
+# ==============================================================================
+"""Implementation of Cluster Resolvers for Slurm workload manager."""
+
+import os
+import re
+import subprocess
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import format_master_url
+from tensorflow.python.training.server_lib import ClusterSpec
+from tensorflow.python.util.tf_export import tf_export
+
+
+def expand_hostlist(hostlist):
+  """Create a list of hosts out of a SLURM hostlist.
+
+  The order of nodes is preserved and no deduplication is done
+  Input: 'n[1-2],m5,o[3-4,6,7-9]')
+  Output: ['n1', 'n2', 'm5', 'o3', 'o4', 'o6', 'o7', 'o8', 'o9']
+  """
+
+  def split_hostlist(hostlist):
+    """Split hostlist at commas outside of range expressions ('[3-5]')."""
+    in_brackets = False
+    cur_host = ''
+    for c in hostlist:
+      if in_brackets:
+        assert c != '['
+        if c == ']':
+          in_brackets = False
+      elif c == '[':
+        in_brackets = True
+      elif c == ',':
+        assert cur_host != ''
+        yield cur_host
+        cur_host = ''
+        continue
+      cur_host += c
+    if cur_host:
+      yield cur_host
+
+  def expand_range_expression(range_exp):
+    """Expand a range expression like '3-5' to values 3,4,5."""
+    for part in range_exp.split(','):
+      sub_range = part.split('-')
+      if len(sub_range) == 1:
+        sub_range = sub_range * 2
+      else:
+        assert len(sub_range) == 2
+      num_digits = len(sub_range[0])
+      for i in range(int(sub_range[0]), int(sub_range[1]) + 1):
+        yield str(i).zfill(num_digits)
+
+  hosts = []
+  try:
+    for part in split_hostlist(hostlist):
+      # Match prefix (anything but a range expression) and range expression
+      # Both are optional
+      m = re.match(r'([^,[\]]*)(\[([^\]]+)\])?$', part)
+      if m is None:
+        raise ValueError('Invalid part: %s' % part)
+      prefix = m.group(1) or ''
+      if m.group(3) is None:
+        hosts.append(prefix)
+      else:
+        hosts.extend(prefix + i for i in expand_range_expression(m.group(3)))
+  except Exception as e:
+    raise ValueError('Invalid hostlist format "%s": %s' % (hostlist, e))
+  return hosts
+
+
+def expand_tasks_per_node(tasks_per_node):
+  """Expands the tasks per node expression from SLURM.
+
+  The order is preserved so it can be matched to the hostlist
+  Input: '3(x2),2,1'
+  Output: [3, 3, 2, 1]
+  """
+  result = []
+  try:
+    for part in tasks_per_node.split(','):
+      m = re.match(r'(\d+)(\(x(\d+)\))?$', part)
+      assert m is not None
+      num_tasks = int(m.group(1))
+      num_repetitions = int(m.group(3) or 1)
+      result.extend([num_tasks] * num_repetitions)
+  except Exception as e:
+    raise ValueError('Invalid tasks-per-node list format "%s": %s' %
+                     (tasks_per_node, e))
+  return result
+
+
+def _get_slurm_var(name):
+  """Gets the SLURM variable from the environment.
+
+  Args:
+    name: Name of the step variable
+
+  Returns:
+    SLURM_<name> from os.environ
+  Raises:
+    RuntimeError if variable is not found
+  """
+  name = 'SLURM_' + name
+  try:
+    return os.environ[name]
+  except KeyError:
+    raise RuntimeError('%s not found in environment. '
+                       'Not running inside a SLURM step?' % name)
+
+
+def _get_num_slurm_tasks():
+  """Returns the number of SLURM tasks of the current job step.
+
+  Returns:
+    The number of tasks as an int
+  """
+  return int(_get_slurm_var('STEP_NUM_TASKS'))
+
+
+def _get_num_nvidia_gpus():
+  """Gets the number of NVIDIA GPUs by using CUDA_VISIBLE_DEVICES and nvidia-smi.
+
+  Returns:
+    Number of GPUs available on the node
+  Raises:
+    RuntimeError if executing nvidia-smi failed
+  """
+  try:
+    return len(os.environ['CUDA_VISIBLE_DEVICES'].split(','))
+  except KeyError:
+    pass  # Ignore and fallback to using nvidia-smi
+  try:
+    output = subprocess.check_output(['nvidia-smi', '--list-gpus'],
+                                     encoding='utf-8')
+    return sum(l.startswith('GPU ') for l in output.strip().split('\n'))
+  except subprocess.CalledProcessError as e:
+    raise RuntimeError('Could not get number of GPUs from nvidia-smi. '
+                       'Maybe it is missing?\nOutput: %s' % e.output)
+
+
+def get_num_gpus():
+  """Returns the number of GPUs visible on the current node.
+
+  Currently only implemented for NVIDIA GPUs.
+  """
+  return _get_num_nvidia_gpus()
+
+
+@tf_export('distribute.cluster_resolver.SlurmClusterResolver')
+class SlurmClusterResolver(ClusterResolver):
+  """ClusterResolver for system with Slurm workload manager.
+
+  This is an implementation of ClusterResolver for Slurm clusters. This allows
+  the specification of jobs and task counts, number of tasks per node, number
+  of GPUs on each node and number of GPUs for each task. It retrieves system
+  attributes by Slurm environment variables, resolves allocated computing node
+  names, constructs a cluster and returns a ClusterResolver object which can be
+  used for distributed TensorFlow.
+  """
+
+  def __init__(self,
+               jobs=None,
+               port_base=8888,
+               gpus_per_node=None,
+               gpus_per_task=None,
+               tasks_per_node=None,
+               auto_set_gpu=True,
+               rpc_layer='grpc'):
+    """Creates a new SlurmClusterResolver object.
+
+    For any parameter not set it will query the environment for the value.
+    It uses those parameters to check which nodes have processes reside on and
+    resolves their hostnames.
+    With the number tasks per node it offsets the port number for each process.
+    With the number of GPUs per node and per task it allocates GPUs to tasks by
+    setting environment variables.
+    Using the resolver works best (and is easier) with homogeneous tasks but
+    heterogeneous tasks (number of tasks varying per node) are also possible as
+    long as the number of GPUs per task stays constant.
+
+    Used environment variables:
+      - SLURM_PROCID
+      - (opt) SLURM_STEP_NUM_TASKS
+      - (opt) SLURM_STEP_NODELIST
+      - (opt) SLURM_STEP_TASKS_PER_NODE
+
+    Args:
+      jobs: Dictionary with job names as key and number of tasks in the job as
+        value. Defaults to as many 'worker's as there are (Slurm) tasks.
+      port_base: The first port number to start with for processes on a node.
+      gpus_per_node: Number of GPUs available on each node. Defaults to the
+        number of GPUs reported by nvidia-smi
+      gpus_per_task: Number of GPUs to be used for each task. Default is to
+        evenly distribute the gpus_per_node to tasks_per_node.
+      tasks_per_node: Number of tasks running on each node. Can be an integer if
+        the number of tasks per node is constant or a dictionary mapping
+        hostnames to number of tasks on that node. If not set the Slurm
+        environment is queried for the correct mapping.
+      auto_set_gpu: Set the visible CUDA devices automatically while resolving
+        the cluster by setting CUDA_VISIBLE_DEVICES environment variable.
+        Defaults to True.
+      rpc_layer: The protocol TensorFlow used to communicate between nodes.
+        Defaults to 'grpc'.
+
+    Returns:
+      A ClusterResolver object which can be used with distributed TensorFlow.
+
+    Raises:
+      RuntimeError: If requested more GPUs per node than available or
+        requested more tasks than assigned tasks or
+        resolving missing values from the environment failed.
+    """
+
+    self._rank = self._resolve_own_rank()
+
+    if jobs is None:
+      jobs = {'worker': self._resolve_num_tasks()}
+
+    self._jobs = jobs
+    self._port_base = port_base
+
+    if tasks_per_node is None:
+      self._task_configuration = self._resolve_task_configuration()
+    elif isinstance(tasks_per_node, dict):
+      # User can pass in an explicit configuration as a dict
+      self._task_configuration = tasks_per_node
+    else:
+      # User can pass a fixed number of tasks per node
+      hostlist = self._resolve_hostlist()
+      self._task_configuration = {
+          host: int(tasks_per_node) for host in hostlist
+      }
+
+    max_tasks_per_node = max(self._task_configuration.values())
+    num_tasks = sum(self._task_configuration.values())
+
+    if gpus_per_node is None:
+      gpus_per_node = get_num_gpus()
+    if gpus_per_task is None:
+      gpus_per_task = gpus_per_node // max_tasks_per_node
+    self._gpus_per_node = gpus_per_node
+    self._gpus_per_task = gpus_per_task
+
+    self._auto_set_gpu = auto_set_gpu
+    self.task_type = None
+    self.task_id = None
+    self.rpc_layer = rpc_layer
+
+    self._gpu_allocation = []
+    self._cluster_allocation = {}
+
+    if max_tasks_per_node * self._gpus_per_task > self._gpus_per_node:
+      raise RuntimeError('Requested more GPUs per node than available.')
+
+    if sum(self._jobs.values()) != num_tasks:
+      raise RuntimeError('Requested {} tasks but only {} were assigned.'.format(
+          sum(self._jobs.values()), num_tasks))
+
+  def _resolve_own_rank(self):
+    """Returns the rank of the current task in range [0, num_tasks)."""
+    return int(_get_slurm_var('PROCID'))
+
+  def _resolve_num_tasks(self):
+    """Returns the number of tasks for the current job step."""
+    return _get_num_slurm_tasks()
+
+  def _resolve_hostlist(self):
+    """Returns a list of hostnames for nodes running the current job step."""
+    return expand_hostlist(_get_slurm_var('STEP_NODELIST'))
+
+  def _resolve_task_configuration(self):
+    """Creates a mapping of hostnames to the number of tasks allocated on it.
+
+    Reads the SLURM environment to determine the nodes involved in the current
+    job step and number of tasks running on each node.
+
+    Returns a dictionary mapping each hostname to the number of tasks.
+    """
+    hostlist = self._resolve_hostlist()
+    tasks_per_node = expand_tasks_per_node(
+        _get_slurm_var('STEP_TASKS_PER_NODE'))
+    return {
+        host: num_tasks for (host, num_tasks) in zip(hostlist, tasks_per_node)
+    }
+
+  def cluster_spec(self):
+    """Returns a ClusterSpec object based on the latest instance group info.
+
+    This returns a ClusterSpec object for use based on information from the
+    specified initialization parameters and Slurm environment variables. The
+    cluster specification is resolved each time this function is called. The
+    resolver extract hostnames of nodes by scontrol and pack tasks in that
+    order until a node a has number of tasks that is equal to specification.
+    GPUs on nodes are allocated to tasks by specification through setting
+    CUDA_VISIBLE_DEVICES environment variable.
+
+    Returns:
+      A ClusterSpec containing host information retrieved from Slurm's
+        environment variables.
+    """
+
+    task_list = []
+    self._gpu_allocation = []
+    self._cluster_allocation = {}
+
+    # Sort to make sure the order is the same for each run
+    for host, num_tasks in sorted(self._task_configuration.items()):
+      for port_offset, gpu_offset in zip(
+          range(num_tasks), range(0, self._gpus_per_node, self._gpus_per_task)):
+
+        host_addr = '%s:%d' % (host, self._port_base + port_offset)
+        task_list.append(host_addr)
+        gpu_id_list = []
+
+        for gpu_id in range(gpu_offset, gpu_offset + self._gpus_per_task):
+          gpu_id_list.append(str(gpu_id))
+
+        self._gpu_allocation.append(','.join(gpu_id_list))
+
+    cluster_rank_offset_start = 0
+    cluster_rank_offset_end = 0
+
+    # Sort to make sure the order is the same for each run
+    for task_type, num_tasks in sorted(self._jobs.items()):
+      cluster_rank_offset_end = cluster_rank_offset_start + num_tasks
+
+      self._cluster_allocation[task_type] = (
+          task_list[cluster_rank_offset_start:cluster_rank_offset_end])
+
+      if cluster_rank_offset_start <= self._rank < cluster_rank_offset_end:
+        self.task_type = task_type
+        self.task_id = self._rank - cluster_rank_offset_start
+
+      cluster_rank_offset_start = cluster_rank_offset_end
+
+    if self._auto_set_gpu:
+      os.environ['CUDA_VISIBLE_DEVICES'] = self._gpu_allocation[self._rank]
+
+    return ClusterSpec(self._cluster_allocation)
+
+  def get_task_info(self):
+    """Returns job name and task_id for the process which calls this.
+
+    This returns the job name and task index for the process which calls this
+    function according to its rank and cluster specification. The job name and
+    task index are set after a cluster is constructed by cluster_spec otherwise
+    defaults to None.
+
+    Returns:
+      A string specifying job name the process belongs to and an integer
+        specifying the task index the process belongs to in that job.
+    """
+    return self.task_type, self.task_id
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Returns the master string for connecting to a TensorFlow master.
+
+    Args:
+      task_type: (Optional) Overrides the default auto-selected task type.
+      task_id: (Optional) Overrides the default auto-selected task index.
+      rpc_layer: (Optional) Overrides the default RPC protocol TensorFlow uses
+        to communicate across nodes.
+
+    Returns:
+      A connection string for connecting to a TensorFlow master.
+    """
+    task_type = task_type if task_type is not None else self.task_type
+    task_id = task_id if task_id is not None else self.task_id
+
+    if task_type is not None and task_id is not None:
+      return format_master_url(
+          self.cluster_spec().task_address(task_type, task_id),
+          rpc_layer or self.rpc_layer)
+
+    return ''
+
+  def num_accelerators(self,
+                       task_type=None,
+                       task_id=None,
+                       config_proto=None):
+    # Unused, since this is set in __init__ manually.
+    del task_type, task_id, config_proto
+    return {'GPU': self._gpus_per_task}

tensorflow/python/distribute/cluster_resolver/slurm_cluster_resolver_test.py

@@ -0,0 +1,231 @@
+# 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.
+# ==============================================================================
+"""Tests for SlurmClusterResolver."""
+
+import os
+
+from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import expand_hostlist
+from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import expand_tasks_per_node
+from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+mock = test.mock
+
+
+class SlurmClusterResolverTest(test.TestCase):
+
+  def test_expand_hostlist(self):
+    self.assertEqual(expand_hostlist('n1'), ['n1'])
+    self.assertEqual(expand_hostlist('n[1,3]'), ['n1', 'n3'])
+    self.assertEqual(expand_hostlist('n[1-3]'), ['n1', 'n2', 'n3'])
+    self.assertEqual(
+        expand_hostlist('n[1-2],m5,o[3-4,6,7-9]'),
+        ['n1', 'n2', 'm5', 'o3', 'o4', 'o6', 'o7', 'o8', 'o9'])
+    self.assertEqual(
+        expand_hostlist('n[0001-0003],m5,o[009-011]'),
+        ['n0001', 'n0002', 'n0003', 'm5', 'o009', 'o010', 'o011'])
+
+  def test_expand_tasks_per_node(self):
+    self.assertEqual(expand_tasks_per_node('2'), [2])
+    self.assertEqual(expand_tasks_per_node('2,1,3'), [2, 1, 3])
+    self.assertEqual(expand_tasks_per_node('3(x2),2,1'), [3, 3, 2, 1])
+    self.assertEqual(
+        expand_tasks_per_node('3(x2),2,11(x4)'), [3, 3, 2, 11, 11, 11, 11])
+    self.assertEqual(
+        expand_tasks_per_node('13(x10)'),
+        [13, 13, 13, 13, 13, 13, 13, 13, 13, 13])
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
+  @mock.patch.dict(
+      os.environ, {
+          'SLURM_PROCID': '0',
+          'SLURM_STEP_NUM_TASKS': '3',
+          'SLURM_STEP_TASKS_PER_NODE': '1(x3)',
+          'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
+          'CUDA_VISIBLE_DEVICES': '0',
+      })
+  def testSimpleRetrievalFromEnv(self):
+    slurm_cluster_resolver = SlurmClusterResolver()
+
+    actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: 't02n13:8888' }
+                         tasks { key: 1 value: 't02n41:8888' }
+                         tasks { key: 2 value: 't02n43:8888' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    self.assertEqual(
+        slurm_cluster_resolver.master('worker', 0, rpc_layer='grpc'),
+        'grpc://t02n13:8888')
+    self.assertEqual(slurm_cluster_resolver.num_accelerators(), {'GPU': 1})
+    self.assertEqual(os.environ['CUDA_VISIBLE_DEVICES'], '0')
+
+  @mock.patch.dict(
+      os.environ, {
+          'SLURM_PROCID': '0',
+          'SLURM_STEP_NUM_TASKS': '3',
+          'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
+      })
+  def testSimpleSuccessfulRetrieval(self):
+    slurm_cluster_resolver = SlurmClusterResolver(
+        jobs={
+            'ps': 1,
+            'worker': 2
+        },
+        port_base=8888,
+        tasks_per_node=1,
+        gpus_per_node=1,
+        gpus_per_task=1,
+        auto_set_gpu=False)
+
+    actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'ps' tasks { value: 't02n13:8888' } }
+    job { name: 'worker' tasks { key: 0 value: 't02n41:8888' }
+                         tasks { key: 1 value: 't02n43:8888' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  @mock.patch.dict(
+      os.environ, {
+          'SLURM_PROCID': '0',
+          'SLURM_STEP_NUM_TASKS': '3',
+          'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
+      })
+  def testSimpleMasterRetrieval(self):
+    slurm_cluster_resolver = SlurmClusterResolver(
+        jobs={
+            'ps': 1,
+            'worker': 2
+        },
+        port_base=8888,
+        tasks_per_node=1,
+        gpus_per_node=1,
+        gpus_per_task=1,
+        auto_set_gpu=False)
+
+    slurm_cluster_resolver.task_type = 'worker'
+    slurm_cluster_resolver.task_id = 1
+    self.assertEqual(slurm_cluster_resolver.master(), 'grpc://t02n43:8888')
+
+    slurm_cluster_resolver.rpc_layer = 'ab'
+    self.assertEqual(slurm_cluster_resolver.master('ps', 0), 'ab://t02n13:8888')
+    self.assertEqual(
+        slurm_cluster_resolver.master('ps', 0, rpc_layer='test'),
+        'test://t02n13:8888')
+
+  @mock.patch.dict(
+      os.environ, {
+          'SLURM_PROCID': '0',
+          'SLURM_STEP_NUM_TASKS': '3',
+          'SLURM_STEP_TASKS_PER_NODE': '1(x3)',
+          'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
+      })
+  def testTaskPerNodeNotSetRetrieval(self):
+    slurm_cluster_resolver = SlurmClusterResolver(
+        jobs={
+            'ps': 1,
+            'worker': 2
+        },
+        port_base=8888,
+        gpus_per_node=1,
+        gpus_per_task=1,
+        auto_set_gpu=False)
+
+    actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'ps' tasks { value: 't02n13:8888' } }
+    job { name: 'worker' tasks { key: 0 value: 't02n41:8888' }
+                         tasks { key: 1 value: 't02n43:8888' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  @mock.patch.dict(
+      os.environ, {
+          'SLURM_PROCID': '1',
+          'SLURM_STEP_NUM_TASKS': '5',
+          'SLURM_STEP_TASKS_PER_NODE': '2(x2),1',
+          'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
+          'CUDA_VISIBLE_DEVICES': '',
+      })
+  def testMultiTaskPerNodeRetrieval(self):
+    slurm_cluster_resolver = SlurmClusterResolver(
+        jobs={
+            'ps': 1,
+            'worker': 4
+        },
+        port_base=8888,
+        gpus_per_node=2,
+        gpus_per_task=1,
+        auto_set_gpu=True)
+
+    actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'ps' tasks { value: 't02n13:8888' } }
+    job { name: 'worker' tasks { key: 0 value: 't02n13:8889' }
+                         tasks { key: 1 value: 't02n41:8888' }
+                         tasks { key: 2 value: 't02n41:8889' }
+                         tasks { key: 3 value: 't02n43:8888' } }
+    """
+
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    assert os.environ['CUDA_VISIBLE_DEVICES'] == '1'
+
+  @mock.patch.dict(
+      os.environ, {
+          'SLURM_PROCID': '1',
+          'SLURM_STEP_NUM_TASKS': '5',
+          'SLURM_STEP_TASKS_PER_NODE': '2(x2),1',
+          'SLURM_STEP_NODELIST': 't02n13,t02n41,t02n43',
+          'CUDA_VISIBLE_DEVICES': '',
+      })
+  def testMultipleGpusPerTaskRetrieval(self):
+    slurm_cluster_resolver = SlurmClusterResolver(
+        jobs={
+            'ps': 1,
+            'worker': 4
+        },
+        port_base=8888,
+        gpus_per_node=4,
+        gpus_per_task=2,
+        auto_set_gpu=True)
+
+    actual_cluster_spec = slurm_cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'ps' tasks { value: 't02n13:8888' } }
+    job { name: 'worker' tasks { key: 0 value: 't02n13:8889' }
+                         tasks { key: 1 value: 't02n41:8888' }
+                         tasks { key: 2 value: 't02n41:8889' }
+                         tasks { key: 3 value: 't02n43:8888' } }
+    """
+
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    assert os.environ['CUDA_VISIBLE_DEVICES'] == '2,3'
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver.py

@@ -0,0 +1,200 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Implementation of Cluster Resolvers for TF_CONFIG Environment Variables."""
+
+
+import json
+import os
+
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import ClusterResolver
+from tensorflow.python.training.server_lib import ClusterSpec
+from tensorflow.python.util.tf_export import tf_export
+
+_TF_CONFIG_ENV = 'TF_CONFIG'
+_SESSION_MASTER_KEY = 'session_master'
+_RPC_LAYER_KEY = 'rpc_layer'
+_TASK_KEY = 'task'
+
+
+def format_master_url(master, rpc_layer=None):
+  if rpc_layer:
+    return '%s://%s' % (rpc_layer, master)
+  else:
+    return master
+
+
+def _load_tf_config():
+  return json.loads(os.environ.get(_TF_CONFIG_ENV, '{}'))
+
+
+def _get_value_in_tfconfig(key, default=None):
+  tf_config = _load_tf_config()
+  return tf_config[key] if key in tf_config else default
+
+
+@tf_export('distribute.cluster_resolver.TFConfigClusterResolver')
+class TFConfigClusterResolver(ClusterResolver):
+  """Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.
+
+  This is an implementation of cluster resolvers when using TF_CONFIG to set
+  information about the cluster. The cluster spec returned will be
+  initialized from the TF_CONFIG environment variable.
+
+  An example to set TF_CONFIG is:
+
+    ```Python
+    os.environ['TF_CONFIG'] = json.dumps({
+      'cluster': {
+          'worker': ["localhost:12345", "localhost:23456"]
+      },
+      'task': {'type': 'worker', 'index': 0}
+    })
+    ```
+
+  However, sometimes the container orchestration framework will set TF_CONFIG
+  for you. In this case, you can just create an instance without passing in any
+  arguments. You can find an example here to let Kuburnetes set TF_CONFIG for
+  you: https://github.com/tensorflow/ecosystem/tree/master/kubernetes. Then you
+  can use it with `tf.distribute.Strategy` as:
+
+    ```Python
+    # `TFConfigClusterResolver` is already the default one in the following
+    # strategy.
+    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
+        cluster_resolver=TFConfigClusterResolver())
+    ```
+  """
+
+  def __init__(self,
+               task_type=None,
+               task_id=None,
+               rpc_layer=None,
+               environment=None):
+    """Creates a new TFConfigClusterResolver.
+
+    Args:
+      task_type: (String, optional) Overrides the task type specified in the
+        TF_CONFIG environment variable.
+      task_id: (Integer, optional) Overrides the task index specified in the
+        TF_CONFIG environment variable.
+      rpc_layer: (String, optional) Overrides the rpc layer TensorFlow uses.
+      environment: (String, optional) Overrides the environment TensorFlow
+        operates in.
+    """
+    self._task_type = task_type
+    self._task_id = task_id
+    self._rpc_layer = rpc_layer
+    self._environment = environment
+
+  @property
+  def task_type(self):
+    if self._task_type is None:
+      task_info = _get_value_in_tfconfig(_TASK_KEY, {})
+      return str(task_info['type']) if 'type' in task_info else None
+    else:
+      return str(self._task_type)
+
+  @property
+  def task_id(self):
+    if self._task_id is None:
+      task_info = _get_value_in_tfconfig(_TASK_KEY, {})
+      return int(task_info['index']) if 'index' in task_info else None
+    else:
+      return int(self._task_id)
+
+  @task_type.setter
+  def task_type(self, task_type):
+    self._task_type = task_type
+
+  @task_id.setter
+  def task_id(self, task_id):
+    self._task_id = task_id
+
+  @property
+  def environment(self):
+    return self._environment
+
+  @property
+  def rpc_layer(self):
+    if self._rpc_layer is None:
+      return _get_value_in_tfconfig(_RPC_LAYER_KEY)
+    else:
+      return self._rpc_layer
+
+  @rpc_layer.setter
+  def rpc_layer(self, rpc_layer):
+    self._rpc_layer = rpc_layer
+
+  def num_accelerators(self,
+                       task_type=None,
+                       task_id=None,
+                       config_proto=None):
+    task_type = self.task_type if task_type is None else task_type
+    task_id = self.task_id if task_id is None else task_id
+    return super(TFConfigClusterResolver, self).num_accelerators(
+        task_type, task_id, config_proto)
+
+  def cluster_spec(self):
+    """Returns a ClusterSpec based on the TF_CONFIG environment variable.
+
+    Returns:
+      A ClusterSpec with information from the TF_CONFIG environment variable.
+    """
+    tf_config = _load_tf_config()
+    if 'cluster' not in tf_config:
+      return ClusterSpec({})
+    return ClusterSpec(tf_config['cluster'])
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Returns the master address to use when creating a TensorFlow session.
+
+    Note: this is only useful for TensorFlow 1.x.
+
+    Args:
+      task_type: (String, optional) Overrides and sets the task_type of the
+        master.
+      task_id: (Integer, optional) Overrides and sets the task id of the
+        master.
+      rpc_layer: (String, optional) Overrides and sets the protocol over which
+        TensorFlow nodes communicate with each other.
+
+    Returns:
+      The address of the master.
+
+    Raises:
+      RuntimeError: If the task_type or task_id is not specified and the
+        `TF_CONFIG` environment variable does not contain a task section.
+    """
+
+    # If `session_master` is set, just use that.
+    session_master = _get_value_in_tfconfig(_SESSION_MASTER_KEY)
+    if session_master is not None:
+      return session_master
+
+    # Return an empty string if we are the only job in the ClusterSpec.
+    cluster_spec = self.cluster_spec()
+    if (not cluster_spec.jobs or
+        (len(cluster_spec.jobs) == 1 and
+         len(cluster_spec.job_tasks(cluster_spec.jobs[0])) == 1)):
+      return ''
+
+    # We try to auto-detect the task type and id, but uses the user-supplied one
+    # where available
+    task_type = task_type if task_type is not None else self.task_type
+    task_id = task_id if task_id is not None else self.task_id
+    rpc_layer = rpc_layer if rpc_layer is not None else self.rpc_layer
+
+    return format_master_url(cluster_spec.task_address(task_type, task_id),
+                             rpc_layer)

tensorflow/python/distribute/cluster_resolver/tfconfig_cluster_resolver_test.py

@@ -0,0 +1,329 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for TFCONFIGClusterResolver."""
+
+import os
+
+from tensorflow.python.client import session
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+mock = test.mock
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class TFConfigClusterResolverTest(test.TestCase):
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
+  def testNormalClusterSpecRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "task": {
+        "type": "ps",
+        "index": 0
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 0 value: 'worker0:2222' }
+                         tasks { key: 1 value: 'worker1:2222' }
+                         tasks { key: 2 value: 'worker2:2222' } }
+    """
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testSparseClusterSpecRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": {"1": "worker1:2222"}
+      },
+      "task": {
+        "type": "worker",
+        "index": 1
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    expected_proto = """
+    job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
+                     tasks { key: 1 value: 'ps1:2222' } }
+    job { name: 'worker' tasks { key: 1 value: 'worker1:2222' } }
+    """
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+  def testAutomaticMasterRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "task": {
+        "type": "ps",
+        "index": 0
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('ps0:2222', cluster_resolver.master())
+
+  def testSpecifiedTaskTypeAndIndexMasterRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "task": {
+        "type": "ps",
+        "index": 0
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('worker1:2222', cluster_resolver.master('worker', 1))
+
+  def testSessionMasterRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "session_master": "sessionmaster:2222",
+      "task": {
+        "type": "ps",
+        "index": 0
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('sessionmaster:2222', cluster_resolver.master())
+
+  def testRpcLayerRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "rpc_layer": "grpc",
+      "task": {
+        "type": "ps",
+        "index": 0
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('grpc://ps0:2222', cluster_resolver.master())
+
+  def testTaskTypeIndexRpcRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "rpc_layer": "grpc",
+      "task": {
+        "type": "ps",
+        "index": 0
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('ps', cluster_resolver.task_type)
+    self.assertEqual(0, cluster_resolver.task_id)
+    self.assertEqual('grpc', cluster_resolver.rpc_layer)
+
+  def testParameterOverrides(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "rpc_layer": "grpc",
+      "task": {
+        "type": "ps",
+        "index": 1
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver(
+        task_type='ps', task_id=0)
+
+    self.assertEqual('grpc://ps0:2222', cluster_resolver.master())
+    self.assertEqual('ps', cluster_resolver.task_type)
+    self.assertEqual(0, cluster_resolver.task_id)
+
+    cluster_resolver.task_type = 'worker'
+    cluster_resolver.task_id = 1
+    cluster_resolver.rpc_layer = 'test'
+
+    self.assertEqual('test://worker1:2222', cluster_resolver.master())
+    self.assertEqual('worker', cluster_resolver.task_type)
+    self.assertEqual(1, cluster_resolver.task_id)
+    self.assertEqual('test', cluster_resolver.rpc_layer)
+
+  def testTaskTypeCastToString(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "123456": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "rpc_layer": "grpc",
+      "task": {
+        "type": 123456,
+        "index": 0
+      }
+    }
+    """
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('123456', cluster_resolver.task_type)
+
+  def testTaskIndexCastToInteger(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+      },
+      "rpc_layer": "grpc",
+      "task": {
+        "type": "ps",
+        "index": "1"
+      }
+    }
+    """
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual(1, cluster_resolver.task_id)
+
+  def testTaskIndexOverride(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "worker": ["worker0:2222", "worker1:2222"]
+      },
+      "task": {
+        "type": "worker",
+        "index": "0"
+      }
+    }
+    """
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver(
+        task_id=1)
+    self.assertEqual(1, cluster_resolver.task_id)
+
+  def testZeroItemsInClusterSpecMasterRead(self):
+    os.environ['TF_CONFIG'] = """
+    {}
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('', cluster_resolver.master())
+
+  def testOneItemInClusterSpecMasterRead(self):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "worker": ["worker0:2222"]
+      }
+    }
+    """
+
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    self.assertEqual('', cluster_resolver.master())
+
+  @mock.patch.object(config, 'list_logical_devices')
+  @mock.patch.object(session.BaseSession, 'list_devices')
+  def testNumAcceleratorsFilterTasksByEnvVar(self, mock_list_devices,
+                                             mock_eager_list_devices):
+    os.environ['TF_CONFIG'] = """
+    {
+      "cluster": {
+        "worker1": ["w10:2222"],
+        "worker2": ["w21:2222", "w22:2222", "w23:2222", "w24:2222"]
+      },
+      "rpc_layer": "grpc",
+      "task": {
+        "type": "worker1",
+        "index": "0"
+      }
+    }
+    """
+
+    devices = [
+        context.LogicalDevice('/job:worker1/task:0/device:TPU:0', 'TPU'),
+        context.LogicalDevice('/job:worker1/task:0/device:TPU:1', 'TPU'),
+        context.LogicalDevice('/job:worker1/task:0/device:GPU:0', 'GPU'),
+        context.LogicalDevice('/job:worker1/task:0/device:GPU:1', 'GPU'),
+        context.LogicalDevice('/job:worker2/task:1/device:TPU:2', 'TPU'),
+        context.LogicalDevice('/job:worker2/task:2/device:TPU:3', 'TPU'),
+        context.LogicalDevice('/job:worker2/task:3/device:GPU:2', 'GPU'),
+        context.LogicalDevice('/job:worker2/task:4/device:GPU:3', 'GPU'),
+    ]
+    device_list = [
+        session._DeviceAttributes(d.name, d.device_type, 1024, 0)
+        for d in devices
+    ]
+    mock_eager_list_devices.return_value = devices
+    mock_list_devices.return_value = device_list
+
+    resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+
+    # By default we read from TF_CONFIG
+    self.assertEqual(resolver.num_accelerators(), {'TPU': 2, 'GPU': 2})
+
+    # Override still works when we want it to
+    self.assertEqual(resolver.num_accelerators(task_type='worker2', task_id=3),
+                     {'GPU': 1})
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/cluster_resolver/tpu/BUILD

@@ -0,0 +1,57 @@
+# Description: OSS only cluster resolvers
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_additional_rpc_deps",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "tpu_cluster_resolver_py",
+    srcs = ["tpu_cluster_resolver.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core/protobuf/tpu:topology_proto_py",
+        "//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
+        "//tensorflow/python/eager:remote",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/tpu:tpu_strategy_util",
+        "//tensorflow/python/tpu:tpu_system_metadata",
+        "//tensorflow/python/tpu/client",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:compat",
+    ] + tf_additional_rpc_deps(),
+)
+
+tf_py_strict_test(
+    name = "tpu_cluster_resolver_py_test",
+    size = "small",
+    srcs = ["tpu_cluster_resolver_test.py"],
+    grpc_enabled = True,
+    main = "tpu_cluster_resolver_test.py",
+    deps = [
+        ":tpu_cluster_resolver_py",
+        "//tensorflow/core/protobuf/tpu:topology_proto_py",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/tpu/client",
+        "//tensorflow/python/training:server_lib",
+        "//tensorflow/python/util:compat",
+    ],
+)

tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py

@@ -0,0 +1,480 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Implementation of Cluster Resolvers for Cloud TPUs."""
+
+import collections
+import re
+
+from tensorflow.core.protobuf.tpu import topology_pb2
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
+from tensorflow.python.eager import remote
+from tensorflow.python.framework import config as framework_config
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.tpu import tpu_strategy_util
+from tensorflow.python.tpu import tpu_system_metadata as tpu_system_metadata_lib
+from tensorflow.python.training import server_lib
+from tensorflow.python.util import compat
+
+try:
+  from cloud_tpu_client import client  # pylint: disable=g-import-not-at-top
+except ImportError:
+  logging.debug(
+      'Falling back to TensorFlow client; we recommended you install the Cloud '
+      'TPU client directly with pip install cloud-tpu-client.')
+  from tensorflow.python.tpu.client import client  # pylint: disable=g-import-not-at-top
+
+
+def is_running_in_gce():
+  return True
+
+
+class _LocalCloudTpuClient(object):
+  """Dummy local Cloud TPU client."""
+
+  def api_available(self):
+    return False
+
+
+_TPU_DEVICE_REGEX = re.compile(
+    r'.*task:(?P<host_id>\d+)/.*device:TPU:(?P<core_id>\d+)$')
+_TPU_CONN_RETRIES = 120
+DeviceDetails = collections.namedtuple(
+    'DeviceDetails', ['device_map', 'total_cores'])
+
+
+def initialize_tpu_system(cluster_resolver=None):
+  """Initialize the TPU devices.
+
+  Args:
+    cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver,
+        which provides information about the TPU cluster.
+  Returns:
+    The tf.tpu.Topology object for the topology of the TPU cluster. If called
+    inside tf.function, it returns the serialized topology object instead.
+
+  Raises:
+    RuntimeError: If running inside a tf.function.
+    NotFoundError: If no TPU devices found in eager mode.
+  """
+  return tpu_strategy_util.initialize_tpu_system_impl(
+      cluster_resolver, TPUClusterResolver)
+
+
+def shutdown_tpu_system(cluster_resolver=None):
+  """Shuts down the TPU devices.
+
+  This will clear all caches, even those that are maintained through sequential
+  calls to tf.tpu.experimental.initialize_tpu_system, such as the compilation
+  cache.
+
+  Args:
+    cluster_resolver: A tf.distribute.cluster_resolver.TPUClusterResolver,
+        which provides information about the TPU cluster.
+
+  Raises:
+    RuntimeError: If no TPU devices found for eager execution or if run in a
+        tf.function.
+  """
+  tpu_strategy_util.shutdown_tpu_system_impl(
+      cluster_resolver, TPUClusterResolver)
+
+
+class TPUClusterResolver(cluster_resolver_lib.ClusterResolver):
+  """Cluster Resolver for Google Cloud TPUs.
+
+  This is an implementation of cluster resolvers for the Google Cloud TPU
+  service.
+
+  TPUClusterResolver supports the following distinct environments:
+  Google Compute Engine
+  Google Kubernetes Engine
+  Google internal
+
+  It can be passed into `tf.distribute.TPUStrategy` to support TF2 training on
+  Cloud TPUs.
+  """
+
+  @staticmethod
+  def connect(tpu=None,
+              zone=None,
+              project=None):
+    """Initializes TPU and returns a TPUClusterResolver.
+
+    This API will connect to remote TPU cluster and initialize the TPU
+    hardwares. Example usage:
+
+    >>> resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
+    ...     tpu='')
+
+    It can be viewed as convenient wrapper of the following code:
+
+    >>> resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+    >>> tf.config.experimental_connect_to_cluster(resolver)
+    >>> tf.tpu.experimental.initialize_tpu_system(resolver)
+
+    Args:
+      tpu: A string corresponding to the TPU to use. It can be the TPU name or
+        TPU worker gRPC address. If not set, it will try automatically resolve
+        the TPU address on Cloud TPUs.
+      zone: Zone where the TPUs are located. If omitted or empty, we will assume
+        that the zone of the TPU is the same as the zone of the GCE VM, which we
+        will try to discover from the GCE metadata service.
+      project: Name of the GCP project containing Cloud TPUs. If omitted or
+        empty, we will try to discover the project name of the GCE VM from the
+        GCE metadata service.
+
+    Returns:
+      An instance of TPUClusterResolver object.
+
+    Raises:
+      NotFoundError: If no TPU devices found in eager mode.
+    """
+    resolver = TPUClusterResolver(tpu, zone, project)
+    remote.connect_to_cluster(resolver)
+    tpu_strategy_util.initialize_tpu_system_impl(resolver, TPUClusterResolver)
+    return resolver
+
+  @staticmethod
+  def _get_device_dict_and_cores(devices):
+    """Returns a dict of hosts to cores and total cores given devices names.
+
+    Returns a namedtuple with two attributes:
+      device_map: A map of host_ids to a list of core_ids.
+      total_cores: The total number of cores within the TPU system.
+
+    Args:
+      devices: A list of devices returned by session.list_devices()
+    """
+    device_map = collections.defaultdict(list)
+    num_cores = 0
+    for device in devices:
+      match = _TPU_DEVICE_REGEX.match(device.name)
+      if match:
+        host_id = match.group('host_id')
+        core_id = match.group('core_id')
+        device_map[host_id].append(core_id)
+        num_cores += 1
+    return DeviceDetails(device_map, num_cores)
+
+  @staticmethod
+  def _verify_and_return_same_core_count(device_dict):
+    """Verifies that every device in device_dict has the same # of cores."""
+    num_cores_per_host_set = (
+        {len(core_ids) for core_ids in device_dict.values()})
+    if len(num_cores_per_host_set) != 1:
+      raise RuntimeError('TPU cores on each device is not the same. This '
+                         'should never happen. Devices: {}'.format(device_dict))
+    return num_cores_per_host_set.pop()
+
+  def __init__(self,
+               tpu=None,
+               zone=None,
+               project=None,
+               job_name='worker',
+               coordinator_name=None,
+               coordinator_address=None,
+               credentials='default',
+               service=None,
+               discovery_url=None):
+    """Creates a new TPUClusterResolver object.
+
+    The ClusterResolver will then use the parameters to query the Cloud TPU APIs
+    for the IP addresses and ports of each Cloud TPU listed.
+
+    Args:
+      tpu: A string corresponding to the TPU to use. It can be the TPU name or
+        TPU worker gRPC address. If not set, it will try automatically resolve
+        the TPU address on Cloud TPUs. If set to "local", it will assume that
+        the TPU is directly connected to the VM instead of over the network.
+      zone: Zone where the TPUs are located. If omitted or empty, we will assume
+        that the zone of the TPU is the same as the zone of the GCE VM, which we
+        will try to discover from the GCE metadata service.
+      project: Name of the GCP project containing Cloud TPUs. If omitted or
+        empty, we will try to discover the project name of the GCE VM from the
+        GCE metadata service.
+      job_name: Name of the TensorFlow job the TPUs belong to.
+      coordinator_name: The name to use for the coordinator. Set to None if the
+        coordinator should not be included in the computed ClusterSpec.
+      coordinator_address: The address of the coordinator (typically an ip:port
+        pair). If set to None, a TF server will be started. If coordinator_name
+        is None, a TF server will not be started even if coordinator_address is
+        None.
+      credentials: GCE Credentials. If None, then we use default credentials
+        from the oauth2client
+      service: The GCE API object returned by the googleapiclient.discovery
+        function. If you specify a custom service object, then the credentials
+        parameter will be ignored.
+      discovery_url: A URL template that points to the location of the discovery
+        service. It should have two parameters {api} and {apiVersion} that when
+        filled in produce an absolute URL to the discovery document for that
+        service. The environment variable 'TPU_API_DISCOVERY_URL' will override
+        this.
+
+    Raises:
+      ImportError: If the googleapiclient is not installed.
+      ValueError: If no TPUs are specified.
+      RuntimeError: If an empty TPU name is specified and this is running in a
+        Google Cloud environment.
+    """
+
+    if tpu != 'local':
+      # Default Cloud environment
+      self._cloud_tpu_client = client.Client(
+          tpu=tpu,
+          zone=zone,
+          project=project,
+          credentials=credentials,
+          service=service,
+          discovery_url=discovery_url)
+      self._tpu = self._cloud_tpu_client.name()
+    else:
+      # Directly connected TPU environment
+      self._cloud_tpu_client = _LocalCloudTpuClient()
+      self._tpu = 'local'
+
+    # By default the task_type is 'worker` and the task_id is 0 (which is the
+    # first worker in the task).
+    self.task_type = job_name
+    self.task_id = 0
+    self._coordinator_name = coordinator_name
+    if (coordinator_name and not coordinator_address):
+      self._start_local_server()
+    else:
+      self._coordinator_address = coordinator_address
+
+    self._tpu_topology = None
+
+  def __enter__(self):
+    self._cloud_tpu_client.enter()
+
+  def __exit__(self, type, value, traceback):  # pylint: disable=redefined-builtin
+    self._cloud_tpu_client.exit(type, value, traceback)
+
+  def master(self, task_type=None, task_id=None, rpc_layer=None):
+    """Get the Master string to be used for the session.
+
+    In the normal case, this returns the grpc path (grpc://1.2.3.4:8470) of
+    first instance in the ClusterSpec returned by the cluster_spec function.
+
+    If a non-TPU name is used when constructing a TPUClusterResolver, that will
+    be returned instead (e.g. If the tpus argument's value when constructing
+    this TPUClusterResolver was 'grpc://10.240.1.2:8470',
+    'grpc://10.240.1.2:8470' will be returned).
+
+    Args:
+      task_type: (Optional, string) The type of the TensorFlow task of the
+        master.
+      task_id: (Optional, integer) The index of the TensorFlow task of the
+        master.
+      rpc_layer: (Optional, string) The RPC protocol TensorFlow should use to
+        communicate with TPUs.
+
+    Returns:
+      string, the connection string to use when creating a session.
+
+    Raises:
+      ValueError: If none of the TPUs specified exists.
+    """
+
+    if self._tpu != 'local':
+      cluster_spec = self.cluster_spec()
+      if task_type is not None and task_id is not None:
+        # task_type and task_id is from the function parameter
+        master = cluster_spec.task_address(task_type, task_id)
+      elif self.task_type is not None and self.task_id is not None:
+        # task_type and task_id is from the object
+        master = cluster_spec.task_address(self.task_type, self.task_id)
+      else:
+        # by default we take the first item in the cluster with the right name
+        job_tasks = cluster_spec.job_tasks(self.task_type)
+        if not job_tasks:
+          raise ValueError('No TPUs with the specified names exist.')
+        master = job_tasks[0]
+      return cluster_resolver_lib.format_master_url(master, 'grpc')
+    else:
+      return ''
+
+  def get_master(self):
+    return self.master()
+
+  def get_job_name(self):
+    return self.task_type
+
+  def get_coordination_service_leader(self):
+    """Returns the location for coordination service.
+
+    The coordination service should be located on TPU worker0.
+
+    Returns:
+      A string indicate the location path.
+    """
+    return '/job:' + self.get_job_name() + '/task:0'
+
+  def get_tpu_system_metadata(self):
+    """Returns the metadata of the TPU system.
+
+    Users can call this method to get some facts of the TPU system, like
+    total number of cores, number of TPU workers and the devices. E.g.
+    ```python
+
+    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
+    tpu_system_metadata = resolver.get_tpu_system_metadata()
+    num_hosts = tpu_system_metadata.num_hosts
+    ```
+
+    Returns:
+      A `tf.tpu.experimental.TPUSystemMetadata` object.
+    """
+    cluster_spec = self.cluster_spec()
+    cluster_def = cluster_spec.as_cluster_def() if cluster_spec else None
+    tpu_system_metadata = (
+        tpu_system_metadata_lib._query_tpu_system_metadata(  # pylint: disable=protected-access
+            self.master(),
+            cluster_def=cluster_def,
+            query_topology=False))
+
+    return tpu_system_metadata
+
+  def cluster_spec(self):
+    """Returns a ClusterSpec object based on the latest TPU information.
+
+    We retrieve the information from the GCE APIs every time this method is
+    called.
+
+    Returns:
+      A ClusterSpec containing host information returned from Cloud TPUs,
+      or None.
+
+    Raises:
+      RuntimeError: If the provided TPU is not healthy.
+    """
+    ############################################################################
+    # There are 6 potential cases this code must handle:
+    #  0. [Local case.] When a TPU is connected directly to the VM.
+    #  1. [Normal case.] We should resolve the TPU name to a set of tasks, and
+    #      a. Create a ClusterSpec that includes the coordinator job
+    #      b. Create a ClusterSpec without the coordinator job.
+    #  2. [GKE / No API Access.] We should not resolve the TPU name to a set of
+    #     tasks and
+    #      a. Create a ClusterSpec with the coordinator
+    #      b. Create a ClusterSpec without the coordinator
+    ############################################################################
+
+    if self._tpu != 'local':
+      network_endpoints = self._cloud_tpu_client.network_endpoints()
+      worker_list = [
+          '%s:%s' % (endpoint['ipAddress'], endpoint['port'])
+          for endpoint in network_endpoints
+      ]
+      cluster_spec = {self.task_type: worker_list}
+      if self._coordinator_address:
+        # {1, 2}.a
+        cluster_spec[self._coordinator_name] = [self._coordinator_address]
+      return server_lib.ClusterSpec(cluster_spec)
+    else:
+      return server_lib.ClusterSpec({})
+
+  def num_accelerators(self,
+                       task_type=None,
+                       task_id=None,
+                       config_proto=None):
+    """Returns the number of TPU cores per worker.
+
+    Connects to the master and list all the devices present in the master,
+    and counts them up. Also verifies that the device counts per host in the
+    cluster is the same before returning the number of TPU cores per host.
+
+    Args:
+      task_type: Unused.
+      task_id: Unused.
+      config_proto: Used to create a connection to a TPU master in order to
+        retrieve the system metadata.
+
+    Raises:
+      RuntimeError: If we cannot talk to a TPU worker after retrying or if the
+        number of TPU devices per host is different.
+    """
+    if self._tpu == 'local':
+      return {
+          'TPU':
+              len([
+                  d for d in framework_config.list_logical_devices()
+                  if d.device_type == 'TPU'
+              ])
+      }
+
+    retry_count = 1
+    # TODO(b/120564445): Replace with standard library for retries.
+    while True:
+      try:
+        device_details = TPUClusterResolver._get_device_dict_and_cores(
+            cluster_resolver_lib.get_accelerator_devices(
+                self.master(), config_proto=config_proto))
+        break
+      except errors.DeadlineExceededError:
+        error_message = ('Failed to connect to master. The TPU might not be '
+                         'ready (e.g. still scheduling) or the master '
+                         'address is incorrect: got (%s)' % self.master())
+        if retry_count <= _TPU_CONN_RETRIES:
+          logging.warning(error_message)
+          logging.warning('Retrying (%d/%d)...', retry_count, _TPU_CONN_RETRIES)
+          retry_count += 1
+        else:
+          raise RuntimeError(error_message)
+
+    if device_details.total_cores:
+      return {
+          'TPU':
+              TPUClusterResolver._verify_and_return_same_core_count(
+                  device_details.device_map)
+      }
+    return {'TPU': 0}
+
+  def set_tpu_topology(self, serialized_tpu_topology):
+    """Sets the tpu topology info stored in this resolver."""
+    self._tpu_topology = topology_pb2.TopologyProto()
+    self._tpu_topology.ParseFromString(serialized_tpu_topology)
+
+  @property
+  def tpu_hardware_feature(self):
+    """Returns the tpu topology info stored."""
+    if self._tpu_topology is None:
+      return self._tpu_topology
+    return self._tpu_topology.tpu_hardware_feature
+
+  @property
+  def environment(self):
+    """Returns the current environment which TensorFlow is running in."""
+    return ''
+
+  def _start_local_server(self):
+    address = compat.as_text(self._cloud_tpu_client.get_local_ip())
+    self._server = server_lib.Server({'local': ['0.0.0.0:0']},
+                                     protocol='grpc',
+                                     config=None,
+                                     start=True)
+    # self._server.target is of the form: grpc://ipaddress:port
+    target = compat.as_bytes(self._server.target)
+    splits = target.split(compat.as_bytes(':'))
+    assert len(splits) == 3, self._server.target
+    assert splits[0] == compat.as_bytes('grpc'), self._server.target
+    self._coordinator_port = compat.as_text(splits[2])
+    self._coordinator_address = '%s:%s' % (
+        address, compat.as_text(self._coordinator_port))
+
+  def __deepcopy__(self, memo):
+    # TODO(b/73668574): Remove this once RunConfig avoids performing deepcopy.
+    return self

tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver_test.py

@@ -0,0 +1,728 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Tests for TPUClusterResolver."""
+
+import os
+
+import six
+from six.moves.urllib.error import URLError
+
+from tensorflow.core.protobuf.tpu import topology_pb2
+from tensorflow.python.client import session
+from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver as resolver
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import server_lib
+from tensorflow.python.util import compat
+
+mock = test.mock
+
+try:
+  from cloud_tpu_client import client  # pylint: disable=g-import-not-at-top
+except ImportError:
+  logging.debug(
+      'Falling back to TensorFlow client; we recommended you install the Cloud '
+      'TPU client directly with pip install cloud-tpu-client.')
+  from tensorflow.python.tpu.client import client  # pylint: disable=g-import-not-at-top
+
+
+class MockRequestClass(object):
+
+  def __init__(self, name, tpu_map):
+    self._name = name
+    self._tpu_map = tpu_map
+
+  def execute(self):
+    if self._name in self._tpu_map:
+      return self._tpu_map[self._name]
+    else:
+      raise KeyError('Resource %s was not found' % self._name)
+
+
+class MockNodeClass(object):
+
+  def __init__(self, tpu_map):
+    self._tpu_map = tpu_map
+
+  def get(self, name):
+    return MockRequestClass(name, self._tpu_map)
+
+
+def mock_request_compute_metadata(*args, **kwargs):
+  del kwargs  # Unused.
+  if args[0] == 'project/project-id':
+    return 'test-project'
+  elif args[0] == 'instance/zone':
+    return 'projects/test-project/locations/us-central1-c'
+  elif args[0] == 'instance/network-interfaces/0/ip':
+    return '10.128.1.2'
+  return ''
+
+
+def mock_is_running_in_gce():
+  return True
+
+
+def mock_is_not_running_in_gce():
+  return False
+
+
+def mock_running_in_gce_urlopen(cls, *args, **kwargs):
+  del cls, args, kwargs  # Unused.
+  mock_response = mock.MagicMock()
+  mock_response.info.return_value = {'Metadata-Flavor': 'Google'}
+  return mock_response
+
+
+def mock_not_running_in_gce_urlopen(cls, *args, **kwargs):
+  del cls, args, kwargs  # Unused.
+  raise URLError(reason='Host does not exist.')
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class TPUClusterResolverTest(test.TestCase):
+
+  def _verifyClusterSpecEquality(self, cluster_spec, expected_proto):
+    """Verifies that the ClusterSpec generates the correct proto.
+
+    We are testing this four different ways to ensure that the ClusterSpec
+    returned by the TPUClusterResolver behaves identically to a normal
+    ClusterSpec when passed into the generic ClusterSpec libraries.
+
+    Args:
+      cluster_spec: ClusterSpec returned by the TPUClusterResolver
+      expected_proto: Expected protobuf
+    """
+    self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def())
+    self.assertProtoEquals(
+        expected_proto,
+        server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def())
+
+  def mock_service_client(self, tpu_map=None):
+
+    if tpu_map is None:
+      tpu_map = {}
+
+    mock_locations = mock.MagicMock()
+    mock_locations.nodes.return_value = MockNodeClass(tpu_map)
+
+    mock_project = mock.MagicMock()
+    mock_project.locations.return_value = mock_locations
+
+    mock_client = mock.MagicMock()
+    mock_client.projects.return_value = mock_project
+
+    return mock_client
+
+  @mock.patch.object(resolver, 'is_running_in_gce', mock_is_running_in_gce)
+  def testCheckRunningInGceWithNoTpuName(self):
+    with self.assertRaisesRegex(ValueError,
+                                'Please provide a TPU Name to connect to.*'):
+      resolver.TPUClusterResolver(tpu='')
+
+  @mock.patch.object(six.moves.urllib.request, 'urlopen',
+                     mock_running_in_gce_urlopen)
+  def testIsRunningInGce(self):
+    self.assertTrue(resolver.is_running_in_gce())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRetrieveProjectAndZoneFromMetadata(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY'
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project=None,
+        zone=None,
+        tpu=['test-tpu-1'],
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map),
+        coordinator_name='coordinator')
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'coordinator'
+      tasks { key: 0 value: '10.128.1.2:%s' }
+    }
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.1.2.3:8470' }
+    }
+    """ % cluster_resolver._coordinator_port
+    self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRetrieveProjectAndZoneFromMetadataNoCoordinator(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY'
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project=None,
+        zone=None,
+        tpu=['test-tpu-1'],
+        coordinator_name=None,
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testNotReadyCloudTpu(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'CREATING'
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project=None,
+        zone=None,
+        tpu='test-tpu-1',
+        coordinator_name=None,
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    with self.assertRaises(RuntimeError):
+      cluster_resolver.cluster_spec()
+
+  def testSimpleSuccessfulRetrieval(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY'
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu=['test-tpu-1'],
+        coordinator_name='coordinator',
+        coordinator_address='10.128.1.5:10203',
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } }
+    job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+
+  def testFailedMetadata(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'health': 'HEALTHY'
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu='nonexistent-tpu',
+        coordinator_name='coordinator',
+        coordinator_address='10.128.1.5:10203',
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    with self.assertRaises(ValueError) as error_context:
+      cluster_resolver.cluster_spec()
+
+    self.assertIn('Could not lookup TPU metadata', str(error_context.exception))
+
+  def testNewNetworkEndpointFormat(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'state': 'READY',
+            'health': 'HEALTHY',
+            'networkEndpoints': [{
+                'ipAddress': '10.2.3.4',
+                'port': 8470,
+            }]
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu='test-tpu-1',
+        coordinator_name='coordinator',
+        coordinator_address='10.128.1.5:10203',
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } }
+    job { name: 'worker' tasks { key: 0 value: '10.2.3.4:8470' } }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    self.assertEqual('grpc://10.2.3.4:8470', cluster_resolver.master())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testPodResolution(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'state': 'READY',
+            'health':
+                'HEALTHY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '10.2.3.4',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.5',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.6',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.7',
+                    'port': 8470,
+                },
+            ]
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        tpu='test-tpu-1',
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map),
+        coordinator_name='coordinator')
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'coordinator',
+      tasks { key: 0 value: '10.128.1.2:%s'}
+    }
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.2.3.4:8470' }
+      tasks { key: 1 value: '10.2.3.5:8470' }
+      tasks { key: 2 value: '10.2.3.6:8470' }
+      tasks { key: 3 value: '10.2.3.7:8470' }
+    }
+    """ % cluster_resolver._coordinator_port
+    self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto))
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.4:8470')
+
+  def testPodResolutionNoCoordinator(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'state': 'READY',
+            'health':
+                'HEALTHY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '10.2.3.4',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.5',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.6',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.7',
+                    'port': 8470,
+                },
+            ]
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu='test-tpu-1',
+        coordinator_name=None,
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.2.3.4:8470' }
+      tasks { key: 1 value: '10.2.3.5:8470' }
+      tasks { key: 2 value: '10.2.3.6:8470' }
+      tasks { key: 3 value: '10.2.3.7:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.4:8470')
+
+  def testGetMasterNoEntries(self):
+    tpu_map = {}
+
+    with self.assertRaises(ValueError):
+      resolver.TPUClusterResolver(
+          project='test-project',
+          zone='us-central1-c',
+          tpu=[],
+          coordinator_name=None,
+          credentials=None,
+          service=self.mock_service_client(tpu_map=tpu_map))
+
+  # TODO(saeta): Convert to parameterized test when included in OSS TF.
+  def verifyShouldResolve(self, tpu, should_resolve):
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu=tpu,
+        coordinator_name=None,
+        credentials=None,
+        service=self.mock_service_client(tpu_map={}))
+    self.assertEqual(should_resolve,
+                     cluster_resolver._cloud_tpu_client.api_available(),
+                     "TPU: '%s'" % tpu)
+
+  def testShouldResolveGrpc(self):
+    self.verifyShouldResolve('grpc://10.1.2.3:8470', False)
+
+  def testShouldResolveName(self):
+    self.verifyShouldResolve('mytpu', True)
+
+  def testShouldResolveList(self):
+    self.verifyShouldResolve(['myothertpu'], True)
+
+  def testShouldResolveGrpcPrefix(self):
+    self.verifyShouldResolve('grpctpu', True)
+
+  def testNoCallComputeMetadata(self):
+    cluster_resolver = resolver.TPUClusterResolver(tpu='grpc://10.1.2.3:8470')
+    self.assertEqual('grpc://10.1.2.3:8470', cluster_resolver.master())
+    self.assertEqual(
+        server_lib.ClusterSpec({
+            'worker': ['10.1.2.3:8470']
+        }).as_dict(),
+        cluster_resolver.cluster_spec().as_dict())
+
+  def testGkeEnvironmentForDonut(self):
+    os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470'
+
+    self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ)
+
+    cluster_resolver = resolver.TPUClusterResolver()
+    self.assertEqual(
+        compat.as_bytes('grpc://10.120.27.5:8470'),
+        compat.as_bytes(cluster_resolver.master()))
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.120.27.5:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+    del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS']
+
+  def testGkeEnvironmentForPod(self):
+    os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = ('grpc://10.120.27.5:8470,'
+                                                     'grpc://10.120.27.6:8470,'
+                                                     'grpc://10.120.27.7:8470,'
+                                                     'grpc://10.120.27.8:8470')
+
+    self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ)
+    cluster_resolver = resolver.TPUClusterResolver()
+    self.assertEqual(
+        compat.as_bytes('grpc://10.120.27.5:8470'),
+        compat.as_bytes(cluster_resolver.master()))
+    actual_cluster_spec = cluster_resolver.cluster_spec()
+    expected_proto = """
+    job {
+      name: 'worker'
+      tasks { key: 0 value: '10.120.27.5:8470' }
+      tasks { key: 1 value: '10.120.27.6:8470' }
+      tasks { key: 2 value: '10.120.27.7:8470' }
+      tasks { key: 3 value: '10.120.27.8:8470' }
+    }
+    """
+    self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto)
+
+    del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS']
+
+  def testRpcDetectionForGrpcString(self):
+    cluster_resolver = resolver.TPUClusterResolver(
+        tpu='grpc://10.1.2.3:8470')
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.1.2.3:8470')
+
+  def testOverrideTaskTypeAndIndexAndGetMaster(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'state': 'READY',
+            'health':
+                'HEALTHY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '10.2.3.4',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.5',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.6',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.7',
+                    'port': 8470,
+                },
+            ]
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu='test-tpu-1',
+        coordinator_name=None,
+        credentials=None,
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.4:8470')
+
+    cluster_resolver.task_type = 'worker'
+    cluster_resolver.task_id = 3
+    self.assertEqual(cluster_resolver.master(), 'grpc://10.2.3.7:8470')
+
+  def testGetDeviceDictAndCoresWithTPUs(self):
+    devices = [
+        '/job:tpu_worker/task:0/device:TPU:0',
+        '/job:tpu_worker/task:1/device:TPU:1',
+        '/job:tpu_worker/task:2/device:TPU:0',
+        '/job:tpu_worker/task:3/device:TPU:1',
+        '/job:tpu_worker/task:0/device:TPU:4',
+        '/job:tpu_worker/task:1/device:TPU:5',
+        '/job:tpu_worker/task:2/device:TPU:4',
+        '/job:tpu_worker/task:3/device:TPU:5',
+    ]
+    device_list = [
+        session._DeviceAttributes(name, 'TPU', 1024, 0) for name in devices
+    ]
+
+    device_details = resolver.TPUClusterResolver._get_device_dict_and_cores(
+        device_list)
+    self.assertEqual(device_details.total_cores, 8)
+    self.assertEqual(device_details.device_map,
+                     {'0': ['0', '4'],
+                      '1': ['1', '5'],
+                      '2': ['0', '4'],
+                      '3': ['1', '5']})
+
+  def testGetDeviceDictAndCoresWithCPUsAndGPUs(self):
+    devices = [
+        '/job:tpu_worker/task:0/device:CPU:0',
+        '/job:tpu_worker/task:1/device:CPU:0',
+        '/job:tpu_worker/task:2/device:CPU:0',
+        '/job:tpu_worker/task:3/device:CPU:0',
+        '/job:tpu_worker/task:0/device:GPU:1',
+        '/job:tpu_worker/task:1/device:GPU:1',
+        '/job:tpu_worker/task:2/device:GPU:1',
+        '/job:tpu_worker/task:3/device:GPU:1',
+    ]
+    device_list = [
+        session._DeviceAttributes(name, 'XLA', 1024, 0) for name in devices
+    ]
+
+    device_dict, num_cores =\
+        resolver.TPUClusterResolver._get_device_dict_and_cores(device_list)
+    self.assertEqual(num_cores, 0)
+    self.assertEqual(device_dict, {})
+
+  def testVerifySameCoreCount(self):
+    self.assertEqual(
+        resolver.TPUClusterResolver
+        ._verify_and_return_same_core_count({0: [0, 1, 2, 3, 4, 5, 6, 7]}), 8)
+    self.assertEqual(
+        resolver.TPUClusterResolver
+        ._verify_and_return_same_core_count({
+            0: [0, 1],
+            1: [2, 3]
+        }), 2)
+    with self.assertRaises(RuntimeError):
+      resolver.TPUClusterResolver._verify_and_return_same_core_count(
+          {
+              0: [0],
+              1: [1, 2]
+          })
+
+  @mock.patch.object(config, 'list_logical_devices')
+  @mock.patch.object(session.BaseSession, 'list_devices')
+  @mock.patch.object(resolver, 'is_running_in_gce', mock_is_not_running_in_gce)
+  def testNumAcceleratorsSuccess(self, mock_list_devices,
+                                 mock_eager_list_devices):
+    devices = [
+        context.LogicalDevice('/job:tpu_worker/task:0/device:TPU:0', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:1/device:TPU:1', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:2/device:TPU:0', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:3/device:TPU:1', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:0/device:TPU:4', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:1/device:TPU:5', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:2/device:TPU:4', 'TPU'),
+        context.LogicalDevice('/job:tpu_worker/task:3/device:TPU:5', 'TPU'),
+    ]
+    device_list = [
+        session._DeviceAttributes(d.name, d.device_type, 1024, 0)
+        for d in devices
+    ]
+    mock_eager_list_devices.return_value = devices
+    mock_list_devices.return_value = device_list
+
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'state': 'READY',
+            'health':
+                'HEALTHY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '10.2.3.4',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.5',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.6',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.7',
+                    'port': 8470,
+                },
+            ]
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu='test-tpu-1',
+        service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(cluster_resolver.num_accelerators(), {'TPU': 2})
+
+  @mock.patch.object(config, 'list_logical_devices')
+  @mock.patch.object(session.BaseSession, 'list_devices')
+  @mock.patch.object(resolver, 'is_running_in_gce', mock_is_not_running_in_gce)
+  def testNumAcceleratorsRetryFailure(self, mock_list_devices,
+                                      mock_eager_list_devices):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': {
+            'health':
+                'HEALTHY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '10.2.3.4',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.5',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.6',
+                    'port': 8470,
+                },
+                {
+                    'ipAddress': '10.2.3.7',
+                    'port': 8470,
+                },
+            ]
+        }
+    }
+
+    cluster_resolver = resolver.TPUClusterResolver(
+        project='test-project',
+        zone='us-central1-c',
+        tpu='test-tpu-1',
+        service=self.mock_service_client(tpu_map=tpu_map))
+    mock_list_devices.side_effect = errors.DeadlineExceededError(
+        None, None, 'timeout')
+    mock_eager_list_devices.side_effect = errors.DeadlineExceededError(
+        None, None, 'timeout')
+    with self.assertRaises(RuntimeError):
+      cluster_resolver.num_accelerators()
+
+  def testLocalTpuResolver(self):
+    cr = resolver.TPUClusterResolver(tpu='local')
+    self.assertEqual(cr.get_master(), '')
+
+  def testTpuTopology(self):
+    cluster_resolver = resolver.TPUClusterResolver(tpu='local')
+    self.assertIsNone(cluster_resolver._tpu_topology)
+
+    # Test set with tpu topology proto.
+    cluster_resolver.set_tpu_topology(
+        serialized_tpu_topology=topology_pb2.TopologyProto(
+            mesh_shape=[1, 1, 1, 1]).SerializeToString())
+    self.assertIsInstance(cluster_resolver.tpu_hardware_feature,
+                          topology_pb2.TPUHardwareFeature)
+
+  def testEnvironment(self):
+    cluster_resolver = resolver.TPUClusterResolver(tpu='local')
+    self.assertEqual(cluster_resolver.environment, '')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/cluster_resolver/tpu_cluster_resolver.py

@@ -0,0 +1,26 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Shim so that direct imports of tpu_cluster_resolver get correct symbols.
+"""
+
+from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import initialize_tpu_system
+from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import is_running_in_gce  # pylint: disable=unused-import
+from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import shutdown_tpu_system
+from tensorflow.python.distribute.cluster_resolver.tpu.tpu_cluster_resolver import TPUClusterResolver
+from tensorflow.python.util.tf_export import tf_export
+
+tf_export('distribute.cluster_resolver.TPUClusterResolver')(TPUClusterResolver)
+tf_export('tpu.experimental.initialize_tpu_system')(initialize_tpu_system)
+tf_export('tpu.experimental.shutdown_tpu_system')(shutdown_tpu_system)

tensorflow/python/distribute/collective_all_reduce_strategy_test.py

@@ -0,0 +1,772 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for CollectiveAllReduceStrategy."""
+
+import copy
+import functools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import collective_all_reduce_strategy
+from tensorflow.python.distribute import collective_util
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import strategy_test_lib
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.distribute.v1 import input_lib as input_lib_v1
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config as tf_config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import device as tf_device
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import gradients
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+CollectiveAllReduceStrategy = (
+    collective_all_reduce_strategy.CollectiveAllReduceStrategy)
+CollectiveAllReduceExtended = (
+    collective_all_reduce_strategy.CollectiveAllReduceExtended)
+_CollectiveAllReduceStrategyExperimental = (
+    collective_all_reduce_strategy._CollectiveAllReduceStrategyExperimental)
+
+
+# TODO(b/231630416): Create more tests to cover the case that strategy uses
+# different number of GPUs than the number of physical devices.
+def create_test_objects(cluster_spec=None,
+                        task_type=None,
+                        task_id=None,
+                        num_gpus=None,
+                        num_tpus=None):
+  if num_gpus is None:
+    num_gpus = context.num_gpus()
+  if num_tpus is None:
+    num_tpus = context.context().list_physical_devices('TPU')
+  if num_tpus:
+    tpu_cluster_resolver.initialize_tpu_system()
+
+  if cluster_spec and task_type and task_id is not None:
+    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
+        cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec),
+        task_type=task_type,
+        task_id=task_id,
+        num_accelerators={'GPU': num_gpus, 'TPU': num_tpus})
+    target = 'grpc://' + cluster_spec[task_type][task_id]
+  else:
+    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
+        server_lib.ClusterSpec({}),
+        num_accelerators={'GPU': num_gpus, 'TPU': num_tpus},
+    )
+    target = ''
+
+  strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy(
+      cluster_resolver=cluster_resolver)
+
+  return strategy, target
+
+
+class CollectiveAllReduceStrategyTestBase(
+    multi_worker_test_base.MultiWorkerTestBase):
+
+  def setUp(self):
+    # We use a different key_base for each test so that collective keys won't be
+    # reused.
+    CollectiveAllReduceStrategy._collective_key_base += 100000
+    super(CollectiveAllReduceStrategyTestBase, self).setUp()
+
+  def _get_test_object(self,
+                       task_type,
+                       task_id,
+                       num_gpus=0,
+                       num_tpus=0,
+                       use_devices_arg=False):
+    strategy, target = create_test_objects(
+        cluster_spec=self._cluster_spec,
+        task_type=task_type,
+        task_id=task_id,
+        num_gpus=num_gpus,
+        num_tpus=num_tpus)
+
+    if use_devices_arg and num_gpus > 0:
+      devices = ['GPU:%d' % i for i in range(num_gpus)]
+      # Temporary workaround to manually set the `_extended` field before device
+      # initialization is exposed as a public interface.
+      strategy._extended = CollectiveAllReduceExtended(
+          container_strategy=strategy,
+          cluster_resolver=None,
+          communication_options=collective_util.Options(),
+          devices=devices)
+      # Manually set the field since the workaround bypasses the base
+      # constructor, resulting in the absence of this field.
+      strategy._extended._retrace_functions_for_each_device = (num_gpus > 1)
+
+    return strategy, target
+
+  def _test_minimize_loss_graph(self, task_type, task_id, num_gpus):
+    distribution, master_target = self._get_test_object(task_type, task_id,
+                                                        num_gpus)
+    with ops.Graph().as_default(), \
+         self.cached_session(target=master_target) as sess, \
+         distribution.scope():
+      initializer = functools.partial(
+          init_ops_v2.GlorotUniform(), (1, 1), dtype=dtypes.float32)
+      kernel = variables.Variable(
+          initial_value=initializer,
+          name='gpu_%d/kernel' % distribution.extended._num_devices_per_worker,
+          trainable=True)
+
+      def loss_fn(x):
+        y = array_ops.reshape(
+            gen_math_ops.mat_mul(x, kernel), []) - constant_op.constant(1.)
+        return y * y
+
+      # TODO(yuefengz, apassos): eager.backprop.implicit_grad is not safe for
+      # multiple graphs (b/111216820).
+      def grad_fn(x):
+        loss = loss_fn(x)
+        var_list = (
+            variables.trainable_variables() + ops.get_collection(
+                ops.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
+        grads = gradients.gradients(loss, var_list)
+        ret = list(zip(grads, var_list))
+        return ret
+
+      def update(v, g):
+        return v.assign_sub(0.05 * g, use_locking=True)
+
+      one = constant_op.constant([[1.]])
+
+      def step():
+        """Perform one optimization step."""
+        # Run forward & backward to get gradients, variables list.
+        g_v = distribution.extended.call_for_each_replica(grad_fn, args=[one])
+        # Update the variables using the gradients and the update() function.
+        before_list = []
+        after_list = []
+        for g, v in g_v:
+          fetched = distribution.extended.read_var(v)
+          before_list.append(fetched)
+          with ops.control_dependencies([fetched]):
+            # TODO(yuefengz): support non-Mirrored variable as destinations.
+            g = distribution.extended.reduce_to(
+                reduce_util.ReduceOp.SUM, g, destinations=v)
+            with ops.control_dependencies(
+                distribution.extended.update(v, update, args=(g,),
+                                             group=False)):
+              after_list.append(distribution.extended.read_var(v))
+        return before_list, after_list
+
+      before_out, after_out = step()
+
+      if (distribution.extended._local_device_type == 'GPU' and
+          context.num_gpus() < distribution.extended._num_devices_per_worker):
+        return True
+
+      sess.run(variables.global_variables_initializer())
+
+      for i in range(10):
+        b, a = sess.run((before_out, after_out))
+        if i == 0:
+          before, = b
+        after, = a
+
+      error_before = abs(before - 1)
+      error_after = abs(after - 1)
+      # Error should go down
+      self.assertLess(error_after, error_before)
+
+  def _test_variable_initialization(self, task_type, task_id, num_gpus):
+    distribution, master_target = self._get_test_object(task_type, task_id,
+                                                        num_gpus)
+    with ops.Graph().as_default(), \
+         self.cached_session(target=master_target) as sess, \
+         distribution.scope():
+
+      def model_fn():
+        x = variable_scope.get_variable(
+            'x',
+            shape=(2, 3),
+            initializer=init_ops.random_uniform_initializer(
+                1.0, 10.0, dtype=dtypes.float32))
+        return array_ops.identity(x)
+
+      x = distribution.extended.call_for_each_replica(model_fn)
+      reduced_x = distribution.reduce(reduce_util.ReduceOp.MEAN, x, axis=None)
+      x = distribution.experimental_local_results(x)[0]
+
+      sess.run(variables.global_variables_initializer())
+
+      x_value, reduced_x_value = sess.run([x, reduced_x])
+      self.assertTrue(
+          np.allclose(x_value, reduced_x_value, atol=1e-5),
+          msg=('x_value = %r, reduced_x_value = %r' % (x_value,
+                                                       reduced_x_value)))
+
+  def _test_input_fn_iterator(self,
+                              task_type,
+                              task_id,
+                              num_gpus,
+                              input_fn,
+                              expected_values,
+                              test_reinitialize=True,
+                              ignore_order=False,
+                              use_devices_arg=False):
+    distribution, master_target = self._get_test_object(
+        task_type, task_id, num_gpus, use_devices_arg=use_devices_arg)
+    devices = distribution.extended.worker_devices
+
+    with ops.Graph().as_default(), \
+         self.cached_session(target=master_target) as sess:
+      iterator = distribution.make_input_fn_iterator(input_fn)
+      sess.run(iterator.initializer)
+
+      for expected_value in expected_values:
+        next_element = iterator.get_next()
+        computed_value = sess.run([distribute_utils.select_replica(
+            r, next_element) for r in range(len(devices))])
+        if ignore_order:
+          self.assertCountEqual(list(expected_value), list(computed_value))
+        else:
+          self.assertEqual(list(expected_value), list(computed_value))
+
+      with self.assertRaises(errors.OutOfRangeError):
+        next_element = iterator.get_next()
+        sess.run([distribute_utils.select_replica(r, next_element)
+                  for r in range(len(devices))])
+
+      # After re-initializing the iterator, should be able to iterate again.
+      if test_reinitialize:
+        sess.run(iterator.initializer)
+
+        for expected_value in expected_values:
+          next_element = iterator.get_next()
+          computed_value = sess.run([
+              distribute_utils.select_replica(r, next_element)
+              for r in range(len(devices))])
+          if ignore_order:
+            self.assertCountEqual(list(expected_value), list(computed_value))
+          else:
+            self.assertEqual(list(expected_value), list(computed_value))
+
+
+class DistributedCollectiveAllReduceStrategyTest(
+    CollectiveAllReduceStrategyTestBase,
+    strategy_test_lib.DistributionTestBase,
+    parameterized.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    """Create a local cluster with 3 workers."""
+    cls._cluster_spec = multi_worker_test_base.create_in_process_cluster(
+        num_workers=3, num_ps=0)
+
+  @combinations.generate(combinations.combine(mode=['graph']))
+  def test_num_replicas_in_sync(self):
+    distribution, _ = create_test_objects(
+        cluster_spec=self._cluster_spec,
+        task_type='worker',
+        task_id=0,
+        num_gpus=2)
+    num_workers = len(self._cluster_spec.get('chief', []) +
+                      self._cluster_spec.get('worker', []))
+    self.assertEqual(2 * num_workers,
+                     distribution.num_replicas_in_sync)
+
+  @combinations.generate(combinations.combine(
+      mode=['graph'],
+      prefetch_to_device=[None, True]))
+  def test_prefetch_to_device_dataset(self, prefetch_to_device):
+    distribution, _ = self._get_test_object(
+        task_type='worker', task_id=0, num_gpus=2)
+    if prefetch_to_device is None:
+      input_options = None
+    else:
+      input_options = distribute_lib.InputOptions(
+          experimental_fetch_to_device=prefetch_to_device)
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = dataset.batch(distribution.num_replicas_in_sync)
+    dataset = distribution.experimental_distribute_dataset(
+        dataset, options=input_options)
+    if isinstance(dataset, input_lib_v1.DistributedDatasetV1):
+      item = dataset.make_initializable_iterator().get_next()
+    else:
+      self.skipTest('unsupported test combination')
+    device_types = {
+        tf_device.DeviceSpec.from_string(tensor.device).device_type for
+        tensor in item.values}
+    self.assertAllEqual(list(device_types), ['GPU'])
+
+  @combinations.generate(combinations.combine(mode=['graph']))
+  def test_prefetch_to_host_dataset(self):
+    distribution, _ = self._get_test_object(
+        task_type='worker', task_id=0, num_gpus=2)
+    input_options = distribute_lib.InputOptions(
+        experimental_fetch_to_device=False)
+    dataset = dataset_ops.Dataset.range(100)
+    dataset = dataset.batch(distribution.num_replicas_in_sync)
+    dataset = distribution.experimental_distribute_dataset(
+        dataset, options=input_options)
+    if isinstance(dataset, input_lib_v1.DistributedDatasetV1):
+      item = dataset.make_initializable_iterator().get_next()
+    else:
+      self.skipTest('unsupported test combination')
+    device_types = {
+        tf_device.DeviceSpec.from_string(tensor.device).device_type for
+        tensor in item.values}
+    self.assertAllEqual(list(device_types), ['CPU'])
+
+  @combinations.generate(
+      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
+  def testMinimizeLossGraph(self, required_gpus):
+    self._run_between_graph_clients(self._test_minimize_loss_graph,
+                                    self._cluster_spec, required_gpus)
+
+  @combinations.generate(
+      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
+  def testVariableInitialization(self, required_gpus):
+    self._run_between_graph_clients(
+        self._test_variable_initialization,
+        self._cluster_spec,
+        num_gpus=required_gpus)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=[0, 1, 2], use_dataset=[True, False]))
+  def testMakeInputFnIterator(self, required_gpus, use_dataset):
+    def _worker_fn(task_type, task_id, required_gpus):
+      if use_dataset:
+        fn = lambda: dataset_ops.Dataset.range(20)
+      else:
+        def fn():
+          dataset = dataset_ops.Dataset.range(20)
+          it = dataset_ops.make_one_shot_iterator(dataset)
+          return it.get_next
+      # We use CPU as the device when required_gpus = 0
+      devices_per_worker = max(1, required_gpus)
+      expected_values = [[i+j for j in range(devices_per_worker)]
+                         for i in range(0, 20, devices_per_worker)]
+
+      input_fn = self._input_fn_to_test_input_context(
+          fn,
+          expected_num_replicas_in_sync=3*devices_per_worker,
+          expected_num_input_pipelines=3,
+          expected_input_pipeline_id=task_id)
+      self._test_input_fn_iterator(
+          task_type,
+          task_id,
+          required_gpus,
+          input_fn,
+          expected_values,
+          test_reinitialize=use_dataset,
+          ignore_order=not use_dataset)
+
+    self._run_between_graph_clients(_worker_fn, self._cluster_spec,
+                                    required_gpus)
+
+  @combinations.generate(combinations.combine(mode=['graph']))
+  def testUpdateConfigProto(self):
+    strategy, _ = self._get_test_object(
+        task_type='worker', task_id=1, num_gpus=2)
+
+    config_proto = config_pb2.ConfigProto(device_filters=['to_be_overridden'])
+    rewrite_options = config_proto.graph_options.rewrite_options
+    rewrite_options.scoped_allocator_opts.enable_op.append('to_be_removed')
+
+    new_config = strategy.update_config_proto(config_proto)
+
+    # Verify group leader
+    self.assertEqual('/job:worker/replica:0/task:0',
+                     new_config.experimental.collective_group_leader)
+
+    # Verify device filters.
+    self.assertEqual(['/job:worker/task:1'], new_config.device_filters)
+
+    # Verify rewrite options.
+    new_rewrite_options = new_config.graph_options.rewrite_options
+    self.assertEqual(rewriter_config_pb2.RewriterConfig.ON,
+                     new_rewrite_options.scoped_allocator_optimization)
+    self.assertEqual(['CollectiveReduce'],
+                     new_rewrite_options.scoped_allocator_opts.enable_op)
+
+
+class DistributedCollectiveAllReduceStrategyTestWithChief(
+    CollectiveAllReduceStrategyTestBase, parameterized.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    """Create a local cluster with 3 workers and 1 chief."""
+    cls._cluster_spec = multi_worker_test_base.create_in_process_cluster(
+        num_workers=3, num_ps=0, has_chief=True)
+
+  @combinations.generate(
+      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
+  def testMinimizeLossGraph(self, required_gpus):
+    self._run_between_graph_clients(self._test_minimize_loss_graph,
+                                    self._cluster_spec, required_gpus)
+
+  @combinations.generate(
+      combinations.combine(mode=['graph'], required_gpus=[0, 1, 2]))
+  def testVariableInitialization(self, required_gpus):
+    self._run_between_graph_clients(
+        self._test_variable_initialization,
+        self._cluster_spec,
+        num_gpus=required_gpus)
+
+
+class SingleWorkerCollectiveAllReduceStrategy(
+    CollectiveAllReduceStrategyTestBase, strategy_test_lib.DistributionTestBase,
+    strategy_test_lib.TwoDeviceDistributionTestBase, parameterized.TestCase):
+
+  @combinations.generate(combinations.combine(mode=['eager']))
+  def testStrategyInitializationError(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        'cluster_resolver and devices cannot be set at the same time'):
+      _ = collective_all_reduce_strategy.CollectiveAllReduceExtended(
+          container_strategy=None,
+          cluster_resolver=multi_worker_test_base.create_in_process_cluster(
+              num_workers=3, num_ps=0),
+          communication_options=collective_util.Options(),
+          devices=['GPU:0', 'GPU:1'])
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph', 'eager'],
+          required_gpus=[0, 1, 2],
+          use_devices_arg=[True, False]))
+  def testMinimizeLoss(self, required_gpus, use_devices_arg):
+    # Collective ops doesn't support strategy with one device.
+    if context.executing_eagerly():
+      strategy, _ = self._get_test_object(
+          None, None, required_gpus, use_devices_arg=use_devices_arg)
+      self._test_minimize_loss_eager(strategy)
+    else:
+      self._test_minimize_loss_graph(None, None, required_gpus)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['eager'], required_gpus=[1, 2], use_devices_arg=[True, False]))
+  def testNumReplicasInSync(self, required_gpus, use_devices_arg):
+    strategy, _ = self._get_test_object(
+        None, None, required_gpus, use_devices_arg=use_devices_arg)
+    self.assertEqual(required_gpus, strategy.num_replicas_in_sync)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['eager'],
+          required_tpus=[0, 1, 2],
+          use_devices_arg=[True, False]))
+  def testMinimizeLossTPU(self, required_tpus, use_devices_arg):
+    strategy, _ = self._get_test_object(
+        None, None, num_tpus=required_tpus, use_devices_arg=use_devices_arg)
+    self._test_minimize_loss_eager(strategy)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph', 'eager'],
+          required_gpus=[0, 1, 2],
+          use_devices_arg=[True, False]))
+  def testCallAndMergeExceptions(self, required_gpus, use_devices_arg):
+    strategy, _ = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    self._test_call_and_merge_exceptions(strategy)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'],
+          required_gpus=2,
+          use_dataset=[True, False],
+          use_devices_arg=[True, False]))
+  def testMakeInputFnIterator(self, required_gpus, use_dataset,
+                              use_devices_arg):
+    if use_dataset:
+      fn = lambda: dataset_ops.Dataset.range(5 * required_gpus)
+    else:
+      def fn():
+        dataset = dataset_ops.Dataset.range(5 * required_gpus)
+        it = dataset_ops.make_one_shot_iterator(dataset)
+        return it.get_next
+
+    expected_values = [
+        range(i, i + required_gpus) for i in range(0, 10, required_gpus)
+    ]
+
+    input_fn = self._input_fn_to_test_input_context(
+        fn,
+        expected_num_replicas_in_sync=required_gpus,
+        expected_num_input_pipelines=1,
+        expected_input_pipeline_id=0)
+    self._test_input_fn_iterator(
+        None,
+        None,
+        required_gpus,
+        input_fn,
+        expected_values,
+        test_reinitialize=use_dataset,
+        ignore_order=not use_dataset)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph', 'eager'],
+          required_gpus=[0, 1, 2],
+          use_devices_arg=[True, False]))
+  def testReduceToCpu(self, required_gpus, use_devices_arg):
+    strategy, _ = self._get_test_object(
+        None, None, required_gpus, use_devices_arg=use_devices_arg)
+    with strategy.scope():
+      result = strategy.extended.call_for_each_replica(_replica_id_f32)
+      reduced = strategy.reduce(reduce_util.ReduceOp.SUM, result, axis=None)
+      expected = sum(range(strategy.num_replicas_in_sync))
+      self.assertEqual(expected, self.evaluate(reduced))
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testAllReduceSum(self, required_gpus, use_devices_arg):
+    distribution, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_all_reduce_sum(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testAllReduceSumGradients(self, required_gpus, use_devices_arg):
+    distribution, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_all_reduce_sum_gradients(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testAllReduceSumGradientTape(self, required_gpus, use_devices_arg):
+    distribution, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_all_reduce_sum_gradient_tape(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testAllReduceMean(self, required_gpus, use_devices_arg):
+    distribution, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_all_reduce_mean(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testAllReduceMeanGradients(self, required_gpus, use_devices_arg):
+    distribution, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_all_reduce_mean_gradients(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testAllReduceMeanGradientTape(self, required_gpus, use_devices_arg):
+    distribution, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_all_reduce_mean_gradient_tape(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph'], required_gpus=2, use_devices_arg=[True, False]))
+  def testNumpyDataset(self, required_gpus, use_devices_arg):
+    strategy, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    self._test_numpy_dataset(
+        strategy, session=self.cached_session(target=target))
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['eager'], required_gpus=2, use_devices_arg=[True, False]))
+  def testReplicateDataset(self, required_gpus, use_devices_arg):
+    strategy, _ = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    dataset_fn = lambda: dataset_ops.Dataset.range(10)
+    expected_values = [[i, i + 1] for i in range(0, 10, 2)]
+    input_fn = self._input_fn_to_test_input_context(
+        dataset_fn,
+        expected_num_replicas_in_sync=required_gpus,
+        expected_num_input_pipelines=1,
+        expected_input_pipeline_id=0)
+    self._test_input_fn_iterable(strategy, input_fn, expected_values)
+
+  @combinations.generate(
+      combinations.combine(mode=['graph'], use_devices_arg=[True, False]))
+  def testDeepCopy(self, use_devices_arg):
+    distribution, _ = self._get_test_object(
+        None, None, use_devices_arg=use_devices_arg)
+    copy.deepcopy(distribution)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph', 'eager'],
+          required_gpus=[0, 1, 2],
+          use_devices_arg=[True, False]))
+  def testSummaryForReplicaZeroOnly(self, required_gpus, use_devices_arg):
+    strategy, target = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    with self.cached_session(target=target):
+      self._test_summary_for_replica_zero_only(strategy)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=['graph', 'eager'],
+          required_gpus=[0, 1, 2],
+          use_devices_arg=[True, False]))
+  def testTrainableVariables(self, required_gpus, use_devices_arg):
+    strategy, _ = self._get_test_object(
+        None, None, num_gpus=required_gpus, use_devices_arg=use_devices_arg)
+    self._test_trainable_variable(strategy)
+
+
+class LogicalDeviceTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(combinations.combine(mode=['eager'], required_gpus=1))
+  def testKeepLogicalDevice(self):
+    gpus = tf_config.list_physical_devices('GPU')
+    if len(gpus) > 1:
+      self.skipTest('Skip logical device test on multi GPUs, since partial GPU '
+                    'virtualization is not permitted.')
+    # Cannot change logical device after the context initialization.
+    context._reset_context()  # pylint: disable=protected-access
+    cluster_spec = multi_worker_test_base.create_cluster_spec(
+        has_chief=False, num_workers=1)
+    resolver = cluster_resolver_lib.SimpleClusterResolver(
+        cluster_spec=multi_worker_util.normalize_cluster_spec(cluster_spec),
+        task_type='worker',
+        task_id=0)
+
+    logical_gpus = len(gpus) * 2
+    for i, device in enumerate(gpus):
+      n = (i + 1) * logical_gpus // len(gpus) - i * logical_gpus // len(gpus)
+      assert n > 0  # guaranteed if count >= len(devices)
+      configs = []
+      for ordinal in range(n):
+        config = context.LogicalDeviceConfiguration(
+            memory_limit=64,
+            experimental_device_ordinal=ordinal)
+        configs.append(config)
+
+      tf_config.set_logical_device_configuration(device, configs)
+
+    collective_all_reduce_strategy.CollectiveAllReduceStrategy(
+        cluster_resolver=resolver)
+    # Since we create two logical GPUs out of the last GPU, there should be one
+    # more logical GPUs than physical GPUs.
+    self.assertLen(tf_config.list_logical_devices('GPU'), logical_gpus)
+    context._reset_context()  # pylint: disable=protected-access
+
+
+@combinations.generate(
+    combinations.combine(
+        strategy=[
+            strategy_combinations.multi_worker_mirrored_2x1_cpu,
+            strategy_combinations.multi_worker_mirrored_2x1_gpu,
+            strategy_combinations.multi_worker_mirrored_2x2_gpu,
+            strategy_combinations.multi_worker_mirrored_2x2_gpu_no_merge_call,
+        ],
+        mode=['eager']))
+class CollectiveAllReduceStrategyV2Test(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if context.context().list_physical_devices('TPU'):
+      self.skipTest('Test not supported on TPUs')
+
+  def test_replica_id_in_sync_group(self, strategy):
+
+    def replica_fn():
+      replica_ctx = distribute_lib.get_replica_context()
+      return replica_ctx.replica_id_in_sync_group, replica_ctx._replica_id
+
+    results = test_util.gather(strategy, strategy.run(replica_fn))
+    self.assertAllEqual(list(range(strategy.extended._num_replicas_in_sync)),
+                        results[0].numpy())
+    self.assertAllEqual(
+        list(range(len(strategy.extended.worker_devices))) *
+        strategy.extended._num_workers, results[1].numpy())
+
+  def test_deep_copy_not_allowed(self, strategy):
+    # Check health is disabled in tests by default. We need to enable it for
+    # this test to simulate the real world.
+    strategy.extended._start_check_health_thread()
+    try:
+      with self.assertRaisesRegex(ValueError, 'cannot be deep copied'):
+        copy.deepcopy(strategy)
+      with self.assertRaisesRegex(ValueError, 'cannot be deep copied'):
+        with ops.Graph().as_default():
+          copy.deepcopy(strategy)
+    finally:
+      strategy.extended._stop_check_health_thread()
+
+
+class ExperimentalCompatibilityTest(test.TestCase):
+
+  def testIsInstance(self):
+    # It's not uncommon for people to special case MultiWorkerMirroredStrategy,
+    # so we need to make sure isinstance check works for combinations between
+    # the experimental and non-experimental endpoints.
+    strategy = CollectiveAllReduceStrategy()
+    experimental_strategy = _CollectiveAllReduceStrategyExperimental()
+    self.assertIsInstance(strategy, CollectiveAllReduceStrategy)
+    self.assertIsInstance(strategy, _CollectiveAllReduceStrategyExperimental)
+    self.assertIsInstance(experimental_strategy, CollectiveAllReduceStrategy)
+    self.assertIsInstance(experimental_strategy,
+                          _CollectiveAllReduceStrategyExperimental)
+
+  def testName(self):
+    self.assertEqual(CollectiveAllReduceStrategy.__name__,
+                     'CollectiveAllReduceStrategy')
+    self.assertEqual(_CollectiveAllReduceStrategyExperimental.__name__,
+                     'CollectiveAllReduceStrategy')
+
+
+def _replica_id_f32():
+  return math_ops.cast(
+      distribute_lib.get_replica_context()
+      .replica_id_in_sync_group, dtypes.float32)
+
+
+if __name__ == '__main__':
+  # TODO(b/172304955): enable logical devices.
+  test_util.main(config_logical_devices=False)

tensorflow/python/distribute/collective_util.py

@@ -0,0 +1,233 @@
+# coding=utf-8
+# 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.
+# ==============================================================================
+"""Utilities for collectives."""
+
+import copy
+import enum
+
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+# TODO(b/170340570): print deprecation warning for CollectiveCommunication.
+@tf_export("distribute.experimental.CommunicationImplementation",
+           "distribute.experimental.CollectiveCommunication")
+class CommunicationImplementation(enum.Enum):
+  """Cross device communication implementation.
+
+  Warning: The alias `tf.distribute.experimental.CollectiveCommunication` is
+  deprecated and will be removed in a future version. Use
+  `tf.distribute.experimental.CommunicationImplementation` instead.
+
+  * `AUTO`: Automatically chosen by Tensorflow.
+  * `RING`: TensorFlow's ring algorithms for all-reduce and
+    all-gather.
+  * `NCCL`: NVIDIA®'s NCCL library. This is now only used for all-reduce on
+    GPUs; all-reduce on CPU, all-gather and broadcast fallbacks to RING.
+  """
+  AUTO = "AUTO"
+  RING = "RING"
+  NCCL = "NCCL"
+  # TODO(ayushd): add ncclAllGather implementation.
+
+
+CollectiveCommunication = CommunicationImplementation
+
+
+@tf_export("distribute.experimental.CommunicationOptions")
+class _OptionsExported(object):
+  """Options for cross device communications like All-reduce.
+
+  This can be passed to methods like
+  `tf.distribute.get_replica_context().all_reduce()` to optimize collective
+  operation performance. Note that these are only hints, which may or may not
+  change the actual behavior. Some options only apply to certain strategy and
+  are ignored by others.
+
+  One common optimization is to break gradients all-reduce into multiple packs
+  so that weight updates can overlap with gradient all-reduce.
+
+  Examples:
+
+  ```python
+  options = tf.distribute.experimental.CommunicationOptions(
+      bytes_per_pack=50 * 1024 * 1024,
+      timeout_seconds=120.0,
+      implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
+  )
+  grads = tf.distribute.get_replica_context().all_reduce(
+      'sum', grads, options=options)
+  optimizer.apply_gradients(zip(grads, vars),
+      experimental_aggregate_gradients=False)
+  ```
+
+  """
+
+  def __new__(cls, *args, **kwargs):
+    # We expose a dummy class so that we can separate internal and public APIs.
+    # Note that __init__ won't be called on the returned object if it's a
+    # different class [1].
+    # [1] https://docs.python.org/3/reference/datamodel.html#object.__new__
+    return Options(*args, **kwargs)
+
+  def __init__(self,
+               bytes_per_pack=0,
+               timeout_seconds=None,
+               implementation=CommunicationImplementation.AUTO):
+    """Creates a CollectiveHints.
+
+    Args:
+      bytes_per_pack: a non-negative integer. Breaks collective operations into
+        packs of certain size. If it's zero, the value is determined
+        automatically. This hint is respected by all multi-replica strategies
+        except `TPUStrategy`.
+      timeout_seconds: a float or None, timeout in seconds. If not None, the
+        collective raises `tf.errors.DeadlineExceededError` if it takes longer
+        than this timeout. Zero disables timeout. This can be useful when
+        debugging hanging issues.  This should only be used for debugging since
+        it creates a new thread for each collective, i.e. an overhead of
+        `timeout_seconds * num_collectives_per_second` more threads. This only
+        works for `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
+      implementation: a
+        `tf.distribute.experimental.CommunicationImplementation`. This is a hint
+        on the preferred communication implementation. Possible values include
+        `AUTO`, `RING`, and `NCCL`. NCCL is generally more performant for GPU,
+        but doesn't work for CPU. This only works for
+        `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
+
+    Raises:
+      ValueError: When arguments have invalid value.
+    """
+    pass
+
+
+class Options(object):
+  """Implementation of OptionsInterface."""
+
+  def __init__(self,
+               bytes_per_pack=0,
+               timeout_seconds=None,
+               implementation=CommunicationImplementation.AUTO):
+    if bytes_per_pack < 0:
+      raise ValueError(
+          f"Argument `bytes_per_pack` must be >=0, Received {bytes_per_pack}.")
+    if isinstance(implementation, str):
+      implementation = CommunicationImplementation(implementation.upper())
+    if not isinstance(implementation, CommunicationImplementation):
+      raise ValueError(
+          "Argument `implementation` must be instance of "
+          "`tf.distribute.experimental.CommunicationImplementation`.")
+    self.bytes_per_pack = bytes_per_pack
+    self.timeout_seconds = timeout_seconds
+    self.implementation = implementation
+
+  __init__.__doc__ = _OptionsExported.__init__.__doc__
+
+  def merge(self, options):
+    """Merges with another options and returns a new one.
+
+    Values specified in the `options` takes precedence if they're not the
+    default.
+
+    Args:
+      options: a `tf.distribute.experimental.CollectiveCommunication`.
+
+    Returns:
+      A new `tf.distribute.experimental.CollectiveCommunication`.
+    """
+    merged = copy.deepcopy(self)
+    if options is None:
+      return merged
+    if options.bytes_per_pack != 0:
+      merged.bytes_per_pack = options.bytes_per_pack
+    if options.timeout_seconds is not None:
+      merged.timeout_seconds = options.timeout_seconds
+    if options.implementation != CommunicationImplementation.AUTO:
+      merged.implementation = options.implementation
+    return merged
+
+  def __str__(self):
+    return (f"Options(bytes_per_pack={self.bytes_per_pack},"
+            f"timeout_seconds={self.timeout_seconds}, "
+            f"implementation={self.implementation})")
+
+
+@tf_export("distribute.experimental.CollectiveHints")
+class Hints(object):
+  """Hints for collective operations like AllReduce.
+
+  This can be passed to methods like
+  `tf.distribute.get_replica_context().all_reduce()` to optimize collective
+  operation performance. Note that these are only hints, which may or may not
+  change the actual behavior. Some options only apply to certain strategy and
+  are ignored by others.
+
+  One common optimization is to break gradients all-reduce into multiple packs
+  so that weight updates can overlap with gradient all-reduce.
+
+  Examples:
+
+  - bytes_per_pack
+
+  ```python
+  hints = tf.distribute.experimental.CollectiveHints(
+      bytes_per_pack=50 * 1024 * 1024)
+  grads = tf.distribute.get_replica_context().all_reduce(
+      'sum', grads, experimental_hints=hints)
+  optimizer.apply_gradients(zip(grads, vars),
+      experimental_aggregate_gradients=False)
+  ```
+
+  - timeout_seconds
+
+  ```python
+  strategy = tf.distribute.MirroredStrategy()
+  hints = tf.distribute.experimental.CollectiveHints(
+      timeout_seconds=120.0)
+  try:
+    strategy.reduce("sum", v, axis=None, experimental_hints=hints)
+  except tf.errors.DeadlineExceededError:
+    do_something()
+  ```
+
+  """
+
+  @deprecation.deprecated(
+      None, "use distribute.experimental.CommunicationOptions instead")
+  def __new__(cls, bytes_per_pack=0, timeout_seconds=None):
+    return Options(
+        bytes_per_pack=bytes_per_pack, timeout_seconds=timeout_seconds)
+
+  def __init__(self, bytes_per_pack=0, timeout_seconds=None):
+    """Creates a CollectiveHints.
+
+    Args:
+      bytes_per_pack: a non-negative integer. Breaks collective operations into
+        packs of certain size. If it's zero, the value is determined
+        automatically. This only applies to all-reduce with
+        `MultiWorkerMirroredStrategy` currently.
+      timeout_seconds: a float or None, timeout in seconds. If not None, the
+        collective raises `tf.errors.DeadlineExceededError` if it takes longer
+        than this timeout. This can be useful when debugging hanging issues.
+        This should only be used for debugging since it creates a new thread for
+        each collective, i.e. an overhead of `timeout_seconds *
+        num_collectives_per_second` more threads.  This only works for
+        `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
+
+    Raises:
+      ValueError: When arguments have invalid value.
+    """
+    pass

tensorflow/python/distribute/collective_util_test.py

@@ -0,0 +1,40 @@
+# 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.
+# ==============================================================================
+"""Test for utilities for collectives."""
+
+from tensorflow.python.distribute import collective_util
+from tensorflow.python.eager import test
+
+
+class OptionsTest(test.TestCase):
+
+  def testCreateOptionsViaExportedAPI(self):
+    options = collective_util._OptionsExported(bytes_per_pack=1)
+    self.assertIsInstance(options, collective_util.Options)
+    self.assertEqual(options.bytes_per_pack, 1)
+    with self.assertRaises(ValueError):
+      collective_util._OptionsExported(bytes_per_pack=-1)
+
+  def testCreateOptionsViaHints(self):
+    with self.assertLogs() as cm:
+      options = collective_util.Hints(50, 1)
+    self.assertTrue(any("is deprecated" in msg for msg in cm.output))
+    self.assertIsInstance(options, collective_util.Options)
+    self.assertEqual(options.bytes_per_pack, 50)
+    self.assertEqual(options.timeout_seconds, 1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/combinations.py

@@ -0,0 +1,652 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""This module customizes `test_combinations` for `tf.distribute.Strategy`.
+
+Additionally it provides `generate()`, `combine()` and `times()` with
+`tf.distribute.Strategy` customizations as a default.
+"""
+
+import collections
+import copy
+import re
+import sys
+import types
+import unittest
+
+from absl import app
+import six
+
+
+from tensorflow.python.client import session
+from tensorflow.python.distribute import collective_all_reduce_strategy
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import combinations as framework_combinations
+from tensorflow.python.framework import config
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_combinations as combinations_lib
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util.tf_export import tf_export
+
+
+# TODO(rchao): Rename `distribution` parameter to `strategy` or
+# `distribute_strategy` in all tests.
+class DistributionParameter(combinations_lib.ParameterModifier):
+  """Transforms arguments of type `NamedDistribution`.
+
+  Convert all arguments of type `NamedDistribution` to the value of their
+  `strategy` property.
+  """
+
+  def modified_arguments(self, kwargs, requested_parameters):
+    # Get the parameter that indicates if we need to set the `_use_policy` flag
+    # on the strategy object. This is a temporary flag for testing the variable
+    # policy rollout.
+    use_var_policy = kwargs.get("use_var_policy", None)
+    distribution_arguments = {}
+    for k, v in kwargs.items():
+      if isinstance(v, NamedDistribution):
+        strategy = v.strategy
+        if use_var_policy:
+          strategy.extended._use_var_policy = use_var_policy
+        distribution_arguments[k] = strategy
+    return distribution_arguments
+
+
+class ClusterParameters(combinations_lib.ParameterModifier):
+  """Adds cluster parameters if a `NamedDistribution` has it.
+
+  It needs to be before DistributionParameter.
+  """
+
+  def modified_arguments(self, kwargs, requested_parameters):
+    strategy = None
+    for _, v in kwargs.items():
+      if isinstance(v, NamedDistribution):
+        if strategy is not None and _num_total_workers(v.has_chief,
+                                                       v.num_workers) > 1:
+          raise ValueError("Only support one NamedDistribution for multi worker"
+                           "tests.")
+        strategy = v
+
+    if strategy:
+      has_chief = strategy.has_chief
+      num_workers = strategy.num_workers
+      runner = strategy.runner
+      share_gpu = strategy.share_gpu
+      num_ps = strategy.num_ps
+      if "has_chief" in kwargs and kwargs["has_chief"] != has_chief:
+        raise ValueError(
+            "both has_chief and strategy specified but are not compatible")
+      if "num_workers" in kwargs and kwargs["num_workers"] != num_workers:
+        raise ValueError(
+            "both num_workers and strategy specified but are not compatible")
+    else:
+      has_chief = kwargs.get("has_chief", False)
+      num_workers = kwargs.get("num_workers", 1)
+      runner = kwargs.get("runner", None)
+      share_gpu = kwargs.get("share_gpu", True)
+      num_ps = kwargs.get("num_ps", 0)
+
+    # Always set cluster parameters if they're requested. So that generate()
+    # works when there's no strategy in the combinations.
+    update = {}
+    if "has_chief" in requested_parameters:
+      update["has_chief"] = has_chief
+    if "num_workers" in requested_parameters:
+      update["num_workers"] = num_workers
+    if "runner" in requested_parameters:
+      update["runner"] = runner
+    if "share_gpu" in requested_parameters:
+      update["share_gpu"] = share_gpu
+    if "num_ps" in requested_parameters:
+      update["num_ps"] = num_ps
+    return update
+
+
+class DistributionCombination(combinations_lib.TestCombination):
+  """Sets up distribution strategy for tests."""
+
+  def should_execute_combination(self, kwargs):
+    distributions = [
+        v for v in kwargs.values() if isinstance(v, NamedDistribution)
+    ]
+    if test_util.is_xla_enabled() and any(d.no_xla for d in distributions):
+      return (
+          False,
+          "n/a: skipping strategy combination with no_xla=True in XLA tests")
+    return (True, None)
+
+  def parameter_modifiers(self):
+    return [
+        DistributionParameter(),
+        combinations_lib.OptionalParameter("use_var_policy"),
+    ]
+
+
+class ClusterCombination(combinations_lib.TestCombination):
+  """Sets up multi worker tests."""
+
+  def parameter_modifiers(self):
+    return [ClusterParameters()]
+
+
+class GPUCombination(combinations_lib.TestCombination):
+  """Enable tests to request GPU hardware and skip non-GPU combinations.
+
+  This class expects test_combinations to be generated with `NamedDistribution`
+  wrapping instances of `tf.distribute.Strategy`.
+
+  Optionally, the `required_gpus` argument is supported.  GPU hardware is
+  required, if its value is `True` or > 0.
+
+  Attributes:
+    GPU_TEST: The environment is considered to have GPU hardware available if
+              the name of the program contains "test_gpu" or "test_xla_gpu".
+  """
+  GPU_TEST = False
+  if sys.argv:
+    GPU_TEST = re.search(r"(test_2?gpu|test_xla_2?gpu)$", sys.argv[0])
+
+  def should_execute_combination(self, kwargs):
+    distributions = [
+        v for v in kwargs.values() if isinstance(v, NamedDistribution)
+    ]
+    required_gpus = kwargs.get("required_gpus", 0)
+    required_physical_gpus = kwargs.get("required_physical_gpus", 0)
+
+    if distributions and required_gpus:
+      raise ValueError("Do not use `required_gpus` and arguments of type "
+                       "NamedDistribution together.")
+
+    number_of_required_gpus = max(
+        [required_gpus] + [required_physical_gpus] +
+        [d.required_physical_gpus or 0 for d in distributions] +
+        [d.required_gpus or 0 for d in distributions])
+    number_of_required_physical_gpus = max(
+        [required_physical_gpus] +
+        [d.required_physical_gpus or 0 for d in distributions])
+
+    if (required_physical_gpus and required_gpus):
+      raise ValueError("Only one of `required_physical_gpus`(number of physical"
+                       " GPUs required) and `required_gpus`(total number of "
+                       "GPUs required) should be set. ")
+    if not number_of_required_gpus and GPUCombination.GPU_TEST:
+      return (False, "Test that doesn't require GPUs.")
+    elif (number_of_required_gpus > 0
+          and context.num_gpus() < number_of_required_gpus):
+      return (False, ("Only {} of {} required GPUs are available.".format(
+          context.num_gpus(), number_of_required_gpus)))
+    elif number_of_required_physical_gpus > len(
+        config.list_physical_devices("GPU")):
+      return (False,
+              ("Only {} of {} required physical GPUs are available.".format(
+                  config.list_physical_devices("GPU"), required_physical_gpus)))
+    else:
+      return (True, None)
+
+  def parameter_modifiers(self):
+    return [combinations_lib.OptionalParameter("required_gpus"),
+            combinations_lib.OptionalParameter("required_physical_gpus")]
+
+
+class TPUCombination(combinations_lib.TestCombination):
+  """Allow to request TPU hardware and skip non-TPU combinations.
+
+  This class expects test_combinations to be generated with `NamedDistribution`
+  wrapping instances of `tf.distribute.Strategy`.
+
+  Optionally, the `required_tpus` parameter is supported.  TPU hardware is
+  required, if its argument is `True` or > 0.
+
+  Optionally, the `use_cloud_tpu` parameter is supported. If TPU hardware is
+  required by `required_tpus`, it specifically must be a Cloud TPU (specified
+  with `--tpu`) if `use_cloud_tpu` is `True`.
+
+  Attributes:
+    TPU_TEST: The environment is considered to have TPU hardware available if
+              the name of the program contains "test_tpu".
+  """
+
+  TPU_TEST = False
+  if sys.argv:
+    TPU_TEST = "test_tpu" in sys.argv[0]
+
+  def should_execute_combination(self, kwargs):
+    distributions = [
+        v for v in kwargs.values() if isinstance(v, NamedDistribution)
+    ]
+    # TODO(isaprykin): Migrate all tests away from using 'required_tpu' in favor
+    # of 'required_tpus'.
+    if "required_tpus" in kwargs and "required_tpu" in kwargs:
+      raise ValueError("Do not use `required_tpu`.  Both `required_tpus` and "
+                       "`required_tpu` were specified.")
+    required_tpus = kwargs.get("required_tpus", None) or kwargs.get(
+        "required_tpu", None)
+
+    if distributions and required_tpus:
+      raise ValueError("Do not use `required_tpus` and arguments of type "
+                       "NamedDistribution together.")
+
+    # TODO(isaprykin): Add support for a particular number of TPUs.  Right now
+    # it's binary.
+    number_of_required_tpus = max([required_tpus or 0] +
+                                  [d.required_tpu or 0 for d in distributions])
+    use_cloud_tpu = any([kwargs.get("use_cloud_tpu")] +
+                        [d.use_cloud_tpu for d in distributions])
+    tpu = hasattr(flags.FLAGS, "tpu") and flags.FLAGS.tpu or ""
+
+    if not number_of_required_tpus and TPUCombination.TPU_TEST:
+      return (False, "Test that doesn't require TPUs.")
+    if number_of_required_tpus and not TPUCombination.TPU_TEST:
+      return (False, "Test requires a TPU, but it's not available.")
+    if use_cloud_tpu and not tpu:
+      return (False, "Test requires a Cloud TPU, but none specified.")
+    if not use_cloud_tpu and tpu:
+      return (False, "Test requires local TPU, but Cloud TPU specified.")
+    return (True, None)
+
+  def parameter_modifiers(self):
+    return [
+        combinations_lib.OptionalParameter("required_tpus"),
+        combinations_lib.OptionalParameter("required_tpu"),
+        combinations_lib.OptionalParameter("use_cloud_tpu"),
+    ]
+
+
+class NamedDistribution(object):
+  """Wraps a `tf.distribute.Strategy` and adds a name for test titles."""
+
+  def __init__(self,
+               name,
+               distribution_fn,
+               required_gpus=None,
+               required_physical_gpus=0,
+               required_tpu=False,
+               use_cloud_tpu=False,
+               has_chief=False,
+               num_workers=1,
+               num_ps=0,
+               share_gpu=True,
+               pool_runner_fn=None,
+               no_xla=False):
+    """Initialize NamedDistribution.
+
+    Args:
+      name: Name that will be a part of the name of the test case.
+      distribution_fn: A callable that creates a `tf.distribute.Strategy`.
+      required_gpus: The number of GPUs that the strategy requires. Only one of
+      `required_gpus` and `required_physical_gpus` should be set.
+      required_physical_gpus: Number of physical GPUs required. Only one of
+      `required_gpus` and `required_physical_gpus` should be set.
+      required_tpu: Whether the strategy requires TPU.
+      use_cloud_tpu: Whether the strategy requires cloud TPU.
+      has_chief: Whether the strategy requires a chief worker.
+      num_workers: The number of workers that the strategy requires.
+      num_ps: The number of parameter servers.
+      share_gpu: Whether to share GPUs among workers.
+      pool_runner_fn: An optional callable that returns a MultiProcessPoolRunner
+        to run the test.
+      no_xla: Whether to skip in XLA tests.
+    """
+    object.__init__(self)
+    self._name = name
+    self._distribution_fn = distribution_fn
+    self.required_gpus = required_gpus
+    self.required_physical_gpus = required_physical_gpus
+    self.required_tpu = required_tpu
+    self.use_cloud_tpu = use_cloud_tpu
+    self.has_chief = has_chief
+    self.num_workers = num_workers
+    self.num_ps = num_ps
+    self.share_gpu = share_gpu
+    self._pool_runner_fn = pool_runner_fn
+    self.no_xla = no_xla
+
+  @property
+  def runner(self):
+    if self._pool_runner_fn is not None:
+      return self._pool_runner_fn()
+    return None
+
+  @property
+  def strategy(self):
+    return self._distribution_fn()
+
+  def __repr__(self):
+    return self._name
+
+
+# This is to allow adding combinations that runs a function both as a
+# tf.function and eagerly.
+#
+# @combinations.generate(
+#   combinations.combine(
+#     tf_function = [combinations.tf_function, combinations.no_tf_function]
+#   )
+# )
+# def testXXX(tf_function):
+#   @tf_function
+#   def foo():
+#     tf.add(1., 1.)
+#
+#   foo()
+tf_function = combinations_lib.NamedObject("TfFunction", def_function.function)
+no_tf_function = combinations_lib.NamedObject("NoTfFunction", lambda f: f)
+
+
+def concat(*combined):
+  """Concats combinations."""
+  result = []
+  for one in combined:
+    result += one
+  return result
+
+
+@tf_export("__internal__.distribute.combinations.generate", v1=[])
+def generate(combinations, test_combinations=()):
+  # pylint: disable=g-doc-args,g-doc-return-or-yield
+  """Distributed adapter of `tf.__internal__.test.combinations.generate`.
+
+  All tests with distributed strategy should use this one instead of
+  `tf.__internal__.test.combinations.generate`. This function has support of
+  strategy combinations, GPU/TPU and multi worker support.
+
+  See `tf.__internal__.test.combinations.generate` for usage.
+  """
+  # pylint: enable=g-doc-args,g-doc-return-or-yield
+  default_combinations = (
+      framework_combinations.EagerGraphCombination(),
+      framework_combinations.TFVersionCombination(),
+      ClusterCombination(),
+      DistributionCombination(),
+      GPUCombination(),
+      TPUCombination(),
+  )
+  # We apply our own decoration to handle multi worker tests before applying
+  # framework.test_combinations.generate. The order is important since we need
+  # framework.test_combinations.generate to apply all parameter modifiers first.
+  combination_decorator = combinations_lib.generate(
+      combinations, test_combinations=default_combinations + test_combinations)
+
+  def decorator(test_method_or_class):
+    if isinstance(test_method_or_class, type):
+      # If it's a test class.
+      class_object = test_method_or_class
+      # Decorate each test method with _multi_worker_test.
+      for name, test_method in six.iteritems(class_object.__dict__.copy()):
+        if (name.startswith(unittest.TestLoader.testMethodPrefix) and
+            isinstance(test_method, types.FunctionType)):
+          setattr(class_object, name, _multi_worker_test(test_method))
+      return combination_decorator(class_object)
+    else:
+      return combination_decorator(_multi_worker_test(test_method_or_class))
+
+  return decorator
+
+
+combine = combinations_lib.combine
+times = combinations_lib.times
+NamedObject = combinations_lib.NamedObject
+
+
+# Identifies whether we're in the main process or worker processes.
+# `_multi_worker_test` decoration behaves differently in the main processes and
+# the worker processes. See the documentation of _multi_worker_test for detail.
+_running_in_worker = False
+
+
+@tf_export("__internal__.distribute.combinations.in_main_process", v1=[])
+def in_main_process():
+  """Whether it's in the main test process.
+
+  This is normally used to prepare the test environment which should only happen
+  in the main process.
+
+  Returns:
+    A boolean.
+  """
+  return not _running_in_worker
+
+
+class TestEnvironment(object):
+  """Holds the test environment information.
+
+  Tests should modify the attributes of the instance returned by `env()` in the
+  main process if needed, and it will be passed to the worker processes each
+  time a test case is run.
+  """
+
+  def __init__(self):
+    self.tf_data_service_dispatcher = None
+    # Note that this includes GPUs that may not be visible to the current
+    # worker.
+    self.total_phsyical_gpus = None
+
+  def __setattr__(self, name, value):
+    if not in_main_process():
+      raise ValueError(
+          "combinations.env() should only be modified in the main process. "
+          "Condition your code on combinations.in_main_process().")
+    super().__setattr__(name, value)
+
+
+_env = TestEnvironment()
+
+
+@tf_export("__internal__.distribute.combinations.env", v1=[])
+def env():
+  """Returns the object holds the test environment information.
+
+  Tests should modify this in the main process if needed, and it will be passed
+  to the worker processes each time a test case is run.
+
+  Returns:
+    a TestEnvironment object.
+  """
+  return _env
+
+
+def _set_total_phsyical_gpus():
+  if in_main_process():
+    env().total_phsyical_gpus = len(
+        context.context().list_physical_devices("GPU"))
+
+
+# This is needed in case CUDA is lazily loaded.
+app.call_after_init(_set_total_phsyical_gpus)
+
+
+_TestResult = collections.namedtuple("_TestResult", ["status", "message"])
+
+
+def _test_runner(test_id, test_env):
+  """Executes the test with the given test_id.
+
+  This is a simple wrapper around TestRunner to be used with
+  multi_process_runner. Similar to test.main(), but it executes only one test
+  specified by test_id and returns whether the test succeeds. If the test fails,
+  the function prints failures and errors to stdout.
+
+  Args:
+    test_id: TestCase.id()
+    test_env: a TestEnvironment object.
+
+  Returns:
+    A boolean indicates whether the test succeeds.
+  """
+  global _running_in_worker, _env
+  # No need to restore the value of _running_in_worker since it should always be
+  # True in worker processes.
+  _running_in_worker = True
+  _env = test_env
+  test = unittest.defaultTestLoader.loadTestsFromName(test_id)
+  runner = unittest.TextTestRunner()
+  result = runner.run(test)
+  # Treat expected failures as failures, so that the main process can get
+  # them and fail as expected. Also treat errors as failures to simplify the
+  # handling.
+  failures = result.failures + result.expectedFailures + result.errors
+  if failures:
+    ret = _TestResult(status="failure", message=failures[0][1])
+  elif result.skipped:
+    ret = _TestResult(status="skipped", message=result.skipped[0][1])
+  else:
+    # Treat unexpectedSuccesses as OK so that the test case in the main process
+    # succeed as well.
+    ret = _TestResult(status="ok", message=None)
+  # Print tracebacks to stdout and multi_process_runner will collect
+  # them and stream back to the main process.
+  if ret.message:
+    print(ret.message)
+  return ret
+
+
+def _multi_worker_test(test_method):
+  """Decorate test_method so that it runs in each worker.
+
+  We use `multi_process_runner` to simulate multiple workers. Since we run the
+  this function in the main process and all worker processes, this decoration
+  behaves differently in the main process and worker procssses. In the main
+  process, it spawns subprocesses and runs the test on each of them; in a worker
+  process, it executes test in the same way as a normal test, e.g.
+  setUp()/tearDown() are called before/after the test.
+
+  Args:
+    test_method: a function which must be a test method.
+
+  Returns:
+    Decorated `test_method`. Note that the decorated function has additional
+    arguments.
+  """
+
+  def decorator(self, has_chief, num_workers, num_ps, share_gpu, runner,
+                **kwargs):
+    if _num_total_workers(has_chief,
+                          num_workers) == 1 or _running_in_worker or (
+                              # Use in-process cluster for PS combinations
+                              # when XLA is enabled.
+                              test_util.is_xla_enabled() and num_ps > 0):
+      # We're in worker process or the test is for single worker. Either case we
+      # execute the test method directly instead of spawning subprocesses.
+
+      # For MultiWorkerMirroredStrategy(CollectiveAllReduceStrategy), install a
+      # session that connects to the local server. This is necessary for multi
+      # worker graph mode tests to work. Those tests cannot use their graphs or
+      # sessions, including the one returned by self.cached_session(). Since
+      # existing tests may already be doing so, we only install the session for
+      # multi worker tests.
+      with _multi_worker_session(kwargs):
+        test_method(self, **kwargs)
+      return
+
+    # We're in the main process. We spawn subprocesses and run the *test* on
+    # each of them. Note that we're not directly executing test_method passed to
+    # _multi_worker_test, because we need setUp()/tearDown() to be called and
+    # all the decorations on the test method. The conceptual call stack is:
+    #   [main process]test.main()
+    #     [main process]test_runner.run(test)
+    #       [main process]wrapper by combinations.generate()
+    #         [main process]_multi_worker_test.decorator()
+    #           # A sub process goes through the same code path as the main
+    #           # process.
+    #           [sub process]_test_runner()
+    #             [sub process]test_runner.run(test)
+    #               [sub process]wrapper by combinations.generate()
+    #                 [sub process]_multi_worker_test.decorator()
+    #                   # _running_in_worker is True
+    #                   [sub process]test_method()
+    test_id = self.id()
+    if runner:
+      results = runner.run(_test_runner, args=(test_id, _env))
+    else:
+      cluster_spec = multi_worker_test_base.create_cluster_spec(
+          has_chief=has_chief,
+          num_workers=num_workers,
+          num_ps=num_ps,
+          has_eval=False)
+      ephemeral_runner = multi_process_runner.MultiProcessRunner(
+          _test_runner,
+          cluster_spec,
+          share_gpu=share_gpu,
+          args=(test_id, _env),
+          dependence_on_chief=has_chief)
+      ephemeral_runner.start()
+      results = ephemeral_runner.join().return_value
+
+    skip_reason = None
+    for result in results:
+      if result.status == "failure":
+        # We can't tell which worker the return value come from, so we fail on
+        # the  first error.
+        self.fail(result.message)
+        break
+      elif result.status == "skipped":
+        # Record the skip reason, but do not actually skip the test in case some
+        # processes fail instead.
+        skip_reason = result.message
+    if skip_reason is not None:
+      self.skipTest(skip_reason)
+
+  argspec = tf_inspect.getfullargspec(test_method)
+  decorator_args = (argspec.args or []) + [
+      "has_chief", "num_workers", "num_ps", "share_gpu", "runner"
+  ]
+  decorator_argspec = argspec._replace(args=decorator_args)
+  return tf_decorator.make_decorator(
+      test_method, decorator, decorator_argspec=decorator_argspec)
+
+
+def _num_total_workers(has_chief, num_workers):
+  """Returns the number of workers including the chief."""
+  if has_chief:
+    return num_workers + 1
+  return num_workers
+
+
+def _multi_worker_session(kwargs):
+  """Returns a context manager that enters a session that is configured for the MultiWorkerMirroredStrategy.
+
+  Args:
+    kwargs: a dict. Keyword arguments passed to the test.
+
+  Returns:
+    A context manager. If MultiWorkerMirroredStrategy is the  one and only one
+    strategy in kwargs and it's in graph mode, it's the session that is
+    configured for that strategy.  Otherwise, it's a no-op context manager.
+  """
+  strategy = None
+  for _, v in kwargs.items():
+    if isinstance(v, distribute_lib.StrategyBase):
+      if strategy is not None:
+        logging.warning(
+            "The test uses multiple strategies. Skipping "
+            "entering a session that is configured for the strategy.")
+        return ops.NullContextmanager()
+      strategy = v
+  if context.executing_eagerly() or not isinstance(
+      strategy, collective_all_reduce_strategy.CollectiveAllReduceStrategy):
+    return ops.NullContextmanager()
+  sess_config = copy.deepcopy(context.context().config)
+  sess_config = strategy.update_config_proto(sess_config)
+  target = strategy.cluster_resolver.master()
+  return session.Session(config=sess_config, target=target).as_default()

tensorflow/python/distribute/combinations_test.py

@@ -0,0 +1,237 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for tensorflow.python.distribute.combinations."""
+
+import importlib
+import os
+import sys
+import unittest
+
+from absl.testing import parameterized
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.eager import context
+from tensorflow.python.framework import combinations as framework_combinations
+from tensorflow.python.platform import test
+
+
+class ClusterCombinationTest(test.TestCase, parameterized.TestCase):
+  # For this test we need to use `framework.test_combinations` because our
+  # `generate` eats the cluster parameters.
+  #
+  # Note that we don't have a standalone combination for ClusterParameters, so
+  # we should use GPUCombination which contains it.
+
+  @framework_combinations.generate(  # pylint: disable=redundant-keyword-arg
+      framework_combinations.combine(distribution=[
+          combinations.NamedDistribution(
+              "HasClusterParams", lambda: None, has_chief=True, num_workers=2),
+      ]),
+      test_combinations=(combinations.ClusterCombination(),))
+  def testClusterParams(self, distribution, has_chief, num_workers):
+    self.assertTrue(has_chief)
+    self.assertEqual(num_workers, 2)
+
+  @framework_combinations.generate(  # pylint: disable=redundant-keyword-arg
+      framework_combinations.combine(distribution=[
+          combinations.NamedDistribution("NoClusterParams", lambda: None),
+      ]),
+      test_combinations=(combinations.ClusterCombination(),))
+  def testClusterParamsHasDefault(self, distribution, has_chief, num_workers):
+    self.assertFalse(has_chief)
+    self.assertEqual(num_workers, 1)
+
+  @framework_combinations.generate(  # pylint: disable=redundant-keyword-arg
+      framework_combinations.combine(v=1),
+      test_combinations=(combinations.ClusterCombination(),))
+  def testClusterParamsNoStrategy(self, v, has_chief, num_workers):
+    self.assertFalse(has_chief)
+    self.assertEqual(num_workers, 1)
+
+  @framework_combinations.generate(  # pylint: disable=redundant-keyword-arg
+      framework_combinations.combine(distribution=[
+          combinations.NamedDistribution(
+              "WithClusterParams", lambda: None, has_chief=True, num_workers=2),
+          combinations.NamedDistribution("WithoutClusterParams", lambda: None),
+      ]),
+      test_combinations=(combinations.ClusterCombination(),))
+  def testClusterParamsAreOptional(self, distribution):
+    # If combinations library doesn't raise an exception, the test is passed.
+    pass
+
+  @framework_combinations.generate(  # pylint: disable=redundant-keyword-arg
+      framework_combinations.combine(
+          ds1=combinations.NamedDistribution(
+              "Strategy1", lambda: None, has_chief=True, num_workers=0),
+          ds2=combinations.NamedDistribution(
+              "Strategy2", lambda: None, has_chief=False, num_workers=1),
+          ds3=combinations.NamedDistribution(
+              "Strategy3", lambda: None, has_chief=True, num_workers=0),
+      ),
+      test_combinations=(combinations.ClusterCombination(),))
+  def testMultipleDistributionSingleWorker(self, ds1, ds2, ds3):
+    # If combinations library doesn't raise an exception, the test is passed.
+    pass
+
+  @combinations.generate(combinations.combine(num_workers=2,))
+  def testUseWithoutStrategy(self):
+    # There's no perfect way to check if the test runs in a subprocess. We
+    # approximate by checking the presence of TF_CONFIG, which is normally not
+    # set to the main process.
+    self.assertNotEqual(os.getenv("TF_CONFIG"), "")
+
+
+@combinations.generate(combinations.combine(num_workers=2))
+class ClusterCombinationTestEnvTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    # Note that test case fixtures are executed in both the main process and
+    # worker processes.
+    super().setUp()
+    if combinations.in_main_process():
+      combinations.env().tf_data_service_dispatcher = "localhost"
+
+  def testTfDataServiceDispatcher(self):
+    self.assertEqual(combinations.env().tf_data_service_dispatcher, "localhost")
+
+  def testUpdateEnvInWorker(self):
+    with self.assertRaises(ValueError):
+      combinations.env().tf_data_service_dispatcher = "localhost"
+
+
+# unittest.expectedFailure doesn't work with parameterized test methods, so we
+# have to decorate the class instead.
+@unittest.expectedFailure
+class ClusterParametersShouldFailTest(test.TestCase, parameterized.TestCase):
+
+  @framework_combinations.generate(  # pylint: disable=redundant-keyword-arg
+      framework_combinations.combine(
+          ds1=combinations.NamedDistribution(
+              "Strategy1", lambda: None, has_chief=True, num_workers=2),
+          ds2=combinations.NamedDistribution(
+              "Strategy2", lambda: None, has_chief=True, num_workers=2),
+      ),
+      test_combinations=(combinations.ClusterCombination(),))
+  def testMultipleDistributionMultiWorker(self, ds1, ds2):
+    # combinations library should raise an exception.
+    pass
+
+
+# Tests that we *actually* run the test method in multiple workers instead of
+# just passing silently. More importantly, it verifies that the test can fail.
+# Note that unittest.expectedFailure doesn't work with parameterized test
+# methods, so we have to decorate the class instead.
+@unittest.expectedFailure
+class CombinationsExpectedFailureTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(distribution=[
+          combinations.NamedDistribution(
+              "OneChiefOneWorker", lambda: None, has_chief=True, num_workers=1),
+          combinations.NamedDistribution(
+              "TwoWorkers", lambda: None, has_chief=False, num_workers=2),
+      ]))
+  def testMultiWorkerCanFail(self, distribution):
+    resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    # This should fail.
+    self.assertIsNone(resolver.task_id)
+
+
+# Tests that we *actually* run the test method in multiple workers instead of
+# just passing silently. More importantly, it verifies that the test can fail.
+# Note that unittest.expectedFailure doesn't work with parameterized test
+# methods, so we have to decorate the class instead.
+@unittest.expectedFailure
+@combinations.generate(
+    combinations.combine(distribution=[
+        combinations.NamedDistribution(
+            "OneChiefOneWorker", lambda: None, has_chief=True, num_workers=1),
+        combinations.NamedDistribution(
+            "TwoWorkers", lambda: None, has_chief=False, num_workers=2),
+    ]))
+class CombinationsOnClassMultiWorkerExpectedFailureTest(test.TestCase,
+                                                        parameterized.TestCase):
+
+  def test(self, distribution):
+    resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+    # This should fail.
+    self.assertIsNone(resolver.task_id)
+
+
+class TfFunctionTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          tf_function_1=combinations.tf_function,
+          tf_function_2=combinations.no_tf_function,
+          mode="eager",
+      ))
+  def testFunc(self, tf_function_1, tf_function_2):
+
+    @tf_function_1
+    def foo():
+      self.assertFalse(context.executing_eagerly())
+
+    @tf_function_2
+    def bar():
+      self.assertTrue(context.executing_eagerly())
+
+    foo()
+    bar()
+
+
+class ModuleInitializingTest(test.TestCase, parameterized.TestCase):
+
+  def testSysArgvClearedIsFine(self):
+    original_argv = list(sys.argv)
+    sys.argv.clear()
+    importlib.reload(combinations)
+    sys.argv = original_argv
+
+
+class ShareGPUTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if combinations.in_main_process():
+      num_gpus = combinations.env().total_phsyical_gpus
+      if num_gpus != 2 and num_gpus != 4:
+        self.skipTest("requires 2 or 4 GPUs")
+
+  # Test cases are annotated with required_gpus only for them to run in gpu
+  # targets, otherwise they will be skipped.
+
+  @combinations.generate(
+      combinations.combine(num_workers=2, required_gpus=1, share_gpu=True))
+  def testShareGPU(self):
+    self.assertLen(context.context().list_physical_devices("GPU"),
+                   combinations.env().total_phsyical_gpus)
+
+  @combinations.generate(combinations.combine(num_workers=2, required_gpus=1))
+  def testShareGPUByDefault(self):
+    self.assertLen(context.context().list_physical_devices("GPU"),
+                   combinations.env().total_phsyical_gpus)
+
+  @combinations.generate(
+      combinations.combine(num_workers=2, required_gpus=1, share_gpu=False))
+  def testNotShareGPU(self):
+    self.assertLen(context.context().list_physical_devices("GPU"),
+                   combinations.env().total_phsyical_gpus / 2)
+
+
+if __name__ == "__main__":
+  test_util.main()

tensorflow/python/distribute/coordinator/BUILD

@@ -0,0 +1,302 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+load("//tensorflow/core/platform:distribute.bzl", "distribute_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "cluster_coordinator",
+    srcs = ["cluster_coordinator.py"],
+    strict_deps = True,
+    deps = [
+        ":coordinator_context",
+        ":metric_utils",
+        ":remote_value",
+        ":utils",
+        ":values",
+        ":watchdog",
+        "//tensorflow/python/eager:cancellation",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:executor",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+        "@six_archive//:six",
+    ],
+)
+
+py_library(
+    name = "coordinator_context",
+    srcs = [
+        "coordinator_context.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":remote_value",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "remote_value",
+    srcs = ["remote_value.py"],
+    strict_deps = True,
+    deps = ["//tensorflow/python/util:tf_export"],
+)
+
+py_library(
+    name = "values",
+    srcs = ["values.py"],
+    strict_deps = True,
+    deps = [
+        ":remote_value",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/distribute:input_lib",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//tensorflow/python/ops:experimental_dataset_ops_gen",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+distribute_py_strict_test(
+    name = "cluster_coordinator_test",
+    srcs = ["cluster_coordinator_test.py"],
+    shard_count = 50,
+    tags = [
+        # "multi_gpu", # TODO(b/287692888): re-enable once the 2gpu test passes.
+        "no_oss",  # TODO(b/214432000): Very flaky under Docker
+        "no_pip",
+        "noasan",  # TODO(b/171040359): Flaky timeout, even if maximum shards
+        "notpu",
+        "notsan",  # TODO(b/171040359): Flaky timeout, even if maximum shards
+    ],
+    xla_enable_strict_auto_jit = False,  # TODO(b/291174864)
+    xla_tags = [
+        "no_cuda_asan",  # Race condition on async test
+    ],
+    deps = [
+        ":cluster_coordinator",
+        ":coordinator_context",
+        ":remote_value",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:distribute_utils",
+        "//tensorflow/python/distribute:input_lib",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute:sharded_variable",
+        "//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
+        "//tensorflow/python/eager:cancellation",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:embedding_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:coordinator",
+        "//tensorflow/python/training:server_lib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "fault_tolerance_test_base",
+    srcs = ["fault_tolerance_test_base.py"],
+    strict_deps = True,
+    deps = [
+        ":cluster_coordinator",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute:test_util",
+        "//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:check_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:coordinator",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "fault_tolerance_test",
+    srcs = ["fault_tolerance_test.py"],
+    shard_count = 41,  # = number of tests, so one shard = one test
+    tags = [
+        "no_oss",  # TODO(b/219580021)
+        "noasan",  # Multi-process runner does not work with test sanitizers
+        "nomac",  # TODO(b/177065434)
+        "nomsan",
+    ],
+    deps = [
+        ":cluster_coordinator",
+        ":fault_tolerance_test_base",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/training:coordinator",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "fault_tolerance_coordination_service_test",
+    srcs = ["fault_tolerance_coordination_service_test.py"],
+    shard_count = 41,
+    tags = [
+        # Inherit tags from fault_tolerance_test
+        "no_oss",  # TODO(b/219580021)
+        "noasan",  # Multi-process runner does not work with test sanitizers
+        "nomac",  # TODO(b/177065434)
+        "nomsan",  # TODO(b/290211396)
+        "notap",
+    ],
+    deps = [
+        ":cluster_coordinator",
+        ":fault_tolerance_test_base",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+    ],
+)
+
+py_library(
+    name = "metric_utils",
+    srcs = ["metric_utils.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/eager:monitoring",
+        "//tensorflow/python/util:tf_contextlib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "metric_utils_test",
+    srcs = ["metric_utils_test.py"],
+    shard_count = 3,
+    deps = [
+        ":cluster_coordinator",
+        ":metric_utils",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/training:coordinator",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:server_lib",
+    ],
+)
+
+py_library(
+    name = "remote_eager_lib",
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+)
+
+py_library(
+    name = "watchdog",
+    srcs = ["watchdog.py"],
+    strict_deps = True,
+    deps = ["@absl_py//absl/logging"],
+)
+
+tf_py_strict_test(
+    name = "watchdog_test",
+    srcs = ["watchdog_test.py"],
+    tags = [
+        "nomac",  # TODO(b/239433962)
+    ],
+    deps = [
+        ":watchdog",
+        "//tensorflow/python/eager:test",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "get_task_states_test",
+    srcs = ["get_task_states_test.py"],
+    shard_count = 3,
+    tags = [
+        "no_oss",  # TODO(b/219580021)
+        "nomac",  # TODO(b/177065434)
+    ],
+    deps = [
+        ":cluster_coordinator",
+        ":utils",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/training:server_lib",
+    ],
+)

tensorflow/python/distribute/coordinator/coordinator_context.py

@@ -0,0 +1,145 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""The execution context for ClusterCoordinator."""
+
+import contextlib
+import threading
+
+from tensorflow.core.framework import attr_value_pb2
+from tensorflow.python.distribute.coordinator import remote_value
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.util import compat
+from tensorflow.python.util.tf_export import tf_export
+
+_dispatch_context = threading.local()
+
+
+def get_current_dispatch_context():
+  try:
+    return _dispatch_context.current
+  except AttributeError:
+    return None
+
+
+@contextlib.contextmanager
+def with_dispatch_context(worker_obj):
+  previous_context = getattr(_dispatch_context, "current", None)
+  _dispatch_context.current = DispatchContext(worker_obj)
+  yield
+  _dispatch_context.current = previous_context
+
+
+class DispatchContext(object):
+  """Context entered when executing a closure on a given worker."""
+
+  def __init__(self, worker_obj):
+    self._worker = worker_obj
+    self._worker_index = worker_obj.worker_index
+
+  @property
+  def worker(self):
+    return self._worker
+
+  @property
+  def worker_index(self):
+    return self._worker_index
+
+  def maybe_get_remote_value(self, ret):
+    return maybe_get_remote_value(ret)
+
+
+def maybe_get_remote_value(val):
+  """Gets the value of `val` if it is a `RemoteValue`."""
+  if isinstance(val, remote_value.RemoteValue):
+    error = val._get_error()  # pylint: disable=protected-access
+    if error:
+      raise AssertionError(
+          "RemoteValue doesn't have a value because it has error %r:%s" %
+          (error, error))
+    elif val._status is not remote_value.RemoteValueStatus.READY:  # pylint: disable=protected-access
+      raise AssertionError("The input RemoteValue has not been executed.")
+    else:
+      return val._get_values()  # pylint: disable=protected-access
+  else:
+    return val
+
+
+@tf_export("distribute.coordinator.experimental_get_current_worker_index",
+           v1=[])
+def get_current_worker_index():
+  """Returns the current worker index, when called within a worker closure.
+
+  Some parameter server training workloads may require the worker to know its
+  index, for example for data sharding for reduced-variance training.
+
+  This method may be used within a `tf.function` that is executed on a worker.
+  That is, either a `dataset_fn` that runs via
+  `ClusterCoordinator.create_per_worker_dataset`, or any other function
+  scheduled via `ClusterCoordinator.schedule`.
+
+  Example (sharding data by worker):
+
+  ```python
+  strategy = tf.distribute.ParameterServerStrategy(
+      cluster_resolver=...)
+  coordinator = (
+      tf.distribute.coordinator.ClusterCoordinator(strategy))
+
+  def dataset_fn(context):
+    dataset = tf.data.Dataset.range(10)
+    worker_index = (
+        tf.distribute.coordinator.experimental_get_current_worker_index()
+    )
+    dataset = dataset.shard(
+        num_shards=num_workers,
+        index=worker_index,
+    )
+    return dataset
+
+  @tf.function
+  def per_worker_dataset_fn():
+    return strategy.distribute_datasets_from_function(dataset_fn)
+
+  per_worker_dataset = coordinator.create_per_worker_dataset(
+      per_worker_dataset_fn)
+  ```
+
+  Raises:
+    RuntimeError: if called from outside a `tf.function` or outside of a remote
+      closure execution context (that is, on a non-worker machine).
+  """
+
+  msg = ("Cannot retrieve the worker index. `get_worker_idx_and_num_workers` "
+         "should be called from within a tf.function being executed on a "
+         "worker. This method should only be called from either a dataset_fn "
+         "that is passed into `ClusterCoordinator.create_per_worker_dataset`, "
+         "or a tf.function that is passed into `ClusterCoordinator.schedule`.")
+  if not ops.inside_function():
+    raise RuntimeError(msg)
+
+  def call_time_worker_index():
+    dispatch_context = get_current_dispatch_context()
+    if not dispatch_context:
+      raise RuntimeError(msg)
+    return dispatch_context.worker_index
+
+  worker_index = ops.get_default_graph().capture_call_time_value(
+      call_time_worker_index, tensor.TensorSpec([], dtype=dtypes.int64))
+  worker_index.op._set_attr(  # pylint: disable=protected-access
+      "_user_specified_name",
+      attr_value_pb2.AttrValue(s=compat.as_bytes("worker_index")))
+  return worker_index

tensorflow/python/distribute/coordinator/fault_tolerance_coordination_service_test.py

@@ -0,0 +1,61 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Fault tolerance tests for coordination service-based failure handling."""
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute.coordinator import cluster_coordinator
+from tensorflow.python.distribute.coordinator import fault_tolerance_test_base
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+
+
+class BaseCoordinationServiceTest(
+    fault_tolerance_test_base.BaseFaultToleranceTest):
+  """Modify some tests to have stronger checks."""
+
+  def setUp(self, num_workers, num_ps):
+    super().setUp(num_workers=num_workers, num_ps=num_ps, use_cs=True)
+
+  def testJoinRaisesUnavailableErrorAtPsFailure(self):
+    self._run_and_kill_ps_task()
+    with self.assertRaises(cluster_coordinator.PSUnavailableError):
+      self.cluster_coord.join()
+
+  def testScheduleRaisesUnavailableErrorAtPsFailure(self):
+    self._run_and_kill_ps_task()
+    with self.assertRaises(cluster_coordinator.PSUnavailableError):
+      self.cluster_coord.schedule(def_function.function(lambda: None))
+
+
+class SingleWorkerCoordinationServiceTest(
+    BaseCoordinationServiceTest, test.TestCase
+):
+
+  def setUp(self):
+    super().setUp(num_workers=1, num_ps=1)
+
+
+class MultiWorkerCoordinationServiceTest(
+    BaseCoordinationServiceTest, test.TestCase
+):
+
+  def setUp(self):
+    super().setUp(num_workers=2, num_ps=2)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  multi_process_runner.test_main()

tensorflow/python/distribute/coordinator/fault_tolerance_test.py

@@ -0,0 +1,145 @@
+# 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.
+# ==============================================================================
+"""Fault tolerance test for parameter server training in TF2."""
+
+import threading
+import time
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import parameter_server_strategy_v2
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver
+from tensorflow.python.distribute.coordinator import cluster_coordinator
+from tensorflow.python.distribute.coordinator import fault_tolerance_test_base
+from tensorflow.python.eager import test
+from tensorflow.python.training import coordinator as thread_coordinator
+from tensorflow.python.training import server_lib
+
+
+class MultiWorkerFaultToleranceTest(
+    fault_tolerance_test_base.BaseFaultToleranceTest, test.TestCase):
+  """Multi worker fault tolerance tests.
+
+  This covers the ordinary cases where multiple workers and PS are used.
+  """
+
+  def setUp(self):
+    super(MultiWorkerFaultToleranceTest, self).setUp(2, 2)
+
+
+class SingleWorkerFaultToleranceTest(
+    fault_tolerance_test_base.BaseFaultToleranceTest, test.TestCase):
+  """Single worker fault tolerance tests.
+
+  This covers the cases that ensure training can continue in a single-worker
+  cluster, even if the only worker can become unavailable at some point and
+  recovered (if there are multiple workers, it is possible that the training
+  succeeds with the workers that did not fail). Realistically single worker
+  is very rarely used, but the tests are important to ensure the correct
+  behaviors.
+  """
+
+  def setUp(self):
+    super(SingleWorkerFaultToleranceTest, self).setUp(1, 1)
+
+
+class InitFaultToleranceTest(test.TestCase):
+  """Test preemptions during strategy init."""
+
+  def setUp(self):
+    super().setUp()
+    self.num_workers = 2
+    self.num_ps = 2
+    self._cluster = multi_worker_test_base.create_multi_process_cluster(
+        num_workers=self.num_workers,
+        num_ps=self.num_ps,
+        rpc_layer="grpc",
+        stream_output=True,
+    )
+    self._cluster_def = self._cluster.cluster_resolver.cluster_spec().as_dict()
+    self._cluster_def["chief"] = [
+        "localhost:%d" % multi_worker_test_base.pick_unused_port()
+    ]
+    self._cluster_resolver = SimpleClusterResolver(
+        server_lib.ClusterSpec(self._cluster_def), rpc_layer="grpc"
+    )
+
+    self.thread_coord = thread_coordinator.Coordinator(
+        clean_stop_exception_types=[]
+    )
+
+  def tearDown(self):
+    super().tearDown()
+    self._cluster.stop()
+    self._cluster = None
+
+  def testWorkerPreemptionDuringInit(self):
+
+    worker_down = threading.Condition()
+
+    def _restart_in_thread(downtime_secs, restart_job):
+      # During initial connection in SetOrUpdateServerDef, there is a step that
+      # waits to receive a `GetStatus` response from each worker, in an attempt
+      # to ensure workers are available before sending `CreateContext` to them.
+      # In high-preemption environments, the time between these two requests can
+      # be enough for a worker to be preempted, at which point the
+      # `CreateContext` request will fail and training cannot start. b/298187302
+      # We solve this by enabling retries on `CreateContext` calls when used
+      # with PSS.
+
+      # This test reproduces this behavior by bringing down a worker during
+      # startup, then waiting long enough for strategy startup to reach the
+      # point where it is just waiting for a `GetStatus` response from that
+      # worker. Then, we restart the first downed worker and bring down another
+      # worker at the same time, so the strategy startup progresses to the
+      # `CreateContext` requests, but not before another worker is down.
+
+      def _restart_fn():
+        with self.thread_coord.stop_on_exception():
+          self._cluster.kill_task(restart_job, 0)
+          with worker_down:
+            worker_down.notify_all()
+          time.sleep(downtime_secs)
+          self._cluster.start_task(restart_job, 0)
+          self._cluster.kill_task(restart_job, 1)
+          time.sleep(downtime_secs)
+          self._cluster.start_task(restart_job, 1)
+
+      restart_thread = threading.Thread(target=_restart_fn)
+      restart_thread.start()
+      return restart_thread
+
+    _restart_in_thread(downtime_secs=2, restart_job="worker")
+
+    with worker_down:
+      worker_down.wait()
+
+    # The strategy's constructor would connect to the cluster.
+    self.strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self._cluster_resolver
+    )
+    self.cluster_coord = cluster_coordinator.ClusterCoordinator(self.strategy)
+
+    # Now do a simple training as a sanity check
+    model = fault_tolerance_test_base.Model(self.cluster_coord)
+    model.schedule_training_functions(2)
+    model.join_training_functions()
+    self.assertEqual(model.iterations.numpy(), 2)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  multi_process_runner.test_main()

tensorflow/python/distribute/coordinator/fault_tolerance_test_base.py

@@ -0,0 +1,697 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Fault tolerance test base class for parameter server training in TF2."""
+
+import gc
+import os
+import sys
+import threading
+import time
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import parameter_server_strategy_v2
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver
+from tensorflow.python.distribute.coordinator import cluster_coordinator
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import coordinator as thread_coordinator
+from tensorflow.python.training import server_lib
+
+
+_RPC_ERROR_FROM_WORKER = "GRPC error information from remote target /job:worker"
+_RPC_ERROR_FROM_PS = "GRPC error information from remote target /job:ps"
+_WORKER_PREEMPTION_THREAD_NAME = "WorkerPreemptionHandler"
+_WORKER_THREAD_PREFIX = "WorkerClosureProcessingLoop"
+
+
+class Model(object):
+
+  def __init__(self, coordinator):
+    self.cluster_coord = coordinator
+    self.strategy = self.cluster_coord.strategy
+    with self.cluster_coord.strategy.scope():
+      self.build()
+
+  def build(self):
+    self.w = variables.Variable(
+        initial_value=random_ops.random_uniform((10, 10)), dtype=dtypes.float32)
+    self.iterations = variables.Variable(initial_value=0, dtype=dtypes.int32)
+    # Allow external control to make the model run its train_fn in an infinite
+    # loop. This allows us to reliably test worker preemption in the middle of
+    # function execution.
+    self.do_infinite_step = variables.Variable(False)
+
+    self.rebuild_iterators()
+
+  def rebuild_iterators(self, use_dataset_fn=True):
+    if use_dataset_fn:
+
+      def dataset_fn():
+        data = random_ops.random_uniform((10, 10))
+        dataset = dataset_ops.DatasetV2.from_tensors([data]).repeat()
+        return dataset
+
+      def distribute_dataset_fn():
+        return self.cluster_coord.strategy.distribute_datasets_from_function(
+            lambda _: dataset_fn())
+
+      self.iterator = iter(
+          self.cluster_coord.create_per_worker_dataset(distribute_dataset_fn))
+      self.iterator2 = iter(
+          self.cluster_coord.create_per_worker_dataset(distribute_dataset_fn))
+    else:
+      data = random_ops.random_uniform((10, 10))
+      dataset = dataset_ops.DatasetV2.from_tensors([data]).repeat()
+
+      self.iterator = iter(
+          self.cluster_coord.create_per_worker_dataset(dataset))
+      self.iterator2 = iter(
+          self.cluster_coord.create_per_worker_dataset(dataset))
+
+  def _train_fn_internal(self, iterator, iterator2):
+    x = math_ops.matmul(array_ops.squeeze(next(iterator)), self.w)
+    x = math_ops.matmul(array_ops.squeeze(next(iterator2)), x)
+    x = math_ops.matmul(random_ops.random_uniform((10, 10)), x)
+    self.w.assign_add(x)
+
+  @def_function.function
+  def train_fn(self, iterator, iterator2):
+    self._train_fn_internal(iterator, iterator2)
+    while self.do_infinite_step:
+      self._train_fn_internal(iterator, iterator2)
+    self.iterations.assign_add(1)
+
+  def schedule_training_functions(self, num_steps):
+    with self.strategy.scope():
+      for _ in range(num_steps):
+        self.cluster_coord.schedule(
+            self.train_fn, args=(self.iterator, self.iterator2))
+
+  def join_training_functions(self):
+    self.do_infinite_step.assign(False)
+    self.cluster_coord.join()
+
+
+class BaseFaultToleranceTest(object):  # pylint: disable=missing-docstring
+
+  def setUp(self, num_workers, num_ps, use_cs=False):
+    super(BaseFaultToleranceTest, self).setUp()
+
+    self._cluster = multi_worker_test_base.create_multi_process_cluster(
+        num_workers=num_workers,
+        num_ps=num_ps,
+        rpc_layer="grpc",
+        stream_output=True,
+    )
+    self._cluster_def = self._cluster.cluster_resolver.cluster_spec().as_dict()
+    self._cluster_def["chief"] = [
+        "localhost:%d" % multi_worker_test_base.pick_unused_port()
+    ]
+    cluster_resolver = SimpleClusterResolver(
+        server_lib.ClusterSpec(self._cluster_def), rpc_layer="grpc"
+    )
+
+    if use_cs:
+      os.environ["TF_PSS_ENABLE_COORDINATION_SERVICE"] = "1"
+    # The strategy's constructor would connect to the cluster.
+    self.strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        cluster_resolver
+    )
+    self.cluster_coord = cluster_coordinator.ClusterCoordinator(self.strategy)
+
+    self.thread_coord = thread_coordinator.Coordinator(
+        clean_stop_exception_types=[]
+    )
+    self.num_workers = num_workers
+    self.num_ps = num_ps
+
+  def tearDown(self):
+    super(BaseFaultToleranceTest, self).tearDown()
+    self._cluster.stop()
+    self._cluster = None
+
+  def _restart(self, downtime_secs, job):
+    """Kills `job` (index: 0) and restarts it after `downtime_secs`.
+
+    Args:
+      downtime_secs: secs before restarting the job.
+      job: a string specifying the job to restart.
+    """
+    self._cluster.kill_task(job, 0)
+    time.sleep(downtime_secs)
+    self.assertFalse(context.check_alive("/job:%s/replica:0/task:0" % job))
+    self._cluster.start_task(job, 0)
+    while not context.check_alive("/job:%s/replica:0/task:0" % job):
+      time.sleep(1)
+
+  def _restart_in_thread(self, downtime_secs, restart_job):
+
+    def _restart_fn():
+      with self.thread_coord.stop_on_exception():
+        self._restart(downtime_secs, restart_job)
+
+    restart_thread = threading.Thread(target=_restart_fn)
+    restart_thread.start()
+    return restart_thread
+
+  def _ensure_threads_closed(self):
+    """Ensures worker and preemption threads are closed."""
+    # Worker and preemption threads should exist before releasing
+    # ClusterCoordinator.
+    running_threads = test_util.get_running_threads()
+    self.assertTrue(
+        test_util.has_thread(_WORKER_THREAD_PREFIX, running_threads))
+    self.assertIn(_WORKER_PREEMPTION_THREAD_NAME, running_threads)
+
+    # Print object graph if ClusterCoordinator may leak.
+    if sys.getrefcount(self.cluster_coord) > 2:
+      try:
+        test_util.show_backref(self.cluster_coord)
+      except:  # pylint: disable=bare-except
+        pass
+
+    # Wait for threads to close.
+    self.cluster_coord = None
+    self.strategy = None
+    gc.collect()
+    time.sleep(1)
+
+    # Verify thread names.
+    running_threads = test_util.get_running_threads()
+    self.assertNotIn(_WORKER_PREEMPTION_THREAD_NAME, running_threads)
+    self.assertFalse(
+        test_util.has_thread(_WORKER_THREAD_PREFIX, running_threads),
+        "Worker thread is not stopped properly.")
+
+  def _create_model_and_run_indefinitely(self):
+    model = Model(self.cluster_coord)
+    model.do_infinite_step.assign(True)
+    model.schedule_training_functions(10)
+    # Model does infinite training step, so at this moment, we expect to have
+    # `self.num_workers` infinite closures inflight, and `10-self.num_workers`
+    # closures in the queue.
+    while (self.cluster_coord._cluster.closure_queue._inflight_closure_count <
+           self.num_workers):
+      time.sleep(0.1)
+    return model
+
+  def testClusterCoordinatorDestroyed(self):
+    self._ensure_threads_closed()
+
+  def testWorkerPreemptionBetweenFunctions(self):
+    model = Model(self.cluster_coord)
+    model.schedule_training_functions(2)
+    model.join_training_functions()
+    self.assertEqual(model.iterations.numpy(), 2)
+
+    self._restart(downtime_secs=2, job="worker")
+
+    model.schedule_training_functions(2)
+    model.join_training_functions()
+    self.assertEqual(model.iterations.numpy(), 4)
+
+  def testWorkerPreemptionMidstFunction(self):
+    model = Model(self.cluster_coord)
+    model.do_infinite_step.assign(True)
+
+    model.schedule_training_functions(4)
+    # Model does infinite training step, so at this moment, we expect to have
+    # `self.num_workers` infinite closures inflight, and `4-self.num_workers`
+    # closures in the queue.
+    while (self.cluster_coord._cluster.closure_queue._inflight_closure_count <
+           self.num_workers):
+      time.sleep(0.1)
+    self.assertFalse(self.cluster_coord.done())
+    self._restart(downtime_secs=2, job="worker")
+    model.join_training_functions()
+    self.assertGreaterEqual(model.iterations.numpy(), 4)
+
+  def testOneWorkerPreemptionWithCancellation(self):
+
+    @def_function.function
+    def normal_function():
+      x = random_ops.random_uniform((2, 10))
+      y = random_ops.random_uniform((10, 2))
+      return math_ops.reduce_mean(math_ops.matmul(x, y))
+
+    @def_function.function
+    def error_function():
+      x = random_ops.random_uniform((2, 10))
+      y = random_ops.random_uniform((10, 2))
+      check_ops.assert_non_positive_v2(
+          math_ops.reduce_sum(math_ops.matmul(x, y)))
+      return x
+
+    @def_function.function
+    def long_function():
+      x = random_ops.random_uniform((1000, 1000))
+      for _ in math_ops.range(10000):
+        a = random_ops.random_uniform((1000, 1000))
+        b = random_ops.random_uniform((1000, 1000))
+        x += math_ops.matmul(a, b)
+      return x
+
+    for _ in range(3):
+      self.cluster_coord.schedule(normal_function)
+    long_function_result = self.cluster_coord.schedule(long_function)
+    self.cluster_coord.schedule(error_function)
+
+    time.sleep(1)  # Let it run a couple steps.
+    self._restart(2, "worker")
+
+    # InvalidArgumentError thrown from the error_function.
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.cluster_coord.join()
+
+    # CancelledError thrown by ClusterCoordinator after cancelling due to user
+    # error.
+    with self.assertRaises(errors.CancelledError):
+      long_function_result.fetch()
+
+    for _ in range(3):
+      self.cluster_coord.schedule(normal_function)
+    self.cluster_coord.join()
+
+    # The cluster is likely still being recovered since `join` returned early
+    # due to the error_function.
+    failure_handler = self.cluster_coord._cluster.failure_handler
+    failure_handler.stop()
+    failure_handler._preemption_handler_thread.join()
+
+  def testHandleDatasetCreationFailureWithDatasetFn(self):
+    model = Model(self.cluster_coord)
+
+    restart_thread = self._restart_in_thread(5, "worker")
+
+    model.schedule_training_functions(3)
+    model.rebuild_iterators()
+    model.schedule_training_functions(3)
+    model.rebuild_iterators()
+    model.schedule_training_functions(3)
+
+    model.join_training_functions()
+
+    self.thread_coord.join([restart_thread])
+    self.assertGreaterEqual(model.iterations.numpy(), 3)
+
+  # TODO(yuefengz): consider using combinations when there is more code
+  # duplication.
+  def testHandleDatasetCreationFailureWithDataset(self):
+    model = Model(self.cluster_coord)
+
+    restart_thread = self._restart_in_thread(5, "worker")
+
+    model.schedule_training_functions(3)
+    model.rebuild_iterators(use_dataset_fn=False)
+    model.schedule_training_functions(3)
+    model.rebuild_iterators(use_dataset_fn=False)
+    model.schedule_training_functions(3)
+
+    model.join_training_functions()
+
+    self.thread_coord.join([restart_thread])
+    self.assertGreaterEqual(model.iterations.numpy(), 3)
+
+  def testWorkerPreemptionErrorType(self):
+
+    @def_function.function
+    def worker_train_fn():
+      x = random_ops.random_uniform((2, 10))
+      y = random_ops.random_uniform((10, 2))
+      return math_ops.reduce_mean(math_ops.matmul(x, y))
+
+    def run_fn():
+      with self.thread_coord.stop_on_exception():
+        with ops.device("/job:worker/replica:0/task:0"):
+          for _ in range(3):
+            for _ in range(3):
+              worker_train_fn()
+            time.sleep(5)
+
+    run_thread = threading.Thread(target=run_fn)
+    run_thread.start()
+    time.sleep(1)  # Let it run a couple steps.
+    self._restart(2, "worker")
+
+    try:
+      self.thread_coord.join([run_thread])
+    except (errors.UnavailableError, errors.AbortedError) as e:
+      logging.info("Got exception %r, error message is %s", e, e)
+
+      self.assertIn(_RPC_ERROR_FROM_WORKER, str(e))  # pylint: disable=g-assert-in-except
+      self.assertNotIn(_RPC_ERROR_FROM_PS, str(e))
+
+      self.assertTrue("failed to connect to all addresses" in str(e) or
+                      "Unable to find a context_id" in str(e) or
+                      "Socket closed" in str(e) or
+                      "Connection reset by peer" in str(e) or
+                      "Transport closed" in str(e))
+
+  def testWorkerPreemptionErrorTypeWithPythonFunction(self):
+
+    def worker_train_fn():
+      x = random_ops.random_uniform((2, 10))
+      y = random_ops.random_uniform((10, 2))
+      return math_ops.reduce_mean(math_ops.matmul(x, y))
+
+    def run_fn():
+      with self.thread_coord.stop_on_exception():
+        with ops.device("/job:worker/replica:0/task:0"):
+          for _ in range(3):
+            for _ in range(3):
+              worker_train_fn()
+            time.sleep(5)
+
+    run_thread = threading.Thread(target=run_fn)
+    run_thread.start()
+    time.sleep(1)  # Let it run a couple steps.
+    self._restart(2, "worker")
+
+    try:
+      self.thread_coord.join([run_thread])
+    except (errors.UnavailableError, errors.AbortedError) as e:
+      logging.info("Got exception %r, error message is %s", e, e)
+
+      self.assertIn(_RPC_ERROR_FROM_WORKER, str(e))  # pylint: disable=g-assert-in-except
+      self.assertNotIn(_RPC_ERROR_FROM_PS, str(e))
+
+      self.assertTrue("failed to connect to all addresses" in str(e) or
+                      "Unable to find a context_id" in str(e) or
+                      "Socket closed" in str(e) or
+                      "Connection reset by peer" in str(e) or
+                      "Transport closed" in str(e))
+
+  def testPSPreemptionErrorType(self):
+
+    with ops.device("/job:ps/replica:0/task:0"):
+      v = variables.Variable(
+          initial_value=random_ops.random_uniform((2, 10)),
+          dtype=dtypes.float32)
+
+    @def_function.function
+    def worker_train_fn():
+      y = random_ops.random_uniform((10, 2))
+      return math_ops.reduce_mean(math_ops.matmul(v, y))
+
+    def run_fn():
+      with self.thread_coord.stop_on_exception():
+        with ops.device("/job:worker/replica:0/task:0"):
+          for _ in range(3):
+            for _ in range(3):
+              worker_train_fn()
+            time.sleep(5)
+
+    run_thread = threading.Thread(target=run_fn)
+    run_thread.start()
+    time.sleep(1)  # Let it run a couple steps.
+
+    # Use a short restart delay to cover the case that RPC channel is reused
+    self._restart(1, "ps")
+
+    try:
+      self.thread_coord.join([run_thread])
+    except (errors.UnavailableError, errors.AbortedError) as e:
+      logging.info("Got exception %r, error message is %s", e, e)
+      self.assertIn(_RPC_ERROR_FROM_PS, str(e))  # pylint: disable=g-assert-in-except
+
+      if isinstance(e, errors.UnavailableError):
+        self.assertTrue("failed to connect to all addresses" in str(e) or
+                        "Socket closed" in str(e) or
+                        "Connection reset by peer" in str(e) or
+                        "Transport closed" in str(e))
+
+      if isinstance(e, errors.AbortedError):
+        self.assertTrue(
+            "RecvTensor expects a different device incarnation" in str(e) or
+            "Unable to find a context_id" in str(e))
+      self._ensure_threads_closed()
+
+  def testTwoWorkersPreempted(self):
+    if self.num_workers < 2:
+      self.skipTest("Worker number is less than 2.")
+    model = self._create_model_and_run_indefinitely()
+
+    self.assertFalse(self.cluster_coord.done())
+    self._cluster.kill_task("worker", 0)
+    self._cluster.kill_task("worker", 1)
+    time.sleep(2)
+    self.assertFalse(context.check_alive("/job:worker/replica:0/task:0"))
+    self.assertFalse(context.check_alive("/job:worker/replica:0/task:1"))
+    self._cluster.start_task("worker", 0)
+    self._cluster.start_task("worker", 1)
+    time.sleep(2)
+    self.assertTrue(context.check_alive("/job:worker/replica:0/task:0"))
+    self.assertTrue(context.check_alive("/job:worker/replica:0/task:1"))
+
+    model.join_training_functions()
+    self.assertGreaterEqual(model.iterations.numpy(), 10)
+
+  def testWorkerContinuousFailure(self):
+    model = self._create_model_and_run_indefinitely()
+
+    self.assertFalse(self.cluster_coord.done())
+    self._cluster.kill_task("worker", 0)
+    time.sleep(2)
+    self.assertFalse(context.check_alive("/job:worker/replica:0/task:0"))
+    self._cluster.start_task("worker", 0)
+    time.sleep(2)
+    self.assertTrue(context.check_alive("/job:worker/replica:0/task:0"))
+    self._cluster.kill_task("worker", 0)
+    time.sleep(2)
+    self.assertFalse(context.check_alive("/job:worker/replica:0/task:0"))
+    self._cluster.start_task("worker", 0)
+    time.sleep(2)
+    self.assertTrue(context.check_alive("/job:worker/replica:0/task:0"))
+
+    model.join_training_functions()
+    self.assertGreaterEqual(model.iterations.numpy(), 10)
+
+  def testPSFailureWhileRecoveryFromWokerFailure(self):
+    model = self._create_model_and_run_indefinitely()
+
+    time.sleep(1)
+    self.assertFalse(self.cluster_coord.done())
+
+    def kill(task):
+      self._cluster.kill_task(task, 0)
+      self.sleep(1)
+      self._cluster.start_task(task, 0)
+
+    kill_thread_1 = threading.Thread(target=kill, args=("worker",))
+    kill_thread_2 = threading.Thread(target=kill, args=("ps",))
+    kill_thread_1.start()
+    kill_thread_2.start()
+    kill_thread_1.join()
+    kill_thread_2.join()
+
+    with self.assertRaises(
+        (errors.UnavailableError, errors.InvalidArgumentError)):
+      model.join_training_functions()
+
+  def testNumpyFetchedAfterWorkerFailure(self):
+
+    with self.strategy.scope():
+      v = variables.Variable(initial_value=0, dtype=dtypes.int32)
+
+    @def_function.function
+    def worker_fn():
+      return v + 1, v - 1
+
+    remote_value = self.cluster_coord.schedule(worker_fn)
+    # Attempt to fetch before killing worker task should succeed.
+    self.assertEqual((1, -1), remote_value.fetch())
+    self._cluster.kill_task("worker", 0)
+    # So should attempt to fetch after killing worker task.
+    self.assertEqual((1, -1), remote_value.fetch())
+
+  def testTensorGotAfterWorkerFailure(self):
+
+    with self.strategy.scope():
+      v = variables.Variable(initial_value=0, dtype=dtypes.int32)
+
+    @def_function.function
+    def worker_fn():
+      return v + 1, v - 1
+
+    remote_value = self.cluster_coord.schedule(worker_fn)
+
+    # Attempt to fetch before killing worker task should succeed.
+    fetched = remote_value.get()[0]
+    self.assertIsInstance(fetched, tensor.Tensor)
+    self.assertEqual(fetched.device, "/job:chief/replica:0/task:0/device:CPU:0")
+    self.assertEqual((1, -1), remote_value.get())
+    remote_value.get()[0].numpy()
+
+    # As well as the remote tensors that point to worker0 or worker1.
+    values = remote_value._values[0]
+    self.assertIsInstance(values, tensor.Tensor)
+    self.assertRegex(values.device,
+                     "/job:worker/replica:0/task:[0-1]/device:CPU:0")
+    self.assertEqual((1, -1), remote_value._values)
+    remote_value._values[0].numpy()
+
+    # Terminate the workers and wait a little so that they are indeed killed.
+    for i in range(self.num_workers):
+      self._cluster.kill_task("worker", i)
+    time.sleep(5)
+
+    # Attempt to fetch after killing worker tasks should succeed as well.
+    remote_value.get()[0].numpy()
+    self.assertEqual((1, -1), remote_value.get())
+
+    # Attempting to copy the tensor from worker now should fail.
+    with self.assertRaises(errors.UnavailableError) as cm:
+      remote_value._values[0].numpy()
+    self.assertIn("failed to connect to all addresses", cm.exception.message)
+    self.assertIn("/job:worker/replica:0/task:", cm.exception.message)
+
+  def testFetchFromPSAfterWorkerFailure(self):
+    # Test for flaky failures when reading from a parameter server while a
+    # worker is recovering.
+    # Place some variables on PSes, kill a worker, and continuously poll one of
+    # those variables.
+
+    model = Model(self.cluster_coord)
+
+    # kill the worker after a delay to make sure variable reading runs while
+    # worker is up, while it's down, and while it restarts
+    def kill_after_delay():
+      time.sleep(3)
+      logging.info("Killing worker 0")
+      self._cluster.kill_task("worker", 0)
+      time.sleep(1)
+      logging.info("Restarting worker 0")
+      self._cluster.start_task("worker", 0)
+
+    kill_thread = threading.Thread(target=kill_after_delay)
+    kill_thread.start()
+
+    model.do_infinite_step.assign(True)
+    model.schedule_training_functions(1)
+
+    num_reads = 0
+    num_reads_after_restart = 0
+    read_interval_secs = 0.1
+    worker_has_stopped = False
+    # limit runtime of the test: stop after doing a few reads after worker
+    # is back up, or after a fixed maximum number of reads
+    while num_reads_after_restart <= 5 and num_reads < 200:
+      worker_up = context.check_alive("/job:worker/replica:0/task:0")
+      if not worker_up:
+        worker_has_stopped = True
+      if worker_up and worker_has_stopped:
+        num_reads_after_restart += 1
+
+      model.join_training_functions()
+      start = time.time()
+      while time.time() < start + read_interval_secs:
+        model.iterations.read_value()
+
+      num_reads += 1
+      # run another epoch
+      model.do_infinite_step.assign(True)
+      model.schedule_training_functions(1)
+
+  def testClusterStateNotDisrupted(self):
+    # This test has side effects and can disrupt other tests, even if the
+    # resource created by it will not be used in following tests.
+    # TODO(b/155209534): enable this test.
+    # self.testPSPreemptionErrorType()
+
+    self.thread_coord = thread_coordinator.Coordinator(
+        clean_stop_exception_types=[])
+    self.testWorkerPreemptionMidstFunction()
+
+    self.thread_coord = thread_coordinator.Coordinator(
+        clean_stop_exception_types=[])
+    self.testWorkerPreemptionErrorType()
+
+    # In previous tests, workers may fail after training is done. But the
+    # following tests start with creating resources where failure is not
+    # handled.
+    # TODO(b/153888707): enable the following two tests.
+    # self.testTwoWorkersPreempted()
+    # self.testWorkerContinuousFailure()
+
+  def _run_and_kill_ps_task(self):
+    self._create_model_and_run_indefinitely()
+    self._cluster.kill_task("ps", 0)
+    while self.cluster_coord._cluster.closure_queue._error is None:
+      time.sleep(1)
+    logging.info("Trying to join, expecting error")
+
+  def testJoinRaisesUnavailableErrorAtPsFailure(self):
+    self._run_and_kill_ps_task()
+    with self.assertRaises((errors.UnavailableError, errors.NotFoundError,
+                            errors.FailedPreconditionError)):
+      self.cluster_coord.join()
+
+  def testScheduleRaisesUnavailableErrorAtPsFailure(self):
+    self._run_and_kill_ps_task()
+    with self.assertRaises((errors.UnavailableError, errors.NotFoundError,
+                            errors.FailedPreconditionError)):
+      self.cluster_coord.schedule(def_function.function(lambda: None))
+
+  def testWorkerExecutionAfterPsFailureRaisesExpectedError(self):
+    model = self._create_model_and_run_indefinitely()
+    for i in range(self.num_ps):
+      self._cluster.kill_task("ps", i)
+    while self.cluster_coord._cluster.closure_queue._error is None:
+      time.sleep(1)
+
+    @def_function.function
+    def trivial_function():
+      return model.iterations + 1
+
+    for i in range(self.num_workers):
+      try:
+        with ops.device("/job:worker/replica:0/task:{}".format(i)):
+          trivial_function()
+      except Exception as e:  # pylint: disable=broad-except
+        if cluster_coordinator._is_ps_failure(e):  # pylint: disable=protected-access
+          if i < self.num_workers - 1:
+            continue
+          return
+      raise AssertionError("Executing a function after PS fails, should "
+                           "result in a PS failure.")
+
+  def testAsyncWaitIsNoOp(self):
+    if self.num_workers < 2:
+      self.skipTest("Worker number is less than 2.")
+    model = self._create_model_and_run_indefinitely()
+
+    self.assertFalse(self.cluster_coord.done())
+    self._cluster.kill_task("worker", 0)
+    time.sleep(2)
+    self.assertFalse(context.check_alive("/job:worker/replica:0/task:0"))
+    # Should pass without exception even with failed remote workers
+    context.async_wait()
+
+    model.join_training_functions()
+    self.assertGreaterEqual(model.iterations.numpy(), 10)
+
+    self._cluster.start_task("worker", 0)

tensorflow/python/distribute/coordinator/get_task_states_test.py

@@ -0,0 +1,157 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Get task states test for parameter server strategy in TF2."""
+
+import time
+
+from tensorflow.core.lib.core import error_codes_pb2
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import parameter_server_strategy_v2
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
+from tensorflow.python.distribute.coordinator import cluster_coordinator
+from tensorflow.python.distribute.coordinator import utils
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.framework import errors
+from tensorflow.python.training import server_lib
+
+_PULL_FREQ_IN_SEC = 2
+_COORDINATION_ERROR_PAYLOAD_KEY = (
+    "type.googleapis.com/tensorflow.CoordinationServiceError"
+)
+
+
+class GetTaskStatesTest(object):  # pylint: disable=missing-docstring
+
+  def setUp(self, num_workers, num_ps):
+    super().setUp()
+
+    self._cluster = multi_worker_test_base.create_multi_process_cluster(
+        num_workers=num_workers, num_ps=num_ps, rpc_layer="grpc")
+    self._cluster_def = self._cluster.cluster_resolver.cluster_spec().as_dict()
+    self._cluster_def["chief"] = [
+        "localhost:%d" % multi_worker_test_base.pick_unused_port()
+    ]
+    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
+        server_lib.ClusterSpec(self._cluster_def), rpc_layer="grpc")
+
+    context.context().configure_coordination_service(
+        service_type="standalone",
+        service_leader="/job:ps/replica:0/task:0",
+        heartbeat_timeout_in_ms=_PULL_FREQ_IN_SEC * 1000,
+        allow_new_incarnation_to_reconnect=True)
+    self.strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        cluster_resolver)
+    self.cluster_coord = cluster_coordinator.ClusterCoordinator(self.strategy)
+
+    self.num_workers = num_workers
+    self.num_ps = num_ps
+
+    self.states = None
+    self.polling_thread = utils.RepeatedTimer(
+        interval=_PULL_FREQ_IN_SEC, function=self.get_task_states)
+
+  def tearDown(self):
+    super().tearDown()
+    self.polling_thread.stop()
+    self._cluster.stop()
+    self._cluster = None
+
+  def get_task_states(self):
+    self.states = context.context().get_task_states([("worker",
+                                                      self.num_workers),
+                                                     ("ps", self.num_ps)])
+
+  def testAllTasksHealthy(self):
+    time.sleep(_PULL_FREQ_IN_SEC * 1.5)
+    self.assertLen(self.states, self.num_workers + self.num_ps)
+    for state in self.states:
+      self.assertIsNone(state)
+
+  def testWorkerPreempted(self):
+    self._cluster.kill_task("worker", 0)
+    time.sleep(_PULL_FREQ_IN_SEC * 2)
+    self.assertLen(self.states, self.num_workers + self.num_ps)
+    self.assertIsInstance(self.states[0], errors.UnavailableError)
+    self.assertIn("/job:worker/replica:0/task:0", self.states[0]._message)
+    self.assertEqual(self.states[0]._error_code, error_codes_pb2.UNAVAILABLE)
+    self.assertIn(_COORDINATION_ERROR_PAYLOAD_KEY,
+                  self.states[0]._experimental_payloads)
+    for i in range(1, self.num_workers + self.num_ps):
+      self.assertIsNone(self.states[i])
+    self._cluster.start_task("worker", 0)
+    context.context().update_server_def(context.get_server_def())
+    time.sleep(_PULL_FREQ_IN_SEC * 2)
+    for state in self.states:
+      self.assertIsNone(state)
+
+  def testPSPreempted(self):
+    self._cluster.kill_task("ps", 1)
+    time.sleep(_PULL_FREQ_IN_SEC * 2)
+    self.assertLen(self.states, self.num_workers + self.num_ps)
+    state_ix = self.num_workers + 1
+    self.assertIsInstance(self.states[state_ix], errors.UnavailableError)
+    self.assertIn("/job:ps/replica:0/task:1", self.states[state_ix]._message)
+    self.assertEqual(self.states[state_ix]._error_code,
+                     error_codes_pb2.UNAVAILABLE)
+    # Simulate the restart of all the tasks.
+    self._cluster.kill_task("ps", 0)
+    for index in range(2, self.num_ps):
+      self._cluster.kill_task("ps", index)
+    for index in range(self.num_workers):
+      self._cluster.kill_task("worker", index)
+    for index in range(self.num_ps):
+      self._cluster.start_task("ps", index)
+    for index in range(self.num_workers):
+      self._cluster.start_task("worker", index)
+    context.context().update_server_def(context.get_server_def())
+    time.sleep(_PULL_FREQ_IN_SEC * 2)
+    self.assertLen(self.states, self.num_workers + self.num_ps)
+    for state in self.states:
+      self.assertIsNone(state)
+
+  def testCoordinationServicePreempted(self):
+    self._cluster.kill_task("ps", 0)
+    time.sleep(_PULL_FREQ_IN_SEC * 2)
+    # `states` is None since Coordination Service is not available.
+    self.assertIsNone(self.states)
+    # Simulate the restart of all the tasks.
+    for index in range(1, self.num_ps):
+      self._cluster.kill_task("ps", index)
+    for index in range(self.num_workers):
+      self._cluster.kill_task("worker", index)
+    for index in range(self.num_ps):
+      self._cluster.start_task("ps", index)
+    for index in range(self.num_workers):
+      self._cluster.start_task("worker", index)
+    context.context().update_server_def(context.get_server_def())
+    time.sleep(_PULL_FREQ_IN_SEC * 2)
+    self.assertLen(self.states, self.num_workers + self.num_ps)
+    for state in self.states:
+      self.assertIsNone(state)
+
+
+class MultiWorkerGetTaskStatesTest(GetTaskStatesTest, test.TestCase):
+  """This covers the cases where multiple workers and PS are used."""
+
+  def setUp(self):
+    super().setUp(2, 2)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  multi_process_runner.test_main()

tensorflow/python/distribute/coordinator/metric_utils.py

@@ -0,0 +1,158 @@
+# 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.
+# ==============================================================================
+"""Metrics collecting utilities for single client training."""
+
+import time
+
+from tensorflow.python.eager import monitoring
+from tensorflow.python.util import tf_contextlib
+
+enable_metrics = True
+_METRICS_MAPPING = {}
+
+
+def _init():
+  """Initialize the metrics mapping."""
+  global _METRICS_MAPPING
+
+  # Define the boundaries for bucketing times of distribution (Sampler) metrics.
+
+  # Closure execution: range from 0.1s to 10000s, i.e. [(0.1, 1), (1, 10), ...]
+  execution_time_buckets = monitoring.ExponentialBuckets(
+      scale=0.1, growth_factor=10, bucket_count=6)
+  # Tracing: same range as execution
+  tracing_time_buckets = execution_time_buckets
+  # Remote value fetch: range from 0.001s (i.e. 1ms) to 1000s
+  fetch_time_buckets = monitoring.ExponentialBuckets(
+      scale=0.001, growth_factor=10, bucket_count=7)
+  # Server def update: range from 1s to 10000s
+  server_update_time_buckets = monitoring.ExponentialBuckets(
+      scale=1, growth_factor=10, bucket_count=5)
+
+  function_tracing_sampler = monitoring.Sampler(
+      '/tensorflow/api/ps_strategy/coordinator/function_tracing',
+      tracing_time_buckets,
+      'Sampler to track the time (in seconds) for tracing functions.')
+
+  closure_execution_sampler = monitoring.Sampler(
+      '/tensorflow/api/ps_strategy/coordinator/closure_execution',
+      execution_time_buckets,
+      'Sampler to track the time (in seconds) for executing closures.')
+
+  remote_value_fetch_sampler = monitoring.Sampler(
+      '/tensorflow/api/ps_strategy/coordinator/remote_value_fetch',
+      fetch_time_buckets,
+      'Sampler to track the time (in seconds) for fetching remote_value.')
+
+  server_def_update_sampler = monitoring.Sampler(
+      '/tensorflow/api/ps_strategy/coordinator/server_def_update',
+      server_update_time_buckets,
+      'Sample to track the time (in seconds) for updating the server def upon '
+      'worker recovery.')
+
+  queued_closure_gauge = monitoring.IntGauge(
+      '/tensorflow/api/ps_strategy/coordinator/queued_closures',
+      'Track how many closures are in the coordinator queue pending execution.')
+
+  inflight_closure_gauge = monitoring.IntGauge(
+      '/tensorflow/api/ps_strategy/coordinator/inflight_closures',
+      'Track how many closures are currently being processed by workers.')
+
+  worker_failure_counter = monitoring.Counter(
+      '/tensorflow/api/ps_strategy/coordinator/recoverable_worker_failure_count',
+      'Track how many recoverable worker failures have been encountered.')
+
+  _METRICS_MAPPING = {
+      'function_tracing': function_tracing_sampler,
+      'closure_execution': closure_execution_sampler,
+      'remote_value_fetch': remote_value_fetch_sampler,
+      'server_def_update': server_def_update_sampler,
+      'queued_closures': queued_closure_gauge,
+      'inflight_closures': inflight_closure_gauge,
+      'worker_failures': worker_failure_counter,
+  }
+
+
+@tf_contextlib.contextmanager
+def monitored_timer(metric_name, state_tracker=None):
+  """Monitor the execution time and collect it into the specified metric."""
+  if not enable_metrics:
+    yield
+  else:
+    if not _METRICS_MAPPING:
+      _init()
+    start_time = time.time()
+    start_state = state_tracker() if state_tracker else None
+    yield
+    duration_sec = time.time() - start_time
+    # If a state_checker is provided, record the metric only if the end state is
+    # different from the start state.
+    if state_tracker is None or state_tracker() != start_state:
+      metric = _METRICS_MAPPING[metric_name]
+      metric.get_cell().add(duration_sec)
+
+
+def monitor_int(metric_name, value):
+  if not enable_metrics:
+    return
+  else:
+    if not _METRICS_MAPPING:
+      _init()
+    metric = _METRICS_MAPPING[metric_name]
+    metric.get_cell().set(value)
+
+
+def monitor_increment_counter(metric_name):
+  if not enable_metrics:
+    return
+  else:
+    if not _METRICS_MAPPING:
+      _init()
+    metric = _METRICS_MAPPING[metric_name]
+    metric.get_cell().increase_by(1)
+
+
+def _get_metric_histogram(histogram_proto):
+  """Convert a histogram proto into a dict.
+
+  Args:
+    histogram_proto: a proto containing a Sampler metric's result histogram.
+
+  Returns:
+    A dict containing summary statistics and the raw histogram values.
+  """
+  ret = dict()
+  ret['min'] = histogram_proto.min
+  ret['max'] = histogram_proto.max
+  ret['num'] = histogram_proto.num
+  ret['sum'] = histogram_proto.sum
+
+  bucket_limits = histogram_proto.bucket_limit
+  bucket_vals = histogram_proto.bucket
+  ret['histogram'] = {}
+  # Add lower limit as 0, since all these metrics are durations
+  bucket_limits.insert(0, 0)
+  for lb, ub, val in zip(bucket_limits[:-1], bucket_limits[1:], bucket_vals):
+    ret['histogram'][(lb, ub)] = val
+  return ret
+
+
+def get_metric_summary(metric_name):
+  """Get summary for the specified metric."""
+  metric = _METRICS_MAPPING[metric_name]
+  result = metric.get_cell().value()
+  if isinstance(metric, monitoring.Sampler):
+    result = _get_metric_histogram(result)
+  return result

tensorflow/python/distribute/coordinator/metric_utils_test.py

@@ -0,0 +1,153 @@
+# 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.
+# ==============================================================================
+"""Tests for metrics collecting in coordinator."""
+
+import time
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import parameter_server_strategy_v2
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
+from tensorflow.python.distribute.coordinator import cluster_coordinator as coordinator_lib
+from tensorflow.python.distribute.coordinator import metric_utils
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.training import coordinator as thread_coordinator
+from tensorflow.python.training.server_lib import ClusterSpec
+
+
+# A function that takes long enough that inflight closures can be measured
+# as nonzero.
+@def_function.function
+def long_function():
+  x = random_ops.random_uniform((1000, 1000))
+  for _ in math_ops.range(100):
+    a = random_ops.random_uniform((1000, 1000))
+    b = random_ops.random_uniform((1000, 1000))
+    x += math_ops.matmul(a, b)
+  return x
+
+
+class MetricUtilsTest(test.TestCase):
+
+  def get_rpc_layer(self):
+    return 'grpc'
+
+  def setUp(self):
+    super().setUp()
+
+    self._cluster = multi_worker_test_base.create_multi_process_cluster(
+        num_workers=1,
+        num_ps=1,
+        rpc_layer=self.get_rpc_layer(),
+        stream_output=True,
+    )
+    self._cluster_def = self._cluster.cluster_resolver.cluster_spec().as_dict()
+    self._cluster_def['chief'] = [
+        'localhost:%d' % multi_worker_test_base.pick_unused_port()
+    ]
+    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
+        ClusterSpec(self._cluster_def), rpc_layer=self.get_rpc_layer()
+    )
+    self.strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        cluster_resolver
+    )
+    self.coordinator = coordinator_lib.ClusterCoordinator(self.strategy)
+
+    self.thread_coord = thread_coordinator.Coordinator(
+        clean_stop_exception_types=[]
+    )
+    metric_utils._init()  # Ensure metrics are reset between tests.
+
+  def tearDown(self):
+    super().tearDown()
+    self._cluster.stop()
+    self._cluster = None
+
+  def _restart(self, downtime_secs, job):
+    """Kills `job` (index: 0) and restarts it after `downtime_secs`.
+
+    Args:
+      downtime_secs: secs before restarting the job.
+      job: a string specifying the job to restart.
+    """
+    self._cluster.kill_task(job, 0)
+    time.sleep(downtime_secs)
+    self.assertFalse(context.check_alive('/job:%s/replica:0/task:0' % job))
+    self._cluster.start_task(job, 0)
+    while not context.check_alive('/job:%s/replica:0/task:0' % job):
+      time.sleep(1)
+
+  def testSimpleMetrics(self):
+    # Add a sleep to increase tracing time
+    @def_function.function
+    def func():
+      time.sleep(0.5)
+      return 3
+
+    self.assertEqual(metric_utils.get_metric_summary('queued_closures'), 0)
+    self.assertEqual(metric_utils.get_metric_summary('inflight_closures'), 0)
+
+    self.coordinator.schedule(func, args=None, kwargs=None)
+    result = self.coordinator.schedule(func, args=None, kwargs=None)
+    self.coordinator.join()
+
+    self.assertEqual(metric_utils.get_metric_summary('queued_closures'), 0)
+    self.assertEqual(metric_utils.get_metric_summary('inflight_closures'), 0)
+    # Tracing, closure execution, and remote_value fetching should be executed
+    # exactly once for running this function.
+    metric_tracing = metric_utils.get_metric_summary('function_tracing')
+    self.assertEqual(metric_tracing['num'], 1)
+    # Tracing time should be longer than the sleep time in Python function.
+    self.assertGreater(metric_tracing['sum'], 0.5)
+    metric_closure = metric_utils.get_metric_summary('closure_execution')
+    self.assertEqual(metric_closure['num'], 2)
+    metric_remote_value = metric_utils.get_metric_summary('remote_value_fetch')
+    self.assertEqual(metric_remote_value['num'], 2)
+
+    self.assertEqual(result.fetch(), 3)
+
+  def testInflightClosures(self):
+    self.coordinator.schedule(long_function)
+    self.coordinator.schedule(long_function)
+    self.assertGreater(metric_utils.get_metric_summary('queued_closures'), 0)
+
+    # inflight closures should be greater than 0 at some point
+    max_inflight = 0
+    while not self.coordinator.done():
+      with self.coordinator._cluster.closure_queue._queue_lock:
+        inflight_metric = metric_utils.get_metric_summary('inflight_closures')
+      max_inflight = max(max_inflight, inflight_metric)
+      time.sleep(0.01)
+    self.assertGreater(max_inflight, 0)
+
+  def testWorkerFailureCount(self):
+    self.coordinator.schedule(long_function)
+    self._restart(downtime_secs=2, job='worker')
+    self.coordinator.schedule(long_function)
+
+    self.coordinator.join()
+    metric_closure = metric_utils.get_metric_summary('closure_execution')
+    self.assertEqual(metric_closure['num'], 2)
+    num_failures = metric_utils.get_metric_summary('worker_failures')
+    self.assertEqual(num_failures, 1)
+    recovery_times = metric_utils.get_metric_summary('server_def_update')
+    self.assertEqual(recovery_times['num'], 1)
+
+
+if __name__ == '__main__':
+  multi_process_runner.test_main()

tensorflow/python/distribute/coordinator/remote_value.py

@@ -0,0 +1,131 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""RemoteValue interface class."""
+
+import enum
+
+from tensorflow.python.util.tf_export import tf_export
+
+
+class RemoteValueStatus(enum.Enum):
+  """The status of a `RemoteValue` object.
+
+  A `RemoteValue` object can have three states:
+    1) not ready: no value, no non-retryable error and not aborted;
+    2) aborted: i.e. the execution of function was aborted because of task
+       failure, but can be retried;
+    3) ready: i.e. has value or has non-tryable error;
+
+  The initial state of a `RemoteValue` is "not ready". When its corresponding
+  closure has
+  been executed at least once, it will become aborted or ready. The state
+  transitions are:
+    1) not ready -> 2) aborted:
+      when the corresponding closure is aborted due to worker failure, and the
+      worker failure is not immediately handled.
+    1) not ready -> 3) ready:
+      when the corresponding closure has been executed successfully.
+    2) aborted -> 3) ready:
+      when the `RemoteValue` is rebuilt by rerunning the corresponding closure
+      and the closure has been executed successfully.
+    3) ready -> 2) aborted:
+      when the corresponding closure had been executed successfully but later
+      the corresponding remote worker failed. This is currently only implemented
+      for resource `RemoteValue` like iterators.
+  """
+  NOT_READY = "NOT_READY"
+  ABORTED = "ABORTED"
+  READY = "READY"
+
+
+@tf_export("distribute.experimental.coordinator.RemoteValue",
+           "distribute.coordinator.RemoteValue", v1=[])
+class RemoteValue(object):
+  """An asynchronously available value of a scheduled function.
+
+  This class is used as the return value of
+  `tf.distribute.experimental.coordinator.ClusterCoordinator.schedule` where
+  the underlying value becomes available at a later time once the function has
+  been executed.
+
+  Using `tf.distribute.experimental.coordinator.RemoteValue` as an input to
+  a subsequent function scheduled with
+  `tf.distribute.experimental.coordinator.ClusterCoordinator.schedule` is
+  currently not supported.
+
+  Example:
+
+  ```python
+  strategy = tf.distribute.experimental.ParameterServerStrategy(
+      cluster_resolver=...)
+  coordinator = (
+      tf.distribute.experimental.coordinator.ClusterCoordinator(strategy))
+
+  with strategy.scope():
+    v1 = tf.Variable(initial_value=0.0)
+    v2 = tf.Variable(initial_value=1.0)
+
+  @tf.function
+  def worker_fn():
+    v1.assign_add(0.1)
+    v2.assign_sub(0.2)
+    return v1.read_value() / v2.read_value()
+
+  result = coordinator.schedule(worker_fn)
+  # Note that `fetch()` gives the actual result instead of a `tf.Tensor`.
+  assert result.fetch() == 0.125
+
+  for _ in range(10):
+    # `worker_fn` will be run on arbitrary workers that are available. The
+    # `result` value will be available later.
+    result = coordinator.schedule(worker_fn)
+  ```
+  """
+
+  def fetch(self):
+    """Wait for the result of `RemoteValue` and return the numpy result.
+
+    This makes the value concrete by copying the remote value to local.
+
+    Returns:
+      The numpy array structure of the actual output of the `tf.function`
+      associated with this `RemoteValue`, previously returned by a
+      `tf.distribute.experimental.coordinator.ClusterCoordinator.schedule` call.
+      This can be a single value, or a structure of values, depending on the
+      output of the `tf.function`.
+
+    Raises:
+      tf.errors.CancelledError: If the function that produces this `RemoteValue`
+        is aborted or cancelled due to failure.
+    """
+    raise NotImplementedError("Must be implemented in subclasses.")
+
+  def get(self):
+    """Wait for the result of `RemoteValue` and return the tensor result.
+
+    This makes the value concrete by copying the remote tensor to local.
+
+    Returns:
+      The actual output (in the form of `tf.Tensor`s) of the `tf.function`
+      associated with this `RemoteValue`, previously returned by a
+      `tf.distribute.experimental.coordinator.ClusterCoordinator.schedule` call.
+      This can be a single Tensor, or a structure of Tensors, depending on the
+      output of the `tf.function`.
+
+    Raises:
+      tf.errors.CancelledError: If the function that produces this `RemoteValue`
+        is aborted or cancelled due to failure.
+    """
+    raise NotImplementedError("Must be implemented in subclasses.")

tensorflow/python/distribute/coordinator/utils.py

@@ -0,0 +1,79 @@
+# 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.
+# ==============================================================================
+"""TF2 parameter server training utilities.
+
+Parameter server training in TF2 is currently under development.
+"""
+import threading
+import time
+
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import server_lib
+
+
+def start_server(cluster_resolver, protocol):
+  """Start a server and block the process from exiting."""
+  # This function is for multi-processing test or users who would like to have
+  # every job run the same binary for simplicity.
+  if not (cluster_resolver.task_type == 'worker' or
+          cluster_resolver.task_type == 'ps'):
+    raise ValueError('Unexpected task_type to start a server: {}'.format(
+        cluster_resolver.task_type))
+
+  server = server_lib.Server(
+      cluster_resolver.cluster_spec().as_cluster_def(),
+      job_name=cluster_resolver.task_type,
+      task_index=cluster_resolver.task_id,
+      protocol=protocol)
+
+  logging.info('TensorFlow server started for job %s, task %d.',
+               cluster_resolver.task_type, cluster_resolver.task_id)
+
+  # Blocking the process that starts a server from exiting.
+  server.join()
+
+
+class RepeatedTimer(object):
+  """Threaded Repeated Timer from http://shortn/_3hMZTFr1Iv."""
+
+  def __init__(self, interval, function, *args):
+    self._timer = None
+    self.interval = interval
+    self.function = function
+    self.args = args
+    self.start_time = time.time()
+    self.is_running = False
+    self.start()
+
+  def _get_duration_sec(self):
+    return int(time.time() - self.start_time)
+
+  def _run(self):
+    self.is_running = False
+    self.start()
+    self.function(*self.args)
+
+  def start(self):
+    if not self.is_running:
+      self._timer = threading.Timer(self.interval, self._run)
+      self._timer.start()
+      self.is_running = True
+
+  def stop(self):
+    duration = self._get_duration_sec()
+    self._timer.cancel()
+    self.is_running = False
+    return duration
+

tensorflow/python/distribute/coordinator/values.py

@@ -0,0 +1,386 @@
+# 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.
+# ==============================================================================
+"""Important value classes relevant to `ClusterCoordinator`.
+
+This is currently under development and the API is subject to change.
+"""
+
+import threading
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops.options import ExternalStatePolicy
+from tensorflow.python.distribute import input_lib
+from tensorflow.python.distribute.coordinator import remote_value
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as tf_function
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import type_spec as type_spec_lib
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+# TODO(yuefengz): create an implementation for resource RemoteValue which needs
+# to remember the closure object while a normal RemoteValue doesn't.
+class RemoteValueImpl(remote_value.RemoteValue):
+  """Implementation of `RemoteValue`."""
+
+  def __init__(self, closure, type_spec):  # pylint: disable=super-init-not-called
+    """Initializes a `RemoteValueImpl`.
+
+    Args:
+      closure: The closure from which the `RemoteValue` is created.
+      type_spec: The type spec for this `RemoteValue` which is used to trace
+        functions that take this `RemoteValue` as input.
+    """
+    self._closure = closure
+    self._type_spec = type_spec
+    self._values = None
+    self._has_fetched_to_local = False
+    self._has_fetched_to_local_lock = threading.Lock()
+    self._fetched_tensors = None
+    self._error = None
+    self._status_available_event = threading.Event()
+    self._status = remote_value.RemoteValueStatus.NOT_READY
+
+  def _set_aborted(self, error):
+    self._status = remote_value.RemoteValueStatus.ABORTED
+    self._values = None
+    self._error = error
+
+    # Wake up any waiting thread and clear the event.
+    self._status_available_event.set()
+
+  def _rebuild_on(self, worker):
+    self._status_available_event.clear()
+    # TODO(yuefengz): we may need to rebuild its inputs as well.
+    self._closure.execute_on(worker)
+
+  def _set_values(self, tensors):
+    self._status = remote_value.RemoteValueStatus.READY
+    self._values = tensors
+    self._error = None
+    self._status_available_event.set()
+
+  def _set_error(self, error):
+    self._status = remote_value.RemoteValueStatus.READY
+    self._values = None
+    self._error = error
+    self._status_available_event.set()
+
+  def _get_values(self):
+    self._status_available_event.wait()
+    return self._values
+
+  def _get_error(self):
+    self._status_available_event.wait()
+    return self._error
+
+  def _wait_and_maybe_error(self):
+    self._status_available_event.wait()
+    if self._status is remote_value.RemoteValueStatus.ABORTED:
+      raise errors.CancelledError(
+          None, None,
+          "The corresponding function is aborted. Please reschedule the "
+          "function.")
+    if self._error is not None:
+      raise self._error
+
+  def fetch(self):
+    # TODO(rchao): Discuss the possibility of letting users perform `numpy`
+    # themselves at API graduation.
+    return nest.map_structure(
+        lambda x: x.numpy() if hasattr(x, "numpy") else x, self.get())
+
+  def _copy_to_local(self):
+    def copy_tensor(composite_tensor_obj):
+      """Copy a remote tensor to local (coordinator)."""
+      if isinstance(composite_tensor_obj, input_lib.DistributedIterator):
+        # A DistributedIterator cannot be copied to local; users should not
+        # access that anyway.
+        return composite_tensor_obj
+
+      with ops.device("/job:%s" % context.get_server_def().job_name):
+        # Copying to local (the coordinator) with `tf.device`.
+        return array_ops.identity(composite_tensor_obj)
+
+    fetched_result = None
+    if self._values is not None:
+      # When `self._values` is `None`, it indicates the associated function
+      # does not have a return value.
+      fetched_result = nest.map_structure(copy_tensor, self._values)
+    return fetched_result
+
+  def get(self):
+    self._wait_and_maybe_error()
+
+    with self._has_fetched_to_local_lock:
+      if not self._has_fetched_to_local:
+        self._fetched_tensors = self._copy_to_local()
+        self._has_fetched_to_local = True
+
+    return self._fetched_tensors
+
+
+class RemoteVariable(RemoteValueImpl):
+  """A RemoteValue that represents a mutable per-worker variable."""
+
+  def get(self):
+    """Retrieve value with no caching to ensure we get the up-to-date value."""
+    self._wait_and_maybe_error()
+    return self._copy_to_local()
+
+
+@tf_export("distribute.experimental.coordinator.PerWorkerValues",
+           "distribute.coordinator.PerWorkerValue", v1=[])
+class PerWorkerValues(composite_tensor.CompositeTensor):
+  """A container that holds a list of values, one value per worker.
+
+  `tf.distribute.experimental.coordinator.PerWorkerValues` contains a collection
+  of values, where each of the values is located on its corresponding worker,
+  and upon being used as one of the `args` or `kwargs` of
+  `tf.distribute.experimental.coordinator.ClusterCoordinator.schedule()`, the
+  value specific to a worker will be passed into the function being executed at
+  that corresponding worker.
+
+  Currently, the only supported path to create an object of
+  `tf.distribute.experimental.coordinator.PerWorkerValues` is through calling
+  `iter` on a `ClusterCoordinator.create_per_worker_dataset`-returned
+  distributed dataset instance. The mechanism to create a custom
+  `tf.distribute.experimental.coordinator.PerWorkerValues` is not yet supported.
+  """
+
+  def __init__(self, values):
+    for v in values:
+      if not isinstance(v, remote_value.RemoteValue):
+        raise AssertionError(
+            "`PerWorkerValues` should only take `RemoteValue`s.")
+    self._values = tuple(values)
+
+  @property
+  def _type_spec(self):
+    return PerWorkerValuesTypeSpec(
+        self._values[0]._type_spec,  # pylint: disable=protected-access
+        type(self))
+
+
+class PerWorkerValuesTypeSpec(type_spec_lib.TypeSpec):
+  """TypeSpec for PerWorkerValues.
+
+  It only support tracing a function using a PerWorkerValues.
+  """
+
+  def __init__(self, value_spec, descendant_type):
+    assert value_spec
+    self._value_spec = value_spec
+    self._descendant_type = descendant_type
+
+  def _serialize(self):
+    return (self._value_spec,)
+
+  @property
+  def value_type(self):
+    return self._descendant_type
+
+  def most_specific_common_supertype(self, others):
+    raise NotImplementedError(
+        "most_specific_common_supertype is not implemented")
+
+  @property
+  def _component_specs(self):
+    return self._value_spec
+
+  def _to_components(self, value):
+    return self._value_spec
+
+  def _from_components(self, value):
+    return value
+
+
+class PerWorkerDatasetFromDatasetFunction(object):
+  """Represents worker-distributed datasets created from dataset function."""
+
+  def __init__(self, dataset_fn, coordinator):
+    """Makes an iterable from datasets created by the given function.
+
+    Args:
+      dataset_fn: A function that returns a `Dataset`.
+      coordinator: a `ClusterCoordinator` object, used to create dataset
+        resources.
+    """
+
+    def disallow_variable_creation(next_creator, **kwargs):
+      raise ValueError("Creating variables in `dataset_fn` is not allowed.")
+
+    if isinstance(dataset_fn, def_function.Function):
+      with variable_scope.variable_creator_scope(disallow_variable_creation):
+        dataset_fn = dataset_fn.get_concrete_function()
+    elif not isinstance(dataset_fn, tf_function.ConcreteFunction):
+      with variable_scope.variable_creator_scope(disallow_variable_creation):
+        dataset_fn = def_function.function(dataset_fn).get_concrete_function()
+    self._dataset_fn = dataset_fn
+    self._coordinator = coordinator
+    self._element_spec = None
+
+  def build(self):
+    """Trigger dataset creation on workers without creating an iterator.
+
+    Returns:
+      A PerWorkerValues object containing a tuple of RemoteValues, themselves
+      containing the built Dataset for each worker
+    """
+    def _create_per_worker_dataset():
+      dataset = self._dataset_fn()
+      return dataset
+
+    # pylint: disable=protected-access
+    per_worker_dataset = self._coordinator._create_per_worker_resources(
+        _create_per_worker_dataset)
+    # hack type_spec of RemoteValues
+    dataset_fn_output_type_spec = self._dataset_fn.structured_outputs._type_spec
+    for dataset_remote_value in per_worker_dataset._values:
+      dataset_remote_value._type_spec = dataset_fn_output_type_spec
+    return per_worker_dataset
+
+  def __iter__(self):
+    # We would like users to create iterators outside `tf.function`s so that we
+    # can track them.
+    if (not context.executing_eagerly() or
+        ops.get_default_graph().building_function):
+      raise RuntimeError(
+          "__iter__() is not supported inside of tf.function or in graph mode.")
+
+    def _create_per_worker_iterator():
+      dataset = self._dataset_fn()
+      return iter(dataset)
+
+    # If PerWorkerDatasetFromDatasetFunction.__iter__ is called multiple
+    # times, for the same object it should only create and register resource
+    # once. Using object id to distinguish different iterator resources.
+    per_worker_iterator = self._coordinator._create_per_worker_resources(
+        _create_per_worker_iterator)
+
+    # Setting type_spec of each RemoteValue so that functions taking these
+    # RemoteValues as inputs can be traced.
+    for iterator_remote_value in per_worker_iterator._values:
+      iterator_remote_value._type_spec = (
+          input_lib.get_iterator_spec_from_dataset(
+              self._coordinator.strategy, self._dataset_fn.structured_outputs))
+
+    return PerWorkerDistributedIterator(per_worker_iterator._values)
+
+  @property
+  def element_spec(self):
+    """The type specification of an element of this dataset.
+
+    This property is subject to change without notice.
+    """
+    if not isinstance(self._dataset_fn, tf_function.ConcreteFunction):
+      raise NotImplementedError(
+          "`element_spec` is not supported when the `dataset_fn` is not "
+          "a `ConcreteFunction`.")
+    return self._dataset_fn.structured_outputs.element_spec
+
+
+def serialize_dataset_to_graph(dataset):
+  dataset = dataset._apply_debug_options()  # pylint: disable=protected-access
+  graph_def = gen_dataset_ops.dataset_to_graph_v2(
+      dataset._variant_tensor,  # pylint: disable=protected-access
+      external_state_policy=ExternalStatePolicy.WARN.value,
+      strip_device_assignment=True)
+  return graph_def
+
+
+class _RemoteDataset(dataset_ops.DatasetSource):
+  """Creates a dataset given a graph def."""
+
+  def __init__(self, graph_def, element_spec):
+    self._elem_spec = element_spec
+    variant_tensor = ged_ops.dataset_from_graph(graph_def)
+    super(_RemoteDataset, self).__init__(variant_tensor)
+
+  @property
+  def element_spec(self):
+    return self._elem_spec
+
+
+def deserialize_dataset_from_graph(graph_def, element_spec):
+  return _RemoteDataset(graph_def, element_spec)
+
+
+class PerWorkerDatasetFromDataset(PerWorkerDatasetFromDatasetFunction):
+  """Represents worker-distributed datasets created from a dataset."""
+
+  def __init__(self, dataset, coordinator):
+    """Makes an iterable from datasets created by the given dataset.
+
+    It creates a dataset_fn which deserializes a dataset from a graph under the
+    hood.
+
+    Args:
+      dataset: A tf.data.Dataset, a DistributedDataset or a
+        DistributedDatasetsFromFunction
+      coordinator: a `ClusterCoordinator` object, used to create dataset
+        resources.
+    """
+    if isinstance(dataset, input_lib.DistributedDataset):
+      original_dataset = dataset._original_dataset
+      serialized = serialize_dataset_to_graph(original_dataset)
+
+      def dataset_fn():
+        deserialized = deserialize_dataset_from_graph(
+            serialized, original_dataset.element_spec)
+        dataset.build(dataset_to_replace=deserialized)
+        return dataset
+    elif isinstance(dataset, input_lib.DistributedDatasetsFromFunction):
+      def dataset_fn():
+        dataset.build()
+        return dataset
+    elif isinstance(dataset, dataset_ops.Dataset):
+      serialized = serialize_dataset_to_graph(dataset)
+
+      def dataset_fn():
+        return deserialize_dataset_from_graph(serialized, dataset.element_spec)
+    else:
+      raise ValueError("Unexpected dataset type!")
+
+    super(PerWorkerDatasetFromDataset, self).__init__(dataset_fn, coordinator)
+
+
+def get_per_worker_dataset(dataset_or_dataset_fn, coordinator):
+  """Returns a per-worker dataset from a dataset or a dataset function."""
+  if callable(dataset_or_dataset_fn):
+    return PerWorkerDatasetFromDatasetFunction(dataset_or_dataset_fn,
+                                               coordinator)
+  else:
+    return PerWorkerDatasetFromDataset(dataset_or_dataset_fn, coordinator)
+
+
+class PerWorkerDistributedIterator(PerWorkerValues):
+  """Distributed iterator for `ClusterCoordinator`."""
+
+  def __next__(self):
+    return self.get_next()
+
+  def get_next(self, name=None):
+    """Returns the next input from the iterator for all replicas."""
+    raise NotImplementedError("Iterating over an `AsyncDistributedIterator` "
+                              "is not supported right now.")

tensorflow/python/distribute/coordinator/watchdog.py

@@ -0,0 +1,63 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Watchdog that monitors activity of ClusterCoordinator."""
+
+import faulthandler
+import os
+import sys
+import threading
+import time
+from absl import logging
+
+
+class WatchDog(object):
+  """A class to dump stack traces if no activity happens in ClusterCoordinator."""
+
+  def __init__(self, timeout=-1, traceback_file=sys.stdout, on_triggered=None):
+    if os.environ.get("TF_CLUSTER_COORDINATOR_WATCH_DOG_TIMEOUT",
+                      "").isnumeric():
+      timeout = int(os.environ["TF_CLUSTER_COORDINATOR_WATCH_DOG_TIMEOUT"])
+    self._timeout = timeout
+    self._last_activity_time = time.time()
+    self._traceback_file = traceback_file
+    self._on_triggered = on_triggered
+    self._stopped = False
+    if timeout > 0:
+      self._watchdog_thread = threading.Thread(
+          target=self._watchdog_function, name="WatchDog", daemon=True)
+      self._watchdog_thread.start()
+
+  def stop(self):
+    self._stopped = True
+
+  def _watchdog_function(self):
+    """The watchdog thread."""
+    logging.info("Starting watchdog thread with timeout %r", self._timeout)
+    while not self._stopped:
+      time.sleep(self._timeout / 10.0)
+      current_time = time.time()
+      if current_time - self._last_activity_time >= self._timeout:
+        logging.warning(
+            "No activity for ClusterCoordinator for %r seconds. "
+            "Dumping stack traces.", self._timeout)
+        if self._on_triggered:
+          self._on_triggered()
+        faulthandler.dump_traceback(file=self._traceback_file)
+        self._traceback_file.write("==== End of stack traces ====\n")
+        self._last_activity_time = current_time
+
+  def report_closure_done(self):
+    if self._timeout > 0:
+      self._last_activity_time = time.time()

tensorflow/python/distribute/coordinator/watchdog_test.py

@@ -0,0 +1,66 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for watchdog.py."""
+
+import os
+import time
+
+from absl.testing import parameterized
+
+from tensorflow.python.distribute.coordinator import watchdog
+from tensorflow.python.eager import test
+
+
+class WatchDogTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(True, False)
+  def testWatchDogTimeout(self, use_env_var):
+    tmp_file = self.create_tempfile()
+    f = open(tmp_file, "w+")
+
+    triggerred_count = [0]
+
+    def on_triggered_fn():
+      triggerred_count[0] += 1
+
+    timeout = 3
+    if use_env_var:
+      os.environ["TF_CLUSTER_COORDINATOR_WATCH_DOG_TIMEOUT"] = str(timeout)
+      wd = watchdog.WatchDog(traceback_file=f, on_triggered=on_triggered_fn)
+    else:
+      wd = watchdog.WatchDog(
+          timeout=timeout, traceback_file=f, on_triggered=on_triggered_fn)
+    time.sleep(6)
+
+    self.assertGreaterEqual(triggerred_count[0], 1)
+    wd.report_closure_done()
+    time.sleep(1)
+    self.assertGreaterEqual(triggerred_count[0], 1)
+    time.sleep(5)
+    self.assertGreaterEqual(triggerred_count[0], 2)
+
+    wd.stop()
+    time.sleep(5)
+    last_triggered_count = triggerred_count[0]
+    time.sleep(10)
+    self.assertEqual(last_triggered_count, triggerred_count[0])
+
+    f.close()
+    with open(tmp_file) as f:
+      self.assertIn("Current thread", f.read())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/cross_device_utils.py

@@ -0,0 +1,730 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Utilities for cross_device_ops."""
+
+import copy
+import threading
+from typing import Callable, List, Optional, Union
+
+from tensorflow.python.distribute import collective_util
+from tensorflow.python.distribute import values as value_lib
+from tensorflow.python.eager import backprop_util
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import collective_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nccl_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.types import core
+
+INSTANCE_KEY_START_NUMBER = 100
+
+
+def aggregate_gradients_using_nccl(replica_grads):
+  """Aggregate gradients using nccl allreduce."""
+  agg_all_g_and_v = []
+  for single_g_and_v in zip(*replica_grads):
+    single_grads = [g for g, _ in single_g_and_v]
+    agg_grads = nccl_ops.all_sum(single_grads)
+    agg_all_g_and_v.append(
+        [(g, v) for g, (_, v) in zip(agg_grads, single_g_and_v)])
+
+  agg_all_g_and_v = list(zip(*agg_all_g_and_v))
+
+  return agg_all_g_and_v
+
+
+def aggregate_gradients_using_hierarchical_copy(avail_devices, replica_grads):
+  """Aggregate gradients using hierarchical copies.
+
+  Args:
+    avail_devices: available GPU devices.
+    replica_grads: List of lists of (gradient, variable) tuples. The outer list
+      is over replicas. The inner list is over individual gradients.
+
+  Returns:
+    The list of (aggregated_gradient, variable), where the gradient has been
+      summed across all replicas and the variable is chosen from the first
+      replica.
+  """
+  # This only works for DGX-1 type of machine topology
+  # Device peer to peer matrix
+  # DMA: 0 1 2 3 4 5 6 7
+  # 0:   Y Y Y Y Y N N N
+  # 1:   Y Y Y Y N Y N N
+  # 2:   Y Y Y Y N N Y N
+  # 3:   Y Y Y Y N N N Y
+  # 4:   Y N N N Y Y Y Y
+  # 5:   N Y N N Y Y Y Y
+  # 6:   N N Y N Y Y Y Y
+  # 7:   N N N Y Y Y Y Y
+  agg_grads = []
+  num_devices = len(avail_devices)
+  # In the special case of DGX-1 machine topology, the two groups have equal
+  # size.
+  group_size = num_devices // 2
+  for i, single_grads in enumerate(zip(*replica_grads)):
+    group_0_main_device = i % num_devices
+    group_1_main_device = (group_0_main_device + group_size) % num_devices
+    if group_0_main_device < group_size:
+      group_0_begin = 0
+      group_1_begin = group_size
+    else:
+      group_0_begin = group_size
+      group_1_begin = 0
+
+    # Aggregate the first group.
+    group_0_device_grads = single_grads[group_0_begin:
+                                        group_0_begin + group_size]
+    with ops.device(avail_devices[group_0_main_device]):
+      group_0_agg_grads, _ = aggregate_single_gradient_using_copy(
+          group_0_device_grads, False, False)
+
+    # Aggregate the second group.
+    group_1_device_grads = single_grads[group_1_begin:
+                                        group_1_begin + group_size]
+    with ops.device(avail_devices[group_1_main_device]):
+      group_1_agg_grads, _ = aggregate_single_gradient_using_copy(
+          group_1_device_grads, False, False)
+
+    # Aggregate between the groups.
+    with ops.device(avail_devices[group_0_main_device]):
+      (agg_total_grads, _), _ = aggregate_single_gradient_using_copy(
+          [group_0_agg_grads, group_1_agg_grads], False, False)
+
+    # Broadcast the result back into the root of each group.
+    with ops.device(avail_devices[group_0_main_device]):
+      group_0_agg_grads_bcast = array_ops.identity(agg_total_grads)
+    with ops.device(avail_devices[group_1_main_device]):
+      group_1_agg_grads_bcast = array_ops.identity(agg_total_grads)
+
+    agg_grads_bcast = []
+    for j in range(len(single_grads)):
+      with ops.device(avail_devices[j]):
+        # Broadcast the result back to each member in the group from the root.
+        if (group_0_main_device < group_size) == (j < group_size):
+          src_device_grad = group_0_agg_grads_bcast
+        else:
+          src_device_grad = group_1_agg_grads_bcast
+        agg_grads_bcast.append(array_ops.identity(src_device_grad))
+
+    agg_grads.append(
+        [(g, v) for g, (_, v) in zip(agg_grads_bcast, single_grads)])
+
+  agg_grads = list(zip(*agg_grads))
+
+  return agg_grads
+
+
+def aggregate_single_gradient_using_copy(grad_and_vars, use_mean,
+                                         check_inf_nan):
+  """Calculate the average gradient for a shared variable across all replicas.
+
+  Note that this function provides a synchronization point across all replicas.
+
+  Args:
+    grad_and_vars: A list or tuple of (gradient, variable) tuples. Each
+      (gradient, variable) pair within the outer list represents the gradient
+      of the variable calculated for a single replica, and the number of pairs
+      equals the number of replicas.
+    use_mean: if True, mean is taken, else sum of gradients is taken.
+    check_inf_nan: check grads for nans and infs.
+
+  Returns:
+    The tuple ([(average_gradient, variable),], has_nan_or_inf) where the
+      gradient has been averaged across all replicas. The variable is chosen
+      from the first replica. The has_nan_or_inf indicates the grads has nan or
+      inf.
+  """
+  grads = [g for g, _ in grad_and_vars]
+  grad = math_ops.add_n(grads)
+
+  if use_mean and len(grads) > 1:
+    grad = array_ops.multiply(grad, 1.0 / len(grads))
+
+  v = grad_and_vars[0][1]
+  if check_inf_nan:
+    has_nan_or_inf = array_ops.logical_not(
+        array_ops.reduce_all(array_ops.is_finite(grads)))
+    return (grad, v), has_nan_or_inf
+  else:
+    return (grad, v), None
+
+
+# TODO(yuefengz): use random key starts to avoid reusing keys?
+class CollectiveKeys(object):
+  """Class that manages collective keys.
+
+  We need to manage three different keys for collective:
+
+  *Group key*: an integer key to identify the set of cooperative devices.
+  Collective ops work under the same set of devices must using the same group
+  key.
+
+  *Instance key*: an integer key to identify the set of same counterpart of
+  tensors on different devices in a device group that need to be all-reduced.
+
+  This class is thread safe.
+  """
+
+  def __init__(self, group_key_start=1):
+    """Initializes the object.
+
+    Args:
+      group_key_start: the starting integer of group key.
+    """
+    self._group_key = group_key_start
+    self._instance_key_table = {}
+    self._lock = threading.Lock()
+    self._known_groups = {}
+
+  def get_group_key(self, devices):
+    """Returns a group key for the list of local devices.
+
+    The same group key is returned if the list of local devices is the same.
+
+    Args:
+      devices: a list of local canonical device strings in a collective group.
+
+    Returns:
+      a group key.
+    """
+    with self._lock:
+      devices_key = ','.join(devices)
+      if devices_key not in self._known_groups:
+        self._known_groups[devices_key] = self._get_new_group_key(devices)
+      return self._known_groups[devices_key]
+
+  def _get_new_group_key(self, devices):
+    """Returns a new group key.
+
+    The caller should store and reuse the same group key for the same set of
+    devices. Calling this method always returns a new group key.
+
+    This method is not thread-safe.
+
+    Args:
+      devices: a list of canonical device strings in a collective group.
+
+    Returns:
+      a new group key.
+    """
+    new_key = self._group_key
+    self._group_key += 1
+    self._instance_key_table[new_key] = {}
+    for device in devices:
+      self._instance_key_table[new_key][device] = INSTANCE_KEY_START_NUMBER
+    return new_key
+
+  def get_instance_key(self, group_key, device):
+    """Returns a new instance key for use in defining a collective op.
+
+    You should call this once per each collective op of a collective instance.
+
+    Args:
+      group_key: the group key returned by get_group_key(). You should not
+        assign the group key yourself.
+      device: a canonical device string. It should be the device this collective
+        op is on.
+
+    Returns:
+      a new instance key.
+
+    Raises:
+      ValueError: when the group key is invalid or the device is not in the
+      group.
+    """
+    with self._lock:
+      group = self._instance_key_table.get(group_key, None)
+      if group is None:
+        raise ValueError(f'Group {group_key} is not found.')
+      if device not in group:
+        raise ValueError(f'Device {device} is not present in group {group_key}')
+      v = group[device]
+      group[device] += 1
+      return v
+
+  def __deepcopy__(self, memo):
+    # distribute_coordinator deep-copies the strategy object, so
+    # CollectiveKeys needs to support deep copy as well.
+    copied = CollectiveKeys()
+    copied._group_key = self._group_key
+    copied._instance_key_table = copy.deepcopy(self._instance_key_table, memo)
+    return copied
+
+
+class CollectiveReplicaLauncher(object):
+  """Launch collectives on one replica."""
+
+  _prefer_unique_instance_key = True
+  _prefer_ordering_token = True
+
+  def __init__(self, group_key: int, group_size: int,
+               collective_keys: CollectiveKeys, device: str,
+               options: collective_util.Options):
+    self._group_key = group_key
+    self._group_size = group_size
+    self._collective_keys = collective_keys
+    self._device = device
+    self._options = options
+    if self._use_ordering_token():
+      with ops.init_scope(), ops.device(device):
+        self._ordering_token = resource_variable_ops.ResourceVariable(0.)
+    else:
+      self._ordering_token = None
+
+  def _control_input(self, control_input: Union[core.TensorLike,
+                                                ops.Operation]):
+    if control_input is not None and not self._use_ordering_token():
+      return ops.control_dependencies([control_input])
+    return ops.NullContextmanager()
+
+  def _use_unique_instance_key(self):
+    if not ops.executing_eagerly_outside_functions():
+      return False
+    return CollectiveReplicaLauncher._prefer_unique_instance_key
+
+  def _use_ordering_token(self):
+    # We rely on auto control dep to insert control edges between NCCL calls,
+    # but for tf1 graph mode auto control dep is not used.
+    if not ops.executing_eagerly_outside_functions():
+      return False
+    return CollectiveReplicaLauncher._prefer_ordering_token
+
+  def _next_instance_key(self):
+    """Returns the next instance key."""
+    if self._use_unique_instance_key():
+      # Assigning instance keys at function building time have issues since
+      # different workers may retrace the function at different times. With
+      # collective V2 we can use capture_call_time_value to use a placeholder as
+      # the instance key and feed it at function call time. In this way we also
+      # don't reuse instance keys, which allows for per-instance cancellation.
+      graph = ops.get_default_graph()
+      # Control flow ops don't work with capture_call_time_value, so we put the
+      # capture in the function graph of that control flow op.
+      while getattr(graph, 'is_control_flow_graph', False):
+        graph = graph.outer_graph
+      if not context.executing_eagerly() and graph.building_function:
+        with graph.as_default():
+          # Capture self._next_instance_key so that when building a function
+          # that calls another tf.function, the instance key assignment is
+          # further delayed until we actually call the function in eager. Note
+          # that capture_call_time_value doesn't automatically propagate the
+          # deferred capture to the outer function.
+          return graph.capture_call_time_value(
+              self._next_instance_key, tensor_spec.TensorSpec([], dtypes.int32))
+      else:
+        instance_key = self._collective_keys.get_instance_key(
+            self._group_key, self._device)
+        with ops.device('CPU:0'):
+          return ops.convert_to_tensor(instance_key, dtype=dtypes.int32)
+    else:
+      return self._collective_keys.get_instance_key(self._group_key,
+                                                    self._device)
+
+  def _get_ordering_token(self):
+    if self._use_ordering_token():
+      return self._ordering_token.handle  # pytype: disable=attribute-error
+
+  def can_order_nccl(self):
+    """Whether this launcher can order NCCL operations."""
+    return self._use_ordering_token()
+
+  def all_reduce(
+      self,
+      input_tensor: core.TensorLike,
+      control_input: Optional[Union[core.TensorLike, ops.Operation]] = None,
+      options: Optional[collective_util.Options] = None) -> core.Tensor:
+    """All-reduce a dense tensor.
+
+    Args:
+      input_tensor: a dense tensor. It must have the same shape on all replicas.
+      control_input: if not None, add control edges between control_input and
+        the all-reduce.
+      options: an optional tf.distribute.experimental.CommunicationOptions. If
+        provided, it overrides the default options.
+
+    Returns:
+      The reduced tensor.
+    """
+    instance_key = self._next_instance_key()
+    options = self._options.merge(options)
+    ordering_token = self._get_ordering_token()
+    with ops.device(self._device), \
+         self._control_input(control_input):
+      return collective_ops.all_reduce_v2(
+          input_tensor,
+          self._group_size,
+          self._group_key,
+          instance_key,
+          communication_hint=options.implementation.value,
+          timeout=options.timeout_seconds,
+          ordering_token=ordering_token)
+
+  def _all_gather(self, input_tensor: core.TensorLike,
+                  options: Optional[collective_util.Options]) -> core.Tensor:
+    """All-gather a dense tensor.
+
+    Args:
+      input_tensor: a dense tensor. It must have the same shape on all replicas.
+      options: an optional tf.distribute.experimental.CommunicationOptions. If
+        provided, it overrides the default options.
+
+    Returns:
+      The reduced tensor.
+    """
+    instance_key = self._next_instance_key()
+    options = self._options.merge(options)
+    ordering_token = self._get_ordering_token()
+    with ops.device(self._device):
+      return collective_ops.all_gather_v2(
+          input_tensor,
+          self._group_size,
+          self._group_key,
+          instance_key,
+          communication_hint=options.implementation.value,
+          timeout=options.timeout_seconds,
+          ordering_token=ordering_token)
+
+  def batch_all_reduce(
+      self,
+      input_tensor_packs: List[List[core.TensorLike]],
+      options: Optional[collective_util.Options] = None) -> core.Tensor:
+    """Batch all-reduce dense tensors.
+
+    This takes a list of batches of tensors. Using multiple batches have the
+    benefit that it doesn't need to wait for all inputs to be ready to start the
+    all-reduce.
+
+    Args:
+      input_tensor_packs: a list of lists of dense tensors.
+      options: an optional tf.distribute.experimental.CommunicationOptions. If
+        provided, it overrides the default options.
+
+    Returns:
+      A flat list of reduced tensors.
+    """
+    options = self._options.merge(options)
+    outputs = []
+    for pack in input_tensor_packs:
+      if context.executing_eagerly():
+        # We don't batch in eager as it sometimes makes the performance worse
+        # due the concat/split ops.
+        for input_tensor in pack:
+          outputs.append(self.all_reduce(input_tensor, None, options))
+      else:
+        # TODO(b/169168846): inserts a parallel all_gather to verify packings
+        # are the same on each replica.
+        with ops.device(self._device):
+          flat_tensors = [array_ops.reshape(t, [-1]) for t in pack]
+          shapes = [array_ops.shape(t) for t in pack]
+          if (options.implementation
+              == collective_util.CommunicationImplementation.NCCL and outputs):
+            control_input = outputs[-1]
+          else:
+            control_input = None
+          reduced = self.all_reduce(
+              array_ops.concat(flat_tensors, axis=0), control_input, options)
+          num_elements = [math_ops.reduce_prod(s) for s in shapes]
+          flat_outputs = array_ops.split(reduced, num_elements, axis=0)
+          for shape, flat_output in zip(shapes, flat_outputs):
+            outputs.append(array_ops.reshape(flat_output, shape))
+
+    return outputs
+
+  def all_gather(
+      self,
+      input_tensor: core.TensorLike,
+      axis: core.TensorLike,
+      options: Optional[collective_util.Options] = None) -> core.Tensor:
+    """All-gather a dense tensor.
+
+    This method must be called inside a tf.function.
+
+    Args:
+      input_tensor: a dense tensor. It must have the same rank on all replicas,
+        and dimensions other than `axis` need to be the same as well.
+      axis: 0-D int32 Tensor. Dimension along which to gather. Must be in the
+        range [0, rank(value)).
+      options: an optional tf.distribute.experimental.CommunicationOptions. If
+        provided, it overrides the default options.
+
+    Returns:
+      The gathered Tensor.
+
+    Raises:
+      RuntimeError: if called in eager mode.
+    """
+    if context.executing_eagerly():
+      raise RuntimeError('all_gather is not supported in eager mode.')
+
+    with ops.device(self._device), \
+         ops.control_dependencies([array_ops.identity(input_tensor)]):
+      # 1. Transpose
+      # E.g. Given an input_tensor with shape [2,2,5,1] and axis to gather is 3,
+      # we use perm_pre=[3 0 1 2] to reshape it to [1,2,2,5], which
+      # brings the 3rd dim first; afterwards we use perm_after=[1,2,3,0] to
+      # place it back.
+      perm_pre = array_ops.concat(
+          ([axis], math_ops.range(axis),
+           math_ops.range(axis + 1, array_ops.rank(input_tensor))),
+          axis=0)
+      input_tensor_t = array_ops.transpose(input_tensor, perm=perm_pre)
+      # 2. Pad
+      gathered_shape = self._all_gather(
+          array_ops.expand_dims_v2(array_ops.shape_v2(input_tensor_t), axis=0),
+          options)
+      first_dims = gathered_shape[:, 0]
+      full_axis_dim = math_ops.reduce_max(first_dims)
+      padded_input_tensor = _pad_util(input_tensor_t, full_axis_dim)
+
+      # 3. Gather
+      gather_padded_out_tensor = self._all_gather(padded_input_tensor, options)
+      # 4. Unpad
+      split_tensors = []
+      for i in range(self._group_size):
+        start_pos = i * full_axis_dim
+        split_tensors.append(gather_padded_out_tensor[start_pos:start_pos +
+                                                      first_dims[i]])
+      out_tensor_t = array_ops.concat(split_tensors, 0)
+
+      # 5. Transpose back
+      perm_after = array_ops.concat(
+          (math_ops.range(1, axis + 1), [0],
+           math_ops.range(axis + 1, array_ops.rank(input_tensor_t))),
+          axis=0)
+      return array_ops.transpose(out_tensor_t, perm=perm_after)
+
+  def all_reduce_indexed_slices(
+      self,
+      input_slices: indexed_slices.IndexedSlices,
+      options: Optional[collective_util.Options] = None
+  ) -> indexed_slices.IndexedSlices:
+    """All-reduce an IndexedSlices.
+
+    This method can be called outside  tf.function.
+
+    Args:
+      input_slices: an IndexedSlices.
+      options: an optional tf.distribute.experimental.CommunicationOptions. If
+        provided, it overrides the default options.
+
+    Returns:
+      The reduced IndexedSlices.
+    """
+
+    # Current CollectiveAllGather implementations require input IndexedSlices to
+    # have consistent length across the board, we handle the reduction of
+    # IndexedSlices as follows:
+    #   1. Gather the lengths of IndexedSlices from all participants.
+    #   2. If they have consistent length, apply all_gather.
+    #   3. Otherwise pad IndexedSlices to be the same length across all
+    #      participants and apply_gather.
+    options = self._options.merge(options)
+    with ops.device(self._device):
+
+      def all_gather_indexed_slices(
+          all_gather_fn: Callable[
+              [core.TensorLike, Optional[collective_util.Options]], core.Tensor]
+      ) -> indexed_slices.IndexedSlices:
+        """Use all_gather_fn to aggregate `IndexedSlices`."""
+        all_values = all_gather_fn(input_slices.values, options)
+        # Add control dependency to order the all-gather.
+        if (options.implementation ==
+            collective_util.CommunicationImplementation.NCCL):
+          control = [all_values]
+        else:
+          control = []
+        with ops.control_dependencies(control):
+          all_indices = all_gather_fn(input_slices.indices, options)
+        return indexed_slices.IndexedSlices(
+            values=all_values,
+            indices=all_indices,
+            dense_shape=input_slices.dense_shape)
+
+      length = array_ops.shape(input_slices.indices)
+      all_lengths = self._all_gather(length, options)
+
+      def all_gather_with_padding(
+          input_tensor: core.TensorLike,
+          options: Optional[collective_util.Options]) -> core.Tensor:
+        """all_gather tensors of different sizes using padding."""
+        max_length = math_ops.reduce_max(all_lengths)
+        padded_tensor = _pad_util(input_tensor, max_length)
+        all_padded_tensors = self._all_gather(padded_tensor, options)
+        split_tensors = []
+        for i in range(self._group_size):
+          start_pos = i * max_length
+          split_tensors.append(all_padded_tensors[start_pos:start_pos +
+                                                  all_lengths[i]])
+        return array_ops.concat(split_tensors, 0)
+
+      return cond.cond(
+          math_ops.equal(
+              math_ops.reduce_max(all_lengths),
+              math_ops.reduce_min(all_lengths)),
+          lambda: all_gather_indexed_slices(self._all_gather),
+          lambda: all_gather_indexed_slices(all_gather_with_padding))
+
+
+def aggregate_tensors_or_indexed_slices(values, accumulation_fn=math_ops.add_n):
+  """Aggregate tensors using `accumulation_fn` and IndexedSlices via concat."""
+  if any(isinstance(v, indexed_slices.IndexedSlices) for v in values):
+    return backprop_util.AggregateIndexedSlicesGradients(values)
+  else:
+    return accumulation_fn(values)
+
+
+def divide_by_n_tensors_or_indexed_slices(value, n):
+  if isinstance(value, indexed_slices.IndexedSlices):
+    value = backprop_util.FlattenNestedIndexedSlices(value)
+    return indexed_slices.IndexedSlices(value.values / n, value.indices,
+                                        value.dense_shape)
+  else:
+    return value / n
+
+
+def copy_tensor_or_indexed_slices_to_device(value, device):
+  """Copies a tensor or IndexedSlices to a device."""
+  with ops.device(device):
+    if isinstance(value, indexed_slices.IndexedSlices):
+      copied_values = array_ops.identity(value.values)
+      copied_indices = array_ops.identity(value.indices)
+      if value.dense_shape is not None:
+        copied_shape = array_ops.identity(value.dense_shape)
+      else:
+        copied_shape = None
+      result = indexed_slices.IndexedSlices(copied_values, copied_indices,
+                                            copied_shape)
+    else:
+      result = array_ops.identity(value)
+  return result
+
+
+def is_indexed_slices(value):
+  if isinstance(value, indexed_slices.IndexedSlices):
+    return True
+  if isinstance(value, value_lib.DistributedValues):
+    return all(
+        isinstance(v, indexed_slices.IndexedSlices) for v in value.values)
+  return False
+
+
+def split_by_sparsity(values):
+  """Split values into dense and sparse values.
+
+  Args:
+    values: a list of tensors or `PerReplica`s.
+
+  Returns:
+    Four lists:
+      a list of dense values, a list of their indices in `values` and
+      a list of sparse values, a list of their indices in `values`.
+  """
+  dense_values = []
+  dense_indices = []
+  sparse_values = []
+  sparse_indices = []
+  for i, v in enumerate(values):
+    if is_indexed_slices(v):
+      sparse_values.append(v)
+      sparse_indices.append(i)
+    else:
+      dense_values.append(v)
+      dense_indices.append(i)
+  return dense_values, dense_indices, sparse_values, sparse_indices
+
+
+def stitch_values(values_and_indices_list):
+  """Stitch values together according to their indices.
+
+  Args:
+    values_and_indices_list: a list of tuples of values and indices indicating
+      the values and positions in the returned list.
+
+  Returns:
+    a stitched list of values.
+  """
+  length = 0
+  for values_and_indices in values_and_indices_list:
+    length += len(values_and_indices[0])
+
+  result = [None] * length
+  for values_and_indices in values_and_indices_list:
+    if values_and_indices and values_and_indices[0]:
+      for v, i in zip(*values_and_indices):
+        assert result[i] is None
+        result[i] = v
+  return result
+
+
+def group_by_size(input_tensors, bytes_per_pack):
+  """Groups `input_tensors` into chunks of `bytes_per_pack`.
+
+  The method preserves the original order of `input_tensors`. The grouping is
+  best effort, each pack could have more or less bytes than `bytes_per_pack`.
+  It only groups values with known shape.
+
+  Args:
+    input_tensors: a list of Tensor.
+    bytes_per_pack: an integer.
+
+  Returns:
+    A list of packs of Tensor. All values are grouped into one pack if
+    `bytes_per_pack` is zero or any of the value has unknown shape.
+  """
+
+  if bytes_per_pack == 0:
+    return [input_tensors]
+  packs = []
+  last_pack_size = 0
+  for value in input_tensors:
+    num_elements = value.shape.num_elements()
+    if num_elements is None:
+      # Can't pack values with unknown shape.
+      logging.warning(
+          'not packing values due to the unknown or inconsistent shape of %s',
+          value)
+      return [input_tensors]
+    size = num_elements * value.dtype.size
+    # Try to keep each pack as close to bytes_per_pack as possible, while each
+    # pack is at least bytes_per_pack large. I.E. we err on the side of having
+    # few but large packs.
+    if not packs or last_pack_size > bytes_per_pack:
+      packs.append([])
+      last_pack_size = 0
+    packs[-1].append(value)
+    last_pack_size += size
+  return packs
+
+
+def _pad_util(input_tensor, full_axis_dim):
+  """Pad the `input_tensor`'s first dimension to be `full_axis_dim`."""
+  missing_axis_dim = full_axis_dim - array_ops.shape_v2(input_tensor)[0]
+  tensor_rank = array_ops.rank(input_tensor)
+  paddings_axis = [[0, missing_axis_dim]]
+  paddings = array_ops.concat([
+      paddings_axis,
+      array_ops.zeros(shape=(tensor_rank - 1, 2), dtype=dtypes.int32)
+  ],
+                              axis=0)
+  padded_input_tensor = array_ops.pad(input_tensor, paddings)
+  return padded_input_tensor

tensorflow/python/distribute/cross_device_utils_test.py

@@ -0,0 +1,179 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for cross_device_utils."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import cross_device_utils
+from tensorflow.python.distribute import device_util
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+
+
+class IndexedSlicesUtilsTest(test.TestCase, parameterized.TestCase):
+
+  def _assert_values_equal(self, left, right):
+    self.assertAllEqual(
+        self.evaluate(ops.convert_to_tensor(left)),
+        self.evaluate(ops.convert_to_tensor(right)))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testAggregateTensors(self):
+    t0 = constant_op.constant([[1., 2.], [0, 0], [3., 4.]])
+    t1 = constant_op.constant([[0., 0.], [5, 6], [7., 8.]])
+    total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]])
+    result = cross_device_utils.aggregate_tensors_or_indexed_slices([t0, t1])
+    self._assert_values_equal(total, result)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testAggregateIndexedSlices(self):
+    t0 = math_ops._as_indexed_slices(
+        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
+    t1 = math_ops._as_indexed_slices(
+        constant_op.constant([[0., 0.], [5, 6], [7., 8.]]))
+    total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]])
+    result = cross_device_utils.aggregate_tensors_or_indexed_slices([t0, t1])
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
+    self._assert_values_equal(total, result)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDivideTensor(self):
+    t = constant_op.constant([[1., 2.], [0, 0], [3., 4.]])
+    n = 2
+    expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]])
+    result = cross_device_utils.divide_by_n_tensors_or_indexed_slices(t, n)
+    self._assert_values_equal(expected, result)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDivideIndexedSlices(self):
+    t = math_ops._as_indexed_slices(
+        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
+    n = 2
+    expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]])
+    result = cross_device_utils.divide_by_n_tensors_or_indexed_slices(t, n)
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
+    self._assert_values_equal(expected, result)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testIsIndexedSlices(self):
+    t = math_ops._as_indexed_slices(
+        constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
+    self.assertTrue(cross_device_utils.is_indexed_slices(t))
+
+  @combinations.generate(combinations.combine(
+      mode=["graph", "eager"],
+      required_gpus=1))
+  def testCopyTensor(self):
+    with ops.device("/cpu:0"):
+      t = constant_op.constant([[1., 2.], [0, 0], [3., 4.]])
+    destination = "/gpu:0"
+    result = cross_device_utils.copy_tensor_or_indexed_slices_to_device(
+        t, destination)
+
+    self._assert_values_equal(t, result)
+    self.assertEqual(device_util.resolve(destination),
+                     device_util.resolve(result.device))
+
+  @combinations.generate(combinations.combine(
+      mode=["graph", "eager"],
+      required_gpus=1))
+  def testCopyIndexedSlices(self):
+    with ops.device("/cpu:0"):
+      t = math_ops._as_indexed_slices(
+          constant_op.constant([[1., 2.], [0, 0], [3., 4.]]))
+    destination = "/gpu:0"
+    result = cross_device_utils.copy_tensor_or_indexed_slices_to_device(
+        t, destination)
+
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
+    self._assert_values_equal(t, result)
+    self.assertEqual(
+        device_util.resolve(destination), device_util.resolve(result.device))
+
+  @combinations.generate(
+      combinations.combine(mode=["graph", "eager"], required_gpus=1))
+  def testCopyIndexedSlicesNoDenseShape(self):
+    with ops.device("/cpu:0"):
+      t = indexed_slices.IndexedSlices(
+          indices=array_ops.identity([0]), values=array_ops.identity([1.]))
+    destination = "/gpu:0"
+    result = cross_device_utils.copy_tensor_or_indexed_slices_to_device(
+        t, destination)
+
+    self.assertIsInstance(result, indexed_slices.IndexedSlices)
+    self.assertAllEqual(t.indices, result.indices)
+    self.assertAllEqual(t.values, result.values)
+    self.assertEqual(
+        device_util.resolve(destination), device_util.resolve(result.device))
+
+
+class GroupBySizeTest(test.TestCase):
+
+  def testPreferLargerPack(self):
+    # Each packs except the last one should be equal or larger than
+    # bytes_per_pack.
+    values = [
+        # size = 2 * 4 * 4 * 4 = 128
+        array_ops.ones([2, 4, 4], dtype=dtypes.float32),
+        # size = 8 * 4 = 32
+        array_ops.ones([8], dtype=dtypes.int32),
+        # size = 10 * 10 * 8 = 800
+        array_ops.ones([10, 10], dtype=dtypes.int64),
+        # size = 1 * 4 = 4
+        array_ops.ones([1], dtype=dtypes.int32),
+    ]
+    packs = cross_device_utils.group_by_size(values, bytes_per_pack=200)
+    self.assertLen(packs, 2)
+    self.assertLen(packs[0], 3)
+    self.assertEqual(packs[0][0].shape, [2, 4, 4])
+    self.assertEqual(packs[0][1].shape, [8])
+    self.assertEqual(packs[0][2].shape, [10, 10])
+    self.assertLen(packs[1], 1)
+    self.assertEqual(packs[1][0].shape, [1])
+
+  def testZeroBytesPerPack(self):
+    values = [
+        array_ops.ones([1], dtype=dtypes.float32),
+        array_ops.ones([2], dtype=dtypes.float32),
+    ]
+    packs = cross_device_utils.group_by_size(values, bytes_per_pack=0)
+    self.assertLen(packs, 1)
+    self.assertLen(packs[0], 2)
+    self.assertEqual(packs[0][0].shape, [1])
+    self.assertEqual(packs[0][1].shape, [2])
+
+  def testUnknownShape(self):
+    def create_placeholder(shape, dtype):
+      with ops.Graph().as_default():
+        return array_ops.placeholder(dtype=dtype, shape=shape)
+
+    values = [
+        array_ops.ones([10, 10], dtype=dtypes.float32),
+        create_placeholder([None, 10], dtype=dtypes.float32),
+    ]
+    packs = cross_device_utils.group_by_size(values, bytes_per_pack=1)
+    self.assertLen(packs, 1)
+    self.assertEqual(packs[0], values)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/custom_training_loop_gradient_test.py

@@ -0,0 +1,147 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for custom training loops."""
+
+from absl.testing import parameterized
+
+from tensorflow.python import tf2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+
+
+def get_dataset_from_tensor_slices(inp_array):
+  dataset = dataset_ops.DatasetV2.from_tensor_slices(inp_array)
+  # TODO(b/138326910): Remove Dataset V1 version once bug resolved.
+  if not tf2.enabled():
+    dataset = dataset_ops.Dataset.from_tensor_slices(inp_array)
+  return dataset
+
+
+class AssertFlattenedMixin(object):
+  """Mixin for specialized asserts."""
+
+  def assert_equal_flattened(self, expected_results, actual_results):
+    """Asserts that flattened results are equal.
+
+    Due to the number of replicas in the strategy, the output may have a
+    different structure and needs to be flattened for comparison.
+
+    Args:
+      expected_results: The results expected as a result of a computation.
+      actual_results: The actual results of a computation.
+    """
+    self.assertEqual(len(expected_results), len(actual_results))
+
+    for i, expected_result in enumerate(expected_results):
+      final_result = []
+      actual_result = actual_results[i]
+      for val in actual_result:
+        final_result.extend(val.numpy())
+      self.assertAllEqual(expected_result, final_result)
+
+
+class GradientTapeTest(test.TestCase, parameterized.TestCase,
+                       AssertFlattenedMixin):
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=strategy_combinations.all_strategies,
+          mode=["eager"]
+      ))
+  def testStepInFunctionGradient(self, distribution):
+    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
+
+    @def_function.function
+    def train_step(x):
+      def computation(x):
+        return math_ops.square(x)
+      with backprop.GradientTape() as tape:
+        tape.watch(x)  # Manually watch non-variable tensors.
+        y = computation(x)
+      grads = tape.gradient(y, x)
+      return grads
+
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    results = []
+    for x in dist_dataset:
+      output = distribution.experimental_local_results(
+          distribution.run(train_step, args=(x,)))
+      results.append(output)
+    self.assert_equal_flattened([[10., 12.], [14., 16.]], results)
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=strategy_combinations.all_strategies,
+          mode=["eager"]
+      ))
+  def testRunInFunctionGradient(self, distribution):
+    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
+
+    @def_function.function
+    def run(x):
+      def train_step(x):
+        def computation(x):
+          return math_ops.square(x)
+        with backprop.GradientTape() as tape:
+          tape.watch(x)  # Manually watch non-variable tensors.
+          y = computation(x)
+        grads = tape.gradient(y, x)
+        return grads
+      return distribution.experimental_local_results(
+          distribution.run(train_step, args=(x,)))
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    results = []
+    for x in dist_dataset:
+      output = run(x)
+      results.append(output)
+    self.assert_equal_flattened([[10., 12.], [14., 16.]], results)
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=strategy_combinations.all_strategies,
+          mode=["eager"],
+          model_in_tf_function=[True, False]
+      ))
+  def testNestedFunction(self, distribution, model_in_tf_function):
+    def model(x):
+      return x * x
+
+    if model_in_tf_function:
+      model = def_function.function(model)
+
+    with distribution.scope():
+      x = variables.Variable(1.0)
+
+      @def_function.function
+      def train_step():
+        def replica_step():
+          with backprop.GradientTape() as tape:
+            y = model(x)
+          return tape.gradient(y, x)
+        return distribution.run(replica_step)
+
+      grads = distribution.experimental_local_results(train_step())
+      self.assertLen(grads, distribution.num_replicas_in_sync)
+      self.assertTrue(all(g is not None for g in grads))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/device_util.py

@@ -0,0 +1,163 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Device-related support functions."""
+
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import config
+from tensorflow.python.framework import device as tf_device
+from tensorflow.python.framework import ops
+
+
+def canonicalize(d, default=None):
+  """Canonicalize device string.
+
+  If d has missing components, the rest would be deduced from the `default`
+  argument or from '/replica:0/task:0/device:CPU:0'. For example:
+    If d = '/cpu:0', default='/job:worker/task:1', it returns
+      '/job:worker/replica:0/task:1/device:CPU:0'.
+    If d = '/cpu:0', default='/job:worker', it returns
+      '/job:worker/replica:0/task:0/device:CPU:0'.
+    If d = '/gpu:0', default=None, it returns
+      '/replica:0/task:0/device:GPU:0'.
+
+  Note: This uses "job:localhost" as the default if executing eagerly.
+
+  Args:
+    d: a device string or tf.config.LogicalDevice
+    default: a string for default device if d doesn't have all components.
+
+  Returns:
+    a canonicalized device string.
+  """
+  if isinstance(d, context.LogicalDevice):
+    d = tf_device.DeviceSpec.from_string(d.name)
+  else:
+    d = tf_device.DeviceSpec.from_string(d)
+
+  assert d.device_type is None or d.device_type == d.device_type.upper(), (
+      "Device type '%s' must be all-caps." % (d.device_type,))
+  # Fill in missing device fields using defaults.
+  result = tf_device.DeviceSpec(
+      replica=0, task=0, device_type="CPU", device_index=0)
+  if ops.executing_eagerly_outside_functions():
+    # Try to deduce job, replica and task in case it's in a multi worker setup.
+    # TODO(b/151452748): Using list_logical_devices is not always safe since it
+    # may return remote devices as well, but we're already doing this elsewhere.
+    host_cpu = tf_device.DeviceSpec.from_string(
+        config.list_logical_devices("CPU")[0].name)
+    if host_cpu.job:
+      result = result.make_merged_spec(host_cpu)
+    else:
+      # The default job is localhost if eager execution is enabled
+      result = result.replace(job="localhost")
+  if default:
+    # Overrides any defaults with values from the default device if given.
+    result = result.make_merged_spec(
+        tf_device.DeviceSpec.from_string(default))
+
+  # Apply `d` last, so that it's values take precedence over the defaults.
+  result = result.make_merged_spec(d)
+  return result.to_string()
+
+
+def canonicalize_without_job_and_task(d):
+  """Partially canonicalize device string.
+
+  This returns device string from `d` without including job and task.
+  This is most useful for parameter server strategy where the device strings are
+  generated on the chief, but executed on workers.
+
+   For example:
+    If d = '/cpu:0', default='/job:worker/task:1', it returns
+      '/replica:0/device:CPU:0'.
+    If d = '/cpu:0', default='/job:worker', it returns
+      '/replica:0/device:CPU:0'.
+    If d = '/gpu:0', default=None, it returns
+      '/replica:0/device:GPU:0'.
+
+  Note: This uses "job:localhost" as the default if executing eagerly.
+
+  Args:
+    d: a device string or tf.config.LogicalDevice
+
+  Returns:
+    a partially canonicalized device string.
+  """
+  canonicalized_device = canonicalize(d)
+  spec = tf_device.DeviceSpec.from_string(canonicalized_device)
+  spec = spec.replace(job=None, task=None, replica=0)
+  return spec.to_string()
+
+
+def resolve(d):
+  """Canonicalize `d` with current device as default."""
+  return canonicalize(d, default=current())
+
+
+class _FakeNodeDef(object):
+  """A fake NodeDef for _FakeOperation."""
+
+  __slots__ = ["op", "name"]
+
+  def __init__(self):
+    self.op = ""
+    self.name = ""
+
+
+class _FakeOperation(object):
+  """A fake Operation object to pass to device functions."""
+
+  def __init__(self):
+    self.device = ""
+    self.type = ""
+    self.name = ""
+    self.node_def = _FakeNodeDef()
+
+  def _set_device(self, device):
+    self.device = ops._device_string(device)  # pylint: disable=protected-access
+
+  def _set_device_from_string(self, device_str):
+    self.device = device_str
+
+
+def current():
+  """Return a string (not canonicalized) for the current device."""
+  # TODO(josh11b): Work out how this function interacts with ops.colocate_with.
+  if ops.executing_eagerly_outside_functions():
+    d = context.context().device_name
+  else:
+    op = _FakeOperation()
+    ops.get_default_graph()._apply_device_functions(op)  # pylint: disable=protected-access
+    d = op.device
+  return d
+
+
+def get_host_for_device(device, device_index=0):
+  """Returns the corresponding host device for the given device."""
+  spec = tf_device.DeviceSpec.from_string(device)
+  return tf_device.DeviceSpec(
+      job=spec.job,
+      replica=spec.replica,
+      task=spec.task,
+      device_type="CPU",
+      device_index=device_index,
+  ).to_string()
+
+
+def local_devices_from_num_gpus(num_gpus):
+  """Returns device strings for local GPUs or CPU."""
+  return (tuple("/device:GPU:%d" % i for i in range(num_gpus)) or
+          ("/device:CPU:0",))

tensorflow/python/distribute/device_util_test.py

@@ -0,0 +1,125 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Tests for device utilities."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+
+
+class DeviceUtilTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(DeviceUtilTest, self).setUp()
+    context._reset_context()  # pylint: disable=protected-access
+
+  @combinations.generate(
+      combinations.combine(mode="graph")
+  )
+  def testCurrentDeviceWithGlobalGraph(self):
+    with ops.device("/cpu:0"):
+      self.assertEqual(device_util.current(), "/device:CPU:0")
+
+    with ops.device("/job:worker"):
+      with ops.device("/cpu:0"):
+        self.assertEqual(device_util.current(), "/job:worker/device:CPU:0")
+
+    with ops.device("/cpu:0"):
+      with ops.device("/gpu:0"):
+        self.assertEqual(device_util.current(), "/device:GPU:0")
+
+  def testCurrentDeviceWithNonGlobalGraph(self):
+    with ops.Graph().as_default():
+      with ops.device("/cpu:0"):
+        self.assertEqual(device_util.current(), "/device:CPU:0")
+
+  def testCurrentDeviceWithEager(self):
+    with context.eager_mode():
+      with ops.device("/cpu:0"):
+        self.assertEqual(device_util.current(),
+                         "/job:localhost/replica:0/task:0/device:CPU:0")
+
+  @combinations.generate(combinations.combine(mode=["graph", "eager"]))
+  def testCanonicalizeWithoutDefaultDevice(self, mode):
+    if mode == "graph":
+      self.assertEqual(
+          device_util.canonicalize("/cpu:0"),
+          "/replica:0/task:0/device:CPU:0")
+    else:
+      self.assertEqual(
+          device_util.canonicalize("/cpu:0"),
+          "/job:localhost/replica:0/task:0/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/cpu:0"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/task:1/cpu:0"),
+        "/job:worker/replica:0/task:1/device:CPU:0")
+
+  @combinations.generate(combinations.combine(mode=["eager"]))
+  def testCanonicalizeWithoutDefaultDeviceCollectiveEnabled(self):
+    cluster_spec = server_lib.ClusterSpec(
+        multi_worker_test_base.create_cluster_spec(
+            has_chief=False, num_workers=1, num_ps=0, has_eval=False))
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_spec.as_cluster_def(),
+        job_name="worker",
+        task_index=0,
+        protocol="grpc",
+        port=0)
+    context.context().enable_collective_ops(server_def)
+    self.assertEqual(
+        device_util.canonicalize("/cpu:0"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+
+  def testCanonicalizeWithDefaultDevice(self):
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/task:1/cpu:0", default="/gpu:0"),
+        "/job:worker/replica:0/task:1/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/job:worker/task:1", default="/gpu:0"),
+        "/job:worker/replica:0/task:1/device:GPU:0")
+    self.assertEqual(
+        device_util.canonicalize("/cpu:0", default="/job:worker"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+    self.assertEqual(
+        device_util.canonicalize(
+            "/job:worker/replica:0/task:1/device:CPU:0",
+            default="/job:chief/replica:0/task:1/device:CPU:0"),
+        "/job:worker/replica:0/task:1/device:CPU:0")
+
+  def testResolveWithDeviceScope(self):
+    with ops.device("/gpu:0"):
+      self.assertEqual(
+          device_util.resolve("/job:worker/task:1/cpu:0"),
+          "/job:worker/replica:0/task:1/device:CPU:0")
+      self.assertEqual(
+          device_util.resolve("/job:worker/task:1"),
+          "/job:worker/replica:0/task:1/device:GPU:0")
+    with ops.device("/job:worker"):
+      self.assertEqual(
+          device_util.resolve("/cpu:0"),
+          "/job:worker/replica:0/task:0/device:CPU:0")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/distribute_config.py

@@ -0,0 +1,41 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""A configure tuple for high-level APIs for running distribution strategies."""
+
+import collections
+
+
+class DistributeConfig(
+    collections.namedtuple(
+        'DistributeConfig',
+        ['train_distribute', 'eval_distribute', 'remote_cluster'])):
+  """A config tuple for distribution strategies.
+
+  Attributes:
+    train_distribute: a `DistributionStrategy` object for training.
+    eval_distribute: an optional `DistributionStrategy` object for
+      evaluation.
+    remote_cluster: a dict, `ClusterDef` or `ClusterSpec` object specifying
+      the cluster configurations. If this is given, the `train_and_evaluate`
+      method will be running as a standalone client which connects to the
+      cluster for training.
+  """
+
+  def __new__(cls,
+              train_distribute=None,
+              eval_distribute=None,
+              remote_cluster=None):
+    return super(DistributeConfig, cls).__new__(cls, train_distribute,
+                                                eval_distribute, remote_cluster)

tensorflow/python/distribute/distribute_coordinator.py

@@ -0,0 +1,872 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""A component for running distributed TensorFlow."""
+
+import copy
+import json
+import os
+import threading
+import time
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.distribute import distribute_coordinator_context
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import coordinator
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import server_lib
+
+
+_thread_local = threading.local()
+
+
+class _TaskType(object):
+  PS = "ps"
+  WORKER = "worker"
+  CHIEF = "chief"
+  EVALUATOR = "evaluator"
+  CLIENT = "client"
+
+
+# TODO(yuefengz): support another mode where the client colocates with one
+# worker.
+class CoordinatorMode(object):
+  """Specify how distribute coordinator runs."""
+  # The default mode where distribute coordinator will run as a standalone
+  # client and connects to remote servers for training.  Each remote server can
+  # use the distribute coordinator binary with task_type set correctly which
+  # will then turn into standard servers.
+  STANDALONE_CLIENT = "standalone_client"
+
+  # The distribute coordinator runs on each worker. It will run a standard
+  # server on each worker and optionally run the `worker_fn` that is configured
+  # to talk to its standard server.
+  INDEPENDENT_WORKER = "independent_worker"
+
+
+class _Barrier(object):
+  """A reusable barrier class for worker synchronization."""
+
+  def __init__(self, num_participants):
+    """Initializes the barrier object.
+
+    Args:
+      num_participants: an integer which is the expected number of calls of
+        `wait` pass to through this barrier.
+    """
+    self._num_participants = num_participants
+    self._counter = 0
+    self._flag = False
+    self._local_sense = threading.local()
+    self._lock = threading.Lock()
+    self._condition = threading.Condition()
+
+  def wait(self):
+    """Waits until all other callers reach the same wait call."""
+    self._local_sense.value = not self._flag
+    with self._lock:
+      self._counter += 1
+      if self._counter == self._num_participants:
+        self._counter = 0
+        self._flag = self._local_sense.value
+    with self._condition:
+      while self._flag != self._local_sense.value:
+        self._condition.wait()
+      self._condition.notify_all()
+
+
+def _get_num_workers(cluster_spec):
+  """Gets number of workers including chief."""
+  if not cluster_spec:
+    return 0
+  return len(cluster_spec.as_dict().get(_TaskType.WORKER, [])) + len(
+      cluster_spec.as_dict().get(_TaskType.CHIEF, []))
+
+
+class _WorkerContext(object):
+  """The worker context class.
+
+  This context object provides configuration information for each task. One
+  context manager with a worker context object will be created per
+  invocation to the `worker_fn` where `get_current_worker_context` can be called
+  to access the worker context object.
+  """
+
+  def __init__(self,
+               strategy,
+               cluster_spec,
+               task_type,
+               task_id,
+               session_config=None,
+               rpc_layer="grpc",
+               worker_barrier=None):
+    """Initialize the worker context object.
+
+    Args:
+      strategy: a `DistributionStrategy` object.
+      cluster_spec: a ClusterSpec object. It can be empty or None in the local
+        training case.
+      task_type: a string indicating the role of the corresponding task, such as
+        "worker" or "ps". It can be None if it is local training or in-graph
+        replicated training.
+      task_id: an integer indicating id of the corresponding task. It can be
+        None if it is local training or in-graph replicated training.
+      session_config: an optional `tf.compat.v1.ConfigProto` object.
+      rpc_layer: optional string specifying the RPC protocol for communication
+        with worker masters. If None or empty, hosts in the `cluster_spec` will
+        be used directly.
+      worker_barrier: optional, the barrier object for worker synchronization.
+    """
+    self._strategy = strategy
+    self._cluster_spec = cluster_spec
+    self._task_type = task_type
+    self._task_id = task_id
+    self._session_config = session_config
+    self._worker_barrier = worker_barrier
+    self._rpc_layer = rpc_layer
+    self._master_target = self._get_master_target()
+    self._num_workers = _get_num_workers(cluster_spec)
+    self._is_chief_node = self._is_chief()
+
+  def _debug_message(self):
+    if self._cluster_spec:
+      return "[cluster_spec: %r, task_type: %r, task_id: %r]" % (
+          self._cluster_spec, self.task_type, self.task_id)
+    else:
+      return "[local]"
+
+  def __enter__(self):
+    old_context = distribute_coordinator_context.get_current_worker_context()
+    if old_context:
+      raise ValueError(
+          "You cannot run distribute coordinator in a `worker_fn`.\t" +
+          self._debug_message())
+    # pylint: disable=protected-access
+    distribute_coordinator_context._worker_context.current = self
+
+  def __exit__(self, unused_exception_type, unused_exception_value,
+               unused_traceback):
+    # pylint: disable=protected-access
+    distribute_coordinator_context._worker_context.current = None
+
+  def _get_master_target(self):
+    """Return the master target for a task."""
+    # If cluster_spec is None or empty, we use local master.
+    if not self._cluster_spec or self._task_type == _TaskType.EVALUATOR:
+      return ""
+
+    # If task_type is None, then it is in-graph replicated training. In this
+    # case we use the chief or first worker's master target.
+    if not self._task_type:
+      if _TaskType.CHIEF in self._cluster_spec.jobs:
+        task_type = _TaskType.CHIEF
+        task_id = 0
+      else:
+        assert _TaskType.WORKER in self._cluster_spec.jobs
+        task_type = _TaskType.WORKER
+        task_id = 0
+    else:
+      task_type = self._task_type
+      task_id = self._task_id
+
+    prefix = ""
+    if self._rpc_layer:
+      prefix = self._rpc_layer + "://"
+    return prefix + self._cluster_spec.job_tasks(task_type)[task_id or 0]
+
+  def _is_chief(self):
+    """Return whether the task is the chief worker."""
+    if (not self._cluster_spec or
+        self._task_type in [_TaskType.CHIEF, _TaskType.EVALUATOR, None]):
+      return True
+
+    # If not local and chief not in the cluster_spec, use the first worker as
+    # chief.
+    if (_TaskType.CHIEF not in self._cluster_spec.jobs and
+        self._task_type == _TaskType.WORKER and self._task_id == 0):
+      return True
+    return False
+
+  def wait_for_other_workers(self):
+    """Waits for other workers to reach the same call to this method.
+
+    Raises:
+      ValueError: if `worker_barrier` is not passed to the __init__ method.
+    """
+    if not self._worker_barrier:
+      # TODO(yuefengz): we should throw an error in independent worker mode.
+      return
+    self._worker_barrier.wait()
+
+  def session_creator(self,
+                      scaffold=None,
+                      config=None,
+                      checkpoint_dir=None,
+                      checkpoint_filename_with_path=None,
+                      max_wait_secs=7200):
+    """Returns a session creator.
+
+    The returned session creator will be configured with the correct master
+    target and session configs. It will also run either init ops or ready ops
+    by querying the `strategy` object when `create_session` is called on it.
+
+    Args:
+      scaffold: A `Scaffold` used for gathering or building supportive ops. If
+        not specified a default one is created. It's used to finalize the graph.
+      config: `ConfigProto` proto used to configure the session.
+      checkpoint_dir: A string. Optional path to a directory where to restore
+        variables.
+      checkpoint_filename_with_path: Full file name path to the checkpoint file.
+        Only one of `checkpoint_dir` or `checkpoint_filename_with_path` can be
+        specified.
+      max_wait_secs: Maximum time to wait for the session to become available.
+
+    Returns:
+      a descendant of SessionCreator.
+    """
+    if config:
+      session_config = copy.deepcopy(config)
+      session_config.MergeFrom(self._session_config)
+    else:
+      session_config = self._session_config
+
+    if not self._strategy or self._strategy.extended.experimental_should_init:
+      logging.info("Creating chief session creator with config: %r", config)
+      return monitored_session.ChiefSessionCreator(
+          scaffold,
+          master=self.master_target,
+          config=session_config,
+          checkpoint_dir=checkpoint_dir,
+          checkpoint_filename_with_path=checkpoint_filename_with_path)
+    else:
+      logging.info("Creating worker session creator with config: %r", config)
+      return monitored_session.WorkerSessionCreator(
+          scaffold,
+          master=self.master_target,
+          config=session_config,
+          max_wait_secs=max_wait_secs)
+
+  @property
+  def session_config(self):
+    return copy.deepcopy(self._session_config)
+
+  @property
+  def has_barrier(self):
+    """Whether the barrier is set or not."""
+    return self._worker_barrier is not None
+
+  @property
+  def distributed_mode(self):
+    """Whether it is distributed training or not."""
+    return bool(self._cluster_spec) and self._task_type != _TaskType.EVALUATOR
+
+  @property
+  def cluster_spec(self):
+    """Returns a copy of the cluster_spec object."""
+    return copy.deepcopy(self._cluster_spec)
+
+  @property
+  def task_type(self):
+    """Returns the role of the corresponding task."""
+    return self._task_type
+
+  @property
+  def task_id(self):
+    """Returns the id or index of the corresponding task."""
+    return self._task_id
+
+  @property
+  def master_target(self):
+    """Returns the session master for the corresponding task to connect to."""
+    return self._master_target
+
+  @property
+  def is_chief(self):
+    """Returns whether the task is a chief node."""
+    return self._is_chief_node
+
+  @property
+  def num_workers(self):
+    """Returns number of workers in the cluster, including chief."""
+    return self._num_workers
+
+  @property
+  def experimental_should_init(self):
+    """Whether to run init ops."""
+    return self._strategy.extended.experimental_should_init
+
+  @property
+  def should_checkpoint(self):
+    """Whether to save checkpoint."""
+    return self._strategy.extended.should_checkpoint
+
+  @property
+  def should_save_summary(self):
+    """Whether to save summaries."""
+    return self._strategy.extended.should_save_summary
+
+
+def _run_single_worker(worker_fn,
+                       strategy,
+                       cluster_spec,
+                       task_type,
+                       task_id,
+                       session_config,
+                       rpc_layer="",
+                       worker_barrier=None,
+                       coord=None):
+  """Runs a single worker by calling `worker_fn` under context."""
+  session_config = copy.deepcopy(session_config)
+  strategy = copy.deepcopy(strategy)
+  # If there is an EVALUATOR task, we run single-machine eval on that task.
+  if task_type == _TaskType.EVALUATOR:
+    # It is possible to not have a strategy object for EVALUATOR task.
+    if strategy:
+      strategy.configure(session_config)
+  else:
+    assert strategy
+    strategy.configure(session_config, cluster_spec, task_type, task_id)
+
+  context = _WorkerContext(
+      strategy,
+      cluster_spec,
+      task_type,
+      task_id,
+      session_config=session_config,
+      rpc_layer=rpc_layer,
+      worker_barrier=worker_barrier)
+  with context:
+    if coord:
+      with coord.stop_on_exception():
+        return worker_fn(strategy)
+    else:
+      return worker_fn(strategy)
+
+
+def _split_cluster_for_evaluator(cluster_spec, task_type):
+  """Split the cluster for evaluator since it needn't talk to other tasks."""
+  # Splitting the cluster is important to prevent the evaluator from talking to
+  # other tasks in the cluster. Since we allow evaluator not to use
+  # distribution strategies and as a result ops in the evaluator task may have
+  # unspecified devices. Those ops may end up on other tasks if we don't split
+  # the cluster.
+  # Note: if you bypass distribute coordinator and bring the cluster yourself,
+  # you can equivalently set device filters to split clusters. This is already
+  # done by distribution strategy's `update_config_proto` method.
+  new_cluster_spec = multi_worker_util.normalize_cluster_spec(
+      cluster_spec).as_dict()
+  if task_type == _TaskType.EVALUATOR:
+    assert _TaskType.EVALUATOR in new_cluster_spec
+    new_cluster_spec = {
+        _TaskType.EVALUATOR: new_cluster_spec[_TaskType.EVALUATOR]
+    }
+  else:
+    new_cluster_spec.pop(_TaskType.EVALUATOR, None)
+  return multi_worker_util.normalize_cluster_spec(new_cluster_spec)
+
+
+def _run_std_server(cluster_spec=None,
+                    task_type=None,
+                    task_id=None,
+                    session_config=None,
+                    rpc_layer=None,
+                    environment=None):
+  """Runs a standard server."""
+  # Check if the Server is already running. If so, assert that no configuration
+  # options have changed, and return the existing Server. This allows us to
+  # call `run_distribute_coordinator` multiple times.
+  if getattr(_thread_local, "server", None) is not None:
+    assert _thread_local.cluster_spec == cluster_spec
+    assert _thread_local.task_type == task_type
+    assert _thread_local.task_id == task_id
+    assert _thread_local.session_config_str == repr(session_config)
+    assert _thread_local.rpc_layer == rpc_layer
+    assert _thread_local.environment == environment
+    return _thread_local.server
+  else:
+    # This method is not thread-safe.
+    _thread_local.server_started = True
+    _thread_local.cluster_spec = cluster_spec
+    _thread_local.task_type = task_type
+    _thread_local.task_id = task_id
+    _thread_local.session_config_str = repr(session_config)
+    _thread_local.rpc_layer = rpc_layer
+    _thread_local.environment = environment
+
+  assert cluster_spec
+  target = cluster_spec.task_address(task_type, task_id)
+  if rpc_layer:
+    target = rpc_layer + "://" + target
+
+  class _FakeServer(object):
+    """A fake server that runs a master session."""
+
+    def start(self):
+      # A tensorflow server starts when a remote session is created.
+      logging.info(
+          "Creating a remote session to start a TensorFlow server, "
+          "target = %r, session_config=%r", target, session_config)
+      session.Session(target=target, config=session_config)
+
+    def join(self):
+      while True:
+        time.sleep(5)
+
+  if environment == "google":
+    server = _FakeServer()
+  else:
+    if session_config:
+      logging.info(
+          "Starting standard TensorFlow server, target = %r, session_config= "
+          "%r", target, session_config)
+    else:
+      logging.info("Starting standard TensorFlow server, target = %r", target)
+    cluster_spec = _split_cluster_for_evaluator(cluster_spec, task_type)
+    server = server_lib.Server(
+        cluster_spec,
+        job_name=task_type,
+        task_index=task_id,
+        config=session_config,
+        protocol=rpc_layer)
+
+  server.start()
+  _thread_local.server = server
+  return server
+
+
+def _run_between_graph_client(worker_fn, strategy, eval_fn, eval_strategy,
+                              cluster_spec, session_config, rpc_layer):
+  """Runs a standalone client for between-graph replication."""
+  coord = coordinator.Coordinator()
+  eval_thread = None
+  if _TaskType.EVALUATOR in cluster_spec.jobs:
+    eval_thread = threading.Thread(
+        target=_run_single_worker,
+        args=(eval_fn, eval_strategy, cluster_spec, _TaskType.EVALUATOR, 0,
+              session_config),
+        kwargs={
+            "rpc_layer": rpc_layer,
+            "coord": coord,
+        })
+    eval_thread.start()
+
+  threads = []
+  worker_barrier = _Barrier(_get_num_workers(cluster_spec))
+  for task_type in [_TaskType.CHIEF, _TaskType.WORKER]:
+    for task_id in range(len(cluster_spec.as_dict().get(task_type, []))):
+      t = threading.Thread(
+          target=_run_single_worker,
+          args=(worker_fn, strategy, cluster_spec, task_type, task_id,
+                session_config),
+          kwargs={
+              "rpc_layer": rpc_layer,
+              "worker_barrier": worker_barrier,
+              "coord": coord,
+          })
+      t.start()
+      threads.append(t)
+
+  if eval_thread:
+    # TODO(yuefengz): is it necessary to join eval thread?
+    threads_to_join = threads + [eval_thread]
+  else:
+    threads_to_join = threads
+  coord.join(threads_to_join)
+
+  # TODO(yuefengz): we probably want to return results from all workers?
+  return None
+
+
+def _run_in_graph_client(worker_fn, strategy, eval_fn, eval_strategy,
+                         cluster_spec, session_config, rpc_layer):
+  """Runs a standalone client for in-graph replication."""
+  coord = coordinator.Coordinator()
+  eval_thread = None
+  if _TaskType.EVALUATOR in cluster_spec.jobs:
+    eval_thread = threading.Thread(
+        target=_run_single_worker,
+        args=(eval_fn, eval_strategy, cluster_spec, _TaskType.EVALUATOR, 0,
+              session_config),
+        kwargs={
+            "rpc_layer": rpc_layer,
+            "coord": coord,
+        })
+    eval_thread.start()
+
+  worker_result = _run_single_worker(
+      worker_fn,
+      strategy,
+      cluster_spec,
+      None,
+      None,
+      session_config,
+      rpc_layer=rpc_layer,
+      coord=coord)
+
+  if eval_thread:
+    coord.join([eval_thread])
+
+  return worker_result
+
+
+def _configure_session_config_for_std_servers(
+    strategy, eval_strategy, session_config, cluster_spec, task_type, task_id):
+  # pylint: disable=g-doc-args
+  """Call strategy's `configure` to mutate the session_config.
+
+  The session_config is currently needed as default config for a TensorFlow
+  server. In the future, we should be able to remove this method and only pass
+  the session config to a client session.
+  """
+  if task_type == _TaskType.EVALUATOR:
+    if eval_strategy:
+      eval_strategy.configure(session_config=session_config)
+  else:
+    # The strategy may be shared in standalone client mode.
+    strategy = copy.deepcopy(strategy)
+    strategy.configure(
+        session_config=session_config,
+        cluster_spec=cluster_spec,
+        task_type=task_type,
+        task_id=task_id)
+  # Remove the device filters specific to the strategy, so that the
+  # TensorFlow server brought up with one strategy can be used by other
+  # strategies. The device filters can be set in the client side as well.
+  del session_config.device_filters[:]
+
+
+def run_standard_tensorflow_server(session_config=None):
+  """Starts a standard TensorFlow server.
+
+  This method parses configurations from "TF_CONFIG" environment variable and
+  starts a TensorFlow server. The "TF_CONFIG" is typically a json string and
+  must have information of the cluster and the role of the server in the
+  cluster. One example is:
+
+  TF_CONFIG='{
+      "cluster": {
+          "worker": ["host1:2222", "host2:2222", "host3:2222"],
+          "ps": ["host4:2222", "host5:2222"]
+      },
+      "task": {"type": "worker", "index": 1}
+  }'
+
+  This "TF_CONFIG" specifies there are 3 workers and 2 ps tasks in the cluster
+  and the current role is worker 1.
+
+  Valid task types are "chief", "worker", "ps" and "evaluator" and you can have
+  at most one "chief" and at most one "evaluator".
+
+  An optional key-value can be specified is "rpc_layer". The default value is
+  "grpc".
+
+  Args:
+    session_config: an optional `tf.compat.v1.ConfigProto` object. Users can
+      pass in the session config object to configure server-local devices.
+
+  Returns:
+    a `tf.distribute.Server` object which has already been started.
+
+  Raises:
+    ValueError: if the "TF_CONFIG" environment is not complete.
+  """
+  tf_config = json.loads(os.environ.get("TF_CONFIG", "{}"))
+  if "cluster" not in tf_config:
+    raise ValueError("\"cluster\" is not found in TF_CONFIG.")
+  cluster_spec = multi_worker_util.normalize_cluster_spec(tf_config["cluster"])
+  if "task" not in tf_config:
+    raise ValueError("\"task\" is not found in TF_CONFIG.")
+  task_env = tf_config["task"]
+  if "type" not in task_env:
+    raise ValueError(
+        "\"task_type\" is not found in the `task` part of TF_CONFIG.")
+  task_type = task_env["type"]
+  task_id = int(task_env.get("index", 0))
+
+  rpc_layer = tf_config.get("rpc_layer", "grpc")
+
+  session_config = session_config or config_pb2.ConfigProto()
+  # Set the collective group leader for collective ops to initialize collective
+  # ops when server starts.
+  if "chief" in cluster_spec.jobs:
+    session_config.experimental.collective_group_leader = (
+        "/job:chief/replica:0/task:0")
+  else:
+    if "worker" not in cluster_spec.jobs:
+      raise ValueError(
+          "You must have `chief` or `worker` jobs in the `cluster_spec`.")
+    session_config.experimental.collective_group_leader = (
+        "/job:worker/replica:0/task:0")
+
+  server = _run_std_server(
+      cluster_spec=cluster_spec,
+      task_type=task_type,
+      task_id=task_id,
+      session_config=session_config,
+      rpc_layer=rpc_layer)
+  server.start()
+  return server
+
+
+# TODO(yuefengz): propagate cluster_spec in the STANDALONE_CLIENT mode.
+# TODO(yuefengz): we may need a smart way to figure out whether the current task
+# is the special task when we support cluster_spec propagation.
+def run_distribute_coordinator(worker_fn,
+                               strategy,
+                               eval_fn=None,
+                               eval_strategy=None,
+                               mode=CoordinatorMode.STANDALONE_CLIENT,
+                               cluster_spec=None,
+                               task_type=None,
+                               task_id=None,
+                               session_config=None,
+                               rpc_layer="grpc"):
+  """Runs the coordinator for distributed TensorFlow.
+
+  This function runs a split coordinator for distributed TensorFlow in its
+  default mode, i.e the STANDALONE_CLIENT mode. Given a `cluster_spec`
+  specifying server addresses and their roles in a cluster, this coordinator
+  will figure out how to set them up, give the underlying function the right
+  targets for master sessions via a scope object and coordinate their training.
+  The cluster consisting of standard servers needs to be brought up either with
+  the standard server binary or with a binary running distribute coordinator
+  with `task_type` set to non-client type which will then turn into standard
+  servers.
+
+  In addition to be the distribute coordinator, this is also the source of
+  configurations for each job in the distributed training. As there are multiple
+  ways to configure a distributed TensorFlow cluster, its context object
+  provides these configurations so that users or higher-level APIs don't have to
+  figure out the configuration for each job by themselves.
+
+  In the between-graph replicated training, this coordinator will create
+  multiple threads and each calls the `worker_fn` which is supposed to create
+  its own graph and connect to one worker master given by its context object. In
+  the in-graph replicated training, it has only one thread calling this
+  `worker_fn`.
+
+  Another mode is the INDEPENDENT_WORKER mode where each server runs a
+  distribute coordinator which will start a standard server and optionally runs
+  `worker_fn` depending whether it is between-graph training or in-graph
+  replicated training.
+
+  The `strategy` object is expected to be a DistributionStrategy object which
+  has implemented methods needed by distributed coordinator such as
+  `configure(session_config, cluster_spec, task_type, task_id)` which configures
+  the strategy object for a specific task and `experimental_should_init`
+  property which instructs the distribute coordinator whether to run init ops
+  for a task. The distribute coordinator will make a copy of the `strategy`
+  object, call its `configure` method and pass it to `worker_fn` as an argument.
+
+  The `worker_fn` defines the training logic and is called under its own
+  worker context which can be accessed to via `get_current_worker_context`. A
+  worker context provides access to configurations for each task, e.g. the
+  task_type, task_id, master target and so on. Since `worker_fn` will be called
+  in a thread and possibly multiple times, caller should be careful when it
+  accesses global data. For example, it is unsafe to define flags in a
+  `worker_fn` or to define different environment variables for different
+  `worker_fn`s.
+
+  The `worker_fn` for the between-graph replication is defined as if there is
+  only one worker corresponding to the `worker_fn` and possibly ps jobs. For
+  example, when training with parameter servers, it assigns variables to
+  parameter servers and all other operations to that worker. In the in-graph
+  replication case, the `worker_fn` has to define operations for all worker
+  jobs. Using a distribution strategy can simplify the `worker_fn` by not having
+  to worry about the replication and device assignment of variables and
+  operations.
+
+  This method is intended to be invoked by high-level APIs so that users don't
+  have to explicitly call it to run this coordinator. For those who don't use
+  high-level APIs, to change a program to use this coordinator, wrap everything
+  in a the program after global data definitions such as commandline flag
+  definition into the `worker_fn` and get task-specific configurations from
+  the worker context.
+
+  The `cluster_spec` can be either passed by the argument or parsed from the
+  "TF_CONFIG" environment variable. Example of a TF_CONFIG:
+  ```
+    cluster = {'chief': ['host0:2222'],
+               'ps': ['host1:2222', 'host2:2222'],
+               'worker': ['host3:2222', 'host4:2222', 'host5:2222']}
+    os.environ['TF_CONFIG'] = json.dumps({'cluster': cluster})
+  ```
+
+  If `cluster_spec` is not given in any format, it becomes local training and
+  this coordinator will connect to a local session.
+
+  For evaluation, if "evaluator" exists in the cluster_spec, a separate thread
+  will be created to call `eval_fn` with its `task_type` set to "evaluator". If
+  `eval_fn` is not defined, fall back to `worker_fn`. This implies that
+  evaluation will be done on a single machine if there is an "evaluator" task.
+  If "evaluator" doesn't exist in the cluster_spec, it entirely depends on the
+  `worker_fn` for how to do evaluation.
+
+  Args:
+    worker_fn: the function to be called. The function should accept a
+      `strategy` object and will be given access to a context object via a
+      context manager scope.
+    strategy: a DistributionStrategy object specifying whether it should
+      run between-graph replicated training or not, whether to run init ops,
+      etc. This object will also be configured given `session_config`,
+      `cluster_spec`, `task_type` and `task_id`.
+    eval_fn: optional function for "evaluator" task. If `eval_fn` is not passed
+      in but a "evaluator" task is found in the `cluster_spec`, the `worker_fn`
+      will be used for this task.
+    eval_strategy: optional DistributionStrategy object for "evaluator" task.
+    mode: in which mode this distribute coordinator runs.
+    cluster_spec: a dict, ClusterDef or ClusterSpec specifying servers and roles
+      in a cluster. If not set or empty, fall back to local training.
+    task_type: the current task type, optional if this is a client.
+    task_id: the current task id, optional if this is a client.
+    session_config: an optional `tf.compat.v1.ConfigProto` object which will be
+      passed to `strategy`'s `configure` method and used to create a session.
+    rpc_layer: optional string, the protocol for RPC, e.g. "grpc".
+
+  Raises:
+    ValueError: if `cluster_spec` is supplied but not a dict or a ClusterDef or
+      a ClusterSpec.
+
+  Returns:
+    In the client job, return the value returned by `worker_fn` if
+    it is in-graph replication or INDEPENDENT_WORKER mode; return None
+    otherwise.
+  """
+  tf_config = json.loads(os.environ.get("TF_CONFIG", "{}"))
+  rpc_layer = tf_config.get("rpc_layer", rpc_layer)
+  environment = tf_config.get("environment", None)
+
+  if not cluster_spec:
+    cluster_spec = tf_config.get("cluster", {})
+    task_env = tf_config.get("task", {})
+    if task_env:
+      task_type = task_env.get("type", task_type)
+      task_id = int(task_env.get("index", task_id))
+
+  if cluster_spec:
+    # TODO(yuefengz): validate cluster_spec.
+    cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec)
+  elif hasattr(strategy.extended, "_cluster_resolver"):
+    cluster_resolver = strategy.extended._cluster_resolver  # pylint: disable=protected-access
+    task_type = cluster_resolver.task_type
+    task_id = cluster_resolver.task_id
+    rpc_layer = cluster_resolver.rpc_layer or rpc_layer
+    environment = cluster_resolver.environment
+    cluster_spec = cluster_resolver.cluster_spec()
+
+  # Setting the session config is necessary for some strategies such as
+  # CollectiveAllReduceStrategy.
+  session_config = session_config or config_pb2.ConfigProto(
+      allow_soft_placement=True)
+
+  if cluster_spec:
+    logging.info(
+        "Running Distribute Coordinator with mode = %r, cluster_spec = %r, "
+        "task_type = %r, task_id = %r, environment = %r, rpc_layer = %r", mode,
+        cluster_spec.as_dict(), task_type, task_id, environment, rpc_layer)
+
+  if not cluster_spec:
+    # `mode` is ignored in the local case.
+    logging.info("Running local Distribute Coordinator.")
+    _run_single_worker(worker_fn, strategy, None, None, None, session_config,
+                       rpc_layer)
+    if eval_fn:
+      _run_single_worker(eval_fn, eval_strategy, None, None, None,
+                         session_config, rpc_layer)
+    else:
+      logging.warning("Skipped evaluation since `eval_fn` is not passed in.")
+  elif mode == CoordinatorMode.STANDALONE_CLIENT:
+    if not eval_fn:
+      logging.warning("`eval_fn` is not passed in. The `worker_fn` will be "
+                      "used if an \"evaluator\" task exists in the cluster.")
+    eval_fn = eval_fn or worker_fn
+    if not eval_strategy:
+      logging.warning("`eval_strategy` is not passed in. No distribution "
+                      "strategy will be used for evaluation.")
+
+    # The client must know the cluster but servers in the cluster don't have to
+    # know the client.
+    if task_type in [_TaskType.CLIENT, None]:
+      if strategy.extended.experimental_between_graph:
+        return _run_between_graph_client(worker_fn, strategy, eval_fn,
+                                         eval_strategy, cluster_spec,
+                                         session_config, rpc_layer)
+      else:
+        return _run_in_graph_client(worker_fn, strategy, eval_fn, eval_strategy,
+                                    cluster_spec, session_config, rpc_layer)
+    else:
+      # If not a client job, run the standard server.
+      _configure_session_config_for_std_servers(strategy, eval_strategy,
+                                                session_config, cluster_spec,
+                                                task_type, task_id)
+      server = _run_std_server(
+          cluster_spec=cluster_spec,
+          task_type=task_type,
+          task_id=task_id,
+          session_config=session_config,
+          rpc_layer=rpc_layer,
+          environment=environment)
+      server.join()
+  else:
+    if mode != CoordinatorMode.INDEPENDENT_WORKER:
+      raise ValueError("Unexpected coordinator mode: %r" % mode)
+
+    if not eval_fn:
+      logging.warning("`eval_fn` is not passed in. The `worker_fn` will be "
+                      "used if an \"evaluator\" task exists in the cluster.")
+    eval_fn = eval_fn or worker_fn
+    if not eval_strategy:
+      logging.warning("`eval_strategy` is not passed in. No distribution "
+                      "strategy will be used for evaluation.")
+
+    # Every one starts a standard server, get session config from `configure`
+    # method.
+    _configure_session_config_for_std_servers(strategy, eval_strategy,
+                                              session_config, cluster_spec,
+                                              task_type, task_id)
+
+    if (task_type != _TaskType.EVALUATOR and
+        not getattr(strategy.extended, "_std_server_started", False)):
+      # Right now, with eager mode, context is configured with a std server at
+      # the very beginning while with graph mode the std server is started when
+      # distribute coordinator is called. We should consolidate these two paths.
+      server = _run_std_server(
+          cluster_spec=cluster_spec,
+          task_type=task_type,
+          task_id=task_id,
+          session_config=session_config,
+          rpc_layer=rpc_layer,
+          environment=environment)
+    if task_type in [_TaskType.CHIEF, _TaskType.WORKER]:
+      if strategy.extended.experimental_between_graph:
+        # All jobs run `worker_fn` if between-graph.
+        return _run_single_worker(worker_fn, strategy, cluster_spec, task_type,
+                                  task_id, session_config, rpc_layer)
+      else:
+        # Only one node runs `worker_fn` if in-graph.
+        context = _WorkerContext(strategy, cluster_spec, task_type, task_id)
+        if context.is_chief:
+          return _run_single_worker(worker_fn, strategy, cluster_spec, None,
+                                    None, session_config, rpc_layer)
+        else:
+          server.join()
+    elif task_type == _TaskType.EVALUATOR:
+      return _run_single_worker(eval_fn, eval_strategy, cluster_spec, task_type,
+                                task_id, session_config, rpc_layer)
+    else:
+      if task_type != _TaskType.PS:
+        raise ValueError("Unexpected task_type: %r" % task_type)
+      server.join()

tensorflow/python/distribute/distribute_coordinator_context.py

@@ -0,0 +1,27 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""The context retrieval method for distribute coordinator."""
+
+import threading
+
+_worker_context = threading.local()
+
+
+def get_current_worker_context():
+  """Returns the current task context."""
+  try:
+    return _worker_context.current
+  except AttributeError:
+    return None

tensorflow/python/distribute/distribute_coordinator_test.py

@@ -0,0 +1,942 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Tests for Distribute Coordinator."""
+
+import contextlib
+import copy
+import json
+import os
+import sys
+import threading
+import time
+
+import six
+
+# pylint: disable=g-import-not-at-top
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.distribute import distribute_coordinator
+from tensorflow.python.distribute import distribute_coordinator_context
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import coordinator
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import session_manager
+
+
+CHIEF = distribute_coordinator._TaskType.CHIEF
+WORKER = distribute_coordinator._TaskType.WORKER
+PS = distribute_coordinator._TaskType.PS
+EVALUATOR = distribute_coordinator._TaskType.EVALUATOR
+
+STANDALONE_CLIENT = distribute_coordinator.CoordinatorMode.STANDALONE_CLIENT
+INDEPENDENT_WORKER = distribute_coordinator.CoordinatorMode.INDEPENDENT_WORKER
+
+NUM_WORKERS = 3
+NUM_PS = 2
+
+original_sys_exit = sys.exit
+
+
+def _bytes_to_str(maybe_bytes):
+  if isinstance(maybe_bytes, six.string_types):
+    return maybe_bytes
+  else:
+    return str(maybe_bytes, "utf-8")
+
+
+def _strip_protocol(target):
+  # cluster_spec expects "host:port" strings.
+  if "//" in target:
+    return target.split("//")[1]
+  else:
+    return target
+
+
+class MockExtended(object):
+
+  def __init__(self,
+               between_graph=False,
+               should_init=None,
+               should_checkpoint=None,
+               should_save_summary=None):
+    self.experimental_between_graph = between_graph
+    self.experimental_should_init = should_init
+    self.should_checkpoint = should_checkpoint
+    self.should_save_summary = should_save_summary
+
+
+class MockStrategy(object):
+
+  def __init__(self,
+               between_graph=False,
+               should_init=None,
+               should_checkpoint=None,
+               should_save_summary=None):
+    self.extended = MockExtended(between_graph, should_init, should_checkpoint,
+                                 should_save_summary)
+
+  def configure(self,
+                session_config=None,
+                cluster_spec=None,
+                task_type=None,
+                task_id=None):
+    if self.extended.experimental_should_init is None:
+      if task_id == 0:
+        self.extended.experimental_should_init = True
+      else:
+        self.extended.experimental_should_init = False
+    if self.extended.should_checkpoint is None:
+      if task_id == 0:
+        self.extended.should_checkpoint = True
+      else:
+        self.extended.should_checkpoint = False
+    if self.extended.should_save_summary is None:
+      if task_id == 0:
+        self.extended.should_save_summary = True
+      else:
+        self.extended.should_save_summary = False
+
+    if session_config:
+      if (cluster_spec and task_type and task_id is not None and
+          self.extended.experimental_between_graph):
+        session_config.intra_op_parallelism_threads += 1
+        if task_type in ["chief", "worker"]:
+          session_config.device_filters.extend(
+              ["/job:%s/task:%d" % (task_type, task_id), "/job:ps"])
+      else:
+        session_config.inter_op_parallelism_threads += 1
+        session_config.device_filters.append("/job:somejob")
+
+
+class MockServer(object):
+
+  def __init__(self):
+    self._joined = False
+    self._started = False
+
+  def start(self):
+    self._started = True
+
+  def join(self):
+    assert not self._joined
+    self._joined = True
+
+  @property
+  def joined(self):
+    return self._joined
+
+  @property
+  def started(self):
+    return self._started
+
+
+class DistributeCoordinatorTestBase(test.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    # We have to create a global in-process cluster because once an in-process
+    # tensorflow server is created, there is no way to terminate it. Please see
+    # multi_worker_test_base.py for more details.
+    # TODO(yuefengz): use the utitliy from multi_worker_test_base.
+    cls._workers, cls._ps = test_util.create_local_cluster(
+        NUM_WORKERS, num_ps=NUM_PS)
+    cls._cluster_spec = {
+        WORKER: [
+            _strip_protocol(_bytes_to_str(w.target)) for w in cls._workers
+        ],
+        PS: [_strip_protocol(_bytes_to_str(ps.target)) for ps in cls._ps]
+    }
+
+  def setUp(self):
+    self._result_correct = 0
+    self._lock = threading.Lock()
+    self._worker_context = {}
+    self._strategy_property = {}
+    self._std_servers = {}
+    self._barrier = distribute_coordinator._Barrier(NUM_WORKERS)
+    self._coord = coordinator.Coordinator()
+
+  @contextlib.contextmanager
+  def _test_session(self, target):
+    config = config_pb2.ConfigProto(allow_soft_placement=True)
+    config.graph_options.optimizer_options.opt_level = -1
+    with session.Session(graph=None, config=config, target=target) as sess:
+      yield sess
+
+  # TODO(yuefengz): use the utitliy from multi_worker_test_base.
+  def _create_cluster_spec(self,
+                           has_chief=False,
+                           num_workers=1,
+                           num_ps=0,
+                           has_eval=False):
+    cluster_spec = {}
+    if has_chief:
+      cluster_spec[CHIEF] = ["localhost:%s" % test_util.pick_unused_port()]
+    if num_workers:
+      cluster_spec[WORKER] = [
+          "localhost:%s" % test_util.pick_unused_port()
+          for _ in range(num_workers)
+      ]
+    if num_ps:
+      cluster_spec[PS] = [
+          "localhost:%s" % test_util.pick_unused_port() for _ in range(num_ps)
+      ]
+    if has_eval:
+      cluster_spec[EVALUATOR] = ["localhost:%s" % test_util.pick_unused_port()]
+    return cluster_spec
+
+  def _in_graph_worker_fn(self, strategy):
+    context = distribute_coordinator_context.get_current_worker_context()
+    self.assertTrue(context is not None)
+    with self._test_session(target=context.master_target) as sess:
+      xs = []
+      expected = 0.0
+      for i in range(context.num_workers):
+        with ops.device("/job:worker/task:%d" % i):
+          x = variable_scope.get_variable("x_%d" % i, initializer=10.0)
+          x_add = x.assign_add(float(i))
+          xs.append(x_add)
+          expected += i + 10.0
+
+      with ops.device("/job:worker/task:0"):
+        result = math_ops.add_n(xs)
+
+      self.evaluate(variables.global_variables_initializer())
+      result_value = sess.run(result)
+    self.assertEqual(result_value, expected)
+    if result_value == expected:
+      self._result_correct += 1
+
+  def _wrapped_worker_fn(self, worker_fn):
+    def wrapped(*args, **kwargs):
+      with self._coord.stop_on_exception():
+        return worker_fn(*args, **kwargs)
+    return wrapped
+
+  def _run_coordinator_in_thread(self, worker_fn, strategy, **kwargs):
+    t = threading.Thread(
+        target=distribute_coordinator.run_distribute_coordinator,
+        args=(self._wrapped_worker_fn(worker_fn), strategy),
+        kwargs=kwargs)
+    t.start()
+    return t
+
+  def _run_multiple_coordinator_in_threads(self, worker_fn, strategy,
+                                           cluster_spec, **kwargs):
+    threads = {}
+    for task_type in cluster_spec.keys():
+      threads[task_type] = []
+      for task_id in range(len(cluster_spec[task_type])):
+        t = self._run_coordinator_in_thread(
+            worker_fn,
+            strategy,
+            cluster_spec=cluster_spec,
+            task_type=task_type,
+            task_id=task_id,
+            **kwargs)
+        threads[task_type].append(t)
+    return threads
+
+  def _join_threads(self, threads):
+    try:
+      self._coord.join(threads)
+    except errors.UnknownError as e:
+      if "Could not start gRPC server" in e.message:
+        self.skipTest("Cannot start std servers.")
+      else:
+        raise
+
+  def _between_graph_worker_fn(self, strategy):
+    context = distribute_coordinator_context.get_current_worker_context()
+    self.assertTrue(context is not None)
+    with self._test_session(target=context.master_target) as sess:
+      with ops.device("/job:ps/task:0"):
+        # TODO(yuefengz): investigate why not using resource variable will make
+        # the test flaky.
+        x = variable_scope.get_variable(
+            "x", initializer=10.0, use_resource=True)
+      with ops.device("/job:ps/task:1"):
+        y = variable_scope.get_variable(
+            "y", initializer=20.0, use_resource=True)
+
+      x_add = x.assign_add(2.0)
+      y_sub = y.assign_sub(2.0)
+      train_op = control_flow_ops.group([x_add, y_sub])
+
+      if context.is_chief:
+        self.evaluate(variables.global_variables_initializer())
+
+      # Synchronize workers after initialization.
+      if context.has_barrier:
+        context.wait_for_other_workers()
+      else:
+        while True:
+          uninit_vars = sess.run(variables.report_uninitialized_variables())
+          # pylint: disable=g-explicit-length-test
+          if len(uninit_vars) == 0:
+            break
+
+      sess.run(train_op)
+
+      # Synchronize workers after one step to make sure they all have finished
+      # training.
+      if context.has_barrier:
+        context.wait_for_other_workers()
+      else:
+        self._barrier.wait()
+
+      x_val, y_val = sess.run([x, y])
+
+      self.assertEqual(x_val, 16.0)
+      self.assertEqual(y_val, 14.0)
+      if x_val == 16.0 and y_val == 14.0:
+        with self._lock:
+          self._result_correct += 1
+
+  def _between_graph_with_monitored_session(self, strategy):
+    context = distribute_coordinator_context.get_current_worker_context()
+    self.assertTrue(context is not None)
+    with ops.device("/job:ps/task:0"):
+      # TODO(yuefengz): investigate why not using resource variable will make
+      # the test flaky.
+      x = variable_scope.get_variable("xx", initializer=10.0, use_resource=True)
+    with ops.device("/job:ps/task:1"):
+      y = variable_scope.get_variable("yy", initializer=20.0, use_resource=True)
+
+    x_add = x.assign_add(2.0)
+    y_sub = y.assign_sub(2.0)
+    train_op = control_flow_ops.group([x_add, y_sub])
+
+    # The monitored session will run init or ready ops.
+    with monitored_session.MonitoredSession() as sess:
+      sess.run(train_op)
+
+      # Synchronize workers after one step to make sure they all have finished
+      # training.
+      if context.has_barrier:
+        context.wait_for_other_workers()
+      else:
+        self._barrier.wait()
+
+      x_val, y_val = sess.run([x, y])
+
+    self.assertEqual(x_val, 16.0)
+    self.assertEqual(y_val, 14.0)
+    if x_val == 16.0 and y_val == 14.0:
+      with self._lock:
+        self._result_correct += 1
+
+  def _dump_worker_context(self, strategy):
+    """Dumps the properties of each worker context.
+
+    It dumps the context properties to a dict mapping from task_type to a list
+    of tuples of master_target, num_workers, is_chief and distribute_mode, where
+    the list is indexed by the task_id.
+
+    Args:
+      strategy: a `DistributionStrategy` object.
+    """
+    context = distribute_coordinator_context.get_current_worker_context()
+    self.assertTrue(context is not None)
+    task_type = str(context.task_type)
+    task_id = context.task_id or 0
+    with self._lock:
+      if task_type not in self._worker_context:
+        self._worker_context[task_type] = []
+      while len(self._worker_context[task_type]) <= task_id:
+        self._worker_context[task_type].append(None)
+      self._worker_context[task_type][task_id] = (context.master_target,
+                                                  context.num_workers,
+                                                  context.is_chief,
+                                                  context.distributed_mode)
+
+  def _dump_strategy_property(self, strategy):
+    context = distribute_coordinator_context.get_current_worker_context()
+    self.assertTrue(context is not None)
+
+    self.assertEqual(context._strategy.extended.experimental_should_init,
+                     strategy.extended.experimental_should_init)
+    self.assertEqual(context.should_checkpoint,
+                     strategy.extended.should_checkpoint)
+    self.assertEqual(context.should_save_summary,
+                     strategy.extended.should_save_summary)
+
+    task_type = str(context.task_type)
+    task_id = context.task_id or 0
+    with self._lock:
+      if task_type not in self._strategy_property:
+        self._strategy_property[task_type] = []
+      while len(self._strategy_property[task_type]) <= task_id:
+        self._strategy_property[task_type].append(None)
+      self._strategy_property[task_type][task_id] = (
+          context._strategy.extended.experimental_should_init,
+          context.should_checkpoint,
+          context.should_save_summary)
+
+  def _run_mock_std_server(self,
+                           session_config=None,
+                           cluster_spec=None,
+                           task_type=None,
+                           task_id=None,
+                           rpc_layer=None,
+                           environment=None):
+    task_type = str(task_type)
+    task_id = task_id or 0
+    with self._lock:
+      if task_type not in self._std_servers:
+        self._std_servers[task_type] = []
+      while len(self._std_servers[task_type]) <= task_id:
+        self._std_servers[task_type].append(None)
+
+      server = MockServer()
+      self._std_servers[task_type][task_id] = server
+    return server
+
+
+class DistributeCoordinatorTestStandaloneMode(DistributeCoordinatorTestBase):
+
+  def testInGraphStandaloneMode(self):
+    """Test it runs in-graph replication in standalone client mode."""
+    distribute_coordinator.run_distribute_coordinator(
+        self._in_graph_worker_fn,
+        MockStrategy(between_graph=False),
+        cluster_spec=self._cluster_spec)
+    self.assertEqual(self._result_correct, 1)
+
+  def testBetweenGraph(self):
+    """Test it runs between-graph replication in standalone client mode."""
+    distribute_coordinator.run_distribute_coordinator(
+        self._between_graph_worker_fn,
+        MockStrategy(between_graph=True),
+        cluster_spec=self._cluster_spec)
+
+    # Each finished worker will increment self._result_correct.
+    self.assertEqual(self._result_correct, NUM_WORKERS)
+
+  @test_util.run_v1_only("MonitoredSession removed from v2")
+  def testBetweenGraphWithMonitoredSession(self):
+    """Test monitored session in standalone client mode."""
+    distribute_coordinator.run_distribute_coordinator(
+        self._between_graph_with_monitored_session,
+        MockStrategy(between_graph=True),
+        cluster_spec=self._cluster_spec)
+
+    # Each finished worker will increment self._result_correct.
+    self.assertEqual(self._result_correct, NUM_WORKERS)
+
+  def testBetweenGraphContext(self):
+    # Dumps the task contexts to the self._worker_context dict.
+    distribute_coordinator.run_distribute_coordinator(
+        self._dump_worker_context,
+        MockStrategy(between_graph=True),
+        cluster_spec=self._cluster_spec)
+
+    # There is only one type of task and there three such tasks.
+    self.assertEqual(len(self._worker_context), 1)
+    self.assertTrue(WORKER in self._worker_context)
+    self.assertEqual(len(self._worker_context[WORKER]), NUM_WORKERS)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(
+        self._worker_context[WORKER][0],
+        (_bytes_to_str(self._workers[0].target), NUM_WORKERS, True, True))
+    self.assertEqual(
+        self._worker_context[WORKER][1],
+        (_bytes_to_str(self._workers[1].target), NUM_WORKERS, False, True))
+    self.assertEqual(
+        self._worker_context[WORKER][2],
+        (_bytes_to_str(self._workers[2].target), NUM_WORKERS, False, True))
+
+  def testBetweenGraphStrategyProperties(self):
+    # Dumps properties of the strategy objects.
+    distribute_coordinator.run_distribute_coordinator(
+        self._dump_strategy_property,
+        MockStrategy(between_graph=True, should_init=True),
+        cluster_spec=self._cluster_spec)
+
+    # There is only one type of task and there three such tasks.
+    self.assertEqual(len(self._strategy_property), 1)
+    self.assertTrue(WORKER in self._strategy_property)
+    self.assertEqual(len(self._strategy_property[WORKER]), NUM_WORKERS)
+
+    # Check whether each task has the right properties of should_init,
+    # should_checkpoint and should_save_summary.
+    self.assertEqual(self._strategy_property[WORKER][0], (True, True, True))
+    self.assertEqual(self._strategy_property[WORKER][1], (True, False, False))
+    self.assertEqual(self._strategy_property[WORKER][2], (True, False, False))
+
+  def testInGraphContext(self):
+    # Dumps the task contexts to the self._worker_context dict.
+    distribute_coordinator.run_distribute_coordinator(
+        self._dump_worker_context,
+        MockStrategy(between_graph=False),
+        cluster_spec=self._cluster_spec)
+
+    # There is only a "None" task in the dumped task context.
+    self.assertEqual(len(self._worker_context), 1)
+    self.assertTrue("None" in self._worker_context)
+    self.assertEqual(len(self._worker_context["None"]), 1)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(
+        self._worker_context["None"][0],
+        (_bytes_to_str(self._workers[0].target), NUM_WORKERS, True, True))
+
+  def testLocalContext(self):
+    # Dumps the task contexts to the self._worker_context dict.
+    distribute_coordinator.run_distribute_coordinator(
+        self._dump_worker_context,
+        MockStrategy(between_graph=False),
+        cluster_spec=None)
+
+    # There is only a "None" task.
+    self.assertEqual(len(self._worker_context), 1)
+    self.assertTrue("None" in self._worker_context)
+    self.assertEqual(len(self._worker_context["None"]), 1)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(self._worker_context["None"][0], ("", 0, True, False))
+
+  def testBetweenGraphContextWithChief(self):
+    # Adds a chief node, so there are NUM_WORKERS + 1 workers in total.
+    cluster_spec = copy.deepcopy(self._cluster_spec)
+    cluster_spec[CHIEF] = ["fake_chief"]
+
+    # Dumps the task contexts to the self._worker_context dict.
+    distribute_coordinator.run_distribute_coordinator(
+        self._dump_worker_context,
+        MockStrategy(between_graph=True),
+        cluster_spec=cluster_spec,
+        rpc_layer="grpc")
+
+    # There are one CHIEF and three workers.
+    self.assertEqual(len(self._worker_context), 2)
+    self.assertTrue(CHIEF in self._worker_context)
+    self.assertTrue(WORKER in self._worker_context)
+    self.assertEqual(len(self._worker_context[CHIEF]), 1)
+    self.assertEqual(len(self._worker_context[WORKER]), NUM_WORKERS)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(self._worker_context[CHIEF][0],
+                     ("grpc://fake_chief", 4, True, True))
+    self.assertEqual(
+        self._worker_context[WORKER][0],
+        (_bytes_to_str(self._workers[0].target), NUM_WORKERS + 1, False, True))
+    self.assertEqual(
+        self._worker_context[WORKER][1],
+        (_bytes_to_str(self._workers[1].target), NUM_WORKERS + 1, False, True))
+    self.assertEqual(
+        self._worker_context[WORKER][2],
+        (_bytes_to_str(self._workers[2].target), NUM_WORKERS + 1, False, True))
+
+  def testInGraphContextWithEval(self):
+    # Adds a EVALUATOR job.
+    cluster_spec = copy.deepcopy(self._cluster_spec)
+    cluster_spec[EVALUATOR] = ["fake_evaluator"]
+
+    # Dumps the task contexts to the self._worker_context dict.
+    distribute_coordinator.run_distribute_coordinator(
+        self._dump_worker_context,
+        MockStrategy(between_graph=False),
+        cluster_spec=cluster_spec,
+        rpc_layer=None)
+
+    # There are one "None" task and one EVALUATOR task.
+    self.assertEqual(len(self._worker_context), 2)
+    self.assertTrue("None" in self._worker_context)
+    self.assertTrue(EVALUATOR in self._worker_context)
+    self.assertEqual(len(self._worker_context["None"]), 1)
+    self.assertEqual(len(self._worker_context[EVALUATOR]), 1)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(self._worker_context["None"][0], (_strip_protocol(
+        _bytes_to_str(self._workers[0].target)), 3, True, True))
+    self.assertEqual(self._worker_context[EVALUATOR][0], ("", 3, True, False))
+
+
+class DistributeCoordinatorTestIndependentWorkerMode(
+    DistributeCoordinatorTestBase):
+
+  def testInGraph(self):
+    cluster_spec = self._create_cluster_spec(num_workers=NUM_WORKERS)
+    threads = self._run_multiple_coordinator_in_threads(
+        self._in_graph_worker_fn,
+        MockStrategy(between_graph=False),
+        cluster_spec,
+        mode=INDEPENDENT_WORKER)
+    self._join_threads([threads[WORKER][0]])
+    self.assertEqual(self._result_correct, 1)
+
+  def testBetweenGraph(self):
+    cluster_spec = self._create_cluster_spec(
+        num_workers=NUM_WORKERS, num_ps=NUM_PS)
+    threads = self._run_multiple_coordinator_in_threads(
+        self._between_graph_worker_fn,
+        MockStrategy(between_graph=True),
+        cluster_spec,
+        mode=INDEPENDENT_WORKER)
+    self._join_threads(threads[WORKER])
+
+    # Each finished worker will increment self._result_correct.
+    self.assertEqual(self._result_correct, NUM_WORKERS)
+
+  @test_util.run_v1_only("MonitoredSession removed from v2")
+  def testBetweenGraphWithMonitoredSession(self):
+    cluster_spec = self._create_cluster_spec(
+        num_workers=NUM_WORKERS, num_ps=NUM_PS)
+    threads = self._run_multiple_coordinator_in_threads(
+        self._between_graph_with_monitored_session,
+        MockStrategy(between_graph=True),
+        cluster_spec,
+        mode=INDEPENDENT_WORKER)
+    self._join_threads(threads[WORKER])
+
+    # Each finished worker will increment self._result_correct.
+    self.assertEqual(self._result_correct, NUM_WORKERS)
+
+  def testBetweenGraphContext(self):
+    cluster_spec = self._create_cluster_spec(num_workers=NUM_WORKERS)
+    # Dumps the task contexts and std server arguments.
+    with test.mock.patch.object(distribute_coordinator, "_run_std_server",
+                                self._run_mock_std_server):
+      threads = self._run_multiple_coordinator_in_threads(
+          self._dump_worker_context,
+          MockStrategy(between_graph=True),
+          cluster_spec,
+          mode=INDEPENDENT_WORKER,
+          rpc_layer=None)
+      self._join_threads(threads[WORKER])
+
+    # There is only one type of task and three such tasks.
+    self.assertEqual(len(self._worker_context), 1)
+    self.assertTrue(WORKER in self._worker_context)
+    self.assertEqual(len(self._worker_context[WORKER]), NUM_WORKERS)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(
+        self._worker_context[WORKER][0],
+        (_bytes_to_str(cluster_spec[WORKER][0]), NUM_WORKERS, True, True))
+    self.assertEqual(
+        self._worker_context[WORKER][1],
+        (_bytes_to_str(cluster_spec[WORKER][1]), NUM_WORKERS, False, True))
+    self.assertEqual(
+        self._worker_context[WORKER][2],
+        (_bytes_to_str(cluster_spec[WORKER][2]), NUM_WORKERS, False, True))
+
+    # Make sure each worker runs a std server.
+    self.assertEqual(len(self._std_servers), 1)
+    self.assertTrue(WORKER in self._std_servers)
+    self.assertEqual(len(self._std_servers[WORKER]), 3)
+    self.assertFalse(self._std_servers[WORKER][0].joined)
+    self.assertFalse(self._std_servers[WORKER][1].joined)
+    self.assertFalse(self._std_servers[WORKER][2].joined)
+
+  def testBetweenGraphStrategyProperties(self):
+    cluster_spec = self._create_cluster_spec(num_workers=NUM_WORKERS)
+    # Dumps properties of the strategy objects.
+    with test.mock.patch.object(distribute_coordinator, "_run_std_server",
+                                self._run_mock_std_server):
+      threads = self._run_multiple_coordinator_in_threads(
+          self._dump_strategy_property,
+          MockStrategy(between_graph=True, should_init=True),
+          cluster_spec,
+          mode=INDEPENDENT_WORKER,
+          rpc_layer=None)
+      self._join_threads(threads[WORKER])
+
+    # There is only one type of task and there three such tasks.
+    self.assertEqual(len(self._strategy_property), 1)
+    self.assertTrue(WORKER in self._strategy_property)
+    self.assertEqual(len(self._strategy_property[WORKER]), NUM_WORKERS)
+
+    # Check whether each task has the right properties of should_init,
+    # should_checkpoint and should_save_summary.
+    self.assertEqual(self._strategy_property[WORKER][0], (True, True, True))
+    self.assertEqual(self._strategy_property[WORKER][1], (True, False, False))
+    self.assertEqual(self._strategy_property[WORKER][2], (True, False, False))
+
+  def testInGraphContext(self):
+    cluster_spec = self._create_cluster_spec(num_workers=NUM_WORKERS)
+    # Dumps the task contexts and std server arguments.
+    with test.mock.patch.object(distribute_coordinator, "_run_std_server",
+                                self._run_mock_std_server):
+      threads = self._run_multiple_coordinator_in_threads(
+          self._dump_worker_context,
+          MockStrategy(between_graph=False),
+          cluster_spec,
+          mode=INDEPENDENT_WORKER,
+          rpc_layer=None)
+      self._join_threads(threads[WORKER])
+
+    # There is only a "None" task in the dumped task context.
+    self.assertEqual(len(self._worker_context), 1)
+    self.assertTrue("None" in self._worker_context)
+    self.assertEqual(len(self._worker_context["None"]), 1)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(
+        self._worker_context["None"][0],
+        (_bytes_to_str(cluster_spec[WORKER][0]), NUM_WORKERS, True, True))
+
+    # Make sure each worker runs a std server.
+    self.assertEqual(len(self._std_servers), 1)
+    self.assertTrue(WORKER in self._std_servers)
+    self.assertEqual(len(self._std_servers[WORKER]), 3)
+    self.assertFalse(self._std_servers[WORKER][0].joined)
+    self.assertTrue(self._std_servers[WORKER][1].joined)
+    self.assertTrue(self._std_servers[WORKER][2].joined)
+
+  def testInGraphContextWithEval(self):
+    # Adds a EVALUATOR job.
+    cluster_spec = self._create_cluster_spec(
+        num_workers=NUM_WORKERS, has_eval=True)
+
+    # Dumps the task contexts and std server arguments.
+    with test.mock.patch.object(distribute_coordinator, "_run_std_server",
+                                self._run_mock_std_server):
+      threads = self._run_multiple_coordinator_in_threads(
+          self._dump_worker_context,
+          MockStrategy(between_graph=False),
+          cluster_spec,
+          mode=INDEPENDENT_WORKER,
+          rpc_layer=None)
+      self._join_threads(threads[WORKER])
+      self._join_threads([threads[EVALUATOR][0]])
+
+    # There are one "None" task and one EVALUATOR task.
+    self.assertEqual(len(self._worker_context), 2)
+    self.assertTrue("None" in self._worker_context)
+    self.assertTrue(EVALUATOR in self._worker_context)
+    self.assertEqual(len(self._worker_context["None"]), 1)
+    self.assertEqual(len(self._worker_context[EVALUATOR]), 1)
+
+    # Check whether each task has the right master_target, num_workers, is_chief
+    # and distributed_mode.
+    self.assertEqual(self._worker_context["None"][0],
+                     (_bytes_to_str(cluster_spec[WORKER][0]), 3, True, True))
+    self.assertEqual(self._worker_context[EVALUATOR][0], ("", 3, True, False))
+
+    # Make sure each worker runs a std server.
+    self.assertEqual(len(self._std_servers), 1)
+    self.assertTrue(WORKER in self._std_servers)
+    self.assertEqual(len(self._std_servers[WORKER]), 3)
+    self.assertFalse(self._std_servers[WORKER][0].joined)
+    self.assertTrue(self._std_servers[WORKER][1].joined)
+    self.assertTrue(self._std_servers[WORKER][2].joined)
+
+  def testRunStdServerInGoogleEnvironment(self):
+    cluster_spec = {"worker": ["fake_worker"], "ps": ["localhost:0"]}
+    tf_config = {"cluster": cluster_spec, "environment": "google"}
+
+    joined = [False]
+
+    def _fake_sleep(_):
+      joined[0] = True
+      original_sys_exit(0)
+
+    def _thread_fn(cluster_spec):
+      distribute_coordinator.run_distribute_coordinator(
+          None,
+          MockStrategy(between_graph=True),
+          mode=INDEPENDENT_WORKER,
+          cluster_spec=cluster_spec,
+          task_type="ps",
+          task_id=0)
+
+    with test.mock.patch.dict(
+        "os.environ",
+        {"TF_CONFIG": json.dumps(tf_config)}), test.mock.patch.object(
+            time, "sleep", _fake_sleep):
+      t = threading.Thread(target=_thread_fn, args=(cluster_spec,))
+      t.start()
+      t.join()
+    self.assertTrue(joined[0])
+
+  def testRpcLayerEnvironmentVariable(self):
+    cluster_spec = {"worker": ["fake_worker"], "ps": ["fake_ps"]}
+    tf_config = {"cluster": cluster_spec, "rpc_layer": "cake"}
+
+    rpc_layer_from_coordinator = [None]
+
+    def _run_mock_server(cluster_spec=None,
+                         task_type=None,
+                         task_id=None,
+                         session_config=None,
+                         rpc_layer=None,
+                         environment=None):
+      del cluster_spec, task_type, task_id, session_config, environment
+      rpc_layer_from_coordinator[0] = rpc_layer
+      return MockServer()
+
+    with test.mock.patch.dict(
+        "os.environ",
+        {"TF_CONFIG": json.dumps(tf_config)}), test.mock.patch.object(
+            distribute_coordinator, "_run_std_server", _run_mock_server):
+      distribute_coordinator.run_distribute_coordinator(
+          None,
+          MockStrategy(between_graph=True),
+          mode=INDEPENDENT_WORKER,
+          cluster_spec=cluster_spec,
+          task_type="ps",
+          task_id=0)
+    self.assertEqual(rpc_layer_from_coordinator[0], "cake")
+
+
+class StrategyConfigureTest(test.TestCase):
+
+  def setUp(self):
+    self._device_filters = []
+    self._intra_op_parallelism_threads = None
+    self._inter_op_parallelism_threads = None
+    super(StrategyConfigureTest, self).setUp()
+
+  def _dump_device_filters(self, *args, **kwargs):
+    session_config = kwargs.get("session_config", None)
+    self._device_filters.extend(session_config.device_filters)
+    self._intra_op_parallelism_threads = (
+        session_config.intra_op_parallelism_threads)
+    self._inter_op_parallelism_threads = (
+        session_config.inter_op_parallelism_threads)
+    return MockServer()
+
+  def _worker_fn(self, strategy):
+    worker_context = distribute_coordinator_context.get_current_worker_context()
+    session_config = worker_context._session_config
+    self._device_filters.extend(session_config.device_filters)
+    self._intra_op_parallelism_threads = (
+        session_config.intra_op_parallelism_threads)
+    self._inter_op_parallelism_threads = (
+        session_config.inter_op_parallelism_threads)
+    return MockServer()
+
+  def test_session_config_in_std_server(self):
+    cluster_spec = {"worker": ["fake_worker"], "ps": ["fake_ps"]}
+    tf_config = {"cluster": cluster_spec}
+
+    with test.mock.patch.dict(
+        "os.environ",
+        {"TF_CONFIG": json.dumps(tf_config)}), test.mock.patch.object(
+            distribute_coordinator, "_run_std_server",
+            self._dump_device_filters):
+      distribute_coordinator.run_distribute_coordinator(
+          lambda _: None,
+          MockStrategy(between_graph=True),
+          mode=INDEPENDENT_WORKER,
+          cluster_spec=cluster_spec,
+          task_type="worker",
+          task_id=0)
+    self.assertEqual(self._intra_op_parallelism_threads, 1)
+    self.assertEqual(self._inter_op_parallelism_threads, 0)
+
+  def test_session_config_in_session_creator(self):
+    cluster_spec = {"worker": ["localhost:0"]}
+    tf_config = {"cluster": cluster_spec}
+
+    # Reset the saved Server state.
+    distribute_coordinator._thread_local = threading.local()  # pylint: disable=protected-access
+
+    with test.mock.patch.dict("os.environ",
+                              {"TF_CONFIG": json.dumps(tf_config)}):
+      distribute_coordinator.run_distribute_coordinator(
+          self._worker_fn,
+          MockStrategy(between_graph=True),
+          mode=INDEPENDENT_WORKER,
+          cluster_spec=cluster_spec,
+          task_type="worker",
+          task_id=0)
+    self.assertEqual(self._device_filters, ["/job:worker/task:0", "/job:ps"])
+    self.assertEqual(self._intra_op_parallelism_threads, 2)
+    self.assertEqual(self._inter_op_parallelism_threads, 0)
+
+  def test_eval_strategy_configure(self):
+    cluster_spec = {"evaluator": ["localhost:0"]}
+    tf_config = {"cluster": cluster_spec}
+
+    with test.mock.patch.dict("os.environ",
+                              {"TF_CONFIG": json.dumps(tf_config)}):
+      distribute_coordinator.run_distribute_coordinator(
+          lambda _: None,
+          MockStrategy(between_graph=False),
+          eval_fn=self._worker_fn,
+          eval_strategy=MockStrategy(between_graph=True),
+          mode=INDEPENDENT_WORKER,
+          cluster_spec=cluster_spec,
+          task_type="evaluator",
+          task_id=0)
+    self.assertEqual(self._device_filters, ["/job:somejob"])
+    self.assertEqual(self._intra_op_parallelism_threads, 0)
+    self.assertEqual(self._inter_op_parallelism_threads, 2)
+
+
+class RunStandardTensorflowServerTest(test.TestCase):
+
+  def test_std_server_arguments(self):
+    cs = {"worker": ["fake_worker"], "ps": ["fake_ps"]}
+    tf_config = {"cluster": cs, "task": {"type": "ps", "id": 0}}
+
+    def _mock_run_std_server(cluster_spec=None,
+                             task_type=None,
+                             task_id=None,
+                             session_config=None,
+                             rpc_layer=None):
+      self.assertEqual(cluster_spec.as_dict(), cs)
+      self.assertEqual(task_type, "ps")
+      self.assertEqual(task_id, 0)
+      self.assertEqual(session_config.experimental.collective_group_leader,
+                       "/job:worker/replica:0/task:0")
+      self.assertEqual(session_config.intra_op_parallelism_threads, 1)
+      self.assertEqual(rpc_layer, "grpc")
+
+      return MockServer()
+
+    with test.mock.patch.dict(
+        "os.environ",
+        {"TF_CONFIG": json.dumps(tf_config)}), test.mock.patch.object(
+            distribute_coordinator, "_run_std_server", _mock_run_std_server):
+      session_config = config_pb2.ConfigProto()
+      session_config.intra_op_parallelism_threads = 1
+      mock_server = distribute_coordinator.run_standard_tensorflow_server(
+          session_config)
+      self.assertTrue(mock_server.started)
+
+
+if __name__ == "__main__":
+  # TODO(yuefengz): find a smart way to terminate std server threads.
+  with test.mock.patch.object(sys, "exit", os._exit):
+    # Reduce `recovery_wait_secs` from 30 seconds so the test completes quickly.
+    orig_init = session_manager.SessionManager.__init__
+
+    def new_init(*args, **kwargs):
+      kwargs.pop("recovery_wait_secs", None)
+      kwargs["recovery_wait_secs"] = 0.5
+      orig_init(*args, **kwargs)
+
+    session_manager.SessionManager.__init__ = new_init
+
+    test.main()

tensorflow/python/distribute/distribute_lib_test.py

@@ -0,0 +1,610 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Test DistributionStrategy, ReplicaContext, and supporting APIs."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import input_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute.cluster_resolver import cluster_resolver as cluster_resolver_lib
+from tensorflow.python.distribute.v1 import input_lib as input_lib_v1
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
+from tensorflow.python.util import nest
+
+
+class _TestReplicaContext(distribute_lib.ReplicaContext):
+
+  def merge_call(self, fn, *args, **kwargs):
+    return kwargs["test_arg"]
+
+
+def _get_test_variable(name, synchronization, aggregation):
+  return {
+      "name": name,
+      "synchronization": synchronization,
+      "aggregation": aggregation
+  }
+
+
+def _test_input_fn(input_context):
+  del input_context
+  return dataset_ops.DatasetV2.from_tensors(1.).repeat()
+
+
+class _TestStrategy(distribute_lib.Strategy):
+
+  def __init__(self):
+    super(_TestStrategy, self).__init__(_TestExtended(self))
+
+
+class _TestExtended(distribute_lib.StrategyExtendedV1):
+
+  def __init__(self, distribute):
+    super(_TestExtended, self).__init__(distribute)
+    worker_device_pairs = [("", ["/device:CPU:0"])]
+    self._input_workers = input_lib.InputWorkers(worker_device_pairs)
+
+  def _call_for_each_replica(self, fn, args, kwargs):
+    with _TestReplicaContext(
+        self._container_strategy(), replica_id_in_sync_group=0):
+      return fn(*args, **kwargs)
+
+  def _create_variable(self, next_creator, **kwargs):
+    return _get_test_variable(kwargs["name"], kwargs["synchronization"],
+                              kwargs["aggregation"])
+
+  def _make_input_fn_iterator(
+      self,
+      input_fn,
+      replication_mode=distribute_lib.InputReplicationMode.PER_WORKER):
+    return input_lib_v1.InputFunctionIterator(input_fn, self._input_workers,
+                                              [distribute_lib.InputContext()],
+                                              self._container_strategy())
+
+  def _distribute_datasets_from_function(self, dataset_fn, options):
+    return dataset_fn(distribute_lib.InputContext())
+
+  def _local_results(self, value):
+    return (value,)
+
+  def _reduce_to(self, reduce_op, value, destinations, options):
+    del reduce_op, destinations, options
+    return value
+
+  def _experimental_run_steps_on_iterator(self, fn, iterator, iterations,
+                                          initial_loop_values=None):
+    # TODO(tomhennigan) This is missing many things (e.g. ctx.run_op).
+    ctx = input_lib.MultiStepContext()
+    for _ in range(iterations):
+      fn(ctx, iterator.get_next())
+    return ctx
+
+  def _update(self, var, fn, args, kwargs, group):
+    # The implementations of _update() and _update_non_slot() are identical
+    # except _update() passes `var` as the first argument to `fn()`.
+    return self._update_non_slot(var, fn, (var,) + tuple(args), kwargs, group)
+
+  def _update_non_slot(self, colocate_with, fn, args, kwargs, group):
+    del colocate_with
+    result = fn(*args, **kwargs)
+    if group:
+      return result
+    else:
+      return nest.map_structure(self._unwrap, result)
+
+  def _get_local_replica_id(self, replica_id_in_sync_group):
+    return replica_id_in_sync_group
+
+
+def _assert_in_default_state(t):
+  t.assertIs(distribute_lib._get_default_replica_context(),
+             distribute_lib.get_replica_context())
+  t.assertIs(None, distribute_lib.get_cross_replica_context())
+  t.assertFalse(distribute_lib.in_cross_replica_context())
+  t.assertIs(
+      distribute_lib._get_default_strategy(), distribute_lib.get_strategy())
+  t.assertFalse(distribute_lib.has_strategy())
+
+
+def _run_in_and_out_of_scope(unbound_test_method):
+  def wrapper(test_case):
+    dist = _TestStrategy()
+    # Running in the default (replica) scope should be supported.
+    _assert_in_default_state(test_case)
+    unbound_test_method(test_case, dist)
+    # As well as running in the strategy scope.
+    with dist.scope():
+      unbound_test_method(test_case, dist)
+    _assert_in_default_state(test_case)
+    # When run under a different strategy the test method should fail.
+    another_strategy = _TestStrategy()
+    msg = "Mixing different .*Strategy objects"
+    with test_case.assertRaisesRegex(RuntimeError, msg):
+      with another_strategy.scope():
+        unbound_test_method(test_case, dist)
+  return wrapper
+
+
+class TestStrategyTest(test.TestCase):
+
+  def testCallForEachReplica(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+
+    def run_fn():
+      replica_context = distribute_lib.get_replica_context()
+      self.assertIsNotNone(replica_context)
+      self.assertIs(None, distribute_lib.get_cross_replica_context())
+      self.assertFalse(distribute_lib.in_cross_replica_context())
+      self.assertTrue(distribute_lib.has_strategy())
+      self.assertIs(dist, distribute_lib.get_strategy())
+      self.assertEqual("foo", replica_context.merge_call(None, test_arg="foo"))
+      expected_value = _get_test_variable(
+          "bar", variable_scope.VariableSynchronization.AUTO,
+          variable_scope.VariableAggregation.NONE)
+      self.assertDictEqual(expected_value,
+                           variable_v1.VariableV1(1.0, name="bar"))
+
+    dist.extended.call_for_each_replica(run_fn)
+    with dist.scope():
+      dist.extended.call_for_each_replica(run_fn)
+    _assert_in_default_state(self)
+
+  def testScope(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    with dist.scope():
+      self.assertIs(None, distribute_lib.get_replica_context())
+      self.assertIs(dist, distribute_lib.get_cross_replica_context())
+      self.assertTrue(distribute_lib.in_cross_replica_context())
+      self.assertTrue(distribute_lib.has_strategy())
+      self.assertIs(dist, distribute_lib.get_strategy())
+      expected_value = _get_test_variable(
+          "baz", variable_scope.VariableSynchronization.AUTO,
+          variable_scope.VariableAggregation.NONE)
+      self.assertDictEqual(expected_value,
+                           variable_v1.VariableV1(1.0, name="baz"))
+    _assert_in_default_state(self)
+
+  def testScopeDeviceNestingError(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    # Open a device scope with dist.scope().
+    dist.extended._default_device = "/device:GPU:0"
+    scope = dist.scope()
+    scope.__enter__()
+    self.assertIs(dist, distribute_lib.get_strategy())
+    with ops.device("/device:CPU:0"):
+      with self.assertRaisesRegex(RuntimeError, "Device scope nesting error"):
+        scope.__exit__(None, None, None)
+    scope.__exit__(None, None, None)
+    _assert_in_default_state(self)
+
+  def testScopeVarCreatorNestingError(self):
+
+    def creator(next_creator, **kwargs):
+      return next_creator(**kwargs)
+
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    scope = dist.scope()
+    scope.__enter__()
+    self.assertIs(dist, distribute_lib.get_strategy())
+    with variable_scope.variable_creator_scope(creator):
+      with self.assertRaisesRegex(RuntimeError,
+                                  "Variable creator scope nesting error"):
+        scope.__exit__(None, None, None)
+    scope.__exit__(None, None, None)
+    _assert_in_default_state(self)
+
+  def testScopeVarScopeNestingError(self):
+    # We create a new graph here to simplify clean-up, since the error
+    # we are triggering happens in the middle of scope.__exit__() and
+    # leaves us in a weird state.
+    with ops.Graph().as_default():
+      _assert_in_default_state(self)
+      dist = _TestStrategy()
+      scope = dist.scope()
+      scope.__enter__()
+      self.assertIs(dist, distribute_lib.get_strategy())
+      with variable_scope.variable_scope("AA"):
+        with self.assertRaisesRegex(RuntimeError,
+                                    "Variable scope nesting error"):
+          scope.__exit__(None, None, None)
+    _assert_in_default_state(self)
+
+  def testSettingSynchronizationAndAggregation(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    with dist.scope():
+      expected_value = _get_test_variable(
+          "baz", variable_scope.VariableSynchronization.ON_WRITE,
+          variable_scope.VariableAggregation.MEAN)
+      self.assertDictEqual(
+          expected_value,
+          variable_v1.VariableV1(
+              1.0,
+              name="baz",
+              synchronization=variable_scope.VariableSynchronization.ON_WRITE,
+              aggregation=variable_scope.VariableAggregation.MEAN))
+    _assert_in_default_state(self)
+
+  def testSetStrategy(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    dist2 = _TestStrategy()
+    distribute_lib.experimental_set_strategy(dist)
+    self.assertIs(None, distribute_lib.get_replica_context())
+    self.assertIs(dist, distribute_lib.get_cross_replica_context())
+    self.assertTrue(distribute_lib.in_cross_replica_context())
+    self.assertTrue(distribute_lib.has_strategy())
+    self.assertIs(dist, distribute_lib.get_strategy())
+    expected_value = _get_test_variable(
+        "baz", variable_scope.VariableSynchronization.AUTO,
+        variable_scope.VariableAggregation.NONE)
+    self.assertDictEqual(expected_value,
+                         variable_v1.VariableV1(1.0, name="baz"))
+    distribute_lib.experimental_set_strategy(dist2)
+    self.assertIs(dist2, distribute_lib.get_strategy())
+    distribute_lib.experimental_set_strategy(None)
+    _assert_in_default_state(self)
+
+  def testSetStrategyInScope(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    with dist.scope():
+      with self.assertRaisesRegex(
+          RuntimeError,
+          "Must not be called inside a `tf.distribute.Strategy` scope"):
+        distribute_lib.experimental_set_strategy(_TestStrategy())
+      with self.assertRaisesRegex(
+          RuntimeError,
+          "Must not be called inside a `tf.distribute.Strategy` scope"):
+        distribute_lib.experimental_set_strategy(dist)
+      with self.assertRaisesRegex(
+          RuntimeError,
+          "Must not be called inside a `tf.distribute.Strategy` scope"):
+        distribute_lib.experimental_set_strategy(None)
+    _assert_in_default_state(self)
+
+  def testSameScopeNesting(self):
+    _assert_in_default_state(self)
+    dist = _TestStrategy()
+    scope_a = dist.scope()
+    with scope_a:
+      self.assertIs(dist, distribute_lib.get_strategy())
+      scope_b = dist.scope()
+      with scope_b:
+        self.assertIs(dist, distribute_lib.get_strategy())
+        with scope_a:
+          self.assertIs(dist, distribute_lib.get_strategy())
+        self.assertIs(dist, distribute_lib.get_strategy())
+      self.assertIs(dist, distribute_lib.get_strategy())
+      dist2 = _TestStrategy()
+      scope2 = dist2.scope()
+      with self.assertRaisesRegex(
+          RuntimeError, "Mixing different tf.distribute.Strategy objects"):
+        with scope2:
+          pass
+    _assert_in_default_state(self)
+    with scope_b:
+      self.assertIs(dist, distribute_lib.get_strategy())
+    _assert_in_default_state(self)
+
+  @_run_in_and_out_of_scope
+  def testMakeInputFnIterator(self, dist):
+    self.assertIsNotNone(dist.make_input_fn_iterator(_test_input_fn))
+
+  @_run_in_and_out_of_scope
+  def testReduce(self, dist):
+    x = constant_op.constant(1.)
+    x_r = dist.reduce(reduce_util.ReduceOp.MEAN, x, axis=None)
+    self.assertEqual(self.evaluate(x), self.evaluate(x_r))
+
+  def testReductions_acceptStringOps(self):
+    dist = _TestStrategy()
+    for op in ("mean", "MEAN", "sum", "SUM"):
+      x = constant_op.constant(1.)
+      y = constant_op.constant(1.)
+      x_r = dist.reduce(op, x, axis=None)
+      self.assertEqual(self.evaluate(x), self.evaluate(x_r))
+      x_r = dist.extended.reduce_to(op, x, "/CPU:0")
+      self.assertEqual(self.evaluate(x), self.evaluate(x_r))
+      x_r, y_r = dist.extended.batch_reduce_to(op,
+                                               ((x, "/CPU:0"), (y, "/CPU:0")))
+      self.assertEqual(self.evaluate(x), self.evaluate(x_r))
+      self.assertEqual(self.evaluate(y), self.evaluate(y_r))
+
+  @_run_in_and_out_of_scope
+  def testReduceMeanAxis(self, dist):
+    x = constant_op.constant([[1., 2.], [3., 4.]])
+    x_r = dist.reduce(reduce_util.ReduceOp.MEAN, x, axis=None)
+    self.assertAllEqual(self.evaluate(x), self.evaluate(x_r))
+    x_r = dist.reduce(reduce_util.ReduceOp.MEAN, x, axis=0)
+    self.assertAllEqual([2., 3.], self.evaluate(x_r))
+    x_r = dist.reduce(reduce_util.ReduceOp.MEAN, x, axis=(0, 1))
+    self.assertEqual(2.5, self.evaluate(x_r))
+
+  @_run_in_and_out_of_scope
+  def testReduceSumAxis(self, dist):
+    x = constant_op.constant([[1., 2.], [3., 4.]])
+    x_r = dist.reduce(reduce_util.ReduceOp.SUM, x, axis=None)
+    self.assertAllEqual(self.evaluate(x), self.evaluate(x_r))
+    x_r = dist.reduce(reduce_util.ReduceOp.SUM, x, axis=0)
+    self.assertAllEqual([4., 6.], self.evaluate(x_r))
+    x_r = dist.reduce(reduce_util.ReduceOp.SUM, x, axis=(0, 1))
+    self.assertEqual(10., self.evaluate(x_r))
+
+  @_run_in_and_out_of_scope
+  def testExperimentalRunStepsOnIterator(self, dist):
+    all_inputs = []
+    dataset = dataset_ops.Dataset.from_tensors(1.).repeat()
+    dist.extended.experimental_run_steps_on_iterator(
+        lambda _, inputs: all_inputs.append(self.evaluate(inputs)),
+        dataset_ops.make_one_shot_iterator(dataset))
+    self.assertEqual(all_inputs, [1.])
+
+  @_run_in_and_out_of_scope
+  def testReduceTo(self, dist):
+    x = constant_op.constant(1.)
+    x_r = dist.extended.reduce_to(reduce_util.ReduceOp.MEAN, x, "/CPU:0")
+    self.assertEqual(self.evaluate(x), self.evaluate(x_r))
+
+  @_run_in_and_out_of_scope
+  def testBatchReduceTo(self, dist):
+    x = constant_op.constant(1.)
+    y = constant_op.constant(1.)
+    x_r, y_r = dist.extended.batch_reduce_to(reduce_util.ReduceOp.MEAN,
+                                             ((x, "/CPU:0"), (y, "/CPU:0")))
+    self.assertEqual(self.evaluate(x), self.evaluate(x_r))
+    self.assertEqual(self.evaluate(y), self.evaluate(y_r))
+
+  @_run_in_and_out_of_scope
+  def testUpdate(self, dist):
+    with dist.scope():
+      v = variables.Variable(1.)
+    t = constant_op.constant(2.)
+
+    def assign_fn(vv, tt):
+      self.assertIs(vv, v)
+      self.assertIs(tt, t)
+    dist.extended.update(v, assign_fn, (t,))
+
+  @_run_in_and_out_of_scope
+  def testUpdateAutoGraph(self, dist):
+    with dist.scope():
+      v = variables.Variable(1.)
+    t = constant_op.constant(2.)
+
+    def assign_fn(unused_vv, unused_tt):
+      self.assertTrue(converter_testing.is_inside_generated_code())
+
+    @def_function.function  # AutoGraph is default-on only within tf.function
+    def test_fn():
+      dist.extended.update(v, assign_fn, (t,))
+
+    test_fn()
+
+  @_run_in_and_out_of_scope
+  def testUpdateNonSlot(self, dist):
+    t = constant_op.constant(2.)
+    update_calls = []
+    dist.extended.update_non_slot(t, lambda: update_calls.append(1))
+    self.assertEqual(len(update_calls), 1)
+
+  @_run_in_and_out_of_scope
+  def testUpdateNonSlotAutoGraph(self, dist):
+    t = constant_op.constant(2.)
+
+    def update_fn():
+      self.assertTrue(converter_testing.is_inside_generated_code())
+
+    @def_function.function  # AutoGraph is default-on only within tf.function
+    def test_fn():
+      dist.extended.update_non_slot(t, update_fn)
+
+    test_fn()
+
+  def testClusterResolverDefaultNotImplemented(self):
+    dist = _TestStrategy()
+    self.assertIsNone(dist.cluster_resolver)
+    base_cluster_spec = server_lib.ClusterSpec({
+        "ps": ["ps0:2222", "ps1:2222"],
+        "worker": ["worker0:2222", "worker1:2222", "worker2:2222"]
+    })
+    cluster_resolver = cluster_resolver_lib.SimpleClusterResolver(
+        base_cluster_spec)
+    dist.extended._cluster_resolver = cluster_resolver
+    self.assertIs(dist.cluster_resolver, cluster_resolver)
+
+
+# _TestStrategy2 is like _TestStrategy, except it doesn't change variable
+# creation.
+class _TestStrategy2(distribute_lib.Strategy):
+
+  def __init__(self):
+    super(_TestStrategy2, self).__init__(_TestExtended2(self))
+
+
+class _TestExtended2(_TestExtended):
+
+  def _create_variable(self, next_creator, **kwargs):
+    return next_creator(**kwargs)
+
+
+class DefaultDistributionStrategyTest(test.TestCase, parameterized.TestCase):
+
+  def testMergeCall(self):
+    _assert_in_default_state(self)
+
+    def merge_fn(dist, s):
+      self.assertIs(distribute_lib._get_default_strategy(), dist)
+      self.assertIs(None, distribute_lib.get_replica_context())
+      self.assertIs(dist, distribute_lib.get_cross_replica_context())
+      self.assertTrue(distribute_lib.in_cross_replica_context())
+      self.assertIs(dist, distribute_lib.get_strategy())
+      self.assertFalse(distribute_lib.has_strategy())
+      return "foo_" + s
+
+    replica_ctx = distribute_lib.get_replica_context()
+    self.assertIs(distribute_lib._get_default_replica_context(), replica_ctx)
+    self.assertEqual("foo_bar", replica_ctx.merge_call(merge_fn, args=("bar",)))
+    _assert_in_default_state(self)
+
+  def testMergeCallAutoGraph(self):
+    _assert_in_default_state(self)
+
+    def merge_fn(_, s):
+      self.assertTrue(converter_testing.is_inside_generated_code())
+      return s
+
+    @def_function.function  # AutoGraph is default-on only within tf.function
+    def test_fn():
+      replica_ctx = distribute_lib.get_replica_context()
+      replica_ctx.merge_call(merge_fn, args=("bar",))
+
+    test_fn()
+
+  def testScopeMostlyNoOp(self):
+    _assert_in_default_state(self)
+
+    test_strategy = _TestStrategy2()
+    with test_strategy.scope():
+      variable_v1.VariableV1(1.0, name="before")
+
+    default_strategy = distribute_lib._get_default_strategy()
+    scope = default_strategy.scope()
+    with scope:
+      _assert_in_default_state(self)
+
+      with test_strategy.scope():
+        with self.assertRaisesRegex(
+            RuntimeError, "Mixing different tf.distribute.Strategy objects"):
+          variable_v1.VariableV1(1.0, name="error")
+
+      with scope:
+        _assert_in_default_state(self)
+
+        with test_strategy.scope():
+          with self.assertRaisesRegex(
+              RuntimeError, "Mixing different tf.distribute.Strategy objects"):
+            variable_v1.VariableV1(1.0, name="also_error")
+
+      _assert_in_default_state(self)
+
+    _assert_in_default_state(self)
+    with test_strategy.scope():
+      variable_v1.VariableV1(1.0, name="after")
+
+  def testExperimentalRunV2(self):
+    default_strategy = distribute_lib._get_default_strategy()
+    dataset = dataset_ops.Dataset.range(10).batch(2)
+    iterator = default_strategy.extended._make_dataset_iterator(dataset)
+    next_val = iterator.get_next()
+
+    def train_step(input_data):
+      return input_data
+
+    for _ in range(2):
+      default_strategy.run(train_step, args=(next_val,))
+
+  @combinations.generate(combinations.combine(mode=["graph", "eager"]))
+  def testDistributedDatasets(self):
+    default_strategy = distribute_lib._get_default_strategy()
+    if context.executing_eagerly():
+      dataset_fn = lambda _: dataset_ops.DatasetV2.range(10).batch(2)
+      dist_dataset = default_strategy.experimental_distribute_dataset(
+          dataset_fn(distribute_lib.InputContext()))
+      next_val = next(iter(dist_dataset))
+    else:
+      dataset_fn = lambda _: dataset_ops.DatasetV1.range(10).batch(2)
+      dist_dataset = default_strategy.experimental_distribute_dataset(
+          dataset_fn(distribute_lib.InputContext()))
+      iterator = dist_dataset.make_initializable_iterator()
+      self.evaluate(iterator.initializer)
+      next_val = iterator.get_next()
+    self.assertAllEqual([0, 1], self.evaluate(next_val))
+
+  @combinations.generate(combinations.combine(mode=["graph", "eager"]))
+  def testDistributedDatasetsFromFunction(self):
+    default_strategy = distribute_lib._get_default_strategy()
+    if context.executing_eagerly():
+      dataset_fn = lambda _: dataset_ops.DatasetV2.range(10).batch(2)
+      dist_dataset_from_func = \
+          default_strategy.distribute_datasets_from_function(
+              dataset_fn)
+      next_val = next(iter(dist_dataset_from_func))
+      self.assertAllEqual([0, 1], self.evaluate(next_val))
+    else:
+      dataset_fn = lambda _: dataset_ops.DatasetV2.range(10).batch(2)
+      dist_dataset_from_func = \
+        default_strategy.distribute_datasets_from_function(
+            dataset_fn)
+      dataset_ops.make_initializable_iterator(dist_dataset_from_func)
+
+  @combinations.generate(combinations.combine(tf_api_version=1))
+  def testV1(self):
+    self.assertIsInstance(
+        distribute_lib.get_strategy(), distribute_lib.StrategyV1)
+
+  @combinations.generate(combinations.combine(tf_api_version=2))
+  def testV2(self):
+    self.assertIsInstance(
+        distribute_lib.get_strategy(), distribute_lib.Strategy)
+
+
+class InputContextTest(test.TestCase):
+
+  def testProperties(self):
+    input_context = distribute_lib.InputContext(
+        num_input_pipelines=2, input_pipeline_id=1, num_replicas_in_sync=6)
+    self.assertEqual(6, input_context.num_replicas_in_sync)
+    self.assertEqual(1, input_context.input_pipeline_id)
+    self.assertEqual(2, input_context.num_input_pipelines)
+
+  def testPerReplicaBatchSize(self):
+    input_context = distribute_lib.InputContext(
+        num_input_pipelines=2, input_pipeline_id=1, num_replicas_in_sync=6)
+    self.assertEqual(2, input_context.get_per_replica_batch_size(12))
+    with self.assertRaises(ValueError):
+      input_context.get_per_replica_batch_size(13)
+
+  def testStr(self):
+    input_context = distribute_lib.InputContext(
+        num_input_pipelines=1, input_pipeline_id=0, num_replicas_in_sync=42)
+    self.assertEqual(
+        "tf.distribute.InputContext(input pipeline id 0, total: 1)",
+        str(input_context))
+    input_context = distribute_lib.InputContext(
+        num_input_pipelines=3, input_pipeline_id=1, num_replicas_in_sync=42)
+    self.assertEqual(
+        "tf.distribute.InputContext(input pipeline id 1, total: 3)",
+        str(input_context))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/distribute_utils.py

@@ -0,0 +1,496 @@
+# 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.
+# ==============================================================================
+"""Class implementing utilities used by tf.distribute.Strategy."""
+
+from collections import abc
+import contextlib
+import threading
+
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import tpu_values as tpu_values_lib
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.distribute.reduce_util import ReduceOp
+from tensorflow.python.eager import context
+from tensorflow.python.eager import record
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope as vs
+from tensorflow.python.ops.losses import losses_impl
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["distribute.get_loss_reduction"])
+def get_loss_reduction():
+  """`tf.distribute.ReduceOp` corresponding to the last loss reduction.
+
+  Returns:
+    `tf.distribute.ReduceOp` corresponding to the last loss reduction for
+    estimator and v1 optimizer use case. `tf.distribute.ReduceOp.SUM` otherwise.
+  """
+  if not distribute_lib.get_strategy()._scale_loss_for_estimator:  # pylint: disable=protected-access
+    # If we are not in Estimator context then return 'SUM'. We do not need to
+    # scale loss in the optimizer.
+    return ReduceOp.SUM
+  last_reduction = ops.get_default_graph()._last_loss_reduction  # pylint: disable=protected-access
+  if (last_reduction == losses_impl.Reduction.SUM or
+      last_reduction == "sum"):  # Check for tf.keras.losses.Reduction.SUM
+    return ReduceOp.SUM
+  return ReduceOp.MEAN
+
+
+def regroup(values, wrap_class=values_lib.PerReplica, always_wrap=False):
+  """Makes a nest per-replica into a nest of PerReplica/Mirrored values.
+
+  Args:
+    values: Values to regroup
+    wrap_class: Class that `values` be wrapped in.
+    always_wrap: Always wrap the `values` in `wrap_class` even if the values
+        are the same except for DistributeVariable.
+  Returns:
+    Wrapped `values`.
+  """
+  v0 = values[0]
+
+  if isinstance(v0, list):
+    for v in values[1:]:
+      assert isinstance(v, list)
+      assert len(v) == len(v0), ("len(v) == %d, len(v0) == %d, v: %s, v0: %s" %
+                                 (len(v), len(v0), v, v0))
+    return [
+        regroup(tuple(v[i] for v in values), wrap_class, always_wrap)
+        for i in range(len(v0))
+    ]
+
+  if isinstance(v0, tuple):
+    for v in values[1:]:
+      assert isinstance(v, tuple)
+      assert len(v) == len(v0), ("Values to regroup had different lengths: "
+                                 f"len(v) == {len(v)}, len(v0) == {len(v0)}, "
+                                 f"v: {v}, v0: {v0}")
+    regrouped_tuple = tuple(
+        regroup(tuple(v[i] for v in values), wrap_class, always_wrap)
+        for i in range(len(v0)))
+    if hasattr(v0, "_fields"):
+      # This tuple is in fact a namedtuple! Create a new namedtuple instance
+      # and initialize it with the regrouped values:
+      assert hasattr(v0, "_make")
+      return v0._make(regrouped_tuple)
+    else:
+      return regrouped_tuple
+
+  if isinstance(v0, abc.Mapping):
+    v0keys = v0.keys()
+    for v in values[1:]:
+      assert isinstance(v, abc.Mapping), ("v[0]: %r  v[i]: %r" % (v0, v))
+      assert set(v.keys()) == set(v0keys), ("v[0].keys: %s  v[i].keys: %s" %
+                                            (set(v0keys), set(v.keys())))
+    # Use the actual type in case it is a class inherited from a dict.
+    return type(v0)({
+        key: regroup(tuple(v[key] for v in values),
+                     wrap_class, always_wrap)
+        for key in v0keys
+    })
+
+  # If exactly the same object across all devices, return it unwrapped.
+  same_id = True
+  for v in values[1:]:
+    if v is not v0:
+      same_id = False
+      break
+  # Consider three cases where same_id is true:
+  # * If v0 is a DistributedVariable (a MirroredVariable or
+  #   SyncOnReadVariable, and same_id means it is the same across all
+  #   devices), we want to return it. We check DistributedVariable
+  #   specifically since it can look like it has a
+  #   _distributed_container member since its members do.
+  if same_id and isinstance(v0, values_lib.DistributedVariable):
+    return v0
+  # * If v0 is a member of a distributed variable, in which case
+  #   value_container(v0) is not v0 itself, we want to
+  #   return the DistributedVariable that contains it using the
+  #   _distributed_container logic below. This case can trigger
+  #   same_id when there is only one device.
+  # * In any other situation, same_id means we return v0 unless `always_wrap` is
+  #   true.
+  if same_id and not always_wrap and value_container(v0) is v0:
+    return v0
+
+  # Detect the case where each device has a parallel component of the
+  # same MirroredVariable (or SyncOnReadVariable). In this case we
+  # want to return the containing MirroredVariable, after a bunch of
+  # sanity checking. In particular, each component should have the
+  # same container, and the devices of the variables should match the
+  # keys of the per-replica dictionary. For _UnreadVariables, use the wrap_class
+  # path, which calls tf.identity on them.
+  if (not isinstance(v0, resource_variable_ops._UnreadVariable) and  # pylint: disable=protected-access
+      value_container(v0) is not v0):
+    # pylint: disable=protected-access
+    assert not isinstance(v0, values_lib.MirroredVariable), (
+        "ids = %s, values = %s" % ([id(v) for v in values], values))
+    distributed_container = value_container(v0)
+    assert distributed_container is not None
+    for v in values[1:]:
+      assert distributed_container is value_container(v)
+    return distributed_container
+  # pylint: enable=protected-access
+
+  return wrap_class(values)
+
+
+def select_replica(replica_id, structured):
+  """Specialize a nest of regular & per-replica values for one replica."""
+
+  def _get(x):
+    # `DistributedValues` would be sliced according to replica unless it is a
+    # `DistributedVariable` because `DistributedVariable` can be handled
+    # directly in the replica context.
+    if (isinstance(x, values_lib.DistributedVariable) or
+        not isinstance(x, values_lib.DistributedValues)):
+      return x
+    else:
+      return x.values[replica_id]
+
+  return nest.map_structure(_get, structured)
+
+
+def select_replica_mirrored(replica_id, structured):
+  """Specialize a nest of regular & mirrored values for one replica."""
+  assert_mirrored(structured)
+  return select_replica(replica_id, structured)
+
+
+def assert_mirrored(structured):
+  """Raises if the structured is not composed of mirrored or regular values."""
+
+  def _assert_mirrored(x):
+    if isinstance(x, values_lib.DistributedValues) and not is_mirrored(x):
+      raise TypeError(
+          "Expected value to be mirrored across replicas: %s in %s." %
+          (x, structured))
+
+  nest.map_structure(_assert_mirrored, structured)
+
+
+def update_regroup(extended, updates, group):
+  """Regroup for an update, with dependencies to ensure all updates execute."""
+  if not group:
+    regrouped = regroup(updates, values_lib.Mirrored)
+    return nest.map_structure(extended._local_results, regrouped)  # pylint: disable=protected-access
+
+  def _make_grouped_mirrored(values):
+    """Convert per-replica list `values` into Mirrored type with grouping."""
+    if len(values) == 1:
+      return values_lib.Mirrored(values)
+
+    # Make sure we run all updates. Without this, something like
+    # session.run(extended.update(...)) may only update one replica.
+    g = control_flow_ops.group(values)
+
+    # If values is just ops, the grouping is enough. Everything in values
+    # should have the same type, since we expect every replica to be performing
+    # the same computation.
+    if not all(tensor_util.is_tf_type(v) for v in values):
+      return g
+
+    # Otherwise we need tensors with the same values as `values`, but
+    # that have a dependency on `g`.
+    with_dep = []
+    for v in values:
+      with ops.device(v.device), ops.control_dependencies([g]):
+        with_dep.append(array_ops.identity(v))
+
+    return values_lib.Mirrored(with_dep)
+
+  return regroup(updates, _make_grouped_mirrored)
+
+
+def value_container(val):
+  """Returns the container that this per-replica `value` belongs to.
+
+  Args:
+    val: A value returned by `call_for_each_replica()` or a variable created in
+      `scope()`.
+
+  Returns:
+    A container that `value` belongs to.
+    If value does not belong to any container (including the case of
+    container having been destroyed), returns the value itself.
+  """
+  # DistributedVariable has _distributed_container defined but we don't want to
+  # return it.
+  container = None
+  if not isinstance(val, values_lib.DistributedVariable):
+    if hasattr(val, "_distributed_container"):
+      container = val._distributed_container()  # pylint: disable=protected-access
+    elif (isinstance(val, composite_tensor.CompositeTensor) and
+          hasattr(val, "handle") and
+          hasattr(val.handle, "_distributed_container")):
+      # For ResourceVariables, the _distributed_container attribute
+      # is added to their handle tensors.
+      container = val.handle._distributed_container()  # pylint: disable=protected-access
+  return container if container is not None else val
+
+
+def is_distributed_variable(v):
+  """Determine if a variable is ds variable or TPU mirrored variable."""
+  return getattr(v, "is_distributed_variable", False)
+
+
+def is_distributed_table(v):
+  """Determine if an object is a DistributedTable."""
+  return getattr(v, "is_distributed_table", False)
+
+
+def _validate_colocate_extended(v, extended):
+  variable_strategy = v._distribute_strategy  # pylint: disable=protected-access
+  if not variable_strategy or variable_strategy.extended is not extended:
+    raise ValueError(
+        "`colocate_vars_with` must only be passed a variable created in this "
+        "tf.distribute.Strategy.scope(), not %s created in scope: %s" %
+        (v, variable_strategy))
+
+
+def validate_colocate_distributed_variable(v, extended):
+  if not isinstance(v, values_lib.DistributedVariable):
+    raise ValueError(
+        "`colocate_vars_with` must only be passed a variable created in this "
+        "tf.distribute.Strategy.scope(), not: %r" % (v,))
+  _validate_colocate_extended(v, extended)
+
+
+def validate_colocate(v, extended):
+  if not hasattr(v, "_distribute_strategy"):
+    raise ValueError(
+        "`colocate_vars_with` must only be passed a variable created in this "
+        "tf.distribute.Strategy.scope(), not: %r" % (v,))
+  _validate_colocate_extended(v, extended)
+
+
+# Variable creation function for sync strategies.
+def _validate_synchronization(kwargs):
+  """Validate that given synchronization value is valid."""
+  synchronization = kwargs.get("synchronization",
+                               vs.VariableSynchronization.AUTO)
+  if synchronization == vs.VariableSynchronization.NONE:
+    raise ValueError(
+        "`NONE` variable synchronization mode is not supported with "
+        "tf.distribute strategy. Please change the `synchronization` for "
+        "variable: " + str(kwargs["name"]))
+  if synchronization not in (vs.VariableSynchronization.ON_READ,
+                             vs.VariableSynchronization.ON_WRITE,
+                             vs.VariableSynchronization.AUTO):
+    raise ValueError(
+        "Invalid variable synchronization mode: %s for variable: %s" %
+        (synchronization, kwargs["name"]))
+  if synchronization == vs.VariableSynchronization.AUTO:
+    return vs.VariableSynchronization.ON_WRITE
+  return synchronization
+
+
+def _validate_aggregation(kwargs):
+  aggregation = kwargs.get("aggregation", vs.VariableAggregation.NONE)
+
+  if aggregation not in (vs.VariableAggregation.NONE,
+                         vs.VariableAggregation.SUM,
+                         vs.VariableAggregation.MEAN,
+                         vs.VariableAggregation.ONLY_FIRST_REPLICA):
+    raise ValueError("Invalid variable aggregation mode: %s for variable: %s" %
+                     (aggregation, kwargs["name"]))
+  return aggregation
+
+
+def create_mirrored_variable(strategy, real_mirrored_creator, class_mapping,
+                             policy_mapping, **kwargs):
+  """Create distributed variables with given synchronization and aggregation."""
+  # Figure out what collections this variable should be added to.
+  # We'll add the MirroredVariable to those collections instead.
+
+  if kwargs.pop("experimental_batch_initialization", None):
+    variable_class_key = "LazyVariableClass"
+  else:
+    variable_class_key = "VariableClass"
+
+  var_collections = kwargs.pop("collections", None)
+  if var_collections is None:
+    var_collections = [ops.GraphKeys.GLOBAL_VARIABLES]
+  kwargs["collections"] = []
+
+  synchronization = _validate_synchronization(kwargs)
+  # Update synchronization in kwargs in case it's AUTO, which is converted to
+  # ON_WRITE.
+  kwargs["synchronization"] = synchronization
+  aggregation = _validate_aggregation(kwargs)
+  use_var_policy = getattr(strategy.extended, "_use_var_policy", False)
+
+  # Ignore user-specified caching device, not needed for mirrored variables.
+  kwargs.pop("caching_device", None)
+
+  # TODO(josh11b,apassos): It would be better if variable initialization
+  # was never recorded on the tape instead of having to do this manually
+  # here.
+  with record.stop_recording():
+    value_list = real_mirrored_creator(**kwargs)
+    # MirroredVariable is recreated during saved_model loading, and its
+    # component variables (value_list) will have None initializer. We
+    # set their initializers to no_op so that consumer like
+    # `global_variables_initializer` wouldn't complain, as it groups all
+    # variables' initializers thus all variables have to have initializers.
+    for v in value_list:
+      # pylint:disable=protected-access
+      if hasattr(v, "_initializer_op") and v._initializer_op is None:
+        v._initializer_op = control_flow_ops.no_op()
+      # pylint:enable=protected-access
+    if use_var_policy:
+      var_policy_cls = policy_mapping.get(synchronization)
+      var_policy = var_policy_cls(aggregation=aggregation)
+      var_cls = class_mapping.get(variable_class_key)
+      result = var_cls(strategy, value_list, aggregation, var_policy=var_policy)
+    else:
+      var_cls = class_mapping.get(synchronization)
+      result = var_cls(strategy, value_list, aggregation)
+
+  # Add the wrapped variable to the requested collections.
+  # The handling of eager mode and the global step matches
+  # ResourceVariable._init_from_args().
+  if not context.executing_eagerly():
+    g = ops.get_default_graph()
+    # If "trainable" is True, next_creator() will add the member variables
+    # to the TRAINABLE_VARIABLES collection, so we manually remove
+    # them and replace with the MirroredVariable. We can't set
+    # "trainable" to False for next_creator() since that causes functions
+    # like implicit_gradients to skip those variables.
+    if kwargs.get("trainable", True):
+      var_collections.append(ops.GraphKeys.TRAINABLE_VARIABLES)
+      l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES)
+      for value in value_list:
+        for i, trainable_variable in enumerate(l):
+          if value is trainable_variable:
+            del l[i]
+            break
+
+    g.add_to_collections(var_collections, result)
+  elif ops.GraphKeys.GLOBAL_STEP in var_collections:
+    ops.add_to_collections(ops.GraphKeys.GLOBAL_STEP, result)
+
+  return result
+
+
+# Utility functions
+# Return True if the Value is Mirrored or the Variable is replicated and kept in
+# sync.
+def is_mirrored(val):
+  return (getattr(val, "_is_mirrored", lambda: False))()
+
+
+def is_sync_on_read(val):
+  return not is_mirrored(val)
+
+
+class CachingScopeLocal(threading.local):
+  """Class for maintaining thread local state for caching scope."""
+
+  def __init__(self):
+    super(CachingScopeLocal, self).__init__()
+    self.new_cache_scope_count = 0
+    self.cache_scope_exited_count = 0
+
+  def enter_scope(self):
+    self.new_cache_scope_count += 1
+
+  def exit_scope(self):
+    self.cache_scope_exited_count += 1
+
+  def in_caching_scope(self):
+    return self.new_cache_scope_count > self.cache_scope_exited_count
+
+
+caching_scope_local = CachingScopeLocal()
+
+
+@contextlib.contextmanager
+def cache_variable_reads():
+  """Scope for caching variable reads for AggregatingVariable.
+
+  The variable reads for AggregatingVariable inside this scope are cached. i.e.
+  the first read of variable reads the value from possibly remote handle, but
+  subsequent reads are returned using local cached value.
+
+  For example:
+  strategy = ParameterServerStrategy...
+  with strategy.scope():
+    # Variable v is of AggregatingVariable type with actual variable residing
+    # on PS.
+    v = tf.Variable(1.0)
+
+  with distribute_utils.cache_variable_reads():
+    v.read_value()  # Reads value 1.0
+    v.assign(constant_op.constant(5.0))  # v changes to 5.0
+    t1 = v.read_value()
+    t2 = v.read_value()  # Both t1 & t2 return cached value 1.0 from local CPU.
+
+  Notes about cache_variable_reads scope:
+  1. Nesting of scope cache_variable_reads() is not supported
+  2. And when caching scope is enabled, the thread enabling the cache and
+    mirrored_run._MirroredReplicaThread threads spawned from it will have
+    caching enabled.
+
+  Yields:
+    A context for caching variables.
+  """
+
+  try:
+    if caching_scope_local.in_caching_scope():
+      # There is nested cache scope, which is not supported.
+      raise ValueError("cache_variable_reads scope cannot be nested")
+    caching_scope_local.enter_scope()
+    yield
+  finally:
+    caching_scope_local.exit_scope()
+
+
+# The following mapping indicates the policy that you must use for a given
+# variable `synchronization` and `aggregation` pair.
+# OnWritePolicy is used for:
+# (synchronization=Auto, aggregation=NONE,SUM,MEAN,ONLY_FIRST_REPLICA)
+# (synchronization=ON_WRITE, aggregation=NONE,SUM,MEAN,ONLY_FIRST_REPLICA)
+# OnReadPolicy is used for:
+# (synchronization=ON_READ, aggregation=NONE,SUM,MEAN,ONLY_FIRST_REPLICA)
+VARIABLE_POLICY_MAPPING = {
+    vs.VariableSynchronization.ON_WRITE: values_lib.OnWritePolicy,
+    vs.VariableSynchronization.ON_READ: values_lib.OnReadPolicy,
+}
+
+VARIABLE_CLASS_MAPPING = {
+    "VariableClass": values_lib.DistributedVariable,
+    vs.VariableSynchronization.ON_WRITE: values_lib.MirroredVariable,
+    vs.VariableSynchronization.ON_READ: values_lib.SyncOnReadVariable,
+}
+
+TPU_VARIABLE_POLICY_MAPPING = {
+    vs.VariableSynchronization.ON_WRITE: tpu_values_lib.TPUOnWritePolicy,
+    vs.VariableSynchronization.ON_READ: tpu_values_lib.TPUOnReadPolicy,
+}
+
+TPU_VARIABLE_CLASS_MAPPING = {
+    "VariableClass": tpu_values_lib.TPUDistributedVariable,
+    "LazyVariableClass": tpu_values_lib.TPULazyDistributedVariable,
+    vs.VariableSynchronization.ON_WRITE: tpu_values_lib.TPUMirroredVariable,
+    vs.VariableSynchronization.ON_READ: tpu_values_lib.TPUSyncOnReadVariable,
+}

tensorflow/python/distribute/distribute_utils_test.py

@@ -0,0 +1,197 @@
+# 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.
+# ==============================================================================
+"""Tests for utility functions in distribute_utils."""
+
+import collections
+import collections.abc
+
+from absl.testing import parameterized
+import wrapt
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import values
+from tensorflow.python.eager import test
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+
+
+def _nested_value(d):
+  return ("a" + d, ["b" + d, {"c": "d" + d, "e": "f" + d}, "g" + d], "h" + d)
+
+
+class RegroupAndSelectDeviceTest(test.TestCase, parameterized.TestCase):
+
+  def _is_per_replica(self, result, expected, klass=values.PerReplica):
+    self.assertIsInstance(result, klass)
+    for i, exp in enumerate(expected):
+      self.assertEqual(exp, result.values[i])
+
+  def testNested(self):
+    result = distribute_utils.regroup((_nested_value("1"), _nested_value("2")))
+    self.assertIsInstance(result, tuple)
+    self.assertLen(result, 3)
+    self._is_per_replica(result[0], ["a1", "a2"])
+    self._is_per_replica(result[2], ["h1", "h2"])
+
+    self.assertIsInstance(result[1], list)
+    self.assertLen(result[1], 3)
+    self._is_per_replica(result[1][0], ["b1", "b2"])
+    self._is_per_replica(result[1][2], ["g1", "g2"])
+
+    self.assertIsInstance(result[1][1], dict)
+    self.assertEqual(set(["c", "e"]), set(result[1][1].keys()))
+    self._is_per_replica(result[1][1]["c"], ["d1", "d2"])
+    self._is_per_replica(result[1][1]["e"], ["f1", "f2"])
+
+    # Also test that we can undo the merge using select_replica()
+    self.assertEqual(_nested_value("1"),
+                     distribute_utils.select_replica(0, result))
+    self.assertEqual(_nested_value("2"),
+                     distribute_utils.select_replica(1, result))
+    # select_device_mirrored() should fail due to non-mirrored values
+    with self.assertRaises(TypeError):
+      distribute_utils.select_replica_mirrored(0, result)
+    with self.assertRaises(TypeError):
+      distribute_utils.select_replica_mirrored(1, result)
+
+  def testRegroupKeepsDictBasedClass(self):
+    class DictBasedClass(dict):
+      """Dummy class inherited from a dict."""
+
+    result = distribute_utils.regroup(
+        (DictBasedClass(a="a1", b="b1"), DictBasedClass(a="a2", b="b2")))
+    self.assertIsInstance(result, DictBasedClass)
+    self._is_per_replica(result["a"], ["a1", "a2"])
+    self._is_per_replica(result["b"], ["b1", "b2"])
+
+  def testRegroupCollectionsMapping(self):
+
+    class CollectionsMappingBasedClass(collections.abc.Mapping):
+      """Class inherited from collections.abc.Mapping."""
+
+      def __init__(self, *args, **kwargs):
+        self._d = dict(*args, **kwargs)
+
+      def __getitem__(self, key):
+        return self._d.__getitem__(key)
+
+      def __iter__(self):
+        return iter(self._d)
+
+      def __len__(self):
+        return len(self._d)
+
+    result = distribute_utils.regroup(
+        (CollectionsMappingBasedClass(a="a1", b="b1"),
+         CollectionsMappingBasedClass(a="a2", b="b2")))
+    self.assertIsInstance(result, CollectionsMappingBasedClass)
+    self._is_per_replica(result["a"], ["a1", "a2"])
+    self._is_per_replica(result["b"], ["b1", "b2"])
+
+  def testWrapClass(self):
+    # Normally a mirrored value would be the same across devices, but
+    # for a test it is convenient to be able to tell the values apart.
+    result = distribute_utils.regroup((_nested_value("1"), _nested_value("2")),
+                                      values.Mirrored)
+    self.assertIsInstance(result, tuple)
+    self.assertLen(result, 3)
+    self._is_per_replica(result[0], ["a1", "a2"], values.Mirrored)
+    self._is_per_replica(result[2], ["h1", "h2"], values.Mirrored)
+
+    self.assertIsInstance(result[1], list)
+    self.assertLen(result[1], 3)
+    self._is_per_replica(result[1][0], ["b1", "b2"], values.Mirrored)
+    self._is_per_replica(result[1][2], ["g1", "g2"], values.Mirrored)
+
+    self.assertIsInstance(result[1][1], dict)
+    self.assertEqual(set(["c", "e"]), set(result[1][1].keys()))
+    self._is_per_replica(result[1][1]["c"], ["d1", "d2"], values.Mirrored)
+    self._is_per_replica(result[1][1]["e"], ["f1", "f2"], values.Mirrored)
+
+    # Also test that we can undo the merge using select_replica()
+    self.assertEqual(_nested_value("1"),
+                     distribute_utils.select_replica(0, result))
+    self.assertEqual(_nested_value("2"),
+                     distribute_utils.select_replica(1, result))
+    # Values are marked as mirrored, so select_device_mirrored() is allowed.
+    self.assertEqual(_nested_value("1"),
+                     distribute_utils.select_replica_mirrored(0, result))
+    self.assertEqual(_nested_value("2"),
+                     distribute_utils.select_replica_mirrored(1, result))
+
+  def testWrapAListOfTwoTuples(self):
+    result = distribute_utils.regroup([("1", "2"), ("3", "4")])
+    self.assertIsInstance(result, tuple)
+    self.assertLen(result, 2)
+    self._is_per_replica(result[0], ("1", "3"), values.PerReplica)
+    self._is_per_replica(result[1], ("2", "4"), values.PerReplica)
+
+  @combinations.generate(
+      combinations.combine(
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.mirrored_strategy_with_one_cpu,
+          ],
+          mode=["graph", "eager"],
+      ))
+  def testMirroredContainer(self, distribution):
+    with distribution.scope():
+      v = variable_v1.VariableV1(
+          1., aggregation=variable_scope.VariableAggregation.SUM)
+    self.assertTrue(distribute_utils.is_distributed_variable(v))
+    self.assertTrue(distribute_utils.is_distributed_variable(
+        distribute_utils.regroup(v.values)))
+
+  def testSameId(self):
+    foo = object()
+    result = distribute_utils.regroup((("a", foo), ("b", foo)))
+    self.assertIsInstance(result, tuple)
+    self.assertLen(result, 2)
+    self._is_per_replica(result[0], ["a", "b"])
+    self.assertIs(foo, result[1])
+
+    # Test select_replica(), should undo the merge done by regroup().
+    result_0 = distribute_utils.select_replica(0, result)
+    self.assertIsInstance(result_0, tuple)
+    self.assertLen(result_0, 2)
+    self.assertEqual("a", result_0[0])
+    self.assertIs(foo, result_0[1])
+    result_1 = distribute_utils.select_replica(1, result)
+    self.assertIsInstance(result_1, tuple)
+    self.assertLen(result_1, 2)
+    self.assertEqual("b", result_1[0])
+    self.assertIs(foo, result_1[1])
+
+  def testOneDevice(self):
+    result = distribute_utils.regroup((_nested_value("1"),))
+    # On one device regroup() and select_replica() are basically identity.
+    self.assertEqual(_nested_value("1"), result)
+    self.assertEqual(_nested_value("1"),
+                     distribute_utils.select_replica(0, result))
+
+  def testWrappedNamedTuple(self):
+    Point = collections.namedtuple("Point", ["x", "y"])
+    point1 = Point(x=0, y=2)
+    point2 = Point(x=1, y=3)
+    wrapped1 = wrapt.ObjectProxy(point1)
+    wrapped2 = wrapt.ObjectProxy(point2)
+    result = distribute_utils.regroup([wrapped1, wrapped2])
+    self.assertEqual(result.x.values, (0, 1))
+    self.assertEqual(result.y.values, (2, 3))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/distributed_table_test.py

@@ -0,0 +1,677 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for distributed_table."""
+
+import copy
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import parameter_server_strategy_v2
+from tensorflow.python.distribute import ps_values
+from tensorflow.python.distribute.coordinator import cluster_coordinator as coordinator_lib
+from tensorflow.python.distribute.coordinator import coordinator_context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.module import module
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.saved_model import load as tf_load
+from tensorflow.python.saved_model import save as tf_save
+
+
+source_combination = combinations.combine(source=["textfile", "keyvaluetensor"])
+
+
+class DistributedTableTest(test.TestCase, parameterized.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(DistributedTableTest, cls).setUpClass()
+    cls.cluster = multi_worker_test_base.create_multi_process_cluster(
+        num_workers=2, num_ps=3, rpc_layer="grpc")
+    cls.cluster_resolver = cls.cluster.cluster_resolver
+
+  @classmethod
+  def tearDownClass(cls):
+    super(DistributedTableTest, cls).tearDownClass()
+    cls.cluster.stop()
+
+  def make_initializer(self, init_source, vals):
+    if init_source == "textfile":
+      file = os.path.join(self.get_temp_dir(), "text_file_initializer")
+      with open(file, "w") as f:
+        f.write("\n".join(str(v) for v in vals) + "\n")
+      return lookup_ops.TextFileInitializer(
+          filename=file,
+          key_dtype=dtypes.int64,
+          key_index=lookup_ops.TextFileIndex.LINE_NUMBER,
+          value_dtype=dtypes.int64,
+          value_index=lookup_ops.TextFileIndex.WHOLE_LINE)
+    elif init_source == "keyvaluetensor":
+      keys_tensor = constant_op.constant(
+          list(range(len(vals))), dtype=dtypes.int64)
+      vals_tensor = constant_op.constant(vals, dtype=dtypes.int64)
+      return lookup_ops.KeyValueTensorInitializer(keys_tensor, vals_tensor)
+    else:
+      raise ValueError("Unrecognized init_source: " + init_source)
+
+  def createStaticHashTable(self,
+                            init_source=None,
+                            vals=None,
+                            default_value=None,
+                            initializer=None):
+    if not initializer:
+      initializer = self.make_initializer(init_source, vals)
+    return lookup_ops.StaticHashTable(
+        initializer=initializer, default_value=default_value)
+
+  def makeDatasetFromTensorWithoutUsingResource(self, input_context, tensor):
+    """Returns a dataset made from `tensor`. To be called in a dataset_fn."""
+    global_batch_size = 24
+    batch_size = input_context.get_per_replica_batch_size(global_batch_size)
+    dataset = dataset_ops.DatasetV2.from_tensors(tensor).repeat().batch(
+        batch_size, drop_remainder=True)
+    dataset = dataset.shard(input_context.num_input_pipelines,
+                            input_context.input_pipeline_id)
+    dataset = dataset.prefetch(2)  # This prefetches 2 batches per device.
+    return dataset
+
+  @combinations.generate(source_combination)
+  def testCreateDistributedTableInScope(self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+
+    coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+      lookuptable = self.createStaticHashTable(
+          init_source=source, vals=[0, 1, 2], default_value=-2)
+
+    self.assertIsInstance(lookuptable, ps_values.DistributedTable)
+    self.assertEqual(self.evaluate(lookuptable.size()), 3)
+
+    # Lookup on the coordinator.
+    output = lookuptable.lookup(
+        constant_op.constant([0, 1, -1], dtype=dtypes.int64))
+    self.assertAllEqual([0, 1, -2], output)
+    self.assertEqual(lookuptable.size(), 3)
+
+  @combinations.generate(source_combination)
+  def testCopyDistributedTable(self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+
+    coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+      lookuptable = self.createStaticHashTable(
+          init_source=source, vals=[0, 1, 2], default_value=-2)
+
+    new_table = copy.copy(lookuptable)
+    # No new coordinator instance or distributed tables are created.
+    self.assertDictEqual(lookuptable.__dict__, new_table.__dict__)
+
+  @combinations.generate(source_combination)
+  def testCreateLookupInDatasetFnUnderScope(self, source):
+    # TODO(wxinyi): Warn the user of the inefficiency of this workflow (i.e.
+    # creating `StaticHashTable` inside a `@tf.function`-wrapped `dataset_fn` to
+    # be distributed with `distribute_datasets_from_function` and
+    # `create_per_worker_dataset`.
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+
+      def dataset_fn(input_context):
+        some_out_of_range_tensor = constant_op.constant(10, dtype=dtypes.int64)
+        lookuptable = self.createStaticHashTable(
+            init_source=source, vals=[0, 1, 2], default_value=-2)
+
+        self.assertNotIsInstance(lookuptable, ps_values.DistributedTable)
+
+        generation_tensor = lookuptable.lookup(some_out_of_range_tensor)
+        dataset = self.makeDatasetFromTensorWithoutUsingResource(
+            input_context, generation_tensor)
+        return dataset
+
+      @def_function.function
+      def per_worker_dataset_fn():
+        return strategy.distribute_datasets_from_function(dataset_fn)
+
+      per_worker_dataset = coordinator.create_per_worker_dataset(
+          per_worker_dataset_fn)
+      per_worker_iterator = iter(per_worker_dataset)
+
+      @def_function.function
+      def worker_fn(iterator):
+        return math_ops.reduce_sum(next(iterator))
+
+      result = []
+      for _ in range(10):
+        result.append(
+            coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
+
+      for r in result:
+        returned_input = r.fetch()
+        self.assertAllClose(-48, returned_input)
+
+  @combinations.generate(source_combination)
+  def testAccessingResourceHandleInDatasetFnWithoutMap(self, source):
+
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+      lookuptable = self.createStaticHashTable(
+          init_source=source, vals=[0, 1, 2], default_value=-2)
+
+    def dataset_fn(input_context):
+      some_out_of_range_tensor = constant_op.constant(10, dtype=dtypes.int64)
+
+      self.assertIsInstance(lookuptable, ps_values.DistributedTable)
+
+      generation_tensor = lookuptable.lookup(some_out_of_range_tensor)
+      dataset = self.makeDatasetFromTensorWithoutUsingResource(
+          input_context, generation_tensor)
+      return dataset
+
+    @def_function.function
+    def per_worker_dataset_fn():
+      return strategy.distribute_datasets_from_function(dataset_fn)
+
+    per_worker_dataset = coordinator.create_per_worker_dataset(
+        per_worker_dataset_fn)
+    per_worker_iterator = iter(per_worker_dataset)
+
+    @def_function.function
+    def worker_fn(iterator):
+      return math_ops.reduce_sum(next(iterator))
+
+    result = []
+    for _ in range(10):
+      result.append(
+          coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
+
+    for r in result:
+      returned_input = r.fetch()
+      self.assertAllClose(-48, returned_input)
+
+  @combinations.generate(
+      combinations.combine(
+          source=["textfile", "keyvaluetensor"],
+          create_datasets_under_scope=[True, False],
+          using_dataset_instance_not_function=[True, False],
+          create_per_worker_dataset_takes_instance=[True, False]))
+  def testCreateTableUnderScopeCombo(self, source,
+                                     create_datasets_under_scope,
+                                     using_dataset_instance_not_function,
+                                     create_per_worker_dataset_takes_instance):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+      lookup_table = self.createStaticHashTable(
+          init_source=source, vals=[0, 1, 2], default_value=-2)
+
+    if using_dataset_instance_not_function:
+
+      def per_worker_dataset_fn():
+        dataset = dataset_ops.DatasetV2.from_tensors(
+            constant_op.constant([0, 1, 3], dtype=dtypes.int64))
+        dataset = dataset.repeat().batch(24, drop_remainder=True).prefetch(2)
+        dataset = dataset.map(lookup_table.lookup)
+
+        return strategy.experimental_distribute_dataset(dataset)
+
+    else:
+
+      def per_worker_dataset_fn():
+        def dataset_fn(input_context):
+          batch_size = input_context.get_per_replica_batch_size(24)
+          dataset = dataset_ops.DatasetV2.from_tensors(
+              constant_op.constant([0, 1, 3], dtype=dtypes.int64))
+          dataset = dataset.repeat().batch(batch_size, drop_remainder=True)
+          dataset = dataset.shard(input_context.num_input_pipelines,
+                                  input_context.input_pipeline_id)
+          dataset = dataset.prefetch(2)  # This prefetches 2 batches per device.
+          dataset = dataset.map(lookup_table.lookup)
+          return dataset
+        return strategy.distribute_datasets_from_function(dataset_fn)
+
+    if create_datasets_under_scope:
+      with strategy.scope():
+        if create_per_worker_dataset_takes_instance:
+          per_worker_dataset = coordinator.create_per_worker_dataset(
+              per_worker_dataset_fn())
+        else:
+          per_worker_dataset = coordinator.create_per_worker_dataset(
+              per_worker_dataset_fn)
+        per_worker_iterator = iter(per_worker_dataset)
+
+    else:
+      if create_per_worker_dataset_takes_instance:
+        per_worker_dataset = coordinator.create_per_worker_dataset(
+            per_worker_dataset_fn())
+      else:
+        per_worker_dataset = coordinator.create_per_worker_dataset(
+            per_worker_dataset_fn)
+      per_worker_iterator = iter(per_worker_dataset)
+
+    @def_function.function
+    def worker_fn(iterator):
+      return math_ops.reduce_sum(next(iterator))
+
+    result = []
+    for _ in range(10):
+      result.append(
+          coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
+
+    for r in result:
+      returned_input = r.fetch()
+      self.assertAllClose(-24, returned_input)
+
+  @combinations.generate(
+      combinations.combine(
+          source=["textfile", "keyvaluetensor"],
+          create_datasets_under_scope=[True, False],
+          using_dataset_instance_not_function=[True, False],
+          create_per_worker_dataset_takes_instance=[True, False]))
+  def testCreateTableInDatasetCombo(self, source, create_datasets_under_scope,
+                                    using_dataset_instance_not_function,
+                                    create_per_worker_dataset_takes_instance):
+
+    if using_dataset_instance_not_function and (
+        not create_per_worker_dataset_takes_instance):
+      # This is the case that uses the `experimental_distribute_dataset` API to
+      # distribute dataset (instead of the `distribute_datasets_from_function`
+      # API), and passes `create_per_worker_dataset` a function that returns
+      # the distributed dataset (instead of passing it the distributed dataset
+      # directly).
+      # TODO(b/201775366): evaluate whether we need to handle this case
+      self.skipTest("Failed to serialize the input pipeline graph")
+
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    if using_dataset_instance_not_function:
+
+      def per_worker_dataset_fn():
+        # If this line is being called under strategy.scope(), it becomes a
+        # DistributedTable. Interestingly, after
+        # `experimental_distribute_dataset` serializes the dataset on chief and
+        # deserializes it on workers, `lookup_table` becomes a
+        # RestoredDistributedTable instead of a DistributedTable. And when it’s
+        # `resource_handle` is being accessed on the worker, it does not detect
+        # a DispatchContext, so it returns the restored resource handle,
+        # which is also the one on the local worker. The LookupTableFindV2 ops
+        # is on the local worker, too.
+        lookup_table = self.createStaticHashTable(
+            init_source=source, vals=[0, 1, 2], default_value=-2)
+
+        if create_datasets_under_scope:
+          self.assertIsInstance(lookup_table, ps_values.DistributedTable)
+
+        dataset = dataset_ops.DatasetV2.from_tensors(
+            constant_op.constant([0, 1, 3], dtype=dtypes.int64))
+        dataset = dataset.repeat().batch(24, drop_remainder=True).prefetch(2)
+        dataset = dataset.map(lookup_table.lookup)
+
+        return strategy.experimental_distribute_dataset(dataset)
+
+    else:
+
+      def per_worker_dataset_fn():
+
+        def dataset_fn(input_context):
+          # When we're wrapping the initialization of a StaticHashTable inside a
+          # `dataset_fn` to be distributed with
+          # `distribute_datasets_from_function`, no matter it's called under
+          # strategy.scope() or not, this call creates a StaticHashTable on
+          # chief instead of a DistributedTable on chief and workers.
+          # And correspondingly, LookupTableFindV2 ops is on chief and there are
+          # send-recv communication for the lookup.
+          lookup_table = self.createStaticHashTable(
+              init_source=source, vals=[0, 1, 2], default_value=-2)
+          if create_datasets_under_scope:
+            self.assertIsInstance(lookup_table, lookup_ops.StaticHashTable)
+            self.assertNotIsInstance(lookup_table, ps_values.DistributedTable)
+
+          batch_size = input_context.get_per_replica_batch_size(24)
+          dataset = dataset_ops.DatasetV2.from_tensors(
+              constant_op.constant([0, 1, 3], dtype=dtypes.int64))
+          dataset = dataset.repeat().batch(batch_size, drop_remainder=True)
+          dataset = dataset.shard(input_context.num_input_pipelines,
+                                  input_context.input_pipeline_id)
+          dataset = dataset.prefetch(2)  # This prefetches 2 batches per device.
+          dataset = dataset.map(lookup_table.lookup)
+          return dataset
+
+        return strategy.distribute_datasets_from_function(dataset_fn)
+
+    if create_datasets_under_scope:
+      with strategy.scope():
+        if create_per_worker_dataset_takes_instance:
+          per_worker_dataset = coordinator.create_per_worker_dataset(
+              per_worker_dataset_fn())
+        else:
+          per_worker_dataset = coordinator.create_per_worker_dataset(
+              per_worker_dataset_fn)
+        per_worker_iterator = iter(per_worker_dataset)
+
+    else:
+      if create_per_worker_dataset_takes_instance:
+        per_worker_dataset = coordinator.create_per_worker_dataset(
+            per_worker_dataset_fn())
+      else:
+        per_worker_dataset = coordinator.create_per_worker_dataset(
+            per_worker_dataset_fn)
+      per_worker_iterator = iter(per_worker_dataset)
+
+    @def_function.function
+    def worker_fn(iterator):
+      return math_ops.reduce_sum(next(iterator))
+
+    result = []
+    for _ in range(10):
+      result.append(
+          coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
+
+    for r in result:
+      returned_input = r.fetch()
+      self.assertAllClose(-24, returned_input)
+
+  @combinations.generate(source_combination)
+  def testAccessingTableInStepFunction(self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+      lookup_table = self.createStaticHashTable(
+          init_source=source, vals=[0, 1, 2], default_value=-2)
+
+    dataset = (
+        dataset_ops.DatasetV2.from_tensors(
+            constant_op.constant([0, 1, 3], dtype=dtypes.int64)).repeat().batch(
+                24, drop_remainder=True).prefetch(2))
+    dataset = dataset.map(lookup_table.lookup)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    distributed_dataset = coordinator.create_per_worker_dataset(
+        distributed_dataset)
+
+    @def_function.function
+    def worker_fn(iterator):
+
+      def replica_fn(inputs):
+        return math_ops.reduce_sum(lookup_table.lookup(inputs))
+
+      all_results = strategy.run(replica_fn, args=(next(iterator),))
+      return all_results
+
+    steps_per_epoch = 10
+    distributed_iterator = iter(distributed_dataset)
+    result = []
+    for _ in range(steps_per_epoch):
+
+      result.append(
+          coordinator.schedule(worker_fn, args=(distributed_iterator,)))
+
+    coordinator.join()
+
+    for r in result:
+      returned_input = r.fetch()
+      self.assertAllClose(-24, returned_input)
+
+  @combinations.generate(source_combination)
+  def testAccessingResourceHandleInDatasetFnWithMapFnDefinedOutside(
+      self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    with strategy.scope():
+      lookuptable = self.createStaticHashTable(
+          init_source=source, vals=[0, 1, 2], default_value=-2)
+
+    def map_fn(vals):
+      return lookuptable.lookup(vals)
+
+    def dataset_fn(input_context):
+      generation_tensor = constant_op.constant([0, 1, 3], dtype=dtypes.int64)
+      dataset = self.makeDatasetFromTensorWithoutUsingResource(
+          input_context, generation_tensor)
+      dataset = dataset.map(map_fn)
+      return dataset
+
+    @def_function.function
+    def per_worker_dataset_fn():
+      return strategy.distribute_datasets_from_function(dataset_fn)
+
+    per_worker_dataset = coordinator.create_per_worker_dataset(
+        per_worker_dataset_fn)
+    per_worker_iterator = iter(per_worker_dataset)
+
+    @def_function.function
+    def worker_fn(iterator):
+      return math_ops.reduce_sum(next(iterator))
+
+    result = []
+    for _ in range(10):
+      # batch_size == 24 and each input is [0, 1, -2]
+      result.append(
+          coordinator.schedule(worker_fn, args=(per_worker_iterator,)))
+
+    for r in result:
+      returned_input = r.fetch()
+      self.assertAllClose(-24, returned_input)
+
+  class Model(module.Module):
+
+    def __init__(self, init_source, filepath):
+      vals = [0, 1, 2]
+      if init_source == "textfile":
+
+        with open(filepath, "w") as f:
+          f.write("\n".join(str(v) for v in vals) + "\n")
+
+        self.initializer = lookup_ops.TextFileInitializer(
+            filepath, dtypes.int64, lookup_ops.TextFileIndex.LINE_NUMBER,
+            dtypes.int64, lookup_ops.TextFileIndex.WHOLE_LINE)
+      else:
+        keys_tensor = constant_op.constant(
+            list(range(len(vals))), dtype=dtypes.int64)
+        vals_tensor = constant_op.constant(vals, dtype=dtypes.int64)
+        self.initializer = lookup_ops.KeyValueTensorInitializer(
+            keys_tensor, vals_tensor)
+
+      self.table = lookup_ops.StaticHashTable(
+          self.initializer, default_value=-2)
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.int64)])
+    def use_table(self, x):
+      return self.table.lookup(x)
+
+  def verifyWorkerLocalInstance(self, coordinator, model):
+    # assert capturing a worker-local resource on each worker
+    for worker in coordinator._cluster.workers:
+      with coordinator_context.with_dispatch_context(worker):
+        captures = model.use_table.get_concrete_function().captured_inputs
+        resource_capture = [t for t in captures if t.dtype == dtypes.resource]
+        self.assertNotEmpty(resource_capture)
+        for capture in resource_capture:
+          self.assertEqual(
+              capture.device,
+              device_util.canonicalize("/CPU:0", default=worker.device_name))
+
+  @combinations.generate(source_combination)
+  def testInModelAndCapture(self, source):
+
+    file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
+
+    model = self.Model(source, file_path)
+    func_captures = model.use_table.get_concrete_function(
+    ).graph.external_captures
+    self.assertLen(func_captures, 2)
+    self.assertTrue(
+        any(model.table.resource_handle is t for t in func_captures))
+    deferred_captures = model.use_table.get_concrete_function(
+    ).graph.deferred_external_captures
+    self.assertEmpty(deferred_captures)
+
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    coordinator = coordinator_lib.ClusterCoordinator(strategy)
+    with strategy.scope():
+      distributed_model = self.Model("value", file_path)
+    func_captures = distributed_model.use_table.get_concrete_function(
+    ).graph.external_captures
+    # One less external_capture, since the table handle becomes a closure in the
+    # deferred_external_capture
+    self.assertLen(func_captures, 1)
+    self.assertFalse(
+        any(model.table.resource_handle is t for t in func_captures))
+    deferred_captures = distributed_model.use_table.get_concrete_function(
+    ).graph.deferred_external_captures
+    self.assertNotEmpty(deferred_captures)
+
+    self.verifyWorkerLocalInstance(coordinator, distributed_model)
+
+  @combinations.generate(source_combination)
+  def testLookupInNestedTFWhileLoop(self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    coordinator = coordinator_lib.ClusterCoordinator(strategy=strategy)
+
+    file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
+    with strategy.scope():
+      model = self.Model(source, file_path)
+
+    @def_function.function
+    def replica_fn(batch_data):
+      replica_result = array_ops.zeros(shape=(), dtype=dtypes.int64)
+      for _ in math_ops.range(10):
+        replica_result += math_ops.reduce_sum(model.use_table(batch_data))
+      return replica_result
+
+    @def_function.function
+    def step_fn(iterator):
+
+      step_result = array_ops.zeros(shape=(), dtype=dtypes.int64)
+      for _ in math_ops.range(10):
+        step_result += strategy.run(replica_fn, args=(next(iterator),))
+
+      return step_result
+
+    dataset = (
+        dataset_ops.DatasetV2.from_tensors(
+            constant_op.constant([0, 1, 3], dtype=dtypes.int64)).repeat().batch(
+                24, drop_remainder=True).prefetch(2))
+    distributed_dataset = coordinator.create_per_worker_dataset(
+        strategy.experimental_distribute_dataset(dataset))
+
+    results = []
+    for _ in range(10):
+      results.append(
+          coordinator.schedule(step_fn, args=(iter(distributed_dataset),)))
+
+    coordinator.join()
+
+    for r in results:
+      self.assertAllClose(-2400, r.fetch())
+
+  @combinations.generate(source_combination)
+  def testDistributeTableSaveAndServe(self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
+    with strategy.scope():
+      model = self.Model(source, file_path)
+
+    model_dir = self.get_temp_dir()
+    tf_save.save(model, model_dir)
+
+    loaded_without_strategy = tf_load.load(model_dir)
+    loaded_func_captures_without_strategy = (
+        loaded_without_strategy.use_table.get_concrete_function().graph
+        .external_captures)
+    loaded_func_deferred_captures_without_strategy = (
+        loaded_without_strategy.use_table.get_concrete_function().graph
+        .deferred_external_captures)
+    self.assertLen(loaded_func_captures_without_strategy, 2)
+    self.assertEmpty(loaded_func_deferred_captures_without_strategy)
+
+    self.assertAllEqual(
+        loaded_without_strategy.use_table(
+            constant_op.constant([0, 1, 3], dtype=dtypes.int64)), [0, 1, -2])
+
+  @combinations.generate(source_combination)
+  def testDistributeTableSaveAndLoadUnderStrategy(self, source):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    coordinator = coordinator_lib.ClusterCoordinator(strategy)
+    file_path = os.path.join(self.get_temp_dir(), "text_file_initializer")
+    with strategy.scope():
+      model = self.Model(source, file_path)
+    model_dir = self.get_temp_dir()
+    tf_save.save(model, model_dir)
+
+    with strategy.scope():
+      loaded = tf_load.load(model_dir)
+
+    loaded_func_captures = (
+        loaded.use_table.get_concrete_function().graph.external_captures)
+    loaded_func_deferred_captures = (
+        loaded.use_table.get_concrete_function().graph
+        .deferred_external_captures)
+    # Compared with loading without strategy, there is one less
+    # external_capture, since the captured table handle has been swapped to a
+    # closure in the deferred_external_capture
+    self.assertLen(loaded_func_captures, 1)
+    self.assertNotEmpty(loaded_func_deferred_captures)
+
+    self.assertIsInstance(loaded.table, ps_values.DistributedTable)
+
+    self.assertLen([
+        t for t in loaded.use_table.get_concrete_function().captured_inputs
+        if t.dtype == dtypes.resource
+    ], 1)
+
+    self.verifyWorkerLocalInstance(coordinator, loaded)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  multi_process_runner.test_main()

tensorflow/python/distribute/distributed_variable_test.py

@@ -0,0 +1,660 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for the distributed variables library."""
+
+import copy
+import os
+
+from absl.testing import parameterized
+from tensorflow.python.checkpoint import checkpoint as trackable_utils
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.distribute import collective_all_reduce_strategy
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute import packed_distributed_variable as packed
+from tensorflow.python.distribute import parameter_server_strategy
+from tensorflow.python.distribute import ps_values
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import test_util as ds_test_util
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.saved_model import save
+from tensorflow.python.saved_model import save_context
+from tensorflow.python.saved_model import save_options
+from tensorflow.python.types import core
+
+
+def _device_str(d):
+  return "/device:GPU:" + str(d)
+
+
+def _nested_value(d):
+  return ("a" + d, ["b" + d, {"c": "d" + d, "e": "f" + d}, "g" + d], "h" + d)
+
+
+def mirrored_and_tpu_strategy_combinations():
+  return combinations.combine(
+      distribution=[
+          strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          strategy_combinations.mirrored_strategy_with_two_gpus_no_merge_call,
+          strategy_combinations.tpu_strategy,
+          strategy_combinations.tpu_strategy_packed_var,
+      ],
+      mode=["graph", "eager"])
+
+
+def checkpoint_test_helper(dvar_test, distribution,
+                           synchronization, aggregation, enable_async_ckpt):
+  # This method is added since `testCheckpointing` cannot be parameterized after
+  # the entire class is parameterized.
+  with distribution.scope():
+    v = variables_lib.Variable(
+        constant_op.constant([1., 2., 3., 4]),
+        synchronization=synchronization,
+        aggregation=aggregation)
+
+  dvar_test.evaluate(v.initializer)
+  before_save = dvar_test.evaluate(v.read_value())
+
+  # Save random weights into checkpoint.
+  checkpoint = trackable_utils.Checkpoint(v=v)
+  ckpt_options = checkpoint_options.CheckpointOptions(
+      experimental_enable_async_checkpoint=enable_async_ckpt)
+  prefix = os.path.join(dvar_test.get_temp_dir(), "ckpt")
+  with dvar_test.test_session():
+    save_path = checkpoint.save(file_prefix=prefix, options=ckpt_options)
+
+  # Assign inverted value.
+  dvar_test.evaluate(v.assign(constant_op.constant([4., 3., 2., 1.])))
+  after_assign = dvar_test.evaluate(v.read_value())
+  dvar_test.assertNotAllClose(before_save, after_assign)
+
+  # Restore from the checkpoint.
+  with dvar_test.test_session():
+    checkpoint.restore(save_path).assert_consumed().run_restore_ops()
+  after_restore = dvar_test.evaluate(v)
+  dvar_test.assertAllClose(before_save, after_restore)
+
+  # Another round of saving/restoring to ensure that the logic of
+  # _copy_trackable_to_cpu works when a copy is already created in object_map.
+  dvar_test.evaluate(v.assign(constant_op.constant([5., 6., 7., 8.])))
+  before_save_1 = dvar_test.evaluate(v.read_value())
+  dvar_test.assertNotAllClose(before_save_1, after_restore)
+  with dvar_test.test_session():
+    save_path = checkpoint.save(file_prefix=prefix, options=ckpt_options)
+
+  dvar_test.evaluate(v.assign(constant_op.constant([8., 7., 6., 5.])))
+  after_assign_1 = dvar_test.evaluate(v.read_value())
+  dvar_test.assertNotAllClose(before_save_1, after_assign_1)
+
+  with dvar_test.test_session():
+    checkpoint.restore(save_path).assert_consumed().run_restore_ops()
+  after_restore_1 = dvar_test.evaluate(v)
+  dvar_test.assertAllClose(before_save_1, after_restore_1)
+
+
+@combinations.generate(
+    combinations.combine(
+        distribution=[
+            strategy_combinations.mirrored_strategy_with_one_cpu,
+            strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+            strategy_combinations.mirrored_strategy_with_two_gpus_no_merge_call,
+            strategy_combinations.tpu_strategy,
+            strategy_combinations.tpu_strategy_packed_var,
+            strategy_combinations.tpu_strategy_spmd,
+            strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
+            strategy_combinations.multi_worker_mirrored_2x1_cpu,
+            strategy_combinations.multi_worker_mirrored_2x1_gpu,
+            strategy_combinations.multi_worker_mirrored_2x2_gpu,
+            strategy_combinations.multi_worker_mirrored_2x2_gpu_no_merge_call,
+        ],
+        synchronization=[
+            variables_lib.VariableSynchronization.ON_READ,
+            variables_lib.VariableSynchronization.ON_WRITE,
+        ],
+        aggregation=[
+            variables_lib.VariableAggregation.MEAN,
+            variables_lib.VariableAggregation.SUM,
+            variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
+        ],
+        mode=["graph", "eager"],
+        use_var_policy=[True, False]))
+class DistributedVariableTest(test.TestCase, parameterized.TestCase):
+
+  def testExtendsVariable(self, distribution, synchronization, aggregation):
+    with distribution.scope():
+      v = variables_lib.Variable(
+          1., synchronization=synchronization, aggregation=aggregation)
+    self.assertIsInstance(v, variables_lib.Variable)
+
+  def testCheckpointing(self, distribution, synchronization, aggregation, mode):
+    if (isinstance(distribution,
+                   collective_all_reduce_strategy.CollectiveAllReduceStrategy)
+        and mode == "graph"):
+      self.skipTest("MWMS combinations tests do not work well in graph mode.")
+
+    checkpoint_test_helper(self, distribution, synchronization, aggregation,
+                           enable_async_ckpt=False)
+
+  def testAsyncCheckpointing(self, distribution, synchronization,
+                             aggregation, mode):
+    if (isinstance(distribution,
+                   collective_all_reduce_strategy.CollectiveAllReduceStrategy)
+        and mode == "graph"):
+      self.skipTest("MWMS combinations tests do not work well in graph mode.")
+
+    checkpoint_test_helper(self, distribution, synchronization, aggregation,
+                           enable_async_ckpt=True)
+
+  def testTraceback(self, distribution, synchronization, aggregation):
+    if context.executing_eagerly():
+      self.skipTest("does not apply to eager")
+    with distribution.scope():
+      variable_scope.get_variable(
+          name="testVar",
+          initializer=1.,
+          use_resource=True,
+          synchronization=synchronization,
+          aggregation=aggregation)
+      with self.assertRaisesRegex(ValueError,
+                                  "Variable testVar already exists"):
+        variable_scope.get_variable(
+            name="testVar",
+            initializer=1.,
+            use_resource=True,
+            synchronization=synchronization,
+            aggregation=aggregation)
+
+  def testSelectReplica(self, distribution, synchronization, aggregation):
+    with distribution.scope():
+      v = variables_lib.Variable(
+          1., synchronization=synchronization, aggregation=aggregation)
+    self.assertIs(v, distribute_utils.select_replica(0, v))
+
+  def testIsTensorLike(self, distribution, synchronization, aggregation):
+    if isinstance(distribution.extended,
+                  tpu_strategy.TPUExtended) and context.executing_eagerly():
+      self.skipTest("TPU doesn't support pure eager")
+
+    with distribution.scope():
+      v = variables_lib.Variable(
+          0., synchronization=synchronization, aggregation=aggregation)
+    # In cross replica context.
+    self.assertIsInstance(v, core.Tensor)
+    # In replica context.
+    distribution.run(lambda v: self.assertIsInstance(v, core.Tensor), args=(v,))
+
+  def testAssignReturnValueIsTensorLike(self, distribution, synchronization,
+                                        aggregation):
+    if isinstance(distribution.extended, tpu_strategy.TPUExtended):
+      if context.executing_eagerly():
+        self.skipTest("TPU doesn't support pure eager")
+      else:
+        self.skipTest("b/152076846")
+
+    with distribution.scope():
+      v = variables_lib.Variable(
+          0., synchronization=synchronization, aggregation=aggregation)
+
+    def assert_is_tensor_like(v):
+      # We can't use Python literals because they are treated as non-distributed
+      # values is not allowed when aggregation is SUM. See
+      # `cross_device_ops.reduce_non_distributed_value`.
+      delta = array_ops.identity(1.)
+      self.assertIsInstance(v.assign(delta), core.Tensor)
+      self.assertIsInstance(v.assign_sub(delta), core.Tensor)
+      self.assertIsInstance(v.assign_add(delta), core.Tensor)
+
+    # In cross replica context we return a PerReplica which is not Tensor like
+    # all the time yet.
+    if (synchronization == variables_lib.VariableSynchronization.ON_READ and
+        aggregation != variables_lib.VariableAggregation.SUM):
+      assert_is_tensor_like(v)
+
+    # In replica context.
+    distribution.run(assert_is_tensor_like, args=(v,))
+
+  def testDeepCopy(self, distribution, synchronization, aggregation):
+    if not context.executing_eagerly():
+      self.skipTest("deepcopy only supported in eager mode")
+
+    with distribution.scope():
+      v = variables_lib.Variable(
+          0., synchronization=synchronization, aggregation=aggregation)
+      in_dist_copy = copy.deepcopy(v)
+
+    out_dist_copy = copy.deepcopy(v)
+
+    def assert_is_deep_copy(v1, v2):
+      self.assertIsInstance(v2, type(v1))
+      self.assertEqual(v1.aggregation, v2.aggregation)
+      self.assertEqual(v1.distribute_strategy, v2.distribute_strategy)
+      if isinstance(v1, ps_values.AggregatingVariable):
+        self.assertIsInstance(v2.get(), type(v1.get()))
+        self.assertNotEqual(id(v1.get()), id(v2.get()))
+      else:
+        if v1._policy:
+          self.assertNotEqual(id(v1._policy), id(v2._policy))  # pylint: disable=protected-access
+        else:
+          self.assertEqual(id(v1._policy), id(v2._policy))  # pylint: disable=protected-access
+        self.assertEqual(len(v1.values), len(v2.values))
+        for (v1v, v2v) in zip(v1.values, v2.values):
+          self.assertEqual(v1v.device, v2v.device)
+          self.assertNotEqual(id(v1v), id(v2v))
+          self.assertAllEqual(
+              self.evaluate(v1.values), self.evaluate(v2.values))
+
+    self.evaluate(variables_lib.global_variables_initializer())
+    if not isinstance(distribution.extended, tpu_strategy.TPUExtended):
+      distribution.run(assert_is_deep_copy, args=(v, in_dist_copy))
+      distribution.run(assert_is_deep_copy, args=(v, out_dist_copy))
+
+  def testAssignSignature(self, distribution, synchronization, aggregation):
+    # This test verifies assign*() can be called in the same way as normal
+    # variables.
+    with distribution.scope():
+      v = variables_lib.Variable(
+          0., synchronization=synchronization, aggregation=aggregation)
+
+      def assign():
+        one = constant_op.constant(1.)
+        v.assign(one, True, "assign", False)
+        # TODO(b/154017756): SyncOnReadVariable.assign() doesn't support passing
+        # value as a keyword argument.
+        v.assign(one, use_locking=True, name="assign", read_value=False)
+        v.assign_add(one, True, "assign", False)
+        v.assign_add(one, use_locking=True, name="assign", read_value=False)
+        v.assign_sub(one, True, "assign", False)
+        v.assign_sub(one, use_locking=True, name="assign", read_value=False)
+        # Return something for graph mode to fetch.
+        return constant_op.constant(1)
+
+      self.evaluate(variables_lib.global_variables_initializer())
+      if not (synchronization == variables_lib.VariableSynchronization.ON_READ
+              and aggregation == variables_lib.VariableAggregation.SUM):
+        self.evaluate(distribution.experimental_local_results(assign()))
+      if not (isinstance(distribution.extended, tpu_strategy.TPUExtended) and
+              context.executing_eagerly()):
+        self.evaluate(
+            distribution.experimental_local_results(distribution.run(assign)))
+
+  def testStrategyExtendedUpdate(self, distribution, synchronization,
+                                 aggregation):
+    if len(distribution.extended.parameter_devices) != 2:
+      self.skipTest("n/a: needs exactly two parameter devices")
+    if (synchronization == variables_lib.VariableSynchronization.ON_WRITE and
+        aggregation != variables_lib.VariableAggregation.NONE):
+      self.skipTest("n/a: doesn't apply to ON_WRITE variable with aggregation")
+    with distribution.scope():
+      v = variables_lib.Variable(
+          0., synchronization=synchronization, aggregation=aggregation)
+    value = values_lib.PerReplica([1., 2.])
+
+    assign_fn = lambda var, value: var.assign(value)
+    self.evaluate(distribution.extended.update(v, assign_fn, args=(value,)))
+    self.assertAllEqual(self.evaluate(v.values), [1., 2.])
+
+    assign_add_fn = lambda var, value: var.assign_add(value)
+    self.evaluate(distribution.extended.update(v, assign_add_fn, args=(value,)))
+    self.assertAllEqual(self.evaluate(v.values), [2., 4.])
+
+    assign_sub_fn = lambda var, value: var.assign_sub(value)
+    self.evaluate(distribution.extended.update(v, assign_sub_fn, args=(value,)))
+    self.assertAllEqual(self.evaluate(v.values), [1., 2.])
+
+    read_assign_fn = lambda var, value: var.assign_add(var.value() + var.
+                                                       read_value())
+    self.evaluate(
+        distribution.extended.update(v, read_assign_fn, args=(value,)))
+    self.assertAllEqual(self.evaluate(v.values), [3., 6.])
+
+  def testSaveNonDistributed(self, distribution, synchronization, aggregation):
+    # This test verifies that the DistributedVariable behave like the primary
+    # variable when saving a non-distributed version of the model (the default).
+    # The test asserts that the function traced under SaveContext has no device
+    # annotations and only reference the primary component of the variable. Note
+    # that please avoid capturing other eager tensors in this test to make the
+    # assertion easy.
+
+    if isinstance(distribution.extended,
+                  parameter_server_strategy.ParameterServerStrategyExtended):
+      self.skipTest("b/148689177: AggregatingVariable doesn't "
+                    "conform to Variable interface well")
+
+    # tf.function requires the return value to be Tensors, which is not always
+    # case for properties and methods of Variable, so we simply discard the
+    # return values.
+    def _discard_return(f):
+      f()
+      return
+
+    def _test(f, v):
+      # This verifies that the function under SaveContext:
+      #   - contains no device annotations.
+      #   - only references the primary component of the variable.
+      g = def_function.function(lambda: _discard_return(f))
+      options = save_options.SaveOptions(
+          experimental_variable_policy=save_options.VariablePolicy.NONE)
+      with save_context.save_context(options):
+        # The graph should contain no device.
+        graph = g.get_concrete_function().graph
+      for op in graph.get_operations():
+        self.assertEqual(op.device, "", msg=str(op))
+      # The function should only capture the primary variable. Note that it
+      # may not have captures, e.g. v.aggregation.
+      captures = list(graph.captures)
+      self.assertLessEqual(len(captures), 1)
+      if graph.captures:
+        self.assertIs(captures[0][0], v._primary.handle)
+
+    def _assert(cond):
+      return control_flow_assert.Assert(cond, [cond])
+
+    with distribution.scope():
+      # We use four variables for convenience reasons. They have no special
+      # meaning.
+      # - v is used whenever possible.
+      # - w is used for scatter and gather, which require the variable to be
+      # non-scalar.
+      # - y is used when the dtype needs to be integer. Note that aggregation
+      # cannot be MEAN for integers.
+      v = variables_lib.Variable(
+          0.,
+          synchronization=synchronization,
+          aggregation=aggregation,
+          trainable=True)
+      w = variables_lib.Variable([0., 0., 0.],
+                                 synchronization=synchronization,
+                                 aggregation=aggregation,
+                                 trainable=True)
+      if aggregation != variables_lib.VariableAggregation.MEAN:
+        y = variables_lib.Variable(
+            0, synchronization=synchronization, aggregation=aggregation)
+
+    # pylint: disable=g-long-lambda
+
+    # tf.Variable properties.
+    _test(lambda: self.assertEqual(v.aggregation, aggregation), v)
+    _test(lambda: self.assertIs(v.constraint, None), v)
+    # TODO(crccw): should we raise an error instead?
+    _test(lambda: self.assertEqual(v.device, v._primary.device), v)
+    _test(lambda: self.assertEqual(v.dtype, dtypes.float32), v)
+    if not context.executing_eagerly():
+      _test(lambda: self.assertIs(v.graph, v._primary.graph), v)
+    if not context.executing_eagerly():
+      _test(lambda: _assert(v.initial_value == 0), v)
+    _test(lambda: self.assertIs(v.initializer, v._primary.initializer), v)
+    _test(lambda: self.assertEqual(v.name, "Variable:0"), v)
+    if not context.executing_eagerly():
+      _test(lambda: self.assertIs(v.op, v._primary.op), v)
+    _test(lambda: self.assertEqual(v.shape, tensor_shape.TensorShape(())), v)
+    _test(lambda: self.assertEqual(v.synchronization, synchronization), v)
+    _test(lambda: self.assertEqual(v.trainable, True), v)
+
+    # tf.Variable methods.
+    _test(lambda: check_ops.assert_equal_v2(v.assign(1.), 1.), v)
+    _test(lambda: check_ops.assert_equal_v2(v.assign_add(1.), 2.), v)
+    _test(lambda: check_ops.assert_equal_v2(v.assign_sub(1.), 1.), v)
+    # TODO(b/148689177): Implement batch_scatter_update.
+    # count_up_to() is skipped since it's deprecated.
+    # eval() is skipped since it shouldn't called in a tf.function.
+    # experimental_ref() is skipped since it's deprecated.
+    # from_proto() is skipped since it shouldn't called in a tf.function.
+    # TODO(b/148689177): Implement gather_nd.
+    _test(
+        lambda: check_ops.assert_equal_v2(v.get_shape(),
+                                          tensor_shape.TensorShape(())), v)
+    # initialized_value() is skipped since it shouldn't called in a tf.function.
+    # load() is skipped since it shouldn't called in a tf.function.
+    _test(lambda: check_ops.assert_equal_v2(v.read_value(), 1.), v)
+    # ref() is skipped since it shouldn't called in a tf.function.
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_add(_make_index_slices(values=[1., 2.], indices=[0, 2])),
+            [1., 0., 2.]), w)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_div(_make_index_slices(values=[4., 2.], indices=[0, 2])),
+            [0.25, 0., 1.]), w)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_max(_make_index_slices(values=[1., 0.5], indices=[1, 2])),
+            [0.25, 1., 1.]), w)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_min(_make_index_slices(values=[1., 0.5], indices=[0, 1])),
+            [0.25, 0.5, 1.]), w)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_mul(_make_index_slices(values=[2., 0.5], indices=[0, 1])),
+            [0.5, 0.25, 1.]), w)
+    # TODO(b/148689177): Implement scatter_nd_*
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_sub(_make_index_slices(values=[2., 0.5], indices=[0, 1])),
+            [-1.5, -0.25, 1.]), w)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            w.scatter_update(
+                _make_index_slices(values=[2., 0.5], indices=[0, 1])),
+            [2., 0.5, 1.]), w)
+    # set_shape() is skipped since ResourceVariable doesn't implement it.
+    # to_proto() is skipped since it shouldn't called in a tf.function.
+    _test(lambda: check_ops.assert_equal_v2(v.value(), 1.), v)
+
+    # DistributedVariable should be treated as ResourceVariable, so it needs to
+    # conform to ResourceVariable interface as well.
+    _test(lambda: self.assertIs(v.handle, v._primary.handle), v)
+
+    # Convert to tensor.
+    _test(lambda: check_ops.assert_equal_v2(ops.convert_to_tensor(v), 1.), v)
+
+    # Control dependency.
+    def _with_control_dep():
+      with ops.control_dependencies([v.assign(1.)]):
+        return array_ops.identity(1)
+
+    _test(_with_control_dep, v)
+
+    # Operator overloads.
+    _test(lambda: check_ops.assert_equal_v2(v.assign(7.), 7.), v)
+    _test(lambda: check_ops.assert_equal_v2(v + 1., 8.), v)
+    _test(lambda: check_ops.assert_equal_v2(3 + v, 10.), v)
+    _test(lambda: check_ops.assert_equal_v2(v + v, 14.), v)
+    _test(lambda: check_ops.assert_equal_v2(v - 2., 5.), v)
+    _test(lambda: check_ops.assert_equal_v2(v - v, 0.), v)
+    _test(lambda: check_ops.assert_equal_v2(v * 2., 14.), v)
+    _test(lambda: check_ops.assert_equal_v2(3 * v, 21.), v)
+    _test(lambda: check_ops.assert_equal_v2(v * v, 49.), v)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            math_ops.cast(v / 2., dtypes.float32), 3.5), v)
+    _test(
+        lambda: check_ops.assert_equal_v2(
+            math_ops.cast(14. / v, dtypes.float32), 2.), v)
+    _test(lambda: _assert(v < 12.), v)
+    _test(lambda: _assert(v <= 12.), v)
+    _test(lambda: _assert(not v > 12.), v)
+    _test(lambda: _assert(not v >= 12.), v)
+    _test(lambda: _assert(not 12. < v), v)
+    _test(lambda: _assert(not 12. <= v), v)
+    _test(lambda: _assert(12. > v), v)
+    _test(lambda: _assert(12. >= v), v)
+    _test(lambda: check_ops.assert_near_v2(pow(v, 3.), 343.), v)
+    _test(lambda: check_ops.assert_near_v2(pow(2., v), 128.), v)
+    _test(lambda: check_ops.assert_equal_v2(abs(v), 7.), v)
+
+    # Operator overloads that only works for integers.
+    if aggregation != variables_lib.VariableAggregation.MEAN:
+      _test(lambda: check_ops.assert_equal_v2(y.assign(7), 7), y)
+      _test(lambda: check_ops.assert_equal_v2(y // 2, 3), y)
+      _test(lambda: check_ops.assert_equal_v2(15 // y, 2), y)
+      _test(lambda: check_ops.assert_equal_v2(y % 2, 1), y)
+      _test(lambda: check_ops.assert_equal_v2(16 % y, 2), y)
+      _test(lambda: check_ops.assert_equal_v2(y & 3, 3), y)
+      _test(lambda: check_ops.assert_equal_v2(3 & y, 3), y)
+      _test(lambda: check_ops.assert_equal_v2(y | 8, 15), y)
+      _test(lambda: check_ops.assert_equal_v2(16 | y, 23), y)
+      _test(lambda: check_ops.assert_equal_v2(y ^ 3, 4), y)
+      _test(lambda: check_ops.assert_equal_v2(11 ^ y, 12), y)
+      _test(lambda: check_ops.assert_equal_v2(-y, -7), y)
+      _test(lambda: check_ops.assert_equal_v2(~y, ~7), y)
+
+    # Index.
+    if isinstance(distribution.extended, tpu_strategy.TPUExtended):
+      # TODO(b/161572567): slice assignment doesn't work for TPU.
+      _test(lambda: check_ops.assert_equal_v2(w[0], 2.), w)
+    else:
+      _test(lambda: check_ops.assert_equal_v2(w[0].assign(1.), [1., 0.5, 1.]),
+            w)
+      _test(lambda: check_ops.assert_equal_v2(w[0], 1.), w)
+
+    # pylint: enable=g-long-lambda
+
+  def testUnsaveable(self, distribution, synchronization, aggregation, mode):
+    if isinstance(distribution.extended,
+                  parameter_server_strategy.ParameterServerStrategyExtended):
+      self.skipTest("n/a: not applicable to AggregatingVariable")
+    if (isinstance(distribution,
+                   collective_all_reduce_strategy.CollectiveAllReduceStrategy)
+        and mode == "graph"):
+      self.skipTest("MWMS combinations tests do not work well in graph mode.")
+    if not distribution.extended._use_merge_call():
+      self.skipTest("Unsupported combination.")
+    with distribution.scope():
+      v = variables_lib.Variable([1., 1.],
+                                 synchronization=synchronization,
+                                 aggregation=aggregation)
+
+    with self.cached_session():
+      self.evaluate(variables_lib.global_variables_initializer())
+
+    export_dir = self.get_temp_dir()
+
+    def _assert_unsaveable(f):
+      # Ignore if it cannot be traced. Certain combinations are not supported or
+      # yet or not allowed.
+      try:
+        f = def_function.function(f).get_concrete_function()
+      except (NotImplementedError, ValueError):
+        return
+      with self.assertRaisesRegex(ValueError, "f_with_input_signature"):
+        save.save(v, export_dir, signatures=f)
+
+    _assert_unsaveable(lambda: v.assign(ops.convert_to_tensor([1., 1.])))
+    _assert_unsaveable(lambda: v.assign_add(ops.convert_to_tensor([1., 1.])))
+    _assert_unsaveable(lambda: v.assign_sub(ops.convert_to_tensor([1., 1.])))
+    _assert_unsaveable(lambda: v.scatter_add(_make_index_slices([1.], [0])))
+    _assert_unsaveable(lambda: v.scatter_sub(_make_index_slices([1.], [0])))
+    _assert_unsaveable(lambda: v.scatter_mul(_make_index_slices([1.], [0])))
+    _assert_unsaveable(lambda: v.scatter_div(_make_index_slices([1.], [0])))
+    _assert_unsaveable(lambda: v.scatter_min(_make_index_slices([1.], [0])))
+    _assert_unsaveable(lambda: v.scatter_max(_make_index_slices([1.], [0])))
+    _assert_unsaveable(lambda: v.scatter_update(_make_index_slices([1.], [0])))
+    # Reading a ON_READ variable should be unsaveable if either:
+    # 1) CollectiveAllReduceStrategy, and aggregation is MEAN/SUM.
+    # 2) aggregation is SUM.
+    if (synchronization == variables_lib.VariableSynchronization.ON_READ and
+        (aggregation == variables_lib.VariableAggregation.SUM or
+         (not distribution.extended._use_merge_call()) or
+         (isinstance(distribution.extended,
+                     collective_all_reduce_strategy.CollectiveAllReduceExtended)
+          and aggregation == variables_lib.VariableAggregation.MEAN))):
+      _assert_unsaveable(v.read_value)
+      _assert_unsaveable(v.value)
+      _assert_unsaveable(lambda: ops.convert_to_tensor(v))
+    else:
+      # Otherwise reading a variable should be saveable.
+
+      @def_function.function
+      def f():
+        v.read_value()
+        v.value()
+        return ops.convert_to_tensor(v)
+
+      with self.cached_session():
+        save.save(v, export_dir, signatures=f.get_concrete_function())
+
+
+@combinations.generate(
+    combinations.combine(
+        distribution=[
+            strategy_combinations.mirrored_strategy_with_one_cpu,
+            strategy_combinations.tpu_strategy,
+        ],
+        mode=["eager"]))
+class PackedDistributedVariableTest(test.TestCase, parameterized.TestCase):
+
+  def testPackedVariable(self, distribution):
+    with distribution.scope():
+      v0 = variables_lib.Variable(0.)
+    self.assertIsNone(v0._packed_var)
+
+    distribution._enable_packed_variable_in_eager_mode = True
+    with distribution.scope():
+      v1 = variables_lib.Variable(0)
+      self.assertIsInstance(v1._packed_var, packed.PackedDistributedVariable)
+
+    devices = v1._devices
+    for i in range(1, len(devices)):
+      with distribute_lib.ReplicaContext(distribution, i):
+        v1.assign(i)
+    val = v1._get()
+    self.assertIsInstance(val, packed.PackedVarAndDevice)
+    self.assertEqual(val.device, devices[0])
+    self.assertEqual(self.evaluate(val.read_value()), 0)
+    for i in range(0, len(devices)):
+      with distribute_lib.ReplicaContext(distribution, i):
+        val = v1._get()
+        self.assertIsInstance(val, packed.PackedVarAndDevice)
+        self.assertEqual(val.device, devices[i])
+        self.assertEqual(self.evaluate(val.read_value()), i)
+
+  def testIgnorePackedVariableInSaveContext(self, distribution):
+    distribution._enable_packed_variable_in_eager_mode = True
+    with distribution.scope():
+      v = variables_lib.Variable(0)
+      self.assertIsInstance(v._packed_variable,
+                            packed.PackedDistributedVariable)
+
+    options = save_options.SaveOptions()
+    with save_context.save_context(options):
+      self.assertIsNone(v._packed_variable)
+
+
+def _make_index_slices(values, indices, dense_shape=None):
+  if dense_shape:
+    dense_shape = array_ops.identity(dense_shape)
+  return indexed_slices.IndexedSlices(
+      array_ops.identity(values), array_ops.identity(indices), dense_shape)
+
+
+if __name__ == "__main__":
+  ds_test_util.main()

tensorflow/python/distribute/experimental/BUILD

@@ -0,0 +1,199 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load(
+    "//tensorflow/dtensor:build_defs.bzl",
+    "GPU_2DEVS_BACKEND",
+    "TPU_V3_DONUT_BACKEND",
+    "dtensor_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//learning/brain/distribute/python/benchmark:__subpackages__",
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "experimental",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":mirrored_strategy",
+        ":multi_worker_mirrored_strategy",
+    ],
+)
+
+py_library(
+    name = "mirrored_strategy",
+    srcs = ["mirrored_strategy.py"],
+    strict_deps = True,
+    deps = [
+        ":dtensor_strategy_extended",
+        ":dtensor_util",
+        "//tensorflow/dtensor/python:config",
+        "//tensorflow/dtensor/python:mesh_util",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/framework:device",
+    ],
+)
+
+dtensor_test(
+    name = "mirrored_strategy_test",
+    srcs = ["mirrored_strategy_test.py"],
+    shard_count = {"tpu": 2},
+    tags = ["no_pip"],
+    deps = [
+        ":dtensor_util",
+        ":mirrored_strategy",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:d_variable",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/dtensor/python:mesh_util",
+        "//tensorflow/dtensor/python/tests:test_util",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:stateless_random_ops",
+        "//tensorflow/python/ops:variables",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "dtensor_util",
+    srcs = ["dtensor_util.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/dtensor/python:accelerator_util",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:input_util",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/python/distribute:cross_device_ops",
+        "//tensorflow/python/distribute:device_util",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion_registry",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:summary_ops_v2",
+    ],
+)
+
+dtensor_test(
+    name = "dtensor_util_test",
+    srcs = ["dtensor_util_test.py"],
+    tags = ["no_pip"],
+    deps = [
+        ":dtensor_util",
+        ":mirrored_strategy",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/dtensor/python/tests:test_util",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "dtensor_strategy_extended",
+    srcs = ["dtensor_strategy_extended.py"],
+    strict_deps = True,
+    deps = [
+        ":dtensor_util",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:config",
+        "//tensorflow/dtensor/python:d_variable",
+        "//tensorflow/dtensor/python:input_util",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/python/data/experimental/ops:distribute",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:distribute_utils",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_library(
+    name = "multi_worker_mirrored_strategy",
+    srcs = ["multi_worker_mirrored_strategy.py"],
+    strict_deps = True,
+    deps = [
+        ":dtensor_strategy_extended",
+        ":dtensor_util",
+        "//tensorflow/dtensor/python:config",
+        "//tensorflow/dtensor/python:mesh_util",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:multi_worker_util",
+        "//tensorflow/python/distribute/cluster_resolver:tfconfig_cluster_resolver_py",
+    ],
+)
+
+dtensor_test(
+    name = "multi_worker_mirrored_strategy_test",
+    srcs = ["multi_worker_mirrored_strategy_test.py"],
+    additional_backends = [
+        GPU_2DEVS_BACKEND,
+        TPU_V3_DONUT_BACKEND,
+    ],
+    disable = [
+        "gpu",  # multi-client gpu is tested via GPU_2DEVS_BACKEND.
+        "tpu",  # multi-client tpu is tested via TPU_V3_DONUT_BACKEND.
+    ],
+    disable_tfrt = [
+        "cpu",  # TODO(b/217969210): Re-enable in TFRT CPU.
+        GPU_2DEVS_BACKEND,  # TODO(b/230679405): Re-enable in TFRT GPU.
+    ],
+    tags = [
+        "no_pip",
+        "no_windows",
+        "nosan",
+    ],  # b/195537906
+    deps = [
+        ":dtensor_util",
+        ":multi_worker_mirrored_strategy",
+        "//tensorflow/dtensor/python:api",
+        "//tensorflow/dtensor/python:d_variable",
+        "//tensorflow/dtensor/python:layout",
+        "//tensorflow/dtensor/python/tests:multi_client_test_util",
+        "//tensorflow/dtensor/python/tests:test_util",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:reduce_util",
+        "//tensorflow/python/distribute/cluster_resolver:tfconfig_cluster_resolver_py",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:stateless_random_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/flags",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/distribute/experimental/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Experimental Distribution Strategy library."""
+
+# pylint: disable=unused-import
+from tensorflow.python.distribute.experimental import mirrored_strategy
+from tensorflow.python.distribute.experimental import multi_worker_mirrored_strategy
+# pylint: enable=unused-import

tensorflow/python/distribute/experimental/dtensor_strategy_extended.py

@@ -0,0 +1,278 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Implement a StrategyExtended based on the DTensor low level API."""
+
+import functools
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import config as d_config
+from tensorflow.dtensor.python import d_variable
+from tensorflow.dtensor.python import input_util
+from tensorflow.dtensor.python import layout
+from tensorflow.python.data.experimental.ops import distribute
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.util import nest
+
+
+class DTensorStrategyExtended(distribute_lib.StrategyExtendedV2):
+  """Strategy extension that support both single and multi worker strategy."""
+  # Note that the unit test for this class is via the strategy interface.
+
+  def __init__(self, container_strategy, mesh):
+    super().__init__(container_strategy)
+    self._mesh = mesh
+    self._num_clients = d_config.num_clients()
+    self._client_id = d_config.client_id()
+
+  def _create_variable(self, next_creator, **kwargs):
+    # Make sure the pop the `use_resource` which is not supported by the
+    # base tf.Variable. The `use_resource` is added by
+    # creator_with_resource_vars in distribute_lib.py
+    kwargs.pop('use_resource', None)
+
+    # Ignore the colocate_with for the mirrored strategy. Each of the device
+    # will get same copy of variable in the DTensor's case.
+    # `colocate_with` is added when user call:
+    # strategy.extended.colocate_vars_with(variable)
+    kwargs.pop('colocate_with', None)
+
+    # Ignore expected_shape, which is from the v1 Variable. Keras was somehow
+    # using the v1 Variable, but didn't specify that value particularly.
+    kwargs.pop('expected_shape', None)
+
+    # Make sure to call DVariable initializer under the scope so that it will
+    # have the proper replicated layout. The initial_value is multi-typed,
+    # eg it can be a tensor, or a python/numpy type, or a callable that
+    # produce tensor/python/numpy types. In all those cases, we need to wrap
+    # them invoke convert_to_tensor() under the scope so that the proper
+    # layout can be assigned.
+
+    # TODO(scottzhu): The layout information should be injected via kwargs, or
+    # lazily set later.
+    initial_value = kwargs.pop('initial_value')
+    dtype = kwargs.get('dtype', None)
+    def new_initial_value():
+      if callable(initial_value):
+        init_var = ops.convert_to_tensor(initial_value(), dtype=dtype)
+      else:
+        init_var = ops.convert_to_tensor(initial_value, dtype=dtype)
+      rank = init_var.shape.rank
+      return d_api.copy_to_mesh(
+          init_var, layout.Layout.replicated(self._mesh, rank))
+
+    return d_variable.DVariable(new_initial_value, **kwargs)
+
+  @property
+  def _num_replicas_in_sync(self):
+    # The mesh should be 1D with batch sharding only.
+    # In the model parallel case, it should only return the size of
+    # batch dimension.
+    return self._mesh.size
+
+  def value_container(self, value):
+    return value
+
+  @property
+  def worker_devices(self):
+    # Note that in either single worker (MirroredStrategy) or multi worker (
+    # MultiWorkerMirroredStrategy), worker_devices refers to the local worker
+    # devices.
+    return tuple(self._mesh.local_devices())
+
+  @property
+  def parameter_devices(self):
+    # Same as the worker_devices.
+    return self.worker_devices
+
+  def _in_multi_worker_mode(self):
+    return d_config.num_clients() > 1
+
+  def _get_local_replica_id(self, replica_id_in_sync_group):
+    return replica_id_in_sync_group
+
+  def _default_device_scope(self):
+    return d_api.default_mesh(self._mesh)
+
+  def _experimental_distribute_dataset(self, dataset, options):
+    # Strategy always assume the user input data is a batched dataset for
+    # experimental_distribute_dataset().
+    # TODO(yuefengz): Add check for whether a dataset is batched for all
+    # strategies.
+
+    # TODO(b/265198795): Support dataset already batched to global batch size.
+    # Since DTensorDataset doesn't support batched dataset that is already
+    # batched global batch size, it only supports dataset that is batched to
+    # local batch size, we need to infer the batch size, and unbatch the dataset
+    # until the b/265198795 is resolved.
+    batch_size = distribute.compute_batch_size(dataset)
+
+    # There are multiple case that the batch is not static, eg partial batch,
+    # or uneven batch, in all those case, it will return -1.
+    if batch_size.numpy() < 0:
+      # When we don't have a static batch size.
+      raise ValueError('DTensor strategy requires a static batch size for now.'
+                       'The dynamic batch size will be supported in future')
+    # Unbatch the dataset for now since the DTensorDataset has some limitation
+    # about the local batch size as well as the mesh size.
+    dataset = dataset.unbatch()
+
+    def _create_batch_layout(tensor_spec):
+      # For unbatched dataset, the new layout need to have +1 rank for
+      # the batched result.
+      rank = len(tensor_spec.shape) + 1
+      return layout.Layout.batch_sharded(
+          self._mesh, batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+          rank=rank)
+
+    layouts = nest.map_structure(_create_batch_layout, dataset.element_spec)
+
+    return input_util.DTensorDataset(
+        dataset=dataset,
+        mesh=self._mesh,
+        layouts=layouts,
+        global_batch_size=batch_size,
+        dataset_already_batched=False,
+        batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+        # TODO(scottzhu): Add prefetch support by inspecting the input dataset.
+        prefetch=None,
+        tf_data_service_config=None
+    )
+
+  def _make_dataset_iterator(self, dataset):
+    raise NotImplementedError(
+        'Strategy.make_dataset_iterator() is deprecated, and only available '
+        'in the V1 API.')
+
+  def _make_input_fn_iterator(self, input_fn, replication_mode):
+    raise NotImplementedError(
+        'Strategy.make_input_fn_iterator() is deprecated, and only available '
+        'in the V1 API.')
+
+  def _distribute_datasets_from_function(self, dataset_fn, options):
+    # TODO(scottzhu): Implement the logic for options in future
+    del options
+    input_context = distribute_lib.InputContext(
+        num_input_pipelines=self._num_clients,
+        input_pipeline_id=self._client_id,
+        num_replicas_in_sync=self._num_replicas_in_sync
+    )
+    dataset = dataset_fn(input_context)
+
+    # Note that the dataset should already batched to local per-relica batch
+    def _create_batch_layout(tensor_spec):
+      rank = len(tensor_spec.shape)
+      return layout.Layout.batch_sharded(
+          self._mesh, batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+          rank=rank)
+
+    layouts = nest.map_structure(_create_batch_layout, dataset.element_spec)
+
+    batch_size = distribute.compute_batch_size(dataset)
+    # There are multiple case that the batch is not static, eg partial batch,
+    # or uneven batch, in all those case, it will return -1.
+    if batch_size.numpy() < 0:
+      # When we don't have a static batch size.
+      raise ValueError('DTensor strategy requires a static batch size for now.'
+                       'The dynamic batch size will be supported in future')
+    global_batch_size = batch_size.numpy() * self._num_replicas_in_sync
+
+    return input_util.DTensorDataset(
+        dataset=dataset,
+        mesh=self._mesh,
+        layouts=layouts,
+        global_batch_size=global_batch_size,
+        dataset_already_batched=True,
+        batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
+        # TODO(scottzhu): Add prefetch support by inspecting the input dataset.
+        prefetch=None,
+        tf_data_service_config=None
+    )
+
+  def _experimental_distribute_values_from_function(self, value_fn):
+    per_replica_values = []
+    # Note that in the multi-worker setting, this function only return the
+    # slide of DistributedValue for the current worker.
+    for i in range(self._mesh.num_local_devices()):
+      # In the case of 2 worker with 2 local devices on each worker,
+      # worker 0 will get 0 and 1 for replica_id.
+      # worker 1 will get 2 and 3 for replica_id.
+      replica_id = d_config.client_id() * self._mesh.num_local_devices() + i
+      per_replica_values.append(value_fn(
+          distribute_lib.ValueContext(replica_id,
+                                      self._num_replicas_in_sync)))
+    # Instead of using the DistributeVariable, return a DTensor instead since
+    # the run() will expect a DTensor instance.
+    result = distribute_utils.regroup(per_replica_values, always_wrap=True)
+    map_fn = functools.partial(dtensor_util.convert_per_replica_to_dtensor,
+                               mesh=self._mesh)
+    return nest.map_structure(map_fn, result)
+
+  def call_for_each_replica(self, fn, args=(), kwargs=None):
+    """Run `fn` once per replica.
+
+    This is a method that expected by the strategy base class in its `run()`.
+
+    Args:
+      fn: function to run (will be run once per replica).
+      args: Tuple or list with positional arguments for `fn`.
+      kwargs: Dict with keyword arguments for `fn`.
+
+    Returns:
+      Merged return value of `fn` across all replicas.
+    """
+    # Comparing to the existing MirroredStrategy, which will run the fn on
+    # each of the replica with individual thread, the DTensor will just run
+    # the fn once with the DTensor inputs, and the distribution will be handled
+    # by the DTensor.
+
+    distribute_lib._require_cross_replica_or_default_context_extended(self)   # pylint: disable=protected-access
+    if kwargs is None:
+      kwargs = {}
+
+    # For any value that is not DTensor, eg normal tf.Tensor or
+    # DistributedValues, we need to convert them into DTensor.
+    map_fn = functools.partial(dtensor_util.convert_inputs_to_dtensor,
+                               mesh=self._mesh)
+    d_args = nest.map_structure(map_fn, args)
+    d_kwargs = nest.map_structure(map_fn, kwargs)
+
+    with self._container_strategy().scope():
+      with dtensor_util.DTensorReplicaContext(self._container_strategy()):
+        dtensor_result = fn(*d_args, **d_kwargs)
+
+    return nest.map_structure(
+        dtensor_util.DTensorDistributedValue,
+        dtensor_result)
+
+  def _gather_to_implementation(self, value, destinations, axis, options):
+    if isinstance(value, dtensor_util.DTensorDistributedValue):
+      value = value.get_dtensor()
+    if not d_api.is_dtensor(value):
+      # This is the current behavior for mirrored strategy, should we raise an
+      # error for unsupported types?
+      return value
+
+    # Unpack the dtensor components and gather the tensors on the axis
+    components = d_api.unpack(value)
+    return array_ops.concat(components, axis=axis)
+
+  def _use_merge_call(self):
+    # This is method for V1 StrategyExtended by still used by
+    # tf.__internal__.distribute.strategy_supports_no_merge_call
+    return False

tensorflow/python/distribute/experimental/dtensor_util.py

@@ -0,0 +1,299 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Utilities for strategies that are backed by DTensor."""
+
+from tensorflow.dtensor.python import accelerator_util
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import input_util
+from tensorflow.dtensor.python import layout
+from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.eager import context
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion_registry
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import summary_ops_v2
+
+
+# Default dimension name used for the mesh created when user provide a list
+# of devices. For mirrored strategy, it should be a 1D mesh with batch dim only.
+DEFAULT_BATCH_MESH_DIM_NAME = "batch"
+
+
+class DTensorDistributedValue(values_lib.DistributedValues):
+  """DistributedValue backed by a DTensor instance.
+
+  This class is useful to align the interface between DTensor and tf.distribute.
+  Most of the tf.distribute API will accept/return DistributedValue, whereas
+  DTensor low level API will only accept DTensor instance. In order to avoid
+  the conversion back and forth between DistributedValue and DTensor, we
+  introduce this class so that it can work with both side.
+  """
+
+  def __init__(self, dtensor):
+    if context.executing_eagerly():
+      if not d_api.is_dtensor(dtensor):
+        raise ValueError("The DTensorDistributedValue can only be built with "
+                         f"DTensor instance, got {type(dtensor)}")
+      super().__init__(d_api.unpack(dtensor))
+    else:
+      # We can't unpack the dtensor instance for now due to graph context.
+      # We will treat the dtensor instance as one global instance and let it
+      # return as a global replica instance.
+      # TODO(feyu): Support unpack in the graph context.
+      super().__init__([dtensor,])
+    self._dtensor = dtensor
+
+  def get_dtensor(self):
+    return self._dtensor
+
+  @property
+  def values(self):
+    # Note that this method exists so that it match the interface for PerReplica
+    # The public API in `tf.types.experimental.distributed.PerReplica` doesn't
+    # define any methods.
+    return self._values
+
+
+def _dtensor_distributed_value_to_tensor(
+    var, dtype=None, name=None, as_ref=False):
+  del name
+  dtensor = var.get_dtensor()
+  if dtype is not None and not dtype.is_compatible_with(dtensor.dtype):
+    raise ValueError(
+        "Incompatible type conversion requested to type {!r} for variable "
+        "of type {!r}".format(dtype.name, dtensor.dtype.name))
+  if as_ref:
+    raise NotImplementedError(
+        "PerReplica doesn't support being used as a reference.")
+  return dtensor
+
+
+# Register a conversion function to provide a useful error message when users
+# try to use PerReplica values in the wrong contexts
+tensor_conversion_registry.register_tensor_conversion_function(
+    DTensorDistributedValue, _dtensor_distributed_value_to_tensor)
+
+
+class DTensorReplicaContext(distribute_lib.ReplicaContext):
+  """ReplicaContext for strategy that is backed by DTensor.
+
+  Since the DTensor is operated in the global context, most of the methods from
+  existing strategy ReplicaContext is not applicable since they need to access
+  local values. For now most of the methods in this class will raise explicit
+  error to user, and we will add more support for local values in future.
+  """
+  _UNSUPPORTED_ERROR_MSG = (
+      "Strategy that is backed by DTensor is run with a global context, and "
+      "doesn't support operations for local context, like any call to merge/"
+      "gather/reduce or local replica ID. Please use any strategy that is not "
+      "backed by DTensor")
+
+  def __init__(self, strategy):
+    # Since DTensor strategy only runs in a global context, and we can't have
+    # a local replica ID in the sync group. For now we pass None to parent, and
+    # raise an explicit error when it is accessed.
+    super().__init__(strategy, replica_id_in_sync_group=None)
+
+  def __enter__(self):
+    # This is a copy of parent class, without any check about whether the
+    # current replica is the first one (since DTensor only has one).
+    distribute_lib._push_per_thread_mode(self._thread_context)  # # pylint: disable=protected-access
+
+    summary_state = summary_ops_v2._summary_state  # pylint: disable=protected-access
+    self._summary_recording_distribution_strategy = (
+        summary_state.is_recording_distribution_strategy)
+    summary_state.is_recording_distribution_strategy = True
+
+  @property
+  def replica_id_in_sync_group(self):
+    # Since there is only one global context for DTensor, we always return a
+    # constant value here. This value is needed by the RNG which try to generate
+    # different seed for different replica.
+    return 0
+
+  @property
+  def _replica_id(self):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def merge_call(self, merge_fn, args=(), kwargs=None):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def all_reduce(self, reduce_op, value, options=None):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def all_gather(self, value, axis, options=None):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+  def _update(self, var, fn, args=(), kwargs=None, group=True):
+    raise NotImplementedError(self._UNSUPPORTED_ERROR_MSG)
+
+
+def initialize_accelerator_system_once(device_type):
+  # Initialize the GPU/TPU before creating the mesh.
+  # Note that this method will also trigger the creation of the pairing
+  # virtual host CPUs, which is needed by dataset and checkpoint.
+  if not accelerator_util.is_initialized():
+    # TODO(feyu): Add a method in accelerator_util to check the initialized
+    # mesh device types.
+    accelerator_util.initialize_accelerator_system(
+        device_type,
+        experimental_reset_context=True)
+
+
+def convert_inputs_to_dtensor(inputs, mesh):
+  """Convert any input types to DTensor instance."""
+  if isinstance(inputs, DTensorDistributedValue):
+    return inputs.get_dtensor()
+  elif isinstance(inputs, values_lib.DistributedValues):
+    return convert_per_replica_to_dtensor(inputs, mesh)
+  elif isinstance(inputs, input_util._DTensorIterator):   # pylint: disable=protected-access
+    return inputs
+  elif tensor_util.is_tensor(inputs):
+    if context.executing_eagerly():
+      if d_api.is_dtensor(inputs):
+        return inputs
+      else:
+        # For a non-dtensor input in eager context, we could choose to replica
+        # them into per-replica and then pack them into dtensor. However, this
+        # will cause an eager/graph discrepancy since we can't do this check in
+        # the graph context. For now, we will ask user to provide a distributed
+        # value for inputs.
+        _raise_unsupported_input_type_error(inputs)
+    else:
+      # For graph context, since we can't check if they are dtensor or not. We
+      # will assume the value is already distributed. This is a critical use
+      # case for keras, where all the inputs are pre-distributed via strategy,
+      # and the train function execute within graph context.
+      return inputs
+  else:
+    # For any other types.
+    _raise_unsupported_input_type_error(inputs)
+
+
+def _raise_unsupported_input_type_error(inputs):
+  raise ValueError("Unsupported input types for MirroredStrategy. "
+                   "Please use `strategy.distribute_dataset` or "
+                   "`strategy.distribute_values_from_function` to "
+                   f"distribute inputs. Received input type: {type(inputs)}")
+
+
+def is_distributed_value(value):
+  return isinstance(
+      value, values_lib.DistributedValues) or d_api.is_dtensor(value)
+
+
+def convert_per_replica_to_dtensor(per_replica_value, mesh):
+  """Convert a PerReplica result to a DTensor instance.
+
+  Args:
+    per_replica_value: A PerReplica instance whose value will be converted
+      to DTensor.
+    mesh: The mesh used for layout creation.
+
+  Returns:
+    A DTensor instance that packed from per_replica_value with batch sharded
+      layout.
+  """
+  values = per_replica_value.values
+  if isinstance(values[0], (float, int)):
+    rank = 0
+  else:
+    rank = len(values[0].shape)
+
+  if rank == 0:
+    result = []
+    # dtensor.pack requires each component to have same rank as the packed
+    # result. When the individual value is scalar, it needs to be expanded into
+    # 1D tensor.
+    for v in values:
+      result.append(array_ops.expand_dims_v2(v, axis=0))
+    rank += 1
+  else:
+    result = list(values)   # dtensor.pack requires a list as input.
+
+  # TODO(scottzhu): Note that the result tensor could be a partial value and
+  # not always batch shard or fully replicaed. See
+  # http://screenshot/6ERkXyX95KqftCw as an example.
+  batch_layout = layout.Layout.batch_sharded(
+      mesh, batch_dim=DEFAULT_BATCH_MESH_DIM_NAME, rank=rank)
+
+  return d_api.pack(result, batch_layout)
+
+
+def dtensor_reduce(strategy, reduce_op, value, axis):
+  """Implement dtensor based strategy.reduce()."""
+  # Due to the limitation of using scalar in DTensor (e.g. the rank 0 tensor
+  # loss the batch shard information), we need to override the default
+  # reduce in addition to the strategy.extend._reduce_to()
+  # Most of the logic here is a mimic of the parent class, except for how
+  # mean and sum are calculated in a global context.
+  distribute_lib._require_cross_replica_or_default_context_extended(  # pylint: disable=protected-access
+      strategy.extended)
+  if isinstance(reduce_op, str):
+    reduce_op = reduce_util.ReduceOp(reduce_op.upper())
+
+  distributed_input = is_distributed_value(value)
+  if not distributed_input and axis is None:
+    # For any value that isn't distributed and doesn't need a reduction within
+    # the replica.
+    destinations = (device_util.current() or
+                    strategy.extended._default_device or  # pylint: disable=protected-access
+                    "/device:CPU:0")
+    devices = cross_device_ops_lib.get_devices_from(destinations)
+    with ops.device(devices[0]):
+      return array_ops.identity(
+          cross_device_ops_lib.reduce_non_distributed_value(
+              reduce_op, value, destinations, strategy.num_replicas_in_sync))
+
+  value = convert_inputs_to_dtensor(value, strategy._mesh)  # pylint: disable=protected-access
+  # At this point, the value is a DTensor instance now.
+  # There will be a final reduction step cross replica. In order to maintain
+  # the shape of each local replica, we need to add a new dim to the front.
+  # E.g. 2 replica with local shape as (4, 5, 6), the global tensor shape
+  # should be (8, 5, 6), we will reshape into (2, 4, 5, 6) and then do a
+  # reduction on axis 0.
+  if reduce_op == reduce_util.ReduceOp.MEAN:
+    reduce_op = math_ops.reduce_mean
+  else:
+    reduce_op = math_ops.reduce_sum
+
+  # TODO(scottzhu): Make sure we handle dynamic/uneven shape in future.
+  if d_api.fetch_layout(value).is_fully_replicated():
+    # In case of fully mirrored dtensor, we only need to do one reduce, and
+    # don't need to care about any per-replica logic.
+    if axis is not None:
+      value = reduce_op(value, axis=axis)
+  else:
+    new_shape = [strategy.num_replicas_in_sync, -1]
+    if len(value.shape) > 1:
+      new_shape.extend(array_ops.shape(value)[1:])
+    value = array_ops.reshape(value, new_shape)
+    if axis is not None:
+      # we do a reduce_sum/mean within each of the replica when axis is not
+      # None. Add 1 to the axis since there is a new dim added by reshape in
+      # front.
+      value = reduce_op(value, axis=axis + 1)
+    value = reduce_op(value, axis=0)
+
+  # Note that we return a DTensor instance here, which should have the same
+  # value as the original MirroredStrategy, but with a different type. User
+  # might want a tf.Tensor for the status quo.
+  return value

tensorflow/python/distribute/experimental/dtensor_util_test.py

@@ -0,0 +1,122 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Test for DTensor related utilities in tf.distribute."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.distribute.experimental import mirrored_strategy
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+
+
+class DTensorDistributedValueTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+    tensor_1 = constant_op.constant([1.0])
+    tensor_2 = constant_op.constant([2.0])
+    self.batch_layout = layout.Layout.batch_sharded(
+        self.mesh, batch_dim='batch', rank=1)
+    self.dtensor = d_api.pack([tensor_1, tensor_2], self.batch_layout)
+
+  @parameterized.named_parameters([
+      ('py_floats', [1.0, 2.0]),
+      ('np_floats', np.array([1.0, 2.0])),
+      ('tf_const', lambda: constant_op.constant([1.0, 2.0])),
+      ('distribute_value', values_lib.PerReplica([1.0, 2.0])),
+  ])
+  def test_input_validation(self, inputs):
+    if callable(inputs):
+      inputs = inputs()
+
+    with self.assertRaisesRegex(ValueError, 'can only be built with DTensor'):
+      dtensor_util.DTensorDistributedValue(inputs)
+
+  def test_unpack(self):
+    v = dtensor_util.DTensorDistributedValue(self.dtensor)
+
+    self.assertIs(self.dtensor, v.get_dtensor())
+
+    per_replica_result = v.values
+    self.assertLen(per_replica_result, 2)
+    self.assertAllClose(per_replica_result[0], constant_op.constant([1.0]))
+    self.assertAllClose(per_replica_result[1], constant_op.constant([2.0]))
+
+  def test_graph_behavior(self):
+
+    @def_function.function
+    def run_fn(input_dtensor):
+      return dtensor_util.DTensorDistributedValue(input_dtensor)
+
+    result = run_fn(self.dtensor)
+    # When it cross the boundary of tf.function, it will be unwrapped and
+    # return a dtensor instance directly.
+    self.assertTrue(d_api.is_dtensor(result))
+    self.assertDTensorEqual(constant_op.constant([1.0, 2.0]),
+                            self.batch_layout, result)
+
+
+class DTensorReplicaContextTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+  def test_unsupported_methods(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    replica_context = dtensor_util.DTensorReplicaContext(strategy)
+
+    expected_error = replica_context._UNSUPPORTED_ERROR_MSG
+
+    self.assertEqual(replica_context.num_replicas_in_sync, 2)
+
+    self.assertEqual(replica_context.replica_id_in_sync_group, 0)
+
+    with self.assertRaisesRegex(NotImplementedError, expected_error):
+      replica_context.merge_call(None)
+
+    with self.assertRaisesRegex(NotImplementedError, expected_error):
+      replica_context.all_reduce(reduce_util.ReduceOp.SUM, None)
+
+    with self.assertRaisesRegex(NotImplementedError, expected_error):
+      replica_context.all_gather([], 0)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/experimental/mirrored_strategy.py

@@ -0,0 +1,99 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Implement a MirroredStrategy based on the DTensor low level API.
+
+This is an experiment to validate the viability of the DTensor API, and expose
+any potential feature gaps between the current API and the need.
+"""
+
+from tensorflow.dtensor.python import config as d_config
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute.experimental import dtensor_strategy_extended
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.framework import device as tf_device
+
+
+class MirroredStrategy(distribute_lib.Strategy):
+  """Synchronous training across multiple replicas on one machine.
+
+  This strategy is typically used for training on one machine with multiple
+  accelerators (GPUs/TPUs).
+
+  For example, a variable created under a `MirroredStrategy` is a distributed
+  variable with layout replicated on each dimension. The variables will be
+  placed on the `mesh` that is specified in the __init__.
+  """
+
+  def __init__(self, devices=None, cross_device_ops=None, *, mesh=None):
+    """Synchronous training across multiple replicas on one machine.
+
+    Args:
+      devices: a list of device strings, such as ['/gpu:0', '/gpu:1']. If both
+        `mesh` and `devices` are None, all the available GPU/TPU will be used.
+        If no accelerators are found, CPU is used.
+      cross_device_ops: optional, a descendant of `CrossDeviceOps`. The value is
+        ignored at the moment, and support will be added later.
+      mesh: optional DTensor mesh for the computation. Note that either `mesh`
+        or `devices` should be provided, and not both. The mesh should be 1D,
+        and will be used to split the input data among that dimension.
+    """
+    self._validate_init_args(mesh, devices)
+    if not mesh:
+      mesh = self._build_mesh_from_device_list(devices)
+
+    extended = dtensor_strategy_extended.DTensorStrategyExtended(
+        container_strategy=self, mesh=mesh)
+    super().__init__(extended)
+    self._mesh = mesh
+    self._devices = devices
+
+  @classmethod
+  def _validate_init_args(cls, mesh, devices):
+    if mesh and devices:
+      raise ValueError('Mesh and devices can not be provided at the same time. '
+                       f'received mesh = {mesh}, devices = {devices}')
+
+    # For mirrored strategy, the mesh should be 1D, and only contains a batch
+    # dimension, we will use that dimension to shard the inputs.
+    if mesh and len(mesh.shape()) != 1:
+      raise ValueError('The mesh for MirroredStrategy must be 1D, received: '
+                       f'{len(mesh.shape())}D')
+
+  @classmethod
+  def _build_mesh_from_device_list(cls, devices):
+    if devices:
+      device_type = tf_device.DeviceSpec.from_string(devices[0]).device_type
+      dtensor_util.initialize_accelerator_system_once(device_type)
+      mesh = mesh_util.create_mesh(
+          mesh_dims=[(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME, len(devices))],
+          devices=devices)
+    else:
+      # Trying to detect if there is any GPU/TPUs attached.
+      device_type = d_config.preferred_device_type()
+      devices = d_config.local_devices(device_type)
+      dtensor_util.initialize_accelerator_system_once(device_type)
+      mesh = mesh_util.create_mesh(
+          mesh_dims=[(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME, len(devices))],
+          device_type=device_type)
+    return mesh
+
+  def reduce(self, reduce_op, value, axis):
+    return dtensor_util.dtensor_reduce(self, reduce_op, value, axis)
+
+  @property
+  def mesh(self):
+    """Returns the mesh used by the strategy."""
+    return self._mesh

tensorflow/python/distribute/experimental/mirrored_strategy_test.py

@@ -0,0 +1,722 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Test for MirroredStrategy backed by DTensor API."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import d_variable
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.distribute.experimental import mirrored_strategy
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.ops import variables
+
+
+class StrategyBaseTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+  @parameterized.named_parameters([
+      ('py_floats', lambda: [1.0, 2.0], True),
+      ('np_floats', lambda: np.array([1.0, 2.0]), True),
+      ('tf_const', lambda: constant_op.constant([1.0, 2.0]), True),
+      ('py_floats_callable', lambda: [1.0, 2.0], False),
+      ('np_floats_callable', lambda: np.array([1.0, 2.0]), False),
+      ('tf_const_callable', lambda: constant_op.constant([1.0, 2.0]), False),
+  ])
+  def test_variable_creation(self, init_value, convert_callable):
+    if convert_callable:
+      init_value = init_value()
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy.scope():
+      v = variables.Variable(init_value)
+
+    self.assertIsInstance(v, d_variable.DVariable)
+    self.assertIsNotNone(v.layout)
+    self.assertEqual(v.layout, layout.Layout.replicated(self.mesh, rank=1))
+
+  def test_variable_creation_with_dtype(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy.scope():
+      v = variables.Variable(
+          0, dtype='int64',
+          aggregation=variables.VariableAggregationV2.ONLY_FIRST_REPLICA)
+    self.assertIsInstance(v, d_variable.DVariable)
+    self.assertEqual(v.dtype, dtypes.int64)
+
+  def test_mesh(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertEqual(strategy.mesh, self.mesh)
+
+  def test_strategy_extension(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertIsInstance(strategy.extended, distribute_lib.StrategyExtendedV2)
+
+  def test_num_replica_in_sync(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertEqual(strategy.num_replicas_in_sync, 2)
+
+  def test_worker_devices(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    worker_devices = strategy.extended.worker_devices
+    self.assertLen(worker_devices, 2)
+    self.assertEqual(worker_devices, tuple(self.mesh.local_devices()))
+
+  def test_parameter_devices(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    parameter_devices = strategy.extended.parameter_devices
+    self.assertLen(parameter_devices, 2)
+    self.assertEqual(parameter_devices, tuple(self.mesh.local_devices()))
+
+  def test_variable_created_in_scope(self):
+    strategy1 = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy1.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    v2 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    strategy2 = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy2.scope():
+      v3 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    self.assertTrue(strategy1.extended.variable_created_in_scope(v1))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v2))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v3))
+    self.assertTrue(strategy2.extended.variable_created_in_scope(v3))
+
+  def test_colocate_vars_with(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with strategy.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+      with strategy.extended.colocate_vars_with(v1):
+        v2 = variables.Variable(constant_op.constant([2.0, 3.0]))
+
+    # We assert the layout for the variable, and make sure they are same.
+    self.assertEqual(v1.layout, v2.layout)
+
+  def test_in_multi_worker_mode(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    self.assertFalse(strategy.extended._in_multi_worker_mode())
+
+  def test_run_with_tensor_inputs(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant(3.0)
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2.0
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.run(replica_fn, args=(tensor_input,))
+
+  def test_run_with_graph_tensor_inputs(self):
+    # Note that this is potentially a sharp edge for the user, since the eager
+    # test case was raising an error, but the graph context will run, by treat
+    # the inputs as a global inputs.
+    # TODO(scottzhu): Mitigate this eager/graph behavior difference in future.
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2.0
+
+    @def_function.function
+    def run_fn():
+      tensor_input = constant_op.constant(3.0)
+      return strategy.run(replica_fn, args=(tensor_input,))
+
+    with strategy.scope():
+      result = run_fn()
+    self.assertEqual(result, constant_op.constant(6.0))
+
+  def test_run_with_unsupported_input_types(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    random_inputs = [123, '456']
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2.0
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.run(replica_fn, args=(random_inputs,))
+
+  def test_run_with_distribute_value_input(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    def value_fn(value_context):
+      return value_context.replica_id_in_sync_group
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, 2)
+    # Note that the scalar value from
+    # experimental_distribute_values_from_function will be up rank to 1D since
+    # batched shared dtensor need at least be 1D. So the result from the
+    # strategy.run is [0], instead of just 0.
+    self.assertAllClose(result.values[0], constant_op.constant([0]))
+    self.assertAllClose(result.values[1], constant_op.constant([2]))
+
+  def test_run_without_input(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([1.0])
+
+    result = strategy.run(replica_fn)
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, 2)
+
+    self.assertAllClose(result.values[0], constant_op.constant([1.0]))
+    self.assertAllClose(result.values[1], constant_op.constant([1.0]))
+
+  def test_nested_structure_output(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    array_value = np.array([3., 2., 1.])
+    def value_fn(ctx):
+      value = array_value[ctx.replica_id_in_sync_group]
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertLen(result.keys(), 2)
+    self.assertIsInstance(result['a'], dtensor_util.DTensorDistributedValue)
+    self.assertAllClose(result['a'].values[0], constant_op.constant([6.0]))
+    self.assertAllClose(result['a'].values[1], constant_op.constant([4.0]))
+
+    self.assertIsInstance(result['b'], dtensor_util.DTensorDistributedValue)
+    self.assertAllClose(result['b'].values[0],
+                        constant_op.constant([8.0, 10.0]))
+    self.assertAllClose(result['b'].values[1], constant_op.constant([6.0, 8.0]))
+
+  def test_inputs_with_dtensor_distribute_values(self):
+
+    @def_function.function
+    def replica_fn_1(inputs):
+      return inputs * 2.0
+
+    @def_function.function
+    def replica_fn_2(inputs):
+      return inputs + 1.0
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant(3.0)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    result_1 = strategy.run(replica_fn_1, args=(d_tensor_input,))
+    self.assertIsInstance(result_1, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_1.values, 2)
+    self.assertAllClose(result_1.values[0], constant_op.constant([6.0]))
+    self.assertAllClose(result_1.values[1], constant_op.constant([6.0]))
+
+    result_2 = strategy.run(replica_fn_2, args=(result_1,))
+    self.assertIsInstance(result_2, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_2.values, 2)
+    self.assertAllClose(result_2.values[0], constant_op.constant([7.0]))
+    self.assertAllClose(result_2.values[1], constant_op.constant([7.0]))
+
+  def test_run_with_nullary_ops(self):
+
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([3.0])
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    result = strategy.run(replica_fn)
+
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertAllClose(result.values[0], constant_op.constant([3.0]))
+    self.assertAllClose(result.values[1], constant_op.constant([3.0]))
+
+  def test_get_replica_context(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    tensor_input = constant_op.constant(3)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    @def_function.function
+    def replica_fn(inputs):
+      replica_context = distribute_lib.get_replica_context()
+      self.assertIsInstance(replica_context, dtensor_util.DTensorReplicaContext)
+      return inputs * replica_context.num_replicas_in_sync
+
+    # Default replica context
+    self.assertIsNotNone(distribute_lib.get_replica_context())
+    with strategy.scope():
+      self.assertIsNone(distribute_lib.get_replica_context())
+
+      result = strategy.run(replica_fn, args=(d_tensor_input,))
+
+    self.assertLen(result.values, 2)
+    self.assertAllClose(result.values[0], constant_op.constant([6]))
+    self.assertAllClose(result.values[1], constant_op.constant([6]))
+
+  def test_gather_non_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant(3.0)
+
+    result = strategy.gather(tensor_input, axis=0)
+    self.assertAllClose(result, tensor_input)
+
+  def test_gather_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    def value_fn(value_context):
+      start = value_context.replica_id_in_sync_group
+      return array_ops.reshape(math_ops.range(start=start, limit=start + 6),
+                               shape=(1, 2, 3))
+
+    distribute_result = strategy.experimental_distribute_values_from_function(
+        value_fn)
+    result = strategy.gather(distribute_result, axis=0)
+    self.assertEqual(result.shape, [2, 2, 3])
+    self.assertAllClose(result, [[[0, 1, 2], [3, 4, 5]],
+                                 [[1, 2, 3], [4, 5, 6]]])
+    result = strategy.gather(distribute_result, axis=1)
+    self.assertEqual(result.shape, [1, 4, 3])
+    self.assertAllClose(result, [[[0, 1, 2], [3, 4, 5], [1, 2, 3], [4, 5, 6]]])
+    result = strategy.gather(distribute_result, axis=2)
+    self.assertEqual(result.shape, [1, 2, 6])
+    self.assertAllClose(result, [[[0, 1, 2, 1, 2, 3], [3, 4, 5, 4, 5, 6]]])
+
+  def test_reduce_mean_non_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=0)
+
+  def test_reduce_sum_non_dtensor_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'):
+      strategy.reduce(reduce_util.ReduceOp.SUM, tensor_input, axis=0)
+
+  def test_reduce_mean_distribute_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      n = value_context.num_replicas_in_sync
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * n * 2.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn)
+    # replica 1 has [[0.0, 1.0],[2.0, 3.0]] and replica 2 has
+    # [[4.0, 5.0],[6.0, 7.0]]
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=None)
+    self.assertAllClose(result, constant_op.constant([[2.0, 3.0], [4.0, 5.0]]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=0)
+    self.assertAllClose(result, constant_op.constant([3.0, 4.0]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=1)
+    self.assertAllClose(result, constant_op.constant([2.5, 4.5]))
+
+  def test_reduce_sum_distribute_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      n = value_context.num_replicas_in_sync
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * n * 2.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn)
+    # replica 1 has [[0.0, 1.0],[2.0, 3.0]] and replica 2 has
+    # [[4.0, 5.0],[6.0, 7.0]]
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=None)
+    self.assertAllClose(result, constant_op.constant([[4.0, 6.0], [8.0, 10.0]]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=0)
+    self.assertAllClose(result, constant_op.constant([12.0, 16.0]))
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=1)
+    self.assertAllClose(result, constant_op.constant([10.0, 18.0]))
+
+  def test_reduce_mean_mirrored_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    with  strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([2.0, 3.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([1.5, 3.5]))
+
+  def test_reduce_sum_mirrored_value(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+
+    with  strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([4.0, 6.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([3.0, 7.0]))
+
+  def test_reduce_value_device(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=None)
+    self.assertIn('CPU:0', result.device)
+
+  def test_experimental_local_results(self):
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([3.0])
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    result = strategy.run(replica_fn)
+    local_result = strategy.experimental_local_results(result)
+
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, 2)
+    self.assertEqual(local_result[0], constant_op.constant([3.0]))
+    self.assertEqual(local_result[1], constant_op.constant([3.0]))
+
+  def test_experimental_local_results_with_inputs(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    array_value = np.array([3., 2.])
+    def value_fn(ctx):
+      value = array_value[ctx.replica_id_in_sync_group]
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    local_result = strategy.experimental_local_results(result)
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, 2)
+    self.assertDictEqual(local_result[0],
+                         {'a': constant_op.constant([6.0]),
+                          'b': constant_op.constant([8.0, 10.0])})
+    self.assertDictEqual(local_result[1],
+                         {'a': constant_op.constant([4.0]),
+                          'b': constant_op.constant([6.0, 8.0])})
+
+
+class InvalidMeshTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2, 1))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch', 'model'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh_2d = self.configTestMesh(mesh_dict)
+
+  def test_invalid_mesh_shape(self):
+    with self.assertRaisesRegex(
+        ValueError, 'The mesh for MirroredStrategy must be 1D, received: 2D'):
+      mirrored_strategy.MirroredStrategy(mesh=self.mesh_2d)
+
+
+class StrategyCreationTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    device_type = test_util.preferred_device_type()
+    if device_type != 'TPU':
+      test_util.reset_logical_devices(device_type, 2)
+    self.device_type = device_type
+
+  def test_explicit_device_list(self):
+    device_list = [f'/{self.device_type}:{i}' for i in range(2)]
+    strategy = mirrored_strategy.MirroredStrategy(devices=device_list)
+    mesh = strategy.mesh
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.shape(), [2,])
+    self.assertEqual(mesh.dim_names, ['batch'])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:0',
+        mesh.local_devices()[0])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:1',
+        mesh.local_devices()[1])
+    # Also make sure the host mesh works since it is required by dataset
+    self.assertIsNotNone(mesh.host_mesh())
+
+  def test_implicit_device_list(self):
+    strategy = mirrored_strategy.MirroredStrategy()
+    mesh = strategy.mesh
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.shape(), [2,])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:0',
+        mesh.local_devices()[0])
+    self.assertIn(
+        f'/job:localhost/replica:0/task:0/device:{self.device_type}:1',
+        mesh.local_devices()[1])
+    # Also make sure the host mesh works since it is required by dataset
+    self.assertIsNotNone(mesh.host_mesh())
+
+  def test_mesh_with_device_list(self):
+    device_list = [f'/{self.device_type}:{i}' for i in range(2)]
+    mesh = mesh_util.create_mesh([('batch', 2)], devices=device_list)
+    with self.assertRaisesRegex(
+        ValueError, 'Mesh and devices can not be provided at the same time'):
+      _ = mirrored_strategy.MirroredStrategy(mesh=mesh, devices=device_list)
+
+
+class StrategyDatasetTest(test_util.DTensorBaseTest):
+
+  def setUp(self):
+    super().setUp()
+    global_ids = test_util.create_device_ids_array((2,))
+    local_ids = np.ravel(global_ids).tolist()
+    mesh_dict = {
+        device: layout.Mesh(['batch'], global_ids, local_ids,
+                            test_util.create_device_list((2,), device))
+        for device in ['TPU', 'GPU', 'CPU']
+    }
+    self.mesh = self.configTestMesh(mesh_dict)
+
+    self.images = stateless_random_ops.stateless_random_uniform(
+        [8, 8, 3], seed=(1, 2), minval=0, maxval=255)
+    self.labels = stateless_random_ops.stateless_random_uniform(
+        [1], seed=(1, 2), minval=0, maxval=10)
+
+    self.dataset = dataset_ops.Dataset.from_tensors(
+        (self.images, self.labels)).repeat()
+
+  def test_create_batched_dataset(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    global_batch_size = 8
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    element = next(iter(distributed_dataset))
+    batched_image, batched_label = element
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # Make sure when unpack the tensor, each of them has enough shards.
+    self.assertLen(d_api.unpack(batched_image), self.mesh.num_local_devices())
+    self.assertLen(d_api.unpack(batched_label), self.mesh.num_local_devices())
+
+  def test_uneven_batched_dataset(self):
+    elements = [[1, 2, 3], [1, 2], [1, 2, 3, 4]]
+    dataset = dataset_ops.Dataset.from_generator(
+        lambda: elements, dtypes.int64).repeat()
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with self.assertRaisesRegex(ValueError, 'requires a static batch size'):
+      strategy.experimental_distribute_dataset(dataset)
+
+  def test_create_partial_batched_dataset(self):
+    # TODO(b/210887657): Support last partial batch.
+    self.skipTest('Test failed due to last partial batch')
+    dataset = dataset_ops.Dataset.from_tensors(
+        (self.images, self.labels)).repeat(30)  # There is a last partial batch
+
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    global_batch_size = 8
+    dataset = dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    expected_element_batch_size = [8, 8, 8, 6]
+    # The last batch with 6 element will fail to produce with StopIteration.
+    iterator = iter(distributed_dataset)
+    for batch_size in expected_element_batch_size:
+      element = next(iterator)
+      batched_image, batched_label = element
+      self.assertEqual(batched_image.shape, [batch_size, 8, 8, 3])
+      self.assertEqual(batched_label.shape, [batch_size, 1])
+
+      # Make sure when unpack the tensor, each of them has enough shards.
+      self.assertLen(d_api.unpack(batched_image), self.mesh.num_local_devices())
+      self.assertLen(d_api.unpack(batched_label), self.mesh.num_local_devices())
+
+  def test_deprecated_strategy_methods(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_dataset_iterator(self.dataset)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_input_fn_iterator(lambda _: self.dataset)
+
+  def test_distribute_dataset_from_fn(self):
+    local_batch_size = 4
+    global_batch_size = 8
+    def dataset_fn(option):
+      del option
+      return dataset_ops.Dataset.from_tensors(
+          (self.images, self.labels)).repeat().batch(
+              local_batch_size, drop_remainder=True).prefetch(2)
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    distributed_dataset = strategy.distribute_datasets_from_function(
+        dataset_fn, None)
+    iterator = iter(distributed_dataset)
+
+    self.assertEqual(distributed_dataset.element_spec,
+                     (tensor_spec.TensorSpec(shape=(8, 8, 8, 3),
+                                             dtype=dtypes.float32, name=None),
+                      tensor_spec.TensorSpec(shape=(8, 1),
+                                             dtype=dtypes.float32, name=None)))
+    self.assertEqual(distributed_dataset.element_spec, iterator.element_spec)
+
+    batched_image, batched_label = next(iterator)
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # Make sure there are two shards when unpack, and each of them has 4 as
+    # batch size
+    unpacked_images = d_api.unpack(batched_image)
+    self.assertLen(unpacked_images, self.mesh.num_local_devices())
+    self.assertEqual(unpacked_images[0].shape, [local_batch_size, 8, 8, 3])
+    self.assertEqual(unpacked_images[1].shape, [local_batch_size, 8, 8, 3])
+
+  def test_distribute_values_from_function(self):
+    array_value = np.array([3., 2., 1.])
+    def value_fn(ctx):
+      return array_value[ctx.replica_id_in_sync_group]
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+    self.assertDTensorEqual(
+        constant_op.constant([3., 2.], dtype=dtypes.float64),
+        layout.Layout.batch_sharded(self.mesh, batch_dim='batch', rank=1),
+        distributed_values)
+
+  def test_distribute_values_from_function_with_nested_structure(self):
+    array_value = np.array([3., 2., 1.])
+    def value_fn(ctx):
+      value = array_value[ctx.replica_id_in_sync_group]
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+    self.assertIsInstance(distributed_values, dict)
+    self.assertDTensorEqual(
+        constant_op.constant([3., 2.], dtype=dtypes.float64),
+        layout.Layout.batch_sharded(self.mesh, batch_dim='batch', rank=1),
+        distributed_values['a'])
+    unpacked_a = d_api.unpack(distributed_values['a'])
+    # Note that this might have a slight behavior difference, the original
+    # mirrored strategy may return scalar for each PerReplica. The DTensor
+    # implementation is more strict and always ensures the PerReplica
+    # value has the same rank as the global-view Tensor.
+    self.assertAllClose(unpacked_a[0], [3.])
+    self.assertAllClose(unpacked_a[1], [2.])
+    self.assertDTensorEqual(
+        constant_op.constant([4., 5., 3., 4.], dtype=dtypes.float64),
+        layout.Layout.batch_sharded(self.mesh, batch_dim='batch', rank=1),
+        distributed_values['b'])
+
+  def test_distribute_dataset_in_tf_function(self):
+    strategy = mirrored_strategy.MirroredStrategy(mesh=self.mesh)
+    local_batch_size = 4
+    global_batch_size = 8
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+
+    @def_function.function
+    def step_fn(iterator):
+      images, labels = next(iterator)
+      del labels
+      return images
+
+    result = strategy.run(step_fn, args=(iter(distributed_dataset),))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, self.mesh.num_local_devices())
+    self.assertEqual(result.values[0].shape, [local_batch_size, 8, 8, 3])
+    self.assertEqual(result.values[1].shape, [local_batch_size, 8, 8, 3])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/distribute/experimental/multi_worker_mirrored_strategy.py

@@ -0,0 +1,156 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Implement a MultiMirroredStrategy based on the DTensor low level API.
+
+This is an experiment to validate the viability of the DTensor API, and expose
+any potential feature gaps between the current API and the need.
+"""
+
+import os
+
+from tensorflow.dtensor.python import config as d_config
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.distribute.experimental import dtensor_strategy_extended
+from tensorflow.python.distribute.experimental import dtensor_util
+
+
+class MultiWorkerMirroredStrategy(distribute_lib.Strategy):
+  """A distribution strategy for synchronous training on multiple workers.
+
+  This strategy implements synchronous distributed training across multiple
+  workers, each with potentially multiple GPUs. Similar to
+  `tf.distribute.MirroredStrategy`, it replicates all variables and computations
+  to each local device. The difference is that it uses a distributed collective
+  implementation (e.g. all-reduce), so that multiple workers can work together.
+  """
+
+  def __init__(self, cluster_resolver=None, communication_options=None, *,
+               mesh=None):
+    """Creates the strategy.
+
+    Args:
+      cluster_resolver: optional
+        `tf.distribute.cluster_resolver.ClusterResolver`. In case neither `mesh`
+        nor `cluster_resolver` are provided,
+        `tf.distribute.cluster_resolver.TFConfigClusterResolver` is used.
+      communication_options: currently ignore.
+      mesh: optional Dtensor global mesh for the computation. Note that either
+        `mesh` or the `cluster_resolver` should be provided. and not both.
+    """
+    self._validate_init_args(mesh, cluster_resolver)
+    if not mesh:
+      if not cluster_resolver:
+        # Use the TFConfigClusterResolver as default
+        cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+      dtensor_env_var = _parse_dtensor_env_var_from_cluster_resolver(
+          cluster_resolver)
+      _config_dtensor_env_var(dtensor_env_var)
+      mesh = _build_distributed_mesh(dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME)
+    extended = dtensor_strategy_extended.DTensorStrategyExtended(
+        container_strategy=self, mesh=mesh)
+    super().__init__(extended)
+    self._mesh = mesh
+    self._cluster_resolver = cluster_resolver
+
+  @classmethod
+  def _validate_init_args(cls, mesh, cluster_resolver):
+    if mesh and cluster_resolver:
+      raise ValueError('Mesh and cluster_resolver can not be provided at the '
+                       f'same time. Received mesh = {mesh}, cluster_resolver = '
+                       f'{cluster_resolver}')
+    if mesh and len(mesh.shape()) != 1:
+      raise ValueError('The mesh for MultiWorkerMirroredStrategy must be 1D, '
+                       f'received: {len(mesh.shape())}D')
+
+  def reduce(self, reduce_op, value, axis):
+    return dtensor_util.dtensor_reduce(self, reduce_op, value, axis)
+
+  @property
+  def mesh(self):
+    """Returns the mesh used by the strategy."""
+    return self._mesh
+
+
+def _parse_dtensor_env_var_from_cluster_resolver(cluster_resolver):
+  """Parse the env vars for Dtensor based on the cluster resolver.
+
+  In the multi-client setting, each of the DTensor jobs need to aware of each
+  other, and the interface to setup those values are via the envvars. The
+  value used by dtensor are different from the existing
+  `MultiWorkerMirroredStrategy`. This function will parse the value from
+  cluster resolver, and populate the corresponding value for DTensor jobs in the
+  `os.environ`.
+
+  Args:
+    cluster_resolver: A `tf.distribute.cluster_resolver.ClusterResolver`
+      instance.
+
+  Returns:
+    A dict of {Str:Str} which contains all the env vars needed by DTensor jobs.
+    The value is for verification purpose.
+
+  Raises:
+    The value parsed from existing cluster spec is not valid.
+  """
+  result = {}
+
+  # Retrieve the number of host, cluster config from the resolver.
+  cluster_spec = multi_worker_util.normalize_cluster_spec(
+      cluster_resolver.cluster_spec())
+  # Export all the necessary envvars for dtensor
+  # Get all the jobs from the cluster spec. Note that the in the normal
+  # setting, it could be multiple worker devices without chief, and the
+  # worker 0 will be the chief, or an explicit chief with multiple worker job.
+  dtensor_jobs = []
+  if 'chief' in cluster_spec.jobs:
+    dtensor_jobs.extend(cluster_spec.job_tasks('chief'))
+  if 'worker' in cluster_spec.jobs:
+    dtensor_jobs.extend(cluster_spec.job_tasks('worker'))
+
+  if None in dtensor_jobs:
+    raise ValueError('Unexpected dtensor job address from cluster spec: '
+                     f'{cluster_spec}')
+  result['DTENSOR_JOBS'] = ','.join(dtensor_jobs)
+  result['DTENSOR_NUM_CLIENTS'] = str(len(dtensor_jobs))
+
+  if cluster_resolver.task_type == 'chief':
+    dtensor_client_id = 0
+  elif cluster_resolver.task_type == 'worker':
+    dtensor_client_id = cluster_resolver.task_id
+    if 'chief' in cluster_spec.jobs:
+      dtensor_client_id += 1
+  result['DTENSOR_CLIENT_ID'] = str(dtensor_client_id)
+  result['DTENSOR_JOB_NAME'] = 'worker'
+
+  return result
+
+
+def _config_dtensor_env_var(dtensor_env_vars):
+  for k, v in dtensor_env_vars.items():
+    os.environ[k] = v
+
+
+def _build_distributed_mesh(batch_dim_name):
+  device_type = d_config.preferred_device_type()
+  local_devices = d_config.local_devices(device_type)
+  number_clients = d_config.num_clients()
+  dtensor_util.initialize_accelerator_system_once(device_type)
+  # This assumes each client has same number of devices.
+  mesh_dims = [(batch_dim_name, len(local_devices) * number_clients)]
+  return mesh_util.create_distributed_mesh(
+      mesh_dims, device_type=device_type)

tensorflow/python/distribute/experimental/multi_worker_mirrored_strategy_test.py

@@ -0,0 +1,647 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Test for MultiWorkerMirroredStrategy backed by DTensor API."""
+
+import json
+import os
+
+from absl import flags
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.dtensor.python import api as d_api
+from tensorflow.dtensor.python import d_variable
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python.tests import multi_client_test_util
+from tensorflow.dtensor.python.tests import test_backend_util
+from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.distribute.experimental import dtensor_util
+from tensorflow.python.distribute.experimental import multi_worker_mirrored_strategy as mwms
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test as tf_test
+
+
+class MultiWorkerMirroredStrategyTest(tf_test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.num_client = flags.FLAGS.num_clients
+    self.num_local_devices = flags.FLAGS.num_local_devices
+
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    self.client_id = int(tf_config['task']['index'])
+
+  def test_strategy_creation_with_default_cluster_resolver(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    mesh = strategy.mesh
+    self.assertIsNotNone(mesh)
+    self.assertLen(mesh.global_device_ids(),
+                   self.num_client * self.num_local_devices)
+    self.assertLen(mesh.local_device_ids(), self.num_local_devices)
+    self.assertIsInstance(strategy._cluster_resolver,
+                          tfconfig_cluster_resolver.TFConfigClusterResolver)
+
+  def test_invalid_init_arguments(self):
+    mesh = object()
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+
+    with self.assertRaisesRegex(
+        ValueError,
+        'Mesh and cluster_resolver can not be provided at the same time'):
+      mwms.MultiWorkerMirroredStrategy(
+          mesh=mesh,
+          cluster_resolver=cluster_resolver)
+
+  def test_parse_dtensor_env_var_from_cluster_resolver(self):
+    cluster_resolver = tfconfig_cluster_resolver.TFConfigClusterResolver()
+
+    dtensor_env_vars = mwms._parse_dtensor_env_var_from_cluster_resolver(
+        cluster_resolver)
+
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    worker_jobs = ','.join(tf_config['cluster']['worker'])
+    client_id = tf_config['task']['index']
+
+    self.assertLen(dtensor_env_vars, 4)
+    self.assertEqual(dtensor_env_vars['DTENSOR_JOBS'], worker_jobs)
+    self.assertEqual(dtensor_env_vars['DTENSOR_NUM_CLIENTS'],
+                     str(self.num_client))
+    self.assertEqual(dtensor_env_vars['DTENSOR_CLIENT_ID'], client_id)
+    self.assertEqual(dtensor_env_vars['DTENSOR_JOB_NAME'], 'worker')
+
+  @parameterized.named_parameters([
+      ('py_floats', lambda: [1.0, 2.0], True),
+      ('np_floats', lambda: np.array([1.0, 2.0]), True),
+      ('tf_const', lambda: constant_op.constant([1.0, 2.0]), True),
+      ('py_floats_callable', lambda: [1.0, 2.0], False),
+      ('np_floats_callable', lambda: np.array([1.0, 2.0]), False),
+      ('tf_const_callable', lambda: constant_op.constant([1.0, 2.0]), False),
+  ])
+  def test_variable_creation(self, init_value, convert_callable):
+    if convert_callable:
+      init_value = init_value()
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+      v = variables.Variable(init_value)
+
+    self.assertIsInstance(v, d_variable.DVariable)
+    self.assertIsNotNone(v.layout)
+    self.assertEqual(v.layout, layout.Layout.replicated(strategy.mesh, rank=1))
+
+  def test_strategy_extension(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertIsInstance(strategy.extended, distribute_lib.StrategyExtendedV2)
+
+  def test_num_replica_in_sync(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertEqual(strategy.num_replicas_in_sync,
+                     self.num_client * self.num_local_devices)
+
+  def test_mesh(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertIsNotNone(strategy.mesh)
+
+  def test_worker_devices(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    worker_devices = strategy.extended.worker_devices
+    self.assertLen(worker_devices, self.num_local_devices)
+    self.assertEqual(worker_devices, tuple(strategy.mesh.local_devices()))
+
+  def test_parameter_devices(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    parameter_devices = strategy.extended.parameter_devices
+    self.assertLen(parameter_devices, self.num_local_devices)
+    self.assertEqual(parameter_devices, tuple(strategy.mesh.local_devices()))
+
+  def test_variable_created_in_scope(self):
+    strategy1 = mwms.MultiWorkerMirroredStrategy()
+    with strategy1.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    v2 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    strategy2 = mwms.MultiWorkerMirroredStrategy()
+    with strategy2.scope():
+      v3 = variables.Variable(constant_op.constant([1.0, 2.0]))
+
+    self.assertTrue(strategy1.extended.variable_created_in_scope(v1))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v2))
+    self.assertFalse(strategy1.extended.variable_created_in_scope(v3))
+    self.assertTrue(strategy2.extended.variable_created_in_scope(v3))
+
+  def test_colocate_vars_with(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with strategy.scope():
+      v1 = variables.Variable(constant_op.constant([1.0, 2.0]))
+      with strategy.extended.colocate_vars_with(v1):
+        v2 = variables.Variable(constant_op.constant([2.0, 3.0]))
+
+    # We assert the layout for the variable, and make sure they are same.
+    self.assertEqual(v1.layout, v2.layout)
+
+  def test_in_multi_worker_mode(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    self.assertTrue(strategy.extended._in_multi_worker_mode())
+
+  def test_run_with_distribute_value_input(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    def value_fn(value_context):
+      return value_context.replica_id_in_sync_group
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      return inputs * 2
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, self.num_local_devices)
+    # Note that the scalar value from
+    # experimental_distribute_values_from_function will be up rank to 1D since
+    # batched shared dtensor need at least be 1D.
+
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result.values[i],
+          constant_op.constant(
+              [(self.client_id * self.num_local_devices + i) * 2]))
+
+  def test_nested_structure_output(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    def value_fn(ctx):
+      value = float(ctx.num_replicas_in_sync)
+      return {'a': value,
+              'b': constant_op.constant([value + 1.0, value + 2.0])}
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    self.assertLen(result.keys(), 2)
+    self.assertIsInstance(result['a'], dtensor_util.DTensorDistributedValue)
+    self.assertLen(result['a'].values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result['a'].values[i],
+          constant_op.constant([strategy.num_replicas_in_sync * 2.0]))
+
+    self.assertIsInstance(result['b'], dtensor_util.DTensorDistributedValue)
+    self.assertLen(result['b'].values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result['b'].values[i],
+          constant_op.constant([(strategy.num_replicas_in_sync + 1.0) * 2.0,
+                                (strategy.num_replicas_in_sync + 2.0) * 2.0]))
+
+  def test_inputs_with_dtensor_distribute_values(self):
+
+    @def_function.function
+    def replica_fn_1(inputs):
+      return inputs * 2.0
+
+    @def_function.function
+    def replica_fn_2(inputs):
+      return inputs + 1.0
+
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant(3.0)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    result_1 = strategy.run(replica_fn_1, args=(d_tensor_input,))
+    self.assertIsInstance(result_1, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_1.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(result_1.values[i], constant_op.constant([6.0]))
+
+    result_2 = strategy.run(replica_fn_2, args=(result_1,))
+    self.assertIsInstance(result_2, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result_2.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(result_2.values[i], constant_op.constant([7.0]))
+
+  def test_get_replica_context(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    tensor_input = constant_op.constant(3)
+    d_tensor_input = strategy.experimental_distribute_values_from_function(
+        lambda _: tensor_input)
+
+    @def_function.function
+    def replica_fn(inputs):
+      replica_context = distribute_lib.get_replica_context()
+      self.assertIsInstance(replica_context, dtensor_util.DTensorReplicaContext)
+      return inputs * replica_context.num_replicas_in_sync
+
+    # Default replica context
+    self.assertIsNotNone(distribute_lib.get_replica_context())
+    with strategy.scope():
+      self.assertIsNone(distribute_lib.get_replica_context())
+
+      result = strategy.run(replica_fn, args=(d_tensor_input,))
+
+    self.assertLen(result.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertAllClose(
+          result.values[i],
+          constant_op.constant([3 * strategy.num_replicas_in_sync]))
+
+  def test_gather_non_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant(3.0)
+
+    result = strategy.gather(tensor_input, axis=0)
+    self.assertAllClose(result, tensor_input)
+
+  def test_gather_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    stride = self.num_client * self.num_local_devices
+
+    def value_fn(value_context):
+      start = value_context.replica_id_in_sync_group * stride
+      return array_ops.reshape(
+          math_ops.range(start=start, limit=start + stride), shape=(1, stride)
+      )
+    distribute_result = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+    # distribute_result is a DTensorDistributedValue.
+    # The shape of the global tensor is [stride, stride],
+    # and each worker gets [stride/2, stride].
+    result = strategy.gather(distribute_result, axis=0)
+
+    start = stride * self.num_local_devices * self.client_id
+    end = start + stride * self.num_local_devices
+    self.assertEqual(result.shape, [self.num_local_devices, stride])
+    self.assertAllClose(
+        result,
+        array_ops.reshape(
+            math_ops.range(start=start, limit=end),
+            shape=(self.num_local_devices, -1),
+        ),
+    )
+
+    result = strategy.gather(distribute_result, axis=1)
+    self.assertEqual(result.shape, [1, self.num_local_devices * stride])
+    self.assertAllClose(
+        result,
+        array_ops.reshape(
+            math_ops.range(start=start, limit=end), shape=(1, -1)
+        ),
+    )
+
+  def test_reduce_mean_non_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'
+    ):
+      strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=0)
+
+  def test_reduce_sum_non_dtensor_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported input types for MirroredStrategy.'
+    ):
+      strategy.reduce(reduce_util.ReduceOp.SUM, tensor_input, axis=0)
+
+  def test_reduce_mean_distribute_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * 4.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+    # replica 0 has [[0.0, 1.0],[2.0, 3.0]] and
+    # replica 1 has [[4.0, 5.0],[6.0, 7.0]]. Each worker has 4 replicas.
+    # For worker 2, it has replica 4 ~ 7.
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=None
+    )
+    # This should be a global reduce and each worker should have same value.
+    # [[14.0, 15.0],[16.0, 17.0]]
+    final = (self.num_local_devices * self.num_client - 1) * 2.0
+    self.assertAllClose(
+        result, constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + final
+    )
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=0
+    )
+    # [15.0, 16.0]
+    self.assertAllClose(result, constant_op.constant([0.0, 1.0]) + final + 1)
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.MEAN, distribute_value, axis=1
+    )
+
+    self.assertAllClose(result, constant_op.constant([0.5, 2.5]) + final)
+
+  def test_reduce_sum_distribute_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    @def_function.function
+    def value_fn(value_context):
+      i = value_context.replica_id_in_sync_group
+      return constant_op.constant([[0.0, 1.0], [2.0, 3.0]]) + i * 4.0
+
+    distribute_value = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+    # replica 0 has [[0.0, 1.0],[2.0, 3.0]] and
+    # replica 1 has [[4.0, 5.0],[6.0, 7.0]]. Each worker has 4 replicas.
+    # For worker 2, it has replica 4 ~ 7.
+    # The shape of the global tensor is [16, 2], and each worker gets [8, 2].
+
+    result = strategy.reduce(
+        reduce_util.ReduceOp.SUM, distribute_value, axis=None
+    )
+    self.assertAllClose(result, [[112.0, 120.0], [128.0, 136.0]])
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, distribute_value, axis=0)
+    self.assertAllClose(result, constant_op.constant([240.0, 256.0]))
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, distribute_value, axis=1)
+    self.assertAllClose(result, constant_op.constant([232.0, 264.0]))
+
+  def test_reduce_mean_mirrored_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([2.0, 3.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([1.5, 3.5]))
+
+  def test_reduce_sum_mirrored_value(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    with strategy.scope():
+      v = variables.Variable(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    self.assertIsInstance(v, d_variable.DVariable)
+
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=None)
+    self.assertAllClose(result, constant_op.constant([[1.0, 2.0], [3.0, 4.0]]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=0)
+    self.assertAllClose(result, constant_op.constant([4.0, 6.0]))
+    result = strategy.reduce(reduce_util.ReduceOp.SUM, v, axis=1)
+    self.assertAllClose(result, constant_op.constant([3.0, 7.0]))
+
+  def test_reduce_value_device(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    tensor_input = constant_op.constant([[3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    result = strategy.reduce(reduce_util.ReduceOp.MEAN, tensor_input, axis=None)
+    self.assertIn('CPU:0', result.device)
+
+  def test_experimental_local_results(self):
+    @def_function.function
+    def replica_fn():
+      return constant_op.constant([3.0])
+
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    result = strategy.run(replica_fn)
+    local_result = strategy.experimental_local_results(result)
+
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertEqual(local_result[i], constant_op.constant([3.0]))
+
+  def test_experimental_local_results_with_inputs(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+
+    def value_fn(ctx):
+      value = float(ctx.num_replicas_in_sync)
+      return {'a': value, 'b': constant_op.constant([value + 1.0, value + 2.0])}
+
+    distributed_values = strategy.experimental_distribute_values_from_function(
+        value_fn
+    )
+
+    @def_function.function
+    def replica_fn(inputs):
+      result = {}
+      for key in inputs:
+        result[key] = inputs[key] * 2.0
+      return result
+
+    result = strategy.run(replica_fn, args=(distributed_values,))
+    local_result = strategy.experimental_local_results(result)
+    self.assertIsInstance(local_result, tuple)
+    self.assertLen(local_result, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertDictEqual(
+          local_result[i],
+          {
+              'a': constant_op.constant([16.0]),
+              'b': constant_op.constant([18.0, 20.0]),
+          },
+      )
+
+
+class StrategyDatasetTest(tf_test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.num_client = flags.FLAGS.num_clients
+    self.num_local_devices = flags.FLAGS.num_local_devices
+
+    tf_config = json.loads(os.environ['TF_CONFIG'])
+    self.client_id = int(tf_config['task']['index'])
+
+    self.images = stateless_random_ops.stateless_random_uniform(
+        [8, 8, 3], seed=(1, 2), minval=0, maxval=255)
+    self.labels = stateless_random_ops.stateless_random_uniform(
+        [1], seed=(1, 2), minval=0, maxval=10)
+
+    self.dataset = dataset_ops.Dataset.from_tensors(
+        (self.images, self.labels)).repeat()
+
+  def test_create_batched_dataset(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    global_batch_size = self.num_client * self.num_local_devices * 2
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+    element = next(iter(distributed_dataset))
+    batched_image, batched_label = element
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # After unpack, it should only get the local shards.
+    self.assertLen(d_api.unpack(batched_image), self.num_local_devices)
+    self.assertLen(d_api.unpack(batched_label), self.num_local_devices)
+
+  def test_uneven_batched_dataset(self):
+    elements = [[1, 2, 3], [1, 2], [1, 2, 3, 4]]
+    dataset = dataset_ops.Dataset.from_generator(
+        lambda: elements, dtypes.int64).repeat()
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with self.assertRaisesRegex(ValueError, 'requires a static batch size'):
+      strategy.experimental_distribute_dataset(dataset)
+
+  def test_deprecated_strategy_methods(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_dataset_iterator(self.dataset)
+
+    with self.assertRaisesRegex(
+        NotImplementedError, 'only available in the V1 API'):
+      strategy.make_input_fn_iterator(lambda _: self.dataset)
+
+  def test_distribute_dataset_from_fn(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    local_batch_size = 4
+    global_batch_size = local_batch_size * strategy.num_replicas_in_sync
+
+    def dataset_fn(option):
+      del option
+      return dataset_ops.Dataset.from_tensors(
+          (self.images, self.labels)).repeat().batch(
+              local_batch_size, drop_remainder=True).prefetch(2)
+
+    distributed_dataset = strategy.distribute_datasets_from_function(
+        dataset_fn, None)
+    iterator = iter(distributed_dataset)
+
+    self.assertEqual(distributed_dataset.element_spec,
+                     (tensor_spec.TensorSpec(shape=(global_batch_size, 8, 8, 3),
+                                             dtype=dtypes.float32, name=None),
+                      tensor_spec.TensorSpec(shape=(global_batch_size, 1),
+                                             dtype=dtypes.float32, name=None)))
+    self.assertEqual(distributed_dataset.element_spec, iterator.element_spec)
+
+    batched_image, batched_label = next(iterator)
+    self.assertEqual(batched_image.shape, [global_batch_size, 8, 8, 3])
+    self.assertEqual(batched_label.shape, [global_batch_size, 1])
+
+    # After unpack, it should only get the local shards.
+    unpacked_images = d_api.unpack(batched_image)
+    self.assertLen(unpacked_images, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertEqual(unpacked_images[i].shape, [local_batch_size, 8, 8, 3])
+
+  def test_distribute_values_from_function(self):
+    array_value = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0])
+    def value_fn(ctx):
+      return array_value[ctx.replica_id_in_sync_group]
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    distributed_values = (
+        strategy.experimental_distribute_values_from_function(
+            value_fn))
+
+    self.assertEqual(d_api.fetch_layout(distributed_values),
+                     layout.Layout.batch_sharded(
+                         strategy.mesh, batch_dim='batch', rank=1))
+    unpacked_value = d_api.unpack(distributed_values)
+    self.assertLen(unpacked_value, self.num_local_devices)
+    start = 1.0 + self.num_local_devices * self.client_id
+    for i in range(self.num_local_devices):
+      self.assertEqual(unpacked_value[i], start + i)
+
+  def test_distribute_dataset_in_tf_function(self):
+    strategy = mwms.MultiWorkerMirroredStrategy()
+    local_batch_size = 4
+    global_batch_size = local_batch_size * strategy.num_replicas_in_sync
+    dataset = self.dataset.batch(global_batch_size).prefetch(2)
+
+    distributed_dataset = strategy.experimental_distribute_dataset(dataset)
+
+    @def_function.function
+    def step_fn(iterator):
+      images, labels = next(iterator)
+      del labels
+      return images
+
+    result = strategy.run(step_fn, args=(iter(distributed_dataset),))
+    self.assertIsInstance(result, dtensor_util.DTensorDistributedValue)
+    self.assertLen(result.values, self.num_local_devices)
+    for i in range(self.num_local_devices):
+      self.assertEqual(result.values[i].shape, [local_batch_size, 8, 8, 3])
+
+
+def client_config_function(config_params):
+  client_id = config_params['client_id']
+  worker_jobs = config_params['worker_jobs']
+  num_devices = config_params['num_devices']
+
+  os.environ['TF_CONFIG'] = json.dumps({
+      'cluster': {
+          'worker': worker_jobs
+      },
+      'task': {'type': 'worker', 'index': f'{client_id}'}
+  })
+
+  if config.list_physical_devices('GPU'):
+    device_type = 'GPU'
+  elif test_util.is_tpu_present():
+    device_type = 'TPU'
+  else:
+    device_type = 'CPU'
+
+  # reset_logical_devices
+  test_util.reset_context()
+  if device_type != 'TPU':
+    # Configure virtual devices. This does not initialize the TensorFlow
+    # context.
+    test_util.reset_logical_devices(device_type, num_devices)
+
+  # Validates the correct number of devices are created.
+  logical_devices = test_util.list_local_logical_devices(device_type)
+  assert len(logical_devices) == num_devices, (
+      logical_devices,
+      f'Test is mis-configured: expecting {num_devices} logical_devices.')
+
+
+if __name__ == '__main__':
+  test_backend_util.handle_test_main(
+      multi_client_test_util.multi_client_main, client_config_function)

tensorflow/python/distribute/experimental/rpc/BUILD

@@ -0,0 +1,64 @@
+# Python bindings for RPC client and server ops.
+
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "rpc_ops",
+    srcs = [
+        "rpc_ops.py",
+    ],
+    deps = [
+        "//tensorflow/distribute/experimental/rpc/kernels:gen_rpc_ops",
+        "//tensorflow/distribute/experimental/rpc/proto:tf_rpc_service_proto_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:none_tensor",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/saved_model:nested_structure_coder",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "rpc_ops_test",
+    size = "medium",
+    srcs = ["rpc_ops_test.py"],
+    shard_count = 7,
+    tags = [
+        "no_mac",  # flaky, b/205156709
+    ],
+    deps = [
+        ":rpc_ops",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:data_flow_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+        "@pypi//portpicker",
+    ],
+)

tensorflow/python/distribute/experimental/rpc/rpc_ops.py

@@ -0,0 +1,505 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Module to expose RPC APIs in tensorflow."""
+
+from typing import Optional, Sequence, Union
+
+import tensorflow.distribute.experimental.rpc.kernels.gen_rpc_ops as gen_rpc_ops
+from tensorflow.distribute.experimental.rpc.proto import tf_rpc_service_pb2 as rpc_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as tf_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import none_tensor
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.types import core as core_tf_types
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+def get_output_specs_from_function(func: tf_function.ConcreteFunction):
+  output_specs = nest.map_structure(type_spec.type_spec_from_value,
+                                    func.structured_outputs)
+  output_specs_proto = nested_structure_coder.encode_structure(output_specs)
+  return output_specs_proto.SerializeToString()
+
+
+def get_input_specs_from_function(func: tf_function.ConcreteFunction):
+  arg_specs, _ = func.structured_input_signature
+  arg_specs_proto = nested_structure_coder.encode_structure(arg_specs)
+  return arg_specs_proto.SerializeToString()
+
+
+@tf_export("distribute.experimental.rpc.Server", v1=[])
+class Server(object):
+  """A Server base class for accepting RPCs for registered tf.functions.
+
+    Functions can be registered on the server and are exposed via RPCs.
+  """
+
+  @staticmethod
+  def create(rpc_layer, address):
+    """Create TF RPC server at given address.
+
+    Args:
+      rpc_layer: Communication layer between client and server. Only "grpc" rpc
+        layer is supported at the moment.
+      address: Address where RPC server is hosted.
+
+    Returns:
+      An instance of `tf.distribute.experimental.rpc.Server` class.
+
+    Raises:
+        A ValueError if rpc_layer other than "grpc" is used. Only GRPC
+        is supported at the moment.
+
+    Example usage:
+
+      >>> import portpicker
+      >>> @tf.function(input_signature=[
+      ...      tf.TensorSpec([], tf.int32),
+      ...      tf.TensorSpec([], tf.int32)])
+      ... def remote_fn(a, b):
+      ...   return tf.add(a, b)
+
+      >>> port = portpicker.pick_unused_port()
+      >>> address = "localhost:{}".format(port)
+      >>> server = tf.distribute.experimental.rpc.Server.create("grpc", address)
+      >>> server.register("addition", remote_fn)
+      >>> server.start()
+
+    """
+    if rpc_layer != "grpc":
+      raise ValueError("Only GRPC backend is supported at the moment.")
+    return GrpcServer(address=address)
+
+  def register(self, method_name: str,
+               func: Union[def_function.Function,
+                           tf_function.ConcreteFunction]):
+    """Method for registering tf.function on server.
+
+    Registered methods can be invoked remotely from clients.
+
+    Args:
+      method_name: Name of the tf.function. Clients use this method_name to make
+        RPCs.
+      func: A `tf.function` or ConcreteFunction to register.
+    """
+    raise NotImplementedError("Please use create_server method to create a"
+                              "concrete subclass of Server.")
+
+  def start(self):
+    """Starts the RPC server on provided address.
+
+     Server listens for new requests from client, once it is started.
+    """
+    raise NotImplementedError("Please use create_server method to create a"
+                              "concrete subclass of Server.")
+
+
+@tf_export("distribute.experimental.rpc.Client", v1=[])
+class Client(object):
+  """Client class for invoking RPCs to the server."""
+
+  @staticmethod
+  def create(rpc_layer, address, name="", timeout_in_ms=0):
+    """Create TF RPC client to connect to the given address.
+
+    Args:
+      rpc_layer: Communication layer between client and server. Only "grpc" rpc
+        layer is supported at the moment.
+      address: Address of the server to connect the RPC client to.
+      name: Name of the RPC Client. You can create multiple clients connecting
+        to same server and distinguish them using different names.
+      timeout_in_ms: The default timeout to use for outgoing RPCs from client. 0
+        indicates no timeout. Exceeding timeout during RPC will raise
+        DeadlineExceeded error.
+
+    Returns:
+      An instance of `tf.distribute.experimental.rpc.Client` with the following
+      dynamically added methods for eagerly created clients:
+        * `Registered methods` e.g. multiply(**args):
+            If Client is created when executing eagerly, client will request the
+            list of registered methods from server during client creation.
+            The convenience methods for RPCs will be dynamically added to the
+            created Client instance.
+
+            For example, when a server has method "multiply" registered, the
+            client object created in eager mode will have 'multiply' method
+            available. Users can use client.multiply(..) to make RPC, instead of
+            client.call("multiply", ...)
+
+            Both "call" and "multiply" methods are non-blocking i.e. they return
+            a StatusOrResult object which should be used to wait for getting
+            value or error.
+
+            Along with the above, blocking versions of the registered
+            methods are also dynamically added to client instance.
+            e.g. multiply_blocking(**args). These methods block till the RPC is
+            finished and return response for successful RPC. Otherwise raise
+            exception.
+
+            These methods are not available when Client is created inside a
+            tf.function.
+
+    Raises:
+        A ValueError if rpc_layer other than "grpc" is used. Only GRPC
+          is supported at the moment.
+        A DeadlineExceeded exception in eager mode if timeout exceeds while
+          creating and listing client methods.
+
+    Example usage:
+      >>> # Have server already started.
+      >>> import portpicker
+      >>> @tf.function(input_signature=[
+      ...      tf.TensorSpec([], tf.int32),
+      ...      tf.TensorSpec([], tf.int32)])
+      ... def remote_fn(a, b):
+      ...   return tf.add(a, b)
+
+      >>> port = portpicker.pick_unused_port()
+      >>> address = "localhost:{}".format(port)
+      >>> server = tf.distribute.experimental.rpc.Server.create("grpc", address)
+      >>> server.register("addition", remote_fn)
+      >>> server.start()
+
+      >>> # Start client
+      >>> client = tf.distribute.experimental.rpc.Client.create("grpc",
+      ...      address=address, name="test_client")
+
+      >>> a = tf.constant(2, dtype=tf.int32)
+      >>> b = tf.constant(3, dtype=tf.int32)
+
+      >>> result = client.call(
+      ...    args=[a, b],
+      ...    method_name="addition",
+      ...    output_specs=tf.TensorSpec((), tf.int32))
+
+      >>> if result.is_ok():
+      ...   result.get_value()
+
+      >>> result = client.addition(a, b)
+
+      >>> if result.is_ok():
+      ...   result.get_value()
+
+      >>> value = client.addition_blocking(a, b)
+    """
+    if rpc_layer != "grpc":
+      raise ValueError("Only GRPC backend is supported at the moment.")
+    if context.executing_eagerly():
+      list_registered_methods = True
+    else:
+      list_registered_methods = False
+    return GrpcClient(
+        address=address,
+        name=name,
+        list_registered_methods=list_registered_methods,
+        timeout_in_ms=timeout_in_ms)
+
+  def call(self,
+           method_name: str,
+           args: Optional[Sequence[core_tf_types.Tensor]] = None,
+           output_specs=None,
+           timeout_in_ms=0):
+    """Method for making RPC calls to remote server.
+
+    This invokes RPC to the server, executing the registered method_name
+    remotely.
+    Args:
+      method_name: Remote registered method to invoke
+      args: List of arguments for the registered method.
+      output_specs: Output specs for the output from method.
+         For example, if tf.function is: @tf.function(input_signature=[
+           tf.TensorSpec([], tf.int32), tf.TensorSpec([], tf.int32) ])
+          def multiply_fn(a, b): return tf.math.multiply(a, b)
+        output_spec is: tf.TensorSpec((), tf.int32)  If you have access to TF
+          Function, the output specs can be generated
+       from tf.function by calling: output_specs =
+         tf.nest.map_structure(tf.type_spec_from_value,
+         tf_function.get_concrete_function().structured_outputs  If output_specs
+         are not provided, flattened list of tensors will be returned in
+         response.
+      timeout_in_ms: Timeout for this call. If 0, default client timeout will be
+        used.
+
+    Returns:
+      An instance of `StatusOrResult` class with the following available
+      methods.
+        * `is_ok()`:
+            Returns True of RPC was successful.
+        * `get_error()`:
+            Returns TF error_code and error message for the RPC.
+        * `get_value()`:
+            Returns the returned value from remote TF function execution
+            when RPC is successful.
+
+      Calling any of the above methods will block till RPC is completed and
+      result is available.
+    """
+    raise NotImplementedError("Must be implemented in inherited classes.")
+
+
+class GrpcServer(Server):
+  """GrpcServer object encapsulates a resource with GRPC server.
+
+    Functions can be registered locally and are exposed via RPCs.
+    Example:
+    ```
+    server = rpc_ops.GrpcServer("host:port")
+    @tf.function
+    def add(a, b):
+      return a + b
+
+    server.register("add", add)
+    server.start()
+    ```
+  """
+
+  def __init__(self, address: str):
+    self._server_handle = gen_rpc_ops.rpc_server(address)
+    if context.executing_eagerly():
+      self._handle_deleter = resource_variable_ops.EagerResourceDeleter(
+          handle=self._server_handle, handle_device=self._server_handle.device)
+    else:
+      raise NotImplementedError("Please create the server outside tf.function.")
+
+  def register(self, method_name: str,
+               func: Union[def_function.Function,
+                           tf_function.ConcreteFunction]):
+    """Method for registering functions."""
+
+    if isinstance(func, def_function.Function):
+      if func.function_spec.arg_names:
+        if func.input_signature is None:
+          raise ValueError("Input signature not specified for the function.")
+      concrete_fn = func.get_concrete_function()
+      gen_rpc_ops.rpc_server_register(
+          self._server_handle,
+          method_name=method_name,
+          captured_inputs=concrete_fn.captured_inputs,
+          input_specs=get_input_specs_from_function(concrete_fn),
+          output_specs=get_output_specs_from_function(concrete_fn),
+          f=concrete_fn)
+    elif isinstance(func, tf_function.ConcreteFunction):
+      gen_rpc_ops.rpc_server_register(
+          self._server_handle,
+          method_name=method_name,
+          captured_inputs=func.captured_inputs,
+          input_specs=get_input_specs_from_function(func),
+          output_specs=get_output_specs_from_function(func),
+          f=func)
+    else:
+      # Python functions
+      # TODO(b/186762191): Add an implementation to support python functions.
+      raise ValueError("Only TF functions are supported with Register method")
+
+  def start(self):
+    """Starts GRPC server."""
+    gen_rpc_ops.rpc_server_start(self._server_handle)
+
+
+class GrpcClient(Client):
+  """Client wrapper to connect to remote RPC server using GRPC.
+
+  If Client is created with (list_registered_methods=True):
+  1. Input and output specs for the methods till this point will be fetched from
+  Server.
+  2. convenience methods are added to invoke registered methods directly from
+  client.
+  For example:
+    For call a server method `add`
+    client.add(a, b) or client.add_async(a, b) can be used instead of
+    client.call(args=[a,b], output_specs=[..])
+
+  Prerequisite for using list_registered_methods=True:
+   1. Server should be already started with the registered methods.
+   2. Client must be created in Eager mode.
+  """
+
+  def __init__(self,
+               address: str,
+               name: str = "",
+               list_registered_methods=False,
+               timeout_in_ms=0):
+    self._client_handle, methods = gen_rpc_ops.rpc_client(
+        shared_name=name,
+        server_address=address,
+        list_registered_methods=list_registered_methods,
+        timeout_in_ms=timeout_in_ms)
+    if context.executing_eagerly():
+      self._handle_deleter = resource_variable_ops.EagerResourceDeleter(
+          handle=self._client_handle, handle_device=self._client_handle.device)
+    else:
+      raise NotImplementedError(
+          "Client creation is supported only in eager mode.")
+    self._server_address = address
+    self._method_registry = {}
+    for method in methods.numpy():
+      m = rpc_pb2.RegisteredMethod()
+      m.ParseFromString(method)
+      output_specs = nested_structure_coder.decode_proto(m.output_specs)
+      input_specs = nested_structure_coder.decode_proto(m.input_specs)
+      self._method_registry[m.method] = output_specs
+      # TODO(ishark): Perhaps doc string can also be taken as input during
+      # function registration.
+      doc_string = "RPC Call for " + m.method + " method to server " + address
+      self._add_method(m.method, output_specs, input_specs, self._client_handle,
+                       doc_string)
+
+  def _add_method(self, method_name, output_specs, input_specs, client_handle,
+                  doc_string):
+    """Method to add RPC methods to the client object."""
+
+    def validate_and_get_flat_inputs(*args):
+      if args is None:
+        args = []
+      if input_specs:
+        nest.assert_same_structure(args, input_specs)
+      flat_inputs = nest.flatten(args)
+      return flat_inputs
+
+    def call_wrapper(*args, timeout_in_ms=0):
+      status_or, deleter = gen_rpc_ops.rpc_call(
+          client_handle,
+          args=validate_and_get_flat_inputs(*args),
+          method_name=method_name,
+          timeout_in_ms=timeout_in_ms)
+      return StatusOrResult(status_or, deleter, output_specs)
+
+    def call_blocking_wrapper(*args, timeout_in_ms=0):
+      status_or, deleter = gen_rpc_ops.rpc_call(
+          client_handle,
+          args=validate_and_get_flat_inputs(*args),
+          method_name=method_name,
+          timeout_in_ms=timeout_in_ms)
+      status_or = StatusOrResult(status_or, deleter, output_specs)
+      if status_or.is_ok():
+        return status_or.get_value()
+      else:
+        error_code, error_msg = status_or.get_error()
+        raise errors.exception_type_from_error_code(error_code.numpy())(
+            None, None, error_msg.numpy())
+
+    setattr(self, method_name, call_wrapper)
+    call_wrapper.__doc__ = doc_string
+
+    blocking_method_name = method_name + "_blocking"
+    setattr(self, blocking_method_name, call_blocking_wrapper)
+    call_blocking_wrapper.__doc__ = doc_string
+
+  def call(self,
+           method_name: str,
+           args: Optional[Sequence[core_tf_types.Tensor]] = None,
+           output_specs=None,
+           timeout_in_ms=0):
+    """Method to invoke remote registered functions on the connected server.
+
+    Server should be started before making an RPC Call.
+
+    Args:
+      method_name: Registered method to invoke on Server.
+      args: Input arguments for the method.
+      output_specs: Output specs for the output from method.
+      timeout_in_ms: Timeout for this call. If 0, default client timeout will be
+       used.
+
+    Returns:
+      StatusOrResult object. This function issues the RPC call to server, it
+      does not block for the duration of RPC. Please call is_ok, get_error or
+      get_value methods on the returned object to blocked till RPC finishes.
+    """
+    if args is None:
+      args = []
+    status_or, deleter = gen_rpc_ops.rpc_call(
+        self._client_handle,
+        args=nest.flatten(args),
+        method_name=method_name,
+        timeout_in_ms=timeout_in_ms)
+    return StatusOrResult(status_or, deleter, output_specs)
+
+
+class StatusOrResult(object):
+  """Class representing result and status from RPC Call."""
+
+  def __init__(self, status_or, deleter, output_specs=None):
+    self._status_or = status_or
+    self._output_specs = output_specs
+    self._deleter = deleter
+    self._error_code: dtypes.int64 = None
+    self._error_message: dtypes.string = None
+
+  def _check_status(self):
+    if self._error_code is None:
+      self._error_code, self._error_message = gen_rpc_ops.rpc_check_status(
+          self._status_or)
+
+  def __del__(self):
+    # Make sure the resource is deleted in the same mode as it was created in.
+    if context.executing_eagerly():
+      with context.eager_mode():
+        gen_rpc_ops.delete_rpc_future_resource(
+            handle=self._status_or, deleter=self._deleter)
+    else:
+      with context.graph_mode():
+        gen_rpc_ops.delete_rpc_future_resource(
+            handle=self._status_or, deleter=self._deleter)
+
+  def is_ok(self):
+    """Returns True if RPC is successful, otherwise returns False.
+
+    This call will block for RPC result.
+    """
+    self._check_status()
+    return math_ops.equal(self._error_code,
+                          constant_op.constant(0, dtype=dtypes.int64))
+
+  def get_error(self):
+    """Returns (TF Error Code, Error Message) from RPC Response.
+
+    This call will block for RPC result.
+    """
+    self._check_status()
+    return self._error_code, self._error_message
+
+  def get_value(self):
+    """Returns the returned response value from RPC Call when RPC is successful.
+
+      The returned value is tensors in the output_specs format as returned from
+      the RPC call
+
+
+    This call will block for RPC result.
+    """
+
+    self._check_status()
+    if self._output_specs is None or isinstance(self._output_specs,
+                                                none_tensor.NoneTensorSpec):
+      flat_output_dtypes = []
+      return_none = True
+    else:
+      return_none = False
+      flat_output_dtypes = [s.dtype for s in nest.flatten(self._output_specs)]
+
+    result = gen_rpc_ops.rpc_get_value(self._status_or, Tout=flat_output_dtypes)
+    if return_none:
+      return None
+    else:
+      return nest.pack_sequence_as(self._output_specs, result)

tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py

@@ -0,0 +1,850 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for rpc_ops.py."""
+
+import threading
+import time
+import numpy as np
+import portpicker
+
+from tensorflow.python.distribute.experimental.rpc import rpc_ops
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function as eager_def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.util import nest
+
+
+@test_util.with_eager_op_as_function
+class RpcOpsTest(test.TestCase):
+
+  def setUp(self):
+    super(RpcOpsTest, self).setUp()
+    cpus = config.list_physical_devices("CPU")
+    # Set 2 virtual CPUs
+    config.set_logical_device_configuration(cpus[0], [
+        context.LogicalDeviceConfiguration(),
+        context.LogicalDeviceConfiguration()
+    ])
+
+  def test_generated_rpc_ops(self):
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    concrete_remote_fn = remote_fn.get_concrete_function()
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.gen_rpc_ops.rpc_server(server_address=address)
+
+    rpc_ops.gen_rpc_ops.rpc_server_register(
+        server_resource,
+        f=concrete_remote_fn,
+        captured_inputs=concrete_remote_fn.captured_inputs,
+        output_specs=rpc_ops.get_output_specs_from_function(concrete_remote_fn),
+        method_name="multiply")
+
+    rpc_ops.gen_rpc_ops.rpc_server_start(server_resource)
+    client_handle, _ = rpc_ops.gen_rpc_ops.rpc_client(
+        server_address=address, timeout_in_ms=5000)
+    future_resource, deleter = rpc_ops.gen_rpc_ops.rpc_call(
+        client_handle, args=[a, b], method_name="multiply", timeout_in_ms=0)
+
+    error_code, _ = rpc_ops.gen_rpc_ops.rpc_check_status(future_resource)
+    self.assertAllEqual(error_code, 0)
+    self.assertAllEqual(
+        rpc_ops.gen_rpc_ops.rpc_get_value(future_resource, Tout=[dtypes.int32]),
+        [6])
+
+    resource_variable_ops.EagerResourceDeleter(
+        handle=server_resource, handle_device=server_resource.device)
+
+    resource_variable_ops.EagerResourceDeleter(
+        handle=client_handle, handle_device=client_handle.device)
+
+    rpc_ops.gen_rpc_ops.delete_rpc_future_resource(future_resource, deleter)
+
+  def test_exported_rpc_api_static_factory(self):
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.Server.create("grpc", address)
+    server_resource.register("multiply", _remote_fn)
+
+    server_resource.start()
+    client = rpc_ops.Client.create("grpc", address=address, name="test_client")
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    mul_or = client.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    # Test empty client name
+    client1 = rpc_ops.Client.create("grpc", address)
+    mul_or = client1.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    # Test without output_spec
+    mul_or = client1.multiply(a, b)
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    self.assertEqual(client1.multiply.__doc__,
+                     "RPC Call for multiply method to server " + address)
+
+  def test_rpc_ops_call_method(self):
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def add_fn(a, b):
+      return math_ops.add(a, b)
+
+    # Register TF function
+    server_resource.register("multiply", _remote_fn)
+
+    # Register concrete Function
+    server_resource.register("add", add_fn.get_concrete_function())
+
+    server_resource.start()
+    client = rpc_ops.GrpcClient(address=address, name="test_client")
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    mul_or = client.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    add_or = client.call(
+        args=[a, b],
+        method_name="add",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+
+    self.assertAllEqual(add_or.is_ok(), True)
+    self.assertAllEqual(add_or.get_value(), 5)
+
+    # Test empty client name
+    client1 = rpc_ops.GrpcClient(address, list_registered_methods=True)
+    mul_or = client1.call(
+        args=[a, b],
+        method_name="multiply",
+        output_specs=tensor_spec.TensorSpec((), dtypes.int32))
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+  def test_rpc_ops_non_blocking_convenience_methods(self):
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    # Register TF function
+    server_resource.register("multiply", _remote_fn)
+
+    server_resource.start()
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+
+    mul_or = client.multiply(a, b)
+    self.assertAllEqual(mul_or.is_ok(), True)
+    self.assertAllEqual(mul_or.get_value(), 6)
+
+    self.assertEqual(client.multiply.__doc__,
+                     "RPC Call for multiply method to server " + address)
+
+  def test_rpc_ops_blocking_convenience_methods(self):
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    # Register TF function
+    server_resource.register("multiply", _remote_fn)
+
+    server_resource.start()
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+    self.assertAllEqual(client.multiply_blocking(a, b), 6)
+
+    self.assertEqual(
+        client.multiply_blocking.__doc__,
+        "RPC Call for multiply method to server " + address)
+
+  def test_output_specs(self):
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def test_dict(val):
+      return {"key": val}
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def is_positive(a):
+      if a > 0:
+        return True
+      return False
+
+    @eager_def_function.function(input_signature=[])
+    def do_nothing():
+      return []
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
+    def test_nested_structure(v):
+      return {"test": (v, [v, v]), "test1": (v,)}
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    server_resource.register("test_dict", test_dict)
+    server_resource.register("is_positive", is_positive)
+    server_resource.register("test_nested_structure", test_nested_structure)
+    server_resource.register("do_nothing", do_nothing)
+
+    server_resource.start()
+
+    client = rpc_ops.GrpcClient(
+        address=address, name="test_client", list_registered_methods=True)
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+
+    result_or = client.test_dict(a)
+    self.assertAllEqual(result_or.is_ok(), True)
+    nest.map_structure(self.assertAllEqual, result_or.get_value(), {"key": 2})
+
+    result_or = client.is_positive(a)
+    self.assertTrue(result_or.is_ok())
+    self.assertTrue(result_or.get_value())
+
+    result_or = client.test_nested_structure(a)
+    self.assertAllEqual(result_or.is_ok(), True)
+    nest.map_structure(self.assertAllEqual, result_or.get_value(), {
+        "test": (2, [2, 2]),
+        "test1": (2,)
+    })
+
+    result_or = client.do_nothing()
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [])
+
+  def test_input_specs(self):
+
+    @eager_def_function.function(input_signature=[{
+        "a": tensor_spec.TensorSpec([], dtypes.int32),
+        "b": tensor_spec.TensorSpec([], dtypes.int32)
+    }])
+    def test_input_dict(value):
+      return math_ops.add(value["a"], value["b"])
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    server_resource.register("test_input_dict", test_input_dict)
+
+    server_resource.start()
+
+    client = rpc_ops.GrpcClient(
+        address=address, name="test_client", list_registered_methods=True)
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+    result_or = client.test_input_dict({"a": a, "b": b})
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), 5)
+
+    with self.assertRaises(TypeError):
+      client.test_input_dict([a, b])
+
+  def test_call_register_ordering(self):
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+
+    # Create client succeeds before server start and registration
+    client = rpc_ops.GrpcClient(address)
+
+    # Create client with list_registered_methods fails before server is started.
+    with self.assertRaises(errors.DeadlineExceededError):
+      rpc_ops.GrpcClient(
+          address,
+          name="client1",
+          list_registered_methods=True,
+          timeout_in_ms=1)
+
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    server = rpc_ops.GrpcServer(address)
+
+    def start_server():
+      # Delay server start to test whether client creation also waits
+      # till server is up.
+      time.sleep(1)
+      server.register("assign_add", assign_add)
+      server.start()
+
+    t = threading.Thread(target=start_server)
+    t.start()
+
+    # Create same "client1" again should succeed.
+    client1_with_listed_methods = rpc_ops.GrpcClient(
+        address, name="client1", list_registered_methods=True)
+
+    result_or = client1_with_listed_methods.assign_add(
+        variables.Variable(2, dtype=dtypes.int64))
+    self.assertAllEqual(result_or.is_ok(), True)
+
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+
+    # Create client with registered methods
+    client2_with_listed_methods = rpc_ops.GrpcClient(
+        address=address, name="client2", list_registered_methods=True)
+
+    result_or = client2_with_listed_methods.assign_add(
+        variables.Variable(2, dtype=dtypes.int64))
+    self.assertAllEqual(result_or.is_ok(), True)
+
+    self.assertAllEqual(v, 6)
+
+    # Register new method after server started.
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError,
+        "All methods must be registered before starting the server"):
+      server.register("read_var", read_var)
+
+  def test_client_timeout(self):
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def add(a, b):
+      return math_ops.add(a, b)
+
+    server = rpc_ops.GrpcServer(address)
+
+    def start_server():
+      # Delay server start to simulate deadline exceeded for 1st RPC call
+      # response. Client waits till server is started, thus it can trigger
+      # deadline exceeded.
+      time.sleep(1)
+      server.register("add", add)
+      server.start()
+
+    t = threading.Thread(target=start_server)
+    t.start()
+
+    def ensure_server_is_ready(client):
+      server_ready = False
+      while not server_ready:
+        result_or = client.call(
+            "add", [constant_op.constant(20),
+                    constant_op.constant(30)])
+        if result_or.is_ok():
+          server_ready = True
+        else:
+          error_code, _ = result_or.get_error()
+          if error_code == errors.UNAVAILABLE:
+            server_ready = False
+          else:
+            server_ready = True
+      return
+
+    # Create client with list_registered_methods fails before server is started.
+    with self.assertRaises(errors.DeadlineExceededError):
+      rpc_ops.GrpcClient(
+          address,
+          name="client1",
+          list_registered_methods=True,
+          timeout_in_ms=1)
+
+    # Create same client again should succeed with
+    # list_registered_methods=False. Default timeout for client is 1 ms.
+    client = rpc_ops.GrpcClient(
+        address, name="client1", list_registered_methods=False, timeout_in_ms=1)
+
+    ensure_server_is_ready(client)
+    # Make explicit RPC call, the timeout of 1 ms should lead to
+    # deadline exceeded error.
+
+    result_or = client.call(
+        "add", [constant_op.constant(20),
+                constant_op.constant(30)],
+        timeout_in_ms=1)
+    self.assertAllEqual(result_or.is_ok(), False)
+    error_code, error_message = result_or.get_error()
+    self.assertAllEqual(error_code, errors.DEADLINE_EXCEEDED, error_message)
+
+    # Specifying reasonable timeout for call should succeed.
+    result_or = client.call(
+        "add", [constant_op.constant(20),
+                constant_op.constant(30)],
+        timeout_in_ms=5000)
+    self.assertAllEqual(result_or.is_ok(), True)
+    error_code, _ = result_or.get_error()
+
+    # Test timeouts for convenience methods
+
+    # Restart server again with delay to simulate deadline exceeded.
+    del server
+    server = rpc_ops.GrpcServer(address)
+    t = threading.Thread(target=start_server)
+    t.start()
+
+    # Client with no default timeout.
+    client = rpc_ops.GrpcClient(
+        address, name="client2", list_registered_methods=True)
+
+    # Succeeds with reasonable timeout.
+    result_or = client.add(
+        constant_op.constant(20), constant_op.constant(30), timeout_in_ms=5000)
+    self.assertAllEqual(result_or.is_ok(), True)
+
+  def test_async_call_op_wrapper(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign_add", assign_add)
+    server.register("read_var", read_var)
+    server.start()
+
+    client = rpc_ops.GrpcClient(address)
+
+    futures = []
+    for _ in range(10):
+      futures.append(
+          client.call("assign_add",
+                      [variables.Variable(2, dtype=dtypes.int64)]))
+
+    for f in futures:
+      f.is_ok()
+
+    result_or = client.call(
+        "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [20])
+
+  def test_rpc_call_op_in_tf_function(self):
+
+    @eager_def_function.function(input_signature=[
+        tensor_spec.TensorSpec([], dtypes.int32),
+        tensor_spec.TensorSpec([], dtypes.int32)
+    ])
+    def _remote_fn(a, b):
+      return math_ops.multiply(a, b)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server_resource = rpc_ops.GrpcServer(address)
+
+    server_resource.register("remote_fn", _remote_fn)
+
+    server_resource.start()
+    client = rpc_ops.GrpcClient(address=address, name="test_client")
+
+    a = variables.Variable(2, dtype=dtypes.int32)
+    b = variables.Variable(3, dtype=dtypes.int32)
+
+    @eager_def_function.function
+    def call_fn():
+      result_or = client.call(
+          args=[a, b],
+          method_name="remote_fn",
+          output_specs=[tensor_spec.TensorSpec([], dtypes.int32)])
+
+      self.assertAllEqual(True, result_or.is_ok())
+      result = result_or.get_value()
+      self.assertEqual(len(result), 1)  # Call returns a list(tensors)
+      # TODO(ishark): Shape for output tensor is unknown currently.
+      # Add attribute for capturing TensorSpec for output and enable
+      # check below:
+      # self.assertIsNotNone(result[0].shape.rank)
+      return result
+
+    self.assertAllEqual(call_fn(), [6])
+
+  def test_resource_deletion(self):
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server_handle = server._server_handle
+
+    # Test Future resource deletion
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    server.register("read_var", read_var)
+
+    server.start()
+    client = rpc_ops.GrpcClient(address)
+
+    client_handle = client._client_handle
+
+    # Check future resource deletion without calling get_value.
+    def _create_and_delete_rpc_future():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      return handle._status_or
+
+    @eager_def_function.function
+    def _create_and_delete_rpc_future_fn():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      return handle._status_or
+
+    for _ in range(2):
+      handle = _create_and_delete_rpc_future()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            handle, ignore_lookup_error=False)
+
+    for _ in range(2):
+      handle = _create_and_delete_rpc_future_fn()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            handle, ignore_lookup_error=False)
+
+    # Check future resource deletion with calling get_value.
+    def _create_and_delete_with_future():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      status_or_handle = handle._status_or
+      handle.get_value()
+      return status_or_handle
+
+    # Check future resource deletion with calling get_value with tf.function.
+    @eager_def_function.function
+    def _create_and_delete_with_future_fn():
+      handle = client.call(
+          "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+      status_or_handle = handle._status_or
+      handle.get_value()
+      return status_or_handle
+
+    for _ in range(2):
+      resource_handle = _create_and_delete_with_future()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            resource_handle, ignore_lookup_error=False)
+
+    for _ in range(2):
+      resource_handle = _create_and_delete_with_future_fn()
+      with self.assertRaises(errors.NotFoundError):
+        resource_variable_ops.destroy_resource_op(
+            resource_handle, ignore_lookup_error=False)
+
+    # Test server client resource gets deleted.
+    del client
+    with self.assertRaises(errors.NotFoundError):
+      resource_variable_ops.destroy_resource_op(
+          client_handle, ignore_lookup_error=False)
+
+    # Test server server resource gets deleted.
+    del server
+    with self.assertRaises(errors.NotFoundError):
+      resource_variable_ops.destroy_resource_op(
+          server_handle, ignore_lookup_error=False)
+
+  def test_rpc_error(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign_add", assign_add)
+    server.register("read_var", read_var)
+    server.start()
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+
+    # confirm it works as expected when arguments are passed.
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    result_or = client.call(
+        "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [2])
+    result_or = client.assign_add(variables.Variable(2, dtype=dtypes.int64))
+    self.assertAllEqual(True, result_or.is_ok())
+
+    result_or = client.read_var()
+    self.assertAllEqual(True, result_or.is_ok())
+    self.assertAllEqual(result_or.get_value(), 4)
+
+    # Fails with invalid argument error when no arguments are passed.
+    result_or = client.call("assign_add")
+    self.assertAllEqual(result_or.is_ok(), False)
+    error_code, _ = result_or.get_error()
+    self.assertAllEqual(error_code, errors.INVALID_ARGUMENT)
+
+    del server
+    with self.assertRaises(errors.DeadlineExceededError):
+      _ = client.assign_add_blocking(
+          variables.Variable(2, dtype=dtypes.int64), timeout_in_ms=1)
+
+  def test_captured_inputs(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(input_signature=[])
+    def read_var():
+      return v.value()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign_add", assign_add)
+    server.register("read_var", read_var)
+
+    server.start()
+
+    client = rpc_ops.GrpcClient(address)
+
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    result_or = client.call("assign_add",
+                            [variables.Variable(2, dtype=dtypes.int64)])
+    self.assertAllEqual(result_or.is_ok(), True)
+    result_or = client.call(
+        "read_var", output_specs=[tensor_spec.TensorSpec([], dtypes.int64)])
+
+    self.assertAllEqual(result_or.is_ok(), True)
+    self.assertAllEqual(result_or.get_value(), [4])
+
+  def test_register_method_twice(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign(a):
+      v.assign(a)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("assign", assign_add)
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "assign is already registered."):
+      # Reusing the same error name.
+      server.register("assign", assign)
+
+  def test_tf_function_register_without_input_signature(self):
+    v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function
+    def assign(a):
+      v.assign(a)
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    with self.assertRaisesRegex(
+        ValueError, "Input signature not specified for the function."):
+      server.register("assign", assign)
+
+    # Register without input signature should work for functions without input
+    # args.
+    @eager_def_function.function
+    def read_var():
+      return v.value()
+
+    server.register("read_var", read_var)
+
+  def test_multi_device_resource(self):
+    elements = np.random.randint(100, size=[200])
+
+    with ops.device("/device:CPU:1"):
+      queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
+
+    @eager_def_function.function()
+    def populate_queue():
+      queue.enqueue_many(elements)
+      queue.close()
+
+    with ops.device("/device:CPU:0"):
+      port = portpicker.pick_unused_port()
+      address = "localhost:{}".format(port)
+      server = rpc_ops.GrpcServer(address)
+      server.register("populate_queue", populate_queue)
+      server.start()
+
+      client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+      client.populate_queue()
+
+    for e in elements:
+      self.assertAllEqual(e, queue.dequeue())
+
+  def test_queue_resource(self):
+    elements = np.random.randint(100, size=[200])
+    queue = data_flow_ops.FIFOQueue(200, dtypes.int64, shapes=[])
+
+    @eager_def_function.function()
+    def populate_queue():
+      queue.enqueue_many(elements)
+      queue.close()
+
+    port = portpicker.pick_unused_port()
+    address = "localhost:{}".format(port)
+    server = rpc_ops.GrpcServer(address)
+    server.register("populate_queue", populate_queue)
+    server.start()
+
+    client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+    client.populate_queue()
+
+    for e in elements:
+      self.assertAllEqual(e, queue.dequeue())
+
+  def test_multi_device_resource_cpu(self):
+    with ops.device("/device:cpu:1"):
+      v = variables.Variable(initial_value=0, dtype=dtypes.int64)
+
+    @eager_def_function.function(
+        input_signature=[tensor_spec.TensorSpec([], dtypes.int64)])
+    def assign_add(a):
+      v.assign_add(a)
+
+    with ops.device("/device:CPU:0"):
+      port = portpicker.pick_unused_port()
+      address = "localhost:{}".format(port)
+      server = rpc_ops.GrpcServer(address)
+      server.register("assign_add", assign_add)
+      server.start()
+
+      client = rpc_ops.GrpcClient(address, list_registered_methods=True)
+      result_or = client.assign_add(variables.Variable(2, dtype=dtypes.int64))
+      self.assertAllEqual(result_or.is_ok(), True)
+
+    self.assertAllEqual(v, 2)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/distribute/failure_handling/BUILD

@@ -0,0 +1,153 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_custom_op_py_strict_library", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "failure_handling_lib",
+    srcs = [
+        "failure_handling.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":check_preemption_py",
+        ":failure_handling_util",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_context",
+        "//tensorflow/python/checkpoint:checkpoint_management",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:multi_worker_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_contextlib",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/tools/docs:doc_controls",
+    ],
+)
+
+py_library(
+    name = "failure_handling_util",
+    srcs = [
+        "failure_handling_util.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/platform:tf_logging",
+        "@pypi//requests",
+        "@six_archive//:six",
+    ],
+)
+
+py_library(
+    name = "preemption_watcher",
+    srcs = ["preemption_watcher.py"],
+    strict_deps = True,
+    deps = [
+        ":failure_handling_util",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:monitoring",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/util:tf_export",
+        "@absl_py//absl/logging",
+    ],
+)
+
+tf_py_strict_test(
+    name = "failure_handler_test",
+    timeout = "long",
+    srcs = ["failure_handler_test.py"],
+    shard_count = 8,
+    tags = [
+        "no_mac",
+        "no_pip",  # TODO(b/266520226)
+        "no_windows",  # TODO(b/197981388)
+    ],
+    deps = [
+        ":failure_handling_lib",
+        ":failure_handling_util",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_management",
+        "//tensorflow/python/distribute:collective_all_reduce_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:multi_worker_util",
+        "//tensorflow/python/distribute:one_device_strategy",
+        "//tensorflow/python/distribute:test_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:server_lib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "gce_failure_handler_test",
+    srcs = ["gce_failure_handler_test.py"],
+    shard_count = 32,
+    tags = [
+        "no_mac",  # Fails on CI but works fine locally.
+        "no_pip",  # TODO(b/266520226)
+        "noasan",  # TODO(b/226154233): Flaky test
+    ],
+    deps = [
+        ":failure_handling_lib",
+        ":failure_handling_util",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_management",
+        "//tensorflow/python/distribute:collective_all_reduce_strategy",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:mirrored_strategy",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:multi_worker_util",
+        "//tensorflow/python/distribute:one_device_strategy",
+        "//tensorflow/python/distribute:test_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/training:server_lib",
+        "@absl_py//absl/testing:parameterized",
+        "@pypi//dill",  # build_cleaner: keep
+    ],
+)
+
+tf_custom_op_py_strict_library(
+    name = "check_preemption_py",
+    kernels = [
+        "//tensorflow/core/distributed_runtime/preemption:check_preemption_op_kernel",
+        "//tensorflow/core/distributed_runtime/preemption:check_preemption_op_op_lib",
+    ],
+    deps = [
+        "//tensorflow/core/distributed_runtime/preemption:gen_check_preemption_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+    ],
+)

tensorflow/python/distribute/failure_handling/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Library imports for PreemptionCheckpointHandler."""
+
+from tensorflow.python.distribute.failure_handling.failure_handling import PreemptionCheckpointHandler
+from tensorflow.python.distribute.failure_handling.failure_handling import TerminationConfig
+from tensorflow.python.distribute.failure_handling.tpu_preemption_watcher import PreemptionWatcher

tensorflow/python/distribute/failure_handling/failure_handler_test.py

@@ -0,0 +1,558 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Tests for PreemptionCheckpointHandler."""
+import os
+import random
+import re
+import signal
+import sys
+import threading
+import time
+
+from absl.testing import parameterized
+
+# pylint:disable=g-direct-tensorflow-import
+from tensorflow.python.checkpoint import checkpoint as tracking_util
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.distribute import collective_all_reduce_strategy
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.distribute import one_device_strategy as one_device_lib
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute.failure_handling import failure_handling
+from tensorflow.python.distribute.failure_handling import failure_handling_util
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.module import module
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import server_lib
+
+
+try:
+  import dill  # pylint:disable=g-import-not-at-top
+
+  _REGISTER_DECORATOR = dill.register
+except ImportError:
+  _REGISTER_DECORATOR = lambda fn, *_: fn
+
+mock = test.mock
+
+
+CLUSTER_SIZE = 4
+EPOCHS_TO_RUN = 8
+STEPS_PER_EPOCH = 15
+MAX_WAIT_TIME = 40
+
+
+def _is_oss():
+  """Returns whether the test is run under OSS."""
+  return len(sys.argv) >= 1 and 'bazel' in sys.argv[0]
+
+
+def _make_checkpoint_manager(checkpoint, checkpoint_dir, cluster_resolver):
+  if not cluster_resolver or not cluster_resolver.cluster_spec().as_dict() or (
+      multi_worker_util.is_chief(
+          cluster_spec=cluster_resolver.cluster_spec(),
+          task_type=cluster_resolver.task_type,
+          task_id=cluster_resolver.task_id)):
+    return checkpoint_management.CheckpointManager(
+        checkpoint, directory=checkpoint_dir, max_to_keep=1)
+  else:
+    return checkpoint_management.CheckpointManager(
+        checkpoint,
+        directory=failure_handling._non_chief_checkpoint_dir(
+            checkpoint_dir, cluster_resolver.task_id),
+        max_to_keep=1)
+
+
+def raise_if_not_all_exit(grace_period, mpr):
+  """Wait for all cluster to exit with a time out."""
+  waiting_time = 0
+  exit_process_count = 0
+  # This addition to mitigate the fact that our step time is too short in test
+  while exit_process_count != CLUSTER_SIZE and waiting_time < max(
+      grace_period + 15, MAX_WAIT_TIME):
+    exit_process_count = 0
+    for worker_id in range(CLUSTER_SIZE):
+      if not mpr.process_exists('worker', worker_id):
+        exit_process_count += 1
+    waiting_time += 1
+    time.sleep(1)
+
+  if waiting_time == max(grace_period + 5, 40):
+    raise RuntimeError('Waited long but at least one worker still exist. '
+                       'Considering size of our model, this should not'
+                       ' happen.')
+
+
+class PreemptionCheckpointTest(test.TestCase, parameterized.TestCase):
+  """Integration test for PreemptionCheckpointHandler."""
+
+  def _maybe_trigger_a_preemption(self, training_started_event,
+                                  trigger_it=False):
+    if not training_started_event:
+      return
+    clear_events = [
+        event for event in training_started_event if not event.is_set()
+    ]
+    if clear_events:
+      if trigger_it:
+        logging.info('Set preemption signal')
+        clear_events[0].set()
+      elif random.randrange(0, 9) > 6:
+        clear_events[0].set()
+
+  def worker_fn(self,
+                checkpoint_dir,
+                cluster_spec,
+                strategy_option,
+                input_arg='checkpoint',
+                training_started_event=None,
+                raise_app_error_on_worker=None,
+                training_restarted=None,
+                training_finished=None,
+                termination_config=failure_handling.TerminationConfig(),
+                api_wrapping_train=True):
+
+    if strategy_option == 'MS':
+      strategy = mirrored_strategy.MirroredStrategy()
+    elif strategy_option == 'OneDevice':
+      if config.list_physical_devices('GPU'):
+        strategy = one_device_lib.OneDeviceStrategy(device='/gpu:0')
+      else:
+        strategy = one_device_lib.OneDeviceStrategy(device='/cpu:0')
+    else:
+      strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy()
+
+    class Model(module.Module):
+
+      def __init__(self):
+        self.v = variables_lib.Variable(
+            0.,
+            synchronization=variables_lib.VariableSynchronization.ON_WRITE,
+            aggregation=variables_lib.VariableAggregation.SUM)
+
+      @def_function.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    with mock.patch.object(failure_handling_util, 'on_gcp', lambda: False):
+
+      with strategy.scope():
+        model = Model()
+        # Named it fh_ckpt because it'd be better that the user have their
+        # regular checkpoint separate from the checkpoint for
+        # PreemptionCheckpointHandler, since we will create CheckpointManager
+        # to manage the checkpoint and only one CheckpointManager should be
+        # active in a particular directory at a time.
+        fh_ckpt = tracking_util.Checkpoint(model=model)
+        if input_arg == 'checkpoint':
+          checkpoint_or_manager = fh_ckpt
+        else:
+          checkpoint_or_manager = _make_checkpoint_manager(
+              fh_ckpt, checkpoint_dir, strategy.cluster_resolver)
+        preemption_handler = (
+            failure_handling.PreemptionCheckpointHandler(
+                strategy.cluster_resolver, checkpoint_or_manager,
+                checkpoint_dir, termination_config))
+
+      def distributed_train_step(current_epoch, current_step):
+
+        @def_function.function
+        def train_step():
+          if cluster_spec and (
+              distribute_lib.get_distribution_strategy(
+              ).cluster_resolver.task_id == raise_app_error_on_worker):
+            raise errors_impl.ResourceExhaustedError(
+                node_def=None, op=None, message='Running out of resources')
+
+          model.v.assign_add(constant_op.constant(1.))
+
+        strategy.run(train_step)
+
+        if current_step == STEPS_PER_EPOCH - 1:
+          logging.info('epoch %d finished', current_epoch)
+
+      logging.info('Start training at %d',
+                   preemption_handler.total_run_calls)
+
+      # If the training process has been restarted, verify that the expected
+      # number of checkpoints have been written.
+      # we also want to check training_finished, because there's a corner case
+      # where the signal is sent quite late and training finishes before the
+      # grace period ends.
+      if training_restarted and training_restarted.is_set(
+      ) and not training_finished.is_set():
+        logging.info('training restarted')
+        match_group = [
+            re.search(r'.*ckpt-(\d+).index', a_file)
+            for a_file in gfile.ListDirectory(checkpoint_dir)
+        ]
+        checkpoint_index = [
+            a_match.group(1) for a_match in match_group if a_match
+        ]
+        if getattr(termination_config, 'grace_period', 0):
+          # Two checkpoints were saved for the extended grace period.
+          self.assertEqual(int(checkpoint_index[0]), 2)
+        else:
+          self.assertEqual(int(checkpoint_index[0]), 1)
+
+      for epoch in range(
+          preemption_handler.total_run_calls // STEPS_PER_EPOCH,
+          EPOCHS_TO_RUN):
+
+        for step in range(
+            preemption_handler.total_run_calls % STEPS_PER_EPOCH,
+            STEPS_PER_EPOCH):
+          if api_wrapping_train:
+            preemption_handler.run(distributed_train_step, epoch, step)
+          else:
+            preemption_handler.save_checkpoint_if_preempted()
+            distributed_train_step(epoch, step)
+        # Add some randomness to when preemption actually happens. We should
+        # trigger it for sure if the training is coming to an end and it hasn't
+        # been triggered yet.
+        if epoch >= EPOCHS_TO_RUN - 2:
+          trigger_it = True
+        else:
+          trigger_it = False
+
+        self._maybe_trigger_a_preemption(training_started_event, trigger_it)
+
+      training_finished.set()
+
+      logging.info('Training finished.')
+
+      self.assertEqual(
+          model.v.numpy(),
+          strategy.num_replicas_in_sync * EPOCHS_TO_RUN * STEPS_PER_EPOCH)
+
+  @combinations.generate(
+      combinations.combine(
+          input_arg=['checkpoint', 'manager'],
+          strategy_option=[
+              'MS',
+              'OneDevice',
+              'MWMS_local',
+              'MWMS_multi_worker',
+          ],
+          api_wrapping_train=[True, False]))
+  def test_preemption_checkpointing(self, input_arg, strategy_option,
+                                    api_wrapping_train):
+    has_chief = False
+
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt')
+    def exit_fn_checking_metric():
+      self.assertGreater(tracking_util._preemption_checkpoint_saved_time_usecs.get_cell().value(), 0)  # pylint: disable=line-too-long
+      sys.exit(42)
+
+    termination_config = failure_handling.TerminationConfig(
+        exit_fn=exit_fn_checking_metric)
+
+    if strategy_option == 'MWMS_multi_worker':
+      cluster_spec = multi_worker_test_base.create_cluster_spec(
+          has_chief=has_chief,
+          num_workers=CLUSTER_SIZE)
+      training_started_event = multi_process_runner.manager().Event()
+      training_restarted = multi_process_runner.manager().Event()
+      training_finished = multi_process_runner.manager().Event()
+
+      mpr = multi_process_runner.MultiProcessRunner(
+          self.worker_fn,
+          cluster_spec,
+          args=(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                [training_started_event], None, training_restarted,
+                training_finished, termination_config),
+          kwargs={'api_wrapping_train': api_wrapping_train},
+          rpc_layer=rpc_layer,
+          return_output=True,
+          dependence_on_chief=has_chief)
+
+      logging.info('Cluster starting.')
+      mpr.start()
+      while not training_started_event.is_set():
+        time.sleep(1)
+
+      logging.info('sending sigterm')
+      killed_worker = random.randrange(0, CLUSTER_SIZE)
+      os.kill(mpr.get_process_id('worker', killed_worker), signal.SIGTERM)
+
+      logging.info('sigterm sent')
+      raise_if_not_all_exit(0, mpr)
+
+      logging.info('restarting workers')
+      training_restarted.set()
+      for worker_id in range(CLUSTER_SIZE):
+        mpr.start_single_process('worker', worker_id, cluster_spec)
+      logging.info('workers restarted')
+
+      mpr.join(timeout=270)
+
+    else:
+      cluster_spec = server_lib.ClusterSpec({})
+
+      training_started_event = threading.Event()
+      training_restarted = threading.Event()
+      training_finished = threading.Event()
+
+      def sending_sigterm(training_started_event):
+        while not training_started_event.is_set():
+          time.sleep(1)
+        logging.info('sending sigterm')
+        training_started_event.set()
+        os.kill(os.getpid(), signal.SIGTERM)
+
+      preemption_sender_thread = threading.Thread(
+          target=sending_sigterm, args=(training_started_event,))
+      preemption_sender_thread.start()
+
+      caught_exit = False
+      try:
+        self.worker_fn(checkpoint_dir, cluster_spec, strategy_option, input_arg,
+                       [training_started_event], None, training_restarted,
+                       training_finished)
+
+      except SystemExit as exit_error:
+        caught_exit = True
+        # We cannot use assertRaise instead, since termination is not always
+        # triggered.
+        self.assertEqual(exit_error.code, 42)  # pylint: disable=g-assert-in-except
+
+      preemption_sender_thread.join(10)
+      if not training_finished.is_set():
+        self.assertTrue(caught_exit)
+
+        logging.info('restarting workers')
+        training_restarted.set()
+        self.worker_fn(checkpoint_dir, cluster_spec, strategy_option, input_arg,
+                       [training_started_event], None, training_restarted,
+                       training_finished)
+
+  def test_error_propagation(self):
+    error_worker = random.randint(0, CLUSTER_SIZE)
+    cluster_spec = multi_worker_test_base.create_cluster_spec(
+        has_chief=False, num_workers=CLUSTER_SIZE)
+    checkpoint_dir = self.get_temp_dir()
+
+    def assert_raise_error():
+      # Asserts that an error raised during a training step on one of the worker
+      # is caught on all workers.
+      with self.assertRaises(errors_impl.ResourceExhaustedError) as error:
+        self.worker_fn(
+            checkpoint_dir,
+            cluster_spec,
+            'MWMS_multi_worker',
+            raise_app_error_on_worker=error_worker)
+      self.assertIn('Running out of resources', str(error.exception))
+
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+
+    mpr = multi_process_runner.MultiProcessRunner(
+        assert_raise_error,
+        cluster_spec,
+        rpc_layer=rpc_layer,
+        return_output=True,
+        dependence_on_chief=False,
+    )
+
+    logging.info('Cluster starting.')
+    mpr.start()
+    mpr.join(timeout=250)
+
+  @combinations.generate(
+      combinations.combine(
+          input_arg=['checkpoint', 'manager'],
+          strategy_option=[
+              'MS',
+              'OneDevice',
+              'MWMS_local',
+              'MWMS_multi_worker',
+          ],
+      )
+  )
+  def test_grace_period_continue_training(self, input_arg, strategy_option):
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt')
+
+    if strategy_option == 'MWMS_multi_worker':
+      grace_period = 5
+      termination_config = failure_handling.TerminationConfig(
+          grace_period=grace_period)
+      has_chief = False
+      cluster_spec = multi_worker_test_base.create_cluster_spec(
+          has_chief=has_chief,
+          num_workers=CLUSTER_SIZE)
+      training_started_event = multi_process_runner.manager().Event()
+      training_restarted = multi_process_runner.manager().Event()
+      training_finished = multi_process_runner.manager().Event()
+
+      mpr = multi_process_runner.MultiProcessRunner(
+          self.worker_fn,
+          cluster_spec,
+          args=(checkpoint_dir, cluster_spec, strategy_option, input_arg,
+                [training_started_event], None, training_restarted,
+                training_finished, termination_config),
+          rpc_layer=rpc_layer,
+          return_output=True,
+          dependence_on_chief=has_chief)
+
+      logging.info('Cluster starting.')
+      mpr.start()
+      while not training_started_event.is_set():
+        time.sleep(1)
+
+      killed_worker = random.randrange(0, CLUSTER_SIZE)
+      logging.info('sending SIGTERM')
+      os.kill(mpr.get_process_id('worker', killed_worker), signal.SIGTERM)
+      logging.info('SIGTERM sent')
+
+      raise_if_not_all_exit(grace_period, mpr)
+
+      logging.info('restarting workers')
+      training_restarted.set()
+      for worker_id in range(CLUSTER_SIZE):
+        mpr.start_single_process('worker', worker_id, cluster_spec)
+      logging.info('workers restarted')
+
+      mpr.join(timeout=250)
+
+    else:
+      # This is because single worker trains super fast with regards to the size
+      # of "model" here. With a longer grace period, the training just finishes
+      # within the grace period so we can't verify the exit behavior.
+      grace_period = 1
+      termination_config = failure_handling.TerminationConfig(
+          grace_period=grace_period)
+      cluster_spec = server_lib.ClusterSpec({})
+
+      training_started_event = threading.Event()
+      training_restarted = threading.Event()
+      training_finished = threading.Event()
+      def sending_sigterm(training_started_event):
+        while not training_started_event.is_set():
+          time.sleep(1)
+        logging.info('sending sigterm')
+        training_started_event.set()
+        os.kill(os.getpid(), signal.SIGTERM)
+
+      preemption_sender_thread = threading.Thread(
+          target=sending_sigterm, args=(training_started_event,))
+      preemption_sender_thread.start()
+
+      caught_exit = False
+      try:
+        self.worker_fn(checkpoint_dir, cluster_spec, strategy_option, input_arg,
+                       [training_started_event], None, training_restarted,
+                       training_finished, termination_config)
+
+      except SystemExit as exit_error:
+        caught_exit = True
+        # We cannot use assertRaise instead, since termination is not always
+        # triggered.
+        self.assertEqual(exit_error.code, 42)  # pylint: disable=g-assert-in-except
+
+      preemption_sender_thread.join(10)
+      if not training_finished.is_set():
+        self.assertTrue(caught_exit)
+
+        logging.info('restarting workers')
+        training_restarted.set()
+        self.worker_fn(checkpoint_dir, cluster_spec, strategy_option, input_arg,
+                       [training_started_event], None, training_restarted,
+                       training_finished, termination_config)
+
+  def test_passed_in_save_fn(self):
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt')
+    gfile.MakeDirs(checkpoint_dir)
+
+    save_fn = lambda: print('Do nothing')
+    termination_config = failure_handling.TerminationConfig(
+        save_fn=save_fn)
+
+    has_chief = False
+    cluster_spec = multi_worker_test_base.create_cluster_spec(
+        has_chief=has_chief,
+        num_workers=CLUSTER_SIZE)
+    training_started_event = multi_process_runner.manager().Event()
+
+    mpr = multi_process_runner.MultiProcessRunner(
+        self.worker_fn,
+        cluster_spec,
+        args=(checkpoint_dir, cluster_spec, 'MWMS_multi_worker', 'checkpoint',
+              [training_started_event], None, None, None, termination_config),
+        rpc_layer=rpc_layer,
+        return_output=True,
+        dependence_on_chief=has_chief)
+
+    logging.info('Cluster starting.')
+    mpr.start()
+    while not training_started_event.is_set():
+      time.sleep(1)
+
+    killed_worker = random.randrange(0, CLUSTER_SIZE)
+    logging.info('sending SIGTERM')
+    os.kill(mpr.get_process_id('worker', killed_worker), signal.SIGTERM)
+    logging.info('SIGTERM sent')
+
+    # 5 is the grace period length
+    raise_if_not_all_exit(5, mpr)
+
+    match_group = [
+        re.search(r'.*ckpt-(\d+).index', a_file)
+        for a_file in gfile.ListDirectory(checkpoint_dir)
+    ]
+
+    # By default, as tested by other test cases, checkpoint will be saved.
+    # This passed in save_fn skips it.
+    self.assertEmpty(match_group)
+
+
+@_REGISTER_DECORATOR(PreemptionCheckpointTest)
+def _save_test_case(pickler, obj):
+  def reconstruct(*args, **kwargs):
+    del args, kwargs
+    return PreemptionCheckpointTest()
+
+  return pickler.save_reduce(reconstruct, (), obj=obj)
+
+
+if __name__ == '__main__':
+  test_util.main()

tensorflow/python/distribute/failure_handling/failure_handling_util.py

@@ -0,0 +1,119 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Util of GCE specifics to ingegrate with WorkerPreemptionHandler."""
+import enum
+import os
+import sys
+import requests
+
+from six.moves.urllib import request
+from tensorflow.python.eager import context
+from tensorflow.python.platform import tf_logging as logging
+
+
+GCP_METADATA_HEADER = {'Metadata-Flavor': 'Google'}
+_GCE_METADATA_URL_ENV_VARIABLE = 'GCE_METADATA_IP'
+_RESTARTABLE_EXIT_CODE = 143
+GRACE_PERIOD_GCE = 3600
+
+
+def gce_exit_fn():
+  sys.exit(_RESTARTABLE_EXIT_CODE)
+
+
+def default_tpu_exit_fn():
+  """Default exit function to run after saving checkpoint for TPUStrategy.
+
+  For TPUStrategy, we want the coordinator to exit after workers are down so
+  that restarted coordinator would not connect to workers scheduled to be
+  preempted. This function achieves so by attempting to get a key-value store
+  from coordination service, which will block until workers are done and then
+  returns with error. Then we have the coordinator sys.exit(42) to re-schedule
+  the job.
+  """
+  logging.info('Waiting for workers to exit...')
+  try:
+    context.context().get_config_key_value('BLOCK_TILL_EXIT')
+  except:  # pylint: disable=bare-except
+    logging.info('Restarting cluster due to preemption.')
+    sys.exit(42)
+
+
+def request_compute_metadata(path):
+  """Returns GCE VM compute metadata."""
+  gce_metadata_endpoint = 'http://' + os.environ.get(
+      _GCE_METADATA_URL_ENV_VARIABLE, 'metadata.google.internal')
+  req = request.Request(
+      '%s/computeMetadata/v1/%s' % (gce_metadata_endpoint, path),
+      headers={'Metadata-Flavor': 'Google'})
+  info = request.urlopen(req).read()
+  if isinstance(info, bytes):
+    return info.decode('utf-8')
+  else:
+    return info
+
+
+def termination_watcher_function_gce():
+  result = request_compute_metadata(
+      'instance/maintenance-event') == 'TERMINATE_ON_HOST_MAINTENANCE'
+  return result
+
+
+def on_gcp():
+  """Detect whether the current running environment is on GCP."""
+  gce_metadata_endpoint = 'http://' + os.environ.get(
+      _GCE_METADATA_URL_ENV_VARIABLE, 'metadata.google.internal')
+
+  try:
+    # Timeout in 5 seconds, in case the test environment has connectivity issue.
+    # There is not default timeout, which means it might block forever.
+    response = requests.get(
+        '%s/computeMetadata/v1/%s' %
+        (gce_metadata_endpoint, 'instance/hostname'),
+        headers=GCP_METADATA_HEADER,
+        timeout=5)
+    return response.status_code == 200
+  except requests.exceptions.RequestException:
+    return False
+
+
+@enum.unique
+class PlatformDevice(enum.Enum):
+  INTERNAL_CPU = 'internal_CPU'
+  INTERNAL_GPU = 'internal_GPU'
+  INTERNAL_TPU = 'internal_TPU'
+  GCE_GPU = 'GCE_GPU'
+  GCE_TPU = 'GCE_TPU'
+  GCE_CPU = 'GCE_CPU'
+  UNSUPPORTED = 'unsupported'
+
+
+def detect_platform():
+  """Returns the platform and device information."""
+  if on_gcp():
+    if context.context().list_logical_devices('GPU'):
+      return PlatformDevice.GCE_GPU
+    elif context.context().list_logical_devices('TPU'):
+      return PlatformDevice.GCE_TPU
+    else:
+      return PlatformDevice.GCE_CPU
+
+  else:
+    if context.context().list_logical_devices('GPU'):
+      return PlatformDevice.INTERNAL_GPU
+    elif context.context().list_logical_devices('TPU'):
+      return PlatformDevice.INTERNAL_TPU
+    else:
+      return PlatformDevice.INTERNAL_CPU

tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py

@@ -0,0 +1,567 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Tests for GCE specifics of PreemptionCheckpointHandler."""
+import os
+import random
+import re
+import sys
+import threading
+import time
+import urllib
+
+from absl.testing import parameterized
+
+# pylint:disable=g-direct-tensorflow-import
+from tensorflow.python.checkpoint import checkpoint as tracking_util
+from tensorflow.python.checkpoint import checkpoint_management
+from tensorflow.python.distribute import collective_all_reduce_strategy
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.distribute import one_device_strategy as one_device_lib
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute.failure_handling import failure_handling
+from tensorflow.python.distribute.failure_handling import failure_handling_util
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.module import module
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import server_lib
+
+
+try:
+  import dill  # pylint:disable=g-import-not-at-top
+
+  _REGISTER_DECORATOR = dill.register
+except ImportError:
+  _REGISTER_DECORATOR = lambda fn, *_: fn
+
+mock = test.mock
+
+
+CLUSTER_SIZE = 4
+EPOCHS_TO_RUN = 5
+STEPS_PER_EPOCH = 6
+_PEER_WATCHER_THREAD_PREFIX = 'PeerTerminationWatcher'
+_LOCAL_WATCHER_THREAD_PREFIX = 'WorkerTerminationSignalWatcher'
+MAX_WAIT_TIME = 30
+
+
+def _is_oss():
+  """Returns whether the test is run under OSS."""
+  return len(sys.argv) >= 1 and 'bazel' in sys.argv[0]
+
+
+def _make_checkpoint_manager(checkpoint, checkpoint_dir, cluster_resolver):
+
+  if not cluster_resolver or not cluster_resolver.cluster_spec().as_dict() or (
+      multi_worker_util.is_chief(
+          cluster_spec=cluster_resolver.cluster_spec(),
+          task_type=cluster_resolver.task_type,
+          task_id=cluster_resolver.task_id)):
+    return checkpoint_management.CheckpointManager(
+        checkpoint, directory=checkpoint_dir, max_to_keep=1)
+  else:
+    return checkpoint_management.CheckpointManager(
+        checkpoint,
+        directory=failure_handling._non_chief_checkpoint_dir(
+            checkpoint_dir, cluster_resolver.task_id),
+        max_to_keep=1)
+
+
+def raise_if_not_all_exit(grace_period, mpr):
+  """Wait for all cluster to exit with a time out."""
+  waiting_time = 0
+  exit_process_count = 0
+  # This addition to mitigate the fact that our step time is too short in test
+  while exit_process_count != CLUSTER_SIZE and waiting_time < max(
+      grace_period + 15, MAX_WAIT_TIME):
+    exit_process_count = 0
+    for worker_id in range(CLUSTER_SIZE):
+      if not mpr.process_exists('worker', worker_id):
+        exit_process_count += 1
+    waiting_time += 1
+    time.sleep(1)
+
+  if waiting_time == max(grace_period + 5, 40):
+    raise RuntimeError('Waited long but at least one worker still exist. '
+                       'Considering size of our model, this should not'
+                       ' happen.')
+
+
+class GceFailureHandlingTest(test.TestCase, parameterized.TestCase):
+  """Integration test for PreemptionCheckpointHandler."""
+
+  def worker_fn(
+      self,
+      checkpoint_dir,
+      cluster_spec,
+      input_arg,
+      strategy_option,
+      maintenance_event=None,
+      training_finished=None,
+      frequent_send=False,
+      training_restarted=None,
+      termination_config=failure_handling.TerminationConfig(grace_period=0),
+      api_wrapping_train=True,
+  ):
+    if strategy_option == 'MS':
+      strategy = mirrored_strategy.MirroredStrategy()
+    elif strategy_option == 'OneDevice':
+      if config.list_physical_devices('GPU'):
+        strategy = one_device_lib.OneDeviceStrategy(device='/gpu:0')
+      else:
+        strategy = one_device_lib.OneDeviceStrategy(device='/cpu:0')
+    else:
+      strategy = collective_all_reduce_strategy.CollectiveAllReduceStrategy()
+
+    def mock_termination_watcher_function_gce(*args, **kwargs):
+      del args, kwargs
+      if not frequent_send:
+        time.sleep(1)
+        if (not maintenance_event.is_set()) and (random.randrange(0, 11) == 5):
+          maintenance_event.set()
+          logging.info('Termination notice available.')
+          return True
+
+      elif frequent_send and not maintenance_event.is_set():
+        logging.info('Termination notice available.')
+        return True
+
+      return False
+
+    with mock.patch.object(
+        failure_handling_util, 'termination_watcher_function_gce',
+        mock_termination_watcher_function_gce), mock.patch.object(
+            failure_handling_util, 'detect_platform',
+            lambda: failure_handling_util.PlatformDevice.GCE_GPU):
+
+      class Model(module.Module):
+
+        def __init__(self):
+          self.v = variables_lib.Variable(
+              0.,
+              synchronization=variables_lib.VariableSynchronization.ON_WRITE,
+              aggregation=variables_lib.VariableAggregation.SUM)
+
+        @def_function.function(input_signature=[])
+        def __call__(self):
+          return self.v.read_value()
+
+      with strategy.scope():
+        model = Model()
+        fh_ckpt = tracking_util.Checkpoint(model=model)
+
+        if input_arg == 'checkpoint':
+          checkpoint_or_manager = fh_ckpt
+        else:
+          checkpoint_or_manager = _make_checkpoint_manager(
+              fh_ckpt, checkpoint_dir, strategy.cluster_resolver)
+        preemption_handler = (
+            failure_handling.PreemptionCheckpointHandler(
+                strategy.cluster_resolver, checkpoint_or_manager,
+                checkpoint_dir, termination_config))
+
+      def distributed_train_step(current_epoch, current_step):
+
+        @def_function.function
+        def train_step():
+          model.v.assign_add(constant_op.constant(1.))
+
+        strategy.run(train_step)
+
+        if current_step == STEPS_PER_EPOCH - 1:
+          logging.info('epoch %d finished', current_epoch)
+
+      logging.info('Start training at %d',
+                   preemption_handler.total_run_calls)
+
+      # If the training process has been restarted, verify that the expected
+      # number of checkpoints have been written.
+      # We also want to check training_finished, because there's a corner case
+      # where the signal is sent quite late and training finishes before the
+      # grace period ends.
+      if training_restarted.is_set() and not training_finished.is_set():
+        logging.info(gfile.ListDirectory(checkpoint_dir))
+        match_group = [
+            re.search(r'.*ckpt-(\d+).index', a_file)
+            for a_file in gfile.ListDirectory(checkpoint_dir)
+        ]
+        checkpoint_index = [
+            a_match.group(1) for a_match in match_group if a_match
+        ]
+        self.assertNotEmpty(checkpoint_index)
+
+        if api_wrapping_train:
+          if termination_config.grace_period > 0:
+            # Two checkpoints were saved for the extended grace period.
+            self.assertEqual(
+                max([int(ckpt_index) for ckpt_index in checkpoint_index]), 2)
+          else:
+            self.assertEqual(
+                max([int(ckpt_index) for ckpt_index in checkpoint_index]), 1)
+
+        else:
+          # Test if arguments to save_checkpoint_if_preempted are passed
+          # successfully.
+          self.assertEqual(
+              max([int(ckpt_index) for ckpt_index in checkpoint_index]),
+              preemption_handler.total_run_calls)
+
+      for epoch in range(preemption_handler.total_run_calls // STEPS_PER_EPOCH,
+                         EPOCHS_TO_RUN):
+
+        for step in range(preemption_handler.total_run_calls % STEPS_PER_EPOCH,
+                          STEPS_PER_EPOCH):
+
+          # Testing two different APIs to save checkpoint.
+          if api_wrapping_train:
+            preemption_handler.run(distributed_train_step, epoch, step)
+
+          else:
+            preemption_handler.save_checkpoint_if_preempted(
+                checkpoint_number=preemption_handler.total_run_calls)
+            distributed_train_step(epoch, step)
+
+      logging.info('Training finished.')
+      training_finished.set()
+
+      self.assertEqual(
+          model.v.numpy(),
+          strategy.num_replicas_in_sync * EPOCHS_TO_RUN * STEPS_PER_EPOCH)
+
+      running_threads = test_util.get_running_threads()
+      if test_util.has_thread(_PEER_WATCHER_THREAD_PREFIX,
+                              running_threads) and test_util.has_thread(
+                                  _LOCAL_WATCHER_THREAD_PREFIX,
+                                  running_threads):
+        try:
+          # Explicitly call __del__ since making it None and gc.collect does
+          # not invoke __del__ here.
+          preemption_handler.__del__()
+
+          time.sleep(2)
+
+          running_threads = test_util.get_running_threads()
+          self.assertFalse(
+              test_util.has_thread(_LOCAL_WATCHER_THREAD_PREFIX,
+                                   running_threads))
+          self.assertFalse(
+              test_util.has_thread(_PEER_WATCHER_THREAD_PREFIX,
+                                   running_threads))
+
+        except urllib.error.URLError as e:
+          if 'Temporary failure in name resolution' in e.message:
+            # This is caused by a weird flakiness that mock.patch does not
+            # correctly patch failure_handling_util.request_compute_metadata, a
+            # real request is attempted, and an error is hit in
+            # failure_handling_util.request_compute_metadata
+            logging.warning('Hit a mock issue.')
+            return
+
+  @combinations.generate(
+      combinations.combine(
+          input_arg=['checkpoint', 'manager'],
+          strategy_option=[
+              'MS',
+              'OneDevice',
+              'MWMS_local',
+              'MWMS_multi_worker',
+          ],
+      )
+  )
+  def test_basic_run(self, input_arg, strategy_option):
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt/')
+
+    if strategy_option == 'MWMS_multi_worker':
+      has_chief = False
+      cluster_spec = multi_worker_test_base.create_cluster_spec(
+          has_chief=has_chief,
+          num_workers=CLUSTER_SIZE)
+      maintenance_event = multi_process_runner.manager().Event()
+      training_finished = multi_process_runner.manager().Event()
+      training_restarted = multi_process_runner.manager().Event()
+
+      mpr = multi_process_runner.MultiProcessRunner(
+          self.worker_fn,
+          cluster_spec,
+          args=(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                maintenance_event, training_finished, False,
+                training_restarted),
+          rpc_layer=rpc_layer,
+          return_output=True,
+          dependence_on_chief=has_chief)
+
+      logging.info('Cluster starting.')
+      mpr.start()
+
+      raise_if_not_all_exit(0, mpr)
+
+      if not training_finished.is_set():
+        logging.info('restarting workers')
+        training_restarted.set()
+        for worker_id in range(CLUSTER_SIZE):
+          mpr.start_single_process('worker', worker_id, cluster_spec)
+        logging.info('workers restarted')
+
+      mpr.join(timeout=250)
+      self.assertTrue(training_finished.is_set())
+
+    else:
+      maintenance_event = threading.Event()
+      training_finished = threading.Event()
+      training_restarted = threading.Event()
+
+      cluster_spec = server_lib.ClusterSpec({})
+      caught_exit = False
+      try:
+        self.worker_fn(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                       maintenance_event, training_finished, False,
+                       training_restarted)
+      except SystemExit as exit_error:
+        caught_exit = True
+        # We cannot use assertRaise instead, since termination is not always
+        # triggered.
+        self.assertEqual(exit_error.code, 143)  # pylint:disable=g-assert-in-except
+
+      if maintenance_event.is_set() and not training_finished.is_set():
+        self.assertTrue(caught_exit)
+
+        logging.info('restarting workers')
+        training_restarted.set()
+        self.worker_fn(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                       maintenance_event, training_finished, False,
+                       training_restarted)
+
+      self.assertTrue(training_finished.is_set())
+
+  @combinations.generate(
+      combinations.combine(
+          grace_period=[0, 11],
+          input_arg=['checkpoint', 'manager'],
+          strategy_option=[
+              'MS',
+              'OneDevice',
+              'MWMS_local',
+              'MWMS_multi_worker',
+          ],
+          api_wrapping_train=[True, False]
+      ))
+  def test_multiple_workers_preempted_consecutively(self, grace_period,
+                                                    input_arg, strategy_option,
+                                                    api_wrapping_train):
+
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt/')
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+
+    termination_config = failure_handling.TerminationConfig(
+        grace_period=grace_period)
+
+    if strategy_option == 'MWMS_multi_worker':
+      has_chief = False
+      cluster_spec = multi_worker_test_base.create_cluster_spec(
+          has_chief=has_chief,
+          num_workers=CLUSTER_SIZE)
+      has_chief = False
+      maintenance_event = multi_process_runner.manager().Event()
+      training_finished = multi_process_runner.manager().Event()
+      training_restarted = multi_process_runner.manager().Event()
+      mpr = multi_process_runner.MultiProcessRunner(
+          self.worker_fn,
+          cluster_spec,
+          args=(
+              checkpoint_dir,
+              cluster_spec,
+              input_arg,
+              strategy_option,
+              maintenance_event,
+              training_finished,
+              True,
+              training_restarted,
+              termination_config,
+          ),
+          kwargs={'api_wrapping_train': api_wrapping_train},
+          rpc_layer=rpc_layer,
+          return_output=True,
+          dependence_on_chief=has_chief,
+      )
+
+      logging.info('Cluster starting.')
+      mpr.start()
+
+      raise_if_not_all_exit(grace_period, mpr)
+
+      maintenance_event.set()
+      logging.info('restarting workers')
+      training_restarted.set()
+      for worker_id in range(CLUSTER_SIZE):
+        mpr.start_single_process('worker', worker_id, cluster_spec)
+      logging.info('workers restarted')
+
+      mpr.join(timeout=250)
+      self.assertTrue(training_finished.is_set())
+
+    else:
+      maintenance_event = threading.Event()
+      training_finished = threading.Event()
+      training_restarted = threading.Event()
+
+      cluster_spec = server_lib.ClusterSpec({})
+
+      caught_exit = False
+      try:
+        self.worker_fn(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                       maintenance_event, training_finished, True,
+                       training_restarted, termination_config)
+      except SystemExit as exit_error:
+        caught_exit = True
+        # We cannot use assertRaise instead, since termination is not always
+        # triggered.
+        self.assertEqual(exit_error.code, 143)  # pylint:disable=g-assert-in-except
+
+      if maintenance_event.is_set() and not training_finished.is_set():
+        self.assertTrue(caught_exit)
+
+        logging.info('restarting workers')
+        training_restarted.set()
+        self.worker_fn(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                       maintenance_event, training_finished, True,
+                       training_restarted, termination_config)
+
+  @combinations.generate(
+      combinations.combine(
+          input_arg=['checkpoint', 'manager'],
+          strategy_option=[
+              'MS',
+              'OneDevice',
+              'MWMS_local',
+              'MWMS_multi_worker',
+          ],
+          api_wrapping_train=[True, False]))
+  def test_grace_period_continue_training(self, input_arg, strategy_option,
+                                          api_wrapping_train):
+    checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt/')
+    grace_period = 11
+    if _is_oss():
+      rpc_layer = 'grpc'
+    else:
+      rpc_layer = 'grpc+loas'
+
+    termination_config = failure_handling.TerminationConfig(
+        grace_period=grace_period)
+
+    if strategy_option == 'multi_worker':
+      has_chief = False
+      cluster_spec = multi_worker_test_base.create_cluster_spec(
+          has_chief=has_chief,
+          num_workers=CLUSTER_SIZE)
+
+      checkpoint_dir = os.path.join(self.get_temp_dir(), 'fh_ckpt/')
+      maintenance_event = multi_process_runner.manager().Event()
+      training_finished = multi_process_runner.manager().Event()
+      training_restarted = multi_process_runner.manager().Event()
+
+      mpr = multi_process_runner.MultiProcessRunner(
+          self.worker_fn,
+          cluster_spec,
+          args=(
+              checkpoint_dir,
+              cluster_spec,
+              input_arg,
+              strategy_option,
+              maintenance_event,
+              training_finished,
+              False,
+              training_restarted,
+              termination_config,
+          ),
+          kwargs={'api_wrapping_train': api_wrapping_train},
+          rpc_layer=rpc_layer,
+          return_output=True,
+          dependence_on_chief=has_chief,
+      )
+
+      logging.info('Cluster starting.')
+      mpr.start()
+
+      while (not maintenance_event.is_set()) and (
+          not training_finished.is_set()):
+        time.sleep(1)
+
+      raise_if_not_all_exit(grace_period, mpr)
+
+      if not training_finished.is_set():
+        logging.info('restarting workers')
+        training_restarted.set()
+        for worker_id in range(CLUSTER_SIZE):
+          mpr.start_single_process('worker', worker_id, cluster_spec)
+        logging.info('workers restarted')
+
+      mpr.join(timeout=250)
+
+      self.assertTrue(training_finished.is_set())
+
+    else:
+      maintenance_event = threading.Event()
+      training_finished = threading.Event()
+      training_restarted = threading.Event()
+
+      cluster_spec = server_lib.ClusterSpec({})
+      caught_exit = False
+      try:
+        self.worker_fn(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                       maintenance_event, training_finished, False,
+                       training_restarted, termination_config)
+      except SystemExit as exit_error:
+        caught_exit = True
+        # We cannot use assertRaise instead, since termination is not always
+        # triggered.
+        self.assertEqual(exit_error.code, 143)  # pylint:disable=g-assert-in-except
+
+      if maintenance_event.is_set() and not training_finished.is_set():
+        self.assertTrue(caught_exit)
+
+        logging.info('restarting workers')
+        training_restarted.set()
+        self.worker_fn(checkpoint_dir, cluster_spec, input_arg, strategy_option,
+                       maintenance_event, training_finished, False,
+                       training_restarted, termination_config)
+        self.assertTrue(training_finished.is_set())
+
+
+@_REGISTER_DECORATOR(GceFailureHandlingTest)
+def _save_test_case(pickler, obj):
+  def reconstruct(*args, **kwargs):
+    del args, kwargs
+    return GceFailureHandlingTest()
+
+  return pickler.save_reduce(reconstruct, (), obj=obj)
+
+
+if __name__ == '__main__':
+  test_util.main()

tensorflow/python/distribute/failure_handling/preemption_watcher.py

@@ -0,0 +1,137 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Provides a utility class for preemption detection and recovery."""
+
+import threading
+
+from absl import logging
+
+from tensorflow.python.distribute.failure_handling.failure_handling_util import detect_platform
+from tensorflow.python.distribute.failure_handling.failure_handling_util import PlatformDevice
+from tensorflow.python.eager import context
+from tensorflow.python.eager import monitoring
+from tensorflow.python.framework.errors import AbortedError
+from tensorflow.python.framework.errors import CancelledError
+from tensorflow.python.framework.errors import InternalError
+from tensorflow.python.framework.errors import UnavailableError
+from tensorflow.python.util.tf_export import tf_export
+
+
+_preemption_watcher_initialization_counter = monitoring.Counter(
+    "/tensorflow/api/distribution_strategy/preemption_watcher_initialized",
+    "Counter for usages of PreemptionWatcher",
+)
+_preemption_handling_counter = monitoring.Counter(
+    "/tensorflow/api/distribution_strategy/preemption_watcher_handled",
+    "Counter for number of preempions caught and handled by PreemptionWatcher",
+)
+
+_PREEMPTION_KEY = "TF_DEFAULT_PREEMPTION_NOTICE_KEY"
+
+
+@tf_export("distribute.experimental.PreemptionWatcher", v1=[])
+class PreemptionWatcher:
+  """Watch preemption signal and store it.
+
+  Notice: Currently only support Borg TPU environment with TPUClusterResolver.
+
+  This class provides a way to monitor the preemption signal during training on
+  TPU. It will start a background thread to watch the training process, trying
+  to fetch preemption message from the coordination service. When preemption
+  happens, the preempted worker will write the preemption message to the
+  coordination service. Thus getting a non-empty preemption message means there
+  is a preemption happened.
+
+  User can use the preemption message as a reliable preemption indicator, and
+  then set the coordinator to reconnect to the TPU worker instead of a fully
+  restart triggered by Borg. For example, a training process with
+  preemption recovery will be like:
+
+  ```python
+  keep_running = True
+  preemption_watcher = None
+  while keep_running:
+    try:
+      # Initialize TPU cluster and stratygy.
+      resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
+      tf.config.experimental_connect_to_cluster(resolver)
+      tf.tpu.experimental.initialize_tpu_system(resolver)
+      strategy = tf.distribute.TPUStrategy(resolver)
+
+      # PreemptionWatcher must be created after connected to cluster.
+      preemption_watcher = tf.distribute.experimental.PreemptionWatcher()
+      train_model(strategy)
+      keep_running = False
+    except Exception as e:
+      if preemption_watcher and preemption_watcher.preemption_message:
+        preemption_watcher.block_until_worker_exit()
+        keep_running = True
+      else:
+        raise e
+  ```
+
+  Attributes:
+    preemption_message: A variable to store the preemption message fetched from
+      the coordination service. If it is not None, then there is a preemption
+      happened.
+    platform: A PlatformDevice to indicate the current job's platform. Refer to
+      failure_handling_util.py for the definition of enum class PlatformDevice.
+  """
+
+  def __init__(self):
+    # TODO(b/254321514): Integrate with GPU and cloud enviornmenmt.
+    self._preemption_message = None
+    self._platform = detect_platform()
+    if self._platform != PlatformDevice.INTERNAL_TPU:
+      logging.warning(
+          "Preemption watcher does not support environment: %s", self._platform
+      )
+    else:
+      _preemption_watcher_initialization_counter.get_cell().increase_by(1)
+      threading.Thread(target=self._watch_preemption_key, daemon=True).start()
+
+  @property
+  def preemption_message(self):
+    """Returns the preemption message."""
+    return self._preemption_message
+
+  def _watch_preemption_key(self):
+    logging.info("Watching preemption signal.")
+    message = context.context().get_config_key_value(_PREEMPTION_KEY)
+    _preemption_handling_counter.get_cell().increase_by(1)
+    logging.info("Preemption signal received.")
+    self._preemption_message = message
+
+  def block_until_worker_exit(self):
+    """Block coordinator until workers exit.
+
+    In some rare cases, another error could be raised during the
+    preemption grace period. This will cause the coordinator to reconnect to the
+    same TPU workers, which will be killed later. It prevents the coordinator to
+    reconnect to new TPU workers, and falls back to a hard restart. To avoid
+    this situation, this method will block the coordinator to reconnect until
+    workers exit. This method will be a no-op for non-TPU platform.
+    """
+    if self._platform != PlatformDevice.INTERNAL_TPU:
+      return
+    try:
+      context.context().get_config_key_value("BLOCK_TILL_EXIT")
+    except InternalError as e:
+      # Ensure that internal error is related to coordination service.
+      if "Coordination service is not enabled." not in e.message:
+        raise
+      logging.info("Workers exited.")
+    except (AbortedError, CancelledError, UnavailableError):
+      logging.info("Workers exited.")

tensorflow/python/distribute/input_lib_type_spec_test.py

@@ -0,0 +1,761 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for the input_lib library which tests iterator type specs."""
+
+import os
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python import tf2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import readers
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.distribute import values
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_tensor as ragged_tensor_lib
+from tensorflow.python.util import nest
+
+
+class DistributedIteratorTest(test.TestCase,
+                              parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          input_type=["dataset"],
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpec(self, input_type, distribution,
+                   enable_get_next_as_optional):
+    if not tf2.enabled():
+      self.skipTest("DistributedIterator has CompositeTensor support in "
+                    "TF 2 only.")
+    dataset = dataset_ops.DatasetV2.range(10).batch(2)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    with distribution.scope():
+      iterator = iter(dist_dataset)
+      _check_type_spec_structure(iterator)
+
+    spec = iterator._type_spec
+    self.assertEqual(spec._input_workers, iterator._input_workers)
+    self.assertEqual(spec._element_spec._value_specs,
+                     (tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64,
+                                             name=None),
+                      tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64,
+                                             name=None)))
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          input_type=["dataset"],
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpecRoundTrip(self, input_type,
+                            distribution, enable_get_next_as_optional):
+    if not tf2.enabled():
+      self.skipTest("DistributedIterator CompositeTensor support is only "
+                    "present in TF 2.0 only.")
+
+    dataset = dataset_ops.DatasetV2.range(10).batch(2)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    with distribution.scope():
+      iterator = iter(dist_dataset)
+      _check_type_spec_structure(iterator)
+
+    spec = iterator._type_spec
+
+    tensor_list = spec._to_components(iterator)
+    re_iterator = spec._from_components(tensor_list)
+
+    self.assertEqual(iterator._input_workers, re_iterator._input_workers)
+    self.assertAllEqual(iterator._iterators, re_iterator._iterators)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          input_type=["dataset"],
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+          ],
+          enable_get_next_as_optional=[True, False],
+          drop_remainder=[True, False],
+          tf_api_version=2,
+      ))
+  def testDoesNotTriggerFunctionTracing(self, input_type, distribution,
+                                        enable_get_next_as_optional,
+                                        drop_remainder):
+    trace_count = [0]
+
+    @def_function.function
+    def f(iterator):
+      trace_count[0] += 1
+      counter = np.int64(0)
+      for _ in range(5):
+        next(iterator)
+        counter += 1
+      return counter
+
+    dataset = dataset_ops.DatasetV2.range(10).batch(
+        2, drop_remainder=drop_remainder)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    with distribution.scope():
+      for _ in range(3):
+        iterator = iter(dist_dataset)
+        _check_type_spec_structure(iterator)
+        counter = f(iterator)
+
+        self.assertEqual(trace_count[0], 1)
+        self.assertEqual(counter, 5)
+
+
+class InputTypeSpecTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          distribution=[
+              strategy_combinations.one_device_strategy,
+              strategy_combinations.mirrored_strategy_with_one_cpu,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+              strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
+              strategy_combinations.multi_worker_mirrored_2x1_cpu,
+              strategy_combinations.multi_worker_mirrored_2x1_gpu,
+              strategy_combinations.multi_worker_mirrored_2x2_gpu,
+          ],
+          tf_api_version=2,
+          enable_get_next_as_optional=[True, False],
+          drop_remainder=[True, False],
+      ))
+  def testInputSignatureForPerReplicaValues(self, distribution,
+                                            enable_get_next_as_optional,
+                                            drop_remainder):
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+    ds = dataset_ops.DatasetV2.from_tensor_slices(
+        np.ones([9, 12]).astype(np.float32)).batch(
+            4, drop_remainder=drop_remainder)
+    ds = distribution.experimental_distribute_dataset(ds)
+    _check_type_spec_structure(iter(ds))
+    element_spec = ds.element_spec
+    iter_element_spec = iter(ds).element_spec
+    nest.assert_same_structure(element_spec, iter_element_spec)
+    self.assertAllEqual(
+        nest.flatten(element_spec), nest.flatten(iter_element_spec))
+
+    @def_function.function(input_signature=[element_spec])
+    def process_inputs(inputs):
+      distribution.run(lambda inputs: inputs, args=(inputs,))
+
+    for x in ds:
+      process_inputs(x)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          distribution=[
+              strategy_combinations.one_device_strategy,
+              strategy_combinations.mirrored_strategy_with_one_cpu,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+              strategy_combinations.central_storage_strategy_with_two_gpus,
+          ],
+      ))
+  def testInputSignatureForNestedPerReplicaValues(self, distribution):
+    a = np.ones((10, 2)) * 5
+    b = np.ones((10, 3)) * 6
+    dataset = dataset_ops.DatasetV2.from_tensor_slices((a, b)).batch(2)
+
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+
+    @def_function.function(input_signature=[dist_dataset.element_spec])
+    def process_inputs(inputs):
+      distribution.run(lambda inputs: inputs, args=(inputs,))
+
+    for x in dist_dataset:
+      process_inputs(x)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          input_type=["dataset"],
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testMostSpecificCompatibleType(self, input_type, distribution,
+                                     enable_get_next_as_optional):
+    if not tf2.enabled():
+      self.skipTest("DistributedIterator has CompositeTensor support in "
+                    "TF 2 only.")
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    ds1 = dataset_ops.DatasetV2.range(10).batch(2).batch(5)
+    ds2 = dataset_ops.DatasetV2.from_tensors(
+        array_ops.zeros([5, 2], dtypes.int64))
+    dist_ds1 = distribution.experimental_distribute_dataset(ds1)
+    dist_ds2 = distribution.experimental_distribute_dataset(ds2)
+
+    with distribution.scope():
+      iter1 = iter(dist_ds1)
+      iter2 = iter(dist_ds2)
+
+    spec1 = iter1._type_spec  # Wrapped TensorSpec has shape [None, None]
+    spec2 = iter2._type_spec  # Wrapped TensorSpec has shape [None, 2]
+
+    self.assertNotEqual(spec1, spec2)
+    self.assertEqual(spec1, spec1.most_specific_compatible_type(spec2))
+    self.assertEqual(spec1, spec2.most_specific_compatible_type(spec1))
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          ],
+          enable_get_next_as_optional=[True, False],
+          experimental_place_dataset_on_device=[True, False],
+          experimental_fetch_to_device=[True, False],
+      ))
+  def testFromFunctionInputSignatureForPerReplicaValuesWithOptions(
+      self, distribution, enable_get_next_as_optional,
+      experimental_place_dataset_on_device, experimental_fetch_to_device):
+
+    if experimental_place_dataset_on_device and experimental_fetch_to_device:
+      self.skipTest("Setting experimental_place_dataset_on_device and "
+                    "experimental_fetch_to_device to `True` is not "
+                    "allowed when using "
+                    "distribute_lib.InputReplicationMode.PER_REPLICA.")
+
+    fname1 = os.path.join(self.get_temp_dir(), "1.txt")
+    _create_text_file(fname1, 5)
+    fname2 = os.path.join(self.get_temp_dir(), "2.txt")
+    _create_text_file(fname2, 9)
+
+    def dataset_fn(input_context):
+      dataset = dataset_ops.DatasetV2.from_tensor_slices([fname1, fname2])
+      dataset = dataset.shard(input_context.num_input_pipelines,
+                              input_context.input_pipeline_id)
+      return readers.TextLineDatasetV2(dataset).map(
+          string_ops.string_to_number).batch(
+              input_context.get_per_replica_batch_size(4))
+
+    options = distribute_lib.InputOptions(
+        experimental_place_dataset_on_device=(
+            experimental_place_dataset_on_device),
+        experimental_fetch_to_device=experimental_fetch_to_device,
+        experimental_replication_mode=(
+            distribute_lib.InputReplicationMode.PER_REPLICA))
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+    ds = distribution.experimental_distribute_datasets_from_function(
+        dataset_fn, options)
+
+    iterator = iter(ds)
+    _check_type_spec_structure(iterator)
+    spec = iterator._type_spec
+    tensor_list = spec._to_components(iterator)
+    re_iterator = spec._from_components(tensor_list)
+
+    _check_type_spec_structure(iter(ds))
+    element_spec = ds.element_spec
+    iter_element_spec = iter(ds).element_spec
+    nest.assert_same_structure(element_spec, iter_element_spec)
+    self.assertAllEqual(
+        nest.flatten(element_spec), nest.flatten(iter_element_spec))
+    self.assertEqual(iterator._input_workers, re_iterator._input_workers)
+    self.assertAllEqual(iterator._iterators, re_iterator._iterators)
+
+    @def_function.function(input_signature=[element_spec])
+    def process_inputs(inputs):
+      distribution.run(lambda inputs: inputs, args=(inputs,))
+
+    for x in ds:
+      process_inputs(x)
+
+
+class DistributedDatasetTypeSpecTest(test.TestCase, parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpecBase(self, distribution, enable_get_next_as_optional):
+
+    def create_dataset():
+      dataset = dataset_ops.DatasetV2.range(10).batch(2)
+      return dataset
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    dist_dataset = distribution.experimental_distribute_dataset(
+        create_dataset())
+
+    spec = dist_dataset._type_spec
+    self.assertEqual(spec._input_workers, dist_dataset._input_workers)
+    self.assertEqual(
+        spec._element_spec._value_specs,
+        (tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64, name=None),
+         tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64, name=None)))
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpecReturnedFromTFFunction(self, distribution,
+                                         enable_get_next_as_optional):
+    # TODO(ishark): This is observed when tensor is copied from one device to
+    # other and since DatasetVariantWrapper does not have a copy
+    # function. Some Context: b/146981184
+    # Try to renable with non-canonicalized input workers, which
+    # helped in PS Strategy for similar error.
+    self.skipTest("Failures observed in Ubuntu presubmit: No unary variant  "
+                  "device copy function found for direction: 1 and Variant "
+                  "type_index:tensorflow::data::(anonymous namespace)::"
+                  "DatasetVariantWrapper")
+
+    @def_function.function
+    def create_dist_dataset():
+      dataset = dataset_ops.DatasetV2.range(10).batch(2)
+      return distribution.experimental_distribute_dataset(dataset)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    dist_dataset = create_dist_dataset()
+
+    spec = dist_dataset._type_spec
+    self.assertEqual(spec._input_workers, dist_dataset._input_workers)
+    self.assertEqual(
+        spec._element_spec._value_specs,
+        (tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64, name=None),
+         tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64, name=None)))
+
+    # Read distributed data to confirm values are correct.
+    iterator = iter(dist_dataset)
+    data = []
+    for it in iterator:
+      data.append(distribution.experimental_local_results(it))
+    self.assertAllEqual(
+        nest.flatten(data),
+        list(dataset_ops.DatasetV2.range(10).batch(1).as_numpy_iterator()))
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpecRaggedTensor(self, distribution, enable_get_next_as_optional):
+    ctx = distribute_lib.InputContext()
+    batch_size = ctx.get_per_replica_batch_size(8)
+    # Use 20 which isn't divisible by 8 to test partial batch behavior.
+    row_lengths = np.mod(np.arange(20), 4).astype(np.int64)
+    ragged_tensor = ragged_tensor_lib.RaggedTensor.from_row_lengths(
+        np.repeat(np.arange(20, dtype=np.float32), row_lengths), row_lengths)
+    dataset = dataset_ops.DatasetV2.from_tensor_slices({
+        "dense": ragged_tensor.to_tensor(),
+        "ragged": ragged_tensor,
+        "sparse": ragged_tensor.to_sparse(),
+    })
+    dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
+    dataset = dataset.batch(batch_size)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    spec = dist_dataset._type_spec
+    self.assertEqual(spec._input_workers, dist_dataset._input_workers)
+    self.assertEqual(
+        spec._element_spec, {
+            "sparse":
+                values.PerReplicaSpec(
+                    sparse_tensor.SparseTensorSpec(
+                        tensor_shape.TensorShape([None, 3]), dtypes.float32),
+                    sparse_tensor.SparseTensorSpec(
+                        tensor_shape.TensorShape([None, 3]), dtypes.float32)),
+            "dense":
+                values.PerReplicaSpec(
+                    tensor_spec.TensorSpec(
+                        shape=(None, 3), dtype=dtypes.float32, name=None),
+                    tensor_spec.TensorSpec(
+                        shape=(None, 3), dtype=dtypes.float32, name=None)),
+            "ragged":
+                values.PerReplicaSpec(
+                    ragged_tensor_lib.RaggedTensorSpec(
+                        tensor_shape.TensorShape([None, None]), dtypes.float32,
+                        1, dtypes.int64),
+                    ragged_tensor_lib.RaggedTensorSpec(
+                        tensor_shape.TensorShape([None, None]), dtypes.float32,
+                        1, dtypes.int64))
+        })
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          ],
+          enable_get_next_as_optional=[True, False],
+          experimental_place_dataset_on_device=[True, False],
+          experimental_fetch_to_device=[True, False]))
+  def testTypeSpecComponents(self, distribution, enable_get_next_as_optional,
+                             experimental_place_dataset_on_device,
+                             experimental_fetch_to_device):
+    dataset = dataset_ops.DatasetV2.range(10).batch(2)
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    options = distribute_lib.InputOptions(
+        experimental_place_dataset_on_device=
+        experimental_place_dataset_on_device,
+        experimental_fetch_to_device=experimental_fetch_to_device)
+
+    dist_dataset = distribution.experimental_distribute_dataset(
+        dataset, options)
+
+    spec = dist_dataset._type_spec
+    self.assertEqual(spec._input_workers, dist_dataset._input_workers)
+    self.assertEqual(
+        spec._element_spec._value_specs,
+        (tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64, name=None),
+         tensor_spec.TensorSpec(shape=(None,), dtype=dtypes.int64, name=None)))
+    components = spec._to_components(dist_dataset)
+    re_dist_dataset = spec._from_components(components)
+
+    self.assertEqual(dist_dataset._input_workers,
+                     re_dist_dataset._input_workers)
+    self.assertAllEqual(dist_dataset._cloned_datasets,
+                        re_dist_dataset._cloned_datasets)
+    self.assertEqual(dist_dataset._element_spec, re_dist_dataset._element_spec)
+    self.assertEqual(dist_dataset._enable_get_next_as_optional,
+                     re_dist_dataset._enable_get_next_as_optional)
+    self.assertEqual(dist_dataset._options, re_dist_dataset._options)
+
+
+class DistributedDatasetsFromFunctionSpecTest(test.TestCase,
+                                              parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_two_gpus,
+              strategy_combinations.mirrored_strategy_with_cpu_1_and_2,
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          ],
+          enable_get_next_as_optional=[True, False],
+          experimental_place_dataset_on_device=[True, False],
+          experimental_fetch_to_device=[True, False],
+      ))
+  def testDistributedDatasetsFromFunctionSpec(
+      self, distribution, enable_get_next_as_optional,
+      experimental_place_dataset_on_device, experimental_fetch_to_device):
+
+    if experimental_place_dataset_on_device and experimental_fetch_to_device:
+      self.skipTest("Setting experimental_place_dataset_on_device and "
+                    "experimental_fetch_to_device to `True` is not "
+                    "allowed when using "
+                    "distribute_lib.InputReplicationMode.PER_REPLICA.")
+
+    fname1 = os.path.join(self.get_temp_dir(), "1.txt")
+    _create_text_file(fname1, 5)
+    fname2 = os.path.join(self.get_temp_dir(), "2.txt")
+    _create_text_file(fname2, 9)
+
+    def dataset_fn(input_context):
+      dataset = dataset_ops.DatasetV2.from_tensor_slices([fname1, fname2])
+      dataset = dataset.shard(input_context.num_input_pipelines,
+                              input_context.input_pipeline_id)
+      return readers.TextLineDatasetV2(dataset).map(
+          string_ops.string_to_number).batch(
+              input_context.get_per_replica_batch_size(4))
+
+    options = distribute_lib.InputOptions(
+        experimental_place_dataset_on_device=
+        experimental_place_dataset_on_device,
+        experimental_fetch_to_device=experimental_fetch_to_device,
+        experimental_replication_mode=(
+            distribute_lib.InputReplicationMode.PER_REPLICA))
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+    ds = distribution.experimental_distribute_datasets_from_function(
+        dataset_fn, options)
+
+    spec = ds._type_spec
+    components = spec._to_components(ds)
+    re_ds = spec._from_components(components)
+
+    element_spec = re_ds.element_spec
+    iter_element_spec = iter(ds).element_spec
+    nest.assert_same_structure(element_spec, iter_element_spec)
+    self.assertAllEqual(
+        nest.flatten(element_spec), nest.flatten(iter_element_spec))
+    self.assertEqual(ds._input_workers, re_ds._input_workers)
+    self.assertEqual(ds._element_spec, re_ds._element_spec)
+
+    @def_function.function(input_signature=[element_spec])
+    def process_inputs(inputs):
+      distribution.run(lambda inputs: inputs, args=(inputs,))
+
+    for x in ds:
+      process_inputs(x)
+
+
+class RaggedTensorDistributedIteratorTest(test.TestCase,
+                                          parameterized.TestCase):
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
+              strategy_combinations.multi_worker_mirrored_2x2_gpu,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpec(self, distribution, enable_get_next_as_optional):
+    ctx = distribute_lib.InputContext()
+    batch_size = ctx.get_per_replica_batch_size(8)
+    # Use 20 which isn't divisible by 8 to test partial batch behavior.
+    row_lengths = np.mod(np.arange(20), 4).astype(np.int64)
+    ragged_tensor = ragged_tensor_lib.RaggedTensor.from_row_lengths(
+        np.repeat(np.arange(20, dtype=np.float32), row_lengths), row_lengths)
+    dataset = dataset_ops.DatasetV2.from_tensor_slices({
+        "dense": ragged_tensor.to_tensor(),
+        "ragged": ragged_tensor,
+        "sparse": ragged_tensor.to_sparse(),
+    })
+    dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
+    dataset = dataset.batch(batch_size)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    dist_dataset = distribution.experimental_distribute_dataset(dataset)
+    with distribution.scope():
+      iterator = iter(dist_dataset)
+      _check_type_spec_structure(iterator)
+
+    spec = iterator._type_spec
+    self.assertEqual(spec._input_workers, iterator._input_workers)
+    self.assertEqual(
+        spec._element_spec, {
+            "sparse":
+                values.PerReplicaSpec(
+                    sparse_tensor.SparseTensorSpec(
+                        tensor_shape.TensorShape([None, 3]), dtypes.float32),
+                    sparse_tensor.SparseTensorSpec(
+                        tensor_shape.TensorShape([None, 3]), dtypes.float32)),
+            "dense":
+                values.PerReplicaSpec(
+                    tensor_spec.TensorSpec(
+                        shape=(None, 3), dtype=dtypes.float32, name=None),
+                    tensor_spec.TensorSpec(
+                        shape=(None, 3), dtype=dtypes.float32, name=None)),
+            "ragged":
+                values.PerReplicaSpec(
+                    ragged_tensor_lib.RaggedTensorSpec(
+                        tensor_shape.TensorShape([None, None]), dtypes.float32,
+                        1, dtypes.int64),
+                    ragged_tensor_lib.RaggedTensorSpec(
+                        tensor_shape.TensorShape([None, None]), dtypes.float32,
+                        1, dtypes.int64))
+        })
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+              strategy_combinations.central_storage_strategy_with_gpu_and_cpu,
+              strategy_combinations.multi_worker_mirrored_2x2_gpu,
+          ],
+
+          enable_get_next_as_optional=[True, False]))
+  def testTypeSpecRoundTrip(self, distribution, enable_get_next_as_optional):
+    ctx = distribute_lib.InputContext()
+    batch_size = ctx.get_per_replica_batch_size(8)
+    # Use 20 which isn't divisible by 8 to test partial batch behavior.
+    row_lengths = np.mod(np.arange(20), 4).astype(np.int64)
+    ragged_tensor = ragged_tensor_lib.RaggedTensor.from_row_lengths(
+        np.repeat(np.arange(20, dtype=np.float32), row_lengths), row_lengths)
+    dataset = dataset_ops.DatasetV2.from_tensor_slices({
+        "dense": ragged_tensor.to_tensor(),
+        "ragged": ragged_tensor,
+        "sparse": ragged_tensor.to_sparse(),
+    })
+    dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
+    dataset = dataset.batch(batch_size)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    if isinstance(distribution,
+                  (tpu_strategy.TPUStrategyV2, tpu_strategy.TPUStrategy)):
+      # TPUStrategy does not support distributed datasets with device prefetch
+      # when using sparse or ragged tensors.
+      options = distribute_lib.InputOptions(experimental_fetch_to_device=False)
+    else:
+      options = None
+
+    dist_dataset = distribution.experimental_distribute_dataset(
+        dataset, options)
+    with distribution.scope():
+      iterator = iter(dist_dataset)
+      _check_type_spec_structure(iterator)
+
+    spec = iterator._type_spec
+
+    tensor_list = spec._to_components(iterator)
+    re_iterator = spec._from_components(tensor_list)
+
+    self.assertEqual(iterator._input_workers, re_iterator._input_workers)
+    self.assertAllEqual(iterator._iterators, re_iterator._iterators)
+
+  @combinations.generate(
+      combinations.combine(
+          mode=["eager"],
+          tf_api_version=2,
+          distribution=[
+              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+              strategy_combinations.tpu_strategy,
+          ],
+          enable_get_next_as_optional=[True, False]))
+  def testDoesNotTriggerFunctionTracing(self, distribution,
+                                        enable_get_next_as_optional):
+    trace_count = [0]
+
+    @def_function.function
+    def f(iterator):
+      trace_count[0] += 1
+      counter = np.int64(0)
+      for _ in range(5):
+        next(iterator)
+        counter += 1
+      return counter
+
+    ctx = distribute_lib.InputContext()
+    batch_size = ctx.get_per_replica_batch_size(8)
+    # Use 20 which isn't divisible by 8 to test partial batch behavior.
+    row_lengths = np.mod(np.arange(50), 4).astype(np.int64)
+    ragged_tensor = ragged_tensor_lib.RaggedTensor.from_row_lengths(
+        np.repeat(np.arange(50, dtype=np.float32), row_lengths), row_lengths)
+    dataset = dataset_ops.DatasetV2.from_tensor_slices({
+        "dense": ragged_tensor.to_tensor(),
+        "ragged": ragged_tensor,
+        "sparse": ragged_tensor.to_sparse(),
+    })
+    dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
+    dataset = dataset.batch(batch_size)
+
+    distribution.extended.experimental_enable_get_next_as_optional = (
+        enable_get_next_as_optional)
+
+    if isinstance(distribution,
+                  (tpu_strategy.TPUStrategyV2, tpu_strategy.TPUStrategy)):
+      # TPUStrategy does not support distributed datasets with device prefetch
+      # when using sparse or ragged tensors.
+      options = distribute_lib.InputOptions(experimental_fetch_to_device=False)
+    else:
+      options = None
+
+    dist_dataset = distribution.experimental_distribute_dataset(
+        dataset, options)
+    with distribution.scope():
+      for _ in range(3):
+        iterator = iter(dist_dataset)
+        _check_type_spec_structure(iterator)
+        counter = f(iterator)
+
+        self.assertEqual(trace_count[0], 1)
+        self.assertEqual(counter, 5)
+
+
+def _check_type_spec_structure(x):
+  """Verifies that `x` has the same structure as its `TypeSpec`."""
+  if isinstance(x, composite_tensor.CompositeTensor):
+    nest.assert_same_structure(x, x._type_spec, expand_composites=True)
+
+
+def _create_text_file(fname, num_lines):
+  with open(fname, "w") as f:
+    for i in range(num_lines):
+      f.write("%d\n" % i)
+
+
+if __name__ == "__main__":
+  test_util.main()

tensorflow/python/distribute/input_ops.py

@@ -0,0 +1,107 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Input-pipeline utilities for Distribution strategies."""
+
+from tensorflow.python.data.experimental.ops import distribute
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops.options import AutoShardPolicy
+from tensorflow.python.data.util import traverse
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+from tensorflow.python.types import data as data_types
+from tensorflow.python.types import distribute as distribute_types
+
+
+# pylint: disable=protected-access
+def auto_shard_dataset(dataset, num_shards, index, num_replicas_in_sync=None):
+  """Shard the input pipeline by sharding the underlying list of files.
+
+  Args:
+    dataset: A `tf.data.Dataset` instance, typically the result of a bunch of
+      dataset transformations.
+    num_shards: A `tf.int64` scalar `tf.Tensor`, representing the number of
+        shards operating in parallel. Same usage as in `tf.data.Dataset.shard`.
+    index: A `tf.int64` scalar `tf.Tensor`, representing the worker index.
+      Same usage as in `tf.data.Dataset.shard`.
+    num_replicas_in_sync: An integer representing the total number of replicas
+      across all workers. This is used in the rewrite when sharding by data.
+
+  Returns:
+    A modified `Dataset` obtained by updating the pipeline sharded by the
+    files. The input dataset will be returned if we cannot automatically
+    determine a good way to shard the input dataset.
+  """
+  if isinstance(dataset, distribute_types.DistributedDatasetInterface):
+    return dataset.auto_shard(num_shards, index)
+  if (dataset.options().experimental_distribute.auto_shard_policy !=
+      AutoShardPolicy.OFF):
+    if num_replicas_in_sync is None:
+      num_replicas_in_sync = 1
+    if isinstance(dataset, data_types.DatasetV1):
+      return distribute._AutoShardDatasetV1(dataset, num_shards, index,
+                                            num_replicas_in_sync)
+    else:
+      return distribute._AutoShardDataset(dataset, num_shards, index,
+                                          num_replicas_in_sync)
+  else:
+    return dataset
+
+
+def _clone_dataset(dataset):
+  """Returns a cloned version of `dataset`."""
+  variant_tensor_ops = traverse.obtain_all_variant_tensor_ops(dataset)
+  remap_dict = _clone_helper(dataset._variant_tensor.op, variant_tensor_ops)
+  new_variant_tensor = remap_dict[dataset._variant_tensor.op].outputs[0]
+  return dataset_ops._VariantDataset(new_variant_tensor, dataset.element_spec)
+
+
+def _get_op_def(op):
+  return op.op_def or op_def_registry.get(op.type)
+
+
+def _clone_helper(op_to_clone, variant_tensor_ops):
+  """Helper method that recursively clones `op_to_clone`.
+
+  Args:
+    op_to_clone: The op we want to clone.
+    variant_tensor_ops: A list of ops that we have to clone along the way.
+
+  Returns:
+    A dictionary mapping old_ops to new_ops created. Includes op_to_clone
+    as a key.
+  """
+  remap_dict = {}
+  for input_tensor in op_to_clone.inputs:
+    input_tensor_op = input_tensor.op
+    if input_tensor_op in variant_tensor_ops:
+      recursive_map = _clone_helper(input_tensor_op, variant_tensor_ops)
+      remap_dict.update(recursive_map)
+  inputs_list = []
+  for input_tensor in op_to_clone.inputs:
+    input_tensor_op = input_tensor.op
+    if input_tensor_op in remap_dict:
+      remapped_input = remap_dict[input_tensor_op].outputs[0]
+      inputs_list.append(remapped_input)
+    else:
+      inputs_list.append(input_tensor_op.outputs[input_tensor.value_index])
+  g = ops.get_default_graph()
+  new_op = g.create_op(
+      op_to_clone.type,
+      inputs_list, [o.dtype for o in op_to_clone.outputs],
+      name=op_to_clone.name,
+      attrs=op_to_clone.node_def.attr,
+      op_def=_get_op_def(op_to_clone))
+  remap_dict[op_to_clone] = new_op
+  return remap_dict

tensorflow/python/distribute/input_ops_test.py

@@ -0,0 +1,326 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for input pipeline modifications for distribution strategies."""
+
+import os
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import readers
+from tensorflow.python.data.util import structure
+from tensorflow.python.distribute import input_ops
+from tensorflow.python.eager import context
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.ops import gen_dataset_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.util import compat
+
+
+class AutoShardDatasetTest(test.TestCase):
+
+  def setUp(self):
+    super(AutoShardDatasetTest, self).setUp()
+    self._num_files = 10
+    self._num_records = 4
+    self._num_shards = 2
+    self._shard_index = 0
+    self._record_bytes = 10
+
+  def _getNext(self, dataset):
+    if context.executing_eagerly():
+      iterator = iter(dataset)
+      return iterator._next_internal  # pylint: disable=protected-access
+    else:
+      iterator = dataset_ops.make_one_shot_iterator(dataset)
+      get_next = iterator.get_next()
+      return lambda: get_next
+
+  def _record(self, r, f):
+    return compat.as_bytes("Record %d of file %d" % (r, f))
+
+  def _text_line(self, r, f):
+    return compat.as_bytes("Text line %d of file %d" % (r, f))
+
+  def _fixed_length_record(self, r, f):
+    return compat.as_bytes(str((r * f) % 10) * self._record_bytes)
+
+  def _createTFRecordFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "tf_record.%d.txt" % i)
+      filenames.append(fn)
+      writer = python_io.TFRecordWriter(fn)
+      for j in range(self._num_records):
+        record = self._record(j, i)
+        writer.write(record)
+      writer.close()
+    return filenames
+
+  def _createTextFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "text_line.%d.txt" % i)
+      filenames.append(fn)
+      contents = []
+      for j in range(self._num_records):
+        contents.append(self._text_line(j, i))
+        if j + 1 != self._num_records or i == 0:
+          contents.append(b"\r\n")
+      contents = b"".join(contents)
+
+      with open(fn, "wb") as f:
+        f.write(contents)
+    return filenames
+
+  def _createFixedLengthRecordFiles(self):
+    filenames = []
+    for i in range(self._num_files):
+      fn = os.path.join(self.get_temp_dir(), "fixed_length_record.%d.txt" % i)
+      filenames.append(fn)
+      with open(fn, "wb") as f:
+        for j in range(self._num_records):
+          f.write(self._fixed_length_record(j, i))
+    return filenames
+
+  def _verifySimpleShardingOutput(self, dataset, record_fn):
+    next_element_fn = self._getNext(dataset)
+    with self.cached_session():
+      for f in range(self._shard_index, self._num_files, self._num_shards):
+        for r in range(self._num_records):
+          self.assertAllEqual(record_fn(r, f), self.evaluate(next_element_fn()))
+      with self.assertRaises(errors.OutOfRangeError):
+        self.evaluate(next_element_fn())
+
+  @test_util.run_in_graph_and_eager_modes
+  def testTFRecordDataset(self):
+    dataset = readers.TFRecordDataset(self._createTFRecordFiles())
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testFlatMap(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testInterleave(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.interleave(
+        readers.TFRecordDataset, cycle_length=4, block_length=self._num_records)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    # Since block_length == num records in each file, the output will still
+    # contain records in order of files.
+    self._verifySimpleShardingOutput(dataset, self._record)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testListfiles(self):
+    filenames = self._createTFRecordFiles()
+    file_pattern = filenames[0].rsplit(os.sep, 1)[0] + "/tf_record.*.txt"
+    dataset = dataset_ops.Dataset.list_files(file_pattern, shuffle=False)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    next_element_fn = self._getNext(dataset)
+    actual, expected = [], []
+    for f in range(self._shard_index, self._num_files, self._num_shards):
+      for r in range(self._num_records):
+        actual.append(self.evaluate(next_element_fn()))
+        expected.append(self._record(r, f))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element_fn())
+    self.assertAllEqual(expected, actual)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testComplexPipeline(self):
+    # Setup a complex input pipeline.
+    batch_size = 2
+    num_epochs = 5
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        self._createTFRecordFiles())
+    dataset = dataset.shuffle(buffer_size=self._num_files)
+    dataset = dataset.flat_map(readers.TFRecordDataset)
+    dataset = dataset.prefetch(buffer_size=batch_size)
+    dataset = dataset.shuffle(2 * self._num_files * self._num_records)
+    dataset = dataset.repeat(num_epochs)
+    dataset = dataset.map(lambda x: x)
+    dataset = dataset.batch(batch_size)
+    dataset = dataset.prefetch(buffer_size=None)
+
+    # Auto shard.
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    # Verify output.
+    next_element_fn = self._getNext(dataset)
+    actual = []
+    num_iterations = (self._num_files * self._num_records * num_epochs) // (
+        self._num_shards * batch_size)
+    for _ in range(num_iterations):
+      actual.extend(self.evaluate(next_element_fn()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element_fn())
+
+    expected = []
+    for f in range(0, self._num_files, self._num_shards):
+      for r in range(self._num_records):
+        expected.append(self._record(r, f))
+    expected *= num_epochs
+
+    self.assertAllEqual(sorted(expected), sorted(actual))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testZip(self):
+    dataset1 = readers.TFRecordDataset(self._createTFRecordFiles())
+    dataset2 = readers.TextLineDataset(self._createTextFiles())
+
+    dataset = dataset_ops.Dataset.zip((dataset1, dataset2))
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    record_fn = lambda r, f: (self._record(r, f), self._text_line(r, f))
+    self._verifySimpleShardingOutput(dataset, record_fn)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testConcat(self):
+    dataset1 = readers.TFRecordDataset(self._createTFRecordFiles())
+    dataset2 = readers.TextLineDataset(self._createTextFiles())
+
+    dataset = dataset1.concatenate(dataset2)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    next_element_fn = self._getNext(dataset)
+    for f in range(self._shard_index, self._num_files, self._num_shards):
+      for r in range(self._num_records):
+        self.assertAllEqual(
+            self._record(r, f), self.evaluate(next_element_fn()))
+    for f in range(self._shard_index, self._num_files, self._num_shards):
+      for r in range(self._num_records):
+        self.assertAllEqual(
+            self._text_line(r, f), self.evaluate(next_element_fn()))
+    with self.assertRaises(errors.OutOfRangeError):
+      self.evaluate(next_element_fn())
+
+  @test_util.run_in_graph_and_eager_modes
+  def testTextLineReader(self):
+    dataset = readers.TextLineDataset(self._createTextFiles())
+
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._text_line)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testTextLineReaderWithFlatMap(self):
+    dataset = readers.TextLineDataset(self._createTextFiles())
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._text_line)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testFixedLengthReaderWithFlatMap(self):
+    dataset = readers.FixedLengthRecordDataset(
+        self._createFixedLengthRecordFiles(), self._record_bytes)
+    dataset = input_ops.auto_shard_dataset(
+        dataset, self._num_shards, self._shard_index)
+
+    self._verifySimpleShardingOutput(dataset, self._fixed_length_record)
+
+
+# A dataset that creates two variant tensors.
+class _TestDataset(dataset_ops.UnaryUnchangedStructureDataset):
+
+  def __init__(self, input_dataset):
+    self._input_dataset = input_dataset
+    temp_variant_tensor = gen_dataset_ops.prefetch_dataset(
+        input_dataset._variant_tensor,
+        buffer_size=1,
+        **self._flat_structure)
+    variant_tensor = gen_dataset_ops.model_dataset(
+        temp_variant_tensor, **self._flat_structure)
+    super(_TestDataset, self).__init__(input_dataset, variant_tensor)
+
+
+class CloneDatasetTest(test.TestCase):
+
+  def _assert_datasets_equal(self, ds1, ds2):
+    # First lets assert the structure is the same.
+    self.assertTrue(
+        structure.are_compatible(ds1.element_spec, ds2.element_spec))
+
+    # Now create iterators on both and assert they produce the same values.
+    it1 = dataset_ops.make_initializable_iterator(ds1)
+    it2 = dataset_ops.make_initializable_iterator(ds2)
+
+    get_next1 = it1.get_next()
+    get_next2 = it2.get_next()
+
+    with self.cached_session():
+      self.evaluate([it1.initializer, it2.initializer])
+      val1, val2 = self.evaluate([get_next1, get_next2])
+      self.assertEqual(val1, val2)
+
+  @test_util.run_deprecated_v1
+  def testOnlySource(self):
+    ds = dataset_ops.Dataset.range(10)
+    cloned_ds = input_ops._clone_dataset(ds)
+    self._assert_datasets_equal(ds, cloned_ds)
+
+  @test_util.run_deprecated_v1
+  def testSimplePipeline(self):
+    ds = dataset_ops.Dataset.range(10).map(math_ops.square)
+    cloned_ds = input_ops._clone_dataset(ds)
+    self._assert_datasets_equal(ds, cloned_ds)
+
+  @test_util.run_deprecated_v1
+  def testConcat(self):
+    ds1 = dataset_ops.Dataset.range(10)
+    ds2 = dataset_ops.Dataset.range(10)
+    ds = ds1.concatenate(ds2)
+    cloned_ds = input_ops._clone_dataset(ds)
+    self._assert_datasets_equal(ds, cloned_ds)
+
+  @test_util.run_deprecated_v1
+  def testZip(self):
+    ds1 = dataset_ops.Dataset.range(10)
+    ds2 = dataset_ops.Dataset.range(10)
+    ds = dataset_ops.Dataset.zip((ds1, ds2))
+    cloned_ds = input_ops._clone_dataset(ds)
+    self._assert_datasets_equal(ds, cloned_ds)
+
+  @test_util.run_deprecated_v1
+  def testMultipleVariantTensors(self):
+    ds = dataset_ops.Dataset.range(10)
+    ds = _TestDataset(ds)
+    cloned_ds = input_ops._clone_dataset(ds)
+    self._assert_datasets_equal(ds, cloned_ds)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/input_util.py

@@ -0,0 +1,155 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Utils to create distributed datasets based on TF version."""
+
+from tensorflow.python import tf2
+from tensorflow.python.distribute import input_lib
+from tensorflow.python.distribute.v1 import input_lib as input_lib_v1
+
+
+def get_distributed_dataset(
+    dataset,
+    input_workers,
+    strategy,
+    num_replicas_in_sync=None,
+    input_context=None,
+    options=None,
+    build=True,
+    replica_order=None,
+):
+  """Returns a distributed dataset from the given tf.data.Dataset instance.
+
+  This is a common function that is used by all strategies to return a
+  distributed dataset. The distributed dataset instance returned is different
+  depending on if we are in a TF 1 or TF 2 context. The distributed dataset
+  instances returned differ from each other in the APIs supported by each of
+  them.
+
+  Args:
+    dataset: a tf.data.Dataset instance.
+    input_workers: an InputWorkers object which specifies devices on which
+      iterators should be created.
+    strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
+      handle last partial batch.
+    num_replicas_in_sync: Optional integer. If this is not None, the value is
+      used to decide how to rebatch datasets into smaller batches so that the
+      total batch size for each step (across all workers and replicas) adds up
+      to `dataset`'s batch size.
+    input_context: `InputContext` for sharding. Only pass this in for between
+      graph multi-worker cases where there is only one `input_worker`. In these
+      cases, we will shard based on the `input_pipeline_id` and
+      `num_input_pipelines` in the `InputContext`.
+    options: Default is None. `tf.distribute.InputOptions` used to control
+      options on how this dataset is distributed.
+    build: whether to build underlying datasets when a DistributedDataset is
+      created. This is only useful for `ParameterServerStrategy` now.
+    replica_order: the order of the replicas, which will be used to reorder the
+      iterators to match the device order.
+
+  Returns:
+    A distributed dataset instance.
+  """
+  if tf2.enabled():
+    return input_lib.DistributedDataset(
+        input_workers,
+        strategy,
+        dataset,
+        num_replicas_in_sync=num_replicas_in_sync,
+        input_context=input_context,
+        build=build,
+        options=options,
+        replica_order=replica_order,
+    )
+  else:
+    return input_lib_v1.DistributedDatasetV1(
+        dataset,
+        input_workers,
+        strategy,
+        num_replicas_in_sync=num_replicas_in_sync,
+        input_context=input_context,
+        options=options)
+
+
+def get_distributed_datasets_from_function(
+    dataset_fn,
+    input_workers,
+    input_contexts,
+    strategy,
+    options=None,
+    build=True,
+    replica_order=None,
+):
+  """Returns a distributed dataset from the given input function.
+
+  This is a common function that is used by all strategies to return a
+  distributed dataset. The distributed dataset instance returned is different
+  depending on if we are in a TF 1 or TF 2 context. The distributed dataset
+  instances returned differ from each other in the APIs supported by each of
+  them.
+
+  Args:
+    dataset_fn: a function that returns a tf.data.Dataset instance.
+    input_workers: an InputWorkers object which specifies devices on which
+      iterators should be created.
+    input_contexts: A list of `InputContext` instances to be passed to call(s)
+      to `dataset_fn`. Length and order should match worker order in
+      `worker_device_pairs`.
+    strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
+      handle last partial batch.
+    options: Default is None. `tf.distribute.InputOptions` used to control
+      options on how this dataset is distributed.
+    build: whether to build underlying datasets when a
+      `DistributedDatasetFromFunction` is created. This is only useful for
+      `ParameterServerStrategy` now.
+    replica_order: the order of the replicas, which will be used to reorder the
+      iterators to match the device order.
+
+  Returns:
+    A distributed dataset instance.
+
+  Raises:
+    ValueError: if `options.experimental_replication_mode` and
+    `options.experimental_place_dataset_on_device` are not consistent
+  """
+  if (options is not None and options.experimental_replication_mode !=
+      input_lib.InputReplicationMode.PER_REPLICA and
+      options.experimental_place_dataset_on_device):
+    raise ValueError(
+        "When `experimental_place_dataset_on_device` is set for dataset "
+        "placement, you must also specify `PER_REPLICA` for the "
+        "replication mode")
+
+  if (options is not None and options.experimental_replication_mode
+      == input_lib.InputReplicationMode.PER_REPLICA and
+      options.experimental_fetch_to_device and
+      options.experimental_place_dataset_on_device):
+    raise ValueError(
+        "`experimental_place_dataset_on_device` can not be set to True "
+        "when experimental_fetch_to_device is True and "
+        "replication mode is set to `PER_REPLICA`")
+
+  if tf2.enabled():
+    return input_lib.DistributedDatasetsFromFunction(
+        input_workers,
+        strategy,
+        input_contexts=input_contexts,
+        dataset_fn=dataset_fn,
+        options=options,
+        build=build,
+        replica_order=replica_order,
+    )
+  else:
+    return input_lib_v1.DistributedDatasetsFromFunctionV1(
+        input_workers, strategy, input_contexts, dataset_fn, options)

tensorflow/python/distribute/integration_test/BUILD

@@ -0,0 +1,90 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+load("//tensorflow/core/platform:distribute.bzl", "distribute_py_strict_test")
+load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+distribute_py_strict_test(
+    name = "saved_model_test",
+    srcs = ["saved_model_test.py"],
+    exec_properties = select({
+        # copybara:uncomment_begin(google-only)
+        # "@bazel_tools//tools/cpp:asan_build": {"mem": "16g"},
+        # copybara:uncomment_end
+        "//conditions:default": None,
+    }),
+    tags = [
+        "no_windows",  # TODO(b/171350360)
+        "nomultivm",  # multi_worker_test_base incompatible with multivm base
+    ],
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/distribute:combinations",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:parameter_server_strategy_v2",
+        "//tensorflow/python/distribute:sharded_variable",
+        "//tensorflow/python/distribute:strategy_combinations",
+        "//tensorflow/python/distribute:test_util",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/ops:lookup_ops",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "mwms_peer_failure_test",
+    size = "medium",
+    srcs = ["mwms_peer_failure_test.py"],
+    shard_count = 2,
+    tags = [
+        "multi_and_single_gpu",
+        "no_oss",  # TODO(b/227372713)
+        "notsan",  # b/195248428
+    ],
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/distribute:collective_all_reduce_strategy",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:test_util",
+        "//tensorflow/python/eager:test",
+    ],
+)
+
+py_library(
+    name = "mwms_peer_failure_test_lib",
+    srcs = ["mwms_peer_failure_test.py"],
+    strict_deps = True,
+    visibility = ["//learning/brain/runtime/python:__pkg__"],
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/distribute:collective_all_reduce_strategy",
+        "//tensorflow/python/distribute:multi_process_runner",
+        "//tensorflow/python/distribute:multi_worker_test_base",
+        "//tensorflow/python/distribute:test_util",
+        "//tensorflow/python/eager:test",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_memory_test",
+    size = "medium",
+    srcs = ["tpu_memory_test.py"],
+    disable_experimental = True,
+    disable_mlir_bridge = True,
+    disable_tfrt = False,
+    disable_v2 = True,
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/platform:flags",
+    ],
+)

tensorflow/python/distribute/integration_test/mwms_peer_failure_test.py

@@ -0,0 +1,229 @@
+# 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.
+# ==============================================================================
+"""This file contains tests that simulate peer failures.
+
+When a peer fails during MultiWorkerMirroredStrategy training. All workers
+should get Unavailable error.
+"""
+
+import os
+
+import tensorflow as tf
+
+from tensorflow.python.distribute import collective_all_reduce_strategy as mwms_lib
+from tensorflow.python.distribute import multi_process_runner
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import test_util
+from tensorflow.python.eager import test
+
+RPC_PROTOCOL = "grpc"
+
+# Put it in top level so it executes in the child processes as well.
+mwms_lib.CollectiveAllReduceExtended._enable_check_health = True
+mwms_lib.CollectiveAllReduceExtended._check_health_interval = 3
+mwms_lib.CollectiveAllReduceExtended._check_health_initial_timeout = 0
+# This is needed for OSS, which issues all RPCs with fail_fast=false by default.
+mwms_lib.CollectiveAllReduceExtended._check_health_timeout = 1
+
+
+def get_attempt(strategy, attempts):
+  task_type = strategy.cluster_resolver.task_type
+  task_id = strategy.cluster_resolver.task_id
+  attempts[(task_type, task_id)] = attempts.get((task_type, task_id), 0) + 1
+  return task_id, attempts[(task_type, task_id)]
+
+
+quick_exit = os._exit  # pylint: disable=protected-access
+
+
+class PeerFailureTest(test.TestCase):
+  # Note that all the tests use auto_restart=True. Currently we rely on the
+  # assumption that an external system restarts failed tasks. If the assumption
+  # is not true, the remaining tasks may still hang instead of fail.
+  #
+  # In these tests we leverage the auto restart feature of MultiProcessRunner.
+  # Failed workers are restarted automatically. In reality there needs to be
+  # some job management system that does the restart, e.g. Kubernetes.
+  #
+  # Worker failures may cause problems if there're more than one collective, and
+  # the failure happens after the first collective. In this case the recovered
+  # worker will be running a different collective with the rest, which causes a
+  # deadlock. Note that collectives are common, e.g. when creating variables the
+  # initial values are broadcasted from the first worker.
+  #
+  # We use a multiprocessing.Manager().dict() object to track the attempts of
+  # each worker. We take different actions in different attempts to simuate the
+  # events in real world. E.g. some tests make a worker fail on the first
+  # attempt only, and asserts that it should recovery.
+
+  def test_creating_variable(self):
+    # This test simulates the case when a worker fails before or during creating
+    # a variable. Creating variables involve broadcasting the initial value from
+    # the first replica to all replicas.
+
+    def worker_fn():
+      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+      with strategy.scope():
+        tf.Variable(1.)
+        # worker-1 dies here.
+        if strategy.cluster_resolver.task_id == 1:
+          quick_exit(1)
+        v = tf.Variable(tf.random.uniform(()))
+        return v.read_value().numpy()
+
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    mpr = multi_process_runner.MultiProcessRunner(
+        worker_fn, cluster_spec, rpc_layer=RPC_PROTOCOL)
+    mpr.start()
+    # TODO(b/151232436): Always raise UnavailableError when a peer fails.
+    with self.assertRaises(
+        (tf.errors.UnavailableError, tf.errors.DeadlineExceededError)):
+      mpr.join(timeout=60)
+
+  def test_reduce_small_tensor(self):
+    # This test simulates the case when a worker fails before or during reducing
+    # a small tensors, e.g. reading a metric.
+    #
+    # Note that this is written for a specific corner case that used to happen
+    # only when all of the following conditions are met:
+    #   - There're two workers.
+    #   - They're reducing a small tensor. The definition of small varies
+    #     per platform.
+    #   - They're reducing a single tensor. Batched all-reduce are not affected.
+    #   - It must be worker-1 that fails.
+    # Under this case, the all-reduce is effectively two send/recv operation,
+    # the first one from worker-0 to worker-1, and the second one vice versa.
+    # The first one blocks the second one. In send/recv, the sending party is
+    # not aware of the failures of the receiving party.
+
+    def worker_fn():
+      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+      value = tf.identity([1.])
+      strategy.reduce("sum", value, axis=None)
+      # worker-1 dies here.
+      if strategy.cluster_resolver.task_id == 1:
+        quick_exit(1)
+      strategy.reduce("sum", value, axis=None)
+
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    mpr = multi_process_runner.MultiProcessRunner(
+        worker_fn, cluster_spec, rpc_layer=RPC_PROTOCOL)
+    mpr.start()
+    # TODO(b/151232436): Always raise UnavailableError when a peer fails.
+    with self.assertRaises(
+        (tf.errors.UnavailableError, tf.errors.DeadlineExceededError)):
+      mpr.join(timeout=60)
+
+
+class PeerFailureRecoverTest(test.TestCase):
+  # Similar to PeerFailureTest but simulates the situation where there's some
+  # external system that automatically restarts failed workers.
+
+  def test_creating_variable(self):
+    # See PeerFailureTest.test_creating_variable
+
+    def worker_fn(attempts):
+      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+      task_id, attempt = get_attempt(strategy, attempts)
+      with strategy.scope():
+        tf.Variable(1.)
+        # worker-1 dies here.
+        if attempt == 1 and task_id == 1:
+          quick_exit(1)
+        v = tf.Variable(tf.random.uniform(()))
+        return v.read_value().numpy()
+
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    attempts = multi_process_runner.manager().dict()
+    mpr = multi_process_runner.MultiProcessRunner(
+        worker_fn,
+        cluster_spec,
+        rpc_layer=RPC_PROTOCOL,
+        args=(attempts,),
+        auto_restart=True)
+    mpr.start()
+    results = mpr.join(timeout=90).return_value
+    self.assertEqual(results[0], results[1])
+
+  def test_reduce_small_tensor(self):
+    # See PeerFailureTest.test_reduce_small_tensor
+
+    def worker_fn(attempts):
+      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+      task_id, attempt = get_attempt(strategy, attempts)
+      value = tf.identity([1.])
+      strategy.reduce("sum", value, axis=None)
+      # worker-1 dies here.
+      if attempt == 1 and task_id == 1:
+        quick_exit(1)
+      return strategy.reduce("sum", value, axis=None).numpy()
+
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    attempts = multi_process_runner.manager().dict()
+    mpr = multi_process_runner.MultiProcessRunner(
+        worker_fn,
+        cluster_spec,
+        rpc_layer=RPC_PROTOCOL,
+        args=(attempts,),
+        auto_restart=True)
+    mpr.start()
+    results = mpr.join(timeout=90).return_value
+    self.assertAllEqual(results, [[2.], [2.]])
+
+  def test_quick_recover(self):
+    # This test simulates the case when a worker fails but recovers quickly
+    # before the next collective.
+    #
+    # It's not guaranteed that the cluster only restarts once when one worker
+    # fails. The external job management system is expected to keep restarting
+    # failed workers.
+
+    def worker_fn(attempts):
+      # Set a long check alive interval to better simulate the case when a
+      # worker fails and recovers during a check alive interval.
+      mwms_lib.CollectiveAllReduceExtended._check_alive_interval = 30
+      mwms_lib.CollectiveAllReduceExtended._check_alive_initial_timeout = 30
+
+      strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
+      task_id, attempt = get_attempt(strategy, attempts)
+
+      @tf.function
+      def replica_fn():
+        ctx = tf.distribute.get_replica_context()
+        # Use a large tensor because small tensor may hang regardless when the
+        # worker recovers.
+        value = tf.ones((64, 64))
+        ctx.all_reduce(tf.distribute.ReduceOp.SUM, [value, value])
+
+      strategy.run(replica_fn)
+      # worker-1 dies here.
+      if attempt == 1 and task_id == 1:
+        quick_exit(1)
+      strategy.run(replica_fn)
+
+    cluster_spec = multi_worker_test_base.create_cluster_spec(num_workers=2)
+    attempts = multi_process_runner.manager().dict()
+    mpr = multi_process_runner.MultiProcessRunner(
+        worker_fn,
+        cluster_spec,
+        rpc_layer=RPC_PROTOCOL,
+        args=(attempts,),
+        auto_restart=True)
+    mpr.start()
+    mpr.join(timeout=90)
+
+
+if __name__ == "__main__":
+  test_util.main()

tensorflow/python/distribute/integration_test/saved_model_test.py

@@ -0,0 +1,698 @@
+# 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.
+# ==============================================================================
+"""This file contains unit tests of past, existing and potential saving issues with tf.distribute.
+
+The tests are written as minimum reproductions in the hope to demonstrate the
+expected behavior in a straightforward fashion.
+
+Test cases ending with _broken are known issues. Assertions in such tests
+describes the current incorrect behavior. If you fix something, you should
+expect some test cases to fail and please update them.
+
+This file is not intended to provide exhaustive test coverage. Exhaustive tests
+using Keras models are in keras*_test.py
+"""
+
+import os
+from absl.testing import parameterized
+
+import tensorflow as tf
+import tensorflow.compat.v1 as tf1
+
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import multi_worker_test_base
+from tensorflow.python.distribute import parameter_server_strategy_v2
+from tensorflow.python.distribute import sharded_variable
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import test_util
+from tensorflow.python.distribute import values
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.ops import lookup_ops
+
+_sixteen_worker_pool = strategy_combinations._deferred_pool_runner(
+    has_chief=True,
+    num_workers=8,
+    initializer=strategy_combinations._get_multi_worker_mirrored_creator(
+        required_gpus=0))
+
+
+@combinations.generate(
+    combinations.combine(
+        strategy=[
+            combinations.NamedDistribution(
+                "MultiWorkerMirrored8x1CPU",
+                strategy_combinations._get_multi_worker_mirrored_creator(
+                    required_gpus=0),
+                has_chief=True,
+                num_workers=8,
+                pool_runner_fn=_sixteen_worker_pool,
+                no_xla=True,
+            ),
+        ],
+        mode=["eager"]))
+class SaveModelForMultipleWorkers(test.TestCase, parameterized.TestCase):
+
+  def test_read_sync_on_read_variable(self, strategy):
+    # TODO(b/178943315): Enable test when the design in b/17894331 is
+    # implemented.
+    self.skipTest(
+        "This test fails today due to issue in multiple workers trying to write"
+        " to same file location: b/178943315")
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_READ,
+            aggregation=tf.VariableAggregation.SUM)
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    export_dir = os.path.join(self._get_tempdir_path_test(),
+                              "test-file-failure")
+    with strategy.scope():
+      m = Model()
+      m.v.assign(1.)
+      # This fails when multiple workers try to write to the same file location.
+      # b/178943315 for tracking this bug.
+      tf.saved_model.save(m, export_dir)
+
+
+@combinations.generate(
+    combinations.combine(
+        strategy=[
+            strategy_combinations.mirrored_strategy_with_two_cpus,
+            strategy_combinations.mirrored_strategy_with_two_gpus,
+            strategy_combinations.tpu_strategy,
+        ],
+        mode="eager",
+    ))
+class SaveAndLoadForServingTest(test.TestCase, parameterized.TestCase):
+  # These test cases cover the case when a model is trained under
+  # tf.distribute.Strategy and used for serving later. Serving usually only uses
+  # one device and this is simulated by loading the model under no strategy
+  # using tf.saved_model.load. Note that tf.keras.models.load_model doesn't use
+  # the saved functions so they're not fit for this test.
+  #
+  # When saving, it's expected that the model is saved as if there's no
+  # tf.distribute.Strategy.  The saved tf.function should be an inference
+  # function on a single device, and the distributed variables are saved as
+  # single variables.
+  #
+  # Currently references to components of a distributed variable are mapped to
+  # the single variable that is saved. This means that if the saved tf.functions
+  # access components of a distributed variable, for example if it triggers
+  # variable aggregation, the outputs are likely incorrect.
+  #
+  # Note that distributed variables have different behavior in the replica
+  # context and the cross-replica context. Saving happens in the cross replica
+  # context or the default strategy's replica context.
+
+  def test_read_sync_on_read_variable(self, strategy):
+    # synchronizaiton=ON_READ variables are typically used in Keras metrics and
+    # batch norm layers.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_READ,
+            aggregation=tf.VariableAggregation.SUM)
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    export_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = Model()
+      # Note that each component is assigned with the value divided by the
+      # number of replicas.
+      m.v.assign(1.)
+      self.assertAllEqual(
+          self.evaluate(strategy.experimental_local_results(m.v)), [0.5, 0.5])
+      tf.saved_model.save(m, export_dir)
+
+    loaded = tf.saved_model.load(export_dir)
+    # The variable already has the aggregated value.
+    self.assertEqual(self.evaluate(loaded.v.read_value()), 1.)
+    self.assertEqual(self.evaluate(loaded()), 1.)
+
+  def test_read_mirrored_variable(self, strategy):
+    # synchronizaiton=ON_WRITE is the default variable created under
+    # tf.distribute.Strategy.scope(). Most model parameters are this kind of
+    # variable. Reading a synchronization=ON_WRITE simply reads the primary
+    # component, so it works as intended.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0., synchronization=tf.VariableSynchronization.ON_WRITE)
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    export_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = Model()
+      m.v.assign(1.)
+      tf.saved_model.save(m, export_dir)
+
+    loaded = tf.saved_model.load(export_dir)
+    self.assertEqual(self.evaluate(loaded()), 1.)
+
+  def test_update_sync_on_read_variable(self, strategy):
+    # It's rare to update aggregation=ON_READ variables in serving, but it's
+    # possible that the SavedModel contains both serving and training graphs,
+    # and the training may contain metrics layers.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_READ,
+            aggregation=tf.VariableAggregation.SUM)
+
+      @tf.function(input_signature=[])
+      def update(self):
+        self.v.assign_add(1.)
+        return self.v.read_value()
+
+    export_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = Model()
+      tf.saved_model.save(m, export_dir)
+
+    loaded = tf.saved_model.load(export_dir)
+    loaded.update()
+    self.assertEqual(self.evaluate(loaded.v), 1.)
+
+  def test_update_mirrored_variable(self, strategy):
+    # It's very rare to update aggregation=ON_WRITE variables in the forward
+    # path, and this test case is mainly for completeness.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_WRITE,
+            aggregation=tf.VariableAggregation.MEAN)
+
+      @tf.function(input_signature=[])
+      def update(self):
+        self.v.assign_add(1.)
+
+    export_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = Model()
+      tf.saved_model.save(m, export_dir)
+
+    loaded = tf.saved_model.load(export_dir)
+    self.assertEqual(self.evaluate(loaded.v), 0.)
+    loaded.update()
+    self.assertEqual(self.evaluate(loaded.v), 1.)
+
+  def test_training_only_device(self, strategy):
+    # tf.distribute APIs may enter device scopes, but the saved model should not
+    # contain device annotations, since devices during training may not be
+    # available when the saved model is used.
+    #
+    # Models trained with MultiWorkerMirroredStrategy is affected the most,
+    # since under MultiWorkerMirroredStrategy the device annotations contain job
+    # and task, which causes error if that job or task is not available even
+    # with soft placement enabled.
+
+    class Model(tf.Module):
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return tf.identity(1.)
+
+    export_dir = self.get_temp_dir()
+    with strategy.scope(), tf.device("GPU:0"):
+      m = Model()
+      tf.saved_model.save(m, export_dir)
+
+    export_dir = self.get_temp_dir()
+    loaded = tf.saved_model.load(export_dir)
+    graph = loaded.signatures["serving_default"].graph
+    for op in graph.get_operations():
+      self.assertEqual(op.device, "")
+    self.assertEqual(loaded().numpy(), 1.)
+
+  def test_model_with_loaded_layer(self, strategy):
+    # When a model is loaded under strategy, we wrap it so that when it's passed
+    # to strategy.run(), the captured variables resolve to the ones of the
+    # current replica. Since the saved tf.function may contain updates to the
+    # variables, we don't allow using the model outside of strategy.run().
+    #
+    # That is to say, a loaded model is different from the original Python one.
+    # We need to test save-load-save-load to make sure things work correctly.
+
+    class Layer(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(1.)
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    class Model(tf.Module):
+
+      def __init__(self, layer):
+        self.layer = layer
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.layer()
+
+    layer_export_dir = self.get_temp_dir()
+    tf.saved_model.save(Layer(), layer_export_dir)
+
+    with strategy.scope():
+      m = Model(tf.saved_model.load(layer_export_dir))
+      export_dir = self.get_temp_dir()
+      # Saving a ConcreteFunction should raise an error.
+      with self.assertRaisesRegex(
+          ValueError, "saving a tf.function with input_signature instead"):
+        tf.saved_model.save(
+            m,
+            export_dir,
+            signatures={
+                "call": m.__call__.get_concrete_function(),
+            })
+      tf.saved_model.save(m, export_dir)
+
+    loaded = tf.saved_model.load(export_dir)
+    self.assertAllEqual(
+        self.evaluate(
+            strategy.experimental_local_results(strategy.run(loaded))),
+        [1., 1.])
+
+
+@combinations.generate(
+    combinations.combine(
+        strategy=[
+            strategy_combinations.mirrored_strategy_with_two_cpus,
+            strategy_combinations.mirrored_strategy_with_two_gpus,
+        ],
+        mode="eager",
+    ))
+class SaveAndLoadForTrainingTest(test.TestCase, parameterized.TestCase):
+  # These test cases cover the case when the user loads a model and continues to
+  # train it. The model could originally be trained with or without
+  # tf.distribute.Strategy.
+  #
+  # tf.distribute does not distinguish whether the model is saved for inference
+  # or for training. The implications are that all issues with serving are
+  # issues with training as well, possibly with different symptoms.
+  #
+  # Note that for Keras models, loading them with tf.keras.models.load_model()
+  # can workaround most issues since Keras loader restructs the layers with
+  # saved configs if possible, in which case the saved graph is not used.
+
+  def test_read_sync_on_read_variable(self, strategy):
+    # Reading a synchronizaiton=ON_READ in the replica context should just read
+    # the local value. Reading it in the cross replica context aggregates the
+    # value from all replicas. Both are true with a loaded model.
+    #
+    # Note that if aggregation=SUM, the value of each replica is the saved value
+    # divided by the number of replicas. In this way if you load a model and
+    # save it again, the values of the variables don't change.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_READ,
+            aggregation=tf.VariableAggregation.SUM)
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    export_dir = self.get_temp_dir()
+    value = strategy.experimental_distribute_values_from_function(
+        lambda ctx: tf.identity([3., 7.][ctx.replica_id_in_sync_group]))
+    with strategy.scope():
+      m = Model()
+      strategy.run(m.v.assign, args=(value,))
+      self.assertAllEqual(
+          self.evaluate(strategy.experimental_local_results(m.v)), [3., 7.])
+      self.assertEqual(self.evaluate(m.v.read_value()), 10.)
+      tf.saved_model.save(m, export_dir)
+      del m
+
+    with strategy.scope():
+      loaded = tf.saved_model.load(export_dir)
+    # It's intended that we don't save the each replica, but just the aggregated
+    # value.
+    self.assertAllEqual(
+        self.evaluate(
+            strategy.experimental_local_results(strategy.run(loaded))),
+        [5., 5.])
+    self.assertEqual(self.evaluate(loaded.v.read_value()), 10.)
+
+    # save and load again.
+    export_dir2 = self.get_temp_dir()
+    tf.saved_model.save(loaded, export_dir2)
+    # loaded.v.read_value() is still 1., both with and without strategy.
+    loaded = tf.saved_model.load(export_dir2)
+    self.assertEqual(self.evaluate(loaded.v.read_value()), 10.)
+    with strategy.scope():
+      loaded = tf.saved_model.load(export_dir2)
+      self.assertEqual(self.evaluate(loaded.v.read_value()), 10.)
+
+  def test_update_sync_on_read_variable(self, strategy):
+    # Updating a synchronizaiton=ON_READ in the replica context should just
+    # update the local value. Updating it in the cross replica context updates
+    # each component of the variable. Both are true with a loaded model.
+    #
+    # Note that if assigning a variable whose aggregation=SUM in the cross
+    # replica context, each replica is assigned with the value divided by the
+    # number of replicas.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_READ,
+            aggregation=tf.VariableAggregation.SUM)
+
+      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
+      def update(self, value):
+        self.v.assign_add(value)
+
+    export_dir = self.get_temp_dir()
+    # TODO(b/157621013): strategy.run doesn't work with tf.function with
+    # input_signature.
+    # value = strategy.experimental_distribute_values_from_function(
+    #     lambda ctx: tf.identity([3., 7.][ctx.replica_id_in_sync_group]))
+    with strategy.scope():
+      m = Model()
+      tf.saved_model.save(m, export_dir)
+      self.evaluate(m.v.assign(10.))
+      self.assertAllEqual(
+          self.evaluate(strategy.experimental_local_results(m.v)), [5., 5.])
+      del m
+      # TODO(b/157621013): strategy.run doesn't work with tf.function with
+      # input_signature.
+      # self.evaluate(strategy.run(m.update, args=(value,)))
+      # self.assertAllEqual(
+      #     self.evaluate(strategy.experimental_local_results(m.v)), [8., 12.])
+
+    with strategy.scope():
+      loaded = tf.saved_model.load(export_dir)
+      self.evaluate(loaded.v.assign(10.))
+      self.assertAllEqual(
+          self.evaluate(strategy.experimental_local_results(loaded.v)),
+          [5., 5.])
+      # TODO(b/157621013): strategy.run doesn't work with tf.function with
+      # input_signature.
+      # self.evaluate(strategy.run(loaded.update, args=(value,)))
+      # self.assertAllEqual(
+      #     self.evaluate(strategy.experimental_local_results(loaded.v)),
+      #     [8., 12.])
+
+  def test_read_mirrored_variable(self, strategy):
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0., synchronization=tf.VariableSynchronization.ON_WRITE)
+
+      @tf.function(input_signature=[])
+      def __call__(self):
+        return self.v.read_value()
+
+    export_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = Model()
+      m.v.assign(1.)
+      tf.saved_model.save(m, export_dir)
+
+    with strategy.scope():
+      loaded = tf.saved_model.load(export_dir)
+    self.assertAllEqual(
+        self.evaluate(
+            strategy.experimental_local_results(strategy.run(loaded))),
+        [1., 1.])
+
+  def test_update_mirrored_variable(self, strategy):
+    # This is also uncommon since most model parameters should be updated by
+    # optimizer, and this test case is for completeness.
+    #
+    # In the cross replica context, assigning to the variable assigns the same
+    # value to all replicas. This is true with the loaded model as well.
+    #
+    # However in replica context, MirroredVariable (synchronization=ON_WRITE)
+    # in a loaded model behaves differently. Updating MirroredVariable only
+    # update the current replica's variable with the current replica's value.
+    # There's no aggregation. This doesn't affect variables that are updated
+    # through optimizer. This is work as intended but can be surprising.
+
+    class Model(tf.Module):
+
+      def __init__(self):
+        self.v = tf.Variable(
+            0.,
+            synchronization=tf.VariableSynchronization.ON_WRITE,
+            aggregation=tf.VariableAggregation.MEAN)
+
+      @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
+      def update(self, value):
+        self.v.assign_add(value)
+
+    export_dir = self.get_temp_dir()
+    # value = strategy.experimental_distribute_values_from_function(
+    #     lambda ctx: tf.identity([1., 2.][ctx.replica_id_in_sync_group]))
+    with strategy.scope():
+      m = Model()
+      tf.saved_model.save(m, export_dir)
+      del m
+
+    with strategy.scope():
+      loaded = tf.saved_model.load(export_dir)
+    self.assertAllEqual(
+        self.evaluate(strategy.experimental_local_results(loaded.v)), [0., 0.])
+    self.evaluate(loaded.v.assign(1.))
+    self.assertAllEqual(
+        self.evaluate(strategy.experimental_local_results(loaded.v)), [1., 1.])
+    # TODO(b/157621013): strategy.run doesn't work with tf.function with
+    # input_signature (Similar to test_update_sync_on_read_variable)
+    # strategy.run(loaded.update, args=(value,))
+    # self.assertAllEqual(
+    #    self.evaluate(strategy.experimental_local_results(loaded.v)), [2., 3.])
+
+  # TODO(crccw): add a test case that trains a saved model with optimizer.
+
+  def test_model_with_loaded_v1_layer_broken(self, strategy):
+    # If a model contains a layer loaded from SavedModel, and if the model is
+    # loaded under tf.distribute.Strategy, it can't be saved and loaded again
+    # under tf.distribute.Strategy.
+    #
+    # Although the error is the same models with TF2 SavedModel, the cause is
+    # different. TF1 models loaded in API contain an initializer, which is
+    # invoked upon loading. Since loading is in the cross-replica context, that
+    # fails.
+    #
+    # Note that these saved model can still be loaded and used without
+    # tf.distribute.Strategy.
+    #
+    # Some tf.hub layers are converted from TF1, by loading TF1 saved model in
+    # TF2 then saved in TF2. This issue disables them to work with
+    # tf.distribute.Strategy.
+    v1_export_dir = self.get_temp_dir()
+
+    with tf1.Graph().as_default(), tf1.Session() as sess:
+      v = tf1.Variable(1., use_resource=True)
+      sess.run(v.initializer)
+      builder = tf1.saved_model.Builder(v1_export_dir)
+      builder.add_meta_graph_and_variables(
+          sess,
+          tags=[tf1.saved_model.tag_constants.TRAINING],
+          main_op=tf1.tables_initializer(),
+          strip_default_attrs=True)
+      builder.save()
+
+    v1_loaded = tf.saved_model.load(v1_export_dir)
+    v2_export_dir = self.get_temp_dir()
+    tf.saved_model.save(v1_loaded, v2_export_dir)
+    with strategy.scope():
+      # TODO(b/150009657): remove after fix.
+      # got error, want no error.
+      with self.assertRaisesRegex(
+          tf.errors.InvalidArgumentError,
+          "from the cross-replica context in an in-replica context"):
+        tf.saved_model.load(v2_export_dir)
+
+
+class PSStrategySaveAndLoadTest(test.TestCase):
+  # Test saved_model saving and loading for parameter server strategy. These
+  # tests are different enough than the tests in `SaveAndLoadForXXX` so we make
+  # a separate test class for them.
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cluster_def = multi_worker_test_base.create_in_process_cluster(
+        num_workers=2, num_ps=2)
+    cls.cluster_resolver = tf.distribute.cluster_resolver.SimpleClusterResolver(
+        tf.train.ClusterSpec(cluster_def))
+
+  def tearDown(self):
+    super().tearDown()
+    context._reset_context()
+
+  def load_and_run_v1(self,
+                      model_dir,
+                      inputs,
+                      signature_key=tf1.saved_model.signature_constants
+                      .DEFAULT_SERVING_SIGNATURE_DEF_KEY):
+    """Load a SavedModel into a TF 1.x-style graph and run `signature_key`."""
+    graph = tf.Graph()
+    with graph.as_default(), tf1.Session() as session:
+      meta_graph_def = tf1.saved_model.load(
+          session, [tf1.saved_model.tag_constants.SERVING], model_dir)
+      signature = meta_graph_def.signature_def[signature_key]
+      feed_dict = {}
+      for arg_name in inputs.keys():
+        input_tensor = session.graph.get_tensor_by_name(
+            signature.inputs[arg_name].name)
+        feed_dict[input_tensor] = inputs[arg_name]
+      output_dict = {}
+      for output_name, output_tensor_info in signature.outputs.items():
+        output_dict[output_name] = session.graph.get_tensor_by_name(
+            output_tensor_info.name)
+      return session.run(output_dict, feed_dict)["output_0"]
+
+  class Model(tf.Module):
+
+    def __init__(self):
+      self.v1 = tf.Variable([0, 0, 0, 0])
+      self.v2 = tf.Variable([1, 1, 1, 1])
+      self.table = lookup_ops.MutableHashTable(
+          key_dtype=tf.int32, value_dtype=tf.int32, default_value=-1)
+
+    def train(self):
+      # simulate a training process to mutate the state of the model.
+      self.v1.assign([2, 2, 2, 2])
+      self.v2.assign([3, 3, 3, 3])
+      self.table.insert(keys=1, values=1)
+
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=(), dtype=tf.dtypes.int32, name="x")
+    ])
+    def __call__(self, x):
+      t = tf.math.add(self.v1, self.v2)
+      return tf.math.add(t, self.table.lookup(x))
+
+  def test_basic(self):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    model_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = self.Model()
+    m.train()
+    tf.saved_model.save(m, model_dir)
+
+    # Load via V2 API.
+    loaded = tf.saved_model.load(model_dir)
+    self.assertRegex(loaded.v1.device, "/job:chief/replica:0/task:0")
+    self.assertRegex(loaded.v2.device, "/job:chief/replica:0/task:0")
+    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
+    loaded.v2.assign([1, 1, 1, 1])
+    self.assertAllEqual(loaded(tf.identity(1)), [4, 4, 4, 4])
+
+    # Load via V1 API.
+    self.assertAllEqual(self.load_and_run_v1(model_dir, {"x": 1}), [6, 6, 6, 6])
+
+  def test_load_to_same_strategy(self):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    model_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = self.Model()
+    m.train()
+    tf.saved_model.save(m, model_dir)
+
+    with strategy.scope():
+      loaded = tf.saved_model.load(model_dir)
+
+    # Make sure that the variables are created on different devices. SavedModel
+    # may load the variables in a different order compared to the creation order
+    # so the devices may not be exactly the same as before.
+    self.assertTrue(("/job:ps/replica:0/task:0" in loaded.v1.device and
+                     "/job:ps/replica:0/task:1" in loaded.v2.device) or
+                    ("/job:ps/replica:0/task:1" in loaded.v1.device and
+                     "/job:ps/replica:0/task:0" in loaded.v2.device))
+    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
+
+  def test_load_to_different_strategy(self):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver)
+    model_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = self.Model()
+    m.train()
+    tf.saved_model.save(m, model_dir)
+
+    del m  # Garbage collect variables before we reset the context.
+    context._reset_context()
+
+    mirrored_strategy = tf.distribute.MirroredStrategy(devices=["CPU:0"])
+    with mirrored_strategy.scope():
+      loaded = tf.saved_model.load(model_dir)
+    self.assertIsInstance(loaded.v1, values.DistributedVariable)
+    self.assertAllEqual(loaded(tf.identity(1)), [6, 6, 6, 6])
+
+  def test_sharded_variable(self):
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver, tf1.fixed_size_partitioner(2))
+    model_dir = self.get_temp_dir()
+    with strategy.scope():
+      m = self.Model()
+      self.assertIsInstance(m.v1, sharded_variable.ShardedVariable)
+    m.train()
+    tf.saved_model.save(m, model_dir)
+
+    self.assertAllEqual(self.load_and_run_v1(model_dir, {"x": 1}), [6, 6, 6, 6])
+
+  def test_load_with_partitioner_works(self):
+    model = self.Model()
+    model_dir = self.get_temp_dir()
+    tf.saved_model.save(model, model_dir)
+
+    strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
+        self.cluster_resolver, tf1.fixed_size_partitioner(2))
+    with strategy.scope():
+      tf.saved_model.load(model_dir)
+
+
+if __name__ == "__main__":
+  test_util.main()

tensorflow/python/distribute/integration_test/tpu_memory_test.py

@@ -0,0 +1,202 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""This file contains integration test for TPUStrategy in regards to memory."""
+
+import gc
+
+import tensorflow as tf
+
+from tensorflow.python.eager import context
+from tensorflow.python.platform import flags
+
+FLAGS = flags.FLAGS
+NUM_CLASS = 10
+
+
+def get_dataset():
+
+  def generate_data(_):
+    image = tf.ones([500, 500, 3], dtype=tf.float32)
+    label = tf.zeros([1], dtype=tf.int32)
+    return image, label
+
+  def preprocess(image, label):
+    label = tf.cast(label, tf.int32)
+    label = tf.one_hot(label, NUM_CLASS)
+    label = tf.reshape(label, [NUM_CLASS])
+    return image, label
+
+  dataset = tf.data.Dataset.range(1)
+  dataset = dataset.repeat()
+  dataset = dataset.map(
+      generate_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+  dataset = dataset.map(
+      preprocess, num_parallel_calls=tf.data.experimental.AUTOTUNE)
+  dataset = dataset.repeat()
+  dataset = dataset.batch(128, drop_remainder=True)
+  return dataset
+
+
+class TpuMemoryTest(tf.test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # Clear all cached tensors
+    context._reset_context()
+    # Run garbage collection to free any tensors from previous
+    # runs.
+    gc.collect()
+
+    # Run a small program and copy the result to CPU.
+    # This causes deferred deallocations to be flushed and new memory to be
+    # allocated in a less fragmented way.
+    # Turning deferred deallocations off no longer seems to work.
+    assert tf.reduce_sum(tf.random.uniform(
+        (1024, 128), dtype=tf.float32)).numpy() > 1.0
+
+    self.resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
+        tpu="", project=None, zone=None)
+
+    tf.config.experimental_connect_to_cluster(self.resolver)
+    tf.tpu.experimental.initialize_tpu_system(self.resolver)
+
+  def testAutoDefragInProgramLoading(self):
+    # This test covers the case when training a large model on TPU. TPU HBM
+    # is not big enough to hold all TPU buffers and preserve stack for the
+    # TPU program. Runtime will automatically unload unused TPU program to
+    # free up space for TPU buffers. Having lots of TPU buffer may also
+    # introduce fragmentation in HBM to prevent us loading a TPU program
+    # properly. Runtime will automatically defrag in order to load a large
+    # TPU program.
+
+    strategy = tf.distribute.TPUStrategy(self.resolver)
+    dataset = get_dataset()
+    iterator = iter(
+        strategy.experimental_distribute_dataset(dataset,
+                                                 tf.distribute.InputOptions()))
+
+    # Create a dummy big model that is close to HBM limit (15G):
+    # Temp HBM: 11G
+    # Sharded variable size: 2G
+    # Unsharded variables size: 4G
+    with strategy.scope():
+      x = tf.keras.layers.Input(shape=(500, 500, 3), name="input")
+      y = tf.keras.layers.Conv2D(
+          384, (15, 15),
+          strides=(2, 2),
+          padding="valid",
+          use_bias=False,
+          kernel_initializer="he_normal",
+          name="conv1")(
+              x)
+      y = tf.keras.layers.BatchNormalization(
+          momentum=0.997, center=True, scale=True)(
+              y)
+      y = tf.keras.layers.Dense(
+          10,
+          activation="softmax",
+          kernel_initializer=tf.random_normal_initializer(stddev=0.01))(
+              y)
+      y = tf.keras.layers.Conv2D(
+          64, (9, 9),
+          strides=(2, 2),
+          padding="valid",
+          use_bias=False,
+          kernel_initializer="he_normal",
+          name="conv2")(
+              y)
+      y = tf.keras.layers.Flatten()(y)
+      y = tf.keras.layers.Dense(
+          1024,
+          activation="softmax",
+          kernel_initializer=tf.random_normal_initializer(stddev=0.01))(
+              y)
+      y = tf.keras.layers.Dense(
+          1024,
+          activation="softmax",
+          kernel_initializer=tf.random_normal_initializer(stddev=0.01))(
+              y)
+      y = tf.keras.layers.Dense(
+          NUM_CLASS,
+          activation="softmax",
+          kernel_initializer=tf.random_normal_initializer(stddev=0.01))(
+              y)
+      model = tf.keras.Model(x, y)
+      optimizer = tf.keras.optimizers.legacy.SGD(learning_rate=0.1)
+      loss_obj = tf.keras.losses.CategoricalCrossentropy(
+          label_smoothing=0.0, reduction=tf.keras.losses.Reduction.NONE)
+      model.compile(optimizer=optimizer, loss=loss_obj)
+
+    @tf.function
+    def train_step(iterator):
+
+      def step_fn(inputs):
+        images, targets = inputs
+        with tf.GradientTape() as tape:
+          outputs = model(images, training=True)
+          loss = model.loss(targets, outputs)
+
+        grads = tape.gradient(loss, model.trainable_variables)
+        model.optimizer.apply_gradients(zip(grads, model.trainable_variables))
+        return loss
+
+      # Using host training loop here to trigger weight-update-sharding. It will
+      # introduce shard variable and unshard variable ops into the graph.
+      # When running unshard variable op, HBM won't have enough space for
+      # unsharded variables: 11G + 2G + 4G > 15G. So Runtime will have to
+      # automatically unload step function to free up space for unshard
+      # variable op.
+      for _ in tf.range(tf.constant(20)):
+        strategy.run(step_fn, args=(next(iterator),))
+
+      # We want to load the step function again after unshard variable op.
+      # However, we won't have enough space due to fragamentation:
+      # 15G - 2G - 4G < 11G. So Runtime will have to automatically defrag
+      # in order to load the program successfully.
+      strategy.run(step_fn, args=(next(iterator),))
+
+      # A dummy result to indicate this @tf.function has finished.
+      return 1.0
+
+    result = train_step(iterator)
+    self.assertAllClose(1.0, result, atol=1e-07)
+
+  def testAutoDefragInBufferAllocation(self):
+
+    with tf.device("TPU:0"):
+      # DF has ~15G HBM. Following 7 buffers will consume most HBM.
+      # pylint: disable=unused-variable
+      buffer_2g_1 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      buffer_2g_2 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      buffer_2g_3 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      buffer_2g_4 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      buffer_2g_5 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      buffer_2g_6 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      buffer_2g_7 = tf.random.uniform((2, 256, 1024, 1024), dtype=tf.float32)
+      #  pylint: enable=unused-variable
+
+      # Deallocate two buffers.
+      del buffer_2g_1, buffer_2g_3
+      gc.collect()
+
+      # The buffer we just deallocated doesn't provide enough contiguous region
+      # for allocating 4G. This allocation will trigger auto-defrag.
+      buffer_4g = tf.random.uniform((4, 256, 1024, 1024), dtype=tf.float32)
+
+    self.assertEndsWith(buffer_4g.device, "device:TPU:0")
+
+
+if __name__ == "__main__":
+  tf.test.main()

tensorflow/python/distribute/load_context.py

@@ -0,0 +1,66 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Context for storing options for loading a SavedModel."""
+
+import contextlib
+import threading
+
+from tensorflow.python.util.tf_export import tf_export
+
+
+class LoadContext(threading.local):
+  """A context for loading a model."""
+
+  def __init__(self):
+    super().__init__()
+    self._entered_load_context = []
+    self._load_options = None
+
+  def set_load_options(self, load_options):
+    self._load_options = load_options
+    self._entered_load_context.append(True)
+
+  def clear_load_options(self):
+    self._load_options = None
+    self._entered_load_context.pop()
+
+  def load_options(self):
+    return self._load_options
+
+  def in_load_context(self):
+    return self._entered_load_context
+
+
+_load_context = LoadContext()
+
+
+@tf_export("__internal__.load_context", v1=[])
+@contextlib.contextmanager
+def load_context(load_options):
+  _load_context.set_load_options(load_options)
+  try:
+    yield
+  finally:
+    _load_context.clear_load_options()
+
+
+def get_load_options():
+  """Returns the load options under a load context."""
+  return _load_context.load_options()
+
+
+def in_load_context():
+  """Returns whether under a load context."""
+  return _load_context.in_load_context()

tensorflow/python/distribute/merge_call_interim.py

@@ -0,0 +1,54 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A module for interm merge-call related internal APIs."""
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__internal__.distribute.strategy_supports_no_merge_call", v1=[])
+def strategy_supports_no_merge_call():
+  """Returns if the current `Strategy` can operate in pure replica context."""
+  if not distribute_lib.has_strategy():
+    return True
+  strategy = distribute_lib.get_strategy()
+  return not strategy.extended._use_merge_call()  # pylint: disable=protected-access
+
+
+@tf_export("__internal__.distribute.interim.maybe_merge_call", v1=[])
+def maybe_merge_call(fn, strategy, *args, **kwargs):
+  """Maybe invoke `fn` via `merge_call` which may or may not be fulfilled.
+
+  The caller of this utility function requests to invoke `fn` via `merge_call`
+  at `tf.distribute.Strategy`'s best efforts. It is `tf.distribute`'s internal
+  whether the request is honored, depending on the `Strategy`. See
+  `tf.distribute.ReplicaContext.merge_call()` for more information.
+
+  This is an interim API which is subject to removal and does not guarantee
+  backward-compatibility.
+
+  Args:
+    fn: the function to be invoked.
+    strategy: the `tf.distribute.Strategy` to call `fn` with.
+    *args: the positional arguments to be passed in to `fn`.
+    **kwargs: the keyword arguments to be passed in to `fn`.
+
+  Returns:
+    The return value of the `fn` call.
+  """
+  if strategy_supports_no_merge_call():
+    return fn(strategy, *args, **kwargs)
+  else:
+    return distribute_lib.get_replica_context().merge_call(
+        fn, args=args, kwargs=kwargs)

tensorflow/python/distribute/metrics_v1_test.py

@@ -0,0 +1,457 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for V1 metrics."""
+from absl.testing import parameterized
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import combinations
+from tensorflow.python.distribute import strategy_combinations
+from tensorflow.python.distribute import strategy_test_lib
+from tensorflow.python.eager import test
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import metrics
+from tensorflow.python.ops import variables
+
+
+def _labeled_dataset_fn():
+  # First four batches of x: labels, predictions -> (labels == predictions)
+  #  0: 0, 0 -> True;   1: 1, 1 -> True;   2: 2, 2 -> True;   3: 3, 0 -> False
+  #  4: 4, 1 -> False;  5: 0, 2 -> False;  6: 1, 0 -> False;  7: 2, 1 -> False
+  #  8: 3, 2 -> False;  9: 4, 0 -> False; 10: 0, 1 -> False; 11: 1, 2 -> False
+  # 12: 2, 0 -> False; 13: 3, 1 -> False; 14: 4, 2 -> False; 15: 0, 0 -> True
+  return dataset_ops.Dataset.range(1000).map(
+      lambda x: {"labels": x % 5, "predictions": x % 3}).batch(
+          4, drop_remainder=True)
+
+
+def _boolean_dataset_fn():
+  # First four batches of labels, predictions: {TP, FP, TN, FN}
+  # with a threshold of 0.5:
+  #   T, T -> TP;  F, T -> FP;   T, F -> FN
+  #   F, F -> TN;  T, T -> TP;   F, T -> FP
+  #   T, F -> FN;  F, F -> TN;   T, T -> TP
+  #   F, T -> FP;  T, F -> FN;   F, F -> TN
+  return dataset_ops.Dataset.from_tensor_slices({
+      "labels": [True, False, True, False],
+      "predictions": [True, True, False, False]}).repeat().batch(
+          3, drop_remainder=True)
+
+
+def _threshold_dataset_fn():
+  # First four batches of labels, predictions: {TP, FP, TN, FN}
+  # with a threshold of 0.5:
+  #   True, 1.0 -> TP;  False, .75 -> FP;   True, .25 -> FN
+  #  False, 0.0 -> TN;   True, 1.0 -> TP;  False, .75 -> FP
+  #   True, .25 -> FN;  False, 0.0 -> TN;   True, 1.0 -> TP
+  #  False, .75 -> FP;   True, .25 -> FN;  False, 0.0 -> TN
+  return dataset_ops.Dataset.from_tensor_slices({
+      "labels": [True, False, True, False],
+      "predictions": [1.0, 0.75, 0.25, 0.]}).repeat().batch(
+          3, drop_remainder=True)
+
+
+def _regression_dataset_fn():
+  return dataset_ops.Dataset.from_tensor_slices({
+      "labels": [1., .5, 1., 0.],
+      "predictions": [1., .75, .25, 0.]}).repeat()
+
+
+def all_combinations():
+  return combinations.combine(
+      distribution=[
+          strategy_combinations.default_strategy,
+          strategy_combinations.one_device_strategy,
+          strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
+          strategy_combinations.mirrored_strategy_with_two_gpus,
+      ],
+      mode=["graph"])
+
+
+def tpu_combinations():
+  return combinations.combine(
+      distribution=[
+          strategy_combinations.tpu_strategy_one_step,
+          strategy_combinations.tpu_strategy
+      ],
+      mode=["graph"])
+
+
+# TODO(josh11b): Test metrics.recall_at_top_k, metrics.average_precision_at_k,
+# metrics.precision_at_k
+class MetricsV1Test(test.TestCase, parameterized.TestCase):
+
+  def _test_metric(self, distribution, dataset_fn, metric_fn, expected_fn):
+    with ops.Graph().as_default(), distribution.scope():
+      iterator = distribution.make_input_fn_iterator(lambda _: dataset_fn())
+      if strategy_test_lib.is_tpu_strategy(distribution):
+        def step_fn(ctx, inputs):
+          value, update = distribution.extended.call_for_each_replica(
+              metric_fn, args=(inputs,))
+          ctx.set_non_tensor_output(name="value", output=value)
+          return distribution.group(update)
+
+        ctx = distribution.extended.experimental_run_steps_on_iterator(
+            step_fn, iterator, iterations=distribution.extended.steps_per_run)
+        update = ctx.run_op
+        value = ctx.non_tensor_outputs["value"]
+        # In each run, we run multiple steps, and each steps consumes as many
+        # batches as number of replicas.
+        batches_per_update = (
+            distribution.num_replicas_in_sync *
+            distribution.extended.steps_per_run)
+      else:
+        value, update = distribution.extended.call_for_each_replica(
+            metric_fn, args=(iterator.get_next(),))
+        update = distribution.group(update)
+        # TODO(josh11b): Once we switch to using a global batch size for input,
+        # replace "distribution.num_replicas_in_sync" with "1".
+        batches_per_update = distribution.num_replicas_in_sync
+
+      self.evaluate(iterator.initializer)
+      self.evaluate(variables.local_variables_initializer())
+
+      batches_consumed = 0
+      for i in range(4):
+        self.evaluate(update)
+        batches_consumed += batches_per_update
+        self.assertAllClose(expected_fn(batches_consumed),
+                            self.evaluate(value),
+                            0.001,
+                            msg="After update #" + str(i+1))
+        if batches_consumed >= 4:  # Consume 4 input batches in total.
+          break
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testMean(self, distribution):
+    def _dataset_fn():
+      return dataset_ops.Dataset.range(1000).map(math_ops.to_float).batch(
+          4, drop_remainder=True)
+
+    def _expected_fn(num_batches):
+      # Mean(0..3) = 1.5, Mean(0..7) = 3.5, Mean(0..11) = 5.5, etc.
+      return num_batches * 2 - 0.5
+
+    self._test_metric(distribution, _dataset_fn, metrics.mean, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testAccuracy(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.accuracy(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [3./4, 3./8, 3./12, 4./16][num_batches - 1]
+
+    self._test_metric(
+        distribution, _labeled_dataset_fn, _metric_fn, _expected_fn)
+
+  # TODO(priyag, jhseu): Enable TPU for this test once scatter_add is added
+  # for TPUMirroredVariable.
+  @combinations.generate(all_combinations())
+  def testMeanPerClassAccuracy(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.mean_per_class_accuracy(
+          labels, predictions, num_classes=5)
+
+    def _expected_fn(num_batches):
+      mean = lambda x: sum(x) / len(x)
+      return [mean([1., 1., 1., 0., 0.]),
+              mean([0.5, 0.5, 0.5, 0., 0.]),
+              mean([1./3, 1./3, 0.5, 0., 0.]),
+              mean([0.5, 1./3, 1./3, 0., 0.])][num_batches - 1]
+
+    self._test_metric(
+        distribution, _labeled_dataset_fn, _metric_fn, _expected_fn)
+
+  # NOTE(priyag): This metric doesn't work on TPUs yet.
+  @combinations.generate(all_combinations())
+  def testMeanIOU(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.mean_iou(
+          labels, predictions, num_classes=5)
+
+    def _expected_fn(num_batches):
+      mean = lambda x: sum(x) / len(x)
+      return [mean([1./2, 1./1, 1./1, 0.]),  # no class 4 in first batch
+              mean([1./4, 1./4, 1./3, 0., 0.]),
+              mean([1./6, 1./6, 1./5, 0., 0.]),
+              mean([2./8, 1./7, 1./7, 0., 0.])][num_batches - 1]
+
+    self._test_metric(
+        distribution, _labeled_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testMeanTensor(self, distribution):
+    def _dataset_fn():
+      dataset = dataset_ops.Dataset.range(1000).map(math_ops.to_float)
+      # Want to produce a fixed, known shape, so drop remainder when batching.
+      dataset = dataset.batch(4, drop_remainder=True)
+      return dataset
+
+    def _expected_fn(num_batches):
+      # Mean(0, 4, ..., 4 * num_batches - 4) == 2 * num_batches - 2
+      # Mean(1, 5, ..., 4 * num_batches - 3) == 2 * num_batches - 1
+      # Mean(2, 6, ..., 4 * num_batches - 2) == 2 * num_batches
+      # Mean(3, 7, ..., 4 * num_batches - 1) == 2 * num_batches + 1
+      first = 2. * num_batches - 2.
+      return [first, first + 1., first + 2., first + 3.]
+
+    self._test_metric(
+        distribution, _dataset_fn, metrics.mean_tensor, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testAUCROC(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.auc(labels, predictions, num_thresholds=8, curve="ROC",
+                         summation_method="careful_interpolation")
+
+    def _expected_fn(num_batches):
+      return [0.5, 7./9, 0.8, 0.75][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testAUCPR(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.auc(labels, predictions, num_thresholds=8, curve="PR",
+                         summation_method="careful_interpolation")
+
+    def _expected_fn(num_batches):
+      return [0.797267, 0.851238, 0.865411, 0.797267][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testFalseNegatives(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.false_negatives(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [1., 1., 2., 3.][num_batches - 1]
+
+    self._test_metric(
+        distribution, _boolean_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testFalseNegativesAtThresholds(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.false_negatives_at_thresholds(labels, predictions, [.5])
+
+    def _expected_fn(num_batches):
+      return [[1.], [1.], [2.], [3.]][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testTrueNegatives(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.true_negatives(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [0., 1., 2., 3.][num_batches - 1]
+
+    self._test_metric(
+        distribution, _boolean_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testTrueNegativesAtThresholds(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.true_negatives_at_thresholds(labels, predictions, [.5])
+
+    def _expected_fn(num_batches):
+      return [[0.], [1.], [2.], [3.]][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testFalsePositives(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.false_positives(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [1., 2., 2., 3.][num_batches - 1]
+
+    self._test_metric(
+        distribution, _boolean_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testFalsePositivesAtThresholds(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.false_positives_at_thresholds(labels, predictions, [.5])
+
+    def _expected_fn(num_batches):
+      return [[1.], [2.], [2.], [3.]][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testTruePositives(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.true_positives(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [1., 2., 3., 3.][num_batches - 1]
+
+    self._test_metric(
+        distribution, _boolean_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testTruePositivesAtThresholds(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.true_positives_at_thresholds(labels, predictions, [.5])
+
+    def _expected_fn(num_batches):
+      return [[1.], [2.], [3.], [3.]][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testPrecision(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.precision(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [0.5, 0.5, 0.6, 0.5][num_batches - 1]
+
+    self._test_metric(
+        distribution, _boolean_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testPrecisionAtThreshold(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.precision_at_thresholds(labels, predictions, [0.5])
+
+    def _expected_fn(num_batches):
+      return [[0.5], [0.5], [0.6], [0.5]][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testRecall(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.recall(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [0.5, 2./3, 0.6, 0.5][num_batches - 1]
+
+    self._test_metric(
+        distribution, _boolean_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testRecallAtThreshold(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.recall_at_thresholds(labels, predictions, [0.5])
+
+    def _expected_fn(num_batches):
+      return [[0.5], [2./3], [0.6], [0.5]][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testMeanSquaredError(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.mean_squared_error(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [0., 1./32, 0.208333, 0.15625][num_batches - 1]
+
+    self._test_metric(
+        distribution, _regression_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations() + tpu_combinations())
+  def testRootMeanSquaredError(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.root_mean_squared_error(labels, predictions)
+
+    def _expected_fn(num_batches):
+      return [0., 0.176777, 0.456435, 0.395285][num_batches - 1]
+
+    self._test_metric(
+        distribution, _regression_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations())
+  def testSensitivityAtSpecificity(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.sensitivity_at_specificity(labels, predictions, 0.8)
+
+    def _expected_fn(num_batches):
+      return [0.5, 2./3, 0.6, 0.5][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+  @combinations.generate(all_combinations())
+  def testSpecificityAtSensitivity(self, distribution):
+    def _metric_fn(x):
+      labels = x["labels"]
+      predictions = x["predictions"]
+      return metrics.specificity_at_sensitivity(labels, predictions, 0.95)
+
+    def _expected_fn(num_batches):
+      return [0., 1./3, 0.5, 0.5][num_batches - 1]
+
+    self._test_metric(
+        distribution, _threshold_dataset_fn, _metric_fn, _expected_fn)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/distribute/mirrored_run.py

@@ -0,0 +1,531 @@
+# 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.
+# ==============================================================================
+"""Class MirroredStrategy implementing tf.distribute.Strategy."""
+
+import contextlib
+import threading
+import weakref
+
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.autograph.core import ag_ctx as autograph_ctx
+from tensorflow.python.autograph.impl import api as autograph
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute import shared_variable_creator
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import device as tf_device
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import summary_ops_v2
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import coordinator
+from tensorflow.python.util import traceback_utils
+
+
+def _is_gpu_device(device):
+  return tf_device.DeviceSpec.from_string(device).device_type == "GPU"
+
+
+def call_for_each_replica(strategy, fn, args=None, kwargs=None):
+  """Call `fn` on each worker devices(replica).
+
+  It's highly recommended to wrap the call to this function inside a
+  `tf.function`, otherwise the performance is poor.
+
+  Args:
+    strategy: `tf.distribute.Strategy`.
+    fn: function to call on each worker devices.
+    args: positional arguments to `fn`.
+    kwargs: keyword arguments to `fn`.
+
+  Returns:
+    Wrapped returned value of `fn` from all replicas.
+  """
+  if args is None:
+    args = ()
+  if kwargs is None:
+    kwargs = {}
+
+  if isinstance(fn, def_function.Function):
+    # Don't lift up the tf.function decoration if `fn` is compiled with XLA
+    # and all devices are GPU. In this case we will use collectives to do
+    # cross-device communication, thus no merge_call is in the path.
+    if fn._jit_compile and all(  # pylint: disable=protected-access
+        [_is_gpu_device(d) for d in strategy.extended.worker_devices]):
+      return _call_for_each_replica(strategy, fn, args, kwargs)
+
+    if strategy not in _cfer_fn_cache:
+      _cfer_fn_cache[strategy] = weakref.WeakKeyDictionary()
+    wrapped = _cfer_fn_cache[strategy].get(fn)
+    if wrapped is None:
+      # We need to wrap fn such that it triggers _call_for_each_replica inside
+      # the tf.function. We use _clone() instead of @tf.function wrapped
+      # call_for_each_replica() because we would like to retain the arguments to
+      # the @tf.function decorator of fn.
+      def wrapped_fn(*args, **kwargs):
+        return call_for_each_replica(strategy, fn.python_function, args, kwargs)
+
+      wrapped = fn._clone(  # pylint: disable=protected-access
+          python_function=wrapped_fn)
+      _cfer_fn_cache[strategy][fn] = wrapped
+    return wrapped(*args, **kwargs)
+
+  if context.executing_eagerly():
+    logging.log_first_n(
+        logging.WARN, "Using %s eagerly has significant "
+        "overhead currently. We will be working on improving "
+        "this in the future, but for now please wrap "
+        "`call_for_each_replica` or `experimental_run` or "
+        "`run` inside a tf.function to get "
+        "the best performance." % strategy.__class__.__name__, 5)
+  else:
+    # When a tf.function is wrapped to trigger _call_for_each_replica (see
+    # the other branch above), AutoGraph stops conversion at
+    # _call_for_each_replica itself (TF library functions are allowlisted).
+    # This makes sure that the Python function that originally passed to
+    # the tf.function is still converted.
+    fn = autograph.tf_convert(fn, autograph_ctx.control_status_ctx())
+
+  return _call_for_each_replica(strategy, fn, args, kwargs)
+
+
+# Per strategy cache for call_for_each_replica def_function.Function objects.
+_cfer_fn_cache = weakref.WeakKeyDictionary()
+
+
+@contextlib.contextmanager
+def _enter_graph(g, eager, creator_stack=None):
+  """Context manager for selecting a graph and maybe eager mode."""
+  if eager:
+    with g.as_default(), context.eager_mode():
+      if creator_stack is not None:
+        g._variable_creator_stack = creator_stack  # pylint: disable=protected-access
+      yield
+  else:
+    with g.as_default():
+      if creator_stack is not None:
+        g._variable_creator_stack = creator_stack  # pylint: disable=protected-access
+      yield
+
+
+@contextlib.contextmanager
+def _maybe_enter_eager_mode(eager):
+  if eager:
+    with context.eager_mode():
+      yield
+  else:
+    yield
+
+
+def _cpu_device(device):
+  cpu_device = tf_device.DeviceSpec.from_string(device)
+  cpu_device = cpu_device.replace(device_type="CPU", device_index=0)
+  return cpu_device.to_string()
+
+
+class _RequestedStop(Exception):  # pylint: disable=g-bad-exception-name
+  pass
+
+
+def _get_thread_local_configuration_callable():
+  if traceback_utils.is_traceback_filtering_enabled():
+    thread_local_callables = {traceback_utils.enable_traceback_filtering}
+  else:
+    thread_local_callables = {traceback_utils.disable_traceback_filtering}
+  return thread_local_callables
+
+
+def _call_for_each_replica(distribution, fn, args, kwargs):
+  """Run `fn` in separate threads, once per replica/worker device.
+
+  Args:
+    distribution: the DistributionStrategy object.
+    fn: function to run (will be run once per replica, each in its own thread).
+    args: positional arguments for `fn`
+    kwargs: keyword arguments for `fn`.
+
+  Returns:
+    Merged return value of `fn` across all replicas.
+
+  Raises:
+    RuntimeError: If fn() calls get_replica_context().merge_call() a different
+        number of times from the available devices.
+  """
+  # TODO(josh11b): Add this option once we add synchronization to variable
+  # creation. Until then, this is pretty unsafe to use.
+  run_concurrently = False
+  if not context.executing_eagerly():
+    # Needed for per-thread device, etc. contexts in graph mode.
+    ops.get_default_graph().switch_to_thread_local()
+
+  coord = coordinator.Coordinator(clean_stop_exception_types=(_RequestedStop,))
+
+  shared_variable_store = {}
+  devices = distribution.extended.worker_devices
+
+  thread_local_callables = _get_thread_local_configuration_callable()
+
+  # TODO(isaprykin): Create these threads once instead of during every call.
+  threads = []
+  for index in range(len(devices)):
+    variable_creator_fn = shared_variable_creator.make_fn(
+        shared_variable_store, index)
+    t = _MirroredReplicaThread(distribution, coord, index, devices,
+                               variable_creator_fn, fn,
+                               distribute_utils.caching_scope_local,
+                               distribute_utils.select_replica(index, args),
+                               distribute_utils.select_replica(index, kwargs),
+                               thread_local_callables)
+    threads.append(t)
+
+  for t in threads:
+    t.start()
+
+  # When `fn` starts `should_run` event is set on _MirroredReplicaThread
+  # (`MRT`) threads. The execution waits until
+  # `MRT.has_paused` is set, which indicates that either `fn` is
+  # complete or a `get_replica_context().merge_call()` is called.  If `fn` is
+  # complete, then `MRT.done` is set to True.  Otherwise, arguments
+  # of `get_replica_context().merge_call` from all paused threads are grouped
+  # and the `merge_fn` is performed.  Results of the
+  # `get_replica_context().merge_call` are then set to `MRT.merge_result`.
+  # Each such `get_replica_context().merge_call` call returns the
+  # `MRT.merge_result` for that thread when `MRT.should_run` event
+  # is reset again. Execution of `fn` resumes.
+
+  try:
+    with coord.stop_on_exception():
+      all_done = False
+      while not all_done and not coord.should_stop():
+        done = []
+        if run_concurrently:
+          for t in threads:
+            t.should_run.set()
+          for t in threads:
+            t.has_paused.wait()
+            t.has_paused.clear()
+            if coord.should_stop():
+              return None
+            done.append(t.done)
+        else:
+          for t in threads:
+            t.should_run.set()
+            t.has_paused.wait()
+            t.has_paused.clear()
+            if coord.should_stop():
+              return None
+            done.append(t.done)
+        if coord.should_stop():
+          return None
+        all_done = all(done)
+        if not all_done:
+          if any(done):
+            raise RuntimeError("Some replicas made a different number of "
+                               "replica_context().merge_call() calls.")
+          # get_replica_context().merge_call() case
+          merge_args = distribute_utils.regroup(
+              tuple(t.merge_args for t in threads))
+          merge_kwargs = distribute_utils.regroup(
+              tuple(t.merge_kwargs for t in threads))
+          # We capture the name_scope of the MRT when we call merge_fn
+          # to ensure that if we have opened a name scope in the MRT,
+          # it will be respected when executing the merge function. We only
+          # capture the name_scope from the first MRT and assume it is
+          # the same for all other MRTs.
+          mtt_captured_name_scope = threads[0].captured_name_scope
+          mtt_captured_var_scope = threads[0].captured_var_scope
+          # Capture and merge the control dependencies from all the threads.
+          mtt_captured_control_deps = set()
+          for t in threads:
+            mtt_captured_control_deps.update(t.captured_control_deps)
+
+          # Control is transferred from _MirroredReplicaThread (MRT) to the main
+          # thread, i.e., here, to perform `merge_fn`, and thus we preserve the
+          # name scope,  control dependencies, etc. from MRT at the time
+          # `merge_call` is made.
+          # One special case is that the `merge_call` is made under an
+          # `tf.init_scope` in the MRT. `tf.init_scope` will clear control
+          # dependencies, pause gradient tape, and enter the lowest context on
+          # the `context_stack` that is not building a graph function. Entering
+          # the lowest context could be one of the two things: installation of a
+          # graph as the default graph or switch into eager mode. If the former
+          # is done and causes `merge_call` to be called in a different graph
+          # from the one in which `call_for_each_replica` is called, we do not
+          # allow this case (see comment in `_merge_call`) and we would not have
+          # arrived here due to the assertion in `_merge_call`. However, if the
+          # latter is done, we want to make sure the main thread enter an eager
+          # mode scope as well so that `merge_fn` does not have trouble
+          # accessing resources defined in MRT under the same context.
+          with ops.name_scope(
+              mtt_captured_name_scope), ops.control_dependencies(
+                  mtt_captured_control_deps), variable_scope.variable_scope(
+                      mtt_captured_var_scope), _maybe_enter_eager_mode(
+                          threads[0].merge_call_entered_in_eager):
+            merge_result = threads[0].merge_fn(distribution, *merge_args,
+                                               **merge_kwargs)
+          for r, t in enumerate(threads):
+            t.merge_result = distribute_utils.select_replica(r, merge_result)
+  finally:
+    for t in threads:
+      t.should_run.set()
+    coord.join(threads)
+
+  return distribute_utils.regroup(tuple(t.main_result for t in threads))
+
+
+class _MirroredReplicaThread(threading.Thread):
+  """A thread that runs() a function on a device."""
+
+  def __init__(self, dist, coord, replica_id, devices, variable_creator_fn, fn,
+               caching_scope, args, kwargs, thread_local_callables=None):
+    super(_MirroredReplicaThread, self).__init__()
+    self.coord = coord
+    self.distribution = dist
+    self.devices = devices
+    self.replica_id = replica_id
+    self.replica_id_in_sync_group = (
+        dist.extended._get_replica_id_in_sync_group(replica_id))  # pylint: disable=protected-access
+
+    self.variable_creator_fn = variable_creator_fn
+    # State needed to run and return the results of `fn`.
+    self.main_fn = fn
+    self.main_args = args
+    self.main_kwargs = kwargs
+    self.main_result = None
+    self.done = False
+    # State needed to run the next merge_call() (if any) requested via
+    # ReplicaContext.
+    self.merge_fn = None
+    self.merge_args = None
+    self.merge_kwargs = None
+    self.merge_result = None
+    self.captured_name_scope = None
+    self.captured_var_scope = None
+    try:
+      self.caching_scope_entered = caching_scope.new_cache_scope_count
+      self.caching_scope_exited = caching_scope.cache_scope_exited_count
+    except AttributeError:
+      self.caching_scope_entered = None
+      self.caching_scope_exited = None
+
+    # We use a thread.Event for the main thread to signal when this
+    # thread should start running (`should_run`), and another for
+    # this thread to transfer control back to the main thread
+    # (`has_paused`, either when it gets to a
+    # `get_replica_context().merge_call` or when `fn` returns). In
+    # either case the event starts cleared, is signaled by calling
+    # set(). The receiving thread waits for the signal by calling
+    # wait() and then immediately clearing the event using clear().
+    self.should_run = threading.Event()
+    self.has_paused = threading.Event()
+    # These fields have to do with inheriting various contexts from the
+    # parent thread:
+    context.ensure_initialized()
+    ctx = context.context()
+    self.in_eager = ctx.executing_eagerly()
+    self.record_thread_local_summary_state()
+    self.record_thread_local_eager_context_state()
+    self.context_device_policy = (
+        pywrap_tfe.TFE_ContextGetDevicePlacementPolicy(
+            ctx._context_handle))  # pylint: disable=protected-access
+    self.graph = ops.get_default_graph()
+    with ops.init_scope():
+      self._init_in_eager = context.executing_eagerly()
+      self._init_graph = ops.get_default_graph()
+    self._variable_creator_stack = self.graph._variable_creator_stack[:]  # pylint: disable=protected-access
+    self._var_scope = variable_scope.get_variable_scope()
+    # Adding a "/" at end lets us re-enter this scope later.
+    self._name_scope = self.graph.get_name_scope()
+    if self._name_scope:
+      self._name_scope += "/"
+    if self.replica_id > 0:
+      if not self._name_scope:
+        self._name_scope = ""
+      self._name_scope += "replica_%d/" % self.replica_id
+
+    self._thread_local_callables = thread_local_callables
+
+  def run(self):
+    self.should_run.wait()
+    self.should_run.clear()
+    try:
+      if self.coord.should_stop():
+        return
+      self.restore_thread_local_summary_state()
+      self.restore_thread_local_callable()
+      self.restore_thread_local_eager_context_state()
+      if (self.caching_scope_entered is not None and
+          self.caching_scope_exited is not None):
+        distribute_utils.caching_scope_local.new_cache_scope_count = self.caching_scope_entered
+        distribute_utils.caching_scope_local.cache_scope_exited_count = self.caching_scope_exited
+      # TODO(josh11b): Use current logical device instead of 0 here.
+      with self.coord.stop_on_exception(), \
+          _enter_graph(self._init_graph, self._init_in_eager), \
+          _enter_graph(self.graph, self.in_eager,
+                       self._variable_creator_stack), \
+          context.device_policy(self.context_device_policy), \
+          _MirroredReplicaContext(self.distribution,
+                                  self.replica_id_in_sync_group), \
+          ops.device(self.devices[self.replica_id]), \
+          ops.name_scope(self._name_scope), \
+          variable_scope.variable_scope(
+              self._var_scope, reuse=self.replica_id > 0), \
+          variable_scope.variable_creator_scope(self.variable_creator_fn):
+        self.main_result = self.main_fn(*self.main_args, **self.main_kwargs)
+        self.done = True
+    finally:
+      self.has_paused.set()
+
+  def record_thread_local_summary_state(self):
+    """Record the thread local summary state in self."""
+    # TODO(slebedev): is this still relevant? the referenced bug is closed.
+    summary_state = summary_ops_v2._summary_state  # pylint: disable=protected-access
+    self._summary_step = summary_state.step
+    self._summary_writer = summary_state.writer
+    self._summary_recording = summary_state.is_recording
+    self._summary_recording_distribution_strategy = (
+        summary_state.is_recording_distribution_strategy)
+
+  def restore_thread_local_summary_state(self):
+    """Restore thread local summary state from self."""
+    # TODO(slebedev): is this still relevant? the referenced bug is closed.
+    summary_state = summary_ops_v2._summary_state  # pylint: disable=protected-access
+    summary_state.step = self._summary_step
+    summary_state.writer = self._summary_writer
+    summary_state.is_recording = self._summary_recording
+    summary_state.is_recording_distribution_strategy = (
+        self._summary_recording_distribution_strategy)
+
+  def record_thread_local_eager_context_state(self):
+    ctx = context.context()
+    eager_context_state = ctx._thread_local_data  # pylint: disable=protected-access
+    self._eager_context_op_callbacks = eager_context_state.op_callbacks
+    # TODO(b/125892694): record other fields in EagerContext.
+
+  def restore_thread_local_eager_context_state(self):
+    ctx = context.context()
+    eager_context_state = ctx._thread_local_data  # pylint: disable=protected-access
+    eager_context_state.op_callbacks = self._eager_context_op_callbacks
+    # TODO(b/125892694): record other fields in EagerContext.
+
+  def restore_thread_local_callable(self):
+    if self._thread_local_callables:
+      for fn in self._thread_local_callables:
+        fn()
+
+
+class _MirroredReplicaContext(distribute_lib.ReplicaContext):
+  """ReplicaContext for synchronized replica."""
+
+  def _merge_call(self, fn, args, kwargs):
+    """`merge_call()` implementation for synchronized replica.
+
+    This pauses the current replica thread and passes `fn` and its arguments to
+    the main thread. The main thread will wait until all replicas pause, then
+    invoke `fn` with grouped arguments. The current replica thread will continue
+    after `fn` completes.
+
+    See `_call_for_each_replica` for the logic in the main thread.
+
+    Args:
+      fn: a function that is called in cross replica context with grouped
+        arguments from each replica. `fn` should returns grouped values.
+      args: positional arguments to `fn`.
+      kwargs: keyword arguments to `fn`.
+
+    Returns:
+      Return value of `fn` for the current replica.
+
+    Raises:
+      RuntimeError: when merge_call happens in a different graph, e.g. in a
+        different tf.function, which is not supported now.
+      _RequestedStop: when stop is requested.
+    """
+    t = threading.current_thread()
+    assert isinstance(t, _MirroredReplicaThread)
+    t.merge_fn = fn
+    t.merge_args = args
+    t.merge_kwargs = kwargs
+    t.captured_name_scope = t.graph.get_name_scope()
+    # Adding a "/" at end lets us re-enter this scope later.
+    if t.captured_name_scope:
+      t.captured_name_scope += "/"
+
+    t.captured_var_scope = variable_scope.get_variable_scope()
+    t.captured_control_deps = t.graph._current_control_dependencies()  # pylint: disable=protected-access
+
+    t.merge_call_entered_in_eager = context.context().executing_eagerly()
+
+    # It is problematic if `merge_call` is called under a different graph other
+    # than the one that `_call_for_each_replica` is called under, there are
+    # 3 cases this can happen:
+    #
+    #   1. The `fn` passed to `_call_for_each_replica` is decorated with
+    #   `tf.function` and there is a `merge_call` in `fn`. Since
+    #   MirroredStrategy traces a separate function per thread (per device),
+    #   and each trace takes a shared lock, the lock is never released by the
+    #   first thread and subsequent replica threads cannot proceed to trace
+    #   their own functions. This issue is addressed by always converting
+    #   `_call_for_each_replica(tf.function(f))` to
+    #   ``tf.function(_call_for_each_replica(f))`.` in
+    #   `MirroredStrategy._call_for_each_replica`.
+    #
+    #   2. The `fn` passed to `_call_for_each_replica` contains a nested
+    #   `tf.function`, and there is a `merge_call` in the nested `tf.function`.
+    #   In this case each thread can successfully trace its own function, but
+    #   since the `merge_fn` passed to `merge_call` is executed in the main
+    #   thread (where `_call_for_each_replica` is executed), it can't access
+    #   the tensors that come from different graphs.
+    #
+    #   3. The `fn` passed to `_call_for_each_replica` contains a control-flow
+    #   statement, and there is a `merge_call` inside the control-flow body,
+    #   `fn` or `_call_for_each_replica` is decorated with `tf.function`.
+    #   Control flow statement creates a separate graph for its body, similar
+    #   to #2, `merge_fn` executed in the main thread can't access the
+    #   tensors that come from different graphs.
+    #
+    #   We raise an error for #2 and #3.
+    if ops.get_default_graph() != t.graph:
+      raise RuntimeError(
+          "`merge_call` called while defining a new graph or a tf.function."
+          " This can often happen if the function `fn` passed to"
+          " `strategy.run()` contains a nested `@tf.function`, and the nested "
+          "`@tf.function` contains a synchronization point, such as aggregating"
+          " gradients (e.g, optimizer.apply_gradients), or if the function `fn`"
+          " uses a control flow statement which contains a synchronization"
+          " point in the body. Such behaviors are not yet supported. Instead,"
+          " please avoid nested `tf.function`s or control flow statements that"
+          " may potentially cross a synchronization boundary, for example,"
+          " wrap the `fn` passed to `strategy.run` or the entire `strategy.run`"
+          " inside a `tf.function` or move the control flow out of `fn`. If"
+          " you are subclassing a `tf.keras.Model`, please avoid decorating"
+          " overridden methods `test_step` and `train_step` in `tf.function`.")
+
+    t.has_paused.set()
+    t.should_run.wait()
+    t.should_run.clear()
+    if t.coord.should_stop():
+      raise _RequestedStop()
+    t.merge_call_entered_in_eager = None
+    return t.merge_result
+
+  @property
+  def devices(self):
+    distribute_lib.require_replica_context(self)
+    return [
+        self._strategy.extended.worker_devices_by_replica[
+            self._replica_id_in_sync_group]
+    ]

tensorflow/python/distribute/mirrored_strategy.py

@@ -0,0 +1,941 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Class MirroredStrategy implementing tf.distribute.Strategy."""
+
+import copy
+
+from tensorflow.python import tf2
+from tensorflow.python.distribute import collective_util
+from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib
+from tensorflow.python.distribute import cross_device_utils
+from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import distribute_utils
+from tensorflow.python.distribute import input_lib
+from tensorflow.python.distribute import input_util
+from tensorflow.python.distribute import mirrored_run
+from tensorflow.python.distribute import multi_worker_util
+from tensorflow.python.distribute import numpy_dataset
+from tensorflow.python.distribute import reduce_util
+from tensorflow.python.distribute import values
+from tensorflow.python.distribute import values_util
+from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
+from tensorflow.python.distribute.v1 import input_lib as input_lib_v1
+from tensorflow.python.eager import context
+from tensorflow.python.eager import record
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import device as tf_device
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+# TODO(josh11b): Replace asserts in this file with if ...: raise ...
+
+
+def _is_device_list_single_worker(devices):
+  """Checks whether the devices list is for single or multi-worker.
+
+  Args:
+    devices: a list of device strings or tf.config.LogicalDevice objects, for
+      either local or for remote devices.
+
+  Returns:
+    a boolean indicating whether these device strings are for local or for
+    remote.
+
+  Raises:
+    ValueError: if device strings are not consistent.
+  """
+  specs = []
+  for d in devices:
+    name = d.name if isinstance(d, context.LogicalDevice) else d
+    specs.append(tf_device.DeviceSpec.from_string(name))
+  num_workers = len({(d.job, d.task, d.replica) for d in specs})
+  all_local = all(d.job in (None, "localhost") for d in specs)
+  any_local = any(d.job in (None, "localhost") for d in specs)
+
+  if any_local and not all_local:
+    raise ValueError("Local device should have only 'localhost' in the job "
+                     "field in device string. "
+                     "E.g. 'job:localhost' in "
+                     "/job:localhost/replica:0/task:0/device:CPU:0"
+                     "Devices cannot have mixed list of device strings "
+                     "containing both localhost and other job types such as "
+                     "worker, ps etc. ")
+
+  if num_workers == 1 and not all_local:
+    if any(d.task is None for d in specs):
+      raise ValueError("Remote device string must have task specified."
+                       "E.g. 'task:0' in "
+                       "/job:worker/replica:0/task:0/device:CPU:0")
+
+  return num_workers == 1
+
+
+def _cluster_spec_to_device_list(cluster_spec, num_gpus_per_worker):
+  """Returns a device list given a cluster spec."""
+  cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec)
+  devices = []
+  for task_type in ("chief", "worker"):
+    for task_id in range(len(cluster_spec.as_dict().get(task_type, []))):
+      if num_gpus_per_worker == 0:
+        devices.append("/job:%s/task:%d/device:CPU:0" % (task_type, task_id))
+      else:
+        devices.extend([
+            "/job:%s/task:%d/device:GPU:%i" % (task_type, task_id, gpu_id)
+            for gpu_id in range(num_gpus_per_worker)
+        ])
+  return devices
+
+
+def _group_device_list(devices):
+  """Groups the devices list by task_type and task_id.
+
+  Args:
+    devices: a list of device strings for remote devices.
+
+  Returns:
+    a dict of list of device strings mapping from task_type to a list of devices
+    for the task_type in the ascending order of task_id.
+  """
+  assert not _is_device_list_single_worker(devices)
+  device_dict = {}
+
+  for d in devices:
+    d_spec = tf_device.DeviceSpec.from_string(d)
+
+    # Create an entry for the task_type.
+    if d_spec.job not in device_dict:
+      device_dict[d_spec.job] = []
+
+    # Fill the device list for task_type until it covers the task_id.
+    while len(device_dict[d_spec.job]) <= d_spec.task:
+      device_dict[d_spec.job].append([])
+
+    device_dict[d_spec.job][d_spec.task].append(d)
+
+  return device_dict
+
+
+def _is_gpu_device(device):
+  return tf_device.DeviceSpec.from_string(device).device_type == "GPU"
+
+
+def _infer_num_gpus_per_worker(devices):
+  """Infers the number of GPUs on each worker.
+
+  Currently to make multi-worker cross device ops work, we need all workers to
+  have the same number of GPUs.
+
+  Args:
+    devices: a list of device strings, can be either local devices or remote
+      devices.
+
+  Returns:
+    number of GPUs per worker.
+
+  Raises:
+    ValueError if workers have different number of GPUs or GPU indices are not
+    consecutive and starting from 0.
+  """
+  if _is_device_list_single_worker(devices):
+    return sum(1 for d in devices if _is_gpu_device(d))
+  else:
+    device_dict = _group_device_list(devices)
+    num_gpus = None
+    for _, devices_in_task in device_dict.items():
+      for device_in_task in devices_in_task:
+        if num_gpus is None:
+          num_gpus = sum(1 for d in device_in_task if _is_gpu_device(d))
+
+        # Verify other workers have the same number of GPUs.
+        elif num_gpus != sum(1 for d in device_in_task if _is_gpu_device(d)):
+          raise ValueError("All workers should have the same number of GPUs.")
+
+        for d in device_in_task:
+          d_spec = tf_device.DeviceSpec.from_string(d)
+          if (d_spec.device_type == "GPU" and
+              d_spec.device_index >= num_gpus):
+            raise ValueError("GPU `device_index` on a worker should be "
+                             "consecutive and start from 0.")
+    return num_gpus
+
+
+def all_local_devices(num_gpus=None):
+  devices = config.list_logical_devices("GPU")
+  if num_gpus is not None:
+    devices = devices[:num_gpus]
+  return devices or config.list_logical_devices("CPU")
+
+
+def all_devices():
+  devices = []
+  tfconfig = tfconfig_cluster_resolver.TFConfigClusterResolver()
+  if tfconfig.cluster_spec().as_dict():
+    devices = _cluster_spec_to_device_list(tfconfig.cluster_spec(),
+                                           context.num_gpus())
+  return devices if devices else all_local_devices()
+
+
+@tf_export("distribute.MirroredStrategy", v1=[])  # pylint: disable=g-classes-have-attributes
+class MirroredStrategy(distribute_lib.Strategy):
+  """Synchronous training across multiple replicas on one machine.
+
+  This strategy is typically used for training on one
+  machine with multiple GPUs. For TPUs, use
+  `tf.distribute.TPUStrategy`. To use `MirroredStrategy` with multiple workers,
+  please refer to `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
+
+  For example, a variable created under a `MirroredStrategy` is a
+  `MirroredVariable`. If no devices are specified in the constructor argument of
+  the strategy then it will use all the available GPUs. If no GPUs are found, it
+  will use the available CPUs. Note that TensorFlow treats all CPUs on a
+  machine as a single device, and uses threads internally for parallelism.
+
+  >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+  >>> with strategy.scope():
+  ...   x = tf.Variable(1.)
+  >>> x
+  MirroredVariable:{
+    0: <tf.Variable ... shape=() dtype=float32, numpy=1.0>,
+    1: <tf.Variable ... shape=() dtype=float32, numpy=1.0>
+  }
+
+  While using distribution strategies, all the variable creation should be done
+  within the strategy's scope. This will replicate the variables across all the
+  replicas and keep them in sync using an all-reduce algorithm.
+
+  Variables created inside a `MirroredStrategy` which is wrapped with a
+  `tf.function` are still `MirroredVariables`.
+
+  >>> x = []
+  >>> @tf.function  # Wrap the function with tf.function.
+  ... def create_variable():
+  ...   if not x:
+  ...     x.append(tf.Variable(1.))
+  ...   return x[0]
+  >>> strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
+  >>> with strategy.scope():
+  ...   _ = create_variable()
+  ...   print(x[0])
+  MirroredVariable:{
+    0: <tf.Variable ... shape=() dtype=float32, numpy=1.0>,
+    1: <tf.Variable ... shape=() dtype=float32, numpy=1.0>
+  }
+
+  `experimental_distribute_dataset` can be used to distribute the dataset across
+  the replicas when writing your own training loop. If you are using `.fit` and
+  `.compile` methods available in `tf.keras`, then `tf.keras` will handle the
+  distribution for you.
+
+  For example:
+
+  ```python
+  my_strategy = tf.distribute.MirroredStrategy()
+  with my_strategy.scope():
+    @tf.function
+    def distribute_train_epoch(dataset):
+      def replica_fn(input):
+        # process input and return result
+        return result
+
+      total_result = 0
+      for x in dataset:
+        per_replica_result = my_strategy.run(replica_fn, args=(x,))
+        total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
+                                           per_replica_result, axis=None)
+      return total_result
+
+    dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
+    for _ in range(EPOCHS):
+      train_result = distribute_train_epoch(dist_dataset)
+  ```
+
+  Args:
+    devices: a list of device strings such as `['/gpu:0', '/gpu:1']`.  If
+      `None`, all available GPUs are used. If no GPUs are found, CPU is used.
+    cross_device_ops: optional, a descendant of `CrossDeviceOps`. If this is not
+      set, `NcclAllReduce()` will be used by default.  One would customize this
+      if NCCL isn't available or if a special implementation that exploits
+      the particular hardware is available.
+  """
+
+  # Only set this in tests.
+  _collective_key_base = 0
+
+  def __init__(self, devices=None, cross_device_ops=None):
+    extended = MirroredExtended(
+        self, devices=devices, cross_device_ops=cross_device_ops)
+    super(MirroredStrategy, self).__init__(extended)
+    distribute_lib.distribution_strategy_gauge.get_cell("V2").set(
+        "MirroredStrategy")
+
+
+@tf_export(v1=["distribute.MirroredStrategy"])
+class MirroredStrategyV1(distribute_lib.StrategyV1):  # pylint: disable=g-missing-docstring
+
+  __doc__ = MirroredStrategy.__doc__
+
+  # Only set this in tests.
+  _collective_key_base = 0
+
+  def __init__(self, devices=None, cross_device_ops=None):
+    extended = MirroredExtended(
+        self, devices=devices, cross_device_ops=cross_device_ops)
+    super(MirroredStrategyV1, self).__init__(extended)
+    distribute_lib.distribution_strategy_gauge.get_cell("V1").set(
+        "MirroredStrategy")
+
+
+# TODO(josh11b): Switch to V2 when we no longer need to support tf.compat.v1.
+class MirroredExtended(distribute_lib.StrategyExtendedV1):
+  """Implementation of MirroredStrategy."""
+
+  def __init__(self, container_strategy, devices=None, cross_device_ops=None):
+    super(MirroredExtended, self).__init__(container_strategy)
+    if context.executing_eagerly():
+      if devices and not _is_device_list_single_worker(devices):
+        raise RuntimeError("In-graph multi-worker training with "
+                           "`MirroredStrategy` is not supported in eager mode.")
+      else:
+        if (
+            tfconfig_cluster_resolver.TFConfigClusterResolver()
+            .cluster_spec()
+            .as_dict()
+        ):
+          # if you are executing in eager mode, only the single machine code
+          # path is supported.
+          logging.info("Initializing local devices since in-graph multi-worker "
+                       "training with `MirroredStrategy` is not supported in "
+                       "eager mode. TF_CONFIG will be ignored when "
+                       "when initializing `MirroredStrategy`.")
+        devices = devices or all_local_devices()
+    else:
+      devices = devices or all_devices()
+
+    assert devices, ("Got an empty `devices` list and unable to recognize "
+                     "any local devices.")
+
+    self._collective_key_base = container_strategy._collective_key_base
+    self._communication_options = collective_util.Options(
+        implementation=collective_util.CommunicationImplementation.NCCL)
+    self._cross_device_ops = cross_device_ops
+    self._initialize_strategy(devices)
+
+    # TODO(b/128995245): Enable last partial batch support in graph mode.
+    if ops.executing_eagerly_outside_functions():
+      self.experimental_enable_get_next_as_optional = True
+
+    # Flag to turn on VariablePolicy.
+    self._use_var_policy = False
+
+  def _use_merge_call(self):
+    # We currently only disable merge_call when XLA is used to compile the `fn`
+    # passed to `strategy.run` and all devices are GPU.
+    return not control_flow_util.GraphOrParentsInXlaContext(
+        ops.get_default_graph()) or not all(
+            [_is_gpu_device(d) for d in self._devices])
+
+  def _initialize_strategy(self, devices):
+    # The _initialize_strategy method is intended to be used by distribute
+    # coordinator as well.
+    assert devices, "Must specify at least one device."
+    devices = tuple(device_util.resolve(d) for d in devices)
+    assert len(set(devices)) == len(devices), (
+        "No duplicates allowed in `devices` argument: %s" % (devices,))
+
+    self._initialize_single_worker(devices)
+
+    self._collective_ops = self._make_collective_ops_with_fallbacks()
+    # If cross_device_ops is not provided, set it to collective op by default.
+    if not self._cross_device_ops:
+      self._cross_device_ops = self._collective_ops
+
+  def _make_collective_ops_with_fallbacks(self):
+    self._collective_keys = cross_device_utils.CollectiveKeys(
+        group_key_start=1 + self._collective_key_base)
+
+    if not ops.executing_eagerly_outside_functions() and any(
+        "gpu" not in d.lower() for d in self._devices):
+      # In TF1/Session, fall back to ReductionToOneDevice() if there are
+      # non-GPU devices or virtual GPUs are used.
+      return cross_device_ops_lib.ReductionToOneDevice()
+
+    # Use ReductionToOneDevice() if mixed devices are used.
+    if any("cpu" in d.lower() for d in self._devices) and any(
+        "gpu" in d.lower() for d in self._devices):
+      return cross_device_ops_lib.ReductionToOneDevice()
+
+    if all("cpu" in d.lower() for d in self._devices):
+      # Use RING collective ops if all devices are CPU.
+      self._communication_options = collective_util.Options(
+          implementation=collective_util.CommunicationImplementation.RING)
+
+    else:
+      physical_gpus = context.context().list_physical_devices(device_type="GPU")
+      logical_gpus = context.context().list_logical_devices(device_type="GPU")
+      # Use RING collective ops if virtual devices are used.
+      if len(physical_gpus) < len(logical_gpus):
+        self._communication_options = collective_util.Options(
+            implementation=collective_util.CommunicationImplementation.RING)
+
+    # If all devices are physical GPU, use NCCL implementation.
+    return cross_device_ops_lib.CollectiveAllReduce(
+        devices=self._devices,
+        group_size=len(self._devices),
+        options=self._communication_options,
+        collective_keys=self._collective_keys)
+
+  def _initialize_single_worker(self, devices):
+    """Initializes the object for single-worker training."""
+    self._devices = tuple(device_util.canonicalize(d) for d in devices)
+    self._input_workers_devices = (
+        (device_util.canonicalize("/device:CPU:0", devices[0]), devices),)
+
+    self._host_input_device = numpy_dataset.SingleDevice(
+        self._input_workers_devices[0][0])
+    device_spec = tf_device.DeviceSpec.from_string(
+        self._input_workers_devices[0][0])
+    # Ensures when we enter strategy.scope() we use the correct default device
+    if device_spec.job is not None and device_spec.job != "localhost":
+      self._default_device = "/job:%s/replica:%d/task:%d" % (
+          device_spec.job, device_spec.replica, device_spec.task)
+
+    logging.info("Using MirroredStrategy with devices %r", devices)
+
+  def _initialize_multi_worker(self, devices):
+    """Initializes the object for multi-worker training."""
+    device_dict = _group_device_list(devices)
+    workers = []
+    worker_devices = []
+    for job in ("chief", "worker"):
+      for task in range(len(device_dict.get(job, []))):
+        worker = "/job:%s/task:%d" % (job, task)
+        workers.append(worker)
+        worker_devices.append((worker, device_dict[job][task]))
+
+    # Setting `_default_device` will add a device scope in the
+    # distribution.scope. We set the default device to the first worker. When
+    # users specify device under distribution.scope by
+    #   with tf.device("/cpu:0"):
+    #     ...
+    # their ops will end up on the cpu device of its first worker, e.g.
+    # "/job:worker/task:0/device:CPU:0". Note this is not used in replica mode.
+    self._default_device = workers[0]
+    self._host_input_device = numpy_dataset.SingleDevice(workers[0])
+
+    self._devices = tuple(devices)
+    self._input_workers_devices = worker_devices
+    self._is_multi_worker_training = True
+
+    if len(workers) > 1:
+      # Grandfather usage in the legacy tests if they're configured properly.
+      if (not isinstance(self._cross_device_ops,
+                         cross_device_ops_lib.ReductionToOneDevice) or
+          self._cross_device_ops._num_between_graph_workers > 1):  # pylint: disable=protected-access
+        raise ValueError(
+            "In-graph multi-worker training with `MirroredStrategy` is not "
+            "supported.")
+      self._inferred_cross_device_ops = self._cross_device_ops
+    else:
+      # TODO(yuefengz): make `select_cross_device_ops` work with device strings
+      # containing job names.
+      self._inferred_cross_device_ops = cross_device_ops_lib.NcclAllReduce()
+
+    logging.info("Using MirroredStrategy with remote devices %r", devices)
+
+  def _input_workers_with_options(self, options=None):
+    if not options:
+      return input_lib.InputWorkers(self._input_workers_devices)
+    if (options.experimental_replication_mode ==
+        distribute_lib.InputReplicationMode.PER_REPLICA):
+      if options.experimental_place_dataset_on_device:
+        self._input_workers_devices = (
+            tuple(
+                (device_util.canonicalize(d, d), (d,)) for d in self._devices))
+      else:
+        self._input_workers_devices = (
+            tuple((device_util.canonicalize("/device:CPU:0", d), (d,))
+                  for d in self._devices))
+      return input_lib.InputWorkers(self._input_workers_devices)
+    else:
+      if not options.experimental_fetch_to_device:
+        return input_lib.InputWorkers([
+            (host_device, (host_device,) * len(compute_devices))
+            for host_device, compute_devices in self._input_workers_devices
+        ])
+      else:
+        return input_lib.InputWorkers(self._input_workers_devices)
+
+  @property
+  def _input_workers(self):
+    return self._input_workers_with_options()
+
+  def _get_variable_creator_initial_value(self,
+                                          replica_id,
+                                          device,
+                                          primary_var,
+                                          **kwargs):
+    """Return the initial value for variables on a replica."""
+    if replica_id == 0:
+      return kwargs["initial_value"]
+    else:
+      assert primary_var is not None
+      assert device is not None
+      assert kwargs is not None
+
+      def initial_value_fn():
+        if context.executing_eagerly() or ops.inside_function():
+          init_value = primary_var.value()
+          return array_ops.identity(init_value)
+        else:
+          with ops.device(device):
+            init_value = primary_var.initial_value
+            return array_ops.identity(init_value)
+
+      return initial_value_fn
+
+  def _create_variable(self, next_creator, **kwargs):
+    """Create a mirrored variable. See `DistributionStrategy.scope`."""
+    colocate_with = kwargs.pop("colocate_with", None)
+    if colocate_with is None:
+      devices = self._devices
+    elif isinstance(colocate_with, numpy_dataset.SingleDevice):
+      with ops.device(colocate_with.device):
+        return next_creator(**kwargs)
+    else:
+      devices = colocate_with._devices  # pylint: disable=protected-access
+
+    def _real_mirrored_creator(**kwargs):  # pylint: disable=g-missing-docstring
+      value_list = []
+      for i, d in enumerate(devices):
+        with ops.device(d):
+          kwargs["initial_value"] = self._get_variable_creator_initial_value(
+              replica_id=i,
+              device=d,
+              primary_var=value_list[0] if value_list else None,
+              **kwargs)
+          if i > 0:
+            # Give replicas meaningful distinct names:
+            var0name = value_list[0].name.split(":")[0]
+            # We append a / to variable names created on replicas with id > 0 to
+            # ensure that we ignore the name scope and instead use the given
+            # name as the absolute name of the variable.
+            kwargs["name"] = "%s/replica_%d/" % (var0name, i)
+          with context.device_policy(context.DEVICE_PLACEMENT_SILENT):
+            # Don't record operations (e.g. other variable reads) during
+            # variable creation.
+            with record.stop_recording():
+              v = next_creator(**kwargs)
+          assert not isinstance(v, values.DistributedVariable)
+          value_list.append(v)
+      return value_list
+
+    return distribute_utils.create_mirrored_variable(
+        self._container_strategy(), _real_mirrored_creator,
+        distribute_utils.VARIABLE_CLASS_MAPPING,
+        distribute_utils.VARIABLE_POLICY_MAPPING, **kwargs)
+
+  def _validate_colocate_with_variable(self, colocate_with_variable):
+    distribute_utils.validate_colocate_distributed_variable(
+        colocate_with_variable, self)
+
+  def _make_dataset_iterator(self, dataset):
+    return input_lib_v1.DatasetIterator(
+        dataset,
+        self._input_workers,
+        self._container_strategy(),
+        num_replicas_in_sync=self._num_replicas_in_sync)
+
+  def _make_input_fn_iterator(
+      self,
+      input_fn,
+      replication_mode=distribute_lib.InputReplicationMode.PER_WORKER):
+    input_contexts = []
+    num_workers = self._input_workers.num_workers
+    for i in range(num_workers):
+      input_contexts.append(distribute_lib.InputContext(
+          num_input_pipelines=num_workers,
+          input_pipeline_id=i,
+          num_replicas_in_sync=self._num_replicas_in_sync))
+    return input_lib_v1.InputFunctionIterator(input_fn, self._input_workers,
+                                              input_contexts,
+                                              self._container_strategy())
+
+  def _experimental_distribute_dataset(self, dataset, options):
+    if (options and options.experimental_replication_mode ==
+        distribute_lib.InputReplicationMode.PER_REPLICA):
+      raise NotImplementedError(
+          "InputReplicationMode.PER_REPLICA "
+          "is only supported in "
+          "`distribute_datasets_from_function`."
+      )
+    return input_util.get_distributed_dataset(
+        dataset,
+        self._input_workers_with_options(options),
+        self._container_strategy(),
+        num_replicas_in_sync=self._num_replicas_in_sync,
+        options=options)
+
+  def _experimental_make_numpy_dataset(self, numpy_input, session):
+    return numpy_dataset.one_host_numpy_dataset(
+        numpy_input, self._host_input_device, session)
+
+  def _distribute_datasets_from_function(self, dataset_fn, options):
+    input_workers = self._input_workers_with_options(options)
+    input_contexts = []
+    num_workers = input_workers.num_workers
+    for i in range(num_workers):
+      input_contexts.append(distribute_lib.InputContext(
+          num_input_pipelines=num_workers,
+          input_pipeline_id=i,
+          num_replicas_in_sync=self._num_replicas_in_sync))
+
+    return input_util.get_distributed_datasets_from_function(
+        dataset_fn, input_workers, input_contexts, self._container_strategy(),
+        options)
+
+  def _experimental_distribute_values_from_function(self, value_fn):
+    per_replica_values = []
+    for replica_id in range(self._num_replicas_in_sync):
+      per_replica_values.append(value_fn(
+          distribute_lib.ValueContext(replica_id,
+                                      self._num_replicas_in_sync)))
+    return distribute_utils.regroup(per_replica_values, always_wrap=True)
+
+  # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed.
+  def _experimental_run_steps_on_iterator(self, fn, iterator, iterations,
+                                          initial_loop_values=None):
+    if initial_loop_values is None:
+      initial_loop_values = {}
+    initial_loop_values = nest.flatten(initial_loop_values)
+
+    ctx = input_lib.MultiStepContext()
+    def body(i, *args):
+      """A wrapper around `fn` to create the while loop body."""
+      del args
+      fn_result = fn(ctx, iterator.get_next())
+      for (name, output) in ctx.last_step_outputs.items():
+        # Convert all outputs to tensors, potentially from `DistributedValues`.
+        ctx.last_step_outputs[name] = self._local_results(output)
+      flat_last_step_outputs = nest.flatten(ctx.last_step_outputs)
+      with ops.control_dependencies([fn_result]):
+        return [i + 1] + flat_last_step_outputs
+
+    # We capture the control_flow_context at this point, before we run `fn`
+    # inside a while_loop. This is useful in cases where we might need to exit
+    # these contexts and get back to the outer context to do some things, for
+    # e.g. create an op which should be evaluated only once at the end of the
+    # loop on the host. One such usage is in creating metrics' value op.
+    self._outer_control_flow_context = (
+        ops.get_default_graph()._get_control_flow_context())  # pylint: disable=protected-access
+
+    cond = lambda i, *args: i < iterations
+    i = constant_op.constant(0)
+    loop_result = while_loop.while_loop(
+        cond,
+        body, [i] + initial_loop_values,
+        name="",
+        parallel_iterations=1,
+        back_prop=False,
+        swap_memory=False,
+        return_same_structure=True)
+    del self._outer_control_flow_context
+
+    ctx.run_op = control_flow_ops.group(loop_result)
+
+    # Convert the last_step_outputs from a list to the original dict structure
+    # of last_step_outputs.
+    last_step_tensor_outputs = loop_result[1:]
+    last_step_tensor_outputs_dict = nest.pack_sequence_as(
+        ctx.last_step_outputs, last_step_tensor_outputs)
+
+    for name, reduce_op in ctx._last_step_outputs_reduce_ops.items():  # pylint: disable=protected-access
+      output = last_step_tensor_outputs_dict[name]
+      # For outputs that have already been reduced, wrap them in a Mirrored
+      # container, else in a PerReplica container.
+      if reduce_op is None:
+        last_step_tensor_outputs_dict[name] = distribute_utils.regroup(output)
+      else:
+        assert len(output) == 1
+        last_step_tensor_outputs_dict[name] = output[0]
+
+    ctx._set_last_step_outputs(last_step_tensor_outputs_dict)  # pylint: disable=protected-access
+    return ctx
+
+  def _broadcast_to(self, tensor, destinations):
+    # This is both a fast path for Python constants, and a way to delay
+    # converting Python values to a tensor until we know what type it
+    # should be converted to. Otherwise we have trouble with:
+    #   global_step.assign_add(1)
+    # since the `1` gets broadcast as an int32 but global_step is int64.
+    if isinstance(tensor, (float, int)):
+      return tensor
+    # TODO(josh11b): In eager mode, use one thread per device, or async mode.
+    if not destinations:
+      # TODO(josh11b): Use current logical device instead of 0 here.
+      destinations = self._devices
+    return self._get_cross_device_ops(tensor).broadcast(tensor, destinations)
+
+  def _call_for_each_replica(self, fn, args, kwargs):
+    return mirrored_run.call_for_each_replica(
+        self._container_strategy(), fn, args, kwargs)
+
+  def _configure(self,
+                 session_config=None,
+                 cluster_spec=None,
+                 task_type=None,
+                 task_id=None):
+    del task_type, task_id
+
+    if session_config:
+      session_config.CopyFrom(self._update_config_proto(session_config))
+
+    if cluster_spec:
+      # TODO(yuefengz): remove the following code once cluster_resolver is
+      # added.
+      num_gpus_per_worker = _infer_num_gpus_per_worker(self._devices)
+      multi_worker_devices = _cluster_spec_to_device_list(
+          cluster_spec, num_gpus_per_worker)
+      self._initialize_multi_worker(multi_worker_devices)
+
+  def _update_config_proto(self, config_proto):
+    updated_config = copy.deepcopy(config_proto)
+    updated_config.isolate_session_state = True
+    return updated_config
+
+  def _get_cross_device_ops(self, value):
+    # Always use CollectiveAllReduce when XLA is enabled, since other cross
+    # device ops don't have as good support on XLA.
+    if not self._use_merge_call():
+      if not isinstance(self._cross_device_ops,
+                        cross_device_ops_lib.CollectiveAllReduce):
+        logging.warning(
+            "Under XLA context, MirroredStrategy uses CollectiveAllReduce op. "
+            "Although %r is provided to initialize MirroredStrategy, it is "
+            "ignored in XLA. Please use CollectiveAllReduce(or default option) "
+            "in the future, since other cross device ops are not well "
+            "supported on XLA.", self._cross_device_ops
+        )
+      return self._collective_ops
+
+    if isinstance(value, values.DistributedValues):
+      value_int32 = True in {
+          dtypes.as_dtype(v.dtype) == dtypes.int32 for v in value.values
+      }
+    else:
+      value_int32 = dtypes.as_dtype(value.dtype) == dtypes.int32
+
+    if value_int32:
+      return cross_device_ops_lib.ReductionToOneDevice()
+    else:
+      return self._cross_device_ops
+
+  def _gather_to_implementation(self, value, destinations, axis, options):
+    if not isinstance(value, values.DistributedValues):
+      # ReductionToOneDevice._gather accepts DistributedValues only.
+      return value
+    return self._get_cross_device_ops(value)._gather(  # pylint: disable=protected-access
+        value,
+        destinations=destinations,
+        axis=axis,
+        options=self._communication_options.merge(options))
+
+  def _reduce_to(self, reduce_op, value, destinations, options):
+    if (distribute_utils.is_mirrored(value) and
+        reduce_op == reduce_util.ReduceOp.MEAN):
+      return value
+    assert not distribute_utils.is_mirrored(value)
+    def get_values(value):
+      if not isinstance(value, values.DistributedValues):
+        # This function handles reducing values that are not PerReplica or
+        # Mirrored values. For example, the same value could be present on all
+        # replicas in which case `value` would be a single value or value could
+        # be 0.
+        return cross_device_ops_lib.reduce_non_distributed_value(
+            reduce_op, value, destinations, self._num_replicas_in_sync)
+
+      if self._use_merge_call() and (
+          not cross_device_ops_lib._devices_match(value, destinations) or  # pylint: disable=protected-access
+          any("cpu" in d.lower()
+              for d in cross_device_ops_lib.get_devices_from(destinations))):
+        return cross_device_ops_lib.ReductionToOneDevice().reduce(
+            reduce_op, value, destinations)
+      return self._get_cross_device_ops(value).reduce(
+          reduce_op,
+          value,
+          destinations=destinations,
+          options=self._communication_options.merge(options))
+
+    return nest.map_structure(get_values, value)
+
+  def _batch_reduce_to(self, reduce_op, value_destination_pairs, options):
+    cross_device_ops = None
+    for value, _ in value_destination_pairs:
+      if cross_device_ops is None:
+        cross_device_ops = self._get_cross_device_ops(value)
+      elif cross_device_ops is not self._get_cross_device_ops(value):
+        raise ValueError("Inputs to batch_reduce_to must be either all on "
+                         "the host or all on the compute devices.")
+    return cross_device_ops.batch_reduce(
+        reduce_op,
+        value_destination_pairs,
+        options=self._communication_options.merge(options))
+
+  def _update(self, var, fn, args, kwargs, group):
+    # TODO(josh11b): In eager mode, use one thread per device.
+    assert isinstance(var, values.DistributedVariable)
+    updates = []
+    for i, v in enumerate(var.values):
+      name = "update_%d" % i
+      with ops.device(v.device), \
+           distribute_lib.UpdateContext(i), \
+           ops.name_scope(name):
+        # If args and kwargs are not mirrored, the value is returned as is.
+        updates.append(
+            fn(v, *distribute_utils.select_replica(i, args),
+               **distribute_utils.select_replica(i, kwargs)))
+    return distribute_utils.update_regroup(self, updates, group)
+
+  def _replica_ctx_all_reduce(self, reduce_op, value, options=None):
+    """Implements `StrategyExtendedV2._replica_ctx_all_reduce`."""
+    # This implementation avoids using `merge_call` and just launches collective
+    # ops in one replica.
+    if options is None:
+      options = collective_util.Options()
+
+    if context.executing_eagerly() or (
+        not tf2.enabled()) or self._use_merge_call():
+      # In eager mode, falls back to the default implementation that uses
+      # `merge_call`. Replica functions are running sequentially in eager mode,
+      # and due to the blocking nature of collective ops, execution will hang if
+      # collective ops are to be launched sequentially.
+      return super()._replica_ctx_all_reduce(reduce_op, value, options)
+
+    replica_context = distribute_lib.get_replica_context()
+    assert replica_context, (
+        "`StrategyExtended._replica_ctx_all_reduce` must be called in a "
+        "replica context")
+    return self._get_cross_device_ops(value)._all_reduce(  # pylint: disable=protected-access
+        reduce_op,
+        value,
+        replica_context._replica_id,  # pylint: disable=protected-access
+        options)
+
+  def _replica_ctx_update(self, var, fn, args, kwargs, group):
+    if self._use_merge_call():
+      return super()._replica_ctx_update(var, fn, args, kwargs, group)
+
+    replica_context = distribute_lib.get_replica_context()
+    assert replica_context
+    replica_id = values_util.get_current_replica_id_as_int()
+    name = "update_%d" % replica_id
+
+    if isinstance(var, values.DistributedVariable):
+      var = var._get_replica(replica_id)  # pylint: disable=protected-access
+
+    with ops.device(var.device), ops.name_scope(name):
+      result = fn(var, *args, **kwargs)
+    return result
+
+  def _update_non_slot(self, colocate_with, fn, args, kwargs, group):
+    assert isinstance(colocate_with, tuple)
+    # TODO(josh11b): In eager mode, use one thread per device.
+    updates = []
+    for i, d in enumerate(colocate_with):
+      name = "update_%d" % i
+      with ops.device(d), distribute_lib.UpdateContext(i), ops.name_scope(name):
+        updates.append(
+            fn(*distribute_utils.select_replica(i, args),
+               **distribute_utils.select_replica(i, kwargs)))
+    return distribute_utils.update_regroup(self, updates, group)
+
+  def read_var(self, replica_local_var):
+    """Read the aggregate value of a replica-local variable."""
+    # pylint: disable=protected-access
+    if distribute_utils.is_sync_on_read(replica_local_var):
+      return replica_local_var._get_cross_replica()
+    assert distribute_utils.is_mirrored(replica_local_var)
+    return array_ops.identity(replica_local_var._get())
+    # pylint: enable=protected-access
+
+  def value_container(self, val):
+    return distribute_utils.value_container(val)
+
+  @property
+  def _num_replicas_in_sync(self):
+    return len(self._devices)
+
+  @property
+  def worker_devices(self):
+    return self._devices
+
+  @property
+  def worker_devices_by_replica(self):
+    return [[d] for d in self._devices]
+
+  @property
+  def parameter_devices(self):
+    return self.worker_devices
+
+  @property
+  def experimental_between_graph(self):
+    return False
+
+  @property
+  def experimental_should_init(self):
+    return True
+
+  @property
+  def should_checkpoint(self):
+    return True
+
+  @property
+  def should_save_summary(self):
+    return True
+
+  def non_slot_devices(self, var_list):
+    del var_list
+    # TODO(josh11b): Should this be the last logical device instead?
+    return self._devices
+
+  # TODO(priyag): Delete this once all strategies use global batch size.
+  @property
+  def _global_batch_size(self):
+    """`make_dataset_iterator` and `make_numpy_iterator` use global batch size.
+
+    `make_input_fn_iterator` assumes per-replica batching.
+
+    Returns:
+      Boolean.
+    """
+    return True
+
+  def _in_multi_worker_mode(self):
+    """Whether this strategy indicates working in multi-worker settings."""
+    return False
+
+  def _get_local_replica_id(self, replica_id_in_sync_group):
+    return replica_id_in_sync_group
+
+  def _get_replica_id_in_sync_group(self, replica_id):
+    return replica_id