diff --git a/BUILD.bazel b/BUILD.bazel index 32fd1ecf6..7eb7125fc 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -662,7 +662,13 @@ cc_test( cc_test( name = "gcs_server_rpc_test", srcs = ["src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc"], + args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"], copts = COPTS, + data = [ + "//:libray_redis_module.so", + "//:redis-cli", + "//:redis-server", + ], deps = [ ":gcs_server_lib", "@com_google_googletest//:gtest_main", diff --git a/src/ray/gcs/gcs_server/actor_info_handler_impl.cc b/src/ray/gcs/gcs_server/actor_info_handler_impl.cc new file mode 100644 index 000000000..bcaaaa14d --- /dev/null +++ b/src/ray/gcs/gcs_server/actor_info_handler_impl.cc @@ -0,0 +1,84 @@ +#include "actor_info_handler_impl.h" +#include +#include "ray/util/logging.h" + +namespace ray { +namespace rpc { + +void DefaultActorInfoHandler::HandleGetActorInfo( + const rpc::GetActorInfoRequest &request, rpc::GetActorInfoReply *reply, + rpc::SendReplyCallback send_reply_callback) { + ActorID actor_id = ActorID::FromBinary(request.actor_id()); + RAY_LOG(DEBUG) << "Getting actor info, actor id = " << actor_id; + + auto on_done = [actor_id, reply, send_reply_callback]( + Status status, const boost::optional &result) { + if (status.ok()) { + assert(result); + reply->mutable_actor_table_data()->CopyFrom(*result); + } else { + RAY_LOG(ERROR) << "Failed to get actor info: " << status.ToString() + << ", actor id = " << actor_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Actors().AsyncGet(actor_id, on_done); + if (!status.ok()) { + on_done(status, boost::none); + } + RAY_LOG(DEBUG) << "Finished getting actor info, actor id = " << actor_id; +} + +void DefaultActorInfoHandler::HandleRegisterActorInfo( + const rpc::RegisterActorInfoRequest &request, rpc::RegisterActorInfoReply *reply, + rpc::SendReplyCallback send_reply_callback) { + ActorID actor_id = ActorID::FromBinary(request.actor_table_data().actor_id()); + RAY_LOG(DEBUG) << "Registering actor info, actor id = " << actor_id; + auto actor_table_data = std::make_shared(); + actor_table_data->CopyFrom(request.actor_table_data()); + auto on_done = [actor_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to register actor info: " << status.ToString() + << ", actor id = " << actor_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Actors().AsyncRegister( + actor_table_data, [send_reply_callback](Status status) { + send_reply_callback(status, nullptr, nullptr); + }); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished registering actor info, actor id = " << actor_id; +} + +void DefaultActorInfoHandler::HandleUpdateActorInfo( + const rpc::UpdateActorInfoRequest &request, rpc::UpdateActorInfoReply *reply, + rpc::SendReplyCallback send_reply_callback) { + ActorID actor_id = ActorID::FromBinary(request.actor_id()); + RAY_LOG(DEBUG) << "Updating actor info, actor id = " << actor_id; + auto actor_table_data = std::make_shared(); + actor_table_data->CopyFrom(request.actor_table_data()); + auto on_done = [actor_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to update actor info: " << status.ToString() + << ", actor id = " << actor_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Actors().AsyncUpdate( + actor_id, actor_table_data, [send_reply_callback](Status status) { + send_reply_callback(status, nullptr, nullptr); + }); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished updating actor info, 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 new file mode 100644 index 000000000..e5c2fa1a0 --- /dev/null +++ b/src/ray/gcs/gcs_server/actor_info_handler_impl.h @@ -0,0 +1,34 @@ +#ifndef RAY_GCS_ACTOR_INFO_HANDLER_IMPL_H +#define RAY_GCS_ACTOR_INFO_HANDLER_IMPL_H + +#include "ray/gcs/redis_gcs_client.h" +#include "ray/rpc/gcs_server/gcs_rpc_server.h" + +namespace ray { +namespace rpc { + +/// This implementation class of `ActorInfoHandler`. +class DefaultActorInfoHandler : public rpc::ActorInfoHandler { + public: + explicit DefaultActorInfoHandler(gcs::RedisGcsClient &gcs_client) + : gcs_client_(gcs_client) {} + + void HandleGetActorInfo(const GetActorInfoRequest &request, GetActorInfoReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleRegisterActorInfo(const RegisterActorInfoRequest &request, + RegisterActorInfoReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleUpdateActorInfo(const UpdateActorInfoRequest &request, + UpdateActorInfoReply *reply, + SendReplyCallback send_reply_callback) override; + + private: + gcs::RedisGcsClient &gcs_client_; +}; + +} // namespace rpc +} // namespace ray + +#endif // RAY_GCS_ACTOR_INFO_HANDLER_IMPL_H diff --git a/src/ray/gcs/gcs_server/gcs_server.cc b/src/ray/gcs/gcs_server/gcs_server.cc index 905d6304a..f4a647584 100644 --- a/src/ray/gcs/gcs_server/gcs_server.cc +++ b/src/ray/gcs/gcs_server/gcs_server.cc @@ -1,4 +1,5 @@ #include "gcs_server.h" +#include "actor_info_handler_impl.h" #include "job_info_handler_impl.h" namespace ray { @@ -20,6 +21,11 @@ void GcsServer::Start() { job_info_service_.reset(new rpc::JobInfoGrpcService(main_service_, *job_info_handler_)); rpc_server_.RegisterService(*job_info_service_); + actor_info_handler_ = InitActorInfoHandler(); + actor_info_service_.reset( + new rpc::ActorInfoGrpcService(main_service_, *actor_info_handler_)); + rpc_server_.RegisterService(*actor_info_service_); + // Run rpc server. rpc_server_.Run(); @@ -40,14 +46,20 @@ void GcsServer::Stop() { void GcsServer::InitBackendClient() { GcsClientOptions options(config_.redis_address, config_.redis_port, - config_.redis_password); + config_.redis_password, config_.is_test); redis_gcs_client_ = std::make_shared(options); auto status = redis_gcs_client_->Connect(main_service_); RAY_CHECK(status.ok()) << "Failed to init redis gcs client as " << status; } std::unique_ptr GcsServer::InitJobInfoHandler() { - return std::unique_ptr(new rpc::DefaultJobInfoHandler()); + return std::unique_ptr( + new rpc::DefaultJobInfoHandler(*redis_gcs_client_)); +} + +std::unique_ptr GcsServer::InitActorInfoHandler() { + return std::unique_ptr( + new rpc::DefaultActorInfoHandler(*redis_gcs_client_)); } } // namespace gcs diff --git a/src/ray/gcs/gcs_server/gcs_server.h b/src/ray/gcs/gcs_server/gcs_server.h index da811b771..5bc7b408e 100644 --- a/src/ray/gcs/gcs_server/gcs_server.h +++ b/src/ray/gcs/gcs_server/gcs_server.h @@ -17,6 +17,7 @@ struct GcsServerConfig { std::string redis_address; uint16_t redis_port = 6379; bool retry_redis = true; + bool is_test = false; }; /// The GcsServer will take over all requests from ServiceBasedGcsClient and transparent @@ -48,6 +49,9 @@ class GcsServer { /// The job info handler virtual std::unique_ptr InitJobInfoHandler(); + /// The actor info handler + virtual std::unique_ptr InitActorInfoHandler(); + private: /// Gcs server configuration GcsServerConfig config_; @@ -58,6 +62,9 @@ class GcsServer { /// Job info handler and service std::unique_ptr job_info_handler_; std::unique_ptr job_info_service_; + /// Actor info handler and service + std::unique_ptr actor_info_handler_; + std::unique_ptr actor_info_service_; /// Backend client std::shared_ptr redis_gcs_client_; }; diff --git a/src/ray/gcs/gcs_server/job_info_handler_impl.cc b/src/ray/gcs/gcs_server/job_info_handler_impl.cc index 7a62b9e22..85da8e6b6 100644 --- a/src/ray/gcs/gcs_server/job_info_handler_impl.cc +++ b/src/ray/gcs/gcs_server/job_info_handler_impl.cc @@ -5,15 +5,54 @@ namespace rpc { void DefaultJobInfoHandler::HandleAddJob(const rpc::AddJobRequest &request, rpc::AddJobReply *reply, rpc::SendReplyCallback send_reply_callback) { - RAY_LOG(DEBUG) << "Received new job ..."; - // TODO(zsl): The detailed implementation will be committed in next PR. + JobID job_id = JobID::FromBinary(request.data().job_id()); + RAY_LOG(DEBUG) << "Adding job, job id = " << job_id + << ", driver pid = " << request.data().driver_pid(); + auto job_table_data = std::make_shared(); + job_table_data->CopyFrom(request.data()); + auto on_done = [job_id, request, reply, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to add job, job id = " << job_id + << ", driver pid = " << request.data().driver_pid(); + } + reply->set_success(status.ok()); + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Jobs().AsyncAdd( + job_table_data, [reply, send_reply_callback](Status status) { + reply->set_success(status.ok()); + send_reply_callback(status, nullptr, nullptr); + }); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished adding job, job id = " << job_id + << ", driver pid = " << request.data().driver_pid(); } void DefaultJobInfoHandler::HandleMarkJobFinished( const rpc::MarkJobFinishedRequest &request, rpc::MarkJobFinishedReply *reply, rpc::SendReplyCallback send_reply_callback) { - RAY_LOG(DEBUG) << "Mark job as finished ..."; - // TODO(zsl): The detailed implementation will be committed in next PR. + JobID job_id = JobID::FromBinary(request.job_id()); + RAY_LOG(DEBUG) << "Marking job state, job id = " << job_id; + auto on_done = [job_id, reply, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to mark job state, job id = " << job_id; + } + reply->set_success(status.ok()); + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Jobs().AsyncMarkFinished( + job_id, [reply, send_reply_callback](Status status) { + reply->set_success(status.ok()); + send_reply_callback(status, nullptr, nullptr); + }); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished marking job state, job id = " << job_id; } } // namespace rpc } // namespace ray diff --git a/src/ray/gcs/gcs_server/job_info_handler_impl.h b/src/ray/gcs/gcs_server/job_info_handler_impl.h index c49d6891d..2e6c21318 100644 --- a/src/ray/gcs/gcs_server/job_info_handler_impl.h +++ b/src/ray/gcs/gcs_server/job_info_handler_impl.h @@ -1,24 +1,30 @@ #ifndef RAY_GCS_JOB_INFO_HANDLER_IMPL_H #define RAY_GCS_JOB_INFO_HANDLER_IMPL_H +#include "ray/gcs/redis_gcs_client.h" #include "ray/rpc/gcs_server/gcs_rpc_server.h" namespace ray { namespace rpc { -/// This class is used to implement `JobInfoHandler`, but only two logs have been printed. -/// The detailed implementation is reflected in the following PR. +/// This implementation class of `JobInfoHandler`. class DefaultJobInfoHandler : public rpc::JobInfoHandler { public: + explicit DefaultJobInfoHandler(gcs::RedisGcsClient &gcs_client) + : gcs_client_(gcs_client) {} + void HandleAddJob(const AddJobRequest &request, AddJobReply *reply, SendReplyCallback send_reply_callback) override; void HandleMarkJobFinished(const MarkJobFinishedRequest &request, MarkJobFinishedReply *reply, SendReplyCallback send_reply_callback) override; + + private: + gcs::RedisGcsClient &gcs_client_; }; } // namespace rpc } // namespace ray -#endif +#endif // RAY_GCS_JOB_INFO_HANDLER_IMPL_H 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 194c2f1ee..dbe71f744 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 @@ -1,93 +1,221 @@ -#include -#include "gmock/gmock.h" #include "gtest/gtest.h" +#include "ray/gcs/gcs_server/actor_info_handler_impl.h" #include "ray/gcs/gcs_server/gcs_server.h" +#include "ray/gcs/gcs_server/job_info_handler_impl.h" #include "ray/rpc/gcs_server/gcs_rpc_client.h" +#include "ray/util/test_util.h" namespace ray { -class GcsServerRpcTest : public ::testing::Test { - protected: - template - class MockedGcsServer : public gcs::GcsServer { - public: - explicit MockedGcsServer(const gcs::GcsServerConfig &config) - : gcs::GcsServer(config) {} - protected: - void InitBackendClient() override {} +static std::string redis_server_executable; +static std::string redis_client_executable; +static std::string libray_redis_module_path; - std::unique_ptr InitJobInfoHandler() override { - return std::unique_ptr(new T); - } - }; -}; +class GcsServerTest : public ::testing::Test { + public: + using CallFunction = std::function &promise)>; -TEST_F(GcsServerRpcTest, JobInfo) { - class MockedJobInfoHandler : public rpc::JobInfoHandler { - public: - void HandleAddJob(const rpc::AddJobRequest &request, rpc::AddJobReply *reply, - rpc::SendReplyCallback send_reply_callback) override { - send_reply_callback(Status::OK(), nullptr, nullptr); - } - - void HandleMarkJobFinished(const rpc::MarkJobFinishedRequest &request, - rpc::MarkJobFinishedReply *reply, - rpc::SendReplyCallback send_reply_callback) override { - send_reply_callback(Status::OK(), nullptr, nullptr); - } - }; - - gcs::GcsServerConfig config; - config.grpc_server_port = 0; - config.grpc_server_name = "MockedGcsServer"; - config.grpc_server_thread_num = 1; - MockedGcsServer server(config); - std::thread([&server] { server.Start(); }).detach(); - - // Wait until server starts listening. - while (server.GetPort() == 0) { - std::this_thread::sleep_for(std::chrono::milliseconds(10)); + static void SetUpTestCase() { + std::string start_redis_command = redis_server_executable + + " --loglevel warning --loadmodule " + + libray_redis_module_path + " --port 6379 &"; + RAY_LOG(INFO) << "Start redis command is: " << start_redis_command; + RAY_CHECK(system(start_redis_command.c_str()) == 0); + usleep(200 * 1000); } - boost::asio::io_context client_io_context; - boost::asio::io_context::work worker(client_io_context); - rpc::ClientCallManager client_call_manager(client_io_context); - rpc::GcsRpcClient client("0.0.0.0", server.GetPort(), client_call_manager); - std::thread([&client_io_context] { client_io_context.run(); }).detach(); + static void TearDownTestCase() { + std::string stop_redis_command = redis_client_executable + " -p 6379 shutdown"; + RAY_LOG(INFO) << "Stop redis command is: " << stop_redis_command; + RAY_CHECK(system(stop_redis_command.c_str()) == 0); + } - rpc::AddJobRequest add_job_request; - std::promise add_job_reply_promise; - auto add_job_reply_future = add_job_reply_promise.get_future(); - client.AddJob(add_job_request, [&add_job_reply_promise](const Status &status, - const rpc::AddJobReply &reply) { - if (status.ok()) { - add_job_reply_promise.set_value(reply); + void SetUp() override { + gcs::GcsServerConfig config; + config.grpc_server_port = 0; + config.grpc_server_name = "MockedGcsServer"; + config.grpc_server_thread_num = 1; + config.redis_address = "127.0.0.1"; + config.is_test = true; + gcs_server_.reset(new gcs::GcsServer(config)); + + thread_io_service_.reset(new std::thread([this] { + std::unique_ptr work( + new boost::asio::io_service::work(io_service_)); + io_service_.run(); + })); + + thread_gcs_server_.reset(new std::thread([this] { gcs_server_->Start(); })); + + // Wait until server starts listening. + while (gcs_server_->GetPort() == 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); } - }); - auto future_status = add_job_reply_future.wait_for(std::chrono::milliseconds(200)); - ASSERT_EQ(future_status, std::future_status::ready); + // Create gcs rpc client + client_call_manager_.reset(new rpc::ClientCallManager(io_service_)); + client_.reset( + new rpc::GcsRpcClient("0.0.0.0", gcs_server_->GetPort(), *client_call_manager_)); + } + + void TearDown() override { + gcs_server_->Stop(); + io_service_.stop(); + thread_io_service_->join(); + thread_gcs_server_->join(); + } + + void TestAddJob(const rpc::AddJobRequest &request) { + auto call_function = [this, request](std::promise &promise) { + client_->AddJob(request, + [&promise](const Status &status, const rpc::AddJobReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + }; + AsyncCall(call_function, timeout_ms_); + } + + void TestMarkJobFinished(const rpc::MarkJobFinishedRequest &request) { + auto call_function = [this, request](std::promise &promise) { + client_->MarkJobFinished( + request, + [&promise](const Status &status, const rpc::MarkJobFinishedReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + }; + AsyncCall(call_function, timeout_ms_); + } + + void TestRegisterActorInfo(const rpc::RegisterActorInfoRequest &request) { + auto call_function = [this, request](std::promise &promise) { + client_->RegisterActorInfo( + request, + [&promise](const Status &status, const rpc::RegisterActorInfoReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + }; + AsyncCall(call_function, timeout_ms_); + } + + void TestUpdateActorInfo(const rpc::UpdateActorInfoRequest &request) { + auto call_function = [this, request](std::promise &promise) { + client_->UpdateActorInfo( + request, + [&promise](const Status &status, const rpc::UpdateActorInfoReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + }; + AsyncCall(call_function, timeout_ms_); + } + + void TestGetActorInfo(const rpc::ActorTableData &expected) { + rpc::GetActorInfoRequest request; + request.set_actor_id(expected.actor_id()); + auto call_function = [this, request, expected](std::promise &promise) { + client_->GetActorInfo( + request, [&promise, &expected](const Status &status, + const rpc::GetActorInfoReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + ASSERT_TRUE(reply.actor_table_data().state() == expected.state()); + }); + }; + AsyncCall(call_function, timeout_ms_); + } + + void AsyncCall(const CallFunction &function, uint64_t timeout_ms) { + std::promise promise_; + auto future = promise_.get_future(); + function(promise_); + auto status = future.wait_for(std::chrono::milliseconds(timeout_ms)); + ASSERT_EQ(status, std::future_status::ready); + } + + rpc::JobTableData GenJobTableData(JobID job_id) { + rpc::JobTableData job_table_data; + job_table_data.set_job_id(job_id.Binary()); + job_table_data.set_is_dead(false); + job_table_data.set_timestamp(std::time(nullptr)); + job_table_data.set_node_manager_address("127.0.0.1"); + job_table_data.set_driver_pid(5667L); + return job_table_data; + } + + rpc::ActorTableData GenActorTableData(const JobID &job_id) { + rpc::ActorTableData actor_table_data; + ActorID actor_id = ActorID::Of(job_id, RandomTaskId(), 0); + actor_table_data.set_actor_id(actor_id.Binary()); + actor_table_data.set_job_id(job_id.Binary()); + actor_table_data.set_state( + rpc::ActorTableData_ActorState::ActorTableData_ActorState_ALIVE); + actor_table_data.set_max_reconstructions(1); + actor_table_data.set_remaining_reconstructions(1); + return actor_table_data; + } + + protected: + // Gcs server + std::unique_ptr gcs_server_; + std::unique_ptr thread_io_service_; + std::unique_ptr thread_gcs_server_; + boost::asio::io_service io_service_; + + // Gcs client + std::unique_ptr client_; + std::unique_ptr client_call_manager_; + + // Timeout waiting for gcs server reply, default is 2s + const uint64_t timeout_ms_ = 2000; +}; + +TEST_F(GcsServerTest, TestActorInfo) { + // Create actor_table_data + JobID job_id = JobID::FromInt(1); + rpc::ActorTableData actor_table_data = GenActorTableData(job_id); + + // Register actor + rpc::RegisterActorInfoRequest register_actor_info_request; + register_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data); + TestRegisterActorInfo(register_actor_info_request); + TestGetActorInfo(actor_table_data); + + // Update actor state + rpc::UpdateActorInfoRequest update_actor_info_request; + actor_table_data.set_state( + rpc::ActorTableData_ActorState::ActorTableData_ActorState_DEAD); + update_actor_info_request.set_actor_id(actor_table_data.actor_id()); + update_actor_info_request.mutable_actor_table_data()->CopyFrom(actor_table_data); + TestUpdateActorInfo(update_actor_info_request); + TestGetActorInfo(actor_table_data); +} + +TEST_F(GcsServerTest, TestJobInfo) { + // Create job_table_data + JobID job_id = JobID::FromInt(1); + rpc::JobTableData job_table_data = GenJobTableData(job_id); + + // Add job + rpc::AddJobRequest add_job_request; + add_job_request.mutable_data()->CopyFrom(job_table_data); + TestAddJob(add_job_request); + + // Mark job finished rpc::MarkJobFinishedRequest mark_job_finished_request; - std::promise mark_job_finished_reply_promise; - auto mark_job_finished_reply_future = mark_job_finished_reply_promise.get_future(); - client.MarkJobFinished( - mark_job_finished_request, - [&mark_job_finished_reply_promise](const Status &status, - const rpc::MarkJobFinishedReply &reply) { - if (status.ok()) { - mark_job_finished_reply_promise.set_value(reply); - } - }); - future_status = mark_job_finished_reply_future.wait_for(std::chrono::milliseconds(200)); - ASSERT_EQ(future_status, std::future_status::ready); - - client_io_context.stop(); - server.Stop(); + mark_job_finished_request.set_job_id(job_table_data.job_id()); + TestMarkJobFinished(mark_job_finished_request); } } // namespace ray int main(int argc, char **argv) { ::testing::InitGoogleTest(&argc, argv); + RAY_CHECK(argc == 4); + ray::redis_server_executable = argv[1]; + ray::redis_client_executable = argv[2]; + ray::libray_redis_module_path = argv[3]; return RUN_ALL_TESTS(); } diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index 6adbdd947..3c4936cb7 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -27,3 +27,41 @@ service JobInfoGcsService { // Mark job as finished to gcs server. rpc MarkJobFinished(MarkJobFinishedRequest) returns (MarkJobFinishedReply); } + +message GetActorInfoRequest { + // ID of this actor. + bytes actor_id = 1; +} + +message GetActorInfoReply { + // Data of actor. + ActorTableData actor_table_data = 1; +} + +message RegisterActorInfoRequest { + // Data of actor. + ActorTableData actor_table_data = 1; +} + +message RegisterActorInfoReply { +} + +message UpdateActorInfoRequest { + // ID of this actor. + bytes actor_id = 1; + // Data of actor. + ActorTableData actor_table_data = 2; +} + +message UpdateActorInfoReply { +} + +// Service for actor info access. +service ActorInfoGcsService { + // Get actor data from GCS Service. + rpc GetActorInfo(GetActorInfoRequest) returns (GetActorInfoReply); + // Register an actor to GCS Service. + rpc RegisterActorInfo(RegisterActorInfoRequest) returns (RegisterActorInfoReply); + // Update actor info in GCS Service. + rpc UpdateActorInfo(UpdateActorInfoRequest) returns (UpdateActorInfoReply); +} diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index 441804f45..cd825ee96 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_client.h @@ -24,7 +24,8 @@ class GcsRpcClient { : client_call_manager_(client_call_manager) { std::shared_ptr channel = grpc::CreateChannel( address + ":" + std::to_string(port), grpc::InsecureChannelCredentials()); - stub_ = JobInfoGcsService::NewStub(channel); + job_info_stub_ = JobInfoGcsService::NewStub(channel); + actor_info_stub_ = ActorInfoGcsService::NewStub(channel); }; /// Add job info to gcs server. @@ -33,7 +34,7 @@ class GcsRpcClient { /// \param callback The callback function that handles reply from server. void AddJob(const AddJobRequest &request, const ClientCallback &callback) { client_call_manager_.CreateCall( - *stub_, &JobInfoGcsService::Stub::PrepareAsyncAddJob, request, callback); + *job_info_stub_, &JobInfoGcsService::Stub::PrepareAsyncAddJob, request, callback); } /// Mark job as finished to gcs server. @@ -44,13 +45,50 @@ class GcsRpcClient { const ClientCallback &callback) { client_call_manager_ .CreateCall( - *stub_, &JobInfoGcsService::Stub::PrepareAsyncMarkJobFinished, request, - callback); + *job_info_stub_, &JobInfoGcsService::Stub::PrepareAsyncMarkJobFinished, + request, callback); + } + + /// Get actor data from GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void GetActorInfo(const GetActorInfoRequest &request, + const ClientCallback &callback) { + client_call_manager_ + .CreateCall( + *actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncGetActorInfo, + request, callback); + } + + /// Register an actor to GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void RegisterActorInfo(const RegisterActorInfoRequest &request, + const ClientCallback &callback) { + client_call_manager_.CreateCall( + *actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncRegisterActorInfo, + request, callback); + } + + /// Update actor info in GCS Service. + /// + /// \param request The request message. + /// \param callback The callback function that handles reply from server. + void UpdateActorInfo(const UpdateActorInfoRequest &request, + const ClientCallback &callback) { + client_call_manager_ + .CreateCall( + *actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncUpdateActorInfo, + request, callback); } private: /// The gRPC-generated stub. - std::unique_ptr stub_; + std::unique_ptr job_info_stub_; + std::unique_ptr actor_info_stub_; /// The `ClientCallManager` used for managing requests. ClientCallManager &client_call_manager_; @@ -59,4 +97,4 @@ class GcsRpcClient { } // namespace rpc } // namespace ray -#endif +#endif // RAY_RPC_GCS_RPC_CLIENT_H diff --git a/src/ray/rpc/gcs_server/gcs_rpc_server.h b/src/ray/rpc/gcs_server/gcs_rpc_server.h index c9f4bf15b..cff996a91 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_server.h @@ -9,6 +9,24 @@ namespace ray { namespace rpc { +#define JOB_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \ + std::unique_ptr HANDLER##_call_factory( \ + new ServerCallFactoryImpl( \ + service_, &JobInfoGcsService::AsyncService::Request##HANDLER, \ + service_handler_, &JobInfoHandler::Handle##HANDLER, cq, main_service_)); \ + server_call_factories_and_concurrencies->emplace_back( \ + std::move(HANDLER##_call_factory), CONCURRENCY); + +#define ACTOR_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \ + std::unique_ptr HANDLER##_call_factory( \ + new ServerCallFactoryImpl( \ + service_, &ActorInfoGcsService::AsyncService::Request##HANDLER, \ + service_handler_, &ActorInfoHandler::Handle##HANDLER, cq, main_service_)); \ + server_call_factories_and_concurrencies->emplace_back( \ + std::move(HANDLER##_call_factory), CONCURRENCY); + class JobInfoHandler { public: virtual ~JobInfoHandler() = default; @@ -38,21 +56,8 @@ class JobInfoGrpcService : public GrpcService { const std::unique_ptr &cq, std::vector, int>> *server_call_factories_and_concurrencies) override { - std::unique_ptr add_job_call_factory( - new ServerCallFactoryImpl( - service_, &JobInfoGcsService::AsyncService::RequestAddJob, service_handler_, - &JobInfoHandler::HandleAddJob, cq, main_service_)); - server_call_factories_and_concurrencies->emplace_back(std::move(add_job_call_factory), - 1); - - std::unique_ptr mark_job_finished_call_factory( - new ServerCallFactoryImpl( - service_, &JobInfoGcsService::AsyncService::RequestMarkJobFinished, - service_handler_, &JobInfoHandler::HandleMarkJobFinished, cq, main_service_)); - server_call_factories_and_concurrencies->emplace_back( - std::move(mark_job_finished_call_factory), 1); + JOB_INFO_SERVICE_RPC_HANDLER(AddJob, 1); + JOB_INFO_SERVICE_RPC_HANDLER(MarkJobFinished, 1); } private: @@ -62,7 +67,53 @@ class JobInfoGrpcService : public GrpcService { JobInfoHandler &service_handler_; }; +class ActorInfoHandler { + public: + virtual ~ActorInfoHandler() = default; + + virtual void HandleGetActorInfo(const GetActorInfoRequest &request, + GetActorInfoReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleRegisterActorInfo(const RegisterActorInfoRequest &request, + RegisterActorInfoReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleUpdateActorInfo(const UpdateActorInfoRequest &request, + UpdateActorInfoReply *reply, + SendReplyCallback send_reply_callback) = 0; +}; + +/// The `GrpcService` for `ActorInfoGcsService`. +class ActorInfoGrpcService : public GrpcService { + public: + /// Constructor. + /// + /// \param[in] handler The service handler that actually handle the requests. + explicit ActorInfoGrpcService(boost::asio::io_service &io_service, + ActorInfoHandler &handler) + : GrpcService(io_service), service_handler_(handler){}; + + protected: + grpc::Service &GetGrpcService() override { return service_; } + + void InitServerCallFactories( + const std::unique_ptr &cq, + std::vector, int>> + *server_call_factories_and_concurrencies) override { + ACTOR_INFO_SERVICE_RPC_HANDLER(GetActorInfo, 1); + ACTOR_INFO_SERVICE_RPC_HANDLER(RegisterActorInfo, 1); + ACTOR_INFO_SERVICE_RPC_HANDLER(UpdateActorInfo, 1); + } + + private: + /// The grpc async service object. + ActorInfoGcsService::AsyncService service_; + /// The service handler that actually handle the requests. + ActorInfoHandler &service_handler_; +}; + } // namespace rpc } // namespace ray -#endif +#endif // RAY_RPC_GCS_RPC_SERVER_H