diff --git a/src/ray/gcs/gcs_server/gcs_node_manager.cc b/src/ray/gcs/gcs_server/gcs_node_manager.cc index 2dd1cd557..a3180e128 100644 --- a/src/ray/gcs/gcs_server/gcs_node_manager.cc +++ b/src/ray/gcs/gcs_server/gcs_node_manager.cc @@ -21,6 +21,18 @@ namespace ray { namespace gcs { +bool GcsNodeResource::IsSubset( + const std::unordered_map &request_resources) const { + for (const auto &request_resource : request_resources) { + auto iter = resources_available_.find(request_resource.first); + if (iter == resources_available_.end() || iter->second < request_resource.second) { + return false; + } + } + return true; +} + +////////////////////////////////////////////////////////////////////////////////////////// GcsNodeManager::NodeFailureDetector::NodeFailureDetector( boost::asio::io_service &io_service, std::shared_ptr gcs_table_storage, @@ -231,6 +243,10 @@ void GcsNodeManager::HandleReportHeartbeat(const rpc::ReportHeartbeatRequest &re ClientID node_id = ClientID::FromBinary(request.heartbeat().client_id()); auto heartbeat_data = std::make_shared(); heartbeat_data->CopyFrom(request.heartbeat()); + + // Update node realtime resources. + UpdateNodeRealtimeResources(node_id, *heartbeat_data); + // Note: To avoid heartbeats being delayed by main thread, make sure heartbeat is always // handled by its own IO service. node_failure_detector_service_.post([this, node_id, heartbeat_data] { @@ -451,5 +467,22 @@ void GcsNodeManager::StartNodeFailureDetector() { node_failure_detector_service_.post([this] { node_failure_detector_->Start(); }); } +void GcsNodeManager::UpdateNodeRealtimeResources( + const ClientID &node_id, const rpc::HeartbeatTableData &heartbeat) { + auto resources_available = MapFromProtobuf(heartbeat.resources_available()); + auto iter = cluster_realtime_resources_.find(node_id); + if (iter != cluster_realtime_resources_.end()) { + iter->second->resources_available_ = resources_available; + } else { + cluster_realtime_resources_[node_id] = + std::make_shared(resources_available); + } +} + +const absl::flat_hash_map> + &GcsNodeManager::GetClusterRealtimeResources() const { + return cluster_realtime_resources_; +} + } // namespace gcs } // namespace ray diff --git a/src/ray/gcs/gcs_server/gcs_node_manager.h b/src/ray/gcs/gcs_server/gcs_node_manager.h index 6a4d96518..3aa27094a 100644 --- a/src/ray/gcs/gcs_server/gcs_node_manager.h +++ b/src/ray/gcs/gcs_server/gcs_node_manager.h @@ -27,6 +27,18 @@ namespace ray { namespace gcs { +/// Node resource information. +class GcsNodeResource { + public: + explicit GcsNodeResource(std::unordered_map resources_available) + : resources_available_(std::move(resources_available)) {} + + bool IsSubset(const std::unordered_map &request_resources) const; + + /// Dynamic resource capacity. + std::unordered_map resources_available_; +}; + /// GcsNodeManager is responsible for managing and monitoring nodes as well as handing /// node and resource related rpc requests. /// This class is not thread-safe. @@ -143,6 +155,14 @@ class GcsNodeManager : public rpc::NodeInfoHandler { /// Start node failure detector. void StartNodeFailureDetector(); + // Update node realtime resources. + void UpdateNodeRealtimeResources(const ClientID &node_id, + const rpc::HeartbeatTableData &heartbeat); + + /// Get cluster realtime resources. + const absl::flat_hash_map> + &GetClusterRealtimeResources() const; + protected: class NodeFailureDetector { public: @@ -239,6 +259,9 @@ class GcsNodeManager : public rpc::NodeInfoHandler { std::shared_ptr gcs_pub_sub_; /// Storage for GCS tables. std::shared_ptr gcs_table_storage_; + /// Cluster realtime resources. + absl::flat_hash_map> + cluster_realtime_resources_; }; } // namespace gcs diff --git a/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc b/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc index 6a493f743..53ce725db 100644 --- a/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc +++ b/src/ray/gcs/gcs_server/test/gcs_node_manager_test.cc @@ -86,6 +86,28 @@ TEST_F(GcsNodeManagerTest, TestListener) { } } +TEST_F(GcsNodeManagerTest, TestGetClusterRealtimeResources) { + boost::asio::io_service io_service; + gcs::GcsNodeManager node_manager(io_service, io_service, gcs_pub_sub_, + gcs_table_storage_); + + auto node_id = ClientID::FromRandom(); + rpc::HeartbeatTableData heartbeat; + (*heartbeat.mutable_resources_available())["CPU"] = 10; + node_manager.UpdateNodeRealtimeResources(node_id, heartbeat); + auto node_resources = node_manager.GetClusterRealtimeResources(); + std::unordered_map request_resources; + request_resources["CPU"] = 9; + ASSERT_TRUE(node_resources[node_id]->IsSubset(request_resources)); + request_resources["CPU"] = 10; + ASSERT_TRUE(node_resources[node_id]->IsSubset(request_resources)); + request_resources["CPU"] = 10.1; + ASSERT_FALSE(node_resources[node_id]->IsSubset(request_resources)); + request_resources.clear(); + request_resources["GPU"] = 1; + ASSERT_FALSE(node_resources[node_id]->IsSubset(request_resources)); +} + } // namespace ray int main(int argc, char **argv) {