Add files using upload-large-folder tool

lib/python3.12/site-packages/ray/autoscaler/v2/autoscaler.py

@@ -0,0 +1,216 @@
+import logging
+from queue import Queue
+from typing import List, Optional
+from urllib.parse import urlsplit
+
+from ray._raylet import GcsClient
+from ray.autoscaler._private.providers import _get_node_provider
+from ray.autoscaler.v2.event_logger import AutoscalerEventLogger
+from ray.autoscaler.v2.instance_manager.cloud_providers.kuberay.cloud_provider import (
+    KubeRayProvider,
+)
+from ray.autoscaler.v2.instance_manager.cloud_providers.read_only.cloud_provider import (  # noqa
+    ReadOnlyProvider,
+)
+from ray.autoscaler.v2.instance_manager.config import (
+    AutoscalingConfig,
+    IConfigReader,
+    Provider,
+)
+from ray.autoscaler.v2.instance_manager.instance_manager import (
+    InstanceManager,
+    InstanceUpdatedSubscriber,
+)
+from ray.autoscaler.v2.instance_manager.instance_storage import InstanceStorage
+from ray.autoscaler.v2.instance_manager.node_provider import (
+    ICloudInstanceProvider,
+    NodeProviderAdapter,
+)
+from ray.autoscaler.v2.instance_manager.ray_installer import RayInstaller
+from ray.autoscaler.v2.instance_manager.reconciler import Reconciler
+from ray.autoscaler.v2.instance_manager.storage import InMemoryStorage
+from ray.autoscaler.v2.instance_manager.subscribers.cloud_instance_updater import (
+    CloudInstanceUpdater,
+)
+from ray.autoscaler.v2.instance_manager.subscribers.ray_stopper import RayStopper
+from ray.autoscaler.v2.instance_manager.subscribers.threaded_ray_installer import (
+    ThreadedRayInstaller,
+)
+from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
+from ray.autoscaler.v2.scheduler import ResourceDemandScheduler
+from ray.autoscaler.v2.sdk import get_cluster_resource_state
+from ray.core.generated.autoscaler_pb2 import AutoscalingState
+
+logger = logging.getLogger(__name__)
+
+
+class Autoscaler:
+    def __init__(
+        self,
+        session_name: str,
+        config_reader: IConfigReader,
+        gcs_client: GcsClient,
+        event_logger: Optional[AutoscalerEventLogger] = None,
+        metrics_reporter: Optional[AutoscalerMetricsReporter] = None,
+    ) -> None:
+        """
+        Args:
+            session_name: The current Ray session name.
+            config_reader: The config reader.
+            gcs_client: The GCS client.
+            event_logger: The event logger for emitting cluster events.
+            metrics_reporter: The metrics reporter for emitting cluster metrics.
+        """
+
+        self._config_reader = config_reader
+
+        config = config_reader.get_cached_autoscaling_config()
+        logger.info(f"Using Autoscaling Config: \n{config.dump()}")
+
+        self._gcs_client = gcs_client
+        self._cloud_instance_provider = None
+        self._instance_manager = None
+        self._ray_stop_errors_queue = Queue()
+        self._ray_install_errors_queue = Queue()
+        self._event_logger = event_logger
+        self._metrics_reporter = metrics_reporter
+
+        self._init_cloud_instance_provider(config, config_reader)
+        self._init_instance_manager(
+            session_name=session_name,
+            config=config,
+            cloud_provider=self._cloud_instance_provider,
+            gcs_client=self._gcs_client,
+        )
+        self._scheduler = ResourceDemandScheduler(self._event_logger)
+
+    def _init_cloud_instance_provider(
+        self, config: AutoscalingConfig, config_reader: IConfigReader
+    ):
+        """
+        Initialize the cloud provider, and its dependencies (the v1 node provider)
+
+        Args:
+            config: The autoscaling config.
+            config_reader: The config reader.
+
+        """
+        provider_config = config.get_provider_config()
+        if provider_config["type"] == "kuberay":
+            provider_config["head_node_type"] = config.get_head_node_type()
+            self._cloud_instance_provider = KubeRayProvider(
+                config.get_config("cluster_name"),
+                provider_config,
+            )
+        elif config.provider == Provider.READ_ONLY:
+            provider_config["gcs_address"] = self._gcs_client.address
+            self._cloud_instance_provider = ReadOnlyProvider(
+                provider_config=provider_config,
+            )
+        else:
+            node_provider_v1 = _get_node_provider(
+                provider_config,
+                config.get_config("cluster_name"),
+            )
+
+            self._cloud_instance_provider = NodeProviderAdapter(
+                v1_provider=node_provider_v1,
+                config_reader=config_reader,
+            )
+
+    def _init_instance_manager(
+        self,
+        session_name: str,
+        cloud_provider: ICloudInstanceProvider,
+        gcs_client: GcsClient,
+        config: AutoscalingConfig,
+    ):
+        """
+        Initialize the instance manager, and its dependencies.
+        """
+
+        instance_storage = InstanceStorage(
+            cluster_id=session_name,
+            storage=InMemoryStorage(),
+        )
+        subscribers: List[InstanceUpdatedSubscriber] = []
+        subscribers.append(CloudInstanceUpdater(cloud_provider=cloud_provider))
+        subscribers.append(
+            RayStopper(gcs_client=gcs_client, error_queue=self._ray_stop_errors_queue)
+        )
+        if not config.disable_node_updaters() and isinstance(
+            cloud_provider, NodeProviderAdapter
+        ):
+            head_node_ip = urlsplit("//" + self._gcs_client.address).hostname
+            assert head_node_ip is not None, "Invalid GCS address format"
+            subscribers.append(
+                ThreadedRayInstaller(
+                    head_node_ip=head_node_ip,
+                    instance_storage=instance_storage,
+                    ray_installer=RayInstaller(
+                        provider=cloud_provider.v1_provider,
+                        config=config,
+                    ),
+                    error_queue=self._ray_install_errors_queue,
+                    # TODO(rueian): Rewrite the ThreadedRayInstaller and its underlying
+                    # NodeUpdater and CommandRunner to use the asyncio, so that we don't
+                    # need to use so many threads. We use so many threads now because
+                    # they are blocking and letting the new cloud machines to wait for
+                    # previous machines to finish installing Ray is quite inefficient.
+                    max_concurrent_installs=config.get_max_num_worker_nodes() or 50,
+                )
+            )
+
+        self._instance_manager = InstanceManager(
+            instance_storage=instance_storage,
+            instance_status_update_subscribers=subscribers,
+        )
+
+    def update_autoscaling_state(
+        self,
+    ) -> Optional[AutoscalingState]:
+        """Update the autoscaling state of the cluster by reconciling the current
+        state of the cluster resources, the cloud providers as well as instance
+        update subscribers with the desired state.
+
+        Returns:
+            AutoscalingState: The new autoscaling state of the cluster or None if
+            the state is not updated.
+
+        Raises:
+            None: No exception.
+        """
+
+        try:
+            ray_stop_errors = []
+            while not self._ray_stop_errors_queue.empty():
+                ray_stop_errors.append(self._ray_stop_errors_queue.get())
+
+            ray_install_errors = []
+            while not self._ray_install_errors_queue.empty():
+                ray_install_errors.append(self._ray_install_errors_queue.get())
+
+            # Get the current state of the ray cluster resources.
+            ray_cluster_resource_state = get_cluster_resource_state(self._gcs_client)
+
+            # Refresh the config from the source
+            self._config_reader.refresh_cached_autoscaling_config()
+            autoscaling_config = self._config_reader.get_cached_autoscaling_config()
+
+            return Reconciler.reconcile(
+                instance_manager=self._instance_manager,
+                scheduler=self._scheduler,
+                cloud_provider=self._cloud_instance_provider,
+                ray_cluster_resource_state=ray_cluster_resource_state,
+                non_terminated_cloud_instances=(
+                    self._cloud_instance_provider.get_non_terminated()
+                ),
+                cloud_provider_errors=self._cloud_instance_provider.poll_errors(),
+                ray_install_errors=ray_install_errors,
+                ray_stop_errors=ray_stop_errors,
+                autoscaling_config=autoscaling_config,
+                metrics_reporter=self._metrics_reporter,
+            )
+        except Exception as e:
+            logger.exception(e)
+            return None

lib/python3.12/site-packages/ray/autoscaler/v2/event_logger.py

