code: def test_get_draining_nodes(ray_start_cluster):
      cluster = ray_start_cluster
      cluster.add_node()
      ray.init(address=cluster.address)
      cluster.add_node(resources={"worker": 1})
      cluster.wait_for_nodes()
  
      @ray.remote
      def get_node_id():
          return ray.get_runtime_context().get_node_id()
  
      worker_node_id = ray.get(
          get_node_id.options(resources={"worker": 1}).remote()
      )
  
      # Initially there is no draining node.
      assert ray._private.state.state.get_draining_nodes() == set()
  
      @ray.remote(num_cpus=1, resources={"worker": 1})
      class Actor:
          def ping(self):
              pass
  
      actor = Actor.remote()
      ray.get(actor.ping.remote())
  
      gcs_client = GcsClient(
          address=ray.get_runtime_context().gcs_address
      )
  
      # Drain the worker node.
      is_accepted = gcs_client.drain_node(
          worker_node_id,
          autoscaler_pb2.DrainNodeReason.Value(
              "DRAIN_NODE_REASON_PREEMPTION"
          ),
          "preemption",
      )
      assert is_accepted
  
      wait_for_condition(
          lambda: ray._private.state.state.get_draining_nodes()
          == {worker_node_id}
      )
  
      # Kill the actor running on the draining worker node so
      # that the worker node becomes idle and can be drained.
      ray.kill(actor)
  
      wait_for_condition(
          lambda: ray._private.state.state.get_draining_nodes() == set()
      )
  
imports: import os
  import time
  import pytest
  import ray
  import ray._private.gcs_utils
  import ray._private.ray_constants
  from ray._raylet import GcsClient
  from ray.core.generated import autoscaler_pb2
  from ray._private.test_utils import convert_actor_state, make_global_state_accessor, wait_for_condition
label_desc: Use image node_id 16 for calling a global function with example usage: test_get_draining_nodes(ray_start_cluster) and returns: ray