code: def testAbortCommunication(self, device, communication):
      if communication == "NCCL":
          self.skipTest("b/171358086: cannot test multi worker NCCL")
      dev0 = "/device:%s:0" % device
      cluster_resolver = cluster_resolver_lib.TFConfigClusterResolver()
      enable_collective_ops_with_barrier(cluster_resolver)
      group_size = 2
      group_key = 100
      instance_key = 100
      in_tensor = constant_op.constant([1.0])
  
      # First perform a normal all-reduce to complete the group and instance
      # resolution.
      with ops.device(dev0):
          collective_ops.all_reduce(
              in_tensor,
              group_size,
              group_key,
              instance_key,
              communication_hint=communication,
          )
  
      if cluster_resolver.task_id == 1:
  
          def abort_fn():
              time.sleep(2)
              context.context().abort_collective_ops(
                  errors.UNAVAILABLE, "peer down"
              )
  
          t = threading.Thread(target=abort_fn)
          t.start()
  
          with self.assertRaisesRegex(
              errors.UnavailableError, "peer down"
          ):
              with ops.device(dev0):
                  collective_ops.all_reduce(
                      in_tensor,
                      group_size,
                      group_key,
                      instance_key,
                      communication_hint=communication,
                  )
  
          # After abortion, subsequent collectives should fail immediately.
          with self.assertRaisesRegex(
              errors.UnavailableError, "peer down"
          ):
              with ops.device(dev0):
                  collective_ops.all_reduce(
                      in_tensor,
                      group_size,
                      group_key,
                      instance_key,
                      communication_hint=communication,
                  )
  
          t.join()
  
      # Enable collective ops again in order to reset the collective executor.
      enable_collective_ops_with_barrier(cluster_resolver)
      with ops.device(dev0):
          collective_ops.all_reduce(
              in_tensor,
              group_size,
              group_key,
              instance_key,
              communication_hint=communication,
          )
  
imports: import copy
  import os
  import threading
  import time
  from absl.testing import parameterized
  from tensorflow.core.protobuf import tensorflow_server_pb2
  from tensorflow.python.distribute import cluster_resolver
  from tensorflow.python.distribute import combinations
  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 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.ops import collective_ops
label_desc: Use image node_id 1 for calling the AbortCollectiveOpsTest obj's underlying member method code with example usage: obj.testAbortCommunication(device, communication) without return types