chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/instance_manager/subscribers/cloud_instance_updater.py

@@ -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
+            )

python/ray/autoscaler/v2/instance_manager/subscribers/cloud_resource_monitor.py

@@ -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

python/ray/autoscaler/v2/instance_manager/subscribers/ray_stopper.py

@@ -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))

python/ray/autoscaler/v2/instance_manager/subscribers/threaded_ray_installer.py

@@ -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),
+            )
+        )