code: def raise_if_not_all_exit(grace_period, mpr):
      """Wait for all cluster to exit with a time out."""
      waiting_time = 0
      exit_process_count = 0
      # This addition to mitigate the fact that our step time is too short in test
      while exit_process_count != CLUSTER_SIZE and waiting_time < max(
          grace_period + 15, MAX_WAIT_TIME
      ):
          exit_process_count = 0
          for worker_id in range(CLUSTER_SIZE):
              if not mpr.process_exists("worker", worker_id):
                  exit_process_count += 1
          waiting_time += 1
          time.sleep(1)
  
      if waiting_time == max(grace_period + 5, 40):
          raise RuntimeError(
              "Waited long but at least one worker still exist. "
              "Considering size of our model, this should not"
              " happen."
          )
  
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 3 for calling a global function with example usage: raise_if_not_all_exit(grace_period, mpr) without return types