[Core] Adjust NodeDeathInfo in raylet (#45533)

In #45357, in the case of preemption, and after the draining deadline, we expect autoscaler to directly SIGKILL the node without sending SIGTERM first. However, autoscaler would send SIGTERM first then SIGKILL. Therefore, this PR does the following:

- At Raylet, distinguish if a SIGTERM is initiated from preemption draining or normal termination based on existing draining request
- Construct node death info accordingly and send UnregisterNode() RPC to GCS

Signed-off-by: Rui Qiao <[email protected]>

Author: ruisearch42

python/ray/tests/test_draining.py

@@ -143,7 +143,11 @@ def ping(self):
     assert worker_node["DeathReasonMessage"] == "preemption"
 
 
-def test_preemption_after_draining_deadline(monkeypatch, ray_start_cluster):
+@pytest.mark.parametrize(
+    "graceful",
+    [True, False],
+)
+def test_preemption_after_draining_deadline(monkeypatch, ray_start_cluster, graceful):
     monkeypatch.setenv("RAY_health_check_failure_threshold", "3")
     monkeypatch.setenv("RAY_health_check_timeout_ms", "100")
     monkeypatch.setenv("RAY_health_check_period_ms", "1000")
@@ -182,9 +186,8 @@ def ping(self):
     )
     assert is_accepted
 
-    # Simulate node provider forcefully terminates the worker node
-    # after the draining deadline.
-    cluster.remove_node(worker_node, False)
+    # Simulate autoscaler terminates the worker node after the draining deadline.
+    cluster.remove_node(worker_node, graceful)
 
     wait_for_condition(
         lambda: {node["NodeID"] for node in ray.nodes() if (node["Alive"])}
@@ -384,7 +387,11 @@ def get_node_id():
     ray.get(obj, timeout=2) == head_node_id
 
 
-def test_draining_reason(ray_start_cluster):
+@pytest.mark.parametrize(
+    "graceful",
+    [False, True],
+)
+def test_draining_reason(ray_start_cluster, graceful):
     cluster = ray_start_cluster
     cluster.add_node(num_cpus=1, resources={"node1": 1})
     ray.init(
@@ -414,8 +421,8 @@ def ping(self):
     )
     assert is_accepted
 
-    # Simulate node provider forcefully terminates the worker node
-    cluster.remove_node(node2, False)
+    # Simulate autoscaler terminates the worker node after the draining deadline.
+    cluster.remove_node(node2, graceful)
     try:
         ray.get(actor.ping.remote())
         raise

src/ray/raylet/main.cc

@@ -439,11 +439,25 @@ int main(int argc, char *argv[]) {
         raylet->Start();
       }));
 
-  auto signal_handler = [shutdown_raylet_gracefully_internal](
+  auto signal_handler = [&raylet, shutdown_raylet_gracefully_internal](
                             const boost::system::error_code &error, int signal_number) {
     ray::rpc::NodeDeathInfo node_death_info;
-    node_death_info.set_reason(ray::rpc::NodeDeathInfo::EXPECTED_TERMINATION);
-    node_death_info.set_reason_message("received SIGTERM");
+    optional<ray::rpc::DrainRayletRequest> drain_request =
+        raylet->node_manager().GetLocalDrainRequest();
+    RAY_LOG(INFO) << "received SIGTERM. Existing local drain request = "
+                  << (drain_request.has_value() ? drain_request->DebugString() : "None");
+    if (drain_request.has_value() &&
+        drain_request->reason() ==
+            ray::rpc::autoscaler::DrainNodeReason::DRAIN_NODE_REASON_PREEMPTION &&
+        drain_request->deadline_timestamp_ms() != 0 &&
+        drain_request->deadline_timestamp_ms() < current_time_ms()) {
+      node_death_info.set_reason(ray::rpc::NodeDeathInfo::AUTOSCALER_DRAIN_PREEMPTED);
+      node_death_info.set_reason_message(drain_request->reason_message());
+    } else {
+      node_death_info.set_reason(ray::rpc::NodeDeathInfo::EXPECTED_TERMINATION);
+      node_death_info.set_reason_message("received SIGTERM");
+    }
+
     shutdown_raylet_gracefully_internal(node_death_info);
   };
   boost::asio::signal_set signals(main_service);

src/ray/raylet/node_manager.cc

@@ -1974,11 +1974,8 @@ void NodeManager::HandleDrainRaylet(rpc::DrainRayletRequest request,
     const bool is_idle =
         cluster_resource_scheduler_->GetLocalResourceManager().IsLocalNodeIdle();
     if (is_idle) {
-      rpc::NodeDeathInfo node_death_info;
-      node_death_info.set_reason(rpc::NodeDeathInfo::AUTOSCALER_DRAIN_IDLE);
-      node_death_info.set_reason_message(request.reason_message());
       cluster_resource_scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-          request.deadline_timestamp_ms(), node_death_info);
+          request);
       reply->set_is_accepted(true);
     } else {
       reply->set_is_accepted(false);
@@ -1989,11 +1986,7 @@ void NodeManager::HandleDrainRaylet(rpc::DrainRayletRequest request,
     // Non-rejectable draining request.
     RAY_CHECK_EQ(request.reason(),
                  rpc::autoscaler::DrainNodeReason::DRAIN_NODE_REASON_PREEMPTION);
-    rpc::NodeDeathInfo node_death_info;
-    node_death_info.set_reason(rpc::NodeDeathInfo::AUTOSCALER_DRAIN_PREEMPTED);
-    node_death_info.set_reason_message(request.reason_message());
-    cluster_resource_scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-        request.deadline_timestamp_ms(), node_death_info);
+    cluster_resource_scheduler_->GetLocalResourceManager().SetLocalNodeDraining(request);
     reply->set_is_accepted(true);
   }
 

src/ray/raylet/node_manager.h

@@ -223,6 +223,11 @@ class NodeManager : public rpc::NodeManagerServiceHandler,
     return mutable_object_provider_;
   }
 
+  /// Get the local drain request.
+  optional<rpc::DrainRayletRequest> GetLocalDrainRequest() const {
+    return cluster_resource_scheduler_->GetLocalResourceManager().GetLocalDrainRequest();
+  }
+
  private:
   void ReleaseWorker(const WorkerID &worker_id) {
     leased_workers_.erase(worker_id);

src/ray/raylet/placement_group_resource_manager_test.cc

@@ -184,9 +184,10 @@ TEST_F(NewPlacementGroupResourceManagerTest, TestNewPrepareBundleDuringDraining)
 
   ASSERT_TRUE(new_placement_group_resource_manager_->PrepareBundles(bundle1_specs));
   // Drain the node, new bundle prepare will fail.
-  rpc::NodeDeathInfo node_death_info;
+  rpc::DrainRayletRequest drain_request;
+  drain_request.set_deadline_timestamp_ms(std::numeric_limits<int64_t>::max());
   cluster_resource_scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-      std::numeric_limits<int64_t>::max(), node_death_info);
+      drain_request);
   ASSERT_FALSE(new_placement_group_resource_manager_->PrepareBundles(bundle2_specs));
   // Prepared bundles can still be committed.
   new_placement_group_resource_manager_->CommitBundles(bundle1_specs);

src/ray/raylet/scheduling/cluster_resource_scheduler.h

@@ -125,6 +125,7 @@ class ClusterResourceScheduler {
                            bool requires_object_store_memory);
 
   LocalResourceManager &GetLocalResourceManager() { return *local_resource_manager_; }
+
   ClusterResourceManager &GetClusterResourceManager() {
     return *cluster_resource_manager_;
   }

src/ray/raylet/scheduling/cluster_resource_scheduler_test.cc

