wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/autoscaler/_private/kuberay/node_provider.py
+558
View File
@@ -0,0 +1,558 @@
import datetime
import json
import logging
import os
from abc import ABC, abstractmethod
from collections import defaultdict
from typing import Any, Dict, List, Optional, Tuple, Union
import requests
from ray._common.network_utils import build_address
from ray.autoscaler._private.constants import WORKER_LIVENESS_CHECK_KEY
from ray.autoscaler._private.util import NodeID, NodeIP, NodeKind, NodeStatus, NodeType
from ray.autoscaler.batching_node_provider import (
BatchingNodeProvider,
NodeData,
ScaleRequest,
)
from ray.autoscaler.tags import (
NODE_KIND_HEAD,
NODE_KIND_WORKER,
STATUS_UP_TO_DATE,
STATUS_UPDATE_FAILED,
TAG_RAY_USER_NODE_TYPE,
)
# Key for KubeRay label that identifies a Ray pod as head or worker.
KUBERAY_LABEL_KEY_KIND = "ray.io/node-type"
# Key for KubeRay label that identifies the worker group (autoscaler node type) of a
# Ray pod.
KUBERAY_LABEL_KEY_TYPE = "ray.io/group"
# These should be synced with:
# https://github.com/ray-project/kuberay/blob/f2d94ffe213dd8f69481b09c474047cb899fa73b/ray-operator/apis/ray/v1/raycluster_types.go#L165-L171 # noqa
# Kind label value indicating the pod is the head.
KUBERAY_KIND_HEAD = "head"
# Kind label value indicating the pod is the worker.
KUBERAY_KIND_WORKER = "worker"
# KubeRay CRD version
KUBERAY_CRD_VER = os.getenv("KUBERAY_CRD_VER", "v1alpha1")
KUBERAY_REQUEST_TIMEOUT_S = int(os.getenv("KUBERAY_REQUEST_TIMEOUT_S", 60))
RAY_HEAD_POD_NAME = os.getenv("RAY_HEAD_POD_NAME")
# https://kubernetes.io/docs/tasks/run-application/access-api-from-pod
# While running in a Pod, your container can create an HTTPS URL for the
# Kubernetes API server by fetching the KUBERNETES_SERVICE_HOST and
# KUBERNETES_SERVICE_PORT_HTTPS environment variables.
KUBERNETES_SERVICE_HOST = os.getenv(
"KUBERNETES_SERVICE_HOST", "https://kubernetes.default"
)
KUBERNETES_SERVICE_PORT = os.getenv("KUBERNETES_SERVICE_PORT_HTTPS", "443")
KUBERNETES_HOST = build_address(KUBERNETES_SERVICE_HOST, KUBERNETES_SERVICE_PORT)
# Key for GKE label that identifies which multi-host replica a pod belongs to
REPLICA_INDEX_KEY = "replicaIndex"
TOKEN_REFRESH_PERIOD = datetime.timedelta(minutes=1)
# Design:
# Each modification the autoscaler wants to make is posted to the API server goal state
# (e.g. if the autoscaler wants to scale up, it increases the number of
# replicas of the worker group it wants to scale, if it wants to scale down
# it decreases the number of replicas and adds the exact pods that should be
# terminated to the scaleStrategy).
# KubeRayNodeProvider inherits from BatchingNodeProvider.
# Thus, the autoscaler's create and terminate requests are batched into a single
# Scale Request object which is submitted at the end of autoscaler update.
# KubeRay node provider converts the ScaleRequest into a RayCluster CR patch
# and applies the patch in the submit_scale_request method.
# To reduce potential for race conditions, KubeRayNodeProvider
# aborts the autoscaler update if the operator has not yet processed workersToDelete -
# see KubeRayNodeProvider.safe_to_scale().
# Once it is confirmed that workersToDelete have been cleaned up, KubeRayNodeProvider
# clears the workersToDelete list.
# Note: Log handlers set up in autoscaling monitor entrypoint.
logger = logging.getLogger(__name__)
def node_data_from_pod(pod: Dict[str, Any]) -> NodeData:
"""Converts a Ray pod extracted from K8s into Ray NodeData.
NodeData is processed by BatchingNodeProvider.
"""
kind, type = kind_and_type(pod)
status = status_tag(pod)
ip = pod_ip(pod)
replica_index = _replica_index_label(pod)
return NodeData(
kind=kind, type=type, replica_index=replica_index, status=status, ip=ip
)
def kind_and_type(pod: Dict[str, Any]) -> Tuple[NodeKind, NodeType]:
"""Determine Ray node kind (head or workers) and node type (worker group name)
from a Ray pod's labels.
"""
labels = pod["metadata"]["labels"]
kind = (
NODE_KIND_HEAD
if labels[KUBERAY_LABEL_KEY_KIND] == KUBERAY_KIND_HEAD
else NODE_KIND_WORKER
)
type = labels[KUBERAY_LABEL_KEY_TYPE]
return kind, type
def _replica_index_label(pod: Dict[str, Any]) -> Optional[str]:
"""Returns the replicaIndex label for a Pod in a multi-host TPU worker group.
The replicaIndex label is set by the GKE TPU Ray webhook and is of
the form {$WORKER_GROUP_NAME-$REPLICA_INDEX} where $REPLICA_INDEX
is an integer from 0 to Replicas-1.
"""
labels = pod["metadata"]["labels"]
return labels.get(REPLICA_INDEX_KEY, None)
def pod_ip(pod: Dict[str, Any]) -> NodeIP:
return pod["status"].get("podIP", "IP not yet assigned")
def status_tag(pod: Dict[str, Any]) -> NodeStatus:
"""Convert pod state to Ray autoscaler node status.
See the doc string of the class
batching_node_provider.NodeData for the semantics of node status.
"""
if (
"containerStatuses" not in pod["status"]
or not pod["status"]["containerStatuses"]
):
return "pending"
state = pod["status"]["containerStatuses"][0]["state"]
if "pending" in state:
return "pending"
if "running" in state:
return STATUS_UP_TO_DATE
if "waiting" in state:
return "waiting"
if "terminated" in state:
return STATUS_UPDATE_FAILED
raise ValueError("Unexpected container state.")
def worker_delete_patch(group_index: str, workers_to_delete: List[NodeID]):
path = f"/spec/workerGroupSpecs/{group_index}/scaleStrategy"
value = {"workersToDelete": workers_to_delete}
return replace_patch(path, value)
def worker_replica_patch(group_index: str, target_replicas: int):
path = f"/spec/workerGroupSpecs/{group_index}/replicas"
value = target_replicas
return replace_patch(path, value)
def replace_patch(path: str, value: Any) -> Dict[str, Any]:
return {"op": "replace", "path": path, "value": value}
def load_k8s_secrets() -> Tuple[Dict[str, str], str]:
"""
Loads secrets needed to access K8s resources.
Returns:
headers: Headers with K8s access token
verify: Path to certificate
"""
with open("/var/run/secrets/kubernetes.io/serviceaccount/token") as secret:
token = secret.read()
headers = {
"Authorization": "Bearer " + token,
}
verify = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
return headers, verify
def url_from_resource(
namespace: str,
path: str,
kuberay_crd_version: str = KUBERAY_CRD_VER,
kubernetes_host: str = KUBERNETES_HOST,
) -> str:
"""Convert resource path to REST URL for Kubernetes API server.
Args:
namespace: The K8s namespace of the resource
path: The part of the resource path that starts with the resource type.
Supported resource types are "pods" and "rayclusters".
kuberay_crd_version: The API version of the KubeRay CRD.
Looks like "v1alpha1", "v1".
kubernetes_host: The host of the Kubernetes API server.
Uses $KUBERNETES_SERVICE_HOST and
$KUBERNETES_SERVICE_PORT to construct the kubernetes_host if not provided.
When set by Kubernetes,
$KUBERNETES_SERVICE_HOST could be an IP address. That's why the https
scheme is added here.
Defaults to "https://kubernetes.default:443".
Returns:
The REST URL for the resource.
"""
if kubernetes_host.startswith("http://"):
raise ValueError("Kubernetes host must be accessed over HTTPS.")
if not kubernetes_host.startswith("https://"):
kubernetes_host = "https://" + kubernetes_host
if path.startswith("pods"):
api_group = "/api/v1"
elif path.startswith("rayclusters"):
api_group = "/apis/ray.io/" + kuberay_crd_version
else:
raise NotImplementedError("Tried to access unknown entity at {}".format(path))
return kubernetes_host + api_group + "/namespaces/" + namespace + "/" + path
def _worker_group_index(raycluster: Dict[str, Any], group_name: str) -> int:
"""Extract worker group index from RayCluster."""
group_names = [
spec["groupName"] for spec in raycluster["spec"].get("workerGroupSpecs", [])
]
return group_names.index(group_name)
def _worker_group_max_replicas(
raycluster: Dict[str, Any], group_index: int
) -> Optional[int]:
"""Extract the maxReplicas of a worker group.
If maxReplicas is unset, return None, to be interpreted as "no constraint".
At time of writing, it should be impossible for maxReplicas to be unset, but it's
better to handle this anyway.
"""
return raycluster["spec"]["workerGroupSpecs"][group_index].get("maxReplicas")
def _worker_group_replicas(raycluster: Dict[str, Any], group_index: int):
# 1 is the default replicas value used by the KubeRay operator
return raycluster["spec"]["workerGroupSpecs"][group_index].get("replicas", 1)
def _worker_group_num_of_hosts(raycluster: Dict[str, Any], group_index: int):
# 1 is the default numOfHosts value used by the KubeRay operator
return raycluster["spec"]["workerGroupSpecs"][group_index].get("numOfHosts", 1)
class IKubernetesHttpApiClient(ABC):
"""
An interface for a Kubernetes HTTP API client.
This interface could be used to mock the Kubernetes API client in tests.
"""
@abstractmethod
def get(self, path: str) -> Dict[str, Any]:
"""Wrapper for REST GET of resource with proper headers."""
pass
@abstractmethod
def patch(
self,
path: str,
payload: Union[List[Dict[str, Any]], Dict[str, Any]],
content_type: str = "application/json-patch+json",
) -> Dict[str, Any]:
"""Wrapper for REST PATCH of resource with proper headers."""
pass
class KubernetesHttpApiClient(IKubernetesHttpApiClient):
def __init__(self, namespace: str, kuberay_crd_version: str = KUBERAY_CRD_VER):
self._kuberay_crd_version = kuberay_crd_version
self._namespace = namespace
self._token_expires_at = datetime.datetime.now() + TOKEN_REFRESH_PERIOD
self._headers, self._verify = None, None
def _get_refreshed_headers_and_verify(self):
if (datetime.datetime.now() >= self._token_expires_at) or (
self._headers is None or self._verify is None
):
logger.info("Refreshing K8s API client token and certs.")
self._headers, self._verify = load_k8s_secrets()
self._token_expires_at = datetime.datetime.now() + TOKEN_REFRESH_PERIOD
return self._headers, self._verify
else:
return self._headers, self._verify
def get(self, path: str) -> Dict[str, Any]:
"""Wrapper for REST GET of resource with proper headers.
Args:
path: The part of the resource path that starts with the resource type.
Returns:
The JSON response of the GET request.
Raises:
HTTPError: If the GET request fails.
"""
url = url_from_resource(
namespace=self._namespace,
path=path,
kuberay_crd_version=self._kuberay_crd_version,
)
headers, verify = self._get_refreshed_headers_and_verify()
result = requests.get(
url,
headers=headers,
timeout=KUBERAY_REQUEST_TIMEOUT_S,
verify=verify,
)
if not result.status_code == 200:
result.raise_for_status()
return result.json()
def patch(
self,
path: str,
payload: Union[List[Dict[str, Any]], Dict[str, Any]],
content_type: str = "application/json-patch+json",
) -> Dict[str, Any]:
"""Wrapper for REST PATCH of resource with proper headers
Args:
path: The part of the resource path that starts with the resource type.
payload: The patch payload, either a JSON Patch list or a
strategic-merge patch object.
content_type: The content type of the merge strategy.
Returns:
The JSON response of the PATCH request.
Raises:
HTTPError: If the PATCH request fails.
"""
url = url_from_resource(
namespace=self._namespace,
path=path,
kuberay_crd_version=self._kuberay_crd_version,
)
headers, verify = self._get_refreshed_headers_and_verify()
result = requests.patch(
url,
json.dumps(payload),
headers={**headers, "Content-type": content_type},
timeout=KUBERAY_REQUEST_TIMEOUT_S,
verify=verify,
)
if not result.status_code == 200:
result.raise_for_status()
return result.json()
class KubeRayNodeProvider(BatchingNodeProvider): # type: ignore
def __init__(
self,
provider_config: Dict[str, Any],
cluster_name: str,
):
logger.info("Creating KubeRayNodeProvider.")
self.namespace = provider_config["namespace"]
self.cluster_name = cluster_name
self.k8s_api_client = KubernetesHttpApiClient(self.namespace)
assert (
provider_config.get(WORKER_LIVENESS_CHECK_KEY, True) is False
), f"To use KubeRayNodeProvider, must set `{WORKER_LIVENESS_CHECK_KEY}:False`."
BatchingNodeProvider.__init__(self, provider_config, cluster_name)
def get_node_data(self) -> Dict[NodeID, NodeData]:
"""Queries K8s for pods in the RayCluster. Converts that pod data into a
map of pod name to Ray NodeData, as required by BatchingNodeProvider.
"""
# Store the raycluster CR
self._raycluster = self._get(f"rayclusters/{self.cluster_name}")
# Get the pods resource version.
# Specifying a resource version in list requests is important for scalability:
# https://kubernetes.io/docs/reference/using-api/api-concepts/#semantics-for-get-and-list
resource_version = self._get_pods_resource_version()
if resource_version:
logger.info(
f"Listing pods for RayCluster {self.cluster_name}"
f" in namespace {self.namespace}"
f" at pods resource version >= {resource_version}."
)
# Filter pods by cluster_name.
label_selector = requests.utils.quote(f"ray.io/cluster={self.cluster_name}")
resource_path = f"pods?labelSelector={label_selector}"
if resource_version:
resource_path += (
f"&resourceVersion={resource_version}"
+ "&resourceVersionMatch=NotOlderThan"
)
pod_list = self._get(resource_path)
fetched_resource_version = pod_list["metadata"]["resourceVersion"]
logger.info(
f"Fetched pod data at resource version" f" {fetched_resource_version}."
)
# Extract node data from the pod list.
node_data_dict = {}
for pod in pod_list["items"]:
# Kubernetes sets metadata.deletionTimestamp immediately after admitting a
# request to delete an object. Full removal of the object may take some time
# after the deletion timestamp is set. See link for details:
# https://kubernetes.io/docs/reference/using-api/api-concepts/#resource-deletion
if "deletionTimestamp" in pod["metadata"]:
# Ignore pods marked for termination.
continue
pod_name = pod["metadata"]["name"]
node_data_dict[pod_name] = node_data_from_pod(pod)
return node_data_dict
def submit_scale_request(self, scale_request: ScaleRequest):
"""Converts the scale request generated by BatchingNodeProvider into
a patch that modifies the RayCluster CR's replicas and/or workersToDelete
fields. Then submits the patch to the K8s API server.
"""
# Transform the scale request into a patch payload.
patch_payload = self._scale_request_to_patch_payload(
scale_request, self._raycluster
)
# Submit the patch to K8s.
logger.info(
"Autoscaler is submitting the following patch to RayCluster "
f"{self.cluster_name} in namespace {self.namespace}."
)
logger.info(patch_payload)
self._submit_raycluster_patch(patch_payload)
def safe_to_scale(self) -> bool:
"""Returns False iff non_terminated_nodes contains any pods in the RayCluster's
workersToDelete lists.
Explanation:
If there are any workersToDelete which are non-terminated,
we should wait for the operator to do its job and delete those
pods. Therefore, we back off the autoscaler update.
If, on the other hand, all of the workersToDelete have already been cleaned up,
then we patch away the workersToDelete lists and return True.
In the future, we may consider having the operator clean up workersToDelete
on it own:
https://github.com/ray-project/kuberay/issues/733
Note (Dmitri):
It is stylistically bad that this function has a side effect.
"""
# Get the list of nodes.
node_set = set(self.node_data_dict.keys())
worker_groups = self._raycluster["spec"].get("workerGroupSpecs", [])
# Accumulates the indices of worker groups with non-empty workersToDelete
non_empty_worker_group_indices = []
for group_index, worker_group in enumerate(worker_groups):
workersToDelete = worker_group.get("scaleStrategy", {}).get(
"workersToDelete", []
)
if workersToDelete:
non_empty_worker_group_indices.append(group_index)
for worker in workersToDelete:
if worker in node_set:
# The operator hasn't removed this worker yet. Abort
# the autoscaler update.
logger.warning(f"Waiting for operator to remove worker {worker}.")
return False
# All required workersToDelete have been removed.
# Clean up the workersToDelete field.
patch_payload = []
for group_index in non_empty_worker_group_indices:
patch = worker_delete_patch(group_index, workers_to_delete=[])
patch_payload.append(patch)
if patch_payload:
logger.info("Cleaning up workers to delete.")
logger.info(f"Submitting patch {patch_payload}.")
self._submit_raycluster_patch(patch_payload)
# It's safe to proceed with the autoscaler update.
return True
def _get_pods_resource_version(self) -> str:
"""
Extract a recent pods resource version by reading the head pod's
metadata.resourceVersion of the response.
"""
if not RAY_HEAD_POD_NAME:
return None
pod_resp = self._get(f"pods/{RAY_HEAD_POD_NAME}")
return pod_resp["metadata"]["resourceVersion"]
def _scale_request_to_patch_payload(
self, scale_request: ScaleRequest, raycluster: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Converts autoscaler scale request into a RayCluster CR patch payload."""
patch_payload = []
# Collect patches for replica counts.
for node_type, target_replicas in scale_request.desired_num_workers.items():
group_index = _worker_group_index(raycluster, node_type)
group_max_replicas = _worker_group_max_replicas(raycluster, group_index)
# Cap the replica count to maxReplicas.
if group_max_replicas is not None and group_max_replicas < target_replicas:
logger.warning(
"Autoscaler attempted to create "
+ "more than maxReplicas pods of type {}.".format(node_type)
)
target_replicas = group_max_replicas
# Check if we need to change the target count.
if target_replicas == _worker_group_replicas(raycluster, group_index):
# No patch required.
continue
# Need to patch replica count. Format the patch and add it to the payload.
patch = worker_replica_patch(group_index, target_replicas)
patch_payload.append(patch)
# Maps node_type to nodes to delete for that group.
deletion_groups = defaultdict(list)
for worker in scale_request.workers_to_delete:
node_type = self.node_tags(worker)[TAG_RAY_USER_NODE_TYPE]
deletion_groups[node_type].append(worker)
for node_type, workers_to_delete in deletion_groups.items():
group_index = _worker_group_index(raycluster, node_type)
patch = worker_delete_patch(group_index, workers_to_delete)
patch_payload.append(patch)
return patch_payload
def _submit_raycluster_patch(self, patch_payload: List[Dict[str, Any]]):
"""Submits a patch to modify a RayCluster CR."""
path = "rayclusters/{}".format(self.cluster_name)
self._patch(path, patch_payload)
def _get(self, path: str) -> Dict[str, Any]:
"""Wrapper for REST GET of resource with proper headers."""
return self.k8s_api_client.get(path)
def _patch(self, path: str, payload: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Wrapper for REST PATCH of resource with proper headers."""
return self.k8s_api_client.patch(path, payload)