wehub-resource-sync
146 lines
5.7 KiB
Python
146 lines
5.7 KiB
Python
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
|