[core] Revamp GetAllNodeInfo to be more efficient and take multiple NodeID's (#55115)

Needed for #55112 as
GetAllNodeInfo needs to be able to take multiple node id's as arguments
so you can get information for multiple nodes in one rpc to the GCS.
Revamping the request here to make it make more sense with this model
where you can select multiple nodes or get all nodes if you don't use
any selectors. Also adding an "optimized path" where if all the
selectors are just node id's, we can just do lookups in the map and add
them to the reply instead of iterating through all of it. Also adding
some extra tests for new behavior.

---------

Signed-off-by: dayshah <dhyey2019@gmail.com>
Signed-off-by: sampan <sampan@anyscale.com>

Author: dayshah

python/ray/includes/common.pxd

@@ -455,7 +455,7 @@ cdef extern from "ray/gcs/gcs_client/accessor.h" nogil:
         void AsyncGetAll(
             const MultiItemPyCallback[CGcsNodeInfo] &callback,
             int64_t timeout_ms,
-            optional[CNodeID] node_id)
+            c_vector[CNodeID] node_ids)
 
     cdef cppclass CNodeResourceInfoAccessor "ray::gcs::NodeResourceInfoAccessor":
         CRayStatus GetAllResourceUsage(

python/ray/includes/gcs_client.pxi

@@ -337,18 +337,18 @@ cdef class InnerGcsClient:
     ) -> Future[Dict[NodeID, gcs_pb2.GcsNodeInfo]]:
         cdef:
             int64_t timeout_ms = round(1000 * timeout) if timeout else -1
-            optional[CNodeID] c_node_id
+            c_vector[CNodeID] c_node_ids
             fut = incremented_fut()
         if node_id:
-            c_node_id = (<NodeID>node_id).native()
+            c_node_ids.push_back((<NodeID>node_id).native())
         with nogil:
             self.inner.get().Nodes().AsyncGetAll(
                 MultiItemPyCallback[CGcsNodeInfo](
                     convert_get_all_node_info,
                     assign_and_decrement_fut,
                     fut),
                 timeout_ms,
-                c_node_id)
+                c_node_ids)
         return asyncio.wrap_future(fut)
 
     #############################################################

python/ray/util/state/state_manager.py

@@ -312,26 +312,33 @@ async def get_all_node_info(
         if filters is None:
             filters = []
 
-        req_filters = GetAllNodeInfoRequest.Filters()
+        node_selectors = []
+        state_filter = None
         for filter in filters:
             key, predicate, value = filter
             if predicate != "=":
                 # We only support EQUAL predicate for source side filtering.
                 continue
 
             if key == "node_id":
-                req_filters.node_id = NodeID(hex_to_binary(value)).binary()
+                node_selector = GetAllNodeInfoRequest.NodeSelector()
+                node_selector.node_id = NodeID(hex_to_binary(value)).binary()
+                node_selectors.append(node_selector)
             elif key == "state":
                 value = value.upper()
                 if value not in GcsNodeInfo.GcsNodeState.keys():
                     raise ValueError(f"Invalid node state for filtering: {value}")
-                req_filters.state = GcsNodeInfo.GcsNodeState.Value(value)
+                state_filter = GcsNodeInfo.GcsNodeState.Value(value)
             elif key == "node_name":
-                req_filters.node_name = value
+                node_selector = GetAllNodeInfoRequest.NodeSelector()
+                node_selector.node_name = value
+                node_selectors.append(node_selector)
             else:
                 continue
 
-        request = GetAllNodeInfoRequest(limit=limit, filters=req_filters)
+        request = GetAllNodeInfoRequest(
+            limit=limit, node_selectors=node_selectors, state_filter=state_filter
+        )
         reply = await self._gcs_node_info_stub.GetAllNodeInfo(request, timeout=timeout)
         return reply
 

src/mock/ray/gcs/gcs_client/accessor.h

@@ -140,7 +140,7 @@ class MockNodeInfoAccessor : public NodeInfoAccessor {
               AsyncGetAll,
               (const MultiItemCallback<rpc::GcsNodeInfo> &callback,
                int64_t timeout_ms,
-               std::optional<NodeID> node_id),
+               const std::vector<NodeID> &node_ids),
               (override));
   MOCK_METHOD(void,
               AsyncSubscribeToNodeChange,

src/ray/gcs/gcs_client/accessor.cc

@@ -594,11 +594,11 @@ Status NodeInfoAccessor::DrainNodes(const std::vector<NodeID> &node_ids,
 
 void NodeInfoAccessor::AsyncGetAll(const MultiItemCallback<rpc::GcsNodeInfo> &callback,
                                    int64_t timeout_ms,
-                                   std::optional<NodeID> node_id) {
+                                   const std::vector<NodeID> &node_ids) {
   RAY_LOG(DEBUG) << "Getting information of all nodes.";
   rpc::GetAllNodeInfoRequest request;
-  if (node_id) {
-    request.mutable_filters()->set_node_id(node_id->Binary());
+  for (const auto &node_id : node_ids) {
+    request.add_node_selectors()->set_node_id(node_id.Binary());
   }
   client_impl_->GetGcsRpcClient().GetAllNodeInfo(
       request,
@@ -683,9 +683,12 @@ Status NodeInfoAccessor::GetAllNoCache(int64_t timeout_ms,
 }
 
 StatusOr<std::vector<rpc::GcsNodeInfo>> NodeInfoAccessor::GetAllNoCacheWithFilters(
-    int64_t timeout_ms, rpc::GetAllNodeInfoRequest_Filters filters) {
+    int64_t timeout_ms,
+    rpc::GcsNodeInfo::GcsNodeState state_filter,
+    rpc::GetAllNodeInfoRequest::NodeSelector node_selector) {
   rpc::GetAllNodeInfoRequest request;
-  *request.mutable_filters() = std::move(filters);
+  *request.add_node_selectors() = std::move(node_selector);
+  request.set_state_filter(state_filter);
   rpc::GetAllNodeInfoReply reply;
   RAY_RETURN_NOT_OK(
       client_impl_->GetGcsRpcClient().SyncGetAllNodeInfo(request, &reply, timeout_ms));

src/ray/gcs/gcs_client/accessor.h

@@ -349,10 +349,11 @@ class NodeInfoAccessor {
   ///
   /// \param callback Callback that will be called after lookup finishes.
   /// \param timeout_ms The timeout for this request.
-  /// \param node_id If not nullopt, only return the node info of the specified node.
+  /// \param node_ids If this is not empty, only return the node info of the specified
+  /// nodes.
   virtual void AsyncGetAll(const MultiItemCallback<rpc::GcsNodeInfo> &callback,
                            int64_t timeout_ms,
-                           std::optional<NodeID> node_id = std::nullopt);
+                           const std::vector<NodeID> &node_ids = {});
 
   /// Subscribe to node addition and removal events from GCS and cache those information.
   ///
@@ -393,7 +394,9 @@ class NodeInfoAccessor {
   ///
   /// \return All nodes that match the given filters from the gcs without the cache.
   virtual StatusOr<std::vector<rpc::GcsNodeInfo>> GetAllNoCacheWithFilters(
-      int64_t timeout_ms, rpc::GetAllNodeInfoRequest_Filters filters);
+      int64_t timeout_ms,
+      rpc::GcsNodeInfo::GcsNodeState state_filter,
+      rpc::GetAllNodeInfoRequest::NodeSelector node_selector);
 
   /// Send a check alive request to GCS for the liveness of some nodes.
   ///

src/ray/gcs/gcs_client/global_state_accessor.cc

@@ -408,16 +408,15 @@ ray::Status GlobalStateAccessor::GetNode(const std::string &node_id_hex_str,
 
   std::vector<rpc::GcsNodeInfo> node_infos;
   while (true) {
-    rpc::GetAllNodeInfoRequest_Filters filters;
-    filters.set_state(rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_ALIVE);
-    filters.set_node_id(node_id_binary);
+    rpc::GetAllNodeInfoRequest::NodeSelector selector;
+    selector.set_node_id(node_id_binary);
     {
       absl::ReaderMutexLock lock(&mutex_);
       auto timeout_ms =
           std::max(end_time_point - current_time_ms(), static_cast<int64_t>(0));
-      RAY_ASSIGN_OR_RETURN(
-          node_infos,
-          gcs_client_->Nodes().GetAllNoCacheWithFilters(timeout_ms, std::move(filters)));
+      RAY_ASSIGN_OR_RETURN(node_infos,
+                           gcs_client_->Nodes().GetAllNoCacheWithFilters(
+                               timeout_ms, rpc::GcsNodeInfo::ALIVE, std::move(selector)));
     }
     if (!node_infos.empty()) {
       *node_info = node_infos[0].SerializeAsString();
@@ -442,16 +441,16 @@ ray::Status GlobalStateAccessor::GetNodeToConnectForDriver(
       current_time_ms() + RayConfig::instance().raylet_start_wait_time_s() * 1000;
 
   std::vector<rpc::GcsNodeInfo> node_infos;
-  rpc::GetAllNodeInfoRequest_Filters filters;
-  filters.set_state(rpc::GcsNodeInfo_GcsNodeState::GcsNodeInfo_GcsNodeState_ALIVE);
-  filters.set_node_ip_address(node_ip_address);
+  rpc::GetAllNodeInfoRequest::NodeSelector selector;
+  selector.set_node_ip_address(node_ip_address);
   while (true) {
     {
       absl::ReaderMutexLock lock(&mutex_);
       auto timeout_ms =
           std::max(end_time_point - current_time_ms(), static_cast<int64_t>(0));
-      RAY_ASSIGN_OR_RETURN(
-          node_infos, gcs_client_->Nodes().GetAllNoCacheWithFilters(timeout_ms, filters));
+      RAY_ASSIGN_OR_RETURN(node_infos,
+                           gcs_client_->Nodes().GetAllNoCacheWithFilters(
+                               timeout_ms, rpc::GcsNodeInfo::ALIVE, selector));
     }
     if (!node_infos.empty()) {
       *node_to_connect = node_infos[0].SerializeAsString();
@@ -464,22 +463,23 @@ ray::Status GlobalStateAccessor::GetNodeToConnectForDriver(
       auto [address, _] = gcs_client_->GetGcsServerAddress();
       gcs_address = std::move(address);
     }
-    filters.set_node_ip_address(gcs_address);
+    selector.set_node_ip_address(gcs_address);
     {
       absl::ReaderMutexLock lock(&mutex_);
       auto timeout_ms = end_time_point - current_time_ms();
-      RAY_ASSIGN_OR_RETURN(
-          node_infos, gcs_client_->Nodes().GetAllNoCacheWithFilters(timeout_ms, filters));
+      RAY_ASSIGN_OR_RETURN(node_infos,
+                           gcs_client_->Nodes().GetAllNoCacheWithFilters(
+                               timeout_ms, rpc::GcsNodeInfo::ALIVE, selector));
     }
     if (node_infos.empty() && node_ip_address == gcs_address) {
-      filters.set_node_ip_address("127.0.0.1");
+      selector.set_node_ip_address("127.0.0.1");
       {
         absl::ReaderMutexLock lock(&mutex_);
         auto timeout_ms =
             std::max(end_time_point - current_time_ms(), static_cast<int64_t>(0));
-        RAY_ASSIGN_OR_RETURN(
-            node_infos,
-            gcs_client_->Nodes().GetAllNoCacheWithFilters(timeout_ms, filters));
+        RAY_ASSIGN_OR_RETURN(node_infos,
+                             gcs_client_->Nodes().GetAllNoCacheWithFilters(
+                                 timeout_ms, rpc::GcsNodeInfo::ALIVE, selector));
       }
     }
     if (!node_infos.empty()) {