code: 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)
  
imports: 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
  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
label_desc: Use image node_id 6 for calling the GetTaskStatesTest obj's underlying member method code with example usage: obj.testPSPreempted() without return types