@@ -362,9 +362,9 @@ TEST_F(ClusterResourceSchedulerTest, NodeAffinitySchedulingStrategyTest) {
   ASSERT_TRUE(resource_scheduler.GetLocalResourceManager().AllocateLocalTaskResources(
       resource_request, task_allocation));
   // Drain the local node so that it's not schedulable for new tasks.
-  rpc::NodeDeathInfo node_death_info;
-  resource_scheduler.GetLocalResourceManager().SetLocalNodeDraining(
-      std::numeric_limits<int64_t>::max(), node_death_info);
+  rpc::DrainRayletRequest drain_request;
+  drain_request.set_deadline_timestamp_ms(std::numeric_limits<int64_t>::max());
+  resource_scheduler.GetLocalResourceManager().SetLocalNodeDraining(drain_request);
 
   scheduling_strategy.mutable_node_affinity_scheduling_strategy()->set_node_id(
       local_node_id.Binary());

src/ray/raylet/scheduling/cluster_task_manager_test.cc

@@ -743,9 +743,9 @@ TEST_F(ClusterTaskManagerTest, DrainingWhileResolving) {
   ASSERT_EQ(pool_.workers.size(), 1);
 
   // Drain the local node.
-  rpc::NodeDeathInfo node_death_info;
-  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-      std::numeric_limits<int64_t>::max(), node_death_info);
+  rpc::DrainRayletRequest drain_request;
+  drain_request.set_deadline_timestamp_ms(std::numeric_limits<int64_t>::max());
+  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(drain_request);
 
   // Arg is resolved.
   missing_objects_.erase(missing_arg);
@@ -1078,9 +1078,9 @@ TEST_F(ClusterTaskManagerTest, NotOKPopWorkerAfterDrainingTest) {
   AddNode(remote_node_id, 5);
 
   // Drain the local node.
-  rpc::NodeDeathInfo node_death_info;
-  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-      std::numeric_limits<int64_t>::max(), node_death_info);
+  rpc::DrainRayletRequest drain_request;
+  drain_request.set_deadline_timestamp_ms(std::numeric_limits<int64_t>::max());
+  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(drain_request);
 
   pool_.callbacks[task1.GetTaskSpecification().GetRuntimeEnvHash()].front()(
       nullptr, PopWorkerStatus::WorkerPendingRegistration, "");
@@ -2637,9 +2637,9 @@ TEST_F(ClusterTaskManagerTest, PopWorkerBeforeDraining) {
   task_manager_.QueueAndScheduleTask(task, false, false, &reply, callback);
 
   // Drain the local node.
-  rpc::NodeDeathInfo node_death_info;
-  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-      std::numeric_limits<int64_t>::max(), node_death_info);
+  rpc::DrainRayletRequest drain_request;
+  drain_request.set_deadline_timestamp_ms(std::numeric_limits<int64_t>::max());
+  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(drain_request);
 
   std::shared_ptr<MockWorker> worker =
       std::make_shared<MockWorker>(WorkerID::FromRandom(), 1234);
@@ -2677,9 +2677,9 @@ TEST_F(ClusterTaskManagerTest, UnscheduleableWhileDraining) {
   AddNode(remote_node_id, 5);
 
   // Drain the local node.
-  rpc::NodeDeathInfo node_death_info;
-  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(
-      std::numeric_limits<int64_t>::max(), node_death_info);
+  rpc::DrainRayletRequest drain_request;
+  drain_request.set_deadline_timestamp_ms(std::numeric_limits<int64_t>::max());
+  scheduler_->GetLocalResourceManager().SetLocalNodeDraining(drain_request);
 
   RayTask spillback_task = CreateTask({{ray::kCPU_ResourceLabel, 1}});
   rpc::RequestWorkerLeaseReply spillback_reply;

src/ray/raylet/scheduling/local_resource_manager.cc

@@ -254,9 +254,8 @@ NodeResources LocalResourceManager::ToNodeResources() const {
   node_resources.available = local_resources_.available.ToNodeResourceSet();
   node_resources.total = local_resources_.total.ToNodeResourceSet();
   node_resources.labels = local_resources_.labels;
-  node_resources.is_draining = is_local_node_draining_;
-  node_resources.draining_deadline_timestamp_ms =
-      local_node_draining_deadline_timestamp_ms_;
+  node_resources.is_draining = IsLocalNodeDraining();
+  node_resources.draining_deadline_timestamp_ms = GetDrainingDeadline();
   return node_resources;
 }
 
@@ -330,8 +329,7 @@ void LocalResourceManager::PopulateResourceViewSyncMessage(
   }
 
   resource_view_sync_message.set_is_draining(IsLocalNodeDraining());
-  resource_view_sync_message.set_draining_deadline_timestamp_ms(
-      local_node_draining_deadline_timestamp_ms_);
+  resource_view_sync_message.set_draining_deadline_timestamp_ms(GetDrainingDeadline());
 
   for (const auto &iter : last_idle_times_) {
     if (iter.second == absl::nullopt) {
@@ -383,7 +381,8 @@ std::optional<syncer::RaySyncMessage> LocalResourceManager::CreateSyncMessage(
 void LocalResourceManager::OnResourceOrStateChanged() {
   if (IsLocalNodeDraining() && IsLocalNodeIdle()) {
     RAY_LOG(INFO) << "The node is drained, continue to shut down raylet...";
-    shutdown_raylet_gracefully_(node_death_info_);
+    rpc::NodeDeathInfo node_death_info = DeathInfoFromDrainRequest();
+    shutdown_raylet_gracefully_(std::move(node_death_info));
   }
 
   ++version_;
@@ -393,6 +392,22 @@ void LocalResourceManager::OnResourceOrStateChanged() {
   resource_change_subscriber_(ToNodeResources());
 }
 
+rpc::NodeDeathInfo LocalResourceManager::DeathInfoFromDrainRequest() {
+  rpc::NodeDeathInfo death_info;
+  RAY_CHECK(drain_request_.has_value());
+  if (drain_request_->reason() ==
+      rpc::autoscaler::DrainNodeReason::DRAIN_NODE_REASON_IDLE_TERMINATION) {
+    death_info.set_reason(rpc::NodeDeathInfo::AUTOSCALER_DRAIN_IDLE);
+    death_info.set_reason_message(drain_request_->reason_message());
+  } else {
+    RAY_CHECK_EQ(drain_request_->reason(),
+                 rpc::autoscaler::DrainNodeReason::DRAIN_NODE_REASON_PREEMPTION);
+    death_info.set_reason(rpc::NodeDeathInfo::AUTOSCALER_DRAIN_PREEMPTED);
+    death_info.set_reason_message(drain_request_->reason_message());
+  }
+  return death_info;
+}
+
 bool LocalResourceManager::ResourcesExist(scheduling::ResourceID resource_id) const {
   return local_resources_.total.Has(resource_id);
 }
@@ -445,11 +460,8 @@ void LocalResourceManager::RecordMetrics() const {
 }
 
 void LocalResourceManager::SetLocalNodeDraining(
-    int64_t draining_deadline_timestamp_ms, const rpc::NodeDeathInfo &node_death_info) {
-  RAY_CHECK_GE(draining_deadline_timestamp_ms, 0);
-  is_local_node_draining_ = true;
-  local_node_draining_deadline_timestamp_ms_ = draining_deadline_timestamp_ms;
-  node_death_info_ = node_death_info;
+    const rpc::DrainRayletRequest &drain_request) {
+  drain_request_ = std::make_optional(drain_request);
   OnResourceOrStateChanged();
 }
 

src/ray/raylet/scheduling/local_resource_manager.h

@@ -31,6 +31,7 @@
 #include "ray/gcs/gcs_client/gcs_client.h"
 #include "ray/util/logging.h"
 #include "src/ray/protobuf/gcs.pb.h"
+#include "src/ray/protobuf/node_manager.pb.h"
 
 namespace ray {
 
@@ -154,10 +155,17 @@ class LocalResourceManager : public syncer::ReporterInterface {
 
   /// Change the local node to the draining state.
   /// After that, no new tasks can be scheduled onto the local node.
-  void SetLocalNodeDraining(int64_t draining_deadline_timestamp_ms,
-                            const rpc::NodeDeathInfo &node_death_info);
+  void SetLocalNodeDraining(const rpc::DrainRayletRequest &drain_request);
 
-  bool IsLocalNodeDraining() const { return is_local_node_draining_; }
+  bool IsLocalNodeDraining() const { return drain_request_.has_value(); }
+
+  /// Get the local drain request.
+  std::optional<rpc::DrainRayletRequest> GetLocalDrainRequest() const {
+    return drain_request_;
+  }
+
+  /// Generate node death info from existing drain request.
+  rpc::NodeDeathInfo DeathInfoFromDrainRequest();
 
  private:
   struct ResourceUsage {
@@ -206,6 +214,12 @@ class LocalResourceManager : public syncer::ReporterInterface {
 
   absl::optional<absl::Time> GetResourceIdleTime() const;
 
+  /// Get the draining deadline if node is in draining state.
+  ///
+  /// \return The draining deadline if node is in draining state, otherwise -1.
+  int64_t GetDrainingDeadline() const {
+    return drain_request_.has_value() ? drain_request_->deadline_timestamp_ms() : -1;
+  }
   /// Identifier of local node.
   scheduling::NodeID local_node_id_;
   /// Resources of local node.
@@ -226,15 +240,8 @@ class LocalResourceManager : public syncer::ReporterInterface {
   // Version of this resource. It will incr by one whenever the state changed.
   int64_t version_ = 0;
 
-  // Whether the local node is being drained or not.
-  bool is_local_node_draining_ = false;
-  // The value is the timestamp when
-  // the node will be force killed.
-  // 0 if there is no deadline.
-  int64_t local_node_draining_deadline_timestamp_ms_ = -1;
-
-  /// This is set when the node is being drained and indicates the reason for draining.
-  rpc::NodeDeathInfo node_death_info_;
+  /// The draining request this node received.
+  std::optional<rpc::DrainRayletRequest> drain_request_;
 
   FRIEND_TEST(ClusterResourceSchedulerTest, SchedulingUpdateTotalResourcesTest);
   FRIEND_TEST(ClusterResourceSchedulerTest, AvailableResourceInstancesOpsTest);