@@ -0,0 +1,185 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+from ray._private.event.event_logger import EventLoggerAdapter
+from ray.autoscaler.v2.utils import ResourceRequestUtil
+from ray.core.generated.autoscaler_pb2 import (
+    ClusterResourceConstraint,
+    GangResourceRequest,
+    ResourceRequest,
+)
+from ray.core.generated.common_pb2 import LabelSelectorOperator
+from ray.core.generated.instance_manager_pb2 import LaunchRequest, TerminationRequest
+
+logger = logging.getLogger(__name__)
+
+
+class AutoscalerEventLogger:
+    """
+    Logs events related to the autoscaler.
+
+    # TODO:
+    - Add more logging for other events.
+    - Rate limit the events if too spammy.
+    """
+
+    def __init__(self, logger: EventLoggerAdapter):
+        self._logger = logger
+
+    def log_cluster_scheduling_update(
+        self,
+        cluster_resources: Dict[str, float],
+        launch_requests: Optional[List[LaunchRequest]] = None,
+        terminate_requests: Optional[List[TerminationRequest]] = None,
+        infeasible_requests: Optional[List[ResourceRequest]] = None,
+        infeasible_gang_requests: Optional[List[GangResourceRequest]] = None,
+        infeasible_cluster_resource_constraints: Optional[
+            List[ClusterResourceConstraint]
+        ] = None,
+    ) -> None:
+        """
+        Log updates to the autoscaler scheduling state.
+
+        Emits:
+        - info logs for node launches and terminations (counts grouped by node type).
+        - an info log summarizing the cluster size after a resize (CPUs/GPUs/TPUs).
+        - warnings describing infeasible single resource requests, infeasible gang
+          (placement group) requests, and infeasible cluster resource constraints.
+
+        Args:
+            cluster_resources: Mapping of resource name to total resources for the
+                current cluster state.
+            launch_requests: Node launch requests issued in this scheduling step.
+            terminate_requests: Node termination requests issued in this scheduling
+                step.
+            infeasible_requests: Resource requests that could not be satisfied by
+                any available node type.
+            infeasible_gang_requests: Gang/placement group requests that could not
+                be scheduled.
+            infeasible_cluster_resource_constraints: Cluster-level resource
+                constraints that could not be satisfied.
+
+        Returns:
+            None
+        """
+
+        # Log any launch events.
+        if launch_requests:
+            launch_type_count = defaultdict(int)
+            for req in launch_requests:
+                launch_type_count[req.instance_type] += req.count
+
+            for idx, (instance_type, count) in enumerate(launch_type_count.items()):
+                log_str = f"Adding {count} node(s) of type {instance_type}."
+                self._logger.info(f"{log_str}")
+                logger.info(f"{log_str}")
+
+        # Log any terminate events.
+        if terminate_requests:
+            termination_by_causes_and_type = defaultdict(int)
+            for req in terminate_requests:
+                termination_by_causes_and_type[(req.cause, req.instance_type)] += 1
+
+            cause_reason_map = {
+                TerminationRequest.Cause.OUTDATED: "outdated",
+                TerminationRequest.Cause.MAX_NUM_NODES: "max number of worker nodes reached",  # noqa
+                TerminationRequest.Cause.MAX_NUM_NODE_PER_TYPE: "max number of worker nodes per type reached",  # noqa
+                TerminationRequest.Cause.IDLE: "idle",
+            }
+
+            for idx, ((cause, instance_type), count) in enumerate(
+                termination_by_causes_and_type.items()
+            ):
+                log_str = f"Removing {count} nodes of type {instance_type} ({cause_reason_map[cause]})."  # noqa
+                self._logger.info(f"{log_str}")
+                logger.info(f"{log_str}")
+
+        # Cluster shape changes.
+        if launch_requests or terminate_requests:
+            num_cpus = cluster_resources.get("CPU", 0)
+            log_str = f"Resized to {int(num_cpus)} CPUs"
+
+            if "GPU" in cluster_resources:
+                log_str += f", {int(cluster_resources['GPU'])} GPUs"
+            if "TPU" in cluster_resources:
+                log_str += f", {int(cluster_resources['TPU'])} TPUs"
+
+            self._logger.info(f"{log_str}.")
+            self._logger.debug(f"Current cluster resources: {dict(cluster_resources)}.")
+
+        # Log any infeasible requests.
+        if infeasible_requests:
+            requests_by_count = ResourceRequestUtil.group_by_count(infeasible_requests)
+            log_str = "No available node types can fulfill resource requests "
+            for idx, req_count in enumerate(requests_by_count):
+                resource_map = ResourceRequestUtil.to_resource_map(req_count.request)
+                log_str += f"{resource_map}*{req_count.count}"
+                if idx < len(requests_by_count) - 1:
+                    log_str += ", "
+
+                # Parse and log label selectors if present
+                if req_count.request.label_selectors:
+                    selector_strs = []
+                    for selector in req_count.request.label_selectors:
+                        for constraint in selector.label_constraints:
+                            op = LabelSelectorOperator.Name(constraint.operator)
+                            values = ",".join(constraint.label_values)
+                            selector_strs.append(
+                                f"{constraint.label_key} {op} [{values}]"
+                            )
+                    if selector_strs:
+                        log_str += (
+                            " with label selectors: [" + "; ".join(selector_strs) + "]"
+                        )
+
+            log_str += (
+                ". Add suitable node types to this cluster to resolve this issue."
+            )
+            self._logger.warning(log_str)
+
+        if infeasible_gang_requests:
+            # Log for each placement group requests.
+            for gang_request in infeasible_gang_requests:
+                log_str = (
+                    "No available node types can fulfill "
+                    "placement group requests (detail={details}): ".format(
+                        details=gang_request.details
+                    )
+                )
+                requests_by_count = ResourceRequestUtil.group_by_count(
+                    gang_request.requests
+                )
+                for idx, req_count in enumerate(requests_by_count):
+                    resource_map = ResourceRequestUtil.to_resource_map(
+                        req_count.request
+                    )
+                    log_str += f"{resource_map}*{req_count.count}"
+                    if idx < len(requests_by_count) - 1:
+                        log_str += ", "
+
+                log_str += (
+                    ". Add suitable node types to this cluster to resolve this issue."
+                )
+                self._logger.warning(log_str)
+
+        if infeasible_cluster_resource_constraints:
+            # We will only have max 1 cluster resource constraint for now since it's
+            # from `request_resources()` sdk, where the most recent call would override
+            # the previous one.
+            for infeasible_constraint in infeasible_cluster_resource_constraints:
+                log_str = "No available node types can fulfill cluster constraint: "
+                for i, requests_by_count in enumerate(
+                    infeasible_constraint.resource_requests
+                ):
+                    resource_map = ResourceRequestUtil.to_resource_map(
+                        requests_by_count.request
+                    )
+                    log_str += f"{resource_map}*{requests_by_count.count}"
+                    if i < len(infeasible_constraint.resource_requests) - 1:
+                        log_str += ", "
+
+                log_str += (
+                    ". Add suitable node types to this cluster to resolve this issue."
+                )
+                self._logger.warning(log_str)

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/cloud_providers/kuberay/cloud_provider.py

@@ -0,0 +1,595 @@
+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
+
+# 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.node_provider import (
+    CloudInstance,
+    CloudInstanceId,
+    CloudInstanceProviderError,
+    ICloudInstanceProvider,
+    LaunchNodeError,
+    NodeKind,
+    TerminateNodeError,
+)
+from ray.autoscaler.v2.schema import NodeType
+
+logger = logging.getLogger(__name__)
+
+
+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],
+        k8s_api_client: Optional[IKubernetesHttpApiClient] = None,
+    ):
+        """
+        Args:
+            cluster_name: The name of the RayCluster resource.
+            provider_config: The namespace of the RayCluster.
+            k8s_api_client: The client to the Kubernetes API server.
+                This could 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
+        )
+
+        # Below are states that are cached locally.
+        self._requests = set()
+        self._launch_errors_queue = []
+        self._terminate_errors_queue = []
+
+        # Below are states that are fetched from the Kubernetes API server.
+        self._ray_cluster = None
+        self._cached_instances: Dict[CloudInstanceId, CloudInstance]
+
+    @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()
+        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
+        self._requests.add(request_id)
+        logger.info("Terminating worker pods: {}".format(ids))
+
+        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
+
+        try:
+            self._submit_scale_request(scale_request)
+        except Exception as e:
+            logger.exception(f"Error terminating nodes: {scale_request}")
+            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
+        self._requests.add(request_id)
+
+        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
+
+        try:
+            self._submit_scale_request(scale_request)
+        except Exception as e:
+            logger.exception(f"Error launching nodes: {scale_request}")
+            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
+
+    ############################
+    # 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._cached_instances = self._fetch_instances()
+
+    @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.
+
+        Returns:
+            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
+        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.
+        """
+        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 _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"]

lib/python3.12/site-packages/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, instance_id: CloudInstanceId) -> 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 []

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/common.py

@@ -0,0 +1,472 @@
+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.
+        """
+        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,
+        }
+
+    @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.
+        """
+        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 timeout_seconds.
+
+        Args:
+            instance: The instance to check.
+            timeout_seconds: 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,
+            },
+            # 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,
+                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,
+            },
+            # 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,
+            },
+            # Whenever a cloud instance disappears from the list of running cloud
+            # instances from the node provider, the instance is marked as stopped. Since
+            # we guarantee 1:1 mapping of a Instance to a cloud instance, 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.
+        """
+        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.
+            instance_status: The status to search for. If None, returns the last
+                status update.
+        """
+        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.
+            instance_status: The status to search for. If None, returns all
+                status updates 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)

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/config.py

@@ -0,0 +1,552 @@
+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", 300
+    )
+    # 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 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:
+        """
+        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")
+        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),
+                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:
+        """
+        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"] = len(available_node_types)
+            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)

lib/python3.12/site-packages/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

lib/python3.12/site-packages/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_verison: 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_verison,
+            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_version is
+        specified, the update will fail if the current storage version does not
+        match the expected version.
+
+        Args:
+            to_delete: A list of instances to be deleted.
+            expected_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

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/node_provider.py

@@ -0,0 +1,529 @@
+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:
+        """
+        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()

lib/python3.12/site-packages/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)}"
+            )

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/reconciler.py

@@ -0,0 +1,1601 @@
+import logging
+import math
+import time
+import uuid
+from collections import defaultdict
+from typing import Dict, List, Optional, Set, Tuple
+
+from ray._common.utils import binary_to_hex
+from ray.autoscaler.v2.instance_manager.common import InstanceUtil
+from ray.autoscaler.v2.instance_manager.config import (
+    AutoscalingConfig,
+    InstanceReconcileConfig,
+    Provider,
+)
+from ray.autoscaler.v2.instance_manager.instance_manager import InstanceManager
+from ray.autoscaler.v2.instance_manager.node_provider import (
+    CloudInstance,
+    CloudInstanceId,
+    CloudInstanceProviderError,
+    ICloudInstanceProvider,
+    LaunchNodeError,
+    TerminateNodeError,
+)
+from ray.autoscaler.v2.instance_manager.subscribers.ray_stopper import RayStopError
+from ray.autoscaler.v2.instance_manager.subscribers.threaded_ray_installer import (
+    RayInstallError,
+)
+from ray.autoscaler.v2.metrics_reporter import AutoscalerMetricsReporter
+from ray.autoscaler.v2.scheduler import IResourceScheduler, SchedulingRequest
+from ray.autoscaler.v2.schema import AutoscalerInstance, NodeType
+from ray.autoscaler.v2.utils import is_head_node
+from ray.core.generated.autoscaler_pb2 import (
+    AutoscalingState,
+    ClusterResourceState,
+    FailedInstanceRequest,
+    NodeState,
+    NodeStatus,
+    PendingInstance,
+    PendingInstanceRequest,
+)
+from ray.core.generated.instance_manager_pb2 import (
+    GetInstanceManagerStateRequest,
+    Instance as IMInstance,
+    InstanceUpdateEvent as IMInstanceUpdateEvent,
+    NodeKind,
+    StatusCode,
+    UpdateInstanceManagerStateRequest,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Reconciler:
+    """
+    A singleton class that reconciles the instance states of the instance manager
+    for autoscaler.
+
+    """
+
+    @staticmethod
+    def reconcile(
+        instance_manager: InstanceManager,
+        scheduler: IResourceScheduler,
+        cloud_provider: ICloudInstanceProvider,
+        ray_cluster_resource_state: ClusterResourceState,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+        autoscaling_config: AutoscalingConfig,
+        cloud_provider_errors: Optional[List[CloudInstanceProviderError]] = None,
+        ray_install_errors: Optional[List[RayInstallError]] = None,
+        ray_stop_errors: Optional[List[RayStopError]] = None,
+        metrics_reporter: Optional[AutoscalerMetricsReporter] = None,
+        _logger: Optional[logging.Logger] = None,
+    ) -> AutoscalingState:
+        """
+        The reconcile method computes InstanceUpdateEvents for the instance manager
+        by:
+
+        1. Reconciling the instance manager's instances with external states like
+        the cloud provider's, the ray cluster's states, the ray installer's results.
+        It performs "passive" status transitions for the instances (where the status
+        transition should only be reflecting the external states of the cloud provider
+        and the ray cluster, and should not be actively changing them)
+
+        2. Stepping the instances to the active states by computing instance status
+        transitions that are needed and updating the instance manager's state.
+        These transitions should be "active" where the transitions have side effects
+        (through InstanceStatusSubscriber) to the cloud provider and the ray cluster.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            ray_cluster_resource_state: The ray cluster's resource state.
+            non_terminated_cloud_instances: The non-terminated cloud instances from
+                the cloud provider.
+            cloud_provider_errors: The errors from the cloud provider.
+            ray_install_errors: The errors from RayInstaller.
+            ray_stop_errors: The errors from RayStopper.
+            metrics_reporter: The metric reporter to report the autoscaler metrics.
+            _logger: The logger (for testing).
+
+        """
+        cloud_provider_errors = cloud_provider_errors or []
+        ray_install_errors = ray_install_errors or []
+        ray_stop_errors = ray_stop_errors or []
+
+        autoscaling_state = AutoscalingState()
+        autoscaling_state.last_seen_cluster_resource_state_version = (
+            ray_cluster_resource_state.cluster_resource_state_version
+        )
+
+        Reconciler._sync_from(
+            instance_manager=instance_manager,
+            ray_nodes=ray_cluster_resource_state.node_states,
+            non_terminated_cloud_instances=non_terminated_cloud_instances,
+            cloud_provider_errors=cloud_provider_errors,
+            ray_install_errors=ray_install_errors,
+            ray_stop_errors=ray_stop_errors,
+            autoscaling_config=autoscaling_config,
+        )
+
+        Reconciler._step_next(
+            autoscaling_state=autoscaling_state,
+            instance_manager=instance_manager,
+            scheduler=scheduler,
+            cloud_provider=cloud_provider,
+            ray_cluster_resource_state=ray_cluster_resource_state,
+            non_terminated_cloud_instances=non_terminated_cloud_instances,
+            autoscaling_config=autoscaling_config,
+            _logger=_logger,
+        )
+
+        Reconciler._report_metrics(
+            instance_manager=instance_manager,
+            autoscaling_config=autoscaling_config,
+            metrics_reporter=metrics_reporter,
+        )
+
+        return autoscaling_state
+
+    @staticmethod
+    def _sync_from(
+        instance_manager: InstanceManager,
+        ray_nodes: List[NodeState],
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+        cloud_provider_errors: List[CloudInstanceProviderError],
+        ray_install_errors: List[RayInstallError],
+        ray_stop_errors: List[RayStopError],
+        autoscaling_config: AutoscalingConfig,
+    ):
+        """
+        Reconcile the instance states of the instance manager from external states like
+        the cloud provider's, the ray cluster's states, the ray installer's results,
+        etc.
+
+        For each instance, we try to figure out if we need to transition the instance
+        status to a new status, and if so, what the new status should be.
+
+        These transitions should be purely "passive", meaning they should only be
+        reflecting the external states of the cloud provider and the ray cluster,
+        and should not be actively changing the states of the cloud provider or the ray
+        cluster.
+
+        More specifically, we will reconcile status transitions for:
+            1.  QUEUED/REQUESTED -> ALLOCATED:
+                When a instance with launch request id (indicating a previous launch
+                request was made) could be assigned to an unassigned cloud instance
+                of the same instance type.
+            2.  REQUESTED -> ALLOCATION_FAILED:
+                When there's an error from the cloud provider for launch failure so
+                that the instance becomes ALLOCATION_FAILED.
+            3.  * -> RAY_RUNNING:
+                When a ray node on a cloud instance joins the ray cluster, we will
+                transition the instance to RAY_RUNNING.
+            4.  * -> TERMINATED:
+                When the cloud instance is already terminated, we will transition the
+                instance to TERMINATED.
+            5.  TERMINATING -> TERMINATION_FAILED:
+                When there's an error from the cloud provider for termination failure.
+            6.  * -> RAY_STOPPED:
+                When ray was stopped on the cloud instance, we will transition the
+                instance to RAY_STOPPED.
+            7.  * -> RAY_INSTALL_FAILED:
+                When there's an error from RayInstaller.
+            8. RAY_STOP_REQUESTED -> RAY_RUNNING:
+                When requested to stop ray, but failed to stop/drain the ray node
+                (e.g. idle termination drain rejected by the node).
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            ray_nodes: The ray cluster's states of ray nodes.
+            non_terminated_cloud_instances: The non-terminated cloud instances from
+                the cloud provider.
+            cloud_provider_errors: The errors from the cloud provider.
+            ray_install_errors: The errors from RayInstaller.
+            ray_stop_errors: The errors from RayStopper.
+
+        """
+
+        # Handle 1 & 2 for cloud instance allocation.
+        Reconciler._handle_cloud_instance_allocation(
+            instance_manager,
+            non_terminated_cloud_instances,
+            cloud_provider_errors,
+        )
+        Reconciler._handle_cloud_instance_terminated(
+            instance_manager, non_terminated_cloud_instances
+        )
+
+        Reconciler._handle_cloud_instance_termination_errors(
+            instance_manager, cloud_provider_errors
+        )
+
+        Reconciler._handle_extra_cloud_instances(
+            instance_manager, non_terminated_cloud_instances, ray_nodes
+        )
+
+        Reconciler._handle_ray_status_transition(
+            instance_manager, ray_nodes, autoscaling_config
+        )
+
+        Reconciler._handle_ray_install_failed(instance_manager, ray_install_errors)
+
+        Reconciler._handle_ray_stop_failed(instance_manager, ray_stop_errors, ray_nodes)
+
+    @staticmethod
+    def _step_next(
+        autoscaling_state: AutoscalingState,
+        instance_manager: InstanceManager,
+        scheduler: IResourceScheduler,
+        cloud_provider: ICloudInstanceProvider,
+        ray_cluster_resource_state: ClusterResourceState,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+        autoscaling_config: AutoscalingConfig,
+        _logger: Optional[logging.Logger] = None,
+    ):
+        """
+        Step the reconciler to the next state by computing instance status transitions
+        that are needed and updating the instance manager's state.
+
+        Specifically, we will:
+            1. Shut down leak cloud instances
+                Leaked cloud instances that are not managed by the instance manager.
+            2. Terminating instances with ray stopped or ray install failure.
+            3. Scale down the cluster:
+              (* -> RAY_STOP_REQUESTED/TERMINATING)
+                b. Extra cloud due to max nodes config.
+                c. Cloud instances with outdated configs.
+            4. Scale up the cluster:
+              (new QUEUED)
+                Create new instances based on the IResourceScheduler's decision for
+                scaling up.
+            5. Request cloud provider to launch new instances.
+              (QUEUED -> REQUESTED)
+            6. Install ray
+              (ALLOCATED -> RAY_INSTALLING)
+                When ray could be installed and launched.
+            7. Handle any stuck instances with timeouts.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            scheduler: The resource scheduler to make scaling decisions.
+            ray_cluster_resource_state: The ray cluster's resource state.
+            non_terminated_cloud_instances: The non-terminated cloud instances from
+                the cloud provider.
+            autoscaling_config: The autoscaling config.
+            _logger: The logger (for testing).
+
+        """
+
+        Reconciler._handle_stuck_instances(
+            instance_manager=instance_manager,
+            reconcile_config=autoscaling_config.get_instance_reconcile_config(),
+            _logger=_logger or logger,
+        )
+
+        Reconciler._scale_cluster(
+            autoscaling_state=autoscaling_state,
+            instance_manager=instance_manager,
+            ray_state=ray_cluster_resource_state,
+            scheduler=scheduler,
+            autoscaling_config=autoscaling_config,
+        )
+
+        Reconciler._handle_instances_launch(
+            instance_manager=instance_manager, autoscaling_config=autoscaling_config
+        )
+
+        Reconciler._terminate_instances(instance_manager=instance_manager)
+        if not autoscaling_config.disable_node_updaters():
+            Reconciler._install_ray(
+                instance_manager=instance_manager,
+                non_terminated_cloud_instances=non_terminated_cloud_instances,
+            )
+
+        Reconciler._fill_autoscaling_state(
+            instance_manager=instance_manager, autoscaling_state=autoscaling_state
+        )
+
+    #######################################################
+    # Utility methods for reconciling instance states.
+    #######################################################
+
+    @staticmethod
+    def _handle_cloud_instance_allocation(
+        instance_manager: InstanceManager,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+        cloud_provider_errors: List[CloudInstanceProviderError],
+    ):
+        im_instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+
+        # Compute intermediate states.
+
+        instances_with_launch_requests: List[IMInstance] = []
+        for instance in im_instances:
+            if instance.status != IMInstance.REQUESTED:
+                continue
+
+            assert (
+                instance.launch_request_id
+            ), "Instance in REQUESTED status should have launch_request_id set."
+            instances_with_launch_requests.append(instance)
+
+        assigned_cloud_instance_ids: Set[CloudInstanceId] = {
+            instance.cloud_instance_id
+            for instance in im_instances
+            if instance.cloud_instance_id
+            and instance.status
+            not in [IMInstance.TERMINATED, IMInstance.ALLOCATION_FAILED]
+        }
+        launch_errors: Dict[str, LaunchNodeError] = {
+            error.request_id: error
+            for error in cloud_provider_errors
+            if isinstance(error, LaunchNodeError)
+        }
+        unassigned_cloud_instances_by_type: Dict[
+            str, List[CloudInstance]
+        ] = defaultdict(list)
+
+        for cloud_instance_id, cloud_instance in non_terminated_cloud_instances.items():
+            if cloud_instance_id not in assigned_cloud_instance_ids:
+                unassigned_cloud_instances_by_type[cloud_instance.node_type].append(
+                    cloud_instance
+                )
+
+        # Sort the request instance by the increasing request time.
+        instances_with_launch_requests.sort(
+            key=lambda instance: InstanceUtil.get_status_transition_times_ns(
+                instance, IMInstance.REQUESTED
+            )
+        )
+
+        # For each instance, try to allocate or fail the allocation.
+        for instance in instances_with_launch_requests:
+            # Try allocate or fail with errors.
+            update_event = Reconciler._try_resolve_pending_allocation(
+                instance, unassigned_cloud_instances_by_type, launch_errors
+            )
+            if not update_event:
+                continue
+
+            updates[instance.instance_id] = update_event
+
+        # Update the instance manager for the events.
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _try_resolve_pending_allocation(
+        im_instance: IMInstance,
+        unassigned_cloud_instances_by_type: Dict[str, List[CloudInstance]],
+        launch_errors: Dict[str, LaunchNodeError],
+    ) -> Optional[IMInstanceUpdateEvent]:
+        """
+        Allocate, or fail the cloud instance allocation for the instance.
+
+        Args:
+            im_instance: The instance to allocate or fail.
+            unassigned_cloud_instances_by_type: The unassigned cloud instances by type.
+            launch_errors: The launch errors from the cloud provider.
+
+        Returns:
+            Instance update to ALLOCATED: if there's a matching unassigned cloud
+                instance with the same type.
+            Instance update to ALLOCATION_FAILED: if the instance allocation failed
+                with errors.
+            None: if there's no update.
+
+        """
+        unassigned_cloud_instance = None
+
+        # Try to allocate an unassigned cloud instance.
+        # TODO(rickyx): We could also look at the launch request id
+        # on the cloud node and the im instance later once all node providers
+        # support request id. For now, we only look at the instance type.
+        if len(unassigned_cloud_instances_by_type.get(im_instance.instance_type, [])):
+            unassigned_cloud_instance = unassigned_cloud_instances_by_type[
+                im_instance.instance_type
+            ].pop()
+
+        if unassigned_cloud_instance:
+            return IMInstanceUpdateEvent(
+                instance_id=im_instance.instance_id,
+                new_instance_status=IMInstance.ALLOCATED,
+                cloud_instance_id=unassigned_cloud_instance.cloud_instance_id,
+                node_kind=unassigned_cloud_instance.node_kind,
+                instance_type=unassigned_cloud_instance.node_type,
+                details=(
+                    "allocated unassigned cloud instance "
+                    f"{unassigned_cloud_instance.cloud_instance_id}"
+                ),
+            )
+
+        # If there's a launch error, transition to ALLOCATION_FAILED.
+        launch_error = launch_errors.get(im_instance.launch_request_id)
+        if launch_error and launch_error.node_type == im_instance.instance_type:
+            return IMInstanceUpdateEvent(
+                instance_id=im_instance.instance_id,
+                new_instance_status=IMInstance.ALLOCATION_FAILED,
+                details=f"launch failed with {str(launch_error)}",
+            )
+        # No update.
+        return None
+
+    @staticmethod
+    def _handle_ray_stop_failed(
+        instance_manager: InstanceManager,
+        ray_stop_errors: List[RayStopError],
+        ray_nodes: List[NodeState],
+    ):
+        """
+        The instance requested to stop ray, but failed to stop/drain the ray node.
+        E.g. connection errors, idle termination drain rejected by the node.
+
+        We will transition the instance back to RAY_RUNNING.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            ray_stop_errors: The errors from RayStopper.
+
+        """
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+
+        ray_stop_errors_by_instance_id = {
+            error.im_instance_id: error for error in ray_stop_errors
+        }
+
+        ray_nodes_by_ray_node_id = {binary_to_hex(n.node_id): n for n in ray_nodes}
+
+        ray_stop_requested_instances = {
+            instance.instance_id: instance
+            for instance in instances
+            if instance.status == IMInstance.RAY_STOP_REQUESTED
+        }
+
+        for instance_id, instance in ray_stop_requested_instances.items():
+            stop_error = ray_stop_errors_by_instance_id.get(instance_id)
+            if not stop_error:
+                continue
+
+            assert instance.node_id
+            ray_node = ray_nodes_by_ray_node_id.get(instance.node_id)
+            assert ray_node is not None and ray_node.status in [
+                NodeStatus.RUNNING,
+                NodeStatus.IDLE,
+            ], (
+                "There should be a running ray node for instance with ray stop "
+                "requested failed."
+            )
+
+            updates[instance_id] = IMInstanceUpdateEvent(
+                instance_id=instance_id,
+                new_instance_status=IMInstance.RAY_RUNNING,
+                details="failed to stop/drain ray",
+                ray_node_id=instance.node_id,
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _handle_ray_install_failed(
+        instance_manager: InstanceManager, ray_install_errors: List[RayInstallError]
+    ):
+
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+
+        # Get all instances with RAY_INSTALLING status.
+        instances_with_ray_installing = {
+            instance.instance_id: instance
+            for instance in instances
+            if instance.status == IMInstance.RAY_INSTALLING
+        }
+
+        install_errors = {error.im_instance_id: error for error in ray_install_errors}
+
+        # For each instance with RAY_INSTALLING status, check if there's any
+        # install error.
+        for instance_id, instance in instances_with_ray_installing.items():
+            install_error = install_errors.get(instance_id)
+            if install_error:
+                updates[instance_id] = IMInstanceUpdateEvent(
+                    instance_id=instance_id,
+                    new_instance_status=IMInstance.RAY_INSTALL_FAILED,
+                    details=(
+                        f"failed to install ray with errors: {install_error.details}"
+                    ),
+                )
+
+        # Update the instance manager for the events.
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _handle_cloud_instance_terminated(
+        instance_manager: InstanceManager,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+    ):
+        """
+        For any IM (instance manager) instance with a cloud node id, if the mapped
+        cloud instance is no longer running, transition the instance to TERMINATED.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            non_terminated_cloud_instances: The non-terminated cloud instances from
+                the cloud provider.
+        """
+        updates = {}
+        instances, version = Reconciler._get_im_instances(instance_manager)
+
+        non_terminated_instances_with_cloud_instance_assigned = {
+            instance.cloud_instance_id: instance
+            for instance in instances
+            if instance.cloud_instance_id and instance.status != IMInstance.TERMINATED
+        }
+
+        for (
+            cloud_instance_id,
+            instance,
+        ) in non_terminated_instances_with_cloud_instance_assigned.items():
+            if cloud_instance_id in non_terminated_cloud_instances.keys():
+                # The cloud instance is still running.
+                continue
+
+            # The cloud instance is terminated.
+            updates[instance.instance_id] = IMInstanceUpdateEvent(
+                instance_id=instance.instance_id,
+                new_instance_status=IMInstance.TERMINATED,
+                details=f"cloud instance {cloud_instance_id} no longer found",
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _handle_cloud_instance_termination_errors(
+        instance_manager: InstanceManager,
+        cloud_provider_errors: List[CloudInstanceProviderError],
+    ):
+        """
+        If any TERMINATING instances have termination errors, transition the instance to
+        TERMINATION_FAILED.
+
+        We will retry the termination for the TERMINATION_FAILED instances in the next
+        reconciler step.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            cloud_provider_errors: The errors from the cloud provider.
+
+        """
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+
+        termination_errors = {
+            error.cloud_instance_id: error
+            for error in cloud_provider_errors
+            if isinstance(error, TerminateNodeError)
+        }
+
+        terminating_instances_by_cloud_instance_id = {
+            instance.cloud_instance_id: instance
+            for instance in instances
+            if instance.status == IMInstance.TERMINATING
+        }
+
+        for cloud_instance_id, failure in termination_errors.items():
+            instance = terminating_instances_by_cloud_instance_id.get(cloud_instance_id)
+            if not instance:
+                # The instance is no longer in TERMINATING status.
+                continue
+
+            updates[instance.instance_id] = IMInstanceUpdateEvent(
+                instance_id=instance.instance_id,
+                new_instance_status=IMInstance.TERMINATION_FAILED,
+                details=f"termination failed: {str(failure)}",
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _get_im_instances(
+        instance_manager: InstanceManager,
+    ) -> Tuple[List[IMInstance], int]:
+        reply = instance_manager.get_instance_manager_state(
+            request=GetInstanceManagerStateRequest()
+        )
+        assert reply.status.code == StatusCode.OK
+        im_state = reply.state
+        return im_state.instances, im_state.version
+
+    @staticmethod
+    def _update_instance_manager(
+        instance_manager: InstanceManager,
+        version: int,
+        updates: Dict[str, IMInstanceUpdateEvent],
+    ) -> None:
+        if not updates:
+            return
+
+        updates = list(updates.values()) or []
+
+        reply = instance_manager.update_instance_manager_state(
+            request=UpdateInstanceManagerStateRequest(
+                expected_version=version,
+                updates=updates,
+            )
+        )
+        # TODO: While it's possible that a version mismatch
+        # happens, or some other failures could happen. But given
+        # the current implementation:
+        #   1. There's only 1 writer (the reconciler) for updating the instance
+        #       manager states, so there shouldn't be version mismatch.
+        #   2. Any failures in one reconciler step should be caught at a higher
+        #       level and be retried in the next reconciler step. If the IM
+        #       fails to be updated, we don't have sufficient info to handle it
+        #       here.
+        assert (
+            reply.status.code == StatusCode.OK
+        ), f"Failed to update instance manager: {reply}"
+
+    @staticmethod
+    def _handle_ray_status_transition(
+        instance_manager: InstanceManager,
+        ray_nodes: List[NodeState],
+        autoscaling_config: AutoscalingConfig,
+    ):
+        """
+        Handle the ray status transition for the instance manager.
+
+        If a new ray node running on the instance, transition it to RAY_RUNNING.
+        If a ray node stopped, transition it to RAY_STOPPED.
+        If a ray node is draining, transition it to RAY_STOPPING.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            ray_nodes: The ray cluster's states of ray nodes.
+        """
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+
+        im_instances_by_cloud_instance_id = {
+            instance.cloud_instance_id: instance
+            for instance in instances
+            if instance.cloud_instance_id
+            and instance.status
+            not in [IMInstance.TERMINATED, IMInstance.ALLOCATION_FAILED]
+        }
+        im_instances_by_ray_node_id = {
+            instance.node_id: instance for instance in instances if instance.node_id
+        }
+
+        for ray_node in ray_nodes:
+            im_instance = None
+            ray_node_id = binary_to_hex(ray_node.node_id)
+            if ray_node_id in im_instances_by_ray_node_id:
+                im_instance = im_instances_by_ray_node_id[ray_node_id]
+            else:
+                if autoscaling_config.provider == Provider.READ_ONLY:
+                    # We will use the node id as the cloud instance id for read-only
+                    # provider.
+                    im_instance = im_instances_by_cloud_instance_id[ray_node_id]
+                elif ray_node.instance_id:
+                    im_instance = im_instances_by_cloud_instance_id[
+                        ray_node.instance_id
+                    ]
+                else:
+                    # This should only happen to a ray node that's not managed by us.
+                    logger.warning(
+                        f"Ray node {ray_node_id} has no instance id. "
+                        "This only happens to a ray node not managed by autoscaler. "
+                        "If not, please file a bug at "
+                        "https://github.com/ray-project/ray"
+                    )
+                    continue
+
+            assert im_instance is not None, (
+                f"Ray node {ray_node_id} has no matching "
+                f"instance with cloud instance id={ray_node.instance_id}. We should "
+                "not see a ray node with cloud instance id not found in IM since "
+                "we have reconciled all cloud instances, and ray nodes by now."
+            )
+
+            reconciled_im_status = Reconciler._reconciled_im_status_from_ray_status(
+                ray_node.status, im_instance.status
+            )
+
+            if reconciled_im_status != im_instance.status:
+                updates[ray_node_id] = IMInstanceUpdateEvent(
+                    instance_id=im_instance.instance_id,
+                    new_instance_status=reconciled_im_status,
+                    details=(
+                        f"ray node {ray_node_id} is "
+                        f"{NodeStatus.Name(ray_node.status)}"
+                    ),
+                    ray_node_id=ray_node_id,
+                )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _reconciled_im_status_from_ray_status(
+        ray_status: NodeStatus, cur_im_status: IMInstance.InstanceStatus
+    ) -> "IMInstance.InstanceStatus":
+        """
+        Reconcile the instance status from the ray node status.
+        Args:
+            ray_status: the current ray node status.
+            cur_im_status: the current IM instance status.
+        Returns:
+            The reconciled IM instance status
+
+        Raises:
+            ValueError: If the ray status is unknown.
+        """
+        reconciled_im_status = None
+        if ray_status in [NodeStatus.RUNNING, NodeStatus.IDLE]:
+            reconciled_im_status = IMInstance.RAY_RUNNING
+        elif ray_status == NodeStatus.DEAD:
+            reconciled_im_status = IMInstance.RAY_STOPPED
+        elif ray_status == NodeStatus.DRAINING:
+            reconciled_im_status = IMInstance.RAY_STOPPING
+        else:
+            raise ValueError(f"Unknown ray status: {ray_status}")
+
+        if (
+            cur_im_status == reconciled_im_status
+            or cur_im_status
+            in InstanceUtil.get_reachable_statuses(reconciled_im_status)
+        ):
+            # No need to reconcile if the instance is already in the reconciled status
+            # or has already transitioned beyond it.
+            return cur_im_status
+
+        return reconciled_im_status
+
+    @staticmethod
+    def _handle_instances_launch(
+        instance_manager: InstanceManager, autoscaling_config: AutoscalingConfig
+    ):
+
+        instances, version = Reconciler._get_im_instances(instance_manager)
+
+        queued_instances = []
+        requested_instances = []
+        running_instances = []
+
+        for instance in instances:
+            if instance.status == IMInstance.QUEUED:
+                queued_instances.append(instance)
+            elif instance.status == IMInstance.REQUESTED:
+                requested_instances.append(instance)
+            elif instance.status == IMInstance.RAY_RUNNING:
+                running_instances.append(instance)
+
+        if not queued_instances:
+            # No QUEUED instances
+            return
+
+        to_launch = Reconciler._compute_to_launch(
+            queued_instances,
+            requested_instances,
+            running_instances,
+            autoscaling_config.get_upscaling_speed(),
+            autoscaling_config.get_max_concurrent_launches(),
+        )
+
+        # Transition the instances to REQUESTED for instance launcher to
+        # launch them.
+        updates = {}
+        new_launch_request_id = str(uuid.uuid4())
+        for instance_type, instances in to_launch.items():
+            for instance in instances:
+                # Reuse launch request id for any QUEUED instances that have been
+                # requested before due to retry.
+                launch_request_id = (
+                    new_launch_request_id
+                    if len(instance.launch_request_id) == 0
+                    else instance.launch_request_id
+                )
+                updates[instance.instance_id] = IMInstanceUpdateEvent(
+                    instance_id=instance.instance_id,
+                    new_instance_status=IMInstance.REQUESTED,
+                    launch_request_id=launch_request_id,
+                    instance_type=instance_type,
+                    details=(
+                        f"requested to launch {instance_type} with request id "
+                        f"{launch_request_id}"
+                    ),
+                )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _compute_to_launch(
+        queued_instances: List[IMInstance],
+        requested_instances: List[IMInstance],
+        running_instances: List[IMInstance],
+        upscaling_speed: float,
+        max_concurrent_launches: int,
+    ) -> Dict[NodeType, List[IMInstance]]:
+        def _group_by_type(instances):
+            instances_by_type = defaultdict(list)
+            for instance in instances:
+                instances_by_type[instance.instance_type].append(instance)
+            return instances_by_type
+
+        # Sort the instances by the time they were queued.
+        def _sort_by_earliest_queued(instance: IMInstance) -> List[int]:
+            queue_times = InstanceUtil.get_status_transition_times_ns(
+                instance, IMInstance.QUEUED
+            )
+            return sorted(queue_times)
+
+        queued_instances_by_type = _group_by_type(queued_instances)
+        running_instances_by_type = _group_by_type(running_instances)
+
+        total_num_requested_to_launch = len(requested_instances)
+        all_to_launch: Dict[NodeType : List[IMInstance]] = defaultdict(list)
+
+        for (
+            instance_type,
+            queued_instances_for_type,
+        ) in queued_instances_by_type.items():
+            running_instances_for_type = running_instances_by_type.get(
+                instance_type, []
+            )
+
+            # Enforce the max allowed pending nodes based on current running nodes
+            num_desired_to_upscale = max(
+                1,
+                math.ceil(upscaling_speed * max(len(running_instances_for_type), 1)),
+            )
+
+            # Enforce global limit, at most we can launch `max_concurrent_launches`
+            num_to_launch = min(
+                max_concurrent_launches - total_num_requested_to_launch,
+                num_desired_to_upscale,
+            )
+
+            # Cap both ends 0 <= num_to_launch <= num_queued
+            num_to_launch = max(0, num_to_launch)
+            num_to_launch = min(len(queued_instances_for_type), num_to_launch)
+
+            to_launch = sorted(queued_instances_for_type, key=_sort_by_earliest_queued)[
+                :num_to_launch
+            ]
+
+            all_to_launch[instance_type].extend(to_launch)
+            total_num_requested_to_launch += num_to_launch
+
+        return all_to_launch
+
+    @staticmethod
+    def _handle_stuck_instances(
+        instance_manager: InstanceManager,
+        reconcile_config: InstanceReconcileConfig,
+        _logger: logging.Logger,
+    ):
+        """
+        Handle stuck instances with timeouts.
+
+        Instances could be stuck in the following status and needs to be updated:
+            - REQUESTED: cloud provider is slow/fails to launch instances.
+            - ALLOCATED: ray fails to be started on the instance.
+            - RAY_INSTALLING: ray fails to be installed on the instance.
+            - TERMINATING: cloud provider is slow/fails to terminate instances.
+
+        Instances could be in the following status which could be unbounded or
+        transient, and we don't have a timeout mechanism to handle them. We would
+        warn if they are stuck for too long:
+            - RAY_STOPPING: ray taking time to drain.
+            - QUEUED: cloud provider is slow to launch instances, resulting in long
+                queue.
+
+            Reconciler should handle below statuses, if not, could be slow
+                reconcilation loop or a bug:
+            - RAY_INSTALL_FAILED
+            - RAY_STOPPED
+            - TERMINATION_FAILED
+
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            reconcile_config: The instance reconcile config.
+            _logger: The logger to log the warning messages. It's used for testing.
+
+        """
+        instances, version = Reconciler._get_im_instances(instance_manager)
+
+        instances_by_status = defaultdict(list)
+        for instance in instances:
+            instances_by_status[instance.status].append(instance)
+
+        im_updates = {}
+
+        # Fail or retry the cloud instance allocation if it's stuck
+        # in the REQUESTED state.
+        for instance in instances_by_status[IMInstance.REQUESTED]:
+            update = Reconciler._handle_stuck_requested_instance(
+                instance,
+                reconcile_config.request_status_timeout_s,
+                reconcile_config.max_num_retry_request_to_allocate,
+            )
+            if update:
+                im_updates[instance.instance_id] = update
+
+        # Leaked ALLOCATED instances should be terminated.
+        # This usually happens when ray fails to be started on the instance, so
+        # it's unable to be RAY_RUNNING after a long time.
+        for instance in instances_by_status[IMInstance.ALLOCATED]:
+            assert (
+                instance.cloud_instance_id
+            ), "cloud instance id should be set on ALLOCATED instance"
+            update = Reconciler._handle_stuck_instance(
+                instance,
+                reconcile_config.allocate_status_timeout_s,
+                new_status=IMInstance.TERMINATING,
+                cloud_instance_id=instance.cloud_instance_id,
+            )
+            if update:
+                im_updates[instance.instance_id] = update
+
+        # Fail the installation if it's stuck in RAY_INSTALLING for too long.
+        # If RAY_INSTALLING is stuck for too long, it's likely that the instance
+        # is not able to install ray, so we should also fail the installation.
+        for instance in instances_by_status[IMInstance.RAY_INSTALLING]:
+            update = Reconciler._handle_stuck_instance(
+                instance,
+                reconcile_config.ray_install_status_timeout_s,
+                new_status=IMInstance.RAY_INSTALL_FAILED,
+            )
+            if update:
+                im_updates[instance.instance_id] = update
+
+        # If we tried to terminate the instance, but it doesn't terminate (disappear
+        # from the cloud provider) after a long time, we fail the termination.
+        # This will trigger another attempt to terminate the instance.
+        for instance in instances_by_status[IMInstance.TERMINATING]:
+            update = Reconciler._handle_stuck_instance(
+                instance,
+                reconcile_config.terminating_status_timeout_s,
+                new_status=IMInstance.TERMINATION_FAILED,
+            )
+            if update:
+                im_updates[instance.instance_id] = update
+
+        # If we tried to stop ray on the instance, but it doesn't stop after a long
+        # time, we will transition it back to RAY_RUNNING as the stop/drain somehow
+        # failed. If it had succeed, we should have transitioned it to RAY_STOPPING
+        # or RAY_STOPPED.
+        for instance in instances_by_status[IMInstance.RAY_STOP_REQUESTED]:
+            update = Reconciler._handle_stuck_instance(
+                instance,
+                reconcile_config.ray_stop_requested_status_timeout_s,
+                new_status=IMInstance.RAY_RUNNING,
+                ray_node_id=instance.node_id,
+            )
+            if update:
+                im_updates[instance.instance_id] = update
+
+        # These statues could be unbounded or transient, and we don't have a timeout
+        # mechanism to handle them. We only warn if they are stuck for too long.
+        for status in [
+            # Ray taking time to drain. We could also have a timeout when Drain protocol
+            # supports timeout.
+            IMInstance.RAY_STOPPING,
+            # These should just be transient, we will terminate instances with this
+            # status in the next reconciler step.
+            IMInstance.RAY_INSTALL_FAILED,
+            IMInstance.RAY_STOPPED,
+            IMInstance.TERMINATION_FAILED,
+            # Instances could be in the QUEUED status for a long time if the cloud
+            # provider is slow to launch instances.
+            IMInstance.QUEUED,
+        ]:
+            Reconciler._warn_stuck_instances(
+                instances_by_status[status],
+                status=status,
+                warn_interval_s=reconcile_config.transient_status_warn_interval_s,
+                logger=_logger,
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, im_updates)
+
+    @staticmethod
+    def _warn_stuck_instances(
+        instances: List[IMInstance],
+        status: IMInstance.InstanceStatus,
+        warn_interval_s: int,
+        logger: logging.Logger,
+    ):
+        """Warn if any instance is stuck in a transient/unbounded status for too
+        long.
+        """
+        for instance in instances:
+            status_times_ns = InstanceUtil.get_status_transition_times_ns(
+                instance, select_instance_status=status
+            )
+            assert len(status_times_ns) >= 1
+            status_time_ns = sorted(status_times_ns)[-1]
+
+            if time.time_ns() - status_time_ns > warn_interval_s * 1e9:
+                logger.warning(
+                    "Instance {}({}) is stuck in {} for {} seconds.".format(
+                        instance.instance_id,
+                        IMInstance.InstanceStatus.Name(instance.status),
+                        IMInstance.InstanceStatus.Name(status),
+                        (time.time_ns() - status_time_ns) // 1e9,
+                    )
+                )
+
+    @staticmethod
+    def _is_head_node_running(instance_manager: InstanceManager) -> bool:
+        """
+        Check if the head node is running and ready.
+
+        If we scale up the cluster before head node is running,
+        it would cause issues when launching the worker nodes.
+
+        There are corner cases when the GCS is up (so the ray cluster resource
+        state is retrievable from the GCS), but the head node's raylet is not
+        running so the head node is missing from the reported nodes. This happens
+        when the head node is still starting up, or the raylet is not running
+        due to some issues, and this would yield false.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+
+        Returns:
+            True if the head node is running and ready, False otherwise.
+        """
+
+        im_instances, _ = Reconciler._get_im_instances(instance_manager)
+
+        for instance in im_instances:
+            if instance.node_kind == NodeKind.HEAD:
+                if instance.status == IMInstance.RAY_RUNNING:
+                    return True
+        return False
+
+    @staticmethod
+    def _scale_cluster(
+        autoscaling_state: AutoscalingState,
+        instance_manager: InstanceManager,
+        ray_state: ClusterResourceState,
+        scheduler: IResourceScheduler,
+        autoscaling_config: AutoscalingConfig,
+    ) -> None:
+        """
+        Scale the cluster based on the resource state and the resource scheduler's
+        decision:
+
+        - It launches new instances if needed.
+        - It terminates extra ray nodes if they should be shut down (preemption
+            or idle termination)
+
+        Args:
+            autoscaling_state: The autoscaling state to reconcile.
+            instance_manager: The instance manager to reconcile.
+            ray_state: The ray cluster's resource state.
+            scheduler: The resource scheduler to make scaling decisions.
+            autoscaling_config: The autoscaling config.
+
+        """
+
+        # Get the current instance states.
+        im_instances, version = Reconciler._get_im_instances(instance_manager)
+
+        im_instances_by_instance_id = {
+            i.instance_id: i for i in im_instances if i.instance_id
+        }
+        autoscaler_instances = []
+        ray_nodes_by_id = {
+            binary_to_hex(node.node_id): node for node in ray_state.node_states
+        }
+
+        for im_instance in im_instances:
+            ray_node = ray_nodes_by_id.get(im_instance.node_id)
+            autoscaler_instances.append(
+                AutoscalerInstance(
+                    ray_node=ray_node,
+                    im_instance=im_instance,
+                    cloud_instance_id=(
+                        im_instance.cloud_instance_id
+                        if im_instance.cloud_instance_id
+                        else None
+                    ),
+                )
+            )
+
+        # TODO(rickyx): We should probably name it as "Planner" or "Scaler"
+        # or "ClusterScaler"
+        sched_request = SchedulingRequest(
+            node_type_configs=autoscaling_config.get_node_type_configs(),
+            max_num_nodes=autoscaling_config.get_max_num_nodes(),
+            resource_requests=ray_state.pending_resource_requests,
+            gang_resource_requests=ray_state.pending_gang_resource_requests,
+            cluster_resource_constraints=ray_state.cluster_resource_constraints,
+            current_instances=autoscaler_instances,
+            idle_timeout_s=autoscaling_config.get_idle_timeout_s(),
+            disable_launch_config_check=(
+                autoscaling_config.disable_launch_config_check()
+            ),
+        )
+
+        # Ask scheduler for updates to the cluster shape.
+        reply = scheduler.schedule(sched_request)
+
+        # Populate the autoscaling state.
+        autoscaling_state.infeasible_resource_requests.extend(
+            reply.infeasible_resource_requests
+        )
+        autoscaling_state.infeasible_gang_resource_requests.extend(
+            reply.infeasible_gang_resource_requests
+        )
+        autoscaling_state.infeasible_cluster_resource_constraints.extend(
+            reply.infeasible_cluster_resource_constraints
+        )
+
+        if not Reconciler._is_head_node_running(instance_manager):
+            # We shouldn't be scaling the cluster until the head node is ready.
+            # This could happen when the head node (i.e. the raylet) is still
+            # pending registration even though GCS is up.
+            # We will wait until the head node is running and ready to avoid
+            # scaling the cluster from min worker nodes constraint.
+            return
+
+        if autoscaling_config.provider == Provider.READ_ONLY:
+            # We shouldn't be scaling the cluster if the provider is read-only.
+            return
+
+        # Scale the clusters if needed.
+        to_launch = reply.to_launch
+        to_terminate = reply.to_terminate
+        updates = {}
+        # Add terminating instances.
+        for terminate_request in to_terminate:
+            instance_id = terminate_request.instance_id
+            if terminate_request.instance_status == IMInstance.ALLOCATED:
+                # The instance is not yet running, so we can't request to stop/drain Ray.
+                # Therefore, we can skip the RAY_STOP_REQUESTED state and directly terminate the node.
+                im_instance_to_terminate = im_instances_by_instance_id[instance_id]
+                updates[terminate_request.instance_id] = IMInstanceUpdateEvent(
+                    instance_id=instance_id,
+                    new_instance_status=IMInstance.TERMINATING,
+                    cloud_instance_id=im_instance_to_terminate.cloud_instance_id,
+                    termination_request=terminate_request,
+                    details=f"terminating ray: {terminate_request.details}",
+                )
+            else:
+                updates[terminate_request.instance_id] = IMInstanceUpdateEvent(
+                    instance_id=instance_id,
+                    new_instance_status=IMInstance.RAY_STOP_REQUESTED,
+                    termination_request=terminate_request,
+                    details=f"draining ray: {terminate_request.details}",
+                )
+
+        # Add new instances.
+        for launch_request in to_launch:
+            for _ in range(launch_request.count):
+                instance_id = InstanceUtil.random_instance_id()
+                updates[instance_id] = IMInstanceUpdateEvent(
+                    instance_id=instance_id,
+                    new_instance_status=IMInstance.QUEUED,
+                    instance_type=launch_request.instance_type,
+                    upsert=True,
+                    details=(
+                        f"queuing new instance of {launch_request.instance_type} "
+                        "from scheduler"
+                    ),
+                )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _terminate_instances(instance_manager: InstanceManager):
+        """
+        Terminate instances with the below statuses:
+            - RAY_STOPPED: ray was stopped on the cloud instance.
+            - RAY_INSTALL_FAILED: ray installation failed on the cloud instance,
+                we will not retry.
+            - TERMINATION_FAILED: cloud provider failed to terminate the instance
+                or timeout for termination happened, we will retry again.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+        """
+
+        im_instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+        for instance in im_instances:
+            if instance.status not in [
+                IMInstance.RAY_STOPPED,
+                IMInstance.RAY_INSTALL_FAILED,
+                IMInstance.TERMINATION_FAILED,
+            ]:
+                continue
+
+            # Terminate the instance.
+            updates[instance.instance_id] = IMInstanceUpdateEvent(
+                instance_id=instance.instance_id,
+                new_instance_status=IMInstance.TERMINATING,
+                cloud_instance_id=instance.cloud_instance_id,
+                details="terminating instance from "
+                f"{IMInstance.InstanceStatus.Name(instance.status)}",
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _install_ray(
+        instance_manager: InstanceManager,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+    ) -> None:
+        """
+        Install ray on the allocated instances when it's ready (cloud instance
+        should be running)
+
+        This is needed if ray installation needs to be performed by
+        the instance manager.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+        """
+        im_instances, version = Reconciler._get_im_instances(instance_manager)
+        updates = {}
+        for instance in im_instances:
+            if instance.status != IMInstance.ALLOCATED:
+                continue
+
+            if instance.node_kind == NodeKind.HEAD:
+                # Skip head node.
+                continue
+
+            cloud_instance = non_terminated_cloud_instances.get(
+                instance.cloud_instance_id
+            )
+
+            assert cloud_instance, (
+                f"Cloud instance {instance.cloud_instance_id} is not found "
+                "in non_terminated_cloud_instances."
+            )
+
+            if not cloud_instance.is_running:
+                # It might still be pending (e.g. setting up ssh)
+                continue
+
+            # Install ray on the running cloud instance
+            updates[instance.instance_id] = IMInstanceUpdateEvent(
+                instance_id=instance.instance_id,
+                new_instance_status=IMInstance.RAY_INSTALLING,
+                details="installing ray",
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _fill_autoscaling_state(
+        instance_manager: InstanceManager,
+        autoscaling_state: AutoscalingState,
+    ) -> None:
+
+        # Use the IM instance version for the autoscaler_state_version
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        autoscaling_state.autoscaler_state_version = version
+
+        # Group instances by status
+        instances_by_status = defaultdict(list)
+        for instance in instances:
+            instances_by_status[instance.status].append(instance)
+
+        # Pending instance requests
+        instances_by_launch_request = defaultdict(list)
+        queued_instances = []
+        for instance in (
+            instances_by_status[IMInstance.REQUESTED]
+            + instances_by_status[IMInstance.QUEUED]
+        ):
+            if instance.launch_request_id:
+                instances_by_launch_request[instance.launch_request_id].append(instance)
+            else:
+                queued_instances.append(instance)
+
+        for _, instances in instances_by_launch_request.items():
+            num_instances_by_type = defaultdict(int)
+            for instance in instances:
+                num_instances_by_type[instance.instance_type] += 1
+
+            # All instances with same request id should have the same
+            # request time.
+            request_update = InstanceUtil.get_last_status_transition(
+                instances[0], IMInstance.REQUESTED
+            )
+            request_time_ns = request_update.timestamp_ns if request_update else 0
+
+            for instance_type, count in num_instances_by_type.items():
+                autoscaling_state.pending_instance_requests.append(
+                    PendingInstanceRequest(
+                        ray_node_type_name=instance_type,
+                        count=int(count),
+                        request_ts=int(request_time_ns // 1e9),
+                    )
+                )
+
+        # Pending instances
+        for instance in (
+            instances_by_status[IMInstance.ALLOCATED]
+            + instances_by_status[IMInstance.RAY_INSTALLING]
+        ):
+
+            status_history = sorted(
+                instance.status_history, key=lambda x: x.timestamp_ns, reverse=True
+            )
+            autoscaling_state.pending_instances.append(
+                PendingInstance(
+                    instance_id=instance.instance_id,
+                    ray_node_type_name=instance.instance_type,
+                    details=status_history[0].details,
+                )
+            )
+
+        # Failed instance requests
+        for instance in instances_by_status[IMInstance.ALLOCATION_FAILED]:
+            request_status_update = InstanceUtil.get_last_status_transition(
+                instance, IMInstance.REQUESTED
+            )
+            failed_status_update = InstanceUtil.get_last_status_transition(
+                instance, IMInstance.ALLOCATION_FAILED
+            )
+            failed_time = (
+                failed_status_update.timestamp_ns if failed_status_update else 0
+            )
+            request_time = (
+                request_status_update.timestamp_ns if request_status_update else 0
+            )
+            autoscaling_state.failed_instance_requests.append(
+                FailedInstanceRequest(
+                    ray_node_type_name=instance.instance_type,
+                    start_ts=int(request_time // 1e9),
+                    failed_ts=int(
+                        failed_time // 1e9,
+                    ),
+                    reason=failed_status_update.details,
+                    count=1,
+                )
+            )
+
+    @staticmethod
+    def _handle_stuck_requested_instance(
+        instance: IMInstance, timeout_s: int, max_num_retry_request_to_allocate: int
+    ) -> Optional[IMInstanceUpdateEvent]:
+        """
+        Fail the cloud instance allocation if it's stuck in the REQUESTED state.
+
+        Args:
+            instance: The instance to handle.
+            timeout_s: The timeout in seconds.
+            max_num_retry_request_to_allocate: The maximum number of times an instance
+                could be requested to allocate.
+
+        Returns:
+            Instance update to ALLOCATION_FAILED: if the instance allocation failed
+                with errors.
+            None: if there's no update.
+
+        """
+        if not InstanceUtil.has_timeout(instance, timeout_s):
+            # Not timeout yet, be patient.
+            return None
+
+        all_request_times_ns = sorted(
+            InstanceUtil.get_status_transition_times_ns(
+                instance, select_instance_status=IMInstance.REQUESTED
+            )
+        )
+
+        # Fail the allocation if we have tried too many times.
+        if len(all_request_times_ns) > max_num_retry_request_to_allocate:
+            return IMInstanceUpdateEvent(
+                instance_id=instance.instance_id,
+                new_instance_status=IMInstance.ALLOCATION_FAILED,
+                details=(
+                    "failed to allocate cloud instance after "
+                    f"{len(all_request_times_ns)} attempts > "
+                    f"max_num_retry_request_to_allocate={max_num_retry_request_to_allocate}"  # noqa
+                ),
+            )
+
+        # Retry the allocation if we could by transitioning to QUEUED again.
+        return IMInstanceUpdateEvent(
+            instance_id=instance.instance_id,
+            new_instance_status=IMInstance.QUEUED,
+            details=f"queue again to launch after timeout={timeout_s}s",
+        )
+
+    @staticmethod
+    def _handle_stuck_instance(
+        instance: IMInstance,
+        timeout_s: int,
+        new_status: IMInstance.InstanceStatus,
+        **update_kwargs: Dict,
+    ) -> Optional[IMInstanceUpdateEvent]:
+        """
+        Fail the instance if it's stuck in the status for too long.
+
+        Args:
+            instance: The instance to handle.
+            timeout_s: The timeout in seconds.
+            new_status: The new status to transition to.
+            update_kwargs: The update kwargs for InstanceUpdateEvent
+
+        Returns:
+            Instance update to the new status: if the instance is stuck in the status
+                for too long.
+            None: if there's no update.
+
+        """
+        if not InstanceUtil.has_timeout(instance, timeout_s):
+            # Not timeout yet, be patient.
+            return None
+
+        return IMInstanceUpdateEvent(
+            instance_id=instance.instance_id,
+            new_instance_status=new_status,
+            details=f"timeout={timeout_s}s at status "
+            f"{IMInstance.InstanceStatus.Name(instance.status)}",
+            **update_kwargs,
+        )
+
+    @staticmethod
+    def _handle_extra_cloud_instances(
+        instance_manager: InstanceManager,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+        ray_nodes: List[NodeState],
+    ):
+        """
+        For extra cloud instances (i.e. cloud instances that are non terminated as
+        returned by cloud provider, but not managed by the instance manager), we
+        will create new IM instances with ALLOCATED status.
+
+        Such instances could either be:
+            1. Leaked instances that are incorrectly started by the cloud instance
+            provider, and they would be terminated eventually if they fail to
+            transition to RAY_RUNNING by stuck instances reconciliation, or they
+            would join the  ray cluster and be terminated when the cluster scales down.
+            2. Instances that are started by the cloud instance provider intentionally
+            but not yet discovered by the instance manager. This could happen for
+               a. Head node that's started before the autoscaler.
+               b. Worker nodes that's started by the cloud provider upon users'
+               actions: i.e. KubeRay scaling up the cluster with ray cluster config
+               change.
+            3. Ray nodes with cloud instance id not in the cloud provider. This could
+            happen if there's delay in the Ray's state (i.e. cloud instance already
+            terminated, but the ray node is still not dead yet).
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            non_terminated_cloud_instances: The non-terminated cloud instances from
+                the cloud provider.
+            ray_nodes: The ray cluster's states of ray nodes.
+        """
+        Reconciler._handle_extra_cloud_instances_from_ray_nodes(
+            instance_manager, ray_nodes
+        )
+        Reconciler._handle_extra_cloud_instances_from_cloud_provider(
+            instance_manager, non_terminated_cloud_instances
+        )
+
+    @staticmethod
+    def _handle_extra_cloud_instances_from_cloud_provider(
+        instance_manager: InstanceManager,
+        non_terminated_cloud_instances: Dict[CloudInstanceId, CloudInstance],
+    ):
+        """
+        For extra cloud instances that are not managed by the instance manager but
+        are running in the cloud provider, we will create new IM instances with
+        ALLOCATED status.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            non_terminated_cloud_instances: The non-terminated cloud instances from
+                the cloud provider.
+        """
+        updates = {}
+
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        cloud_instance_ids_managed_by_im = {
+            instance.cloud_instance_id
+            for instance in instances
+            if instance.cloud_instance_id
+            and instance.status
+            not in [IMInstance.TERMINATED, IMInstance.ALLOCATION_FAILED]
+        }
+
+        # Find the extra cloud instances that are not managed by the instance manager.
+        for cloud_instance_id, cloud_instance in non_terminated_cloud_instances.items():
+            if cloud_instance_id in cloud_instance_ids_managed_by_im:
+                continue
+            updates[cloud_instance_id] = IMInstanceUpdateEvent(
+                instance_id=InstanceUtil.random_instance_id(),  # Assign a new id.
+                cloud_instance_id=cloud_instance_id,
+                new_instance_status=IMInstance.ALLOCATED,
+                node_kind=cloud_instance.node_kind,
+                instance_type=cloud_instance.node_type,
+                details=(
+                    "allocated unmanaged cloud instance :"
+                    f"{cloud_instance.cloud_instance_id} "
+                    f"({NodeKind.Name(cloud_instance.node_kind)}) from cloud provider"
+                ),
+                upsert=True,
+            )
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _handle_extra_cloud_instances_from_ray_nodes(
+        instance_manager: InstanceManager, ray_nodes: List[NodeState]
+    ):
+        """
+        For extra cloud instances reported by Ray but not managed by the instance
+        manager, we will create new IM instances with ALLOCATED status.
+
+        Args:
+            instance_manager: The instance manager to reconcile.
+            ray_nodes: The ray cluster's states of ray nodes.
+        """
+        updates = {}
+
+        instances, version = Reconciler._get_im_instances(instance_manager)
+        cloud_instance_ids_managed_by_im = {
+            instance.cloud_instance_id
+            for instance in instances
+            if instance.cloud_instance_id
+            and not instance.node_id
+            and instance.status
+            not in [IMInstance.TERMINATED, IMInstance.ALLOCATION_FAILED]
+        }
+        ray_node_ids_managed_by_im = {
+            instance.node_id for instance in instances if instance.node_id
+        }
+
+        for ray_node in ray_nodes:
+            if not ray_node.instance_id:
+                continue
+
+            ray_node_id = binary_to_hex(ray_node.node_id)
+            if ray_node_id in ray_node_ids_managed_by_im:
+                continue
+
+            cloud_instance_id = ray_node.instance_id
+            if cloud_instance_id in cloud_instance_ids_managed_by_im:
+                continue
+
+            is_head = is_head_node(ray_node)
+            updates[ray_node_id] = IMInstanceUpdateEvent(
+                instance_id=InstanceUtil.random_instance_id(),  # Assign a new id.
+                cloud_instance_id=cloud_instance_id,
+                new_instance_status=IMInstance.ALLOCATED,
+                node_kind=NodeKind.HEAD if is_head else NodeKind.WORKER,
+                ray_node_id=ray_node_id,
+                instance_type=ray_node.ray_node_type_name,
+                details=(
+                    "allocated unmanaged worker cloud instance from ray node: "
+                    f"{ray_node_id}"
+                ),
+                upsert=True,
+            )
+
+        Reconciler._update_instance_manager(instance_manager, version, updates)
+
+    @staticmethod
+    def _report_metrics(
+        instance_manager: InstanceManager,
+        autoscaling_config: AutoscalingConfig,
+        metrics_reporter: Optional[AutoscalerMetricsReporter] = None,
+    ):
+        if not metrics_reporter:
+            return
+
+        instances, _ = Reconciler._get_im_instances(instance_manager)
+        node_type_configs = autoscaling_config.get_node_type_configs()
+
+        metrics_reporter.report_instances(instances, node_type_configs)
+        metrics_reporter.report_resources(instances, node_type_configs)

lib/python3.12/site-packages/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

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/subscribers/cloud_instance_updater.py

@@ -0,0 +1,93 @@
+import logging
+import uuid
+from collections import defaultdict
+from typing import List
+
+from ray.autoscaler.v2.instance_manager.instance_manager import (
+    InstanceUpdatedSubscriber,
+)
+from ray.autoscaler.v2.instance_manager.node_provider import ICloudInstanceProvider
+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,
+    ) -> None:
+        self._cloud_provider = cloud_provider
+
+    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
+        ]
+        self._launch_new_instances(new_requests)
+        self._terminate_instances(new_terminations)
+
+    def _terminate_instances(self, new_terminations: List[InstanceUpdateEvent]):
+        """
+        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 _launch_new_instances(self, new_requests: List[InstanceUpdateEvent]):
+        """
+        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
+            )

lib/python3.12/site-packages/ray/autoscaler/v2/instance_manager/subscribers/ray_stopper.py

@@ -0,0 +1,154 @@
+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.
+            ray_node_id: The ray node id to drain.
+            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.
+            ray_node_id: The ray node id to stop.
+        """
+        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))

lib/python3.12/site-packages/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),
+            )
+        )