diff --git a/src/ray/gcs/gcs_server/task_info_handler_impl.cc b/src/ray/gcs/gcs_server/task_info_handler_impl.cc index 2c62466c3..3b3a7de43 100644 --- a/src/ray/gcs/gcs_server/task_info_handler_impl.cc +++ b/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(); + 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 diff --git a/src/ray/gcs/gcs_server/task_info_handler_impl.h b/src/ray/gcs/gcs_server/task_info_handler_impl.h index 0d210ac2f..7f66f257a 100644 --- a/src/ray/gcs/gcs_server/task_info_handler_impl.h +++ b/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_; }; diff --git a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc b/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc index 7b3c5848b..78a86d2af 100644 --- a/src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc +++ b/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 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 &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_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 diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index d98dcef0a..505512e4f 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/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); } diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index 0ef5550c0..d1051bfd5 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/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> job_info_grpc_client_; diff --git a/src/ray/rpc/gcs_server/gcs_rpc_server.h b/src/ray/rpc/gcs_server/gcs_rpc_server.h index c4bc63b0c..a5b8e4b86 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/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: