class_docstr: Integration test for PreemptionCheckpointHandler.
class_docstr_tok: ["Integration","test","for","PreemptionCheckpointHandler","."]
code: def test_grace_period_continue_training(
      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":
          grace_period = 5
          termination_config = failure_handling.TerminationConfig(
              grace_period=grace_period
          )
          has_chief = False
          cluster_spec = multi_worker_test_base.create_cluster_spec(
              has_chief=has_chief, num_workers=CLUSTER_SIZE
          )
          training_started_event = (
              multi_process_runner.manager().Event()
          )
          training_restarted = multi_process_runner.manager().Event()
          training_finished = multi_process_runner.manager().Event()
  
          mpr = multi_process_runner.MultiProcessRunner(
              self.worker_fn,
              cluster_spec,
              args=(
                  checkpoint_dir,
                  cluster_spec,
                  strategy_option,
                  input_arg,
                  [training_started_event],
                  None,
                  training_restarted,
                  training_finished,
                  termination_config,
              ),
              rpc_layer=rpc_layer,
              return_output=True,
              dependence_on_chief=has_chief,
          )
  
          logging.info("Cluster starting.")
          mpr.start()
          while not training_started_event.is_set():
              time.sleep(1)
  
          killed_worker = random.randrange(0, CLUSTER_SIZE)
          logging.info("sending SIGTERM")
          os.kill(
              mpr.get_process_id("worker", killed_worker),
              signal.SIGTERM,
          )
          logging.info("SIGTERM sent")
  
          raise_if_not_all_exit(grace_period, mpr)
  
          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)
  
      else:
          # This is because single worker trains super fast with regards to the size
          # of "model" here. With a longer grace period, the training just finishes
          # within the grace period so we can't verify the exit behavior.
          grace_period = 1
          termination_config = failure_handling.TerminationConfig(
              grace_period=grace_period
          )
          cluster_spec = server_lib.ClusterSpec({})
  
          training_started_event = threading.Event()
          training_restarted = threading.Event()
          training_finished = threading.Event()
  
          def sending_sigterm(training_started_event):
              while not training_started_event.is_set():
                  time.sleep(1)
              logging.info("sending sigterm")
              training_started_event.set()
              os.kill(os.getpid(), signal.SIGTERM)
  
          preemption_sender_thread = threading.Thread(
              target=sending_sigterm, args=(training_started_event,)
          )
          preemption_sender_thread.start()
  
          caught_exit = False
          try:
              self.worker_fn(
                  checkpoint_dir,
                  cluster_spec,
                  strategy_option,
                  input_arg,
                  [training_started_event],
                  None,
                  training_restarted,
                  training_finished,
                  termination_config,
              )
  
          except SystemExit as exit_error:
              caught_exit = True
              # We cannot use assertRaise instead, since termination is not always
              # triggered.
              self.assertEqual(
                  exit_error.code, 42
              )  # pylint: disable=g-assert-in-except
  
          preemption_sender_thread.join(10)
          if not training_finished.is_set():
              self.assertTrue(caught_exit)
  
              logging.info("restarting workers")
              training_restarted.set()
              self.worker_fn(
                  checkpoint_dir,
                  cluster_spec,
                  strategy_option,
                  input_arg,
                  [training_started_event],
                  None,
                  training_restarted,
                  training_finished,
                  termination_config,
              )
  
imports: import os
  import random
  import re
  import signal
  import sys
  import threading
  import time
  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 distribute_lib
  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.framework import errors_impl
  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 5 for calling the PreemptionCheckpointTest obj's underlying member method code with example usage: obj.test_grace_period_continue_training(input_arg, strategy_option) without return types