[core] Fix idempotency issues in RequestWorkerLease for scheduled leases (#58265)

## Description
> Briefly describe what this PR accomplishes and why it's needed.

Using the ip tables script created in #58241 we found a bug in
RequestWorkerLease where a RAY_CHECK was being triggered here:
https://github.com/ray-project/ray/blob/66c08b47a195bcfac6878a234dc804142e488fc2/src/ray/raylet/lease_dependency_manager.cc#L222-L223
The issue is that transient network errors can happen ANYTIME, including
when the server logic is executing and has not yet replied back to the
client. Our original testing framework using an env variable to drop the
request or reply when it's being sent, hence this was missed. The issue
specifically is that RequestWorkerLease could be in the process of
pulling the lease dependencies to it's local plasma store, and the retry
can arrive triggering this check. Created a cpp unit test that
specifically triggers this RAY_CHECK without this change and is fixed. I
decided to store the callbacks instead of replacing the older one with
the new one due to the possibility of message reordering where the new
one could arrive before the old one.

---------

Signed-off-by: joshlee <[email protected]>

Author: Sparks0219

src/mock/ray/raylet/local_lease_manager.h

@@ -79,6 +79,17 @@ class MockLocalLeaseManager : public LocalLeaseManagerInterface {
   MOCK_METHOD(size_t, GetNumLeaseSpilled, (), (const, override));
   MOCK_METHOD(size_t, GetNumWaitingLeaseSpilled, (), (const, override));
   MOCK_METHOD(size_t, GetNumUnschedulableLeaseSpilled, (), (const, override));
+  MOCK_METHOD(bool,
+              IsLeaseQueued,
+              (const SchedulingClass &scheduling_class, const LeaseID &lease_id),
+              (const, override));
+  MOCK_METHOD(bool,
+              AddReplyCallback,
+              (const SchedulingClass &scheduling_class,
+               const LeaseID &lease_id,
+               rpc::SendReplyCallback send_reply_callback,
+               rpc::RequestWorkerLeaseReply *reply),
+              (override));
 };
 
 }  // namespace ray::raylet

src/ray/gcs/gcs_actor_scheduler.cc

@@ -103,11 +103,12 @@ void GcsActorScheduler::ScheduleByGcs(std::shared_ptr<GcsActor> actor) {
   RayLease lease(
       actor->GetLeaseSpecification(),
       owner_node.has_value() ? actor->GetOwnerNodeID().Binary() : std::string());
-  cluster_lease_manager_.QueueAndScheduleLease(std::move(lease),
-                                               /*grant_or_reject=*/false,
-                                               /*is_selected_based_on_locality=*/false,
-                                               /*reply=*/reply.get(),
-                                               send_reply_callback);
+  cluster_lease_manager_.QueueAndScheduleLease(
+      std::move(lease),
+      /*grant_or_reject=*/false,
+      /*is_selected_based_on_locality=*/false,
+      {ray::raylet::internal::ReplyCallback(std::move(send_reply_callback),
+                                            reply.get())});
 }
 
 void GcsActorScheduler::ScheduleByRaylet(std::shared_ptr<GcsActor> actor) {

src/ray/raylet/local_lease_manager.cc

@@ -36,11 +36,13 @@ namespace {
 void ReplyCancelled(const std::shared_ptr<internal::Work> &work,
                     rpc::RequestWorkerLeaseReply::SchedulingFailureType failure_type,
                     const std::string &scheduling_failure_message) {
-  auto reply = work->reply_;
-  reply->set_canceled(true);
-  reply->set_failure_type(failure_type);
-  reply->set_scheduling_failure_message(scheduling_failure_message);
-  work->send_reply_callback_(Status::OK(), nullptr, nullptr);
+  for (const auto &reply_callback : work->reply_callbacks_) {
+    auto reply = reply_callback.reply_;
+    reply->set_canceled(true);
+    reply->set_failure_type(failure_type);
+    reply->set_scheduling_failure_message(scheduling_failure_message);
+    reply_callback.send_reply_callback_(Status::OK(), nullptr, nullptr);
+  }
 }
 }  // namespace
 
@@ -547,8 +549,7 @@ bool LocalLeaseManager::PoppedWorkerHandler(
     bool is_detached_actor,
     const rpc::Address &owner_address,
     const std::string &runtime_env_setup_error_message) {
-  const auto &reply = work->reply_;
-  const auto &send_reply_callback = work->send_reply_callback_;
+  const auto &reply_callbacks = work->reply_callbacks_;
   const bool canceled = work->GetState() == internal::WorkStatus::CANCELLED;
   const auto &lease = work->lease_;
   bool granted = false;
@@ -662,12 +663,7 @@ bool LocalLeaseManager::PoppedWorkerHandler(
     RAY_LOG(DEBUG) << "Granting lease " << lease_id << " to worker "
                    << worker->WorkerId();
 
-    Grant(worker,
-          leased_workers_,
-          work->allocated_instances_,
-          lease,
-          reply,
-          send_reply_callback);
+    Grant(worker, leased_workers_, work->allocated_instances_, lease, reply_callbacks);
     erase_from_leases_to_grant_queue_fn(work, scheduling_class);
     granted = true;
   }
@@ -677,11 +673,11 @@ bool LocalLeaseManager::PoppedWorkerHandler(
 
 void LocalLeaseManager::Spillback(const NodeID &spillback_to,
                                   const std::shared_ptr<internal::Work> &work) {
-  auto send_reply_callback = work->send_reply_callback_;
-
   if (work->grant_or_reject_) {
-    work->reply_->set_rejected(true);
-    send_reply_callback(Status::OK(), nullptr, nullptr);
+    for (const auto &reply_callback : work->reply_callbacks_) {
+      reply_callback.reply_->set_rejected(true);
+      reply_callback.send_reply_callback_(Status::OK(), nullptr, nullptr);
+    }
     return;
   }
 
@@ -701,13 +697,15 @@ void LocalLeaseManager::Spillback(const NodeID &spillback_to,
   RAY_CHECK(node_info_ptr)
       << "Spilling back to a node manager, but no GCS info found for node "
       << spillback_to;
-  auto reply = work->reply_;
-  reply->mutable_retry_at_raylet_address()->set_ip_address(
-      node_info_ptr->node_manager_address());
-  reply->mutable_retry_at_raylet_address()->set_port(node_info_ptr->node_manager_port());
-  reply->mutable_retry_at_raylet_address()->set_node_id(spillback_to.Binary());
-
-  send_reply_callback(Status::OK(), nullptr, nullptr);
+  for (const auto &reply_callback : work->reply_callbacks_) {
+    auto reply = reply_callback.reply_;
+    reply->mutable_retry_at_raylet_address()->set_ip_address(
+        node_info_ptr->node_manager_address());
+    reply->mutable_retry_at_raylet_address()->set_port(
+        node_info_ptr->node_manager_port());
+    reply->mutable_retry_at_raylet_address()->set_node_id(spillback_to.Binary());
+    reply_callback.send_reply_callback_(Status::OK(), nullptr, nullptr);
+  }
 }
 
 void LocalLeaseManager::LeasesUnblocked(const std::vector<LeaseID> &ready_ids) {
@@ -969,8 +967,7 @@ void LocalLeaseManager::Grant(
     absl::flat_hash_map<LeaseID, std::shared_ptr<WorkerInterface>> &leased_workers,
     const std::shared_ptr<TaskResourceInstances> &allocated_instances,
     const RayLease &lease,
-    rpc::RequestWorkerLeaseReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
+    const std::vector<internal::ReplyCallback> &reply_callbacks) {
   const auto &lease_spec = lease.GetLeaseSpecification();
 
   if (lease_spec.IsActorCreationTask()) {
@@ -982,11 +979,14 @@ void LocalLeaseManager::Grant(
   worker->GrantLease(lease);
 
   // Pass the contact info of the worker to use.
-  reply->set_worker_pid(worker->GetProcess().GetId());
-  reply->mutable_worker_address()->set_ip_address(worker->IpAddress());
-  reply->mutable_worker_address()->set_port(worker->Port());
-  reply->mutable_worker_address()->set_worker_id(worker->WorkerId().Binary());
-  reply->mutable_worker_address()->set_node_id(self_node_id_.Binary());
+  for (const auto &reply_callback : reply_callbacks) {
+    reply_callback.reply_->set_worker_pid(worker->GetProcess().GetId());
+    reply_callback.reply_->mutable_worker_address()->set_ip_address(worker->IpAddress());
+    reply_callback.reply_->mutable_worker_address()->set_port(worker->Port());
+    reply_callback.reply_->mutable_worker_address()->set_worker_id(
+        worker->WorkerId().Binary());
+    reply_callback.reply_->mutable_worker_address()->set_node_id(self_node_id_.Binary());
+  }
 
   RAY_CHECK(!leased_workers.contains(lease_spec.LeaseId()));
   leased_workers[lease_spec.LeaseId()] = worker;
@@ -1000,26 +1000,29 @@ void LocalLeaseManager::Grant(
   } else {
     allocated_resources = worker->GetAllocatedInstances();
   }
-  ::ray::rpc::ResourceMapEntry *resource;
   for (auto &resource_id : allocated_resources->ResourceIds()) {
-    bool first = true;  // Set resource name only if at least one of its
-                        // instances has available capacity.
     auto instances = allocated_resources->Get(resource_id);
-    for (size_t inst_idx = 0; inst_idx < instances.size(); inst_idx++) {
-      if (instances[inst_idx] > 0.) {
-        if (first) {
-          resource = reply->add_resource_mapping();
-          resource->set_name(resource_id.Binary());
-          first = false;
+    for (const auto &reply_callback : reply_callbacks) {
+      ::ray::rpc::ResourceMapEntry *resource = nullptr;
+      for (size_t inst_idx = 0; inst_idx < instances.size(); inst_idx++) {
+        if (instances[inst_idx] > 0.) {
+          // Set resource name only if at least one of its instances has available
+          // capacity.
+          if (resource == nullptr) {
+            resource = reply_callback.reply_->add_resource_mapping();
+            resource->set_name(resource_id.Binary());
+          }
+          auto rid = resource->add_resource_ids();
+          rid->set_index(inst_idx);
+          rid->set_quantity(instances[inst_idx].Double());
         }
-        auto rid = resource->add_resource_ids();
-        rid->set_index(inst_idx);
-        rid->set_quantity(instances[inst_idx].Double());
       }
     }
   }
-  // Send the result back.
-  send_reply_callback(Status::OK(), nullptr, nullptr);
+  // Send the result back to the clients.
+  for (const auto &reply_callback : reply_callbacks) {
+    reply_callback.send_reply_callback_(Status::OK(), nullptr, nullptr);
+  }
 }
 
 void LocalLeaseManager::ClearWorkerBacklog(const WorkerID &worker_id) {
@@ -1258,5 +1261,41 @@ void LocalLeaseManager::DebugStr(std::stringstream &buffer) const {
   }
 }
 
+bool LocalLeaseManager::IsLeaseQueued(const SchedulingClass &scheduling_class,
+                                      const LeaseID &lease_id) const {
+  if (waiting_leases_index_.contains(lease_id)) {
+    return true;
+  }
+  auto leases_to_grant_it = leases_to_grant_.find(scheduling_class);
+  if (leases_to_grant_it != leases_to_grant_.end()) {
+    for (const auto &work : leases_to_grant_it->second) {
+      if (work->lease_.GetLeaseSpecification().LeaseId() == lease_id) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+bool LocalLeaseManager::AddReplyCallback(const SchedulingClass &scheduling_class,
+                                         const LeaseID &lease_id,
+                                         rpc::SendReplyCallback send_reply_callback,
+                                         rpc::RequestWorkerLeaseReply *reply) {
+  if (leases_to_grant_.contains(scheduling_class)) {
+    for (const auto &work : leases_to_grant_[scheduling_class]) {
+      if (work->lease_.GetLeaseSpecification().LeaseId() == lease_id) {
+        work->reply_callbacks_.emplace_back(std::move(send_reply_callback), reply);
+        return true;
+      }
+    }
+  }
+  auto it = waiting_leases_index_.find(lease_id);
+  if (it != waiting_leases_index_.end()) {
+    (*it->second)->reply_callbacks_.emplace_back(std::move(send_reply_callback), reply);
+    return true;
+  }
+  return false;
+}
+
 }  // namespace raylet
 }  // namespace ray

src/ray/raylet/local_lease_manager.h

@@ -189,6 +189,14 @@ class LocalLeaseManager : public LocalLeaseManagerInterface {
     return num_unschedulable_lease_spilled_;
   }
 
+  bool IsLeaseQueued(const SchedulingClass &scheduling_class,
+                     const LeaseID &lease_id) const override;
+
+  bool AddReplyCallback(const SchedulingClass &scheduling_class,
+                        const LeaseID &lease_id,
+                        rpc::SendReplyCallback send_reply_callback,
+                        rpc::RequestWorkerLeaseReply *reply) override;
+
  private:
   struct SchedulingClassInfo;
 
@@ -248,8 +256,7 @@ class LocalLeaseManager : public LocalLeaseManagerInterface {
       absl::flat_hash_map<LeaseID, std::shared_ptr<WorkerInterface>> &leased_workers_,
       const std::shared_ptr<TaskResourceInstances> &allocated_instances,
       const RayLease &lease,
-      rpc::RequestWorkerLeaseReply *reply,
-      rpc::SendReplyCallback send_reply_callback);
+      const std::vector<internal::ReplyCallback> &reply_callbacks);
 
   void Spillback(const NodeID &spillback_to, const std::shared_ptr<internal::Work> &work);
 

src/ray/raylet/node_manager.cc

@@ -1806,7 +1806,7 @@ void NodeManager::HandleRequestWorkerLease(rpc::RequestWorkerLeaseRequest reques
                                            rpc::SendReplyCallback send_reply_callback) {
   auto lease_id = LeaseID::FromBinary(request.lease_spec().lease_id());
   // If the lease is already granted, this is a retry and forward the address of the
-  // already leased worker to use.
+  // already leased worker to use
   if (leased_workers_.contains(lease_id)) {
     const auto &worker = leased_workers_[lease_id];
     RAY_LOG(DEBUG) << "Lease " << lease_id
@@ -1844,9 +1844,6 @@ void NodeManager::HandleRequestWorkerLease(rpc::RequestWorkerLeaseRequest reques
     actor_id = lease.GetLeaseSpecification().ActorId();
   }
 
-  const auto &lease_spec = lease.GetLeaseSpecification();
-  worker_pool_.PrestartWorkers(lease_spec, request.backlog_size());
-
   auto send_reply_callback_wrapper =
       [this, is_actor_creation_task, actor_id, reply, send_reply_callback](
           Status status, std::function<void()> success, std::function<void()> failure) {
@@ -1877,11 +1874,34 @@ void NodeManager::HandleRequestWorkerLease(rpc::RequestWorkerLeaseRequest reques
         send_reply_callback(status, std::move(success), std::move(failure));
       };
 
-  cluster_lease_manager_.QueueAndScheduleLease(std::move(lease),
-                                               request.grant_or_reject(),
-                                               request.is_selected_based_on_locality(),
-                                               reply,
-                                               std::move(send_reply_callback_wrapper));
+  if (cluster_lease_manager_.IsLeaseQueued(
+          lease.GetLeaseSpecification().GetSchedulingClass(), lease_id)) {
+    RAY_CHECK(cluster_lease_manager_.AddReplyCallback(
+        lease.GetLeaseSpecification().GetSchedulingClass(),
+        lease_id,
+        std::move(send_reply_callback_wrapper),
+        reply));
+    return;
+  }
+
+  if (local_lease_manager_.IsLeaseQueued(
+          lease.GetLeaseSpecification().GetSchedulingClass(), lease_id)) {
+    RAY_CHECK(local_lease_manager_.AddReplyCallback(
+        lease.GetLeaseSpecification().GetSchedulingClass(),
+        lease_id,
+        std::move(send_reply_callback_wrapper),
+        reply));
+    return;
+  }
+
+  const auto &lease_spec = lease.GetLeaseSpecification();
+  worker_pool_.PrestartWorkers(lease_spec, request.backlog_size());
+
+  cluster_lease_manager_.QueueAndScheduleLease(
+      std::move(lease),
+      request.grant_or_reject(),
+      request.is_selected_based_on_locality(),
+      {internal::ReplyCallback(std::move(send_reply_callback_wrapper), reply)});
 }
 
 void NodeManager::HandlePrestartWorkers(rpc::PrestartWorkersRequest request,
@@ -2197,8 +2217,7 @@ void NodeManager::HandleDrainRaylet(rpc::DrainRayletRequest request,
       cluster_lease_manager_.QueueAndScheduleLease(work->lease_,
                                                    work->grant_or_reject_,
                                                    work->is_selected_based_on_locality_,
-                                                   work->reply_,
-                                                   work->send_reply_callback_);
+                                                   work->reply_callbacks_);
     }
   }
 }