# 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 remote execution."""

import os
import random
import time

from absl.testing import parameterized
import numpy as np
import portpicker

from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute.cluster_resolver.cluster_resolver import SimpleClusterResolver
from tensorflow.python.eager import cancellation
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import executor
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
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_ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import data_flow_ops
from tensorflow.python.ops import functional_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.training import server_lib
from tensorflow.python.training.server_lib import ClusterSpec
from tensorflow.python.util import compat


class SingleWorkerTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super(SingleWorkerTest, self).setUp()

    workers, _ = test_util.create_local_cluster(1, 0)
    remote.connect_to_remote_host(workers[0].target)

  def tearDown(self):
    super(SingleWorkerTest, self).tearDown()

    # Clear the current device scope to avoid polluting other test cases.
    ops.device(None).__enter__()
    # Reset the context to avoid polluting other test cases.
    context._reset_context()

  def testMultiDeviceFunctionBasic(self):

    @def_function.function
    def basic(i):
      with ops.device('/job:localhost/replica:0/task:0/cpu:0'):
        a = constant_op.constant([2]) + i
      with ops.device('/job:worker/replica:0/task:0/cpu:0'):
        b = constant_op.constant([1])

      return a + b

    self.assertAllEqual(basic(constant_op.constant([2])).numpy(), [5])
    self.assertAllEqual(basic(constant_op.constant([1])).numpy(), [4])

  def testMultiDeviceFunctionVariable(self):
    with ops.device('/job:worker/replica:0/task:0/cpu:0'):
      variable_b = variables.Variable(1)

    # Add a sync point to avoid the out-of-order issue of eager async execution
    # (b/155789951).
    context.async_wait()

    @def_function.function
    def with_variable(i):
      return i + variable_b

    self.assertAllEqual(with_variable(constant_op.constant([2])).numpy(), [3])

  def testMultiDeviceFunctionRemoteOutput(self):
    with ops.device('/job:worker/replica:0/task:0/cpu:0'):
      variable_b = variables.Variable(1)

    @def_function.function
    def remote_output(i):
      with ops.device('/job:worker/replica:0/task:0/cpu:0'):
        c = variable_b + 1
      return i + variable_b, c

    rets = remote_output(constant_op.constant([1]))
    self.assertAllEqual(rets[0].numpy(), [2])
    self.assertAllEqual(rets[1].numpy(), 2)
    self.assertEqual(rets[0].backing_device,
                     '/job:localhost/replica:0/task:0/device:CPU:0')
    self.assertEqual(rets[1].backing_device,
                     '/job:worker/replica:0/task:0/device:CPU:0')

  def testStreaming(self):
    """A mini stress test for streaming - issuing many RPCs back to back."""
    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
      x = array_ops.ones([2, 2])
      y = array_ops.zeros([2, 2])
      num_iters = 200
      for _ in range(num_iters):
        y = x + y
        # Ask for y's shape after every 10 additions on average.
        # This exercises waiting for remote shape logic in TensorHandle.
        if random.randint(1, 10) == 1:
          _ = y.shape
    np.testing.assert_array_equal(
        [[num_iters, num_iters], [num_iters, num_iters]], y.numpy())

  def testTwoExecutors(self):
    # Run an op on the main executor that by default uses StreamingEnqueue to
    # schedule the op to run on the remote async executor. This op produces an
    # error, i.e., division by zero, but will not be immediately caught due to
    # streaming enqueue.
    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
      a = constant_op.constant(3)
      b = constant_op.constant(0)
      math_ops.div(a, b)

    # Run another op using another executor that disables streaming enqueue,
    # which would run the op using the tf_compute thread pool in the remote
    # worker. Since the op is not run in the same remotes async executor, it
    # will not carry back that error produced by the op above, even though this
    # op is executed synchronously.
    with context.executor_scope(
        executor.new_executor(
            enable_async=False, enable_streaming_enqueue=False)):
      with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
        c = constant_op.constant(4)
        d = constant_op.constant(2)
        self.assertEqual(math_ops.div(c, d).numpy(), 2)

    # Sync on the context to force to catch the error produced by the first op.
    with self.assertRaises(errors.InvalidArgumentError) as cm:
      context.async_wait()
    self.assertIn('division by zero', cm.exception.message)

  def testShapeError_OpByOp(self):
    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
      x = array_ops.ones([2, 3])
      y = array_ops.zeros([2, 2])
      with self.assertRaises(errors.InvalidArgumentError) as cm:
        math_ops.matmul(x, y)

    self.assertIn('Dimensions must be equal', cm.exception.message)

  def testShapeError_Function(self):

    @def_function.function
    def matmul_func(x, y):
      return math_ops.matmul(x, y)

    x = array_ops.ones([2, 3])
    y = array_ops.zeros([2, 2])

    with ops.device('job:worker/replica:0/task:0/device:CPU:0'):
      with self.assertRaises(ValueError) as cm:
        matmul_func(x, y)

    self.assertIn('Dimensions must be equal', cm.exception.args[0])

  def testClientVarible(self):
    var = variables.Variable(initial_value=0)

    @def_function.function
    def func():
      with ops.device('/job:localhost/task:0'):
        read = var.read_value()
      return read + 1

    with ops.device('/job:worker/task:0'):
      self.assertAllEqual(func(), 1)

  def testRemoteCall(self):

    @def_function.function(
        input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
    def _remote_fn(x):
      return constant_op.constant(1) + x

    remote_fn = _remote_fn.get_concrete_function()

    @def_function.function
    def func(x):
      return functional_ops.remote_call(
          args=[x],
          Tout=[dtypes.int32],
          f=remote_fn,
          target='/job:worker/task:0')

    with ops.device('/job:localhost/task:0'):
      self.assertAllEqual(func(constant_op.constant(1)), [2])

  def testOperationTimeout(self):
    context._reset_context()
    context.context().operation_timeout_in_ms = 10
    workers, _ = test_util.create_local_cluster(1, 0)
    remote.connect_to_remote_host(workers[0].target)

    q = data_flow_ops.FIFOQueue(1, dtypes.int32)

    @def_function.function
    def f():
      return q.dequeue()

    with self.assertRaises(errors.DeadlineExceededError):
      with ops.device('/job:worker/replica:0/task:0'):
        f()
      # If streaming RPC is enabled, fetch remote errors before end of execution
      context.async_wait()


class RemoteAsyncTest(test.TestCase):

  def setUp(self):
    super(RemoteAsyncTest, self).setUp()

    workers, _ = test_util.create_local_cluster(1, 0)
    remote.connect_to_remote_host(workers[0].target)

  def tearDown(self):
    super(RemoteAsyncTest, self).tearDown()

    # Reset the context to avoid polluting other test cases.
    context._reset_context()

  def test_out_of_range_with_while_loop(self):

    with ops.device('/job:worker/task:0'):
      dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0])
      dataset = dataset.batch(1, drop_remainder=False)
      iterator = iter(dataset)
      v = variables.Variable(1.0)

    @def_function.function
    def train_step(iterator):
      i = next(iterator)
      v.assign_add(math_ops.reduce_mean(i))

    while True:
      try:
        with ops.device('/job:worker/task:0'):
          train_step(iterator)
      except (errors.OutOfRangeError, errors.InternalError):
        context.async_clear_error()
        break

    self.assertAllEqual(v.numpy(), 4.0)

  def test_out_of_range_with_for_loop(self):

    with ops.device('/job:worker/task:0'):
      dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0])
      dataset = dataset.batch(1, drop_remainder=False)
      iterator = iter(dataset)
      v = variables.Variable(1.0)

    @def_function.function
    def train_step(iterator):
      i = next(iterator)
      v.assign_add(math_ops.reduce_mean(i))

    num_steps = 3
    for i in range(num_steps):
      try:
        with ops.device('/job:worker/task:0'):
          train_step(iterator)
        if i == num_steps - 1:
          context.async_wait()
      except errors.OutOfRangeError:
        context.async_clear_error()
        break

    self.assertAllEqual(v.numpy(), 4.0)

  def test_out_of_range_with_async_scope(self):

    with ops.device('/job:worker/task:0'):
      dataset = dataset_ops.Dataset.from_tensor_slices([1.0, 2.0])
      dataset = dataset.batch(1, drop_remainder=False)
      iterator = iter(dataset)
      v = variables.Variable(1.0)

    @def_function.function
    def train_step(iterator):
      i = next(iterator)
      v.assign_add(math_ops.reduce_mean(i))

    num_steps = 3
    try:
      with context.async_scope():
        for _ in range(num_steps):
          with ops.device('/job:worker/task:0'):
            train_step(iterator)
    except errors.OutOfRangeError:
      context.async_clear_error()

    self.assertAllEqual(v.numpy(), 4.0)


class MultiWorkersTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super(MultiWorkersTest, self).setUp()

    workers, _ = test_util.create_local_cluster(3, 0)
    remote.connect_to_remote_host(
        [workers[0].target, workers[1].target, workers[2].target])

  def tearDown(self):
    super(MultiWorkersTest, self).tearDown()

    # Clear the current device scope to avoid polluting other test cases.
    ops.device(None).__enter__()
    # Reset the context to avoid polluting other test cases.
    context._reset_context()

  def testReturnRemoteArgument(self):

    @def_function.function
    def local_func(i):
      return i

    with ops.device('/job:worker/replica:0/task:0'):
      x = constant_op.constant([2, 1])

    with ops.device('/job:worker/replica:0/task:1'):
      self.assertAllEqual(local_func(x), [2, 1])

  def testMultiDeviceFunctionAmbiguousDevice(self):

    @def_function.function
    def ambiguous_device(i):
      with ops.device('/job:worker'):
        # Multiple worker tasks, thus ambiguous device found error will be
        # raised.
        return i + constant_op.constant([2])

    with self.assertRaises(errors.InvalidArgumentError) as cm:
      ambiguous_device(constant_op.constant([2])).numpy()

    self.assertIn('the output node must match exactly one device',
                  cm.exception.message)

  # Note that the following tests for remote function cancellation only works
  # when non-streaming RPC. We need to disable streaming explicitly and restore
  # this config to its initial value at the end of each test case.
  def testCancelRemoteFunctionBeforeExecution(self):
    remote_async_env_var = 'TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE'
    default_streaming = os.environ.get(remote_async_env_var)
    os.environ[remote_async_env_var] = str(False)

    q = data_flow_ops.FIFOQueue(1, dtypes.int32)

    @def_function.function
    def f():
      return q.dequeue()

    c_mgr = cancellation.CancellationManager()
    cancelable_func = c_mgr.get_cancelable_function(f.get_concrete_function())

    c_mgr.start_cancel()
    with self.assertRaises(errors.CancelledError):
      with ops.device('/job:worker/replica:0/task:1'):
        cancelable_func()

    if default_streaming is None:
      del os.environ[remote_async_env_var]
    else:
      os.environ[remote_async_env_var] = default_streaming

  def testCancelRemoteFunctionDuringExecution(self):
    remote_async_env_var = 'TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE'
    default_streaming = os.environ.get(remote_async_env_var)
    os.environ[remote_async_env_var] = str(False)

    q = data_flow_ops.FIFOQueue(1, dtypes.int32)

    @def_function.function
    def f():
      return q.dequeue()

    c_mgr = cancellation.CancellationManager()
    cancelable_func = c_mgr.get_cancelable_function(f.get_concrete_function())

    def cancel_thread():
      time.sleep(0.5)
      c_mgr.start_cancel()

    t = self.checkedThread(cancel_thread)
    t.start()
    with self.assertRaises(errors.CancelledError):
      with ops.device('/job:worker/replica:0/task:1'):
        cancelable_func()
    t.join()

    if default_streaming is None:
      del os.environ[remote_async_env_var]
    else:
      os.environ[remote_async_env_var] = default_streaming

  def testMultiDeviceFunctionOnLocalDevice(self):
    with ops.device('/job:worker/replica:0/task:1'):
      variable_b = variables.Variable(1.0)

    @def_function.function
    def remote_function(i):
      with ops.device('/job:worker/replica:0/task:0'):
        a = i + variable_b
      c = a + 1.0
      return c

    self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])

  def testMultiDeviceFunctionExecutionOrderingWithPackedInput(self):
    shape = [2]
    with ops.device('/job:worker/replica:0/task:2/device:CPU:0'):
      # Send 20 remote requests to simulate heavy load on worker:2.
      unused_values = []
      for _ in range(20):
        unused_values.append(array_ops.zeros(shape))
      func_input = array_ops.zeros(shape)

    packed_input = ops.pack_eager_tensors([func_input])

    @def_function.function
    def func(packed_input):
      # When worker:2 receives the component function request, packed_input
      # should be ready on worker:2.
      with ops.device('/job:worker/replica:0/task:2/device:CPU:0'):
        ret = packed_input + constant_op.constant(1.0)
      return ret + constant_op.constant(1.0)

    # Run the function on a worker:1
    with ops.device('/job:worker/replica:0/task:1/device:CPU:0'):
      self.assertAllEqual(func(packed_input).numpy(),
                          array_ops.ones(shape).numpy() * 2)

  def testMultiDeviceFunctionWithPackedVariable(self):
    with ops.device('/job:worker/replica:0/task:0/device:CPU:0'):
      var0 = resource_variable_ops.ResourceVariable(1.0)
    with ops.device('/job:worker/replica:0/task:1/device:CPU:0'):
      var1 = resource_variable_ops.ResourceVariable(2.0)

    packed_var = ops.pack_eager_tensors([var0.handle, var1.handle])
    self.assertEqual(packed_var.device,
                     '/job:localhost/replica:0/task:0/device:COMPOSITE:0')
    self.assertEqual(packed_var.backing_device,
                     '/job:localhost/replica:0/task:0/device:COMPOSITE:0')

    @def_function.function
    def add_variables():
      with ops.device('/job:worker/replica:0/task:0/device:CPU:0'):
        read0 = resource_variable_ops.read_variable_op(
            packed_var, dtype=dtypes.float32)
      with ops.device('/job:worker/replica:0/task:1/device:CPU:0'):
        read1 = resource_variable_ops.read_variable_op(
            packed_var, dtype=dtypes.float32)

      return read0 + read1

    # Run the function on a remote device
    with ops.device('/job:worker/replica:0/task:0'):
      self.assertAllEqual(add_variables().numpy(), 3.0)

    # Run the function on a local worker
    self.assertAllEqual(add_variables().numpy(), 3.0)

  def testMultiDeviceFunctionOnRemoteDeviceWithWait(self):
    with ops.device('/job:worker/replica:0/task:1'):
      variable_b = variables.Variable([1.0])

    @def_function.function
    def remote_function(i):
      x = array_ops.ones([1000, 1000])
      for _ in range(1, 1000):
        x = x * x
      variable_b.assign_add(i)
      a = 1.0 + variable_b
      return a

    @def_function.function
    def remote_function2(i):
      variable_b.assign_add(i)
      a = 1.0 + variable_b
      return a

    # Runs first function:
    # - on remote device
    # - needs remote input
    # - is side impacting
    # - runs much slower
    with ops.device('/job:worker/replica:0/task:0'):
      remote_function(constant_op.constant([2.0]))

    # Runs second function:
    # - on remote device
    # - is side impacting
    # There should be a sync point here and the next function will be executed
    # only after the first function has completed.
    with ops.device('/job:worker/replica:0/task:2'):
      self.assertAllEqual(remote_function2(constant_op.constant([3.0])), [7.0])

  def testMultiDeviceFunctionOnRemoteDevice(self):
    with ops.device('/job:worker/replica:0/task:1'):
      variable_b = variables.Variable(1.0)

    @def_function.function
    def remote_function(i):
      with ops.device('/job:worker/replica:0/task:0'):
        a = i + variable_b
      c = a + 1.0
      return c

    with ops.device('/job:worker/replica:0/task:0'):
      self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])

    if test_util.is_gpu_available():
      with ops.device('/job:worker/replica:0/task:0/device:GPU:0'):
        self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])

  def testMultiDeviceFunctionRemoteOutput(self):
    with ops.device('/job:worker/replica:0/task:1/cpu:0'):
      variable_b = variables.Variable(1)

    @def_function.function
    def remote_output(i):
      with ops.device('/job:worker/replica:0/task:1/cpu:0'):
        c = variable_b + 1
      return i + variable_b, c

    with ops.device('/job:worker/replica:0/task:0/cpu:0'):
      rets = remote_output(constant_op.constant([1]))
    self.assertEqual(rets[0].backing_device,
                     '/job:worker/replica:0/task:0/device:CPU:0')
    self.assertEqual(rets[1].backing_device,
                     '/job:worker/replica:0/task:1/device:CPU:0')
    self.assertAllEqual(rets[0].numpy(), [2])
    self.assertAllEqual(rets[1].numpy(), 2)

  def testMultiDeviceWhileLoopOnRemoteDevice(self):
    with ops.device('/job:worker/replica:0/task:1'):
      variable_b = variables.Variable(1.0)

    @def_function.function
    def remote_function(i):

      def body(i, _):
        with ops.device('/job:worker/replica:0/task:0'):
          a = i + variable_b
        return a + 1.0, 1

      return while_loop.while_loop_v2(lambda _, d: d < 1, body, [i, 0])[0]

    with ops.device('/job:worker/replica:0/task:0'):
      self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])

    if test_util.is_gpu_available():
      with ops.device('/job:worker/replica:0/task:0/device:GPU:0'):
        self.assertAllEqual(remote_function(constant_op.constant([1.0])), [3.0])

  def testSimpleParameterServer(self):

    with ops.device('/job:worker/task:2/device:CPU:0'):
      v1 = variables.Variable(initial_value=0)
      v2 = variables.Variable(initial_value=10)

    @def_function.function
    def worker_fn():
      v1.assign_add(1)
      v2.assign_sub(2)
      return v1.read_value() + v2.read_value()

    with ops.device('/job:worker/task:0/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 9)

    with ops.device('/job:worker/task:1/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 8)


_GRPC_PREFIX = 'grpc://'


class MultiJobsTest(test.TestCase, parameterized.TestCase):

  def setUp(self):
    super(MultiJobsTest, self).setUp()

    workers, ps = test_util.create_local_cluster(num_workers=2, num_ps=2)
    cluster = {
        'my_worker': [_strip_prefix(t.target, _GRPC_PREFIX) for t in workers],
        'my_ps': [_strip_prefix(t.target, _GRPC_PREFIX) for t in ps],
    }
    self._cluster = server_lib.ClusterSpec(cluster)
    self._cluster_resolver = SimpleClusterResolver(
        cluster_spec=self._cluster, master=ps[0].target)

  def tearDown(self):
    super(MultiJobsTest, self).tearDown()

    # Clear the current device scope to avoid polluting other test cases.
    ops.device(None).__enter__()
    # Reset the context to avoid polluting other test cases.
    context._reset_context()

  def testMultipleDeviceFoundCheck(self):
    remote.connect_to_cluster(self._cluster)

    @def_function.function
    def func():
      with ops.device('cpu:0'):
        # Multiple CPU:0 devices match would be found, but the CPU:0 from the
        # parent device scope should be picked.
        x = test_ops.device_placement_op()
        y = string_ops.string_upper(x)
        packed_var_0 = array_ops_stack.stack([x, y], 0)
        return packed_var_0

    with ops.device('/job:my_worker/task:1'):
      output = self.evaluate(func())
      self.assertEqual(
          compat.as_bytes('/job:my_worker/replica:0/task:1/device:CPU:0'),
          output[0])
      self.assertIn(compat.as_bytes('/JOB:MY_WORKER'), output[1])
    with ops.device('/job:my_ps/task:1'):
      output = self.evaluate(func())
      self.assertEqual(
          compat.as_bytes('/job:my_ps/replica:0/task:1/device:CPU:0'),
          output[0])
      self.assertIn(compat.as_bytes('/JOB:MY_PS'), output[1])

  def testSimpleParameterServer(self):
    remote.connect_to_cluster(self._cluster)

    with ops.device('/job:my_ps/task:0/device:CPU:0'):
      v1 = variables.Variable(initial_value=0)
      v2 = variables.Variable(initial_value=10)

    @def_function.function
    def worker_fn():
      v1.assign_add(1)
      v2.assign_sub(2)
      return v1.read_value() + v2.read_value()

    with ops.device('/job:my_worker/task:0/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 9)

    with ops.device('/job:my_worker/task:1/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 8)

  def testResetClusterWithDifferentJobNames(self):
    addr = 'localhost:%s' % portpicker.pick_unused_port()
    cluster = server_lib.ClusterSpec({'localhost': [addr]})
    remote.connect_to_cluster(cluster, job_name='localhost')
    with ops.device('/job:localhost/task:0/device:CPU:0'):
      v1 = variables.Variable(initial_value=0)
      v1.assign_add(1)

    # Replace job name from 'localhost' to 'worker' in the cluster.
    addr = 'localhost:%s' % portpicker.pick_unused_port()
    cluster = server_lib.ClusterSpec({'worker': [addr]})
    remote.connect_to_cluster(cluster, job_name='worker')

    with ops.device('/job:worker/task:0/device:CPU:0'):
      v2 = variables.Variable(initial_value=0)
      v2.assign_add(1)

  # TODO(b/152224115): Re-enable this test.
  def DISABLED_testSimpleParameterServerWithDeviceFilters(self):
    cluster_device_filters = server_lib.ClusterDeviceFilters()
    for i in range(2):
      cluster_device_filters.set_device_filters('my_worker', i, ['/job:my_ps'])
      cluster_device_filters.set_device_filters('my_ps', i, ['/job:my_worker'])
    remote.connect_to_cluster(
        self._cluster, cluster_device_filters=cluster_device_filters)

    with ops.device('/job:my_ps/task:0/device:CPU:0'):
      v1 = variables.Variable(initial_value=0)
    with ops.device('/job:my_ps/task:1/device:CPU:0'):
      v2 = variables.Variable(initial_value=10)

    @def_function.function
    def worker_fn():
      v1.assign_add(1)
      v2.assign_sub(2)
      return v1.read_value() + v2.read_value()

    with ops.device('/job:my_worker/task:0/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 9)
    with ops.device('/job:my_worker/task:1/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 8)

    # The following remote call would fail because the ps nodes cannot see each
    # other due to the device filters.
    with self.assertRaises(errors.InvalidArgumentError) as cm:
      with ops.device('/job:my_ps/task:0/device:CPU:0'):
        worker_fn().numpy()
    self.assertIn('/job:my_ps/replica:0/task:1/device:CPU:0 unknown device',
                  cm.exception.message)

    with self.assertRaises(errors.InvalidArgumentError) as cm:
      with ops.device('/job:my_ps/task:1/device:CPU:0'):
        worker_fn().numpy()
    self.assertIn('/job:my_ps/replica:0/task:0/device:CPU:0 unknown device',
                  cm.exception.message)

    with ops.device('/job:my_worker/task:0/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 7)
    with ops.device('/job:my_worker/task:1/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 6)
    # Explicitly delete variables to avoid triggering errors when being GC'ed in
    # subsequent tests.
    del v1, v2

  def testConnectWithClusterResolver(self):
    remote.connect_to_cluster(self._cluster_resolver)

    v1 = variables.Variable(initial_value=0)
    v2 = variables.Variable(initial_value=10)

    @def_function.function
    def worker_fn():
      v1.assign_add(1)
      v2.assign_sub(2)
      return v1.read_value() + v2.read_value()

    with ops.device('/job:my_worker/task:0/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 9)

    with ops.device('/job:my_worker/task:1/device:CPU:0'):
      self.assertAllEqual(worker_fn(), 8)

  def testConnectToClusterTwiceOk(self):
    remote.connect_to_cluster(self._cluster_resolver)
    remote.connect_to_cluster(self._cluster_resolver)

  def testConnectToClusterOnMismatchedDevice(self):
    remote.connect_to_cluster(self._cluster_resolver)

    # enter into another device scope.
    ops.device('/job:my_worker/task:0/device:CPU:0').__enter__()

    with self.assertRaises(ValueError):
      remote.connect_to_cluster(self._cluster_resolver)

  def testConnectToClusterWithLocalMaster(self):
    local_resolver = SimpleClusterResolver(ClusterSpec({}), master='local')
    remote.connect_to_cluster(local_resolver)

  def testConnectToClusterInGraphModeWillFail(self):
    ops.disable_eager_execution()
    with self.assertRaises(ValueError):
      remote.connect_to_cluster(self._cluster_resolver)
    ops.enable_eager_execution()

  def testConnectToClusterWithoutLocalGpu(self):
    # Only remote workers have GPU devices
    context.context().set_visible_devices([], 'GPU')
    # Ensure that no default device is set in eager context
    remote.connect_to_cluster(self._cluster_resolver,
                              make_master_device_default=False)
    self.assertEmpty(context.get_device_name())

    v1 = variables.Variable(initial_value=0)
    v1.assign_add(1)
    self.assertAllEqual(v1.read_value(), 1)

  # TODO(b/249134783): Add a test for task failures by introducing an Op for
  # reporting errors.
  def testGetTaskStatesAllOK(self):
    context.context().configure_coordination_service(
        service_type='standalone', service_leader='/job:my_ps/replica:0/task:0')
    remote.connect_to_cluster(self._cluster)
    context.context().ensure_initialized()

    states = context.context().get_task_states([('my_worker', 2), ('my_ps', 2)])
    self.assertLen(states, 4)
    for state in states:
      self.assertIsNone(state)


def _strip_prefix(s, prefix):
  return s[len(prefix):] if s.startswith(prefix) else s


if __name__ == '__main__':
  test.main()