chore: import upstream snapshot with attribution

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
+76
View File
@@ -0,0 +1,76 @@
load("@rules_python//python:defs.bzl", "py_library")
# --------------------------------------------------------------------
# Tests from the python/ray/autoscaler/v2/tests directory.
# Covers all tests starting with `test_`.
# Please keep these sorted alphabetically.
# --------------------------------------------------------------------
load("//bazel:python.bzl", "py_test_module_list")
py_library(
name = "conftest",
srcs = ["tests/conftest.py"],
deps = ["//python/ray/tests:conftest"],
)
# Unit test. (These should not have ray started)
py_test_module_list(
size = "small",
files = [
"tests/test_allocation_timeout_recovery.py",
"tests/test_config.py",
"tests/test_event_logger.py",
"tests/test_instance_manager.py",
"tests/test_instance_storage.py",
"tests/test_instance_util.py",
"tests/test_ippr_provider.py",
"tests/test_metrics_reporter.py",
"tests/test_node_provider.py",
"tests/test_priority_selection.py",
"tests/test_ray_installer.py",
"tests/test_reconciler.py",
"tests/test_scheduler.py",
"tests/test_schema.py",
"tests/test_storage.py",
"tests/test_subscribers.py",
"tests/test_threaded_ray_installer.py",
"tests/test_utils.py",
],
tags = ["team:core"],
deps = [
":conftest",
"//:ray_lib",
],
)
# Integration test.
py_test_module_list(
size = "medium",
files = [
"tests/test_autoscaler.py",
"tests/test_sdk.py",
],
tags = [
"exclusive",
"team:core",
],
deps = [
":conftest",
"//:ray_lib",
],
)
py_test_module_list(
size = "large",
files = [
"tests/test_e2e.py",
],
tags = [
"exclusive",
"team:core",
],
deps = [
":conftest",
"//:ray_lib",
],
)
+234
View File
@@ -0,0 +1,234 @@
import logging
from queue import Queue
from typing import List, Optional
from urllib.parse import urlsplit
from ray._raylet import GcsClient
from ray.autoscaler._private.providers import _get_node_provider
from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
from ray.autoscaler.v2.instance_manager.cloud_providers.kuberay.cloud_provider import (
KubeRayProvider,
)
from ray.autoscaler.v2.instance_manager.cloud_providers.read_only.cloud_provider import ( # noqa
ReadOnlyProvider,
)
from ray.autoscaler.v2.instance_manager.config import (
AutoscalingConfig,
IConfigReader,
Provider,
)
from ray.autoscaler.v2.instance_manager.instance_manager import (
InstanceManager,
InstanceUpdatedSubscriber,
)
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.autoscaler.v2.instance_manager.node_provider import (
ICloudInstanceProvider,
NodeProviderAdapter,
)
from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
from ray.autoscaler.v2.instance_manager.reconciler import Reconciler
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import (
CloudInstanceUpdater,
)
from ray.autoscaler.v2.instance_manager.subscribers.cloud_resource_monitor import (
CloudResourceMonitor,
)
from ray.autoscaler.v2.instance_manager.subscribers.ray_stopper import RayStopper
from ray.autoscaler.v2.instance_manager.subscribers.threaded_ray_installer import (
ThreadedRayInstaller,
)
from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
from ray.autoscaler.v2.scheduler import ResourceDemandScheduler
from ray.autoscaler.v2.sdk import get_cluster_resource_state
from ray.core.generated.autoscaler_pb2 import AutoscalingState
from ray.exceptions import AuthenticationError
logger = logging.getLogger(__name__)
class Autoscaler:
def __init__(
self,
session_name: str,
config_reader: IConfigReader,
gcs_client: GcsClient,
event_logger: Optional[AutoscalerEventLogger] = None,
metrics_reporter: Optional[AutoscalerMetricsReporter] = None,
) -> None:
"""Initialize the autoscaler.
Args:
session_name: The current Ray session name.
config_reader: The config reader.
gcs_client: The GCS client.
event_logger: The event logger for emitting cluster events.
metrics_reporter: The metrics reporter for emitting cluster metrics.
"""
self._config_reader = config_reader
config = config_reader.get_cached_autoscaling_config()
logger.info(f"Using Autoscaling Config: \n{config.dump()}")
self._gcs_client = gcs_client
self._cloud_instance_provider = None
self._instance_manager = None
self._ray_stop_errors_queue = Queue()
self._ray_install_errors_queue = Queue()
self._event_logger = event_logger
self._metrics_reporter = metrics_reporter
self._init_cloud_instance_provider(config, config_reader)
self._cloud_resource_monitor = None
self._init_instance_manager(
session_name=session_name,
config=config,
cloud_provider=self._cloud_instance_provider,
gcs_client=self._gcs_client,
)
self._scheduler = ResourceDemandScheduler(self._event_logger)
def _init_cloud_instance_provider(
self, config: AutoscalingConfig, config_reader: IConfigReader
):
"""
Initialize the cloud provider, and its dependencies (the v1 node provider)
Args:
config: The autoscaling config.
config_reader: The config reader.
"""
provider_config = config.get_provider_config()
if provider_config["type"] == "kuberay":
provider_config["head_node_type"] = config.get_head_node_type()
self._cloud_instance_provider = KubeRayProvider(
config.get_config("cluster_name"),
provider_config,
gcs_client=self._gcs_client,
)
elif config.provider == Provider.READ_ONLY:
provider_config["gcs_address"] = self._gcs_client.address
self._cloud_instance_provider = ReadOnlyProvider(
provider_config=provider_config,
)
else:
node_provider_v1 = _get_node_provider(
provider_config,
config.get_config("cluster_name"),
)
self._cloud_instance_provider = NodeProviderAdapter(
v1_provider=node_provider_v1,
config_reader=config_reader,
)
def _init_instance_manager(
self,
session_name: str,
cloud_provider: ICloudInstanceProvider,
gcs_client: GcsClient,
config: AutoscalingConfig,
):
"""
Initialize the instance manager, and its dependencies.
"""
instance_storage = InstanceStorage(
cluster_id=session_name,
storage=InMemoryStorage(),
)
subscribers: List[InstanceUpdatedSubscriber] = []
subscribers.append(
CloudInstanceUpdater(
cloud_provider=cloud_provider,
metrics_reporter=self._metrics_reporter,
)
)
subscribers.append(
RayStopper(gcs_client=gcs_client, error_queue=self._ray_stop_errors_queue)
)
if not config.disable_node_updaters() and isinstance(
cloud_provider, NodeProviderAdapter
):
head_node_ip = urlsplit("//" + self._gcs_client.address).hostname
assert head_node_ip is not None, "Invalid GCS address format"
subscribers.append(
ThreadedRayInstaller(
head_node_ip=head_node_ip,
instance_storage=instance_storage,
ray_installer=RayInstaller(
provider=cloud_provider.v1_provider,
config=config,
),
error_queue=self._ray_install_errors_queue,
# TODO(rueian): Rewrite the ThreadedRayInstaller and its underlying
# NodeUpdater and CommandRunner to use the asyncio, so that we don't
# need to use so many threads. We use so many threads now because
# they are blocking and letting the new cloud machines to wait for
# previous machines to finish installing Ray is quite inefficient.
max_concurrent_installs=config.get_max_num_worker_nodes() or 50,
)
)
self._cloud_resource_monitor = CloudResourceMonitor()
subscribers.append(self._cloud_resource_monitor)
self._instance_manager = InstanceManager(
instance_storage=instance_storage,
instance_status_update_subscribers=subscribers,
)
def update_autoscaling_state(
self,
) -> Optional[AutoscalingState]:
"""Update the autoscaling state of the cluster by reconciling the current
state of the cluster resources, the cloud providers as well as instance
update subscribers with the desired state.
Returns:
AutoscalingState: The new autoscaling state of the cluster or None if
the state is not updated.
Raises:
None: No exception.
"""
try:
ray_stop_errors = []
while not self._ray_stop_errors_queue.empty():
ray_stop_errors.append(self._ray_stop_errors_queue.get())
ray_install_errors = []
while not self._ray_install_errors_queue.empty():
ray_install_errors.append(self._ray_install_errors_queue.get())
# Get the current state of the ray cluster resources.
ray_cluster_resource_state = get_cluster_resource_state(self._gcs_client)
# Refresh the config from the source
self._config_reader.refresh_cached_autoscaling_config()
autoscaling_config = self._config_reader.get_cached_autoscaling_config()
return Reconciler.reconcile(
instance_manager=self._instance_manager,
scheduler=self._scheduler,
cloud_provider=self._cloud_instance_provider,
cloud_resource_monitor=self._cloud_resource_monitor,
ray_cluster_resource_state=ray_cluster_resource_state,
non_terminated_cloud_instances=(
self._cloud_instance_provider.get_non_terminated()
),
cloud_provider_errors=self._cloud_instance_provider.poll_errors(),
ray_install_errors=ray_install_errors,
ray_stop_errors=ray_stop_errors,
autoscaling_config=autoscaling_config,
metrics_reporter=self._metrics_reporter,
)
except AuthenticationError as e:
logger.warning(f"AuthenticationError detected, restarting autoscaler: {e}")
raise
except Exception as e:
logger.exception(e)
return None
+186
View File
@@ -0,0 +1,186 @@
import logging
from collections import defaultdict
from typing import Dict, List, Optional
from ray._private.event.event_logger import EventLoggerAdapter
from ray.autoscaler.v2.utils import ResourceRequestUtil
from ray.core.generated.autoscaler_pb2 import (
ClusterResourceConstraint,
GangResourceRequest,
ResourceRequest,
)
from ray.core.generated.common_pb2 import LabelSelectorOperator
from ray.core.generated.instance_manager_pb2 import LaunchRequest, TerminationRequest
logger = logging.getLogger(__name__)
class AutoscalerEventLogger:
"""
Logs events related to the autoscaler.
# TODO:
- Add more logging for other events.
- Rate limit the events if too spammy.
"""
def __init__(self, logger: EventLoggerAdapter, log_cluster_shape: bool = True):
self._logger = logger
self._log_cluster_shape = log_cluster_shape
def log_cluster_scheduling_update(
self,
cluster_resources: Dict[str, float],
launch_requests: Optional[List[LaunchRequest]] = None,
terminate_requests: Optional[List[TerminationRequest]] = None,
infeasible_requests: Optional[List[ResourceRequest]] = None,
infeasible_gang_requests: Optional[List[GangResourceRequest]] = None,
infeasible_cluster_resource_constraints: Optional[
List[ClusterResourceConstraint]
] = None,
) -> None:
"""
Log updates to the autoscaler scheduling state.
Emits:
- info logs for node launches and terminations (counts grouped by node type).
- an info log summarizing the cluster size after a resize (CPUs/GPUs/TPUs).
- warnings describing infeasible single resource requests, infeasible gang
(placement group) requests, and infeasible cluster resource constraints.
Args:
cluster_resources: Mapping of resource name to total resources for the
current cluster state.
launch_requests: Node launch requests issued in this scheduling step.
terminate_requests: Node termination requests issued in this scheduling
step.
infeasible_requests: Resource requests that could not be satisfied by
any available node type.
infeasible_gang_requests: Gang/placement group requests that could not
be scheduled.
infeasible_cluster_resource_constraints: Cluster-level resource
constraints that could not be satisfied.
Returns:
None
"""
# Log any launch events.
if self._log_cluster_shape and launch_requests:
launch_type_count = defaultdict(int)
for req in launch_requests:
launch_type_count[req.instance_type] += req.count
for idx, (instance_type, count) in enumerate(launch_type_count.items()):
log_str = f"Adding {count} node(s) of type {instance_type}."
self._logger.info(f"{log_str}")
logger.info(f"{log_str}")
# Log any terminate events.
if self._log_cluster_shape and terminate_requests:
termination_by_causes_and_type = defaultdict(int)
for req in terminate_requests:
termination_by_causes_and_type[(req.cause, req.instance_type)] += 1
cause_reason_map = {
TerminationRequest.Cause.OUTDATED: "outdated",
TerminationRequest.Cause.MAX_NUM_NODES: "max number of worker nodes reached", # noqa
TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE: "max number of worker nodes per type reached", # noqa
TerminationRequest.Cause.IDLE: "idle",
}
for idx, ((cause, instance_type), count) in enumerate(
termination_by_causes_and_type.items()
):
log_str = f"Removing {count} nodes of type {instance_type} ({cause_reason_map[cause]})." # noqa
self._logger.info(f"{log_str}")
logger.info(f"{log_str}")
# Cluster shape changes.
if self._log_cluster_shape and (launch_requests or terminate_requests):
num_cpus = cluster_resources.get("CPU", 0)
log_str = f"Resized to {int(num_cpus)} CPUs"
if "GPU" in cluster_resources:
log_str += f", {int(cluster_resources['GPU'])} GPUs"
if "TPU" in cluster_resources:
log_str += f", {int(cluster_resources['TPU'])} TPUs"
self._logger.info(f"{log_str}.")
self._logger.debug(f"Current cluster resources: {dict(cluster_resources)}.")
# Log any infeasible requests.
if infeasible_requests:
requests_by_count = ResourceRequestUtil.group_by_count(infeasible_requests)
log_str = "No available node types can fulfill resource requests "
for idx, req_count in enumerate(requests_by_count):
resource_map = ResourceRequestUtil.to_resource_map(req_count.request)
log_str += f"{resource_map}*{req_count.count}"
if idx < len(requests_by_count) - 1:
log_str += ", "
# Parse and log label selectors if present
if req_count.request.label_selectors:
selector_strs = []
for selector in req_count.request.label_selectors:
for constraint in selector.label_constraints:
op = LabelSelectorOperator.Name(constraint.operator)
values = ",".join(constraint.label_values)
selector_strs.append(
f"{constraint.label_key} {op} [{values}]"
)
if selector_strs:
log_str += (
" with label selectors: [" + "; ".join(selector_strs) + "]"
)
log_str += (
". Add suitable node types to this cluster to resolve this issue."
)
self._logger.warning(log_str)
if infeasible_gang_requests:
# Log for each placement group requests.
for gang_request in infeasible_gang_requests:
log_str = (
"No available node types can fulfill "
"placement group requests (detail={details}): ".format(
details=gang_request.details
)
)
requests_by_count = ResourceRequestUtil.group_by_count(
gang_request.requests
)
for idx, req_count in enumerate(requests_by_count):
resource_map = ResourceRequestUtil.to_resource_map(
req_count.request
)
log_str += f"{resource_map}*{req_count.count}"
if idx < len(requests_by_count) - 1:
log_str += ", "
log_str += (
". Add suitable node types to this cluster to resolve this issue."
)
self._logger.warning(log_str)
if infeasible_cluster_resource_constraints:
# We will only have max 1 cluster resource constraint for now since it's
# from `request_resources()` sdk, where the most recent call would override
# the previous one.
for infeasible_constraint in infeasible_cluster_resource_constraints:
log_str = "No available node types can fulfill cluster constraint: "
for i, requests_by_count in enumerate(
infeasible_constraint.resource_requests
):
resource_map = ResourceRequestUtil.to_resource_map(
requests_by_count.request
)
log_str += f"{resource_map}*{requests_by_count.count}"
if i < len(infeasible_constraint.resource_requests) - 1:
log_str += ", "
log_str += (
". Add suitable node types to this cluster to resolve this issue."
)
self._logger.warning(log_str)
@@ -0,0 +1,768 @@
import copy
import logging
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple
import requests
from ray._raylet import RAY_INTERNAL_NAMESPACE_PREFIX, GcsClient
# TODO(rickyx): We should eventually remove these imports
# when we deprecate the v1 kuberay node provider.
from ray.autoscaler._private.kuberay.node_provider import (
KUBERAY_KIND_HEAD,
KUBERAY_KIND_WORKER,
KUBERAY_LABEL_KEY_KIND,
KUBERAY_LABEL_KEY_TYPE,
RAY_HEAD_POD_NAME,
IKubernetesHttpApiClient,
KubernetesHttpApiClient,
_worker_group_index,
_worker_group_max_replicas,
_worker_group_num_of_hosts,
_worker_group_replicas,
worker_delete_patch,
worker_replica_patch,
)
from ray.autoscaler.v2.instance_manager.cloud_providers.kuberay.ippr_provider import (
KubeRayIPPRProvider,
)
from ray.autoscaler.v2.instance_manager.node_provider import (
CloudInstance,
CloudInstanceId,
CloudInstanceProviderError,
ICloudInstanceProvider,
LaunchNodeError,
NodeKind,
TerminateNodeError,
)
from ray.autoscaler.v2.schema import IPPRSpecs, IPPRStatus, NodeType
logger = logging.getLogger(__name__)
# Annotation the KubeRay operator acts on to terminate the cluster.
NO_DRIVER_TTL_EXPIRED_ANNOTATION = "ray.io/no-driver-ttl-expired"
AUTOSCALER_OPTIONS_KEY = "autoscalerOptions"
NO_DRIVER_TIMEOUT_SECONDS_KEY = "noDriverTimeoutSeconds"
class KubeRayProvider(ICloudInstanceProvider):
"""
This class is a thin wrapper around the Kubernetes API client. It modifies
the RayCluster resource spec on the Kubernetes API server to scale the cluster:
It launches new instances/nodes by submitting patches to the Kubernetes API
to update the RayCluster CRD.
"""
def __init__(
self,
cluster_name: str,
provider_config: Dict[str, Any],
gcs_client: GcsClient,
k8s_api_client: Optional[IKubernetesHttpApiClient] = None,
):
"""
Initializes a new KubeRayProvider.
Args:
cluster_name: The name of the RayCluster resource.
provider_config: The configuration for the RayCluster.
gcs_client: The client to the GCS server. Will be used for resizing raylets.
k8s_api_client: The client to the Kubernetes
API server. This can be used to mock the Kubernetes API server for testing.
"""
self._cluster_name = cluster_name
self._namespace = provider_config["namespace"]
self._k8s_api_client = k8s_api_client or KubernetesHttpApiClient(
namespace=self._namespace
)
self._gcs_client = gcs_client
# Below are states that are cached locally.
self._requests = set()
self._launch_errors_queue = []
self._terminate_errors_queue = []
# Below are states for idle-cluster termination tracking.
# Monotonic timestamp when no driver was first observed; None resets it.
self._no_driver_observed_since: Optional[float] = None
# Latest GCS job end time seen; a newer one means a driver came and went.
self._last_seen_job_end_time = 0
# No-driver timeout (seconds) from the CR; None disables the feature.
self._no_driver_timeout_seconds: Optional[float] = None
# Below are states that are fetched from the Kubernetes API server.
self._ray_cluster = None
self._cached_instances: Dict[CloudInstanceId, CloudInstance]
self._ippr_provider = KubeRayIPPRProvider(
gcs_client=gcs_client, k8s_api_client=self._k8s_api_client
)
@dataclass
class ScaleRequest:
"""Represents a scale request that contains the current states and go-to states
for the ray cluster.
This class will be converted to patches to be submitted to the Kubernetes API
server:
- For launching new instances, it will adjust the `replicas` field in the
workerGroupSpecs.
- For terminating instances, it will adjust the `workersToDelete` field in the
workerGroupSpecs.
"""
# The desired number of workers for each node type.
desired_num_workers: Dict[NodeType, int] = field(default_factory=dict)
# The workers to delete for each node type.
workers_to_delete: Dict[NodeType, List[CloudInstanceId]] = field(
default_factory=dict
)
# The worker groups with empty workersToDelete field.
# This is needed since we will also need to clear the workersToDelete field
# for the worker groups that have finished deletes.
worker_groups_without_pending_deletes: Set[NodeType] = field(
default_factory=set
)
# The worker groups that still have workers to be deleted.
worker_groups_with_pending_deletes: Set[NodeType] = field(default_factory=set)
################################
# Interface for ICloudInstanceProvider
################################
def get_non_terminated(self) -> Dict[CloudInstanceId, CloudInstance]:
self._sync_with_api_server()
self._evaluate_no_driver_termination()
return copy.deepcopy(dict(self._cached_instances))
def terminate(self, ids: List[CloudInstanceId], request_id: str) -> None:
if request_id in self._requests:
# This request is already processed.
logger.warning(f"Request {request_id} is already processed for: {ids}")
return
logger.info("Terminating worker pods: {}".format(ids))
scale_request = None
try:
scale_request = self._initialize_scale_request(
to_launch={}, to_delete_instances=ids
)
if scale_request.worker_groups_with_pending_deletes:
errors_msg = (
"There are workers to be deleted from: "
f"{scale_request.worker_groups_with_pending_deletes}. "
"Waiting for them to be deleted before adding new workers "
" to be deleted"
)
logger.warning(errors_msg)
self._add_terminate_errors(
ids,
request_id,
details=errors_msg,
)
return
self._submit_scale_request(scale_request)
# Only add to processed requests if successful
self._requests.add(request_id)
except Exception as e:
logger.exception(f"Error terminating nodes: {scale_request or 'N/A'}")
self._add_terminate_errors(ids, request_id, details=str(e), e=e)
def launch(self, shape: Dict[NodeType, int], request_id: str) -> None:
if request_id in self._requests:
# This request is already processed.
return
scale_request = None
try:
scale_request = self._initialize_scale_request(
to_launch=shape, to_delete_instances=[]
)
if scale_request.worker_groups_with_pending_deletes:
error_msg = (
"There are workers to be deleted from: "
f"{scale_request.worker_groups_with_pending_deletes}. "
"Waiting for them to be deleted before creating new workers."
)
logger.warning(error_msg)
self._add_launch_errors(
shape,
request_id,
details=error_msg,
)
return
self._submit_scale_request(scale_request)
# Only add to processed requests if successful
self._requests.add(request_id)
except Exception as e:
logger.exception(f"Error launching nodes: {scale_request or 'N/A'}")
self._add_launch_errors(shape, request_id, details=str(e), e=e)
def poll_errors(self) -> List[CloudInstanceProviderError]:
errors = []
errors += self._launch_errors_queue
errors += self._terminate_errors_queue
self._launch_errors_queue = []
self._terminate_errors_queue = []
return errors
def get_ippr_specs(self) -> IPPRSpecs:
"""Return the cached, validated IPPR specs for the cluster.
The IPPR specs are refreshed during the provider's periodic sync with the
API server by reading the RayCluster annotation and validating it against
the IPPR schema.
"""
return self._ippr_provider.get_ippr_specs()
def get_ippr_statuses(self) -> Dict[str, IPPRStatus]:
"""Return the latest per-pod IPPR statuses keyed by pod name.
These statuses are refreshed from the current pod list during the provider's
periodic sync with the API server.
"""
return self._ippr_provider.get_ippr_statuses()
def do_ippr_requests(self, resizes: List[IPPRStatus]) -> None:
"""Execute IPPR resize requests via the underlying IPPR provider.
Args:
resizes: The list of per-pod IPPR actions produced by the scheduler.
"""
self._ippr_provider.do_ippr_requests(resizes)
############################
# Private
############################
def _initialize_scale_request(
self, to_launch: Dict[NodeType, int], to_delete_instances: List[CloudInstanceId]
) -> "KubeRayProvider.ScaleRequest":
"""
Initialize the scale request based on the current state of the cluster and
the desired state (to launch, to delete).
Args:
to_launch: The desired number of workers to launch for each node type.
to_delete_instances: The instances to delete.
Returns:
The scale request.
"""
# Update the cached states.
self._sync_with_api_server()
ray_cluster = self.ray_cluster
cur_instances = self.instances
# Get the worker groups that have pending deletes and the worker groups that
# have finished deletes, and the set of workers included in the workersToDelete
# field of any worker group.
(
worker_groups_with_pending_deletes,
worker_groups_without_pending_deletes,
worker_to_delete_set,
) = self._get_workers_delete_info(ray_cluster, set(cur_instances.keys()))
observed_workers_dict = defaultdict(int)
for instance in cur_instances.values():
if instance.node_kind != NodeKind.WORKER:
continue
if instance.cloud_instance_id in worker_to_delete_set:
continue
observed_workers_dict[instance.node_type] += 1
# Calculate the desired number of workers by type.
num_workers_dict = defaultdict(int)
worker_groups = ray_cluster["spec"].get("workerGroupSpecs", [])
for worker_group in worker_groups:
node_type = worker_group["groupName"]
# Handle the case where users manually increase `minReplicas`
# to scale up the number of worker Pods. In this scenario,
# `replicas` will be smaller than `minReplicas`.
# num_workers_dict should account for multi-host replicas when
# `numOfHosts`` is set.
num_of_hosts = worker_group.get("numOfHosts", 1)
replicas = (
max(worker_group["replicas"], worker_group["minReplicas"])
* num_of_hosts
)
# The `replicas` field in worker group specs can be updated by users at any time.
# However, users should only increase the field (manually upscaling the worker group), not decrease it,
# because downscaling the worker group requires specifying which workers to delete explicitly in the `workersToDelete` field.
# Since we don't have a way to enforce this, we need to fix unexpected decreases on the `replicas` field by using actual observations.
# For example, if the user manually decreases the `replicas` field to 0 without specifying which workers to delete,
# we should fix the `replicas` field back to the number of observed workers excluding the workers to be deleted,
# otherwise, we won't have a correct `replicas` matches the actual number of workers eventually.
num_workers_dict[node_type] = max(
replicas, observed_workers_dict[node_type]
)
# Add to launch nodes.
for node_type, count in to_launch.items():
num_workers_dict[node_type] += count
to_delete_instances_by_type = defaultdict(list)
# Update the number of workers with to_delete_instances
# and group them by type.
for to_delete_id in to_delete_instances:
to_delete_instance = cur_instances.get(to_delete_id, None)
if to_delete_instance is None:
# This instance has already been deleted.
continue
if to_delete_instance.node_kind == NodeKind.HEAD:
# Not possible to delete head node.
continue
if to_delete_instance.cloud_instance_id in worker_to_delete_set:
# If the instance is already in the workersToDelete field of
# any worker group, skip it.
continue
num_workers_dict[to_delete_instance.node_type] -= 1
assert num_workers_dict[to_delete_instance.node_type] >= 0
to_delete_instances_by_type[to_delete_instance.node_type].append(
to_delete_instance
)
scale_request = KubeRayProvider.ScaleRequest(
desired_num_workers=num_workers_dict,
workers_to_delete=to_delete_instances_by_type,
worker_groups_without_pending_deletes=worker_groups_without_pending_deletes,
worker_groups_with_pending_deletes=worker_groups_with_pending_deletes,
)
return scale_request
def _submit_scale_request(
self, scale_request: "KubeRayProvider.ScaleRequest"
) -> None:
"""Submits a scale request to the Kubernetes API server.
This method will convert the scale request to patches and submit the patches
to the Kubernetes API server.
Args:
scale_request: The scale request.
Raises:
Exception: An exception is raised if the Kubernetes API server returns an
error.
"""
# Get the current ray cluster spec.
patch_payload = []
raycluster = self.ray_cluster
# Collect patches for replica counts.
for node_type, num_workers 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)
group_num_of_hosts = _worker_group_num_of_hosts(raycluster, group_index)
# the num_workers from the scale request is multiplied by numOfHosts, so we need to divide it back.
target_replicas = num_workers // group_num_of_hosts
# 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.
for (
node_type,
workers_to_delete_of_type,
) in scale_request.workers_to_delete.items():
group_index = _worker_group_index(raycluster, node_type)
worker_ids_to_delete = [
worker.cloud_instance_id for worker in workers_to_delete_of_type
]
patch = worker_delete_patch(group_index, worker_ids_to_delete)
patch_payload.append(patch)
# Clear the workersToDelete field for the worker groups that have been deleted.
for node_type in scale_request.worker_groups_without_pending_deletes:
if node_type in scale_request.workers_to_delete:
# This node type is still being deleted.
continue
group_index = _worker_group_index(raycluster, node_type)
patch = worker_delete_patch(group_index, [])
patch_payload.append(patch)
if len(patch_payload) == 0:
# No patch required.
return
logger.info(f"Submitting a scale request: {scale_request}")
self._patch(f"rayclusters/{self._cluster_name}", patch_payload)
def _add_launch_errors(
self,
shape: Dict[NodeType, int],
request_id: str,
details: str,
e: Optional[Exception] = None,
) -> None:
"""
Adds launch errors to the error queue.
Args:
shape: The shape of the nodes that failed to launch.
request_id: The request id of the launch request.
details: The details of the error.
e: The exception that caused the error.
"""
for node_type, count in shape.items():
self._launch_errors_queue.append(
LaunchNodeError(
node_type=node_type,
timestamp_ns=time.time_ns(),
count=count,
request_id=request_id,
details=details,
cause=e,
)
)
def _add_terminate_errors(
self,
ids: List[CloudInstanceId],
request_id: str,
details: str,
e: Optional[Exception] = None,
) -> None:
"""
Adds terminate errors to the error queue.
Args:
ids: The ids of the nodes that failed to terminate.
request_id: The request id of the terminate request.
details: The details of the error.
e: The exception that caused the error.
"""
for id in ids:
self._terminate_errors_queue.append(
TerminateNodeError(
cloud_instance_id=id,
timestamp_ns=time.time_ns(),
request_id=request_id,
details=details,
cause=e,
)
)
def _sync_with_api_server(self) -> None:
"""Fetches the RayCluster resource from the Kubernetes API server."""
self._ray_cluster = self._get(f"rayclusters/{self._cluster_name}")
self._refresh_no_driver_timeout_seconds()
self._ippr_provider.validate_and_set_ippr_specs(self._ray_cluster)
self._cached_instances = self._fetch_instances()
self._ippr_provider.sync_with_raylets()
def _refresh_no_driver_timeout_seconds(self) -> None:
"""Reads noDriverTimeoutSeconds from the RayCluster CR."""
opts = self._ray_cluster["spec"].get(AUTOSCALER_OPTIONS_KEY, {})
secs = opts.get(NO_DRIVER_TIMEOUT_SECONDS_KEY)
self._no_driver_timeout_seconds = float(secs) if secs is not None else None
@property
def ray_cluster(self) -> Dict[str, Any]:
return copy.deepcopy(self._ray_cluster)
@property
def instances(self) -> Dict[CloudInstanceId, CloudInstance]:
return copy.deepcopy(self._cached_instances)
@staticmethod
def _get_workers_delete_info(
ray_cluster_spec: Dict[str, Any], node_set: Set[CloudInstanceId]
) -> Tuple[Set[NodeType], Set[NodeType], Set[CloudInstanceId]]:
"""
Gets the worker groups that have pending deletes and the worker groups that
have finished deletes.
Args:
ray_cluster_spec: The RayCluster CR spec dict.
node_set: The set of currently known cloud instance IDs.
Returns:
A tuple of:
- worker_groups_with_pending_deletes: The worker groups that have pending
deletes.
- worker_groups_with_finished_deletes: The worker groups that have finished
deletes.
- worker_to_delete_set: A set of Pods that are included in the
workersToDelete field of any worker group.
"""
worker_groups_with_pending_deletes = set()
worker_groups_with_deletes = set()
worker_to_delete_set = set()
worker_groups = ray_cluster_spec["spec"].get("workerGroupSpecs", [])
for worker_group in worker_groups:
workersToDelete = worker_group.get("scaleStrategy", {}).get(
"workersToDelete", []
)
if not workersToDelete:
# No workers to delete in this group.
continue
node_type = worker_group["groupName"]
worker_groups_with_deletes.add(node_type)
for worker in workersToDelete:
worker_to_delete_set.add(worker)
if worker in node_set:
worker_groups_with_pending_deletes.add(node_type)
worker_groups_with_finished_deletes = (
worker_groups_with_deletes - worker_groups_with_pending_deletes
)
return (
worker_groups_with_pending_deletes,
worker_groups_with_finished_deletes,
worker_to_delete_set,
)
def _fetch_instances(self) -> Dict[CloudInstanceId, CloudInstance]:
"""
Fetches the pods from the Kubernetes API server and convert them to Ray
CloudInstance.
Returns:
A dict of CloudInstanceId to CloudInstance.
"""
# 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_head_pod_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.
cloud_instances = {}
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"]
cloud_instance = self._cloud_instance_from_pod(pod)
if cloud_instance:
cloud_instances[pod_name] = cloud_instance
self._ippr_provider.sync_ippr_status_from_pods(pod_list["items"])
return cloud_instances
@staticmethod
def _cloud_instance_from_pod(pod: Dict[str, Any]) -> Optional[CloudInstance]:
"""
Convert a pod to a Ray CloudInstance.
Args:
pod: The pod resource dict.
Returns:
The CloudInstance representing the pod, or None if the pod is not a
tracked Ray node (e.g. a redis-cleanup pod).
"""
labels = pod["metadata"]["labels"]
if labels[KUBERAY_LABEL_KEY_KIND] == KUBERAY_KIND_HEAD:
kind = NodeKind.HEAD
type = labels[KUBERAY_LABEL_KEY_TYPE]
elif labels[KUBERAY_LABEL_KEY_KIND] == KUBERAY_KIND_WORKER:
kind = NodeKind.WORKER
type = labels[KUBERAY_LABEL_KEY_TYPE]
else:
# Other ray nodes types defined by KubeRay.
# e.g. this could also be `redis-cleanup`
# We will not track these nodes.
return None
# TODO: we should prob get from the pod's env var (RAY_CLOUD_INSTANCE_ID)
# directly.
cloud_instance_id = pod["metadata"]["name"]
return CloudInstance(
cloud_instance_id=cloud_instance_id,
node_type=type,
node_kind=kind,
is_running=KubeRayProvider._is_running(pod),
)
@staticmethod
def _is_running(pod) -> bool:
"""Convert pod state to Ray NodeStatus
A cloud instance is considered running if the pod is in the running state,
else it could be pending/containers-terminated.
When it disappears from the list, it is considered terminated.
"""
if (
"containerStatuses" not in pod["status"]
or not pod["status"]["containerStatuses"]
):
return False
state = pod["status"]["containerStatuses"][0]["state"]
if "running" in state:
return True
return False
def _get(self, remote_path: str) -> Dict[str, Any]:
"""Get a resource from the Kubernetes API server."""
return self._k8s_api_client.get(remote_path)
def _patch(self, remote_path: str, payload: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Patch a resource on the Kubernetes API server."""
return self._k8s_api_client.patch(remote_path, payload)
def _evaluate_no_driver_termination(self) -> None:
"""Patches the no-driver-TTL annotation once no driver held for the timeout.
Detached actors do not count as a driver.
"""
# Feature disabled or a driver is attached: reset the anchor.
if self._no_driver_timeout_seconds is None:
self._no_driver_observed_since = None
return
has_active_driver, latest_job_end_time = self._driver_status()
if has_active_driver:
self._no_driver_observed_since = None
return
# A driver finished since the last check: it was attached during the
# no-driver window, so restart the timer.
if latest_job_end_time > self._last_seen_job_end_time:
self._last_seen_job_end_time = latest_job_end_time
self._no_driver_observed_since = None
# Anchor on the first loop with no driver, then dispatch once the
# no-driver window reaches the timeout.
now = time.monotonic()
if self._no_driver_observed_since is None:
self._no_driver_observed_since = now
if now - self._no_driver_observed_since < self._no_driver_timeout_seconds:
return
self._set_no_driver_annotation()
def _driver_status(self) -> Tuple[bool, int]:
"""Returns whether a non-internal driver is alive and the latest job end time.
Fails closed: a failed GCS query reports a driver as present.
"""
try:
jobs = self._gcs_client.get_all_job_info(
skip_submission_job_info_field=True,
skip_is_running_tasks_field=True,
)
except Exception:
logger.exception(
"Failed to query GCS job table; treating as drivers attached."
)
return True, self._last_seen_job_end_time
has_active_driver = False
latest_job_end_time = 0
for job in jobs.values():
# Ray-internal drivers (e.g. the dashboard) are not user activity.
if job.config.ray_namespace.startswith(RAY_INTERNAL_NAMESPACE_PREFIX):
continue
if job.is_dead:
latest_job_end_time = max(latest_job_end_time, job.end_time)
else:
has_active_driver = True
return has_active_driver, latest_job_end_time
def _set_no_driver_annotation(self) -> None:
"""Sets `ray.io/no-driver-ttl-expired=true` on the RayCluster CR.
Idempotent via the CR cached this reconcile loop; PATCH errors are swallowed.
"""
annotations = self._ray_cluster.get("metadata", {}).get("annotations", {})
if annotations.get(NO_DRIVER_TTL_EXPIRED_ANNOTATION) == "true":
return
path = f"rayclusters/{self._cluster_name}"
# Merge patch covers missing and present annotations in one call.
payload = {
"metadata": {"annotations": {NO_DRIVER_TTL_EXPIRED_ANNOTATION: "true"}}
}
try:
self._k8s_api_client.patch(
path,
payload,
content_type="application/merge-patch+json",
)
except Exception:
logger.exception(
"Failed to PATCH %s=true on RayCluster %s",
NO_DRIVER_TTL_EXPIRED_ANNOTATION,
self._cluster_name,
)
return
logger.info(
"Set %s=true on RayCluster %s.",
NO_DRIVER_TTL_EXPIRED_ANNOTATION,
self._cluster_name,
)
def _get_head_pod_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"]
@@ -0,0 +1,736 @@
import json
import logging
import re
import time
from decimal import Decimal
from typing import Any, Dict, List, Optional, Tuple, Union
import jsonschema
from ray._raylet import GcsClient
from ray.autoscaler._private.kuberay.node_provider import (
KUBERAY_KIND_HEAD,
KUBERAY_KIND_WORKER,
KUBERAY_LABEL_KEY_KIND,
KUBERAY_LABEL_KEY_TYPE,
IKubernetesHttpApiClient,
replace_patch,
)
from ray.autoscaler._private.kuberay.utils import parse_quantity
from ray.autoscaler.v2.schema import (
IPPRGroupSpec,
IPPRSpecs,
IPPRSpecsSchema,
IPPRStatus,
)
logger = logging.getLogger(__name__)
class KubeRayIPPRProvider:
"""Implements in-place pod resize (IPPR) operations for KubeRay pods.
This provider is responsible for:
- Validating and materializing IPPR specs from the RayCluster annotation
(``ray.io/ippr``) into typed structures (``IPPRSpecs``/``IPPRGroupSpec``).
- Tracking per-pod resize status (``IPPRStatus``) from Kubernetes pods and
computing the desired resize actions.
- Issuing Kubernetes Pod Resize API requests and keeping a shadow annotation
(``ray.io/ippr-status``) to track progress and temporary caps.
- Synchronizing successful resource changes with the Raylet so Ray's local
resource view matches Kubernetes.
Attributes:
_gcs_client: Ray GCS client used to fetch Raylet node information.
_k8s_api_client: Kubernetes HTTP client for patching pods.
_ippr_specs: Validated per-group IPPR specs (limits and timeouts).
_ippr_statuses: Latest per-pod IPPR statuses indexed by pod name.
_container_resources: Snapshot of container resource requests/limits
from both pod spec and pod status, per pod name, used to compute
patch diffs.
"""
def __init__(
self,
gcs_client: GcsClient,
k8s_api_client: IKubernetesHttpApiClient,
):
"""Create a new IPPR provider.
Args:
gcs_client: Ray GCS client for resolving Raylet addresses.
k8s_api_client: Kubernetes HTTP client to issue patch requests.
"""
self._gcs_client = gcs_client
self._k8s_api_client = k8s_api_client
self._ippr_specs: IPPRSpecs = IPPRSpecs(groups={})
self._ippr_statuses: Dict[str, IPPRStatus] = {}
self._container_resources: Dict[str, Any] = {}
def validate_and_set_ippr_specs(
self, ray_cluster: Optional[Dict[str, Any]]
) -> None:
"""Validate and load IPPR specs from a RayCluster CR.
Reads the ``ray.io/ippr`` annotation, validates it against
``IPPRSpecsSchema``, and converts it to typed ``IPPRSpecs`` with per-group
``IPPRGroupSpec`` entries. Minimal resources are derived from the group's
pod template; maximums and timeout come from the annotation. If the
annotation is removed, clear any previously loaded IPPR specs.
Args:
ray_cluster: The RayCluster custom resource as a dict. If missing or
lacking the annotation, this method is a no-op.
Raises:
ValueError: If the Ray pod template is incompatible with IPPR (e.g.,
missing required requests, using unsupported resizePolicy restarts,
or conflicting ``rayStartParams``).
Example:
import json
ray_cluster = {
"metadata": {
"name": "example-raycluster",
"annotations": {
"ray.io/ippr": json.dumps(
{
"groups": {
"headgroup": {
"max-cpu": "4",
"max-memory": "8Gi",
"resize-timeout": 300,
},
"small-workers": {
"max-cpu": 2,
"max-memory": "4Gi",
"resize-timeout": 120,
},
}
}
),
},
},
"spec": {
"headGroupSpec": {
"rayStartParams": {},
"template": {
"spec": {
"containers": [
{
"name": "ray-head",
"resources": {
"requests": {
"cpu": "1",
"memory": "2Gi",
},
"limits": {
"cpu": "2",
"memory": "4Gi",
},
},
"resizePolicy": [
{
"resourceName": "cpu",
"restartPolicy": "NotRequired",
},
{
"resourceName": "memory",
"restartPolicy": "NotRequired",
},
],
}
],
}
},
},
"workerGroupSpecs": [
{
"groupName": "small-workers",
"rayStartParams": {},
"template": {
"spec": {
"containers": [
{
"name": "ray-worker",
"resources": {
"requests": {
"cpu": "500m",
"memory": "1Gi",
},
},
}
],
}
},
}
],
},
}
provider.validate_and_set_ippr_specs(ray_cluster)
"""
if not ray_cluster:
return
specs_str = ray_cluster["metadata"].get("annotations", {}).get("ray.io/ippr")
if not specs_str:
self._ippr_specs = IPPRSpecs(groups={})
return
ippr_specs_raw = json.loads(specs_str)
jsonschema.validate(instance=ippr_specs_raw, schema=IPPRSpecsSchema)
# Validate and build typed spec per group
worker_groups = {
worker_group_spec["groupName"]: worker_group_spec
for worker_group_spec in ray_cluster["spec"].get("workerGroupSpecs", [])
}
worker_groups["headgroup"] = ray_cluster["spec"]["headGroupSpec"]
groups = {
group_name: _build_ippr_group_spec(group_spec, worker_groups[group_name])
for group_name, group_spec in ippr_specs_raw.get("groups", {}).items()
if group_name in worker_groups
}
self._ippr_specs = IPPRSpecs(groups=groups)
def sync_with_raylets(self) -> None:
"""Propagate completed K8s resizes to Raylets via GCS.
For any pod whose K8s resize has completed, update the corresponding Raylet's local resource
instances via GCS gRPC and clear the pending timestamp on the pod's
``ray.io/ippr-status`` annotation.
Three situations we can have exceptions are:
1. K8s API is not available.
2. GCS is not available.
3. Raylet is not available.
If a raylet is truly dead, its pod will also be deleted eventually.
All of the above exceptions can only be resolved in the future reconcile loops.
"""
for ippr_status in self._ippr_statuses.values():
if not ippr_status.need_sync_with_raylet():
continue
try:
self._gcs_client.resize_raylet_resource_instances(
ippr_status.raylet_id,
{
"CPU": ippr_status.current_cpu,
"memory": ippr_status.current_memory,
},
)
self._patch_ippr_status(ippr_status, resizing_at=None)
logger.info(f"Pod {ippr_status.cloud_instance_id} resized successfully")
except Exception as e:
logger.error(
f"Failed to resize pod {ippr_status.cloud_instance_id}: {e}. "
"If this persists, check GCS (e.g. Head/GCS logs and Raylet reachability) "
"and Kubernetes (e.g. API errors, pod events, ray.io/ippr-status) for "
"related request failures."
)
def sync_ippr_status_from_pods(self, pods: List[Dict[str, Any]]) -> None:
"""Refresh internal IPPR statuses and container resources from pods.
Parses pods to produce up-to-date ``IPPRStatus`` objects and stores
relevant request/limit snapshots for later patch computations.
Args:
pods: List of Kubernetes Pod resources for the Ray cluster.
"""
self._ippr_statuses = {}
self._container_resources = {}
if not self._ippr_specs.groups:
return
for pod in pods:
if "deletionTimestamp" in pod["metadata"]:
continue
labels = pod["metadata"].get("labels", {})
kind = labels.get(KUBERAY_LABEL_KEY_KIND)
if kind not in (KUBERAY_KIND_HEAD, KUBERAY_KIND_WORKER):
continue
ippr_group_spec = self._ippr_specs.groups.get(
labels.get(KUBERAY_LABEL_KEY_TYPE)
)
ippr_status, container_resource = _get_ippr_status_from_pod(
ippr_group_spec, pod
)
if ippr_status:
self._ippr_statuses[ippr_status.cloud_instance_id] = ippr_status
if ippr_status and container_resource:
self._container_resources[
ippr_status.cloud_instance_id
] = container_resource
def do_ippr_requests(self, resizes: List[IPPRStatus]) -> None:
"""Issue Kubernetes Pod Resize requests for the given targets.
If any dimension downsizes, attempt to first adjust the Raylet's local
resources via gRPC to avoid temporary overcommit in Ray's scheduler.
The raylet request uses ``min(desired, current)`` per resource so mixed
resizes (e.g. CPU up and memory down) do not advertise increases before
Kubernetes applies them; the reply is merged so upsizing targets are kept
for the pod patch.
Args:
resizes: List of IPPRStatus with desired resources and metadata.
"""
for resize in resizes:
logger.info(
f"Resizing pod {resize.cloud_instance_id} to cpu={resize.desired_cpu} memory={resize.desired_memory} from cpu={resize.current_cpu} memory={resize.current_memory}"
)
if resize.raylet_id and (
resize.desired_cpu < resize.current_cpu
or resize.desired_memory < resize.current_memory
):
# For any downsize, update Raylet first. Cap each resource at
# current so we never tell the scheduler about an upsize until
# sync_with_raylets runs after K8s applies the pod resize.
try:
updated_total_resources = (
self._gcs_client.resize_raylet_resource_instances(
resize.raylet_id,
{
"CPU": min(resize.desired_cpu, resize.current_cpu),
"memory": min(
resize.desired_memory, resize.current_memory
),
},
)
)
if (
"CPU" not in updated_total_resources
or "memory" not in updated_total_resources
):
raise RuntimeError(
f"CPU or memory not found in the response of resizing raylet resources: {updated_total_resources}"
)
if resize.desired_cpu < resize.current_cpu:
resize.desired_cpu = float(updated_total_resources["CPU"])
if resize.desired_memory < resize.current_memory:
resize.desired_memory = int(updated_total_resources["memory"])
except Exception as e:
logger.error(
f"Skip failed downsizing on pod {resize.cloud_instance_id}: {e}"
)
continue
self._patch_ippr_resize(resize)
def get_ippr_specs(self) -> IPPRSpecs:
"""Return the current validated IPPR specs."""
return self._ippr_specs
def get_ippr_statuses(self) -> Dict[str, IPPRStatus]:
"""Return the latest per-pod IPPR statuses keyed by pod name."""
return self._ippr_statuses
def _patch_ippr_resize(self, resize: IPPRStatus) -> None:
patch = self._make_ippr_patch(resize)
self._k8s_api_client.patch(
"pods/{}/resize".format(resize.cloud_instance_id), patch
)
self._patch_ippr_status(resize, resizing_at=int(time.time()))
def _patch_ippr_status(
self, resize: IPPRStatus, resizing_at: Optional[int]
) -> None:
"""Save the IPPR status to the pod annotation ``ray.io/ippr-status``.
The annotation is used to track the IPPR status of the pod across reconcile loops.
Args:
resize: The IPPR status to save.
resizing_at: Timestamp while a resize is in progress; pass ``None``
to clear after the resize completes (e.g. from ``sync_with_raylets``).
"""
self._k8s_api_client.patch(
"pods/{}".format(resize.cloud_instance_id),
{
"metadata": {
"annotations": {
"ray.io/ippr-status": json.dumps(
{
"raylet-id": resize.raylet_id,
"resizing-at": resizing_at,
"suggested-max-cpu": resize.suggested_max_cpu,
"suggested-max-memory": resize.suggested_max_memory,
"last-failed-at": resize.last_failed_at,
"last-failed-reason": resize.last_failed_reason,
}
)
}
}
},
content_type="application/strategic-merge-patch+json",
)
def _make_ippr_patch(self, resize: IPPRStatus) -> List[Dict[str, Any]]:
patch = []
path_prefix = "/spec/containers/0/resources"
container_resource = self._container_resources[resize.cloud_instance_id]
# When limits are present, preserve the existing gap (limits - requests)
# by adjusting requests proportionally so QoS doesn't change.
for resource_name, desired in (
("cpu", resize.desired_cpu),
("memory", resize.desired_memory),
):
if container_resource["status"]["limits"].get(resource_name):
# Gap between status limits and requests for each resource.
diff = _resource_gap(
container_resource["status"]["limits"],
container_resource["status"]["requests"],
resource_name,
)
patch.append(
replace_patch(f"{path_prefix}/limits/{resource_name}", desired)
)
patch.append(
replace_patch(
f"{path_prefix}/requests/{resource_name}",
_request_from_desired(desired, diff),
)
)
else:
# No limits configured: adjust requests only.
patch.append(
replace_patch(f"{path_prefix}/requests/{resource_name}", desired)
)
return patch
def _build_ippr_group_spec(
group_spec: Dict[str, Any], worker_group_spec: Dict[str, Any]
) -> IPPRGroupSpec:
# Disallow per-pod overrides that conflict with IPPR's dynamic sizing.
ray_start_params = worker_group_spec.get("rayStartParams", {})
if "num-cpus" in ray_start_params or "memory" in ray_start_params:
raise ValueError(
"should not have 'num-cpus' or 'memory' in rayStartParams if IPPR is used"
)
container_spec = worker_group_spec["template"]["spec"]["containers"][0]
pod_spec_requests = container_spec.get("resources", {}).get("requests", {})
# Pod template must declare baseline CPU/memory requests for IPPR.
if "cpu" not in pod_spec_requests or "memory" not in pod_spec_requests:
raise ValueError(
"should have 'cpu' and 'memory' in resource requests as the resources lower bounds if IPPR is used"
)
for policy in container_spec.get("resizePolicy", []):
resource_name = policy.get("resourceName")
if resource_name != "cpu" and resource_name != "memory":
continue
restart = policy.get("restartPolicy")
# IPPR requires NotRequired so that K8s won't restart the container
# during in-place resource updates.
if restart is not None and restart != "NotRequired":
raise ValueError("IPPR only supports restartPolicy=NotRequired")
# pod_spec_limits are the initial resource limits specified for the pod.
# we use it together with pod_spec_requests to derive the lower bounds for IPPR.
pod_spec_limits = container_spec.get("resources", {}).get("limits", {})
return IPPRGroupSpec(
min_cpu=_resource_value(pod_spec_requests, pod_spec_limits, "cpu", float),
min_memory=_resource_value(pod_spec_requests, pod_spec_limits, "memory", int),
max_cpu=float(parse_quantity(group_spec.get("max-cpu"))),
max_memory=int(parse_quantity(group_spec.get("max-memory"))),
resize_timeout=int(group_spec.get("resize-timeout")),
)
def _get_ippr_status_from_pod(
ippr_group_spec: Optional[IPPRGroupSpec],
pod: Dict[str, Any],
) -> Tuple[Optional[IPPRStatus], Optional[Dict[str, Any]]]:
"""Build IPPRStatus and container resource snapshots from a Pod.
Returns a tuple of (ippr_status, container_resource_snapshot). The snapshot
contains both spec and status requests/limits used to construct resize
patches that preserve the current QoS gap.
IPPRStatus includes the updated desired resources based on the failure messages of the previous resize request carried on the pod.
This function doesn't have any external side effects.
"""
if not ippr_group_spec:
return (None, None)
container_status = {}
other_container_resources = []
for status in pod.get("status", {}).get("containerStatuses", []):
if status["name"] == pod["spec"]["containers"][0]["name"]:
container_status = status
else:
# We need to substract other containers' resources when adjusting the
# the new resource requests based on the capactity report from the Kubelet.
requests = status.get("resources", {}).get("requests")
if requests:
other_container_resources.append(requests)
pod_spec_requests = (
pod["spec"]["containers"][0].get("resources", {}).get("requests", {})
)
pod_spec_limits = (
pod["spec"]["containers"][0].get("resources", {}).get("limits", {})
)
pod_status_requests = container_status.get("resources", {}).get(
"requests", pod_spec_requests
)
pod_status_limits = container_status.get("resources", {}).get(
"limits", pod_spec_limits
)
ippr_status = IPPRStatus(
cloud_instance_id=pod["metadata"]["name"],
spec=ippr_group_spec,
desired_cpu=_resource_value(pod_spec_requests, pod_spec_limits, "cpu", float),
desired_memory=_resource_value(
pod_spec_requests, pod_spec_limits, "memory", int
),
current_cpu=_resource_value(
pod_status_requests, pod_status_limits, "cpu", float
),
current_memory=_resource_value(
pod_status_requests, pod_status_limits, "memory", int
),
)
ippr_status = _restore_ippr_status_from_annotation(ippr_status, pod)
ippr_status = _apply_resize_conditions(
ippr_status=ippr_status,
pod=pod,
pod_status_requests=pod_status_requests,
pod_status_limits=pod_status_limits,
other_container_resources=other_container_resources,
)
ippr_status = _handle_failed_or_timed_out_ippr(ippr_status)
return (
ippr_status,
{
"spec": {
"requests": pod_spec_requests,
"limits": pod_spec_limits,
},
"status": {
"requests": pod_status_requests,
"limits": pod_status_limits,
},
},
)
def _restore_ippr_status_from_annotation(
ippr_status: IPPRStatus, pod: Dict[str, Any]
) -> IPPRStatus:
"""Restore previously persisted IPPR status fields from pod annotations."""
pod_ippr_status_json = (
pod["metadata"].get("annotations", {}).get("ray.io/ippr-status")
)
if not pod_ippr_status_json:
return ippr_status
pod_ippr_status = json.loads(pod_ippr_status_json)
ippr_status.raylet_id = pod_ippr_status.get("raylet-id")
ippr_status.resizing_at = pod_ippr_status.get("resizing-at")
ippr_status.suggested_max_cpu = pod_ippr_status.get("suggested-max-cpu")
ippr_status.suggested_max_memory = pod_ippr_status.get("suggested-max-memory")
ippr_status.last_failed_at = pod_ippr_status.get("last-failed-at")
ippr_status.last_failed_reason = pod_ippr_status.get("last-failed-reason")
return ippr_status
def _apply_resize_conditions(
ippr_status: IPPRStatus,
pod: Dict[str, Any],
pod_status_requests: Dict[str, Any],
pod_status_limits: Dict[str, Any],
other_container_resources: List[Dict[str, Any]],
) -> IPPRStatus:
"""Parse pod resize conditions and queue any follow-up suggestions."""
for condition in pod.get("status", {}).get("conditions", []):
if condition["type"] == "PodResizePending" and condition["status"] == "True":
ippr_status.k8s_resize_message = condition.get("message")
ippr_status.k8s_resize_status = condition.get("reason", "").lower()
ippr_status = _apply_resize_suggestion(
ippr_status=ippr_status,
pod_status_requests=pod_status_requests,
pod_status_limits=pod_status_limits,
other_container_resources=other_container_resources,
)
break
if condition["type"] == "PodResizeInProgress" and condition["status"] == "True":
ippr_status.k8s_resize_message = condition.get("message")
ippr_status.k8s_resize_status = "inprogress"
if condition.get("reason") == "Error":
ippr_status.k8s_resize_status = "error"
break
return ippr_status
def _apply_resize_suggestion(
ippr_status: IPPRStatus,
pod_status_requests: Dict[str, Any],
pod_status_limits: Dict[str, Any],
other_container_resources: List[Dict[str, Any]],
) -> IPPRStatus:
"""Parse Kubelet capacity reports and queue a suggested follow-up resize.
A suggested follow-up resize is taken from the failure message of the previous resize request on the pod.
For example, a message, "Node didn't have enough resource: cpu, requested: 8000, used: 5000, capacity: 9000",
means the pod can only have CPU request up to 4. For our suggested follow-up resize, we will set the suggested_max_cpu to
4 + (orignal cpu limit - original cpu request) to keep the gap.
"""
report = None
if ippr_status.k8s_resize_message and ippr_status.k8s_resize_status == "deferred":
report = re.search(
r"Node didn't have enough resource: (cpu|memory), requested: (\d+), used: (\d+), capacity: (\d+)",
ippr_status.k8s_resize_message,
)
elif (
ippr_status.k8s_resize_message and ippr_status.k8s_resize_status == "infeasible"
):
report = re.search(
r"Node didn't have enough capacity: (cpu|memory), requested: (\d+), ()capacity: (\d+)",
ippr_status.k8s_resize_message,
)
if not report:
return ippr_status
# Example (resize to max-cpu = 8, max-memory = 20Gi):
# Initial pod status:
# - CPU limit is 4 cores
# - CPU request is 1 core (gap = 3 cores)
# - Mem limit is 8Gi
# - Mem request is 2Gi (gap = 6Gi)
# Initial resize request will be:
# - desired_cpu=8 cores
# - desired_memory=20Gi
# The actual resize patch will be:
# - CPU limit is 8 cores (upsize from 4)
# - CPU request is 8 - 3 = 5 cores (upsize from 1, keep the gap 3 cores)
# - Mem limit is 20Gi (upsize from 8Gi)
# - Mem request is 20 6 = 14Gi (upsize from 2Gi, keep the gap 6Gi)
# If Kubelet reports (the deferred case):
# - CPU: used=5, capacity=9 → remaining_cpu = 4 cores
# - Mem: used=6Gi, capacity=10Gi → remaining_mem = 4Gi
# The suggestions used in the next patch will be:
# - suggested_max_cpu = remaining_cpu + cpu_gap = 4 + 3 = 7 cores
# - suggested_max_memory = remaining_mem + mem_gap = 4Gi + 6Gi = 10Gi
# The actual resize patch will be:
# - CPU limit is 7 cores
# - CPU request is 7 - 3 = 4 cores (aligned with the kubelet's report, and keep the gap 3 cores)
# - Mem limit is 10Gi
# - Mem request is 10Gi - 6Gi = 4Gi (aligned with the kubelet's report, and keep the gap 6Gi)
used = int(
report.group(3) or "0"
) # this field is the used resource request capacity of the k8s node excluding the current pod.
capacity = int(
report.group(4)
) # this field is the total resource request capacity of the k8s node.
max_request = (
capacity - used
) # so this max_request is the remaining resource request capacity that this pod can still request.
resource_name = report.group(1)
suggested = _suggested_resize_limit(
resource_name,
max_request,
pod_status_requests,
pod_status_limits,
other_container_resources,
)
if resource_name == "cpu":
ippr_status.suggested_max_cpu = float(suggested)
if ippr_status.queue_resize_request(desired_cpu=ippr_status.suggested_max_cpu):
logger.info(
f"Apply resize suggestions for {ippr_status.cloud_instance_id} to cpu={ippr_status.suggested_max_cpu}"
)
else:
ippr_status.suggested_max_memory = int(suggested)
if ippr_status.queue_resize_request(
desired_memory=ippr_status.suggested_max_memory,
):
logger.info(
f"Apply resize suggestions for {ippr_status.cloud_instance_id} to memory={ippr_status.suggested_max_memory}"
)
return ippr_status
def _suggested_resize_limit(
resource_name: str,
max_request: int,
pod_status_requests: Dict[str, Any],
pod_status_limits: Dict[str, Any],
other_container_resources: List[Dict[str, Any]],
) -> Union[float, int]:
gap = _resource_gap(pod_status_limits, pod_status_requests, resource_name)
other_requests = sum(
parse_quantity(requests.get(resource_name, "0"))
for requests in other_container_resources
)
if resource_name == "cpu":
available = Decimal(str(max_request)) / 1000
return float(available + gap - other_requests)
else:
available = Decimal(str(max_request))
return int(available + gap - other_requests)
def _resource_value(
requests: Dict[str, Any],
limits: Dict[str, Any],
resource_name: str,
value_type: Union[type[float], type[int]],
) -> Union[float, int]:
return value_type(
parse_quantity(limits.get(resource_name) or requests.get(resource_name))
)
def _resource_gap(
limits: Dict[str, Any],
requests: Dict[str, Any],
resource_name: str,
) -> Decimal:
return parse_quantity(
limits.get(resource_name) or requests.get(resource_name)
) - parse_quantity(requests.get(resource_name))
def _request_from_desired(
desired: Union[float, int], gap: Decimal
) -> Union[float, int]:
requested = Decimal(str(desired)) - gap
return type(desired)(requested)
def _handle_failed_or_timed_out_ippr(ippr_status: IPPRStatus) -> IPPRStatus:
"""Record terminal IPPR failures and queue a revert to current resources."""
if not (ippr_status.is_errored() or ippr_status.is_timeout()):
return ippr_status
if ippr_status.last_failed_at is None:
if ippr_status.is_errored():
ippr_status.record_failure(
reason=ippr_status.k8s_resize_message or "Pod resize failed"
)
else:
ippr_status.record_failure(reason="Pod resize timed out")
if ippr_status.queue_resize_request(
desired_cpu=ippr_status.current_cpu,
desired_memory=ippr_status.current_memory,
):
logger.info(
f"Revert failed or stuck IPPR for {ippr_status.cloud_instance_id} to cpu={ippr_status.current_cpu} memory={ippr_status.current_memory}"
)
return ippr_status
@@ -0,0 +1,73 @@
from typing import Dict, List
from ray._common.utils import binary_to_hex
from ray._raylet import GcsClient
from ray.autoscaler._private.util import format_readonly_node_type
from ray.autoscaler.v2.instance_manager.node_provider import (
CloudInstance,
CloudInstanceId,
CloudInstanceProviderError,
ICloudInstanceProvider,
NodeKind,
)
from ray.autoscaler.v2.sdk import get_cluster_resource_state
from ray.autoscaler.v2.utils import is_head_node
from ray.core.generated.autoscaler_pb2 import NodeStatus
class ReadOnlyProvider(ICloudInstanceProvider):
"""
A read only provider that use the ray node states from the GCS as the
cloud instances.
This is used for laptop mode / manual cluster setup modes, in order to
provide status reporting in the same way for users.
"""
def __init__(self, provider_config: dict):
self._provider_config = provider_config
self._gcs_address = provider_config["gcs_address"]
self._gcs_client = GcsClient(address=self._gcs_address)
def get_non_terminated(self) -> Dict[str, CloudInstance]:
cluster_resource_state = get_cluster_resource_state(self._gcs_client)
cloud_instances = {}
for gcs_node_state in cluster_resource_state.node_states:
if gcs_node_state.status == NodeStatus.DEAD:
# Skip dead nodes.
continue
# Use node's node id if instance id is not available
cloud_instance_id = (
gcs_node_state.instance_id
if gcs_node_state.instance_id
else binary_to_hex(gcs_node_state.node_id)
)
# TODO: we should add a field to the proto to indicate if the node is head
# or not.
is_head = is_head_node(gcs_node_state)
cloud_instances[cloud_instance_id] = CloudInstance(
cloud_instance_id=cloud_instance_id,
node_kind=NodeKind.HEAD if is_head else NodeKind.WORKER,
node_type=format_readonly_node_type(
binary_to_hex(gcs_node_state.node_id) # Legacy behavior.
),
is_running=True,
request_id="",
)
return cloud_instances
def terminate(self, ids: List[CloudInstanceId], request_id: str) -> None:
raise NotImplementedError("Cannot terminate instances in read-only mode.")
def launch(
self, shape: Dict[CloudInstanceId, int], request_id: CloudInstanceId
) -> None:
raise NotImplementedError("Cannot launch instances in read-only mode.")
def poll_errors(self) -> List[CloudInstanceProviderError]:
return []
@@ -0,0 +1,511 @@
import time
import uuid
from typing import Dict, List, Optional, Set
from ray.core.generated.instance_manager_pb2 import Instance, InstanceUpdateEvent
class InstanceUtil:
"""
A helper class to group updates and operations on an Instance object defined
in instance_manager.proto
"""
# Memoized reachable from sets, where the key is the instance status, and
# the value is the set of instance status that is reachable from the key
# instance status.
_reachable_from: Optional[
Dict["Instance.InstanceStatus", Set["Instance.InstanceStatus"]]
] = None
@staticmethod
def new_instance(
instance_id: str,
instance_type: str,
status: Instance.InstanceStatus,
details: str = "",
) -> Instance:
"""
Returns a new instance with the given status.
Args:
instance_id: The instance id.
instance_type: The instance type.
status: The status of the new instance.
details: The details of the status transition.
Returns:
The newly-created instance.
"""
instance = Instance()
instance.version = 0 # it will be populated by the underlying storage.
instance.instance_id = instance_id
instance.instance_type = instance_type
instance.status = status
InstanceUtil._record_status_transition(instance, status, details)
return instance
@staticmethod
def random_instance_id() -> str:
"""
Returns a random instance id.
"""
return str(uuid.uuid4())
@staticmethod
def is_cloud_instance_allocated(instance_status: Instance.InstanceStatus) -> bool:
"""
Returns True if the instance is in a status where there could exist
a cloud instance allocated by the cloud provider.
"""
assert instance_status != Instance.UNKNOWN
return instance_status in {
Instance.ALLOCATED,
Instance.RAY_INSTALLING,
Instance.RAY_RUNNING,
Instance.RAY_STOPPING,
Instance.RAY_STOP_REQUESTED,
Instance.RAY_STOPPED,
Instance.TERMINATING,
Instance.RAY_INSTALL_FAILED,
Instance.TERMINATION_FAILED,
Instance.ALLOCATION_TIMEOUT,
}
@staticmethod
def is_ray_running(instance_status: Instance.InstanceStatus) -> bool:
"""
Returns True if the instance is in a status where the ray process is
running on the cloud instance.
i.e. RAY_RUNNING, RAY_STOP_REQUESTED, RAY_STOPPING
"""
assert instance_status != Instance.UNKNOWN
if instance_status in InstanceUtil.get_reachable_statuses(
Instance.RAY_STOPPING
):
return False
if instance_status in InstanceUtil.get_reachable_statuses(Instance.RAY_RUNNING):
return True
return False
@staticmethod
def is_ray_pending(instance_status: Instance.InstanceStatus) -> bool:
"""
Returns True if the instance is in a status where the ray process is
pending to be started on the cloud instance.
"""
assert instance_status != Instance.UNKNOWN
# Not gonna be in a RAY_RUNNING status.
if Instance.RAY_RUNNING not in InstanceUtil.get_reachable_statuses(
instance_status
):
return False
# Already running ray.
if instance_status in InstanceUtil.get_reachable_statuses(Instance.RAY_RUNNING):
return False
return True
def is_ray_running_reachable(instance_status: Instance.InstanceStatus) -> bool:
"""
Returns True if the instance is in a status where it may transition
to RAY_RUNNING status.
"""
return Instance.RAY_RUNNING in InstanceUtil.get_reachable_statuses(
instance_status
)
@staticmethod
def set_status(
instance: Instance,
new_instance_status: Instance.InstanceStatus,
details: str = "",
) -> bool:
"""Transitions the instance to the new state.
Args:
instance: The instance to update.
new_instance_status: The new status to transition to.
details: The details of the transition.
Returns:
True if the status transition is successful, False otherwise.
"""
if (
new_instance_status
not in InstanceUtil.get_valid_transitions()[instance.status]
):
return False
instance.status = new_instance_status
InstanceUtil._record_status_transition(instance, new_instance_status, details)
return True
@staticmethod
def _record_status_transition(
instance: Instance, status: Instance.InstanceStatus, details: str
):
"""Records the status transition.
Args:
instance: The instance to update.
status: The new status to transition to.
details: The details of the status transition.
"""
now_ns = time.time_ns()
instance.status_history.append(
Instance.StatusHistory(
instance_status=status,
timestamp_ns=now_ns,
details=details,
)
)
@staticmethod
def has_timeout(instance: Instance, timeout_s: int) -> bool:
"""
Returns True if the instance has been in the current status for more
than the given timeout.
Args:
instance: The instance to check.
timeout_s: The timeout in seconds.
Returns:
True if the instance has been in the current status for more than
the timeout_s seconds.
"""
cur_status = instance.status
status_times_ns = InstanceUtil.get_status_transition_times_ns(
instance, select_instance_status=cur_status
)
assert len(status_times_ns) >= 1, (
f"instance {instance.instance_id} has {len(status_times_ns)} "
f"{Instance.InstanceStatus.Name(cur_status)} status"
)
status_time_ns = sorted(status_times_ns)[-1]
if time.time_ns() - status_time_ns <= (timeout_s * 1e9):
return False
return True
@staticmethod
def get_valid_transitions() -> Dict[
"Instance.InstanceStatus", Set["Instance.InstanceStatus"]
]:
return {
# This is the initial status of a new instance.
Instance.QUEUED: {
# Cloud provider requested to launch a node for the instance.
# This happens when the a launch request is made to the node provider.
Instance.REQUESTED,
# Allocation request canceled before being requested.
# This happens when max_workers config is reduced or other termination
# triggers occur while the instance is still queued.
Instance.TERMINATED,
},
# When in this status, a launch request to the node provider is made.
Instance.REQUESTED: {
# Cloud provider allocated a cloud instance for the instance.
# This happens when the cloud instance first appears in the list of
# running cloud instances from the cloud instance provider.
Instance.ALLOCATED,
# Retry the allocation, become queueing again.
Instance.QUEUED,
# Cloud provider fails to allocate one. Either as a timeout or
# the launch request fails immediately.
Instance.ALLOCATION_FAILED,
},
# When in this status, the cloud instance is allocated and running. This
# happens when the cloud instance is present in node provider's list of
# running cloud instances.
Instance.ALLOCATED: {
# Ray needs to be install and launch on the provisioned cloud instance.
# This happens when the cloud instance is allocated, and the autoscaler
# is responsible for installing and launching ray on the cloud instance.
# For node provider that manages the ray installation and launching,
# this state is skipped.
Instance.RAY_INSTALLING,
# Ray is already installed on the provisioned cloud
# instance. It could be any valid ray status.
Instance.RAY_RUNNING,
# The cloud provider timed out for allocating running cloud instance.
# The CloudResourceMonitor subscriber will lower this node-type's priority
# in feature schedules.
Instance.ALLOCATION_TIMEOUT,
Instance.RAY_STOPPING,
Instance.RAY_STOPPED,
# Instance is requested to be stopped, e.g. instance leaked: no matching
# Instance with the same type is found in the autoscaler's state.
Instance.TERMINATING,
# cloud instance somehow failed.
Instance.TERMINATED,
},
# Ray process is being installed and started on the cloud instance.
# This status is skipped for node provider that manages the ray
# installation and launching. (e.g. Ray-on-Spark)
Instance.RAY_INSTALLING: {
# Ray installed and launched successfully, reported by the ray cluster.
# Similar to the Instance.ALLOCATED -> Instance.RAY_RUNNING transition,
# where the ray process is managed by the node provider.
Instance.RAY_RUNNING,
# Ray installation failed. This happens when the ray process failed to
# be installed and started on the cloud instance.
Instance.RAY_INSTALL_FAILED,
# Wen the ray node is reported as stopped by the ray cluster.
# This could happen that the ray process was stopped quickly after start
# such that a ray running node wasn't discovered and the RAY_RUNNING
# transition was skipped.
Instance.RAY_STOPPED,
# A cloud instance is being terminated (when the instance itself is no
# longer needed, e.g. instance is outdated, autoscaler is scaling down)
Instance.TERMINATING,
# cloud instance somehow failed during the installation process.
Instance.TERMINATED,
},
# Ray process is installed and running on the cloud instance. When in this
# status, a ray node must be present in the ray cluster.
Instance.RAY_RUNNING: {
# Ray is requested to be stopped.
Instance.RAY_STOP_REQUESTED,
# Ray is stopping (currently draining),
# e.g. idle termination.
Instance.RAY_STOPPING,
# Ray is already stopped, as reported by the ray cluster.
Instance.RAY_STOPPED,
# A cloud instance is being terminated (when the instance itself is no
# longer needed, e.g. instance is outdated, autoscaler is scaling down)
Instance.TERMINATING,
# cloud instance somehow failed.
Instance.TERMINATED,
},
# Ray process should be stopped on the cloud instance. The RayStopper
# subscriber will listen to this status and stop the ray process.
Instance.RAY_STOP_REQUESTED: {
# Ray is stopping on the cloud instance.
Instance.RAY_STOPPING,
# Ray stopped already.
Instance.RAY_STOPPED,
# Ray stop request failed (e.g. idle node no longer idle),
# ray is still running.
Instance.RAY_RUNNING,
# cloud instance somehow failed.
Instance.TERMINATED,
},
# An instance has been allocated to a cloud instance, but the cloud
# provider timed out for allocating running cloud instance, e.g. the
# a kubernetes pod remains pending due to insufficient resources.
Instance.ALLOCATION_TIMEOUT: {
# Instance is requested to be stopped
Instance.TERMINATING,
# Cloud instance already disappeared; skip termination request.
# This transition is allowed to avoid unnecessary termination attempts
# when the cloud instance has already disappeared (e.g., manually deleted
# or terminated by another process). While this helps avoid unnecessary
# retries, it's important to monitor this transition as it may indicate
# underlying issues with the allocation or termination process itself.
Instance.TERMINATED,
},
# When in this status, the ray process is requested to be stopped to the
# ray cluster, but not yet present in the dead ray node list reported by
# the ray cluster.
Instance.RAY_STOPPING: {
# Ray is stopped, and the ray node is present in the dead ray node list
# reported by the ray cluster.
Instance.RAY_STOPPED,
# A cloud instance is being terminated (when the instance itself is no
# longer needed, e.g. instance is outdated, autoscaler is scaling down)
Instance.TERMINATING,
# cloud instance somehow failed.
Instance.TERMINATED,
},
# When in this status, the ray process is stopped, and the ray node is
# present in the dead ray node list reported by the ray cluster.
Instance.RAY_STOPPED: {
# A cloud instance is being terminated (when the instance itself is no
# longer needed, e.g. instance is outdated, autoscaler is scaling down)
Instance.TERMINATING,
# cloud instance somehow failed.
Instance.TERMINATED,
},
# When in this status, the cloud instance is requested to be stopped to
# the node provider.
Instance.TERMINATING: {
# When a cloud instance no longer appears in the list of running cloud
# instances from the node provider.
Instance.TERMINATED,
# When the cloud instance failed to be terminated.
Instance.TERMINATION_FAILED,
},
# When in this status, the cloud instance failed to be terminated by the
# node provider. We will keep retrying.
Instance.TERMINATION_FAILED: {
# Retry the termination, become terminating again.
Instance.TERMINATING,
# Cloud instance already disappeared; skip termination request.
Instance.TERMINATED,
},
# An instance is marked as terminated when:
# 1. A cloud instance disappears from the list of running cloud instances
# from the node provider (follows from TERMINATING or other running states).
# 2. An allocation request is canceled before cloud resources are allocated
# (follows from QUEUED).
# This is a terminal state.
Instance.TERMINATED: set(), # Terminal state.
# When in this status, the cloud instance failed to be allocated by the
# node provider.
Instance.ALLOCATION_FAILED: set(), # Terminal state.
Instance.RAY_INSTALL_FAILED: {
# Autoscaler requests to shutdown the instance when ray install failed.
Instance.TERMINATING,
# cloud instance somehow failed.
Instance.TERMINATED,
},
# Initial state before the instance is created. Should never be used.
Instance.UNKNOWN: set(),
}
@staticmethod
def get_status_transitions(
instance: Instance,
select_instance_status: Optional["Instance.InstanceStatus"] = None,
) -> List["Instance.StatusHistory"]:
"""
Returns the status history of the instance.
Args:
instance: The instance.
select_instance_status: The go-to status to search for, i.e. select
only status history when the instance transitions into the status.
If None, returns all status updates.
Returns:
The list of status updates matching ``select_instance_status``,
or all status updates when ``select_instance_status`` is None.
"""
history = []
for status_update in instance.status_history:
if (
select_instance_status
and status_update.instance_status != select_instance_status
):
continue
history.append(status_update)
return history
@staticmethod
def get_last_status_transition(
instance: Instance,
select_instance_status: Optional["Instance.InstanceStatus"] = None,
) -> Optional["Instance.StatusHistory"]:
"""
Returns the last status transition of the instance.
Args:
instance: The instance.
select_instance_status: The status to search for. If None, returns
the last status update.
Returns:
The last matching status update, or None if no status updates match.
"""
history = InstanceUtil.get_status_transitions(instance, select_instance_status)
history.sort(key=lambda x: x.timestamp_ns)
if history:
return history[-1]
return None
@staticmethod
def get_status_transition_times_ns(
instance: Instance,
select_instance_status: Optional["Instance.InstanceStatus"] = None,
) -> List[int]:
"""
Returns a list of timestamps of the instance status update.
Args:
instance: The instance.
select_instance_status: The status to search for. If None, returns
all status update timestamps.
Returns:
The list of timestamps of the instance status updates.
"""
return [
e.timestamp_ns
for e in InstanceUtil.get_status_transitions(
instance, select_instance_status
)
]
@classmethod
def get_reachable_statuses(
cls,
instance_status: Instance.InstanceStatus,
) -> Set["Instance.InstanceStatus"]:
"""
Returns the set of instance status that is reachable from the given
instance status following the status transitions.
This method is memoized.
Args:
instance_status: The instance status to start from.
Returns:
The set of instance status that is reachable from the given instance
status.
"""
if cls._reachable_from is None:
cls._compute_reachable()
return cls._reachable_from[instance_status]
@staticmethod
def get_log_str_for_update(instance: Instance, update: InstanceUpdateEvent) -> str:
"""Returns a log string for the given instance update."""
if update.upsert:
return (
f"New instance "
f"{Instance.InstanceStatus.Name(update.new_instance_status)} (id="
f"{instance.instance_id}, type={instance.instance_type}, "
f"cloud_instance_id={instance.cloud_instance_id}, "
f"ray_id={instance.node_id}): {update.details}"
)
return (
f"Update instance "
f"{Instance.InstanceStatus.Name(instance.status)}->"
f"{Instance.InstanceStatus.Name(update.new_instance_status)} (id="
f"{instance.instance_id}, type={instance.instance_type}, "
f"cloud_instance_id={instance.cloud_instance_id}, "
f"ray_id={instance.node_id}): {update.details}"
)
@classmethod
def _compute_reachable(cls):
"""
Computes and memorize the from status sets for each status machine with
a DFS search.
"""
valid_transitions = cls.get_valid_transitions()
def dfs(graph, start, visited):
"""
Regular DFS algorithm to find all reachable nodes from a given node.
"""
for next_node in graph[start]:
if next_node not in visited:
# We delay adding the visited set here so we could capture
# the self loop.
visited.add(next_node)
dfs(graph, next_node, visited)
return visited
# Initialize the graphs
cls._reachable_from = {}
for status in Instance.InstanceStatus.values():
# All nodes reachable from 'start'
visited = set()
cls._reachable_from[status] = dfs(valid_transitions, status, visited)
@@ -0,0 +1,561 @@
import copy
import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional
import yaml
from ray._common.utils import binary_to_hex
from ray._private.ray_constants import env_integer
from ray._raylet import GcsClient
from ray.autoscaler._private.constants import (
AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
DEFAULT_UPSCALING_SPEED,
DISABLE_LAUNCH_CONFIG_CHECK_KEY,
DISABLE_NODE_UPDATERS_KEY,
)
from ray.autoscaler._private.kuberay.autoscaling_config import AutoscalingConfigProducer
from ray.autoscaler._private.monitor import BASE_READONLY_CONFIG
from ray.autoscaler._private.util import (
format_readonly_node_type,
hash_launch_conf,
hash_runtime_conf,
prepare_config,
validate_config,
)
from ray.autoscaler.v2.schema import NodeType
from ray.autoscaler.v2.sdk import get_cluster_resource_state
from ray.autoscaler.v2.utils import is_head_node
logger = logging.getLogger(__name__)
class Provider(Enum):
UNKNOWN = 0
ALIYUN = 1
AWS = 2
AZURE = 3
GCP = 4
KUBERAY = 5
LOCAL = 6
READ_ONLY = 7
class IConfigReader(ABC):
"""An interface for reading Autoscaling config.
A utility class that reads autoscaling configs from various sources:
- File
- In-memory dict
- Remote config service (e.g. KubeRay's config)
Example:
reader = FileConfigReader("path/to/config.yaml")
# Get the recently cached config.
config = reader.get_cached_autoscaling_config()
...
# Refresh the cached config.
reader.refresh_cached_autoscaling_config()
config = reader.get_cached_autoscaling_config()
"""
@abstractmethod
def get_cached_autoscaling_config(self) -> "AutoscalingConfig":
"""Returns the recently read autoscaling config.
Returns:
AutoscalingConfig: The recently read autoscaling config.
"""
pass
@abstractmethod
def refresh_cached_autoscaling_config(self):
"""Read the config from the source."""
pass
@dataclass(frozen=True)
class InstanceReconcileConfig:
# The timeout for waiting for a REQUESTED instance to be ALLOCATED.
request_status_timeout_s: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_REQUEST_STATUS_TIMEOUT_S", 10 * 60
)
# The timeout for waiting for a ALLOCATED instance to be RAY_RUNNING.
allocate_status_timeout_s: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_ALLOCATE_STATUS_TIMEOUT_S", 60 * 60
)
# The timeout for waiting for a RAY_INSTALLING instance to be RAY_RUNNING.
ray_install_status_timeout_s: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_RAY_INSTALL_STATUS_TIMEOUT_S", 30 * 60
)
# The timeout for waiting for a TERMINATING instance to be TERMINATED.
terminating_status_timeout_s: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_TERMINATING_STATUS_TIMEOUT_S", 300
)
# The timeout for waiting for a RAY_STOP_REQUESTED instance
# to be RAY_STOPPING or RAY_STOPPED.
ray_stop_requested_status_timeout_s: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_RAY_STOP_REQUESTED_STATUS_TIMEOUT_S", 300
)
# The interval for raise a warning when an instance in transient status
# is not updated for a long time.
transient_status_warn_interval_s: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_TRANSIENT_STATUS_WARN_INTERVAL_S", 90
)
# The number of times to retry requesting to allocate an instance.
max_num_retry_request_to_allocate: int = env_integer(
"RAY_AUTOSCALER_RECONCILE_MAX_NUM_RETRY_REQUEST_TO_ALLOCATE", 3
)
@dataclass
class NodeTypeConfig:
"""
NodeTypeConfig is the helper class to provide node type specific configs.
This maps to subset of the `available_node_types` field in the
autoscaling config.
"""
# Node type name
name: NodeType
# The minimal number of worker nodes to be launched for this node type.
min_worker_nodes: int
# The maximal number of worker nodes can be launched for this node type.
max_worker_nodes: int
# Idle timeout seconds for worker nodes of this node type.
idle_timeout_s: Optional[float] = None
# The priority of the worker group. Higher value means the group will be scaled up first if everything else is equal.
priority: int = 0
# The total resources on the node.
resources: Dict[str, float] = field(default_factory=dict)
# The labels on the node.
labels: Dict[str, str] = field(default_factory=dict)
# The node config's launch config hash. It's calculated from the auth
# config, and the node's config in the `AutoscalingConfig` for the node
# type when launching the node. It's used to detect config changes.
launch_config_hash: str = ""
def __post_init__(self):
assert self.min_worker_nodes <= self.max_worker_nodes
assert self.min_worker_nodes >= 0
class AutoscalingConfig:
"""
AutoscalingConfig is the helper class to provide autoscaling
related configs.
# TODO(rickyx):
1. Move the config validation logic here.
2. Deprecate the ray-schema.json for validation because it's
static thus not possible to validate the config with interdependency
of each other.
"""
def __init__(
self,
configs: Dict[str, Any],
skip_content_hash: bool = False,
) -> None:
"""Initialize the autoscaling config.
Args:
configs: The raw configs dict.
skip_content_hash: Whether to skip file mounts/ray command hash
calculation.
"""
self._sync_continuously = False
self.update_configs(configs, skip_content_hash)
def update_configs(self, configs: Dict[str, Any], skip_content_hash: bool) -> None:
self._configs = prepare_config(configs)
validate_config(self._configs)
if skip_content_hash:
return
self._calculate_hashes()
self._sync_continuously = self._configs.get(
"generate_file_mounts_contents_hash", True
)
def _calculate_hashes(self) -> None:
logger.info("Calculating hashes for file mounts and ray commands.")
self._runtime_hash, self._file_mounts_contents_hash = hash_runtime_conf(
self._configs.get("file_mounts", {}),
self._configs.get("cluster_synced_files", []),
[
self._configs.get("worker_setup_commands", []),
self._configs.get("worker_start_ray_commands", []),
],
generate_file_mounts_contents_hash=self._configs.get(
"generate_file_mounts_contents_hash", True
),
)
def get_cloud_node_config(self, ray_node_type: NodeType) -> Dict[str, Any]:
return copy.deepcopy(
self.get_node_type_specific_config(ray_node_type, "node_config") or {}
)
def get_docker_config(self, ray_node_type: NodeType) -> Dict[str, Any]:
"""
Return the docker config for the specified node type.
If it's a head node, the image will be chosen in the following order:
1. Node specific docker image.
2. The 'docker' config's 'head_image' field.
3. The 'docker' config's 'image' field.
If it's a worker node, the image will be chosen in the following order:
1. Node specific docker image.
2. The 'docker' config's 'worker_image' field.
3. The 'docker' config's 'image' field.
"""
# TODO(rickyx): It's unfortunate we have multiple fields in ray-schema.json
# that can specify docker images. We should consolidate them.
docker_config = copy.deepcopy(self._configs.get("docker", {}))
node_specific_docker_config = self._configs["available_node_types"][
ray_node_type
].get("docker", {})
# Override the global docker config with node specific docker config.
docker_config.update(node_specific_docker_config)
if self._configs.get("head_node_type") == ray_node_type:
if "head_image" in docker_config:
logger.info(
"Overwriting image={} by head_image({}) for head node docker.".format( # noqa: E501
docker_config["image"], docker_config["head_image"]
)
)
docker_config["image"] = docker_config["head_image"]
else:
if "worker_image" in docker_config:
logger.info(
"Overwriting image={} by worker_image({}) for worker node docker.".format( # noqa: E501
docker_config["image"], docker_config["worker_image"]
)
)
docker_config["image"] = docker_config["worker_image"]
# These fields should be merged.
docker_config.pop("head_image", None)
docker_config.pop("worker_image", None)
return docker_config
def get_worker_start_ray_commands(self) -> List[str]:
return self._configs.get("worker_start_ray_commands", [])
def get_head_setup_commands(self) -> List[str]:
return self._configs.get("head_setup_commands", [])
def get_head_start_ray_commands(self) -> List[str]:
return self._configs.get("head_start_ray_commands", [])
def get_worker_setup_commands(self, ray_node_type: NodeType) -> List[str]:
"""
Return the worker setup commands for the specified node type.
If the node type specific worker setup commands are not specified,
return the global worker setup commands.
"""
worker_setup_command = self.get_node_type_specific_config(
ray_node_type, "worker_setup_commands"
)
if worker_setup_command is None:
# Return global worker setup commands if node type specific
# worker setup commands are not specified.
logger.info(
"Using global worker setup commands for {}".format(ray_node_type)
)
return self._configs.get("worker_setup_commands", [])
return worker_setup_command
def get_initialization_commands(self, ray_node_type: NodeType) -> List[str]:
"""
Return the initialization commands for the specified node type.
If the node type specific initialization commands are not specified,
return the global initialization commands.
"""
initialization_command = self.get_node_type_specific_config(
ray_node_type, "initialization_commands"
)
if initialization_command is None:
logger.info(
"Using global initialization commands for {}".format(ray_node_type)
)
return self._configs.get("initialization_commands", [])
return initialization_command
def get_node_type_specific_config(
self, ray_node_type: NodeType, config_name: str
) -> Optional[Any]:
node_specific_config = self._configs["available_node_types"].get(
ray_node_type, {}
)
return node_specific_config.get(config_name, None)
def get_node_resources(self, ray_node_type: NodeType) -> Dict[str, float]:
return copy.deepcopy(
self.get_node_type_specific_config(ray_node_type, "resources") or {}
)
def get_node_labels(self, ray_node_type: NodeType) -> Dict[str, str]:
return copy.deepcopy(
self.get_node_type_specific_config(ray_node_type, "labels") or {}
)
def get_config(self, config_name, default=None) -> Any:
return self._configs.get(config_name, default)
def get_provider_instance_type(self, ray_node_type: NodeType) -> str:
provider = self.provider
node_config = (
self.get_node_type_specific_config(ray_node_type, "node_config") or {}
)
if provider in [Provider.AWS, Provider.ALIYUN]:
return node_config.get("InstanceType", "")
elif provider == Provider.AZURE:
return node_config.get("azure_arm_parameters", {}).get("vmSize", "")
elif provider == Provider.GCP:
return node_config.get("machineType", "")
elif provider in [Provider.KUBERAY, Provider.LOCAL, Provider.UNKNOWN]:
return ""
else:
raise ValueError(f"Unknown provider {provider}")
def get_node_type_configs(self) -> Dict[NodeType, NodeTypeConfig]:
"""
Returns the node type configs from the `available_node_types` field.
Returns:
Dict[NodeType, NodeTypeConfig]: The node type configs.
"""
available_node_types = self._configs.get("available_node_types", {})
if not available_node_types:
return None
node_type_configs = {}
auth_config = self._configs.get("auth", {})
head_node_type = self.get_head_node_type()
assert head_node_type
for node_type, node_config in available_node_types.items():
launch_config_hash = hash_launch_conf(
node_config.get("node_config", {}), auth_config
)
max_workers_nodes = node_config.get("max_workers", 0)
if head_node_type == node_type:
max_workers_nodes += 1
node_type_configs[node_type] = NodeTypeConfig(
name=node_type,
min_worker_nodes=node_config.get("min_workers", 0),
max_worker_nodes=max_workers_nodes,
idle_timeout_s=node_config.get("idle_timeout_s", None),
priority=node_config.get("priority", 0),
resources=node_config.get("resources", {}),
labels=node_config.get("labels", {}),
launch_config_hash=launch_config_hash,
)
return node_type_configs
def get_head_node_type(self) -> NodeType:
"""
Returns the head node type.
If there is only one node type, return the only node type as the head
node type.
If there are multiple node types, return the head node type specified
in the config.
"""
available_node_types = self._configs.get("available_node_types", {})
if len(available_node_types) == 1:
return list(available_node_types.keys())[0]
return self._configs.get("head_node_type")
def get_max_num_worker_nodes(self) -> Optional[int]:
return self.get_config("max_workers", None)
def get_max_num_nodes(self) -> Optional[int]:
max_num_workers = self.get_max_num_worker_nodes()
if max_num_workers is not None:
return max_num_workers + 1 # For head node
return None
def get_raw_config_mutable(self) -> Dict[str, Any]:
return self._configs
def get_upscaling_speed(self) -> float:
return self.get_config("upscaling_speed", DEFAULT_UPSCALING_SPEED)
def get_max_concurrent_launches(self) -> int:
return AUTOSCALER_MAX_CONCURRENT_LAUNCHES
def disable_node_updaters(self) -> bool:
provider_config = self._configs.get("provider", {})
return provider_config.get(DISABLE_NODE_UPDATERS_KEY, False)
def get_idle_timeout_s(self) -> Optional[float]:
"""
Returns the idle timeout in seconds if present in config, otherwise None.
"""
idle_timeout_s = self.get_config("idle_timeout_minutes", None)
return idle_timeout_s * 60 if idle_timeout_s is not None else None
def disable_launch_config_check(self) -> bool:
provider_config = self.get_provider_config()
return provider_config.get(DISABLE_LAUNCH_CONFIG_CHECK_KEY, True)
def get_instance_reconcile_config(self) -> InstanceReconcileConfig:
# TODO(rickyx): we need a way to customize these configs,
# either extending the current ray-schema.json, or just use another
# schema validation paths.
return InstanceReconcileConfig()
def get_provider_config(self) -> Dict[str, Any]:
return self._configs.get("provider", {})
def dump(self) -> str:
return yaml.safe_dump(self._configs)
@property
def provider(self) -> Provider:
provider_str = self._configs.get("provider", {}).get("type", "")
if provider_str == "local":
return Provider.LOCAL
elif provider_str == "aws":
return Provider.AWS
elif provider_str == "azure":
return Provider.AZURE
elif provider_str == "gcp":
return Provider.GCP
elif provider_str == "aliyun":
return Provider.ALIYUN
elif provider_str == "kuberay":
return Provider.KUBERAY
elif provider_str == "readonly":
return Provider.READ_ONLY
else:
return Provider.UNKNOWN
@property
def runtime_hash(self) -> str:
if not hasattr(self, "_runtime_hash"):
self._calculate_hashes()
return self._runtime_hash
@property
def file_mounts_contents_hash(self) -> str:
if not hasattr(self, "_file_mounts_contents_hash"):
self._calculate_hashes()
return self._file_mounts_contents_hash
class FileConfigReader(IConfigReader):
"""A class that reads cluster config from a yaml file."""
def __init__(self, config_file: str, skip_content_hash: bool = True) -> None:
"""Initialize the file config reader.
Args:
config_file: The path to the config file.
skip_content_hash: Whether to skip file mounts/ray command
hash calculation. Default to True.
"""
self._config_file_path = Path(config_file).resolve()
self._skip_content_hash = skip_content_hash
self._cached_config = self._read()
def _read(self) -> AutoscalingConfig:
with open(self._config_file_path) as f:
config = yaml.safe_load(f.read())
return AutoscalingConfig(config, skip_content_hash=self._skip_content_hash)
def get_cached_autoscaling_config(self) -> AutoscalingConfig:
"""
Returns:
AutoscalingConfig: The autoscaling config.
"""
return self._cached_config
def refresh_cached_autoscaling_config(self):
self._cached_config = self._read()
class KubeRayConfigReader(IConfigReader):
"""A class that reads cluster config from a K8s RayCluster CR."""
def __init__(self, config_producer: AutoscalingConfigProducer):
self._config_producer = config_producer
self._cached_config = self._generate_configs_from_k8s()
def _generate_configs_from_k8s(self) -> AutoscalingConfig:
return AutoscalingConfig(self._config_producer())
def get_cached_autoscaling_config(self) -> AutoscalingConfig:
"""
Returns:
AutoscalingConfig: The autoscaling config.
"""
return self._cached_config
def refresh_cached_autoscaling_config(self):
"""
Reads the configs from the K8s RayCluster CR.
This reads from the K8s API server every time to pick up changes.
"""
self._cached_config = self._generate_configs_from_k8s()
class ReadOnlyProviderConfigReader(IConfigReader):
"""A class that reads cluster config for a read-only provider.
This is used for laptop mode / manual cluster setup modes, in order to
provide status reporting in the same way for users."""
def __init__(self, gcs_address: str):
self._configs = BASE_READONLY_CONFIG
self._gcs_client = GcsClient(address=gcs_address)
def refresh_cached_autoscaling_config(self) -> AutoscalingConfig:
# Update the config with node types from GCS.
ray_cluster_resource_state = get_cluster_resource_state(self._gcs_client)
# Format each node type's config from the running nodes.
available_node_types = {}
head_node_type = None
for node_state in ray_cluster_resource_state.node_states:
node_type = node_state.ray_node_type_name
if not node_type:
node_type = format_readonly_node_type(binary_to_hex(node_state.node_id))
if is_head_node(node_state):
head_node_type = node_type
if node_type not in available_node_types:
available_node_types[node_type] = {
"resources": dict(node_state.total_resources),
"min_workers": 0,
"max_workers": 0 if is_head_node(node_state) else 1,
"node_config": {},
}
elif not is_head_node(node_state):
available_node_types[node_type]["max_workers"] += 1
if available_node_types:
self._configs["available_node_types"].update(available_node_types)
self._configs["max_workers"] = sum(
v["max_workers"] for v in available_node_types.values()
)
assert head_node_type, "Head node type should be found."
self._configs["head_node_type"] = head_node_type
# Don't idle terminated nodes in read-only mode.
self._configs.pop("idle_timeout_minutes", None)
def get_cached_autoscaling_config(self) -> AutoscalingConfig:
return AutoscalingConfig(self._configs, skip_content_hash=True)
@@ -0,0 +1,271 @@
import logging
from abc import ABC, abstractmethod
from typing import List, Optional
from ray.autoscaler.v2.instance_manager.common import InstanceUtil
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.core.generated.instance_manager_pb2 import (
GetInstanceManagerStateReply,
GetInstanceManagerStateRequest,
Instance,
InstanceUpdateEvent,
NodeKind,
StatusCode,
UpdateInstanceManagerStateReply,
UpdateInstanceManagerStateRequest,
)
logger = logging.getLogger(__name__)
class InstanceUpdatedSubscriber(ABC):
"""Subscribers to instance status changes."""
@abstractmethod
def notify(self, events: List[InstanceUpdateEvent]) -> None:
pass
class InstanceManager:
"""
See `InstanceManagerService` in instance_manager.proto
This handles updates to an instance, or inserts a new instance if
it's an insert update. We should only be inserting new instances
of the below statuses:
1. ALLOCATED: For unmanaged instance not initialized by InstanceManager,
e.g. head node
2. QUEUED: For new instance being queued to launch.
3. TERMINATING: For leaked cloud instance that needs to be terminated.
For full status transitions, see:
https://docs.google.com/document/d/1NzQjA8Mh-oMc-QxXOa529oneWCoA8sDiVoNkBqqDb4U/edit#heading=h.k9a1sp4qpqj4
Not thread safe, should be used as a singleton.
"""
def __init__(
self,
instance_storage: InstanceStorage,
instance_status_update_subscribers: Optional[List[InstanceUpdatedSubscriber]],
):
self._instance_storage = instance_storage
self._status_update_subscribers = instance_status_update_subscribers or []
def update_instance_manager_state(
self, request: UpdateInstanceManagerStateRequest
) -> UpdateInstanceManagerStateReply:
"""
Updates the instance manager state.
If there's any failure, no updates would be made and the reply
would contain the latest version of the instance manager state,
and the error info.
Args:
request: The request to update the instance manager state.
Returns:
The reply to the request.
"""
# Handle updates
ids_to_updates = {update.instance_id: update for update in request.updates}
to_update_instances, version = self._instance_storage.get_instances(
instance_ids=ids_to_updates.keys()
)
if request.expected_version >= 0 and request.expected_version != version:
err_str = (
f"Version mismatch: expected: {request.expected_version}, "
f"actual: {version}"
)
logger.warning(err_str)
return self._get_update_im_state_reply(
StatusCode.VERSION_MISMATCH,
version,
err_str,
)
# Handle instances states update.
to_upsert_instances = []
for instance_id, update in ids_to_updates.items():
if instance_id in to_update_instances:
instance = self._update_instance(
to_update_instances[instance_id], update
)
else:
instance = self._create_instance(update)
to_upsert_instances.append(instance)
# Updates the instance storage.
result = self._instance_storage.batch_upsert_instances(
updates=to_upsert_instances,
expected_storage_version=version,
)
if not result.success:
if result.version != version:
err_str = (
f"Version mismatch: expected: {version}, actual: {result.version}"
)
logger.warning(err_str)
return self._get_update_im_state_reply(
StatusCode.VERSION_MISMATCH, result.version, err_str
)
else:
err_str = "Failed to update instance storage."
logger.error(err_str)
return self._get_update_im_state_reply(
StatusCode.UNKNOWN_ERRORS, result.version, err_str
)
# Successful updates.
for subscriber in self._status_update_subscribers:
subscriber.notify(request.updates)
return self._get_update_im_state_reply(StatusCode.OK, result.version)
def get_instance_manager_state(
self, request: GetInstanceManagerStateRequest
) -> GetInstanceManagerStateReply:
"""
Gets the instance manager state.
Args:
request: The request to get the instance manager state.
Returns:
The reply to the request.
"""
reply = GetInstanceManagerStateReply()
instances, version = self._instance_storage.get_instances()
reply.state.instances.extend(instances.values())
reply.state.version = version
reply.status.code = StatusCode.OK
return reply
#########################################
# Private methods
#########################################
@staticmethod
def _get_update_im_state_reply(
status_code: StatusCode, version: int, error_message: str = ""
) -> UpdateInstanceManagerStateReply:
"""
Returns a UpdateInstanceManagerStateReply with the given status code and
version.
Args:
status_code: The status code.
version: The version.
error_message: The error message if any.
Returns:
The reply.
"""
reply = UpdateInstanceManagerStateReply()
reply.status.code = status_code
reply.version = version
if error_message:
reply.status.message = error_message
return reply
@staticmethod
def _apply_update(instance: Instance, update: InstanceUpdateEvent):
"""
Apply status specific update to the instance.
Args:
instance: The instance to update.
update: The update to apply.
"""
if update.new_instance_status == Instance.ALLOCATED:
assert (
update.cloud_instance_id
), "ALLOCATED update must have cloud_instance_id"
assert update.node_kind in [
NodeKind.WORKER,
NodeKind.HEAD,
], "ALLOCATED update must have node_kind as WORKER or HEAD"
assert update.instance_type, "ALLOCATED update must have instance_type"
assert (
update.cloud_instance_id
), "ALLOCATED update must have cloud_instance_id"
instance.cloud_instance_id = update.cloud_instance_id
instance.node_kind = update.node_kind
instance.instance_type = update.instance_type
instance.node_id = update.ray_node_id
elif update.new_instance_status == Instance.RAY_RUNNING:
assert update.ray_node_id, "RAY_RUNNING update must have ray_node_id"
instance.node_id = update.ray_node_id
elif update.new_instance_status == Instance.REQUESTED:
assert (
update.launch_request_id
), "REQUESTED update must have launch_request_id"
assert update.instance_type, "REQUESTED update must have instance_type"
instance.launch_request_id = update.launch_request_id
instance.instance_type = update.instance_type
elif update.new_instance_status == Instance.TERMINATING:
assert (
update.cloud_instance_id
), "TERMINATING update must have cloud instance id"
@staticmethod
def _create_instance(update: InstanceUpdateEvent) -> Instance:
"""
Create a new instance from the given update.
"""
assert update.upsert, "upsert must be true for creating new instance."
assert update.new_instance_status in [
# For unmanaged instance not initialized by InstanceManager,
# e.g. head node
Instance.ALLOCATED,
# For new instance being queued to launch.
Instance.QUEUED,
# For leaked cloud instance that needs to be terminated.
Instance.TERMINATING,
], (
"Invalid status for new instance, must be one of "
"[ALLOCATED, QUEUED, TERMINATING]"
)
# Create a new instance first for common fields.
instance = InstanceUtil.new_instance(
instance_id=update.instance_id,
instance_type=update.instance_type,
status=update.new_instance_status,
details=update.details,
)
# Apply the status specific updates.
logger.info(InstanceUtil.get_log_str_for_update(instance, update))
InstanceManager._apply_update(instance, update)
return instance
@staticmethod
def _update_instance(instance: Instance, update: InstanceUpdateEvent) -> Instance:
"""
Update the instance with the given update.
Args:
instance: The instance to update.
update: The update to apply.
Returns:
The updated instance.
"""
logger.info(InstanceUtil.get_log_str_for_update(instance, update))
assert InstanceUtil.set_status(instance, update.new_instance_status), (
"Invalid status transition from "
f"{Instance.InstanceStatus.Name(instance.status)} to "
f"{Instance.InstanceStatus.Name(update.new_instance_status)}"
)
InstanceManager._apply_update(instance, update)
return instance
@@ -0,0 +1,151 @@
import copy
import logging
from typing import Dict, List, Optional, Set, Tuple
from ray.autoscaler.v2.instance_manager.storage import Storage, StoreStatus
from ray.core.generated.instance_manager_pb2 import Instance
logger = logging.getLogger(__name__)
class InstanceStorage:
"""Instance storage stores the states of instances in the storage."""
def __init__(
self,
cluster_id: str,
storage: Storage,
) -> None:
self._storage = storage
self._cluster_id = cluster_id
self._table_name = f"instance_table@{cluster_id}"
def batch_upsert_instances(
self,
updates: List[Instance],
expected_storage_version: Optional[int] = None,
) -> StoreStatus:
"""Upsert instances into the storage. If the instance already exists,
it will be updated. Otherwise, it will be inserted. If the
expected_storage_version is specified, the update will fail if the
current storage version does not match the expected version.
Note the version of the upserted instances will be set to the current
storage version.
Args:
updates: A list of instances to be upserted.
expected_storage_version: The expected storage version.
Returns:
StoreStatus: A tuple of (success, storage_version).
"""
mutations = {}
version = self._storage.get_version()
# handle version mismatch
if expected_storage_version and expected_storage_version != version:
return StoreStatus(False, version)
for instance in updates:
instance = copy.deepcopy(instance)
# the instance version is set to 0, it will be
# populated by the storage entry's verion on read
instance.version = 0
mutations[instance.instance_id] = instance.SerializeToString()
result, version = self._storage.batch_update(
self._table_name, mutations, {}, expected_storage_version
)
return StoreStatus(result, version)
def upsert_instance(
self,
instance: Instance,
expected_instance_version: Optional[int] = None,
expected_storage_version: Optional[int] = None,
) -> StoreStatus:
"""Upsert an instance in the storage.
If the expected_instance_version is specified, the update will fail
if the current instance version does not match the expected version.
Similarly, if the expected_storage_version is
specified, the update will fail if the current storage version does not
match the expected version.
Note the version of the upserted instances will be set to the current
storage version.
Args:
instance: The instance to be updated.
expected_instance_version: The expected instance version.
expected_storage_version: The expected storage version.
Returns:
StoreStatus: A tuple of (success, storage_version).
"""
instance = copy.deepcopy(instance)
# the instance version is set to 0, it will be
# populated by the storage entry's verion on read
instance.version = 0
result, version = self._storage.update(
self._table_name,
key=instance.instance_id,
value=instance.SerializeToString(),
expected_entry_version=expected_instance_version,
expected_storage_version=expected_storage_version,
insert_only=False,
)
return StoreStatus(result, version)
def get_instances(
self,
instance_ids: List[str] = None,
status_filter: Set[int] = None,
) -> Tuple[Dict[str, Instance], int]:
"""Get instances from the storage.
Args:
instance_ids: A list of instance ids to be retrieved. If empty, all
instances will be retrieved.
status_filter: Only instances with the specified status will be returned.
Returns:
Tuple[Dict[str, Instance], int]: A tuple of (instances, version).
The instances is a dictionary of (instance_id, instance) pairs.
"""
instance_ids = instance_ids or []
status_filter = status_filter or set()
pairs, version = self._storage.get(self._table_name, instance_ids)
instances = {}
for instance_id, (instance_data, entry_version) in pairs.items():
instance = Instance()
instance.ParseFromString(instance_data)
instance.version = entry_version
if status_filter and instance.status not in status_filter:
continue
instances[instance_id] = instance
return instances, version
def batch_delete_instances(
self, instance_ids: List[str], expected_storage_version: Optional[int] = None
) -> StoreStatus:
"""Delete instances from the storage. If the expected_storage_version
is specified, the update will fail if the current storage version does
not match the expected version.
Args:
instance_ids: A list of instance ids to be deleted.
expected_storage_version: The expected storage version.
Returns:
StoreStatus: A tuple of (success, storage_version).
"""
version = self._storage.get_version()
if expected_storage_version and expected_storage_version != version:
return StoreStatus(False, version)
result = self._storage.batch_update(
self._table_name, {}, instance_ids, expected_storage_version
)
return result
@@ -0,0 +1,530 @@
import logging
import math
import time
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from queue import Queue
from typing import Any, Dict, List, Optional
from ray.autoscaler._private.constants import (
AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
AUTOSCALER_MAX_LAUNCH_BATCH,
)
from ray.autoscaler._private.util import hash_launch_conf
from ray.autoscaler.node_provider import NodeProvider as NodeProviderV1
from ray.autoscaler.tags import (
NODE_KIND_HEAD,
NODE_KIND_UNMANAGED,
NODE_KIND_WORKER,
STATUS_UNINITIALIZED,
TAG_RAY_LAUNCH_CONFIG,
TAG_RAY_LAUNCH_REQUEST,
TAG_RAY_NODE_KIND,
TAG_RAY_NODE_NAME,
TAG_RAY_NODE_STATUS,
TAG_RAY_USER_NODE_TYPE,
)
from ray.autoscaler.v2.instance_manager.config import IConfigReader
from ray.autoscaler.v2.schema import NodeType
from ray.core.generated.instance_manager_pb2 import NodeKind
logger = logging.getLogger(__name__)
# Type Alias. This is a **unique identifier** for a cloud instance in the cluster.
# The provider should guarantee that this id is unique across the cluster,
# such that:
# - When a cloud instance is created and running, no other cloud instance in the
# cluster has the same id.
# - When a cloud instance is terminated, no other cloud instance in the cluster will
# be assigned the same id later.
CloudInstanceId = str
@dataclass
class CloudInstance:
"""
A class that represents a cloud instance in the cluster, with necessary metadata
of the cloud instance.
"""
# The cloud instance id.
cloud_instance_id: CloudInstanceId
# The node type of the cloud instance.
node_type: NodeType
# The node kind, i.e head or worker.
node_kind: NodeKind
# If the cloud instance is already running.
is_running: bool
# Update request id from which the cloud instance is launched.
# This could be None if the cloud instance couldn't be associated with requests
# by the cloud provider: e.g. cloud provider doesn't support per-instance
# extra metadata.
# This is fine for now since the reconciler should be able to know how
# to handle cloud instances w/o request ids.
# TODO: make this a required field.
request_id: Optional[str] = None
class CloudInstanceProviderError(Exception):
"""
An base error class that represents an error that happened in the cloud instance
provider.
"""
# The timestamp of the error occurred in nanoseconds.
timestamp_ns: int
def __init__(self, msg, timestamp_ns) -> None:
super().__init__(msg)
self.timestamp_ns = timestamp_ns
class LaunchNodeError(CloudInstanceProviderError):
# The node type that failed to launch.
node_type: NodeType
# Number of nodes that failed to launch.
count: int
# A unique id that identifies from which update request the error originates.
request_id: str
def __init__(
self,
node_type: NodeType,
count: int,
request_id: str,
timestamp_ns: int,
details: str = "",
cause: Optional[Exception] = None,
) -> None:
msg = (
f"Failed to launch {count} nodes of type {node_type} with "
f"request id {request_id}: {details}"
)
super().__init__(msg, timestamp_ns=timestamp_ns)
self.node_type = node_type
self.count = count
self.request_id = request_id
if cause:
self.__cause__ = cause
def __repr__(self) -> str:
return (
f"LaunchNodeError(node_type={self.node_type}, count={self.count}, "
f"request_id={self.request_id}): {self.__cause__}"
)
class TerminateNodeError(CloudInstanceProviderError):
# The cloud instance id of the node that failed to terminate.
cloud_instance_id: CloudInstanceId
# A unique id that identifies from which update request the error originates.
request_id: str
def __init__(
self,
cloud_instance_id: CloudInstanceId,
request_id: str,
timestamp_ns: int,
details: str = "",
cause: Optional[Exception] = None,
) -> None:
msg = (
f"Failed to terminate node {cloud_instance_id} with "
f"request id {request_id}: {details}"
)
super().__init__(msg, timestamp_ns=timestamp_ns)
self.cloud_instance_id = cloud_instance_id
self.request_id = request_id
if cause:
self.__cause__ = cause
def __repr__(self) -> str:
return (
f"TerminateNodeError(cloud_instance_id={self.cloud_instance_id}, "
f"request_id={self.request_id}): {self.__cause__}"
)
class ICloudInstanceProvider(ABC):
"""
The interface for a cloud instance provider.
This interface is a minimal interface that should be implemented by the
various cloud instance providers (e.g. AWS, and etc).
The cloud instance provider is responsible for managing the cloud instances in the
cluster. It provides the following main functionalities:
- Launch new cloud instances.
- Terminate existing running instances.
- Get the non-terminated cloud instances in the cluster.
- Poll the errors that happened for the updates to the cloud instance provider.
Below properties of the cloud instance provider are assumed with this interface:
1. Eventually consistent
The cloud instance provider is expected to be eventually consistent with the
cluster state. For example, when a cloud instance is request to be terminated
or launched, the provider may not immediately reflect the change in its state.
However, the provider is expected to eventually reflect the change in its state.
2. Asynchronous
The provider could also be asynchronous, where the termination/launch
request may not immediately return the result of the request.
3. Unique cloud instance ids
Cloud instance ids are expected to be unique across the cluster.
4. Idempotent updates
For the update APIs (e.g. ensure_min_nodes, terminate), the provider may use the
request ids to provide idempotency.
Usage:
```
provider: ICloudInstanceProvider = ...
# Update the cluster with a desired shape.
provider.launch(
shape={
"worker_nodes": 10,
"ray_head": 1,
},
request_id="1",
)
# Get the non-terminated nodes of the cloud instance provider.
running = provider.get_non_terminated()
# Poll the errors
errors = provider.poll_errors()
# Terminate nodes.
provider.terminate(
ids=["cloud_instance_id_1", "cloud_instance_id_2"],
request_id="2",
)
# Process the state of the provider.
...
```
"""
@abstractmethod
def get_non_terminated(self) -> Dict[CloudInstanceId, CloudInstance]:
"""Get the non-terminated cloud instances in the cluster.
Returns:
A dictionary of the non-terminated cloud instances in the cluster.
The key is the cloud instance id, and the value is the cloud instance.
"""
pass
@abstractmethod
def terminate(self, ids: List[CloudInstanceId], request_id: str) -> None:
"""
Terminate the cloud instances asynchronously.
This method is expected to be idempotent, i.e. if the same request id is used
to terminate the same cloud instances, this should be a no-op if
the cloud instances are already terminated or being terminated.
Args:
ids: the cloud instance ids to terminate.
request_id: a unique id that identifies the request.
"""
pass
@abstractmethod
def launch(
self,
shape: Dict[NodeType, int],
request_id: str,
) -> None:
"""Launch the cloud instances asynchronously.
Args:
shape: A map from node type to number of nodes to launch.
request_id: a unique id that identifies the update request.
"""
pass
@abstractmethod
def poll_errors(self) -> List[CloudInstanceProviderError]:
"""
Poll the errors that happened since the last poll.
This method would also clear the errors that happened since the last poll.
Returns:
The errors that happened since the last poll.
"""
pass
@dataclass(frozen=True)
class CloudInstanceLaunchRequest:
"""
The arguments to launch a node.
"""
# The node type to launch.
node_type: NodeType
# Number of nodes to launch.
count: int
# A unique id that identifies the request.
request_id: str
@dataclass(frozen=True)
class CloudInstanceTerminateRequest:
"""
The arguments to terminate a node.
"""
# The cloud instance id of the node to terminate.
cloud_instance_id: CloudInstanceId
# A unique id that identifies the request.
request_id: str
class NodeProviderAdapter(ICloudInstanceProvider):
"""
Warps a NodeProviderV1 to a ICloudInstanceProvider.
TODO(rickyx):
The current adapter right now consists of two sets of APIs:
- v1: the old APIs that are used by the autoscaler, where
we forward the calls to the NodeProviderV1.
- v2: the new APIs that are used by the autoscaler v2, this is
defined in the ICloudInstanceProvider interface.
We should eventually remove the v1 APIs and only use the v2 APIs.
It's currently left as a TODO since changing the v1 APIs would
requires a lot of changes in the cluster launcher codebase.
"""
def __init__(
self,
v1_provider: NodeProviderV1,
config_reader: IConfigReader,
max_launch_batch_per_type: int = AUTOSCALER_MAX_LAUNCH_BATCH,
max_concurrent_launches: int = AUTOSCALER_MAX_CONCURRENT_LAUNCHES,
) -> None:
"""Initialize the node provider adapter.
Args:
v1_provider: The v1 node provider to wrap.
config_reader: The config reader to read the autoscaling config.
max_launch_batch_per_type: The maximum number of nodes to launch per
node type in a single batch.
max_concurrent_launches: The maximum number of concurrent launches.
"""
super().__init__()
self._v1_provider = v1_provider
self._config_reader = config_reader
# Executor to async launching and terminating nodes.
self._main_executor = ThreadPoolExecutor(
max_workers=1, thread_name_prefix="ray::NodeProviderAdapter"
)
# v1 legacy rate limiting on the node provider launch calls.
self._max_launch_batch_per_type = max_launch_batch_per_type
max_batches = math.ceil(
max_concurrent_launches / float(max_launch_batch_per_type)
)
self._node_launcher_executors = ThreadPoolExecutor(
max_workers=max_batches,
thread_name_prefix="ray::NodeLauncherPool",
)
# Queue to retrieve new errors occur in the multi-thread executors
# temporarily.
self._errors_queue = Queue()
@property
def v1_provider(self) -> NodeProviderV1:
return self._v1_provider
def get_non_terminated(self) -> Dict[CloudInstanceId, CloudInstance]:
nodes = {}
cloud_instance_ids = self._v1_non_terminated_nodes({})
# Filter out nodes that are not running.
# This is efficient since the provider is expected to cache the
# running status of the nodes.
for cloud_instance_id in cloud_instance_ids:
node_tags = self._v1_node_tags(cloud_instance_id)
node_kind_tag = node_tags.get(TAG_RAY_NODE_KIND, NODE_KIND_UNMANAGED)
if node_kind_tag == NODE_KIND_UNMANAGED:
# Filter out unmanaged nodes.
continue
elif node_kind_tag == NODE_KIND_WORKER:
node_kind = NodeKind.WORKER
elif node_kind_tag == NODE_KIND_HEAD:
node_kind = NodeKind.HEAD
else:
raise ValueError(f"Invalid node kind: {node_kind_tag}")
nodes[cloud_instance_id] = CloudInstance(
cloud_instance_id=cloud_instance_id,
node_type=node_tags.get(TAG_RAY_USER_NODE_TYPE, ""),
is_running=self._v1_is_running(cloud_instance_id),
request_id=node_tags.get(TAG_RAY_LAUNCH_REQUEST, ""),
node_kind=node_kind,
)
return nodes
def poll_errors(self) -> List[CloudInstanceProviderError]:
errors = []
while not self._errors_queue.empty():
errors.append(self._errors_queue.get_nowait())
return errors
def launch(
self,
shape: Dict[NodeType, int],
request_id: str,
) -> None:
self._main_executor.submit(self._do_launch, shape, request_id)
def terminate(self, ids: List[CloudInstanceId], request_id: str) -> None:
self._main_executor.submit(self._do_terminate, ids, request_id)
###########################################
# Private APIs
###########################################
def _do_launch(
self,
shape: Dict[NodeType, int],
request_id: str,
) -> None:
"""
Launch the cloud instances by calling into the v1 base node provider.
Args:
shape: The requested to launch node type and number of nodes.
request_id: The request id that identifies the request.
"""
for node_type, count in shape.items():
# Keep submitting the launch requests to the launch pool in batches.
while count > 0:
to_launch = min(count, self._max_launch_batch_per_type)
self._node_launcher_executors.submit(
self._launch_nodes_by_type,
node_type,
to_launch,
request_id,
)
count -= to_launch
def _do_terminate(self, ids: List[CloudInstanceId], request_id: str) -> None:
"""
Terminate the cloud instances by calling into the v1 base node provider.
If errors happen during the termination, the errors will be put into the
errors queue.
Args:
ids: The cloud instance ids to terminate.
request_id: The request id that identifies the request.
"""
try:
self._v1_terminate_nodes(ids)
except Exception as e:
for id in ids:
error = TerminateNodeError(id, request_id, int(time.time_ns()))
error.__cause__ = e
self._errors_queue.put(error)
def _launch_nodes_by_type(
self,
node_type: NodeType,
count: int,
request_id: str,
) -> None:
"""
Launch nodes of the given node type.
Args:
node_type: The node type to launch.
count: Number of nodes to launch.
request_id: A unique id that identifies the request.
Raises:
ValueError: If the node type is invalid.
LaunchNodeError: If the launch failed and raised by the underlying provider.
"""
# Check node type is valid.
try:
config = self._config_reader.get_cached_autoscaling_config()
launch_config = config.get_cloud_node_config(node_type)
resources = config.get_node_resources(node_type)
labels = config.get_node_labels(node_type)
# This is to be compatible with the v1 node launcher.
# See more in https://github.com/ray-project/ray/blob/6f5a189bc463e52c51a70f8aea41fb2950b443e8/python/ray/autoscaler/_private/node_launcher.py#L78-L85 # noqa
# TODO: this should be synced with what's stored in the IM, it should
# probably be made as a metadata field in the cloud instance. This is
# another incompatibility with KubeRay.
launch_hash = hash_launch_conf(launch_config, config.get_config("auth", {}))
node_tags = {
TAG_RAY_NODE_NAME: "ray-{}-worker".format(
config.get_config("cluster_name", "")
),
TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
TAG_RAY_NODE_STATUS: STATUS_UNINITIALIZED,
TAG_RAY_LAUNCH_CONFIG: launch_hash,
TAG_RAY_LAUNCH_REQUEST: request_id,
TAG_RAY_USER_NODE_TYPE: node_type,
}
logger.info("Launching {} nodes of type {}.".format(count, node_type))
self._v1_provider.create_node_with_resources_and_labels(
launch_config, node_tags, count, resources, labels
)
logger.info("Launched {} nodes of type {}.".format(count, node_type))
except Exception as e:
logger.info(
"Failed to launch {} nodes of type {}: {}".format(count, node_type, e)
)
error = LaunchNodeError(node_type, count, request_id, int(time.time_ns()))
error.__cause__ = e
self._errors_queue.put(error)
###########################################
# V1 Legacy APIs
###########################################
"""
Below are the necessary legacy APIs from the V1 node provider.
These are needed as of now to provide the needed features
for V2 node provider.
The goal is to eventually remove these APIs and only use the
V2 APIs by modifying the individual node provider to inherit
from ICloudInstanceProvider.
"""
def _v1_terminate_nodes(
self, ids: List[CloudInstanceId]
) -> Optional[Dict[str, Any]]:
return self._v1_provider.terminate_nodes(ids)
def _v1_non_terminated_nodes(
self, tag_filters: Dict[str, str]
) -> List[CloudInstanceId]:
return self._v1_provider.non_terminated_nodes(tag_filters)
def _v1_is_running(self, node_id: CloudInstanceId) -> bool:
return self._v1_provider.is_running(node_id)
def _v1_post_process(self) -> None:
self._v1_provider.post_process()
def _v1_node_tags(self, node_id: CloudInstanceId) -> Dict[str, str]:
return self._v1_provider.node_tags(node_id)
def _v1_safe_to_scale(self) -> bool:
return self._v1_provider.safe_to_scale()
@@ -0,0 +1,96 @@
import logging
import subprocess
from ray.autoscaler._private.updater import (
STATUS_UP_TO_DATE,
TAG_RAY_NODE_STATUS,
NodeUpdater,
)
from ray.autoscaler._private.util import with_envs, with_head_node_ip
from ray.autoscaler.node_provider import NodeProvider as NodeProviderV1
from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
from ray.core.generated.instance_manager_pb2 import Instance
logger = logging.getLogger(__name__)
class RayInstaller(object):
"""
RayInstaller is responsible for installing ray on the target instance.
"""
def __init__(
self,
provider: NodeProviderV1,
config: AutoscalingConfig,
process_runner=subprocess,
) -> None:
self._provider = provider
self._config = config
self._process_runner = process_runner
def install_ray(self, instance: Instance, head_node_ip: str) -> None:
"""
Install ray on the target instance synchronously.
TODO:(rickyx): This runs in another thread, and errors are silently
ignored. We should propagate the error to the main thread.
"""
setup_commands = self._config.get_worker_setup_commands(instance.instance_type)
ray_start_commands = self._config.get_worker_start_ray_commands()
docker_config = self._config.get_docker_config(instance.instance_type)
logger.info(
f"Creating new (spawn_updater) updater thread for node"
f" {instance.cloud_instance_id}."
)
provider_instance_type_name = self._config.get_provider_instance_type(
instance.instance_type
)
updater = NodeUpdater(
node_id=instance.cloud_instance_id,
provider_config=self._config.get_config("provider"),
provider=self._provider,
auth_config=self._config.get_config("auth"),
cluster_name=self._config.get_config("cluster_name"),
file_mounts=self._config.get_config("file_mounts"),
initialization_commands=with_head_node_ip(
self._config.get_initialization_commands(instance.instance_type),
head_node_ip,
),
setup_commands=with_head_node_ip(setup_commands, head_node_ip),
# This will prepend envs to the begin of the ray start commands, e.g.
# `RAY_HEAD_IP=<head_node_ip> \
# RAY_CLOUD_INSTANCE_ID=<instance_id> \
# ray start --head ...`
# See src/ray/common/constants.h for ENV name definitions.
ray_start_commands=with_envs(
ray_start_commands,
{
"RAY_HEAD_IP": head_node_ip,
"RAY_CLOUD_INSTANCE_ID": instance.cloud_instance_id,
"RAY_NODE_TYPE_NAME": instance.instance_type,
"RAY_CLOUD_INSTANCE_TYPE_NAME": provider_instance_type_name,
},
),
runtime_hash=self._config.runtime_hash,
file_mounts_contents_hash=self._config.file_mounts_contents_hash,
is_head_node=False,
cluster_synced_files=self._config.get_config("cluster_synced_files"),
rsync_options={
"rsync_exclude": self._config.get_config("rsync_exclude"),
"rsync_filter": self._config.get_config("rsync_filter"),
},
use_internal_ip=True,
docker_config=docker_config,
node_resources=self._config.get_node_resources(instance.instance_type),
node_labels=self._config.get_node_labels(instance.instance_type),
process_runner=self._process_runner,
)
updater.run()
# check if the updater was successful by checking the node tags
# since the updater could hide exceptions and just set the status tag
tags = self._provider.node_tags(instance.cloud_instance_id)
if tags.get(TAG_RAY_NODE_STATUS) != STATUS_UP_TO_DATE:
raise Exception(
f"Ray installation failed with unexpected status: {tags.get(TAG_RAY_NODE_STATUS)}"
)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,180 @@
import copy
from abc import ABCMeta, abstractmethod
from collections import defaultdict, namedtuple
from threading import Lock
from typing import Dict, List, Optional, Tuple
StoreStatus = namedtuple("StoreStatus", ["success", "version"])
VersionedValue = namedtuple("VersionedValue", ["value", "version"])
class Storage(metaclass=ABCMeta):
"""Interface for a storage backend that stores the state of nodes in the cluster.
The storage is thread-safe.
The storage is versioned, which means that each successful stage change to the
storage will bump the version number. The version number can be used to
implement optimistic concurrency control.
Each entry in the storage table is also versioned. The version number of an entry
is the last version number when the entry is updated.
"""
@abstractmethod
def batch_update(
self,
table: str,
mutation: Optional[Dict[str, str]] = None,
deletion: Optional[List[str]] = None,
expected_storage_version: Optional[int] = None,
) -> StoreStatus:
"""Batch update the storage table. This method is atomic.
Args:
table: The name of the table.
mutation: A dictionary of key-value pairs to be updated.
deletion: A list of keys to be deleted.
expected_storage_version: The expected storage version. The
update will fail if the version does not match the
current storage version.
Returns:
StoreStatus: A tuple of (success, version). If the update is
successful, returns (True, new_version).
Otherwise, returns (False, current_version).
"""
raise NotImplementedError("batch_update() has to be implemented")
@abstractmethod
def update(
self,
table: str,
key: str,
value: str,
expected_entry_version: Optional[int] = None,
insert_only: bool = False,
) -> StoreStatus:
"""Update a single entry in the storage table.
Args:
table: The name of the table.
key: The key of the entry.
value: The value of the entry.
expected_entry_version: The expected version of the entry.
The update will fail if the version does not match the current
version of the entry.
insert_only: If True, the update will
fail if the entry already exists.
Returns:
StoreStatus: A tuple of (success, version). If the update is
successful, returns (True, new_version). Otherwise,
returns (False, current_version).
"""
raise NotImplementedError("update() has to be implemented")
@abstractmethod
def get_all(self, table: str) -> Tuple[Dict[str, Tuple[str, int]], int]:
raise NotImplementedError("get_all() has to be implemented")
@abstractmethod
def get(
self, table: str, keys: List[str]
) -> Tuple[Dict[str, Tuple[str, int]], int]:
"""Get a list of entries from the storage table.
Args:
table: The name of the table.
keys: A list of keys to be retrieved. If the list is empty,
all entries in the table will be returned.
Returns:
Tuple[Dict[str, VersionedValue], int]: A tuple of
(entries, storage_version). The entries is a dictionary of
(key, (value, entry_version)) pairs. The entry_version is the
version of the entry when it was last updated. The
storage_version is the current storage version.
"""
raise NotImplementedError("get() has to be implemented")
@abstractmethod
def get_version(self) -> int:
"""Get the current storage version.
Returns:
int: The current storage version.
"""
raise NotImplementedError("get_version() has to be implemented")
class InMemoryStorage(Storage):
"""An in-memory implementation of the Storage interface. This implementation
is not durable"""
def __init__(self):
self._version = 0
self._tables = defaultdict(dict)
self._lock = Lock()
def batch_update(
self,
table: str,
mutation: Dict[str, str] = None,
deletion: List[str] = None,
expected_version: Optional[int] = None,
) -> StoreStatus:
mutation = mutation if mutation else {}
deletion = deletion if deletion else []
with self._lock:
if expected_version is not None and expected_version != self._version:
return StoreStatus(False, self._version)
self._version += 1
key_value_pairs_with_version = {
key: VersionedValue(value, self._version)
for key, value in mutation.items()
}
self._tables[table].update(key_value_pairs_with_version)
for deleted_key in deletion:
self._tables[table].pop(deleted_key, None)
return StoreStatus(True, self._version)
def update(
self,
table: str,
key: str,
value: str,
expected_entry_version: Optional[int] = None,
expected_storage_version: Optional[int] = None,
insert_only: bool = False,
) -> StoreStatus:
with self._lock:
if (
expected_storage_version is not None
and expected_storage_version != self._version
):
return StoreStatus(False, self._version)
if insert_only and key in self._tables[table]:
return StoreStatus(False, self._version)
_, version = self._tables[table].get(key, (None, -1))
if expected_entry_version is not None and expected_entry_version != version:
return StoreStatus(False, self._version)
self._version += 1
self._tables[table][key] = VersionedValue(value, self._version)
return StoreStatus(True, self._version)
def get_all(self, table: str) -> Tuple[Dict[str, VersionedValue], int]:
with self._lock:
return (copy.deepcopy(self._tables[table]), self._version)
def get(self, table: str, keys: List[str]) -> Tuple[Dict[str, VersionedValue], int]:
if not keys:
return self.get_all(table)
with self._lock:
result = {}
for key in keys:
if key in self._tables.get(table, {}):
result[key] = self._tables[table][key]
return StoreStatus(result, self._version)
def get_version(self) -> int:
return self._version
@@ -0,0 +1,118 @@
import logging
import uuid
from collections import defaultdict
from typing import List, Optional
from ray.autoscaler.v2.instance_manager.instance_manager import (
InstanceUpdatedSubscriber,
)
from ray.autoscaler.v2.instance_manager.node_provider import ICloudInstanceProvider
from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
from ray.core.generated.instance_manager_pb2 import Instance, InstanceUpdateEvent
logger = logging.getLogger(__name__)
class CloudInstanceUpdater(InstanceUpdatedSubscriber):
"""CloudInstanceUpdater is responsible for launching
new instances and terminating cloud instances
It requests the cloud instance provider to launch new instances when
there are new instance requests (with REQUESTED status change).
It requests the cloud instance provider to terminate instances when
there are new instance terminations (with TERMINATING status change).
The cloud instance APIs are async and non-blocking.
"""
def __init__(
self,
cloud_provider: ICloudInstanceProvider,
metrics_reporter: Optional[AutoscalerMetricsReporter] = None,
) -> None:
self._cloud_provider = cloud_provider
self._metrics_reporter = metrics_reporter
def notify(self, events: List[InstanceUpdateEvent]) -> None:
new_requests = [
event for event in events if event.new_instance_status == Instance.REQUESTED
]
new_terminations = [
event
for event in events
if event.new_instance_status == Instance.TERMINATING
]
terminated_instances = [
event
for event in events
if event.new_instance_status == Instance.TERMINATED
and event.cloud_instance_id
]
self._launch_new_instances(new_requests)
self._terminate_instances(new_terminations)
self._count_stopped_instances(terminated_instances)
def _terminate_instances(self, new_terminations: List[InstanceUpdateEvent]) -> None:
"""
Terminate cloud instances through cloud provider.
Args:
new_terminations: List of new instance terminations.
"""
if not new_terminations:
logger.debug("No instances to terminate.")
return
# Terminate the instances.
cloud_instance_ids = [event.cloud_instance_id for event in new_terminations]
# This is an async call.
self._cloud_provider.terminate(
ids=cloud_instance_ids, request_id=str(uuid.uuid4())
)
def _count_stopped_instances(self, terminated_instances: List[InstanceUpdateEvent]):
"""
Record successfully terminated cloud instances.
Args:
terminated_instances: List of terminated cloud instances.
Returns:
None.
"""
if not terminated_instances or not self._metrics_reporter:
return
self._metrics_reporter.inc_stopped_nodes(len(terminated_instances))
def _launch_new_instances(self, new_requests: List[InstanceUpdateEvent]) -> None:
"""
Launches new instances by requesting the cloud provider.
Args:
new_requests: List of new instance requests.
"""
if not new_requests:
logger.debug("No instances to launch.")
return
# Group new requests by launch request id.
requests_by_launch_request_id = defaultdict(list)
for event in new_requests:
assert (
event.launch_request_id
), "Launch request id should have been set by the reconciler"
requests_by_launch_request_id[event.launch_request_id].append(event)
for launch_request_id, events in requests_by_launch_request_id.items():
request_shape = defaultdict(int)
for event in events:
request_shape[event.instance_type] += 1
# Make requests to the cloud provider.
self._cloud_provider.launch(
shape=request_shape, request_id=launch_request_id
)
@@ -0,0 +1,102 @@
import logging
import time
from typing import Dict, List
from ray.autoscaler._private.constants import RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S
from ray.autoscaler.v2.instance_manager.instance_manager import (
InstanceUpdatedSubscriber,
)
from ray.autoscaler.v2.schema import NodeType
from ray.core.generated.instance_manager_pb2 import Instance, InstanceUpdateEvent
logger = logging.getLogger(__name__)
class CloudResourceMonitor(InstanceUpdatedSubscriber):
"""CloudResourceMonitor records the availability of all node types.
In the Spot scenario, the resources in the cluster change dynamically.
When scaling up, it is necessary to know which node types are most
likely to have resources, in order to decide which type of node to request.
During scaling up, if resource of a node type is requested but fail to
allocate, that type is considered unavailable at that timestamp.This class
records the last timestamp at which a node type is unavailable,allowing the
autoscaler to skip such node types when making future scaling decisions.
"""
def __init__(
self,
) -> None:
self._last_unavailable_timestamp: Dict[NodeType, float] = {}
def allocation_timeout(self, failed_event: InstanceUpdateEvent):
unavailable_timestamp = time.time()
self._last_unavailable_timestamp[
failed_event.instance_type
] = unavailable_timestamp
logger.info(
f"Cloud Resource Type {failed_event.instance_type} is "
f"unavailable at timestamp={unavailable_timestamp}. "
f"We will lower its priority in feature schedules."
)
def allocation_succeeded(self, succeeded_event: InstanceUpdateEvent):
if succeeded_event.instance_type in self._last_unavailable_timestamp:
self._last_unavailable_timestamp.pop(succeeded_event.instance_type)
logger.info(
f"Cloud Resource Type {succeeded_event.instance_type} is "
f"available at timestamp={time.time()}. We will higher its priority in "
f"feature schedules."
)
def notify(self, events: List[InstanceUpdateEvent]) -> None:
for event in events:
if event.new_instance_status == Instance.ALLOCATION_TIMEOUT:
self.allocation_timeout(event)
elif (
event.new_instance_status == Instance.RAY_RUNNING
and event.instance_type
):
self.allocation_succeeded(event)
def get_resource_availabilities(self) -> Dict[NodeType, float]:
"""Calculate the availability scores of node types.
Higher values indicate a higher likelihood of resource allocation.
"""
resource_availability_scores: Dict[NodeType, float] = {}
if self._last_unavailable_timestamp:
max_ts = max(self._last_unavailable_timestamp.values())
for node_type in self._last_unavailable_timestamp:
resource_availability_scores[node_type] = (
1 - self._last_unavailable_timestamp[node_type] / max_ts
)
return resource_availability_scores
def get_recoverable_resource_availabilities(self) -> Dict[NodeType, float]:
"""Calculate a continuous recovery score from 0.0 to 1.0.
score = 0.0 if (current_time - last_unavailable_timestamp) < safety_floor
else min(1.0, (current_time - last_unavailable_timestamp) /
RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S)
"""
assert (
RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S > 0
), "RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S must be positive"
recovery_scores: Dict[NodeType, float] = {}
current_time = time.time()
# Safety floor is 10s or 10% of recovery window.
# This ensures that we don't immediately retry a failed node type
# and be stuck in a retry loop.
safety_floor = min(10, RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S * 0.1)
for node_type, last_ts in self._last_unavailable_timestamp.items():
diff = current_time - last_ts
if diff < safety_floor:
recovery_scores[node_type] = 0.0
else:
recovery_scores[node_type] = min(
1.0, diff / RAY_AUTOSCALER_AVAILABILITY_RECOVERY_S
)
return recovery_scores
@@ -0,0 +1,158 @@
import logging
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from queue import Queue
from typing import List
from ray._common.utils import hex_to_binary
from ray._raylet import GcsClient
from ray.autoscaler.v2.instance_manager.instance_manager import (
InstanceUpdatedSubscriber,
)
from ray.core.generated.autoscaler_pb2 import DrainNodeReason
from ray.core.generated.instance_manager_pb2 import (
Instance,
InstanceUpdateEvent,
TerminationRequest,
)
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class RayStopError:
# Instance manager's instance id.
im_instance_id: str
class RayStopper(InstanceUpdatedSubscriber):
"""RayStopper is responsible for stopping ray on instances.
It will drain the ray node if it's for idle termination.
For other terminations, it will stop the ray node. (e.g. scale down, etc.)
If any failures happen when stopping/draining the node, we will not retry
and rely on the reconciler to handle the failure.
TODO: we could also surface the errors back to the reconciler for
quicker failure detection.
"""
def __init__(self, gcs_client: GcsClient, error_queue: Queue) -> None:
self._gcs_client = gcs_client
self._error_queue = error_queue
self._executor = ThreadPoolExecutor(max_workers=1)
def notify(self, events: List[InstanceUpdateEvent]) -> None:
for event in events:
if event.new_instance_status == Instance.RAY_STOP_REQUESTED:
fut = self._executor.submit(self._stop_or_drain_ray, event)
def _log_on_error(fut):
try:
fut.result()
except Exception:
logger.exception("Error stopping/drain ray.")
fut.add_done_callback(_log_on_error)
def _stop_or_drain_ray(self, event: InstanceUpdateEvent) -> None:
"""
Stops or drains the ray node based on the termination request.
"""
assert event.HasField("termination_request"), "Termination request is required."
termination_request = event.termination_request
ray_node_id = termination_request.ray_node_id
instance_id = event.instance_id
if termination_request.cause == TerminationRequest.Cause.IDLE:
reason = DrainNodeReason.DRAIN_NODE_REASON_IDLE_TERMINATION
reason_str = "Termination of node that's idle for {} seconds.".format(
termination_request.idle_duration_ms / 1000
)
self._drain_ray_node(
self._gcs_client,
self._error_queue,
ray_node_id,
instance_id,
reason,
reason_str,
)
return
# If it's not an idle termination, we stop the ray node.
self._stop_ray_node(
self._gcs_client, self._error_queue, ray_node_id, instance_id
)
@staticmethod
def _drain_ray_node(
gcs_client: GcsClient,
error_queue: Queue,
ray_node_id: str,
instance_id: str,
reason: DrainNodeReason,
reason_str: str,
):
"""
Drains the ray node.
Args:
gcs_client: The gcs client to use.
error_queue: Queue to put errors on when draining fails.
ray_node_id: The ray node id to drain.
instance_id: The instance id corresponding to the ray node.
reason: The reason to drain the node.
reason_str: The reason message to drain the node.
"""
try:
accepted, reject_msg_str = gcs_client.drain_node(
node_id=ray_node_id,
reason=reason,
reason_message=reason_str,
# TODO: we could probably add a deadline here that's derived
# from the stuck instance reconciliation configs.
deadline_timestamp_ms=0,
)
logger.info(
f"Drained ray on {ray_node_id}(success={accepted}, "
f"msg={reject_msg_str})"
)
if not accepted:
error_queue.put_nowait(RayStopError(im_instance_id=instance_id))
except Exception:
logger.exception(f"Error draining ray on {ray_node_id}")
error_queue.put_nowait(RayStopError(im_instance_id=instance_id))
@staticmethod
def _stop_ray_node(
gcs_client: GcsClient,
error_queue: Queue,
ray_node_id: str,
instance_id: str,
):
"""
Stops the ray node.
Args:
gcs_client: The gcs client to use.
error_queue: Queue to put errors on when stopping fails.
ray_node_id: The ray node id to stop.
instance_id: The instance id corresponding to the ray node.
"""
try:
drained = gcs_client.drain_nodes(node_ids=[hex_to_binary(ray_node_id)])
success = len(drained) > 0
logger.info(
f"Stopping ray on {ray_node_id}(instance={instance_id}): "
f"success={success})"
)
if not success:
error_queue.put_nowait(RayStopError(im_instance_id=instance_id))
except Exception:
logger.exception(
f"Error stopping ray on {ray_node_id}(instance={instance_id})"
)
error_queue.put_nowait(RayStopError(im_instance_id=instance_id))
@@ -0,0 +1,95 @@
import dataclasses
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from queue import Queue
from typing import List
from ray.autoscaler.v2.instance_manager.instance_manager import (
InstanceUpdatedSubscriber,
)
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
from ray.core.generated.instance_manager_pb2 import (
Instance,
InstanceUpdateEvent,
NodeKind,
)
logger = logging.getLogger(__name__)
@dataclasses.dataclass(frozen=True)
class RayInstallError:
# Instance manager's instance id.
im_instance_id: str
# Error details.
details: str
class ThreadedRayInstaller(InstanceUpdatedSubscriber):
"""ThreadedRayInstaller is responsible for install ray on new nodes."""
def __init__(
self,
head_node_ip: str,
instance_storage: InstanceStorage,
ray_installer: RayInstaller,
error_queue: Queue,
max_install_attempts: int = 3,
install_retry_interval: int = 10,
max_concurrent_installs: int = 50,
) -> None:
self._head_node_ip = head_node_ip
self._instance_storage = instance_storage
self._ray_installer = ray_installer
self._max_concurrent_installs = max_concurrent_installs
self._max_install_attempts = max_install_attempts
self._install_retry_interval = install_retry_interval
self._error_queue = error_queue
self._ray_installation_executor = ThreadPoolExecutor(
max_workers=self._max_concurrent_installs
)
def notify(self, events: List[InstanceUpdateEvent]) -> None:
for event in events:
if event.new_instance_status == Instance.RAY_INSTALLING:
self._install_ray_on_new_nodes(event.instance_id)
def _install_ray_on_new_nodes(self, instance_id: str) -> None:
allocated_instance, _ = self._instance_storage.get_instances(
instance_ids={instance_id},
status_filter={Instance.RAY_INSTALLING},
)
for instance in allocated_instance.values():
assert instance.node_kind == NodeKind.WORKER
self._ray_installation_executor.submit(
self._install_ray_on_single_node, instance
)
def _install_ray_on_single_node(self, instance: Instance) -> None:
assert instance.status == Instance.RAY_INSTALLING
# install with exponential backoff
backoff_factor = 1
last_exception = None
for _ in range(self._max_install_attempts):
try:
self._ray_installer.install_ray(instance, self._head_node_ip)
return
except Exception as e:
logger.info(
f"Ray installation failed on instance {instance.cloud_instance_id}: {e}"
)
last_exception = e
logger.warning("Failed to install ray, retrying...")
time.sleep(self._install_retry_interval * backoff_factor)
backoff_factor *= 2
self._error_queue.put_nowait(
RayInstallError(
im_instance_id=instance.instance_id,
details=str(last_exception),
)
)
@@ -0,0 +1,110 @@
from collections import defaultdict
from typing import Dict, List
from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
from ray.autoscaler.v2.instance_manager.common import InstanceUtil
from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
from ray.autoscaler.v2.schema import NodeType
from ray.core.generated.instance_manager_pb2 import Instance as IMInstance
class AutoscalerMetricsReporter:
def __init__(self, prom_metrics: AutoscalerPrometheusMetrics) -> None:
self._prom_metrics = prom_metrics
def inc_stopped_nodes(self, count: int) -> None:
self._prom_metrics.stopped_nodes.inc(count)
def report_instances(
self,
instances: List[IMInstance],
node_type_configs: Dict[NodeType, NodeTypeConfig],
):
"""
Record autoscaler metrics for:
- pending_nodes: Nodes that are launching/pending ray start
- active_nodes: Active nodes (nodes running ray)
- recently_failed_nodes: Nodes that are being terminated.
"""
# map of instance type to a dict of status to count.
status_count_by_type: Dict[NodeType, Dict[str, int]] = {}
def _new_status_count() -> Dict[str, int]:
return {
"pending": 0,
"running": 0,
"terminating": 0,
"terminated": 0,
}
# initialize the status count by type.
for instance_type in node_type_configs.keys():
status_count_by_type[instance_type] = _new_status_count()
for instance in instances:
status_count = status_count_by_type.get(instance.instance_type)
if status_count is None:
status_count = _new_status_count()
status_count_by_type[instance.instance_type] = status_count
if InstanceUtil.is_ray_pending(instance.status):
status_count["pending"] += 1
elif InstanceUtil.is_ray_running(instance.status):
status_count["running"] += 1
elif instance.status == IMInstance.TERMINATING:
status_count["terminating"] += 1
elif instance.status == IMInstance.TERMINATED:
status_count["terminated"] += 1
for instance_type, status_count in status_count_by_type.items():
self._prom_metrics.pending_nodes.labels(
SessionName=self._prom_metrics.session_name, NodeType=instance_type
).set(status_count["pending"])
self._prom_metrics.active_nodes.labels(
SessionName=self._prom_metrics.session_name, NodeType=instance_type
).set(status_count["running"])
self._prom_metrics.recently_failed_nodes.labels(
SessionName=self._prom_metrics.session_name, NodeType=instance_type
).set(status_count["terminating"])
def report_resources(
self,
instances: List[IMInstance],
node_type_configs: Dict[NodeType, NodeTypeConfig],
):
"""
Record autoscaler metrics for:
- pending_resources: Pending resources
- cluster_resources: Cluster resources (resources running on the cluster)
"""
# pending resources.
pending_resources = defaultdict(float)
cluster_resources = defaultdict(float)
def _add_resources(resource_map, node_type_configs, node_type, count):
node_resources = node_type_configs[node_type].resources
for resource_name, resource_value in node_resources.items():
resource_map[resource_name] += resource_value * count
for instance in instances:
if instance.instance_type not in node_type_configs:
continue
if InstanceUtil.is_ray_pending(instance.status):
_add_resources(
pending_resources, node_type_configs, instance.instance_type, 1
)
elif InstanceUtil.is_ray_running(instance.status):
_add_resources(
cluster_resources, node_type_configs, instance.instance_type, 1
)
for resource_name, resource_value in pending_resources.items():
self._prom_metrics.pending_resources.labels(
SessionName=self._prom_metrics.session_name, resource=resource_name
).set(resource_value)
for resource_name, resource_value in cluster_resources.items():
self._prom_metrics.cluster_resources.labels(
SessionName=self._prom_metrics.session_name, resource=resource_name
).set(resource_value)
+348
View File
@@ -0,0 +1,348 @@
"""Autoscaler monitoring loop daemon.
See autoscaler._private/monitor.py for the legacy implementation. All the legacy flags
are supported here, but the new implementation uses the new autoscaler v2.
"""
import argparse
import logging
import os
import sys
import time
from typing import Optional
import ray
import ray._private.ray_constants as ray_constants
from ray._common.network_utils import (
build_address,
get_localhost_ip,
is_localhost,
parse_address,
)
from ray._common.ray_constants import (
LOGGING_ROTATE_BACKUP_COUNT,
LOGGING_ROTATE_BYTES,
)
from ray._common.usage.usage_lib import record_extra_usage_tag
from ray._private import logging_utils
from ray._private.event.event_logger import get_event_logger
from ray._private.ray_logging import setup_component_logger
from ray._private.worker import SCRIPT_MODE
from ray._raylet import GcsClient
from ray.autoscaler._private.constants import (
AUTOSCALER_METRIC_PORT,
AUTOSCALER_UPDATE_INTERVAL_S,
)
from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
from ray.autoscaler.v2.autoscaler import Autoscaler
from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
from ray.autoscaler.v2.instance_manager.config import (
FileConfigReader,
IConfigReader,
Provider,
ReadOnlyProviderConfigReader,
)
from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
from ray.core.generated.autoscaler_pb2 import AutoscalingState
from ray.core.generated.event_pb2 import Event as RayEvent
from ray.core.generated.usage_pb2 import TagKey
try:
import prometheus_client
except ImportError:
prometheus_client = None
logger = logging.getLogger(__name__)
class AutoscalerMonitor:
"""Autoscaling monitor.
This process periodically collects stats from the GCS and triggers
autoscaler updates.
TODO:
We should also handle autoscaler failures properly in the future.
Right now, we don't restart autoscaler if it fails (internal reconciliation
however, should not fail the autoscaler process).
With the Reconciler able to handle extra cloud instances, we could in fact
recover the autoscaler process from reconciliation.
"""
def __init__(
self,
address: str,
config_reader: IConfigReader,
log_dir: Optional[str] = None,
monitor_ip: Optional[str] = None,
):
# Record v2 usage (we do this as early as possible to capture usage)
record_autoscaler_v2_usage(GcsClient(address))
self.gcs_address = address
worker = ray._private.worker.global_worker
# TODO: eventually plumb ClusterID through to here
self.gcs_client = GcsClient(address=self.gcs_address)
if monitor_ip:
monitor_addr = build_address(monitor_ip, AUTOSCALER_METRIC_PORT)
self.gcs_client.internal_kv_put(
b"AutoscalerMetricsAddress", monitor_addr.encode(), True, None
)
self._session_name = self._get_session_name(self.gcs_client)
logger.info(f"session_name: {self._session_name}")
worker.set_mode(SCRIPT_MODE)
head_node_ip = parse_address(self.gcs_address)[0]
self.autoscaler = None
if log_dir:
try:
ray_event_logger = get_event_logger(
RayEvent.SourceType.AUTOSCALER, log_dir
)
self.event_logger = AutoscalerEventLogger(
ray_event_logger,
log_cluster_shape=config_reader.get_cached_autoscaling_config().provider
!= Provider.READ_ONLY,
)
except Exception:
self.event_logger = None
else:
self.event_logger = None
prom_metrics = AutoscalerPrometheusMetrics(session_name=self._session_name)
self.metric_reporter = AutoscalerMetricsReporter(prom_metrics)
if monitor_ip and prometheus_client:
# If monitor_ip wasn't passed in, then don't attempt to start the
# metric server to keep behavior identical to before metrics were
# introduced
try:
logger.info(
"Starting autoscaler metrics server on port {}".format(
AUTOSCALER_METRIC_PORT
)
)
kwargs = (
{"addr": get_localhost_ip()} if is_localhost(head_node_ip) else {}
)
prometheus_client.start_http_server(
port=AUTOSCALER_METRIC_PORT,
registry=prom_metrics.registry,
**kwargs,
)
except Exception:
logger.exception(
"An exception occurred while starting the metrics server."
)
elif not prometheus_client:
logger.warning(
"`prometheus_client` not found, so metrics will not be exported."
)
self.autoscaler = Autoscaler(
session_name=self._session_name,
config_reader=config_reader,
gcs_client=self.gcs_client,
event_logger=self.event_logger,
metrics_reporter=self.metric_reporter,
)
@staticmethod
def _get_session_name(gcs_client: GcsClient) -> Optional[str]:
"""Obtain the session name from the GCS.
If the GCS doesn't respond, session name is considered None.
In this case, the metrics reported from the monitor won't have
the correct session name.
"""
session_name = gcs_client.internal_kv_get(
b"session_name",
ray_constants.KV_NAMESPACE_SESSION,
timeout=10,
)
if session_name:
session_name = session_name.decode()
return session_name
@staticmethod
def _report_autoscaling_state(
gcs_client: GcsClient, autoscaling_state: AutoscalingState
):
"""Report the autoscaling state to the GCS."""
try:
gcs_client.report_autoscaling_state(autoscaling_state.SerializeToString())
except Exception:
logger.exception("Error reporting autoscaling state to GCS.")
def _run(self):
"""Run the monitor loop."""
while True:
autoscaling_state = self.autoscaler.update_autoscaling_state()
if autoscaling_state:
# report autoscaling state
self._report_autoscaling_state(self.gcs_client, autoscaling_state)
else:
logger.warning("No autoscaling state to report.")
# Wait for a autoscaler update interval before processing the next
# round of messages.
time.sleep(AUTOSCALER_UPDATE_INTERVAL_S)
def run(self):
try:
self._run()
except Exception:
logger.exception("Error in monitor loop")
raise
def record_autoscaler_v2_usage(gcs_client: GcsClient) -> None:
"""
Record usage for autoscaler v2.
"""
try:
record_extra_usage_tag(TagKey.AUTOSCALER_VERSION, "v2", gcs_client)
except Exception:
logger.exception("Error recording usage for autoscaler v2.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description=("Parse GCS server for the monitor to connect to.")
)
parser.add_argument(
"--gcs-address", required=False, type=str, help="The address (ip:port) of GCS."
)
parser.add_argument(
"--autoscaling-config",
required=False,
type=str,
help="the path to the autoscaling config file",
)
parser.add_argument(
"--logging-level",
required=False,
type=str,
default=ray_constants.LOGGER_LEVEL,
choices=ray_constants.LOGGER_LEVEL_CHOICES,
help=ray_constants.LOGGER_LEVEL_HELP,
)
parser.add_argument(
"--logging-format",
required=False,
type=str,
default=ray_constants.LOGGER_FORMAT,
help=ray_constants.LOGGER_FORMAT_HELP,
)
parser.add_argument(
"--logging-filename",
required=False,
type=str,
default=ray_constants.MONITOR_LOG_FILE_NAME,
help="Specify the name of log file, "
"log to stdout if set empty, default is "
f'"{ray_constants.MONITOR_LOG_FILE_NAME}"',
)
parser.add_argument(
"--logs-dir",
required=True,
type=str,
help="Specify the path of the temporary directory used by Ray processes.",
)
parser.add_argument(
"--logging-rotate-bytes",
required=False,
type=int,
default=LOGGING_ROTATE_BYTES,
help="Specify the max bytes for rotating "
"log file, default is "
f"{LOGGING_ROTATE_BYTES} bytes.",
)
parser.add_argument(
"--logging-rotate-backup-count",
required=False,
type=int,
default=LOGGING_ROTATE_BACKUP_COUNT,
help="Specify the backup count of rotated log file, default is "
f"{LOGGING_ROTATE_BACKUP_COUNT}.",
)
parser.add_argument(
"--monitor-ip",
required=False,
type=str,
default=None,
help="The IP address of the machine hosting the monitor process.",
)
parser.add_argument(
"--stdout-filepath",
required=False,
type=str,
default="",
help="The filepath to dump monitor stdout.",
)
parser.add_argument(
"--stderr-filepath",
required=False,
type=str,
default="",
help="The filepath to dump monitor stderr.",
)
args = parser.parse_args()
# Disable log rotation for windows platform.
logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0
logging_rotation_backup_count = (
args.logging_rotate_backup_count if sys.platform != "win32" else 1
)
setup_component_logger(
logging_level=args.logging_level,
logging_format=args.logging_format,
log_dir=args.logs_dir,
filename=args.logging_filename,
max_bytes=logging_rotation_bytes,
backup_count=logging_rotation_backup_count,
)
# Setup stdout/stderr redirect files if redirection enabled.
logging_utils.redirect_stdout_stderr_if_needed(
args.stdout_filepath,
args.stderr_filepath,
logging_rotation_bytes,
logging_rotation_backup_count,
)
logger.info(
f"Starting autoscaler v2 monitor using ray installation: {ray.__file__}"
)
logger.info(f"Ray version: {ray.__version__}")
logger.info(f"Ray commit: {ray.__commit__}")
logger.info(f"AutoscalerMonitor started with command: {sys.argv}")
gcs_address = args.gcs_address
if gcs_address is None:
raise ValueError("--gcs-address must be set!")
if not args.autoscaling_config:
logger.info("No autoscaling config provided: use read only node provider.")
config_reader = ReadOnlyProviderConfigReader(gcs_address)
else:
autoscaling_config = os.path.expanduser(args.autoscaling_config)
config_reader = FileConfigReader(
config_file=autoscaling_config, skip_content_hash=True
)
monitor = AutoscalerMonitor(
gcs_address,
config_reader,
log_dir=args.logs_dir,
monitor_ip=args.monitor_ip,
)
monitor.run()
File diff suppressed because it is too large Load Diff
+621
View File
@@ -0,0 +1,621 @@
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Tuple
from ray.autoscaler.v2.instance_manager.common import InstanceUtil
from ray.core.generated.autoscaler_pb2 import NodeState, NodeStatus
from ray.core.generated.instance_manager_pb2 import Instance
# TODO(rickyx): once we have graceful shutdown, we could populate
# the failure detail with the actual termination message. As of now,
# we will use a more generic message to include cases such as:
# (idle termination, node death, crash, preemption, etc)
NODE_DEATH_CAUSE_RAYLET_DIED = "NodeTerminated"
# e.g., cpu_4_ondemand.
NodeType = str
@dataclass
class ResourceUsage:
# Resource name.
resource_name: str = ""
# Total resource.
total: float = 0.0
# Resource used.
used: float = 0.0
@dataclass
class NodeUsage:
# The node resource usage.
usage: List[ResourceUsage]
# How long the node has been idle.
idle_time_ms: int
@dataclass
class NodeInfo:
# The instance type name, e.g. p3.2xlarge
instance_type_name: str
# ray node type name.
ray_node_type_name: str
# Cloud instance id.
instance_id: str
# Ip address of the node when alive.
ip_address: str
# The status of the node. Optional for pending nodes.
node_status: Optional[str] = None
# ray node id in hex. None if still pending.
node_id: Optional[str] = None
# Resource usage breakdown if node is running.
resource_usage: Optional[NodeUsage] = None
# Failure detail if the node failed.
failure_detail: Optional[str] = None
# Descriptive details.
details: Optional[str] = None
# Activity on the node.
node_activity: Optional[List[str]] = None
# Ray node labels.
labels: Optional[Dict[str, str]] = None
def total_resources(self) -> Dict[str, float]:
if self.resource_usage is None:
return {}
return {r.resource_name: r.total for r in self.resource_usage.usage}
def available_resources(self) -> Dict[str, float]:
if self.resource_usage is None:
return {}
return {r.resource_name: r.total - r.used for r in self.resource_usage.usage}
def used_resources(self) -> Dict[str, float]:
if self.resource_usage is None:
return {}
return {r.resource_name: r.used for r in self.resource_usage.usage}
@dataclass
class LaunchRequest:
class Status(Enum):
FAILED = "FAILED"
PENDING = "PENDING"
# The instance type name, e.g. p3.2xlarge
instance_type_name: str
# ray node type name.
ray_node_type_name: str
# count.
count: int
# State: (e.g. PENDING, FAILED)
state: Status
# When the launch request was made in unix timestamp in secs.
request_ts_s: int
# When the launch request failed unix timestamp in secs if failed.
failed_ts_s: Optional[int] = None
# Request details, e.g. error reason if the launch request failed.
details: Optional[str] = None
@dataclass
class ResourceRequestByCount:
# Bundles in the demand.
bundle: Dict[str, float]
# Number of bundles with the same shape.
count: int
def __str__(self) -> str:
return f"[{self.count} {self.bundle}]"
@dataclass
class ResourceDemand:
# The bundles in the demand with shape and count info.
bundles_by_count: List[ResourceRequestByCount]
@dataclass
class PlacementGroupResourceDemand(ResourceDemand):
# Details string (parsed into below information)
details: str
# Placement group's id.
pg_id: Optional[str] = None
# Strategy, e.g. STRICT_SPREAD
strategy: Optional[str] = None
# Placement group's state, e.g. PENDING
state: Optional[str] = None
def __post_init__(self):
if not self.details:
return
# Details in the format of <pg_id>:<strategy>|<state>, parse
# it into the above fields.
pattern = r"^.*:.*\|.*$"
match = re.match(pattern, self.details)
if not match:
return
pg_id, details = self.details.split(":")
strategy, state = details.split("|")
self.pg_id = pg_id
self.strategy = strategy
self.state = state
@dataclass
class RayTaskActorDemand(ResourceDemand):
pass
@dataclass
class ClusterConstraintDemand(ResourceDemand):
pass
@dataclass
class ResourceDemandSummary:
# Placement group demand.
placement_group_demand: List[PlacementGroupResourceDemand] = field(
default_factory=list
)
# Ray task actor demand.
ray_task_actor_demand: List[RayTaskActorDemand] = field(default_factory=list)
# Cluster constraint demand.
cluster_constraint_demand: List[ClusterConstraintDemand] = field(
default_factory=list
)
@dataclass
class Stats:
# How long it took to get the GCS request.
gcs_request_time_s: float = 0.0
# How long it took to get all live instances from node provider.
none_terminated_node_request_time_s: Optional[float] = None
# How long for autoscaler to process the scaling decision.
autoscaler_iteration_time_s: Optional[float] = None
# The last seen autoscaler state version from Ray.
autoscaler_version: Optional[str] = None
# The last seen cluster state resource version.
cluster_resource_state_version: Optional[str] = None
# Request made time unix timestamp: when the data was pulled from GCS.
request_ts_s: Optional[int] = None
@dataclass
class ClusterStatus:
# Healthy nodes information (non-idle)
active_nodes: List[NodeInfo] = field(default_factory=list)
# Idle node information
idle_nodes: List[NodeInfo] = field(default_factory=list)
# Pending launches.
pending_launches: List[LaunchRequest] = field(default_factory=list)
# Failed launches.
failed_launches: List[LaunchRequest] = field(default_factory=list)
# Pending nodes.
pending_nodes: List[NodeInfo] = field(default_factory=list)
# Failures
failed_nodes: List[NodeInfo] = field(default_factory=list)
# Resource usage summary for entire cluster.
cluster_resource_usage: List[ResourceUsage] = field(default_factory=list)
# Demand summary.
resource_demands: ResourceDemandSummary = field(
default_factory=ResourceDemandSummary
)
# Query metics
stats: Stats = field(default_factory=Stats)
def total_resources(self) -> Dict[str, float]:
return {r.resource_name: r.total for r in self.cluster_resource_usage}
def available_resources(self) -> Dict[str, float]:
return {r.resource_name: r.total - r.used for r in self.cluster_resource_usage}
# TODO(rickyx): we don't show infeasible requests as of now.
# (They will just be pending forever as part of the demands)
# We should show them properly in the future.
IPPRSpecsSchema = {
# JSON schema for IPPR (In-Place Pod Resize) specs provided via the
# Kubernetes annotation `ray.io/ippr` on a RayCluster CR.
#
# Structure:
# {
# "groups": {
# "<groupName>": {
# "max-cpu": string|number, # K8s quantity (e.g. "2", "1500m")
# "max-memory": string|integer, # K8s quantity (e.g. "8Gi", 2147483648)
# "resize-timeout": integer # Seconds to wait for a pod resize to
# # complete before considering it timed out
# },
# ...
# }
# }
#
# Notes:
# - The set of valid <groupName> keys corresponds to the RayCluster
# `workerGroupSpecs[].groupName` plus the implicit "headgroup".
# - The minimal CPU/memory values (min_*) are derived from the pod template
# requests/limits and are not part of this schema.
"type": "object",
"properties": {
"groups": {
"type": "object",
"additionalProperties": {
"type": "object",
"properties": {
"max-cpu": {"type": ["string", "number"]},
"max-memory": {"type": ["string", "integer"]},
"resize-timeout": {"type": "integer"},
},
"required": ["max-cpu", "max-memory", "resize-timeout"],
},
}
},
}
@dataclass
class IPPRGroupSpec:
"""Per-group IPPR limits and baseline resources.
This mirrors a single Ray group (worker group or head group). The minimal
resources are derived from the pod template's container resources (either
requests or limits if present), and the maximal resources and timeout are
provided by the IPPR spec annotation validated by ``IPPRSpecsSchema``.
Attributes:
min_cpu: Baseline CPU in cores derived from the pod template (float, e.g., 1.5).
max_cpu: Maximum CPU in cores allowed for in-place resize for this group.
min_memory: Baseline memory in bytes derived from the pod template.
max_memory: Maximum memory in bytes allowed for in-place resize for this group.
resize_timeout: Timeout in seconds for a single resize operation before
it is considered timed out.
"""
min_cpu: float
max_cpu: float
min_memory: int
max_memory: int
resize_timeout: int
@dataclass
class IPPRSpecs:
"""Typed, validated IPPR specs across Ray groups.
Attributes:
groups: Mapping from Ray group name (e.g., worker group ``groupName`` or
``"headgroup"``) to its ``IPPRGroupSpec``.
"""
groups: Dict[str, IPPRGroupSpec]
@dataclass
class IPPRStatus:
"""Represents the current and target resources for a pod under IPPR.
This structure is the working state used by the autoscaler to decide if and
when to apply in-place pod resizes and when to synchronize resource changes
with the Raylet.
Attributes:
cloud_instance_id: Cloud instance identifier for the pod (K8s pod name).
spec: The group-level limits and baselines for this pod.
current_cpu: Current CPU allocation in cores from the pod status.
current_memory: Current memory allocation in bytes from the pod status.
desired_cpu: Target CPU allocation in cores.
desired_memory: Target memory allocation in bytes.
resizing_at: Unix timestamp (seconds) when a resize request was issued to
Kubernetes, or None if not pending/needed.
k8s_resize_status: Lower-cased status from pod conditions for IPPR, e.g.
"inprogress", "deferred", "infeasible", "error"; None indicates no
active resize and is treated as finished.
k8s_resize_message: Message from the pod condition describing the resize
state or failure, if any.
suggested_max_cpu: Current-iteration suggested CPU cap (cores), computed from
node remaining capacity plus existing gap when resize is deferred/infeasible.
suggested_max_memory: Current-iteration suggested memory cap (bytes), computed
from node remaining capacity plus existing gap when resize is deferred/infeasible.
last_failed_at: Unix timestamp (seconds) when this pod most recently hit a
terminal IPPR failure. Once set, the autoscaler stops sending further
IPPR requests for this pod.
last_failed_reason: Human-readable reason for the terminal IPPR failure.
raylet_id: Raylet node id (hex) running in this pod, used to sync Ray's
internal resource view when K8s successfully changed pod resources.
"""
cloud_instance_id: str
spec: IPPRGroupSpec
current_cpu: float
current_memory: int
desired_cpu: float
desired_memory: int
resizing_at: Optional[int] = None
k8s_resize_status: Optional[str] = None
k8s_resize_message: Optional[str] = None
suggested_max_cpu: Optional[float] = None
suggested_max_memory: Optional[int] = None
last_failed_at: Optional[int] = None
last_failed_reason: Optional[str] = None
raylet_id: Optional[str] = None
def queue_resize_request(
self,
desired_cpu: Optional[float] = None,
desired_memory: Optional[int] = None,
) -> bool:
"""Queue the new desired resources and reset resize tracking state.
Queues the new desired CPU/memory if provided, associates the Raylet id,
and marks the resize state as "new" so the scheduler can identify the IPPR
action before the next iteration.
Args:
desired_cpu: Optional new desired CPU in cores.
desired_memory: Optional new desired memory in bytes.
Returns:
bool: True if the resize request is queued.
"""
updated = False
if desired_cpu is not None and desired_cpu != self.desired_cpu:
self.desired_cpu = desired_cpu
updated = True
if desired_memory is not None and desired_memory != self.desired_memory:
self.desired_memory = desired_memory
updated = True
if updated:
self.resizing_at = None
self.k8s_resize_status = "new"
self.k8s_resize_message = None
return updated
def has_resize_request_to_send(self) -> bool:
"""Whether this pod should be sent an IPPR request now.
Returns True if there is a Raylet id and the status is marked as "new".
"""
return self.raylet_id is not None and self.k8s_resize_status == "new"
def is_in_progress(self) -> bool:
"""Whether a resize is on going or about to be issued.
True if a resize was already issued to K8s (``resizing_at`` set), or if
the status is newly queued for resize (``has_resize_request_to_send``).
"""
return self.resizing_at is not None or self.has_resize_request_to_send()
def is_k8s_resize_finished(self) -> bool:
"""Whether the Kubernetes-side resize is considered finished.
We treat ``k8s_resize_status is None`` as finished. After K8s completes a
resize, the provider clears ``k8s_resize_status`` as None.
"""
return self.k8s_resize_status is None
def is_raylet_synced(self) -> bool:
"""Whether the Raylet's internal resources have been updated.
After the Raylet has been synchronized,
the provider clears ``resizing_at``.
"""
return self.resizing_at is None
def need_sync_with_raylet(self) -> bool:
"""Whether the Raylet's internal resources need to be updated.
Returns True when all of the following hold:
- We know the ``raylet_id`` for this pod.
- A resize was previously issued to Kubernetes (``resizing_at`` set).
- Kubernetes has finished applying the resize (``k8s_resize_status`` is None).
- The pod's current resources equal the desired resources.
In this case, the provider will call GCS's resize_raylet_resource_instances to update the raylet's local
resource instances and then clear ``resizing_at`` on the pod annotation.
"""
return (
self.raylet_id is not None
and self.resizing_at is not None
and self.k8s_resize_status is None
and self.desired_cpu == self.current_cpu
and self.desired_memory == self.current_memory
)
def max_cpu(self) -> float:
"""Effective maximum CPU cores allowed for this pod.
Preference order:
1) ``suggested_max_cpu`` (discovered limit)
2) ``spec.max_cpu`` (static limit)
"""
if self.suggested_max_cpu is not None:
return self.suggested_max_cpu
return self.spec.max_cpu
def max_memory(self) -> int:
"""Effective maximum memory bytes allowed for this pod.
Preference order:
1) ``suggested_max_memory`` (discovered limit)
2) ``spec.max_memory`` (static limit)
"""
if self.suggested_max_memory is not None:
return self.suggested_max_memory
return self.spec.max_memory
def can_resize_up(self) -> bool:
"""Whether the pod can still be scaled up within allowed limits.
Only returns True when no resize is in progress and the current
CPU/memory are below the effective max limits. Pods that have already
hit a terminal IPPR failure are permanently excluded from future IPPR.
"""
return (
self.last_failed_at is None
and self.is_k8s_resize_finished()
and self.is_raylet_synced()
and (
self.current_cpu < self.max_cpu()
or self.current_memory < self.max_memory()
)
)
def is_timeout(self) -> bool:
"""Whether an in-flight resize has exceeded the group's timeout.
Returns True when a resize was issued and now current time exceeds
``resizing_at + spec.resize_timeout``.
"""
return (
self.resizing_at is not None
and not self.need_sync_with_raylet()
and (self.resizing_at + self.spec.resize_timeout) < time.time()
)
def is_errored(self) -> bool:
"""Whether the last resize attempt reported an error from Kubernetes."""
return self.k8s_resize_status == "error"
def record_failure(self, reason: str, failed_at: Optional[int] = None) -> None:
"""Record the IPPR failure."""
self.last_failed_at = int(time.time()) if failed_at is None else failed_at
self.last_failed_reason = reason
@dataclass
class AutoscalerInstance:
"""
AutoscalerInstance represents an instance that's managed by the autoscaler.
This includes two states:
1. the instance manager state: information of the underlying cloud instance.
2. the ray node state, e.g. resources, ray node status.
The two states are linked by the cloud instance id, which should be set
when the ray node is started.
"""
# The cloud instance id. It could be None if the instance hasn't been assigned
# a cloud instance id, e.g. the instance is still in QUEUED or REQUESTED status.
cloud_instance_id: Optional[str] = None
# The ray node state status. It could be None when no ray node is running
# or has run on the cloud instance: for example, ray is still being installed
# or the instance manager hasn't had a cloud instance assigned (e.g. QUEUED,
# REQUESTED).
ray_node: Optional[NodeState] = None
# The instance manager instance state. It would be None when the ray_node is not
# None.
# It could be None iff:
# 1. There's a ray node, but the instance manager hasn't discovered the
# cloud instance that's running this ray process yet. This could happen since
# the instance manager only discovers instances periodically.
#
# 2. There was a ray node running on the cloud instance, which was already stopped
# and removed from the instance manager state. But the ray state is still lagging
# behind.
#
# 3. There is a ray node that's unmanaged by the instance manager.
#
im_instance: Optional[Instance] = None
# | cloud_instance_id | ray_node | im_instance |
# |-------------------|----------|-------------|
# | None | None | None | Not possible.
# | None | None | not None | OK. An instance hasn't had ray running on it yet. # noqa E501
# | None | Not None | None | OK. Possible if the ray node is not started by autoscaler. # noqa E501
# | None | Not None | not None | Not possible - no way to link im instance with ray node. # noqa E501
# | not None | None | None | Not possible since cloud instance id is either part of im state or ray node. # noqa E501
# | not None | None | not None | OK. e.g. An instance that's not running ray yet. # noqa E501
# | not None | Not None | None | OK. See scenario 1, 2, 3 above.
# | not None | Not None | not None | OK. An instance that's running ray.
def validate(self) -> Tuple[bool, str]:
"""Validate the autoscaler instance state.
Returns:
A tuple of (valid, error_msg) where:
- valid is whether the state is valid
- error_msg is the error message for the validation results.
"""
state_combinations = {
# (cloud_instance_id is None, ray_node is None, im_instance is None): (valid, error_msg) # noqa E501
(True, True, True): (False, "Not possible"),
(True, True, False): (True, ""),
(True, False, True): (
True,
"There's a ray node w/o cloud instance id, must be started not "
"by autoscaler",
),
(True, False, False): (
False,
"Not possible - no way to link im instance with ray node",
),
(False, True, True): (
False,
"Not possible since cloud instance id is either part of "
"im state or ray node",
),
(False, True, False): (True, ""),
(False, False, True): (True, ""),
(False, False, False): (True, ""),
}
valid, error_msg = state_combinations[
(
self.cloud_instance_id is None,
self.ray_node is None,
self.im_instance is None,
)
]
if not valid:
return valid, error_msg
if self.im_instance is not None and self.ray_node is None:
# We don't see a ray node, but tracking an im instance.
if self.cloud_instance_id is None:
if InstanceUtil.is_cloud_instance_allocated(self.im_instance.status):
return (
False,
"instance should be in a status where cloud instance "
"is not allocated.",
)
else:
if not InstanceUtil.is_cloud_instance_allocated(
self.im_instance.status
):
return (
False,
"instance should be in a status where cloud instance is "
"allocated.",
)
if self.ray_node is not None:
if self.cloud_instance_id != self.ray_node.instance_id:
return False, "cloud instance id doesn't match."
if self.im_instance is not None and self.cloud_instance_id is not None:
if self.cloud_instance_id != self.im_instance.cloud_instance_id:
return False, "cloud instance id doesn't match."
return True, ""
def is_ray_running(self) -> bool:
"""Whether the ray node is running."""
return self.ray_node is not None and self.ray_node.status in [
NodeStatus.RUNNING,
NodeStatus.IDLE,
]
def is_ray_stop(self) -> bool:
"""Whether the ray node is stopped."""
return self.ray_node is None or self.ray_node.status in [
NodeStatus.DEAD,
]
+121
View File
@@ -0,0 +1,121 @@
import time
from collections import Counter
from typing import List, NamedTuple
from ray._raylet import GcsClient
from ray.autoscaler.v2.schema import ClusterStatus, Stats
from ray.autoscaler.v2.utils import ClusterStatusParser
from ray.core.generated.autoscaler_pb2 import (
ClusterResourceState,
GetClusterResourceStateReply,
GetClusterStatusReply,
)
DEFAULT_RPC_TIMEOUT_S = 10
class ResourceRequest(NamedTuple):
resources: dict
label_selector: dict
def request_cluster_resources(
gcs_address: str,
to_request: List[dict],
timeout: int = DEFAULT_RPC_TIMEOUT_S,
):
"""Request resources from the autoscaler.
This will add a cluster resource constraint to GCS. GCS will asynchronously
pass the constraint to the autoscaler, and the autoscaler will try to provision the
requested minimal bundles in `to_request`.
If the cluster already has `to_request` resources, this will be an no-op.
Future requests submitted through this API will overwrite the previous requests.
Args:
gcs_address: The GCS address to query.
to_request: A list of resource requests to request the cluster to have.
Each resource request is a tuple of resources and a label_selector
to apply per-bundle. e.g.: [{"resources": {"CPU": 1, "GPU": 1}, "label_selector": {"accelerator-type": "A100"}}]
timeout: Timeout in seconds for the request to be timeout
"""
assert len(gcs_address) > 0, "GCS address is not specified."
# Convert bundle dicts to ResourceRequest tuples.
normalized: List[ResourceRequest] = []
for r in to_request:
assert isinstance(
r, dict
), f"Internal Error: Expected a dict, but got {type(r)}"
resources = r.get("resources", {})
selector = r.get("label_selector", {})
normalized.append(ResourceRequest(resources, selector))
to_request = normalized
# Aggregate bundle by shape
def keyfunc(r):
return (
frozenset(r.resources.items()),
frozenset(r.label_selector.items()),
)
grouped_requests = Counter(keyfunc(r) for r in to_request)
bundles: List[dict] = []
label_selectors: List[dict] = []
counts: List[int] = []
for (bundle, selector), count in grouped_requests.items():
bundles.append(dict(bundle))
label_selectors.append(dict(selector))
counts.append(count)
GcsClient(gcs_address).request_cluster_resource_constraint(
bundles, label_selectors, counts, timeout_s=timeout
)
def get_cluster_status(
gcs_address: str, timeout: int = DEFAULT_RPC_TIMEOUT_S
) -> ClusterStatus:
"""
Get the cluster status from the autoscaler.
Args:
gcs_address: The GCS address to query.
timeout: Timeout in seconds for the request to be timeout
Returns:
A ClusterStatus object.
"""
assert len(gcs_address) > 0, "GCS address is not specified."
req_time = time.time()
str_reply = GcsClient(gcs_address).get_cluster_status(timeout_s=timeout)
reply_time = time.time()
reply = GetClusterStatusReply()
reply.ParseFromString(str_reply)
# TODO(rickyx): To be more accurate, we could add a timestamp field from the reply.
return ClusterStatusParser.from_get_cluster_status_reply(
reply,
stats=Stats(gcs_request_time_s=reply_time - req_time, request_ts_s=req_time),
)
def get_cluster_resource_state(gcs_client: GcsClient) -> ClusterResourceState:
"""
Get the cluster resource state from GCS.
Args:
gcs_client: The GCS client to query.
Returns:
A ClusterResourceState object
Raises:
Exception: If the request times out or failed.
"""
str_reply = gcs_client.get_cluster_resource_state()
reply = GetClusterResourceStateReply()
reply.ParseFromString(str_reply)
return reply.cluster_resource_state
@@ -0,0 +1 @@
from ray.tests.conftest import * # noqa
@@ -0,0 +1,544 @@
# coding: utf-8
"""
Test recovery behavior for the ALLOCATION_TIMEOUT scenario.
Verifies that when an instance reaches ALLOCATION_TIMEOUT, the autoscaler can:
1. Terminate the timed-out old instance.
2. Launch a replacement instance.
3. Avoid a QUEUED->REQUESTED->QUEUED loop.
Test design principles:
- Pure Python, mocking only the k8s client.
- Validate instance state rather than log output.
- Run reconcile multiple times to verify the state does not get stuck.
"""
import time
from typing import Any, Dict, List
import pytest
from ray.autoscaler.v2.instance_manager.config import InstanceReconcileConfig, Provider
from ray.autoscaler.v2.instance_manager.instance_manager import InstanceManager
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.autoscaler.v2.instance_manager.node_provider import (
CloudInstance,
ICloudInstanceProvider,
)
from ray.autoscaler.v2.instance_manager.reconciler import Reconciler
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import (
CloudInstanceUpdater,
)
from ray.autoscaler.v2.instance_manager.subscribers.cloud_resource_monitor import (
CloudResourceMonitor,
)
from ray.autoscaler.v2.scheduler import (
ResourceDemandScheduler,
)
from ray.autoscaler.v2.tests.util import create_instance
from ray.core.generated.autoscaler_pb2 import (
ClusterResourceState,
NodeState,
NodeStatus,
)
from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
s_to_ns = 1 * 1_000_000_000
class MockAutoscalingConfig:
"""Mock autoscaling config for testing"""
def __init__(self, configs=None):
if configs is None:
configs = {}
self._configs = configs
def get_node_type_configs(self):
return self._configs.get("node_type_configs", {})
def get_max_num_worker_nodes(self):
return self._configs.get("max_num_worker_nodes")
def get_max_num_nodes(self):
n = self._configs.get("max_num_worker_nodes")
return n + 1 if n is not None else None
def get_upscaling_speed(self):
return self._configs.get("upscaling_speed", 0.0)
def get_max_concurrent_launches(self):
return self._configs.get("max_concurrent_launches", 100)
def get_instance_reconcile_config(self):
return self._configs.get("instance_reconcile_config", InstanceReconcileConfig())
def disable_node_updaters(self):
return self._configs.get("disable_node_updaters", True)
def disable_launch_config_check(self):
return self._configs.get("disable_launch_config_check", False)
def get_idle_timeout_s(self):
return self._configs.get("idle_timeout_s", 999)
def get_provider_instance_type(self, ray_node_type):
return ""
@property
def provider(self):
return Provider.UNKNOWN
class EventCapturingSubscriber:
"""
Subscriber that captures events for verification and delegates to CloudInstanceUpdater.
This wraps the real CloudInstanceUpdater to capture events for test assertions
while using the actual launch/terminate logic.
"""
def __init__(self, cloud_provider: ICloudInstanceProvider):
self.cloud_provider = cloud_provider
self.updater = CloudInstanceUpdater(cloud_provider=cloud_provider)
self.events = []
def notify(self, events):
self.events.extend(events)
# Delegate to real updater which calls cloud_provider.launch/terminate
self.updater.notify(events)
def clear(self):
self.events.clear()
def events_by_id(self, instance_id):
return [e for e in self.events if e.instance_id == instance_id]
class MockK8sClient:
"""
Mock Kubernetes API client that simulates RayCluster behavior.
Tracks:
- replicas: current replica count
- workers_to_delete: list of worker pod names to delete
- patch_history: history of all patches for verification
"""
def __init__(self, max_replicas=3, initial_replicas=3, worker_pod_names=None):
self.max_replicas = max_replicas
self.replicas = initial_replicas
self.workers_to_delete: List[str] = []
self.patch_history: List[Dict] = []
self._resource_version = 100
# Worker pod names - defaults to worker-001, worker-002, worker-003
self.worker_pod_names = worker_pod_names or [
"worker-001",
"worker-002",
"worker-003",
]
def get(self, path: str) -> Dict[str, Any]:
"""Handle GET requests"""
if "rayclusters" in path:
return {
"metadata": {
"name": "test-ray-cluster",
"namespace": "default",
"resourceVersion": str(self._resource_version),
},
"spec": {
"workerGroupSpecs": [
{
"groupName": "default-worker-group",
"replicas": self.replicas,
"minReplicas": 0,
"maxReplicas": self.max_replicas,
"scaleStrategy": {
"workersToDelete": list(self.workers_to_delete)
},
"numOfHosts": 1,
}
]
},
}
elif "pods" in path:
# Return pods based on current replicas (excluding workers_to_delete)
# Use worker_pod_names to match the cloud_instances in the test
items = []
for i in range(min(self.replicas, len(self.worker_pod_names))):
pod_name = self.worker_pod_names[i]
if pod_name not in self.workers_to_delete:
# worker-003 is in ALLOCATION_TIMEOUT state, so it should not be running
if pod_name == "worker-003":
# Pod is pending/failed - not running
container_state = {"waiting": {"reason": "ContainerCreating"}}
else:
container_state = {"running": {}}
items.append(
{
"metadata": {
"name": pod_name,
"labels": {
"ray.io/cluster": "test-ray-cluster",
"ray.io/node-type": "worker",
"ray.io/group": "default-worker-group",
},
},
"status": {
"containerStatuses": [{"state": container_state}]
},
}
)
return {
"metadata": {"resourceVersion": str(self._resource_version)},
"items": items,
}
return {}
def patch(self, path: str, payload: List[Dict]) -> Dict[str, Any]:
"""Handle PATCH requests and update internal state"""
self.patch_history.append({"path": path, "payload": payload})
self._resource_version += 1
for op in payload:
if op["op"] == "replace" and "replicas" in op["path"]:
self.replicas = op["value"]
elif op["op"] == "replace" and "scaleStrategy" in op["path"]:
self.workers_to_delete = op["value"].get("workersToDelete", [])
return {}
class TestAllocationTimeoutRecovery:
"""Test ALLOCATION_TIMEOUT instance recovery"""
@staticmethod
def _add_instances(instance_storage, instances):
for instance in instances:
ok, _ = instance_storage.upsert_instance(instance)
assert ok
def test_no_queued_requested_loop(self):
"""
Minimal reproduction path:
Preconditions:
- maxReplicas = 3
- 3 worker instances: 2 ALLOCATED (healthy), 1 ALLOCATION_TIMEOUT
- idle_worker_nodes = 1
Steps:
1. Run Reconciler.reconcile() multiple times to simulate repeated
reconciler cycles.
Expected results:
1. The new instance does not enter a QUEUED->REQUESTED->QUEUED loop.
2. The new instance can eventually transition into REQUESTED.
3. Terminate is submitted to Kubernetes before launch.
"""
# ===== Preconditions =====
instance_storage = InstanceStorage(
cluster_id="test_cluster_id",
storage=InMemoryStorage(),
)
cloud_resource_monitor = CloudResourceMonitor()
# Mock K8s client: maxReplicas=3, initial_replicas=3
# Worker pod names match the cloud_instance_id in cloud_instances
mock_k8s = MockK8sClient(
max_replicas=3,
initial_replicas=3,
worker_pod_names=["worker-001", "worker-002", "worker-003"],
)
# Create real cloud provider with mock k8s client
from unittest.mock import MagicMock
from ray.autoscaler.v2.instance_manager.cloud_providers.kuberay.cloud_provider import (
KubeRayProvider,
)
cloud_provider = KubeRayProvider(
cluster_name="test-ray-cluster",
provider_config={"namespace": "default"},
gcs_client=MagicMock(),
k8s_api_client=mock_k8s,
)
# Use EventCapturingSubscriber that wraps real CloudInstanceUpdater
# and captures events for test verification
mock_subscriber = EventCapturingSubscriber(cloud_provider=cloud_provider)
instance_manager = InstanceManager(
instance_storage=instance_storage,
instance_status_update_subscribers=[mock_subscriber],
)
# Create 2 ALLOCATED instances and 1 ALLOCATION_TIMEOUT instance.
current_time = time.time_ns()
timeout_time = (
current_time - 200 * s_to_ns
) # 200s ago, beyond the timeout threshold.
from ray._common.utils import binary_to_hex
instances = [
# Head node
create_instance(
"head",
status=Instance.RAY_RUNNING,
cloud_instance_id="head-001",
node_kind=NodeKind.HEAD,
instance_type="head",
ray_node_id=binary_to_hex(b"head"),
),
# Worker 1: healthy
create_instance(
"worker-1",
status=Instance.ALLOCATED,
instance_type="default-worker-group",
cloud_instance_id="worker-001",
node_kind=NodeKind.WORKER,
ray_node_id=binary_to_hex(b"wkr1"),
),
# Worker 2: healthy
create_instance(
"worker-2",
status=Instance.ALLOCATED,
instance_type="default-worker-group",
cloud_instance_id="worker-002",
node_kind=NodeKind.WORKER,
ray_node_id=binary_to_hex(b"wkr2"),
),
# Worker 3: ALLOCATION_TIMEOUT (startup timed out)
create_instance(
"worker-3",
status=Instance.ALLOCATION_TIMEOUT,
instance_type="default-worker-group",
cloud_instance_id="worker-003",
node_kind=NodeKind.WORKER,
status_times=[(Instance.ALLOCATION_TIMEOUT, timeout_time)],
),
]
TestAllocationTimeoutRecovery._add_instances(instance_storage, instances)
# Mock ray nodes: head + 2 healthy workers.
# ray_node_id must match the instance node_id.
# available_resources=0 means resources are fully consumed, so scale-up is needed.
ray_nodes = [
NodeState(
node_id=b"head",
status=NodeStatus.RUNNING,
instance_id="head-001",
total_resources={"CPU": 0},
available_resources={"CPU": 0},
),
NodeState(
node_id=b"wkr1",
status=NodeStatus.RUNNING,
instance_id="worker-001",
total_resources={"CPU": 4},
available_resources={"CPU": 0}, # Resources are fully consumed.
),
NodeState(
node_id=b"wkr2",
status=NodeStatus.RUNNING,
instance_id="worker-002",
total_resources={"CPU": 4},
available_resources={"CPU": 0}, # Resources are fully consumed.
),
]
# Mock cloud instances
# worker-003 must be included because the cloud instance still exists
# while the instance is in ALLOCATION_TIMEOUT.
# is_running=False means the pod is not running normally.
cloud_instances = {
"head-001": CloudInstance("head-001", "head", True, NodeKind.HEAD),
"worker-001": CloudInstance(
"worker-001", "default-worker-group", True, NodeKind.WORKER
),
"worker-002": CloudInstance(
"worker-002", "default-worker-group", True, NodeKind.WORKER
),
"worker-003": CloudInstance(
"worker-003",
"default-worker-group",
False,
NodeKind.WORKER, # is_running=False
),
}
# Use the real ResourceDemandScheduler.
scheduler = ResourceDemandScheduler()
# Add pending_resource_requests to simulate resource demand.
# Request 4 CPU while current available CPU is 0, so one more worker is needed.
from ray.core.generated.autoscaler_pb2 import (
ResourceRequest,
ResourceRequestByCount,
)
ray_cluster_resource_state = ClusterResourceState(
node_states=ray_nodes,
pending_resource_requests=[
ResourceRequestByCount(
request=ResourceRequest(resources_bundle={"CPU": 4}), count=1
),
],
)
# Mock autoscaling config. Use MockAutoscalingConfig to avoid schema validation.
# Build real NodeTypeConfig objects.
from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
node_type_configs = {
"default-worker-group": NodeTypeConfig(
name="default-worker-group",
min_worker_nodes=0,
max_worker_nodes=3,
resources={"CPU": 4},
labels={},
launch_config_hash="hash1",
idle_timeout_s=None,
),
"head": NodeTypeConfig(
name="head",
min_worker_nodes=0,
max_worker_nodes=1,
resources={"CPU": 0},
labels={},
launch_config_hash="hash1",
idle_timeout_s=None,
),
}
autoscaling_config = MockAutoscalingConfig(
configs={
"node_type_configs": node_type_configs,
"max_num_worker_nodes": 3,
"upscaling_speed": 1.0,
"max_concurrent_launches": 100,
"instance_reconcile_config": InstanceReconcileConfig(
request_status_timeout_s=10,
allocate_status_timeout_s=300,
),
"disable_node_updaters": True,
"disable_launch_config_check": True,
"idle_timeout_s": 999,
}
)
# ===== Steps: run reconcile multiple times =====
# Run 3 reconcile cycles to simulate repeated reconciler execution.
for cycle in range(3):
# Clear subscriber events so each iteration is checked independently.
mock_subscriber.events.clear()
Reconciler.reconcile(
instance_manager=instance_manager,
scheduler=scheduler,
cloud_provider=cloud_provider,
cloud_resource_monitor=cloud_resource_monitor,
ray_cluster_resource_state=ray_cluster_resource_state,
non_terminated_cloud_instances=cloud_instances,
cloud_provider_errors=[],
ray_install_errors=[],
autoscaling_config=autoscaling_config,
)
# Fetch current instance state to ensure repeated reconcile calls do not fail.
instance_storage.get_instances()
# ===== Expected result checks =====
# Get final instance states.
all_instances, _ = instance_storage.get_instances()
# 1. Verify the ALLOCATION_TIMEOUT instance transitions to TERMINATING.
worker_3 = all_instances.get("worker-3")
assert worker_3 is not None, "Expected worker-3 instance to exist"
assert worker_3.status == Instance.TERMINATING, (
f"Expected worker-3 status to be TERMINATING, "
f"got {Instance.InstanceStatus.Name(worker_3.status)}"
)
# 2. Verify at least one new instance is created (QUEUED or REQUESTED).
new_instances = [
i
for i in all_instances.values()
if i.status in (Instance.QUEUED, Instance.REQUESTED)
]
assert len(new_instances) >= 1, (
f"Expected at least 1 new instance (QUEUED/REQUESTED), "
f"got {len(new_instances)}"
)
# 3. Verify K8s patch history shows terminate before launch.
# Each patch may contain both replicas and scaleStrategy updates.
# Key checks:
# a) There should be two patches (terminate + launch).
# b) The first patch should contain workersToDelete.
# c) The second patch should increase replicas.
# Verify there are two patches (terminate + launch).
assert len(mock_k8s.patch_history) == 2, (
f"Expected 2 patches (terminate + launch), got {len(mock_k8s.patch_history)}. "
"If only 1 patch, launch may have failed due to maxReplicas bug."
)
# Verify the first patch contains workersToDelete.
first_patch = mock_k8s.patch_history[0]
first_patch_payload_str = str(first_patch.get("payload", []))
assert "workersToDelete" in first_patch_payload_str, (
f"Expected first patch to contain workersToDelete (terminate before launch). "
f"First patch: {first_patch}"
)
# Verify workersToDelete contains the timed-out instance.
for op in first_patch["payload"]:
if "scaleStrategy" in op.get("path", ""):
workers_to_delete = op["value"].get("workersToDelete", [])
assert (
"worker-003" in workers_to_delete
), f"Expected worker-003 in workersToDelete, got {workers_to_delete}"
break
# Verify the second patch increases replicas.
second_patch = mock_k8s.patch_history[1]
for op in second_patch["payload"]:
if "replicas" in op.get("path", ""):
assert (
op["value"] == 3
), f"Expected replicas=3 after launch, got {op['value']}"
break
# 4. Verify final replicas do not exceed maxReplicas.
assert (
mock_k8s.replicas <= 3
), f"Expected replicas <= 3, got {mock_k8s.replicas}"
# 5. Verify no instance enters a REQUESTED->QUEUED->REQUESTED loop.
# Check whether status_history contains a repeated transition pattern.
for instance in all_instances.values():
status_sequence = [h.instance_status for h in instance.status_history]
# Check for a REQUESTED->QUEUED->REQUESTED pattern.
for i in range(len(status_sequence) - 2):
assert not (
status_sequence[i] == Instance.REQUESTED
and status_sequence[i + 1] == Instance.QUEUED
and status_sequence[i + 2] == Instance.REQUESTED
), (
f"Instance {instance.instance_id} entered REQUESTED->QUEUED->REQUESTED loop. "
f"Status sequence: {[Instance.InstanceStatus.Name(s) for s in status_sequence]}"
)
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
@@ -0,0 +1,262 @@
import logging
import os
import subprocess
import sys
import time
from unittest.mock import MagicMock, patch
import pytest
import ray
from ray._common.test_utils import wait_for_condition
from ray._raylet import GcsClient
from ray.autoscaler._private.fake_multi_node.node_provider import FAKE_HEAD_NODE_ID
from ray.autoscaler.v2.autoscaler import Autoscaler
from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
from ray.autoscaler.v2.monitor import AutoscalerMonitor
from ray.autoscaler.v2.sdk import get_cluster_status, request_cluster_resources
from ray.autoscaler.v2.tests.util import MockEventLogger
from ray.cluster_utils import Cluster
logger = logging.getLogger(__name__)
DEFAULT_AUTOSCALING_CONFIG = {
"cluster_name": "fake_multinode",
"max_workers": 8,
"provider": {
"type": "fake_multinode",
},
"available_node_types": {
"ray.head.default": {
"resources": {
"CPU": 0,
},
"max_workers": 0,
"node_config": {},
},
"ray.worker.cpu": {
"resources": {
"CPU": 1,
},
"min_workers": 0,
"max_workers": 10,
"node_config": {},
},
"ray.worker.gpu": {
"resources": {
"GPU": 1,
},
"min_workers": 0,
"max_workers": 10,
"node_config": {},
},
},
"head_node_type": "ray.head.default",
"upscaling_speed": 0,
"idle_timeout_minutes": 0.08, # ~5 seconds
}
@pytest.fixture(scope="function")
def make_autoscaler():
ctx = {}
def _make_autoscaler(config):
head_node_kwargs = {
"env_vars": {
"RAY_CLOUD_INSTANCE_ID": FAKE_HEAD_NODE_ID,
"RAY_OVERRIDE_NODE_ID_FOR_TESTING": FAKE_HEAD_NODE_ID,
"RAY_NODE_TYPE_NAME": "ray.head.default",
},
"num_cpus": config["available_node_types"]["ray.head.default"]["resources"][
"CPU"
],
}
cluster = Cluster(
initialize_head=True, head_node_args=head_node_kwargs, connect=True
)
ctx["cluster"] = cluster
mock_config_reader = MagicMock()
gcs_address = cluster.address
# Configs for the node provider
config["provider"]["gcs_address"] = gcs_address
config["provider"]["head_node_id"] = FAKE_HEAD_NODE_ID
config["provider"]["launch_multiple"] = True
os.environ["RAY_FAKE_CLUSTER"] = "1"
mock_config_reader.get_cached_autoscaling_config.return_value = (
AutoscalingConfig(configs=config, skip_content_hash=True)
)
gcs_address = gcs_address
gcs_client = GcsClient(gcs_address)
event_logger = AutoscalerEventLogger(MockEventLogger(logger))
autoscaler = Autoscaler(
session_name="test",
config_reader=mock_config_reader,
gcs_client=gcs_client,
event_logger=event_logger,
)
return autoscaler
yield _make_autoscaler
try:
ray.shutdown()
ctx["cluster"].shutdown()
except Exception:
logger.exception("Error during teardown")
# Run ray stop to clean up everything
subprocess.run(
["ray", "stop", "--force"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
)
def test_basic_scaling(make_autoscaler):
config = DEFAULT_AUTOSCALING_CONFIG
autoscaler = make_autoscaler(DEFAULT_AUTOSCALING_CONFIG)
gcs_address = autoscaler._gcs_client.address
# Resource requests
print("=================== Test scaling up constraint 1/2====================")
request_cluster_resources(
gcs_address, [{"resources": {"CPU": 1}}, {"resources": {"GPU": 1}}]
)
def verify():
autoscaler.update_autoscaling_state()
cluster_state = get_cluster_status(gcs_address)
assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 3
return True
wait_for_condition(verify, retry_interval_ms=2000)
# Test scaling down shouldn't happen
print("=================== Test scaling down constraint 2/2 ====================")
idle_timeout_s = config["idle_timeout_minutes"] * 60
time.sleep(idle_timeout_s + 2)
wait_for_condition(verify, retry_interval_ms=2000)
# Test scaling down.
print("=================== Test scaling down idle ====================")
request_cluster_resources(gcs_address, [])
def verify():
autoscaler.update_autoscaling_state()
cluster_state = get_cluster_status(gcs_address)
assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 1
return True
wait_for_condition(verify, retry_interval_ms=2000)
# Test scaling up again with tasks
print("=================== Test scaling up with tasks ====================")
@ray.remote
def task():
time.sleep(999)
task.options(num_cpus=1).remote()
task.options(num_cpus=0, num_gpus=1).remote()
def verify():
autoscaler.update_autoscaling_state()
cluster_state = get_cluster_status(gcs_address)
assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 3
return True
wait_for_condition(verify, retry_interval_ms=2000)
# Test with placement groups
print(
"=================== Test scaling up with placement groups ===================="
)
# Spread to create another 2 nodes.
ray.util.placement_group(
name="pg", strategy="STRICT_SPREAD", bundles=[{"CPU": 0.5}, {"CPU": 0.5}]
)
def verify():
autoscaler.update_autoscaling_state()
cluster_state = get_cluster_status(gcs_address)
assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 5
return True
wait_for_condition(verify, retry_interval_ms=2000)
# Pack with feasible ones
ray.util.placement_group(
name="pg_feasible", strategy="STRICT_PACK", bundles=[{"CPU": 0.5}, {"CPU": 0.5}]
)
def verify():
autoscaler.update_autoscaling_state()
cluster_state = get_cluster_status(gcs_address)
assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 6
return True
wait_for_condition(verify, retry_interval_ms=2000)
# Pack with infeasible request
ray.util.placement_group(
name="pg_infeasible", strategy="STRICT_PACK", bundles=[{"CPU": 1}, {"CPU": 1}]
)
def verify():
autoscaling_state = autoscaler.update_autoscaling_state()
cluster_state = get_cluster_status(gcs_address)
assert len(cluster_state.active_nodes + cluster_state.idle_nodes) == 6
assert len(autoscaling_state.infeasible_gang_resource_requests) == 1
return True
wait_for_condition(verify, retry_interval_ms=2000)
# Test report autoscaling state
def verify():
autoscaling_state = autoscaler.update_autoscaling_state()
assert len(autoscaling_state.infeasible_gang_resource_requests) == 1
# For now, we just track that it's called ok.
autoscaler._gcs_client.report_autoscaling_state(
autoscaling_state.SerializeToString()
)
return True
wait_for_condition(verify, retry_interval_ms=2000)
class TestAutoscalerMonitor(AutoscalerMonitor):
"""Lightweight wrapper for testing _run() without full init."""
def __init__(self, gcs_address, gcs_client, autoscaler):
self.gcs_address = gcs_address
self.gcs_client = gcs_client
self.autoscaler = autoscaler
self._session_name = "test"
def test_raise_AuthenticationError_v2(make_autoscaler):
autoscaler = make_autoscaler(DEFAULT_AUTOSCALING_CONFIG)
gcs_client = autoscaler._gcs_client
gcs_address = gcs_client.address
monitor = TestAutoscalerMonitor(gcs_address, gcs_client, autoscaler)
def flaky():
raise ray.exceptions.AuthenticationError("WrongClusterID")
with patch.object(autoscaler, "update_autoscaling_state", side_effect=flaky):
with pytest.raises(ray.exceptions.AuthenticationError):
monitor._run()
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,251 @@
# coding: utf-8
import os
import sys
import tempfile
import pytest # noqa
from ray._common.utils import binary_to_hex
from ray._private.test_utils import get_test_config_path
from ray.autoscaler import AUTOSCALER_DIR_PATH
from ray.autoscaler._private.util import format_readonly_node_type
from ray.autoscaler.v2.instance_manager import config as config_mod
from ray.autoscaler.v2.instance_manager.config import (
FileConfigReader,
Provider,
ReadOnlyProviderConfigReader,
)
@pytest.mark.parametrize(
"skip_hash",
[True, False],
)
def test_simple(skip_hash):
config = FileConfigReader(
get_test_config_path("test_multi_node.yaml"), skip_content_hash=skip_hash
).get_cached_autoscaling_config()
assert config.get_cloud_node_config("head_node") == {"InstanceType": "m5.large"}
assert config.get_docker_config("head_node") == {
"image": "anyscale/ray-ml:latest",
"container_name": "ray_container",
"pull_before_run": True,
}
assert config.get_worker_start_ray_commands
def test_complex():
config = FileConfigReader(
get_test_config_path("test_ray_complex.yaml")
).get_cached_autoscaling_config()
assert config.get_head_setup_commands() == [
"echo a",
"echo b",
"echo ${echo hi}",
"echo head",
]
assert config.get_head_start_ray_commands() == [
"ray stop",
"ray start --head --autoscaling-config=~/ray_bootstrap_config.yaml",
]
assert config.get_worker_setup_commands("worker_nodes") == [
"echo a",
"echo b",
"echo ${echo hi}",
"echo worker",
]
assert config.get_worker_start_ray_commands() == [
"ray stop",
"ray start --address=$RAY_HEAD_IP",
]
assert config.get_worker_setup_commands("worker_nodes1") == [
"echo worker1",
]
assert config.get_docker_config("head_node") == {
"image": "anyscale/ray-ml:head-default",
"container_name": "ray_container",
"pull_before_run": True,
}
assert config.get_docker_config("default") == {
"image": "anyscale/ray-ml:worker-default",
"container_name": "ray_container",
"pull_before_run": True,
}
assert config.get_docker_config("worker_nodes") == {
"image": "anyscale/ray-ml:worker-default",
"container_name": "ray_container",
"pull_before_run": True,
}
assert config.get_docker_config("worker_nodes1") == {
"image": "anyscale/ray-ml:worker_nodes1",
"container_name": "ray_container",
"pull_before_run": True,
}
assert config.get_initialization_commands("worker_nodes") == ["echo what"]
assert config.get_initialization_commands("worker_nodes1") == ["echo init"]
assert config.get_node_resources("worker_nodes1") == {"CPU": 2}
assert config.get_node_resources("worker_nodes") == {}
assert config.get_node_labels("worker_nodes1") == {"foo": "bar"}
assert config.get_config("cluster_name") == "test-cli"
assert config.get_config("non-existing", "default") == "default"
assert config.get_config("non-existing") is None
def test_multi_provider_instance_type():
def load_config(file):
path = os.path.join(AUTOSCALER_DIR_PATH, file)
return FileConfigReader(path).get_cached_autoscaling_config()
aws_config = load_config("aws/defaults.yaml")
assert aws_config.get_provider_instance_type("ray.head.default") == "m5.large"
gcp_config = load_config("gcp/defaults.yaml")
# NOTE: Why is this underscore....
assert gcp_config.get_provider_instance_type("ray_head_default") == "n1-standard-2"
aliyun_config = load_config("aliyun/defaults.yaml")
assert (
aliyun_config.get_provider_instance_type("ray.head.default") == "ecs.n4.large"
)
azure_config = load_config("azure/defaults.yaml")
assert (
azure_config.get_provider_instance_type("ray.head.default") == "Standard_D2s_v3"
)
# TODO(rickyx):
# We don't have kuberay and local config yet.
def test_node_type_configs():
config = FileConfigReader(
get_test_config_path("test_ray_complex.yaml")
).get_cached_autoscaling_config()
node_type_configs = config.get_node_type_configs()
assert config.get_max_num_worker_nodes() == 10
assert len(node_type_configs) == 4
assert node_type_configs["head_node"].max_worker_nodes == 1
assert node_type_configs["head_node"].min_worker_nodes == 0
assert node_type_configs["head_node"].resources == {}
assert node_type_configs["head_node"].labels == {}
assert node_type_configs["default"].max_worker_nodes == 2
assert node_type_configs["default"].min_worker_nodes == 0
assert node_type_configs["default"].resources == {}
assert node_type_configs["default"].labels == {}
assert node_type_configs["worker_nodes"].max_worker_nodes == 2
assert node_type_configs["worker_nodes"].min_worker_nodes == 1
assert node_type_configs["worker_nodes"].resources == {}
assert node_type_configs["worker_nodes"].labels == {}
assert node_type_configs["worker_nodes1"].max_worker_nodes == 2
assert node_type_configs["worker_nodes1"].min_worker_nodes == 1
assert node_type_configs["worker_nodes1"].resources == {"CPU": 2}
assert node_type_configs["worker_nodes1"].labels == {"foo": "bar"}
def test_read_config():
# Make a temp config file from aws/defaults.yaml
with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
# Write "aws/defaults.yaml" to the temp file
with open(
os.path.join(AUTOSCALER_DIR_PATH, "aws/defaults.yaml"), "r"
) as default_file:
f.write(default_file.read())
config_reader = FileConfigReader(f.name)
# Check that the config is read correctly
assert config_reader.get_cached_autoscaling_config().provider == Provider.AWS
# Now override the file with a different provider
with open(f.name, "w") as f:
# Replace the file with "gcp/defaults.yaml"
with open(
os.path.join(AUTOSCALER_DIR_PATH, "gcp/defaults.yaml"), "r"
) as default_file:
f.write(default_file.read())
# Still the same.
assert config_reader.get_cached_autoscaling_config().provider == Provider.AWS
# Reload
config_reader.refresh_cached_autoscaling_config()
assert config_reader.get_cached_autoscaling_config().provider == Provider.GCP
def test_readonly_node_type_name_and_fallback(monkeypatch):
class _DummyNodeState:
def __init__(self, ray_node_type_name, node_id, total_resources):
self.ray_node_type_name = ray_node_type_name
self.node_id = node_id
self.total_resources = total_resources
class _DummyClusterState:
def __init__(self, node_states):
self.node_states = node_states
# Avoid real GCS usage.
monkeypatch.setattr(config_mod, "GcsClient", lambda address: object())
# Build a cluster with:
# - 1 named head type
# - 2 named worker types of the same type (aggregation check)
# - 1 worker type without name (fallback to node_id-based type)
unnamed_worker_id = b"\xab"
fallback_name = format_readonly_node_type(binary_to_hex(unnamed_worker_id))
nodes = [
_DummyNodeState(
"ray.head.default", b"\x01", {"CPU": 1, "node:__internal_head__": 1}
),
_DummyNodeState("worker.custom", b"\x02", {"CPU": 2}),
_DummyNodeState("worker.custom", b"\x03", {"CPU": 2}),
_DummyNodeState("worker.custom", b"\x04", {"CPU": 2}),
_DummyNodeState("", unnamed_worker_id, {"CPU": 3}),
]
monkeypatch.setattr(
config_mod,
"get_cluster_resource_state",
lambda _gc: _DummyClusterState(nodes),
)
reader = ReadOnlyProviderConfigReader("dummy:0")
reader.refresh_cached_autoscaling_config()
cfg = reader.get_cached_autoscaling_config()
node_types = cfg.get_config("available_node_types")
# Head assertions
assert "ray.head.default" in node_types
assert node_types["ray.head.default"]["max_workers"] == 0
assert cfg.get_head_node_type() == "ray.head.default"
# Preferred name aggregation
assert "worker.custom" in node_types
assert node_types["worker.custom"]["max_workers"] == 3
# Fallback for unnamed worker
assert fallback_name in node_types
assert node_types[fallback_name]["max_workers"] == 1
# Global max_workers should be the sum of all worker-type max_workers,
# NOT the count of node type names.
# Here: 3 distinct types (head, worker.custom, fallback), but
# 4 actual workers (3 x worker.custom + 1 x fallback).
# The old buggy `len(available_node_types)` returned 3 instead of 4.
assert cfg.get_config("max_workers") == 4
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,131 @@
import logging
import os
import sys
import pytest
from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
from ray.autoscaler.v2.tests.util import MockEventLogger
from ray.autoscaler.v2.utils import ResourceRequestUtil
from ray.core.generated.autoscaler_pb2 import (
ClusterResourceConstraint,
GangResourceRequest,
)
from ray.core.generated.instance_manager_pb2 import LaunchRequest, TerminationRequest
# coding: utf-8
OUTDATED = TerminationRequest.Cause.OUTDATED
IDLE = TerminationRequest.Cause.IDLE
MAX_NUM_NODE_PER_TYPE = TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE
MAX_NUM_NODES = TerminationRequest.Cause.MAX_NUM_NODES
def launch_request(instance_type: str, count: int) -> LaunchRequest:
return LaunchRequest(
instance_type=instance_type,
count=count,
)
def termination_request(
instance_type: str, cause: TerminationRequest.Cause
) -> TerminationRequest:
return TerminationRequest(
instance_id="",
instance_type=instance_type,
cause=cause,
)
logger = logging.getLogger(__name__)
def test_log_scheduling_updates():
mock_logger = MockEventLogger(logger)
event_logger = AutoscalerEventLogger(mock_logger)
launch_requests = [
launch_request("m4.large", 2),
launch_request("m4.xlarge", 2),
]
terminate_requests = [
termination_request("m4.large", IDLE),
termination_request("m4.xlarge", OUTDATED),
]
infeasible_requests = [
ResourceRequestUtil.make({"CPU": 4, "GPU": 1}),
] * 100 + [ResourceRequestUtil.make({"CPU": 4})]
gang_resource_requests = [
[
ResourceRequestUtil.make({"CPU": 4, "GPU": 1}),
ResourceRequestUtil.make({"CPU": 4, "GPU": 1}),
]
]
cluster_resource_constraints = [
ResourceRequestUtil.make({"CPU": 1, "GPU": 1}),
] * 100
event_logger.log_cluster_scheduling_update(
launch_requests=launch_requests,
terminate_requests=terminate_requests,
infeasible_requests=infeasible_requests,
infeasible_gang_requests=[
GangResourceRequest(requests=reqs) for reqs in gang_resource_requests
],
infeasible_cluster_resource_constraints=[
ClusterResourceConstraint(
resource_requests=ResourceRequestUtil.group_by_count(
cluster_resource_constraints
)
)
],
cluster_resources={"CPU": 5, "GPU": 5, "TPU": 2},
)
assert mock_logger.get_logs("info") == [
"Adding 2 node(s) of type m4.large.",
"Adding 2 node(s) of type m4.xlarge.",
"Removing 1 nodes of type m4.large (idle).",
"Removing 1 nodes of type m4.xlarge (outdated).",
"Resized to 5 CPUs, 5 GPUs, 2 TPUs.",
]
expect_lines = [
"No available node types can fulfill resource requests", # noqa
"No available node types can fulfill placement group requests", # noqa
"No available node types can fulfill cluster constraint", # noqa
]
for expect_line, actual_line in zip(expect_lines, mock_logger.get_logs("error")):
assert expect_line in actual_line
assert mock_logger.get_logs("error") == []
assert mock_logger.get_logs("debug") == [
"Current cluster resources: {'CPU': 5, 'GPU': 5, 'TPU': 2}."
]
def test_log_scheduling_updates_without_cluster_shape():
mock_logger = MockEventLogger(logger)
event_logger = AutoscalerEventLogger(mock_logger, log_cluster_shape=False)
event_logger.log_cluster_scheduling_update(
launch_requests=[launch_request("m4.large", 1)],
terminate_requests=[termination_request("m4.xlarge", OUTDATED)],
infeasible_requests=[ResourceRequestUtil.make({"CPU": 4})],
cluster_resources={"CPU": 5},
)
assert mock_logger.get_logs("info") == []
assert mock_logger.get_logs("warning") == [
"No available node types can fulfill resource requests {'CPU': 4.0}*1. Add suitable node types to this cluster to resolve this issue."
]
assert mock_logger.get_logs("debug") == []
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,414 @@
import os
import sys
import unittest
from collections import defaultdict
from unittest.mock import MagicMock
# coding: utf-8
import pytest
from ray.autoscaler.v2.instance_manager.instance_manager import InstanceManager
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage, StoreStatus
from ray.autoscaler.v2.tests.util import MockSubscriber
from ray.core.generated.instance_manager_pb2 import (
GetInstanceManagerStateRequest,
Instance,
InstanceUpdateEvent,
NodeKind,
StatusCode,
UpdateInstanceManagerStateRequest,
)
class InstanceManagerTest(unittest.TestCase):
def test_instances_version_mismatch(self):
ins_storage = MagicMock()
subscriber = MockSubscriber()
im = InstanceManager(
ins_storage, instance_status_update_subscribers=[subscriber]
)
# Version mismatch on reading from the storage.
ins_storage.get_instances.return_value = ({}, 1)
update = InstanceUpdateEvent(
instance_id="id-1",
new_instance_status=Instance.QUEUED,
instance_type="type-1",
upsert=True,
)
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=0,
updates=[update],
)
)
assert reply.status.code == StatusCode.VERSION_MISMATCH
assert len(subscriber.events) == 0
# Version OK.
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=1,
updates=[update],
)
)
assert reply.status.code == StatusCode.OK
assert len(subscriber.events) == 1
assert subscriber.events[0].new_instance_status == Instance.QUEUED
# Version mismatch when writing to the storage (race happens)
ins_storage.batch_upsert_instances.return_value = StoreStatus(
False, 2 # No longer 1
)
subscriber.clear()
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=1,
updates=[update],
)
)
assert reply.status.code == StatusCode.VERSION_MISMATCH
assert len(subscriber.events) == 0
# Non-version mismatch error.
ins_storage.batch_upsert_instances.return_value = StoreStatus(
False, 1 # Still 1
)
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=1,
updates=[update],
)
)
assert reply.status.code == StatusCode.UNKNOWN_ERRORS
assert len(subscriber.events) == 0
def test_get_and_updates(self):
ins_storage = InstanceStorage(
"cluster-id",
InMemoryStorage(),
)
subscriber = MockSubscriber()
im = InstanceManager(
ins_storage, instance_status_update_subscribers=[subscriber]
)
# Empty storage.
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert reply.status.code == StatusCode.OK
assert list(reply.state.instances) == []
# Launch nodes.
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=0,
updates=[
InstanceUpdateEvent(
instance_type="type-1",
instance_id="id-1",
new_instance_status=Instance.QUEUED,
upsert=True,
),
InstanceUpdateEvent(
instance_type="type-2",
instance_id="id-2",
new_instance_status=Instance.QUEUED,
upsert=True,
),
InstanceUpdateEvent(
instance_type="type-2",
instance_id="id-3",
new_instance_status=Instance.QUEUED,
upsert=True,
),
],
)
)
assert reply.status.code == StatusCode.OK
assert len(subscriber.events) == 3
for e in subscriber.events:
assert e.new_instance_status == Instance.QUEUED
# Get launched nodes.
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert reply.status.code == StatusCode.OK
assert len(reply.state.instances) == 3
instance_ids = [ins.instance_id for ins in reply.state.instances]
types_count = defaultdict(int)
for ins in reply.state.instances:
types_count[ins.instance_type] += 1
assert ins.status == Instance.QUEUED
assert types_count["type-1"] == 1
assert types_count["type-2"] == 2
# Update node status.
subscriber.clear()
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=1,
updates=[
InstanceUpdateEvent(
instance_id=instance_ids[0],
new_instance_status=Instance.REQUESTED,
instance_type="type-1",
launch_request_id="l1",
),
InstanceUpdateEvent(
instance_id=instance_ids[1],
new_instance_status=Instance.REQUESTED,
launch_request_id="l1",
instance_type="type-1",
),
],
)
)
assert reply.status.code == StatusCode.OK
assert len(subscriber.events) == 2
for e in subscriber.events:
assert e.new_instance_status == Instance.REQUESTED
# Get updated nodes.
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert reply.status.code == StatusCode.OK
assert len(reply.state.instances) == 3
types_count = defaultdict(int)
for ins in reply.state.instances:
types_count[ins.instance_type] += 1
if ins.instance_id in [instance_ids[0], instance_ids[1]]:
assert ins.status == Instance.REQUESTED
else:
assert ins.status == Instance.QUEUED
# Invalid instances status update.
subscriber.clear()
with pytest.raises(AssertionError):
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=2,
updates=[
InstanceUpdateEvent(
instance_id=instance_ids[2],
# Not requested yet.
new_instance_status=Instance.RAY_RUNNING,
),
],
)
)
assert len(subscriber.events) == 0
# Invalid versions.
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=0, # Invalid version, outdated.
updates=[
InstanceUpdateEvent(
instance_id=instance_ids[2],
new_instance_status=Instance.REQUESTED,
instance_type="type-2",
),
],
)
)
assert reply.status.code == StatusCode.VERSION_MISMATCH
assert len(subscriber.events) == 0
def test_insert(self):
ins_storage = InstanceStorage(
"cluster-id",
InMemoryStorage(),
)
subscriber = MockSubscriber()
im = InstanceManager(
ins_storage, instance_status_update_subscribers=[subscriber]
)
im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=0,
updates=[
InstanceUpdateEvent(
instance_type="type-1",
instance_id="id-1",
new_instance_status=Instance.QUEUED,
upsert=True,
),
InstanceUpdateEvent(
instance_id="id-2",
new_instance_status=Instance.TERMINATING,
cloud_instance_id="cloud-id-2",
upsert=True,
),
InstanceUpdateEvent(
instance_id="id-3",
new_instance_status=Instance.ALLOCATED,
cloud_instance_id="cloud-id-3",
node_kind=NodeKind.WORKER,
instance_type="type-3",
upsert=True,
),
],
)
)
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert len(reply.state.instances) == 3
instance_by_ids = {ins.instance_id: ins for ins in reply.state.instances}
assert instance_by_ids["id-1"].status == Instance.QUEUED
assert instance_by_ids["id-1"].instance_type == "type-1"
assert instance_by_ids["id-2"].status == Instance.TERMINATING
assert instance_by_ids["id-3"].status == Instance.ALLOCATED
assert instance_by_ids["id-3"].cloud_instance_id == "cloud-id-3"
version = reply.state.version
# With non-upsert flags.
with pytest.raises(AssertionError):
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=version,
updates=[
InstanceUpdateEvent(
instance_type="type-1",
instance_id="id-999",
new_instance_status=Instance.QUEUED,
),
],
)
)
# With invalid statuses
all_statuses = set(Instance.InstanceStatus.values())
non_insertable_statuses = all_statuses - {
Instance.QUEUED,
Instance.TERMINATING,
Instance.ALLOCATED,
}
for status in non_insertable_statuses:
subscriber.clear()
with pytest.raises(AssertionError):
reply = im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=version,
updates=[
InstanceUpdateEvent(
instance_id="id-999",
new_instance_status=status,
),
],
)
)
assert len(subscriber.events) == 0
def test_apply_update(self):
ins_storage = InstanceStorage(
"cluster-id",
InMemoryStorage(),
)
subscriber = MockSubscriber()
im = InstanceManager(
ins_storage, instance_status_update_subscribers=[subscriber]
)
# Insert a new instance.
im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=0,
updates=[
InstanceUpdateEvent(
instance_type="type-1",
instance_id="id-1",
new_instance_status=Instance.QUEUED,
upsert=True,
),
],
)
)
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert len(reply.state.instances) == 1
assert reply.state.instances[0].status == Instance.QUEUED
assert reply.state.instances[0].instance_type == "type-1"
# Request
im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=1,
updates=[
InstanceUpdateEvent(
instance_id="id-1",
new_instance_status=Instance.REQUESTED,
launch_request_id="l1",
instance_type="type-1",
),
],
)
)
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert len(reply.state.instances) == 1
assert reply.state.instances[0].status == Instance.REQUESTED
assert reply.state.instances[0].launch_request_id == "l1"
# ALLOCATED
im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=2,
updates=[
InstanceUpdateEvent(
instance_id="id-1",
new_instance_status=Instance.ALLOCATED,
cloud_instance_id="cloud-id-1",
node_kind=NodeKind.WORKER,
instance_type="type-1",
),
],
)
)
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert len(reply.state.instances) == 1
assert reply.state.instances[0].status == Instance.ALLOCATED
assert reply.state.instances[0].cloud_instance_id == "cloud-id-1"
# RAY_RUNNING
im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=3,
updates=[
InstanceUpdateEvent(
instance_id="id-1",
new_instance_status=Instance.RAY_RUNNING,
ray_node_id="ray-node-1",
),
],
)
)
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert len(reply.state.instances) == 1
assert reply.state.instances[0].status == Instance.RAY_RUNNING
assert reply.state.instances[0].node_id == "ray-node-1"
# TERMINATED
im.update_instance_manager_state(
UpdateInstanceManagerStateRequest(
expected_version=4,
updates=[
InstanceUpdateEvent(
instance_id="id-1",
new_instance_status=Instance.TERMINATED,
),
],
)
)
reply = im.get_instance_manager_state(GetInstanceManagerStateRequest())
assert len(reply.state.instances) == 1
assert reply.state.instances[0].status == Instance.TERMINATED
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,220 @@
# coding: utf-8
import copy
import os
import sys
from unittest import mock
import pytest # noqa
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
from ray.core.generated.instance_manager_pb2 import Instance
def create_instance(instance_id, status=Instance.QUEUED, version=0):
return Instance(
instance_id=instance_id,
status=status,
version=version,
)
@mock.patch("time.time", mock.MagicMock(return_value=1))
def test_upsert():
storage = InstanceStorage(
cluster_id="test_cluster",
storage=InMemoryStorage(),
)
instance1 = create_instance("instance1")
instance2 = create_instance("instance2")
instance3 = create_instance("instance3")
assert (True, 1) == storage.batch_upsert_instances(
[instance1, instance2],
expected_storage_version=None,
)
instance1.version = 1
instance2.version = 1
entries, storage_version = storage.get_instances()
assert storage_version == 1
assert entries == {
"instance1": instance1,
"instance2": instance2,
}
assert (False, 1) == storage.batch_upsert_instances(
[create_instance("instance1"), create_instance("instance2")],
expected_storage_version=0,
)
instance2.status = Instance.ALLOCATED
assert (True, 2) == storage.batch_upsert_instances(
[instance3, instance2],
expected_storage_version=1,
)
instance1.version = 1
instance2.version = 2
instance3.version = 2
entries, storage_version = storage.get_instances()
assert storage_version == 2
assert entries == {
"instance1": instance1,
"instance2": instance2,
"instance3": instance3,
}
@mock.patch("time.time", mock.MagicMock(return_value=1))
def test_update():
storage = InstanceStorage(
cluster_id="test_cluster",
storage=InMemoryStorage(),
)
instance1 = create_instance("instance1")
instance2 = create_instance("instance2")
assert (True, 1) == storage.upsert_instance(instance=instance1)
assert (True, 2) == storage.upsert_instance(instance=instance2)
assert (
{
"instance1": create_instance("instance1", version=1),
"instance2": create_instance("instance2", version=2),
},
2,
) == storage.get_instances()
# failed because instance version is not correct
assert (False, 2) == storage.upsert_instance(
instance=instance1,
expected_instance_version=0,
)
# failed because storage version is not correct
assert (False, 2) == storage.upsert_instance(
instance=instance1,
expected_storage_version=0,
)
assert (True, 3) == storage.upsert_instance(
instance=instance2,
expected_storage_version=2,
)
assert (
{
"instance1": create_instance("instance1", version=1),
"instance2": create_instance("instance2", version=3),
},
3,
) == storage.get_instances()
assert (True, 4) == storage.upsert_instance(
instance=instance1,
expected_instance_version=1,
)
assert (
{
"instance1": create_instance("instance1", version=4),
"instance2": create_instance("instance2", version=3),
},
4,
) == storage.get_instances()
@mock.patch("time.time", mock.MagicMock(return_value=1))
def test_delete():
storage = InstanceStorage(
cluster_id="test_cluster",
storage=InMemoryStorage(),
)
instance1 = create_instance("instance1")
instance2 = create_instance("instance2")
instance3 = create_instance("instance3")
assert (True, 1) == storage.batch_upsert_instances(
[instance1, instance2, instance3],
expected_storage_version=None,
)
assert (False, 1) == storage.batch_delete_instances(
instance_ids=["instance1"], expected_storage_version=0
)
assert (True, 2) == storage.batch_delete_instances(instance_ids=["instance1"])
assert (
{
"instance2": create_instance("instance2", version=1),
"instance3": create_instance("instance3", version=1),
},
2,
) == storage.get_instances()
assert (True, 3) == storage.batch_delete_instances(
instance_ids=["instance2"], expected_storage_version=2
)
assert (
{
"instance3": create_instance("instance3", version=1),
},
3,
) == storage.get_instances()
@mock.patch("time.time", mock.MagicMock(return_value=1))
def test_get_instances():
storage = InstanceStorage(
cluster_id="test_cluster",
storage=InMemoryStorage(),
)
instance1 = create_instance("instance1", version=1)
instance2 = create_instance("instance2", status=Instance.ALLOCATED, version=1)
instance3 = create_instance("instance3", status=Instance.TERMINATING, version=1)
assert (True, 1) == storage.batch_upsert_instances(
[copy.deepcopy(instance1), copy.deepcopy(instance2), copy.deepcopy(instance3)],
expected_storage_version=None,
)
assert (
{
"instance1": instance1,
"instance2": instance2,
"instance3": instance3,
},
1,
) == storage.get_instances()
assert (
{
"instance1": instance1,
"instance2": instance2,
},
1,
) == storage.get_instances(instance_ids=["instance1", "instance2"])
assert ({"instance2": instance2}, 1) == storage.get_instances(
instance_ids=["instance1", "instance2"], status_filter={Instance.ALLOCATED}
)
assert (
{
"instance2": instance2,
},
1,
) == storage.get_instances(status_filter={Instance.ALLOCATED})
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,336 @@
import os
import sys
import unittest
# coding: utf-8
from unittest.mock import patch
import pytest
from ray.autoscaler.v2.instance_manager.common import InstanceUtil
from ray.core.generated.instance_manager_pb2 import Instance
class InstanceUtilTest(unittest.TestCase):
def test_basic(self):
# New instance.
instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
assert instance.instance_id == "i-123"
assert instance.instance_type == "type_1"
assert instance.status == Instance.QUEUED
# Set status.
assert InstanceUtil.set_status(instance, Instance.REQUESTED)
assert instance.status == Instance.REQUESTED
# Set status with invalid status.
assert not InstanceUtil.set_status(instance, Instance.RAY_RUNNING)
assert not InstanceUtil.set_status(instance, Instance.UNKNOWN)
def test_transition_graph(self):
# Assert on each edge in the graph.
all_status = set(Instance.InstanceStatus.values())
g = InstanceUtil.get_valid_transitions()
assert g[Instance.QUEUED] == {Instance.REQUESTED, Instance.TERMINATED}
all_status.remove(Instance.QUEUED)
assert g[Instance.REQUESTED] == {
Instance.ALLOCATED,
Instance.QUEUED,
Instance.ALLOCATION_FAILED,
}
all_status.remove(Instance.REQUESTED)
assert g[Instance.ALLOCATED] == {
Instance.RAY_INSTALLING,
Instance.RAY_RUNNING,
Instance.RAY_STOPPING,
Instance.RAY_STOPPED,
Instance.TERMINATING,
Instance.TERMINATED,
Instance.ALLOCATION_TIMEOUT,
}
all_status.remove(Instance.ALLOCATED)
assert g[Instance.RAY_INSTALLING] == {
Instance.RAY_RUNNING,
Instance.RAY_INSTALL_FAILED,
Instance.RAY_STOPPED,
Instance.TERMINATING,
Instance.TERMINATED,
}
all_status.remove(Instance.RAY_INSTALLING)
assert g[Instance.RAY_RUNNING] == {
Instance.RAY_STOP_REQUESTED,
Instance.RAY_STOPPING,
Instance.RAY_STOPPED,
Instance.TERMINATING,
Instance.TERMINATED,
}
all_status.remove(Instance.RAY_RUNNING)
assert g[Instance.ALLOCATION_TIMEOUT] == {
Instance.TERMINATING,
Instance.TERMINATED,
}
all_status.remove(Instance.ALLOCATION_TIMEOUT)
assert g[Instance.RAY_STOP_REQUESTED] == {
Instance.RAY_STOPPING,
Instance.RAY_STOPPED,
Instance.TERMINATED,
Instance.RAY_RUNNING,
}
all_status.remove(Instance.RAY_STOP_REQUESTED)
assert g[Instance.RAY_STOPPING] == {
Instance.RAY_STOPPED,
Instance.TERMINATING,
Instance.TERMINATED,
}
all_status.remove(Instance.RAY_STOPPING)
assert g[Instance.RAY_STOPPED] == {Instance.TERMINATED, Instance.TERMINATING}
all_status.remove(Instance.RAY_STOPPED)
assert g[Instance.TERMINATING] == {
Instance.TERMINATED,
Instance.TERMINATION_FAILED,
}
all_status.remove(Instance.TERMINATING)
assert g[Instance.TERMINATION_FAILED] == {
Instance.TERMINATING,
Instance.TERMINATED,
}
all_status.remove(Instance.TERMINATION_FAILED)
assert g[Instance.TERMINATED] == set()
all_status.remove(Instance.TERMINATED)
assert g[Instance.ALLOCATION_FAILED] == set()
all_status.remove(Instance.ALLOCATION_FAILED)
assert g[Instance.RAY_INSTALL_FAILED] == {
Instance.TERMINATED,
Instance.TERMINATING,
}
all_status.remove(Instance.RAY_INSTALL_FAILED)
assert g[Instance.UNKNOWN] == set()
all_status.remove(Instance.UNKNOWN)
assert len(all_status) == 0
@patch("time.time_ns")
def test_status_time(self, mock_time):
mock_time.return_value = 1
instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
# OK
assert (
InstanceUtil.get_status_transition_times_ns(instance, Instance.QUEUED)[0]
== 1
)
# No filter.
assert InstanceUtil.get_status_transition_times_ns(
instance,
) == [1]
# Missing status returns empty list
assert (
InstanceUtil.get_status_transition_times_ns(instance, Instance.REQUESTED)
== []
)
# Multiple status.
mock_time.return_value = 2
InstanceUtil.set_status(instance, Instance.REQUESTED)
mock_time.return_value = 3
InstanceUtil.set_status(instance, Instance.QUEUED)
mock_time.return_value = 4
InstanceUtil.set_status(instance, Instance.REQUESTED)
assert InstanceUtil.get_status_transition_times_ns(
instance, Instance.QUEUED
) == [1, 3]
@patch("time.time_ns")
def test_get_last_status_transition(self, mock_time):
mock_time.return_value = 1
instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
assert (
InstanceUtil.get_last_status_transition(instance).instance_status
== Instance.QUEUED
)
assert InstanceUtil.get_last_status_transition(instance).timestamp_ns == 1
mock_time.return_value = 2
InstanceUtil.set_status(instance, Instance.REQUESTED)
assert (
InstanceUtil.get_last_status_transition(instance).instance_status
== Instance.REQUESTED
)
assert InstanceUtil.get_last_status_transition(instance).timestamp_ns == 2
mock_time.return_value = 3
InstanceUtil.set_status(instance, Instance.QUEUED)
assert (
InstanceUtil.get_last_status_transition(instance).instance_status
== Instance.QUEUED
)
assert InstanceUtil.get_last_status_transition(instance).timestamp_ns == 3
assert (
InstanceUtil.get_last_status_transition(
instance, select_instance_status=Instance.REQUESTED
).instance_status
== Instance.REQUESTED
)
assert (
InstanceUtil.get_last_status_transition(
instance, select_instance_status=Instance.REQUESTED
).timestamp_ns
== 2
)
assert (
InstanceUtil.get_last_status_transition(
instance, select_instance_status=Instance.RAY_RUNNING
)
is None
)
def test_is_cloud_instance_allocated(self):
all_status = set(Instance.InstanceStatus.values())
instance = InstanceUtil.new_instance("i-123", "type_1", Instance.QUEUED)
positive_status = {
Instance.ALLOCATED,
Instance.RAY_INSTALLING,
Instance.RAY_INSTALL_FAILED,
Instance.RAY_RUNNING,
Instance.RAY_STOP_REQUESTED,
Instance.RAY_STOPPING,
Instance.RAY_STOPPED,
Instance.TERMINATING,
Instance.TERMINATION_FAILED,
Instance.ALLOCATION_TIMEOUT,
}
for s in positive_status:
instance.status = s
assert InstanceUtil.is_cloud_instance_allocated(instance.status)
all_status.remove(s)
# Unknown not possible.
all_status.remove(Instance.UNKNOWN)
for s in all_status:
instance.status = s
assert not InstanceUtil.is_cloud_instance_allocated(instance.status)
def test_is_ray_running(self):
all_statuses = set(Instance.InstanceStatus.values())
positive_statuses = {
Instance.RAY_RUNNING,
Instance.RAY_STOP_REQUESTED,
Instance.RAY_STOPPING,
}
all_statuses.remove(Instance.UNKNOWN)
for s in positive_statuses:
assert InstanceUtil.is_ray_running(s)
all_statuses.remove(s)
for s in all_statuses:
assert not InstanceUtil.is_ray_running(s)
def test_is_ray_pending(self):
all_statuses = set(Instance.InstanceStatus.values())
all_statuses.remove(Instance.UNKNOWN)
positive_statuses = {
Instance.QUEUED,
Instance.REQUESTED,
Instance.RAY_INSTALLING,
Instance.ALLOCATED,
}
for s in positive_statuses:
assert InstanceUtil.is_ray_pending(s), Instance.InstanceStatus.Name(s)
all_statuses.remove(s)
for s in all_statuses:
assert not InstanceUtil.is_ray_pending(s), Instance.InstanceStatus.Name(s)
def test_is_ray_running_reachable(self):
all_status = set(Instance.InstanceStatus.values())
positive_status = {
Instance.QUEUED,
Instance.REQUESTED,
Instance.ALLOCATED,
Instance.RAY_INSTALLING,
Instance.RAY_RUNNING,
Instance.RAY_STOP_REQUESTED,
}
for s in positive_status:
assert InstanceUtil.is_ray_running_reachable(
s
), Instance.InstanceStatus.Name(s)
all_status.remove(s)
# Unknown not possible.
all_status.remove(Instance.UNKNOWN)
for s in all_status:
assert not InstanceUtil.is_ray_running_reachable(
s
), Instance.InstanceStatus.Name(s)
def test_reachable_from(self):
def add_reachable_from(reachable, src, transitions):
reachable[src] = set()
for dst in transitions[src]:
reachable[src].add(dst)
reachable[src] |= (
reachable[dst] if reachable[dst] is not None else set()
)
expected_reachable = {s: None for s in Instance.InstanceStatus.values()}
# Error status and terminal status.
expected_reachable[Instance.ALLOCATION_FAILED] = set()
expected_reachable[Instance.UNKNOWN] = set()
expected_reachable[Instance.TERMINATED] = set()
transitions = InstanceUtil.get_valid_transitions()
# Recursively build the reachable set from terminal statuses.
add_reachable_from(expected_reachable, Instance.TERMINATION_FAILED, transitions)
add_reachable_from(expected_reachable, Instance.TERMINATING, transitions)
# Add TERMINATION_FAILED again since it's also reachable from TERMINATING.
add_reachable_from(expected_reachable, Instance.TERMINATION_FAILED, transitions)
add_reachable_from(expected_reachable, Instance.RAY_STOPPED, transitions)
add_reachable_from(expected_reachable, Instance.RAY_STOPPING, transitions)
add_reachable_from(expected_reachable, Instance.RAY_STOP_REQUESTED, transitions)
add_reachable_from(expected_reachable, Instance.RAY_RUNNING, transitions)
# Add RAY_STOP_REQUESTED again since it's also reachable from RAY_RUNNING.
add_reachable_from(expected_reachable, Instance.RAY_STOP_REQUESTED, transitions)
add_reachable_from(expected_reachable, Instance.RAY_INSTALL_FAILED, transitions)
add_reachable_from(expected_reachable, Instance.RAY_INSTALLING, transitions)
add_reachable_from(expected_reachable, Instance.ALLOCATED, transitions)
add_reachable_from(expected_reachable, Instance.REQUESTED, transitions)
add_reachable_from(expected_reachable, Instance.QUEUED, transitions)
# Add REQUESTED again since it's also reachable from QUEUED.
add_reachable_from(expected_reachable, Instance.REQUESTED, transitions)
add_reachable_from(expected_reachable, Instance.ALLOCATION_TIMEOUT, transitions)
for s, expected in expected_reachable.items():
assert InstanceUtil.get_reachable_statuses(s) == expected, (
f"reachable_from({s}) = {InstanceUtil.get_reachable_statuses(s)} "
f"!= {expected}"
)
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,219 @@
import os
import sys
from typing import List
import pytest
from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
from ray.autoscaler.v2.tests.util import create_instance
from ray.core.generated.instance_manager_pb2 import Instance
def _get_metrics(metrics, name) -> List[float]:
sample_values = []
for x in metrics:
for sample in x.samples:
if sample.name == name:
sample_values.append(sample.value)
return sample_values
def test_report_nodes_resources():
"""
Test that the metrics reporter reports the correct number of nodes and resources
"""
reporter = AutoscalerMetricsReporter(
AutoscalerPrometheusMetrics(session_name="test")
)
node_type_configs = {
"type_1": NodeTypeConfig(
name="type_1",
max_worker_nodes=10,
min_worker_nodes=1,
resources={"CPU": 1},
),
"type_2": NodeTypeConfig(
name="type_2",
max_worker_nodes=10,
min_worker_nodes=1,
resources={"GPU": 1},
),
}
_i = 0
def id():
nonlocal _i
_i += 1
return f"i-{_i}"
terminating_type_1 = create_instance(
id(), status=Instance.TERMINATING, instance_type="type_1"
)
terminating_type_2 = create_instance(
id(), status=Instance.TERMINATING, instance_type="type_2"
)
instances = [
# Active = 3
create_instance(id(), status=Instance.RAY_RUNNING, instance_type="type_1"),
create_instance(
id(), status=Instance.RAY_STOP_REQUESTED, instance_type="type_1"
),
create_instance(id(), status=Instance.RAY_STOPPING, instance_type="type_1"),
create_instance(id(), status=Instance.RAY_RUNNING, instance_type="type_2"),
# Pending
create_instance(id(), status=Instance.QUEUED, instance_type="type_1"),
create_instance(id(), status=Instance.REQUESTED, instance_type="type_1"),
create_instance(id(), status=Instance.RAY_INSTALLING, instance_type="type_1"),
create_instance(id(), status=Instance.ALLOCATED, instance_type="type_1"),
create_instance(id(), status=Instance.RAY_INSTALLING, instance_type="type_2"),
create_instance(id(), status=Instance.ALLOCATED, instance_type="type_2"),
# Terminating
terminating_type_1,
terminating_type_2,
]
reporter.report_instances(instances, node_type_configs)
assert _get_metrics(
reporter._prom_metrics.active_nodes.labels(
SessionName="test", NodeType="type_1"
).collect(),
"autoscaler_active_nodes",
) == [3]
assert _get_metrics(
reporter._prom_metrics.pending_nodes.labels(
SessionName="test", NodeType="type_1"
).collect(),
"autoscaler_pending_nodes",
) == [4]
assert _get_metrics(
reporter._prom_metrics.recently_failed_nodes.labels(
SessionName="test", NodeType="type_1"
).collect(),
"autoscaler_recently_failed_nodes",
) == [1]
# Test that resources are reported correctly
reporter.report_resources(instances, node_type_configs)
assert _get_metrics(
reporter._prom_metrics.cluster_resources.labels(
SessionName="test", resource="CPU"
).collect(),
"autoscaler_cluster_resources",
) == [3]
assert _get_metrics(
reporter._prom_metrics.cluster_resources.labels(
SessionName="test", resource="GPU"
).collect(),
"autoscaler_cluster_resources",
) == [1]
assert _get_metrics(
reporter._prom_metrics.pending_resources.labels(
SessionName="test", resource="CPU"
).collect(),
"autoscaler_pending_resources",
) == [4]
assert _get_metrics(
reporter._prom_metrics.pending_resources.labels(
SessionName="test", resource="GPU"
).collect(),
"autoscaler_pending_resources",
) == [2]
def test_report_nodes_resources_handles_deleted_node_type():
reporter = AutoscalerMetricsReporter(
AutoscalerPrometheusMetrics(session_name="test_deleted")
)
node_type_configs = {
"current_type": NodeTypeConfig(
name="current_type",
max_worker_nodes=10,
min_worker_nodes=1,
resources={"CPU": 2},
),
}
instances = [
create_instance(
"current-running",
status=Instance.RAY_RUNNING,
instance_type="current_type",
),
create_instance(
"deleted-terminated",
status=Instance.TERMINATED,
instance_type="deleted_type",
),
create_instance(
"deleted-terminating",
status=Instance.TERMINATING,
instance_type="deleted_type",
),
create_instance(
"deleted-pending",
status=Instance.QUEUED,
instance_type="deleted_type",
),
]
reporter.report_instances(instances, node_type_configs)
reporter.report_resources(instances, node_type_configs)
assert _get_metrics(
reporter._prom_metrics.active_nodes.labels(
SessionName="test_deleted", NodeType="current_type"
).collect(),
"autoscaler_active_nodes",
) == [1]
assert _get_metrics(
reporter._prom_metrics.pending_nodes.labels(
SessionName="test_deleted", NodeType="current_type"
).collect(),
"autoscaler_pending_nodes",
) == [0]
assert _get_metrics(
reporter._prom_metrics.recently_failed_nodes.labels(
SessionName="test_deleted", NodeType="current_type"
).collect(),
"autoscaler_recently_failed_nodes",
) == [0]
assert _get_metrics(
reporter._prom_metrics.active_nodes.labels(
SessionName="test_deleted", NodeType="deleted_type"
).collect(),
"autoscaler_active_nodes",
) == [0]
assert _get_metrics(
reporter._prom_metrics.pending_nodes.labels(
SessionName="test_deleted", NodeType="deleted_type"
).collect(),
"autoscaler_pending_nodes",
) == [1]
assert _get_metrics(
reporter._prom_metrics.recently_failed_nodes.labels(
SessionName="test_deleted", NodeType="deleted_type"
).collect(),
"autoscaler_recently_failed_nodes",
) == [1]
assert _get_metrics(
reporter._prom_metrics.cluster_resources.labels(
SessionName="test_deleted", resource="CPU"
).collect(),
"autoscaler_cluster_resources",
) == [2]
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,328 @@
import os
import sys
import time
import pytest
from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
from ray.autoscaler.v2.instance_manager.subscribers.cloud_resource_monitor import (
CloudResourceMonitor,
)
from ray.autoscaler.v2.scheduler import (
ResourceDemandScheduler,
ResourceRequestSource,
SchedulingNode,
SchedulingNodeStatus,
SchedulingRequest,
)
from ray.core.generated.autoscaler_pb2 import ResourceRequest as PBResourceRequest
from ray.core.generated.instance_manager_pb2 import (
Instance,
InstanceUpdateEvent,
NodeKind,
)
def test_recovery_scoring():
monitor = CloudResourceMonitor()
node_type = "gpu-node"
# Mock failure
event = InstanceUpdateEvent(
instance_type=node_type, new_instance_status=Instance.ALLOCATION_TIMEOUT
)
monitor.notify([event])
# Immediately after failure, score should be 0.0
scores = monitor.get_recoverable_resource_availabilities()
assert scores[node_type] == 0.0
# After safety floor (e.g., 11s, default safety floor is 10s)
monitor._last_unavailable_timestamp[node_type] = time.time() - 11
scores = monitor.get_recoverable_resource_availabilities()
assert 0.0 < scores[node_type] < 0.1
# Halfway through recovery window (600s / 2 = 300s)
monitor._last_unavailable_timestamp[node_type] = time.time() - 300
scores = monitor.get_recoverable_resource_availabilities()
# 300 / 600 = 0.5
assert pytest.approx(scores[node_type], 0.01) == 0.5
# After recovery window
monitor._last_unavailable_timestamp[node_type] = time.time() - 601
scores = monitor.get_recoverable_resource_availabilities()
assert scores[node_type] == 1.0
def test_scheduler_priority_tie_breaking():
# Two node types with identical resources
resources = {"CPU": 4}
node_type_1 = "high-priority"
node_type_2 = "low-priority"
config_1 = NodeTypeConfig(
name=node_type_1,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=10,
)
config_2 = NodeTypeConfig(
name=node_type_2,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=0,
)
node_1 = SchedulingNode.from_node_config(
config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
node_2 = SchedulingNode.from_node_config(
config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
request = PBResourceRequest(resources_bundle={"CPU": 1})
# Utilization and Availability are equal (all perfect)
# Priority should break the tie
best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
[request], [node_2, node_1], ResourceRequestSource.PENDING_DEMAND, {}, {}
)
assert best_node.node_type == node_type_1
def test_schedule_context_propagation():
resources = {"CPU": 4}
node_type = "gpu-node"
config = NodeTypeConfig(
name=node_type,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=7,
)
# Mock cloud availabilities
cloud_availabilities = {node_type: 0.1} # failure recency
recoverable_availabilities = {node_type: 0.5}
req = SchedulingRequest(
node_type_configs={node_type: config},
cloud_resource_availabilities=cloud_availabilities,
recoverable_resource_availabilities=recoverable_availabilities,
disable_launch_config_check=True,
)
ctx = ResourceDemandScheduler.ScheduleContext.from_schedule_request(req)
# Check if a new node created from this context has the correct priority
node_pools = [
SchedulingNode.from_node_config(
ctx.get_node_type_configs()[nt],
status=SchedulingNodeStatus.TO_LAUNCH,
node_kind=NodeKind.WORKER,
)
for nt, num_available in ctx.get_node_type_available().items()
]
assert len(node_pools) == 1
node = node_pools[0]
assert node.priority == 7
# Dynamic scores are in the context, not the node
assert ctx.get_recoverable_resource_availabilities()[node_type] == 0.5
assert ctx.get_cloud_resource_availabilities()[node_type] == 0.1
def test_scheduler_availability_over_priority():
# High priority node is recovering (score 0.5)
# Low priority node is available (score 1.0)
resources = {"CPU": 4}
node_type_1 = "high-priority"
node_type_2 = "low-priority"
config_1 = NodeTypeConfig(
name=node_type_1,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=10,
)
config_2 = NodeTypeConfig(
name=node_type_2,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=0,
)
node_1 = SchedulingNode.from_node_config(
config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
node_2 = SchedulingNode.from_node_config(
config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
request = PBResourceRequest(resources_bundle={"CPU": 1})
# Recoverable Availability is higher for node_2, so it should be chosen despite lower priority
best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
[request],
[node_1, node_2],
ResourceRequestSource.PENDING_DEMAND,
cloud_resource_availabilities={},
recoverable_resource_availabilities={node_type_1: 0.5, node_type_2: 1.0},
)
assert best_node.node_type == node_type_2
def test_scheduler_failure_recency_tie_breaking():
# Same priority, same recoverable availability (1.0)
# One has an older failure than the other.
resources = {"CPU": 4}
node_type_1 = "older-failure"
node_type_2 = "newer-failure"
config_1 = NodeTypeConfig(
name=node_type_1,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=5,
)
config_2 = NodeTypeConfig(
name=node_type_2,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=5,
)
node_1 = SchedulingNode.from_node_config(
config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
node_2 = SchedulingNode.from_node_config(
config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
request = PBResourceRequest(resources_bundle={"CPU": 1})
best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
[request],
[node_1, node_2],
ResourceRequestSource.PENDING_DEMAND,
cloud_resource_availabilities={node_type_1: 0.9, node_type_2: 0.1},
recoverable_resource_availabilities={node_type_1: 1.0, node_type_2: 1.0},
)
assert best_node.node_type == node_type_1
def test_recovery_integration():
monitor = CloudResourceMonitor()
node_type_1 = "high-priority"
node_type_2 = "low-priority"
# Mock failure for node_1 (high priority)
event = InstanceUpdateEvent(
instance_type=node_type_1, new_instance_status=Instance.ALLOCATION_TIMEOUT
)
monitor.notify([event])
# Score should be 0.0 for node_1 immediately
scores = monitor.get_recoverable_resource_availabilities()
assert scores[node_type_1] == 0.0
# Setup scheduler structures
resources = {"CPU": 4}
config_1 = NodeTypeConfig(
name=node_type_1,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=10,
)
config_2 = NodeTypeConfig(
name=node_type_2,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources,
priority=0,
)
node_1 = SchedulingNode.from_node_config(
config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
node_2 = SchedulingNode.from_node_config(
config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
request = PBResourceRequest(resources_bundle={"CPU": 1})
# Pass scores from monitor to scheduler
best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
[request],
[node_1, node_2],
ResourceRequestSource.PENDING_DEMAND,
cloud_resource_availabilities={},
recoverable_resource_availabilities=scores,
)
# Node 2 should be chosen because Node 1 is recovering, despite Node 1 having higher priority
assert best_node.node_type == node_type_2
def test_scheduler_utilization_over_priority():
# Node 1: 4 CPUs, priority 10
# Node 2: 2 CPUs, priority 0
# Request: 2 CPUs
# Node 2 should be selected because it fits perfectly (utilization score is higher),
# even though Node 1 has higher priority.
resources_1 = {"CPU": 4}
resources_2 = {"CPU": 2}
node_type_1 = "large-node"
node_type_2 = "small-node"
config_1 = NodeTypeConfig(
name=node_type_1,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources_1,
priority=10,
)
config_2 = NodeTypeConfig(
name=node_type_2,
min_worker_nodes=0,
max_worker_nodes=10,
resources=resources_2,
priority=0,
)
node_1 = SchedulingNode.from_node_config(
config_1, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
node_2 = SchedulingNode.from_node_config(
config_2, status=SchedulingNodeStatus.TO_LAUNCH, node_kind=NodeKind.WORKER
)
request = PBResourceRequest(resources_bundle={"CPU": 2})
best_node, infeasible, remaining_nodes = ResourceDemandScheduler._sched_best_node(
[request],
[node_1, node_2],
ResourceRequestSource.PENDING_DEMAND,
cloud_resource_availabilities={},
recoverable_resource_availabilities={},
)
assert best_node.node_type == node_type_2
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,68 @@
# coding: utf-8
import os
import sys
import unittest
import pytest # noqa
from ray._private.test_utils import load_test_config
from ray.autoscaler.tags import TAG_RAY_NODE_KIND
from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
from ray.core.generated.instance_manager_pb2 import Instance
from ray.tests.autoscaler_test_utils import MockProcessRunner, MockProvider
class RayInstallerTest(unittest.TestCase):
def setUp(self):
self.base_provider = MockProvider()
self.config = AutoscalingConfig(load_test_config("test_ray_complex.yaml"))
self.runner = MockProcessRunner()
self.ray_installer = RayInstaller(self.base_provider, self.config, self.runner)
def test_install_succeeded(self):
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
self.ray_installer.install_ray(
Instance(
instance_id="0", instance_type="worker_nodes1", cloud_instance_id="0"
),
head_node_ip="1.2.3.4",
)
def test_install_failed(self):
# creation failed because no such node.
with self.assertRaisesRegex(KeyError, "0"):
assert not self.ray_installer.install_ray(
Instance(
instance_id="0",
instance_type="worker_nodes1",
cloud_instance_id="0",
),
head_node_ip="1.2.3.4",
)
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
self.runner.fail_cmds = [
"echo" # this is the command used in the test_ray_complex.yaml
]
self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
# creation failed because setup command failed.
with self.assertRaisesRegex(Exception, "unexpected status"):
self.ray_installer.install_ray(
Instance(
instance_id="0",
instance_type="worker_nodes1",
cloud_instance_id="0",
),
head_node_ip="1.2.3.4",
)
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,232 @@
import os
import sys
import time
import pytest
from ray.autoscaler.v2.schema import (
AutoscalerInstance,
ClusterStatus,
IPPRGroupSpec,
IPPRStatus,
)
from ray.core.generated.autoscaler_pb2 import NodeState, NodeStatus
from ray.core.generated.instance_manager_pb2 import Instance
def test_cluster_status_default_stats():
status = ClusterStatus()
assert status.active_nodes == []
assert status.idle_nodes == []
assert status.pending_launches == []
assert status.failed_launches == []
assert status.pending_nodes == []
assert status.failed_nodes == []
assert status.cluster_resource_usage == []
assert status.stats.gcs_request_time_s == 0.0
assert status.stats.request_ts_s is None
def _make_ippr_status() -> IPPRStatus:
return IPPRStatus(
cloud_instance_id="ray-worker-1",
spec=IPPRGroupSpec(
min_cpu=1.0,
max_cpu=4.0,
min_memory=2,
max_memory=8,
resize_timeout=10,
),
current_cpu=1.0,
current_memory=2,
desired_cpu=1.0,
desired_memory=2,
)
def test_autoscaler_instance():
i = AutoscalerInstance()
assert not i.validate()[0], "Empty instance should be invalid"
i = AutoscalerInstance(im_instance=Instance(status=Instance.QUEUED))
assert i.validate()[0], "Instance with only im_instance should be valid"
i = AutoscalerInstance(
ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
cloud_instance_id="i-123",
)
assert i.validate()[0], i.validate()[1]
i = AutoscalerInstance(
ray_node=NodeState(status=NodeStatus.RUNNING),
)
assert not i.validate()[0], "Missing cloud node id."
i = AutoscalerInstance(
im_instance=Instance(status=Instance.QUEUED),
ray_node=NodeState(status=NodeStatus.RUNNING),
)
assert not i.validate()[
0
], "cloud node id is required to link the ray node with im state"
i = AutoscalerInstance(
cloud_instance_id="i-123",
)
assert not i.validate()[
0
], "cloud instance id is not possible without im or ray node"
i = AutoscalerInstance(
im_instance=Instance(status=Instance.ALLOCATED, cloud_instance_id="i-123"),
cloud_instance_id="i-123",
)
assert i.validate()[0]
i = AutoscalerInstance(
im_instance=Instance(status=Instance.ALLOCATED, cloud_instance_id="i-123"),
cloud_instance_id="i-124", # mismatch.
)
assert not i.validate()[0], "cloud instance id should match"
i = AutoscalerInstance(
im_instance=Instance(status=Instance.QUEUED, cloud_instance_id="i-123"),
cloud_instance_id="i-123",
)
assert not i.validate()[0], "cloud instance id is not possible with queued state"
i = AutoscalerInstance(
im_instance=Instance(status=Instance.ALLOCATED, cloud_instance_id="i-123"),
)
assert not i.validate()[0], "cloud instance id should also be set"
i = AutoscalerInstance(
ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
cloud_instance_id="i-123",
)
assert i.validate()[0]
i = AutoscalerInstance(
ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
cloud_instance_id="i-124", # mismatch.
)
assert not i.validate()[0]
i = AutoscalerInstance(
im_instance=Instance(status=Instance.RAY_RUNNING, cloud_instance_id="i-123"),
ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
cloud_instance_id="i-123",
)
assert i.validate()[0]
i = AutoscalerInstance(
im_instance=Instance(status=Instance.RAY_RUNNING, cloud_instance_id="i-123"),
ray_node=NodeState(status=NodeStatus.RUNNING, instance_id="i-123"),
cloud_instance_id="i-124", # mismatch.
)
assert not i.validate()[0]
def test_ippr_status_queue_resize_request():
status = _make_ippr_status()
status.resizing_at = 123
status.k8s_resize_status = "deferred"
status.k8s_resize_message = "pending"
status.raylet_id = "abc"
assert status.queue_resize_request(desired_cpu=2.0, desired_memory=4)
assert status.desired_cpu == 2.0
assert status.desired_memory == 4
assert status.resizing_at is None
assert status.k8s_resize_status == "new"
assert status.k8s_resize_message is None
assert not status.queue_resize_request(desired_cpu=2.0, desired_memory=4)
def test_ippr_status_request_and_progress_helpers():
status = _make_ippr_status()
assert status.has_resize_request_to_send() is False
assert status.is_in_progress() is False
status.raylet_id = "abc"
assert status.is_k8s_resize_finished()
assert not status.has_resize_request_to_send()
assert not status.is_in_progress()
status.queue_resize_request(desired_cpu=2.0)
assert status.has_resize_request_to_send()
assert status.is_in_progress()
assert not status.is_k8s_resize_finished()
status.resizing_at = int(time.time())
status.k8s_resize_status = None
assert status.is_in_progress()
assert status.is_k8s_resize_finished()
def test_ippr_status_need_sync_with_raylet():
status = _make_ippr_status()
assert status.need_sync_with_raylet() is False
status.raylet_id = "abc"
status.resizing_at = int(time.time())
status.k8s_resize_status = None
assert status.need_sync_with_raylet()
status.desired_cpu = 2.0
assert not status.need_sync_with_raylet()
def test_ippr_status_limits_and_can_resize_up():
status = _make_ippr_status()
assert status.can_resize_up()
status.raylet_id = "abc"
assert status.max_cpu() == 4.0
assert status.max_memory() == 8
assert status.can_resize_up()
status.suggested_max_cpu = 1.5
status.suggested_max_memory = 3
assert status.max_cpu() == 1.5
assert status.max_memory() == 3
status.current_cpu = 1.5
status.current_memory = 3
assert not status.can_resize_up()
status.current_cpu = 1.0
status.last_failed_at = 1
assert not status.can_resize_up()
def test_ippr_status_failure_and_timeout_helpers():
status = _make_ippr_status()
status.raylet_id = "abc"
status.desired_cpu = 2.0
status.desired_memory = 4
status.resizing_at = int(time.time()) - 20
assert status.is_timeout()
status.k8s_resize_status = "error"
assert status.is_errored()
status.record_failure("resize failed", failed_at=123)
assert status.last_failed_at == 123
assert status.last_failed_reason == "resize failed"
# K8s resize finished and resources match desired, but raylet sync is still
# pending: do not treat as timeout (provider will call GCS to sync).
status2 = _make_ippr_status()
status2.raylet_id = "abc"
status2.resizing_at = int(time.time()) - 20
status2.k8s_resize_status = None
assert status2.need_sync_with_raylet()
assert not status2.is_timeout()
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
+995
View File
@@ -0,0 +1,995 @@
import os
import sys
import time
# coding: utf-8
from dataclasses import dataclass
from typing import Callable, List, Optional, Tuple
import pytest
import ray
import ray._private.ray_constants as ray_constants
from ray._common.test_utils import wait_for_condition
from ray._private import authentication_test_utils
from ray.autoscaler.v2.schema import (
ClusterStatus,
LaunchRequest,
NodeInfo,
ResourceRequestByCount,
)
from ray.autoscaler.v2.sdk import (
get_cluster_status,
request_cluster_resources,
)
from ray.autoscaler.v2.tests.util import (
get_available_resources,
get_cluster_resource_state,
get_total_resources,
report_autoscaling_state,
)
from ray.core.generated import autoscaler_pb2, autoscaler_pb2_grpc
from ray.core.generated.autoscaler_pb2 import ClusterResourceState, NodeStatus
from ray.core.generated.common_pb2 import LabelSelectorOperator
from ray.util.state.api import list_nodes
def _autoscaler_state_service_stub():
"""Get the grpc stub for the autoscaler state service"""
from ray._private.grpc_utils import init_grpc_channel
gcs_address = ray.get_runtime_context().gcs_address
gcs_channel = init_grpc_channel(gcs_address, ray_constants.GLOBAL_GRPC_OPTIONS)
return autoscaler_pb2_grpc.AutoscalerStateServiceStub(gcs_channel)
def get_node_ids() -> Tuple[str, List[str]]:
"""Get the node ids of the head node and a worker node"""
head_node_id = None
nodes = list_nodes()
worker_node_ids = []
for node in nodes:
if node.is_head_node:
head_node_id = node.node_id
else:
worker_node_ids += [node.node_id]
return head_node_id, worker_node_ids
def assert_cluster_resource_constraints(
state: ClusterResourceState, expected_bundles: List[dict], expected_count: List[int]
):
"""
Assert a GetClusterResourceStateReply has cluster_resource_constraints that
matches with the expected resources.
"""
# We only have 1 constraint for now.
assert len(state.cluster_resource_constraints) == 1
resource_requests = state.cluster_resource_constraints[0].resource_requests
assert len(resource_requests) == len(expected_bundles) == len(expected_count)
# Sort all the bundles by bundle's resource names
resource_requests = sorted(
resource_requests,
key=lambda bundle_by_count: "".join(
bundle_by_count.request.resources_bundle.keys()
),
)
expected = zip(expected_bundles, expected_count)
expected = sorted(
expected, key=lambda bundle_count: "".join(bundle_count[0].keys())
)
for actual_bundle_count, expected_bundle_count in zip(resource_requests, expected):
assert (
dict(actual_bundle_count.request.resources_bundle)
== expected_bundle_count[0]
)
assert actual_bundle_count.count == expected_bundle_count[1]
@dataclass
class ExpectedNodeState:
node_id: str
node_status: NodeStatus
idle_time_check_cb: Optional[Callable] = None
labels: Optional[dict] = None
def assert_node_states(
state: ClusterResourceState, expected_nodes: List[ExpectedNodeState]
):
"""
Assert a GetClusterResourceStateReply has node states that
matches with the expected nodes.
"""
assert len(state.node_states) == len(expected_nodes)
# Sort all the nodes by node's node_id
node_states = sorted(state.node_states, key=lambda node: node.node_id)
expected_nodes = sorted(expected_nodes, key=lambda node: node.node_id)
for actual_node, expected_node in zip(node_states, expected_nodes):
assert actual_node.status == expected_node.node_status
if expected_node.idle_time_check_cb:
assert expected_node.idle_time_check_cb(actual_node.idle_duration_ms)
if expected_node.labels:
assert sorted(actual_node.dynamic_labels) == sorted(expected_node.labels)
@dataclass
class ExpectedNodeInfo:
node_id: Optional[str] = None
node_status: Optional[str] = None
idle_time_check_cb: Optional[Callable] = None
instance_id: Optional[str] = None
ray_node_type_name: Optional[str] = None
instance_type_name: Optional[str] = None
ip_address: Optional[str] = None
details: Optional[str] = None
# Check those resources are included in the actual node info.
total_resources: Optional[dict] = None
available_resources: Optional[dict] = None
def assert_nodes(actual_nodes: List[NodeInfo], expected_nodes: List[ExpectedNodeInfo]):
assert len(actual_nodes) == len(expected_nodes)
# Sort the nodes by id.
actual_nodes = sorted(actual_nodes, key=lambda node: node.node_id)
expected_nodes = sorted(expected_nodes, key=lambda node: node.node_id)
for actual_node, expected_node in zip(actual_nodes, expected_nodes):
if expected_node.node_id is not None:
assert actual_node.node_id == expected_node.node_id
if expected_node.node_status is not None:
assert actual_node.node_status == expected_node.node_status
if expected_node.instance_id is not None:
assert actual_node.instance_id == expected_node.instance_id
if expected_node.ray_node_type_name is not None:
assert actual_node.ray_node_type_name == expected_node.ray_node_type_name
if expected_node.instance_type_name is not None:
assert actual_node.instance_type_name == expected_node.instance_type_name
if expected_node.ip_address is not None:
assert actual_node.ip_address == expected_node.ip_address
if expected_node.details is not None:
assert expected_node.details in actual_node.details
if expected_node.idle_time_check_cb:
assert expected_node.idle_time_check_cb(
actual_node.resource_usage.idle_time_ms
)
if expected_node.total_resources:
for resource_name, total in expected_node.total_resources.items():
assert (
total
== get_total_resources(actual_node.resource_usage.usage)[
resource_name
]
)
if expected_node.available_resources:
for resource_name, available in expected_node.available_resources.items():
assert (
available
== get_available_resources(actual_node.resource_usage.usage)[
resource_name
]
)
def assert_launches(
cluster_status: ClusterStatus,
expected_pending_launches: List[LaunchRequest],
expected_failed_launches: List[LaunchRequest],
):
def assert_launches(actuals, expects):
for actual, expect in zip(actuals, expects):
assert actual.instance_type_name == expect.instance_type_name
assert actual.ray_node_type_name == expect.ray_node_type_name
assert actual.count == expect.count
assert actual.state == expect.state
assert actual.request_ts_s == expect.request_ts_s
assert len(cluster_status.pending_launches) == len(expected_pending_launches)
assert len(cluster_status.failed_launches) == len(expected_failed_launches)
actual_pending = sorted(
cluster_status.pending_launches, key=lambda launch: launch.ray_node_type_name
)
expected_pending = sorted(
expected_pending_launches, key=lambda launch: launch.ray_node_type_name
)
assert_launches(actual_pending, expected_pending)
actual_failed = sorted(
cluster_status.failed_launches, key=lambda launch: launch.ray_node_type_name
)
expected_failed = sorted(
expected_failed_launches, key=lambda launch: launch.ray_node_type_name
)
assert_launches(actual_failed, expected_failed)
@dataclass
class GangResourceRequest:
# Resource bundles.
bundles: List[dict]
# List of detail information about the request
details: List[str]
def assert_gang_requests(
state: ClusterResourceState, expected: List[GangResourceRequest]
):
"""
Assert a GetClusterResourceStateReply has gang requests that
matches with the expected requests.
"""
assert len(state.pending_gang_resource_requests) == len(expected)
# Sort all the requests by request's details
requests = sorted(
state.pending_gang_resource_requests, key=lambda request: request.details
)
expected = sorted(expected, key=lambda request: "".join(request.details))
for actual_request, expected_request in zip(requests, expected):
# Assert the detail contains the expected details
for detail_str in expected_request.details:
assert detail_str in actual_request.details
def test_request_cluster_resources_basic(shutdown_only):
ctx = ray.init(num_cpus=1)
stub = _autoscaler_state_service_stub()
gcs_address = ctx.address_info["gcs_address"]
# Request one
request_cluster_resources(gcs_address, [{"resources": {"CPU": 1}}])
def verify():
state = get_cluster_resource_state(stub)
assert_cluster_resource_constraints(state, [{"CPU": 1}], [1])
return True
wait_for_condition(verify)
# Request another overrides the previous request
request_cluster_resources(
gcs_address, [{"resources": {"CPU": 2, "GPU": 1}}, {"resources": {"CPU": 1}}]
)
def verify():
state = get_cluster_resource_state(stub)
assert_cluster_resource_constraints(
state, [{"CPU": 2, "GPU": 1}, {"CPU": 1}], [1, 1]
)
return True
# Request multiple is aggregated by shape.
request_cluster_resources(gcs_address, [{"resources": {"CPU": 1}}] * 100)
def verify():
state = get_cluster_resource_state(stub)
assert_cluster_resource_constraints(state, [{"CPU": 1}], [100])
return True
wait_for_condition(verify)
def test_request_cluster_resources_with_label_selectors(shutdown_only):
ctx = ray.init(num_cpus=1)
stub = _autoscaler_state_service_stub()
gcs_address = ctx.address_info["gcs_address"]
# Define two bundles, each with its own label_selector, to request.
bundles = [
{"CPU": 1},
{"GPU": 1, "CPU": 2},
]
bundle_label_selectors = [
{"region": "us-west1"},
{"accelerator-type": "!in(A100)"},
]
to_request = [
{"resources": b, "label_selector": s}
for b, s in zip(bundles, bundle_label_selectors)
]
# Send the request for these resource bundles
request_cluster_resources(gcs_address, to_request)
def verify():
state = get_cluster_resource_state(stub)
# Validate shape and resource request count
assert_cluster_resource_constraints(state, bundles, [1, 1])
# Check that requests carry expected label selectors
requests = state.cluster_resource_constraints[0].resource_requests
# First resource request
label_selectors_0 = requests[0].request.label_selectors
selector_0 = label_selectors_0[0]
constraints_0 = {
c.label_key: list(c.label_values) for c in selector_0.label_constraints
}
assert constraints_0 == {"region": ["us-west1"]}
assert (
selector_0.label_constraints[0].operator
== LabelSelectorOperator.LABEL_OPERATOR_IN
)
# Second resource request
label_selectors_1 = requests[1].request.label_selectors
selector_1 = label_selectors_1[0]
constraints_1 = {
c.label_key: list(c.label_values) for c in selector_1.label_constraints
}
assert constraints_1 == {"accelerator-type": ["A100"]}
assert (
selector_1.label_constraints[0].operator
== LabelSelectorOperator.LABEL_OPERATOR_NOT_IN
)
return True
wait_for_condition(verify)
def test_node_info_basic(shutdown_only, monkeypatch):
with monkeypatch.context() as m:
m.setenv("RAY_CLOUD_INSTANCE_ID", "instance-id")
m.setenv("RAY_NODE_TYPE_NAME", "node-type-name")
m.setenv("RAY_CLOUD_INSTANCE_TYPE_NAME", "instance-type-name")
ctx = ray.init(num_cpus=1)
ip = ctx.address_info["node_ip_address"]
stub = _autoscaler_state_service_stub()
def verify():
state = get_cluster_resource_state(stub)
assert len(state.node_states) == 1
node = state.node_states[0]
assert node.instance_id == "instance-id"
assert node.ray_node_type_name == "node-type-name"
assert node.node_ip_address == ip
assert node.instance_type_name == "instance-type-name"
assert (
state.cluster_session_name
== ray._private.worker.global_worker.node.session_name
)
return True
wait_for_condition(verify)
def test_pg_pending_gang_requests_basic(shutdown_only):
ray.init(num_cpus=1)
# Create a pg that's pending.
pg = ray.util.placement_group([{"CPU": 1}] * 3, strategy="STRICT_SPREAD")
try:
ray.get(pg.ready(), timeout=2)
except TimeoutError:
pass
pg_id = pg.id.hex()
stub = _autoscaler_state_service_stub()
def verify():
state = get_cluster_resource_state(stub)
assert_gang_requests(
state,
[
GangResourceRequest(
[{"CPU": 1}] * 3, details=[pg_id, "STRICT_SPREAD", "PENDING"]
)
],
)
return True
wait_for_condition(verify)
def test_pg_usage_labels(shutdown_only):
ray.init(num_cpus=1)
# Create a pg
pg = ray.util.placement_group([{"CPU": 1}])
ray.get(pg.ready())
# Check the labels
stub = _autoscaler_state_service_stub()
head_node_id, _ = get_node_ids()
pg_id = pg.id.hex()
def verify():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
head_node_id,
NodeStatus.RUNNING,
labels={f"_PG_{pg_id}": ""},
),
],
)
return True
wait_for_condition(verify)
def test_node_state_lifecycle_basic(ray_start_cluster):
start_s = time.perf_counter()
cluster = ray_start_cluster
cluster.add_node(num_cpus=0)
ray.init(address=cluster.address)
node = cluster.add_node(num_cpus=1)
stub = _autoscaler_state_service_stub()
# We don't have node id from `add_node` unfortunately.
def nodes_up():
nodes = list_nodes()
assert len(nodes) == 2
return True
wait_for_condition(nodes_up)
head_node_id, worker_node_ids = get_node_ids()
node_id = worker_node_ids[0]
def verify_cluster_idle():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
),
ExpectedNodeState(
head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
),
],
)
return True
wait_for_condition(verify_cluster_idle)
# Schedule a task running
@ray.remote(num_cpus=0.1)
def f():
while True:
pass
t = f.remote()
def verify_cluster_busy():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
node_id, NodeStatus.RUNNING, lambda idle_ms: idle_ms == 0
),
ExpectedNodeState(
head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
),
],
)
return True
wait_for_condition(verify_cluster_busy)
# Kill the task
ray.cancel(t, force=True)
wait_for_condition(verify_cluster_idle)
# Kill the node.
cluster.remove_node(node)
# Sleep for a bit so head node should be idle longer than this.
time.sleep(3)
def verify_cluster_no_node():
state = get_cluster_resource_state(stub)
now_s = time.perf_counter()
test_dur_ms = (now_s - start_s) * 1000
assert_node_states(
state,
[
ExpectedNodeState(node_id, NodeStatus.DEAD),
ExpectedNodeState(
head_node_id,
NodeStatus.IDLE,
lambda idle_ms: idle_ms > 3 * 1000 and idle_ms < test_dur_ms,
),
],
)
return True
wait_for_condition(verify_cluster_no_node)
# We test that a node with only workers blocked on get
# is considered idle.
def test_idle_node_blocked(ray_start_cluster):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1)
ray.init(address=cluster.address)
stub = _autoscaler_state_service_stub()
# We don't have node id from `add_node` unfortunately.
def nodes_up():
nodes = list_nodes()
assert len(nodes) == 1
return True
wait_for_condition(nodes_up)
head_node_id = get_node_ids()
def verify_cluster_idle():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
),
],
)
return True
wait_for_condition(verify_cluster_idle)
# Unschedulable
@ray.remote(num_cpus=10000)
def f():
pass
# Schedule a task running
@ray.remote(num_cpus=1)
def g():
ray.get(f.remote())
t = g.remote()
def verify_cluster_busy():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
head_node_id, NodeStatus.RUNNING, lambda idle_ms: idle_ms == 0
),
],
)
return True
wait_for_condition(verify_cluster_busy)
for _ in range(10):
time.sleep(0.5)
verify_cluster_busy()
# Kill the task
ray.cancel(t, force=True)
wait_for_condition(verify_cluster_idle)
def test_idle_node_no_resource(ray_start_cluster):
cluster = ray_start_cluster
cluster.add_node(num_cpus=1)
ray.init(address=cluster.address)
stub = _autoscaler_state_service_stub()
# We don't have node id from `add_node` unfortunately.
def nodes_up():
nodes = list_nodes()
assert len(nodes) == 1
return True
wait_for_condition(nodes_up)
head_node_id = get_node_ids()
def verify_cluster_idle():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
head_node_id, NodeStatus.IDLE, lambda idle_ms: idle_ms > 0
),
],
)
return True
wait_for_condition(verify_cluster_idle)
# Schedule a task running
@ray.remote(num_cpus=0)
def f():
while True:
pass
t = f.remote()
def verify_cluster_busy():
state = get_cluster_resource_state(stub)
assert_node_states(
state,
[
ExpectedNodeState(
head_node_id, NodeStatus.RUNNING, lambda idle_ms: idle_ms == 0
),
],
)
return True
wait_for_condition(verify_cluster_busy)
# Kill the task
ray.cancel(t, force=True)
wait_for_condition(verify_cluster_idle)
def test_get_cluster_status_resources(ray_start_cluster):
cluster = ray_start_cluster
# Head node
cluster.add_node(num_cpus=1, _system_config={"enable_autoscaler_v2": True})
ray.init(address=cluster.address)
# Worker node
cluster.add_node(num_cpus=2)
@ray.remote(num_cpus=1)
class Actor:
def loop(self):
while True:
pass
# Schedule tasks to use all resources.
@ray.remote(num_cpus=1)
def loop():
while True:
pass
[loop.remote() for _ in range(2)]
actor = Actor.remote()
actor.loop.remote()
def verify_cpu_resources_all_used():
cluster_status = get_cluster_status(cluster.address)
total_cluster_resources = get_total_resources(
cluster_status.cluster_resource_usage
)
assert total_cluster_resources["CPU"] == 3.0
available_cluster_resources = get_available_resources(
cluster_status.cluster_resource_usage
)
assert available_cluster_resources["CPU"] == 0.0
return True
wait_for_condition(verify_cpu_resources_all_used)
# Schedule more tasks should show up as task demands
[loop.remote() for _ in range(2)]
def verify_task_demands():
resource_demands = get_cluster_status(cluster.address).resource_demands
assert len(resource_demands.ray_task_actor_demand) == 1
assert resource_demands.ray_task_actor_demand[0].bundles_by_count == [
ResourceRequestByCount(
bundle={"CPU": 1.0},
count=2,
)
]
return True
wait_for_condition(verify_task_demands)
# Request resources through SDK
request_cluster_resources(
gcs_address=cluster.address, to_request=[{"resources": {"GPU": 1, "CPU": 2}}]
)
def verify_cluster_constraint_demand():
resource_demands = get_cluster_status(cluster.address).resource_demands
assert len(resource_demands.cluster_constraint_demand) == 1
assert resource_demands.cluster_constraint_demand[0].bundles_by_count == [
ResourceRequestByCount(
bundle={"GPU": 1.0, "CPU": 2.0},
count=1,
)
]
return True
wait_for_condition(verify_cluster_constraint_demand)
# Try to schedule some PGs
pg1 = ray.util.placement_group([{"CPU": 1}] * 3)
def verify_pg_demands():
resource_demands = get_cluster_status(cluster.address).resource_demands
assert len(resource_demands.placement_group_demand) == 1
assert resource_demands.placement_group_demand[0].bundles_by_count == [
ResourceRequestByCount(
bundle={"CPU": 1.0},
count=3,
)
]
assert resource_demands.placement_group_demand[0].pg_id == pg1.id.hex()
assert resource_demands.placement_group_demand[0].strategy == "PACK"
assert resource_demands.placement_group_demand[0].state == "PENDING"
return True
wait_for_condition(verify_pg_demands)
def test_get_cluster_status(ray_start_cluster):
# This test is to make sure the grpc stub is working.
# TODO(rickyx): Add e2e tests for the autoscaler state service in a separate PR
# to validate the data content.
cluster = ray_start_cluster
# Head node
cluster.add_node(num_cpus=1, _system_config={"enable_autoscaler_v2": True})
ray.init(address=cluster.address)
# Worker node
cluster.add_node(num_cpus=2)
head_node_id, worker_node_ids = get_node_ids()
def verify_nodes():
cluster_status = get_cluster_status(cluster.address)
assert_nodes(
cluster_status.idle_nodes,
[
ExpectedNodeInfo(
worker_node_ids[0],
"IDLE",
lambda idle_ms: idle_ms > 0,
total_resources={"CPU": 2.0},
available_resources={"CPU": 2.0},
),
ExpectedNodeInfo(
head_node_id,
"IDLE",
lambda idle_ms: idle_ms > 0,
total_resources={"CPU": 1.0},
available_resources={"CPU": 1.0},
),
],
)
return True
wait_for_condition(verify_nodes)
# Schedule a task running
@ray.remote(num_cpus=2)
def f():
while True:
pass
f.remote()
def verify_nodes_busy():
cluster_status = get_cluster_status(cluster.address)
assert_nodes(
cluster_status.idle_nodes,
[
ExpectedNodeInfo(head_node_id, "IDLE", lambda idle_ms: idle_ms > 0),
],
)
assert_nodes(
cluster_status.active_nodes,
[
ExpectedNodeInfo(
worker_node_ids[0],
"RUNNING",
lambda idle_ms: idle_ms == 0,
total_resources={"CPU": 2.0},
available_resources={"CPU": 0.0},
),
],
)
return True
wait_for_condition(verify_nodes_busy)
stub = _autoscaler_state_service_stub()
state = autoscaler_pb2.AutoscalingState(
last_seen_cluster_resource_state_version=0,
# since the autoscaler will also update the autoscaler_state_version periodically,
# we need to use a large number here, such as 10, to override it to avoid flaky test.
autoscaler_state_version=10,
pending_instance_requests=[
autoscaler_pb2.PendingInstanceRequest(
instance_type_name="m5.large",
ray_node_type_name="worker",
count=2,
request_ts=1000,
)
],
failed_instance_requests=[
autoscaler_pb2.FailedInstanceRequest(
instance_type_name="m5.large",
ray_node_type_name="worker",
count=2,
start_ts=1000,
failed_ts=2000,
reason="insufficient quota",
)
],
pending_instances=[
autoscaler_pb2.PendingInstance(
instance_id="instance-id",
instance_type_name="m5.large",
ray_node_type_name="worker",
ip_address="10.10.10.10",
details="launching",
)
],
)
report_autoscaling_state(stub, autoscaling_state=state)
def verify_autoscaler_state():
# TODO(rickyx): Add infeasible asserts.
cluster_status = get_cluster_status(cluster.address)
assert len(cluster_status.pending_launches) == 1
assert_launches(
cluster_status,
expected_pending_launches=[
LaunchRequest(
instance_type_name="m5.large",
ray_node_type_name="worker",
count=2,
state=LaunchRequest.Status.PENDING,
request_ts_s=1000,
)
],
expected_failed_launches=[
LaunchRequest(
instance_type_name="m5.large",
ray_node_type_name="worker",
count=2,
state=LaunchRequest.Status.FAILED,
request_ts_s=1000,
failed_ts_s=2000,
details="insufficient quota",
)
],
)
assert_nodes(
cluster_status.pending_nodes,
[
ExpectedNodeInfo(
instance_id="instance-id",
ray_node_type_name="worker",
details="launching",
ip_address="10.10.10.10",
)
],
)
return True
wait_for_condition(verify_autoscaler_state)
@pytest.mark.parametrize(
"env_val,enabled",
[
("1", True),
("0", False),
("", False),
],
)
def test_is_autoscaler_v2_enabled(shutdown_only, monkeypatch, env_val, enabled):
def reset_autoscaler_v2_enabled_cache():
import ray.autoscaler.v2.utils as u
u.cached_is_autoscaler_v2 = None
reset_autoscaler_v2_enabled_cache()
with monkeypatch.context() as m:
m.setenv("RAY_enable_autoscaler_v2", env_val)
ray.init()
def verify():
assert ray.autoscaler.v2.utils.is_autoscaler_v2() == enabled
return True
wait_for_condition(verify)
@pytest.mark.parametrize(
"token_state,setup_token,should_fail",
[
("valid", lambda: None, False),
("invalid", lambda: _setup_invalid_token(), True),
],
)
def test_autoscaler_api_with_token_auth(
setup_cluster_with_token_auth,
cleanup_auth_token_env,
token_state,
setup_token,
should_fail,
):
"""Parametrized test for autoscaler API with different token states.
Tests request_cluster_resources with valid, invalid, and missing tokens.
"""
# Setup token state (this changes the client-side token)
setup_token()
if should_fail:
# API call should fail with invalid token
with pytest.raises(Exception) as exc_info:
request_cluster_resources(
ray.get_runtime_context().gcs_address,
[{"resources": {"CPU": 1}, "label_selector": {}}],
)
# Verify it's an authentication error
error_str = str(exc_info.value).lower()
assert (
"unauthenticated" in error_str or "invalidauthtoken" in error_str
), f"request_cluster_resources with {token_state} token should return auth error, got: {exc_info.value}"
else:
# API call should succeed with valid token
request_cluster_resources(
ray.get_runtime_context().gcs_address,
[{"resources": {"CPU": 1}, "label_selector": {}}],
)
# Verify the request was successful using the autoscaler state service stub
stub = _autoscaler_state_service_stub()
state = get_cluster_resource_state(stub)
assert (
len(state.cluster_resource_constraints) > 0
), f"request_cluster_resources with {token_state} token should succeed"
def _setup_invalid_token():
"""Helper to set up an invalid authentication token."""
invalid_token = "invalid_token_value"
authentication_test_utils.set_env_auth_token(invalid_token)
authentication_test_utils.reset_auth_token_state()
def _clear_token():
"""Helper to clear authentication token sources."""
authentication_test_utils.clear_auth_token_sources()
authentication_test_utils.reset_auth_token_state()
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,87 @@
# coding: utf-8
import os
import sys
import pytest # noqa
from ray.autoscaler.v2.instance_manager.storage import (
InMemoryStorage,
StoreStatus,
VersionedValue,
)
@pytest.mark.parametrize("storage", [InMemoryStorage()])
def test_storage(storage):
assert storage.get_version() == 0
assert storage.get_all(table="test_table") == ({}, 0)
assert storage.get(table="test_table", keys=[]) == ({}, 0)
assert storage.get(table="test_table", keys=["key1"]) == ({}, 0)
assert storage.batch_update(
table="test_table", mutation={"key1": "value1"}
) == StoreStatus(
True,
1,
)
assert storage.get_version() == 1
assert storage.get_all(table="test_table") == (
{"key1": VersionedValue("value1", 1)},
1,
)
assert storage.get(table="test_table", keys=[]) == (
{"key1": VersionedValue("value1", 1)},
1,
)
assert storage.batch_update(
table="test_table", mutation={"key1": "value2"}, expected_version=0
) == StoreStatus(False, 1)
assert storage.batch_update(
table="test_table", mutation={"key1": "value2"}, expected_version=1
) == StoreStatus(True, 2)
assert storage.get_all(table="test_table") == (
{"key1": VersionedValue("value2", 2)},
2,
)
assert storage.batch_update(
table="test_table",
mutation={"key2": "value3", "key3": "value4"},
deletion=["key1"],
expected_version=2,
) == StoreStatus(True, 3)
assert storage.get_all(table="test_table") == (
{"key2": VersionedValue("value3", 3), "key3": VersionedValue("value4", 3)},
3,
)
assert storage.get(table="test_table", keys=["key2", "key1"]) == (
{"key2": VersionedValue("value3", 3)},
3,
)
assert storage.update(
table="test_table", key="key2", value="value5"
) == StoreStatus(True, 4)
assert storage.update(
table="test_table", key="key2", value="value5", insert_only=True
) == StoreStatus(False, 4)
assert storage.update(
table="test_table", key="key2", value="value5", expected_entry_version=3
) == StoreStatus(False, 4)
assert storage.update(
table="test_table", key="key2", value="value6", expected_entry_version=4
) == StoreStatus(True, 5)
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,392 @@
# coding: utf-8
import os
import sys
from queue import Queue
from unittest import mock
import pytest
from ray._common.test_utils import wait_for_condition
from ray._common.utils import binary_to_hex, hex_to_binary
from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
from ray.autoscaler.v2.instance_manager.cloud_providers.read_only.cloud_provider import (
ReadOnlyProvider,
)
from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import (
CloudInstanceUpdater,
)
from ray.autoscaler.v2.instance_manager.subscribers.ray_stopper import ( # noqa
RayStopper,
)
from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
from ray.core.generated.autoscaler_pb2 import DrainNodeReason
from ray.core.generated.instance_manager_pb2 import (
Instance,
InstanceUpdateEvent,
TerminationRequest,
)
def _get_stopped_nodes_total(metrics_reporter: AutoscalerMetricsReporter) -> float:
total_samples = [
sample.value
for metric in metrics_reporter._prom_metrics.stopped_nodes.collect()
for sample in metric.samples
if sample.name == "autoscaler_stopped_nodes_total"
]
assert len(total_samples) == 1
return total_samples[0]
class TestRayStopper:
def test_no_op(self):
mock_gcs_client = mock.MagicMock()
ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=Queue())
ray_stopper.notify(
[
InstanceUpdateEvent(
instance_id="test_id",
new_instance_status=Instance.REQUESTED,
)
]
)
assert mock_gcs_client.drain_node.call_count == 0
@pytest.mark.parametrize(
"drain_accepted",
[True, False, None],
ids=["drain_accepted", "drain_rejected", "drain_error"],
)
def test_idle_termination(self, drain_accepted):
mock_gcs_client = mock.MagicMock()
if drain_accepted is None:
mock_gcs_client.drain_node.side_effect = Exception("error")
else:
mock_gcs_client.drain_node.return_value = (
drain_accepted,
f"accepted={str(drain_accepted)}",
)
error_queue = Queue()
ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=error_queue)
ray_stopper.notify(
[
InstanceUpdateEvent(
instance_id="test_id",
new_instance_status=Instance.RAY_STOP_REQUESTED,
termination_request=TerminationRequest(
cause=TerminationRequest.Cause.IDLE,
idle_duration_ms=1000,
ray_node_id="0000",
),
)
]
)
def verify():
mock_gcs_client.drain_node.assert_has_calls(
[
mock.call(
node_id="0000",
reason=DrainNodeReason.DRAIN_NODE_REASON_IDLE_TERMINATION,
reason_message=(
"Termination of node that's idle for 1.0 seconds."
),
deadline_timestamp_ms=0,
)
]
)
if drain_accepted:
assert error_queue.empty()
else:
error = error_queue.get_nowait()
assert error.im_instance_id == "test_id"
return True
wait_for_condition(verify)
@pytest.mark.parametrize(
"stop_accepted",
[True, False],
ids=["stop_accepted", "stop_rejected"],
)
def test_preemption(self, stop_accepted):
mock_gcs_client = mock.MagicMock()
mock_gcs_client.drain_nodes.return_value = [0] if stop_accepted else []
error_queue = Queue()
ray_stopper = RayStopper(gcs_client=mock_gcs_client, error_queue=error_queue)
ray_stopper.notify(
[
InstanceUpdateEvent(
instance_id="i-1",
new_instance_status=Instance.RAY_STOP_REQUESTED,
termination_request=TerminationRequest(
cause=TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE,
max_num_nodes_per_type=10,
ray_node_id=binary_to_hex(hex_to_binary(b"1111")),
),
),
InstanceUpdateEvent(
instance_id="i-2",
new_instance_status=Instance.RAY_STOP_REQUESTED,
termination_request=TerminationRequest(
cause=TerminationRequest.Cause.MAX_NUM_NODES,
max_num_nodes=100,
ray_node_id=binary_to_hex(hex_to_binary(b"2222")),
),
),
]
)
def verify():
mock_gcs_client.drain_nodes.assert_has_calls(
[
mock.call(
node_ids=[hex_to_binary(b"1111")],
),
mock.call(
node_ids=[hex_to_binary(b"2222")],
),
]
)
if stop_accepted:
assert error_queue.empty()
else:
error_in_ids = set()
while not error_queue.empty():
error = error_queue.get_nowait()
error_in_ids.add(error.im_instance_id)
assert error_in_ids == {"i-1", "i-2"}
return True
wait_for_condition(verify)
class TestCloudInstanceUpdater:
def test_launch_no_op(self):
mock_provider = mock.MagicMock()
launcher = CloudInstanceUpdater(mock_provider)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.RAY_RUNNING,
launch_request_id="1",
instance_type="type-1",
),
]
)
mock_provider.launch.assert_not_called()
def test_launch_new_instances(self):
mock_provider = mock.MagicMock()
launcher = CloudInstanceUpdater(mock_provider)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.REQUESTED,
launch_request_id="1",
instance_type="type-1",
),
InstanceUpdateEvent(
new_instance_status=Instance.REQUESTED,
launch_request_id="1",
instance_type="type-1",
),
InstanceUpdateEvent(
new_instance_status=Instance.REQUESTED,
launch_request_id="2",
instance_type="type-1",
),
InstanceUpdateEvent(
new_instance_status=Instance.REQUESTED,
launch_request_id="2",
instance_type="type-2",
),
]
)
def verify():
mock_provider.launch.assert_has_calls(
[
mock.call(shape={"type-1": 2}, request_id="1"),
mock.call(shape={"type-1": 1, "type-2": 1}, request_id="2"),
]
)
return True
wait_for_condition(verify)
def test_multi_notify(self):
mock_provider = mock.MagicMock()
launcher = CloudInstanceUpdater(mock_provider)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.REQUESTED,
launch_request_id="1",
instance_type="type-1",
),
]
)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.REQUESTED,
launch_request_id="2",
instance_type="type-1",
),
]
)
def verify():
assert mock_provider.launch.call_count == 2
mock_provider.launch.assert_has_calls(
[
mock.call(shape={"type-1": 1}, request_id="1"),
mock.call(shape={"type-1": 1}, request_id="2"),
]
)
return True
wait_for_condition(verify)
def test_terminate_no_op(self):
mock_provider = mock.MagicMock()
launcher = CloudInstanceUpdater(mock_provider)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.RAY_RUNNING,
instance_id="1",
cloud_instance_id="c1",
),
]
)
def verify():
mock_provider.terminate.assert_not_called()
return True
wait_for_condition(verify)
def test_terminate_instances(self):
mock_provider = mock.MagicMock()
metrics_reporter = AutoscalerMetricsReporter(
AutoscalerPrometheusMetrics(session_name="test")
)
launcher = CloudInstanceUpdater(mock_provider, metrics_reporter)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATING,
instance_id="1",
cloud_instance_id="c1",
),
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATING,
instance_id="2",
cloud_instance_id="c2",
),
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATING,
instance_id="3",
cloud_instance_id="c3",
),
]
)
def verify():
mock_provider.terminate.assert_called_once_with(
ids=["c1", "c2", "c3"], request_id=mock.ANY
)
assert _get_stopped_nodes_total(metrics_reporter) == 0
return True
wait_for_condition(verify)
def test_count_stopped_instances_on_terminated(self):
mock_provider = mock.MagicMock()
metrics_reporter = AutoscalerMetricsReporter(
AutoscalerPrometheusMetrics(session_name="test")
)
launcher = CloudInstanceUpdater(mock_provider, metrics_reporter)
launcher.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATED,
instance_id="1",
cloud_instance_id="c1",
),
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATED,
instance_id="2",
cloud_instance_id="c2",
),
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATED,
instance_id="3",
),
]
)
def verify():
mock_provider.terminate.assert_not_called()
assert _get_stopped_nodes_total(metrics_reporter) == 2
return True
wait_for_condition(verify)
class TestReadOnlyProvider:
def test_terminate_raises_not_implemented_with_correct_interface(self):
"""ReadOnlyProvider.terminate() must accept ids and request_id kwargs.
Regression test for:
TypeError: ReadOnlyProvider.terminate() got an unexpected keyword
argument 'ids'
CloudInstanceUpdater calls provider.terminate(ids=..., request_id=...)
which matches the ICloudInstanceProvider interface. ReadOnlyProvider
must accept the same signature even though it raises NotImplementedError.
"""
provider = object.__new__(ReadOnlyProvider) # skip __init__ (needs GCS)
with pytest.raises(NotImplementedError):
provider.terminate(ids=["node-1", "node-2"], request_id="req-1")
def test_terminate_via_cloud_instance_updater_raises_not_implemented(self):
"""Verify the full call path: CloudInstanceUpdater -> ReadOnlyProvider.
When a TERMINATING event arrives, CloudInstanceUpdater calls
provider.terminate(ids=..., request_id=...). With ReadOnlyProvider this
should surface as NotImplementedError, not TypeError.
"""
provider = object.__new__(ReadOnlyProvider)
updater = CloudInstanceUpdater(cloud_provider=provider)
with pytest.raises(NotImplementedError):
updater.notify(
[
InstanceUpdateEvent(
new_instance_status=Instance.TERMINATING,
instance_id="i-1",
cloud_instance_id="cloud-node-1",
),
]
)
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,119 @@
# coding: utf-8
import os
import sys
import unittest
from queue import Queue
from unittest.mock import patch
import pytest # noqa
from ray._private.test_utils import load_test_config
from ray.autoscaler.tags import TAG_RAY_NODE_KIND
from ray.autoscaler.v2.instance_manager.config import AutoscalingConfig
from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
from ray.autoscaler.v2.instance_manager.subscribers.threaded_ray_installer import (
ThreadedRayInstaller,
)
from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
from ray.tests.autoscaler_test_utils import MockProcessRunner, MockProvider
class ThreadedRayInstallerTest(unittest.TestCase):
def setUp(self):
self.base_provider = MockProvider()
self.config = AutoscalingConfig(load_test_config("test_ray_complex.yaml"))
self.runner = MockProcessRunner()
self.ray_installer = RayInstaller(self.base_provider, self.config, self.runner)
self.instance_storage = InstanceStorage(
cluster_id="test_cluster_id",
storage=InMemoryStorage(),
)
self.error_queue = Queue()
self.threaded_ray_installer = ThreadedRayInstaller(
head_node_ip="127.0.0.1",
instance_storage=self.instance_storage,
ray_installer=self.ray_installer,
error_queue=self.error_queue,
)
def test_install_ray_on_new_node_version_mismatch(self):
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
instance = Instance(
instance_id="0",
instance_type="worker_nodes1",
cloud_instance_id="0",
status=Instance.RAY_INSTALLING,
node_kind=NodeKind.WORKER,
)
success, verison = self.instance_storage.upsert_instance(instance)
assert success
self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
self.threaded_ray_installer._install_ray_on_single_node(instance)
instances, _ = self.instance_storage.get_instances(
instance_ids={instance.instance_id}
)
assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
assert instances[instance.instance_id].version == verison
@patch.object(RayInstaller, "install_ray")
def test_install_ray_on_new_node_install_failed(self, mock_method):
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
instance = Instance(
instance_id="0",
instance_type="worker_nodes1",
cloud_instance_id="0",
status=Instance.RAY_INSTALLING,
node_kind=NodeKind.WORKER,
)
success, verison = self.instance_storage.upsert_instance(instance)
assert success
instance.version = verison
mock_method.side_effect = RuntimeError("Installation failed")
self.threaded_ray_installer._install_retry_interval = 0
self.threaded_ray_installer._max_install_attempts = 1
self.threaded_ray_installer._install_ray_on_single_node(instance)
instances, _ = self.instance_storage.get_instances(
instance_ids={instance.instance_id}
)
# Make sure the instance status is not updated by the ThreadedRayInstaller
# since it should be updated by the Reconciler.
assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
# Make sure the error is added to the error queue.
error = self.error_queue.get()
assert error.im_instance_id == instance.instance_id
assert "Installation failed" in error.details
def test_install_ray_on_new_nodes(self):
self.base_provider.create_node({}, {TAG_RAY_NODE_KIND: "worker_nodes1"}, 1)
instance = Instance(
instance_id="0",
instance_type="worker_nodes1",
cloud_instance_id="0",
status=Instance.RAY_INSTALLING,
node_kind=NodeKind.WORKER,
)
success, verison = self.instance_storage.upsert_instance(instance)
assert success
instance.version = verison
self.runner.respond_to_call("json .Config.Env", ["[]" for i in range(1)])
self.threaded_ray_installer._install_ray_on_new_nodes(instance.instance_id)
self.threaded_ray_installer._ray_installation_executor.shutdown(wait=True)
instances, _ = self.instance_storage.get_instances(
instance_ids={instance.instance_id}
)
# Make sure the instance status is not updated by the ThreadedRayInstaller
# since it should be updated by the Reconciler.
assert instances[instance.instance_id].status == Instance.RAY_INSTALLING
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
@@ -0,0 +1,607 @@
# coding: utf-8
import os
import sys
from typing import Dict
import pytest # noqa
from google.protobuf.json_format import ParseDict
from ray.autoscaler.v2.schema import (
ClusterConstraintDemand,
ClusterStatus,
LaunchRequest,
NodeInfo,
NodeUsage,
PlacementGroupResourceDemand,
RayTaskActorDemand,
ResourceDemandSummary,
ResourceRequestByCount,
ResourceUsage,
Stats,
)
from ray.autoscaler.v2.utils import (
ClusterStatusFormatter,
ClusterStatusParser,
ResourceRequestUtil,
)
from ray.core.generated.autoscaler_pb2 import GetClusterStatusReply
def _gen_cluster_status_reply(data: Dict):
return ParseDict(data, GetClusterStatusReply())
class TestResourceRequestUtil:
@staticmethod
def test_combine_requests_with_affinity():
AFFINITY = ResourceRequestUtil.PlacementConstraintType.AFFINITY
ANTI_AFFINITY = ResourceRequestUtil.PlacementConstraintType.ANTI_AFFINITY
rqs = [
ResourceRequestUtil.make({"CPU": 1}, [(AFFINITY, "1", "1")]), # 1
ResourceRequestUtil.make({"CPU": 2}, [(AFFINITY, "1", "1")]), # 1
ResourceRequestUtil.make({"CPU": 1}, [(AFFINITY, "2", "2")]), # 2
ResourceRequestUtil.make({"CPU": 1}, [(AFFINITY, "2", "2")]), # 2
ResourceRequestUtil.make({"CPU": 1}, [(ANTI_AFFINITY, "2", "2")]), # 3
ResourceRequestUtil.make({"CPU": 1}, [(ANTI_AFFINITY, "2", "2")]), # 4
ResourceRequestUtil.make({"CPU": 1}), # 5
]
rq_result = ResourceRequestUtil.combine_requests_with_affinity(rqs)
assert len(rq_result) == 5
actual = ResourceRequestUtil.to_dict_list(rq_result)
expected = [
ResourceRequestUtil.to_dict(
ResourceRequestUtil.make(
{"CPU": 3}, # Combined
[
(AFFINITY, "1", "1"),
],
)
),
ResourceRequestUtil.to_dict(
ResourceRequestUtil.make(
{"CPU": 2}, # Combined
[
(AFFINITY, "2", "2"),
],
)
),
ResourceRequestUtil.to_dict(
ResourceRequestUtil.make(
{"CPU": 1},
[(ANTI_AFFINITY, "2", "2")],
)
),
ResourceRequestUtil.to_dict(
ResourceRequestUtil.make(
{"CPU": 1},
[(ANTI_AFFINITY, "2", "2")],
)
),
ResourceRequestUtil.to_dict(
ResourceRequestUtil.make(
{"CPU": 1},
)
),
]
actual_str_serialized = [str(x) for x in actual]
expected_str_serialized = [str(x) for x in expected]
assert sorted(actual_str_serialized) == sorted(expected_str_serialized)
def test_cluster_status_parser_cluster_resource_state():
test_data = {
"cluster_resource_state": {
"node_states": [
{
"node_id": b"1" * 4,
"instance_id": "instance1",
"ray_node_type_name": "head_node",
"available_resources": {
"CPU": 0.5,
"GPU": 2.0,
},
"total_resources": {
"CPU": 1,
"GPU": 2.0,
},
"status": "RUNNING",
"node_ip_address": "10.10.10.10",
"instance_type_name": "m5.large",
},
{
"node_id": b"2" * 4,
"instance_id": "instance2",
"ray_node_type_name": "worker_node",
"available_resources": {},
"total_resources": {
"CPU": 1,
"GPU": 2.0,
},
"status": "DEAD",
"node_ip_address": "22.22.22.22",
"instance_type_name": "m5.large",
},
{
"node_id": b"3" * 4,
"instance_id": "instance3",
"ray_node_type_name": "worker_node",
"available_resources": {
"CPU": 1.0,
"GPU": 2.0,
},
"total_resources": {
"CPU": 1,
"GPU": 2.0,
},
"idle_duration_ms": 100,
"status": "IDLE",
"node_ip_address": "22.22.22.22",
"instance_type_name": "m5.large",
},
],
"pending_gang_resource_requests": [
{
"requests": [
{
"resources_bundle": {"CPU": 1, "GPU": 1},
"placement_constraints": [
{
"anti_affinity": {
"label_name": "_PG_1x1x",
"label_value": "",
}
}
],
},
],
"details": "1x1x:STRICT_SPREAD|PENDING",
},
{
"requests": [
{
"resources_bundle": {"GPU": 2},
"placement_constraints": [
{
"affinity": {
"label_name": "_PG_2x2x",
"label_value": "",
}
}
],
},
],
"details": "2x2x:STRICT_PACK|PENDING",
},
],
"pending_resource_requests": [
{
"request": {
"resources_bundle": {"CPU": 1, "GPU": 1},
"placement_constraints": [],
},
"count": 1,
},
],
"cluster_resource_constraints": [
{
"resource_requests": [
{
"request": {
"resources_bundle": {"GPU": 2, "CPU": 100},
"placement_constraints": [],
},
"count": 1,
},
]
}
],
"cluster_resource_state_version": 10,
},
"autoscaling_state": {},
}
reply = _gen_cluster_status_reply(test_data)
stats = Stats(gcs_request_time_s=0.1)
cluster_status = ClusterStatusParser.from_get_cluster_status_reply(reply, stats)
# Assert on health nodes
assert len(cluster_status.idle_nodes) + len(cluster_status.active_nodes) == 2
assert cluster_status.active_nodes[0].instance_id == "instance1"
assert cluster_status.active_nodes[0].ray_node_type_name == "head_node"
cluster_status.active_nodes[0].resource_usage.usage.sort(
key=lambda x: x.resource_name
)
assert cluster_status.active_nodes[0].resource_usage == NodeUsage(
usage=[
ResourceUsage(resource_name="CPU", total=1.0, used=0.5),
ResourceUsage(resource_name="GPU", total=2.0, used=0.0),
],
idle_time_ms=0,
)
assert cluster_status.idle_nodes[0].instance_id == "instance3"
assert cluster_status.idle_nodes[0].ray_node_type_name == "worker_node"
cluster_status.idle_nodes[0].resource_usage.usage.sort(
key=lambda x: x.resource_name
)
assert cluster_status.idle_nodes[0].resource_usage == NodeUsage(
usage=[
ResourceUsage(resource_name="CPU", total=1.0, used=0.0),
ResourceUsage(resource_name="GPU", total=2.0, used=0.0),
],
idle_time_ms=100,
)
# Assert on dead nodes
assert len(cluster_status.failed_nodes) == 1
assert cluster_status.failed_nodes[0].instance_id == "instance2"
assert cluster_status.failed_nodes[0].ray_node_type_name == "worker_node"
assert cluster_status.failed_nodes[0].resource_usage is None
# Assert on resource demands from tasks
assert len(cluster_status.resource_demands.ray_task_actor_demand) == 1
assert cluster_status.resource_demands.ray_task_actor_demand[
0
].bundles_by_count == [
ResourceRequestByCount(
bundle={"CPU": 1, "GPU": 1},
count=1,
)
]
# Assert on resource demands from placement groups
assert len(cluster_status.resource_demands.placement_group_demand) == 2
assert sorted(
cluster_status.resource_demands.placement_group_demand, key=lambda x: x.pg_id
) == [
PlacementGroupResourceDemand(
bundles_by_count=[
ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=1)
],
strategy="STRICT_SPREAD",
pg_id="1x1x",
state="PENDING",
details="1x1x:STRICT_SPREAD|PENDING",
),
PlacementGroupResourceDemand(
bundles_by_count=[ResourceRequestByCount(bundle={"GPU": 2}, count=1)],
strategy="STRICT_PACK",
pg_id="2x2x",
state="PENDING",
details="2x2x:STRICT_PACK|PENDING",
),
]
# Assert on resource constraints
assert len(cluster_status.resource_demands.cluster_constraint_demand) == 1
assert cluster_status.resource_demands.cluster_constraint_demand[
0
].bundles_by_count == [
ResourceRequestByCount(bundle={"GPU": 2, "CPU": 100}, count=1)
]
# Assert on the cluster_resource_usage
assert sorted(
cluster_status.cluster_resource_usage, key=lambda x: x.resource_name
) == [
ResourceUsage(resource_name="CPU", total=2.0, used=0.5),
ResourceUsage(resource_name="GPU", total=4.0, used=0.0),
]
# Assert on the node stats
assert cluster_status.stats.cluster_resource_state_version == "10"
assert cluster_status.stats.gcs_request_time_s == 0.1
def test_cluster_status_parser_autoscaler_state():
test_data = {
"cluster_resource_state": {},
"autoscaling_state": {
"pending_instance_requests": [
{
"instance_type_name": "m5.large",
"ray_node_type_name": "head_node",
"count": 1,
"request_ts": 29999,
},
{
"instance_type_name": "m5.large",
"ray_node_type_name": "worker_node",
"count": 2,
"request_ts": 19999,
},
],
"pending_instances": [
{
"instance_type_name": "m5.large",
"ray_node_type_name": "head_node",
"instance_id": "instance1",
"ip_address": "10.10.10.10",
"details": "Starting Ray",
},
],
"failed_instance_requests": [
{
"instance_type_name": "m5.large",
"ray_node_type_name": "worker_node",
"count": 2,
"reason": "Insufficient capacity",
"start_ts": 10000,
"failed_ts": 20000,
}
],
"autoscaler_state_version": 10,
},
}
reply = _gen_cluster_status_reply(test_data)
stats = Stats(gcs_request_time_s=0.1)
cluster_status = ClusterStatusParser.from_get_cluster_status_reply(reply, stats)
# Assert on the pending requests
assert len(cluster_status.pending_launches) == 2
assert cluster_status.pending_launches[0].instance_type_name == "m5.large"
assert cluster_status.pending_launches[0].ray_node_type_name == "head_node"
assert cluster_status.pending_launches[0].count == 1
assert cluster_status.pending_launches[0].request_ts_s == 29999
assert cluster_status.pending_launches[1].instance_type_name == "m5.large"
assert cluster_status.pending_launches[1].ray_node_type_name == "worker_node"
assert cluster_status.pending_launches[1].count == 2
assert cluster_status.pending_launches[1].request_ts_s == 19999
# Assert on the failed requests
assert len(cluster_status.failed_launches) == 1
assert cluster_status.failed_launches[0].instance_type_name == "m5.large"
assert cluster_status.failed_launches[0].ray_node_type_name == "worker_node"
assert cluster_status.failed_launches[0].count == 2
assert cluster_status.failed_launches[0].details == "Insufficient capacity"
assert cluster_status.failed_launches[0].request_ts_s == 10000
assert cluster_status.failed_launches[0].failed_ts_s == 20000
# Assert on the pending nodes
assert len(cluster_status.pending_nodes) == 1
assert cluster_status.pending_nodes[0].instance_type_name == "m5.large"
assert cluster_status.pending_nodes[0].ray_node_type_name == "head_node"
assert cluster_status.pending_nodes[0].instance_id == "instance1"
assert cluster_status.pending_nodes[0].ip_address == "10.10.10.10"
assert cluster_status.pending_nodes[0].details == "Starting Ray"
# Assert on stats
assert cluster_status.stats.autoscaler_version == "10"
assert cluster_status.stats.gcs_request_time_s == 0.1
def test_cluster_status_formatter():
state = ClusterStatus(
idle_nodes=[
NodeInfo(
instance_id="instance1",
instance_type_name="m5.large",
ray_node_type_name="head_node",
ip_address="127.0.0.1",
node_status="RUNNING",
node_id="fffffffffffffffffffffffffffffffffffffffffffffffffff00001",
resource_usage=NodeUsage(
usage=[
ResourceUsage(resource_name="CPU", total=1.0, used=0.5),
ResourceUsage(resource_name="GPU", total=2.0, used=0.0),
ResourceUsage(
resource_name="object_store_memory",
total=10282.0,
used=5555.0,
),
],
idle_time_ms=0,
),
),
NodeInfo(
instance_id="instance2",
instance_type_name="m5.large",
ray_node_type_name="worker_node",
ip_address="127.0.0.2",
node_status="RUNNING",
node_id="fffffffffffffffffffffffffffffffffffffffffffffffffff00002",
resource_usage=NodeUsage(
usage=[
ResourceUsage(resource_name="CPU", total=1.0, used=0),
ResourceUsage(resource_name="GPU", total=2.0, used=0),
],
idle_time_ms=0,
),
),
NodeInfo(
instance_id="instance3",
instance_type_name="m5.large",
ray_node_type_name="worker_node",
ip_address="127.0.0.2",
node_status="RUNNING",
node_id="fffffffffffffffffffffffffffffffffffffffffffffffffff00003",
resource_usage=NodeUsage(
usage=[
ResourceUsage(resource_name="CPU", total=1.0, used=0.0),
],
idle_time_ms=0,
),
),
],
pending_launches=[
LaunchRequest(
instance_type_name="m5.large",
count=2,
ray_node_type_name="worker_node",
state=LaunchRequest.Status.PENDING,
request_ts_s=10000,
),
LaunchRequest(
instance_type_name="g5n.large",
count=1,
ray_node_type_name="worker_node_gpu",
state=LaunchRequest.Status.PENDING,
request_ts_s=20000,
),
],
failed_launches=[
LaunchRequest(
instance_type_name="m5.large",
count=2,
ray_node_type_name="worker_node",
state=LaunchRequest.Status.FAILED,
details="Insufficient capacity",
request_ts_s=10000,
failed_ts_s=20000,
),
],
pending_nodes=[
NodeInfo(
instance_id="instance4",
instance_type_name="m5.large",
ray_node_type_name="worker_node",
ip_address="127.0.0.3",
details="Starting Ray",
),
],
failed_nodes=[
NodeInfo(
instance_id="instance5",
instance_type_name="m5.large",
ray_node_type_name="worker_node",
ip_address="127.0.0.5",
node_status="DEAD",
),
],
cluster_resource_usage=[
ResourceUsage(resource_name="CPU", total=3.0, used=0.5),
ResourceUsage(resource_name="GPU", total=4.0, used=0.0),
ResourceUsage(
resource_name="object_store_memory", total=10282.0, used=5555.0
),
],
resource_demands=ResourceDemandSummary(
placement_group_demand=[
PlacementGroupResourceDemand(
pg_id="1x1x",
strategy="STRICT_SPREAD",
state="PENDING",
details="1x1x:STRICT_SPREAD|PENDING",
bundles_by_count=[
ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=1)
],
),
PlacementGroupResourceDemand(
pg_id="2x2x",
strategy="STRICT_PACK",
state="PENDING",
details="2x2x:STRICT_PACK|PENDING",
bundles_by_count=[
ResourceRequestByCount(bundle={"GPU": 2}, count=1)
],
),
PlacementGroupResourceDemand(
pg_id="3x3x",
strategy="STRICT_PACK",
state="PENDING",
details="3x3x:STRICT_PACK|PENDING",
bundles_by_count=[
ResourceRequestByCount(bundle={"GPU": 2}, count=1)
],
),
],
ray_task_actor_demand=[
RayTaskActorDemand(
bundles_by_count=[
ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=1)
]
),
RayTaskActorDemand(
bundles_by_count=[
ResourceRequestByCount(bundle={"CPU": 1, "GPU": 1}, count=10)
]
),
],
cluster_constraint_demand=[
ClusterConstraintDemand(
bundles_by_count=[
ResourceRequestByCount(bundle={"GPU": 2, "CPU": 100}, count=2)
]
),
],
),
stats=Stats(
gcs_request_time_s=0.1,
none_terminated_node_request_time_s=0.2,
autoscaler_iteration_time_s=0.3,
autoscaler_version="10",
cluster_resource_state_version="20",
request_ts_s=775303535,
),
)
actual = ClusterStatusFormatter.format(state, verbose=True)
expected = """======== Autoscaler status: 1994-07-27 10:05:35 ========
GCS request time: 0.100000s
Node Provider non_terminated_nodes time: 0.200000s
Autoscaler iteration time: 0.300000s
Node status
--------------------------------------------------------
Active:
(no active nodes)
Idle:
1 head_node
2 worker_node
Pending:
worker_node, 1 launching
worker_node_gpu, 1 launching
instance4: worker_node, starting ray
Recent failures:
worker_node: LaunchFailed (latest_attempt: 02:46:40) - Insufficient capacity
worker_node: NodeTerminated (instance_id: instance5)
Resources
--------------------------------------------------------
Total Usage:
0.5/3.0 CPU
0.0/4.0 GPU
5.42KiB/10.04KiB object_store_memory
From request_resources:
{'GPU': 2, 'CPU': 100}: 2 from request_resources()
Pending Demands:
{'CPU': 1, 'GPU': 1}: 11+ pending tasks/actors
{'CPU': 1, 'GPU': 1} * 1 (STRICT_SPREAD): 1+ pending placement groups
{'GPU': 2} * 1 (STRICT_PACK): 2+ pending placement groups
Node: instance1 (head_node)
Id: fffffffffffffffffffffffffffffffffffffffffffffffffff00001
Usage:
0.5/1.0 CPU
0.0/2.0 GPU
5.42KiB/10.04KiB object_store_memory
Activity:
(no activity)
Node: instance2 (worker_node)
Id: fffffffffffffffffffffffffffffffffffffffffffffffffff00002
Usage:
0/1.0 CPU
0/2.0 GPU
Activity:
(no activity)
Node: instance3 (worker_node)
Id: fffffffffffffffffffffffffffffffffffffffffffffffffff00003
Usage:
0.0/1.0 CPU
Activity:
(no activity)"""
assert actual == expected
if __name__ == "__main__":
if os.environ.get("PARALLEL_CI"):
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
else:
sys.exit(pytest.main(["-sv", __file__]))
+206
View File
@@ -0,0 +1,206 @@
import abc
import operator
import time
from abc import abstractmethod
from collections import defaultdict
from typing import Dict, List, Optional, Tuple
import ray
from ray.autoscaler.v2.schema import AutoscalerInstance, ClusterStatus, ResourceUsage
from ray.autoscaler.v2.sdk import get_cluster_status
from ray.core.generated import autoscaler_pb2
from ray.core.generated.instance_manager_pb2 import Instance, NodeKind
class MockEventLogger:
def __init__(self, logger) -> None:
self._logs = defaultdict(list)
self._logger = logger
def info(self, s):
self._logger.info(s)
self._logs["info"].append(s)
def warning(self, s):
self._logger.warning(s)
self._logs["warning"].append(s)
def error(self, s):
self._logger.error(s)
self._logs["error"].append(s)
def debug(self, s):
self._logger.debug(s)
self._logs["debug"].append(s)
def get_logs(self, level: str) -> List[str]:
return self._logs[level]
class MockSubscriber:
def __init__(self):
self.events = []
def notify(self, events):
self.events.extend(events)
def clear(self):
self.events.clear()
def events_by_id(self, instance_id):
return [e for e in self.events if e.instance_id == instance_id]
def make_autoscaler_instance(
im_instance: Optional[Instance] = None,
ray_node: Optional[autoscaler_pb2.NodeState] = None,
cloud_instance_id: Optional[str] = None,
) -> AutoscalerInstance:
if cloud_instance_id:
if im_instance:
im_instance.cloud_instance_id = cloud_instance_id
if ray_node:
ray_node.instance_id = cloud_instance_id
return AutoscalerInstance(
im_instance=im_instance,
ray_node=ray_node,
cloud_instance_id=cloud_instance_id,
)
def get_cluster_resource_state(stub) -> autoscaler_pb2.ClusterResourceState:
request = autoscaler_pb2.GetClusterResourceStateRequest(
last_seen_cluster_resource_state_version=0
)
return stub.GetClusterResourceState(request).cluster_resource_state
class FakeCounter:
def dec(self, *args, **kwargs):
pass
def create_instance(
instance_id,
status=Instance.UNKNOWN,
instance_type="worker_nodes1",
status_times: List[Tuple["Instance.InstanceStatus", int]] = None,
launch_request_id="",
version=0,
cloud_instance_id="",
ray_node_id="",
node_kind=NodeKind.WORKER,
):
if not status_times:
status_times = [(status, time.time_ns())]
return Instance(
instance_id=instance_id,
status=status,
version=version,
instance_type=instance_type,
launch_request_id=launch_request_id,
status_history=[
Instance.StatusHistory(instance_status=status, timestamp_ns=ts)
for status, ts in status_times
],
cloud_instance_id=cloud_instance_id,
node_id=ray_node_id,
node_kind=node_kind,
)
def report_autoscaling_state(stub, autoscaling_state: autoscaler_pb2.AutoscalingState):
request = autoscaler_pb2.ReportAutoscalingStateRequest(
autoscaling_state=autoscaling_state
)
stub.ReportAutoscalingState(request)
def get_total_resources(usages: List[ResourceUsage]) -> Dict[str, float]:
"""Returns a map of resource name to total resource."""
return {r.resource_name: r.total for r in usages}
def get_available_resources(usages: List[ResourceUsage]) -> Dict[str, float]:
"""Returns a map of resource name to available resource."""
return {r.resource_name: r.total - r.used for r in usages}
def get_used_resources(usages: List[ResourceUsage]) -> Dict[str, float]:
"""Returns a map of resource name to used resource."""
return {r.resource_name: r.used for r in usages}
"""
Test utils for e2e autoscaling states checking.
"""
class Check(abc.ABC):
@abstractmethod
def check(self, status: ClusterStatus):
pass
def __repr__(self) -> str:
return self.__str__()
class CheckFailure(RuntimeError):
pass
class NodeCountCheck(Check):
def __init__(self, count: int):
self.count = count
def check(self, status: ClusterStatus):
healthy_nodes = len(status.active_nodes) + len(status.idle_nodes)
if healthy_nodes != self.count:
raise CheckFailure(f"Expected {self.count} nodes, got {healthy_nodes}")
def __str__(self) -> str:
return f"NodeCountCheck: {self.count}"
class TotalResourceCheck(Check):
def __init__(
self, resources: Dict[str, float], op: operator = operator.eq, enforce_all=False
):
self.resources = resources
self.op = op
self.enforce_all = enforce_all
def check(self, status: ClusterStatus):
actual = status.total_resources()
if self.enforce_all and len(actual) != len(self.resources):
raise CheckFailure(
f"Expected {len(self.resources)} resources, got {len(actual)}"
)
for k, v in self.resources.items():
if k not in actual and v:
raise CheckFailure(f"Expected resource {k} not found")
if not self.op(v, actual.get(k, 0)):
raise CheckFailure(
f"Expected resource {k} {self.op} {v}, got {actual.get(k, 0)}"
)
def __str__(self) -> str:
return f"TotalResourceCheck({self.op}): {self.resources}"
def check_cluster(
targets: List[Check],
) -> bool:
gcs_address = ray.get_runtime_context().gcs_address
cluster_status = get_cluster_status(gcs_address)
for target in targets:
target.check(cluster_status)
return True
File diff suppressed because it is too large Load Diff