[Core] Improve Large-Scale Resource View Synchronization Through Sync Message Batching

src/ray/common/BUILD.bazel

@@ -16,6 +16,7 @@ ray_cc_library(
     hdrs = ["test_utils.h"],
     deps = [
         ":asio",
+        ":function_descriptor",
         ":id",
         ":placement_group",
         ":ray_object",

src/ray/common/ray_config_def.h

@@ -1017,3 +1017,15 @@ RAY_CONFIG(bool, start_python_gc_manager_thread, true)
 // Whether to enable the feature of outputting error log if the task is
 // still retryable.
 RAY_CONFIG(bool, enable_output_error_log_if_still_retry, true)
+
+// The maximum batch size for ray_syncer_bidi_reactor to sync messages to other nodes.
+// If `gcs_resource_broadcast_max_batch_size == 1`, the ray_syncer_bidi_reactor will
+// disable batching. Otherwise, it will start send with maximum of
+// `gcs_resource_broadcast_max_batch_size` messages before timeout.
+RAY_CONFIG(size_t, gcs_resource_broadcast_max_batch_size, 1)
+
+// The maximum period that the GCS will wait for resource update messages before
+// broadcasting the batch. If `gcs_resource_broadcast_max_batch_size` messages arrive
+// before the timeout, the batch will be broadcasted eagerly. This flag only applies if
+// `gcs_resource_broadcast_max_batch_size != 1`.
+RAY_CONFIG(uint64_t, gcs_resource_broadcast_max_batch_delay_ms, 0)

src/ray/gcs/gcs_server.cc

@@ -607,9 +607,23 @@ GcsServer::StorageType GcsServer::GetStorageType() const {
 }
 
 void GcsServer::InitRaySyncer(const GcsInitData &gcs_init_data) {
+  if (RayConfig::instance().gcs_resource_broadcast_max_batch_delay_ms() > 0 &&
+      RayConfig::instance().gcs_resource_broadcast_max_batch_size() == 1) {
+    RAY_LOG(WARNING) << "Configuration inconsistency detected: "
+                        "gcs_resource_broadcast_max_batch_delay_ms > 0 "
+                     << "but gcs_resource_broadcast_max_batch_size = 1. The delay "
+                        "setting will have no effect "
+                     << "since only one message can be batched at a time. Consider "
+                        "increasing batch_size or "
+                     << "setting delay_ms to 0.";
+  }
+
   ray_syncer_ = std::make_unique<syncer::RaySyncer>(
       io_context_provider_.GetIOContext<syncer::RaySyncer>(),
       kGCSNodeID.Binary(),
+      /* batch_size */ RayConfig::instance().gcs_resource_broadcast_max_batch_size(),
+      /* batch_delay_ms */
+      RayConfig::instance().gcs_resource_broadcast_max_batch_delay_ms(),
       [this](const NodeID &node_id) {
         gcs_healthcheck_manager_->MarkNodeHealthy(node_id);
       });
@@ -976,6 +990,7 @@ void GcsServer::TryGlobalGC() {
     std::string serialized_msg;
     RAY_CHECK(commands_sync_message.SerializeToString(&serialized_msg));
     msg->set_sync_message(std::move(serialized_msg));
+
     ray_syncer_->BroadcastMessage(std::move(msg));
   }
 }

src/ray/protobuf/ray_syncer.proto

@@ -60,6 +60,11 @@ message RaySyncMessage {
   bytes node_id = 4;
 }
 
+message RaySyncMessageBatch {
+  // Batch of RaySyncMessages.
+  repeated RaySyncMessage messages = 1;
+}
+
 service RaySyncer {
-  rpc StartSync(stream RaySyncMessage) returns (stream RaySyncMessage);
+  rpc StartSync(stream RaySyncMessageBatch) returns (stream RaySyncMessageBatch);
 }

src/ray/ray_syncer/ray_syncer.cc

@@ -30,11 +30,15 @@ namespace ray::syncer {
 
 RaySyncer::RaySyncer(instrumented_io_context &io_context,
                      const std::string &local_node_id,
+                     size_t max_batch_size,
+                     uint64_t max_batch_delay_ms,
                      RpcCompletionCallback on_rpc_completion)
     : io_context_(io_context),
       local_node_id_(local_node_id),
       node_state_(std::make_unique<NodeState>()),
       timer_(PeriodicalRunner::Create(io_context)),
+      max_batch_size_(max_batch_size),
+      max_batch_delay_ms_(max_batch_delay_ms),
       on_rpc_completion_(std::move(on_rpc_completion)) {
   stopped_ = std::make_shared<bool>(false);
 }
@@ -79,7 +83,7 @@ void RaySyncer::Connect(const std::string &node_id,
   boost::asio::dispatch(
       io_context_.get_executor(), std::packaged_task<void()>([=]() {
         auto stub = ray::rpc::syncer::RaySyncer::NewStub(channel);
-        auto *reactor = new RayClientBidiReactor(
+        auto reactor = std::make_shared<RayClientBidiReactor>(
             /* remote_node_id */ node_id,
             /* local_node_id */ GetLocalNodeID(),
             /* io_context */ io_context_,
@@ -90,7 +94,7 @@ void RaySyncer::Connect(const std::string &node_id,
               const std::string &remote_node_id = bidi_reactor->GetRemoteNodeID();
               auto iter = sync_reactors_.find(remote_node_id);
               if (iter != sync_reactors_.end()) {
-                if (iter->second != bidi_reactor) {
+                if (iter->second.get() != bidi_reactor) {
                   // The client is already reconnected.
                   return;
                 }
@@ -109,14 +113,17 @@ void RaySyncer::Connect(const std::string &node_id,
                 node_state_->RemoveNode(remote_node_id);
               }
             },
-            /* stub */ std::move(stub));
+            /* stub */ std::move(stub),
+            /* max_batch_size */ max_batch_size_,
+            /* max_batch_delay_ms */ max_batch_delay_ms_);
+        reactor->SetSelfRef(reactor);
         Connect(reactor);
         reactor->StartCall();
       }))
       .get();
 }
 
-void RaySyncer::Connect(RaySyncerBidiReactor *reactor) {
+void RaySyncer::Connect(std::shared_ptr<RaySyncerBidiReactor> reactor) {
   // Bind rpc completion callback.
   if (on_rpc_completion_) {
     reactor->SetRpcCompletionCallbackForOnce(on_rpc_completion_);
@@ -217,19 +224,20 @@ void RaySyncer::BroadcastMessage(std::shared_ptr<const RaySyncMessage> message)
 }
 
 ServerBidiReactor *RaySyncerService::StartSync(grpc::CallbackServerContext *context) {
-  auto reactor = new RayServerBidiReactor(
+  auto reactor = std::make_shared<RayServerBidiReactor>(
       context,
       syncer_.GetIOContext(),
       syncer_.GetLocalNodeID(),
-      /*message_processor=*/[this](auto msg) mutable { syncer_.BroadcastMessage(msg); },
+      /*message_processor=*/
+      [this](auto msg) mutable { syncer_.BroadcastMessage(msg); },
       /*cleanup_cb=*/
       [this](RaySyncerBidiReactor *bidi_reactor, bool reconnect) mutable {
         // No need to reconnect for server side.
         RAY_CHECK(!reconnect);
         const auto &node_id = bidi_reactor->GetRemoteNodeID();
         auto iter = syncer_.sync_reactors_.find(node_id);
         if (iter != syncer_.sync_reactors_.end()) {
-          if (iter->second != bidi_reactor) {
+          if (iter->second.get() != bidi_reactor) {
             // There is a new connection to the node, no need to clean up.
             // This can happen when there is transient network error and the client
             // reconnects. The sequence of events are:
@@ -245,22 +253,26 @@ ServerBidiReactor *RaySyncerService::StartSync(grpc::CallbackServerContext *cont
         RAY_LOG(INFO).WithField(NodeID::FromBinary(node_id)) << "Connection is broken.";
         syncer_.node_state_->RemoveNode(node_id);
       },
-      /*auth_token=*/auth_token_);
+      /*auth_token=*/auth_token_,
+      /*max_batch_size=*/syncer_.max_batch_size_,
+      /*max_batch_delay_ms=*/syncer_.max_batch_delay_ms_);
   RAY_LOG(DEBUG).WithField(NodeID::FromBinary(reactor->GetRemoteNodeID()))
       << "Get connection";
 
+  reactor->SetSelfRef(reactor);
+
   // If the reactor has already called Finish() (e.g., due to authentication failure),
   // skip registration. The reactor will clean itself up via OnDone().
   if (reactor->IsFinished()) {
-    return reactor;
+    return reactor.get();
   }
 
   // Disconnect exiting connection if there is any.
   // This can happen when there is transient network error
   // and the client reconnects.
   syncer_.Disconnect(reactor->GetRemoteNodeID());
   syncer_.Connect(reactor);
-  return reactor;
+  return reactor.get();
 }
 
 }  // namespace ray::syncer

src/ray/ray_syncer/ray_syncer.h

@@ -37,6 +37,7 @@ namespace ray::syncer {
 using ray::rpc::syncer::CommandsSyncMessage;
 using ray::rpc::syncer::MessageType;
 using ray::rpc::syncer::RaySyncMessage;
+using ray::rpc::syncer::RaySyncMessageBatch;
 using ray::rpc::syncer::ResourceViewSyncMessage;
 
 /// The interface for a reporter. Reporter is defined to be a local module which would
@@ -93,8 +94,13 @@ class RaySyncer {
   /// \param io_context The io context for this component.
   /// \param node_id The id of current node.
   /// \param on_rpc_completion A callback which invokes after a sync rpc succeeds.
+  /// \param max_batch_size The max number of messages to be sent in a batch.
+  /// \param max_batch_delay_ms The max delay in milliseconds to wait before sending a
+  /// batch.
   RaySyncer(instrumented_io_context &io_context,
             const std::string &node_id,
+            size_t max_batch_size,
+            uint64_t max_batch_delay_ms,
             RpcCompletionCallback on_rpc_completion = {});
   ~RaySyncer();
 
@@ -155,7 +161,7 @@ class RaySyncer {
   std::vector<std::string> GetAllConnectedNodeIDs() const;
 
  private:
-  void Connect(RaySyncerBidiReactor *connection);
+  void Connect(std::shared_ptr<RaySyncerBidiReactor> connection);
 
   std::shared_ptr<bool> stopped_;
 
@@ -169,14 +175,20 @@ class RaySyncer {
   const std::string local_node_id_;
 
   /// Manage connections. Here the key is the NodeID in binary form.
-  absl::flat_hash_map<std::string, RaySyncerBidiReactor *> sync_reactors_;
+  absl::flat_hash_map<std::string, std::shared_ptr<RaySyncerBidiReactor>> sync_reactors_;
 
   /// The local node state
   std::unique_ptr<NodeState> node_state_;
 
   /// Timer is used to do broadcasting.
   std::shared_ptr<PeriodicalRunner> timer_;
 
+  /// The max number of messages to be sent in a batch.
+  const size_t max_batch_size_;
+
+  /// The max delay in milliseconds to wait before sending a batch.
+  const uint64_t max_batch_delay_ms_;
+
   /// Sync message observer, which is a callback on received message response for
   /// [RaySyncerBidiReactor], so should be passed to each of them.
   RpcCompletionCallback on_rpc_completion_;
@@ -204,7 +216,7 @@ class RaySyncerService : public ray::rpc::syncer::RaySyncer::CallbackService {
       std::optional<ray::rpc::AuthenticationToken> auth_token = std::nullopt)
       : syncer_(syncer), auth_token_(std::move(auth_token)) {}
 
-  grpc::ServerBidiReactor<RaySyncMessage, RaySyncMessage> *StartSync(
+  grpc::ServerBidiReactor<RaySyncMessageBatch, RaySyncMessageBatch> *StartSync(
       grpc::CallbackServerContext *context) override;
 
  private:

src/ray/ray_syncer/ray_syncer_bidi_reactor.h

@@ -28,6 +28,7 @@ namespace ray::syncer {
 using ray::rpc::syncer::CommandsSyncMessage;
 using ray::rpc::syncer::MessageType;
 using ray::rpc::syncer::RaySyncMessage;
+using ray::rpc::syncer::RaySyncMessageBatch;
 using ray::rpc::syncer::ResourceViewSyncMessage;
 
 /// This is the base class for the bidi-streaming call and defined the method
@@ -106,13 +107,23 @@ class RaySyncerBidiReactor {
   /// Return true if it's disconnected.
   std::shared_ptr<bool> IsDisconnected() const { return disconnected_; }
 
+  // Set self-owned shared_ptr used to guard lifetime for async callbacks via weak_ptr
+  // lock.
+  void SetSelfRef(std::shared_ptr<RaySyncerBidiReactor> self) {
+    self_ref_ = std::move(self);
+  }
+
   // Node id which is communicating with the current reactor.
   std::string remote_node_id_;
 
  protected:
   /// Sync message observer, which is a callback on received message response.
   RpcCompletionCallback on_rpc_completion_;
 
+  /// Self-owned shared_ptr for lifetime guarding in async callbacks (timer/dispatch). Set
+  /// by caller/factory, released in OnDone.
+  std::shared_ptr<RaySyncerBidiReactor> self_ref_;
+
  private:
   virtual void DoDisconnect() = 0;
   std::shared_ptr<bool> disconnected_ = std::make_shared<bool>(false);