Gcs pull resource reports

BUILD.bazel

@@ -1232,6 +1232,19 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "gcs_resource_report_poller_test",
+    srcs = [
+        "src/ray/gcs/gcs_server/test/gcs_resource_report_poller_test.cc",
+    ],
+    copts = COPTS,
+    deps = [
+        ":gcs_server_lib",
+        ":gcs_test_util_lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "service_based_gcs_client_lib",
     srcs = glob(

src/ray/common/ray_config_def.h

@@ -229,6 +229,13 @@ RAY_CONFIG(uint32_t, maximum_gcs_destroyed_actor_cached_count, 100000)
 RAY_CONFIG(uint32_t, maximum_gcs_dead_node_cached_count, 1000)
 /// The interval at which the gcs server will print debug info.
 RAY_CONFIG(int64_t, gcs_dump_debug_log_interval_minutes, 1)
+// The interval at which the gcs server will pull a new resource.
+RAY_CONFIG(int, gcs_resource_report_poll_period_ms, 100)
+// The number of concurrent polls to polls to GCS.
+RAY_CONFIG(uint64_t, gcs_max_concurrent_resource_pulls, 100)
+// Feature flag to turn on resource report polling. Polling and raylet pushing are
+// mutually exlusive.
+RAY_CONFIG(bool, pull_based_resource_reporting, true)
 
 /// Duration to sleep after failing to put an object in plasma because it is full.
 RAY_CONFIG(uint32_t, object_store_full_delay_ms, 10)

src/ray/gcs/gcs_server/gcs_resource_manager.cc

@@ -161,12 +161,10 @@ void GcsResourceManager::HandleGetAllAvailableResources(
   ++counts_[CountType::GET_ALL_AVAILABLE_RESOURCES_REQUEST];
 }
 
-void GcsResourceManager::HandleReportResourceUsage(
-    const rpc::ReportResourceUsageRequest &request, rpc::ReportResourceUsageReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
-  NodeID node_id = NodeID::FromBinary(request.resources().node_id());
+void GcsResourceManager::UpdateFromResourceReport(const rpc::ResourcesData &data) {
+  NodeID node_id = NodeID::FromBinary(data.node_id());
   auto resources_data = std::make_shared<rpc::ResourcesData>();
-  resources_data->CopyFrom(request.resources());
+  resources_data->CopyFrom(data);
 
   // We use `node_resource_usages_` to filter out the nodes that report resource
   // information for the first time. `UpdateNodeResourceUsage` will modify
@@ -177,13 +175,19 @@ void GcsResourceManager::HandleReportResourceUsage(
     SetAvailableResources(node_id, ResourceSet(resource_changed));
   }
 
-  UpdateNodeResourceUsage(node_id, request);
+  UpdateNodeResourceUsage(node_id, data);
 
   if (resources_data->should_global_gc() || resources_data->resources_total_size() > 0 ||
       resources_data->resources_available_changed() ||
       resources_data->resource_load_changed()) {
     resources_buffer_[node_id] = *resources_data;
   }
+}
+
+void GcsResourceManager::HandleReportResourceUsage(
+    const rpc::ReportResourceUsageRequest &request, rpc::ReportResourceUsageReply *reply,
+    rpc::SendReplyCallback send_reply_callback) {
+  UpdateFromResourceReport(request.resources());
 
   GCS_RPC_SEND_REPLY(send_reply_callback, reply, Status::OK());
   ++counts_[CountType::REPORT_RESOURCE_USAGE_REQUEST];
@@ -238,26 +242,24 @@ void GcsResourceManager::HandleGetAllResourceUsage(
   ++counts_[CountType::GET_ALL_RESOURCE_USAGE_REQUEST];
 }
 
-void GcsResourceManager::UpdateNodeResourceUsage(
-    const NodeID node_id, const rpc::ReportResourceUsageRequest &request) {
+void GcsResourceManager::UpdateNodeResourceUsage(const NodeID &node_id,
+                                                 const rpc::ResourcesData &resources) {
   auto iter = node_resource_usages_.find(node_id);
   if (iter == node_resource_usages_.end()) {
     auto resources_data = std::make_shared<rpc::ResourcesData>();
-    resources_data->CopyFrom(request.resources());
+    resources_data->CopyFrom(resources);
     node_resource_usages_[node_id] = *resources_data;
   } else {
-    if (request.resources().resources_total_size() > 0) {
-      (*iter->second.mutable_resources_total()) = request.resources().resources_total();
+    if (resources.resources_total_size() > 0) {
+      (*iter->second.mutable_resources_total()) = resources.resources_total();
     }
-    if (request.resources().resources_available_changed()) {
-      (*iter->second.mutable_resources_available()) =
-          request.resources().resources_available();
+    if (resources.resources_available_changed()) {
+      (*iter->second.mutable_resources_available()) = resources.resources_available();
     }
-    if (request.resources().resource_load_changed()) {
-      (*iter->second.mutable_resource_load()) = request.resources().resource_load();
+    if (resources.resource_load_changed()) {
+      (*iter->second.mutable_resource_load()) = resources.resource_load();
     }
-    (*iter->second.mutable_resource_load_by_shape()) =
-        request.resources().resource_load_by_shape();
+    (*iter->second.mutable_resource_load_by_shape()) = resources.resource_load_by_shape();
   }
 }
 

src/ray/gcs/gcs_server/gcs_resource_manager.h

@@ -136,8 +136,12 @@ class GcsResourceManager : public rpc::NodeResourceInfoHandler {
   ///
   /// \param node_id Node id.
   /// \param request Request containing resource usage.
-  void UpdateNodeResourceUsage(const NodeID node_id,
-                               const rpc::ReportResourceUsageRequest &request);
+  void UpdateNodeResourceUsage(const NodeID &node_id,
+                               const rpc::ResourcesData &resources);
+
+  /// Process a new resource report from a node, independent of the rpc handler it came
+  /// from.
+  void UpdateFromResourceReport(const rpc::ResourcesData &data);
 
   /// Update the placement group load information so that it will be reported through
   /// heartbeat.

src/ray/gcs/gcs_server/gcs_resource_report_poller.cc

@@ -0,0 +1,147 @@
+#include "ray/gcs/gcs_server/gcs_resource_report_poller.h"
+
+namespace ray {
+namespace gcs {
+
+GcsResourceReportPoller::GcsResourceReportPoller(
+    std::shared_ptr<GcsResourceManager> gcs_resource_manager,
+    std::shared_ptr<rpc::NodeManagerClientPool> raylet_client_pool,
+    std::function<void(const rpc::ResourcesData &)> handle_resource_report,
+    std::function<int64_t(void)> get_current_time_milli,
+    std::function<void(
+        const rpc::Address &, std::shared_ptr<rpc::NodeManagerClientPool> &,
+        std::function<void(const Status &, const rpc::RequestResourceReportReply &)>)>
+        request_report)
+    : ticker_(polling_service_),
+      max_concurrent_pulls_(RayConfig::instance().gcs_max_concurrent_resource_pulls()),
+      inflight_pulls_(0),
+      gcs_resource_manager_(gcs_resource_manager),
+      raylet_client_pool_(raylet_client_pool),
+      handle_resource_report_(handle_resource_report),
+      get_current_time_milli_(get_current_time_milli),
+      request_report_(request_report),
+      poll_period_ms_(RayConfig::instance().gcs_resource_report_poll_period_ms()) {}
+
+GcsResourceReportPoller::~GcsResourceReportPoller() { Stop(); }
+
+void GcsResourceReportPoller::Initialize(const GcsInitData &gcs_init_data) {
+  for (const auto &pair : gcs_init_data.Nodes()) {
+    HandleNodeAdded(pair.second);
+  }
+}
+
+void GcsResourceReportPoller::Start() {
+  polling_thread_.reset(new std::thread{[this]() {
+    SetThreadName("resource_report_poller");
+    boost::asio::io_service::work work(polling_service_);
+
+    polling_service_.run();
+    RAY_LOG(DEBUG) << "GCSResourceReportPoller has stopped. This should only happen if "
+                      "the cluster has stopped";
+  }});
+  ticker_.RunFnPeriodically([this] { TryPullResourceReport(); }, 10);
+}
+
+void GcsResourceReportPoller::Stop() {
+  if (polling_thread_ != nullptr) {
+    // TODO (Alex): There's technically a race condition here if we start and stop the
+    // thread in rapid succession.
+    polling_service_.stop();
+    if (polling_thread_->joinable()) {
+      polling_thread_->join();
+    }
+  }
+}
+
+void GcsResourceReportPoller::HandleNodeAdded(const rpc::GcsNodeInfo &node_info) {
+  absl::MutexLock guard(&mutex_);
+
+  rpc::Address address;
+  address.set_raylet_id(node_info.node_id());
+  address.set_ip_address(node_info.node_manager_address());
+  address.set_port(node_info.node_manager_port());
+
+  auto state =
+      std::make_shared<PullState>(NodeID::FromBinary(node_info.node_id()),
+                                  std::move(address), -1, get_current_time_milli_());
+
+  const auto &node_id = state->node_id;
+
+  RAY_CHECK(!nodes_.count(node_id)) << "Node with id: " << node_id << " was added twice!";
+
+  nodes_[node_id] = state;
+  to_pull_queue_.push_front(state);
+  RAY_LOG(DEBUG) << "Node was added with id: " << node_id;
+
+  polling_service_.post([this]() { TryPullResourceReport(); });
+}
+
+void GcsResourceReportPoller::HandleNodeRemoved(const rpc::GcsNodeInfo &node_info) {
+  NodeID node_id = NodeID::FromBinary(node_info.node_id());
+  {
+    absl::MutexLock guard(&mutex_);
+    nodes_.erase(node_id);
+    RAY_CHECK(!nodes_.count(node_id));
+    RAY_LOG(DEBUG) << "Node removed (node_id: " << node_id
+                   << ")# of remaining nodes: " << nodes_.size();
+  }
+}
+
+void GcsResourceReportPoller::TryPullResourceReport() {
+  absl::MutexLock guard(&mutex_);
+  int64_t cur_time = get_current_time_milli_();
+
+  RAY_LOG(DEBUG) << "Trying to pull inflight_pulls " << inflight_pulls_ << "/"
+                 << max_concurrent_pulls_ << ", queue size: " << to_pull_queue_.size();
+  while (inflight_pulls_ < max_concurrent_pulls_ && !to_pull_queue_.empty()) {
+    auto to_pull = to_pull_queue_.front();
+    if (cur_time < to_pull->next_pull_time) {
+      break;
+    }
+
+    to_pull_queue_.pop_front();
+
+    if (!nodes_.count(to_pull->node_id)) {
+      RAY_LOG(DEBUG)
+          << "Update finished, but node was already removed from the cluster. Ignoring.";
+      continue;
+    }
+
+    PullResourceReport(to_pull);
+  }
+}
+
+void GcsResourceReportPoller::PullResourceReport(const std::shared_ptr<PullState> state) {
+  inflight_pulls_++;
+
+  request_report_(
+      state->address, raylet_client_pool_,
+      [this, state](const Status &status, const rpc::RequestResourceReportReply &reply) {
+        if (status.ok()) {
+          // TODO (Alex): This callback is always posted onto the main thread. Since most
+          // of the work is in the callback we should move this callback's execution to
+          // the polling thread. We will need to implement locking once we switch threads.
+          handle_resource_report_(reply.resources());
+        } else {
+          RAY_LOG(INFO) << "Couldn't get resource request from raylet " << state->node_id
+                        << ": " << status.ToString();
+        }
+        polling_service_.post([this, state]() { NodeResourceReportReceived(state); });
+      });
+}
+
+void GcsResourceReportPoller::NodeResourceReportReceived(
+    const std::shared_ptr<PullState> state) {
+  absl::MutexLock guard(&mutex_);
+  inflight_pulls_--;
+
+  // Schedule the next pull. The scheduling `TryPullResourceReport` loop will handle
+  // validating that this node is still in the cluster.
+  state->next_pull_time = get_current_time_milli_() + poll_period_ms_;
+  to_pull_queue_.push_back(state);
+
+  polling_service_.post([this] { TryPullResourceReport(); });
+}
+
+}  // namespace gcs
+}  // namespace ray