class_docstr: Test preemptions during strategy init.
code: 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)
  
imports: 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
  from tensorflow.python.training import server_lib
label_desc: Use image node_id 3 for calling the InitFaultToleranceTest obj's underlying member method code with example usage: obj.testWorkerPreemptionDuringInit() and returns: restart_thread