diff --git a/src/ray/gcs/gcs_server/node_info_handler_impl.cc b/src/ray/gcs/gcs_server/node_info_handler_impl.cc index 7f8c1edc0..82e1cb586 100644 --- a/src/ray/gcs/gcs_server/node_info_handler_impl.cc +++ b/src/ray/gcs/gcs_server/node_info_handler_impl.cc @@ -110,5 +110,85 @@ void DefaultNodeInfoHandler::HandleReportBatchHeartbeat( << request.heartbeat_batch().batch_size(); } +void DefaultNodeInfoHandler::HandleGetResources(const GetResourcesRequest &request, + GetResourcesReply *reply, + SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.node_id()); + RAY_LOG(DEBUG) << "Getting node resources, node id = " << node_id; + + auto on_done = [node_id, reply, send_reply_callback]( + Status status, + const boost::optional &result) { + if (status.ok()) { + if (result) { + for (auto &resource : *result) { + (*reply->mutable_resources())[resource.first] = *resource.second; + } + } + } else { + RAY_LOG(ERROR) << "Failed to get node resources: " << status.ToString() + << ", node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Nodes().AsyncGetResources(node_id, on_done); + if (!status.ok()) { + on_done(status, boost::none); + } + + RAY_LOG(DEBUG) << "Finished getting node resources, node id = " << node_id; +} + +void DefaultNodeInfoHandler::HandleUpdateResources( + const UpdateResourcesRequest &request, UpdateResourcesReply *reply, + SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.node_id()); + gcs::NodeInfoAccessor::ResourceMap resources; + for (auto resource : request.resources()) { + resources[resource.first] = std::make_shared(resource.second); + } + + RAY_LOG(DEBUG) << "Updating node resources, node id = " << node_id; + auto on_done = [node_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to update node resources: " << status.ToString() + << ", node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = gcs_client_.Nodes().AsyncUpdateResources(node_id, resources, on_done); + if (!status.ok()) { + on_done(status); + } + + RAY_LOG(DEBUG) << "Finished updating node resources, node id = " << node_id; +} + +void DefaultNodeInfoHandler::HandleDeleteResources( + const DeleteResourcesRequest &request, DeleteResourcesReply *reply, + SendReplyCallback send_reply_callback) { + ClientID node_id = ClientID::FromBinary(request.node_id()); + auto resource_names = VectorFromProtobuf(request.resource_name_list()); + RAY_LOG(DEBUG) << "Deleting node resources, node id = " << node_id; + + auto on_done = [node_id, send_reply_callback](Status status) { + if (!status.ok()) { + RAY_LOG(ERROR) << "Failed to delete node resources: " << status.ToString() + << ", node id = " << node_id; + } + send_reply_callback(status, nullptr, nullptr); + }; + + Status status = + gcs_client_.Nodes().AsyncDeleteResources(node_id, resource_names, on_done); + if (!status.ok()) { + on_done(status); + } + + RAY_LOG(DEBUG) << "Finished deleting node resources, node id = " << node_id; +} + } // namespace rpc } // namespace ray diff --git a/src/ray/gcs/gcs_server/node_info_handler_impl.h b/src/ray/gcs/gcs_server/node_info_handler_impl.h index e779044c4..3310ec14f 100644 --- a/src/ray/gcs/gcs_server/node_info_handler_impl.h +++ b/src/ray/gcs/gcs_server/node_info_handler_impl.h @@ -32,6 +32,17 @@ class DefaultNodeInfoHandler : public rpc::NodeInfoHandler { ReportBatchHeartbeatReply *reply, SendReplyCallback send_reply_callback) override; + void HandleGetResources(const GetResourcesRequest &request, GetResourcesReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleUpdateResources(const UpdateResourcesRequest &request, + UpdateResourcesReply *reply, + SendReplyCallback send_reply_callback) override; + + void HandleDeleteResources(const DeleteResourcesRequest &request, + DeleteResourcesReply *reply, + SendReplyCallback send_reply_callback) override; + private: gcs::RedisGcsClient &gcs_client_; }; 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 7683e095d..3bea2a317 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 @@ -208,6 +208,44 @@ class GcsServerTest : public RedisServiceManagerForTest { return WaitReady(promise.get_future(), timeout_ms_); } + bool UpdateResources(const rpc::UpdateResourcesRequest &request) { + std::promise promise; + client_->UpdateResources(request, [&promise](const Status &status, + const rpc::UpdateResourcesReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + + bool DeleteResources(const rpc::DeleteResourcesRequest &request) { + std::promise promise; + client_->DeleteResources(request, [&promise](const Status &status, + const rpc::DeleteResourcesReply &reply) { + RAY_CHECK_OK(status); + promise.set_value(true); + }); + return WaitReady(promise.get_future(), timeout_ms_); + } + + std::map GetResources(const std::string &node_id) { + rpc::GetResourcesRequest request; + request.set_node_id(node_id); + std::map resources; + std::promise promise; + client_->GetResources(request, + [&resources, &promise](const Status &status, + const rpc::GetResourcesReply &reply) { + RAY_CHECK_OK(status); + for (auto resource : reply.resources()) { + resources[resource.first] = resource.second; + } + promise.set_value(true); + }); + EXPECT_TRUE(WaitReady(promise.get_future(), timeout_ms_)); + return resources; + } + bool AddObjectLocation(const rpc::AddObjectLocationRequest &request) { std::promise promise; client_->AddObjectLocation( @@ -390,6 +428,25 @@ TEST_F(GcsServerTest, TestNodeInfo) { ASSERT_TRUE(node_info_list.size() == 1); ASSERT_TRUE(node_info_list[0].state() == rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_DEAD); + + // Update node resources + rpc::UpdateResourcesRequest update_resources_request; + update_resources_request.set_node_id(node_id.Binary()); + rpc::ResourceTableData resource_table_data; + resource_table_data.set_resource_capacity(1.0); + std::string resource_name = "CPU"; + (*update_resources_request.mutable_resources())[resource_name] = resource_table_data; + ASSERT_TRUE(UpdateResources(update_resources_request)); + auto resources = GetResources(node_id.Binary()); + ASSERT_TRUE(resources.size() == 1); + + // Delete node resources + rpc::DeleteResourcesRequest delete_resources_request; + delete_resources_request.set_node_id(node_id.Binary()); + delete_resources_request.add_resource_name_list(resource_name); + ASSERT_TRUE(DeleteResources(delete_resources_request)); + resources = GetResources(node_id.Binary()); + ASSERT_TRUE(resources.size() == 0); } TEST_F(GcsServerTest, TestObjectInfo) { diff --git a/src/ray/protobuf/gcs_service.proto b/src/ray/protobuf/gcs_service.proto index 1916d0783..3e373b2d1 100644 --- a/src/ray/protobuf/gcs_service.proto +++ b/src/ray/protobuf/gcs_service.proto @@ -133,6 +133,30 @@ message ReportBatchHeartbeatRequest { message ReportBatchHeartbeatReply { } +message GetResourcesRequest { + bytes node_id = 1; +} + +message GetResourcesReply { + map resources = 1; +} + +message UpdateResourcesRequest { + bytes node_id = 1; + map resources = 2; +} + +message UpdateResourcesReply { +} + +message DeleteResourcesRequest { + bytes node_id = 1; + repeated string resource_name_list = 2; +} + +message DeleteResourcesReply { +} + // Service for node info access. service NodeInfoGcsService { // Register a node to GCS Service. @@ -146,6 +170,12 @@ service NodeInfoGcsService { // Report batch heartbeat to GCS Service. rpc ReportBatchHeartbeat(ReportBatchHeartbeatRequest) returns (ReportBatchHeartbeatReply); + // Get node's resources from GCS Service. + rpc GetResources(GetResourcesRequest) returns (GetResourcesReply); + // Update resources of a node in GCS Service. + rpc UpdateResources(UpdateResourcesRequest) returns (UpdateResourcesReply); + // Delete resources of a node in GCS Service. + rpc DeleteResources(DeleteResourcesRequest) returns (DeleteResourcesReply); } message GetObjectLocationsRequest { diff --git a/src/ray/rpc/gcs_server/gcs_rpc_client.h b/src/ray/rpc/gcs_server/gcs_rpc_client.h index 4ff5e5528..c64fc34be 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_client.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_client.h @@ -84,6 +84,18 @@ class GcsRpcClient { VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, ReportBatchHeartbeat, request, callback, node_info_grpc_client_) + /// Get node's resources from GCS Service. + VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, GetResources, request, callback, + node_info_grpc_client_) + + /// Update resources of a node in GCS Service. + VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, UpdateResources, request, callback, + node_info_grpc_client_) + + /// Delete resources of a node in GCS Service. + VOID_RPC_CLIENT_METHOD(NodeInfoGcsService, DeleteResources, request, callback, + node_info_grpc_client_) + /// Get object's locations from GCS Service. VOID_RPC_CLIENT_METHOD(ObjectInfoGcsService, GetObjectLocations, request, callback, object_info_grpc_client_) diff --git a/src/ray/rpc/gcs_server/gcs_rpc_server.h b/src/ray/rpc/gcs_server/gcs_rpc_server.h index fe1e650fb..7998a7630 100644 --- a/src/ray/rpc/gcs_server/gcs_rpc_server.h +++ b/src/ray/rpc/gcs_server/gcs_rpc_server.h @@ -169,6 +169,18 @@ class NodeInfoHandler { virtual void HandleReportBatchHeartbeat(const ReportBatchHeartbeatRequest &request, ReportBatchHeartbeatReply *reply, SendReplyCallback send_reply_callback) = 0; + + virtual void HandleGetResources(const GetResourcesRequest &request, + GetResourcesReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleUpdateResources(const UpdateResourcesRequest &request, + UpdateResourcesReply *reply, + SendReplyCallback send_reply_callback) = 0; + + virtual void HandleDeleteResources(const DeleteResourcesRequest &request, + DeleteResourcesReply *reply, + SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcService` for `NodeInfoGcsService`. @@ -193,6 +205,9 @@ class NodeInfoGrpcService : public GrpcService { NODE_INFO_SERVICE_RPC_HANDLER(GetAllNodeInfo, 1); NODE_INFO_SERVICE_RPC_HANDLER(ReportHeartbeat, 1); NODE_INFO_SERVICE_RPC_HANDLER(ReportBatchHeartbeat, 1); + NODE_INFO_SERVICE_RPC_HANDLER(GetResources, 1); + NODE_INFO_SERVICE_RPC_HANDLER(UpdateResources, 1); + NODE_INFO_SERVICE_RPC_HANDLER(DeleteResources, 1); } private: