mirrors/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/__init__.py
View File
python/ray/autoscaler/_private/kuberay/autoscaling_config.py
+555
View File
@@ -0,0 +1,555 @@
import decimal
import json
import logging
import time
from itertools import chain
from typing import Any, Dict, Optional
import requests
from ray._private.label_utils import (
validate_node_label_syntax,
)
from ray.autoscaler._private.constants import (
DISABLE_LAUNCH_CONFIG_CHECK_KEY,
DISABLE_NODE_UPDATERS_KEY,
FOREGROUND_NODE_LAUNCH_KEY,
WORKER_LIVENESS_CHECK_KEY,
)
from ray.autoscaler._private.kuberay import node_provider, utils
from ray.autoscaler._private.util import validate_config
from ray.util.debug import log_once
logger = logging.getLogger(__name__)
AUTOSCALER_OPTIONS_KEY = "autoscalerOptions"
IDLE_SECONDS_KEY = "idleTimeoutSeconds"
UPSCALING_KEY = "upscalingMode"
UPSCALING_VALUE_AGGRESSIVE = "Aggressive"
UPSCALING_VALUE_DEFAULT = "Default"
UPSCALING_VALUE_CONSERVATIVE = "Conservative"
MAX_RAYCLUSTER_FETCH_TRIES = 5
RAYCLUSTER_FETCH_RETRY_S = 5
GKE_TPU_TOPOLOGY_LABEL = "cloud.google.com/gke-tpu-topology"
GKE_TPU_ACCELERATOR_LABEL = "cloud.google.com/gke-tpu-accelerator"
# Logical group name for the KubeRay head group.
# Used as the name of the "head node type" by the autoscaler.
_HEAD_GROUP_NAME = "headgroup"
class AutoscalingConfigProducer:
"""Produces an autoscaling config by reading data from the RayCluster CR.
Used to fetch the autoscaling config at the beginning of each autoscaler iteration.
In the context of Ray deployment on Kubernetes, the autoscaling config is an
internal interface.
The autoscaling config carries the strict subset of RayCluster CR data required by
the autoscaler to make scaling decisions; in particular, the autoscaling config does
not carry pod configuration data.
This class is the only public object in this file.
"""
def __init__(self, ray_cluster_name, ray_cluster_namespace):
self.kubernetes_api_client = node_provider.KubernetesHttpApiClient(
namespace=ray_cluster_namespace
)
self._ray_cr_path = f"rayclusters/{ray_cluster_name}"
def __call__(self):
ray_cr = self._fetch_ray_cr_from_k8s_with_retries()
autoscaling_config = _derive_autoscaling_config_from_ray_cr(ray_cr)
return autoscaling_config
def _fetch_ray_cr_from_k8s_with_retries(self) -> Dict[str, Any]:
"""Fetch the RayCluster CR by querying the K8s API server.
Retry on HTTPError for robustness, in particular to protect autoscaler
initialization.
"""
for i in range(1, MAX_RAYCLUSTER_FETCH_TRIES + 1):
try:
return self.kubernetes_api_client.get(self._ray_cr_path)
except requests.HTTPError as e:
if i < MAX_RAYCLUSTER_FETCH_TRIES:
logger.exception(
"Failed to fetch RayCluster CR from K8s. Retrying."
)
time.sleep(RAYCLUSTER_FETCH_RETRY_S)
else:
raise e from None
# This branch is inaccessible. Raise to satisfy mypy.
raise AssertionError
def _derive_autoscaling_config_from_ray_cr(ray_cr: Dict[str, Any]) -> Dict[str, Any]:
provider_config = _generate_provider_config(ray_cr["metadata"]["namespace"])
available_node_types = _generate_available_node_types_from_ray_cr_spec(
ray_cr["spec"]
)
# The autoscaler expects a global max workers field. We set it to the sum of
# node type max workers.
global_max_workers = sum(
node_type["max_workers"] for node_type in available_node_types.values()
)
# Legacy autoscaling fields carry no information but are required for compatibility.
legacy_autoscaling_fields = _generate_legacy_autoscaling_config_fields()
# Process autoscaler options.
autoscaler_options = ray_cr["spec"].get(AUTOSCALER_OPTIONS_KEY, {})
if IDLE_SECONDS_KEY in autoscaler_options:
idle_timeout_minutes = autoscaler_options[IDLE_SECONDS_KEY] / 60.0
else:
idle_timeout_minutes = 1.0
if autoscaler_options.get(UPSCALING_KEY) == UPSCALING_VALUE_CONSERVATIVE:
upscaling_speed = 1 # Rate-limit upscaling if "Conservative" is set by user.
# This elif is redudant but included for clarity.
elif autoscaler_options.get(UPSCALING_KEY) == UPSCALING_VALUE_DEFAULT:
upscaling_speed = 1000 # i.e. big, no rate-limiting by default
# This elif is redudant but included for clarity.
elif autoscaler_options.get(UPSCALING_KEY) == UPSCALING_VALUE_AGGRESSIVE:
upscaling_speed = 1000
else:
upscaling_speed = 1000
autoscaling_config = {
"provider": provider_config,
"cluster_name": ray_cr["metadata"]["name"],
"head_node_type": _HEAD_GROUP_NAME,
"available_node_types": available_node_types,
"max_workers": global_max_workers,
# Should consider exposing `idleTimeoutMinutes` in the RayCluster CRD,
# under an `autoscaling` field.
"idle_timeout_minutes": idle_timeout_minutes,
# Should consider exposing `upscalingSpeed` in the RayCluster CRD,
# under an `autoscaling` field.
"upscaling_speed": upscaling_speed,
**legacy_autoscaling_fields,
}
# Make sure the config is readable by the autoscaler.
validate_config(autoscaling_config)
return autoscaling_config
def _generate_provider_config(ray_cluster_namespace: str) -> Dict[str, Any]:
"""Generates the `provider` field of the autoscaling config, which carries data
required to instantiate the KubeRay node provider.
"""
return {
"type": "kuberay",
"namespace": ray_cluster_namespace,
DISABLE_NODE_UPDATERS_KEY: True,
DISABLE_LAUNCH_CONFIG_CHECK_KEY: True,
FOREGROUND_NODE_LAUNCH_KEY: True,
WORKER_LIVENESS_CHECK_KEY: False,
}
def _generate_legacy_autoscaling_config_fields() -> Dict[str, Any]:
"""Generates legacy autoscaling config fields required for compatibiliy."""
return {
"file_mounts": {},
"cluster_synced_files": [],
"file_mounts_sync_continuously": False,
"initialization_commands": [],
"setup_commands": [],
"head_setup_commands": [],
"worker_setup_commands": [],
"head_start_ray_commands": [],
"worker_start_ray_commands": [],
"auth": {},
}
def _generate_available_node_types_from_ray_cr_spec(
ray_cr_spec: Dict[str, Any],
) -> Dict[str, Any]:
"""Formats autoscaler "available_node_types" field based on the Ray CR's group
specs.
"""
headGroupSpec = ray_cr_spec["headGroupSpec"]
return {
_HEAD_GROUP_NAME: _node_type_from_group_spec(headGroupSpec, is_head=True),
**{
worker_group_spec["groupName"]: _node_type_from_group_spec(
worker_group_spec, is_head=False
)
for worker_group_spec in ray_cr_spec["workerGroupSpecs"]
},
}
def _node_type_from_group_spec(
group_spec: Dict[str, Any], is_head: bool
) -> Dict[str, Any]:
"""Converts CR group spec to autoscaler node type."""
group_name = _HEAD_GROUP_NAME if is_head else group_spec["groupName"]
if is_head:
# The head node type has no workers because the head is not a worker.
min_workers = max_workers = 0
else:
# `minReplicas` and `maxReplicas` are required fields for each workerGroupSpec.
# numOfHosts specifies the number of workers per replica in KubeRay v1.1+.
min_workers = group_spec["minReplicas"] * group_spec.get("numOfHosts", 1)
max_workers = group_spec["maxReplicas"] * group_spec.get("numOfHosts", 1)
resources = _get_ray_resources_from_group_spec(group_spec, is_head)
labels = _get_labels_from_group_spec(group_spec, group_name)
node_type = {
"min_workers": min_workers,
"max_workers": max_workers,
# `node_config` is a legacy field required for compatibility.
# Pod config data is required by the operator but not by the autoscaler.
"node_config": {},
"resources": resources,
"labels": labels,
}
idle_timeout_s = group_spec.get(IDLE_SECONDS_KEY)
if idle_timeout_s is not None:
node_type["idle_timeout_s"] = float(idle_timeout_s)
return node_type
def _get_ray_resources_from_group_spec(
group_spec: Dict[str, Any], is_head: bool
) -> Dict[str, int]:
"""
Infers Ray resources from group `Resources` field, rayStartCommands, or K8s limits.
The resources extracted are used in autoscaling calculations.
"""
# Set resources from top-level group 'Resources' field if it exists.
group_resources = group_spec.get("resources", {})
ray_start_params = group_spec.get("rayStartParams", {})
# In KubeRay, Ray container is always the first application container of a Ray Pod.
k8s_resources = group_spec["template"]["spec"]["containers"][0].get("resources", {})
group_name = _HEAD_GROUP_NAME if is_head else group_spec["groupName"]
num_cpus = _get_num_cpus(
group_resources, ray_start_params, k8s_resources, group_name
)
num_gpus = _get_num_gpus(
group_resources, ray_start_params, k8s_resources, group_name
)
custom_resource_dict = _get_custom_resources(
group_resources, ray_start_params, group_name
)
num_tpus = _get_num_tpus(group_resources, custom_resource_dict, k8s_resources)
memory = _get_memory(group_resources, ray_start_params, k8s_resources)
# It's not allowed to use object store memory as a resource request, so we don't
# add that to the autoscaler's resources annotations.
resources = {}
assert isinstance(num_cpus, int)
resources["CPU"] = num_cpus
if num_gpus is not None:
resources["GPU"] = num_gpus
if num_tpus is not None:
# Add TPU Ray resource if not already added by ray_start_params,
# but specified in k8s_resource_limits.
if "TPU" not in custom_resource_dict:
resources["TPU"] = num_tpus
"""Add TPU head resource, similar to the GCP node_provider.
Sets the Ray resource TPU-{...}-head to ensure the Ray autoscaler
has sufficient resources to make scaling decisions.
TPU worker groups treat each TPU podslice as a replica, with `NumOfHosts`
specifying the number of workers per slice. Each replica of a TPU worker
group has one TPU head.
For example, a v4-16 worker group with 2 replicas should have the following
resource labels on worker 0 of each replica:
worker 0: resources = {"TPU": 4, "TPU-v4-16-head": 1}
"""
if (
"nodeSelector" in group_spec["template"]["spec"]
and GKE_TPU_TOPOLOGY_LABEL in group_spec["template"]["spec"]["nodeSelector"]
and GKE_TPU_ACCELERATOR_LABEL
in group_spec["template"]["spec"]["nodeSelector"]
):
topology = group_spec["template"]["spec"]["nodeSelector"][
GKE_TPU_TOPOLOGY_LABEL
]
accelerator = group_spec["template"]["spec"]["nodeSelector"][
GKE_TPU_ACCELERATOR_LABEL
]
accelerator_type = utils.tpu_node_selectors_to_type(topology, accelerator)
if accelerator_type:
resources[f"TPU-{accelerator_type}-head"] = 1
else:
logger.error(
f"Pods using TPUs require both `{GKE_TPU_TOPOLOGY_LABEL}` and `{GKE_TPU_ACCELERATOR_LABEL}` node selectors. "
"See https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html#configuring-ray-pods-for-tpu-usage "
"and https://cloud.google.com/kubernetes-engine/docs/how-to/tpus."
)
if memory is not None:
resources["memory"] = memory
resources.update(custom_resource_dict)
return resources
def _get_labels_from_group_spec(
group_spec: Dict[str, Any], group_name: str = ""
) -> Dict[str, str]:
"""
Parses Ray node labels for the autoscaling config based on the following
priority:
1. Top-level `labels` field in the group spec.
2. `labels` field in `rayStartParams`.
Args:
group_spec: The group specification dictionary.
group_name: The name of the group (used in warning messages).
Returns:
A dictionary of labels for the node type.
"""
labels_dict = {}
ray_start_params = group_spec.get("rayStartParams", {})
labels_str = ray_start_params.get("labels")
# Use a unique log_once key per group to ensure each group's warning is shown.
log_once_key = (
f"raystartparams_labels_warning_{group_name}"
if group_name
else "raystartparams_labels_warning"
)
if labels_str and log_once(log_once_key):
if group_name:
logger.warning(
f"Ignoring labels: {labels_str} set in rayStartParams for group "
f"'{group_name}'. Group labels are supported in the top-level "
"Labels field starting in KubeRay v1.5"
)
else:
logger.warning(
f"Ignoring labels: {labels_str} set in rayStartParams. "
"Group labels are supported in the top-level Labels field "
"starting in KubeRay v1.5"
)
# Check for top-level structured Labels field.
if "labels" in group_spec and isinstance(group_spec.get("labels"), dict):
labels_dict = group_spec.get("labels")
# Validate node labels follow expected Kubernetes label syntax.
validate_node_label_syntax(labels_dict)
return labels_dict
def _get_num_cpus(
group_resources: Dict[str, str],
ray_start_params: Dict[str, str],
k8s_resources: Dict[str, Dict[str, str]],
group_name: str,
) -> int:
"""Get CPU annotation from `resources` field, ray_start_params or k8s_resources,
with priority for `resources` field.
"""
if "CPU" in group_resources:
if "num-cpus" in ray_start_params:
logger.warning(
f"'CPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
f"Using the value from 'resources': {group_resources['CPU']}."
)
return _round_up_k8s_quantity(group_resources["CPU"])
if "num-cpus" in ray_start_params:
return int(ray_start_params["num-cpus"])
elif "cpu" in k8s_resources.get("limits", {}):
cpu_quantity: str = k8s_resources["limits"]["cpu"]
return _round_up_k8s_quantity(cpu_quantity)
elif "cpu" in k8s_resources.get("requests", {}):
cpu_quantity: str = k8s_resources["requests"]["cpu"]
return _round_up_k8s_quantity(cpu_quantity)
else:
# Getting the number of CPUs is important, so raise an error if we can't do it.
raise ValueError(
f"Autoscaler failed to detect `CPU` resources for group {group_name}."
"\nSet the `--num-cpus` rayStartParam and/or "
"the CPU resource limit for the Ray container."
)
def _get_memory(
group_resources: Dict[str, str],
ray_start_params: Dict[str, str],
k8s_resources: Dict[str, Dict[str, str]],
) -> Optional[int]:
"""Get memory resource annotation from `resources` field, ray_start_params or k8s_resources,
with priority for `resources` field.
"""
if "memory" in group_resources:
if "memory" in ray_start_params:
logger.warning(
f"'memory' specified in both the top-level 'resources' field and in 'rayStartParams'. "
f"Using the value from 'resources': {group_resources['memory']}."
)
return _round_up_k8s_quantity(group_resources["memory"])
if "memory" in ray_start_params:
return int(ray_start_params["memory"])
elif "memory" in k8s_resources.get("limits", {}):
memory_quantity: str = k8s_resources["limits"]["memory"]
return _round_up_k8s_quantity(memory_quantity)
elif "memory" in k8s_resources.get("requests", {}):
memory_quantity: str = k8s_resources["requests"]["memory"]
return _round_up_k8s_quantity(memory_quantity)
return None
def _get_num_gpus(
group_resources: Dict[str, str],
ray_start_params: Dict[str, str],
k8s_resources: Dict[str, Dict[str, str]],
group_name: str,
) -> Optional[int]:
"""Get GPU resource annotation from `resources` field, ray_start_params or k8s_resources,
with priority for `resources` field.
"""
if "GPU" in group_resources:
if "num-gpus" in ray_start_params:
logger.warning(
f"'GPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
f"Using the value from 'resources': {group_resources['GPU']}."
)
return _round_up_k8s_quantity(group_resources["GPU"])
elif "num-gpus" in ray_start_params:
return int(ray_start_params["num-gpus"])
else:
for key, resource_quantity in chain(
k8s_resources.get("limits", {}).items(),
k8s_resources.get("requests", {}).items(),
):
# e.g. nvidia.com/gpu
if key.endswith("gpu"):
# Typically, this is a string representing an interger, e.g. "1".
# Convert to int, making no assumptions on the resource_quantity,
# besides that it's valid as a K8s resource quantity.
num_gpus = _round_up_k8s_quantity(resource_quantity)
if num_gpus > 0:
# Only one GPU type supported for now, break out on first
# "/gpu" match.
return num_gpus
return None
def _get_num_tpus(
group_resources: Dict[str, str],
custom_resource_dict: Dict[str, int],
k8s_resources: Dict[str, Dict[str, str]],
) -> Optional[int]:
"""Get TPU custom resource annotation from `resources` field, custom_resource_dict in ray_start_params,
or k8s_resources, with priority for `resources` field.
"""
if "TPU" in group_resources:
return _round_up_k8s_quantity(group_resources["TPU"])
elif "TPU" in custom_resource_dict:
return custom_resource_dict["TPU"]
else:
for typ in ["limits", "requests"]:
tpu_resource_quantity = k8s_resources.get(typ, {}).get("google.com/tpu")
if tpu_resource_quantity is not None:
# Typically, this is a string representing an integer, e.g. "1".
# Convert to int, making no assumptions on the tpu_resource_quantity,
# besides that it's valid as a K8s resource quantity.
num_tpus = _round_up_k8s_quantity(tpu_resource_quantity)
if num_tpus > 0:
return num_tpus
return None
def _round_up_k8s_quantity(quantity: str) -> int:
"""Rounds a Kubernetes resource quantity up to the nearest integer.
Args:
quantity: Resource quantity as a string in the canonical K8s form.
Returns:
The quantity, rounded up, as an integer.
"""
resource_decimal: decimal.Decimal = utils.parse_quantity(quantity)
rounded = resource_decimal.to_integral_value(rounding=decimal.ROUND_UP)
return int(rounded)
def _get_custom_resources(
group_resources: Dict[str, str], ray_start_params: Dict[str, Any], group_name: str
) -> Dict[str, int]:
"""Format custom resources based on the group `resources` field or `resources` Ray start param.
Currently, the value of the rayStartParam `resources` field must
be formatted as follows:
'"{\"Custom1\": 1, \"Custom2\": 5}"'.
This method first converts the input to a correctly formatted
json string and then loads that json string to a dict.
"""
# If the top-level `resources` field is defined, use it as the exclusive source.
if group_resources:
if "resources" in ray_start_params:
logger.warning(
f"custom resources specified in both the top-level 'resources' field and in 'rayStartParams'. "
f"Using the values from 'resources': {group_resources}."
)
standard_keys = {"CPU", "GPU", "TPU", "memory"}
try:
custom_resources = {
k: _round_up_k8s_quantity(v)
for k, v in group_resources.items()
if k not in standard_keys
}
except Exception as e:
logger.error(
f"Error reading `resource` for group {group_name}."
" For the correct format, refer to example configuration at "
"https://github.com/ray-project/ray/blob/master/python/"
"ray/autoscaler/kuberay/ray-cluster.complete.yaml."
)
raise e
return custom_resources
# Otherwise, check rayStartParams.
if "resources" not in ray_start_params:
return {}
resources_string = ray_start_params["resources"]
try:
# Drop the extra pair of quotes and remove the backslash escapes.
# resources_json should be a json string.
resources_json = resources_string[1:-1].replace("\\", "")
# Load a dict from the json string.
resources = json.loads(resources_json)
assert isinstance(resources, dict)
for key, value in resources.items():
assert isinstance(key, str)
assert isinstance(value, int)
except Exception as e:
logger.error(
f"Error reading `resource` rayStartParam for group {group_name}."
" For the correct format, refer to example configuration at "
"https://github.com/ray-project/ray/blob/master/python/"
"ray/autoscaler/kuberay/ray-cluster.complete.yaml."
)
raise e
return resources
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)
python/ray/autoscaler/_private/kuberay/run_autoscaler.py
+143
View File
@@ -0,0 +1,143 @@
import logging
import os
import subprocess
import time
import ray
from ray._common.network_utils import build_address
from ray._common.ray_constants import (
LOGGING_ROTATE_BACKUP_COUNT,
LOGGING_ROTATE_BYTES,
)
from ray._common.utils import env_integer, try_to_create_directory
from ray._private import ray_constants
from ray._private.ray_logging import setup_component_logger
from ray._private.services import get_node_ip_address
from ray._private.utils import get_all_node_info_until_retrieved
from ray._raylet import GcsClient
from ray.autoscaler._private.kuberay.autoscaling_config import AutoscalingConfigProducer
from ray.autoscaler._private.monitor import Monitor
from ray.autoscaler.v2.instance_manager.config import KubeRayConfigReader
from ray.autoscaler.v2.utils import is_autoscaler_v2
from ray.core.generated.gcs_service_pb2 import GetAllNodeInfoRequest
logger = logging.getLogger(__name__)
BACKOFF_S = 5
def _get_log_dir(gcs_client: GcsClient) -> str:
head_node_selector = GetAllNodeInfoRequest.NodeSelector()
head_node_selector.is_head_node = True
# We need to wait until head node's raylet is registered in GCS.
node_infos = get_all_node_info_until_retrieved(
gcs_client,
node_selectors=[head_node_selector],
)
node_info = next(iter(node_infos))
temp_dir = getattr(node_info, "temp_dir", None)
if temp_dir is None:
raise Exception(
"Node temp_dir was not found in NodeInfo. did the head node's raylet start successfully?"
)
return os.path.join(temp_dir, ray._private.ray_constants.SESSION_LATEST, "logs")
def run_kuberay_autoscaler(cluster_name: str, cluster_namespace: str):
"""Wait until the Ray head container is ready. Then start the autoscaler."""
head_ip = get_node_ip_address()
ray_address = build_address(head_ip, 6379)
while True:
try:
# Autoscaler Ray version might not exactly match GCS version, so skip the
# version check when checking GCS status.
subprocess.check_call(
[
"ray",
"health-check",
"--address",
ray_address,
"--skip-version-check",
]
)
logger.info("The Ray head is ready. Starting the autoscaler.")
break
except subprocess.CalledProcessError:
logger.warning(
f"The Ray head is not ready. Will check again in {BACKOFF_S} seconds."
)
time.sleep(BACKOFF_S)
gcs_client = GcsClient(ray_address)
log_dir = _get_log_dir(gcs_client)
# The Ray head container sets up the log directory. Thus, we set up logging
# only after the Ray head is ready.
_setup_logging(log_dir)
# autoscaling_config_producer reads the RayCluster CR from K8s and uses the CR
# to output an autoscaling config.
autoscaling_config_producer = AutoscalingConfigProducer(
cluster_name, cluster_namespace
)
if is_autoscaler_v2(fetch_from_server=True, gcs_client=gcs_client):
from ray.autoscaler.v2.monitor import AutoscalerMonitor as MonitorV2
MonitorV2(
address=gcs_client.address,
config_reader=KubeRayConfigReader(autoscaling_config_producer),
log_dir=log_dir,
monitor_ip=head_ip,
).run()
else:
Monitor(
address=gcs_client.address,
# The `autoscaling_config` arg can be a dict or a `Callable: () -> dict`.
# In this case, it's a callable.
autoscaling_config=autoscaling_config_producer,
monitor_ip=head_ip,
# Let the autoscaler process exit after it hits 5 exceptions.
# (See ray.autoscaler._private.constants.AUTOSCALER_MAX_NUM_FAILURES.)
# Kubernetes will then restart the autoscaler container.
retry_on_failure=False,
).run()
def _setup_logging(log_dir: str) -> None:
"""Log to autoscaler log file
(typically, /tmp/ray/session_latest/logs/monitor.*)
Also log to pod stdout (logs viewable with `kubectl logs <head-pod> -c autoscaler`).
Args:
log_dir: The path to the log directory.
"""
# The director should already exist, but try (safely) to create it just in case.
try_to_create_directory(log_dir)
# Write logs at info level to monitor.log.
max_bytes = env_integer("RAY_ROTATION_MAX_BYTES", LOGGING_ROTATE_BYTES)
backup_count = env_integer("RAY_ROTATION_BACKUP_COUNT", LOGGING_ROTATE_BACKUP_COUNT)
setup_component_logger(
logging_level=ray_constants.LOGGER_LEVEL,
logging_format=ray_constants.LOGGER_FORMAT,
log_dir=log_dir,
filename=ray_constants.MONITOR_LOG_FILE_NAME, # monitor.log
max_bytes=max_bytes,
backup_count=backup_count,
)
# For the autoscaler, the root logger _also_ needs to write to stderr, not just
# ray_constants.MONITOR_LOG_FILE_NAME.
level = logging.getLevelName(ray_constants.LOGGER_LEVEL.upper())
stderr_handler = logging._StderrHandler()
stderr_handler.setFormatter(logging.Formatter(ray_constants.LOGGER_FORMAT))
stderr_handler.setLevel(level)
logging.root.setLevel(level)
logging.root.addHandler(stderr_handler)
# The stdout handler was set up in the Ray CLI entry point.
# See ray.scripts.scripts::cli().
python/ray/autoscaler/_private/kuberay/utils.py
+112
View File
@@ -0,0 +1,112 @@
# Source:
# https://github.com/kubernetes-client/python/blob/master/kubernetes/utils/quantity.py
from decimal import Decimal, InvalidOperation
from typing import Optional, Union
from ray._private.accelerators.tpu import (
get_num_chips_from_topology,
get_tpu_cores_per_chip,
)
# Mapping used to get generation for TPU-{accelerator}-head resource
# https://cloud.google.com/kubernetes-engine/docs/how-to/tpus#run
gke_tpu_accelerator_to_generation = {
"tpu-v4-podslice": "v4",
"tpu-v5-lite-device": "v5litepod",
"tpu-v5-lite-podslice": "v5litepod",
"tpu-v5p-slice": "v5p",
"tpu-v6e-slice": "v6e",
"tpu7x": "v7x",
}
def parse_quantity(quantity: Union[str, int, float, Decimal]) -> Decimal:
"""Parse kubernetes canonical form quantity like 200Mi to a decimal number.
Supported SI suffixes:
base1024: Ki | Mi | Gi | Ti | Pi | Ei
base1000: n | u | m | "" | k | M | G | T | P | E
See
https://github.com/kubernetes/apimachinery/blob/master/pkg/api/resource/quantity.go
Args:
quantity: kubernetes canonical form quantity.
Returns:
The parsed quantity as a decimal number.
Raises:
ValueError: On invalid or unknown input
"""
if isinstance(quantity, (int, float, Decimal)):
return Decimal(quantity)
exponents = {
"n": -3,
"u": -2,
"m": -1,
"K": 1,
"k": 1,
"M": 2,
"G": 3,
"T": 4,
"P": 5,
"E": 6,
}
quantity = str(quantity)
number = quantity
suffix = None
if len(quantity) >= 2 and quantity[-1] == "i":
if quantity[-2] in exponents:
number = quantity[:-2]
suffix = quantity[-2:]
elif len(quantity) >= 1 and quantity[-1] in exponents:
number = quantity[:-1]
suffix = quantity[-1:]
try:
number = Decimal(number)
except InvalidOperation:
raise ValueError("Invalid number format: {}".format(number))
if suffix is None:
return number
if suffix.endswith("i"):
base = 1024
elif len(suffix) == 1:
base = 1000
else:
raise ValueError("{} has unknown suffix".format(quantity))
# handle SI inconsistency
if suffix == "ki":
raise ValueError("{} has unknown suffix".format(quantity))
if suffix[0] not in exponents:
raise ValueError("{} has unknown suffix".format(quantity))
exponent = Decimal(exponents[suffix[0]])
return number * (base**exponent)
def tpu_node_selectors_to_type(topology: str, accelerator: str) -> Optional[str]:
"""Convert Kubernetes gke-tpu nodeSelectors to TPU accelerator_type
for a kuberay TPU worker group.
Args:
topology: value of the cloud.google.com/gke-tpu-topology Kubernetes
nodeSelector, describes the physical topology of the TPU podslice.
accelerator: value of the cloud.google.com/gke-tpu-accelerator nodeSelector,
the name of the TPU accelerator, e.g. tpu-v4-podslice
Returns:
A string, accelerator_type, e.g. "v4-8".
"""
if topology and accelerator:
generation = gke_tpu_accelerator_to_generation.get(accelerator)
if not generation:
return None
num_chips = get_num_chips_from_topology(topology)
num_cores = num_chips * get_tpu_cores_per_chip(generation)
return f"{generation}-{num_cores}"
return None