code: def testRestartWorker(
      self, use_same_port, work_dir, fault_tolerant_mode
  ):
      cluster = data_service_test_base.TestCluster(
          num_workers=1,
          work_dir=work_dir,
          fault_tolerant_mode=fault_tolerant_mode,
      )
      num_elements = 2 * multiprocessing.cpu_count() + 100
      ds = self.make_distributed_range_dataset(num_elements, cluster)
      iterator = iter(ds)
      # Read halfway through the dataset.
      midpoint = num_elements // 2
      for i in range(midpoint):
          self.assertEqual(i, next(iterator).numpy())
  
      # Stop the original worker and start a new one.
      cluster.workers[0].restart(use_same_port=use_same_port)
  
      # There may have been some elements prefetched from the first worker
      # before it was stopped.
      while True:
          val = next(iterator).numpy()
          if val == 0:
              break
  
      # The dataset starts over now that we read from the new worker.
      # TODO(b/157086991): Iterate until end of sequence when we support
      # detecting lost workers.
      for i in range(1, num_elements // 2):
          val = next(iterator).numpy()
          self.assertEqual(i, val)
  
imports: import multiprocessing
  import threading
  import time
  from absl.testing import parameterized
  from tensorflow.python.data.experimental.kernel_tests.service import test_base
  from tensorflow.python.data.experimental.ops import data_service_ops
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import array_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 14 for calling the FaultToleranceTest obj's underlying member method code with example usage: obj.testRestartWorker(use_same_port, work_dir, fault_tolerant_mode) without return types