wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-16 11:21:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add gcs server job info & actor info handler (#6469)

Browse Source
This commit is contained in:
fangfengbin
2019-12-20 14:28:04 +08:00
committed by Hao Chen
parent de22cdb233
commit 3c0164419b
11 changed files with 545 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
BUILD.bazel
+6
View File
@@ -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",
src/ray/gcs/gcs_server/actor_info_handler_impl.cc
+84
View File
@@ -0,0 +1,84 @@
#include "actor_info_handler_impl.h"
#include <assert.h>
#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<ActorTableData> &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<ActorTableData>();
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<ActorTableData>();
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
src/ray/gcs/gcs_server/actor_info_handler_impl.h
+34
View File
@@ -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
src/ray/gcs/gcs_server/gcs_server.cc
+14 -2
View File
@@ -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<RedisGcsClient>(options);
auto status = redis_gcs_client_->Connect(main_service_);
RAY_CHECK(status.ok()) << "Failed to init redis gcs client as " << status;
}
std::unique_ptr<rpc::JobInfoHandler> GcsServer::InitJobInfoHandler() {
return std::unique_ptr<rpc::DefaultJobInfoHandler>(new rpc::DefaultJobInfoHandler());
return std::unique_ptr<rpc::DefaultJobInfoHandler>(
new rpc::DefaultJobInfoHandler(*redis_gcs_client_));
}
std::unique_ptr<rpc::ActorInfoHandler> GcsServer::InitActorInfoHandler() {
return std::unique_ptr<rpc::DefaultActorInfoHandler>(
new rpc::DefaultActorInfoHandler(*redis_gcs_client_));
}
} // namespace gcs
src/ray/gcs/gcs_server/gcs_server.h
+7
View File
@@ -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<rpc::JobInfoHandler> InitJobInfoHandler();
/// The actor info handler
virtual std::unique_ptr<rpc::ActorInfoHandler> InitActorInfoHandler();
private:
/// Gcs server configuration
GcsServerConfig config_;
@@ -58,6 +62,9 @@ class GcsServer {
/// Job info handler and service
std::unique_ptr<rpc::JobInfoHandler> job_info_handler_;
std::unique_ptr<rpc::JobInfoGrpcService> job_info_service_;
/// Actor info handler and service
std::unique_ptr<rpc::ActorInfoHandler> actor_info_handler_;
std::unique_ptr<rpc::ActorInfoGrpcService> actor_info_service_;
/// Backend client
std::shared_ptr<RedisGcsClient> redis_gcs_client_;
};
src/ray/gcs/gcs_server/job_info_handler_impl.cc
+43 -4
View File
@@ -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<JobTableData>();
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
src/ray/gcs/gcs_server/job_info_handler_impl.h
+9 -3
View File
@@ -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
src/ray/gcs/gcs_server/test/gcs_server_rpc_test.cc
+199 -71
View File
@@ -1,93 +1,221 @@
#include <boost/asio/basic_socket.hpp>
#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 <typename T>
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<rpc::JobInfoHandler> InitJobInfoHandler() override {
return std::unique_ptr<rpc::JobInfoHandler>(new T);
}
};
};
class GcsServerTest : public ::testing::Test {
public:
using CallFunction = std::function<void(std::promise<bool> &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<MockedJobInfoHandler> 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<rpc::AddJobReply> 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<boost::asio::io_service::work> 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<bool> &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<bool> &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<bool> &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<bool> &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<bool> &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<bool> 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::GcsServer> gcs_server_;
std::unique_ptr<std::thread> thread_io_service_;
std::unique_ptr<std::thread> thread_gcs_server_;
boost::asio::io_service io_service_;
// Gcs client
std::unique_ptr<rpc::GcsRpcClient> client_;
std::unique_ptr<rpc::ClientCallManager> 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<rpc::MarkJobFinishedReply> 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();
}
src/ray/protobuf/gcs_service.proto
+38
View File
@@ -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);
}
src/ray/rpc/gcs_server/gcs_rpc_client.h
+44 -6
View File
@@ -24,7 +24,8 @@ class GcsRpcClient {
: client_call_manager_(client_call_manager) {
std::shared_ptr<grpc::Channel> 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<AddJobReply> &callback) {
client_call_manager_.CreateCall<JobInfoGcsService, AddJobRequest, AddJobReply>(
*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<MarkJobFinishedReply> &callback) {
client_call_manager_
.CreateCall<JobInfoGcsService, MarkJobFinishedRequest, MarkJobFinishedReply>(
*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<GetActorInfoReply> &callback) {
client_call_manager_
.CreateCall<ActorInfoGcsService, GetActorInfoRequest, GetActorInfoReply>(
*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<RegisterActorInfoReply> &callback) {
client_call_manager_.CreateCall<ActorInfoGcsService, RegisterActorInfoRequest,
RegisterActorInfoReply>(
*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<UpdateActorInfoReply> &callback) {
client_call_manager_
.CreateCall<ActorInfoGcsService, UpdateActorInfoRequest, UpdateActorInfoReply>(
*actor_info_stub_, &ActorInfoGcsService::Stub::PrepareAsyncUpdateActorInfo,
request, callback);
}
private:
/// The gRPC-generated stub.
std::unique_ptr<JobInfoGcsService::Stub> stub_;
std::unique_ptr<JobInfoGcsService::Stub> job_info_stub_;
std::unique_ptr<ActorInfoGcsService::Stub> 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
src/ray/rpc/gcs_server/gcs_rpc_server.h
+67 -16
View File
@@ -9,6 +9,24 @@
namespace ray {
namespace rpc {
#define JOB_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
std::unique_ptr<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<JobInfoGcsService, JobInfoHandler, HANDLER##Request, \
HANDLER##Reply>( \
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<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<ActorInfoGcsService, ActorInfoHandler, HANDLER##Request, \
HANDLER##Reply>( \
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<grpc::ServerCompletionQueue> &cq,
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
*server_call_factories_and_concurrencies) override {
std::unique_ptr<ServerCallFactory> add_job_call_factory(
new ServerCallFactoryImpl<JobInfoGcsService, JobInfoHandler, AddJobRequest,
AddJobReply>(
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<ServerCallFactory> mark_job_finished_call_factory(
new ServerCallFactoryImpl<JobInfoGcsService, JobInfoHandler,
MarkJobFinishedRequest, MarkJobFinishedReply>(
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<grpc::ServerCompletionQueue> &cq,
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, 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