Gcs pull resource reports

python/ray/tests/test_multi_node_2.py

@@ -101,7 +101,7 @@ def verify_load_metrics(monitor, expected_resource_usage=None, timeout=30):
     pg_demands = [{"GPU": 2}, {"extra_resource": 2}]
     strategy = "STRICT_PACK"
     pg = placement_group(pg_demands, strategy=strategy)
-    pg.ready()
+    # ray.get(pg.ready())
     time.sleep(2)  # wait for placemnt groups to propogate.
 
     # Disable event clearing for test.

src/ray/common/ray_config_def.h

@@ -228,6 +228,11 @@ 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)
+// Feature flag to turn on resource report polling. Polling and raylet pushing or 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

@@ -159,12 +159,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
@@ -175,13 +173,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];
@@ -236,26 +240,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

@@ -134,8 +134,11 @@ 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,130 @@
+#include "ray/gcs/gcs_server/gcs_resource_report_poller.h"
+
+namespace ray {
+namespace gcs {
+
+GcsResourceReportPoller::GcsResourceReportPoller(
+    uint64_t max_concurrent_pulls,
+    std::shared_ptr<GcsResourceManager> gcs_resource_manager,
+    std::shared_ptr<rpc::NodeManagerClientPool> raylet_client_pool)
+    : max_concurrent_pulls_(max_concurrent_pulls),
+      inflight_pulls_(0),
+      gcs_resource_manager_(gcs_resource_manager),
+      raylet_client_pool_(raylet_client_pool),
+      poll_period_ms_(boost::posix_time::milliseconds(
+          RayConfig::instance().gcs_resource_report_poll_period_ms())) {}
+
+GcsResourceReportPoller::~GcsResourceReportPoller() {
+  Stop();
+}
+
+void GcsResourceReportPoller::Start() {
+  polling_thread_ = std::unique_ptr<std::thread>(new std::thread{[&]() {
+    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";
+  }});
+}
+
+void GcsResourceReportPoller::Stop() {
+  polling_service_.stop();
+  if (polling_thread_->joinable()) {
+    polling_thread_->join();
+  }
+}
+
+void GcsResourceReportPoller::HandleNodeAdded(
+    std::shared_ptr<rpc::GcsNodeInfo> node_info) {
+  absl::MutexLock guard(&mutex_);
+  const auto node_id = NodeID::FromBinary(node_info->node_id());
+
+  RAY_CHECK(!nodes_.count(node_id)) << "Node with id: " << node_id << " was added twice!";
+
+  auto &state = nodes_[node_id];
+  state.node_id = node_id;
+
+  state.address.set_raylet_id(node_info->node_id());
+  state.address.set_ip_address(node_info->node_manager_address());
+  state.address.set_port(node_info->node_manager_port());
+
+  state.last_pull_time = absl::GetCurrentTimeNanos();
+
+  state.next_pull_timer = std::unique_ptr<boost::asio::deadline_timer>(
+      new boost::asio::deadline_timer(polling_service_));
+
+
+  polling_service_.post([&, node_id]() { TryPullResourceReport(node_id); });
+}
+
+void GcsResourceReportPoller::HandleNodeRemoved(
+    std::shared_ptr<rpc::GcsNodeInfo> node_info) {
+  NodeID node_id = NodeID::FromBinary(node_info->node_id());
+
+  {
+    absl::MutexLock guard(&mutex_);
+    nodes_.erase(node_id);
+  }
+}
+
+void GcsResourceReportPoller::TryPullResourceReport(const NodeID &node_id) {
+  absl::MutexLock guard(&mutex_);
+
+  to_pull_queue_.push_back(node_id);
+
+  while (inflight_pulls_ < max_concurrent_pulls_ && !to_pull_queue_.empty()) {
+    const NodeID &to_pull = to_pull_queue_.front();
+    to_pull_queue_.pop_front();
+
+    auto it = nodes_.find(to_pull);
+    if (it == nodes_.end()) {
+      RAY_LOG(DEBUG)
+          << "Update finished, but node was already removed from the cluster. Ignoring.";
+      continue;
+    }
+    auto &state = it->second;
+    PullResourceReport(state);
+  }
+}
+
+void GcsResourceReportPoller::PullResourceReport(PullState &state) {
+  inflight_pulls_++;
+  const auto &node_id = state.node_id;
+  auto raylet_client = raylet_client_pool_->GetOrConnectByAddress(state.address);
+  raylet_client->RequestResourceReport(
+      [&, node_id](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.
+          gcs_resource_manager_->UpdateFromResourceReport(reply.resources());
+          polling_service_.post([this, node_id] { NodeResourceReportReceived(node_id); });
+        } else {
+          RAY_LOG(INFO) << "Couldn't get resource request from raylet: " << node_id << ". " << status.ToString();
+        }
+      });
+}
+
+void GcsResourceReportPoller::NodeResourceReportReceived(const NodeID &node_id) {
+  absl::MutexLock guard(&mutex_);
+  inflight_pulls_--;
+  auto it = nodes_.find(node_id);
+  if (it == nodes_.end()) {
+    RAY_LOG(DEBUG)
+        << "Update finished, but node was already removed from the cluster. Ignoring.";
+    return;
+  }
+
+  auto &state = it->second;
+  state.next_pull_timer->expires_from_now(poll_period_ms_);
+  state.next_pull_timer->async_wait([&, node_id](const boost::system::error_code &error) {
+    if (!error) {
+      TryPullResourceReport(node_id);
+    }
+    RAY_LOG(INFO) << "GcsResourceReportPoller timer failed: " << error.message() << ".";
+  });
+}
+
+}  // namespace gcs
+}  // namespace ray

src/ray/gcs/gcs_server/gcs_resource_report_poller.h

@@ -0,0 +1,97 @@
+#include "ray/gcs/gcs_server/gcs_resource_manager.h"
+#include "ray/rpc/node_manager/node_manager_client_pool.h"
+
+namespace ray {
+namespace gcs {
+
+/// Polls raylets on a separate thread to update GCS's view of the cluster's resource
+/// utilization. This class creates and manages a polling thread. All public methods are
+/// thread safe.
+class GcsResourceReportPoller {
+  /*
+  This class roughly polls each node independently (with the exception of max
+  concurrency). The process for polling a single node is as follows:
+
+  A new node joins the cluster.
+  1. (Main thread) Begin tracking the node, and begin the polling process.
+
+  Main polling procedure.
+  2. (Polling thread) Enqueue the node to be pulled.
+  3. (Polling thread) Node is popped off the back of the queue and RequestResourceReport
+  is sent to the raylet.
+  4. (Main thread) The raylet responds and the resource manager is updated. This section
+  is _not_ thread safe (i.e. should not modify the resource report poller state).
+  5. (Polling thread) The RequestResourceReport continuation runs, scheduling the next
+  pull time.
+  6. (Polling thread) The next pull time occurs, and step 2 is repeated.
+
+  The node leaves the cluster.
+  7. Untrack the node. The next time the main polling procedure comes across the node, it
+  should be dropped from the system.
+   */
+
+ public:
+  GcsResourceReportPoller(uint64_t max_concurrent_pulls,
+                          std::shared_ptr<GcsResourceManager> gcs_resource_manager,
+                          std::shared_ptr<rpc::NodeManagerClientPool> raylet_client_pool);
+
+  ~GcsResourceReportPoller();
+
+  /// Start a thread to poll for resource updates.
+  void Start();
+
+  /// Stop polling for resource updates.
+  void Stop();
+
+  /// Event handler when a new node joins the cluster.
+  void HandleNodeAdded(std::shared_ptr<rpc::GcsNodeInfo> node_info);
+
+  /// Event handler when a node leaves the cluster.
+  void HandleNodeRemoved(std::shared_ptr<rpc::GcsNodeInfo> node_info);
+
+ private:
+  // An asio service which does the polling work.
+  boost::asio::io_context polling_service_;
+  // The associated thread it runs on.
+  std::unique_ptr<std::thread> polling_thread_;
+
+  // The maximum number of pulls that can occur at once.
+  const uint64_t max_concurrent_pulls_;
+  // The number of ongoing pulls.
+  uint64_t inflight_pulls_;
+  // The resource manager which maintains GCS's view of the cluster's resource
+  // utilization.
+  std::shared_ptr<GcsResourceManager> gcs_resource_manager_;
+  // The shared, thread safe pool of raylet clients, which we use to minimize connections.
+  std::shared_ptr<rpc::NodeManagerClientPool> raylet_client_pool_;
+  // The minimum delay between two pull requests to the same thread.
+  const boost::posix_time::milliseconds poll_period_ms_;
+
+  struct PullState {
+    NodeID node_id;
+    rpc::Address address;
+    int64_t last_pull_time;
+    std::unique_ptr<boost::asio::deadline_timer> next_pull_timer;
+  };
+
+  // A global lock for internal operations. This lock is shared between the main thread
+  // and polling thread, so we should be mindful about how long we hold it.
+  absl::Mutex mutex_;
+  // All the state regarding how to and when to send a new pull request to a raylet.
+  std::unordered_map<NodeID, PullState> nodes_;
+  // The set of all nodes which we are allowed to pull from. We can't necessarily pull
+  // from this list immediately because we limit the number of concurrent pulls.
+  std::deque<NodeID> to_pull_queue_;
+
+  /// Try to pull from the node. We may not be able to if it violates max concurrent
+  /// pulls. This method is thread safe.
+  void TryPullResourceReport(const NodeID &node_id);
+  /// Pull resource report without validation. This method is NOT thread safe.
+  void PullResourceReport(PullState &state);
+  /// A resource report was successfully pulled (and the resource manager was already
+  /// updated). This method is thread safe.
+  void NodeResourceReportReceived(const NodeID &node_id);
+};
+
+}  // namespace gcs
+}  // namespace ray

src/ray/gcs/gcs_server/gcs_server.cc

@@ -99,6 +99,9 @@ void GcsServer::DoStart(const GcsInitData &gcs_init_data) {
   // Init stats handler.
   InitStatsHandler();
 
+  // Init resource report polling.
+  InitResourceReportPolling();
+
   // Install event listeners.
   InstallEventListeners();
 
@@ -130,6 +133,8 @@ void GcsServer::Stop() {
 
     gcs_heartbeat_manager_->Stop();
 
+    gcs_resource_report_poller_->Stop();
+
     is_stopped_ = true;
     RAY_LOG(INFO) << "GCS server stopped.";
   }
@@ -281,6 +286,14 @@ void GcsServer::InitTaskInfoHandler() {
   rpc_server_.RegisterService(*task_info_service_);
 }
 
+void GcsServer::InitResourceReportPolling() {
+  gcs_resource_report_poller_ = std::unique_ptr<GcsResourceReportPoller>(
+      new GcsResourceReportPoller(100, gcs_resource_manager_, raylet_client_pool_));
+  if (config_.pull_based_resource_reporting) {
+    gcs_resource_report_poller_->Start();
+  }
+}
+
 void GcsServer::InitStatsHandler() {
   RAY_CHECK(gcs_table_storage_);
   stats_handler_.reset(new rpc::DefaultStatsHandler(gcs_table_storage_));
@@ -307,6 +320,7 @@ void GcsServer::InstallEventListeners() {
     gcs_placement_group_manager_->SchedulePendingPlacementGroups();
     gcs_actor_manager_->SchedulePendingActors();
     gcs_heartbeat_manager_->AddNode(NodeID::FromBinary(node->node_id()));
+    gcs_resource_report_poller_->HandleNodeAdded(node);
   });
   gcs_node_manager_->AddNodeRemovedListener(
       [this](std::shared_ptr<rpc::GcsNodeInfo> node) {
@@ -317,6 +331,7 @@ void GcsServer::InstallEventListeners() {
         gcs_placement_group_manager_->OnNodeDead(node_id);
         gcs_actor_manager_->OnNodeDead(node_id);
         raylet_client_pool_->Disconnect(NodeID::FromBinary(node->node_id()));
+        gcs_resource_report_poller_->HandleNodeRemoved(node);
       });
 
   // Install worker event listener.