diff --git a/src/ray/gcs/gcs_server/gcs_server.cc b/src/ray/gcs/gcs_server/gcs_server.cc index 52aff19b8..d5e567775 100644 --- a/src/ray/gcs/gcs_server/gcs_server.cc +++ b/src/ray/gcs/gcs_server/gcs_server.cc @@ -3,6 +3,7 @@ #include "job_info_handler_impl.h" #include "node_info_handler_impl.h" #include "object_info_handler_impl.h" +#include "stats_handler_impl.h" #include "task_info_handler_impl.h" namespace ray { @@ -44,6 +45,10 @@ void GcsServer::Start() { new rpc::TaskInfoGrpcService(main_service_, *task_info_handler_)); rpc_server_.RegisterService(*task_info_service_); + stats_handler_ = InitStatsHandler(); + stats_service_.reset(new rpc::StatsGrpcService(main_service_, *stats_handler_)); + rpc_server_.RegisterService(*stats_service_); + // Run rpc server. rpc_server_.Run(); @@ -95,5 +100,10 @@ std::unique_ptr GcsServer::InitTaskInfoHandler() { new rpc::DefaultTaskInfoHandler(*redis_gcs_client_)); } +std::unique_ptr GcsServer::InitStatsHandler() { + return std::unique_ptr( + new rpc::DefaultStatsHandler(*redis_gcs_client_)); +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_server.h b/src/ray/gcs/gcs_server/gcs_server.h index 4b33feea0..e65aab5c2 100644 --- a/src/ray/gcs/gcs_server/gcs_server.h +++ b/src/ray/gcs/gcs_server/gcs_server.h @@ -59,6 +59,9 @@ class GcsServer { /// The task info handler virtual std::unique_ptr InitTaskInfoHandler(); + /// The stats handler + virtual std::unique_ptr InitStatsHandler(); + private: /// Gcs server configuration GcsServerConfig config_; @@ -81,6 +84,9 @@ class GcsServer { /// Task info handler and service std::unique_ptr task_info_handler_; std::unique_ptr task_info_service_; + /// Stats handler and service + std::unique_ptr stats_handler_; + std::unique_ptr stats_service_; /// Backend client std::shared_ptr redis_gcs_client_; }; diff --git a/src/ray/gcs/gcs_server/stats_handler_impl.cc b/src/ray/gcs/gcs_server/stats_handler_impl.cc new file mode 100644 index 000000000..79bf54611 --- /dev/null +++ b/src/ray/gcs/gcs_server/stats_handler_impl.cc @@ -0,0 +1,32 @@ +#include "stats_handler_impl.h" + +namespace ray { +namespace rpc { + +void DefaultStatsHandler::HandleAddProfileData(const AddProfileDataRequest &request, + AddProfileDataReply *reply, + SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.profile_data().component_id()); + RAY_LOG(DEBUG) << "Adding profile data, component type = " + << request.profile_data().component_type() << ", node id = " << node_id; + auto profile_table_data = std::make_shared(); + profile_table_data->CopyFrom(request.profile_data()); + auto on_done = [node_id, request, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to add profile data, component type = " + << request.profile_data().component_type() + << ", node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Stats().AsyncAddProfileData(profile_table_data, on_done); + if (!status.ok()) { + on_done(status); + } + RAY_LOG(DEBUG) << "Finished adding profile data, component type = " + << request.profile_data().component_type() << ", node id = " << node_id; +} + +} // namespace rpc +} // namespace ray diff --git a/src/ray/gcs/gcs_server/stats_handler_impl.h b/src/ray/gcs/gcs_server/stats_handler_impl.h new file mode 100644 index 000000000..fa7b41009 --- /dev/null +++ b/src/ray/gcs/gcs_server/stats_handler_impl.h @@ -0,0 +1,27 @@ +#ifndef RAY_GCS_STATS_HANDLER_IMPL_H +#define RAY_GCS_STATS_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 `StatsHandler`. +class DefaultStatsHandler : public rpc::StatsHandler { + public: + explicit DefaultStatsHandler(gcs::RedisGcsClient &gcs_client) + : gcs_client_(gcs_client) {} + + void HandleAddProfileData(const AddProfileDataRequest &request, + AddProfileDataReply *reply, + SendReplyCallback send_reply_callback) override; + + private: + gcs::RedisGcsClient &gcs_client_; +}; + +} // namespace rpc +} // namespace ray + +#endif // RAY_GCS_STATS_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 51f35b2a1..de689d527 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 @@ -345,6 +345,16 @@ class GcsServerTest : public RedisServiceManagerForTest { return WaitReady(promise.get_future(), timeout_ms_); } + bool AddProfileData(const rpc::AddProfileDataRequest &request) { + std::promise promise; + client_->AddProfileData( + request, [&promise](const Status &status, const rpc::AddProfileDataReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + bool WaitReady(const std::future &future, uint64_t timeout_ms) { auto status = future.wait_for(std::chrono::milliseconds(timeout_ms)); return status == std::future_status::ready; @@ -595,6 +605,14 @@ TEST_F(GcsServerTest, TestTaskInfo) { ASSERT_TRUE(AttemptTaskReconstruction(attempt_task_reconstruction_request)); } +TEST_F(GcsServerTest, TestStats) { + rpc::ProfileTableData profile_table_data; + profile_table_data.set_component_id(ClientID::FromRandom().Binary()); + rpc::AddProfileDataRequest add_profile_data_request; + add_profile_data_request.mutable_profile_data()->CopyFrom(profile_table_data); + ASSERT_TRUE(AddProfileData(add_profile_data_request)); +} + } // namespace ray int main(int argc, char **argv) { diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index a435da8ff..fac438e58 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -22,9 +22,9 @@ message MarkJobFinishedReply { // Service for job info access. service JobInfoGcsService { - // Add job to gcs server. + // Add job to GCS Service. rpc AddJob(AddJobRequest) returns (AddJobReply); - // Mark job as finished to gcs server. + // Mark job as finished to GCS Service. rpc MarkJobFinished(MarkJobFinishedRequest) returns (MarkJobFinishedReply); } @@ -269,3 +269,16 @@ service TaskInfoGcsService { rpc AttemptTaskReconstruction(AttemptTaskReconstructionRequest) returns (AttemptTaskReconstructionReply); } + +message AddProfileDataRequest { + ProfileTableData profile_data = 1; +} + +message AddProfileDataReply { +} + +// Service for stats access. +service StatsGcsService { + // Add profile data to GCS Service. + rpc AddProfileData(AddProfileDataRequest) returns (AddProfileDataReply); +} diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index a9d996f41..6d6042dde 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_client.h @@ -26,6 +26,8 @@ class GcsRpcClient { new GrpcClient(address, port, client_call_manager)); task_info_grpc_client_ = std::unique_ptr>( new GrpcClient(address, port, client_call_manager)); + stats_grpc_client_ = std::unique_ptr>( + new GrpcClient(address, port, client_call_manager)); }; /// Add job info to gcs server. @@ -109,6 +111,9 @@ class GcsRpcClient { VOID_RPC_CLIENT_METHOD(TaskInfoGcsService, AttemptTaskReconstruction, task_info_grpc_client_, ) + /// Add profile data to GCS Service. + VOID_RPC_CLIENT_METHOD(StatsGcsService, AddProfileData, stats_grpc_client_, ) + private: /// The gRPC-generated stub. std::unique_ptr> job_info_grpc_client_; @@ -116,6 +121,7 @@ class GcsRpcClient { std::unique_ptr> node_info_grpc_client_; std::unique_ptr> object_info_grpc_client_; std::unique_ptr> task_info_grpc_client_; + std::unique_ptr> stats_grpc_client_; }; } // namespace rpc diff --git a/src/ray/rpc/gcs_server/gcs_rpc_server.h b/src/ray/rpc/gcs_server/gcs_rpc_server.h index aa9a1c98e..364aa7ad5 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_server.h @@ -24,6 +24,9 @@ namespace rpc { #define TASK_INFO_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \ RPC_SERVICE_HANDLER(TaskInfoGcsService, HANDLER, CONCURRENCY) +#define STATS_SERVICE_RPC_HANDLER(HANDLER, CONCURRENCY) \ + RPC_SERVICE_HANDLER(StatsGcsService, HANDLER, CONCURRENCY) + class JobInfoGcsServiceHandler { public: virtual ~JobInfoGcsServiceHandler() = default; @@ -296,11 +299,48 @@ class TaskInfoGrpcService : public GrpcService { TaskInfoGcsServiceHandler &service_handler_; }; +class StatsGcsServiceHandler { + public: + virtual ~StatsGcsServiceHandler() = default; + + virtual void HandleAddProfileData(const AddProfileDataRequest &request, + AddProfileDataReply *reply, + SendReplyCallback send_reply_callback) = 0; +}; + +/// The `GrpcService` for `StatsGcsService`. +class StatsGrpcService : public GrpcService { + public: + /// Constructor. + /// + /// \param[in] handler The service handler that actually handle the requests. + explicit StatsGrpcService(boost::asio::io_service &io_service, + StatsGcsServiceHandler &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 { + STATS_SERVICE_RPC_HANDLER(AddProfileData, 1); + } + + private: + /// The grpc async service object. + StatsGcsService::AsyncService service_; + /// The service handler that actually handle the requests. + StatsGcsServiceHandler &service_handler_; +}; + using JobInfoHandler = JobInfoGcsServiceHandler; using ActorInfoHandler = ActorInfoGcsServiceHandler; using NodeInfoHandler = NodeInfoGcsServiceHandler; using ObjectInfoHandler = ObjectInfoGcsServiceHandler; using TaskInfoHandler = TaskInfoGcsServiceHandler; +using StatsHandler = StatsGcsServiceHandler; } // namespace rpc } // namespace ray