Add actor checkpoint methods to gcs server actor info handler (#6663)

This commit is contained in:
fangfengbin
2020-01-02 19:31:54 +08:00
committed by Hao Chen
parent a7e9d63979
commit a13781d70e
6 changed files with 217 additions and 26 deletions
@@ -1,5 +1,4 @@
#include "actor_info_handler_impl.h"
#include <assert.h>
#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<ActorTableData> &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<ActorCheckpointData>();
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<ActorCheckpointData> &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<ActorCheckpointIdData> &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
@@ -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_;
};
@@ -106,6 +106,49 @@ class GcsServerTest : public RedisServiceManagerForTest {
return actor_table_data;
}
bool AddActorCheckpoint(const rpc::AddActorCheckpointRequest &request) {
std::promise<bool> 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<bool> 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<bool> 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<bool> 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) {
+30
View File
@@ -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 {
+20 -24
View File
@@ -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<ReportHeartbeatReply> &callback) {
client_call_manager_
.CreateCall<NodeInfoGcsService, ReportHeartbeatRequest, ReportHeartbeatReply>(
*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<ReportBatchHeartbeatReply> &callback) {
client_call_manager_.CreateCall<NodeInfoGcsService, ReportBatchHeartbeatRequest,
ReportBatchHeartbeatReply>(
*node_info_stub_, &NodeInfoGcsService::Stub::PrepareAsyncReportBatchHeartbeat,
request, callback);
}
private:
/// The gRPC-generated stub.
std::unique_ptr<GrpcClient<JobInfoGcsService>> job_info_grpc_client_;
+15
View File
@@ -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: