[rpc] refactor RPC server code (#6661)

* refactor RPC client

* remove unused code

* format

* fix

* resolve comments

* format

* update

* refactor rpc server

* update

* format

* fix

* fix

* Update src/ray/rpc/worker/core_worker_server.h

Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com>

* resolve comments

* format

* update

* update

* add a comment

* fix

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
This commit is contained in:
Zhijun Fu
2020-01-07 22:03:42 +08:00
committed by GitHub
parent aaeb3c44a5
commit 72335dbe46
11 changed files with 168 additions and 231 deletions
+14 -18
View File
@@ -26,23 +26,19 @@
/// The set of gRPC handlers and their associated level of concurrency. If you want to
/// add a new call to the worker gRPC server, do the following:
/// 1) Add the rpc to the CoreWorkerService in core_worker.proto, e.g., "ExampleCall"
/// 2) Add a new handler to the macro below: "RAY_CORE_WORKER_RPC_HANDLER(ExampleCall, 1)"
/// 3) Add a method to the CoreWorker class below: "CoreWorker::HandleExampleCall"
#define RAY_CORE_WORKER_RPC_HANDLERS \
RAY_CORE_WORKER_RPC_HANDLER(AssignTask, 5) \
RAY_CORE_WORKER_RPC_HANDLER(PushTask, 9999) \
RAY_CORE_WORKER_RPC_HANDLER(DirectActorCallArgWaitComplete, 100) \
RAY_CORE_WORKER_RPC_HANDLER(GetObjectStatus, 9999) \
RAY_CORE_WORKER_RPC_HANDLER(WaitForObjectEviction, 9999) \
RAY_CORE_WORKER_RPC_HANDLER(KillActor, 9999) \
RAY_CORE_WORKER_RPC_HANDLER(GetCoreWorkerStats, 100)
/// 2) Add a new macro to RAY_CORE_WORKER_DECLARE_RPC_HANDLERS
/// in core_worker_server.h,
// e.g. "DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(ExampleCall)"
/// 3) Add a new macro to RAY_CORE_WORKER_RPC_HANDLERS in core_worker_server.h, e.g.
/// "RPC_SERVICE_HANDLER(CoreWorkerService, ExampleCall, 1)"
/// 4) Add a method to the CoreWorker class below: "CoreWorker::HandleExampleCall"
namespace ray {
/// The root class that contains all the core and language-independent functionalities
/// of the worker. This class is supposed to be used to implement app-language (Java,
/// Python, etc) workers.
class CoreWorker {
class CoreWorker : public rpc::CoreWorkerServiceHandler {
// Callback that must be implemented and provided by the language-specific worker
// frontend to execute tasks and return their results.
using TaskExecutionCallback = std::function<Status(
@@ -81,7 +77,7 @@ class CoreWorker {
std::function<Status()> check_signals = nullptr,
bool ref_counting_enabled = false);
~CoreWorker();
virtual ~CoreWorker();
void Disconnect();
@@ -400,22 +396,22 @@ class CoreWorker {
/// Implements gRPC server handler.
void HandleAssignTask(const rpc::AssignTaskRequest &request,
rpc::AssignTaskReply *reply,
rpc::SendReplyCallback send_reply_callback);
rpc::SendReplyCallback send_reply_callback) override;
/// Implements gRPC server handler.
void HandlePushTask(const rpc::PushTaskRequest &request, rpc::PushTaskReply *reply,
rpc::SendReplyCallback send_reply_callback);
rpc::SendReplyCallback send_reply_callback) override;
/// Implements gRPC server handler.
void HandleDirectActorCallArgWaitComplete(
const rpc::DirectActorCallArgWaitCompleteRequest &request,
rpc::DirectActorCallArgWaitCompleteReply *reply,
rpc::SendReplyCallback send_reply_callback);
rpc::SendReplyCallback send_reply_callback) override;
/// Implements gRPC server handler.
void HandleGetObjectStatus(const rpc::GetObjectStatusRequest &request,
rpc::GetObjectStatusReply *reply,
rpc::SendReplyCallback send_reply_callback);
rpc::SendReplyCallback send_reply_callback) override;
/// Implements gRPC server handler.
void HandleWaitForObjectEviction(const rpc::WaitForObjectEvictionRequest &request,
@@ -424,12 +420,12 @@ class CoreWorker {
/// Implements gRPC server handler.
void HandleKillActor(const rpc::KillActorRequest &request, rpc::KillActorReply *reply,
rpc::SendReplyCallback send_reply_callback);
rpc::SendReplyCallback send_reply_callback) override;
/// Get statistics from core worker.
void HandleGetCoreWorkerStats(const rpc::GetCoreWorkerStatsRequest &request,
rpc::GetCoreWorkerStatsReply *reply,
rpc::SendReplyCallback send_reply_callback);
rpc::SendReplyCallback send_reply_callback) override;
///
/// Public methods related to async actor call. This should only be used when
+7 -9
View File
@@ -665,9 +665,8 @@ void ObjectManager::WaitComplete(const UniqueID &wait_id) {
}
/// Implementation of ObjectManagerServiceHandler
void ObjectManager::HandlePushRequest(const rpc::PushRequest &request,
rpc::PushReply *reply,
rpc::SendReplyCallback send_reply_callback) {
void ObjectManager::HandlePush(const rpc::PushRequest &request, rpc::PushReply *reply,
rpc::SendReplyCallback send_reply_callback) {
ObjectID object_id = ObjectID::FromBinary(request.object_id());
ClientID client_id = ClientID::FromBinary(request.client_id());
@@ -712,9 +711,8 @@ ray::Status ObjectManager::ReceiveObjectChunk(const ClientID &client_id,
return status;
}
void ObjectManager::HandlePullRequest(const rpc::PullRequest &request,
rpc::PullReply *reply,
rpc::SendReplyCallback send_reply_callback) {
void ObjectManager::HandlePull(const rpc::PullRequest &request, rpc::PullReply *reply,
rpc::SendReplyCallback send_reply_callback) {
ObjectID object_id = ObjectID::FromBinary(request.object_id());
ClientID client_id = ClientID::FromBinary(request.client_id());
RAY_LOG(DEBUG) << "Received pull request from client " << client_id << " for object ["
@@ -735,9 +733,9 @@ void ObjectManager::HandlePullRequest(const rpc::PullRequest &request,
send_reply_callback(Status::OK(), nullptr, nullptr);
}
void ObjectManager::HandleFreeObjectsRequest(const rpc::FreeObjectsRequest &request,
rpc::FreeObjectsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
void ObjectManager::HandleFreeObjects(const rpc::FreeObjectsRequest &request,
rpc::FreeObjectsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
std::vector<ObjectID> object_ids;
for (const auto &e : request.object_ids()) {
object_ids.emplace_back(ObjectID::FromBinary(e));
+7 -7
View File
@@ -81,25 +81,25 @@ class ObjectManager : public ObjectManagerInterface,
/// \param request Push request including the object chunk data
/// \param reply Reply to the sender
/// \param send_reply_callback Callback of the request
void HandlePushRequest(const rpc::PushRequest &request, rpc::PushReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
void HandlePush(const rpc::PushRequest &request, rpc::PushReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Handle pull request from remote object manager
///
/// \param request Pull request
/// \param reply Reply
/// \param send_reply_callback Callback of request
void HandlePullRequest(const rpc::PullRequest &request, rpc::PullReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
void HandlePull(const rpc::PullRequest &request, rpc::PullReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Handle free objects request
///
/// \param request Free objects request
/// \param reply Reply
/// \param send_reply_callback
void HandleFreeObjectsRequest(const rpc::FreeObjectsRequest &request,
rpc::FreeObjectsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
void HandleFreeObjects(const rpc::FreeObjectsRequest &request,
rpc::FreeObjectsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Send object to remote object manager
///
+9 -9
View File
@@ -1496,7 +1496,7 @@ void NodeManager::WaitForTaskArgsRequests(std::pair<ScheduleFn, Task> &work) {
}
};
void NodeManager::HandleWorkerLeaseRequest(const rpc::RequestWorkerLeaseRequest &request,
void NodeManager::HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest &request,
rpc::RequestWorkerLeaseReply *reply,
rpc::SendReplyCallback send_reply_callback) {
rpc::Task task_message;
@@ -2968,9 +2968,9 @@ std::string compact_tag_string(const opencensus::stats::ViewDescriptor &view,
return result.str();
}
void NodeManager::HandlePinObjectIDsRequest(const rpc::PinObjectIDsRequest &request,
rpc::PinObjectIDsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
void NodeManager::HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request,
rpc::PinObjectIDsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
if (!object_pinning_enabled_) {
send_reply_callback(Status::OK(), nullptr, nullptr);
return;
@@ -3044,9 +3044,9 @@ void NodeManager::HandlePinObjectIDsRequest(const rpc::PinObjectIDsRequest &requ
send_reply_callback(Status::OK(), nullptr, nullptr);
}
void NodeManager::HandleNodeStatsRequest(const rpc::GetNodeStatsRequest &request,
rpc::GetNodeStatsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
void NodeManager::HandleGetNodeStats(const rpc::GetNodeStatsRequest &request,
rpc::GetNodeStatsReply *reply,
rpc::SendReplyCallback send_reply_callback) {
for (const auto &driver : worker_pool_.GetAllDrivers()) {
auto worker_stats = reply->add_workers_stats();
worker_stats->set_pid(driver->Pid());
@@ -3094,11 +3094,11 @@ void NodeManager::HandleNodeStatsRequest(const rpc::GetNodeStatsRequest &request
}
}
}
// As a result of the HandleNodeStatsRequest, we are collecting information from all
// As a result of the HandleGetNodeStats, we are collecting information from all
// workers on this node. This is done by calling GetCoreWorkerStats on each worker. In
// order to send up-to-date information back, we wait until all workers have replied,
// and return the information from HandleNodesStatsRequest. The caller of
// HandleNodeStatsRequest should set a timeout so that the rpc finishes even if not all
// HandleGetNodeStats should set a timeout so that the rpc finishes even if not all
// workers have replied.
auto all_workers = worker_pool_.GetAllWorkers();
for (const auto &worker : all_workers) {
+7 -7
View File
@@ -534,7 +534,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
bool success);
/// Handle a `WorkerLease` request.
void HandleWorkerLeaseRequest(const rpc::RequestWorkerLeaseRequest &request,
void HandleRequestWorkerLease(const rpc::RequestWorkerLeaseRequest &request,
rpc::RequestWorkerLeaseReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
@@ -549,14 +549,14 @@ class NodeManager : public rpc::NodeManagerServiceHandler {
rpc::SendReplyCallback send_reply_callback) override;
/// Handle a `PinObjectIDs` request.
void HandlePinObjectIDsRequest(const rpc::PinObjectIDsRequest &request,
rpc::PinObjectIDsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
void HandlePinObjectIDs(const rpc::PinObjectIDsRequest &request,
rpc::PinObjectIDsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Handle a `NodeStats` request.
void HandleNodeStatsRequest(const rpc::GetNodeStatsRequest &request,
rpc::GetNodeStatsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
void HandleGetNodeStats(const rpc::GetNodeStatsRequest &request,
rpc::GetNodeStatsReply *reply,
rpc::SendReplyCallback send_reply_callback) override;
/// Push an error to the driver if this node is full of actors and so we are
/// unable to schedule new tasks or actors at all.
+37 -61
View File
@@ -9,54 +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 JOB_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
RPC_SERVICE_HANDLER(JobInfoGcsService, HANDLER, 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);
#define ACTOR_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
RPC_SERVICE_HANDLER(ActorInfoGcsService, HANDLER, CONCURRENCY)
#define NODE_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
std::unique_ptr<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<NodeInfoGcsService, NodeInfoHandler, HANDLER##Request, \
HANDLER##Reply>( \
service_, &NodeInfoGcsService::AsyncService::Request##HANDLER, \
service_handler_, &NodeInfoHandler::Handle##HANDLER, cq, main_service_)); \
server_call_factories_and_concurrencies->emplace_back( \
std::move(HANDLER##_call_factory), CONCURRENCY);
#define NODE_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
RPC_SERVICE_HANDLER(NodeInfoGcsService, HANDLER, CONCURRENCY)
#define OBJECT_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
std::unique_ptr<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<ObjectInfoGcsService, ObjectInfoHandler, \
HANDLER##Request, HANDLER##Reply>( \
service_, &ObjectInfoGcsService::AsyncService::Request##HANDLER, \
service_handler_, &ObjectInfoHandler::Handle##HANDLER, cq, main_service_)); \
server_call_factories_and_concurrencies->emplace_back( \
std::move(HANDLER##_call_factory), CONCURRENCY);
#define OBJECT_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
RPC_SERVICE_HANDLER(ObjectInfoGcsService, HANDLER, CONCURRENCY)
#define TASK_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
std::unique_ptr<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<TaskInfoGcsService, TaskInfoHandler, HANDLER##Request, \
HANDLER##Reply>( \
service_, &TaskInfoGcsService::AsyncService::Request##HANDLER, \
service_handler_, &TaskInfoHandler::Handle##HANDLER, cq, main_service_)); \
server_call_factories_and_concurrencies->emplace_back( \
std::move(HANDLER##_call_factory), CONCURRENCY);
#define TASK_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \
RPC_SERVICE_HANDLER(TaskInfoGcsService, HANDLER, CONCURRENCY)
class JobInfoHandler {
class JobInfoGcsServiceHandler {
public:
virtual ~JobInfoHandler() = default;
virtual ~JobInfoGcsServiceHandler() = default;
virtual void HandleAddJob(const AddJobRequest &request, AddJobReply *reply,
SendReplyCallback send_reply_callback) = 0;
@@ -73,7 +43,7 @@ class JobInfoGrpcService : public GrpcService {
///
/// \param[in] handler The service handler that actually handle the requests.
explicit JobInfoGrpcService(boost::asio::io_service &io_service,
JobInfoHandler &handler)
JobInfoGcsServiceHandler &handler)
: GrpcService(io_service), service_handler_(handler){};
protected:
@@ -91,12 +61,12 @@ class JobInfoGrpcService : public GrpcService {
/// The grpc async service object.
JobInfoGcsService::AsyncService service_;
/// The service handler that actually handle the requests.
JobInfoHandler &service_handler_;
JobInfoGcsServiceHandler &service_handler_;
};
class ActorInfoHandler {
class ActorInfoGcsServiceHandler {
public:
virtual ~ActorInfoHandler() = default;
virtual ~ActorInfoGcsServiceHandler() = default;
virtual void HandleGetActorInfo(const GetActorInfoRequest &request,
GetActorInfoReply *reply,
@@ -130,7 +100,7 @@ class ActorInfoGrpcService : public GrpcService {
///
/// \param[in] handler The service handler that actually handle the requests.
explicit ActorInfoGrpcService(boost::asio::io_service &io_service,
ActorInfoHandler &handler)
ActorInfoGcsServiceHandler &handler)
: GrpcService(io_service), service_handler_(handler){};
protected:
@@ -152,12 +122,12 @@ class ActorInfoGrpcService : public GrpcService {
/// The grpc async service object.
ActorInfoGcsService::AsyncService service_;
/// The service handler that actually handle the requests.
ActorInfoHandler &service_handler_;
ActorInfoGcsServiceHandler &service_handler_;
};
class NodeInfoHandler {
class NodeInfoGcsServiceHandler {
public:
virtual ~NodeInfoHandler() = default;
virtual ~NodeInfoGcsServiceHandler() = default;
virtual void HandleRegisterNode(const RegisterNodeRequest &request,
RegisterNodeReply *reply,
@@ -199,7 +169,7 @@ class NodeInfoGrpcService : public GrpcService {
///
/// \param[in] handler The service handler that actually handle the requests.
explicit NodeInfoGrpcService(boost::asio::io_service &io_service,
NodeInfoHandler &handler)
NodeInfoGcsServiceHandler &handler)
: GrpcService(io_service), service_handler_(handler){};
protected:
@@ -223,12 +193,12 @@ class NodeInfoGrpcService : public GrpcService {
/// The grpc async service object.
NodeInfoGcsService::AsyncService service_;
/// The service handler that actually handle the requests.
NodeInfoHandler &service_handler_;
NodeInfoGcsServiceHandler &service_handler_;
};
class ObjectInfoHandler {
class ObjectInfoGcsServiceHandler {
public:
virtual ~ObjectInfoHandler() = default;
virtual ~ObjectInfoGcsServiceHandler() = default;
virtual void HandleGetObjectLocations(const GetObjectLocationsRequest &request,
GetObjectLocationsReply *reply,
@@ -243,14 +213,14 @@ class ObjectInfoHandler {
SendReplyCallback send_reply_callback) = 0;
};
/// The `GrpcService` for `ObjectInfoHandler`.
/// The `GrpcService` for `ObjectInfoGcsServiceHandler`.
class ObjectInfoGrpcService : public GrpcService {
public:
/// Constructor.
///
/// \param[in] handler The service handler that actually handle the requests.
explicit ObjectInfoGrpcService(boost::asio::io_service &io_service,
ObjectInfoHandler &handler)
ObjectInfoGcsServiceHandler &handler)
: GrpcService(io_service), service_handler_(handler){};
protected:
@@ -269,12 +239,12 @@ class ObjectInfoGrpcService : public GrpcService {
/// The grpc async service object.
ObjectInfoGcsService::AsyncService service_;
/// The service handler that actually handle the requests.
ObjectInfoHandler &service_handler_;
ObjectInfoGcsServiceHandler &service_handler_;
};
class TaskInfoHandler {
class TaskInfoGcsServiceHandler {
public:
virtual ~TaskInfoHandler() = default;
virtual ~TaskInfoGcsServiceHandler() = default;
virtual void HandleAddTask(const AddTaskRequest &request, AddTaskReply *reply,
SendReplyCallback send_reply_callback) = 0;
@@ -294,7 +264,7 @@ class TaskInfoGrpcService : public GrpcService {
///
/// \param[in] handler The service handler that actually handle the requests.
explicit TaskInfoGrpcService(boost::asio::io_service &io_service,
TaskInfoHandler &handler)
TaskInfoGcsServiceHandler &handler)
: GrpcService(io_service), service_handler_(handler){};
protected:
@@ -313,9 +283,15 @@ class TaskInfoGrpcService : public GrpcService {
/// The grpc async service object.
TaskInfoGcsService::AsyncService service_;
/// The service handler that actually handle the requests.
TaskInfoHandler &service_handler_;
TaskInfoGcsServiceHandler &service_handler_;
};
using JobInfoHandler = JobInfoGcsServiceHandler;
using ActorInfoHandler = ActorInfoGcsServiceHandler;
using NodeInfoHandler = NodeInfoGcsServiceHandler;
using ObjectInfoHandler = ObjectInfoGcsServiceHandler;
using TaskInfoHandler = TaskInfoGcsServiceHandler;
} // namespace rpc
} // namespace ray
+15
View File
@@ -13,6 +13,21 @@
namespace ray {
namespace rpc {
#define RPC_SERVICE_HANDLER(SERVICE, HANDLER, CONCURRENCY) \
std::unique_ptr<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<SERVICE, SERVICE##Handler, HANDLER##Request, \
HANDLER##Reply>( \
service_, &SERVICE::AsyncService::Request##HANDLER, service_handler_, \
&SERVICE##Handler::Handle##HANDLER, cq, main_service_)); \
server_call_factories_and_concurrencies->emplace_back( \
std::move(HANDLER##_call_factory), CONCURRENCY);
// Define a void RPC client method.
#define DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(METHOD) \
virtual void Handle##METHOD(const rpc::METHOD##Request &request, \
rpc::METHOD##Reply *reply, \
rpc::SendReplyCallback send_reply_callback) = 0;
class GrpcService;
/// Class that represents an gRPC server.
+16 -54
View File
@@ -9,6 +9,14 @@
namespace ray {
namespace rpc {
/// NOTE: See src/ray/core_worker/core_worker.h on how to add a new grpc handler.
#define RAY_NODE_MANAGER_RPC_HANDLERS \
RPC_SERVICE_HANDLER(NodeManagerService, RequestWorkerLease, 100) \
RPC_SERVICE_HANDLER(NodeManagerService, ReturnWorker, 100) \
RPC_SERVICE_HANDLER(NodeManagerService, ForwardTask, 100) \
RPC_SERVICE_HANDLER(NodeManagerService, PinObjectIDs, 100) \
RPC_SERVICE_HANDLER(NodeManagerService, GetNodeStats, 1)
/// Interface of the `NodeManagerService`, see `src/ray/protobuf/node_manager.proto`.
class NodeManagerServiceHandler {
public:
@@ -23,7 +31,7 @@ class NodeManagerServiceHandler {
/// \param[out] reply The reply message.
/// \param[in] send_reply_callback The callback to be called when the request is done.
virtual void HandleWorkerLeaseRequest(const RequestWorkerLeaseRequest &request,
virtual void HandleRequestWorkerLease(const RequestWorkerLeaseRequest &request,
RequestWorkerLeaseReply *reply,
SendReplyCallback send_reply_callback) = 0;
@@ -35,13 +43,13 @@ class NodeManagerServiceHandler {
ForwardTaskReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandlePinObjectIDsRequest(const PinObjectIDsRequest &request,
PinObjectIDsReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandlePinObjectIDs(const PinObjectIDsRequest &request,
PinObjectIDsReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandleNodeStatsRequest(const GetNodeStatsRequest &request,
GetNodeStatsReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandleGetNodeStats(const GetNodeStatsRequest &request,
GetNodeStatsReply *reply,
SendReplyCallback send_reply_callback) = 0;
};
/// The `GrpcService` for `NodeManagerService`.
@@ -62,53 +70,7 @@ class NodeManagerGrpcService : public GrpcService {
const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
*server_call_factories_and_concurrencies) override {
// Initialize the factory for requests.
std::unique_ptr<ServerCallFactory> request_worker_lease_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
RequestWorkerLeaseRequest, RequestWorkerLeaseReply>(
service_, &NodeManagerService::AsyncService::RequestRequestWorkerLease,
service_handler_, &NodeManagerServiceHandler::HandleWorkerLeaseRequest, cq,
main_service_));
std::unique_ptr<ServerCallFactory> release_worker_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
ReturnWorkerRequest, ReturnWorkerReply>(
service_, &NodeManagerService::AsyncService::RequestReturnWorker,
service_handler_, &NodeManagerServiceHandler::HandleReturnWorker, cq,
main_service_));
std::unique_ptr<ServerCallFactory> forward_task_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
ForwardTaskRequest, ForwardTaskReply>(
service_, &NodeManagerService::AsyncService::RequestForwardTask,
service_handler_, &NodeManagerServiceHandler::HandleForwardTask, cq,
main_service_));
std::unique_ptr<ServerCallFactory> pin_object_ids_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
PinObjectIDsRequest, PinObjectIDsReply>(
service_, &NodeManagerService::AsyncService::RequestPinObjectIDs,
service_handler_, &NodeManagerServiceHandler::HandlePinObjectIDsRequest, cq,
main_service_));
std::unique_ptr<ServerCallFactory> node_stats_call_factory(
new ServerCallFactoryImpl<NodeManagerService, NodeManagerServiceHandler,
GetNodeStatsRequest, GetNodeStatsReply>(
service_, &NodeManagerService::AsyncService::RequestGetNodeStats,
service_handler_, &NodeManagerServiceHandler::HandleNodeStatsRequest, cq,
main_service_));
// Set accept concurrency.
server_call_factories_and_concurrencies->emplace_back(
std::move(request_worker_lease_call_factory), 100);
server_call_factories_and_concurrencies->emplace_back(
std::move(release_worker_call_factory), 100);
server_call_factories_and_concurrencies->emplace_back(
std::move(forward_task_call_factory), 100);
server_call_factories_and_concurrencies->emplace_back(
std::move(pin_object_ids_call_factory), 100);
server_call_factories_and_concurrencies->emplace_back(
std::move(node_stats_call_factory), 1);
RAY_NODE_MANAGER_RPC_HANDLERS
}
private:
@@ -10,6 +10,11 @@
namespace ray {
namespace rpc {
#define RAY_OBJECT_MANAGER_RPC_HANDLERS \
RPC_SERVICE_HANDLER(ObjectManagerService, Push, 5) \
RPC_SERVICE_HANDLER(ObjectManagerService, Pull, 5) \
RPC_SERVICE_HANDLER(ObjectManagerService, FreeObjects, 2)
/// Implementations of the `ObjectManagerGrpcService`, check interface in
/// `src/ray/protobuf/object_manager.proto`.
class ObjectManagerServiceHandler {
@@ -21,15 +26,15 @@ class ObjectManagerServiceHandler {
/// \param[in] request The request message.
/// \param[out] reply The reply message.
/// \param[in] send_reply_callback The callback to be called when the request is done.
virtual void HandlePushRequest(const PushRequest &request, PushReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandlePush(const PushRequest &request, PushReply *reply,
SendReplyCallback send_reply_callback) = 0;
/// Handle a `Pull` request
virtual void HandlePullRequest(const PullRequest &request, PullReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandlePull(const PullRequest &request, PullReply *reply,
SendReplyCallback send_reply_callback) = 0;
/// Handle a `FreeObjects` request
virtual void HandleFreeObjectsRequest(const FreeObjectsRequest &request,
FreeObjectsReply *reply,
SendReplyCallback send_reply_callback) = 0;
virtual void HandleFreeObjects(const FreeObjectsRequest &request,
FreeObjectsReply *reply,
SendReplyCallback send_reply_callback) = 0;
};
/// The `GrpcService` for `ObjectManagerGrpcService`.
@@ -50,33 +55,7 @@ class ObjectManagerGrpcService : public GrpcService {
const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
*server_call_factories_and_concurrencies) override {
// Initialize the factory for `Push` requests.
std::unique_ptr<ServerCallFactory> push_call_factory(
new ServerCallFactoryImpl<ObjectManagerService, ObjectManagerServiceHandler,
PushRequest, PushReply>(
service_, &ObjectManagerService::AsyncService::RequestPush, service_handler_,
&ObjectManagerServiceHandler::HandlePushRequest, cq, main_service_));
server_call_factories_and_concurrencies->emplace_back(std::move(push_call_factory),
5);
// Initialize the factory for `Pull` requests.
std::unique_ptr<ServerCallFactory> pull_call_factory(
new ServerCallFactoryImpl<ObjectManagerService, ObjectManagerServiceHandler,
PullRequest, PullReply>(
service_, &ObjectManagerService::AsyncService::RequestPull, service_handler_,
&ObjectManagerServiceHandler::HandlePullRequest, cq, main_service_));
server_call_factories_and_concurrencies->emplace_back(std::move(pull_call_factory),
5);
// Initialize the factory for `FreeObjects` requests.
std::unique_ptr<ServerCallFactory> free_objects_call_factory(
new ServerCallFactoryImpl<ObjectManagerService, ObjectManagerServiceHandler,
FreeObjectsRequest, FreeObjectsReply>(
service_, &ObjectManagerService::AsyncService::RequestFreeObjects,
service_handler_, &ObjectManagerServiceHandler::HandleFreeObjectsRequest, cq,
main_service_));
server_call_factories_and_concurrencies->emplace_back(
std::move(free_objects_call_factory), 2);
RAY_OBJECT_MANAGER_RPC_HANDLERS
}
private:
-28
View File
@@ -1,28 +0,0 @@
#include "ray/rpc/worker/core_worker_server.h"
#include "ray/core_worker/core_worker.h"
namespace ray {
namespace rpc {
#define RAY_CORE_WORKER_RPC_HANDLER(HANDLER, CONCURRENCY) \
std::unique_ptr<ServerCallFactory> HANDLER##_call_factory( \
new ServerCallFactoryImpl<CoreWorkerService, CoreWorker, HANDLER##Request, \
HANDLER##Reply>( \
service_, &CoreWorkerService::AsyncService::Request##HANDLER, core_worker_, \
&CoreWorker::Handle##HANDLER, cq, main_service_)); \
server_call_factories_and_concurrencies->emplace_back( \
std::move(HANDLER##_call_factory), CONCURRENCY);
CoreWorkerGrpcService::CoreWorkerGrpcService(boost::asio::io_service &main_service,
CoreWorker &core_worker)
: GrpcService(main_service), core_worker_(core_worker){};
void CoreWorkerGrpcService::InitServerCallFactories(
const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
*server_call_factories_and_concurrencies) {
RAY_CORE_WORKER_RPC_HANDLERS
}
} // namespace rpc
} // namespace ray
+43 -4
View File
@@ -13,6 +13,41 @@ class CoreWorker;
namespace rpc {
/// NOTE: See src/ray/core_worker/core_worker.h on how to add a new grpc handler.
#define RAY_CORE_WORKER_RPC_HANDLERS \
RPC_SERVICE_HANDLER(CoreWorkerService, AssignTask, 5) \
RPC_SERVICE_HANDLER(CoreWorkerService, PushTask, 9999) \
RPC_SERVICE_HANDLER(CoreWorkerService, DirectActorCallArgWaitComplete, 100) \
RPC_SERVICE_HANDLER(CoreWorkerService, GetObjectStatus, 9999) \
RPC_SERVICE_HANDLER(CoreWorkerService, WaitForObjectEviction, 9999) \
RPC_SERVICE_HANDLER(CoreWorkerService, KillActor, 9999) \
RPC_SERVICE_HANDLER(CoreWorkerService, GetCoreWorkerStats, 100)
#define RAY_CORE_WORKER_DECLARE_RPC_HANDLERS \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(AssignTask) \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(PushTask) \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(DirectActorCallArgWaitComplete) \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(GetObjectStatus) \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(WaitForObjectEviction) \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(KillActor) \
DECLARE_VOID_RPC_SERVICE_HANDLER_METHOD(GetCoreWorkerStats)
/// Interface of the `CoreWorkerServiceHandler`, see `src/ray/protobuf/core_worker.proto`.
class CoreWorkerServiceHandler {
public:
/// Handlers. For all of the following handlers, the implementations can
/// handle the request asynchronously. When handling is done, the
/// `send_reply_callback` should be called. See
/// src/ray/rpc/node_manager/node_manager_client.h and
/// src/ray/protobuf/node_manager.proto for a description of the
/// functionality of each handler.
///
/// \param[in] request The request message.
/// \param[out] reply The reply message.
/// \param[in] send_reply_callback The callback to be called when the request is done.
RAY_CORE_WORKER_DECLARE_RPC_HANDLERS
};
/// The `GrpcServer` for `CoreWorkerService`.
class CoreWorkerGrpcService : public GrpcService {
public:
@@ -20,7 +55,9 @@ class CoreWorkerGrpcService : public GrpcService {
///
/// \param[in] main_service See super class.
/// \param[in] handler The service handler that actually handle the requests.
CoreWorkerGrpcService(boost::asio::io_service &main_service, CoreWorker &core_worker);
CoreWorkerGrpcService(boost::asio::io_service &main_service,
CoreWorkerServiceHandler &service_handler)
: GrpcService(main_service), service_handler_(service_handler) {}
protected:
grpc::Service &GetGrpcService() override { return service_; }
@@ -28,14 +65,16 @@ class CoreWorkerGrpcService : public GrpcService {
void InitServerCallFactories(
const std::unique_ptr<grpc::ServerCompletionQueue> &cq,
std::vector<std::pair<std::unique_ptr<ServerCallFactory>, int>>
*server_call_factories_and_concurrencies) override;
*server_call_factories_and_concurrencies) override {
RAY_CORE_WORKER_RPC_HANDLERS
}
private:
/// The grpc async service object.
CoreWorkerService::AsyncService service_;
/// The core worker that actually handles the requests.
CoreWorker &core_worker_;
/// The service handler that actually handles the requests.
CoreWorkerServiceHandler &service_handler_;
};
} // namespace rpc