chore: import upstream snapshot with attribution

python/ray/serve/_private/cluster_node_info_cache.py

@@ -0,0 +1,145 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Dict, FrozenSet, List, Optional, Set, Tuple, Union
+
+import ray
+from ray._common.utils import binary_to_hex
+from ray._raylet import GcsClient
+from ray.serve._private.constants import RAY_GCS_RPC_TIMEOUT_S, SERVE_LOGGER_NAME
+
+logger = logging.getLogger(SERVE_LOGGER_NAME)
+
+
+class ClusterNodeInfoCache(ABC):
+    """Provide access to cached node information in the cluster."""
+
+    def __init__(self, gcs_client: GcsClient):
+        self._gcs_client = gcs_client
+        self._cached_alive_nodes = None
+        self._cached_node_labels = dict()
+        self._cached_total_resources_per_node = dict()
+        self._cached_available_resources_per_node = dict()
+        # Track alive node IDs to detect cluster membership changes and skip
+        # rebuilding labels / total resources when nothing changed.
+        self._alive_node_id_set: FrozenSet[str] = frozenset()
+
+    def update(self):
+        """Update the cache by fetching latest node information from GCS.
+
+        This should be called once in each update cycle.
+        Within an update cycle, everyone will see the same
+        cached node info avoiding any potential issues
+        caused by inconsistent node info seen by different components.
+        """
+        nodes = self._gcs_client.get_all_node_info(timeout=RAY_GCS_RPC_TIMEOUT_S)
+        alive_nodes = [
+            (node_id.hex(), node.node_name, node.instance_id)
+            for (node_id, node) in nodes.items()
+            if node.state == ray.core.generated.gcs_pb2.GcsNodeInfo.ALIVE
+        ]
+
+        # Sort on NodeID to ensure the ordering is deterministic across the cluster.
+        alive_nodes.sort()
+        self._cached_alive_nodes = alive_nodes
+
+        # Detect whether the set of alive nodes has changed. Rebuild labels
+        # and total resources only when it has, since they are static per-node
+        # properties that don't change while a node stays alive.
+        current_alive_ids = frozenset(node_id for node_id, _, _ in alive_nodes)
+        if current_alive_ids != self._alive_node_id_set:
+            self._alive_node_id_set = current_alive_ids
+            self._cached_node_labels = {
+                node_id.hex(): dict(node.labels)
+                for (node_id, node) in nodes.items()
+                if node_id.hex() in current_alive_ids
+            }
+            self._cached_total_resources_per_node = {
+                node_id.hex(): dict(node.resources_total)
+                for (node_id, node) in nodes.items()
+                if node_id.hex() in current_alive_ids
+            }
+
+        # Fetch available resources using the existing GCS client rather than
+        # the legacy GlobalStateAccessor path (which opens a second connection
+        # and performs redundant protobuf deserialization).
+        self._cached_available_resources_per_node = (
+            self._fetch_available_resources_per_node()
+        )
+
+    def _fetch_available_resources_per_node(self) -> Dict[str, Dict[str, float]]:
+        """Fetch available resources per alive node via get_all_resource_usage()."""
+        try:
+            reply = self._gcs_client.get_all_resource_usage(
+                timeout=RAY_GCS_RPC_TIMEOUT_S
+            )
+        except Exception:
+            logger.warning(
+                "Failed to fetch resource usage from GCS. "
+                "Available resources cache will be stale.",
+                exc_info=True,
+            )
+            return self._cached_available_resources_per_node
+
+        return {
+            node_id: dict(resource_data.resources_available)
+            for resource_data in reply.resource_usage_data.batch
+            if (node_id := binary_to_hex(resource_data.node_id))
+            in self._alive_node_id_set
+        }
+
+    def get_alive_nodes(self) -> List[Tuple[str, str, str]]:
+        """Get IDs, IPs, and Instance IDs for all live nodes in the cluster.
+
+        Returns a list of (node_id: str, node_ip: str, instance_id: str).
+        The node_id can be passed into the Ray SchedulingPolicy API.
+        """
+        return self._cached_alive_nodes
+
+    def get_total_resources_per_node(self) -> Dict[str, Dict]:
+        """Get total resources for alive nodes."""
+        return self._cached_total_resources_per_node
+
+    def get_alive_node_ids(self) -> Set[str]:
+        """Get IDs of all live nodes in the cluster."""
+        return {node_id for node_id, _, _ in self.get_alive_nodes()}
+
+    @abstractmethod
+    def get_draining_nodes(self) -> Dict[str, int]:
+        """Get draining nodes in the cluster and their deadlines."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_node_az(self, node_id: str) -> Optional[str]:
+        """Get availability zone of a node."""
+        raise NotImplementedError
+
+    def get_active_node_ids(self) -> Set[str]:
+        """Get IDs of all active nodes in the cluster.
+
+        A node is active if it's schedulable for new tasks and actors.
+        """
+        return self.get_alive_node_ids() - set(self.get_draining_nodes())
+
+    def get_available_resources_per_node(self) -> Dict[str, Union[float, Dict]]:
+        """Get available resources per node.
+
+        Returns a map from (node_id -> Dict of resources).
+        """
+
+        return self._cached_available_resources_per_node
+
+    def get_node_labels(self, node_id: str) -> Dict[str, str]:
+        """Get the labels for a specific node from the cache."""
+        return self._cached_node_labels.get(node_id, {})
+
+
+class DefaultClusterNodeInfoCache(ClusterNodeInfoCache):
+    def __init__(self, gcs_client: GcsClient):
+        super().__init__(gcs_client)
+
+    def get_draining_nodes(self) -> Dict[str, int]:
+        return dict()
+
+    def get_node_az(self, node_id: str) -> Optional[str]:
+        """Get availability zone of a node."""
+        return None