class_docstr: Integration test for PreemptionCheckpointHandler.
class_docstr_tok: ["Integration","test","for","PreemptionCheckpointHandler","."]
code: def test_basic_run(self, input_arg, strategy_option):
      if _is_oss():
          rpc_layer = "grpc"
      else:
          rpc_layer = "grpc+loas"
      checkpoint_dir = os.path.join(self.get_temp_dir(), "fh_ckpt/")
  
      if strategy_option == "MWMS_multi_worker":
          has_chief = False
          cluster_spec = multi_worker_test_base.create_cluster_spec(
              has_chief=has_chief, num_workers=CLUSTER_SIZE
          )
          maintenance_event = multi_process_runner.manager().Event()
          training_finished = multi_process_runner.manager().Event()
          training_restarted = multi_process_runner.manager().Event()
  
          mpr = multi_process_runner.MultiProcessRunner(
              self.worker_fn,
              cluster_spec,
              args=(
                  checkpoint_dir,
                  cluster_spec,
                  input_arg,
                  strategy_option,
                  maintenance_event,
                  training_finished,
                  False,
                  training_restarted,
              ),
              rpc_layer=rpc_layer,
              return_output=True,
              dependence_on_chief=has_chief,
          )
  
          logging.info("Cluster starting.")
          mpr.start()
  
          raise_if_not_all_exit(0, mpr)
  
          if not training_finished.is_set():
              logging.info("restarting workers")
              training_restarted.set()
              for worker_id in range(CLUSTER_SIZE):
                  mpr.start_single_process(
                      "worker", worker_id, cluster_spec
                  )
              logging.info("workers restarted")
  
          mpr.join(timeout=250)
          self.assertTrue(training_finished.is_set())
  
      else:
          maintenance_event = threading.Event()
          training_finished = threading.Event()
          training_restarted = threading.Event()
  
          cluster_spec = server_lib.ClusterSpec({})
          caught_exit = False
          try:
              self.worker_fn(
                  checkpoint_dir,
                  cluster_spec,
                  input_arg,
                  strategy_option,
                  maintenance_event,
                  training_finished,
                  False,
                  training_restarted,
              )
          except SystemExit as exit_error:
              caught_exit = True
              # We cannot use assertRaise instead, since termination is not always
              # triggered.
              self.assertEqual(
                  exit_error.code, 143
              )  # pylint:disable=g-assert-in-except
  
          if (
              maintenance_event.is_set()
              and not training_finished.is_set()
          ):
              self.assertTrue(caught_exit)
  
              logging.info("restarting workers")
              training_restarted.set()
              self.worker_fn(
                  checkpoint_dir,
                  cluster_spec,
                  input_arg,
                  strategy_option,
                  maintenance_event,
                  training_finished,
                  False,
                  training_restarted,
              )
  
          self.assertTrue(training_finished.is_set())
  
imports: import os
  import random
  import re
  import sys
  import threading
  import time
  import urllib
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.distribute import collective_all_reduce_strategy
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import mirrored_strategy
  from tensorflow.python.distribute import multi_process_runner
  from tensorflow.python.distribute import multi_worker_test_base
  from tensorflow.python.distribute import multi_worker_util
  from tensorflow.python.distribute import one_device_strategy
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute.failure_handling import failure_handling
  from tensorflow.python.distribute.failure_handling import failure_handling_util
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.module import module
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import gfile
  from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.training import server_lib
label_desc: Use image node_id 2 for calling the GceFailureHandlingTest obj's underlying member method code with example usage: obj.test_basic_run(input_arg, strategy_option) without return types