Add task lease method to task info handler (#6710)

* add task lease methods to task info handler

* rebase master

src/ray/gcs/gcs_server/task_info_handler_impl.cc

@@ -69,5 +69,30 @@ void DefaultTaskInfoHandler::HandleDeleteTasks(const DeleteTasksRequest &request
   RAY_LOG(DEBUG) << "Finished deleting tasks, task id list size = " << task_ids.size();
 }
 
+void DefaultTaskInfoHandler::HandleAddTaskLease(const AddTaskLeaseRequest &request,
+                                                AddTaskLeaseReply *reply,
+                                                SendReplyCallback send_reply_callback) {
+  TaskID task_id = TaskID::FromBinary(request.task_lease_data().task_id());
+  ClientID node_id = ClientID::FromBinary(request.task_lease_data().node_manager_id());
+  RAY_LOG(DEBUG) << "Adding task lease, task id = " << task_id
+                 << ", node id = " << node_id;
+  auto task_lease_data = std::make_shared<TaskLeaseData>();
+  task_lease_data->CopyFrom(request.task_lease_data());
+  auto on_done = [task_id, node_id, request, send_reply_callback](Status status) {
+    if (!status.ok()) {
+      RAY_LOG(ERROR) << "Failed to add task lease, task id = " << task_id
+                     << ", node id = " << node_id;
+    }
+    send_reply_callback(status, nullptr, nullptr);
+  };
+
+  Status status = gcs_client_.Tasks().AsyncAddTaskLease(task_lease_data, on_done);
+  if (!status.ok()) {
+    on_done(status);
+  }
+  RAY_LOG(DEBUG) << "Finished adding task lease, task id = " << task_id
+                 << ", node id = " << node_id;
+}
+
 }  // namespace rpc
 }  // namespace ray

src/ray/gcs/gcs_server/task_info_handler_impl.h

@@ -22,6 +22,9 @@ class DefaultTaskInfoHandler : public rpc::TaskInfoHandler {
   void HandleDeleteTasks(const DeleteTasksRequest &request, DeleteTasksReply *reply,
                          SendReplyCallback send_reply_callback) override;
 
+  void HandleAddTaskLease(const AddTaskLeaseRequest &request, AddTaskLeaseReply *reply,
+                          SendReplyCallback send_reply_callback) override;
+
  private:
   gcs::RedisGcsClient &gcs_client_;
 };

src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc

@@ -324,6 +324,16 @@ class GcsServerTest : public RedisServiceManagerForTest {
     return WaitReady(promise.get_future(), timeout_ms_);
   }
 
+  bool AddTaskLease(const rpc::AddTaskLeaseRequest &request) {
+    std::promise<bool> promise;
+    client_->AddTaskLease(
+        request, [&promise](const Status &status, const rpc::AddTaskLeaseReply &reply) {
+          RAY_CHECK_OK(status);
+          promise.set_value(true);
+        });
+    return WaitReady(promise.get_future(), timeout_ms_);
+  }
+
   bool WaitReady(const std::future<bool> &future, uint64_t timeout_ms) {
     auto status = future.wait_for(std::chrono::milliseconds(timeout_ms));
     return status == std::future_status::ready;
@@ -370,6 +380,14 @@ class GcsServerTest : public RedisServiceManagerForTest {
     return task_table_data;
   }
 
+  rpc::TaskLeaseData GenTaskLeaseData(const std::string &task_id,
+                                      const std::string &node_id) {
+    rpc::TaskLeaseData task_lease_data;
+    task_lease_data.set_task_id(task_id);
+    task_lease_data.set_node_manager_id(node_id);
+    return task_lease_data;
+  }
+
  protected:
   // Gcs server
   std::unique_ptr<gcs::GcsServer> gcs_server_;
@@ -546,6 +564,14 @@ TEST_F(GcsServerTest, TestTaskInfo) {
   ASSERT_TRUE(DeleteTasks(delete_tasks_request));
   result = GetTask(task_id.Binary());
   ASSERT_TRUE(!result.has_task());
+
+  // Add task lease
+  ClientID node_id = ClientID::FromRandom();
+  rpc::TaskLeaseData task_lease_data =
+      GenTaskLeaseData(task_id.Binary(), node_id.Binary());
+  rpc::AddTaskLeaseRequest add_task_lease_request;
+  add_task_lease_request.mutable_task_lease_data()->CopyFrom(task_lease_data);
+  ASSERT_TRUE(AddTaskLease(add_task_lease_request));
 }
 
 }  // namespace ray

src/ray/protobuf/gcs_service.proto

@@ -241,6 +241,13 @@ message DeleteTasksRequest {
 message DeleteTasksReply {
 }
 
+message AddTaskLeaseRequest {
+  TaskLeaseData task_lease_data = 1;
+}
+
+message AddTaskLeaseReply {
+}
+
 // Service for task info access.
 service TaskInfoGcsService {
   // Add a task to GCS Service.
@@ -249,4 +256,6 @@ service TaskInfoGcsService {
   rpc GetTask(GetTaskRequest) returns (GetTaskReply);
   // Delete tasks from GCS Service.
   rpc DeleteTasks(DeleteTasksRequest) returns (DeleteTasksReply);
+  // Add a task lease to GCS Service.
+  rpc AddTaskLease(AddTaskLeaseRequest) returns (AddTaskLeaseReply);
 }

src/ray/rpc/gcs_server/gcs_rpc_client.h

@@ -102,6 +102,9 @@ class GcsRpcClient {
   /// Delete tasks from GCS Service.
   VOID_RPC_CLIENT_METHOD(TaskInfoGcsService, DeleteTasks, task_info_grpc_client_, )
 
+  /// Add a task lease to GCS Service.
+  VOID_RPC_CLIENT_METHOD(TaskInfoGcsService, AddTaskLease, task_info_grpc_client_, )
+
  private:
   /// The gRPC-generated stub.
   std::unique_ptr<GrpcClient<JobInfoGcsService>> job_info_grpc_client_;

src/ray/rpc/gcs_server/gcs_rpc_server.h

@@ -255,6 +255,10 @@ class TaskInfoGcsServiceHandler {
   virtual void HandleDeleteTasks(const DeleteTasksRequest &request,
                                  DeleteTasksReply *reply,
                                  SendReplyCallback send_reply_callback) = 0;
+
+  virtual void HandleAddTaskLease(const AddTaskLeaseRequest &request,
+                                  AddTaskLeaseReply *reply,
+                                  SendReplyCallback send_reply_callback) = 0;
 };
 
 /// The `GrpcService` for `TaskInfoGcsService`.
@@ -277,6 +281,7 @@ class TaskInfoGrpcService : public GrpcService {
     TASK_INFO_SERVICE_RPC_HANDLER(AddTask, 1);
     TASK_INFO_SERVICE_RPC_HANDLER(GetTask, 1);
     TASK_INFO_SERVICE_RPC_HANDLER(DeleteTasks, 1);
+    TASK_INFO_SERVICE_RPC_HANDLER(AddTaskLease, 1);
   }
 
  private: