chore: import upstream snapshot with attribution

python/ray/autoscaler/v2/instance_manager/cloud_providers/kuberay/cloud_provider.py

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

python/ray/autoscaler/v2/instance_manager/cloud_providers/kuberay/ippr_provider.py

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

python/ray/autoscaler/v2/instance_manager/cloud_providers/read_only/cloud_provider.py

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

python/ray/autoscaler/v2/instance_manager/common.py

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

python/ray/autoscaler/v2/instance_manager/config.py

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

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

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

python/ray/autoscaler/v2/instance_manager/instance_storage.py

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

python/ray/autoscaler/v2/instance_manager/node_provider.py

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

python/ray/autoscaler/v2/instance_manager/ray_installer.py

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

python/ray/autoscaler/v2/instance_manager/storage.py

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

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