code: def test_preemption_handling(
      cluster_setup,
      tmp_path,
      error_at_level: str,
  ):
      """Integration test for node preemption handling in Ray Train/Tune.
      Even though `max_failures=0`, preemption errors should still be retried.
      """
      cluster, nodes, node_ids = cluster_setup
      # node 1 = coordinator, node 2 = worker, node 3 = worker
      coordinator_node, worker_node, _ = nodes
      coordinator_node_id, worker_node_id, _ = node_ids
  
      num_workers = 2
      tmp_path.joinpath("markers").mkdir()
  
      def train_fn(config):
          checkpoint = train.get_checkpoint()
          start_iter = 0
          if checkpoint:
              start_iter = load_dict_checkpoint(checkpoint)["iter"] + 1
              print(f"Restored at iter = {start_iter}")
  
          for iter in range(start_iter, 6):
              with create_dict_checkpoint({"iter": iter}) as checkpoint:
                  ray.train.report(
                      {"iter": iter}, checkpoint=checkpoint
                  )
  
              if iter == 2:
                  # Write a "done marker" to tell the driver to simulate a preemption.
                  tmp_path.joinpath(
                      "markers",
                      str(ray.train.get_context().get_world_rank()),
                  ).touch()
                  # Await execution.
                  time.sleep(120)
  
      def launch_training():
          trainer = DataParallelTrainer(
              train_loop_per_worker=train_fn,
              scaling_config=ScalingConfig(
                  num_workers=num_workers,
                  trainer_resources={"coordinator": 1},
                  resources_per_worker={
                      "cpu": 1
                  },  # worker2 and worker3
              ),
              run_config=RunConfig(
                  storage_path=str(tmp_path),
                  name="test_preemption_error",
                  failure_config=train.FailureConfig(
                      fail_fast=False, max_failures=0
                  ),
              ),
          )
          result = trainer.fit()
          assert result.metrics["iter"] == 5
  
      t = threading.Thread(target=launch_training)
      t.start()
  
      # Wait until the workers are ready for preemption (after a few checkpoints).
      while (
          len(list(tmp_path.joinpath("markers").glob("*")))
          < num_workers
      ):
          time.sleep(0.5)
  
      if error_at_level == "coordinator":
          node, node_id = coordinator_node, coordinator_node_id
      elif error_at_level == "worker":
          node, node_id = worker_node, worker_node_id
      else:
          raise NotImplementedError(
              f"Invalid error_at_level = {error_at_level}"
          )
  
      # Preempt a node.
      gcs_client = GcsClient(
          address=ray.get_runtime_context().gcs_address
      )
      print("Draining node...")
      is_accepted = gcs_client.drain_node(
          node_id,
          autoscaler_pb2.DrainNodeReason.Value(
              "DRAIN_NODE_REASON_PREEMPTION"
          ),
          "preemption",
      )
      assert is_accepted
      print("Killing node...")
      cluster.remove_node(node, allow_graceful=True)
      print("Adding new node..")  # so that the job can be rescheduled
      # New node can replace a preempted coordinator or worker
      # NOTE: `cluster.add_node` only works in the main thread.
      cluster.add_node(resources={"coordinator": 1, "cpu": 1})
      t.join()
  
imports: import gc
  import threading
  import time
  from tempfile import TemporaryDirectory
  import pytest
  import ray
  from ray import train, tune
  from ray._private.test_utils import wait_for_condition
  from ray._raylet import GcsClient
  from ray.cluster_utils import Cluster
  from ray.core.generated import autoscaler_pb2
  from ray.train import Checkpoint, FailureConfig, RunConfig, ScalingConfig
  from ray.train.data_parallel_trainer import DataParallelTrainer
  from ray.train.trainer import BaseTrainer, TrainingFailedError
  from ray.tune import Tuner, TuneConfig, TuneError
  from ray.tests.conftest import *
  from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
label_desc: Use image node_id 10 for calling a global function with example usage: test_preemption_handling(cluster_setup, tmp_path, error_at_level) without return types