mirror of
https://github.com/wassname/ray.git
synced 2026-08-12 12:20:11 +08:00
Add actor checkpoint methods to gcs server actor info handler (#6663)
This commit is contained in:
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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_;
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user