diff --git a/src/ray/gcs/gcs_server/actor_info_handler_impl.cc b/src/ray/gcs/gcs_server/actor_info_handler_impl.cc index 628332e3d..37088d882 100644 --- a/src/ray/gcs/gcs_server/actor_info_handler_impl.cc +++ b/src/ray/gcs/gcs_server/actor_info_handler_impl.cc @@ -1,5 +1,4 @@ #include "actor_info_handler_impl.h" -#include #include "ray/util/logging.h" namespace ray { @@ -14,7 +13,7 @@ void DefaultActorInfoHandler::HandleGetActorInfo( auto on_done = [actor_id, reply, send_reply_callback]( Status status, const boost::optional &result) { if (status.ok()) { - assert(result); + RAY_DCHECK(result); reply->mutable_actor_table_data()->CopyFrom(*result); } else { RAY_LOG(ERROR) << "Failed to get actor info: " << status.ToString() @@ -74,5 +73,83 @@ void DefaultActorInfoHandler::HandleUpdateActorInfo( RAY_LOG(DEBUG) << "Finished updating actor info, actor id = " << actor_id; } +void DefaultActorInfoHandler::HandleAddActorCheckpoint( + const AddActorCheckpointRequest &request, AddActorCheckpointReply *reply, + SendReplyCallback send_reply_callback) { + ActorID actor_id = ActorID::FromBinary(request.checkpoint_data().actor_id()); + ActorCheckpointID checkpoint_id = + ActorCheckpointID::FromBinary(request.checkpoint_data().checkpoint_id()); + RAY_LOG(DEBUG) << "Adding actor checkpoint, actor id = " << actor_id + << ", checkpoint id = " << checkpoint_id; + auto actor_checkpoint_data = std::make_shared(); + actor_checkpoint_data->CopyFrom(request.checkpoint_data()); + auto on_done = [actor_id, checkpoint_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to add actor checkpoint: " << status.ToString() + << ", actor id = " << actor_id + << ", checkpoint id = " << checkpoint_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Actors().AsyncAddCheckpoint(actor_checkpoint_data, on_done); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished adding actor checkpoint, actor id = " << actor_id + << ", checkpoint id = " << checkpoint_id; +} + +void DefaultActorInfoHandler::HandleGetActorCheckpoint( + const GetActorCheckpointRequest &request, GetActorCheckpointReply *reply, + SendReplyCallback send_reply_callback) { + ActorCheckpointID checkpoint_id = + ActorCheckpointID::FromBinary(request.checkpoint_id()); + RAY_LOG(DEBUG) << "Getting actor checkpoint, checkpoint id = " << checkpoint_id; + auto on_done = [checkpoint_id, reply, send_reply_callback]( + Status status, const boost::optional &result) { + if (status.ok()) { + RAY_DCHECK(result); + reply->mutable_checkpoint_data()->CopyFrom(*result); + } else { + RAY_LOG(ERROR) << "Failed to get actor checkpoint: " << status.ToString() + << ", checkpoint id = " << checkpoint_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Actors().AsyncGetCheckpoint(checkpoint_id, on_done); + if (!status.ok()) { + on_done(status, boost::none); + } + RAY_LOG(DEBUG) << "Finished getting actor checkpoint, checkpoint id = " + << checkpoint_id; +} + +void DefaultActorInfoHandler::HandleGetActorCheckpointID( + const GetActorCheckpointIDRequest &request, GetActorCheckpointIDReply *reply, + SendReplyCallback send_reply_callback) { + ActorID actor_id = ActorID::FromBinary(request.actor_id()); + RAY_LOG(DEBUG) << "Getting actor checkpoint id, actor id = " << actor_id; + auto on_done = [actor_id, reply, send_reply_callback]( + Status status, + const boost::optional &result) { + if (status.ok()) { + RAY_DCHECK(result); + reply->mutable_checkpoint_id_data()->CopyFrom(*result); + } else { + RAY_LOG(ERROR) << "Failed to get actor checkpoint id: " << status.ToString() + << ", actor id = " << actor_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Actors().AsyncGetCheckpointID(actor_id, on_done); + if (!status.ok()) { + on_done(status, boost::none); + } + RAY_LOG(DEBUG) << "Finished getting actor checkpoint id, actor id = " << actor_id; +} + } // namespace rpc } // namespace ray diff --git a/src/ray/gcs/gcs_server/actor_info_handler_impl.h b/src/ray/gcs/gcs_server/actor_info_handler_impl.h index e5c2fa1a0..37163c482 100644 --- a/src/ray/gcs/gcs_server/actor_info_handler_impl.h +++ b/src/ray/gcs/gcs_server/actor_info_handler_impl.h @@ -24,6 +24,18 @@ class DefaultActorInfoHandler : public rpc::ActorInfoHandler { UpdateActorInfoReply *reply, SendReplyCallback send_reply_callback) override; + void HandleAddActorCheckpoint(const AddActorCheckpointRequest &request, + AddActorCheckpointReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleGetActorCheckpoint(const GetActorCheckpointRequest &request, + GetActorCheckpointReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleGetActorCheckpointID(const GetActorCheckpointIDRequest &request, + GetActorCheckpointIDReply *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 e82ba53aa..7683e095d 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 @@ -106,6 +106,49 @@ class GcsServerTest : public RedisServiceManagerForTest { return actor_table_data; } + bool AddActorCheckpoint(const rpc::AddActorCheckpointRequest &request) { + std::promise promise; + client_->AddActorCheckpoint( + request, + [&promise](const Status &status, const rpc::AddActorCheckpointReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + + rpc::ActorCheckpointData GetActorCheckpoint(const std::string &checkpoint_id) { + rpc::GetActorCheckpointRequest request; + request.set_checkpoint_id(checkpoint_id); + rpc::ActorCheckpointData checkpoint_data; + std::promise promise; + client_->GetActorCheckpoint( + request, [&checkpoint_data, &promise](const Status &status, + const rpc::GetActorCheckpointReply &reply) { + RAY_CHECK_OK(status); + checkpoint_data.CopyFrom(reply.checkpoint_data()); + promise.set_value(true); + }); + EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_)); + return checkpoint_data; + } + + rpc::ActorCheckpointIdData GetActorCheckpointID(const std::string &actor_id) { + rpc::GetActorCheckpointIDRequest request; + request.set_actor_id(actor_id); + rpc::ActorCheckpointIdData checkpoint_id_data; + std::promise promise; + client_->GetActorCheckpointID( + request, [&checkpoint_id_data, &promise]( + const Status &status, const rpc::GetActorCheckpointIDReply &reply) { + RAY_CHECK_OK(status); + checkpoint_id_data.CopyFrom(reply.checkpoint_id_data()); + promise.set_value(true); + }); + EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_)); + return checkpoint_id_data; + } + bool RegisterNode(const rpc::RegisterNodeRequest &request) { std::promise promise; client_->RegisterNode( @@ -280,6 +323,24 @@ TEST_F(GcsServerTest, TestActorInfo) { result = GetActorInfo(actor_table_data.actor_id()); ASSERT_TRUE(result.state() == rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD); + + // Add actor checkpoint + ActorCheckpointID checkpoint_id = ActorCheckpointID::FromRandom(); + rpc::ActorCheckpointData checkpoint; + checkpoint.set_actor_id(actor_table_data.actor_id()); + checkpoint.set_checkpoint_id(checkpoint_id.Binary()); + checkpoint.set_execution_dependency(checkpoint_id.Binary()); + + rpc::AddActorCheckpointRequest add_actor_checkpoint_request; + add_actor_checkpoint_request.mutable_checkpoint_data()->CopyFrom(checkpoint); + ASSERT_TRUE(AddActorCheckpoint(add_actor_checkpoint_request)); + rpc::ActorCheckpointData checkpoint_result = GetActorCheckpoint(checkpoint_id.Binary()); + ASSERT_TRUE(checkpoint_result.actor_id() == actor_table_data.actor_id()); + ASSERT_TRUE(checkpoint_result.checkpoint_id() == checkpoint_id.Binary()); + rpc::ActorCheckpointIdData checkpoint_id_result = + GetActorCheckpointID(actor_table_data.actor_id()); + ASSERT_TRUE(checkpoint_id_result.actor_id() == actor_table_data.actor_id()); + ASSERT_TRUE(checkpoint_id_result.checkpoint_ids_size() == 1); } TEST_F(GcsServerTest, TestJobInfo) { diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index 77883f76d..1916d0783 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -56,6 +56,29 @@ message UpdateActorInfoRequest { message UpdateActorInfoReply { } +message AddActorCheckpointRequest { + ActorCheckpointData checkpoint_data = 1; +} + +message AddActorCheckpointReply { +} + +message GetActorCheckpointRequest { + bytes checkpoint_id = 1; +} + +message GetActorCheckpointReply { + ActorCheckpointData checkpoint_data = 1; +} + +message GetActorCheckpointIDRequest { + bytes actor_id = 1; +} + +message GetActorCheckpointIDReply { + ActorCheckpointIdData checkpoint_id_data = 1; +} + // Service for actor info access. service ActorInfoGcsService { // Get actor data from GCS Service. @@ -64,6 +87,13 @@ service ActorInfoGcsService { rpc RegisterActorInfo(RegisterActorInfoRequest) returns (RegisterActorInfoReply); // Update actor info in GCS Service. rpc UpdateActorInfo(UpdateActorInfoRequest) returns (UpdateActorInfoReply); + // Add actor checkpoint data to GCS Service. + rpc AddActorCheckpoint(AddActorCheckpointRequest) returns (AddActorCheckpointReply); + // Get actor checkpoint data from GCS Service. + rpc GetActorCheckpoint(GetActorCheckpointRequest) returns (GetActorCheckpointReply); + // Get actor checkpoint id data from GCS Service. + rpc GetActorCheckpointID(GetActorCheckpointIDRequest) + returns (GetActorCheckpointIDReply); } message RegisterNodeRequest { diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index 9e730c4da..4ff5e5528 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_client.h @@ -52,6 +52,18 @@ class GcsRpcClient { VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, UpdateActorInfo, request, callback, actor_info_grpc_client_) + /// Add actor checkpoint data to GCS Service. + VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, AddActorCheckpoint, request, callback, + actor_info_grpc_client_) + + /// Get actor checkpoint data from GCS Service. + VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, GetActorCheckpoint, request, callback, + actor_info_grpc_client_) + + /// Get actor checkpoint id data from GCS Service. + VOID_RPC_CLIENT_METHOD(ActorInfoGcsService, GetActorCheckpointID, request, callback, + actor_info_grpc_client_) + /// Register a node to GCS Service. VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, RegisterNode, request, callback, node_info_grpc_client_) @@ -64,6 +76,14 @@ class GcsRpcClient { VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, GetAllNodeInfo, request, callback, node_info_grpc_client_) + /// Report heartbeat of a node to GCS Service. + VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, ReportHeartbeat, request, callback, + node_info_grpc_client_) + + /// Report batch heartbeat to GCS Service. + VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, ReportBatchHeartbeat, request, callback, + node_info_grpc_client_) + /// Get object's locations from GCS Service. VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, GetObjectLocations, request, callback, object_info_grpc_client_) @@ -76,30 +96,6 @@ class GcsRpcClient { VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, RemoveObjectLocation, request, callback, object_info_grpc_client_) - /// Report heartbeat of a node to GCS Service. - /// - /// \param request The request message. - /// \param callback The callback function that handles reply from server. - void ReportHeartbeat(const ReportHeartbeatRequest &request, - const ClientCallback &callback) { - client_call_manager_ - .CreateCall( - *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportHeartbeat, - request, callback); - } - - /// Report batch heartbeat to GCS Service. - /// - /// \param request The request message. - /// \param callback The callback function that handles reply from server. - void ReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request, - const ClientCallback &callback) { - client_call_manager_.CreateCall( - *node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportBatchHeartbeat, - request, callback); - } - 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 bd3b5958f..fe1e650fb 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_server.h @@ -100,6 +100,18 @@ class ActorInfoHandler { virtual void HandleUpdateActorInfo(const UpdateActorInfoRequest &request, UpdateActorInfoReply *reply, SendReplyCallback send_reply_callback) = 0; + + virtual void HandleAddActorCheckpoint(const AddActorCheckpointRequest &request, + AddActorCheckpointReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleGetActorCheckpoint(const GetActorCheckpointRequest &request, + GetActorCheckpointReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleGetActorCheckpointID(const GetActorCheckpointIDRequest &request, + GetActorCheckpointIDReply *reply, + SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcService` for `ActorInfoGcsService`. @@ -122,6 +134,9 @@ class ActorInfoGrpcService : public GrpcService { ACTOR_INFO_SERVICE_RPC_HANDLER(GetActorInfo, 1); ACTOR_INFO_SERVICE_RPC_HANDLER(RegisterActorInfo, 1); ACTOR_INFO_SERVICE_RPC_HANDLER(UpdateActorInfo, 1); + ACTOR_INFO_SERVICE_RPC_HANDLER(AddActorCheckpoint, 1); + ACTOR_INFO_SERVICE_RPC_HANDLER(GetActorCheckpoint, 1); + ACTOR_INFO_SERVICE_RPC_HANDLER(GetActorCheckpointID, 1); } private: