[core] change GcsClient.get_all_node_info return type to proto type. (ray-project#46057)

Signed-off-by: Ruiyang Wang <rywang014@gmail.com>

Author: rynewang

python/ray/_private/gcs_aio_client.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Dict, List, Optional
 from concurrent.futures import ThreadPoolExecutor
-from ray._raylet import GcsClient
+from ray._raylet import GcsClient, JobID
 from ray.core.generated import (
     gcs_pb2,
 )
@@ -146,7 +146,7 @@ async def internal_kv_keys(
 
     async def get_all_job_info(
         self, timeout: Optional[float] = None
-    ) -> Dict[bytes, gcs_pb2.JobTableData]:
+    ) -> Dict[JobID, gcs_pb2.JobTableData]:
         """
         Return dict key: bytes of job_id; value: JobTableData pb message.
         """

python/ray/_private/usage/usage_lib.py

@@ -556,7 +556,7 @@ def get_total_num_nodes_to_report(gcs_client, timeout=None) -> Optional[int]:
         result = gcs_client.get_all_node_info(timeout=timeout)
         total_num_nodes = 0
         for node_id, node_info in result.items():
-            if node_info["state"] == gcs_pb2.GcsNodeInfo.GcsNodeState.ALIVE:
+            if node_info.state == gcs_pb2.GcsNodeInfo.GcsNodeState.ALIVE:
                 total_num_nodes += 1
         return total_num_nodes
     except Exception as e:

python/ray/_raylet.pyx

@@ -194,7 +194,7 @@ from ray.util.scheduling_strategies import (
 import ray._private.ray_constants as ray_constants
 import ray.cloudpickle as ray_pickle
 from ray.core.generated.common_pb2 import ActorDiedErrorContext
-from ray.core.generated.gcs_pb2 import JobTableData
+from ray.core.generated.gcs_pb2 import JobTableData, GcsNodeInfo
 from ray.core.generated.gcs_service_pb2 import GetAllResourceUsageReply
 from ray._private.async_compat import (
     sync_to_async,
@@ -2880,27 +2880,26 @@ cdef class GcsClient:
                 c_uri, expiration_s, timeout_ms))
 
     @_auto_reconnect
-    def get_all_node_info(self, timeout=None):
+    def get_all_node_info(self, timeout=None) -> Dict[NodeID, GcsNodeInfo]:
         cdef:
             int64_t timeout_ms = round(1000 * timeout) if timeout else -1
-            CGcsNodeInfo node_info
-            c_vector[CGcsNodeInfo] node_infos
+            CGcsNodeInfo c_node_info
+            c_vector[CGcsNodeInfo] c_node_infos
+            c_vector[c_string] serialized_node_infos
         with nogil:
-            check_status(self.inner.get().GetAllNodeInfo(timeout_ms, node_infos))
+            check_status(self.inner.get().GetAllNodeInfo(timeout_ms, c_node_infos))
+            for c_node_info in c_node_infos:
+                serialized_node_infos.push_back(c_node_info.SerializeAsString())
 
         result = {}
-        for node_info in node_infos:
-            c_resources = PythonGetResourcesTotal(node_info)
-            result[node_info.node_id()] = {
-                "node_name": node_info.node_name(),
-                "state": node_info.state(),
-                "labels": PythonGetNodeLabels(node_info),
-                "resources": {key.decode(): value for key, value in c_resources}
-            }
+        for serialized in serialized_node_infos:
+            node_info = GcsNodeInfo()
+            node_info.ParseFromString(serialized)
+            result[NodeID.from_binary(node_info.node_id)] = node_info
         return result
 
     @_auto_reconnect
-    def get_all_job_info(self, timeout=None) -> Dict[bytes, JobTableData]:
+    def get_all_job_info(self, timeout=None) -> Dict[JobID, JobTableData]:
         # Ideally we should use json_format.MessageToDict(job_info),
         # but `job_info` is a cpp pb message not a python one.
         # Manually converting each and every protobuf field is out of question,
@@ -2918,7 +2917,7 @@ cdef class GcsClient:
         for serialized in serialized_job_infos:
             job_info = JobTableData()
             job_info.ParseFromString(serialized)
-            result[job_info.job_id] = job_info
+            result[JobID.from_binary(job_info.job_id)] = job_info
         return result
 
     @_auto_reconnect

python/ray/includes/common.pxd

@@ -506,7 +506,7 @@ cdef extern from "src/ray/protobuf/gcs.pb.h" nogil:
 
     cdef cppclass CJobConfig "ray::rpc::JobConfig":
         c_string ray_namespace() const
-        const c_string &SerializeAsString()
+        const c_string &SerializeAsString() const
 
     cdef cppclass CNodeDeathInfo "ray::rpc::NodeDeathInfo":
         int reason() const
@@ -526,6 +526,7 @@ cdef extern from "src/ray/protobuf/gcs.pb.h" nogil:
         int runtime_env_agent_port() const
         CNodeDeathInfo death_info() const
         void ParseFromString(const c_string &serialized)
+        const c_string& SerializeAsString() const
 
     cdef enum CGcsNodeState "ray::rpc::GcsNodeInfo_GcsNodeState":
         ALIVE "ray::rpc::GcsNodeInfo_GcsNodeState_ALIVE",
@@ -534,7 +535,7 @@ cdef extern from "src/ray/protobuf/gcs.pb.h" nogil:
         c_string job_id() const
         c_bool is_dead() const
         CJobConfig config() const
-        const c_string &SerializeAsString()
+        const c_string &SerializeAsString() const
 
     cdef cppclass CPythonFunction "ray::rpc::PythonFunction":
         void set_key(const c_string &key)
@@ -570,7 +571,7 @@ cdef extern from "src/ray/protobuf/gcs.pb.h" nogil:
     cdef cppclass CActorTableData "ray::rpc::ActorTableData":
         CAddress address() const
         void ParseFromString(const c_string &serialized)
-        const c_string &SerializeAsString()
+        const c_string &SerializeAsString() const
 
 cdef extern from "ray/common/task/task_spec.h" nogil:
     cdef cppclass CConcurrencyGroup "ray::ConcurrencyGroup":

python/ray/includes/unique_ids.pxi

@@ -230,6 +230,12 @@ cdef class JobID(BaseID):
         check_id(id, CJobID.Size())
         self.data = CJobID.FromBinary(<c_string>id)
 
+    @classmethod
+    def from_binary(cls, id_bytes):
+        if not isinstance(id_bytes, bytes):
+            raise TypeError("Expect bytes, got " + str(type(id_bytes)))
+        return cls(id_bytes)
+
     cdef CJobID native(self):
         return <CJobID>self.data
 

python/ray/serve/_private/cluster_node_info_cache.py

@@ -26,25 +26,21 @@ def update(self):
         """
         nodes = self._gcs_client.get_all_node_info(timeout=RAY_GCS_RPC_TIMEOUT_S)
         alive_nodes = [
-            (ray.NodeID.from_binary(node_id).hex(), node["node_name"].decode("utf-8"))
+            (node_id.hex(), node.node_name)
             for (node_id, node) in nodes.items()
-            if node["state"] == ray.core.generated.gcs_pb2.GcsNodeInfo.ALIVE
+            if node.state == ray.core.generated.gcs_pb2.GcsNodeInfo.ALIVE
         ]
 
         # Sort on NodeID to ensure the ordering is deterministic across the cluster.
         sorted(alive_nodes)
         self._cached_alive_nodes = alive_nodes
         self._cached_node_labels = {
-            ray.NodeID.from_binary(node_id).hex(): {
-                label_name.decode("utf-8"): label_value.decode("utf-8")
-                for label_name, label_value in node["labels"].items()
-            }
-            for (node_id, node) in nodes.items()
+            node_id.hex(): dict(node.labels) for (node_id, node) in nodes.items()
         }
 
         # Node resources
         self._cached_total_resources_per_node = {
-            ray.NodeID.from_binary(node_id).hex(): node["resources"]
+            node_id.hex(): dict(node.resources_total)
             for (node_id, node) in nodes.items()
         }
 

python/ray/tests/test_state_api.py

@@ -3584,7 +3584,7 @@ def f(signal):
     ray.get(signal.wait.remote())
 
     client = ray.worker.global_worker.gcs_client
-    job_id = ray.worker.global_worker.current_job_id.binary()
+    job_id = ray.worker.global_worker.current_job_id
     all_job_info = client.get_all_job_info()
     assert len(all_job_info) == 1
     assert job_id in all_job_info